Skip to content
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
102 commits
Select commit Hold shift + click to select a range
b8b6465
initial commit
linoytsaban Apr 8, 2025
d728168
initial commit
linoytsaban Apr 8, 2025
0fa0993
initial commit
linoytsaban Apr 10, 2025
5ecf0ed
initial commit
linoytsaban Apr 10, 2025
911c30e
initial commit
linoytsaban Apr 10, 2025
b7fffee
initial commit
linoytsaban Apr 10, 2025
fde5eeb
Merge branch 'main' into hi-dream
linoytsaban Apr 10, 2025
02de3ce
Update examples/dreambooth/train_dreambooth_lora_hidream.py
linoytsaban Apr 14, 2025
4e08343
Merge branch 'main' into hi-dream
linoytsaban Apr 14, 2025
5257b46
move prompt embeds, pooled embeds outside
linoytsaban Apr 14, 2025
e9b4ad2
Update examples/dreambooth/train_dreambooth_lora_hidream.py
linoytsaban Apr 14, 2025
677bab1
Update examples/dreambooth/train_dreambooth_lora_hidream.py
linoytsaban Apr 14, 2025
dafc5fe
Merge branch 'main' into hi-dream
linoytsaban Apr 14, 2025
139568a
Merge remote-tracking branch 'origin/hi-dream' into hi-dream
linoytsaban Apr 14, 2025
31aa0a2
fix import
linoytsaban Apr 14, 2025
de1654a
fix import and tokenizer 4, text encoder 4 loading
linoytsaban Apr 14, 2025
33385c9
te
linoytsaban Apr 14, 2025
d993e16
prompt embeds
linoytsaban Apr 14, 2025
c296b6f
fix naming
linoytsaban Apr 14, 2025
aa6b6e2
shapes
linoytsaban Apr 14, 2025
ecc1c18
initial commit to add HiDreamImageLoraLoaderMixin
linoytsaban Apr 14, 2025
c439c89
fix init
linoytsaban Apr 14, 2025
fc97a54
Merge branch 'huggingface:main' into hi-dream
linoytsaban Apr 14, 2025
22e9ae8
add tests
linoytsaban Apr 14, 2025
8e2e1f1
Merge remote-tracking branch 'origin/hi-dream' into hi-dream
linoytsaban Apr 14, 2025
3653dcc
loader
linoytsaban Apr 14, 2025
32567bb
Merge branch 'main' into hi-dream
linoytsaban Apr 14, 2025
26f289d
Merge branch 'main' into hi-dream
linoytsaban Apr 14, 2025
fcf6eaa
fix model input
linoytsaban Apr 15, 2025
0fdc7dd
add code example to readme
linoytsaban Apr 15, 2025
759d204
Merge remote-tracking branch 'origin/hi-dream' into hi-dream
linoytsaban Apr 15, 2025
62f2f15
fix default max length of text encoders
linoytsaban Apr 15, 2025
3cb1a4c
prints
linoytsaban Apr 15, 2025
82bcd44
nullify training cond in unpatchify for temp fix to incompatible shap…
linoytsaban Apr 15, 2025
e956980
Merge branch 'main' into hi-dream
linoytsaban Apr 15, 2025
764bce6
Merge branch 'main' into hi-dream
sayakpaul Apr 16, 2025
75aa8bd
smol fix
linoytsaban Apr 16, 2025
47e861f
unpatchify
linoytsaban Apr 16, 2025
a461d38
unpatchify
linoytsaban Apr 16, 2025
2a1336c
Merge branch 'main' into hi-dream
linoytsaban Apr 16, 2025
1b20d15
Merge remote-tracking branch 'origin/hi-dream' into hi-dream
linoytsaban Apr 16, 2025
b31b595
fix validation
linoytsaban Apr 16, 2025
640bead
Merge branch 'main' into hi-dream
linoytsaban Apr 16, 2025
466c9c0
flip pred and loss
linoytsaban Apr 17, 2025
04be03a
Merge remote-tracking branch 'origin/hi-dream' into hi-dream
linoytsaban Apr 17, 2025
baa9784
Merge branch 'main' into hi-dream
linoytsaban Apr 17, 2025
efb22e2
Merge branch 'main' into hi-dream
linoytsaban Apr 17, 2025
6043d9d
fix shift!!!
linoytsaban Apr 17, 2025
0a2adfd
Merge remote-tracking branch 'origin/hi-dream' into hi-dream
linoytsaban Apr 17, 2025
92306d9
Merge branch 'main' into hi-dream
linoytsaban Apr 17, 2025
13e6f0d
revert unpatchify changes (for now)
linoytsaban Apr 17, 2025
c3a4047
Merge branch 'main' into hi-dream
linoytsaban Apr 18, 2025
ae39434
smol fix
linoytsaban Apr 18, 2025
c8932ed
Apply style fixes
github-actions[bot] Apr 18, 2025
c8ac7d5
workaround moe training
linoytsaban Apr 18, 2025
bf7ace6
workaround moe training
linoytsaban Apr 18, 2025
2fcb17d
remove prints
linoytsaban Apr 18, 2025
80f13be
to reduce some memory, keep vae in `weight_dtype` same as we have for…
linoytsaban Apr 18, 2025
b8039c9
refactor to align with HiDream refactor
linoytsaban Apr 18, 2025
c331597
refactor to align with HiDream refactor
linoytsaban Apr 18, 2025
c32cccc
refactor to align with HiDream refactor
linoytsaban Apr 18, 2025
6e070b8
add support for cpu offloading of text encoders
linoytsaban Apr 18, 2025
d77e42a
Apply style fixes
github-actions[bot] Apr 18, 2025
5c8c339
adjust lr and rank for train example
linoytsaban Apr 18, 2025
abfb389
fix copies
linoytsaban Apr 18, 2025
a5fe6be
Apply style fixes
github-actions[bot] Apr 18, 2025
d562776
Merge branch 'main' into hi-dream
linoytsaban Apr 19, 2025
2798d40
update README
linoytsaban Apr 19, 2025
ab960c2
update README
linoytsaban Apr 19, 2025
52d9421
update README
linoytsaban Apr 19, 2025
fc5eb48
fix license
linoytsaban Apr 19, 2025
a012914
keep prompt2,3,4 as None in validation
linoytsaban Apr 19, 2025
d5b9ecc
remove reverse ode comment
linoytsaban Apr 19, 2025
f04a13a
Update examples/dreambooth/train_dreambooth_lora_hidream.py
linoytsaban Apr 19, 2025
0b75081
Update examples/dreambooth/train_dreambooth_lora_hidream.py
linoytsaban Apr 19, 2025
4db988f
vae offload change
linoytsaban Apr 19, 2025
13192a3
fix text encoder offloading
linoytsaban Apr 19, 2025
408dfdb
Apply style fixes
github-actions[bot] Apr 19, 2025
3383446
cleaner to_kwargs
linoytsaban Apr 20, 2025
73ab201
fix module name in copied from
linoytsaban Apr 20, 2025
62019cc
Merge branch 'main' into hi-dream
linoytsaban Apr 21, 2025
a07ee59
add requirements
linoytsaban Apr 21, 2025
a751cc2
fix offloading
linoytsaban Apr 21, 2025
120b821
fix offloading
linoytsaban Apr 21, 2025
9d71d3b
fix offloading
linoytsaban Apr 21, 2025
2798d77
update transformers version in reqs
linoytsaban Apr 21, 2025
363d29b
try AutoTokenizer
linoytsaban Apr 21, 2025
8f24e8c
try AutoTokenizer
linoytsaban Apr 21, 2025
b31bdf0
Apply style fixes
github-actions[bot] Apr 21, 2025
6c77651
empty commit
linoytsaban Apr 21, 2025
246978e
Merge remote-tracking branch 'origin/hi-dream' into hi-dream
linoytsaban Apr 21, 2025
9b6ef43
Delete tests/lora/test_lora_layers_hidream.py
linoytsaban Apr 21, 2025
fb3ac74
change tokenizer_4 to load with AutoTokenizer as well
linoytsaban Apr 21, 2025
82a4037
make text_encoder_four and tokenizer_four configurable
linoytsaban Apr 21, 2025
9de10cb
save model card
linoytsaban Apr 21, 2025
418f6a3
save model card
linoytsaban Apr 21, 2025
36c1ada
revert T5
linoytsaban Apr 21, 2025
bd399b1
fix test
linoytsaban Apr 21, 2025
a8e1a0b
Merge branch 'main' into hi-dream
linoytsaban Apr 21, 2025
bd275c0
Merge branch 'main' into hi-dream
linoytsaban Apr 21, 2025
e49be97
Merge branch 'main' into hi-dream
sayakpaul Apr 22, 2025
ed97dba
remove non diffusers lumina2 conversion
linoytsaban Apr 22, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 127 additions & 0 deletions examples/dreambooth/README_hidream.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
# DreamBooth training example for HiDream Image

