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