back to basics

app.py

@@ -1,176 +1,54 @@
 import os
-import re
-import json
-import sys
-import time
-import threading
-import traceback
-
 import gradio as gr
 import numpy as np
-import soundfile as sf
 import torch
+import soundfile as sf
 import spaces
-from huggingface_hub import login, snapshot_download
-
-# --------- 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
-try:
-    torch.set_float32_matmul_precision("high")
-except Exception:
-    pass
 
+from huggingface_hub import login
 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")
 
-# --------- 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 = []
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN:
     try:
-        parts.append(f"torch={torch.__version__}")
-        try:
-            import triton  # type: ignore
-            parts.append(f"triton={getattr(triton, '__version__', 'unknown')}")
-        except Exception:
-            parts.append("triton=not_importable")
-        parts.append(f"cuda.is_available={torch.cuda.is_available()}")
-        if torch.cuda.is_available():
-            parts.append(f"cuda.version={torch.version.cuda}")
-            try:
-                free, total = torch.cuda.mem_get_info()
-                parts.append(f"mem_free={free/1e9:.2f}GB/{total/1e9:.2f}GB")
-            except Exception:
-                pass
+        login(token=HF_TOKEN)
+        print("✅ Logged to Hugging Face Hub.")
     except Exception as e:
-        parts.append(f"env_diag_error={e}")
-    return " | ".join(parts)
+        print("⚠️ HF login failed:", e)
+
+_pardi = None
+_sampling_rate = 24000
 
 def _normalize_text(s: str, lang_hint: str = "fr") -> str:
-    s = (s or "").strip()
+    s = (s or "").strip().lower()
     try:
-        import re as _re
+        import re
         from num2words import num2words
-        def repl(m):
-            try:
-                return num2words(int(m.group()), lang=lang_hint)
-            except Exception:
-                return m.group()
-        s = _re.sub(r"\d+", repl, s)
+        def repl(m): return num2words(int(m.group()), lang=lang_hint)
+        s = re.sub(r"\d+", repl, s)
     except Exception:
         pass
     return s
 
+def _load_model(device: str = "cuda"):
+    global _pardi, _sampling_rate
+    if _pardi is None:
+        _pardi = PardiSpeech.from_pretrained(MODEL_REPO_ID, map_location=device)
+        _sampling_rate = getattr(_pardi, "sampling_rate", 24000)
+        print(f"✅ PardiSpeech loaded on {device} (sr={_sampling_rate}).")
+    return _pardi
+
 def _to_mono_float32(arr: np.ndarray) -> np.ndarray:
-    arr = np.asarray(arr)
+    arr = arr.astype(np.float32)
     if arr.ndim == 2:
         arr = arr.mean(axis=1)
-    return arr.astype(np.float32)
-
-def _extract_repo_ids_from_config(config_path: str):
-    repo_ids = set()
-    preview = None
-    try:
-        with open(config_path, "r", encoding="utf-8") as f:
-            cfg = json.load(f)
-        pattern = re.compile(r"^[\w\-]+\/[\w\.\-]+$")  # org/name
-        def rec(obj):
-            if isinstance(obj, dict):
-                for v in obj.values(): rec(v)
-            elif isinstance(obj, list):
-                for v in obj: rec(v)
-            elif isinstance(obj, str):
-                if pattern.match(obj): repo_ids.add(obj)
-        rec(cfg)
-        try:
-            subset_keys = list(cfg)[:5] if isinstance(cfg, dict) else []
-            preview = json.dumps({k: cfg[k] for k in subset_keys}, ensure_ascii=False)[:600]
-        except Exception:
-            pass
-    except Exception:
-        pass
-    return sorted(repo_ids), preview
-
-def _prefetch_and_load_cpu():
-    """Exécuté dans un thread au démarrage du Space (hors worker GPU)."""
-    try:
-        _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:
-        _MODEL["err"] = e
-        _log(f"[EXC@preload] {type(e).__name__}: {e}")
-        _log(traceback.format_exc())
+    return arr
 
