updates

Author: sayakpaul

src/diffusers/loaders/lora_pipeline.py

@@ -1863,7 +1863,9 @@ def load_lora_weights(
                 "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)
+        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(
@@ -2385,29 +2387,32 @@ def _maybe_expand_transformer_param_shape_or_error_(
     def _maybe_expand_lora_state_dict(cls, transformer, lora_state_dict):
         expanded_module_names = set()
         transformer_state_dict = transformer.state_dict()
-        lora_module_names = set([k.replace(".lora_A.weight", "") for k in lora_state_dict if "lora_A" in k])
-        lora_module_names = sorted(lora_module_names)
+        lora_module_names = sorted({k.replace(".lora_A.weight", "") for k in lora_state_dict if "lora_A" in k})
         is_peft_loaded = getattr(transformer, "peft_config", None) is not None
 
         for k in lora_module_names:
-            base_param_name = f"{k.replace(f'{cls.transformer_name}.', '')}.base_layer.weight" if is_peft_loaded else f"{k.replace(f'{cls.transformer_name}.', '')}.weight"
+            base_param_name = (
+                f"{k.replace(f'{cls.transformer_name}.', '')}.base_layer.weight"
+                if is_peft_loaded
+                else f"{k.replace(f'{cls.transformer_name}.', '')}.weight"
+            )
             base_weight_param = transformer_state_dict[base_param_name]
             lora_A_param = lora_state_dict[f"{k}.lora_A.weight"]
-            # lora_B_param = lora_state_dict[f"{k}.lora_B.weight"]
 
             if base_weight_param.shape[1] > lora_A_param.shape[1]:
-                shape = (lora_A_param.shape[0], base_weight_param.shape[1])
-                expanded_state_dict_weight = torch.zeros(shape, device=base_weight_param.device)
-                expanded_state_dict_weight[:, :lora_A_param.shape[1]].copy_(lora_A_param)
+                # could be made more advanced with `repeats`.
+                # have tried zero-padding but that doesn't work, either.
+                expanded_state_dict_weight = torch.cat([lora_A_param, lora_A_param], dim=1)
                 lora_state_dict[f"{k}.lora_A.weight"] = expanded_state_dict_weight
                 expanded_module_names.add(k)
 
         if expanded_module_names:
-            logger.info(f"Found some LoRA modules for which the weights were expanded: {expanded_module_names}. Please open an issue if you think this was unexpected - https://github.com/huggingface/diffusers/issues/new.")
+            logger.info(
+                f"Found some LoRA modules for which the weights were expanded: {expanded_module_names}. Please open an issue if you think this was unexpected - https://github.com/huggingface/diffusers/issues/new."
+            )
         return lora_state_dict
 
 
-
 # The reason why we subclass from `StableDiffusionLoraLoaderMixin` here is because Amused initially
 # relied on `StableDiffusionLoraLoaderMixin` for its LoRA support.
 class AmusedLoraLoaderMixin(StableDiffusionLoraLoaderMixin):