Add regression tests for FlowMatchEuler, DDIM guard, and LoRA state dict utils

tests/others/test_state_dict_utils.py

@@ -0,0 +1,82 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+
+from diffusers.utils.state_dict_utils import (
+    StateDictType,
+    convert_state_dict_to_diffusers,
+    convert_state_dict_to_peft,
+    convert_unet_state_dict_to_peft,
+)
+
+
+class StateDictUtilsTest(unittest.TestCase):
+    """
+    Unit tests for LoRA state-dict conversion utilities. Every LoRA load path goes through these
+    functions; incorrect key remapping silently produces wrong weights.
+    """
+
+    def test_convert_diffusers_old_to_peft(self):
+        weight = torch.randn(4, 4)
+        state_dict = {
+            "attn1.to_out_lora.down.weight": weight,
+            "attn1.to_out_lora.up.weight": weight,
+        }
+        converted = convert_state_dict_to_peft(state_dict)
+        self.assertIn("attn1.out_proj.lora_A.weight", converted)
+        self.assertIn("attn1.out_proj.lora_B.weight", converted)
+        torch.testing.assert_close(converted["attn1.out_proj.lora_A.weight"], weight)
+        torch.testing.assert_close(converted["attn1.out_proj.lora_B.weight"], weight)
+
+    def test_convert_diffusers_new_is_noop(self):
+        weight = torch.randn(4, 4)
+        state_dict = {
+            "attn1.q_proj.lora_linear_layer.down.weight": weight,
+            "attn1.q_proj.lora_linear_layer.up.weight": weight,
+        }
+        converted = convert_state_dict_to_diffusers(state_dict)
+        self.assertEqual(set(converted.keys()), set(state_dict.keys()))
+
+    def test_convert_peft_to_diffusers(self):
+        weight = torch.randn(4, 4)
+        state_dict = {
+            "attn1.q_proj.lora_A.weight": weight,
+            "attn1.q_proj.lora_B.weight": weight,
+        }
+        converted = convert_state_dict_to_diffusers(state_dict, original_type=StateDictType.PEFT)
+        self.assertIn("attn1.q_proj.lora_linear_layer.down.weight", converted)
+        self.assertIn("attn1.q_proj.lora_linear_layer.up.weight", converted)
+
+    def test_convert_unet_state_dict_to_peft(self):
+        weight = torch.randn(4, 4)
+        state_dict = {
+            "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q_lora.down.weight": weight,
+            "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q_lora.up.weight": weight,
+        }
+        converted = convert_unet_state_dict_to_peft(state_dict)
+        self.assertIn("down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.lora_A.weight", converted)
+        self.assertIn("down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.lora_B.weight", converted)
+
+    def test_convert_state_dict_to_peft_unrecognized_raises(self):
+        state_dict = {"some.unrelated.weight": torch.randn(2, 2)}
+        with self.assertRaises(ValueError):
+            convert_state_dict_to_peft(state_dict)
+
+    def test_convert_state_dict_to_diffusers_unrecognized_raises(self):
+        state_dict = {"some.unrelated.weight": torch.randn(2, 2)}
+        with self.assertRaises(ValueError):
+            convert_state_dict_to_diffusers(state_dict)

tests/schedulers/test_scheduler_ddim.py

@@ -73,6 +73,20 @@ def test_timestep_spacing(self):
         for timestep_spacing in ["trailing", "leading"]:
             self.check_over_configs(timestep_spacing=timestep_spacing)
 
+    def test_set_timesteps_num_inference_steps_exceeds_train_timesteps_raises(self):
+        # Guard against inverted comparison (num_inference_steps < num_train_timesteps) which would
+        # reject all valid inference schedules and accept invalid ones.
+        scheduler_class = self.scheduler_classes[0]
+        scheduler = scheduler_class(**self.get_scheduler_config())
+        with self.assertRaises(ValueError):
+            scheduler.set_timesteps(scheduler.config.num_train_timesteps + 1)
+
+    def test_set_timesteps_num_inference_steps_at_limit_succeeds(self):
+        scheduler_class = self.scheduler_classes[0]
+        scheduler = scheduler_class(**self.get_scheduler_config())
+        scheduler.set_timesteps(scheduler.config.num_train_timesteps)
+        self.assertEqual(scheduler.num_inference_steps, scheduler.config.num_train_timesteps)
+
     def test_rescale_betas_zero_snr(self):
         for rescale_betas_zero_snr in [True, False]:
             self.check_over_configs(rescale_betas_zero_snr=rescale_betas_zero_snr)

tests/schedulers/test_scheduler_flow_match_euler_discrete.py

