From 79a19f30b8b2719d436a959b958ea244caf6aa78 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 22 Jun 2026 02:06:18 +0000 Subject: [PATCH] Add regression tests for FlowMatchEuler, DDIM guard, and LoRA state dict utils - FlowMatchEulerDiscreteScheduler: contract tests for validation, dynamic shifting, integer-timestep rejection, and stochastic sampling - DDIMScheduler: lock down set_timesteps guard against inverted comparison - state_dict_utils: unit tests for LoRA key remapping and error paths Co-authored-by: Simon Lynch --- tests/others/test_state_dict_utils.py | 82 +++++++++++ tests/schedulers/test_scheduler_ddim.py | 14 ++ ...est_scheduler_flow_match_euler_discrete.py | 129 ++++++++++++++++++ 3 files changed, 225 insertions(+) create mode 100644 tests/others/test_state_dict_utils.py create mode 100644 tests/schedulers/test_scheduler_flow_match_euler_discrete.py diff --git a/tests/others/test_state_dict_utils.py b/tests/others/test_state_dict_utils.py new file mode 100644 index 0000000..208b7d0 --- /dev/null +++ b/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) diff --git a/tests/schedulers/test_scheduler_ddim.py b/tests/schedulers/test_scheduler_ddim.py index 13b353a..dc6e688 100644 --- a/tests/schedulers/test_scheduler_ddim.py +++ b/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) diff --git a/tests/schedulers/test_scheduler_flow_match_euler_discrete.py b/tests/schedulers/test_scheduler_flow_match_euler_discrete.py new file mode 100644 index 0000000..38b014c --- /dev/null +++ b/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)