Support EDM-style training in DreamBooth LoRA SDXL script #7126

sayakpaul · 2024-02-28T05:54:07Z

Command example:

CUDA_VISIBLE_DEVICES=1 accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path="playgroundai/playground-v2.5-1024px-aesthetic"  \
  --instance_data_dir="dog" \
  --output_dir="dog-playground-lora" \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-4 \
  --use_8bit_adam \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=25 \
  --seed="0" \
  --push_to_hub

WandB: https://wandb.ai/sayakpaul/dreambooth-lora-playground/runs/sxe4bavp

HuggingFaceDocBuilderDev · 2024-02-28T06:01:59Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

patil-suraj · 2024-02-28T05:56:33Z

examples/dreambooth/train_dreambooth_lora_playground.py

+    def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
+        sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
+        schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
+        timesteps = timesteps.to(accelerator.device)
+
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < n_dim:
+            sigma = sigma.unsqueeze(-1)
+        return sigma


For later:

We could think of making this more general by allowing to sample sigmas as presented in the paper, cf https://github.com/NVlabs/edm/blob/main/training/loss.py#L74

examples/dreambooth/train_dreambooth_lora_playground.py

patil-suraj · 2024-02-28T06:02:32Z

examples/dreambooth/train_dreambooth_lora_playground.py

+                    )[0]
+
+                model_pred = model_pred * (-sigmas) + noisy_model_input
+                weighing = sigmas**-2.0


For later: this could be made configurable, as there are multiple weighing alternatives. In EDM they use
https://github.com/NVlabs/edm/blob/main/training/loss.py#L75

examples/dreambooth/train_dreambooth_lora_playground.py

sayakpaul · 2024-02-28T06:26:26Z

@patil-suraj could you give this another look? Results are still blank: https://wandb.ai/sayakpaul/dreambooth-lora-playground/runs/i7aq50g0.

Experimenting with a lower LR.

examples/dreambooth/train_dreambooth_lora_playground.py

patil-suraj · 2024-02-28T07:06:40Z

Can't seem to find anything else, will also try to run the script and see what's going on.

patil-suraj · 2024-02-28T09:01:39Z

Got a good run: https://wandb.ai/psuraj/dreambooth-lora-playground/runs/j34izml0?workspace=user-psuraj (still going on)

What fixed it:

Load fp32 variant of the vae .
don't use autocast during generation
disable loss weighing.

examples/dreambooth/train_dreambooth_lora_playground.py

Co-authored-by: Suraj Patil <surajp815@gmail.com>

sayakpaul · 2024-02-28T09:13:26Z

Applied the changes, @patil-suraj. Could you try another run?

patil-suraj · 2024-02-28T09:19:35Z

Started a new run here https://wandb.ai/psuraj/dreambooth-lora-playground/runs/ef2qkmre

sayakpaul · 2024-02-28T13:04:17Z

@patil-suraj ready for a review. Feel free to test the script too :)

@pcuenca feel free to give this a review as well.

pcuenca

Looks good in general. I'd maybe try to avoid hardcoded references to the string "playgroundai" to make decisions, if possible.

pcuenca · 2024-02-28T14:58:48Z

examples/dreambooth/README_sdxl.md

+
+It's now possible to perform EDM-style training as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://arxiv.org/abs/2206.00364). 
+
+For the SDXL model, simple set:


Suggested change

For the SDXL model, simple set:

For the standard SDXL model, simply set:

Does it work with SDXL out of the box? 🤯

There's a test that you can check but I haven't done a full-blown training run.

For LoRA it might not work, but can def be fine-tuned with EDM.

@patil-suraj elaborate?

examples/dreambooth/train_dreambooth_lora_sdxl.py

patil-suraj

Looking good, left some comments. +1 to what pedro said.

patil-suraj · 2024-02-28T15:38:40Z

examples/dreambooth/README_sdxl.md

+
+It's now possible to perform EDM-style training as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://arxiv.org/abs/2206.00364). 
+
+For the SDXL model, simple set:


For LoRA it might not work, but can def be fine-tuned with EDM.

examples/dreambooth/train_dreambooth_lora_sdxl.py

Co-authored-by: Suraj Patil <surajp815@gmail.com>

sayakpaul · 2024-02-28T16:18:38Z

@patil-suraj ready for another review.

examples/dreambooth/train_dreambooth_lora_sdxl.py

Co-authored-by: Suraj Patil <surajp815@gmail.com>

sayakpaul · 2024-02-29T01:33:06Z

@pcuenca @patil-suraj I have addressed all your comments. Would appreciate another review.

I am going to run with the command from the OP one more time and also with regular SDXL with --do_edm_style_training.

sayakpaul · 2024-02-29T03:07:46Z

Started a regular SDXL run with EDM:

CUDA_VISIBLE_DEVICES=1 accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0"  \
  --instance_data_dir="dog"\
  --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
  --output_dir="lora-sdxl-dog" \
  --mixed_precision="fp16" \
  --use_8bit_adam \
  --do_edm_style_training \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=25 \
  --seed="0"

