#!/usr/bin/env bash
# Run Wan2.2 smoke test on local GPU

set -eo pipefail

echo "=== Wan2.2 Smoke Test ==="
echo ""

# Activate GPU environment
source tools/activate-gpu-env.sh
echo ""

# Check if DiffSynth-Studio is installed
if ! python3 -c "import diffsynth" 2>/dev/null; then
  echo "DiffSynth-Studio not found. Install it first:"
  echo "  git clone https://github.com/modelscope/DiffSynth-Studio.git ~/src/DiffSynth-Studio"
  echo "  cd ~/src/DiffSynth-Studio && pip install -e ."
  echo ""
  exit 1
fi

# Check if model is downloaded
WAN22_MODEL="${WAN22_MODEL:-$HOME/models/Wan2.2-TI2V-5B}"
if [ ! -d "$WAN22_MODEL" ]; then
  echo "Wan2.2 model not found at $WAN22_MODEL"
  echo ""
  echo "Download it first:"
  echo "  huggingface-cli login"
  echo "  huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir $WAN22_MODEL"
  echo ""
  exit 1
fi

# Run dry-run first
echo "Running metadata validation (dry-run)..."
python3 tools/run_wan22_train.py \
  --training-data-dir materials/training-data/iiw-english-smoke-video-only \
  --output-dir materials/training-data/iiw-english-smoke-video-only/wan22_checkpoints \
  --model-variant ti2v-5b \
  --lora-rank 16 \
  --epochs 1 \
  --dataset-repeat 20 \
  --learning-rate 2e-5 \
  --num-frames 81 \
  --height 480 \
  --width 832 \
  --gradient-accumulation-steps 4 \
  --dry-run

echo ""
echo "=== Dry-run successful ==="
echo ""
echo "To run actual training:"
echo "  python3 tools/run_wan22_train.py \\"
echo "    --training-data-dir materials/training-data/iiw-english-smoke-video-only \\"
echo "    --output-dir materials/training-data/iiw-english-smoke-video-only/wan22_checkpoints \\"
echo "    --diffsynth-path \$HOME/src/DiffSynth-Studio \\"
echo "    --wan22-model $WAN22_MODEL \\"
echo "    --model-variant ti2v-5b \\"
echo "    --lora-rank 16 \\"
echo "    --epochs 1 \\"
echo "    --dataset-repeat 20 \\"
echo "    --learning-rate 2e-5 \\"
echo "    --num-frames 81 \\"
echo "    --height 480 \\"
echo "    --width 832 \\"
echo "    --gradient-accumulation-steps 4"