src/model_loader.py

"""
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()]