From dc1c6d25ae3cf53f8a6a179b3c8d34397fb46d36 Mon Sep 17 00:00:00 2001 From: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Date: Mon, 19 Dec 2022 22:00:53 +0900 Subject: [PATCH] Develop (#66) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme * Fix lora inject, added weight self apply lora (#39) * Develop (#34) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme Co-authored-by: brian6091 * release : version 0.0.4, now able to tune rank, now add loras dynamically * readme : add brain6091's discussions * fix:inject lora in to_out module list * feat: added weight self apply lora * chore: add import copy * fix: readded r Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 Co-authored-by: SimoRyu * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf303f216a54a9c3f5b0cf0ac46b44b9b913. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR * Fix save_steps, max_train_steps, and logging (#45) * v 0.0.5 (#42) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme * Fix lora inject, added weight self apply lora (#39) * Develop (#34) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme Co-authored-by: brian6091 * release : version 0.0.4, now able to tune rank, now add loras dynamically * readme : add brain6091's discussions * fix:inject lora in to_out module list * feat: added weight self apply lora * chore: add import copy * fix: readded r Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 Co-authored-by: SimoRyu * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf303f216a54a9c3f5b0cf0ac46b44b9b913. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR Co-authored-by: brian6091 Co-authored-by: Davide Paglieri * Fix save_steps, max_train_steps, and logging Corrected indenting so checking save_steps, max_train_steps, and updating logs are performed every step instead at the end of an epoch. Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 Co-authored-by: Davide Paglieri * Enable resuming (#52) * v 0.0.5 (#42) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme * Fix lora inject, added weight self apply lora (#39) * Develop (#34) * Add parameter to control rank of decomposition (#28) * ENH: allow controlling rank of approximation * Training script accepts lora_rank * feat : statefully monkeypatch different loras + example ipynb + readme Co-authored-by: brian6091 * release : version 0.0.4, now able to tune rank, now add loras dynamically * readme : add brain6091's discussions * fix:inject lora in to_out module list * feat: added weight self apply lora * chore: add import copy * fix: readded r Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 Co-authored-by: SimoRyu * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf303f216a54a9c3f5b0cf0ac46b44b9b913. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR Co-authored-by: brian6091 Co-authored-by: Davide Paglieri * Enable resume training unet/text encoder (#48) * Enable resume training unet/text encoder New flags --resume_text_encoder --resume_unet accept the paths to .pt files to resume. Make sure to change the output directory from the previous training session, or else .pt files will be overwritten since training does not resume from previous global step. * Load weights from .pt with inject_trainable_lora Adds new loras argument to inject_trainable_lora function which accepts path to a .pt file containing previously trained weights. Co-authored-by: Simo Ryu <35953539+cloneofsimo@users.noreply.github.com> Co-authored-by: brian6091 Co-authored-by: Davide Paglieri * feat : low-rank pivotal tuning * feat : pivotal tuning * v 0.0.6 * Learning rate switching & fix indent (#57) * Learning rate switching & fix indent Make learning rates switch from training textual inversion to unet/text encoder after unfreeze_lora_step. I think this is how it was explained in the paper linked(?) Either way, it might be useful to add another parameter to activate unet/text encoder training at a certain step instead of at unfreeze_lora_step. This would let the user have more control. Also fix indenting to make save_steps and logging work properly. * Fix indent fix accelerator.wait_for_everyone() indent according to original dreambooth training * Re:Fix indent (#58) Fix indenting of accelerator.wait_for_everyone() according to original dreambooth training Co-authored-by: brian6091 Co-authored-by: Davide Paglieri Co-authored-by: hdeezy <82070413+hdeezy@users.noreply.github.com> --- train_lora_dreambooth.py | 2 +- train_lora_w_ti.py | 150 ++++++++++++++++++++------------------- 2 files changed, 78 insertions(+), 74 deletions(-) diff --git a/train_lora_dreambooth.py b/train_lora_dreambooth.py index 43ec4d5..3ed2253 100644 --- a/train_lora_dreambooth.py +++ b/train_lora_dreambooth.py @@ -922,7 +922,7 @@ def collate_fn(examples): if global_step >= args.max_train_steps: break - accelerator.wait_for_everyone() + accelerator.wait_for_everyone() # Create the pipeline using using the trained modules and save it. if accelerator.is_main_process: diff --git a/train_lora_w_ti.py b/train_lora_w_ti.py index a97f985..372c2a6 100644 --- a/train_lora_w_ti.py +++ b/train_lora_w_ti.py @@ -178,7 +178,6 @@ def __getitem__(self, index): placeholder_string = self.placeholder_token text = random.choice(self.templates).format(placeholder_string) - example["instance_prompt_ids"] = self.tokenizer( text, padding="do_not_pad", @@ -899,10 +898,15 @@ def collate_fn(examples): # optimizer = accelerator.prepare(optimizer) for step, batch in enumerate(train_dataloader): - + + # freeze unet and text encoder during ti training if global_step < args.unfreeze_lora_step: optimizer.param_groups[0]["lr"] = 0.0 optimizer.param_groups[1]["lr"] = 0.0 + else: # begin learning with unet and text encoder + optimizer.param_groups[0]["lr"] = args.learning_rate + optimizer.param_groups[1]["lr"] = args.learning_rate_text + optimizer.param_groups[2]["lr"] = 0.0 # stop learning ti # Convert images to latent space latents = vae.encode( @@ -987,86 +991,86 @@ def collate_fn(examples): global_step += 1 - # Checks if the accelerator has performed an optimization step behind the scenes - if accelerator.sync_gradients: - if args.save_steps and global_step - last_save >= args.save_steps: - if accelerator.is_main_process: - # newer versions of accelerate allow the 'keep_fp32_wrapper' arg. without passing - # it, the models will be unwrapped, and when they are then used for further training, - # we will crash. pass this, but only to newer versions of accelerate. fixes - # https://github.com/huggingface/diffusers/issues/1566 - accepts_keep_fp32_wrapper = "keep_fp32_wrapper" in set( - inspect.signature(accelerator.unwrap_model).parameters.keys() - ) - extra_args = ( - {"keep_fp32_wrapper": True} if accepts_keep_fp32_wrapper else {} - ) - pipeline = StableDiffusionPipeline.from_pretrained( - args.pretrained_model_name_or_path, - unet=accelerator.unwrap_model(unet, **extra_args), - text_encoder=accelerator.unwrap_model( - text_encoder, **extra_args - ), - revision=args.revision, - ) - - filename_unet = ( - f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.pt" - ) - filename_text_encoder = f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.text_encoder.pt" - print(f"save weights {filename_unet}, {filename_text_encoder}") - save_lora_weight(pipeline.unet, filename_unet) - - save_lora_weight( - pipeline.text_encoder, - filename_text_encoder, - target_replace_module=["CLIPAttention"], - ) - - for _up, _down in extract_lora_ups_down(pipeline.unet): - print("First Unet Layer's Up Weight is now : ", _up.weight.data) - print( - "First Unet Layer's Down Weight is now : ", - _down.weight.data, + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + if args.save_steps and global_step - last_save >= args.save_steps: + if accelerator.is_main_process: + # newer versions of accelerate allow the 'keep_fp32_wrapper' arg. without passing + # it, the models will be unwrapped, and when they are then used for further training, + # we will crash. pass this, but only to newer versions of accelerate. fixes + # https://github.com/huggingface/diffusers/issues/1566 + accepts_keep_fp32_wrapper = "keep_fp32_wrapper" in set( + inspect.signature(accelerator.unwrap_model).parameters.keys() ) - break - - for _up, _down in extract_lora_ups_down( - pipeline.text_encoder, - target_replace_module=["CLIPAttention"], - ): - print( - "First Text Encoder Layer's Up Weight is now : ", - _up.weight.data, + extra_args = ( + {"keep_fp32_wrapper": True} if accepts_keep_fp32_wrapper else {} ) - print( - "First Text Encoder Layer's Down Weight is now : ", - _down.weight.data, + pipeline = StableDiffusionPipeline.from_pretrained( + args.pretrained_model_name_or_path, + unet=accelerator.unwrap_model(unet, **extra_args), + text_encoder=accelerator.unwrap_model( + text_encoder, **extra_args + ), + revision=args.revision, ) - break - filename_ti = ( - f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.ti.pt" - ) + filename_unet = ( + f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.pt" + ) + filename_text_encoder = f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.text_encoder.pt" + print(f"save weights {filename_unet}, {filename_text_encoder}") + save_lora_weight(pipeline.unet, filename_unet) + + save_lora_weight( + pipeline.text_encoder, + filename_text_encoder, + target_replace_module=["CLIPAttention"], + ) - save_progress( - pipeline.text_encoder, - placeholder_token_id, - accelerator, - args, - filename_ti, - ) + for _up, _down in extract_lora_ups_down(pipeline.unet): + print("First Unet Layer's Up Weight is now : ", _up.weight.data) + print( + "First Unet Layer's Down Weight is now : ", + _down.weight.data, + ) + break + + for _up, _down in extract_lora_ups_down( + pipeline.text_encoder, + target_replace_module=["CLIPAttention"], + ): + print( + "First Text Encoder Layer's Up Weight is now : ", + _up.weight.data, + ) + print( + "First Text Encoder Layer's Down Weight is now : ", + _down.weight.data, + ) + break + + filename_ti = ( + f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.ti.pt" + ) + + save_progress( + pipeline.text_encoder, + placeholder_token_id, + accelerator, + args, + filename_ti, + ) - last_save = global_step + last_save = global_step - logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} - progress_bar.set_postfix(**logs) - accelerator.log(logs, step=global_step) + logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} + progress_bar.set_postfix(**logs) + accelerator.log(logs, step=global_step) - if global_step >= args.max_train_steps: - break + if global_step >= args.max_train_steps: + break - accelerator.wait_for_everyone() + accelerator.wait_for_everyone() # Create the pipeline using using the trained modules and save it. if accelerator.is_main_process: