[Examples] Add a training script for SDXL DreamBooth LoRA by sayakpaul · Pull Request #4016 · huggingface/diffusers

sayakpaul · 2023-07-10T05:43:09Z

What does this PR do?

Adds a training script to support fine-tuning SDXL with DreamBooth and LoRA. On top of #3896. Since the conflicts are brutal this PR exists.

Notes

No support for text encoder yet as the modifications are reasonably sized.
The training artifacts are available here: https://huggingface.co/diffusers/lora-trained-xl (private, only visible to the diffusers team members for now).

Used the following command to do a test run:

export MODEL_NAME="diffusers/stable-diffusion-xl-base-0.9"
export INSTANCE_DIR="dog"
export CLASS_DIR="dog-class"
export OUTPUT_DIR="lora-trained-xl"

accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-5 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=100 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --push_to_hub

The dog dataset was downloaded using the following code:

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir, repo_type="dataset",
    ignore_patterns=".gitattributes",
)

TODOs

Docs
Tests

sayakpaul · 2023-07-10T05:47:29Z

examples/dreambooth/train_dreambooth_lora_sdxl.py

+            repo_id = create_repo(
+                repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
+            ).repo_id


I think it's okay to create the repos with public visibility here since we're NOT distributing the SDXL weights here. It's just the LoRA weights. So, that should be fine w.r.t privacy and access concerns.

do they contain the training checkpoints? or just for the lora weights?

Checkpoints (which are intermediate LoRA parameters) as well as the final LoRA parameters. Example: https://huggingface.co/diffusers/lora-trained-xl-potato-head/tree/main (private diffusers team members).

HuggingFaceDocBuilderDev · 2023-07-10T05:53:28Z

The documentation is not available anymore as the PR was closed or merged.

sayakpaul · 2023-07-10T07:10:27Z

examples/dreambooth/README.md

+## Stable Diffusion XL
+
+We support fine-tuning of the UNet shipped in [Stable Diffusion XL](https://github.com/Stability-AI/generative-models/blob/main/assets/sdxl_report.pdf) with DreamBooth and LoRA via the `train_dreambooth_lora_sdxl.py` script. Please refer to the docs [here](./README_sdxl.md). 


A separate README here makes sense here I think.

src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

examples/dreambooth/train_dreambooth_lora_sdxl.py

patrickvonplaten

Very nice! For now we don't allow training the text encoder I assume no?

patrickvonplaten · 2023-07-10T10:06:50Z

@williamberman could you maybe also take a look here?

sayakpaul · 2023-07-10T10:26:18Z

Very nice! For now we don't allow training the text encoder I assume no?

Yeah that is right since the results are already quite nice especially if someone uses the Refiner. So, let's wait for a bit for the community and then incorporating that support won't be a big deal.

…Pipeline.

pcuenca

Awesome!

pcuenca · 2023-07-10T10:20:42Z

examples/dreambooth/README_sdxl.md

nit: license header :)

We don't include licensing in the REAMDEs, no?

examples/dreambooth/README_sdxl.md

examples/dreambooth/train_dreambooth_lora_sdxl.py

pcuenca · 2023-07-10T10:29:53Z

examples/dreambooth/train_dreambooth_lora_sdxl.py

+            repo_id = create_repo(
+                repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
+            ).repo_id


do they contain the training checkpoints? or just for the lora weights?

examples/dreambooth/README.md

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

n00mkrad · 2023-07-10T17:37:36Z

Training is cool, but do we have a way of applying these LoRAs with SDXL yet?

williamberman · 2023-07-10T17:44:06Z

examples/dreambooth/train_dreambooth_lora_sdxl.py

+    if args.enable_xformers_memory_efficient_attention:
+        if is_xformers_available():
+            import xformers
+
+            xformers_version = version.parse(xformers.__version__)
+            if xformers_version == version.parse("0.0.16"):
+                logger.warn(
+                    "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
+                )
+            unet.enable_xformers_memory_efficient_attention()
+        else:
+            raise ValueError("xformers is not available. Make sure it is installed correctly")


Should we be still including xformers in new training scripts given we have native flash attention in pytorch now?

