File size: 2,311 Bytes
6f2c7f0 |
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 |
#!/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() |