Merge pull request #1853 from bghira/bugfix/lora-alpha-metadata

integrate lora_alpha fully with diffusers metadata collection

Author: bghira

simpletuner/helpers/models/auraflow/pipeline.py

@@ -386,6 +386,8 @@ def save_lora_weights(
         weight_name: str = None,
         save_function: Callable = None,
         safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: Optional[dict] = None,
+        controlnet_lora_adapter_metadata: Optional[dict] = None,
     ):
         r"""
         Save the LoRA parameters corresponding to the transformer and optionally controlnet.
@@ -409,6 +411,7 @@ def save_lora_weights(
                 Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
         """
         state_dict = {}
+        lora_adapter_metadata = {}
 
         if not (transformer_lora_layers or controlnet_lora_layers):
             raise ValueError("You must pass at least one of `transformer_lora_layers` or `controlnet_lora_layers`.")
@@ -419,6 +422,12 @@ def save_lora_weights(
         if controlnet_lora_layers:
             state_dict.update(cls.pack_weights(controlnet_lora_layers, cls.controlnet_name))
 
+        if transformer_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(transformer_lora_adapter_metadata, cls.transformer_name))
+
+        if controlnet_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(controlnet_lora_adapter_metadata, cls.controlnet_name))
+
         # Save the model
         cls.write_lora_layers(
             state_dict=state_dict,
@@ -427,6 +436,7 @@ def save_lora_weights(
             weight_name=weight_name,
             save_function=save_function,
             safe_serialization=safe_serialization,
+            lora_adapter_metadata=lora_adapter_metadata,
         )
 
     # Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora

simpletuner/helpers/models/chroma/pipeline.py

@@ -658,6 +658,9 @@ def save_lora_weights(
         weight_name: str = None,
         save_function: Callable = None,
         safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_lora_adapter_metadata: Optional[dict] = None,
+        controlnet_lora_adapter_metadata: Optional[dict] = None,
     ):
         r"""
         Save the LoRA parameters of the UNet, text encoder and optionally the ControlNet.
@@ -684,6 +687,7 @@ def save_lora_weights(
                 Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
         """
         state_dict = {}
+        lora_adapter_metadata = {}
 
         if not (transformer_lora_layers or text_encoder_lora_layers or controlnet_lora_layers):
             raise ValueError(
@@ -700,6 +704,16 @@ def save_lora_weights(
             controlnet_prefix = "controlnet"
             state_dict.update(cls.pack_weights(controlnet_lora_layers, controlnet_prefix))
 
+        if transformer_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(transformer_lora_adapter_metadata, cls.transformer_name))
+
+        if text_encoder_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_lora_adapter_metadata, cls.text_encoder_name))
+
+        if controlnet_lora_adapter_metadata:
+            controlnet_prefix = "controlnet"
+            lora_adapter_metadata.update(cls.pack_weights(controlnet_lora_adapter_metadata, controlnet_prefix))
+
         # Save the model
         cls.write_lora_layers(
             state_dict=state_dict,
@@ -708,6 +722,7 @@ def save_lora_weights(
             weight_name=weight_name,
             save_function=save_function,
             safe_serialization=safe_serialization,
+            lora_adapter_metadata=lora_adapter_metadata,
         )
 
     # Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.fuse_lora with unet->transformer

simpletuner/helpers/models/flux/pipeline.py

@@ -812,6 +812,9 @@ def save_lora_weights(
         weight_name: str = None,
         save_function: Callable = None,
         safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_lora_adapter_metadata: Optional[dict] = None,
+        controlnet_lora_adapter_metadata: Optional[dict] = None,
     ):
         r"""
         Save the LoRA parameters corresponding to the UNet, text encoder, and optionally controlnet.
@@ -838,6 +841,7 @@ def save_lora_weights(
                 Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
         """
         state_dict = {}
+        lora_adapter_metadata = {}
 
         if not (transformer_lora_layers or text_encoder_lora_layers or controlnet_lora_layers):
             raise ValueError(
@@ -854,6 +858,16 @@ def save_lora_weights(
             controlnet_prefix = "controlnet"
             state_dict.update(cls.pack_weights(controlnet_lora_layers, controlnet_prefix))
 
+        if transformer_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(transformer_lora_adapter_metadata, cls.transformer_name))
+
+        if text_encoder_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_lora_adapter_metadata, cls.text_encoder_name))
+
+        if controlnet_lora_adapter_metadata:
+            controlnet_prefix = "controlnet"
+            lora_adapter_metadata.update(cls.pack_weights(controlnet_lora_adapter_metadata, controlnet_prefix))
+
         # Save the model
         cls.write_lora_layers(
             state_dict=state_dict,
@@ -862,6 +876,7 @@ def save_lora_weights(
             weight_name=weight_name,
             save_function=save_function,
             safe_serialization=safe_serialization,
+            lora_adapter_metadata=lora_adapter_metadata,
         )
 
     # Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.fuse_lora with unet->transformer

simpletuner/helpers/models/hidream/pipeline.py

@@ -680,6 +680,11 @@ def save_lora_weights(
         save_function: Callable = None,
         safe_serialization: bool = True,
         transformer_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_2_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_3_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_4_lora_adapter_metadata: Optional[dict] = None,
+        controlnet_lora_adapter_metadata: Optional[dict] = None,
     ):
         r"""
         Save the LoRA parameters corresponding to the transformer, text encoders, and optionally controlnet.
@@ -711,6 +716,16 @@ def save_lora_weights(
                 Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
             transformer_lora_adapter_metadata:
                 LoRA adapter metadata associated with the transformer to be serialized with the state dict.
+            text_encoder_lora_adapter_metadata:
+                LoRA adapter metadata associated with the text encoder to be serialized with the state dict.
+            text_encoder_2_lora_adapter_metadata:
+                LoRA adapter metadata associated with the second text encoder to be serialized with the state dict.
+            text_encoder_3_lora_adapter_metadata:
+                LoRA adapter metadata associated with the third text encoder to be serialized with the state dict.
+            text_encoder_4_lora_adapter_metadata:
+                LoRA adapter metadata associated with the fourth text encoder to be serialized with the state dict.
+            controlnet_lora_adapter_metadata:
+                LoRA adapter metadata associated with the controlnet to be serialized with the state dict.
         """
         state_dict = {}
         lora_adapter_metadata = {}
@@ -746,6 +761,21 @@ def save_lora_weights(
         if transformer_lora_adapter_metadata is not None:
             lora_adapter_metadata.update(cls.pack_weights(transformer_lora_adapter_metadata, cls.transformer_name))
 
+        if text_encoder_lora_adapter_metadata is not None:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_lora_adapter_metadata, "text_encoder"))
+
+        if text_encoder_2_lora_adapter_metadata is not None:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_2_lora_adapter_metadata, "text_encoder_2"))
+
+        if text_encoder_3_lora_adapter_metadata is not None:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_3_lora_adapter_metadata, "text_encoder_3"))
+
+        if text_encoder_4_lora_adapter_metadata is not None:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_4_lora_adapter_metadata, "text_encoder_4"))
+
+        if controlnet_lora_adapter_metadata is not None:
+            lora_adapter_metadata.update(cls.pack_weights(controlnet_lora_adapter_metadata, cls.controlnet_name))
+
         # Save the model
         cls.write_lora_layers(
             state_dict=state_dict,

simpletuner/helpers/models/pixart/model.py

@@ -393,6 +393,8 @@ def save_lora_weights(self, output_dir: str, **kwargs):
             save_directory=output_dir,
             transformer_lora_layers=None,  # No transformer LoRA
             controlnet_lora_layers=controlnet_lora_layers,
+            transformer_lora_adapter_metadata=kwargs.get("transformer_lora_adapter_metadata"),
+            controlnet_lora_adapter_metadata=kwargs.get("controlnet_lora_adapter_metadata"),
         )
 
     def custom_model_card_schedule_info(self):