[DreamBooth](https://arxiv.org/abs/2208.12242) is a method to personalize text2image models like stable diffusion given just a few (3~5) images of a subject.

The `train_dreambooth_lora_hidream.py` script shows how to implement the training procedure with [LoRA](https://huggingface.co/docs/peft/conceptual_guides/adapter#low-rank-adaptation-lora) and adapt it for [HiDream Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/).


This will also allow us to push the trained model parameters to the Hugging Face Hub platform.

## Running locally with PyTorch

### Installing the dependencies

Before running the scripts, make sure to install the library's training dependencies:

**Important**

To make sure you can successfully run the latest versions of the example scripts, we highly recommend **installing from source** and keeping the install up to date as we update the example scripts frequently and install some example-specific requirements. To do this, execute the following steps in a new virtual environment:

```bash
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install -e .
```

Then cd in the `examples/dreambooth` folder and run
```bash
pip install -r requirements_sana.txt
```

And initialize an [🤗Accelerate](https://github.com/huggingface/accelerate/) environment with:

```bash
accelerate config
```

Or for a default accelerate configuration without answering questions about your environment

```bash
accelerate config default
```

Or if your environment doesn't support an interactive shell (e.g., a notebook)

```python
from accelerate.utils import write_basic_config
write_basic_config()
```

When running `accelerate config`, if we specify torch compile mode to True there can be dramatic speedups.
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

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",
)
```

This will also allow us to push the trained LoRA parameters to the Hugging Face Hub platform.

Now, we can launch training using:

```bash
export MODEL_NAME="HiDream-ai/HiDream-I1-Dev"
export INSTANCE_DIR="dog"
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 \
--output_dir=$OUTPUT_DIR \
--mixed_precision="bf16" \
--instance_prompt="a photo of sks dog" \
--resolution=1024 \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--use_8bit_adam \
--learning_rate=1e-4 \
--report_to="wandb" \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=500 \
--validation_prompt="A photo of sks dog in a bucket" \
--validation_epochs=25 \
--seed="0" \
--push_to_hub
```

For using `push_to_hub`, make you're logged into your Hugging Face account:

```bash
huggingface-cli login
```

To better track our training experiments, we're using the following flags in the command above:

* `report_to="wandb` will ensure the training runs are tracked on [Weights and Biases](https://wandb.ai/site). To use it, be sure to install `wandb` with `pip install wandb`. Don't forget to call `wandb login <your_api_key>` before training if you haven't done it before.
* `validation_prompt` and `validation_epochs` to allow the script to do a few validation inference runs. This allows us to qualitatively check if the training is progressing as expected.

## Notes

Additionally, we welcome you to explore the following CLI arguments:

* `--lora_layers`: The transformer modules to apply LoRA training on. Please specify the layers in a comma seperated. E.g. - "to_k,to_q,to_v" will result in lora training of attention layers only.
* `--system_prompt`: A custom system prompt to provide additional personality to the model.
* `--max_sequence_length`: Maximum sequence length to use for text embeddings.


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.

Refer to the [official documentation](https://huggingface.co/docs/diffusers/main/en/api/pipelines/) of the `HiDreamImagePipeline` to know more about the model.
Loading
Loading