#!/bin/bash

export CUDA_VISIBLE_DEVICES="0,1,2,3"

PRETRAINED_MODEL="/gemini/space/hsd/project/pretrained_model/huggingface/hub/models--stabilityai--stable-diffusion-3-medium-diffusers/snapshots/ea42f8cef0f178587cf766dc8129abd379c90671"
LORA_PATH="/gemini/space/gzy_new/models/Sida/sd3-lora-finetuned-batch-4/checkpoint-500000"
RECTIFIED_WEIGHTS="/gemini/space/gzy_new/models/Sida/rectified-noise-batch-2/checkpoint-220000/sit_weights"

CAPTIONS_JSONL="/gemini/space/hsd/project/dataset/cc3m-wds/validation/metadata.jsonl"
SAMPLE_DIR="./sd3_lora_rn_pair_samples"

NUM_INFERENCE_STEPS=40
GUIDANCE_SCALE=7.0
HEIGHT=512
WIDTH=512
PER_PROC_BATCH_SIZE=1
IMAGES_PER_CAPTION=1
MAX_SAMPLES=500
GLOBAL_SEED=42
MIXED_PRECISION="fp16"
NUM_SIT_LAYERS=1

ARGS=(
  --pretrained_model_name_or_path "$PRETRAINED_MODEL"
  --captions_jsonl "$CAPTIONS_JSONL"
  --sample_dir "$SAMPLE_DIR"
  --num_inference_steps $NUM_INFERENCE_STEPS
  --guidance_scale $GUIDANCE_SCALE
  --height $HEIGHT
  --width $WIDTH
  --per_proc_batch_size $PER_PROC_BATCH_SIZE
  --images_per_caption $IMAGES_PER_CAPTION
  --max_samples $MAX_SAMPLES
  --global_seed $GLOBAL_SEED
  --num_sit_layers $NUM_SIT_LAYERS
  --mixed_precision $MIXED_PRECISION
  --rectified_weights "$RECTIFIED_WEIGHTS"
)

if [ -n "$LORA_PATH" ]; then
  ARGS+=(--lora_path "$LORA_PATH")
fi

torchrun --nproc_per_node=4 --master_port=25923 sample_sd3_lora_rn_pair_ddp.py "${ARGS[@]}" --stage lora
torchrun --nproc_per_node=4 --master_port=25924 sample_sd3_lora_rn_pair_ddp.py "${ARGS[@]}" --stage rn
torchrun --nproc_per_node=4 --master_port=25925 sample_sd3_lora_rn_pair_ddp.py "${ARGS[@]}" --stage pair

echo "Sampling done. Output at: $SAMPLE_DIR"
# nohup bash run_sd3_lora_rn_pair_sampling.sh > run_sd3_lora_rn_pair_sampling.log 2>&1 &