File size: 2,935 Bytes
546ff88 0f961a6 546ff88 0f961a6 546ff88 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
"""
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()
|