Spaces:
Running
Running
| """ | |
| Called during Docker build to pre-download DeepFace AI models. | |
| Tries Hugging Face Hub first (fast & reliable within HF infrastructure), | |
| then falls back to DeepFace's built-in downloader (GitHub releases). | |
| Models are baked into the image so no runtime download is needed. | |
| """ | |
| import os | |
| import sys | |
| from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout | |
| # Use DEEPFACE_HOME from env (set in Dockerfile) or fallback to user home | |
| DEEPFACE_HOME = os.environ.get("DEEPFACE_HOME", os.path.expanduser("~")) | |
| WEIGHTS_DIR = os.path.join(DEEPFACE_HOME, ".deepface", "weights") | |
| os.makedirs(WEIGHTS_DIR, exist_ok=True) | |
| HF_REPO = "Hodfa71/deepface-weights" | |
| MODEL_FILES = [ | |
| "vgg_face_weights.h5", | |
| "age_model_weights.h5", | |
| "gender_model_weights.h5", | |
| "facial_expression_model_weights.h5", | |
| "race_model_single_batch.h5", | |
| "retinaface.h5", | |
| ] | |
| def try_hf_hub(): | |
| """Download weights from Hodfa71/deepface-weights on Hugging Face Hub.""" | |
| try: | |
| from huggingface_hub import hf_hub_download, list_repo_files | |
| # Check the repo exists and has files | |
| files = list(list_repo_files(HF_REPO)) | |
| if not files: | |
| print(f" HF repo {HF_REPO} is empty or not found.", flush=True) | |
| return False | |
| for filename in MODEL_FILES: | |
| dest = os.path.join(WEIGHTS_DIR, filename) | |
| if os.path.exists(dest): | |
| print(f" {filename} already present, skipping.", flush=True) | |
| continue | |
| if filename not in files: | |
| print(f" {filename} not in HF repo, skipping.", flush=True) | |
| continue | |
| print(f" Downloading {filename} from HF Hub...", flush=True) | |
| with ThreadPoolExecutor(max_workers=1) as executor: | |
| future = executor.submit( | |
| hf_hub_download, | |
| repo_id=HF_REPO, | |
| filename=filename, | |
| local_dir=WEIGHTS_DIR, | |
| local_dir_use_symlinks=False, | |
| ) | |
| future.result(timeout=300) # 5 minute cap per file | |
| print(f" {filename} ready.", flush=True) | |
| return True | |
| except Exception as exc: | |
| print(f" HF Hub download failed: {exc}", flush=True) | |
| return False | |
| def try_deepface_build(): | |
| """Download weights via DeepFace (from GitHub releases).""" | |
| try: | |
| from deepface import DeepFace | |
| for model_name in ["Age", "Gender", "Emotion", "Race"]: | |
| print(f" Building {model_name} via DeepFace (GitHub)...", flush=True) | |
| DeepFace.build_model(model_name) | |
| print(f" {model_name} ready.", flush=True) | |
| return True | |
| except Exception as exc: | |
| print(f" DeepFace build failed: {exc}", flush=True) | |
| return False | |
| print("=== Downloading AI model weights ===", flush=True) | |
| print(f"Trying HF Hub ({HF_REPO})...", flush=True) | |
| if try_hf_hub(): | |
| print("=== HF Hub download successful ===", flush=True) | |
| elif try_deepface_build(): | |
| print("=== GitHub download successful ===", flush=True) | |
| else: | |
| print("WARNING: All download attempts failed.", flush=True) | |
| print("Models will attempt to download on first request.", flush=True) | |
| sys.exit(0) | |