[training ] add Kontext i2i training

[sayakpaul]

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]

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

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

[linoytsaban]

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

[sayakpaul]

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)

[linoytsaban]

indeed it will

[sayakpaul]

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

[linoytsaban]

yes good catch!

[sayakpaul]

@bot /style

[linoytsaban]

thanks @sayakpaul 🚀