Upload download_models.py with huggingface_hub

download_models.py

@@ -0,0 +1,98 @@
+"""Download all model assets needed by the whisper-ensemble pipeline.
+
+Run once before using ``pipeline.py``::
+
+    python download_models.py
+
+This will populate the local ``./models`` directory with:
+
+  * ``models/ast-finetuned-audioset/``           (MIT AST router, BSD-3)
+  * ``models/sound-effect-captioning-whisper/``  (laion HF repo snapshot)
+  * ``models/music-whisper/``                    (laion HF repo snapshot)
+  * ``models/voice-tagging-whisper/``            (laion HF repo snapshot)
+  * ``models/BUD-E-Whisper_V1.2/``               (laion HF repo snapshot)
+  * ``models/whisper-small-processor/``          (openai/whisper-small processor only)
+
+The total download is ~5.8 GB. Subsequent runs are no-ops (they skip
+repos that are already mirrored on disk).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent
+MODELS_DIR = ROOT / "models"
+MODELS_DIR.mkdir(parents=True, exist_ok=True)
+
+# ----------------------------------------------------------------------
+# Hugging Face repos to mirror locally
+# ----------------------------------------------------------------------
+HF_REPOS = {
+    # AudioSet router: Audio Spectrogram Transformer fine-tuned on AS-2M
+    # (BSD-3-clause, ~0.459 mAP). The model config ships id2label for all
+    # 527 AudioSet classes, so no extra label CSV is needed.
+    "ast-finetuned-audioset":          "MIT/ast-finetuned-audioset-10-10-0.4593",
+    # Routed Whisper-Small captioners
+    "sound-effect-captioning-whisper": "laion/sound-effect-captioning-whisper",
+    "music-whisper":                   "laion/music-whisper",
+    "voice-tagging-whisper":           "laion/voice-tagging-whisper",
+    "BUD-E-Whisper_V1.2":              "laion/BUD-E-Whisper_V1.2",
+    # Whisper-Small processor (feature extractor + tokenizer fallback,
+    # used by sound-effect / voice-tagging which don't ship one).
+    "whisper-small-processor":         "openai/whisper-small",
+}
+
+
+def download_hf_repos() -> None:
+    """Snapshot all Hugging Face repos used by the pipeline locally."""
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as e:
+        raise SystemExit(
+            "huggingface_hub is required. Install with:\n"
+            "    pip install -r requirements.txt"
+        ) from e
+
+    for local_name, repo_id in HF_REPOS.items():
+        local_dir = MODELS_DIR / local_name
+        if (local_dir / "config.json").exists():
+            print(f"[skip] {repo_id} already mirrored at {local_dir}")
+            continue
+        print(f"[hf]   {repo_id} -> {local_dir}")
+        snapshot_download(
+            repo_id=repo_id,
+            local_dir=str(local_dir),
+            local_dir_use_symlinks=False,
+            allow_patterns=[
+                "*.json",
+                "*.txt",
+                "*.safetensors",
+                "*.bin",
+                "*.model",
+                "*.tiktoken",
+                "tokenizer*",
+                "vocab*",
+                "merges*",
+                "preprocessor*",
+                "generation_config*",
+                "special_tokens*",
+                "added_tokens*",
+                "normalizer*",
+            ],
+        )
+
+
+def main() -> None:
+    print("=" * 72)
+    print("whisper-ensemble: downloading model assets")
+    print("=" * 72)
+    download_hf_repos()
+    print("\nAll downloads complete.")
+    print(f"Models directory: {MODELS_DIR}")
+
+
+if __name__ == "__main__":
+    main()