Garbage results: https://wandb.ai/sayakpaul/dreambooth-lora-sd-xl/runs/bup8u1yc. Let me further tweak around some things.

sayakpaul · 2024-02-29T13:18:52Z

@pcuenca @patil-suraj the script now should work out of the box when do_edm_style_training specified for SDXL:

CUDA_VISIBLE_DEVICES=1 accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0"  \
  --instance_data_dir="dog"\
  --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
  --output_dir="lora-sdxl-dog" \
  --mixed_precision="fp16" \
  --use_8bit_adam \
  --do_edm_style_training \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=25 \
  --seed="0"

Feel free to train one yourselves. Here are my results: https://wandb.ai/sayakpaul/dreambooth-lora-sd-xl/runs/dz77sffl

Please review the changes so that we can ship this beast!

patil-suraj

Thanks for addressing the comments. The script is in a very good state for EDM. ~~I would just suggest to verify the euler bit before adding it here or maybe even do it in another PR.~~ (saw the other comment, all good)

Also are the vae weights loaded and kept in fp32 ?

examples/dreambooth/train_dreambooth_lora_sdxl.py

patil-suraj · 2024-03-01T09:14:31Z

examples/dreambooth/train_dreambooth_lora_sdxl.py

+                    # There might be other alternatives for weighting as well:
+                    # https://github.com/huggingface/diffusers/pull/7126#discussion_r1505404686
+                    if "EDM" not in scheduler_type:
+                        weighting = (sigmas**-2.0).float()


We should verify if this works with euler

It is: https://wandb.ai/sayakpaul/dreambooth-lora-sd-xl/runs/dz77sffl. When do_edm_style_training is True and the scheduler is not EDM*, we are using EulerDiscrete. The run is from that setting.

Does that work?

Sounds good!

sayakpaul · 2024-03-01T10:48:33Z

Also are the vae weights loaded and kept in fp32 ?

Yes, that is the case. I have addressed your other comment as well, @patil-suraj. LMK.

patil-suraj

Looks great now! Feel free to merge

sayakpaul · 2024-03-03T03:59:12Z

I keep getting a PermissionError: [Errno 13] Permission denied when trying to access the dog folder. Logged in to huggingface and running as administrator. All folders have full read/write access.

That is an issue quite unrelated to this PR.

bghira · 2024-03-06T23:40:37Z

examples/dreambooth/train_dreambooth_lora_sdxl.py

+    if args.do_edm_style_training and args.snr_gamma is not None:
+        raise ValueError("Min-SNR formulation is not supported when conducting EDM-style training.")


do this earlier, so it doesn't load the model yet.

It's at the beginning:

diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py

Line 948 in 59433ca

if args.do_edm_style_training and args.snr_gamma is not None:

way before the model loading code.

… of #7126) (#7182) * add edm style training * style * finish adding edm training feature * import fix * fix latents mean * minor adjustments * add edm to readme * style * fix autocast and scheduler config issues when using edm * style --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

sayakpaul added 2 commits February 28, 2024 08:46

add: dreambooth lora script for Playground v2.5

f3013e2

fix: kwarg

cb5f0d7

sayakpaul requested a review from patil-suraj February 28, 2024 05:54

patil-suraj reviewed Feb 28, 2024

View reviewed changes

address suraj's comments.

84047c2

sayakpaul commented Feb 28, 2024

View reviewed changes

examples/dreambooth/train_dreambooth_lora_playground.py Outdated Show resolved Hide resolved

patil-suraj reviewed Feb 28, 2024

View reviewed changes

sayakpaul and others added 3 commits February 28, 2024 14:40

Apply suggestions from code review

746b625

Co-authored-by: Suraj Patil <surajp815@gmail.com>

Merge branch 'main' into playground-dreambooth-lora

dcf36b4

apply suraj's suggestion

0d1547f

sayakpaul added 6 commits February 28, 2024 17:41

incorporate changes in the canonical script./

c9dd1d0

tracker naming

e2b7144

fix: schedule determination

db025fc

add: two simple tests

80ef425

remove playground script

3ae0e28

note about edm-style training

2f49630

sayakpaul marked this pull request as ready for review February 28, 2024 13:03

sayakpaul requested a review from patil-suraj February 28, 2024 13:11

pcuenca approved these changes Feb 28, 2024

View reviewed changes

patil-suraj reviewed Feb 28, 2024

View reviewed changes

sayakpaul and others added 3 commits February 28, 2024 21:24

address pedro's comments.

c8ed8af

address part of Suraj's comments.

e02e6f6

Apply suggestions from code review

a8bea31

Co-authored-by: Suraj Patil <surajp815@gmail.com>

sayakpaul added 2 commits February 28, 2024 21:35

use mse_loss.

00156a3

add comments for preconditioning.

3621e18

