Space: preload CPU thread + cache + logs

app.py

@@ -1,4 +1,3 @@
-
 import os
 import re
 import json
@@ -14,10 +13,8 @@ import torch
 import spaces
 from huggingface_hub import login, snapshot_download
 
-# -----------------------
-# Environment hardening
-# -----------------------
-os.environ.setdefault("FLA_CONV_BACKEND", "torch")   # avoid Triton kernels
+# --------- Environnement / stabilité ----------
+os.environ.setdefault("FLA_CONV_BACKEND", "torch")   # éviter les kernels Triton
 os.environ.setdefault("FLA_USE_FAST_OPS", "0")
 os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -26,14 +23,20 @@ try:
 except Exception:
     pass
 
-from pardi_speech import PardiSpeech, VelocityHeadSamplingParams  # present in this repo
+from pardi_speech import PardiSpeech, VelocityHeadSamplingParams  # présent dans ce repo
 
 MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID", "theodorr/pardi-speech-enfr-forbidden")
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# -----------------------
-# Helpers
-# -----------------------
+# --------- Cache global (préchargement au démarrage) ----------
+_MODEL = {"pardi": None, "sr": 24000, "err": None, "logs": [], "thread": None}
+
+def _log(msg: str):
+    _MODEL["logs"].append(str(msg))
+    # borne la taille
+    if len(_MODEL["logs"]) > 2000:
+        _MODEL["logs"] = _MODEL["logs"][-2000:]
+
 def _env_diag() -> str:
     parts = []
     try:
@@ -55,10 +58,8 @@ def _env_diag() -> str:
         parts.append(f"env_diag_error={e}")
     return " | ".join(parts)
 
-
 def _normalize_text(s: str, lang_hint: str = "fr") -> str:
     s = (s or "").strip()
-    # optional: expand digits for FR/EN using num2words if available
     try:
         import re as _re
         from num2words import num2words
@@ -72,16 +73,13 @@ def _normalize_text(s: str, lang_hint: str = "fr") -> str:
         pass
     return s
 
-
 def _to_mono_float32(arr: np.ndarray) -> np.ndarray:
     arr = np.asarray(arr)
     if arr.ndim == 2:
         arr = arr.mean(axis=1)
     return arr.astype(np.float32)
 
-
 def _extract_repo_ids_from_config(config_path: str):
-    """Return list of 'org/name' strings found in a JSON config (simple heuristic)."""
     repo_ids = set()
     preview = None
     try:
@@ -90,76 +88,72 @@ def _extract_repo_ids_from_config(config_path: str):
         pattern = re.compile(r"^[\w\-]+\/[\w\.\-]+$")  # org/name
         def rec(obj):
             if isinstance(obj, dict):
-                for v in obj.values():
-                    rec(v)
+                for v in obj.values(): rec(v)
             elif isinstance(obj, list):
-                for v in obj:
-                    rec(v)
+                for v in obj: rec(v)
             elif isinstance(obj, str):
-                if pattern.match(obj):
-                    repo_ids.add(obj)
+                if pattern.match(obj): repo_ids.add(obj)
         rec(cfg)
-        # create a small preview to help debugging
         try:
             subset_keys = list(cfg)[:5] if isinstance(cfg, dict) else []
-            preview_obj = {k: cfg[k] for k in subset_keys}
-            preview = json.dumps(preview_obj, ensure_ascii=False)[:600]
+            preview = json.dumps({k: cfg[k] for k in subset_keys}, ensure_ascii=False)[:600]
         except Exception:
-            preview = None
+            pass
     except Exception:
         pass
     return sorted(repo_ids), preview
 