-# 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, logs_acc):
-    def L(msg): logs_acc.append(str(msg))
-    if torch.cuda.is_available():
-        L("[move] moving model to cuda...")
-        try:
-            m = m.to("cuda")  # type: ignore[attr-defined]
-            L("[move] cuda OK")
-        except Exception as e:
-            L(f"[move] .to('cuda') failed: {e}. Keeping on CPU.")
-    else:
-        L("[move] cuda not available, keep CPU")
-    return m
-
-# --------- UI callback (GPU) ----------
-@spaces.GPU(duration=200)
+@spaces.GPU(duration=120)
 def synthesize(
     text: str,
-    debug: bool,
-    adv_sampling: bool,   # Velocity Head sampling
     ref_audio,
     ref_text: str,
     steps: int,
@@ -179,112 +57,83 @@ def synthesize(
     temperature: float,
     max_seq_len: int,
     seed: int,
-    lang_hint: str,
+    lang_hint: str
 ):
-    logs = []
-    def LOG(msg: str):
-        logs.append(str(msg))
-        joined = "\n".join(logs + _MODEL["logs"][-50:])  # mêle quelques logs de préchargement
-        if len(joined) > 12000:
-            joined = joined[-12000:]
-        return joined
-
-    try:
-        if HF_TOKEN:
-            try:
-                login(token=HF_TOKEN)
-                yield None, LOG("✅ HF login ok")
-            except Exception as e:
-                yield None, LOG(f"⚠️ HF login failed: {e}")
-
-        yield None, LOG("[env] " + _env_diag())
-        torch.manual_seed(int(seed))
-        os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
-
-        # 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("Preload timeout (>180s)")
-            time.sleep(2.0)
-
-        if _MODEL["err"]:
-            raise _MODEL["err"]
-
-        pardi = _MODEL["pardi"]
-        sr_out = _MODEL["sr"]
-
-        # 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={sr_out}")
-
-        # ---- Texte + prefix optionnel ----
-        txt = _normalize_text(text or "", lang_hint=lang_hint)
-        yield None, LOG(f"[text] {txt[:120]}{'...' if len(txt) > 120 else ''}")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(int(seed))
 
-        steps = int(min(max(1, int(steps)), 16))
-        max_seq_len = int(min(max(50, int(max_seq_len)), 600))
+    pardi = _load_model(device)
+    txt = _normalize_text(text, lang_hint=lang_hint)
 
-        prefix = None
-        if ref_audio is not None:
-            yield None, LOG("[prefix] encoding reference audio...")
-            if isinstance(ref_audio, str):
-                wav, sr = sf.read(ref_audio)
-            else:
-                sr, wav = ref_audio
-            wav = _to_mono_float32(wav)
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            wav_t = torch.from_numpy(wav).to(device).unsqueeze(0)
-            with torch.inference_mode():
-                prefix_tokens = pardi.patchvae.encode(wav_t)  # type: ignore[attr-defined]
-            prefix = (ref_text or "", prefix_tokens[0])
-            yield None, LOG("[prefix] done.")
+    cache = pardi.tts.audio_decoder.init_cache(int(max_seq_len), device)
 
-        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}")
+    # --- IMPORTANT : signature de VelocityHeadSamplingParams ---
+    # Dans ton notebook d’inférence, la classe attend (cfg_ref, cfg, num_steps) SANS 'temperature'.
+    # On essaie d’abord sans temperature, puis fallback si la classe en accepte une.
+    try:
+        vel_params = VelocityHeadSamplingParams(
+            cfg_ref=float(cfg_ref),
+            cfg=float(cfg),
+            num_steps=int(steps)
+        )
+    except TypeError:
+        vel_params = VelocityHeadSamplingParams(
+            cfg_ref=float(cfg_ref),
+            cfg=float(cfg),
+            num_steps=int(steps),
+            temperature=float(temperature)
+        )
 
-        # ---- Chemin rapide (comme le notebook) ----
+    # Prefix optionnel
+    prefix = None
+    if ref_audio is not None:
+        if isinstance(ref_audio, str):
+            wav, sr = sf.read(ref_audio)
+        else:
+            sr, wav = ref_audio
+        wav = _to_mono_float32(np.array(wav))
+        wav_t = torch.from_numpy(wav).to(device)
+        import torchaudio
+        if sr != pardi.sampling_rate:
+            wav_t = torchaudio.functional.resample(wav_t, sr, pardi.sampling_rate)
+        wav_t = wav_t.unsqueeze(0)
         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)
-            else:
-                wavs, _ = pardi.text_to_speech([txt], prefix, max_seq_len=int(max_seq_len))
+            prefix_tokens = pardi.patchvae.encode(wav_t)
+        prefix = (ref_text or "", prefix_tokens[0])
 
-        wav = wavs[0].detach().cpu().numpy().astype(np.float32)
-        yield (sr_out, wav), LOG("[ok] done.")
+    print(f"[debug] has_prefix={prefix is not None}, steps={steps}, cfg={cfg}, cfg_ref={cfg_ref}, T={temperature}, max_seq_len={max_seq_len}, seed={seed}")
 
+    try:
+        with torch.inference_mode():
+            wavs, _ = pardi.text_to_speech(
+                [txt],
+                prefix,
+                max_seq_len=int(max_seq_len),
+                velocity_head_sampling_params=vel_params,
+                cache=cache
+            )
     except Exception as e:
-        tb = traceback.format_exc()
-        yield None, LOG(f"[EXC] {type(e).__name__}: {e}\n{tb}")
+        import traceback, sys
+        print("❌ text_to_speech failed:", e, file=sys.stderr)
+        traceback.print_exc()
+        raise gr.Error(f"Synthèse échouée: {type(e).__name__}: {e}")
+
+    wav = wavs[0].detach().cpu().numpy()
+    return (_sampling_rate, wav)
 
-# --------- 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"
-            "Chemin rapide par défaut (comme le notebook)."
+            "## 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"
+            "Paramètres avancés: *num_steps*, *CFG*, *température*, *max_seq_len*, *seed*."
         )
+
         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)")
+            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")
@@ -293,26 +142,21 @@ def build_demo():
             with gr.Row():
                 temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Température")
                 max_seq_len = gr.Slider(50, 1200, value=300, step=10, label="max_seq_len (tokens audio)")
