diff --git a/examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py b/examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py index 8ad3a28b..6001af48 100644 --- a/examples/pytorch/stable_diffusion/lora/finetune_stable_diffusion_lora.py +++ b/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 diff --git a/examples/pytorch/stable_diffusion/lora/run_train_lora.sh b/examples/pytorch/stable_diffusion/lora/run_train_lora.sh index 876a2475..bf62f833 100644 --- a/examples/pytorch/stable_diffusion/lora/run_train_lora.sh +++ b/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 \ diff --git a/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py b/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py index 99351fef..7c6644bd 100644 --- a/modelscope/trainers/multi_modal/lora_diffusion/lora_diffusion_trainer.py +++ b/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)