fix stuff

Author: sayakpaul

src/diffusers/loaders/lora_pipeline.py

@@ -2312,15 +2312,17 @@ def _maybe_expand_transformer_param_shape_or_error_(
 
         # Expand transformer parameter shapes if they don't match lora
         has_param_with_shape_update = False
+        is_peft_loaded = getattr(transformer, "peft_config", None) is not None
         for name, module in transformer.named_modules():
             if isinstance(module, torch.nn.Linear):
                 module_weight = module.weight.data
                 module_bias = module.bias.data if module.bias is not None else None
                 bias = module_bias is not None
 
-                lora_A_weight_name = f"{name}.lora_A.weight"
-                lora_B_weight_name = f"{name}.lora_B.weight"
-                if lora_A_weight_name not in state_dict.keys():
+                lora_base_name = name.replace(".base_layer", "") if is_peft_loaded else name
+                lora_A_weight_name = f"{lora_base_name}.lora_A.weight"
+                lora_B_weight_name = f"{lora_base_name}.lora_B.weight"
+                if lora_A_weight_name not in state_dict:
                     continue
 
                 in_features = state_dict[lora_A_weight_name].shape[1]

tests/lora/test_lora_layers_flux.py

@@ -467,18 +467,14 @@ def test_lora_expanding_shape_with_normal_lora(self):
             pipe.load_lora_weights(lora_state_dict, "adapter-2")
 
         self.assertTrue(check_if_lora_correctly_set(pipe.transformer), "Lora not correctly set in denoiser")
+        self.assertTrue("The following LoRA modules were zero padded to match the state dict of" in cap_logger.out)
         self.assertTrue(pipe.get_active_adapters() == ["adapter-2"])
 
         lora_output_2 = pipe(**inputs, generator=torch.manual_seed(0))[0]
-
-        self.assertTrue("The following LoRA modules were zero padded to match the state dict of" in cap_logger.out)
         self.assertFalse(np.allclose(lora_output, lora_output_2, atol=1e-3, rtol=1e-3))
 
         # Test the opposite case where the first lora has the correct input features and the second lora has expanded input features.
-        # This should raise a runtime error on input shapes being incompatible. But it doesn't. This is because PEFT renames the
-        # original layers as `base_layer` and the lora layers with the adapter names. This makes our logic to check if a lora
-        # weight is compatible with the current model inadequate. This should be addressed when attempting support for
-        # https://github.com/huggingface/diffusers/issues/10180 (TODO)
+        # This should raise a runtime error on input shapes being incompatible.
         components, _, _ = self.get_dummy_components(FlowMatchEulerDiscreteScheduler)
         # Change the transformer config to mimic a real use case.
         num_channels_without_control = 4
@@ -501,24 +497,23 @@ def test_lora_expanding_shape_with_normal_lora(self):
             "transformer.x_embedder.lora_A.weight": normal_lora_A.weight,
             "transformer.x_embedder.lora_B.weight": normal_lora_B.weight,
         }
-
-        with CaptureLogger(logger) as cap_logger:
-            pipe.load_lora_weights(lora_state_dict, "adapter-1")
+        pipe.load_lora_weights(lora_state_dict, "adapter-1")
 
         self.assertTrue(check_if_lora_correctly_set(pipe.transformer), "Lora not correctly set in denoiser")
         self.assertTrue(pipe.transformer.x_embedder.weight.data.shape[1] == in_features)
         self.assertTrue(pipe.transformer.config.in_channels == in_features)
-        self.assertFalse(cap_logger.out.startswith("Expanding the nn.Linear input/output features for module"))
 
         lora_state_dict = {
             "transformer.x_embedder.lora_A.weight": shape_expander_lora_A.weight,
             "transformer.x_embedder.lora_B.weight": shape_expander_lora_B.weight,
         }
 
-        # We should check for input shapes being incompatible here.
+        # We should check for input shapes being incompatible here. But because above mentioned issue is
+        # not a supported use case, and because of the PEFT renaming, we will currently have a shape
+        # mismatch error.
         self.assertRaisesRegex(
             RuntimeError,
-            "x_embedder.lora_A.weight",
+            "size mismatch for x_embedder.lora_A.adapter-2.weight",
             pipe.load_lora_weights,
             lora_state_dict,
             "adapter-2",