#!/bin/bash

export CUDA_HOME=/usr/local/cuda
export HF_HUB_CACHE=/tiamat-vePFS/share_data/hf_home/hub

export CUDA_VISIBLE_DEVICES=0
export LORA_RANK=16
export BATCH_SIZE=2
export DATA_CONFIG="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/MergeModel/6_Portrait/6_Portrait.toml"
export OUTPUT_DIR="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/RecraftModel/6_Portrait_makeanything/"
export OUTNAME=test
export LOG_DIR=$OUTPUT_DIR/logs
export MAX_TRAIN_EPOCHS=50
export MAX_TRAIN_STEPS=500
export SAVE_EVERY_N_EPOCHS=1
export SAVE_EVERY_N_STEPS=50

accelerate launch --config_file="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/accelerate_config.yaml" \
  --main_process_port=23322 --mixed_precision=bf16 --num_cpu_threads_per_process=1 flux_train_recraft.py \
  --pretrained_model_name_or_path="/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/MergeModel/6_Portrait/6_Portrait.safetensors" \
  --clip_l=/tiamat-NAS/hailong/storage_backup/models/stabilityai/stable-diffusion-3-medium/text_encoders/clip_l.safetensors \
  --t5xxl=/tiamat-NAS/hailong/storage_backup/models/stabilityai/stable-diffusion-3-medium/text_encoders/t5xxl_fp16.safetensors \
  --ae=/tiamat-vePFS/share_data/storage/huggingface/models/black-forest-labs/FLUX.1-dev/ae.safetensors \
  --max_token_length=225 --apply_t5_attn_mask \
  --dataset_config=$DATA_CONFIG \
  --cache_latents --vae_batch_size=4 --cache_latents_to_disk --cache_text_encoder_outputs --cache_text_encoder_outputs_to_disk --skip_cache_check \
  --enable_bucket --min_bucket_reso=256 --max_bucket_reso=2304 --bucket_reso_steps=64 --bucket_no_upscale \
  --output_dir=$OUTPUT_DIR --output_name=$OUTNAME --logging_dir=$LOG_DIR --log_with=tensorboard --log_prefix="" --log_tracker_name=tryon \
  --save_every_n_epochs=$SAVE_EVERY_N_EPOCHS --save_every_n_steps=$SAVE_EVERY_N_STEPS --save_precision=bf16 --save_model_as=safetensors \
  --max_train_epochs=$MAX_TRAIN_EPOCHS --max_train_steps=$MAX_TRAIN_STEPS \
  --initial_epoch=0 --initial_step=0 \
  --train_batch_size=$BATCH_SIZE --max_data_loader_n_workers=2 --persistent_data_loader_workers --mixed_precision=bf16 --fp8_base \
  --mem_eff_attn --sdpa --gradient_checkpointing --gradient_accumulation_steps=1 \
  --seed=42 --clip_skip=2 --noise_offset=0.0375 --loss_type=l2 --adaptive_noise_scale=0.0375 \
  --learning_rate=1 --unet_lr=1 --text_encoder_lr=1 \
  --optimizer_type=Prodigy --optimizer_args "weight_decay=0.01" "betas=.9,.99" "decouple=True" "use_bias_correction=True" "d_coef=0.5" "d0=1e-4" \
  --lr_scheduler=cosine_with_restarts --lr_scheduler_num_cycles=1 --lr_decay_steps=160 --lr_scheduler_min_lr_ratio=0.1 \
  --network_module=networks.lora_flux --network_dim=$LORA_RANK --network_alpha=$LORA_RANK --network_train_unet_only \
  --guidance_scale=1 --timestep_sampling=flux_shift --discrete_flow_shift 3.1582 \
  --model_prediction_type=raw \
  --sample_every_n_steps 5 \
  --sample_at_first \
  --sample_images "/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/MergeModel/test/image.txt" \
  --sample_prompts "/tiamat-NAS/songyiren/FYP/liucheng/sd-scripts/MergeModel/test/caption.txt" \
  --frame_num 9