-                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")
-            adv_sampling = gr.Checkbox(value=False, label="Sampling avancé (Velocity Head)")
+                seed = gr.Number(value=0, precision=0, label="Seed (reproductibilité)")
+            lang_hint = gr.Dropdown(choices=["fr", "en"], value="fr", label="Langue (normalisation)")
 
         btn = gr.Button("Synthétiser")
         out_audio = gr.Audio(label="Sortie audio", type="numpy")
-        logs_box = gr.Textbox(label="Logs (live)", lines=28)
 
         demo.queue(default_concurrency_limit=1, max_size=32)
+
         btn.click(
             fn=synthesize,
-            inputs=[text, debug, adv_sampling, ref_audio, ref_text, steps, cfg, cfg_ref, temperature, max_seq_len, seed, lang_hint],
-            outputs=[out_audio, logs_box],
-            api_name="synthesize",
+            inputs=[text, ref_audio, ref_text, steps, cfg, cfg_ref, temperature, max_seq_len, seed, lang_hint],
+            outputs=[out_audio]
         )
     return demo
 
 if __name__ == "__main__":
-    build_demo().launch(ssr_mode=False)
-# retrigger 2025-10-30T15:17:49+01:00
-# retrigger 2025-10-30T16:37:47+01:00
+    demo = build_demo()
+    demo.launch()

app.py.bak

@@ -1,54 +1,176 @@
 import os
+import re
+import json
+import sys
+import time
+import threading
+import traceback
+
 import gradio as gr
 import numpy as np
-import torch
 import soundfile as sf
+import torch
 import spaces
+from huggingface_hub import login, snapshot_download
+
+# --------- 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
+try:
+    torch.set_float32_matmul_precision("high")
+except Exception:
+    pass
 
-from huggingface_hub import login
 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")
-if HF_TOKEN:
+
+# --------- 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:
-        login(token=HF_TOKEN)
-        print("✅ Logged to Hugging Face Hub.")
+        parts.append(f"torch={torch.__version__}")
+        try:
+            import triton  # type: ignore
+            parts.append(f"triton={getattr(triton, '__version__', 'unknown')}")
+        except Exception:
+            parts.append("triton=not_importable")
+        parts.append(f"cuda.is_available={torch.cuda.is_available()}")
+        if torch.cuda.is_available():
+            parts.append(f"cuda.version={torch.version.cuda}")
+            try:
+                free, total = torch.cuda.mem_get_info()
+                parts.append(f"mem_free={free/1e9:.2f}GB/{total/1e9:.2f}GB")
+            except Exception:
+                pass
     except Exception as e:
-        print("⚠️ HF login failed:", e)
-
-_pardi = None
-_sampling_rate = 24000
+        parts.append(f"env_diag_error={e}")
+    return " | ".join(parts)
 
 def _normalize_text(s: str, lang_hint: str = "fr") -> str:
-    s = (s or "").strip().lower()
+    s = (s or "").strip()
     try:
-        import re
+        import re as _re
         from num2words import num2words
-        def repl(m): return num2words(int(m.group()), lang=lang_hint)
-        s = re.sub(r"\d+", repl, s)
+        def repl(m):
+            try:
+                return num2words(int(m.group()), lang=lang_hint)
+            except Exception:
+                return m.group()
+        s = _re.sub(r"\d+", repl, s)
     except Exception:
         pass
     return s
 
-def _load_model(device: str = "cuda"):
-    global _pardi, _sampling_rate
-    if _pardi is None:
-        _pardi = PardiSpeech.from_pretrained(MODEL_REPO_ID, map_location=device)
-        _sampling_rate = getattr(_pardi, "sampling_rate", 24000)
-        print(f"✅ PardiSpeech loaded on {device} (sr={_sampling_rate}).")
-    return _pardi
-
 def _to_mono_float32(arr: np.ndarray) -> np.ndarray:
-    arr = arr.astype(np.float32)
+    arr = np.asarray(arr)
     if arr.ndim == 2:
         arr = arr.mean(axis=1)
-    return arr
+    return arr.astype(np.float32)
+
+def _extract_repo_ids_from_config(config_path: str):
+    repo_ids = set()
+    preview = None
+    try:
+        with open(config_path, "r", encoding="utf-8") as f:
+            cfg = json.load(f)
+        pattern = re.compile(r"^[\w\-]+\/[\w\.\-]+$")  # org/name
+        def rec(obj):
+            if isinstance(obj, dict):
+                for v in obj.values(): rec(v)
+            elif isinstance(obj, list):
+                for v in obj: rec(v)
+            elif isinstance(obj, str):
+                if pattern.match(obj): repo_ids.add(obj)
+        rec(cfg)
+        try:
+            subset_keys = list(cfg)[:5] if isinstance(cfg, dict) else []
+            preview = json.dumps({k: cfg[k] for k in subset_keys}, ensure_ascii=False)[:600]
+        except Exception:
+            pass
+    except Exception:
+        pass
+    return sorted(repo_ids), preview
+
+def _prefetch_and_load_cpu():
+    """Exécuté dans un thread au démarrage du Space (hors worker GPU)."""
+    try:
+        _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:
+        _MODEL["err"] = e
+        _log(f"[EXC@preload] {type(e).__name__}: {e}")
+        _log(traceback.format_exc())
 
