Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import importlib.util | |
| import json | |
| import sys | |
| from pathlib import Path | |
| from typing import Any, Dict | |
| import torch | |
| def load_aasist_config(config_path: str | Path) -> Dict[str, Any]: | |
| """Load AASIST JSON-style config and return the model_config section if present.""" | |
| config_path = Path(config_path) | |
| with config_path.open("r", encoding="utf-8") as f: | |
| cfg = json.load(f) | |
| # Official AASIST configs usually contain a top-level "model_config". | |
| return cfg.get("model_config", cfg) | |
| def import_aasist_model_class(aasist_root: str | Path): | |
| """Dynamically import external/aasist/models/AASIST.py.""" | |
| aasist_root = Path(aasist_root) | |
| model_file = aasist_root / "models" / "AASIST.py" | |
| if not model_file.exists(): | |
| raise FileNotFoundError( | |
| f"Could not find AASIST.py at {model_file}. " | |
| "Clone the official repo with: git clone https://github.com/clovaai/aasist.git external/aasist" | |
| ) | |
| # AASIST.py may import modules relative to the repo root. | |
| sys.path.insert(0, str(aasist_root)) | |
| spec = importlib.util.spec_from_file_location("external_aasist_model", model_file) | |
| if spec is None or spec.loader is None: | |
| raise ImportError(f"Could not import {model_file}") | |
| module = importlib.util.module_from_spec(spec) | |
| spec.loader.exec_module(module) | |
| if not hasattr(module, "Model"): | |
| raise AttributeError(f"{model_file} does not define class Model") | |
| return module.Model | |
| def build_aasist_backbone( | |
| aasist_root: str | Path, | |
| aasist_config: str | Path, | |
| checkpoint: str | Path | None = None, | |
| device: str | torch.device = "cpu", | |
| ) -> torch.nn.Module: | |
| """Build the official AASIST backbone and optionally load weights.""" | |
| Model = import_aasist_model_class(aasist_root) | |
| model_config = load_aasist_config(aasist_config) | |
| backbone = Model(model_config).to(device) | |
| if checkpoint is not None: | |
| state = torch.load(checkpoint, map_location=device) | |
| if isinstance(state, dict) and "model" in state: | |
| state = state["model"] | |
| if isinstance(state, dict) and "state_dict" in state: | |
| state = state["state_dict"] | |
| backbone.load_state_dict(state, strict=False) | |
| return backbone | |