CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
    --deepspeed_config_file "ds_config.json" \
    train_flux_kontext_lora.py \
    --pretrained_model_name_or_path="flux-kontext-ckpt" \
    --dataset_name="raulc0399/open_pose_controlnet" \
    --output_dir="long-context-flux" \
    --mixed_precision="bf16" \
    --train_batch_size=1 \
    --rank=16 \
    --gradient_accumulation_steps=4 \
    --gradient_checkpointing \
    --learning_rate=1e-4 \
    --report_to="wandb" \
    --lr_scheduler="constant" \
    --lr_warmup_steps=0 \
    --max_train_steps=5000 \
    --validation_image="openpose.png" \
    --validation_prompt="A couple, 4k photo, highly detailed" \
    --offload \
    --seed="0" \
    --push_to_hub \
    --use_8bit_adam