@@ -0,0 +1,129 @@
+# Copyright 2025 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteSchedulerOutput
+
+
+class FlowMatchEulerDiscreteSchedulerTest(unittest.TestCase):
+    """
+    Contract tests for FlowMatchEulerDiscreteScheduler — the default scheduler for SD3, Flux, Wan,
+    Qwen-Image, and many other pipelines. Pipeline tests only smoke-test it; these tests lock down
+    validation and stepping behavior that is easy to break silently.
+    """
+
+    scheduler_class = FlowMatchEulerDiscreteScheduler
+
+    def get_default_config(self, **kwargs):
+        config = {
+            "num_train_timesteps": 1000,
+            "shift": 1.0,
+        }
+        config.update(**kwargs)
+        return config
+
+    def test_instantiation_with_defaults(self):
+        scheduler = self.scheduler_class(**self.get_default_config())
+        self.assertEqual(scheduler.config.num_train_timesteps, 1000)
+        self.assertEqual(scheduler.config.shift, 1.0)
+
+    def test_mutually_exclusive_sigma_schedules_raises(self):
+        with self.assertRaises(ValueError):
+            self.scheduler_class(
+                **self.get_default_config(use_karras_sigmas=True, use_exponential_sigmas=True),
+            )
+
+    def test_invalid_time_shift_type_raises(self):
+        with self.assertRaises(ValueError):
+            self.scheduler_class(**self.get_default_config(time_shift_type="quadratic"))
+
+    def test_set_timesteps_endpoints(self):
+        scheduler = self.scheduler_class(**self.get_default_config())
+        for nfe in [1, 2, 4, 8, 16]:
+            scheduler.set_timesteps(num_inference_steps=nfe)
+            self.assertEqual(scheduler.timesteps.shape, (nfe,))
+            self.assertEqual(scheduler.sigmas.shape, (nfe + 1,))
+            self.assertAlmostEqual(scheduler.timesteps[0].item(), 1000.0, places=4)
+            self.assertAlmostEqual(scheduler.sigmas[-1].item(), 0.0, places=4)
+
+    def test_set_timesteps_dynamic_shifting_requires_mu(self):
+        scheduler = self.scheduler_class(**self.get_default_config(use_dynamic_shifting=True))
+        with self.assertRaises(ValueError):
+            scheduler.set_timesteps(num_inference_steps=4)
+
+    def test_set_timesteps_dynamic_shifting_with_mu(self):
+        scheduler = self.scheduler_class(**self.get_default_config(use_dynamic_shifting=True))
+        scheduler.set_timesteps(num_inference_steps=4, mu=0.5)
+        self.assertEqual(scheduler.num_inference_steps, 4)
+
+    def test_set_timesteps_sigmas_timesteps_length_mismatch_raises(self):
+        scheduler = self.scheduler_class(**self.get_default_config())
+        with self.assertRaises(ValueError):
+            scheduler.set_timesteps(sigmas=[0.9, 0.5, 0.1], timesteps=[900.0, 500.0])
+
+    def test_set_timesteps_custom_sigmas(self):
+        scheduler = self.scheduler_class(**self.get_default_config(shift=1.0))
+        custom = [0.9, 0.7, 0.4, 0.1]
+        scheduler.set_timesteps(sigmas=custom)
+        self.assertEqual(scheduler.num_inference_steps, 4)
+        self.assertEqual(scheduler.timesteps.shape, (4,))
+        self.assertEqual(scheduler.sigmas.shape, (5,))
+        self.assertAlmostEqual(scheduler.sigmas[-1].item(), 0.0, places=6)
+
+    def test_step_shape_preserved(self):
+        scheduler = self.scheduler_class(**self.get_default_config())
+        scheduler.set_timesteps(num_inference_steps=4)
+
+        sample = torch.randn(2, 16, 8, 8)
+        model_output = torch.randn_like(sample)
+        timestep = scheduler.timesteps[0:1]
+
+        output = scheduler.step(model_output, timestep, sample)
+        self.assertIsInstance(output, FlowMatchEulerDiscreteSchedulerOutput)
+        self.assertEqual(output.prev_sample.shape, sample.shape)
+        self.assertEqual(output.prev_sample.dtype, model_output.dtype)
+
+    def test_step_rejects_integer_timestep_index(self):
+        # Pipelines must pass scheduler.timesteps values, not enumerate() indices.
+        scheduler = self.scheduler_class(**self.get_default_config())
+        scheduler.set_timesteps(num_inference_steps=4)
+
+        sample = torch.randn(1, 4, 4, 4)
+        model_output = torch.randn_like(sample)
+
+        with self.assertRaises(ValueError):
+            scheduler.step(model_output, 0, sample)
+
+        with self.assertRaises(ValueError):
+            scheduler.step(model_output, torch.tensor([0], dtype=torch.long), sample)
+
+    def test_step_stochastic_sampling(self):
+        scheduler = self.scheduler_class(**self.get_default_config(stochastic_sampling=True))
+        scheduler.set_timesteps(num_inference_steps=4)
+
+        sample = torch.randn(1, 4, 4, 4)
+        model_output = torch.randn_like(sample)
+        generator = torch.Generator().manual_seed(0)
+
+        prev_sample = scheduler.step(
+            model_output,
+            scheduler.timesteps[0:1],
+            sample,
+            generator=generator,
+        ).prev_sample
+        self.assertEqual(prev_sample.shape, sample.shape)