factor out the overlaps in save_lora_weights().

Author: sayakpaul

src/diffusers/loaders/lora_base.py

@@ -1060,6 +1060,44 @@ def save_function(weights, filename):
         save_function(state_dict, save_path)
         logger.info(f"Model weights saved in {save_path}")
 
+    @classmethod
+    def _save_lora_weights(
+        cls,
+        save_directory: Union[str, os.PathLike],
+        lora_layers: Dict[str, Dict[str, Union[torch.nn.Module, torch.Tensor]]],
+        lora_metadata: Dict[str, Optional[dict]],
+        is_main_process: bool = True,
+        weight_name: str = None,
+        save_function: Callable = None,
+        safe_serialization: bool = True,
+    ):
+        """
+        Helper method to pack and save LoRA weights and metadata. This method centralizes the saving logic for all
+        pipeline types.
+        """
+        state_dict = {}
+        final_lora_adapter_metadata = {}
+
+        # Pack the weights for each component (e.g., 'unet', 'text_encoder')
+        for prefix, layers in lora_layers.items():
+            state_dict.update(cls.pack_weights(layers, prefix))
+
+        # Pack the metadata for each component
+        for prefix, metadata in lora_metadata.items():
+            if metadata:
+                final_lora_adapter_metadata.update(_pack_dict_with_prefix(metadata, prefix))
+
+        # Call the existing writer function
+        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=final_lora_adapter_metadata if final_lora_adapter_metadata else None,
+        )
+
     @classmethod
     def _optionally_disable_offloading(cls, _pipeline):
         return _func_optionally_disable_offloading(_pipeline=_pipeline)

src/diffusers/loaders/lora_pipeline.py

@@ -509,35 +509,29 @@ def save_lora_weights(
             text_encoder_lora_adapter_metadata:
                 LoRA adapter metadata associated with the text encoder to be serialized with the state dict.
         """
-        state_dict = {}
-        lora_adapter_metadata = {}
-
-        if not (unet_lora_layers or text_encoder_lora_layers):
-            raise ValueError("You must pass at least one of `unet_lora_layers` and `text_encoder_lora_layers`.")
+        lora_layers = {}
+        lora_metadata = {}
 
         if unet_lora_layers:
-            state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name))
+            lora_layers[cls.unet_name] = unet_lora_layers
+            lora_metadata[cls.unet_name] = unet_lora_adapter_metadata
 
         if text_encoder_lora_layers:
-            state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name))
+            lora_layers[cls.text_encoder_name] = text_encoder_lora_layers
+            lora_metadata[cls.text_encoder_name] = text_encoder_lora_adapter_metadata
 
-        if unet_lora_adapter_metadata:
-            lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name))
+        if not lora_layers:
+            raise ValueError("You must pass at least one of `unet_lora_layers` or `text_encoder_lora_layers`.")
 
-        if text_encoder_lora_adapter_metadata:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
-            )
-
-        # Save the model
-        cls.write_lora_layers(
-            state_dict=state_dict,
+        # Delegate to the base helper method
+        cls._save_lora_weights(
             save_directory=save_directory,
+            lora_layers=lora_layers,
+            lora_metadata=lora_metadata,
             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(
@@ -1003,44 +997,34 @@ def save_lora_weights(
             text_encoder_2_lora_adapter_metadata:
                 LoRA adapter metadata associated with the second text encoder to be serialized with the state dict.
         """
-        state_dict = {}
-        lora_adapter_metadata = {}
-
-        if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
-            raise ValueError(
-                "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`."
-            )
+        lora_layers = {}
+        lora_metadata = {}
 
         if unet_lora_layers:
-            state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name))
+            lora_layers[cls.unet_name] = unet_lora_layers
+            lora_metadata[cls.unet_name] = unet_lora_adapter_metadata
 
         if text_encoder_lora_layers:
-            state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder"))
+            lora_layers["text_encoder"] = text_encoder_lora_layers
+            lora_metadata["text_encoder"] = text_encoder_lora_adapter_metadata
 
         if text_encoder_2_lora_layers:
