Add textual inversion XL for Gaudi

[dsocek]

What does this PR do?

This PR adds a script to run Textual Inversion fine-tuning on Gaudi for XL diffusion models, such as "stabilityai/stable-diffusion-xl-base-1.0", "cagliostrolab/animagine-xl-3.1", and so on. README doc is updated accordingly.

[atakaha]

do we need --allow_tf32?

[dsocek]

do we need --allow_tf32?

No, I removed in the PR amendment

[emascarenhas]

@dsocek , Is this still good to review/merge or does it require a sync, make style, fast_tests check?

[dsocek]

@dsocek , Is this still good to review/merge or does it require a sync, make style, fast_tests check?

@emascarenhas this PR is was good shape, but I need to add tests, rebase and test on 1.16/1.17. I can work on this today.

[dsocek]

@emascarenhas turns out guidance is not working properly even in the original diffusers. This PR will need to be slightly reworked, but good news is I managed find a fix. Created HF Diffusers PR here: huggingface/diffusers#9010

[dsocek]

Updated PR (ready for review)

cc: @emascarenhas @libinta @regisss

• Fixed guidance on G2
• Improved training logic to avoid OOMs when doing validation
• Added support for training 2nd text encoder

Fully tested with and without validation enabled on Gaudi2, all tests pass.

Train set (cats)

Classic textual inversion training set (6 images) used to learn new token guidance:

Test Output

Training command:

python textual_inversion_sdxl.py \
  --pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \
  --train_data_dir ./cat \
  --learnable_property object \
  --placeholder_token "<cat-toy>" \
  --initializer_token toy \
  --resolution 768 \
  --train_batch_size 1 \
  --gradient_accumulation_steps 4 \
  --max_train_steps 500 \
  --learning_rate 5.0e-04 \
  --scale_lr \
  --lr_scheduler constant \
  --lr_warmup_steps 0 \
  --output_dir /tmp/textual_inversion_cat_sdxl \
  --save_as_full_pipeline \
  --gaudi_config_name Habana/stable-diffusion \
  --throughput_warmup_steps 3

Inference script:

from optimum.habana.diffusers import GaudiStableDiffusionXLPipeline
import torch

model_id = "/tmp/textual_inversion_cat_sdxl"
pipe = GaudiStableDiffusionXLPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    use_habana=True,
    use_hpu_graphs=True,
    gaudi_config="Habana/stable-diffusion",
)

prompt = "A <cat-toy> backpack"
# engage 1st text encoder
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
image.save(f"cat-backpack.png")

# engage 2nd text encoder
image = pipe(prompt="", prompt_2=prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
image.save(f"cat-backpack_p2.png")

Output (4 images generated using 1st text encoder):

Output (4 images generated using 2nd text encoder):

[dsocek]

@regisss my equivalent PR fix is now merged in main diffusers branch (huggingface/diffusers#9010). May be good to also update optimum-habana with the fix too. Could you please review this PR when you get a chance?

[imangohari1]

@dsocek
Some initial thoughts here:

• Could you rebase/sync on top of OH main?
• Make sure make style is done.
• Could you share the results of test_diffusers.py after this changes? I think that won't change since this is a new code.
• We probably need to add a test for this in test_diffusers.py. If so, I could help with that if needed.

[imangohari1]

I am not sure why this is crashing on my end Daniel.

[imangohari1]

Suggested change

	image = pipe(prompt=prompt_1, prompt_2=prompt_2, num_inference_steps=50, guidance_scale=7.5).images[0]
	image = pipe(prompt=prompt, prompt_2=prompt_2, num_inference_steps=50, guidance_scale=7.5).images[0]

or line 141 prompt_1 = "A <cat-toy> backpack"

[imangohari1]

@dsocek
Thanks for working on this on both HF and OH. I've tested this pr and seems working.
I added some CI tests in attached patch.
I've tested it on both G2 and G1 and it passed.
Please review and apply with git am < 0001*.

Will take a final look afterwards.
Thank you.
0001-fea-ci-Added-text-inversion-sdxl.-minor-fix.patch

[dsocek]

@dsocek Thanks for working on this on both HF and OH. I've tested this pr and seems working. I added some CI tests in attached patch. I've tested it on both G2 and G1 and it passed. Please review and apply with git am < 0001*.

Will take a final look afterwards. Thank you. 0001-fea-ci-Added-text-inversion-sdxl.-minor-fix.patch

Thanks @imangohari1, I added your CI test implementation and also rebased to latest main

[imangohari1]

LGTM!
@regisss Could you pls. do a final review on this PR. it's been on a queue for a while.
Thank you.

[emascarenhas]

Please sync your PR with main/upstream and fix any merge conflicts. Thank you.

[dsocek]

Please sync your PR with main/upstream and fix any merge conflicts. Thank you.

@emascarenhas rebased and resolved conflicts.

@imangohari1 this good PR has been neglected for far too long.. Lets merge!

[emascarenhas]

Please sync your PR with main/upstream and fix any merge conflicts. Thank you.

@emascarenhas rebased and resolved conflicts.

@imangohari1 this good PR has been neglected for far too long.. Lets merge!

Thanks @dsocek , I'm adding it to our queue to review. Can you confirm you also ran fast tests and relevant slow tests and those did not generate new failures.

[dsocek]

@emascarenhas fast tests and relevant slow test PASS with the current PR:

# Setup
make test_installs
export GAUDI2_CI=1

# Run fast tests
export RUN_SLOW=false
python -m pytest tests/test_diffusers.py

# Output of fast tests
... ============== 112 passed, 48 skipped, 280 warnings in 1277.10s (0:21:17)

# Run relevant slow tests
export RUN_SLOW=true
python -m pytest tests/test_diffusers.py -v -s -k "test_textual_inversion_sdxl"

# Output of slow tests
... ============== 1 passed, 159 deselected, 5 warnings in 293.07s (0:04:53) 

[HuggingFaceDocBuilderDev]

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.