load metadata.

Author: sayakpaul

src/diffusers/loaders/lora_base.py

@@ -64,6 +64,7 @@
 
 LORA_WEIGHT_NAME = "pytorch_lora_weights.bin"
 LORA_WEIGHT_NAME_SAFE = "pytorch_lora_weights.safetensors"
+LORA_ADAPTER_METADATA_KEY = "lora_adapter_metadata"
 
 
 def fuse_text_encoder_lora(text_encoder, lora_scale=1.0, safe_fusing=False, adapter_names=None):
@@ -208,7 +209,6 @@ def _fetch_state_dict(
     subfolder,
     user_agent,
     allow_pickle,
-    load_with_metadata=False,
 ):
     model_file = None
     if not isinstance(pretrained_model_name_or_path_or_dict, dict):
@@ -226,8 +226,6 @@ def _fetch_state_dict(
                         file_extension=".safetensors",
                         local_files_only=local_files_only,
                     )
-                if load_with_metadata and not weight_name.endswith(".safetensors"):
-                    raise ValueError("`load_with_metadata` cannot be set to True when not using safetensors.")
 
                 model_file = _get_model_file(
                     pretrained_model_name_or_path_or_dict,
@@ -242,10 +240,7 @@ def _fetch_state_dict(
                     user_agent=user_agent,
                 )
                 state_dict = safetensors.torch.load_file(model_file, device="cpu")
-                if load_with_metadata:
-                    state_dict = _maybe_populate_state_dict_with_metadata(
-                        state_dict, model_file, metadata_key="lora_adapter_metadata"
-                    )
+                state_dict = _maybe_populate_state_dict_with_metadata(state_dict, model_file)
 
             except (IOError, safetensors.SafetensorError) as e:
                 if not allow_pickle:

src/diffusers/loaders/lora_pipeline.py

@@ -4727,7 +4727,6 @@ def lora_state_dict(
                     - A [torch state
                       dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
 
-            load_with_metadata: TODO
             cache_dir (`Union[str, os.PathLike]`, *optional*):
                 Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
                 is not used.
@@ -4762,7 +4761,6 @@ def lora_state_dict(
         subfolder = kwargs.pop("subfolder", None)
         weight_name = kwargs.pop("weight_name", None)
         use_safetensors = kwargs.pop("use_safetensors", None)
-        load_with_metadata = kwargs.pop("load_with_metadata", False)
 
         allow_pickle = False
         if use_safetensors is None:
@@ -4787,7 +4785,6 @@ def lora_state_dict(
             subfolder=subfolder,
             user_agent=user_agent,
             allow_pickle=allow_pickle,
-            load_with_metadata=load_with_metadata,
         )
         if any(k.startswith("diffusion_model.") for k in state_dict):
             state_dict = _convert_non_diffusers_wan_lora_to_diffusers(state_dict)
@@ -4861,7 +4858,6 @@ def load_lora_weights(
             raise ValueError("PEFT backend is required for this method.")
 
         low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
-        load_with_metadata = kwargs.get("load_with_metadata", False)
         if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
             raise ValueError(
                 "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
@@ -4888,7 +4884,6 @@ def load_lora_weights(
             adapter_name=adapter_name,
             _pipeline=self,
             low_cpu_mem_usage=low_cpu_mem_usage,
-            load_with_metadata=load_with_metadata,
             hotswap=hotswap,
         )
 
@@ -4902,54 +4897,25 @@ def load_lora_into_transformer(
         _pipeline=None,
         low_cpu_mem_usage=False,
         hotswap: bool = False,
-        load_with_metadata: bool = False,
     ):
         """
-                This will load the LoRA layers specified in `state_dict` into `transformer`.
-
-                Parameters:
-                    state_dict (`dict`):
-                        A standard state dict containing the lora layer parameters. The keys can either be indexed
-                        directly into the unet or prefixed with an additional `unet` which can be used to distinguish
-                        between text encoder lora layers.
-                    transformer (`WanTransformer3DModel`):
-                        The Transformer model to load the LoRA layers into.
-                    adapter_name (`str`, *optional*):
-                        Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
-                        `default_{i}` where i is the total number of adapters being loaded.
-                    low_cpu_mem_usage (`bool`, *optional*):
-                        Speed up model loading by only loading the pretrained LoRA weights and not initializing the
-                        random weights.
-        <<<<<<< HEAD
-                    hotswap : (`bool`, *optional*)
-                        Defaults to `False`. Whether to substitute an existing (LoRA) adapter with the newly loaded
-                        adapter in-place. This means that, instead of loading an additional adapter, this will take the
-                        existing adapter weights and replace them with the weights of the new adapter. This can be
-                        faster and more memory efficient. However, the main advantage of hotswapping is that when the
-                        model is compiled with torch.compile, loading the new adapter does not require recompilation of
-                        the model. When using hotswapping, the passed `adapter_name` should be the name of an already
-                        loaded adapter.
-
-                        If the new adapter and the old adapter have different ranks and/or LoRA alphas (i.e. scaling),
-                        you need to call an additional method before loading the adapter:
-
-                        ```py
-                        pipeline = ...  # load diffusers pipeline
-                        max_rank = ...  # the highest rank among all LoRAs that you want to load
-                        # call *before* compiling and loading the LoRA adapter
-                        pipeline.enable_lora_hotswap(target_rank=max_rank)
-                        pipeline.load_lora_weights(file_name)
-                        # optionally compile the model now
-                        ```
-
-                        Note that hotswapping adapters of the text encoder is not yet supported. There are some further
-                        limitations to this technique, which are documented here:
-                        https://huggingface.co/docs/peft/main/en/package_reference/hotswap
-                    load_with_metadata: TODO
-        =======
-                    hotswap (`bool`, *optional*):
-                        See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
-        >>>>>>> main
+        This will load the LoRA layers specified in `state_dict` into `transformer`.
+
+        Parameters:
+            state_dict (`dict`):
+                A standard state dict containing the lora layer parameters. The keys can either be indexed directly
+                into the unet or prefixed with an additional `unet` which can be used to distinguish between text
+                encoder lora layers.
+            transformer (`WanTransformer3DModel`):
+                The Transformer model to load the LoRA layers into.
+            adapter_name (`str`, *optional*):
+                Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
+                `default_{i}` where i is the total number of adapters being loaded.
+            low_cpu_mem_usage (`bool`, *optional*):
+                Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
+                weights.
+            hotswap (`bool`, *optional*):
+                See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
         """
         if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
             raise ValueError(
@@ -4965,7 +4931,6 @@ def load_lora_into_transformer(
             _pipeline=_pipeline,
             low_cpu_mem_usage=low_cpu_mem_usage,
             hotswap=hotswap,
-            load_with_metadata=load_with_metadata,
         )
 
     @classmethod

src/diffusers/loaders/peft.py

@@ -120,7 +120,6 @@ def load_lora_adapter(
         pretrained_model_name_or_path_or_dict,
         prefix="transformer",
         hotswap: bool = False,
-        load_with_metadata: bool = False,
         **kwargs,
     ):
         r"""
@@ -190,7 +189,6 @@ def load_lora_adapter(
                 limitations to this technique, which are documented here:
                 https://huggingface.co/docs/peft/main/en/package_reference/hotswap
 
-            load_with_metadata: TODO
         """
         from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
         from peft.tuners.tuners_utils import BaseTunerLayer
@@ -233,7 +231,6 @@ def load_lora_adapter(
             subfolder=subfolder,
             user_agent=user_agent,
             allow_pickle=allow_pickle,
-            load_with_metadata=load_with_metadata,
         )
         if network_alphas is not None and prefix is None:
             raise ValueError("`network_alphas` cannot be None when `prefix` is None.")
@@ -280,7 +277,6 @@ def load_lora_adapter(
                 rank,
                 network_alpha_dict=network_alphas,
                 peft_state_dict=state_dict,
-                load_with_metadata=load_with_metadata,
                 prefix=prefix,
             )
             _maybe_raise_error_for_ambiguity(lora_config_kwargs)

src/diffusers/utils/peft_utils.py

@@ -148,12 +148,16 @@ def unscale_lora_layers(model, weight: Optional[float] = None):
 
 
 def get_peft_kwargs(
-    rank_dict, network_alpha_dict, peft_state_dict, is_unet=True, prefix=None, load_with_metadata=False
+    rank_dict,
+    network_alpha_dict,
+    peft_state_dict,
+    is_unet=True,
+    prefix=None,
 ):
-    if load_with_metadata:
-        if "lora_adapter_metadata" not in peft_state_dict:
-            raise ValueError("Couldn't find 'lora_adapter_metadata' key in the `peft_state_dict`.")
-        metadata = peft_state_dict["lora_adapter_metadata"]
+    from ..loaders.lora_base import LORA_ADAPTER_METADATA_KEY
+
+    if LORA_ADAPTER_METADATA_KEY in peft_state_dict:
+        metadata = peft_state_dict[LORA_ADAPTER_METADATA_KEY]
         if prefix is not None:
             metadata = {k.replace(f"{prefix}.", ""): v for k, v in metadata.items()}
         return metadata

src/diffusers/utils/state_dict_utils.py

@@ -350,9 +350,15 @@ def state_dict_all_zero(state_dict, filter_str=None):
     return all(torch.all(param == 0).item() for param in state_dict.values())
 
 
-def _maybe_populate_state_dict_with_metadata(state_dict, model_file, metadata_key):
+def _maybe_populate_state_dict_with_metadata(state_dict, model_file):
+    if not model_file.endswith(".safetensors"):
+        return state_dict
+
     import safetensors.torch
 
+    from ..loaders.lora_base import LORA_ADAPTER_METADATA_KEY
+
+    metadata_key = LORA_ADAPTER_METADATA_KEY
     with safetensors.torch.safe_open(model_file, framework="pt", device="cpu") as f:
         if hasattr(f, "metadata"):
             metadata = f.metadata()
@@ -361,6 +367,4 @@ def _maybe_populate_state_dict_with_metadata(state_dict, model_file, metadata_ke
                 if not (len(metadata_keys) == 1 and metadata_keys[0] == "format"):
                     peft_metadata = {k: v for k, v in metadata.items() if k != "format"}
                     state_dict["lora_adapter_metadata"] = json.loads(peft_metadata[metadata_key])
-        else:
-            raise ValueError("Metadata couldn't be parsed from the safetensors file.")
     return state_dict

tests/lora/test_lora_layers_wan.py

@@ -26,6 +26,7 @@
     WanPipeline,
     WanTransformer3DModel,
 )
+from diffusers.loaders.lora_base import LORA_ADAPTER_METADATA_KEY
 from diffusers.utils.testing_utils import floats_tensor, require_peft_backend, skip_mps, torch_device
 
 
@@ -162,9 +163,9 @@ def test_adapter_metadata_is_loaded_correctly(self):
             pipe.unload_lora_weights()
             state_dict = pipe.lora_state_dict(tmpdir, load_with_metadata=True)
 
-            self.assertTrue("lora_metadata" in state_dict)
+            self.assertTrue(LORA_ADAPTER_METADATA_KEY in state_dict)
 
-            parsed_metadata = state_dict["lora_metadata"]
+            parsed_metadata = state_dict[LORA_ADAPTER_METADATA_KEY]
             parsed_metadata = {k[len("transformer.") :]: v for k, v in parsed_metadata.items()}
             check_if_dicts_are_equal(parsed_metadata, metadata)