Update lemas_tts/api.py

lemas_tts/api.py

@@ -1,10 +1,10 @@
+import os
 import random
 import sys
 from pathlib import Path
 import re, regex
 import soundfile as sf
 import tqdm
-from cached_path import cached_path
 from hydra.utils import get_class
 from omegaconf import OmegaConf
 
@@ -36,11 +36,47 @@ def _find_repo_root(start: Path) -> Path:
     return start
 
 
+def _find_pretrained_root(start: Path) -> Path:
+    """
+    Locate the `pretrained_models` root, with support for:
+    1) Explicit env override (LEMAS_PRETRAINED_ROOT)
+    2) Hugging Face Spaces model mount under /models
+    3) Local source tree (searching upwards from this file)
+    """
+    # 1) Explicit override
+    env_root = os.environ.get("LEMAS_PRETRAINED_ROOT")
+    if env_root:
+        p = Path(env_root)
+        if p.is_dir():
+            return p
+
+    # 2) HF Spaces model mount: /models/<model_id>/pretrained_models
+    models_dir = Path("/models")
+    if models_dir.is_dir():
+        # Try the expected model name first
+        specific = models_dir / "LEMAS-Project__LEMAS-TTS"
+        if (specific / "pretrained_models").is_dir():
+            return specific / "pretrained_models"
+        # Otherwise, pick the first model that has a pretrained_models subdir
+        for child in models_dir.iterdir():
+            if child.is_dir() and (child / "pretrained_models").is_dir():
+                return child / "pretrained_models"
+
+    # 3) Local repo layout
+    repo_root = _find_repo_root(start)
+    if (repo_root / "pretrained_models").is_dir():
+        return repo_root / "pretrained_models"
+
+    cwd = Path.cwd()
+    if (cwd / "pretrained_models").is_dir():
+        return cwd / "pretrained_models"
+
+    # Fallback: assume under repo root even if directory is missing
+    return repo_root / "pretrained_models"
+
+
 REPO_ROOT = _find_repo_root(THIS_FILE)
-# Local pretrained root (used when running from a repo / Space that bundles weights)
-PRETRAINED_ROOT = REPO_ROOT / "pretrained_models"
-# Remote pretrained root on Hugging Face Hub (fallback when local files are absent)
-HF_PRETRAINED_ROOT = "hf://LEMAS-Project/LEMAS-TTS/pretrained_models"
+PRETRAINED_ROOT = _find_pretrained_root(THIS_FILE)
 CKPTS_ROOT = PRETRAINED_ROOT / "ckpts"
 
 class TTS: