propagate changes.

Author: sayakpaul

src/diffusers/loaders/lora_base.py

@@ -353,8 +353,11 @@ def _load_lora_into_text_encoder(
         raise ValueError("At the moment, hotswapping is not supported for text encoders, please pass `hotswap=False`.")
 
     # Load the layers corresponding to text encoder and make necessary adjustments.
+    if LORA_ADAPTER_METADATA_KEY in state_dict:
+        metadata = state_dict[LORA_ADAPTER_METADATA_KEY]
     if prefix is not None:
         state_dict = {k[len(f"{prefix}.") :]: v for k, v in state_dict.items() if k.startswith(f"{prefix}.")}
+    state_dict[LORA_ADAPTER_METADATA_KEY] = metadata
 
     if len(state_dict) > 0:
         logger.info(f"Loading {prefix}.")
@@ -382,7 +385,7 @@ def _load_lora_into_text_encoder(
             alpha_keys = [k for k in network_alphas.keys() if k.startswith(prefix) and k.split(".")[0] == prefix]
             network_alphas = {k.replace(f"{prefix}.", ""): v for k, v in network_alphas.items() if k in alpha_keys}
 
-        lora_config_kwargs = get_peft_kwargs(rank, network_alphas, state_dict, is_unet=False)
+        lora_config_kwargs = get_peft_kwargs(rank, network_alphas, state_dict, is_unet=False, prefix=prefix)
 
         if "use_dora" in lora_config_kwargs:
             if lora_config_kwargs["use_dora"]:

src/diffusers/loaders/lora_pipeline.py

@@ -644,6 +644,9 @@ def load_lora_weights(
         if not is_correct_format:
             raise ValueError("Invalid LoRA checkpoint.")
 
+        from .lora_base import LORA_ADAPTER_METADATA_KEY
+
+        print(f"{LORA_ADAPTER_METADATA_KEY in state_dict=} before UNet")
         self.load_lora_into_unet(
             state_dict,
             network_alphas=network_alphas,
@@ -653,6 +656,7 @@ def load_lora_weights(
             low_cpu_mem_usage=low_cpu_mem_usage,
             hotswap=hotswap,
         )
+        print(f"{LORA_ADAPTER_METADATA_KEY in state_dict=} before text encoder.")
         self.load_lora_into_text_encoder(
             state_dict,
             network_alphas=network_alphas,
@@ -664,6 +668,7 @@ def load_lora_weights(
             low_cpu_mem_usage=low_cpu_mem_usage,
             hotswap=hotswap,
         )
+        print(f"{LORA_ADAPTER_METADATA_KEY in state_dict=} before text encoder 2.")
         self.load_lora_into_text_encoder(
             state_dict,
             network_alphas=network_alphas,
@@ -732,6 +737,7 @@ def lora_state_dict(
         """
         # Load the main state dict first which has the LoRA layers for either of
         # UNet and text encoder or both.
+
         cache_dir = kwargs.pop("cache_dir", None)
         force_download = kwargs.pop("force_download", False)
         proxies = kwargs.pop("proxies", None)
@@ -914,6 +920,9 @@ def save_lora_weights(
         weight_name: str = None,
         save_function: Callable = None,
         safe_serialization: bool = True,
+        unet_lora_adapter_metadata=None,
+        text_encoder_lora_adapter_metadata=None,
+        text_encoder_2_lora_adapter_metadata=None,
     ):
         r"""
         Save the LoRA parameters corresponding to the UNet and text encoder.
@@ -939,8 +948,12 @@ def save_lora_weights(
                 `DIFFUSERS_SAVE_MODE`.
             safe_serialization (`bool`, *optional*, defaults to `True`):
                 Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
+            unet_lora_adapter_metadata: TODO
+            text_encoder_lora_adapter_metadata: TODO
+            text_encoder_2_lora_adapter_metadata: TODO
         """
         state_dict = {}
+        lora_adapter_metadata = {}
 
         if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
             raise ValueError(
@@ -956,13 +969,23 @@ def save_lora_weights(
         if text_encoder_2_lora_layers:
             state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
 
+        if unet_lora_adapter_metadata is not None:
+            lora_adapter_metadata.update(cls.pack_weights(unet_lora_adapter_metadata, cls.unet_name))
+
+        if text_encoder_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_lora_adapter_metadata, cls.text_encoder_name))
+
+        if text_encoder_2_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_2_lora_adapter_metadata, "text_encoder_2"))
+
         cls.write_lora_layers(
             state_dict=state_dict,
             save_directory=save_directory,
             is_main_process=is_main_process,
             weight_name=weight_name,
             save_function=save_function,
             safe_serialization=safe_serialization,
+            lora_adapter_metadata=lora_adapter_metadata,
         )
 
     def fuse_lora(

src/diffusers/loaders/peft.py

@@ -193,7 +193,7 @@ def load_lora_adapter(
         from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
         from peft.tuners.tuners_utils import BaseTunerLayer
 
-        from .lora_base import LORA_ADAPTER_METADATA_KEY
+        from ..loaders.lora_base import LORA_ADAPTER_METADATA_KEY
 
         cache_dir = kwargs.pop("cache_dir", None)
         force_download = kwargs.pop("force_download", False)
@@ -234,15 +234,14 @@ def load_lora_adapter(
             user_agent=user_agent,
             allow_pickle=allow_pickle,
         )
+        if LORA_ADAPTER_METADATA_KEY in state_dict:
+            metadata = state_dict[LORA_ADAPTER_METADATA_KEY]
         if network_alphas is not None and prefix is None:
             raise ValueError("`network_alphas` cannot be None when `prefix` is None.")
 
         if prefix is not None:
-            state_dict = {k[len(f"{prefix}.") :]: v for k, v in state_dict.items() if k.startswith(f"{prefix}.")}
-
-        metadata = state_dict.pop(LORA_ADAPTER_METADATA_KEY, None)
-        if metadata is not None:
-            state_dict[LORA_ADAPTER_METADATA_KEY] = metadata
+            state_dict = {k.removeprefix(f"{prefix}."): v for k, v in state_dict.items() if k.startswith(f"{prefix}.")}
+        state_dict[LORA_ADAPTER_METADATA_KEY] = metadata
 
         if len(state_dict) > 0:
             if adapter_name in getattr(self, "peft_config", {}) and not hotswap:

src/diffusers/utils/peft_utils.py

@@ -158,10 +158,12 @@ def get_peft_kwargs(
 
     if LORA_ADAPTER_METADATA_KEY in peft_state_dict:
         metadata = peft_state_dict[LORA_ADAPTER_METADATA_KEY]
-        if metadata:
-            if prefix is not None:
-                metadata = {k.removeprefix(prefix + "."): v for k, v in metadata.items()}
-            return metadata
+    else:
+        metadata = None
+    if metadata:
+        if prefix is not None:
+            metadata = {k.removeprefix(f"{prefix}."): v for k, v in metadata.items() if k.startswith(f"{prefix}.")}
+        return metadata
 
     rank_pattern = {}
     alpha_pattern = {}

tests/lora/test_lora_layers_wan.py

@@ -13,10 +13,8 @@
 # limitations under the License.
 
 import sys
-import tempfile
 import unittest
 
-import numpy as np
 import torch
 from transformers import AutoTokenizer, T5EncoderModel
 
@@ -26,13 +24,7 @@
     WanPipeline,
     WanTransformer3DModel,
 )
-from diffusers.loaders.lora_base import LORA_ADAPTER_METADATA_KEY
-from diffusers.utils.testing_utils import (
-    check_if_dicts_are_equal,
-    floats_tensor,
-    require_peft_backend,
-    torch_device,
-)
+from diffusers.utils.testing_utils import floats_tensor, require_peft_backend
 
 
 sys.path.append(".")
@@ -145,64 +137,3 @@ def test_simple_inference_with_text_lora_fused(self):
     @unittest.skip("Text encoder LoRA is not supported in Wan.")
     def test_simple_inference_with_text_lora_save_load(self):
         pass
-
-    def test_lora_adapter_metadata_is_loaded_correctly(self):
-        # TODO: Will write the test in utils.py eventually.
-        scheduler_cls = self.scheduler_classes[0]
-        components, _, denoiser_lora_config = self.get_dummy_components(scheduler_cls)
-        pipe = self.pipeline_class(**components)
-
-        pipe, _ = self.check_if_adapters_added_correctly(
-            pipe, text_lora_config=None, denoiser_lora_config=denoiser_lora_config
-        )
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
-            lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
-            metadata = denoiser_lora_config.to_dict()
-            self.pipeline_class.save_lora_weights(
-                save_directory=tmpdir,
-                transformer_lora_adapter_metadata=metadata,
-                **lora_state_dicts,
-            )
-            pipe.unload_lora_weights()
-            state_dict = pipe.lora_state_dict(tmpdir)
-
-            self.assertTrue(LORA_ADAPTER_METADATA_KEY in state_dict)
-
-            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)
-
-    def test_lora_adapter_metadata_save_load_inference(self):
-        # Will write the test in utils.py eventually.
-        scheduler_cls = self.scheduler_classes[0]
-        components, _, denoiser_lora_config = self.get_dummy_components(scheduler_cls)
-        pipe = self.pipeline_class(**components).to(torch_device)
-        _, _, inputs = self.get_dummy_inputs(with_generator=False)
-
-        output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
-        self.assertTrue(output_no_lora.shape == self.output_shape)
-
-        pipe, _ = self.check_if_adapters_added_correctly(
-            pipe, text_lora_config=None, denoiser_lora_config=denoiser_lora_config
-        )
-        output_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            modules_to_save = self._get_modules_to_save(pipe, has_denoiser=True)
-            lora_state_dicts = self._get_lora_state_dicts(modules_to_save)
-            metadata = denoiser_lora_config.to_dict()
-            self.pipeline_class.save_lora_weights(
-                save_directory=tmpdir,
-                transformer_lora_adapter_metadata=metadata,
-                **lora_state_dicts,
-            )
-            pipe.unload_lora_weights()
-            pipe.load_lora_weights(tmpdir)
-
-            output_lora_pretrained = pipe(**inputs, generator=torch.manual_seed(0))[0]
-
-            self.assertTrue(
-                np.allclose(output_lora, output_lora_pretrained, atol=1e-3, rtol=1e-3), "Lora outputs should match."
-            )