#!/bin/bash
# ============================================================
# RunPod Secure Cloud 一键引导：Qwen-Image-Edit-2511 脸部 LoRA 训练
# 用法：pod 启动后，在 web terminal 里执行：
#   cd /workspace && unzip dataset.zip -d train_dataset && bash bootstrap.sh
# ============================================================
set -e

echo "===== [1/4] 安装 musubi-tuner ====="
cd /workspace
if [ ! -d musubi-tuner ]; then
  git clone --depth 1 https://github.com/kohya-ss/musubi-tuner.git
fi
cd musubi-tuner
pip install -e . --quiet
pip install --quiet huggingface_hub hf_transfer
export HF_HUB_ENABLE_HF_TRANSFER=1

echo "===== [2/4] 下载模型（Qwen-Image-Edit-2511，机房网速约2-5分钟）====="
mkdir -p /workspace/models/qwen-edit-2511/{transformer,text_encoder}
hf download Qwen/Qwen-Image-Edit-2511 \
  --include "transformer/*" \
  --local-dir /workspace/models/qwen-edit-2511
hf download Qwen/Qwen-Image-Edit-2511 \
  --include "text_encoder/*" \
  --local-dir /workspace/models/qwen-edit-2511
hf download Qwen/Qwen-Image-Edit-2511 \
  --include "vae/diffusion_pytorch_model.safetensors" \
  --local-dir /workspace/models/qwen-edit-2511
# VAE 路径兼容（musubi 参数直接指到文件）
find /workspace/models/qwen-edit-2511 -name "*.safetensors" | head -20

echo "===== [3/4] 数据集配置 ====="
mkdir -p /workspace/train_dataset /workspace/config /workspace/ckpt
cat > /workspace/config/dataset.toml <<'EOF'
[general]
resolution = 1024
caption_extension = ".txt"
batch_size = 1
enable_bucket = true
bucket_no_upscale = false

[[datasets]]
image_directory = "/workspace/train_dataset"
num_repeats = 1
EOF
cp /workspace/sample_prompts.txt /workspace/config/ 2>/dev/null || true

echo "===== [3.5] 续炼 checkpoint 就位 ====="
# 续炼起点：本地 scp 上传的 myface_lora-000060.safetensors 在 /workspace/
# （1.1GB，dim32 版本），移到 /workspace/ckpt/ 供 --network_weights 使用
if [ -f /workspace/myface_lora-000060.safetensors ]; then
  cp /workspace/myface_lora-000060.safetensors /workspace/ckpt/
  echo "checkpoint 就位: $(ls -la /workspace/ckpt/myface_lora-000060.safetensors)"
else
  echo "WARNING: /workspace/myface_lora-000060.safetensors 不存在！续炼会失败（除非从头训练）"
fi

echo "===== [4/4] 预缓存（VAE latent + TextEncoder）====="
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)
echo "DIT: $DIT"
echo "TE:  $TE"
echo "VAE: $VAE"

python src/musubi_tuner/qwen_image_cache_latents.py \
  --dataset_config /workspace/config/dataset.toml \
  --vae "$VAE" --device cuda --model_version edit-2511

python src/musubi_tuner/qwen_image_cache_text_encoder_outputs.py \
  --dataset_config /workspace/config/dataset.toml \
  --text_encoder "$TE" --fp8_vl --device cuda --model_version edit-2511

echo ""
echo "===== 全部就绪！启动训练执行：bash train.sh ====="