up

Author: sayakpaul

src/diffusers/loaders/lora_conversion_utils.py

@@ -1978,13 +1978,11 @@ def get_alpha_scales(down_weight, alpha_key):
                 )
 
         if any("head.head" in k for k in original_state_dict):
-            if any(
-                f"head.head.{lora_down_key}.weight" in k and f"head.head.{lora_up_key}.weight" in k
-                for k in original_state_dict
-            ):
+            if any(f"head.head.{lora_down_key}.weight" in k for k in state_dict):
                 converted_state_dict["proj_out.lora_A.weight"] = original_state_dict.pop(
                     f"head.head.{lora_down_key}.weight"
                 )
+            if any(f"head.head.{lora_up_key}.weight" in k for k in state_dict):
                 converted_state_dict["proj_out.lora_B.weight"] = original_state_dict.pop(
                     f"head.head.{lora_up_key}.weight"
                 )
@@ -1995,7 +1993,7 @@ def get_alpha_scales(down_weight, alpha_key):
             # This is my (sayakpaul) assumption that this particular key belongs to the down matrix.
             # Since for this particular LoRA, we don't have the corresponding up matrix, I will use
             # an identity.
-            if any("head.head.diff" in k for k in state_dict):
+            if any("head.head" in k and k.endswith(".diff") for k in state_dict):
                 if f"head.head.{lora_down_key}.weight" in state_dict:
                     logger.info(
                         f"The state dict seems to be have both `head.head.diff` and `head.head.{lora_down_key}.weight` keys, which is unexpected."

src/diffusers/loaders/lora_pipeline.py

@@ -4074,6 +4074,7 @@ def load_lora_weights(
             raise ValueError("Invalid LoRA checkpoint.")
 
         load_into_transformer_2 = kwargs.pop("load_into_transformer_2", False)
+        print(f"{load_into_transformer_2=}")
         if load_into_transformer_2:
             if not hasattr(self, "transformer_2"):
                 raise AttributeError(