preload_neutts.py

import os
import sys  # <--- ADD THIS LINE
import logging

# --- ADD THESE LINES TO FIX THE IMPORT PATH ---
# The Dockerfile places the cloned repo at /app/neutts-air
NEUTTS_PATH = "/app/neutts-air"
if NEUTTS_PATH not in sys.path:
    sys.path.insert(0, NEUTTS_PATH)
# ---------------------------------------------

# Set cache environment variables before any other imports
CACHE_DIR = "/app/cache"
os.environ['HF_HOME'] = CACHE_DIR
os.environ['HUGGINGFACE_HUB_CACHE'] = CACHE_DIR

# Now this import will work correctly
from neuttsair.neutts import NeuTTSAir

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def preload_model():
    """
    Downloads and caches the NeuTTSAir model and its dependencies
    to the directory specified by HF_HOME.
    """
    logger.info(f"Pre-loading NeuTTS Air model to cache: {CACHE_DIR}...")
    try:
        # Instantiating the class triggers the download from Hugging Face Hub
        NeuTTSAir(
            backbone_repo="neuphonic/neutts-air",
            backbone_device="cpu",
            codec_repo="neuphonic/neucodec",
            codec_device="cpu"
        )
        logger.info("✅ NeuTTS Air model pre-loading completed successfully!")
    except Exception as e:
        logger.error(f"❌ Error during model pre-loading: {e}")
        # We raise the exception to fail the build if the model can't be downloaded
        raise e

if __name__ == "__main__":
    preload_model()