-            state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
+            lora_layers["text_encoder_2"] = text_encoder_2_lora_layers
+            lora_metadata["text_encoder_2"] = text_encoder_2_lora_adapter_metadata
 
-        if unet_lora_adapter_metadata is not None:
-            lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name))
-
-        if text_encoder_lora_adapter_metadata:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
-            )
-
-        if text_encoder_2_lora_adapter_metadata:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2")
+        if not lora_layers:
+            raise ValueError(
+                "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers`, or `text_encoder_2_lora_layers`."
             )
 
-        cls.write_lora_layers(
-            state_dict=state_dict,
+        cls._save_lora_weights(
             save_directory=save_directory,
+            lora_layers=lora_layers,
+            lora_metadata=lora_metadata,
             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(
@@ -1466,46 +1450,34 @@ def save_lora_weights(
             text_encoder_2_lora_adapter_metadata:
                 LoRA adapter metadata associated with the second text encoder to be serialized with the state dict.
         """
-        state_dict = {}
-        lora_adapter_metadata = {}
-
-        if not (transformer_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
-            raise ValueError(
-                "You must pass at least one of `transformer_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`."
-            )
+        lora_layers = {}
+        lora_metadata = {}
 
         if transformer_lora_layers:
-            state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
+            lora_layers[cls.transformer_name] = transformer_lora_layers
+            lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata
 
         if text_encoder_lora_layers:
-            state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder"))
+            lora_layers["text_encoder"] = text_encoder_lora_layers
+            lora_metadata["text_encoder"] = text_encoder_lora_adapter_metadata
 
         if text_encoder_2_lora_layers:
-            state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
+            lora_layers["text_encoder_2"] = text_encoder_2_lora_layers
+            lora_metadata["text_encoder_2"] = text_encoder_2_lora_adapter_metadata
 
-        if transformer_lora_adapter_metadata is not None:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
-            )
-
-        if text_encoder_lora_adapter_metadata:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
-            )
-
-        if text_encoder_2_lora_adapter_metadata:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2")
+        if not lora_layers:
+            raise ValueError(
+                "You must pass at least one of `transformer_lora_layers`, `text_encoder_lora_layers`, or `text_encoder_2_lora_layers`."
             )
 
-        cls.write_lora_layers(
-            state_dict=state_dict,
+        cls._save_lora_weights(
             save_directory=save_directory,
+            lora_layers=lora_layers,
+            lora_metadata=lora_metadata,
             is_main_process=is_main_process,
             weight_name=weight_name,
             save_function=save_function,
             safe_serialization=safe_serialization,
-            lora_adapter_metadata=lora_adapter_metadata,
         )
 
     # Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.fuse_lora with unet->transformer
@@ -2434,37 +2406,29 @@ def save_lora_weights(
             text_encoder_lora_adapter_metadata:
                 LoRA adapter metadata associated with the text encoder to be serialized with the state dict.
         """
-        state_dict = {}
-        lora_adapter_metadata = {}
-
-        if not (transformer_lora_layers or text_encoder_lora_layers):
-            raise ValueError("You must pass at least one of `transformer_lora_layers` and `text_encoder_lora_layers`.")
+        lora_layers = {}
+        lora_metadata = {}
 
         if transformer_lora_layers:
-            state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
+            lora_layers[cls.transformer_name] = transformer_lora_layers
+            lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata
 
         if text_encoder_lora_layers:
-            state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name))
-
-        if transformer_lora_adapter_metadata:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
-            )
+            lora_layers[cls.text_encoder_name] = text_encoder_lora_layers
+            lora_metadata[cls.text_encoder_name] = text_encoder_lora_adapter_metadata
 
-        if text_encoder_lora_adapter_metadata:
-            lora_adapter_metadata.update(
-                _pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
-            )
+        if not lora_layers:
+            raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.")
 
-        # Save the model
-        cls.write_lora_layers(
-            state_dict=state_dict,
+        # Delegate to the base helper method
+        cls._save_lora_weights(
             save_directory=save_directory,
+            lora_layers=lora_layers,
+            lora_metadata=lora_metadata,
             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(