quality

206f2c7

patil-suraj reviewed Feb 28, 2024

View reviewed changes

examples/dreambooth/train_dreambooth_lora_sdxl.py Outdated Show resolved Hide resolved

sayakpaul and others added 3 commits February 29, 2024 06:54

Update examples/dreambooth/train_dreambooth_lora_sdxl.py

f7fc1f6

Co-authored-by: Suraj Patil <surajp815@gmail.com>

Merge branch 'main' into playground-dreambooth-lora

d96d8ea

tackle v-pred.

dde7595

Empty-Commit

128b877

sayakpaul changed the title ~~Support DreamBooth LoRA for Playground~~ Support EDM-style training in DreamBooth LoRA SDXL script Feb 29, 2024

Merge branch 'main' into playground-dreambooth-lora

e052fe1

support edm for sdxl too.

65c382c

patil-suraj reviewed Mar 1, 2024

View reviewed changes

sayakpaul added 2 commits March 1, 2024 16:15

Merge branch 'main' into playground-dreambooth-lora

0f046d8

address suraj's comments.

10c55bb

patil-suraj approved these changes Mar 2, 2024

View reviewed changes

sayakpaul added 2 commits March 2, 2024 13:52

Merge branch 'main' into playground-dreambooth-lora

e3210fc

Empty-Commit

fe75b46

linoytsaban mentioned this pull request Mar 2, 2024

[Advanced DreamBooth LoRA SDXL] Support EDM-style training (follow up of #7126) #7182

Merged

sayakpaul merged commit ccb93dc into main Mar 3, 2024
10 checks passed

sayakpaul deleted the playground-dreambooth-lora branch March 3, 2024 03:59

bghira reviewed Mar 6, 2024

View reviewed changes

sayakpaul mentioned this pull request Apr 2, 2024

7529 do not disable autocast for cuda devices #7530

Merged

6 tasks

axel578 mentioned this pull request Apr 25, 2024

Can't apply loras on playground 2.5 comfyanonymous/ComfyUI#3347

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support EDM-style training in DreamBooth LoRA SDXL script #7126

Support EDM-style training in DreamBooth LoRA SDXL script #7126

sayakpaul commented Feb 28, 2024 •

edited

Loading

HuggingFaceDocBuilderDev commented Feb 28, 2024

patil-suraj Feb 28, 2024

patil-suraj Feb 28, 2024

sayakpaul commented Feb 28, 2024 •

edited

Loading

patil-suraj commented Feb 28, 2024

patil-suraj commented Feb 28, 2024

sayakpaul commented Feb 28, 2024

patil-suraj commented Feb 28, 2024

sayakpaul commented Feb 28, 2024

pcuenca left a comment

pcuenca Feb 28, 2024

sayakpaul Feb 28, 2024

patil-suraj Feb 28, 2024

sayakpaul Feb 28, 2024

patil-suraj left a comment •

edited

Loading

patil-suraj Feb 28, 2024

sayakpaul commented Feb 28, 2024

sayakpaul commented Feb 29, 2024

sayakpaul commented Feb 29, 2024

sayakpaul commented Feb 29, 2024

patil-suraj left a comment •

edited

Loading

patil-suraj Mar 1, 2024

sayakpaul Mar 1, 2024

patil-suraj Mar 2, 2024

sayakpaul commented Mar 1, 2024

patil-suraj left a comment

sayakpaul commented Mar 3, 2024

bghira Mar 6, 2024

sayakpaul Mar 7, 2024


		It's now possible to perform EDM-style training as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://arxiv.org/abs/2206.00364).

		For the SDXL model, simple set:

	For the SDXL model, simple set:
	For the standard SDXL model, simply set:

		if args.do_edm_style_training and args.snr_gamma is not None:
		raise ValueError("Min-SNR formulation is not supported when conducting EDM-style training.")

Support EDM-style training in DreamBooth LoRA SDXL script #7126

Support EDM-style training in DreamBooth LoRA SDXL script #7126

Conversation

sayakpaul commented Feb 28, 2024 • edited Loading

HuggingFaceDocBuilderDev commented Feb 28, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Feb 28, 2024 • edited Loading

patil-suraj commented Feb 28, 2024

patil-suraj commented Feb 28, 2024

sayakpaul commented Feb 28, 2024

patil-suraj commented Feb 28, 2024

sayakpaul commented Feb 28, 2024

pcuenca left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

patil-suraj left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Feb 28, 2024

sayakpaul commented Feb 29, 2024

sayakpaul commented Feb 29, 2024

sayakpaul commented Feb 29, 2024

patil-suraj left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Mar 1, 2024

patil-suraj left a comment

Choose a reason for hiding this comment

sayakpaul commented Mar 3, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Feb 28, 2024 •

edited

Loading

sayakpaul commented Feb 28, 2024 •

edited

Loading

patil-suraj left a comment •

edited

Loading

patil-suraj left a comment •

edited

Loading