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

Hi @Nerogar!
For context - the self.training/ if training conditions in the HiDreamImageTransformer2DModel implementation come from the original implementation and refer to the case where aux loss is used and the experts are trained as well.
For LoRA training in this PR we currently do not support aux loss and MoE training, and hence we added
--fore_inference_output

diffusers/src/diffusers/models/transformers/transformer_hidream_image.py

Line 320 in 73ab201

if self.training and self.alpha > 0.0 and not self._force_inference_output:

Specifically we encountered an issue with unpatchify when initially working on this PR, with the resulting shape
of the output not as we'd expect, resulting with an error when trying to calculate the loss due to mismatch of the shapes between target and model_pred

[linoytsaban]

@bot /style

[github-actions]

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

[sayakpaul]

Thank you!

[vladmandic]

follow-up issue ##11383