The community still seems to use xformers quite a bit. So, let's keep its support.

williamberman

c'est magnifique!

…e#4016) * add dreambooth lora script for SDXL incorporating latest changes. * remove use_auth_token=True. * add: documentation * remove unneeded cli. * increase the number of training steps in the readme. * add LoraLoaderMixin to the subclassing mix. * add sdxl lora dreambooth test. * add: inference code sample. * add: refiner output. * add LoraLoaderMixin to the mix of classes of StableDiffusionXLImg2ImgPipeline. * change default resolution of DreamBoothDataset. * better sdxl report path. * Apply suggestions from code review Co-authored-by: Pedro Cuenca <pedro@huggingface.co> --------- Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

sayakpaul added 2 commits July 10, 2023 10:46

add dreambooth lora script for SDXL incorporating latest changes.

449ec65

remove use_auth_token=True.

5603a3a

sayakpaul commented Jul 10, 2023

View reviewed changes

sayakpaul added 5 commits July 10, 2023 11:49

add: documentation

cf98ea3

remove unneeded cli.

d77a876

increase the number of training steps in the readme.

5fd8948

add LoraLoaderMixin to the subclassing mix.

c462b54

add sdxl lora dreambooth test.

089bf77

sayakpaul marked this pull request as ready for review July 10, 2023 07:05

sayakpaul requested review from patrickvonplaten and pcuenca July 10, 2023 07:05

sayakpaul commented Jul 10, 2023

View reviewed changes

sayakpaul mentioned this pull request Jul 10, 2023

[WIP][Examples] Support LoRA DreamBooth training with SD XL #3896

Closed

add: inference code sample.

17e67aa

patrickvonplaten reviewed Jul 10, 2023

View reviewed changes

src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py Show resolved Hide resolved

patrickvonplaten reviewed Jul 10, 2023

View reviewed changes

examples/dreambooth/train_dreambooth_lora_sdxl.py Show resolved Hide resolved

patrickvonplaten reviewed Jul 10, 2023

View reviewed changes

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

patrickvonplaten approved these changes Jul 10, 2023

View reviewed changes

patrickvonplaten requested a review from williamberman July 10, 2023 10:06

add: refiner output.

92ae4b3

sayakpaul added 2 commits July 10, 2023 15:59

add LoraLoaderMixin to the mix of classes of StableDiffusionXLImg2Img…

ae49893

…Pipeline.

change default resolution of DreamBoothDataset.

785c106

pcuenca approved these changes Jul 10, 2023

View reviewed changes

mcmonkey4eva reviewed Jul 10, 2023

View reviewed changes

examples/dreambooth/README.md Outdated Show resolved Hide resolved

sayakpaul and others added 2 commits July 10, 2023 16:14

better sdxl report path.

891f32a

Apply suggestions from code review

82b4cb5

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

williamberman reviewed Jul 10, 2023

View reviewed changes

williamberman approved these changes Jul 10, 2023

View reviewed changes

sayakpaul merged commit 3d74dc2 into main Jul 11, 2023

sayakpaul deleted the dreambooth/sd-xl-3 branch July 11, 2023 02:08

okotaku mentioned this pull request Jul 30, 2023

[Examples] Support train_text_to_image_lora_sdxl.py #4365

Merged

6 tasks

		## Stable Diffusion XL

		We support fine-tuning of the UNet shipped in [Stable Diffusion XL](https://github.com/Stability-AI/generative-models/blob/main/assets/sdxl_report.pdf) with DreamBooth and LoRA via the `train_dreambooth_lora_sdxl.py` script. Please refer to the docs [here](./README_sdxl.md). No newline at end of file

Conversation

sayakpaul commented Jul 10, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Notes

TODOs

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Jul 10, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented Jul 10, 2023

Uh oh!

sayakpaul commented Jul 10, 2023

Uh oh!

pcuenca left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

n00mkrad commented Jul 10, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

williamberman left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

sayakpaul commented Jul 10, 2023 •

edited

Loading

HuggingFaceDocBuilderDev commented Jul 10, 2023 •

edited

Loading