[HiDream LoRA] optimizations + small updates (#11381)

* 1. add pre-computation of prompt embeddings when custom prompts are used as well
2. save model card even if model is not pushed to hub
3. remove scheduler initialization from code example - not necessary anymore (it's now if the base model's config)
4. add skip_final_inference - to allow to run with validation, but skip the final loading of the pipeline with the lora weights to reduce memory reqs

* pre encode validation prompt as well

* Update examples/dreambooth/train_dreambooth_lora_hidream.py

Co-authored-by: Sayak Paul <[email protected]>

* Update examples/dreambooth/train_dreambooth_lora_hidream.py

Co-authored-by: Sayak Paul <[email protected]>

* Update examples/dreambooth/train_dreambooth_lora_hidream.py

Co-authored-by: Sayak Paul <[email protected]>

* pre encode validation prompt as well

* Apply style fixes

* empty commit

* change default trained modules

* empty commit

* address comments + change encoding of validation prompt (before it was only pre-encoded if custom prompts are provided, but should be pre-encoded either way)

* Apply style fixes

* empty commit

* fix validation_embeddings definition

* fix final inference condition

* fix pipeline deletion in last inference

* Apply style fixes

* empty commit

* layers

* remove readme remarks on only pre-computing when instance prompt is provided and change example to 3d icons

* smol fix

* empty commit

---------

Co-authored-by: Sayak Paul <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

examples/dreambooth/README_hidream.md

@@ -51,54 +51,41 @@ When running `accelerate config`, if we specify torch compile mode to True there
 Note also that we use PEFT library as backend for LoRA training, make sure to have `peft>=0.14.0` installed in your environment.
 
 
-### Dog toy example
+### 3d icon example
 
-Now let's get our dataset. For this example we will use some dog images: https://huggingface.co/datasets/diffusers/dog-example.
-
-Let's first download it locally:
-
-```python
-from huggingface_hub import snapshot_download
-
-local_dir = "./dog"
-snapshot_download(
-    "diffusers/dog-example",
-    local_dir=local_dir, repo_type="dataset",
-    ignore_patterns=".gitattributes",
-)
-```
+For this example we will use some 3d icon images: https://huggingface.co/datasets/linoyts/3d_icon.
 
 This will also allow us to push the trained LoRA parameters to the Hugging Face Hub platform.
 
 Now, we can launch training using:
 > [!NOTE]
 > The following training configuration prioritizes lower memory consumption by using gradient checkpointing, 
-> 8-bit Adam optimizer, latent caching, offloading, no validation. 
-> Additionally, when provided with 'instance_prompt' only and no 'caption_column' (used for custom prompts for each image)
-> text embeddings are pre-computed to save memory.
-
+> 8-bit Adam optimizer, latent caching, offloading, no validation.
+> all text embeddings are pre-computed to save memory.
 ```bash
 export MODEL_NAME="HiDream-ai/HiDream-I1-Dev"
-export INSTANCE_DIR="dog"
+export INSTANCE_DIR="linoyts/3d_icon"
 export OUTPUT_DIR="trained-hidream-lora"
 
 accelerate launch train_dreambooth_lora_hidream.py \
   --pretrained_model_name_or_path=$MODEL_NAME  \
-  --instance_data_dir=$INSTANCE_DIR \
+  --dataset_name=$INSTANCE_DIR \
   --output_dir=$OUTPUT_DIR \
   --mixed_precision="bf16" \
-  --instance_prompt="a photo of sks dog" \
+  --instance_prompt="3d icon" \
+  --caption_column="prompt"\
+  --validation_prompt="a 3dicon, a llama eating ramen" \
   --resolution=1024 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=4 \
   --use_8bit_adam \
-  --rank=16 \
+  --rank=8 \
   --learning_rate=2e-4 \
   --report_to="wandb" \
-  --lr_scheduler="constant" \
-  --lr_warmup_steps=0 \
+  --lr_scheduler="constant_with_warmup" \
+  --lr_warmup_steps=100 \
   --max_train_steps=1000 \
-  --cache_latents \
+  --cache_latents\
   --gradient_checkpointing \
   --validation_epochs=25 \
   --seed="0" \
@@ -128,6 +115,5 @@ We provide several options for optimizing memory optimization:
 * `--offload`: When enabled, we will offload the text encoder and VAE to CPU, when they are not used.
 * `cache_latents`: When enabled, we will pre-compute the latents from the input images with the VAE and remove the VAE from memory once done.
 * `--use_8bit_adam`: When enabled, we will use the 8bit version of AdamW provided by the `bitsandbytes` library.
-* `--instance_prompt` and no `--caption_column`: when only an instance prompt is provided, we will pre-compute the text embeddings and remove the text encoders from memory once done.
 
 Refer to the [official documentation](https://huggingface.co/docs/diffusers/main/en/api/pipelines/) of the `HiDreamImagePipeline` to know more about the model.