-@spaces.GPU(duration=120)
+# 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, logs_acc):
+    def L(msg): logs_acc.append(str(msg))
+    if torch.cuda.is_available():
+        L("[move] moving model to cuda...")
+        try:
+            m = m.to("cuda")  # type: ignore[attr-defined]
+            L("[move] cuda OK")
+        except Exception as e:
+            L(f"[move] .to('cuda') failed: {e}. Keeping on CPU.")
+    else:
+        L("[move] cuda not available, keep CPU")
+    return m
+
+# --------- UI callback (GPU) ----------
+@spaces.GPU(duration=200)
 def synthesize(
     text: str,
+    debug: bool,
+    adv_sampling: bool,   # Velocity Head sampling
     ref_audio,
     ref_text: str,
     steps: int,
@@ -57,83 +179,112 @@ def synthesize(
     temperature: float,
     max_seq_len: int,
     seed: int,
-    lang_hint: str
+    lang_hint: str,
 ):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    torch.manual_seed(int(seed))
+    logs = []
+    def LOG(msg: str):
+        logs.append(str(msg))
+        joined = "\n".join(logs + _MODEL["logs"][-50:])  # mêle quelques logs de préchargement
+        if len(joined) > 12000:
+            joined = joined[-12000:]
+        return joined
 
-    pardi = _load_model(device)
-    txt = _normalize_text(text, lang_hint=lang_hint)
+    try:
+        if HF_TOKEN:
+            try:
+                login(token=HF_TOKEN)
+                yield None, LOG("✅ HF login ok")
+            except Exception as e:
+                yield None, LOG(f"⚠️ HF login failed: {e}")
 
-    cache = pardi.tts.audio_decoder.init_cache(int(max_seq_len), device)
+        yield None, LOG("[env] " + _env_diag())
+        torch.manual_seed(int(seed))
+        os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
 
-    # --- IMPORTANT : signature de VelocityHeadSamplingParams ---
-    # Dans ton notebook d’inférence, la classe attend (cfg_ref, cfg, num_steps) SANS 'temperature'.
-    # On essaie d’abord sans temperature, puis fallback si la classe en accepte une.
-    try:
-        vel_params = VelocityHeadSamplingParams(
-            cfg_ref=float(cfg_ref),
-            cfg=float(cfg),
-            num_steps=int(steps)
-        )
-    except TypeError:
-        vel_params = VelocityHeadSamplingParams(
-            cfg_ref=float(cfg_ref),
-            cfg=float(cfg),
-            num_steps=int(steps),
-            temperature=float(temperature)
-        )
+        # 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("Preload timeout (>180s)")
+            time.sleep(2.0)
 
-    # Prefix optionnel
-    prefix = None
-    if ref_audio is not None:
-        if isinstance(ref_audio, str):
-            wav, sr = sf.read(ref_audio)
-        else:
-            sr, wav = ref_audio
-        wav = _to_mono_float32(np.array(wav))
-        wav_t = torch.from_numpy(wav).to(device)
-        import torchaudio
-        if sr != pardi.sampling_rate:
-            wav_t = torchaudio.functional.resample(wav_t, sr, pardi.sampling_rate)
-        wav_t = wav_t.unsqueeze(0)
-        with torch.inference_mode():
-            prefix_tokens = pardi.patchvae.encode(wav_t)
-        prefix = (ref_text or "", prefix_tokens[0])
+        if _MODEL["err"]:
+            raise _MODEL["err"]
 
-    print(f"[debug] has_prefix={prefix is not None}, steps={steps}, cfg={cfg}, cfg_ref={cfg_ref}, T={temperature}, max_seq_len={max_seq_len}, seed={seed}")
+        pardi = _MODEL["pardi"]
+        sr_out = _MODEL["sr"]
 
-    try:
+        # 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={sr_out}")
+
+        # ---- Texte + prefix optionnel ----
+        txt = _normalize_text(text or "", lang_hint=lang_hint)
+        yield None, LOG(f"[text] {txt[:120]}{'...' if len(txt) > 120 else ''}")
+
+        steps = int(min(max(1, int(steps)), 16))
+        max_seq_len = int(min(max(50, int(max_seq_len)), 600))
+
+        prefix = None
+        if ref_audio is not None:
+            yield None, LOG("[prefix] encoding reference audio...")
+            if isinstance(ref_audio, str):
+                wav, sr = sf.read(ref_audio)
+            else:
+                sr, wav = ref_audio
+            wav = _to_mono_float32(wav)
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            wav_t = torch.from_numpy(wav).to(device).unsqueeze(0)
+            with torch.inference_mode():
+                prefix_tokens = pardi.patchvae.encode(wav_t)  # type: ignore[attr-defined]
+            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}, "
+                        f"T={temperature}, max_seq_len={max_seq_len}, seed={seed}, adv_sampling={adv_sampling}")
+
+        # ---- Chemin rapide (comme le notebook) ----
         with torch.inference_mode():
