fix tests

Author: sayakpaul

src/diffusers/loaders/lora_pipeline.py

@@ -4889,7 +4889,7 @@ def load_lora_weights(
             adapter_name=adapter_name,
             _pipeline=self,
             low_cpu_mem_usage=low_cpu_mem_usage,
-            load_with_metdata=load_with_metdata,
+            load_with_metadata=load_with_metdata,
         )
 
     @classmethod

src/diffusers/loaders/peft.py

@@ -236,15 +236,12 @@ def load_lora_adapter(
         if network_alphas is not None and prefix is None:
             raise ValueError("`network_alphas` cannot be None when `prefix` is None.")
 
-        if load_with_metadata is not None and not use_safetensors:
-            raise ValueError("`load_with_metadata` cannot be specified when not using `use_safetensors`.")
-
         if prefix is not None:
-            metadata = state_dict.pop("_metadata", None)
+            metadata = state_dict.pop("lora_metadata", None)
             state_dict = {k[len(f"{prefix}.") :]: v for k, v in state_dict.items() if k.startswith(f"{prefix}.")}
 
             if metadata is not None:
-                state_dict["_metadata"] = metadata
+                state_dict["lora_metadata"] = metadata
 
         if len(state_dict) > 0:
             if adapter_name in getattr(self, "peft_config", {}) and not hotswap:

src/diffusers/utils/peft_utils.py

@@ -151,9 +151,9 @@ def get_peft_kwargs(
     rank_dict, network_alpha_dict, peft_state_dict, is_unet=True, prefix=None, load_with_metadata=False
 ):
     if load_with_metadata:
-        if "_metadata" not in peft_state_dict:
+        if "lora_metadata" not in peft_state_dict:
             raise ValueError("Couldn't find '_metadata' key in the `peft_state_dict`.")
-        metadata = peft_state_dict["_metadata"]
+        metadata = peft_state_dict["lora_metadata"]
         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

@@ -360,7 +360,7 @@ def _maybe_populate_state_dict_with_metadata(state_dict, model_file, metadata_ke
                 metadata_keys = list(metadata.keys())
                 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["_metadata"] = json.loads(peft_metadata[metadata_key])
+                    state_dict["lora_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

@@ -162,9 +162,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("_metadata" in state_dict)
+            self.assertTrue("lora_metadata" in state_dict)
 
-            parsed_metadata = state_dict["_metadata"]
+            parsed_metadata = state_dict["lora_metadata"]
             parsed_metadata = {k[len("transformer.") :]: v for k, v in parsed_metadata.items()}
             check_if_dicts_are_equal(parsed_metadata, metadata)