Skip to content

Commit dc1c6d2

Browse files
cloneofsimobrian6091DavidePaglierihdon96
authored
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 <brian6091@gmail.com> * 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 <brian6091@gmail.com> Co-authored-by: SimoRyu <cloneofsimo@korea.ac.kr> * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf3. * 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 <brian6091@gmail.com> * 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 <brian6091@gmail.com> Co-authored-by: SimoRyu <cloneofsimo@korea.ac.kr> * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf3. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * 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 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * 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 <brian6091@gmail.com> * 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 <brian6091@gmail.com> Co-authored-by: SimoRyu <cloneofsimo@korea.ac.kr> * Revert "Fix lora inject, added weight self apply lora (#39)" (#40) This reverts commit fececf3. * fix : rank bug in monkeypatch * fix cli fix * visualizatio on effect of LR Co-authored-by: brian6091 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * 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 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> * 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 <brian6091@gmail.com> Co-authored-by: Davide Paglieri <paglieridavide@gmail.com> Co-authored-by: hdeezy <82070413+hdeezy@users.noreply.github.com>
1 parent b6b9986 commit dc1c6d2

2 files changed

Lines changed: 78 additions & 74 deletions

File tree

‎train_lora_dreambooth.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -922,7 +922,7 @@ def collate_fn(examples):
922922
if global_step >= args.max_train_steps:
923923
break
924924

925-
accelerator.wait_for_everyone()
925+
accelerator.wait_for_everyone()
926926

927927
# Create the pipeline using using the trained modules and save it.
928928
if accelerator.is_main_process:

‎train_lora_w_ti.py‎

Lines changed: 77 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,6 @@ def __getitem__(self, index):
178178

179179
placeholder_string = self.placeholder_token
180180
text = random.choice(self.templates).format(placeholder_string)
181-
182181
example["instance_prompt_ids"] = self.tokenizer(
183182
text,
184183
padding="do_not_pad",
@@ -899,10 +898,15 @@ def collate_fn(examples):
899898

900899
# optimizer = accelerator.prepare(optimizer)
901900
for step, batch in enumerate(train_dataloader):
902-
901+
902+
# freeze unet and text encoder during ti training
903903
if global_step < args.unfreeze_lora_step:
904904
optimizer.param_groups[0]["lr"] = 0.0
905905
optimizer.param_groups[1]["lr"] = 0.0
906+
else: # begin learning with unet and text encoder
907+
optimizer.param_groups[0]["lr"] = args.learning_rate
908+
optimizer.param_groups[1]["lr"] = args.learning_rate_text
909+
optimizer.param_groups[2]["lr"] = 0.0 # stop learning ti
906910

907911
# Convert images to latent space
908912
latents = vae.encode(
@@ -987,86 +991,86 @@ def collate_fn(examples):
987991

988992
global_step += 1
989993

990-
# Checks if the accelerator has performed an optimization step behind the scenes
991-
if accelerator.sync_gradients:
992-
if args.save_steps and global_step - last_save >= args.save_steps:
993-
if accelerator.is_main_process:
994-
# newer versions of accelerate allow the 'keep_fp32_wrapper' arg. without passing
995-
# it, the models will be unwrapped, and when they are then used for further training,
996-
# we will crash. pass this, but only to newer versions of accelerate. fixes
997-
# https://github.com/huggingface/diffusers/issues/1566
998-
accepts_keep_fp32_wrapper = "keep_fp32_wrapper" in set(
999-
inspect.signature(accelerator.unwrap_model).parameters.keys()
1000-
)
1001-
extra_args = (
1002-
{"keep_fp32_wrapper": True} if accepts_keep_fp32_wrapper else {}
1003-
)
1004-
pipeline = StableDiffusionPipeline.from_pretrained(
1005-
args.pretrained_model_name_or_path,
1006-
unet=accelerator.unwrap_model(unet, **extra_args),
1007-
text_encoder=accelerator.unwrap_model(
1008-
text_encoder, **extra_args
1009-
),
1010-
revision=args.revision,
1011-
)
1012-
1013-
filename_unet = (
1014-
f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.pt"
1015-
)
1016-
filename_text_encoder = f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.text_encoder.pt"
1017-
print(f"save weights {filename_unet}, {filename_text_encoder}")
1018-
save_lora_weight(pipeline.unet, filename_unet)
1019-
1020-
save_lora_weight(
1021-
pipeline.text_encoder,
1022-
filename_text_encoder,
1023-
target_replace_module=["CLIPAttention"],
1024-
)
1025-
1026-
for _up, _down in extract_lora_ups_down(pipeline.unet):
1027-
print("First Unet Layer's Up Weight is now : ", _up.weight.data)
1028-
print(
1029-
"First Unet Layer's Down Weight is now : ",
1030-
_down.weight.data,
994+
# Checks if the accelerator has performed an optimization step behind the scenes
995+
if accelerator.sync_gradients:
996+
if args.save_steps and global_step - last_save >= args.save_steps:
997+
if accelerator.is_main_process:
998+
# newer versions of accelerate allow the 'keep_fp32_wrapper' arg. without passing
999+
# it, the models will be unwrapped, and when they are then used for further training,
1000+
# we will crash. pass this, but only to newer versions of accelerate. fixes
1001+
# https://github.com/huggingface/diffusers/issues/1566
1002+
accepts_keep_fp32_wrapper = "keep_fp32_wrapper" in set(
1003+
inspect.signature(accelerator.unwrap_model).parameters.keys()
10311004
)
1032-
break
1033-
1034-
for _up, _down in extract_lora_ups_down(
1035-
pipeline.text_encoder,
1036-
target_replace_module=["CLIPAttention"],
1037-
):
1038-
print(
1039-
"First Text Encoder Layer's Up Weight is now : ",
1040-
_up.weight.data,
1005+
extra_args = (
1006+
{"keep_fp32_wrapper": True} if accepts_keep_fp32_wrapper else {}
10411007
)
1042-
print(
1043-
"First Text Encoder Layer's Down Weight is now : ",
1044-
_down.weight.data,
1008+
pipeline = StableDiffusionPipeline.from_pretrained(
1009+
args.pretrained_model_name_or_path,
1010+
unet=accelerator.unwrap_model(unet, **extra_args),
1011+
text_encoder=accelerator.unwrap_model(
1012+
text_encoder, **extra_args
1013+
),
1014+
revision=args.revision,
10451015
)
1046-
break
10471016

1048-
filename_ti = (
1049-
f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.ti.pt"
1050-
)
1017+
filename_unet = (
1018+
f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.pt"
1019+
)
1020+
filename_text_encoder = f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.text_encoder.pt"
1021+
print(f"save weights {filename_unet}, {filename_text_encoder}")
1022+
save_lora_weight(pipeline.unet, filename_unet)
1023+
1024+
save_lora_weight(
1025+
pipeline.text_encoder,
1026+
filename_text_encoder,
1027+
target_replace_module=["CLIPAttention"],
1028+
)
10511029

1052-
save_progress(
1053-
pipeline.text_encoder,
1054-
placeholder_token_id,
1055-
accelerator,
1056-
args,
1057-
filename_ti,
1058-
)
1030+
for _up, _down in extract_lora_ups_down(pipeline.unet):
1031+
print("First Unet Layer's Up Weight is now : ", _up.weight.data)
1032+
print(
1033+
"First Unet Layer's Down Weight is now : ",
1034+
_down.weight.data,
1035+
)
1036+
break
1037+
1038+
for _up, _down in extract_lora_ups_down(
1039+
pipeline.text_encoder,
1040+
target_replace_module=["CLIPAttention"],
1041+
):
1042+
print(
1043+
"First Text Encoder Layer's Up Weight is now : ",
1044+
_up.weight.data,
1045+
)
1046+
print(
1047+
"First Text Encoder Layer's Down Weight is now : ",
1048+
_down.weight.data,
1049+
)
1050+
break
1051+
1052+
filename_ti = (
1053+
f"{args.output_dir}/lora_weight_e{epoch}_s{global_step}.ti.pt"
1054+
)
1055+
1056+
save_progress(
1057+
pipeline.text_encoder,
1058+
placeholder_token_id,
1059+
accelerator,
1060+
args,
1061+
filename_ti,
1062+
)
10591063

1060-
last_save = global_step
1064+
last_save = global_step
10611065

1062-
logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
1063-
progress_bar.set_postfix(**logs)
1064-
accelerator.log(logs, step=global_step)
1066+
logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
1067+
progress_bar.set_postfix(**logs)
1068+
accelerator.log(logs, step=global_step)
10651069

1066-
if global_step >= args.max_train_steps:
1067-
break
1070+
if global_step >= args.max_train_steps:
1071+
break
10681072

1069-
accelerator.wait_for_everyone()
1073+
accelerator.wait_for_everyone()
10701074

10711075
# Create the pipeline using using the trained modules and save it.
10721076
if accelerator.is_main_process:

0 commit comments

Comments
 (0)