[training ] add Kontext i2i training by sayakpaul · Pull Request #11858 · huggingface/diffusers (original) (raw)

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.