From e2fa053f486b43a04919840cca3bdb72f777511c Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 27 Dec 2023 08:03:42 +0530 Subject: [PATCH 1/7] add to dreambooth lora. --- examples/dreambooth/train_dreambooth_lora.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py index 67132d6d88df..3219385b1f8f 100644 --- a/examples/dreambooth/train_dreambooth_lora.py +++ b/examples/dreambooth/train_dreambooth_lora.py @@ -54,7 +54,7 @@ ) from diffusers.loaders import LoraLoaderMixin from diffusers.optimization import get_scheduler -from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils import check_min_version, is_wandb_available, convert_state_dict_to_diffusers from diffusers.utils.import_utils import is_xformers_available @@ -853,9 +853,9 @@ def save_model_hook(models, weights, output_dir): for model in models: if isinstance(model, type(accelerator.unwrap_model(unet))): - unet_lora_layers_to_save = get_peft_model_state_dict(model) + unet_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) elif isinstance(model, type(accelerator.unwrap_model(text_encoder))): - text_encoder_lora_layers_to_save = get_peft_model_state_dict(model) + text_encoder_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) else: raise ValueError(f"unexpected save model: {model.__class__}") @@ -1285,11 +1285,11 @@ def compute_text_embeddings(prompt): unet = accelerator.unwrap_model(unet) unet = unet.to(torch.float32) - unet_lora_state_dict = get_peft_model_state_dict(unet) + unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) if args.train_text_encoder: text_encoder = accelerator.unwrap_model(text_encoder) - text_encoder_state_dict = get_peft_model_state_dict(text_encoder) + text_encoder_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(text_encoder)) else: text_encoder_state_dict = None From 56ff55736b6aabb972fc1364b932db8e168185ee Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 27 Dec 2023 08:04:28 +0530 Subject: [PATCH 2/7] add: t2i lora. --- examples/text_to_image/train_text_to_image_lora.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/text_to_image/train_text_to_image_lora.py b/examples/text_to_image/train_text_to_image_lora.py index 2efbaf298d2e..55fc377b141c 100644 --- a/examples/text_to_image/train_text_to_image_lora.py +++ b/examples/text_to_image/train_text_to_image_lora.py @@ -44,7 +44,7 @@ from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, StableDiffusionPipeline, UNet2DConditionModel from diffusers.optimization import get_scheduler from diffusers.training_utils import compute_snr -from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils import check_min_version, is_wandb_available, convert_state_dict_to_diffusers from diffusers.utils.import_utils import is_xformers_available @@ -809,7 +809,7 @@ def collate_fn(examples): accelerator.save_state(save_path) unwrapped_unet = accelerator.unwrap_model(unet) - unet_lora_state_dict = get_peft_model_state_dict(unwrapped_unet) + unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unwrapped_unet)) StableDiffusionPipeline.save_lora_weights( save_directory=save_path, @@ -876,7 +876,7 @@ def collate_fn(examples): unet = unet.to(torch.float32) unwrapped_unet = accelerator.unwrap_model(unet) - unet_lora_state_dict = get_peft_model_state_dict(unwrapped_unet) + unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unwrapped_unet)) StableDiffusionPipeline.save_lora_weights( save_directory=args.output_dir, unet_lora_layers=unet_lora_state_dict, From b79556293358751dbee26f140333516b7db24252 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 27 Dec 2023 08:05:59 +0530 Subject: [PATCH 3/7] add: sdxl t2i lora. --- .../text_to_image/train_text_to_image_lora_sdxl.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py index d95fcbbba033..5d4478cc0bc6 100644 --- a/examples/text_to_image/train_text_to_image_lora_sdxl.py +++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py @@ -52,7 +52,7 @@ from diffusers.loaders import LoraLoaderMixin from diffusers.optimization import get_scheduler from diffusers.training_utils import compute_snr -from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils import check_min_version, is_wandb_available, convert_state_dict_to_diffusers from diffusers.utils.import_utils import is_xformers_available @@ -651,11 +651,11 @@ def save_model_hook(models, weights, output_dir): for model in models: if isinstance(model, type(accelerator.unwrap_model(unet))): - unet_lora_layers_to_save = get_peft_model_state_dict(model) + unet_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) elif isinstance(model, type(accelerator.unwrap_model(text_encoder_one))): - text_encoder_one_lora_layers_to_save = get_peft_model_state_dict(model) + text_encoder_one_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) elif isinstance(model, type(accelerator.unwrap_model(text_encoder_two))): - text_encoder_two_lora_layers_to_save = get_peft_model_state_dict(model) + text_encoder_two_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) else: raise ValueError(f"unexpected save model: {model.__class__}") @@ -1160,14 +1160,14 @@ def compute_time_ids(original_size, crops_coords_top_left): accelerator.wait_for_everyone() if accelerator.is_main_process: unet = accelerator.unwrap_model(unet) - unet_lora_state_dict = get_peft_model_state_dict(unet) + unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) if args.train_text_encoder: text_encoder_one = accelerator.unwrap_model(text_encoder_one) text_encoder_two = accelerator.unwrap_model(text_encoder_two) - text_encoder_lora_layers = get_peft_model_state_dict(text_encoder_one) - text_encoder_2_lora_layers = get_peft_model_state_dict(text_encoder_two) + text_encoder_lora_layers = convert_state_dict_to_diffusers(get_peft_model_state_dict(text_encoder_one)) + text_encoder_2_lora_layers = convert_state_dict_to_diffusers(get_peft_model_state_dict(text_encoder_two)) else: text_encoder_lora_layers = None text_encoder_2_lora_layers = None From 1aedb794bdc941b18a76d504f26dfd5e98cdbea9 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 27 Dec 2023 08:06:15 +0530 Subject: [PATCH 4/7] style --- examples/dreambooth/train_dreambooth_lora.py | 6 ++++-- examples/text_to_image/train_text_to_image_lora.py | 6 ++++-- .../text_to_image/train_text_to_image_lora_sdxl.py | 10 +++++++--- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py index 3219385b1f8f..f0bc78da3f49 100644 --- a/examples/dreambooth/train_dreambooth_lora.py +++ b/examples/dreambooth/train_dreambooth_lora.py @@ -54,7 +54,7 @@ ) from diffusers.loaders import LoraLoaderMixin from diffusers.optimization import get_scheduler -from diffusers.utils import check_min_version, is_wandb_available, convert_state_dict_to_diffusers +from diffusers.utils import check_min_version, convert_state_dict_to_diffusers, is_wandb_available from diffusers.utils.import_utils import is_xformers_available @@ -855,7 +855,9 @@ def save_model_hook(models, weights, output_dir): if isinstance(model, type(accelerator.unwrap_model(unet))): unet_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) elif isinstance(model, type(accelerator.unwrap_model(text_encoder))): - text_encoder_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) + text_encoder_lora_layers_to_save = convert_state_dict_to_diffusers( + get_peft_model_state_dict(model) + ) else: raise ValueError(f"unexpected save model: {model.__class__}") diff --git a/examples/text_to_image/train_text_to_image_lora.py b/examples/text_to_image/train_text_to_image_lora.py index 55fc377b141c..6d285362434d 100644 --- a/examples/text_to_image/train_text_to_image_lora.py +++ b/examples/text_to_image/train_text_to_image_lora.py @@ -44,7 +44,7 @@ from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, StableDiffusionPipeline, UNet2DConditionModel from diffusers.optimization import get_scheduler from diffusers.training_utils import compute_snr -from diffusers.utils import check_min_version, is_wandb_available, convert_state_dict_to_diffusers +from diffusers.utils import check_min_version, convert_state_dict_to_diffusers, is_wandb_available from diffusers.utils.import_utils import is_xformers_available @@ -809,7 +809,9 @@ def collate_fn(examples): accelerator.save_state(save_path) unwrapped_unet = accelerator.unwrap_model(unet) - unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unwrapped_unet)) + unet_lora_state_dict = convert_state_dict_to_diffusers( + get_peft_model_state_dict(unwrapped_unet) + ) StableDiffusionPipeline.save_lora_weights( save_directory=save_path, diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py index 5d4478cc0bc6..3dab86a48492 100644 --- a/examples/text_to_image/train_text_to_image_lora_sdxl.py +++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py @@ -52,7 +52,7 @@ from diffusers.loaders import LoraLoaderMixin from diffusers.optimization import get_scheduler from diffusers.training_utils import compute_snr -from diffusers.utils import check_min_version, is_wandb_available, convert_state_dict_to_diffusers +from diffusers.utils import check_min_version, convert_state_dict_to_diffusers, is_wandb_available from diffusers.utils.import_utils import is_xformers_available @@ -653,9 +653,13 @@ def save_model_hook(models, weights, output_dir): if isinstance(model, type(accelerator.unwrap_model(unet))): unet_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) elif isinstance(model, type(accelerator.unwrap_model(text_encoder_one))): - text_encoder_one_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) + text_encoder_one_lora_layers_to_save = convert_state_dict_to_diffusers( + get_peft_model_state_dict(model) + ) elif isinstance(model, type(accelerator.unwrap_model(text_encoder_two))): - text_encoder_two_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model)) + text_encoder_two_lora_layers_to_save = convert_state_dict_to_diffusers( + get_peft_model_state_dict(model) + ) else: raise ValueError(f"unexpected save model: {model.__class__}") From 2f96aba4c0949ffde89e8f21469a6663de04888e Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 27 Dec 2023 08:07:25 +0530 Subject: [PATCH 5/7] lcm lora sdxl. --- .../train_lcm_distill_lora_sdxl.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py index 2733eb146cd3..3c9482773489 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py @@ -51,7 +51,7 @@ UNet2DConditionModel, ) from diffusers.optimization import get_scheduler -from diffusers.utils import check_min_version, is_wandb_available +from diffusers.utils import check_min_version, convert_state_dict_to_diffusers, is_wandb_available from diffusers.utils.import_utils import is_xformers_available @@ -113,7 +113,7 @@ def log_validation(vae, args, accelerator, weight_dtype, step, unet=None, is_fin if unet is None: raise ValueError("Must provide a `unet` when doing intermediate validation.") unet = accelerator.unwrap_model(unet) - state_dict = get_peft_model_state_dict(unet) + state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) to_load = state_dict else: to_load = args.output_dir @@ -819,7 +819,7 @@ def save_model_hook(models, weights, output_dir): unet_ = accelerator.unwrap_model(unet) # also save the checkpoints in native `diffusers` format so that it can be easily # be independently loaded via `load_lora_weights()`. - state_dict = get_peft_model_state_dict(unet_) + state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet_)) StableDiffusionXLPipeline.save_lora_weights(output_dir, unet_lora_layers=state_dict) for _, model in enumerate(models): @@ -1332,7 +1332,7 @@ def compute_time_ids(original_size, crops_coords_top_left): accelerator.wait_for_everyone() if accelerator.is_main_process: unet = accelerator.unwrap_model(unet) - unet_lora_state_dict = get_peft_model_state_dict(unet) + unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet)) StableDiffusionXLPipeline.save_lora_weights(args.output_dir, unet_lora_layers=unet_lora_state_dict) if args.push_to_hub: From a53543da78459ec26679410f5ffe4f2f1b514f6d Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 27 Dec 2023 20:14:15 +0530 Subject: [PATCH 6/7] unwrap --- .../consistency_distillation/train_lcm_distill_lora_sdxl.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py index 3c9482773489..263588ea97e0 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py @@ -1184,7 +1184,7 @@ def compute_time_ids(original_size, crops_coords_top_left): # solver timestep. # With the adapters disabled, the `unet` is the regular teacher model. - unet.disable_adapters() + accelerator.unwrap_model(unet).disable_adapters() with torch.no_grad(): # 1. Get teacher model prediction on noisy_model_input z_{t_{n + k}} and conditional embedding c cond_teacher_output = unet( @@ -1248,7 +1248,7 @@ def compute_time_ids(original_size, crops_coords_top_left): x_prev = solver.ddim_step(pred_x0, pred_noise, index).to(unet.dtype) # re-enable unet adapters to turn the `unet` into a student unet. - unet.enable_adapters() + accelerator.unwrap_model(unet).disable_adapters() # 9. Get target LCM prediction on x_prev, w, c, t_n (timesteps) # Note that we do not use a separate target network for LCM-LoRA distillation. From fadfac1e2c9becf510d154fe98bb8022c29671ac Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 27 Dec 2023 20:27:19 +0530 Subject: [PATCH 7/7] fix: enable_adapters(). --- .../consistency_distillation/train_lcm_distill_lora_sdxl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py index 263588ea97e0..95e7b2dbaa27 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py @@ -1248,7 +1248,7 @@ def compute_time_ids(original_size, crops_coords_top_left): x_prev = solver.ddim_step(pred_x0, pred_noise, index).to(unet.dtype) # re-enable unet adapters to turn the `unet` into a student unet. - accelerator.unwrap_model(unet).disable_adapters() + accelerator.unwrap_model(unet).enable_adapters() # 9. Get target LCM prediction on x_prev, w, c, t_n (timesteps) # Note that we do not use a separate target network for LCM-LoRA distillation.