#!/bin/bash
# ============================================================
# 正式训练（RunPod Secure Cloud 4090 24G）
# 用法：nohup bash train.sh > /workspace/train.log 2>&1 &
# 特性：崩溃自动重启（进程退出后自动拉起续跑，直到完成或 epoch120 checkpoint 出现）
# ============================================================
cd /workspace/musubi-tuner
MODELS=/workspace/models/qwen-edit-2511
DIT=$MODELS/transformer/$(ls $MODELS/transformer | grep '00001-of' | head -1)
TE=$MODELS/text_encoder/$(ls $MODELS/text_encoder | grep '00001-of' | head -1)
VAE=$(find $MODELS/vae -name "*.safetensors" | head -1)

TRAIN_CMD="accelerate launch --num_cpu_threads_per_process 2 --mixed_precision bf16 \
  src/musubi_tuner/qwen_image_train_network.py \
  --dit \"$DIT\" --vae \"$VAE\" --text_encoder \"$TE\" \
  --model_version edit-2511 \
  --dataset_config /workspace/config/dataset.toml \
  --sdpa --mixed_precision bf16 \
  --timestep_sampling shift --weighting_scheme none --discrete_flow_shift 2.2 \
  --optimizer_type adamw8bit --learning_rate 1e-4 \
  --gradient_checkpointing \
  --network_module networks.lora_qwen_image --network_dim 32 \\
  --network_weights /workspace/ckpt/myface_lora-000060.safetensors \
  --fp8_base --fp8_scaled --fp8_vl --blocks_to_swap 8 \
  --max_train_epochs 50 --save_every_n_epochs 10 --seed 42 \
  --sample_prompts /workspace/config/sample_prompts.txt --sample_every_n_epochs 10 \
  --output_dir /workspace/ckpt --output_name myface_v7"

# 崩溃自动重启：最多 20 次（防止死循环烧钱）
for attempt in $(seq 1 20); do
  echo "===== [尝试 $attempt/20] 启动训练 $(date) ====="
  eval "$TRAIN_CMD"
  rc=$?
  echo "===== 训练退出 rc=$rc $(date) ====="
  # 完成检测：epoch120 checkpoint 已存在 → 训练成功结束
  if ls /workspace/ckpt/myface_v7-000050.safetensors >/dev/null 2>&1; then
    echo "===== epoch120 checkpoint 已生成，训练完成 ====="
    break
  fi
  # 中途崩溃 → 等待 30s 后自动重启续跑
  echo "===== 训练中断（rc=$rc），30 秒后自动重启续跑 ====="
  sleep 30
done
echo "===== train.sh 结束 $(date) ====="