#!/usr/bin/env python3
"""
Setup script for HuggingFace Spaces deployment
Downloads required models and assets for MIMO
"""

import os
import sys
from pathlib import Path
from huggingface_hub import snapshot_download, hf_hub_download

def setup_hf_spaces():
    """Setup models and assets for HF Spaces"""

    print("🚀 Setting up MIMO for HuggingFace Spaces...")

    # Create directories
    os.makedirs("./models", exist_ok=True)
    os.makedirs("./assets", exist_ok=True)

    try:
        # 1. Download MIMO models
        print("📥 Downloading MIMO models...")
        snapshot_download(
            repo_id="menyifang/MIMO",
            cache_dir="./models",
            allow_patterns=["*.pth", "*.json", "*.md"]
        )
        print("✅ MIMO models downloaded")

        # 2. Download base models
        print("📥 Downloading Stable Diffusion v1.5...")
        snapshot_download(
            repo_id="runwayml/stable-diffusion-v1-5",
            cache_dir="./models/stable-diffusion-v1-5"
        )
        print("✅ Stable Diffusion downloaded")

        print("📥 Downloading VAE...")
        snapshot_download(
            repo_id="stabilityai/sd-vae-ft-mse",
            cache_dir="./models/sd-vae-ft-mse"
        )
        print("✅ VAE downloaded")

        print("📥 Downloading image encoder...")
        snapshot_download(
            repo_id="lambdalabs/sd-image-variations-diffusers",
            cache_dir="./models/image_encoder",
            allow_patterns=["image_encoder/**"]
        )
        print("✅ Image encoder downloaded")

        # 3. Download human segmenter
        print("📥 Downloading human segmenter...")
        hf_hub_download(
            repo_id="menyifang/MIMO",
            filename="matting_human.pb",
            cache_dir="./assets",
            local_dir="./assets"
        )
        print("✅ Human segmenter downloaded")

        # 4. Create minimal assets
        print("📁 Creating asset directories...")
        os.makedirs("./assets/test_image", exist_ok=True)
        os.makedirs("./assets/masks", exist_ok=True)

        print("🎉 HuggingFace Spaces setup complete!")
        return True

    except Exception as e:
        print(f"❌ Setup failed: {e}")
        return False

if __name__ == "__main__":
    setup_hf_spaces()