#!/bin/bash

export CUDA_VISIBLE_DEVICES=0,1
NODE='node-name'
NAME='train'

# ** Training Parameters **
DATASET="coco"
BATCHSIZE=2
NUM_SAMPLES=10
NUM_ITER=10

# ** Save Parameter**
LOGDIR="./data"
DATADIR="/mnt/data1/jaayeon/data" # change data path

echo "Train SD1.5 start"
# srun -w $NODE --job-name=$NAME \
# python train.py --model sd1.5 --dataset $DATASET --batch_size $BATCHSIZE --separate_gpus \
#  --datadir $DATADIR --logdir $LOGDIR --num_iter $NUM_ITER --num_eval $NUM_SAMPLES \
#  --n_dc_tokens 4 --apply_dc True False False --use_dc_t --epochs 2 --dweight 0.01

echo "Train SDXL start"
# srun -w $NODE --job-name=$NAME \
# python train.py --model sdxl --dataset $DATASET --batch_size $BATCHSIZE --separate_gpus \
#  --datadir $DATADIR --logdir $LOGDIR --num_iter $NUM_ITER --num_eval $NUM_SAMPLES \
#  --n_dc_tokens 4 --apply_dc True True False --epochs 2 --dweight 0.01

echo "Train SD3 start"
# srun -w $NODE --job-name=$NAME 
python train.py --model sd3 --dataset $DATASET --batch_size $BATCHSIZE --separate_gpus \
 --datadir $DATADIR --logdir $LOGDIR --num_iter $NUM_ITER --num_eval $NUM_SAMPLES \
 --n_dc_tokens 4 --n_dc_layers 5 --epochs 2 --use_dc_t --dweight 0 

echo "All epochs completed!"