-            wavs, _ = pardi.text_to_speech(
-                [txt],
-                prefix,
-                max_seq_len=int(max_seq_len),
-                velocity_head_sampling_params=vel_params,
-                cache=cache
-            )
-    except Exception as e:
-        import traceback, sys
-        print("❌ text_to_speech failed:", e, file=sys.stderr)
-        traceback.print_exc()
-        raise gr.Error(f"Synthèse échouée: {type(e).__name__}: {e}")
+            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)
+            else:
+                wavs, _ = pardi.text_to_speech([txt], prefix, max_seq_len=int(max_seq_len))
 
-    wav = wavs[0].detach().cpu().numpy()
-    return (_sampling_rate, wav)
+        wav = wavs[0].detach().cpu().numpy().astype(np.float32)
+        yield (sr_out, wav), LOG("[ok] done.")
 
+    except Exception as e:
+        tb = traceback.format_exc()
+        yield None, LOG(f"[EXC] {type(e).__name__}: {e}\n{tb}")
+
+# --------- 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"
-            "Paramètres avancés: *num_steps*, *CFG*, *température*, *max_seq_len*, *seed*."
+            "### 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"
+            "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)")
+            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")
@@ -142,22 +293,26 @@ def build_demo():
             with gr.Row():
                 temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Température")
                 max_seq_len = gr.Slider(50, 1200, value=300, step=10, label="max_seq_len (tokens audio)")
-                seed = gr.Number(value=0, precision=0, label="Seed (reproductibilité)")
-            lang_hint = gr.Dropdown(choices=["fr", "en"], value="fr", label="Langue (normalisation)")
+                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")
+            adv_sampling = gr.Checkbox(value=False, label="Sampling avancé (Velocity Head)")
 
         btn = gr.Button("Synthétiser")
         out_audio = gr.Audio(label="Sortie audio", type="numpy")
+        logs_box = gr.Textbox(label="Logs (live)", lines=28)
 
         demo.queue(default_concurrency_limit=1, max_size=32)
-
         btn.click(
             fn=synthesize,
-            inputs=[text, ref_audio, ref_text, steps, cfg, cfg_ref, temperature, max_seq_len, seed, lang_hint],
-            outputs=[out_audio]
+            inputs=[text, debug, adv_sampling, ref_audio, ref_text, steps, cfg, cfg_ref, temperature, max_seq_len, seed, lang_hint],
+            outputs=[out_audio, logs_box],
+            api_name="synthesize",
         )
     return demo
 
 if __name__ == "__main__":
-    demo = build_demo()
-    demo.launch()
-# retrigger 2025-10-29T16:27:55+01:00
+    build_demo().launch(ssr_mode=False)
+# retrigger 2025-10-30T15:17:49+01:00
+# retrigger 2025-10-30T16:37:47+01:00

tts/model/simple_gla.py

@@ -1,304 +1,291 @@
-"""
-Patched Simple GLA decoder for HF Spaces (ZeroGPU) — safe PyTorch-only paths.
-
-- Forces FLA (flash-linear-attention) to avoid fused/Triton kernels during __init__ & forward
-- Adds tolerant construction of SimpleGatedLinearAttention (backend="torch", fused=False)
-- Falls back to a no-op GLA stub if FLA construction fails (for demo resilience)
-- Keeps cache handling defensive to avoid NoneType unpack errors
-
-Drop-in replacement for: tts/model/simple_gla.py
-"""
-
 import os
-from typing import Optional, Dict, Any, Tuple, List, Union
-
-# ---- Force safe runtime defaults (no Triton / no compile) ----
-os.environ.setdefault("FLA_CONV_BACKEND", "torch")
-os.environ.setdefault("FLA_USE_FAST_OPS", "0")
-os.environ.setdefault("FLA_DISABLE_TRITON", "1")     # ignored if not recognized
-os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
-os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
 
 import torch
 import torch.nn.functional as F
-from torch import nn
-
 from einops import rearrange
+from fla.layers.simple_gla import SimpleGatedLinearAttention
+from fla.models.utils import Cache
+from sympy import num_digits
+from torch import nn
 
-# ---------- Try importing FLA; otherwise define a stub ----------
-try:
-    from fla.layers.simple_gla import SimpleGatedLinearAttention  # type: ignore
-    from fla.models.utils import Cache  # type: ignore
-    _FLA_AVAILABLE = True
-except Exception:
-    _FLA_AVAILABLE = False
-
-    class SimpleGatedLinearAttention(nn.Module):  # minimal stub (identity)
-        def __init__(self, *args, **kwargs):
-            super().__init__()
-
-        def forward(self, x, past_key_values=None, use_cache: bool = False, **kwargs):
-            # Match tuple output convention used by callers: (x, kv)
-            return x, None
-
-    # Fallback Cache typing
-    Cache = Dict[str, Any]  # type: ignore
-
-
-# Local layers / utils
 from tts.layers.attention import CrossAttention
 from tts.layers.ffn import SwiGLU
 
 from .cache_utils import FLACache
 from .config import SimpleGLADecoderConfig
 from .registry import register_decoder
+from .shortconv import ShortConvBlock
 
+if "GRAD_CKPT" in os.environ:
 
