Merge pull request #408 from XDUWQ/lora_rank

Add lora_rank parameter for lora stable diffusion

examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py

@@ -19,6 +19,12 @@ class StableDiffusionLoraArguments(TrainingArgs):
             'help': 'The pipeline prompt.',
         })
 
+    lora_rank: int = field(
+        default=4,
+        metadata={
+            'help': 'The rank size of lora intermediate linear.',
+        })
+
 
 training_args = StableDiffusionLoraArguments(
     task='text-to-image-synthesis').parse_cli()
@@ -59,6 +65,7 @@ kwargs = dict(
     work_dir=training_args.work_dir,
     train_dataset=train_dataset,
     eval_dataset=validation_dataset,
+    lora_rank=args.lora_rank,
     cfg_modify_fn=cfg_modify_fn)
 
 # build trainer and training

examples/pytorch/stable_diffusion/lora/run_train_lora.sh

@@ -5,6 +5,7 @@ PYTHONPATH=. torchrun examples/pytorch/stable_diffusion/lora/finetune_stable_dif
     --work_dir './tmp/lora_diffusion' \
     --train_dataset_name 'buptwq/lora-stable-diffusion-finetune' \
     --max_epochs 100 \
+    --lora_rank 4 \
     --save_ckpt_strategy 'by_epoch' \
     --logging_interval 1 \
     --train.dataloader.workers_per_gpu 0 \

modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py

@@ -34,6 +34,14 @@ class LoraDiffusionTrainer(EpochBasedTrainer):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        """Lora trainers for fine-tuning stable diffusion
+
+        Args:
+            lora_rank: The rank size of lora intermediate linear.
+
+        """
+        lora_rank = kwargs.pop('lora_rank', 4)
+
         # set lora save checkpoint processor
         ckpt_hook = list(
             filter(lambda hook: isinstance(hook, CheckpointHook),
@@ -59,7 +67,8 @@ class LoraDiffusionTrainer(EpochBasedTrainer):
 
             lora_attn_procs[name] = LoRAAttnProcessor(
                 hidden_size=hidden_size,
-                cross_attention_dim=cross_attention_dim)
+                cross_attention_dim=cross_attention_dim,
+                rank=lora_rank)
 
         self.model.unet.set_attn_processor(lora_attn_procs)