From db03fc96ff6be8dd9b0c959da5c1bf8e3e9c83e3 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 15:51:37 +0530 Subject: [PATCH 01/19] unbloat lora utils in tests --- tests/lora/test_lora_layers_auraflow.py | 12 -- tests/lora/test_lora_layers_cogvideox.py | 12 -- tests/lora/test_lora_layers_cogview4.py | 12 -- tests/lora/test_lora_layers_hunyuanvideo.py | 12 -- tests/lora/test_lora_layers_ltx_video.py | 12 -- tests/lora/test_lora_layers_lumina2.py | 11 -- tests/lora/test_lora_layers_mochi.py | 14 +- tests/lora/test_lora_layers_sana.py | 12 -- tests/lora/test_lora_layers_wan.py | 12 -- tests/lora/utils.py | 204 ++++++-------------- 10 files changed, 59 insertions(+), 254 deletions(-) diff --git a/tests/lora/test_lora_layers_auraflow.py b/tests/lora/test_lora_layers_auraflow.py index d119feae20d0..c9821cecb066 100644 --- a/tests/lora/test_lora_layers_auraflow.py +++ b/tests/lora/test_lora_layers_auraflow.py @@ -119,18 +119,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") def test_simple_inference_with_text_lora_and_scale(self): pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index 565d6db69727..f276832e26a5 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -152,22 +152,10 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") def test_simple_inference_with_text_lora_and_scale(self): pass - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @unittest.skip("Not supported in CogVideoX.") def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): pass diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index b7367d9b0946..0b57ece158cc 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -171,18 +171,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in CogView4.") def test_simple_inference_with_text_lora_and_scale(self): pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index 19e31f320d0a..6fa4a1c92a18 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -177,22 +177,10 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") def test_simple_inference_with_text_lora_and_scale(self): pass - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @nightly @require_torch_accelerator diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index 88949227cf94..101d7ab4cba1 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -130,18 +130,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") def test_simple_inference_with_text_lora_and_scale(self): pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py index d7096e79b93c..638309097f95 100644 --- a/tests/lora/test_lora_layers_lumina2.py +++ b/tests/lora/test_lora_layers_lumina2.py @@ -117,22 +117,11 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") def test_simple_inference_with_text_lora_and_scale(self): pass @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @skip_mps @pytest.mark.xfail( condition=torch.device(torch_device).type == "cpu" and is_torch_version(">=", "2.5"), diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index 501a4b35f48e..8d53506b026f 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -121,22 +121,10 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in Mochi.") def test_simple_inference_with_text_lora_and_scale(self): pass - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora_save_load(self): - pass - - @unittest.skip("Not supported in CogVideoX.") + @unittest.skip("Not supported in Mochi.") def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): pass diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py index 24beb46b95ff..0c2b6a0a9b00 100644 --- a/tests/lora/test_lora_layers_sana.py +++ b/tests/lora/test_lora_layers_sana.py @@ -121,18 +121,6 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in SANA.") def test_simple_inference_with_text_lora_and_scale(self): pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index fe26a56e77cf..d103c1a2a63b 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -126,18 +126,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in Wan.") def test_simple_inference_with_text_lora_and_scale(self): pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora_save_load(self): - pass diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 91ca188137e7..521d647fbd8f 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -343,28 +343,6 @@ def test_simple_inference(self): output_no_lora = pipe(**inputs)[0] self.assertTrue(output_no_lora.shape == self.output_shape) - def test_simple_inference_with_text_lora(self): - """ - Tests a simple inference with lora attached on the text encoder - and makes sure it works as expected - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) - @require_peft_version_greater("0.13.1") def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" @@ -520,44 +498,18 @@ def test_simple_inference_with_text_lora_and_scale(self): "Lora + 0 scale should lead to same result as no LoRA", ) - def test_simple_inference_with_text_lora_fused(self): + @parameterized.expand([("fused",), ("unloaded",), ("save_load",)]) + def test_lora_text_encoder_actions(self, action): """ - Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model - and makes sure it works as expected + Tests various actions (fusing, unloading, saving/loading) on a pipeline + with LoRA applied to the text encoder(s). """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - pipe.fuse_lora() - # Fusing should still keep the LoRA layers - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - ouput_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(ouput_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" + if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): + pytest.skip( + "Test not supported for {self.__class__.__name__} since there is not text encoder in the LoRA loadable modules." ) - - def test_simple_inference_with_text_lora_unloaded(self): - """ - Tests a simple inference with lora attached to text encoder, then unloads the lora weights - and makes sure it works as expected - """ for scheduler_cls in self.scheduler_classes: + # 1. Setup pipeline and get base output without LoRA components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) @@ -567,66 +519,69 @@ def test_simple_inference_with_text_lora_unloaded(self): output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue(output_no_lora.shape == self.output_shape) + # 2. Add LoRA adapters pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.unload_lora_weights() - # unloading should remove the LoRA layers - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder" - ) - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly unloaded in text encoder 2", - ) - - ouput_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - np.allclose(ouput_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), - "Fused lora should change the output", + not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" ) - def test_simple_inference_with_text_lora_save_load(self): - """ - Tests a simple usecase where users could use saving utilities for LoRA. - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + # 3. Perform the specified action and assert the outcome + if action == "fused": + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) + # Fusing should still keep the LoRA layers + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder)) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder_2)) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + not np.allclose(output_after_action, output_no_lora, atol=1e-3, rtol=1e-3), + "Fused LoRA should produce a different output from the base model.", + ) - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + elif action == "unloaded": + pipe.unload_lora_weights() + # Unloading should remove the LoRA layers + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder)) + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder_2)) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts + output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_after_action, output_no_lora, atol=1e-3, rtol=1e-3), + "Output after unloading LoRA should match the original output.", ) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) + elif action == "save_load": + with tempfile.TemporaryDirectory() as tmpdirname: + modules_to_save = self._get_modules_to_save(pipe) + lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + self.pipeline_class.save_lora_weights( + save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts + ) + self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + # Unload and then load the weights back + pipe.unload_lora_weights() + pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] + for module_name, module in modules_to_save.items(): + self.assertTrue( + check_if_lora_correctly_set(module), + f"LoRA not set correctly in {module_name} after loading.", + ) - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) + output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_lora, output_after_action, atol=1e-3, rtol=1e-3), + "Loading from a saved checkpoint should yield the same result as the original LoRA.", + ) def test_simple_inference_with_partial_text_lora(self): """ @@ -690,49 +645,6 @@ def test_simple_inference_with_partial_text_lora(self): "Removing adapters should change the output", ) - def test_simple_inference_save_pretrained_with_text_lora(self): - """ - Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - - pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) - pipe_from_pretrained.to(torch_device) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe_from_pretrained.text_encoder), - "Lora not correctly set in text encoder", - ) - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe_from_pretrained.text_encoder_2), - "Lora not correctly set in text encoder 2", - ) - - images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - def test_simple_inference_with_text_denoiser_lora_save_load(self): """ Tests a simple usecase where users could use saving utilities for LoRA for Unet + text encoder From 27fe7c5b2344ac3356d6c7aa22358378edaeb21c Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 15:53:55 +0530 Subject: [PATCH 02/19] update --- tests/lora/test_lora_layers_cogview4.py | 32 --------------------- tests/lora/test_lora_layers_hunyuanvideo.py | 5 ---- tests/lora/utils.py | 14 --------- 3 files changed, 51 deletions(-) diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index 0b57ece158cc..f94f3d5a521d 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -13,10 +13,8 @@ # limitations under the License. import sys -import tempfile import unittest -import numpy as np import torch from parameterized import parameterized from transformers import AutoTokenizer, GlmModel @@ -27,7 +25,6 @@ require_peft_backend, require_torch_accelerator, skip_mps, - torch_device, ) @@ -119,35 +116,6 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): def test_simple_inference_with_text_denoiser_lora_unfused(self): super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - def test_simple_inference_save_pretrained(self): - """ - Tests a simple usecase where users could use saving utilities for LoRA through save_pretrained - """ - for scheduler_cls in self.scheduler_classes: - components, _, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - - pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) - pipe_from_pretrained.to(torch_device) - - images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - @parameterized.expand([("block_level", True), ("leaf_level", False)]) @require_torch_accelerator def test_group_offloading_inference_denoiser(self, offload_type, use_stream): diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index 6fa4a1c92a18..c6bc6b9cc205 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -156,11 +156,6 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): def test_simple_inference_with_text_denoiser_lora_unfused(self): super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - # TODO(aryan): Fix the following test - @unittest.skip("This test fails with an error I haven't been able to debug yet.") - def test_simple_inference_save_pretrained(self): - pass - @unittest.skip("Not supported in HunyuanVideo.") def test_simple_inference_with_text_denoiser_block_scale(self): pass diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 521d647fbd8f..0062a5fbe421 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -329,20 +329,6 @@ def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_co ) return pipe, denoiser - def test_simple_inference(self): - """ - Tests a simple inference and makes sure it works as expected - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - _, _, inputs = self.get_dummy_inputs() - output_no_lora = pipe(**inputs)[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - @require_peft_version_greater("0.13.1") def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" From fc88ac4814782191f2a59cf63f8433e5bae303ef Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 16:08:01 +0530 Subject: [PATCH 03/19] move private methods to the bottom --- tests/lora/utils.py | 110 ++++++++++++++++++++++---------------------- 1 file changed, 55 insertions(+), 55 deletions(-) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 0062a5fbe421..3c194d92cdd8 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -251,61 +251,6 @@ def get_dummy_tokens(self): prepared_inputs["input_ids"] = inputs return prepared_inputs - def _get_lora_state_dicts(self, modules_to_save): - state_dicts = {} - for module_name, module in modules_to_save.items(): - if module is not None: - state_dicts[f"{module_name}_lora_layers"] = get_peft_model_state_dict(module) - return state_dicts - - def _get_lora_adapter_metadata(self, modules_to_save): - metadatas = {} - for module_name, module in modules_to_save.items(): - if module is not None: - metadatas[f"{module_name}_lora_adapter_metadata"] = module.peft_config["default"].to_dict() - return metadatas - - def _get_modules_to_save(self, pipe, has_denoiser=False): - modules_to_save = {} - lora_loadable_modules = self.pipeline_class._lora_loadable_modules - - if ( - "text_encoder" in lora_loadable_modules - and hasattr(pipe, "text_encoder") - and getattr(pipe.text_encoder, "peft_config", None) is not None - ): - modules_to_save["text_encoder"] = pipe.text_encoder - - if ( - "text_encoder_2" in lora_loadable_modules - and hasattr(pipe, "text_encoder_2") - and getattr(pipe.text_encoder_2, "peft_config", None) is not None - ): - modules_to_save["text_encoder_2"] = pipe.text_encoder_2 - - if has_denoiser: - if "unet" in lora_loadable_modules and hasattr(pipe, "unet"): - modules_to_save["unet"] = pipe.unet - - if "transformer" in lora_loadable_modules and hasattr(pipe, "transformer"): - modules_to_save["transformer"] = pipe.transformer - - return modules_to_save - - def _get_exclude_modules(self, pipe): - from diffusers.utils.peft_utils import _derive_exclude_modules - - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - denoiser = "unet" if self.unet_kwargs is not None else "transformer" - modules_to_save = {k: v for k, v in modules_to_save.items() if k == denoiser} - denoiser_lora_state_dict = self._get_lora_state_dicts(modules_to_save)[f"{denoiser}_lora_layers"] - pipe.unload_lora_weights() - denoiser_state_dict = pipe.unet.state_dict() if self.unet_kwargs is not None else pipe.transformer.state_dict() - exclude_modules = _derive_exclude_modules( - denoiser_state_dict, denoiser_lora_state_dict, adapter_name="default" - ) - return exclude_modules - def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"): if text_lora_config is not None: if "text_encoder" in self.pipeline_class._lora_loadable_modules: @@ -2408,3 +2353,58 @@ def test_group_offloading_inference_denoiser(self, offload_type, use_stream): # materializes the test methods on invocation which cannot be overridden. return self._test_group_offloading_inference_denoiser(offload_type, use_stream) + + def _get_lora_state_dicts(self, modules_to_save): + state_dicts = {} + for module_name, module in modules_to_save.items(): + if module is not None: + state_dicts[f"{module_name}_lora_layers"] = get_peft_model_state_dict(module) + return state_dicts + + def _get_lora_adapter_metadata(self, modules_to_save): + metadatas = {} + for module_name, module in modules_to_save.items(): + if module is not None: + metadatas[f"{module_name}_lora_adapter_metadata"] = module.peft_config["default"].to_dict() + return metadatas + + def _get_modules_to_save(self, pipe, has_denoiser=False): + modules_to_save = {} + lora_loadable_modules = self.pipeline_class._lora_loadable_modules + + if ( + "text_encoder" in lora_loadable_modules + and hasattr(pipe, "text_encoder") + and getattr(pipe.text_encoder, "peft_config", None) is not None + ): + modules_to_save["text_encoder"] = pipe.text_encoder + + if ( + "text_encoder_2" in lora_loadable_modules + and hasattr(pipe, "text_encoder_2") + and getattr(pipe.text_encoder_2, "peft_config", None) is not None + ): + modules_to_save["text_encoder_2"] = pipe.text_encoder_2 + + if has_denoiser: + if "unet" in lora_loadable_modules and hasattr(pipe, "unet"): + modules_to_save["unet"] = pipe.unet + + if "transformer" in lora_loadable_modules and hasattr(pipe, "transformer"): + modules_to_save["transformer"] = pipe.transformer + + return modules_to_save + + def _get_exclude_modules(self, pipe): + from diffusers.utils.peft_utils import _derive_exclude_modules + + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) + denoiser = "unet" if self.unet_kwargs is not None else "transformer" + modules_to_save = {k: v for k, v in modules_to_save.items() if k == denoiser} + denoiser_lora_state_dict = self._get_lora_state_dicts(modules_to_save)[f"{denoiser}_lora_layers"] + pipe.unload_lora_weights() + denoiser_state_dict = pipe.unet.state_dict() if self.unet_kwargs is not None else pipe.transformer.state_dict() + exclude_modules = _derive_exclude_modules( + denoiser_state_dict, denoiser_lora_state_dict, adapter_name="default" + ) + return exclude_modules From fb182693966255f9408a511847e42a0e3c06a673 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 16:14:25 +0530 Subject: [PATCH 04/19] more reorg. --- tests/lora/utils.py | 300 +++++++++++++++++++++++--------------------- 1 file changed, 155 insertions(+), 145 deletions(-) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 3c194d92cdd8..63e790da363f 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -129,151 +129,6 @@ class PeftLoraLoaderMixinTests: text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"] denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): - if self.unet_kwargs and self.transformer_kwargs: - raise ValueError("Both `unet_kwargs` and `transformer_kwargs` cannot be specified.") - if self.has_two_text_encoders and self.has_three_text_encoders: - raise ValueError("Both `has_two_text_encoders` and `has_three_text_encoders` cannot be True.") - - scheduler_cls = self.scheduler_cls if scheduler_cls is None else scheduler_cls - rank = 4 - lora_alpha = rank if lora_alpha is None else lora_alpha - - torch.manual_seed(0) - if self.unet_kwargs is not None: - unet = UNet2DConditionModel(**self.unet_kwargs) - else: - transformer = self.transformer_cls(**self.transformer_kwargs) - - scheduler = scheduler_cls(**self.scheduler_kwargs) - - torch.manual_seed(0) - vae = self.vae_cls(**self.vae_kwargs) - - text_encoder = self.text_encoder_cls.from_pretrained( - self.text_encoder_id, subfolder=self.text_encoder_subfolder - ) - tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id, subfolder=self.tokenizer_subfolder) - - if self.text_encoder_2_cls is not None: - text_encoder_2 = self.text_encoder_2_cls.from_pretrained( - self.text_encoder_2_id, subfolder=self.text_encoder_2_subfolder - ) - tokenizer_2 = self.tokenizer_2_cls.from_pretrained( - self.tokenizer_2_id, subfolder=self.tokenizer_2_subfolder - ) - - if self.text_encoder_3_cls is not None: - text_encoder_3 = self.text_encoder_3_cls.from_pretrained( - self.text_encoder_3_id, subfolder=self.text_encoder_3_subfolder - ) - tokenizer_3 = self.tokenizer_3_cls.from_pretrained( - self.tokenizer_3_id, subfolder=self.tokenizer_3_subfolder - ) - - text_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=self.text_encoder_target_modules, - init_lora_weights=False, - use_dora=use_dora, - ) - - denoiser_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=self.denoiser_target_modules, - init_lora_weights=False, - use_dora=use_dora, - ) - - pipeline_components = { - "scheduler": scheduler, - "vae": vae, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - } - # Denoiser - if self.unet_kwargs is not None: - pipeline_components.update({"unet": unet}) - elif self.transformer_kwargs is not None: - pipeline_components.update({"transformer": transformer}) - - # Remaining text encoders. - if self.text_encoder_2_cls is not None: - pipeline_components.update({"tokenizer_2": tokenizer_2, "text_encoder_2": text_encoder_2}) - if self.text_encoder_3_cls is not None: - pipeline_components.update({"tokenizer_3": tokenizer_3, "text_encoder_3": text_encoder_3}) - - # Remaining stuff - init_params = inspect.signature(self.pipeline_class.__init__).parameters - if "safety_checker" in init_params: - pipeline_components.update({"safety_checker": None}) - if "feature_extractor" in init_params: - pipeline_components.update({"feature_extractor": None}) - if "image_encoder" in init_params: - pipeline_components.update({"image_encoder": None}) - - return pipeline_components, text_lora_config, denoiser_lora_config - - @property - def output_shape(self): - raise NotImplementedError - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 10 - num_channels = 4 - sizes = (32, 32) - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "A painting of a squirrel eating a burger", - "num_inference_steps": 5, - "guidance_scale": 6.0, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - # Copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb - def get_dummy_tokens(self): - max_seq_length = 77 - - inputs = torch.randint(2, 56, size=(1, max_seq_length), generator=torch.manual_seed(0)) - - prepared_inputs = {} - prepared_inputs["input_ids"] = inputs - return prepared_inputs - - def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"): - if text_lora_config is not None: - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, adapter_name=adapter_name) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - if denoiser_lora_config is not None: - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, adapter_name=adapter_name) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - else: - denoiser = None - - if text_lora_config is not None and self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, adapter_name=adapter_name) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - return pipe, denoiser - @require_peft_version_greater("0.13.1") def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" @@ -2354,6 +2209,161 @@ def test_group_offloading_inference_denoiser(self, offload_type, use_stream): return self._test_group_offloading_inference_denoiser(offload_type, use_stream) + def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): + if self.unet_kwargs and self.transformer_kwargs: + raise ValueError("Both `unet_kwargs` and `transformer_kwargs` cannot be specified.") + if self.has_two_text_encoders and self.has_three_text_encoders: + raise ValueError("Both `has_two_text_encoders` and `has_three_text_encoders` cannot be True.") + + scheduler_cls = self.scheduler_cls if scheduler_cls is None else scheduler_cls + rank = 4 + lora_alpha = rank if lora_alpha is None else lora_alpha + + torch.manual_seed(0) + if self.unet_kwargs is not None: + unet = UNet2DConditionModel(**self.unet_kwargs) + else: + transformer = self.transformer_cls(**self.transformer_kwargs) + + scheduler = scheduler_cls(**self.scheduler_kwargs) + + torch.manual_seed(0) + vae = self.vae_cls(**self.vae_kwargs) + + text_encoder = self.text_encoder_cls.from_pretrained( + self.text_encoder_id, subfolder=self.text_encoder_subfolder + ) + tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id, subfolder=self.tokenizer_subfolder) + + if self.text_encoder_2_cls is not None: + text_encoder_2 = self.text_encoder_2_cls.from_pretrained( + self.text_encoder_2_id, subfolder=self.text_encoder_2_subfolder + ) + tokenizer_2 = self.tokenizer_2_cls.from_pretrained( + self.tokenizer_2_id, subfolder=self.tokenizer_2_subfolder + ) + + if self.text_encoder_3_cls is not None: + text_encoder_3 = self.text_encoder_3_cls.from_pretrained( + self.text_encoder_3_id, subfolder=self.text_encoder_3_subfolder + ) + tokenizer_3 = self.tokenizer_3_cls.from_pretrained( + self.tokenizer_3_id, subfolder=self.tokenizer_3_subfolder + ) + + text_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=self.text_encoder_target_modules, + init_lora_weights=False, + use_dora=use_dora, + ) + + denoiser_lora_config = LoraConfig( + r=rank, + lora_alpha=lora_alpha, + target_modules=self.denoiser_target_modules, + init_lora_weights=False, + use_dora=use_dora, + ) + + pipeline_components = { + "scheduler": scheduler, + "vae": vae, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + # Denoiser + if self.unet_kwargs is not None: + pipeline_components.update({"unet": unet}) + elif self.transformer_kwargs is not None: + pipeline_components.update({"transformer": transformer}) + + # Remaining text encoders. + if self.text_encoder_2_cls is not None: + pipeline_components.update({"tokenizer_2": tokenizer_2, "text_encoder_2": text_encoder_2}) + if self.text_encoder_3_cls is not None: + pipeline_components.update({"tokenizer_3": tokenizer_3, "text_encoder_3": text_encoder_3}) + + # Remaining stuff + init_params = inspect.signature(self.pipeline_class.__init__).parameters + if "safety_checker" in init_params: + pipeline_components.update({"safety_checker": None}) + if "feature_extractor" in init_params: + pipeline_components.update({"feature_extractor": None}) + if "image_encoder" in init_params: + pipeline_components.update({"image_encoder": None}) + + return pipeline_components, text_lora_config, denoiser_lora_config + + @property + def output_shape(self): + raise NotImplementedError + + def get_dummy_inputs(self, with_generator=True): + batch_size = 1 + sequence_length = 10 + num_channels = 4 + sizes = (32, 32) + + generator = torch.manual_seed(0) + noise = floats_tensor((batch_size, num_channels) + sizes) + input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) + + pipeline_inputs = { + "prompt": "A painting of a squirrel eating a burger", + "num_inference_steps": 5, + "guidance_scale": 6.0, + "output_type": "np", + } + if with_generator: + pipeline_inputs.update({"generator": generator}) + + return noise, input_ids, pipeline_inputs + + # Copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb + def get_dummy_tokens(self): + max_seq_length = 77 + + inputs = torch.randint(2, 56, size=(1, max_seq_length), generator=torch.manual_seed(0)) + + prepared_inputs = {} + prepared_inputs["input_ids"] = inputs + return prepared_inputs + + def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"): + if text_lora_config is not None: + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, adapter_name=adapter_name) + self.assertTrue( + check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" + ) + + if denoiser_lora_config is not None: + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, adapter_name=adapter_name) + self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + else: + denoiser = None + + if text_lora_config is not None and self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder_2.add_adapter(text_lora_config, adapter_name=adapter_name) + self.assertTrue( + check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" + ) + return pipe, denoiser + + def _setup_pipeline_and_get_base_output(self, scheduler_cls): + components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + pipe = self.pipeline_class(**components) + pipe = pipe.to(torch_device) + pipe.set_progress_bar_config(disable=None) + _, _, inputs = self.get_dummy_inputs(with_generator=False) + + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + return pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config + def _get_lora_state_dicts(self, modules_to_save): state_dicts = {} for module_name, module in modules_to_save.items(): From be922aebbc40b48dd2cec4470190a6f9e19631a4 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 16:27:18 +0530 Subject: [PATCH 05/19] parameterize more. --- tests/lora/utils.py | 191 ++++++++++++++++++-------------------------- 1 file changed, 79 insertions(+), 112 deletions(-) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 63e790da363f..8930352e58cc 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -244,15 +244,23 @@ def test_low_cpu_mem_usage_with_loading(self): "Loading from saved checkpoints with `low_cpu_mem_usage` should give same results.", ) - def test_simple_inference_with_text_lora_and_scale(self): + def test_simple_inference_with_partial_text_lora(self): """ - Tests a simple inference with lora attached on the text encoder + scale argument + Tests a simple inference with lora attached on the text encoder + with different ranks and some adapters removed and makes sure it works as expected """ - attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) + components, _, _ = self.get_dummy_components(scheduler_cls) + # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). + text_lora_config = LoraConfig( + r=4, + rank_pattern={self.text_encoder_target_modules[i]: i + 1 for i in range(3)}, + lora_alpha=4, + target_modules=self.text_encoder_target_modules, + init_lora_weights=False, + use_dora=False, + ) pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) @@ -263,25 +271,39 @@ def test_simple_inference_with_text_lora_and_scale(self): pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) + state_dict = {} + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + # Gather the state dict for the PEFT model, excluding `layers.4`, to ensure `load_lora_into_text_encoder` + # supports missing layers (PR#8324). + state_dict = { + f"text_encoder.{module_name}": param + for module_name, param in get_peft_model_state_dict(pipe.text_encoder).items() + if "text_model.encoder.layers.4" not in module_name + } - attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + if self.has_two_text_encoders or self.has_three_text_encoders: + if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: + state_dict.update( + { + f"text_encoder_2.{module_name}": param + for module_name, param in get_peft_model_state_dict(pipe.text_encoder_2).items() + if "text_model.encoder.layers.4" not in module_name + } + ) + output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", + not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" ) - attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} - output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + # Unload lora and load it back using the pipe.load_lora_weights machinery + pipe.unload_lora_weights() + pipe.load_lora_weights(state_dict) + output_partial_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), - "Lora + 0 scale should lead to same result as no LoRA", + not np.allclose(output_partial_lora, output_lora, atol=1e-3, rtol=1e-3), + "Removing adapters should change the output", ) @parameterized.expand([("fused",), ("unloaded",), ("save_load",)]) @@ -369,66 +391,57 @@ def test_lora_text_encoder_actions(self, action): "Loading from a saved checkpoint should yield the same result as the original LoRA.", ) - def test_simple_inference_with_partial_text_lora(self): + @parameterized.expand( + [ + ("text_encoder_only",), + ("text_and_denoiser",), + ] + ) + def test_lora_scaling(self, lora_components_to_add): """ - Tests a simple inference with lora attached on the text encoder - with different ranks and some adapters removed - and makes sure it works as expected + Tests inference with LoRA scaling applied via attention_kwargs + for different LoRA configurations. """ + if lora_components_to_add == "text_encoder_only": + if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): + pytest.skip( + "Test not supported for {self.__class__.__name__} since there is not text encoder in the LoRA loadable modules." + ) + attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) + for scheduler_cls in self.scheduler_classes: - components, _, _ = self.get_dummy_components(scheduler_cls) - # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). - text_lora_config = LoraConfig( - r=4, - rank_pattern={self.text_encoder_target_modules[i]: i + 1 for i in range(3)}, - lora_alpha=4, - target_modules=self.text_encoder_target_modules, - init_lora_weights=False, - use_dora=False, + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) ) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - - state_dict = {} - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - # Gather the state dict for the PEFT model, excluding `layers.4`, to ensure `load_lora_into_text_encoder` - # supports missing layers (PR#8324). - state_dict = { - f"text_encoder.{module_name}": param - for module_name, param in get_peft_model_state_dict(pipe.text_encoder).items() - if "text_model.encoder.layers.4" not in module_name - } - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - state_dict.update( - { - f"text_encoder_2.{module_name}": param - for module_name, param in get_peft_model_state_dict(pipe.text_encoder_2).items() - if "text_model.encoder.layers.4" not in module_name - } - ) + # Add LoRA components based on the parameterization + if lora_components_to_add == "text_encoder_only": + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + elif lora_components_to_add == "text_and_denoiser": + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + else: + raise ValueError(f"Unknown `lora_components_to_add`: {lora_components_to_add}") + # 1. Test base LoRA output output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" + self.assertFalse( + np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "LoRA should change the output." ) - # Unload lora and load it back using the pipe.load_lora_weights machinery - pipe.unload_lora_weights() - pipe.load_lora_weights(state_dict) + # 2. Test with a scale of 0.5 + attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} + output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + self.assertFalse( + np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), + "Using a LoRA scale should change the output.", + ) - output_partial_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + # 3. Test with a scale of 0.0, which should be identical to no LoRA + attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} + output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] self.assertTrue( - not np.allclose(output_partial_lora, output_lora, atol=1e-3, rtol=1e-3), - "Removing adapters should change the output", + np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), + "Using a LoRA scale of 0.0 should be the same as no LoRA.", ) def test_simple_inference_with_text_denoiser_lora_save_load(self): @@ -469,52 +482,6 @@ def test_simple_inference_with_text_denoiser_lora_save_load(self): "Loading from saved checkpoints should give same results.", ) - def test_simple_inference_with_text_denoiser_lora_and_scale(self): - """ - Tests a simple inference with lora attached on the text encoder + Unet + scale argument - and makes sure it works as expected - """ - attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) - - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) - - attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), - "Lora + scale should change the output", - ) - - attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} - output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - - self.assertTrue( - np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), - "Lora + 0 scale should lead to same result as no LoRA", - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - pipe.text_encoder.text_model.encoder.layers[0].self_attn.q_proj.scaling["default"] == 1.0, - "The scaling parameter has not been correctly restored!", - ) - def test_simple_inference_with_text_lora_denoiser_fused(self): """ Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model From 66f922cf76ad3eba93b32ba67151a5a1739f5f00 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 16:29:03 +0530 Subject: [PATCH 06/19] unify more. --- tests/lora/test_lora_layers_auraflow.py | 4 ---- tests/lora/test_lora_layers_cogvideox.py | 4 ---- tests/lora/test_lora_layers_cogview4.py | 4 ---- tests/lora/test_lora_layers_hunyuanvideo.py | 4 ---- tests/lora/test_lora_layers_ltx_video.py | 4 ---- tests/lora/test_lora_layers_lumina2.py | 4 ---- tests/lora/test_lora_layers_mochi.py | 4 ---- tests/lora/test_lora_layers_sana.py | 4 ---- tests/lora/test_lora_layers_wan.py | 4 ---- tests/lora/utils.py | 7 +------ 10 files changed, 1 insertion(+), 42 deletions(-) diff --git a/tests/lora/test_lora_layers_auraflow.py b/tests/lora/test_lora_layers_auraflow.py index c9821cecb066..82b723db9942 100644 --- a/tests/lora/test_lora_layers_auraflow.py +++ b/tests/lora/test_lora_layers_auraflow.py @@ -118,7 +118,3 @@ def test_modify_padding_mode(self): @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") def test_simple_inference_with_partial_text_lora(self): pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_text_lora_and_scale(self): - pass diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index f276832e26a5..383613ca0793 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -152,10 +152,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - @unittest.skip("Not supported in CogVideoX.") def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): pass diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index f94f3d5a521d..c60e538353c7 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -138,7 +138,3 @@ def test_modify_padding_mode(self): @unittest.skip("Text encoder LoRA is not supported in CogView4.") def test_simple_inference_with_partial_text_lora(self): pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_text_lora_and_scale(self): - pass diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index c6bc6b9cc205..6109e6a4de7d 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -172,10 +172,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - @nightly @require_torch_accelerator diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index 101d7ab4cba1..b8976ea056ac 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -129,7 +129,3 @@ def test_modify_padding_mode(self): @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") def test_simple_inference_with_partial_text_lora(self): pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_text_lora_and_scale(self): - pass diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py index 638309097f95..9bdfffb90258 100644 --- a/tests/lora/test_lora_layers_lumina2.py +++ b/tests/lora/test_lora_layers_lumina2.py @@ -117,10 +117,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") @skip_mps @pytest.mark.xfail( diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index 8d53506b026f..7f4f06646045 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -121,10 +121,6 @@ def test_modify_padding_mode(self): def test_simple_inference_with_partial_text_lora(self): pass - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - @unittest.skip("Not supported in Mochi.") def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): pass diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py index 0c2b6a0a9b00..d2c4fc25796e 100644 --- a/tests/lora/test_lora_layers_sana.py +++ b/tests/lora/test_lora_layers_sana.py @@ -120,7 +120,3 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Text encoder LoRA is not supported in SANA.") def test_simple_inference_with_partial_text_lora(self): pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_text_lora_and_scale(self): - pass diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index d103c1a2a63b..ef20e7e5aca9 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -125,7 +125,3 @@ def test_modify_padding_mode(self): @unittest.skip("Text encoder LoRA is not supported in Wan.") def test_simple_inference_with_partial_text_lora(self): pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_text_lora_and_scale(self): - pass diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 8930352e58cc..813940fdff7c 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -391,12 +391,7 @@ def test_lora_text_encoder_actions(self, action): "Loading from a saved checkpoint should yield the same result as the original LoRA.", ) - @parameterized.expand( - [ - ("text_encoder_only",), - ("text_and_denoiser",), - ] - ) + @parameterized.expand([("text_encoder_only",), ("text_and_denoiser",)]) def test_lora_scaling(self, lora_components_to_add): """ Tests inference with LoRA scaling applied via attention_kwargs From 705a9fd27e6150d931bb0893a99b2256a05312a5 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 16:59:29 +0530 Subject: [PATCH 07/19] unbloat more. --- tests/lora/test_lora_layers_cogvideox.py | 17 +- tests/lora/test_lora_layers_cogview4.py | 17 +- tests/lora/test_lora_layers_hunyuanvideo.py | 18 +- tests/lora/test_lora_layers_ltx_video.py | 18 +- tests/lora/test_lora_layers_mochi.py | 24 +- tests/lora/test_lora_layers_wan.py | 18 +- tests/lora/utils.py | 279 ++++++-------------- 7 files changed, 172 insertions(+), 219 deletions(-) diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index 383613ca0793..70fa0b48e1f6 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -123,8 +123,21 @@ def get_dummy_inputs(self, with_generator=True): def test_simple_inference_with_text_lora_denoiser_fused_multi(self): super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): + super().test_lora_actions(action, components_to_add, expected_atol=9e-3) def test_lora_scale_kwargs_match_fusion(self): super().test_lora_scale_kwargs_match_fusion(expected_atol=9e-3, expected_rtol=9e-3) diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index c60e538353c7..e9f8302ac5d6 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -113,8 +113,21 @@ def get_dummy_inputs(self, with_generator=True): def test_simple_inference_with_text_lora_denoiser_fused_multi(self): super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): + super().test_lora_actions(action, components_to_add, expected_atol=9e-3) @parameterized.expand([("block_level", True), ("leaf_level", False)]) @require_torch_accelerator diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index 6109e6a4de7d..fdaeb3adae36 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -19,6 +19,7 @@ import numpy as np import pytest import torch +from parameterized import parameterized from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast from diffusers import ( @@ -153,8 +154,21 @@ def get_dummy_inputs(self, with_generator=True): def test_simple_inference_with_text_lora_denoiser_fused_multi(self): super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): + super().test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in HunyuanVideo.") def test_simple_inference_with_text_denoiser_block_scale(self): diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index b8976ea056ac..adb987d968cb 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -16,6 +16,7 @@ import unittest import torch +from parameterized import parameterized from transformers import AutoTokenizer, T5EncoderModel from diffusers import ( @@ -111,8 +112,21 @@ def get_dummy_inputs(self, with_generator=True): def test_simple_inference_with_text_lora_denoiser_fused_multi(self): super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): + super().test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in LTXVideo.") def test_simple_inference_with_text_denoiser_block_scale(self): diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index 7f4f06646045..cf5ac4be35fa 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -16,14 +16,11 @@ import unittest import torch +from parameterized import parameterized from transformers import AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel -from diffusers.utils.testing_utils import ( - floats_tensor, - require_peft_backend, - skip_mps, -) +from diffusers.utils.testing_utils import floats_tensor, require_peft_backend, skip_mps sys.path.append(".") @@ -102,8 +99,21 @@ def get_dummy_inputs(self, with_generator=True): def test_simple_inference_with_text_lora_denoiser_fused_multi(self): super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): + super().test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in Mochi.") def test_simple_inference_with_text_denoiser_block_scale(self): diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index ef20e7e5aca9..b7c7072c696e 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -16,6 +16,7 @@ import unittest import torch +from parameterized import parameterized from transformers import AutoTokenizer, T5EncoderModel from diffusers import ( @@ -107,8 +108,21 @@ def get_dummy_inputs(self, with_generator=True): def test_simple_inference_with_text_lora_denoiser_fused_multi(self): super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): + super().test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in Wan.") def test_simple_inference_with_text_denoiser_block_scale(self): diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 813940fdff7c..27f0a2886313 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -306,91 +306,121 @@ def test_simple_inference_with_partial_text_lora(self): "Removing adapters should change the output", ) - @parameterized.expand([("fused",), ("unloaded",), ("save_load",)]) - def test_lora_text_encoder_actions(self, action): + def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3): """ - Tests various actions (fusing, unloading, saving/loading) on a pipeline - with LoRA applied to the text encoder(s). + A unified test for various LoRA actions (fusing, unloading, saving/loading, etc.) + on different combinations of model components. """ - if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): - pytest.skip( - "Test not supported for {self.__class__.__name__} since there is not text encoder in the LoRA loadable modules." - ) + # Skip text_encoder tests if the pipeline doesn't support it + if lora_components_to_add == "text_encoder_only" and not any( + "text_encoder" in k for k in self.pipeline_class._lora_loadable_modules + ): + pytest.skip(f"Test not supported for {self.__class__.__name__} without a LoRA-compatible text encoder.") + for scheduler_cls in self.scheduler_classes: - # 1. Setup pipeline and get base output without LoRA - components, text_lora_config, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + # 1. Setup pipeline and get base output + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + # 2. Add LoRA adapters based on the parameterization + if lora_components_to_add == "text_encoder_only": + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + elif lora_components_to_add == "text_and_denoiser": + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + else: + raise ValueError(f"Unknown `lora_components_to_add`: {lora_components_to_add}") + modules_to_save = self._get_modules_to_save( + pipe, has_denoiser=lora_components_to_add != "text_encoder_only" + ) - # 2. Add LoRA adapters - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - not np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "Lora should change the output" - ) + self.assertTrue(not np.allclose(output_lora, output_no_lora, atol=expected_atol, rtol=1e-3)) # 3. Perform the specified action and assert the outcome if action == "fused": pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) - # Fusing should still keep the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder)) - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder_2)) - + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - not np.allclose(output_after_action, output_no_lora, atol=1e-3, rtol=1e-3), + not np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=1e-3), "Fused LoRA should produce a different output from the base model.", ) elif action == "unloaded": pipe.unload_lora_weights() - # Unloading should remove the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder)) - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertFalse(check_if_lora_correctly_set(pipe.text_encoder_2)) - + for module_name, module in modules_to_save.items(): + self.assertTrue( + not check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}" + ) output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - np.allclose(output_after_action, output_no_lora, atol=1e-3, rtol=1e-3), + np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=1e-3), "Output after unloading LoRA should match the original output.", ) + elif action == "unfused": + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + output_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + + pipe.unfuse_lora() + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + output_unfused = pipe(**inputs, generator=torch.manual_seed(0))[0] + + self.assertTrue( + np.allclose(output_fused, output_unfused, atol=expected_atol, rtol=1e-3), + "Output after unfusing should match the fused output.", + ) + elif action == "save_load": with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe) + has_denoiser = lora_components_to_add == "text_and_denoiser" + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=has_denoiser) lora_state_dicts = self._get_lora_state_dicts(modules_to_save) + self.pipeline_class.save_lora_weights( save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts ) self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - # Unload and then load the weights back pipe.unload_lora_weights() pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) - for module_name, module in modules_to_save.items(): self.assertTrue( - check_if_lora_correctly_set(module), - f"LoRA not set correctly in {module_name} after loading.", + check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}" ) output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - np.allclose(output_lora, output_after_action, atol=1e-3, rtol=1e-3), - "Loading from a saved checkpoint should yield the same result as the original LoRA.", + np.allclose(output_lora, output_after_action, atol=expected_atol, rtol=1e-3), + "Loading from a saved checkpoint should yield the same result.", ) + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3): + for cls in inspect.getmro(self.__class__): + if "test_lora_actions" in cls.__dict__ and cls is not PeftLoraLoaderMixinTests: + # Skip this test if it is overwritten by child class. We need to do this because parameterized + # materializes the test methods on invocation which cannot be overridden. + return + self._test_lora_actions(action, lora_components_to_add, expected_atol) + @parameterized.expand([("text_encoder_only",), ("text_and_denoiser",)]) def test_lora_scaling(self, lora_components_to_add): """ @@ -426,8 +456,8 @@ def test_lora_scaling(self, lora_components_to_add): # 2. Test with a scale of 0.5 attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - self.assertFalse( - np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), + self.assertTrue( + not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), "Using a LoRA scale should change the output.", ) @@ -439,161 +469,6 @@ def test_lora_scaling(self, lora_components_to_add): "Using a LoRA scale of 0.0 should be the same as no LoRA.", ) - def test_simple_inference_with_text_denoiser_lora_save_load(self): - """ - Tests a simple usecase where users could use saving utilities for LoRA for Unet + text encoder - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights( - save_directory=tmpdirname, safe_serialization=False, **lora_state_dicts - ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) - pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_from_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora, images_lora_from_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - - def test_simple_inference_with_text_lora_denoiser_fused(self): - """ - Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model - and makes sure it works as expected - with unet - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) - - # Fusing should still keep the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - output_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_fused, output_no_lora, atol=1e-3, rtol=1e-3), "Fused lora should change the output" - ) - - def test_simple_inference_with_text_denoiser_lora_unloaded(self): - """ - Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights - and makes sure it works as expected - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.unload_lora_weights() - # unloading should remove the LoRA layers - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly unloaded in text encoder" - ) - self.assertFalse(check_if_lora_correctly_set(denoiser), "Lora not correctly unloaded in denoiser") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertFalse( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly unloaded in text encoder 2", - ) - - output_unloaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_unloaded, output_no_lora, atol=1e-3, rtol=1e-3), - "Fused lora should change the output", - ) - - def test_simple_inference_with_text_denoiser_lora_unfused( - self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3 - ): - """ - Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights - and makes sure it works as expected - """ - for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - pipe, denoiser = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - output_fused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - output_unfused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # unloading should remove the LoRA layers - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") - - self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" - ) - - # Fuse and unfuse should lead to the same results - self.assertTrue( - np.allclose(output_fused_lora, output_unfused_lora, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - def test_simple_inference_with_text_denoiser_multi_adapter(self): """ Tests a simple inference with lora attached to text encoder and unet, attaches From 395f9d93639b34a1121bc730de0b09792d33060f Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 17:47:46 +0530 Subject: [PATCH 08/19] update --- tests/lora/test_lora_layers_lumina2.py | 8 +- tests/lora/utils.py | 291 ++++++++----------------- 2 files changed, 91 insertions(+), 208 deletions(-) diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py index 9bdfffb90258..b6f04feaea91 100644 --- a/tests/lora/test_lora_layers_lumina2.py +++ b/tests/lora/test_lora_layers_lumina2.py @@ -126,11 +126,9 @@ def test_simple_inference_with_partial_text_lora(self): ) def test_lora_fuse_nan(self): for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 27f0a2886313..00913d7c99b3 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -133,10 +133,9 @@ class PeftLoraLoaderMixinTests: def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: inject_adapter_in_model(text_lora_config, pipe.text_encoder, low_cpu_mem_usage=True) @@ -196,17 +195,11 @@ def test_low_cpu_mem_usage_with_loading(self): """Tests if we can load LoRA state dict with low_cpu_mem_usage.""" for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] with tempfile.TemporaryDirectory() as tmpdirname: @@ -251,7 +244,7 @@ def test_simple_inference_with_partial_text_lora(self): and makes sure it works as expected """ for scheduler_cls in self.scheduler_classes: - components, _, _ = self.get_dummy_components(scheduler_cls) + pipe, inputs, output_no_lora, _, _ = self._setup_pipeline_and_get_base_output(scheduler_cls) # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). text_lora_config = LoraConfig( r=4, @@ -261,13 +254,6 @@ def test_simple_inference_with_partial_text_lora(self): init_lora_weights=False, use_dora=False, ) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) @@ -475,13 +461,9 @@ def test_simple_inference_with_text_denoiser_multi_adapter(self): multiple adapters and set them """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -552,11 +534,9 @@ def test_wrong_adapter_name_raises_error(self): adapter_name = "adapter-1" scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) pipe, _ = self.add_adapters_to_pipeline( pipe, text_lora_config, denoiser_lora_config, adapter_name=adapter_name @@ -574,11 +554,9 @@ def test_wrong_adapter_name_raises_error(self): def test_multiple_wrong_adapter_name_raises_error(self): adapter_name = "adapter-1" scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) pipe, _ = self.add_adapters_to_pipeline( pipe, text_lora_config, denoiser_lora_config, adapter_name=adapter_name @@ -604,13 +582,9 @@ def test_simple_inference_with_text_denoiser_block_scale(self): one adapter and set different weights for different blocks (i.e. block lora) """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") @@ -661,13 +635,9 @@ def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): multiple adapters and set different weights for different blocks (i.e. block lora) """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -795,12 +765,9 @@ def all_possible_dict_opts(unet, value): return opts - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(self.scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - + pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + self.scheduler_classes[0] + ) pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet @@ -824,13 +791,9 @@ def test_simple_inference_with_text_denoiser_multi_adapter_delete_adapter(self): multiple adapters and set/delete them """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -918,13 +881,9 @@ def test_simple_inference_with_text_denoiser_multi_adapter_weighted(self): multiple adapters and set them """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -996,11 +955,9 @@ def test_simple_inference_with_text_denoiser_multi_adapter_weighted(self): ) def test_lora_fuse_nan(self): for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1058,11 +1015,9 @@ def test_get_adapters(self): are the expected results """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1087,10 +1042,9 @@ def test_get_list_adapters(self): are the expected results """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) # 1. dicts_to_be_checked = {} @@ -1162,14 +1116,9 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi( and makes sure it works as expected - with unet and multi-adapter case """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1243,14 +1192,9 @@ def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expec for lora_scale in [1.0, 0.8]: for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") @@ -1296,19 +1240,10 @@ def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expec @require_peft_version_greater(peft_version="0.9.0") def test_simple_inference_with_dora(self): for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components( - scheduler_cls, use_dora=True + pipe, inputs, output_no_dora_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls, use_dora=True) ) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_dora_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - output_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertFalse( @@ -1319,10 +1254,7 @@ def test_simple_inference_with_dora(self): def test_missing_keys_warning(self): scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. - components, _, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe, _, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output(scheduler_cls) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet denoiser.add_adapter(denoiser_lora_config) @@ -1356,10 +1288,7 @@ def test_missing_keys_warning(self): def test_unexpected_keys_warning(self): scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. - components, _, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe, _, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output(scheduler_cls) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet denoiser.add_adapter(denoiser_lora_config) @@ -1392,11 +1321,9 @@ def test_simple_inference_with_text_denoiser_lora_unfused_torch_compile(self): and makes sure it works as expected """ for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) @@ -1430,13 +1357,8 @@ def set_pad_mode(network, mode="circular"): def test_logs_info_when_no_lora_keys_found(self): scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. - components, _, _ = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - _, _, inputs = self.get_dummy_inputs(with_generator=False) - original_out = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, output_no_lora, _, _ = self._setup_pipeline_and_get_base_output(scheduler_cls) + output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] no_op_state_dict = {"lora_foo": torch.tensor(2.0), "lora_bar": torch.tensor(3.0)} logger = logging.get_logger("diffusers.loaders.peft") @@ -1448,7 +1370,7 @@ def test_logs_info_when_no_lora_keys_found(self): denoiser = getattr(pipe, "unet") if self.unet_kwargs is not None else getattr(pipe, "transformer") self.assertTrue(cap_logger.out.startswith(f"No LoRA keys associated to {denoiser.__class__.__name__}")) - self.assertTrue(np.allclose(original_out, out_after_lora_attempt, atol=1e-5, rtol=1e-5)) + self.assertTrue(np.allclose(output_no_lora, out_after_lora_attempt, atol=1e-5, rtol=1e-5)) # test only for text encoder for lora_module in self.pipeline_class._lora_loadable_modules: @@ -1476,15 +1398,9 @@ def test_set_adapters_match_attention_kwargs(self): attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) lora_scale = 0.5 @@ -1514,9 +1430,7 @@ def test_set_adapters_match_attention_kwargs(self): ) self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe.unload_lora_weights() pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) for module_name, module in modules_to_save.items(): @@ -1540,10 +1454,7 @@ def test_set_adapters_match_attention_kwargs(self): def test_lora_B_bias(self): # Currently, this test is only relevant for Flux Control LoRA as we are not # aware of any other LoRA checkpoint that has its `lora_B` biases trained. - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe, inputs, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output(self.scheduler_classes[0]) # keep track of the bias values of the base layers to perform checks later. bias_values = {} @@ -1577,13 +1488,9 @@ def test_lora_B_bias(self): self.assertFalse(np.allclose(lora_bias_false_output, lora_bias_true_output, atol=1e-3, rtol=1e-3)) def test_correct_lora_configs_with_different_ranks(self): - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - original_output = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, original_output, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(self.scheduler_classes[0]) + ) if self.unet_kwargs is not None: pipe.unet.add_adapter(denoiser_lora_config, "adapter-1") @@ -1729,10 +1636,7 @@ def check_module(denoiser): self.assertTrue(module._diffusers_hook.get_hook(_PEFT_AUTOCAST_DISABLE_HOOK) is not None) # 1. Test forward with add_adapter - components, _, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device, dtype=compute_dtype) - pipe.set_progress_bar_config(disable=None) + pipe, inputs, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output(self.scheduler_classes[0]) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet denoiser.add_adapter(denoiser_lora_config) @@ -1747,7 +1651,6 @@ def check_module(denoiser): ) check_module(denoiser) - _, _, inputs = self.get_dummy_inputs(with_generator=False) pipe(**inputs, generator=torch.manual_seed(0))[0] # 2. Test forward with load_lora_weights @@ -1759,10 +1662,9 @@ def check_module(denoiser): ) self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - components, _, _ = self.get_dummy_components(self.scheduler_classes[0]) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device, dtype=compute_dtype) - pipe.set_progress_bar_config(disable=None) + pipe, inputs, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + self.scheduler_classes[0] + ) pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet @@ -1774,17 +1676,14 @@ def check_module(denoiser): ) check_module(denoiser) - _, _, inputs = self.get_dummy_inputs(with_generator=False) pipe(**inputs, generator=torch.manual_seed(0))[0] @parameterized.expand([4, 8, 16]) def test_lora_adapter_metadata_is_loaded_correctly(self, lora_alpha): scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components( + pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls, lora_alpha=lora_alpha ) - pipe = self.pipeline_class(**components) - pipe, _ = self.add_adapters_to_pipeline( pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config ) @@ -1829,15 +1728,9 @@ def test_lora_adapter_metadata_is_loaded_correctly(self, lora_alpha): @parameterized.expand([4, 8, 16]) def test_lora_adapter_metadata_save_load_inference(self, lora_alpha): scheduler_cls = self.scheduler_classes[0] - components, text_lora_config, denoiser_lora_config = self.get_dummy_components( + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls, lora_alpha=lora_alpha ) - pipe = self.pipeline_class(**components).to(torch_device) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) - pipe, _ = self.add_adapters_to_pipeline( pipe, text_lora_config=text_lora_config, denoiser_lora_config=denoiser_lora_config ) @@ -1892,7 +1785,6 @@ def test_lora_exclude_modules(self): _, _, inputs = self.get_dummy_inputs(with_generator=False) output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(output_no_lora.shape == self.output_shape) # only supported for `denoiser` now pipe_cp = copy.deepcopy(pipe) @@ -1931,13 +1823,9 @@ def test_lora_exclude_modules(self): def test_inference_load_delete_load_adapters(self): "Tests if `load_lora_weights()` -> `delete_adapters()` -> `load_lora_weights()` works." for scheduler_cls in self.scheduler_classes: - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - _, _, inputs = self.get_dummy_inputs(with_generator=False) - - output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) if "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config) @@ -1982,10 +1870,7 @@ def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): onload_device = torch_device offload_device = torch.device("cpu") - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(self.scheduler_classes[0]) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe, inputs, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output(self.scheduler_classes[0]) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet denoiser.add_adapter(denoiser_lora_config) @@ -1997,17 +1882,14 @@ def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): self.pipeline_class.save_lora_weights( save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts ) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) - components, _, _ = self.get_dummy_components(self.scheduler_classes[0]) - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe, inputs, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + self.scheduler_classes[0] + ) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - check_if_lora_correctly_set(denoiser) - _, _, inputs = self.get_dummy_inputs(with_generator=False) + pipe.unload_lora_weights() + pipe.load_lora_weights(tmpdirname) + self.assertTrue(check_if_lora_correctly_set(denoiser)) # Test group offloading with load_lora_weights denoiser.enable_group_offload( @@ -2025,11 +1907,11 @@ def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): pipe.unload_lora_weights() group_offload_hook_2 = _get_top_level_group_offload_hook(denoiser) self.assertTrue(group_offload_hook_2 is not None) - output_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] # noqa: F841 + _ = pipe(**inputs, generator=torch.manual_seed(0))[0] # noqa: F841 # Add the lora again and check if group offloading works - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - check_if_lora_correctly_set(denoiser) + pipe.load_lora_weights(tmpdirname) + self.assertTrue(check_if_lora_correctly_set(denoiser)) group_offload_hook_3 = _get_top_level_group_offload_hook(denoiser) self.assertTrue(group_offload_hook_3 is not None) output_3 = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -2191,14 +2073,17 @@ def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_co ) return pipe, denoiser - def _setup_pipeline_and_get_base_output(self, scheduler_cls): - components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls) + def _setup_pipeline_and_get_base_output(self, scheduler_cls, lora_alpha=4, use_dora=False): + components, text_lora_config, denoiser_lora_config = self.get_dummy_components( + scheduler_cls, lora_alpha=lora_alpha, use_dora=use_dora + ) pipe = self.pipeline_class(**components) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) _, _, inputs = self.get_dummy_inputs(with_generator=False) output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(output_no_lora.shape == self.output_shape) return pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config def _get_lora_state_dicts(self, modules_to_save): From 9fd5c9a716df13d1569b86b83985d43b6426f87e Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 17:53:36 +0530 Subject: [PATCH 09/19] updatye --- tests/lora/test_lora_layers_sdxl.py | 20 ++++++++++++++++---- tests/lora/utils.py | 12 ++++++------ 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/tests/lora/test_lora_layers_sdxl.py b/tests/lora/test_lora_layers_sdxl.py index 267650056aad..3e33a9e71a4e 100644 --- a/tests/lora/test_lora_layers_sdxl.py +++ b/tests/lora/test_lora_layers_sdxl.py @@ -22,6 +22,7 @@ import numpy as np import torch from packaging import version +from parameterized import parameterized from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer from diffusers import ( @@ -117,7 +118,20 @@ def tearDown(self): def test_multiple_wrong_adapter_name_raises_error(self): super().test_multiple_wrong_adapter_name_raises_error() - def test_simple_inference_with_text_denoiser_lora_unfused(self): + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): if torch.cuda.is_available(): expected_atol = 9e-2 expected_rtol = 9e-2 @@ -125,9 +139,7 @@ def test_simple_inference_with_text_denoiser_lora_unfused(self): expected_atol = 1e-3 expected_rtol = 1e-3 - super().test_simple_inference_with_text_denoiser_lora_unfused( - expected_atol=expected_atol, expected_rtol=expected_rtol - ) + super().test_lora_actions(expected_atol=expected_atol, expected_rtol=expected_rtol) def test_simple_inference_with_text_lora_denoiser_fused_multi(self): if torch.cuda.is_available(): diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 00913d7c99b3..7e1f0faccf2d 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -292,7 +292,7 @@ def test_simple_inference_with_partial_text_lora(self): "Removing adapters should change the output", ) - def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3): + def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, expected_rtol=1e-3): """ A unified test for various LoRA actions (fusing, unloading, saving/loading, etc.) on different combinations of model components. @@ -321,7 +321,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3) ) output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(not np.allclose(output_lora, output_no_lora, atol=expected_atol, rtol=1e-3)) + self.assertTrue(not np.allclose(output_lora, output_no_lora, atol=expected_atol, rtol=expected_rtol)) # 3. Perform the specified action and assert the outcome if action == "fused": @@ -330,7 +330,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3) self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - not np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=1e-3), + not np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=expected_rtol), "Fused LoRA should produce a different output from the base model.", ) @@ -342,7 +342,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3) ) output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=1e-3), + np.allclose(output_after_action, output_no_lora, atol=expected_atol, rtol=expected_rtol), "Output after unloading LoRA should match the original output.", ) @@ -358,7 +358,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3) output_unfused = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - np.allclose(output_fused, output_unfused, atol=expected_atol, rtol=1e-3), + np.allclose(output_fused, output_unfused, atol=expected_atol, rtol=expected_rtol), "Output after unfusing should match the fused output.", ) @@ -382,7 +382,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3) output_after_action = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( - np.allclose(output_lora, output_after_action, atol=expected_atol, rtol=1e-3), + np.allclose(output_lora, output_after_action, atol=expected_atol, rtol=expected_rtol), "Loading from a saved checkpoint should yield the same result.", ) From 1d6aa8a23533066d75cb3631ec41297ae48a317c Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 17:56:54 +0530 Subject: [PATCH 10/19] update --- tests/lora/utils.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 7e1f0faccf2d..44c6f31b1968 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -211,7 +211,7 @@ def test_low_cpu_mem_usage_with_loading(self): self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=False) + pipe.load_lora_weights(tmpdirname, low_cpu_mem_usage=False) for module_name, module in modules_to_save.items(): self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") @@ -224,7 +224,7 @@ def test_low_cpu_mem_usage_with_loading(self): # Now, check for `low_cpu_mem_usage.` pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin"), low_cpu_mem_usage=True) + pipe.load_lora_weights(tmpdirname, low_cpu_mem_usage=True) for module_name, module in modules_to_save.items(): self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") @@ -374,7 +374,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin"))) pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.bin")) + pipe.load_lora_weights(tmpdirname) for module_name, module in modules_to_save.items(): self.assertTrue( check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}" @@ -1428,10 +1428,8 @@ def test_set_adapters_match_attention_kwargs(self): self.pipeline_class.save_lora_weights( save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts ) - - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) pipe.unload_lora_weights() - pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) + pipe.load_lora_weights(os.path.join(tmpdirname)) for module_name, module in modules_to_save.items(): self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") @@ -1661,7 +1659,6 @@ def check_module(denoiser): save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts ) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) pipe, inputs, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output( self.scheduler_classes[0] ) @@ -1851,7 +1848,6 @@ def test_inference_load_delete_load_adapters(self): modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) lora_state_dicts = self._get_lora_state_dicts(modules_to_save) self.pipeline_class.save_lora_weights(save_directory=tmpdirname, **lora_state_dicts) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))) # First, delete adapter and compare. pipe.delete_adapters(pipe.get_active_adapters()[0]) From e56b8d709459cdfc08c268f19ac05931f37e01c3 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 18:24:16 +0530 Subject: [PATCH 11/19] update --- tests/lora/test_lora_layers_auraflow.py | 12 - tests/lora/test_lora_layers_cogvideox.py | 16 +- tests/lora/test_lora_layers_cogview4.py | 12 - tests/lora/test_lora_layers_flux.py | 24 +- tests/lora/test_lora_layers_hunyuanvideo.py | 12 - tests/lora/test_lora_layers_ltx_video.py | 12 - tests/lora/test_lora_layers_lumina2.py | 12 - tests/lora/test_lora_layers_mochi.py | 16 +- tests/lora/test_lora_layers_sana.py | 12 - tests/lora/test_lora_layers_sd3.py | 12 +- tests/lora/test_lora_layers_wan.py | 12 - tests/lora/utils.py | 422 ++++---------------- 12 files changed, 85 insertions(+), 489 deletions(-) diff --git a/tests/lora/test_lora_layers_auraflow.py b/tests/lora/test_lora_layers_auraflow.py index 82b723db9942..45a4c4834a16 100644 --- a/tests/lora/test_lora_layers_auraflow.py +++ b/tests/lora/test_lora_layers_auraflow.py @@ -103,18 +103,6 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - @unittest.skip("Not supported in AuraFlow.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in AuraFlow.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in AuraFlow.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in AuraFlow.") - def test_simple_inference_with_partial_text_lora(self): - pass diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index 70fa0b48e1f6..d08f25a6875e 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -149,22 +149,8 @@ def test_group_offloading_inference_denoiser(self, offload_type, use_stream): # The reason for this can be found here: https://github.com/huggingface/diffusers/pull/11804#issuecomment-3013325338 super()._test_group_offloading_inference_denoiser(offload_type, use_stream) - @unittest.skip("Not supported in CogVideoX.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in CogVideoX.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in CogVideoX.") def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in CogVideoX.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Not supported in CogVideoX.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass + # TODO: skip them properly diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index e9f8302ac5d6..df9e8fbd5f3f 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -136,18 +136,6 @@ def test_group_offloading_inference_denoiser(self, offload_type, use_stream): # The reason for this can be found here: https://github.com/huggingface/diffusers/pull/11804#issuecomment-3013325338 super()._test_group_offloading_inference_denoiser(offload_type, use_stream) - @unittest.skip("Not supported in CogView4.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in CogView4.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in CogView4.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in CogView4.") - def test_simple_inference_with_partial_text_lora(self): - pass diff --git a/tests/lora/test_lora_layers_flux.py b/tests/lora/test_lora_layers_flux.py index 336ac2246fd2..4cf40e16c60a 100644 --- a/tests/lora/test_lora_layers_flux.py +++ b/tests/lora/test_lora_layers_flux.py @@ -263,21 +263,11 @@ def test_lora_expansion_works_for_extra_keys(self): "LoRA should lead to different results.", ) - @unittest.skip("Not supported in Flux.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Flux.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in Flux.") def test_modify_padding_mode(self): pass - @unittest.skip("Not supported in Flux.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass + # TODO: skip them properly class FluxControlLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): @@ -791,21 +781,11 @@ def test_lora_unload_with_parameter_expanded_shapes_and_no_reset(self): self.assertTrue(pipe.transformer.x_embedder.weight.data.shape[1] == in_features * 2) self.assertTrue(pipe.transformer.config.in_channels == in_features * 2) - @unittest.skip("Not supported in Flux.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Flux.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in Flux.") def test_modify_padding_mode(self): pass - @unittest.skip("Not supported in Flux.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass + # TODO: skip them properly @slow diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index fdaeb3adae36..be4bf12575d1 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -170,22 +170,10 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): def test_lora_actions(self, action, components_to_add): super().test_lora_actions(action, components_to_add, expected_atol=9e-3) - @unittest.skip("Not supported in HunyuanVideo.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in HunyuanVideo.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in HunyuanVideo.") def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in HunyuanVideo.") - def test_simple_inference_with_partial_text_lora(self): - pass - @nightly @require_torch_accelerator diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index adb987d968cb..eee0ca7784ee 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -128,18 +128,6 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): def test_lora_actions(self, action, components_to_add): super().test_lora_actions(action, components_to_add, expected_atol=9e-3) - @unittest.skip("Not supported in LTXVideo.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in LTXVideo.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in LTXVideo.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in LTXVideo.") - def test_simple_inference_with_partial_text_lora(self): - pass diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py index b6f04feaea91..63eb0e41b65e 100644 --- a/tests/lora/test_lora_layers_lumina2.py +++ b/tests/lora/test_lora_layers_lumina2.py @@ -101,22 +101,10 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - @unittest.skip("Not supported in Lumina2.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Lumina2.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in Lumina2.") def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") - def test_simple_inference_with_partial_text_lora(self): - pass - @unittest.skip("Text encoder LoRA is not supported in Lumina2.") @skip_mps @pytest.mark.xfail( diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index cf5ac4be35fa..3cab8dfd2c2c 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -115,22 +115,8 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): def test_lora_actions(self, action, components_to_add): super().test_lora_actions(action, components_to_add, expected_atol=9e-3) - @unittest.skip("Not supported in Mochi.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Mochi.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in Mochi.") def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in Mochi.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Not supported in Mochi.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass + # TODO: skip them properly diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py index d2c4fc25796e..5606d06a9945 100644 --- a/tests/lora/test_lora_layers_sana.py +++ b/tests/lora/test_lora_layers_sana.py @@ -108,15 +108,3 @@ def get_dummy_inputs(self, with_generator=True): @unittest.skip("Not supported in SANA.") def test_modify_padding_mode(self): pass - - @unittest.skip("Not supported in SANA.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in SANA.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in SANA.") - def test_simple_inference_with_partial_text_lora(self): - pass diff --git a/tests/lora/test_lora_layers_sd3.py b/tests/lora/test_lora_layers_sd3.py index 8a8f2a676df1..344fa0aedfa5 100644 --- a/tests/lora/test_lora_layers_sd3.py +++ b/tests/lora/test_lora_layers_sd3.py @@ -114,17 +114,7 @@ def test_sd3_lora(self): lora_filename = "lora_peft_format.safetensors" pipe.load_lora_weights(lora_model_id, weight_name=lora_filename) - @unittest.skip("Not supported in SD3.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in SD3.") - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - pass - - @unittest.skip("Not supported in SD3.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass + # TODO: skip them properly @unittest.skip("Not supported in SD3.") def test_modify_padding_mode(self): diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index b7c7072c696e..62fddb96a7e0 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -124,18 +124,6 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): def test_lora_actions(self, action, components_to_add): super().test_lora_actions(action, components_to_add, expected_atol=9e-3) - @unittest.skip("Not supported in Wan.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Wan.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - @unittest.skip("Not supported in Wan.") def test_modify_padding_mode(self): pass - - @unittest.skip("Text encoder LoRA is not supported in Wan.") - def test_simple_inference_with_partial_text_lora(self): - pass diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 44c6f31b1968..44b4f8333e83 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -243,6 +243,8 @@ def test_simple_inference_with_partial_text_lora(self): with different ranks and some adapters removed and makes sure it works as expected """ + if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): + pytest.skip("Test not supported.") for scheduler_cls in self.scheduler_classes: pipe, inputs, output_no_lora, _, _ = self._setup_pipeline_and_get_base_output(scheduler_cls) # Verify `StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder` handles different ranks per module (PR#8324). @@ -455,81 +457,6 @@ def test_lora_scaling(self, lora_components_to_add): "Using a LoRA scale of 0.0 should be the same as no LoRA.", ) - def test_simple_inference_with_text_denoiser_multi_adapter(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set them - """ - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls) - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_1, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) - - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) - - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_no_lora, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter outputs should be different.", - ) - - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - def test_wrong_adapter_name_raises_error(self): adapter_name = "adapter-1" @@ -581,6 +508,8 @@ def test_simple_inference_with_text_denoiser_block_scale(self): Tests a simple inference with lora attached to text encoder and unet, attaches one adapter and set different weights for different blocks (i.e. block lora) """ + if self.unet_kwargs is None: + pytest.skip("Test is only supposed for UNets.") for scheduler_cls in self.scheduler_classes: pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( self._setup_pipeline_and_get_base_output(scheduler_cls) @@ -629,78 +558,10 @@ def test_simple_inference_with_text_denoiser_block_scale(self): "output with no lora and output with lora disabled should give same results", ) - def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set different weights for different blocks (i.e. block lora) - """ - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls) - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - scales_1 = {"text_encoder": 2, "unet": {"down": 5}} - scales_2 = {"unet": {"down": 5, "mid": 5}} - - pipe.set_adapters("adapter-1", scales_1) - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters("adapter-2", scales_2) - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1, scales_2]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - # a mismatching number of adapter_names and adapter_weights should raise an error - with self.assertRaises(ValueError): - pipe.set_adapters(["adapter-1", "adapter-2"], [scales_1]) - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): """Tests that any valid combination of lora block scales can be used in pipe.set_adapter""" + if self.unet_kwargs is None: + pytest.skip("Test not supported.") def updown_options(blocks_with_tf, layers_per_block, value): """ @@ -785,168 +646,6 @@ def all_possible_dict_opts(unet, value): pipe.set_adapters("adapter-1", scale_dict) # test will fail if this line throws an error - def test_simple_inference_with_text_denoiser_multi_adapter_delete_adapter(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set/delete them - """ - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls) - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.delete_adapters("adapter-1") - output_deleted_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_deleted_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - pipe.delete_adapters("adapter-2") - output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - pipe.set_adapters(["adapter-1", "adapter-2"]) - pipe.delete_adapters(["adapter-1", "adapter-2"]) - - output_deleted_adapters = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_deleted_adapters, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - - def test_simple_inference_with_text_denoiser_multi_adapter_weighted(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - multiple adapters and set them - """ - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls) - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # Fuse and unfuse should lead to the same results - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3), - "Adapter 1 and 2 should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 1 and mixed adapters should give different results", - ) - - self.assertFalse( - np.allclose(output_adapter_2, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Adapter 2 and mixed adapters should give different results", - ) - - pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) - output_adapter_mixed_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse( - np.allclose(output_adapter_mixed_weighted, output_adapter_mixed, atol=1e-3, rtol=1e-3), - "Weighted adapter and mixed adapter should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - @skip_mps @pytest.mark.xfail( condition=torch.device(torch_device).type == "cpu" and is_torch_version(">=", "2.5"), @@ -1817,48 +1516,69 @@ def test_lora_exclude_modules(self): "Lora outputs should match.", ) - def test_inference_load_delete_load_adapters(self): - "Tests if `load_lora_weights()` -> `delete_adapters()` -> `load_lora_weights()` works." + @parameterized.expand( + [ + ("simple",), + ("weighted",), + ("block_lora",), + ("delete_adapter",), + ("load_delete_load",), + ] + ) + def test_multi_adapter_scenarios(self, scenario): + """ + A unified test for various multi-adapter LoRA scenarios, including weighting, + block scaling, and adapter deletion. + """ for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls) - ) + pipe, inputs, output_no_lora, _ = self._setup_multi_adapter_pipeline(scheduler_cls) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + # Run inference with each adapter individually and mixed + pipe.set_adapters("adapter-1") + output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) + pipe.set_adapters("adapter-2") + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.set_adapters(["adapter-1", "adapter-2"]) + output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - with tempfile.TemporaryDirectory() as tmpdirname: - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) - lora_state_dicts = self._get_lora_state_dicts(modules_to_save) - self.pipeline_class.save_lora_weights(save_directory=tmpdirname, **lora_state_dicts) + # --- Assert base multi-adapter behavior --- + self.assertFalse(np.allclose(output_no_lora, output_adapter_1, atol=1e-3, rtol=1e-3)) + self.assertFalse(np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3)) + self.assertFalse(np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3)) + + # --- Scenario-specific logic --- + if scenario == "weighted": + pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) + output_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse(np.allclose(output_weighted, output_adapter_mixed, atol=1e-3, rtol=1e-3)) + + elif scenario == "block_lora": + scales = {"unet": {"down": 0.8, "mid": 0.5, "up": 0.2}} + pipe.set_adapters(["adapter-1", "adapter-2"], [scales, scales]) + output_block_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse(np.allclose(output_block_lora, output_adapter_mixed, atol=1e-3, rtol=1e-3)) + + elif scenario == "delete_adapter": + pipe.delete_adapters("adapter-1") + output_after_delete = pipe(**inputs, generator=torch.manual_seed(0))[0] + # After deleting adapter-1, the output should be the same as using only adapter-2 + self.assertTrue(np.allclose(output_after_delete, output_adapter_2, atol=1e-3, rtol=1e-3)) + + elif scenario == "load_delete_load": + with tempfile.TemporaryDirectory() as tmpdirname: + pipe.save_lora_weights(save_directory=tmpdirname, unet_lora_layers=pipe.unet.lora_state_dict()) - # First, delete adapter and compare. - pipe.delete_adapters(pipe.get_active_adapters()[0]) - output_no_adapter = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse(np.allclose(output_adapter_1, output_no_adapter, atol=1e-3, rtol=1e-3)) - self.assertTrue(np.allclose(output_no_lora, output_no_adapter, atol=1e-3, rtol=1e-3)) + # Delete the currently active adapter ("adapter-2") + pipe.delete_adapters(pipe.get_active_adapters()) + output_after_delete = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue(np.allclose(output_after_delete, output_no_lora, atol=1e-3, rtol=1e-3)) - # Then load adapter and compare. - pipe.load_lora_weights(tmpdirname) - output_lora_loaded = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(np.allclose(output_adapter_1, output_lora_loaded, atol=1e-3, rtol=1e-3)) + # Load the weights back + pipe.load_lora_weights(tmpdirname, adapter_name="adapter-new") + output_after_load = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse(np.allclose(output_after_load, output_no_lora, atol=1e-3, rtol=1e-3)) def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): from diffusers.hooks.group_offloading import _get_top_level_group_offload_hook @@ -2082,6 +1802,24 @@ def _setup_pipeline_and_get_base_output(self, scheduler_cls, lora_alpha=4, use_d self.assertTrue(output_no_lora.shape == self.output_shape) return pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config + def _setup_multi_adapter_pipeline(self, scheduler_cls): + """ + Helper to set up a pipeline with two LoRA adapters ("adapter-1", "adapter-2") + attached to the text encoder and denoiser. + """ + pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls) + ) + + # Add adapter-1 + pipe, denoiser = self.add_adapters_to_pipeline( + pipe, text_lora_config, denoiser_lora_config, adapter_name="adapter-1" + ) + # Add adapter-2 + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config, adapter_name="adapter-2") + + return pipe, inputs, output_no_lora, denoiser + def _get_lora_state_dicts(self, modules_to_save): state_dicts = {} for module_name, module in modules_to_save.items(): From f86ccac4c7a34ca09b66d11edda22a23c31b9f1b Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 1 Jul 2025 18:27:41 +0530 Subject: [PATCH 12/19] don't need the compilation test, we test it elsewhere. --- tests/lora/utils.py | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 44b4f8333e83..4a0e2687b2e8 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -17,7 +17,6 @@ import os import re import tempfile -import unittest from itertools import product import numpy as np @@ -1013,28 +1012,6 @@ def test_unexpected_keys_warning(self): self.assertTrue(".diffusers_cat" in cap_logger.out) - @unittest.skip("This is failing for now - need to investigate") - def test_simple_inference_with_text_denoiser_lora_unfused_torch_compile(self): - """ - Tests a simple inference with lora attached to text encoder and unet, then unloads the lora weights - and makes sure it works as expected - """ - for scheduler_cls in self.scheduler_classes: - pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( - scheduler_cls - ) - - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - - pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) - pipe.text_encoder = torch.compile(pipe.text_encoder, mode="reduce-overhead", fullgraph=True) - - if self.has_two_text_encoders or self.has_three_text_encoders: - pipe.text_encoder_2 = torch.compile(pipe.text_encoder_2, mode="reduce-overhead", fullgraph=True) - - # Just makes sure it works.. - _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - def test_modify_padding_mode(self): def set_pad_mode(network, mode="circular"): for _, module in network.named_modules(): From 810726c4c521aa6fffd062df5feccdd4d5168a9d Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 2 Jul 2025 19:07:17 +0530 Subject: [PATCH 13/19] update --- tests/lora/test_lora_layers_cogvideox.py | 7 +- tests/lora/test_lora_layers_cogview4.py | 5 +- tests/lora/test_lora_layers_hunyuanvideo.py | 7 +- tests/lora/test_lora_layers_ltx_video.py | 7 +- tests/lora/test_lora_layers_mochi.py | 7 +- tests/lora/test_lora_layers_sdxl.py | 9 +- tests/lora/test_lora_layers_wan.py | 7 +- tests/lora/utils.py | 353 +++++++------------- 8 files changed, 142 insertions(+), 260 deletions(-) diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index d08f25a6875e..902fbaf631e0 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -120,8 +120,9 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) + def test_lora_set_adapters_scenarios(self, scenario): + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3) @parameterized.expand( [ @@ -137,7 +138,7 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): ] ) def test_lora_actions(self, action, components_to_add): - super().test_lora_actions(action, components_to_add, expected_atol=9e-3) + super()._test_lora_actions(action, components_to_add, expected_atol=9e-3) def test_lora_scale_kwargs_match_fusion(self): super().test_lora_scale_kwargs_match_fusion(expected_atol=9e-3, expected_rtol=9e-3) diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index df9e8fbd5f3f..201e0984ff1e 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -110,9 +110,6 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - @parameterized.expand( [ # Test actions on text_encoder LoRA only @@ -127,7 +124,7 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): ] ) def test_lora_actions(self, action, components_to_add): - super().test_lora_actions(action, components_to_add, expected_atol=9e-3) + super()._test_lora_actions(action, components_to_add, expected_atol=9e-3) @parameterized.expand([("block_level", True), ("leaf_level", False)]) @require_torch_accelerator diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index be4bf12575d1..c82bde75fad2 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -151,8 +151,9 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) + def test_lora_set_adapters_scenarios(self, scenario): + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3, expected_rtol=9e-3) @parameterized.expand( [ @@ -168,7 +169,7 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): ] ) def test_lora_actions(self, action, components_to_add): - super().test_lora_actions(action, components_to_add, expected_atol=9e-3) + super()._test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in HunyuanVideo.") def test_modify_padding_mode(self): diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index eee0ca7784ee..20dc2d5f2c45 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -109,8 +109,9 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) + def test_lora_set_adapters_scenarios(self, scenario): + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3) @parameterized.expand( [ @@ -126,7 +127,7 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): ] ) def test_lora_actions(self, action, components_to_add): - super().test_lora_actions(action, components_to_add, expected_atol=9e-3) + super()._test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in LTXVideo.") def test_modify_padding_mode(self): diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index 3cab8dfd2c2c..d6dd59b4b182 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -96,8 +96,9 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) + def test_lora_set_adapters_scenarios(self, scenario): + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3) @parameterized.expand( [ @@ -113,7 +114,7 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): ] ) def test_lora_actions(self, action, components_to_add): - super().test_lora_actions(action, components_to_add, expected_atol=9e-3) + super()._test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in Mochi.") def test_modify_padding_mode(self): diff --git a/tests/lora/test_lora_layers_sdxl.py b/tests/lora/test_lora_layers_sdxl.py index 3e33a9e71a4e..eee1502b5849 100644 --- a/tests/lora/test_lora_layers_sdxl.py +++ b/tests/lora/test_lora_layers_sdxl.py @@ -139,9 +139,10 @@ def test_lora_actions(self, action, components_to_add): expected_atol = 1e-3 expected_rtol = 1e-3 - super().test_lora_actions(expected_atol=expected_atol, expected_rtol=expected_rtol) + super()._test_lora_actions(action, components_to_add, expected_atol=expected_atol, expected_rtol=expected_rtol) - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",), ("fused_multi",)]) + def test_lora_set_adapters_scenarios(self, scenario): if torch.cuda.is_available(): expected_atol = 9e-2 expected_rtol = 9e-2 @@ -149,8 +150,8 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): expected_atol = 1e-3 expected_rtol = 1e-3 - super().test_simple_inference_with_text_lora_denoiser_fused_multi( - expected_atol=expected_atol, expected_rtol=expected_rtol + super()._test_lora_set_adapters_scenarios( + scenario=scenario, expected_atol=expected_atol, expected_rtol=expected_rtol ) def test_lora_scale_kwargs_match_fusion(self): diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index 62fddb96a7e0..69f2cb663999 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -105,8 +105,9 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) + def test_lora_set_adapters_scenarios(self, scenario): + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3) @parameterized.expand( [ @@ -122,7 +123,7 @@ def test_simple_inference_with_text_lora_denoiser_fused_multi(self): ] ) def test_lora_actions(self, action, components_to_add): - super().test_lora_actions(action, components_to_add, expected_atol=9e-3) + super()._test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in Wan.") def test_modify_padding_mode(self): diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 4a0e2687b2e8..e214f386aab2 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -38,7 +38,6 @@ floats_tensor, is_torch_version, require_peft_backend, - require_peft_version_greater, require_torch_accelerator, require_transformers_version_greater, skip_mps, @@ -128,7 +127,6 @@ class PeftLoraLoaderMixinTests: text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"] denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - @require_peft_version_greater("0.13.1") def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" for scheduler_cls in self.scheduler_classes: @@ -188,7 +186,6 @@ def test_low_cpu_mem_usage_with_injection(self): output_lora = pipe(**inputs)[0] self.assertTrue(output_lora.shape == self.output_shape) - @require_peft_version_greater("0.13.1") @require_transformers_version_greater("4.45.2") def test_low_cpu_mem_usage_with_loading(self): """Tests if we can load LoRA state dict with low_cpu_mem_usage.""" @@ -353,7 +350,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") output_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.unfuse_lora() + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) for module_name, module in modules_to_save.items(): self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") output_unfused = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -401,6 +398,9 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, ] ) def test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3): + """Tests to check if different LoRA actions like fusion, loading-unloading, etc. + work as expected. + """ for cls in inspect.getmro(self.__class__): if "test_lora_actions" in cls.__dict__ and cls is not PeftLoraLoaderMixinTests: # Skip this test if it is overwritten by child class. We need to do this because parameterized @@ -414,10 +414,13 @@ def test_lora_scaling(self, lora_components_to_add): Tests inference with LoRA scaling applied via attention_kwargs for different LoRA configurations. """ + # TODO: + # Currently, the following fails: + # tests/lora/test_lora_layers_sdxl.py::StableDiffusionXLLoRATests::test_lora_scaling_1_text_and_denoiser if lora_components_to_add == "text_encoder_only": if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): pytest.skip( - "Test not supported for {self.__class__.__name__} since there is not text encoder in the LoRA loadable modules." + f"Test not supported for {self.__class__.__name__} since there is not text encoder in the LoRA loadable modules." ) attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) @@ -426,44 +429,54 @@ def test_lora_scaling(self, lora_components_to_add): self._setup_pipeline_and_get_base_output(scheduler_cls) ) + # Create a deep copy to ensure a clean state for each iteration + lora_pipe = copy.deepcopy(pipe) + # Add LoRA components based on the parameterization if lora_components_to_add == "text_encoder_only": - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + lora_pipe, _ = self.add_adapters_to_pipeline(lora_pipe, text_lora_config, denoiser_lora_config=None) elif lora_components_to_add == "text_and_denoiser": - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - else: - raise ValueError(f"Unknown `lora_components_to_add`: {lora_components_to_add}") + lora_pipe, _ = self.add_adapters_to_pipeline(lora_pipe, text_lora_config, denoiser_lora_config) # 1. Test base LoRA output - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + output_lora = lora_pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertFalse( np.allclose(output_lora, output_no_lora, atol=1e-3, rtol=1e-3), "LoRA should change the output." ) # 2. Test with a scale of 0.5 attention_kwargs = {attention_kwargs_name: {"scale": 0.5}} - output_lora_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] - self.assertTrue( - not np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), + output_lora_scale = lora_pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + self.assertFalse( + np.allclose(output_lora, output_lora_scale, atol=1e-3, rtol=1e-3), "Using a LoRA scale should change the output.", ) # 3. Test with a scale of 0.0, which should be identical to no LoRA attention_kwargs = {attention_kwargs_name: {"scale": 0.0}} - output_lora_0_scale = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] + output_lora_0_scale = lora_pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0] self.assertTrue( np.allclose(output_no_lora, output_lora_0_scale, atol=1e-3, rtol=1e-3), "Using a LoRA scale of 0.0 should be the same as no LoRA.", ) + # 4. Final check to ensure the scaling parameter is restored + if "text_encoder" in self.pipeline_class._lora_loadable_modules: + # Get the underlying LoRA layer to check its scaling factor + lora_layer = lora_pipe.text_encoder.text_model.encoder.layers[0].self_attn.q_proj + if hasattr(lora_layer, "scaling"): + self.assertEqual( + lora_layer.scaling.get("default", 1.0), + 1.0, + "The scaling parameter was not correctly restored!", + ) + def test_wrong_adapter_name_raises_error(self): adapter_name = "adapter-1" - scheduler_cls = self.scheduler_classes[0] pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls ) - pipe, _ = self.add_adapters_to_pipeline( pipe, text_lora_config, denoiser_lora_config, adapter_name=adapter_name ) @@ -483,7 +496,6 @@ def test_multiple_wrong_adapter_name_raises_error(self): pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls ) - pipe, _ = self.add_adapters_to_pipeline( pipe, text_lora_config, denoiser_lora_config, adapter_name=adapter_name ) @@ -502,61 +514,6 @@ def test_multiple_wrong_adapter_name_raises_error(self): pipe.set_adapters(adapter_name) _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - def test_simple_inference_with_text_denoiser_block_scale(self): - """ - Tests a simple inference with lora attached to text encoder and unet, attaches - one adapter and set different weights for different blocks (i.e. block lora) - """ - if self.unet_kwargs is None: - pytest.skip("Test is only supposed for UNets.") - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls) - ) - - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - - weights_1 = {"text_encoder": 2, "unet": {"down": 5}} - pipe.set_adapters("adapter-1", weights_1) - output_weights_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - weights_2 = {"unet": {"up": 5}} - pipe.set_adapters("adapter-1", weights_2) - output_weights_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse( - np.allclose(output_weights_1, output_weights_2, atol=1e-3, rtol=1e-3), - "LoRA weights 1 and 2 should give different results", - ) - self.assertFalse( - np.allclose(output_no_lora, output_weights_1, atol=1e-3, rtol=1e-3), - "No adapter and LoRA weights 1 should give different results", - ) - self.assertFalse( - np.allclose(output_no_lora, output_weights_2, atol=1e-3, rtol=1e-3), - "No adapter and LoRA weights 2 should give different results", - ) - - pipe.disable_lora() - output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), - "output with no lora and output with lora disabled should give same results", - ) - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): """Tests that any valid combination of lora block scales can be used in pipe.set_adapter""" if self.unet_kwargs is None: @@ -629,7 +586,6 @@ def all_possible_dict_opts(unet, value): self.scheduler_classes[0] ) pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet denoiser.add_adapter(denoiser_lora_config, "adapter-1") @@ -657,15 +613,7 @@ def test_lora_fuse_nan(self): scheduler_cls ) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config, "adapter-1") # corrupt one LoRA weight with `inf` values with torch.no_grad(): @@ -716,7 +664,6 @@ def test_get_adapters(self): pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls ) - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet @@ -805,87 +752,7 @@ def test_get_list_adapters(self): self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - @require_peft_version_greater(peft_version="0.6.2") - def test_simple_inference_with_text_lora_denoiser_fused_multi( - self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3 - ): - """ - Tests a simple inference with lora attached into text encoder + fuses the lora weights into base model - and makes sure it works as expected - with unet and multi-adapter case - """ - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls) - ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - pipe.text_encoder.add_adapter(text_lora_config, "adapter-2") - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - denoiser.add_adapter(denoiser_lora_config, "adapter-2") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2" - ) - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2") - - # set them to multi-adapter inference mode - pipe.set_adapters(["adapter-1", "adapter-2"]) - outputs_all_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1"]) - outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - # Fusing should still keep the LoRA layers so output should remain the same - outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertTrue( - np.allclose(outputs_lora_1, outputs_lora_1_fused, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers") - - self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers") - - if self.has_two_text_encoders or self.has_three_text_encoders: - if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers" - ) - - pipe.fuse_lora( - components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-2", "adapter-1"] - ) - self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - # Fusing should still keep the LoRA layers - output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_all_lora_fused, outputs_all_lora, atol=expected_atol, rtol=expected_rtol), - "Fused lora should not change the output", - ) - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3): + def test_lora_scale_kwargs_match_fusion(self, expected_atol=1e-3, expected_rtol=1e-3): attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class) for lora_scale in [1.0, 0.8]: @@ -893,25 +760,7 @@ def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expec pipe, inputs, output_no_lora, text_lora_config, denoiser_lora_config = ( self._setup_pipeline_and_get_base_output(scheduler_cls) ) - - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder_2), - "Lora not correctly set in text encoder 2", - ) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config, "adapter-1") pipe.set_adapters(["adapter-1"]) attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}} @@ -935,7 +784,6 @@ def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expec "LoRA should change the output", ) - @require_peft_version_greater(peft_version="0.9.0") def test_simple_inference_with_dora(self): for scheduler_cls in self.scheduler_classes: pipe, inputs, output_no_dora_lora, text_lora_config, denoiser_lora_config = ( @@ -953,10 +801,7 @@ def test_missing_keys_warning(self): scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. pipe, _, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output(scheduler_cls) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + pipe, _ = self.add_adapters_to_pipeline(pipe, None, denoiser_lora_config) with tempfile.TemporaryDirectory() as tmpdirname: modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) @@ -987,10 +832,7 @@ def test_unexpected_keys_warning(self): scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. pipe, _, _, _, denoiser_lora_config = self._setup_pipeline_and_get_base_output(scheduler_cls) - - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config) - self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.") + pipe, _ = self.add_adapters_to_pipeline(pipe, None, denoiser_lora_config) with tempfile.TemporaryDirectory() as tmpdirname: modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True) @@ -1104,9 +946,11 @@ def test_set_adapters_match_attention_kwargs(self): self.pipeline_class.save_lora_weights( save_directory=tmpdirname, safe_serialization=True, **lora_state_dicts ) - pipe.unload_lora_weights() + # We should not have to set up the pipeline again. `unload_lora_weights()` should work. + # TODO: investigate later. + # pipe.unload_lora_weights() + pipe, _, _, _, _ = self._setup_pipeline_and_get_base_output(scheduler_cls) pipe.load_lora_weights(os.path.join(tmpdirname)) - for module_name, module in modules_to_save.items(): self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") @@ -1124,7 +968,6 @@ def test_set_adapters_match_attention_kwargs(self): "Loading from saved checkpoints should give same results as set_adapters().", ) - @require_peft_version_greater("0.13.2") def test_lora_B_bias(self): # Currently, this test is only relevant for Flux Control LoRA as we are not # aware of any other LoRA checkpoint that has its `lora_B` biases trained. @@ -1267,7 +1110,6 @@ def initialize_pipeline(storage_dtype=None, compute_dtype=torch.float32): pipe_float8_e4m3_bf16 = initialize_pipeline(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) pipe_float8_e4m3_bf16(**inputs, generator=torch.manual_seed(0))[0] - @require_peft_version_greater("0.14.0") def test_layerwise_casting_peft_input_autocast_denoiser(self): r""" A test that checks if layerwise casting works correctly with PEFT layers and forward pass does not fail. This @@ -1442,7 +1284,6 @@ def test_lora_unload_add_adapter(self): ) _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - @require_peft_version_greater("0.13.2") def test_lora_exclude_modules(self): """ Test to check if `exclude_modules` works or not. It works in the following way: @@ -1493,20 +1334,7 @@ def test_lora_exclude_modules(self): "Lora outputs should match.", ) - @parameterized.expand( - [ - ("simple",), - ("weighted",), - ("block_lora",), - ("delete_adapter",), - ("load_delete_load",), - ] - ) - def test_multi_adapter_scenarios(self, scenario): - """ - A unified test for various multi-adapter LoRA scenarios, including weighting, - block scaling, and adapter deletion. - """ + def _test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expected_rtol=1e-3): for scheduler_cls in self.scheduler_classes: pipe, inputs, output_no_lora, _ = self._setup_multi_adapter_pipeline(scheduler_cls) @@ -1520,42 +1348,96 @@ def test_multi_adapter_scenarios(self, scenario): pipe.set_adapters(["adapter-1", "adapter-2"]) output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] + # Also disable the LoRA + pipe.disable_lora() + output_lora_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] + # --- Assert base multi-adapter behavior --- - self.assertFalse(np.allclose(output_no_lora, output_adapter_1, atol=1e-3, rtol=1e-3)) - self.assertFalse(np.allclose(output_adapter_1, output_adapter_2, atol=1e-3, rtol=1e-3)) - self.assertFalse(np.allclose(output_adapter_1, output_adapter_mixed, atol=1e-3, rtol=1e-3)) + self.assertFalse(np.allclose(output_no_lora, output_adapter_1, atol=expected_atol, rtol=expected_rtol)) + self.assertTrue(np.allclose(output_no_lora, output_lora_disabled, atol=expected_atol, rtol=expected_rtol)) + self.assertFalse(np.allclose(output_adapter_1, output_adapter_2, atol=expected_atol, rtol=expected_rtol)) + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) # --- Scenario-specific logic --- if scenario == "weighted": pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) output_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse(np.allclose(output_weighted, output_adapter_mixed, atol=1e-3, rtol=1e-3)) + self.assertFalse( + np.allclose(output_weighted, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) elif scenario == "block_lora": scales = {"unet": {"down": 0.8, "mid": 0.5, "up": 0.2}} - pipe.set_adapters(["adapter-1", "adapter-2"], [scales, scales]) - output_block_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse(np.allclose(output_block_lora, output_adapter_mixed, atol=1e-3, rtol=1e-3)) + pipe.set_adapters("adapter-1", scales) + output_block_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_block_lora_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) + + text_encoder_modules = [k for k in self.pipeline_class._lora_loadable_modules if "text_encoder" in k] + if text_encoder_modules: + text_encoder_module_name = text_encoder_modules[0] + scales_2 = {text_encoder_module_name: 2, "unet": {"down": 5}} + pipe.set_adapters("adapter-1", scales_2) + output_block_lora_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_block_lora_2, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) elif scenario == "delete_adapter": + pipe.set_adapters("adapter-2") + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] pipe.delete_adapters("adapter-1") output_after_delete = pipe(**inputs, generator=torch.manual_seed(0))[0] # After deleting adapter-1, the output should be the same as using only adapter-2 - self.assertTrue(np.allclose(output_after_delete, output_adapter_2, atol=1e-3, rtol=1e-3)) + self.assertTrue( + np.allclose(output_after_delete, output_adapter_2, atol=expected_atol, rtol=expected_rtol) + ) - elif scenario == "load_delete_load": - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_lora_weights(save_directory=tmpdirname, unet_lora_layers=pipe.unet.lora_state_dict()) + elif scenario == "fused_multi": + # 1. Fuse a single adapter + pipe.set_adapters("adapter-1") + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) + self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - # Delete the currently active adapter ("adapter-2") - pipe.delete_adapters(pipe.get_active_adapters()) - output_after_delete = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue(np.allclose(output_after_delete, output_no_lora, atol=1e-3, rtol=1e-3)) + output_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_adapter_1, output_lora_1_fused, atol=expected_atol, rtol=expected_rtol) + ) + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + + # 2. Fuse both adapters + pipe.set_adapters(["adapter-1", "adapter-2"]) + pipe.fuse_lora( + components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1", "adapter-2"] + ) + self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_all_lora_fused, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",), ("fused_multi",)]) + def test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expected_rtol=1e-3): + """ + A unified test for various multi-adapter (and single-adapter) LoRA scenarios, including weighting, + block scaling, and adapter deletion. + """ + for cls in inspect.getmro(self.__class__): + if "test_lora_set_adapters_scenarios" in cls.__dict__ and cls is not PeftLoraLoaderMixinTests: + # Skip this test if it is overwritten by child class. We need to do this because parameterized + # materializes the test methods on invocation which cannot be overridden. + return - # Load the weights back - pipe.load_lora_weights(tmpdirname, adapter_name="adapter-new") - output_after_load = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse(np.allclose(output_after_load, output_no_lora, atol=1e-3, rtol=1e-3)) + if scenario == "block_lora" and self.unet_kwargs is None: + pytest.skip(f"Test not supported for {scenario=}.") + + self._test_lora_set_adapters_scenarios(scenario, expected_atol, expected_rtol) def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): from diffusers.hooks.group_offloading import _get_top_level_group_offload_hook @@ -1744,12 +1626,9 @@ def get_dummy_tokens(self): return prepared_inputs def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"): - if text_lora_config is not None: - if "text_encoder" in self.pipeline_class._lora_loadable_modules: - pipe.text_encoder.add_adapter(text_lora_config, adapter_name=adapter_name) - self.assertTrue( - check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder" - ) + if text_lora_config is not None and "text_encoder" in self.pipeline_class._lora_loadable_modules: + pipe.text_encoder.add_adapter(text_lora_config, adapter_name=adapter_name) + self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder") if denoiser_lora_config is not None: denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet @@ -1758,7 +1637,7 @@ def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_co else: denoiser = None - if text_lora_config is not None and self.has_two_text_encoders or self.has_three_text_encoders: + if text_lora_config is not None and (self.has_two_text_encoders or self.has_three_text_encoders): if "text_encoder_2" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder_2.add_adapter(text_lora_config, adapter_name=adapter_name) self.assertTrue( From f72ada150e6ad845e68b6e6509af9c392addaf24 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 3 Jul 2025 13:41:28 +0530 Subject: [PATCH 14/19] fix --- tests/lora/test_lora_layers_hunyuanvideo.py | 2 +- tests/lora/utils.py | 32 +++++++++------------ 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index c82bde75fad2..db345e291545 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -153,7 +153,7 @@ def get_dummy_inputs(self, with_generator=True): @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) def test_lora_set_adapters_scenarios(self, scenario): - super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3, expected_rtol=9e-3) + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3) @parameterized.expand( [ diff --git a/tests/lora/utils.py b/tests/lora/utils.py index e214f386aab2..8f155da512a0 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -252,7 +252,6 @@ def test_simple_inference_with_partial_text_lora(self): init_lora_weights=False, use_dora=False, ) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) state_dict = {} @@ -317,7 +316,6 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, modules_to_save = self._get_modules_to_save( pipe, has_denoiser=lora_components_to_add != "text_encoder_only" ) - output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue(not np.allclose(output_lora, output_no_lora, atol=expected_atol, rtol=expected_rtol)) @@ -348,9 +346,11 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules) for module_name, module in modules_to_save.items(): self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") output_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") for module_name, module in modules_to_save.items(): self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") output_unfused = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -384,6 +384,14 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, "Loading from a saved checkpoint should yield the same result.", ) + elif action == "disable": + pipe.disable_lora() + output_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_no_lora, output_disabled, atol=1e-3, rtol=1e-3), + "output with no lora and output with lora disabled should give same results", + ) + @parameterized.expand( [ # Test actions on text_encoder LoRA only @@ -395,6 +403,7 @@ def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, ("unloaded", "text_and_denoiser"), ("unfused", "text_and_denoiser"), ("save_load", "text_and_denoiser"), + ("disable", "text_and_denoiser"), ] ) def test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3): @@ -1349,12 +1358,12 @@ def _test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expect output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] # Also disable the LoRA - pipe.disable_lora() - output_lora_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] + # pipe.disable_lora() + # output_lora_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] # --- Assert base multi-adapter behavior --- self.assertFalse(np.allclose(output_no_lora, output_adapter_1, atol=expected_atol, rtol=expected_rtol)) - self.assertTrue(np.allclose(output_no_lora, output_lora_disabled, atol=expected_atol, rtol=expected_rtol)) + # self.assertTrue(np.allclose(output_no_lora, output_lora_disabled, atol=expected_atol, rtol=expected_rtol)) self.assertFalse(np.allclose(output_adapter_1, output_adapter_2, atol=expected_atol, rtol=expected_rtol)) self.assertFalse( np.allclose(output_adapter_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) @@ -1400,8 +1409,6 @@ def _test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expect # 1. Fuse a single adapter pipe.set_adapters("adapter-1") pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) - self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - output_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] self.assertTrue( np.allclose(output_adapter_1, output_lora_1_fused, atol=expected_atol, rtol=expected_rtol) @@ -1552,7 +1559,6 @@ def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=No init_lora_weights=False, use_dora=use_dora, ) - denoiser_lora_config = LoraConfig( r=rank, lora_alpha=lora_alpha, @@ -1615,16 +1621,6 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - # Copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb - def get_dummy_tokens(self): - max_seq_length = 77 - - inputs = torch.randint(2, 56, size=(1, max_seq_length), generator=torch.manual_seed(0)) - - prepared_inputs = {} - prepared_inputs["input_ids"] = inputs - return prepared_inputs - def add_adapters_to_pipeline(self, pipe, text_lora_config=None, denoiser_lora_config=None, adapter_name="default"): if text_lora_config is not None and "text_encoder" in self.pipeline_class._lora_loadable_modules: pipe.text_encoder.add_adapter(text_lora_config, adapter_name=adapter_name) From 1200b247aaa466dccc2e52e06374c6caaf9d7836 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 3 Jul 2025 14:41:17 +0530 Subject: [PATCH 15/19] fix --- tests/lora/test_lora_layers_hunyuanvideo.py | 5 ++++- tests/lora/test_lora_layers_sd.py | 14 ++++++++++++++ tests/lora/utils.py | 8 -------- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index db345e291545..07de084d736f 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -153,7 +153,10 @@ def get_dummy_inputs(self, with_generator=True): @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) def test_lora_set_adapters_scenarios(self, scenario): - super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3) + expected_atol = 9e-3 + if scenario == "weighted": + expected_atol = 1e-3 + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=expected_atol) @parameterized.expand( [ diff --git a/tests/lora/test_lora_layers_sd.py b/tests/lora/test_lora_layers_sd.py index a81128fa446b..1aaaf4076f76 100644 --- a/tests/lora/test_lora_layers_sd.py +++ b/tests/lora/test_lora_layers_sd.py @@ -20,6 +20,7 @@ import torch import torch.nn as nn from huggingface_hub import hf_hub_download +from parameterized import parameterized from safetensors.torch import load_file from transformers import CLIPTextModel, CLIPTokenizer @@ -208,6 +209,19 @@ def test_integration_move_lora_dora_cpu(self): if "lora_" in name: self.assertNotEqual(param.device, torch.device("cpu")) + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) + def test_lora_set_adapters_scenarios(self, scenario): + if torch.cuda.is_available(): + expected_atol = 9e-2 + expected_rtol = 9e-2 + else: + expected_atol = 1e-3 + expected_rtol = 1e-3 + + super()._test_lora_set_adapters_scenarios( + scenario=scenario, expected_atol=expected_atol, expected_rtol=expected_rtol + ) + @slow @nightly diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 8f155da512a0..b4510a5a80df 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -423,9 +423,6 @@ def test_lora_scaling(self, lora_components_to_add): Tests inference with LoRA scaling applied via attention_kwargs for different LoRA configurations. """ - # TODO: - # Currently, the following fails: - # tests/lora/test_lora_layers_sdxl.py::StableDiffusionXLLoRATests::test_lora_scaling_1_text_and_denoiser if lora_components_to_add == "text_encoder_only": if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): pytest.skip( @@ -1357,13 +1354,8 @@ def _test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expect pipe.set_adapters(["adapter-1", "adapter-2"]) output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - # Also disable the LoRA - # pipe.disable_lora() - # output_lora_disabled = pipe(**inputs, generator=torch.manual_seed(0))[0] - # --- Assert base multi-adapter behavior --- self.assertFalse(np.allclose(output_no_lora, output_adapter_1, atol=expected_atol, rtol=expected_rtol)) - # self.assertTrue(np.allclose(output_no_lora, output_lora_disabled, atol=expected_atol, rtol=expected_rtol)) self.assertFalse(np.allclose(output_adapter_1, output_adapter_2, atol=expected_atol, rtol=expected_rtol)) self.assertFalse( np.allclose(output_adapter_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) From 36bc333e1c6a846d509f6da85131e847e8edd933 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 3 Jul 2025 14:45:04 +0530 Subject: [PATCH 16/19] update --- tests/lora/test_lora_layers_wan.py | 13 +----- tests/lora/test_lora_layers_wanvace.py | 62 +++++++++----------------- 2 files changed, 23 insertions(+), 52 deletions(-) diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index 69f2cb663999..9933ad7ebd4b 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -19,17 +19,8 @@ from parameterized import parameterized from transformers import AutoTokenizer, T5EncoderModel -from diffusers import ( - AutoencoderKLWan, - FlowMatchEulerDiscreteScheduler, - WanPipeline, - WanTransformer3DModel, -) -from diffusers.utils.testing_utils import ( - floats_tensor, - require_peft_backend, - skip_mps, -) +from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanPipeline, WanTransformer3DModel +from diffusers.utils.testing_utils import floats_tensor, require_peft_backend, skip_mps sys.path.append(".") diff --git a/tests/lora/test_lora_layers_wanvace.py b/tests/lora/test_lora_layers_wanvace.py index 740c00f941ed..dfa04505743c 100644 --- a/tests/lora/test_lora_layers_wanvace.py +++ b/tests/lora/test_lora_layers_wanvace.py @@ -21,18 +21,13 @@ import pytest import safetensors.torch import torch +from parameterized import parameterized from PIL import Image from transformers import AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanVACEPipeline, WanVACETransformer3DModel from diffusers.utils.import_utils import is_peft_available -from diffusers.utils.testing_utils import ( - floats_tensor, - require_peft_backend, - require_peft_version_greater, - skip_mps, - torch_device, -) +from diffusers.utils.testing_utils import floats_tensor, require_peft_backend, skip_mps, torch_device if is_peft_available(): @@ -120,44 +115,30 @@ def get_dummy_inputs(self, with_generator=True): return noise, input_ids, pipeline_inputs - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - - @unittest.skip("Not supported in Wan VACE.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in Wan VACE.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",)]) + def test_lora_set_adapters_scenarios(self, scenario): + super()._test_lora_set_adapters_scenarios(scenario, expected_atol=9e-3) + + @parameterized.expand( + [ + # Test actions on text_encoder LoRA only + ("fused", "text_encoder_only"), + ("unloaded", "text_encoder_only"), + ("save_load", "text_encoder_only"), + # Test actions on both text_encoder and denoiser LoRA + ("fused", "text_and_denoiser"), + ("unloaded", "text_and_denoiser"), + ("unfused", "text_and_denoiser"), + ("save_load", "text_and_denoiser"), + ] + ) + def test_lora_actions(self, action, components_to_add): + super()._test_lora_actions(action, components_to_add, expected_atol=9e-3) @unittest.skip("Not supported in Wan VACE.") def test_modify_padding_mode(self): pass - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_partial_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora_and_scale(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora_fused(self): - pass - - @unittest.skip("Text encoder LoRA is not supported in Wan VACE.") - def test_simple_inference_with_text_lora_save_load(self): - pass - @pytest.mark.xfail( condition=True, reason="RuntimeError: Input type (float) and bias type (c10::BFloat16) should be the same", @@ -166,7 +147,6 @@ def test_simple_inference_with_text_lora_save_load(self): def test_layerwise_casting_inference_denoiser(self): super().test_layerwise_casting_inference_denoiser() - @require_peft_version_greater("0.13.2") def test_lora_exclude_modules_wanvace(self): scheduler_cls = self.scheduler_classes[0] exclude_module_name = "vace_blocks.0.proj_out" From b48fde5b79795c12515b05f84c76fd0606a8c195 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 3 Jul 2025 14:48:11 +0530 Subject: [PATCH 17/19] increase the number of workers. --- .github/workflows/pr_tests.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 34a344528e3e..83a7c69e8bc3 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -265,11 +265,11 @@ jobs: - name: Run fast PyTorch LoRA tests with PEFT run: | python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + python -m pytest -n 6 --max-worker-restart=0 --dist=loadfile \ -s -v \ --make-reports=tests_peft_main \ tests/lora/ - python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile \ + python -m pytest -n 6 --max-worker-restart=0 --dist=loadfile \ -s -v \ --make-reports=tests_models_lora_peft_main \ tests/models/ -k "lora" From fe1af35af72801e2ab7675a0e0a7062b05d90576 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 4 Jul 2025 19:28:35 +0530 Subject: [PATCH 18/19] update --- tests/lora/utils.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index b4510a5a80df..6c9e77eb4b6d 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -127,6 +127,35 @@ class PeftLoraLoaderMixinTests: text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"] denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] + def test_simple_inference_save_pretrained_with_text_lora(self): + """ + Tests a simple usecase where users could use saving utilities for text encoder (only) + LoRA through save_pretrained. + """ + if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): + pytest.skip("Test not supported.") + for scheduler_cls in self.scheduler_classes: + pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + scheduler_cls + ) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + + with tempfile.TemporaryDirectory() as tmpdirname: + pipe.save_pretrained(tmpdirname) + pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) + pipe_from_pretrained.to(torch_device) + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=False) + + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + + images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results.", + ) + def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" for scheduler_cls in self.scheduler_classes: From 9445c4bf4ac65f5924020a292aff7accab0a3f1f Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 4 Jul 2025 21:04:43 +0530 Subject: [PATCH 19/19] update --- tests/lora/utils.py | 465 +++++++++++++++++++++----------------------- 1 file changed, 227 insertions(+), 238 deletions(-) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 6c9e77eb4b6d..0633c8574f8f 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -127,42 +127,12 @@ class PeftLoraLoaderMixinTests: text_encoder_target_modules = ["q_proj", "k_proj", "v_proj", "out_proj"] denoiser_target_modules = ["to_q", "to_k", "to_v", "to_out.0"] - def test_simple_inference_save_pretrained_with_text_lora(self): - """ - Tests a simple usecase where users could use saving utilities for text encoder (only) - LoRA through save_pretrained. - """ - if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): - pytest.skip("Test not supported.") - for scheduler_cls in self.scheduler_classes: - pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( - scheduler_cls - ) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) - images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - with tempfile.TemporaryDirectory() as tmpdirname: - pipe.save_pretrained(tmpdirname) - pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) - pipe_from_pretrained.to(torch_device) - modules_to_save = self._get_modules_to_save(pipe, has_denoiser=False) - - for module_name, module in modules_to_save.items(): - self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") - - images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), - "Loading from saved checkpoints should give same results.", - ) - def test_low_cpu_mem_usage_with_injection(self): """Tests if we can inject LoRA state dict with low_cpu_mem_usage.""" for scheduler_cls in self.scheduler_classes: pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls ) - if "text_encoder" in self.pipeline_class._lora_loadable_modules: inject_adapter_in_model(text_lora_config, pipe.text_encoder, low_cpu_mem_usage=True) self.assertTrue( @@ -223,7 +193,6 @@ def test_low_cpu_mem_usage_with_loading(self): pipe, inputs, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls ) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] @@ -262,6 +231,33 @@ def test_low_cpu_mem_usage_with_loading(self): "Loading from saved checkpoints with `low_cpu_mem_usage` should give same results.", ) + def test_simple_inference_save_pretrained_with_text_lora(self): + """ + Tests a simple usecase where users could use saving utilities for text encoder (only) + LoRA through save_pretrained. + """ + if not any("text_encoder" in k for k in self.pipeline_class._lora_loadable_modules): + pytest.skip("Test not supported.") + for scheduler_cls in self.scheduler_classes: + pipe, inputs, _, text_lora_config, _ = self._setup_pipeline_and_get_base_output(scheduler_cls) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config=None) + images_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + + with tempfile.TemporaryDirectory() as tmpdirname: + pipe.save_pretrained(tmpdirname) + pipe_from_pretrained = self.pipeline_class.from_pretrained(tmpdirname) + pipe_from_pretrained.to(torch_device) + modules_to_save = self._get_modules_to_save(pipe, has_denoiser=False) + + for module_name, module in modules_to_save.items(): + self.assertTrue(check_if_lora_correctly_set(module), f"Lora not correctly set in {module_name}") + + images_lora_save_pretrained = pipe_from_pretrained(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(images_lora, images_lora_save_pretrained, atol=1e-3, rtol=1e-3), + "Loading from saved checkpoints should give same results.", + ) + def test_simple_inference_with_partial_text_lora(self): """ Tests a simple inference with lora attached on the text encoder @@ -318,6 +314,106 @@ def test_simple_inference_with_partial_text_lora(self): "Removing adapters should change the output", ) + def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): + """Tests that any valid combination of lora block scales can be used in pipe.set_adapter""" + if self.unet_kwargs is None: + pytest.skip("Test not supported.") + + def updown_options(blocks_with_tf, layers_per_block, value): + """ + Generate every possible combination for how a lora weight dict for the up/down part can be. + E.g. 2, {"block_1": 2}, {"block_1": [2,2,2]}, {"block_1": 2, "block_2": [2,2,2]}, ... + """ + num_val = value + list_val = [value] * layers_per_block + + node_opts = [None, num_val, list_val] + node_opts_foreach_block = [node_opts] * len(blocks_with_tf) + + updown_opts = [num_val] + for nodes in product(*node_opts_foreach_block): + if all(n is None for n in nodes): + continue + opt = {} + for b, n in zip(blocks_with_tf, nodes): + if n is not None: + opt["block_" + str(b)] = n + updown_opts.append(opt) + return updown_opts + + def all_possible_dict_opts(unet, value): + """ + Generate every possible combination for how a lora weight dict can be. + E.g. 2, {"unet: {"down": 2}}, {"unet: {"down": [2,2,2]}}, {"unet: {"mid": 2, "up": [2,2,2]}}, ... + """ + + down_blocks_with_tf = [i for i, d in enumerate(unet.down_blocks) if hasattr(d, "attentions")] + up_blocks_with_tf = [i for i, u in enumerate(unet.up_blocks) if hasattr(u, "attentions")] + + layers_per_block = unet.config.layers_per_block + + text_encoder_opts = [None, value] + text_encoder_2_opts = [None, value] + mid_opts = [None, value] + down_opts = [None] + updown_options(down_blocks_with_tf, layers_per_block, value) + up_opts = [None] + updown_options(up_blocks_with_tf, layers_per_block + 1, value) + + opts = [] + + for t1, t2, d, m, u in product(text_encoder_opts, text_encoder_2_opts, down_opts, mid_opts, up_opts): + if all(o is None for o in (t1, t2, d, m, u)): + continue + opt = {} + if t1 is not None: + opt["text_encoder"] = t1 + if t2 is not None: + opt["text_encoder_2"] = t2 + if all(o is None for o in (d, m, u)): + # no unet scaling + continue + opt["unet"] = {} + if d is not None: + opt["unet"]["down"] = d + if m is not None: + opt["unet"]["mid"] = m + if u is not None: + opt["unet"]["up"] = u + opts.append(opt) + + return opts + + pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( + self.scheduler_classes[0] + ) + pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") + denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet + denoiser.add_adapter(denoiser_lora_config, "adapter-1") + + if self.has_two_text_encoders or self.has_three_text_encoders: + lora_loadable_components = self.pipeline_class._lora_loadable_modules + if "text_encoder_2" in lora_loadable_components: + pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") + + for scale_dict in all_possible_dict_opts(pipe.unet, value=1234): + # test if lora block scales can be set with this scale_dict + if not self.has_two_text_encoders and "text_encoder_2" in scale_dict: + del scale_dict["text_encoder_2"] + + pipe.set_adapters("adapter-1", scale_dict) # test will fail if this line throws an error + + def test_simple_inference_with_dora(self): + for scheduler_cls in self.scheduler_classes: + pipe, inputs, output_no_dora_lora, text_lora_config, denoiser_lora_config = ( + self._setup_pipeline_and_get_base_output(scheduler_cls, use_dora=True) + ) + pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) + output_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] + + self.assertFalse( + np.allclose(output_dora_lora, output_no_dora_lora, atol=1e-3, rtol=1e-3), + "DoRA lora should change the output", + ) + def _test_lora_actions(self, action, lora_components_to_add, expected_atol=1e-3, expected_rtol=1e-3): """ A unified test for various LoRA actions (fusing, unloading, saving/loading, etc.) @@ -506,6 +602,104 @@ def test_lora_scaling(self, lora_components_to_add): "The scaling parameter was not correctly restored!", ) + def _test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expected_rtol=1e-3): + for scheduler_cls in self.scheduler_classes: + pipe, inputs, output_no_lora, _ = self._setup_multi_adapter_pipeline(scheduler_cls) + + # Run inference with each adapter individually and mixed + pipe.set_adapters("adapter-1") + output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + + pipe.set_adapters("adapter-2") + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + + pipe.set_adapters(["adapter-1", "adapter-2"]) + output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] + + # --- Assert base multi-adapter behavior --- + self.assertFalse(np.allclose(output_no_lora, output_adapter_1, atol=expected_atol, rtol=expected_rtol)) + self.assertFalse(np.allclose(output_adapter_1, output_adapter_2, atol=expected_atol, rtol=expected_rtol)) + self.assertFalse( + np.allclose(output_adapter_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) + + # --- Scenario-specific logic --- + if scenario == "weighted": + pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) + output_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_weighted, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) + + elif scenario == "block_lora": + scales = {"unet": {"down": 0.8, "mid": 0.5, "up": 0.2}} + pipe.set_adapters("adapter-1", scales) + output_block_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_block_lora_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) + + text_encoder_modules = [k for k in self.pipeline_class._lora_loadable_modules if "text_encoder" in k] + if text_encoder_modules: + text_encoder_module_name = text_encoder_modules[0] + scales_2 = {text_encoder_module_name: 2, "unet": {"down": 5}} + pipe.set_adapters("adapter-1", scales_2) + output_block_lora_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertFalse( + np.allclose(output_block_lora_2, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) + + elif scenario == "delete_adapter": + pipe.set_adapters("adapter-2") + output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] + pipe.delete_adapters("adapter-1") + output_after_delete = pipe(**inputs, generator=torch.manual_seed(0))[0] + # After deleting adapter-1, the output should be the same as using only adapter-2 + self.assertTrue( + np.allclose(output_after_delete, output_adapter_2, atol=expected_atol, rtol=expected_rtol) + ) + + elif scenario == "fused_multi": + # 1. Fuse a single adapter + pipe.set_adapters("adapter-1") + pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) + output_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_adapter_1, output_lora_1_fused, atol=expected_atol, rtol=expected_rtol) + ) + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + + # 2. Fuse both adapters + pipe.set_adapters(["adapter-1", "adapter-2"]) + pipe.fuse_lora( + components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1", "adapter-2"] + ) + self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] + self.assertTrue( + np.allclose(output_all_lora_fused, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) + ) + pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) + self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") + + @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",), ("fused_multi",)]) + def test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expected_rtol=1e-3): + """ + A unified test for various multi-adapter (and single-adapter) LoRA scenarios, including weighting, + block scaling, and adapter deletion. + """ + for cls in inspect.getmro(self.__class__): + if "test_lora_set_adapters_scenarios" in cls.__dict__ and cls is not PeftLoraLoaderMixinTests: + # Skip this test if it is overwritten by child class. We need to do this because parameterized + # materializes the test methods on invocation which cannot be overridden. + return + + if scenario == "block_lora" and self.unet_kwargs is None: + pytest.skip(f"Test not supported for {scenario=}.") + + self._test_lora_set_adapters_scenarios(scenario, expected_atol, expected_rtol) + def test_wrong_adapter_name_raises_error(self): adapter_name = "adapter-1" scheduler_cls = self.scheduler_classes[0] @@ -549,93 +743,6 @@ def test_multiple_wrong_adapter_name_raises_error(self): pipe.set_adapters(adapter_name) _ = pipe(**inputs, generator=torch.manual_seed(0))[0] - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - """Tests that any valid combination of lora block scales can be used in pipe.set_adapter""" - if self.unet_kwargs is None: - pytest.skip("Test not supported.") - - def updown_options(blocks_with_tf, layers_per_block, value): - """ - Generate every possible combination for how a lora weight dict for the up/down part can be. - E.g. 2, {"block_1": 2}, {"block_1": [2,2,2]}, {"block_1": 2, "block_2": [2,2,2]}, ... - """ - num_val = value - list_val = [value] * layers_per_block - - node_opts = [None, num_val, list_val] - node_opts_foreach_block = [node_opts] * len(blocks_with_tf) - - updown_opts = [num_val] - for nodes in product(*node_opts_foreach_block): - if all(n is None for n in nodes): - continue - opt = {} - for b, n in zip(blocks_with_tf, nodes): - if n is not None: - opt["block_" + str(b)] = n - updown_opts.append(opt) - return updown_opts - - def all_possible_dict_opts(unet, value): - """ - Generate every possible combination for how a lora weight dict can be. - E.g. 2, {"unet: {"down": 2}}, {"unet: {"down": [2,2,2]}}, {"unet: {"mid": 2, "up": [2,2,2]}}, ... - """ - - down_blocks_with_tf = [i for i, d in enumerate(unet.down_blocks) if hasattr(d, "attentions")] - up_blocks_with_tf = [i for i, u in enumerate(unet.up_blocks) if hasattr(u, "attentions")] - - layers_per_block = unet.config.layers_per_block - - text_encoder_opts = [None, value] - text_encoder_2_opts = [None, value] - mid_opts = [None, value] - down_opts = [None] + updown_options(down_blocks_with_tf, layers_per_block, value) - up_opts = [None] + updown_options(up_blocks_with_tf, layers_per_block + 1, value) - - opts = [] - - for t1, t2, d, m, u in product(text_encoder_opts, text_encoder_2_opts, down_opts, mid_opts, up_opts): - if all(o is None for o in (t1, t2, d, m, u)): - continue - opt = {} - if t1 is not None: - opt["text_encoder"] = t1 - if t2 is not None: - opt["text_encoder_2"] = t2 - if all(o is None for o in (d, m, u)): - # no unet scaling - continue - opt["unet"] = {} - if d is not None: - opt["unet"]["down"] = d - if m is not None: - opt["unet"]["mid"] = m - if u is not None: - opt["unet"]["up"] = u - opts.append(opt) - - return opts - - pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( - self.scheduler_classes[0] - ) - pipe.text_encoder.add_adapter(text_lora_config, "adapter-1") - denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet - denoiser.add_adapter(denoiser_lora_config, "adapter-1") - - if self.has_two_text_encoders or self.has_three_text_encoders: - lora_loadable_components = self.pipeline_class._lora_loadable_modules - if "text_encoder_2" in lora_loadable_components: - pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1") - - for scale_dict in all_possible_dict_opts(pipe.unet, value=1234): - # test if lora block scales can be set with this scale_dict - if not self.has_two_text_encoders and "text_encoder_2" in scale_dict: - del scale_dict["text_encoder_2"] - - pipe.set_adapters("adapter-1", scale_dict) # test will fail if this line throws an error - @skip_mps @pytest.mark.xfail( condition=torch.device(torch_device).type == "cpu" and is_torch_version(">=", "2.5"), @@ -725,7 +832,6 @@ def test_get_list_adapters(self): pipe, _, _, text_lora_config, denoiser_lora_config = self._setup_pipeline_and_get_base_output( scheduler_cls ) - # 1. dicts_to_be_checked = {} if "text_encoder" in self.pipeline_class._lora_loadable_modules: @@ -738,7 +844,6 @@ def test_get_list_adapters(self): else: pipe.transformer.add_adapter(denoiser_lora_config, "adapter-1") dicts_to_be_checked.update({"transformer": ["adapter-1"]}) - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) # 2. @@ -753,12 +858,11 @@ def test_get_list_adapters(self): else: pipe.transformer.add_adapter(denoiser_lora_config, "adapter-2") dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) - self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) - # 3. pipe.set_adapters(["adapter-1", "adapter-2"]) + # 3. dicts_to_be_checked = {} if "text_encoder" in self.pipeline_class._lora_loadable_modules: dicts_to_be_checked = {"text_encoder": ["adapter-1", "adapter-2"]} @@ -767,11 +871,7 @@ def test_get_list_adapters(self): dicts_to_be_checked.update({"unet": ["adapter-1", "adapter-2"]}) else: dicts_to_be_checked.update({"transformer": ["adapter-1", "adapter-2"]}) - - self.assertDictEqual( - pipe.get_list_adapters(), - dicts_to_be_checked, - ) + self.assertDictEqual(pipe.get_list_adapters(), dicts_to_be_checked) # 4. dicts_to_be_checked = {} @@ -819,19 +919,6 @@ def test_lora_scale_kwargs_match_fusion(self, expected_atol=1e-3, expected_rtol= "LoRA should change the output", ) - def test_simple_inference_with_dora(self): - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_dora_lora, text_lora_config, denoiser_lora_config = ( - self._setup_pipeline_and_get_base_output(scheduler_cls, use_dora=True) - ) - pipe, _ = self.add_adapters_to_pipeline(pipe, text_lora_config, denoiser_lora_config) - output_dora_lora = pipe(**inputs, generator=torch.manual_seed(0))[0] - - self.assertFalse( - np.allclose(output_dora_lora, output_no_dora_lora, atol=1e-3, rtol=1e-3), - "DoRA lora should change the output", - ) - def test_missing_keys_warning(self): scheduler_cls = self.scheduler_classes[0] # Skip text encoder check for now as that is handled with `transformers`. @@ -1369,104 +1456,6 @@ def test_lora_exclude_modules(self): "Lora outputs should match.", ) - def _test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expected_rtol=1e-3): - for scheduler_cls in self.scheduler_classes: - pipe, inputs, output_no_lora, _ = self._setup_multi_adapter_pipeline(scheduler_cls) - - # Run inference with each adapter individually and mixed - pipe.set_adapters("adapter-1") - output_adapter_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - - pipe.set_adapters(["adapter-1", "adapter-2"]) - output_adapter_mixed = pipe(**inputs, generator=torch.manual_seed(0))[0] - - # --- Assert base multi-adapter behavior --- - self.assertFalse(np.allclose(output_no_lora, output_adapter_1, atol=expected_atol, rtol=expected_rtol)) - self.assertFalse(np.allclose(output_adapter_1, output_adapter_2, atol=expected_atol, rtol=expected_rtol)) - self.assertFalse( - np.allclose(output_adapter_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) - ) - - # --- Scenario-specific logic --- - if scenario == "weighted": - pipe.set_adapters(["adapter-1", "adapter-2"], [0.5, 0.6]) - output_weighted = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_weighted, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) - ) - - elif scenario == "block_lora": - scales = {"unet": {"down": 0.8, "mid": 0.5, "up": 0.2}} - pipe.set_adapters("adapter-1", scales) - output_block_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_block_lora_1, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) - ) - - text_encoder_modules = [k for k in self.pipeline_class._lora_loadable_modules if "text_encoder" in k] - if text_encoder_modules: - text_encoder_module_name = text_encoder_modules[0] - scales_2 = {text_encoder_module_name: 2, "unet": {"down": 5}} - pipe.set_adapters("adapter-1", scales_2) - output_block_lora_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertFalse( - np.allclose(output_block_lora_2, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) - ) - - elif scenario == "delete_adapter": - pipe.set_adapters("adapter-2") - output_adapter_2 = pipe(**inputs, generator=torch.manual_seed(0))[0] - pipe.delete_adapters("adapter-1") - output_after_delete = pipe(**inputs, generator=torch.manual_seed(0))[0] - # After deleting adapter-1, the output should be the same as using only adapter-2 - self.assertTrue( - np.allclose(output_after_delete, output_adapter_2, atol=expected_atol, rtol=expected_rtol) - ) - - elif scenario == "fused_multi": - # 1. Fuse a single adapter - pipe.set_adapters("adapter-1") - pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"]) - output_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_adapter_1, output_lora_1_fused, atol=expected_atol, rtol=expected_rtol) - ) - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - - # 2. Fuse both adapters - pipe.set_adapters(["adapter-1", "adapter-2"]) - pipe.fuse_lora( - components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1", "adapter-2"] - ) - self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0] - self.assertTrue( - np.allclose(output_all_lora_fused, output_adapter_mixed, atol=expected_atol, rtol=expected_rtol) - ) - pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules) - self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}") - - @parameterized.expand([("simple",), ("weighted",), ("block_lora",), ("delete_adapter",), ("fused_multi",)]) - def test_lora_set_adapters_scenarios(self, scenario, expected_atol=1e-3, expected_rtol=1e-3): - """ - A unified test for various multi-adapter (and single-adapter) LoRA scenarios, including weighting, - block scaling, and adapter deletion. - """ - for cls in inspect.getmro(self.__class__): - if "test_lora_set_adapters_scenarios" in cls.__dict__ and cls is not PeftLoraLoaderMixinTests: - # Skip this test if it is overwritten by child class. We need to do this because parameterized - # materializes the test methods on invocation which cannot be overridden. - return - - if scenario == "block_lora" and self.unet_kwargs is None: - pytest.skip(f"Test not supported for {scenario=}.") - - self._test_lora_set_adapters_scenarios(scenario, expected_atol, expected_rtol) - def _test_group_offloading_inference_denoiser(self, offload_type, use_stream): from diffusers.hooks.group_offloading import _get_top_level_group_offload_hook