-def _force_safe_fla_impl(m: SimpleGatedLinearAttention) -> None:
-    """Force SimpleGatedLinearAttention to use non-fused, PyTorch-only kernels.
-
-    On Hugging Face Spaces (ZeroGPU) + Python 3.10/Triton 3.1, fused kernels can hang at import/first call.
-    We harden the module to avoid fused/triton implementations.
-    """
-    # Prefer explicit mode to avoid backend auto-selection
-    try:
-        if hasattr(m, "mode"):
-            m.mode = "chunk"  # safer than "recurrent" fused paths
-    except Exception:
-        pass
-
-    # For recent versions exposing implementation switches:
-    for attr, val in (("recurrent_impl", "naive"),
-                      ("chunk_impl", "naive"),
-                      ("fused", False),
-                      ("backend", "torch")):
-        if hasattr(m, attr):
-            try:
-                setattr(m, attr, val)
-            except Exception:
-                pass
-
-
-def _make_tmix(dim: int, num_heads: int) -> SimpleGatedLinearAttention:
-    """
-    Construct SimpleGatedLinearAttention using the safest available signature.
-    Falls back gracefully if kwargs are not supported by the installed FLA version.
-    """
-    # Try most explicit signature first
-    try:
-        tmix = SimpleGatedLinearAttention(
-            hidden_size=dim,
-            num_heads=num_heads,
-            causal=True,
-            backend="torch",   # key to avoid Triton
-            fused=False,
-        )
-        _force_safe_fla_impl(tmix)
-        return tmix
-    except TypeError:
-        pass
-    except Exception:
-        # If constructing with explicit kwargs fails for another reason,
-        # we will try progressively simpler signatures below.
-        pass
-
-    # Try without fused/backends but keep causal if supported
-    try:
-        tmix = SimpleGatedLinearAttention(
-            hidden_size=dim,
-            num_heads=num_heads,
-            causal=True,
-        )
-        _force_safe_fla_impl(tmix)
-        return tmix
-    except TypeError:
-        pass
-    except Exception:
-        pass
-
-    # Try minimal signature
-    try:
-        tmix = SimpleGatedLinearAttention(
-            hidden_size=dim,
-            num_heads=num_heads,
-        )
-        _force_safe_fla_impl(tmix)
-        return tmix
-    except Exception:
-        # Last resort: identity stub
-        return SimpleGatedLinearAttention()
-
-
-def _cache_for_layer(cache: Optional[Cache], idx: int) -> Optional[Cache]:
-    """
-    Extract per-layer cache if present; return None if structure is not compatible.
-    FLA expects either:
-      - cache["layers"][i]["conv_state"] being a tuple/list
-      - or a top-level cache dict with "conv_state" key
-    """
-    if isinstance(cache, dict):
-        if "layers" in cache and isinstance(cache["layers"], (list, tuple)):
-            if idx < len(cache["layers"]) and isinstance(cache["layers"][idx], dict):
-                # Layer-specific cache entry
-                c = cache["layers"][idx]
-                # Validate conv_state shape
-                if isinstance(c.get("conv_state", None), (list, tuple)):
-                    return c
-                # If not valid, ignore layer cache to prevent NoneType errors
-                return None
-        # Some layouts put conv states directly at top-level
-        if isinstance(cache.get("conv_state", None), (list, tuple)):
-            return cache
-    return None
+    def maybe_grad_ckpt(f):
+        def grad_ckpt_f(*args, **kwargs):
+            return torch.utils.checkpoint.checkpoint(
+                f, *args, **kwargs, use_reentrant=False
+            )
 
+        return grad_ckpt_f
+else:
 
-class SimpleGLABlock(nn.Module):
-    """One decoder block with GLA time-mixing + feed-forward + (optional) norm/shortconv."""
+    def maybe_grad_ckpt(f):
+        return f
 
+
+class SimpleGLABlock(nn.Module):
     def __init__(
         self,
         dim: int,
         num_heads: int,
-        layer_idx: int = 0,
-        expand_k: float = 1.0,
-        expand_v: float = 1.0,
-        use_short_conv: bool = False,
-        ffn_expansion_factor: int = 4,
+        layer_idx: int,
+        expand_k: float,
+        expand_v: float,
+        use_short_conv: bool,
+        ffn_expansion_factor: int,
     ):
         super().__init__()
-        # Time-mixing (GLA) — robust construction
-        self.tmix = _make_tmix(dim=dim, num_heads=num_heads)
-
-        # Feed-forward
-        hidden_ff = int(dim * ffn_expansion_factor)
-        self.cmix = SwiGLU(dim, hidden_ff)
-
-        # Norms
+        self.tmix = SimpleGatedLinearAttention(
+            hidden_size=dim,
+            num_heads=num_heads,
+            layer_idx=layer_idx,
+        )
+        self.cmix = SwiGLU(dim, ffn_expansion_factor)
         self.norm1 = nn.LayerNorm(dim)
         self.norm2 = nn.LayerNorm(dim)
 
