fixes

Author: sayakpaul

src/diffusers/loaders/lora_base.py

@@ -344,6 +344,10 @@ def _load_lora_into_text_encoder(
 
     # Safe prefix to check with.
     if any(text_encoder_name in key for key in keys):
+        # adapter_name
+        if adapter_name is None:
+            adapter_name = get_adapter_name(text_encoder)
+
         # Load the layers corresponding to text encoder and make necessary adjustments.
         text_encoder_keys = [k for k in keys if k.startswith(prefix) and k.split(".")[0] == prefix]
         text_encoder_lora_state_dict = {
@@ -358,20 +362,6 @@ def _load_lora_into_text_encoder(
             # convert state dict
             text_encoder_lora_state_dict = convert_state_dict_to_peft(text_encoder_lora_state_dict)
 
-            if any("position_embedding" in k for k in text_encoder_lora_state_dict):
-                # TODO: this copying is a big shot in the dark.
-                # https://huggingface.co/sayakpaul/different-lora-from-civitai/tree/main?show_file_info=RM_Artistify_v1.0M.safetensors
-                # only has LoRA keys for the position embedding but not the LoRA embedding keys.
-                text_encoder_lora_state_dict[
-                    "text_model.embeddings.position_embedding.lora_embedding_A.weight"
-                ] = text_encoder_lora_state_dict["text_model.embeddings.position_embedding.lora_A.weight"].clone()
-                text_encoder_lora_state_dict[
-                    "text_model.embeddings.position_embedding.lora_embedding_B.weight"
-                ] = text_encoder_lora_state_dict["text_model.embeddings.position_embedding.lora_B.weight"].clone()
-                rank["text_model.embeddings.position_embedding.lora_B.weight"] = text_encoder_lora_state_dict[
-                    "text_model.embeddings.position_embedding.lora_B.weight"
-                ].shape[1]
-
             for name, _ in text_encoder_attn_modules(text_encoder):
                 for module in ("out_proj", "q_proj", "k_proj", "v_proj"):
                     rank_key = f"{name}.{module}.lora_B.weight"
@@ -414,10 +404,6 @@ def _load_lora_into_text_encoder(
 
             lora_config = LoraConfig(**lora_config_kwargs)
 
-            # adapter_name
-            if adapter_name is None:
-                adapter_name = get_adapter_name(text_encoder)
-
             is_model_cpu_offload, is_sequential_cpu_offload = _func_optionally_disable_offloading(_pipeline)
 
             # inject LoRA layers and load the state dict

src/diffusers/loaders/lora_conversion_utils.py

@@ -17,7 +17,7 @@
 
 import torch
 
-from ..utils import is_peft_version, logging
+from ..utils import is_peft_version, logging, state_dict_all_zero
 
 
 logger = logging.get_logger(__name__)
@@ -755,29 +755,67 @@ def _convert(original_key, diffusers_key, state_dict, new_state_dict):
     state_dict = {k.replace("diffusion_model.", "lora_unet_"): v for k, v in state_dict.items()}
     state_dict = {k.replace("text_encoders.clip_l.transformer.", "lora_te_"): v for k, v in state_dict.items()}
     has_t5xxl = any(k.startswith("text_encoders.t5xxl.transformer.") for k in state_dict)
+
+    if any("position_embedding" in k for k in state_dict):
+        zero_status_pe = state_dict_all_zero(state_dict, "position_embedding")
+        if zero_status_pe:
+            logger.info(
+                "The `position_embedding` LoRA params are all zeros which make them ineffective. "
+                "So, we will purge them out of the curret state dict to make loading possible."
+            )
+            current_pe_lora_keys = [k for k in state_dict if "position_embedding" in k]
+            for k in current_pe_lora_keys:
+                state_dict.pop(k)
+        else:
+            raise NotImplementedError(
+                "The state_dict has position_embedding LoRA params and we currently do not support them. "
+                "Open an issue if you need this supported - https://github.com/huggingface/diffusers/issues/new."
+            )
+
     if has_t5xxl:
-        logger.info(
-            "T5-xxl keys found in the state dict, which are currently unsupported. We will filter them out."
-            "Open an issue if this is a problem - https://github.com/huggingface/diffusers/issues/new."
-        )
+        zero_status_t5 = state_dict_all_zero(state_dict, "text_encoders.t5xxl")
+        if zero_status_t5:
+            logger.info(
+                "The `t5xxl` LoRA params are all zeros which make them ineffective. "
+                "So, we will purge them out of the curret state dict to make loading possible."
+            )
+        else:
+            logger.info(
+                "T5-xxl keys found in the state dict, which are currently unsupported. We will filter them out."
+                "Open an issue if this is a problem - https://github.com/huggingface/diffusers/issues/new."
+            )
         state_dict = {k: v for k, v in state_dict.items() if not k.startswith("text_encoders.t5xxl.transformer.")}
 
     any_diffb_keys = any("diff_b" in k and k.startswith(("lora_unet_", "lora_te_")) for k in state_dict)
     if any_diffb_keys:
-        logger.info(
-            "`diff_b` keys found in the state dict which are currently unsupported. "
-            "So, we will filter out those keys. Open an issue if this is a problem - "
-            "https://github.com/huggingface/diffusers/issues/new."
-        )
+        zero_status_diff_b = state_dict_all_zero(state_dict, "diff_b")
+        if zero_status_diff_b:
+            logger.info(
+                "The `diff_b` LoRA params are all zeros which make them ineffective. "
+                "So, we will purge them out of the curret state dict to make loading possible."
+            )
+        else:
+            logger.info(
+                "`diff_b` keys found in the state dict which are currently unsupported. "
+                "So, we will filter out those keys. Open an issue if this is a problem - "
+                "https://github.com/huggingface/diffusers/issues/new."
+            )
         state_dict = {k: v for k, v in state_dict.items() if "diff_b" not in k}
 
     any_norm_diff_keys = any("norm" in k and "diff" in k for k in state_dict)
     if any_norm_diff_keys:
-        logger.info(
-            "Normalization diff keys found in the state dict which are currently unsupported. "
-            "So, we will filter out those keys. Open an issue if this is a problem - "
-            "https://github.com/huggingface/diffusers/issues/new."
-        )
+        zero_status_diff = state_dict_all_zero(state_dict, "diff")
+        if zero_status_diff:
+            logger.info(
+                "The `diff` LoRA params are all zeros which make them ineffective. "
+                "So, we will purge them out of the curret state dict to make loading possible."
+            )
+        else:
+            logger.info(
+                "Normalization diff keys found in the state dict which are currently unsupported. "
+                "So, we will filter out those keys. Open an issue if this is a problem - "
+                "https://github.com/huggingface/diffusers/issues/new."
+            )
         state_dict = {k: v for k, v in state_dict.items() if "norm" not in k and "diff" not in k}
 
     limit_substrings = ["lora_down", "lora_up"]

src/diffusers/utils/__init__.py

@@ -123,6 +123,7 @@
     convert_state_dict_to_kohya,
     convert_state_dict_to_peft,
     convert_unet_state_dict_to_peft,
+    state_dict_all_zero,
 )
 from .typing_utils import _get_detailed_type, _is_valid_type
 

src/diffusers/utils/state_dict_utils.py

@@ -17,9 +17,14 @@
 
 import enum
 
+from .import_utils import is_torch_available
 from .logging import get_logger
 
 
+if is_torch_available():
+    import torch
+
+
 logger = get_logger(__name__)
 
 
@@ -64,8 +69,8 @@ class StateDictType(enum.Enum):
     ".lora_linear_layer.down": ".lora_A",
     "text_projection.lora.down.weight": "text_projection.lora_A.weight",
     "text_projection.lora.up.weight": "text_projection.lora_B.weight",
-    "position_embedding.lora.down.weight": "position_embedding.lora_A.weight",
-    "position_embedding.lora.up.weight": "position_embedding.lora_B.weight",
+    "position_embedding.lora.down.weight": "position_embedding.lora_embedding_A",
+    "position_embedding.lora.up.weight": "position_embedding.lora_embedding_B",
 }
 
 DIFFUSERS_OLD_TO_PEFT = {
@@ -335,3 +340,12 @@ def convert_state_dict_to_kohya(state_dict, original_type=None, **kwargs):
             kohya_ss_state_dict[alpha_key] = torch.tensor(len(weight))
 
     return kohya_ss_state_dict
+
+
+def state_dict_all_zero(state_dict, filter_str=None):
+    if filter_str is not None:
+        if isinstance(filter_str, str):
+            filter_str = [filter_str]
+        state_dict = {k: v for k, v in state_dict.items() if any(f in k for f in filter_str)}
+
+    return all(torch.all(param == 0).item() for param in state_dict.values())