setup_spaces.py

#!/usr/bin/env python3
"""
Setup script for Hugging Face Spaces deployment
This script handles the installation of complex dependencies like nvdiffrast and PyTorch3D
"""

import os
import sys
import subprocess
import shutil
from pathlib import Path

def run_command(command, cwd=None):
    """Run a shell command and return the result"""
    print(f"Running: {command}")
    result = subprocess.run(command, shell=True, cwd=cwd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Error running command: {command}")
        print(f"Error output: {result.stderr}")
        return False
    print(f"Success: {command}")
    return True

def setup_nvdiffrast():
    """Install nvdiffrast"""
    print("Setting up nvdiffrast...")
    
    # Create packages directory if it doesn't exist
    packages_dir = Path("packages")
    packages_dir.mkdir(exist_ok=True)
    
    # Clone nvdiffrast
    if not (packages_dir / "nvdiffrast").exists():
        if not run_command("git clone https://github.com/NVlabs/nvdiffrast.git", cwd=packages_dir):
            return False
    
    # Install nvdiffrast
    nvdiffrast_dir = packages_dir / "nvdiffrast"
    if not run_command("pip install .", cwd=nvdiffrast_dir):
        return False
    
    return True

def setup_pytorch3d():
    """Install PyTorch3D"""
    print("Setting up PyTorch3D...")
    
    packages_dir = Path("packages")
    
    # Clone PyTorch3D
    if not (packages_dir / "pytorch3d").exists():
        if not run_command("git clone https://github.com/facebookresearch/pytorch3d.git", cwd=packages_dir):
            return False
    
    # Install PyTorch3D
    pytorch3d_dir = packages_dir / "pytorch3d"
    
    # Set environment variable for Linux
    env = os.environ.copy()
    env['DISTUTILS_USE_SDK'] = '1'
    
    # Try to install with pip first
    if not run_command("pip install .", cwd=pytorch3d_dir):
        # Fallback to setup.py
        if not run_command("python setup.py install", cwd=pytorch3d_dir):
            return False
    
    return True

def setup_fashion_clip():
    """Setup Fashion-CLIP"""
    print("Setting up Fashion-CLIP...")
    
    packages_dir = Path("packages")
    
    # Clone Fashion-CLIP if not already present
    if not (packages_dir / "fashion-clip").exists():
        if not run_command("git clone https://github.com/patrickjohncyh/fashion-clip.git", cwd=packages_dir):
            return False
    
    # Install Fashion-CLIP dependencies
    fashion_clip_dir = packages_dir / "fashion-clip"
    dependencies = ["appdirs", "boto3", "annoy", "validators", "transformers", "datasets"]
    
    for dep in dependencies:
        if not run_command(f"pip install {dep}", cwd=fashion_clip_dir):
            print(f"Warning: Failed to install {dep}")
    
    return True

def install_basic_dependencies():
    """Install basic Python dependencies"""
    print("Installing basic dependencies...")
    
    # Install PyTorch with CUDA support
    if not run_command("pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118"):
        return False
    
    # Install other dependencies
    if not run_command("pip install -r requirements.txt"):
        return False
    
    # Install CLIP
    if not run_command("pip install git+https://github.com/openai/CLIP.git"):
        return False
    
    return True

def create_symlinks():
    """Create necessary symlinks for the packages"""
    print("Creating package symlinks...")
    
    # Add packages to Python path
    packages_dir = Path("packages")
    
    # Create __init__.py files if they don't exist
    for package_dir in packages_dir.iterdir():
        if package_dir.is_dir():
            init_file = package_dir / "__init__.py"
            if not init_file.exists():
                init_file.touch()
    
    return True

def main():
    """Main setup function"""
    print("Starting setup for Hugging Face Spaces deployment...")
    
    # Install basic dependencies first
    if not install_basic_dependencies():
        print("Failed to install basic dependencies")
        sys.exit(1)
    
    # Setup complex dependencies
    if not setup_nvdiffrast():
        print("Failed to setup nvdiffrast")
        sys.exit(1)
    
    if not setup_pytorch3d():
        print("Failed to setup PyTorch3D")
        sys.exit(1)
    
    if not setup_fashion_clip():
        print("Failed to setup Fashion-CLIP")
        sys.exit(1)
    
    # Create symlinks
    if not create_symlinks():
        print("Failed to create symlinks")
        sys.exit(1)
    
    print("Setup completed successfully!")
    print("You can now run the Gradio interface with: python app.py")

if __name__ == "__main__":
    main()