Skip to content

[training ] add Kontext i2i training#11858

Merged
sayakpaul merged 14 commits intomainfrom
kontext-i2i-training
Jul 8, 2025
Merged

[training ] add Kontext i2i training#11858
sayakpaul merged 14 commits intomainfrom
kontext-i2i-training

Conversation

@sayakpaul
Copy link
Copy Markdown
Member

@sayakpaul sayakpaul commented Jul 4, 2025

What does this PR do?

Test command:

accelerate launch train_dreambooth_lora_flux_kontext.py \
  --pretrained_model_name_or_path=black-forest-labs/FLUX.1-Kontext-dev  \
  --output_dir="kontext-i2i" \
  --dataset_name="kontext-community/relighting" \
  --image_column="output" --cond_image_column="file_name" --caption_column="instruction" \
  --mixed_precision="bf16" \
  --resolution=1024 \
  --train_batch_size=1 \
  --guidance_scale=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --optimizer="adamw" \
  --use_8bit_adam \
  --cache_latents \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --seed="0"

I haven't finished it fully.

Additionally, I have taken the liberty to modify our training script to precompute the text embeddings when we have train_dataset.custom_instance_prompts. These are better called custom_instruction_prompts, IMO. So, in a future PR, we could switch to better variable names.

@sayakpaul sayakpaul requested a review from linoytsaban July 4, 2025 04:20
sayakpaul
sayakpaul commented Jul 4, 2025
View reviewed changes
examples/dreambooth/train_dreambooth_lora_flux_kontext.py
Comment on lines +962 to +963
to_tensor = transforms.ToTensor()
normalize = transforms.Normalize([0.5], [0.5])
Copy link
Copy Markdown
Member Author

@sayakpaul sayakpaul Jul 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should initialize one-time only. All deterministic transformations should be initialized only once. Future PR.

@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

HuggingFaceDocBuilderDev commented Jul 4, 2025

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.

linoytsaban
linoytsaban reviewed Jul 7, 2025
View reviewed changes
examples/dreambooth/train_dreambooth_lora_flux_kontext.py
Copy link
Copy Markdown
Collaborator

@linoytsaban linoytsaban Jul 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

diffusers/examples/dreambooth/train_dreambooth_lora_flux_kontext.py

Lines 1380 to 1393 in bce55a9

target_modules = [
"attn.to_k",
"attn.to_q",
"attn.to_v",
"attn.to_out.0",
"attn.add_k_proj",
"attn.add_q_proj",
"attn.add_v_proj",
"attn.to_add_out",
"ff.net.0.proj",
"ff.net.2",
"ff_context.net.0.proj",
"ff_context.net.2",
]

let's add proj_out,proj_mlp too here, seems to improve results and other trainers target these as well

Copy link
Copy Markdown
Member Author

@sayakpaul sayakpaul Jul 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But proj_out will also include the final output layer also right? 👁️

diffusers/src/diffusers/models/transformers/transformer_flux.py

Line 286 in 2c30287

self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)

Copy link
Copy Markdown
Collaborator

@linoytsaban linoytsaban Jul 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

indeed it will

Copy link
Copy Markdown
Member Author

@sayakpaul sayakpaul Jul 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, maybe let's just add proj_mlp for now given #11874?

Copy link
Copy Markdown
Collaborator

@linoytsaban linoytsaban Jul 8, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes good catch!

linoytsaban
linoytsaban reviewed Jul 7, 2025
View reviewed changes
linoytsaban
linoytsaban reviewed Jul 7, 2025
View reviewed changes
@sayakpaul
Copy link
Copy Markdown
Member Author

sayakpaul commented Jul 7, 2025 •
edited
Loading

@bot /style

@sayakpaul sayakpaul requested a review from linoytsaban July 7, 2025 14:30
linoytsaban
linoytsaban approved these changes Jul 8, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@linoytsaban linoytsaban left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks @sayakpaul 🚀

sayakpaul and others added 4 commits July 8, 2025 14:14
@sayakpaul
add proj_mlp to the mix
53bf41d
@linoytsaban
Update README_flux.md
9ac1bde 
add note on installing from commit `05e7a854d0a5661f5b433f6dd5954c224b104f0b`
@sayakpaul
fix
d2931e0
@sayakpaul
fix
2b4827f
@sayakpaul sayakpaul merged commit 01240fe into main Jul 8, 2025
11 of 12 checks passed
@sayakpaul sayakpaul deleted the kontext-i2i-training branch July 8, 2025 15:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@linoytsaban linoytsaban linoytsaban approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@sayakpaul @HuggingFaceDocBuilderDev @linoytsaban