Create preload_neutts.py

preload_neutts.py

@@ -0,0 +1,36 @@
+import os
+import logging
+
+# 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 import the model class
+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()