Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import os | |
| import shutil | |
| import urllib.request | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| import spaces | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from src.demo_professional import create_demo | |
| PROJECT_ROOT = Path(__file__).resolve().parent | |
| EXTERNAL_ROOT = PROJECT_ROOT / "external" / "aasist" | |
| AASIST_COMMIT = os.getenv( | |
| "AURALGUARD_AASIST_COMMIT", | |
| "a04c9863f63d44471dde8a6abcb3b082b07cd1d1", | |
| ) | |
| AASIST_BASE_URL = f"https://raw.githubusercontent.com/clovaai/aasist/{AASIST_COMMIT}" | |
| MODEL_REPO = os.getenv("AURALGUARD_MODEL_REPO") | |
| CHECKPOINT_FILE = os.getenv("AURALGUARD_CHECKPOINT_FILE", "best.pt") | |
| MODEL_REVISION = os.getenv("AURALGUARD_MODEL_REVISION", "main") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| def _download(url: str, destination: Path) -> None: | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| if destination.exists(): | |
| return | |
| with urllib.request.urlopen(url) as response, destination.open("wb") as out: | |
| shutil.copyfileobj(response, out) | |
| def ensure_aasist() -> tuple[Path, Path]: | |
| model_path = EXTERNAL_ROOT / "models" / "AASIST.py" | |
| config_path = EXTERNAL_ROOT / "config" / "AASIST.conf" | |
| license_path = EXTERNAL_ROOT / "LICENSE" | |
| _download(f"{AASIST_BASE_URL}/models/AASIST.py", model_path) | |
| _download(f"{AASIST_BASE_URL}/config/AASIST.conf", config_path) | |
| _download(f"{AASIST_BASE_URL}/LICENSE", license_path) | |
| return EXTERNAL_ROOT, config_path | |
| def get_checkpoint() -> Path: | |
| if not MODEL_REPO: | |
| raise RuntimeError( | |
| "AURALGUARD_MODEL_REPO is not set. Add the Hugging Face model repository " | |
| "ID in the Space Settings under Variables." | |
| ) | |
| path = hf_hub_download( | |
| repo_id=MODEL_REPO, | |
| filename=CHECKPOINT_FILE, | |
| repo_type="model", | |
| revision=MODEL_REVISION, | |
| token=HF_TOKEN, | |
| ) | |
| return Path(path) | |
| def build_args(checkpoint: Path, aasist_root: Path, aasist_config: Path): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| return SimpleNamespace( | |
| checkpoint=str(checkpoint), | |
| aasist_root=str(aasist_root), | |
| aasist_config=str(aasist_config), | |
| sample_rate=16000, | |
| duration_sec=4.0, | |
| feature_dim=160, | |
| device=device, | |
| share=False, | |
| ) | |
| aasist_root, aasist_config = ensure_aasist() | |
| checkpoint = get_checkpoint() | |
| demo = create_demo(build_args(checkpoint, aasist_root, aasist_config)) | |
| if __name__ == "__main__": | |
| demo.launch() | |