fix test_set_adapters_match_attention_kwargs

Author: sayakpaul

src/diffusers/models/__init__.py

@@ -114,7 +114,7 @@
     _import_structure["transformers.transformer_hunyuan_video_framepack"] = ["HunyuanVideoFramepackTransformer3DModel"]
     _import_structure["transformers.transformer_hunyuanimage"] = ["HunyuanImageTransformer2DModel"]
     _import_structure["transformers.transformer_joyimage"] = [
-        "JoyImageEditTransformer3DModel",
+        "JoyImageEditTransformer3DModel"
     ]
     _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"]
     _import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"]

tests/lora/utils.py

@@ -1962,6 +1962,9 @@ def test_set_adapters_match_attention_kwargs(self):
             "Lora + scale should match the output of `set_adapters()`.",
         )
 
+        needs_lora_repair = self._needs_text_encoder_lora_repair()
+        captured_lora = self._capture_text_encoder_lora_tensors(pipe) if needs_lora_repair else {}
+
         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)
@@ -1975,6 +1978,9 @@ def test_set_adapters_match_attention_kwargs(self):
             pipe.set_progress_bar_config(disable=None)
             pipe.load_lora_weights(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors"))
 
+            if needs_lora_repair:
+                self._restore_text_encoder_lora_tensors(pipe, captured_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}")