-        # (Optional) short conv placeholder
-        self.use_short_conv = use_short_conv
-
     def forward(
         self,
-        x: torch.Tensor,
-        cache: Optional[Cache] = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        # Extract a valid cache view for this layer (if any)
-        pkv = _cache_for_layer(cache, idx=getattr(self, "layer_idx", 0))
-
-        # Some FLA versions want explicit flags
-        use_cache_flag = isinstance(pkv, dict) and isinstance(pkv.get("conv_state", None), (list, tuple))
-
-        y, _ = self.tmix(
-            self.norm1(x),
-            past_key_values=pkv,
-            use_cache=use_cache_flag,
+        x,
+        freqs: torch.Tensor | None = None,
+        text_freqs: torch.Tensor | None = None,
+        cache: Cache | None = None,
+    ):
+        x = (
+            self.tmix(
+                self.norm1(x),
+                past_key_values=cache,
+                use_cache=cache is not None,
+            )[0]
+            + x
         )
-        x = y + x
         x = self.cmix(self.norm2(x)) + x
         return x
 
 
 class DecoderBlockWithOptionalCrossAttention(nn.Module):
-    """Wrap a GLABlock and add cross-attention (encoder-decoder attention) if provided."""
-
-    def __init__(self, decoder_block: nn.Module, crossatt: Optional[nn.Module] = None):
+    def __init__(self, decoder_block: nn.Module, crossatt: nn.Module | None = None):
         super().__init__()
+
         self.decoder_block = decoder_block
         self.crossatt = crossatt
 
     def forward(
         self,
         x: torch.Tensor,
-        encoder_output: Optional[torch.Tensor] = None,
-        text_freqs: Optional[torch.Tensor] = None,
-        cache: Optional[Cache] = None,
-        crossatt_mask: Optional[torch.Tensor] = None,
+        encoder_output: torch.Tensor | None = None,
+        freqs: torch.Tensor | None = None,
+        text_freqs: torch.Tensor | None = None,
+        cache: Cache | None = None,
+        selfatt_mask: torch.Tensor | None = None,
+        crossatt_mask: torch.Tensor | list[torch.Tensor] | None = None,
     ) -> torch.Tensor:
-        if self.crossatt is not None and encoder_output is not None:
-            # Standard cross-attention (keys/values from encoder_output)
-            x = self.crossatt(
+        x = self.decoder_block(
+            x,
+            freqs=freqs,
+            cache=cache,
+        )
+        if type(crossatt_mask) is list:
+            crossatt_mask = crossatt_mask[self.decoder_block.tmix.layer_idx]
+        if self.crossatt is not None:
+            x = x + self.crossatt(
                 x,
-                context=encoder_output,
+                k=encoder_output,
+                text_freqs=text_freqs,
                 mask=crossatt_mask,
+                cache=cache,
             )
-        x = self.decoder_block(x, cache=cache)
+
         return x
 
 
 @register_decoder("simple_gla")
 class SimpleGLADecoder(nn.Module):
     config = SimpleGLADecoderConfig
-    """Decoder composed of a stack of SimpleGLABlock (+ optional cross-attention)."""
 
-    def __init__(self, config: SimpleGLADecoderConfig):
+    def __init__(self, cfg: SimpleGLADecoderConfig):
         super().__init__()
-        self.config = config
-
-        dim = getattr(config, "hidden_size", getattr(config, "dim", 512))
-        num_heads = getattr(config, "num_heads", 8)
-        num_layers = getattr(config, "num_layers", 12)
-        ffn_expansion_factor = getattr(config, "ffn_expansion_factor", 4)
-        expand_k = getattr(config, "expand_k", 1.0)
-        expand_v = getattr(config, "expand_v", 1.0)
-        use_short_conv = getattr(config, "use_short_conv", False)
-        cross_every = getattr(config, "cross_every", 1)  # add cross-att every N layers (1 = every layer)
-
-        decoder_layers: List[nn.Module] = []
-        for i in range(num_layers):
-            block = SimpleGLABlock(
-                dim=dim,
-                num_heads=num_heads,
-                layer_idx=i,
-                expand_k=expand_k,
-                expand_v=expand_v,
-                use_short_conv=use_short_conv,
-                ffn_expansion_factor=ffn_expansion_factor,
+
+        assert cfg.dim % cfg.num_heads == 0, "num_heads should divide dim"
+        assert cfg.blind_crossatt + (cfg.listen_read_crossatt is not None) < 2, (
+            "at most one specialized cross-attention"
+        )
+
+        self.head_dim = cfg.dim // cfg.num_heads
+        self.num_heads = cfg.num_heads
+
+        def simple_gla_block(i):
+            conv_layers = [] if cfg.conv_layers is None else cfg.conv_layers
+            if i in conv_layers:
+                return ShortConvBlock(
+                    dim=cfg.dim,
+                    kernel_size=4,
+                    ffn_expansion_factor=cfg.ffn_expansion_factor,
+                    layer_idx=i,
+                    use_fast_conv1d=True,
+                )
+
+            else:
+                return SimpleGLABlock(
+                    dim=cfg.dim,
+                    num_heads=cfg.num_heads,
+                    layer_idx=i,
+                    expand_k=cfg.expand_k,
+                    expand_v=cfg.expand_v,
+                    use_short_conv=cfg.use_short_conv,
+                    ffn_expansion_factor=cfg.ffn_expansion_factor,
+                )
+
+        def crossatt_block(i):
+            if i in cfg.crossatt_layer_idx:
+                return CrossAttention(
+                    dim=cfg.dim,
+                    num_heads=cfg.crossatt_num_heads,
+                    dropout=cfg.crossatt_dropout,
+                    layer_idx=i,
+                )
+            else:
+                return None
+
+        self.decoder_layers = nn.ModuleList(
+            [
+                DecoderBlockWithOptionalCrossAttention(
+                    simple_gla_block(i),
+                    crossatt_block(i),
+                )
+                for i in range(cfg.num_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        encoder_output: torch.Tensor,
+        decoder_input: torch.Tensor,
+        crossatt_mask: torch.Tensor | list[torch.Tensor] | None = None,
+        text_ids: torch.Tensor | None = None,
+        cache: FLACache | None = None,
+    ):
+        x = decoder_input
+        text_freqs = None
+
+        for layer in self.decoder_layers:
+            x = maybe_grad_ckpt(layer)(
+                x,
+                encoder_output,
+                text_freqs=text_freqs,
+                cache=cache,
+                crossatt_mask=crossatt_mask,
             )
-            cross = None
-            if cross_every and (i % cross_every == 0):
-                # CrossAttention(dim, num_heads=num_heads) -> module expects (x, context, mask)
-                cross = CrossAttention(dim, num_heads=num_heads)
-            decoder_layers.append(DecoderBlockWithOptionalCrossAttention(block, cross))
+        return x
 
-        self.decoder_layers = nn.ModuleList(decoder_layers)
+    def init_cache(self, max_seq_len, device):
+        return FLACache(num_states=len(self.decoder_layers) + 1)
+
+    def init_initial_state(self, batch_size=1, scale=1e-2, device="cpu"):
+        return tuple(
+            nn.Parameter(
+                torch.randn(
+                    batch_size,
+                    self.num_heads,
+                    self.head_dim,
+                    self.head_dim,
+                    device=device,
+                )
+                * scale
+            )
+            for _ in range(len(self.decoder_layers))
+        )
+    def init_initial_state_lora(self, lora:int=1, batch_size: int = 1, scale: float=1e-2, device: str="cpu"):
+        return tuple(
+            (
+                nn.Parameter(
+                torch.randn(
+                    batch_size,
+                    self.num_heads,
+                    self.head_dim,
+                    lora,
+                    device=device,
+                )
+                * scale
+            ),
+                 nn.Parameter(
+                    torch.randn(
+                        batch_size,
+                        self.num_heads,
+                        lora,
+                        self.head_dim,
+                        device=device,
+                    )
+                    * scale
+                )
+            )
+            for _ in range(len(self.decoder_layers))
+        )
+
+    def _get_query(self, audio_inputs: torch.Tensor, layer_idx: int):
+        assert self.decoder_layers[layer_idx].crossatt is not None
+        x = audio_inputs
+        for _, layer in zip(range(layer_idx - 1), self.decoder_layers):
+            x = layer(x, None)
+        return self.decoder_layers[layer_idx].crossatt._query(x)
+
+    def forward_first_n_layers(
+        self,
+        encoder_output: torch.Tensor,
+        decoder_input: torch.Tensor,
+        n_first_layers: int,
+        crossatt_mask: torch.Tensor | None = None,
+        cache: FLACache | None = None,
+    ):
+        x = decoder_input
+        if self.text_freqs_embd is not None:
+            text_freqs = torch.arange(encoder_output.shape[1], device=x.device)[None, :]
+            text_freqs = self.text_freqs_embd(text_freqs)
+        else:
+            text_freqs = None
+
+        for layer in self.decoder_layers[:n_first_layers]:
+            x = maybe_grad_ckpt(layer)(
+                x,
+                encoder_output,
+                text_freqs=text_freqs,
+                cache=cache,
+                crossatt_mask=crossatt_mask,
+            )
+        return x
 
-    # Backward compatibility with code expecting "prefill" API
     def prefill(
         self,
-        encoder_output: Optional[torch.Tensor],
+        encoder_output: torch.Tensor,
         decoder_input: torch.Tensor,
-        cache: Optional[Cache],
-        text_freqs: Optional[torch.Tensor] = None,
-        crossatt_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        return self(
-            encoder_output=encoder_output,
-            decoder_input=decoder_input,
-            cache=cache,
-            text_freqs=text_freqs,
-            crossatt_mask=crossatt_mask,
-        )
+        crossatt_mask: torch.Tensor | None = None,
+        cache: FLACache | None = None,
+    ):
+        return self(encoder_output, decoder_input, cache=cache, crossatt_mask=crossatt_mask)
 
-    def forward(
+    def decode_one(
         self,
-        encoder_output: Optional[torch.Tensor],
+        encoder_output: torch.Tensor,
         decoder_input: torch.Tensor,
-        cache: Optional[Cache],
-        text_freqs: Optional[torch.Tensor] = None,
-        crossatt_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
+        cache: Cache,
+        text_freqs: torch.Tensor | None = None,
+        crossatt_mask: torch.Tensor | None = None,
+    ):
         x = decoder_input
-        for idx, layer in enumerate(self.decoder_layers):
-            layer_cache = _cache_for_layer(cache, idx)
+        for layer in self.decoder_layers:
             x = layer(
                 x,
-                encoder_output=encoder_output,
+                encoder_output,
                 text_freqs=text_freqs,
-                cache=layer_cache,
+                cache=cache,
                 crossatt_mask=crossatt_mask,
             )
         return x