app.py

"""
Three-View-Style-Embedder - Hugging Face Space App
Running on ZeroGPU with automatic model downloading.
"""
import os
import sys
from pathlib import Path

# Ensure dependencies are available
try:
    import spaces
    import huggingface_hub
except ImportError:
    print("Installing dependencies...")
    os.system("pip install spaces huggingface_hub")
    import spaces
    from huggingface_hub import hf_hub_download
else:
    from huggingface_hub import hf_hub_download

from inference_utils import StyleEmbedderApp

# Configuration
REPO_ID = "ij/Three-View-Style-Embedder-Combined"
MODEL_FILENAME = "best_model.pt"
EMBEDDINGS_FILENAME = "artist_embeddings.npz"
YOLO_WEIGHTS_FILENAME = "yolov5x_anime.pt"

def download_files():
    print(f"Downloading models from {REPO_ID}...")
    
    # Download Main Model
    try:
        model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
        print(f"Model downloaded to {model_path}")
    except Exception as e:
        raise RuntimeError(f"Failed to download model: {e}")

    # Download Embeddings
    try:
        embeddings_path = hf_hub_download(repo_id=REPO_ID, filename=EMBEDDINGS_FILENAME)
        print(f"Embeddings downloaded to {embeddings_path}")
    except Exception as e:
        raise RuntimeError(f"Failed to download embeddings: {e}")

    # Download YOLO Weights (Optional/Check)
    yolo_path = None
    try:
        # Check if local file exists (if uploaded with repo code)
        if os.path.exists(YOLO_WEIGHTS_FILENAME):
            yolo_path = os.path.abspath(YOLO_WEIGHTS_FILENAME)
            print(f"Using local YOLO weights: {yolo_path}")
        else:
            # Try downloading
            print("Downloading YOLO weights...")
            yolo_path = hf_hub_download(repo_id=REPO_ID, filename=YOLO_WEIGHTS_FILENAME)
            print(f"YOLO weights downloaded to {yolo_path}")
    except Exception as e:
        print(f"Warning: Failed to download YOLO weights ({e}). Will attempt to use default/local if available.")
        # If download fails, we pass None to let StyleEmbedderApp try its default logic
        yolo_path = None
        
    return model_path, embeddings_path, yolo_path

def main():
    # Download large files
    model_path, embeddings_path, yolo_path = download_files()
    
    # Initialize App
    # Model loading is lazy - happens inside @spaces.GPU decorated function
    # This avoids CUDA initialization in main process
    print("Initializing Application...")
    app = StyleEmbedderApp(
        checkpoint_path=model_path,
        embeddings_path=embeddings_path,
        device='cuda',  # Will be used when model loads (inside @spaces.GPU context)
        detector_device='cpu',  # Always use CPU for detector to avoid CUDA init
        yolo_weights=yolo_path
    )
    
    # Launch UI
    print("Launching UI...")
    demo = app.create_ui()
    demo.launch()

if __name__ == "__main__":
    main()