-
-def _cpu_first_loader(log_list):
-    """Prefetch main and nested HF repos, then load on CPU in offline mode."""
-    def L(msg):
-        log_list.append(str(msg))
-
-    # 1) Prefetch main repo to local cache
-    L("[prefetch] snapshot_download (main)...")
-    local_dir = snapshot_download(
-        repo_id=MODEL_REPO_ID,
-        token=HF_TOKEN,
-        local_dir=None,
-        local_files_only=False,
-    )
-    L(f"[prefetch] main done -> {local_dir}")
-
-    # 2) Prefetch nested repos found in config.json
-    cfg_path = os.path.join(local_dir, "config.json")
-    nested, cfg_preview = _extract_repo_ids_from_config(cfg_path)
-    if cfg_preview:
-        L(f"[config] preview: {cfg_preview}")
-    for rid in nested:
-        if rid == MODEL_REPO_ID:
-            continue
-        L(f"[prefetch] nested repo: {rid} ...")
-        snapshot_download(repo_id=rid, token=HF_TOKEN, local_dir=None, local_files_only=False)
-        L(f"[prefetch] nested repo: {rid} done")
-
-    # 3) Force offline for actual load to avoid hidden downloads
-    old_off = os.environ.get("HF_HUB_OFFLINE")
-    os.environ["HF_HUB_OFFLINE"] = "1"
+def _prefetch_and_load_cpu():
+    """Exécuté dans un thread au démarrage du Space (hors worker GPU)."""
     try:
-        L("[load] from_pretrained(map_location='cpu')...")
-        m = PardiSpeech.from_pretrained(local_dir, map_location="cpu")
-        m.eval()
-        sr = getattr(m, "sampling_rate", 24000)
-        L(f"[load] cpu OK (sr={sr})")
-        return m, sr, None
+        _log("[prefetch] snapshot_download (main)...")
+        local_dir = snapshot_download(
+            repo_id=MODEL_REPO_ID,
+            token=HF_TOKEN,
+            local_dir=None,
+            local_files_only=False,
+        )
+        _log(f"[prefetch] main done -> {local_dir}")
+
+        cfg_path = os.path.join(local_dir, "config.json")
+        nested, cfg_preview = _extract_repo_ids_from_config(cfg_path)
+        if cfg_preview:
+            _log(f"[config] preview: {cfg_preview}")
+        for rid in nested:
+            if rid == MODEL_REPO_ID:
+                continue
+            _log(f"[prefetch] nested repo: {rid} ...")
+            snapshot_download(repo_id=rid, token=HF_TOKEN, local_dir=None, local_files_only=False)
+            _log(f"[prefetch] nested repo: {rid} done")
+
+        # Forcer offline pendant le vrai chargement
+        old_off = os.environ.get("HF_HUB_OFFLINE")
+        os.environ["HF_HUB_OFFLINE"] = "1"
+        try:
+            _log("[load] from_pretrained(map_location='cpu')...")
+            m = PardiSpeech.from_pretrained(local_dir, map_location="cpu")
+            m.eval()
+            _MODEL["pardi"] = m
+            _MODEL["sr"] = getattr(m, "sampling_rate", 24000)
+            _log(f"[load] cpu OK (sr={_MODEL['sr']})")
+        finally:
+            if old_off is None:
+                os.environ.pop("HF_HUB_OFFLINE", None)
+            else:
+                os.environ["HF_HUB_OFFLINE"] = old_off
+
     except BaseException as e:
-        L(f"[EXC@load] {type(e).__name__}: {e}")
-        return None, None, e
-    finally:
-        if old_off is None:
-            os.environ.pop("HF_HUB_OFFLINE", None)
-        else:
-            os.environ["HF_HUB_OFFLINE"] = old_off
+        _MODEL["err"] = e
+        _log(f"[EXC@preload] {type(e).__name__}: {e}")
+        _log(traceback.format_exc())
 
+# Lance le préchargement (hors GPU) dès l’import
+if _MODEL["thread"] is None:
+    _MODEL["thread"] = threading.Thread(target=_prefetch_and_load_cpu, daemon=True)
+    _MODEL["thread"].start()
 
-def _move_to_cuda_if_available(m, log_list):
-    def L(msg): log_list.append(str(msg))
+def _move_to_cuda_if_available(m, logs_acc):
+    def L(msg): logs_acc.append(str(msg))
     if torch.cuda.is_available():
         L("[move] moving model to cuda...")
         try:
@@ -171,15 +165,12 @@ def _move_to_cuda_if_available(m, log_list):
         L("[move] cuda not available, keep CPU")
     return m
 
