# accelerate launch train_instruct_pix2pix.py  \
#     --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" \
#     --train_data_dir="/home/AURA/dataset" \
#     --train_data_prompt="/home/AURA/dataset/json-text2layout/train_prompt_text2layout_single.json" \
#     --use_ema \
#     --resolution=512 \
#     --train_batch_size=32 \
#     --gradient_accumulation_steps=1 \
#     --gradient_checkpointing  \
#     --max_train_steps=10000  \
#     --checkpointing_steps=500  \
#     --learning_rate=5e-05 \
#     --lr_warmup_steps=0  \
#     --mixed_precision=fp16  \
#     --seed=42 \
#     --report_to=wandb \
#     --validation_image \
#         "/home/AURA/dataset/test/40759/organ_40759.png" \
#         "/home/AURA/dataset/test/40759/organ_40759.png" \
#         "/home/AURA/dataset/test/40775/organ_40775.png" \
#         "/home/AURA/dataset/test/40775/organ_40775.png" \
#     --validation_prompt \
#         "A Chest X-ray photo with mild Effusion on biliteral lung." \
#         "A Chest X-ray photo with moderate Effusion on right lower lung." \
#         "A Chest X-ray photo with severe Cardiomegaly on heart." \
#         "A Chest X-ray photo with mild Pneumothorax on left upper lung."
  

accelerate launch train_instruct_pix2pix_rib.py  \
    --pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4" \
    --train_data_dir="/home/shuhan/blobdata_sd/AURAD_dataset" \
    --train_data_prompt="/home/shuhan/blobdata_sd/AURAD_dataset/train_prompt_text2layout_total_resample.json" \
    --use_ema \
    --resolution=512 \
    --train_batch_size=32 \
    --gradient_accumulation_steps=1 \
    --gradient_checkpointing  \
    --max_train_steps=10000  \
    --checkpointing_steps=500  \
    --learning_rate=5e-05 \
    --lr_warmup_steps=0  \
    --mixed_precision=fp16  \
    --seed=42 \
    --report_to=wandb \
    --validation_json="/home/shuhan/blobdata_sd/AURAD_dataset/log_validation_text2layout.json" \
    --validation_epochs=10