fixes

Author: sayakpaul

src/diffusers/loaders/lora_pipeline.py

@@ -1807,7 +1807,7 @@ def load_lora_weights(
             raise ValueError("Invalid LoRA checkpoint.")
 
         transformer_lora_state_dict = {
-            k: state_dict.pop(k)
+            k: state_dict.get(k)
             for k in list(state_dict.keys())
             if k.startswith(self.transformer_name) and "lora" in k
         }
@@ -1819,29 +1819,33 @@ def load_lora_weights(
         }
 
         transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
-        has_param_with_expanded_shape = self._maybe_expand_transformer_param_shape_or_error_(
-            transformer, transformer_lora_state_dict, transformer_norm_state_dict
-        )
+        has_param_with_expanded_shape = False
+        if len(transformer_lora_state_dict) > 0:
+            has_param_with_expanded_shape = self._maybe_expand_transformer_param_shape_or_error_(
+                transformer, transformer_lora_state_dict, transformer_norm_state_dict
+            )
 
         if has_param_with_expanded_shape:
             logger.info(
                 "The LoRA weights contain parameters that have different shapes that expected by the transformer. "
                 "As a result, the state_dict of the transformer has been expanded to match the LoRA parameter shapes. "
                 "To get a comprehensive list of parameter names that were modified, enable debug logging."
             )
-        transformer_lora_state_dict = self._maybe_expand_lora_state_dict(
-            transformer=transformer, lora_state_dict=transformer_lora_state_dict
-        )
-
         if len(transformer_lora_state_dict) > 0:
-            self.load_lora_into_transformer(
-                transformer_lora_state_dict,
-                network_alphas=network_alphas,
-                transformer=transformer,
-                adapter_name=adapter_name,
-                _pipeline=self,
-                low_cpu_mem_usage=low_cpu_mem_usage,
+            transformer_lora_state_dict = self._maybe_expand_lora_state_dict(
+                transformer=transformer, lora_state_dict=transformer_lora_state_dict
             )
+            for k in transformer_lora_state_dict:
+                state_dict.update({k: transformer_lora_state_dict[k]})
+
+        self.load_lora_into_transformer(
+            state_dict,
+            network_alphas=network_alphas,
+            transformer=transformer,
+            adapter_name=adapter_name,
+            _pipeline=self,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+        )
 
         if len(transformer_norm_state_dict) > 0:
             transformer._transformer_norm_layers = self._load_norm_into_transformer(

src/diffusers/loaders/peft.py

@@ -254,7 +254,7 @@ def load_lora_adapter(self, pretrained_model_name_or_path_or_dict, prefix="trans
 
         if prefix is not None:
             state_dict = {k.replace(f"{prefix}.", ""): v for k, v in state_dict.items() if k.startswith(f"{prefix}.")}
-
+        print(f"{len(state_dict)=}")
         if len(state_dict) > 0:
             if adapter_name in getattr(self, "peft_config", {}):
                 raise ValueError(

tests/lora/utils.py

@@ -1923,7 +1923,8 @@ def test_logs_info_when_no_lora_keys_found(self):
             pipe.load_lora_weights(no_op_state_dict)
         out_after_lora_attempt = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
-        self.assertTrue(cap_logger.out.startswith("No LoRA keys found in the provided state dict"))
+        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))
 
         # test only for text encoder
@@ -1943,7 +1944,9 @@ def test_logs_info_when_no_lora_keys_found(self):
                         no_op_state_dict, network_alphas=None, text_encoder=text_encoder, prefix=prefix
                     )
 
-                self.assertTrue(cap_logger.out.startswith("No LoRA keys found in the provided state dict"))
+                self.assertTrue(
+                    cap_logger.out.startswith(f"No LoRA keys associated to {text_encoder.__class__.__name__}")
+                )
 
     def test_set_adapters_match_attention_kwargs(self):
         """Test to check if outputs after `set_adapters()` and attention kwargs match."""