simpletuner/helpers/models/pixart/pipeline.py

@@ -207,9 +207,12 @@ def save_lora_weights(
         weight_name: str = None,
         save_function: Callable = None,
         safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: Optional[dict] = None,
+        controlnet_lora_adapter_metadata: Optional[dict] = None,
     ):
         """Save LoRA weights for both transformer and controlnet."""
         state_dict = {}
+        lora_adapter_metadata = {}
 
         # Pack transformer weights (only the non-replaced blocks)
         if transformer_lora_layers:
@@ -220,6 +223,12 @@ def save_lora_weights(
         if controlnet_lora_layers:
             state_dict.update(controlnet_lora_layers)  # they're already packed
 
+        if transformer_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(transformer_lora_adapter_metadata, cls.transformer_name))
+
+        if controlnet_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(controlnet_lora_adapter_metadata, cls.controlnet_name))
+
         # Save the model
         cls.write_lora_layers(
             state_dict=state_dict,
@@ -228,6 +237,7 @@ def save_lora_weights(
             weight_name=weight_name,
             save_function=save_function,
             safe_serialization=safe_serialization,
+            lora_adapter_metadata=lora_adapter_metadata,
         )
 
     def load_lora_weights(

simpletuner/helpers/models/sd3/pipeline.py

@@ -779,6 +779,10 @@ def save_lora_weights(
         weight_name: str = None,
         save_function: Callable = None,
         safe_serialization: bool = True,
+        transformer_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_2_lora_adapter_metadata: Optional[dict] = None,
+        controlnet_lora_adapter_metadata: Optional[dict] = None,
     ):
         r"""
         Save the LoRA parameters corresponding to the transformer, text encoders, and optionally controlnet.
@@ -808,6 +812,7 @@ def save_lora_weights(
                 Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
         """
         state_dict = {}
+        lora_adapter_metadata = {}
 
         if not (transformer_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers or controlnet_lora_layers):
             raise ValueError(
@@ -826,6 +831,18 @@ def save_lora_weights(
         if controlnet_lora_layers:
             state_dict.update(cls.pack_weights(controlnet_lora_layers, cls.controlnet_name))
 
+        if transformer_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(transformer_lora_adapter_metadata, cls.transformer_name))
+
+        if text_encoder_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_lora_adapter_metadata, "text_encoder"))
+
+        if text_encoder_2_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_2_lora_adapter_metadata, "text_encoder_2"))
+
+        if controlnet_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(controlnet_lora_adapter_metadata, cls.controlnet_name))
+
         # Save the model
         cls.write_lora_layers(
             state_dict=state_dict,
@@ -834,6 +851,7 @@ def save_lora_weights(
             weight_name=weight_name,
             save_function=save_function,
             safe_serialization=safe_serialization,
+            lora_adapter_metadata=lora_adapter_metadata,
         )
 
     def fuse_lora(

simpletuner/helpers/models/sdxl/pipeline.py

@@ -371,6 +371,10 @@ def save_lora_weights(
         weight_name: str = None,
         save_function: Callable = None,
         safe_serialization: bool = True,
+        unet_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_lora_adapter_metadata: Optional[dict] = None,
+        text_encoder_2_lora_adapter_metadata: Optional[dict] = None,
+        controlnet_lora_adapter_metadata: Optional[dict] = None,
     ):
         r"""
         Save the LoRA parameters corresponding to the UNet, text encoders, and optionally controlnet.
@@ -394,6 +398,7 @@ def save_lora_weights(
                 Whether to save the model using `safetensors` or the traditional PyTorch way.
         """
         state_dict = {}
+        lora_adapter_metadata = {}
 
         if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers or controlnet_lora_layers):
             raise ValueError(
@@ -413,13 +418,26 @@ def save_lora_weights(
         if controlnet_lora_layers:
             state_dict.update(cls.pack_weights(controlnet_lora_layers, cls.controlnet_name))
 
+        if unet_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(unet_lora_adapter_metadata, "unet"))
+
+        if text_encoder_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_lora_adapter_metadata, "text_encoder"))
+
+        if text_encoder_2_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(text_encoder_2_lora_adapter_metadata, "text_encoder_2"))
+
+        if controlnet_lora_adapter_metadata:
+            lora_adapter_metadata.update(cls.pack_weights(controlnet_lora_adapter_metadata, cls.controlnet_name))
+
         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(