export MODEL_DIR="" # your flux path
export OUTPUT_DIR=""  # your save path
export CONFIG="./default_config.yaml"
export TRAIN_DATA="" # your data jsonl file
export LOG_PATH="$OUTPUT_DIR/log"

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --config_file $CONFIG train_kontext_color.py \
    --pretrained_model_name_or_path $MODEL_DIR \
    --lora_num=1 \
    --cond_size=512 \
    --ranks 128 \
    --network_alphas 128 \
    --output_dir=$OUTPUT_DIR \
    --logging_dir=$LOG_PATH \
    --mixed_precision="bf16" \
    --train_data_dir=$TRAIN_DATA \
    --learning_rate=1e-4 \
    --train_batch_size=1 \
    --num_train_epochs=1 \
    --validation_steps=100 \
    --checkpointing_steps=1000 \
    --validation_images "./kontext_color_test/img_1.png" \
    --spatial_test_images "./kontext_color_test/color_1.png"  \
    --validation_prompt "Let this woman have red purple and blue hair" \
    --num_validation_images=1