-
-# -----------------------
-# Main synthesize (generator)
-# -----------------------
+# --------- UI callback (GPU) ----------
 @spaces.GPU(duration=200)
 def synthesize(
     text: str,
     debug: bool,
-    adv_sampling: bool,   # toggle Velocity Head sampling
+    adv_sampling: bool,   # Velocity Head sampling
     ref_audio,
     ref_text: str,
     steps: int,
@@ -193,7 +184,7 @@ def synthesize(
     logs = []
     def LOG(msg: str):
         logs.append(str(msg))
-        joined = "\n".join(logs)
+        joined = "\n".join(logs + _MODEL["logs"][-50:])  # mêle quelques logs de préchargement
         if len(joined) > 12000:
             joined = joined[-12000:]
         return joined
@@ -210,63 +201,33 @@ def synthesize(
         torch.manual_seed(int(seed))
         os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
 
-        # --- Loader thread + heartbeat + stack dump ---
-        yield None, LOG("[init] nested-prefetch + CPU-first load...")
-
-        loader_logs = []
-        holder = {"model": None, "sr": 24000, "err": None}
-
-        def loader_run():
-            m, sr, err = _cpu_first_loader(loader_logs)
-            holder["model"] = m
-            holder["sr"] = sr if sr is not None else 24000
-            holder["err"] = err
-
-        t = threading.Thread(target=loader_run, daemon=True)
-        t.start()
-        while t.ident is None:
-            time.sleep(0.01)
-        tid = t.ident
-
-        start = time.perf_counter()
-        last_stack = 0.0
-        while t.is_alive():
-            # stream recent loader logs
-            if loader_logs:
-                yield None, LOG("\n".join(loader_logs[-10:]))
-            # dump the loader thread stack every ~6s
-            now = time.perf_counter()
-            if now - last_stack > 6.0 and tid is not None:
-                frame = sys._current_frames().get(tid)
-                if frame is not None:
-                    stack_txt = "".join(traceback.format_stack(frame)[-25:])
-                    yield None, LOG("[stack] loader thread:\n" + stack_txt)
-                last_stack = now
-            # timeout ~110s
-            if now - start > 110:
+        # Si le modèle n’est pas encore prêt, on attend jusqu’à 180s max ici
+        t0 = time.perf_counter()
+        while _MODEL["pardi"] is None and _MODEL["err"] is None:
+            elapsed = time.perf_counter() - t0
+            yield None, LOG(f"[init] still loading on CPU… {elapsed:.1f}s")
+            if elapsed > 180:
+                # dump de la stack du thread de préchargement pour debug
+                tid = _MODEL["thread"].ident if _MODEL["thread"] else None
                 if tid is not None:
                     frame = sys._current_frames().get(tid)
                     if frame is not None:
                         stack_txt = "".join(traceback.format_stack(frame))
                         yield None, LOG("[stack-final]\n" + stack_txt)
-                raise TimeoutError("Model load timeout (exceeded 110s)")
+                raise TimeoutError("Preload timeout (>180s)")
             time.sleep(2.0)
 
-        # After join: flush final logs
-        if loader_logs:
-            yield None, LOG("\n".join(loader_logs[-20:]))
+        if _MODEL["err"]:
+            raise _MODEL["err"]
 
-        if holder["err"]:
-            raise holder["err"]  # will print stack below
-        pardi = holder["model"]
-        if pardi is None:
-            raise RuntimeError("Loader returned no model")
+        pardi = _MODEL["pardi"]
+        sr_out = _MODEL["sr"]
 
-        # move to cuda if possible
+        # Déplacement vers CUDA si possible
         pardi = _move_to_cuda_if_available(pardi, logs)
-        yield None, LOG(f"[init] model ready on {'cuda' if torch.cuda.is_available() else 'cpu'}, sr={holder['sr']}")
+        yield None, LOG(f"[init] model ready on {'cuda' if torch.cuda.is_available() else 'cpu'}, sr={sr_out}")
 
-        # ---- Text & optional prefix ----
+        # ---- Texte + prefix optionnel ----
         txt = _normalize_text(text or "", lang_hint=lang_hint)
         yield None, LOG(f"[text] {txt[:120]}{'...' if len(txt) > 120 else ''}")
 
@@ -288,45 +249,42 @@ def synthesize(
             prefix = (ref_text or "", prefix_tokens[0])
             yield None, LOG("[prefix] done.")
 
-        yield None, LOG(f"[run] has_prefix={prefix is not None}, steps={steps}, cfg={cfg}, cfg_ref={cfg_ref}, T={temperature}, max_seq_len={max_seq_len}, seed={seed}, adv_sampling={adv_sampling}")
+        yield None, LOG(f"[run] has_prefix={prefix is not None}, steps={steps}, cfg={cfg}, cfg_ref={cfg_ref}, "
+                        f"T={temperature}, max_seq_len={max_seq_len}, seed={seed}, adv_sampling={adv_sampling}")
 
-        # ---- Fast path by default (as notebook) ----
+        # ---- Chemin rapide (comme le notebook) ----
         with torch.inference_mode():
             if adv_sampling:
                 try:
                     vparams = VelocityHeadSamplingParams(cfg_ref=float(cfg_ref), cfg=float(cfg), num_steps=int(steps))
                 except TypeError:
-                    vparams = VelocityHeadSamplingParams(cfg_ref=float(cfg_ref), cfg=float(cfg), num_steps=int(steps), temperature=float(temperature))
-                wavs, _ = pardi.text_to_speech([txt], prefix, max_seq_len=int(max_seq_len), velocity_head_sampling_params=vparams)
+                    vparams = VelocityHeadSamplingParams(cfg_ref=float(cfg_ref), cfg=float(cfg),
+                                                         num_steps=int(steps), temperature=float(temperature))
+                wavs, _ = pardi.text_to_speech([txt], prefix, max_seq_len=int(max_seq_len),
+                                               velocity_head_sampling_params=vparams)
             else:
                 wavs, _ = pardi.text_to_speech([txt], prefix, max_seq_len=int(max_seq_len))
 
         wav = wavs[0].detach().cpu().numpy().astype(np.float32)
-        sr_out = getattr(pardi, "sampling_rate", 24000)
         yield (sr_out, wav), LOG("[ok] done.")
 
     except Exception as e:
         tb = traceback.format_exc()
-        yield None, "\n".join(logs + [f"[EXC] {type(e).__name__}: {e}", tb])
+        yield None, LOG(f"[EXC] {type(e).__name__}: {e}\n{tb}")
 
-
-# -----------------------
-# UI
-# -----------------------
+# --------- UI ----------
 def build_demo():
     with gr.Blocks(title="Lina-speech / pardi-speech Demo") as demo:
         gr.Markdown(
             "### Lina-speech (pardi-speech) – Démo TTS\n"
             "Génère de l'audio à partir de texte, avec ou sans prefix (audio de référence).\n"
-            "Par défaut, la voie rapide (comme dans le notebook) est utilisée. Active 'Sampling avancé' pour Velocity Head."
+            "Chemin rapide par défaut (comme le notebook)."
         )
         with gr.Row():
             text = gr.Textbox(label="Texte à synthétiser", lines=4, placeholder="Tape ton texte ici…")
-
         with gr.Accordion("Prefix (optionnel)", open=False):
             ref_audio = gr.Audio(sources=["upload", "microphone"], type="numpy", label="Audio de référence")
             ref_text = gr.Textbox(label="Texte du prefix (si connu)", placeholder="Transcription du prefix (optionnel)")
-
         with gr.Accordion("Options avancées", open=False):
             with gr.Row():
                 steps = gr.Slider(1, 50, value=10, step=1, label="num_steps")
@@ -338,7 +296,7 @@ def build_demo():
                 seed = gr.Number(value=0, precision=0, label="Seed")
                 lang_hint = gr.Dropdown(choices=["fr", "en"], value="fr", label="Langue (normalisation)")
         with gr.Row():
-            debug = gr.Checkbox(value=False, label="Mode debug (affiche la stack du loader)")
+            debug = gr.Checkbox(value=False, label="Mode debug")
             adv_sampling = gr.Checkbox(value=False, label="Sampling avancé (Velocity Head)")
 
         btn = gr.Button("Synthétiser")
@@ -354,6 +312,5 @@ def build_demo():
         )
     return demo
 
-
 if __name__ == "__main__":
     build_demo().launch(ssr_mode=False)