""" Model Loader for VITS TTS Models Loads trained models from the models directory. Models are expected to be in the models/ directory after training. """ import os import logging from pathlib import Path from typing import Optional, List logger = logging.getLogger(__name__) # Model directory MODELS_DIR = Path(os.environ.get("MODELS_DIR", "models")) def _ensure_models_available(): """ Internal function to ensure model files are available. Called during engine initialization. """ if MODELS_DIR.exists() and any(MODELS_DIR.iterdir()): return True # Models need to be loaded - this happens during Docker build logger.info("Initializing model directory...") MODELS_DIR.mkdir(exist_ok=True) try: from huggingface_hub import snapshot_download snapshot_download( repo_id="Harshil748/VoiceAPI-Models", local_dir=MODELS_DIR, local_dir_use_symlinks=False, ignore_patterns=["*.md", ".gitattributes"], ) logger.info("Models initialized successfully") return True except Exception as e: logger.warning(f"Could not initialize models: {e}") return False def get_model_path(voice_key: str) -> Optional[Path]: """Get path to a model directory""" model_dir = MODELS_DIR / voice_key if model_dir.exists(): return model_dir return None def list_available_models() -> List[str]: """List all available trained models""" if not MODELS_DIR.exists(): return [] return [d.name for d in MODELS_DIR.iterdir() if d.is_dir()]