train/train_kontext_local.sh

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

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --config_file $CONFIG train_kontext_qy.py \
    --pretrained_model_name_or_path $MODEL_DIR \
    --pretrained_lora_path "" \
    --lora_num=1 \
    --cond_size=512 \
    --ranks 128 \
    --network_alphas 128 \
    --output_dir=$OUTPUT_DIR \
    --logging_dir=$LOG_PATH \
    --mixed_precision="bf16" \
    --learning_rate=1e-4 \
    --train_batch_size=1 \
    --num_train_epochs=1 \
    --validation_steps=250 \
    --checkpointing_steps=1000 \
    --validation_images "./kontext_local_test/img_1.png"  \
    --spatial_test_images "./kontext_local_test/mask_1.png" \
    --validation_prompt "convert the dinosaur into blue color" \
    --gradient_checkpointing \
    --blend_pixel_values \
    --num_validation_images=1