#!/bin/sh

OUT_PATH="/home/ubuntu/out"
DATASET_PATH="/home/ubuntu/dataset"
PROMPT_PATH="/home/ubuntu/sotediffusion-prompt.txt"
WUER_PATH="/home/ubuntu/models/wuerstchen3"

CUDA_VISIBLE_DEVICES="1,2,3,4,5,6,7" accelerate launch --mixed_precision no --multi_gpu /home/ubuntu/sd-scripts/stable_cascade_train_stage_c.py \
--mixed_precision no \
--save_precision float \
--sdpa \
--gradient_checkpointing \
--train_text_encoder \
--resolution "1024,1024" \
--train_batch_size 12 \
--learning_rate 2e-6 \
--learning_rate_te1 1e-7 \
--lr_scheduler constant_with_warmup \
--lr_warmup_steps 100 \
--optimizer_type adamw8bit \
--token_warmup_min 1 \
--token_warmup_step 0 \
--caption_separator ", " \
--caption_dropout_rate 0.1 \
--caption_tag_dropout_rate 0 \
--caption_dropout_every_n_epochs 0 \
--dataset_repeats 1 \
--save_state \
--save_every_n_steps 1024 \
--sample_every_n_steps 1024 \
--max_token_length 225 \
--max_grad_norm 0 \
--max_train_epochs 10 \
--caption_extension ".txt" \
--max_data_loader_n_workers 32 \
--persistent_data_loader_workers \
--enable_bucket \
--min_bucket_reso 256 \
--max_bucket_reso 4096 \
--bucket_reso_steps 128 \
--bucket_no_upscale \
--log_with wandb \
--output_name sotediffusion-wr3_3b \
--train_data_dir $DATASET_PATH \
--in_json $DATASET_PATH.json \
--output_dir $OUT_PATH \
--logging_dir $OUT_PATH/logs \
--stage_c_checkpoint_path $WUER_PATH/sotediffusion-wr3_3b-base.safetensors \
--text_model_checkpoint_path $WUER_PATH/sotediffusion-wr3_text_model.safetensors \
--effnet_checkpoint_path $WUER_PATH/effnet_encoder.safetensors \
--previewer_checkpoint_path $WUER_PATH/previewer.safetensors \
--sample_prompts $PROMPT_PATH