setup_model.py

#!/usr/bin/env python3
"""
Download and cache TripoSR model for HF Spaces
Runs on first startup to fetch the model from HuggingFace
"""

import os
import sys
from pathlib import Path

def setup_triposr_model():
    """Download TripoSR model and setup cache directory"""
    
    # Setup cache directory
    cache_dir = Path.home() / ".cache" / "triposr"
    cache_dir.mkdir(parents=True, exist_ok=True)
    
    checkpoint_dir = cache_dir / "checkpoints"
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    
    # Check if model already exists
    model_ckpt = checkpoint_dir / "model.ckpt"
    config_file = checkpoint_dir / "config.yaml"
    
    if model_ckpt.exists() and config_file.exists():
        print("[INFO] TripoSR model already cached, skipping download")
        return str(checkpoint_dir)
    
    print("[INFO] Downloading TripoSR model from HuggingFace...")
    print("[INFO] This may take 2-5 minutes on first run...")
    
    try:
        from huggingface_hub import snapshot_download
        
        # Download from stabilityai/TripoSR
        model_dir = snapshot_download(
            repo_id="stabilityai/TripoSR",
            cache_dir=str(cache_dir),
            local_dir=str(checkpoint_dir),
            force_download=False,
            resume_download=True
        )
        
        print(f"[OK] TripoSR model downloaded to: {model_dir}")
        return str(checkpoint_dir)
        
    except ImportError:
        print("[ERROR] huggingface_hub not available")
        return None
    except Exception as e:
        print(f"[ERROR] Failed to download model: {e}")
        return None

if __name__ == "__main__":
    setup_triposr_model()