[LoRA] add LoRA support to HiDream and fine-tuning script

[linoytsaban]

add lora training for Hi Dream Image

• trains transformer layers (text encoder training not supported at the moment)
• MoE training disabled for LoRA fine-tuning (potentially re-visited in the future)

memory optimizations:

1. cpu offloading support for vae and text encoders
2. latent caching
3. vae loaded in mixed precision

example -

--pretrained_model_name_or_path=HiDream-ai/HiDream-I1-Full --dataset_name=Norod78/Yarn-art-style --output_dir=trained-hidream-lora --mixed_precision=bf16 --instance_prompt="a dog, yarn art style" --caption_column=text --resolution=512 --train_batch_size=1 --gradient_accumulation_steps=1 --optimizer=prodigy --rank=16 --learning_rate=1.0 --report_to=wandb --lr_scheduler=constant --lr_warmup_steps=0 --max_train_steps=1000 --validation_epochs=25 --validation_prompt="yoda, yarn art style" --seed=0

[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.

[sayakpaul]

Hundred percent not a blocker for this PR. But we could allow users to drop out the other text encoders by zeroing those embeds.

[sayakpaul]

Also, in encode_prompt(),

prompt_embeds= [t5_prompt_embeds, llama3_prompt_embeds]

Would prompt_embeds = prompt_embeds.to(accelerator.device) work here.

[sayakpaul]

Thanks a lot for this!

I have left some clarification questions. LMK if they make sense.

[hlky]

Just a few comments, will review again when out of draft, thanks @linoytsaban!

[linoytsaban]

@bot /style

[github-actions]

Style fixes have been applied. View the workflow run here.

[sayakpaul]

Thank you!

[vladmandic]

follow-up issue ##11383