#!/bin/bash

gpu_ids=0
dataset_config="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/Train/9F_Mix_S.toml"
output_dir="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/Output/9F_Hydra_S_new"
output_name="9F_Hydra_S"
# sample_prompts="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/Samples/9F_Mix_S.txt"

cache_dir="/tiamat-NAS/songyiren/FYP/liucheng/cache/"
export HF_HOME=$cache_dir

accelerate launch \
  --config_file "/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/accelerate_config.yaml" \
  --num_cpu_threads_per_process 1 \
  --gpu_ids $gpu_ids \
  flux_train_network_asylora.py \
  --dataset_config $dataset_config \
  --pretrained_model_name_or_path "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/unet/flux1-dev-fp8.safetensors" \
  --ae "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/vae/ae.safetensors" \
  --clip_l "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/clip/clip_l.safetensors" \
  --t5xxl "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/clip/t5xxl_fp8_e4m3fn.safetensors" \
  --optimizer_type came \
  --max_grad_norm 1.0 \
  --lr_scheduler constant \
  --lr_warmup_steps 0 \
  --lr_scheduler_num_cycles 1 \
  --lr_scheduler_power 1.0 \
  --min_snr_gamma 5 \
  --output_name $output_name \
  --output_dir $output_dir \
  --network_dim 64 \
  --network_alpha 1.0 \
  --learning_rate 1e-4 \
  --max_train_steps 50000 \
  --apply_t5_attn_mask \
  --cache_latents_to_disk \
  --cache_text_encoder_outputs \
  --cache_text_encoder_outputs_to_disk \
  --weighting_scheme logit_normal \
  --logit_mean 0 \
  --logit_std 1.0 \
  --mode_scale 1.29 \
  --timestep_sampling shift \
  --sigmoid_scale 1.0 \
  --model_prediction_type raw \
  --guidance_scale 1.0 \
  --discrete_flow_shift 3.1582 \
  --fp8_base \
  --highvram \
  --gradient_checkpointing \
  --seed 42 \
  --save_precision bf16 \
  --save_every_n_epochs 5 \
  --network_module networks.asylora_flux \
  --network_train_unet_only \
  --vae_batch_size 1 \
  --save_model_as safetensors \
  --max_data_loader_n_workers 0 \
  --mixed_precision bf16 \
  --skip_cache_check \
  --gradient_accumulation_steps 1 \
  --lora_ups_num 10 \
  --log_config