Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """Download every model the service uses into the image, at build time. | |
| Run by the Dockerfile. Without it, each model downloaded the first time it was used after the Space | |
| started, and with a 15-minute sleep the Space starts often: the first film after a wake waited | |
| minutes for gigabytes of weights. HOME and HF_HOME point into the image (see the Dockerfile), so | |
| the caches written here are the ones the service reads at run time. | |
| Downloads only, wherever a download can be had without building the model: the build machine has | |
| far less memory than the GPU it runs on, and loading emotion2vec here got the build OOM-killed. | |
| Each model is fetched in its own process, so whatever one leaves in memory is gone before the next. | |
| A failure is reported, not fatal: a model missing from the image still downloads on first use. | |
| """ | |
| import os | |
| import subprocess | |
| import sys | |
| os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") # the build has no GPU | |
| def whisper(size): | |
| from faster_whisper.utils import download_model | |
| download_model(size) | |
| def hub(repo, **kw): | |
| from huggingface_hub import snapshot_download | |
| snapshot_download(repo, **kw) | |
| def hsemotion(): | |
| from hsemotion_onnx.facial_emotions import get_model_path | |
| get_model_path("enet_b0_8_best_afew") | |
| def easyocr_models(): | |
| # Builds the reader (small: the text detector and the English recogniser), which is how | |
| # EasyOCR downloads its weights. | |
| import easyocr | |
| easyocr.Reader(["en"], gpu=False, verbose=False) | |
| def memorability_models(): | |
| from huggingface_hub import hf_hub_download | |
| hf_hub_download("PerceptCLIP/PerceptCLIP_Memorability", "modeling.py") | |
| hf_hub_download("PerceptCLIP/PerceptCLIP_Memorability", "perceptCLIP_Memorability.pth") | |
| hub("openai/clip-vit-large-patch14", allow_patterns=["*.json", "*.txt", "*.safetensors"]) | |
| hub("sentence-transformers/all-MiniLM-L6-v2") | |
| def deepface_emotion(): | |
| # DeepFace's emotion model is small; building it is how DeepFace downloads it. | |
| import numpy as np | |
| from deepface import DeepFace | |
| DeepFace.analyze(np.zeros((64, 64, 3), dtype=np.uint8), actions=["emotion"], enforce_detection=False, silent=True) | |
| JOBS = { | |
| "whisper medium": lambda: whisper("medium"), | |
| "whisper small": lambda: whisper("small"), | |
| "whisper medium.en (Stage Read)": lambda: whisper("medium.en"), | |
| "whisper small.en (Stage Read)": lambda: whisper("small.en"), | |
| # funasr with hub="hf" reads this from the Hugging Face cache | |
| "emotion2vec": lambda: hub("emotion2vec/emotion2vec_plus_large"), | |
| "HSEmotion": hsemotion, | |
| "EasyOCR": easyocr_models, | |
| "GoEmotions": lambda: hub("SamLowe/roberta-base-go_emotions"), | |
| "memorability (CLIP, PerceptCLIP, MiniLM)": memorability_models, | |
| "DeepFace emotion": deepface_emotion, | |
| "wav2vec2 (Stage Read)": lambda: hub("audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim"), | |
| } | |
| if __name__ == "__main__": | |
| if len(sys.argv) > 1: | |
| JOBS[sys.argv[1]]() | |
| sys.exit(0) | |
| failed = [] | |
| for name in JOBS: | |
| print(f"[prefetch] {name}...", flush=True) | |
| if subprocess.run([sys.executable, __file__, name]).returncode == 0: | |
| print(f"[prefetch] {name} done", flush=True) | |
| else: | |
| print(f"[prefetch] {name} FAILED (it will download on first use instead)", flush=True) | |
| failed.append(name) | |
| print(f"[prefetch] finished; {len(failed)} failed: {failed}", flush=True) | |