#!/usr/bin/env bash
set -euo pipefail

# ============================================================
# ComfyUI + Wan 2.1 T2V 1.3B setup for Google Colab
# venv: none
# ============================================================

COMFY_DIR="/content/ComfyUI"
WAN_REPO="Comfy-Org/Wan_2.1_ComfyUI_repackaged"

echo "=== 1. Install basic packages ==="

apt-get update
apt-get install -y git

echo "=== 2. Install Hugging Face CLI ==="

python -m pip install -U huggingface_hub

echo "=== 3. Clone ComfyUI ==="

if [ ! -d "$COMFY_DIR/.git" ]; then
    git clone https://github.com/Comfy-Org/ComfyUI.git "$COMFY_DIR"
else
    echo "ComfyUI already exists. Skip clone."
fi

cd "$COMFY_DIR"

echo "=== 4. Install ComfyUI dependencies ==="

python -m pip install -r requirements.txt

echo "=== 5. Create model directories ==="

mkdir -p \
    "$COMFY_DIR/models/diffusion_models" \
    "$COMFY_DIR/models/text_encoders" \
    "$COMFY_DIR/models/vae"

TMP_DIR="/content/wan_models"
mkdir -p "$TMP_DIR"

echo "=== 6. Download Wan 2.1 T2V 1.3B diffusion model ==="

hf download \
    "$WAN_REPO" \
    split_files/diffusion_models/wan2.1_t2v_1.3B_fp16.safetensors \
    --local-dir "$TMP_DIR"

echo "=== 7. Download UMT5 XXL FP8 text encoder ==="

hf download \
    "$WAN_REPO" \
    split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors \
    --local-dir "$TMP_DIR"

echo "=== 8. Download Wan 2.1 VAE ==="

hf download \
    "$WAN_REPO" \
    split_files/vae/wan_2.1_vae.safetensors \
    --local-dir "$TMP_DIR"

echo "=== 9. Install models into ComfyUI ==="

cp \
    "$TMP_DIR/split_files/diffusion_models/wan2.1_t2v_1.3B_fp16.safetensors" \
    "$COMFY_DIR/models/diffusion_models/"

cp \
    "$TMP_DIR/split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" \
    "$COMFY_DIR/models/text_encoders/"

cp \
    "$TMP_DIR/split_files/vae/wan_2.1_vae.safetensors" \
    "$COMFY_DIR/models/vae/"

echo "=== 10. Verify models ==="

ls -lh \
    "$COMFY_DIR/models/diffusion_models/wan2.1_t2v_1.3B_fp16.safetensors" \
    "$COMFY_DIR/models/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" \
    "$COMFY_DIR/models/vae/wan_2.1_vae.safetensors"

echo
echo "============================================================"
echo "Setup complete!"
echo
echo "ComfyUI:"
echo "  $COMFY_DIR"
echo
echo "Start with:"
echo "  cd $COMFY_DIR"
echo "  python main.py --listen 0.0.0.0 --port 8188"
echo "============================================================"