update the test script.

Author: ParagEkbote

examples/dreambooth/train_dreambooth_lora_hidream.py

@@ -54,6 +54,7 @@
 )
 from diffusers.optimization import get_scheduler
 from diffusers.training_utils import (
+    _collate_lora_metadata,
     cast_training_params,
     compute_density_for_timestep_sampling,
     compute_loss_weighting_for_sd3,
@@ -420,6 +421,13 @@ def parse_args(input_args=None):
 
     parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers")
 
+    parser.add_argument(
+        "--lora_alpha",
+        type=int,
+        default=4,
+        help="LoRA alpha to be used for additional scaling.",
+    )
+
     parser.add_argument(
         "--with_prior_preservation",
         default=False,
@@ -1163,7 +1171,7 @@ def main(args):
     # now we will add new LoRA weights the transformer layers
     transformer_lora_config = LoraConfig(
         r=args.rank,
-        lora_alpha=args.rank,
+        lora_alpha=args.lora_alpha,
         lora_dropout=args.lora_dropout,
         init_lora_weights="gaussian",
         target_modules=target_modules,
@@ -1180,10 +1188,12 @@ def save_model_hook(models, weights, output_dir):
         if accelerator.is_main_process:
             transformer_lora_layers_to_save = None
 
+            modules_to_save = {}
             for model in models:
                 if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
                     model = unwrap_model(model)
                     transformer_lora_layers_to_save = get_peft_model_state_dict(model)
+                    modules_to_save["transformer"] = model
                 else:
                     raise ValueError(f"unexpected save model: {model.__class__}")
 
@@ -1194,6 +1204,7 @@ def save_model_hook(models, weights, output_dir):
             HiDreamImagePipeline.save_lora_weights(
                 output_dir,
                 transformer_lora_layers=transformer_lora_layers_to_save,
+                **_collate_lora_metadata(modules_to_save),
             )
 
     def load_model_hook(models, input_dir):
@@ -1496,6 +1507,7 @@ def compute_text_embeddings(prompt, text_encoding_pipeline):
     # We need to initialize the trackers we use, and also store our configuration.
     # The trackers initializes automatically on the main process.
     if accelerator.is_main_process:
+        modules_to_save = {}
         tracker_name = "dreambooth-hidream-lora"
         accelerator.init_trackers(tracker_name, config=vars(args))
 
@@ -1737,6 +1749,7 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
             else:
                 transformer = transformer.to(weight_dtype)
         transformer_lora_layers = get_peft_model_state_dict(transformer)
+        modules_to_save["transformer"] = transformer
 
         HiDreamImagePipeline.save_lora_weights(
             save_directory=args.output_dir,