Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-# Абмежуем «шум» патокаў, каб пазбегнуць thrashing
 os.environ.setdefault("OMP_NUM_THREADS", "1")
 os.environ.setdefault("MKL_NUM_THREADS", "1")
 os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
@@ -25,7 +24,7 @@ from huggingface_hub import hf_hub_download
 from scipy.io.wavfile import write
 
 # ---------------------------------------------------------
-# 1) Клануем і падключаем coqui-ai-TTS (fork з падтрымкай BE)
+# 1) coqui-ai-TTS fork
 # ---------------------------------------------------------
 REPO_URL = "https://github.com/tuteishygpt/coqui-ai-TTS.git"
 REPO_DIR = "coqui-ai-TTS"
@@ -42,7 +41,7 @@ from TTS.tts.models.xtts import Xtts
 from TTS.tts.layers.xtts.tokenizer import VoiceBpeTokenizer, split_sentence
 
 # ---------------------------------------------------------
-# 2) Файлы мадэлі
+# 2) мадэльныя файлы
 # ---------------------------------------------------------
 repo_id = "archivartaunik/BE_XTTS_V2_10ep250k"
 model_dir = "./model"
@@ -59,7 +58,7 @@ for fname in ("model.pth", "config.json", "vocab.json", "voice.wav"):
         hf_hub_download(repo_id, filename=fname, local_dir=model_dir)
 
 # ---------------------------------------------------------
-# 3) Загрузка мадэлі і токенайзера
+# 3) загрузка мадэлі
 # ---------------------------------------------------------
 config = XttsConfig()
 config.load_json(config_file)
@@ -73,7 +72,6 @@ XTTS_MODEL.load_checkpoint(
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# GPU/CPU налады
 torch.set_num_threads(1)
 if device.startswith("cuda"):
     torch.backends.cuda.matmul.allow_tf32 = True
@@ -90,12 +88,12 @@ XTTS_MODEL.tokenizer = tokenizer
 # =========================================================
 # 4) Streaming-канфіг
 # =========================================================
-MIN_BUFFER_S = 0.02        # агульны мікс-бфер (стабільны для плэера)
-RUNTIME_FIRST_CHUNK_S = 0.015  # канкрэтна для 1-га чанка ў генерацыі (хутчэй)
+MIN_BUFFER_S = 0.03            # бяспечны выхадны буфер для плэера
+RUNTIME_FIRST_CHUNK_S = 0.02   # унутраны чанк у генерацыі
 FADE_S       = 0.004
 TOKENS_PER_STEP = 1
 ENABLE_TEXT_SPLITTING = True
-FIRST_SEGMENT_LIMIT = 120   # кароткі першы сегмент (хутчэйшы першы чанк)
+FIRST_SEGMENT_LIMIT = 160      # стабільная прасадыя для 1-га сегмента
 
 # -------------------- утыліты аўдыя ----------------------
 def _seconds_to_samples(sec: float, sr: int) -> int:
@@ -216,7 +214,7 @@ def init_stream_support():
 init_stream_support()
 
 # ---------------------------------------------------------
-# 5) Пастаянны кэш латэнтаў (CPU) + GPU-кэш (без дадатковых запытаў)
+# 5) пастаянны кэш латэнтаў (CPU) + GPU-кэш
 # ---------------------------------------------------------
 PERSIST_LATENTS_DIR = pathlib.Path("./latents_cache")
 PERSIST_LATENTS_DIR.mkdir(parents=True, exist_ok=True)
@@ -229,8 +227,8 @@ class LatentsMeta:
     sound_norm_refs: bool
     xtts_git: str | None = None
 
-LATENT_CACHE: dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}        # RAM CPU
-GPU_LATENT_CACHE: dict[Tuple[str, str], Tuple[torch.Tensor, torch.Tensor]] = {}  # RAM GPU
+LATENT_CACHE: dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
+GPU_LATENT_CACHE: dict[Tuple[str, str], Tuple[torch.Tensor, torch.Tensor]] = {}
 
 def _latents_key(path: str | None, meta: LatentsMeta) -> str:
     if path and os.path.exists(path):
@@ -250,18 +248,11 @@ def _latents_disk_path(key: str) -> pathlib.Path:
     return PERSIST_LATENTS_DIR / f"{key}.pt"
 
 def _save_latents_to_disk(key: str, gpt_cond_latent: torch.Tensor, speaker_embedding: torch.Tensor):
-    torch.save(
-        {
-            "gpt_cond_latent": gpt_cond_latent.cpu(),
-            "speaker_embedding": speaker_embedding.cpu(),
-        },
-        _latents_disk_path(key),
-    )
+    torch.save({"gpt_cond_latent": gpt_cond_latent.cpu(), "speaker_embedding": speaker_embedding.cpu()}, _latents_disk_path(key))
 
 def _load_latents_from_disk(key: str) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
     p = _latents_disk_path(key)
-    if not p.exists():
-        return None
+    if not p.exists(): return None
     obj = torch.load(p, map_location="cpu")
     return obj["gpt_cond_latent"], obj["speaker_embedding"]
 
@@ -276,7 +267,6 @@ def _compute_latents_cpu(path: str | None) -> Tuple[torch.Tensor, torch.Tensor]:
     return g.cpu(), s.cpu()
 
 def _latents_for(path: str | None, *, to_device: Optional[str] = None) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Вяртае латэнты з RAM/дыска; калі няма — палічыць на CPU і захаваць. Пры патрэбе — кэшуе і на GPU."""
     meta = LatentsMeta(
         model_id=repo_id,
         gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
@@ -286,21 +276,17 @@ def _latents_for(path: str | None, *, to_device: Optional[str] = None) -> Tuple[
     )
     key = _latents_key(path, meta)
 
-    # 1) CPU RAM
     if key in LATENT_CACHE:
         g, s = LATENT_CACHE[key]
     else:
-        # 2) дыск
         loaded = _load_latents_from_disk(key)
         if loaded is None:
-            # 3) палічыць на CPU і захаваць
             g, s = _compute_latents_cpu(path)
             _save_latents_to_disk(key, g, s)
         else:
             g, s = loaded
         LATENT_CACHE[key] = (g, s)
 
-    # 4) GPU-кэш (калі патрэбны)
     if to_device and to_device.startswith("cuda"):
         dev_key = (key, to_device)
         if dev_key in GPU_LATENT_CACHE:
@@ -309,19 +295,16 @@ def _latents_for(path: str | None, *, to_device: Optional[str] = None) -> Tuple[
         s2 = s.to(to_device, non_blocking=True)
         GPU_LATENT_CACHE[dev_key] = (g2, s2)
         return g2, s2
-
     return g, s
 
-# Аўтаматычна падлічым латэнты на CPU для default voice (адзін раз, без GPU-запытаў)
-DEFAULT_VOICE_PRECOMPUTE = True
-if DEFAULT_VOICE_PRECOMPUTE:
-    try:
-        _ = _latents_for(default_voice_file)  # CPU-разлік і захаванне, калі няма
-    except Exception as e:
-        print(f"[warn] precompute default voice latents failed: {e}")
+# аўтападлік для default voice (CPU) — без дадатковых запытаў
+try:
+    _ = _latents_for(default_voice_file)
+except Exception as e:
+    print(f"[warn] precompute default voice latents failed: {e}")
 
 # ---------------------------------------------------------
-# 6) Хэлперы: буферы + base64
+# 6) буферы + base64
 # ---------------------------------------------------------
 def _merge_for_file(chunks: List[np.ndarray]) -> np.ndarray:
     if not chunks: return np.zeros((0,), dtype=np.float32)
@@ -347,7 +330,7 @@ def _pcm_f32_to_b64(x: np.ndarray) -> str:
     return base64.b64encode(x.tobytes()).decode("ascii")
 
 # ---------------------------------------------------------
-# 7) Хуткі падзел тэксту: lightweight + fallback
+# 7) падзел тэксту: хуткі + fallback
 # ---------------------------------------------------------
 _SENT_END = re.compile(r"([\.!\?…]+[»\")\]]*\s+)")
 _WS = re.compile(r"\s+")
@@ -361,8 +344,7 @@ def _fast_split(text: str, limit: int) -> List[str]:
         end = m.end()
         parts.append(text[start:end].strip())
         start = end
-    if start < len(text):
-        parts.append(text[start:].strip())
+    if start < len(text): parts.append(text[start:].strip())
     chunks = []
     cur = ""
     for s in parts:
@@ -373,8 +355,7 @@ def _fast_split(text: str, limit: int) -> List[str]:
             if len(s) <= limit:
                 cur = s
             else:
-                w = _WS.split(s)
-                acc = ""
+                w = _WS.split(s); acc = ""
                 for tok in w:
                     if len(acc) + 1 + len(tok) <= limit:
                         acc = (acc + " " + tok).strip() if acc else tok
@@ -400,9 +381,7 @@ def _split_text_smart(text_in: str, lang_short: str, chunk_limit: int) -> List[s
         text_for_rest = tail
     else:
         text_for_rest = text_in
-
-    if not text_for_rest:
-        return parts or [text_in]
+    if not text_for_rest: return parts or [text_in]
 
     rest = _fast_split(text_for_rest, chunk_limit)
     if not rest or sum(len(x) for x in rest) < int(0.6 * len(text_for_rest)):
@@ -412,11 +391,10 @@ def _split_text_smart(text_in: str, lang_short: str, chunk_limit: int) -> List[s
             if rest2: rest = rest2
         except Exception:
             pass
-
     return parts + (rest or [text_for_rest])
 
 # ---------------------------------------------------------
-# 8) Асноўная функцыя TTS — стрим + фінальны файл + лагі
+# 8) TTS — стрим + фінальны файл + лагі
 # ---------------------------------------------------------
 @spaces.GPU(duration=60)
 def text_to_speech(belarusian_story, speaker_audio_file=None):
@@ -432,7 +410,6 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
     if not belarusian_story or str(belarusian_story).strip() == "":
         raise gr.Error("Увядзі хоць нейкі тэкст 🙂")
 
-    # Голас па змаўчанні
     if not speaker_audio_file or (
         not isinstance(speaker_audio_file, str)
         and getattr(speaker_audio_file, "name", "") == ""
@@ -443,13 +420,13 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
     lang_short = "be"
     chunk_limit = getattr(XTTS_MODEL.tokenizer, "char_limits", {}).get(lang_short, 250)
 
-    # --- Latents з CPU-/дыск-кэшу; GPU-кэш (без extra запытаў) ---
+    # Latents (кэш CPU/GPU)
     t_lat0 = time.perf_counter()
     to_dev = "cuda:0" if torch.cuda.is_available() else None
     gpt_cond_latent, speaker_embedding = _latents_for(speaker_audio_file, to_device=to_dev)
     t_lat1 = time.perf_counter()
 
-    # --- Split (хуткі) ---
+    # Split
     t_split0 = time.perf_counter()
     texts = _split_text_smart(text_in, lang_short, chunk_limit) if ENABLE_TEXT_SPLITTING else [text_in]
     if not texts: texts = [text_in]
@@ -463,10 +440,8 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
         "server_unaccounted_before_first_chunk_s": None,
         "file_write_s": None,
     }
-    # пачатковыя метрыкі
     yield ("", None, None, json.dumps(server_metrics))
 
-    # --- Генерацыя і стрим ---
     full_audio_chunks: List[np.ndarray] = []
     first_chunk_seen = False
     t_gen0 = time.perf_counter()
@@ -475,13 +450,12 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
         gen = XTTS_MODEL.generate(
             text=part, do_stream=True, language=lang_short,
             gpt_cond_latent=gpt_cond_latent, speaker_embedding=speaker_embedding,
-            min_buffer_s=RUNTIME_FIRST_CHUNK_S,           # 1-ы чанк хутчэй
+            min_buffer_s=RUNTIME_FIRST_CHUNK_S,
             tokens_per_step=TOKENS_PER_STEP,
-            stream_chunk_size_s=RUNTIME_FIRST_CHUNK_S,    # важна для native stream
+            stream_chunk_size_s=RUNTIME_FIRST_CHUNK_S,
             temperature=0.1, length_penalty=1.0, repetition_penalty=10.0,
             top_k=10, top_p=0.3,
         )
-        # На выхад у плэер — больш стабільны буфер MIN_BUFFER_S
         for buf in _chunker(gen, sampling_rate, MIN_BUFFER_S):
             if not first_chunk_seen:
                 t_first = time.perf_counter()
@@ -496,7 +470,6 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
                 yield (_pcm_f32_to_b64(buf), None, None, None)
             full_audio_chunks.append(buf)
 
-    # --- Фінал: WAV ---
     if not full_audio_chunks:
         yield ("__STOP__", None, None, json.dumps(server_metrics)); return
 
@@ -515,14 +488,14 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
     yield ("__STOP__", tmp.name, tmp.name, json.dumps(server_metrics))
 
 # ---------------------------------------------------------
-# 9) UI: лагі ў секундах + Play Final
+# 9) UI (лагі ў секундах + Play Final; без underrun’аў)
 # ---------------------------------------------------------
 examples = [
     ["Прывітанне! Гэта праверка жывога струменя беларускага TTS.", "Nestarka.wav"],
 ]
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Belarusian TTS — Streaming (GPU-кэш latents, хуткі 1-ы чанк) + фінальны файл")
+    gr.Markdown("## Belarusian TTS — Streaming (стабільны старт) + фінальны файл")
 
     with gr.Row():
         inp_text = gr.Textbox(lines=5, label="Тэкст на беларускай мове")
@@ -552,6 +525,9 @@ with gr.Blocks() as demo:
   const AC = window.AudioContext || window.webkitAudioContext;
   if (!AC) return;
 
+  const PRIME_CHUNKS = 2;       // мін. к-ць чанкаў перад стартаваннем гуку
+  let primeCounter = 0;
+
   function toSec(ms) {{ return (ms/1000); }}
   function fmtS(x) {{ return (x===null||x===undefined) ? "n/a" : x.toFixed(3) + " s"; }}
 
@@ -600,7 +576,7 @@ with gr.Blocks() as demo:
 
   if (!window.__wa) {{
     const ctx = new AC({{ sampleRate }});
-    const bufferSize = 1024;
+    const bufferSize = 2048; // большы буфер = менш underrun’аў
     const node = ctx.createScriptProcessor(bufferSize, 0, 1);
     let queue = [];
     let playing = false;
@@ -642,18 +618,31 @@ with gr.Blocks() as demo:
       get eos() {{ return eos; }},
       set eos(v) {{ eos = v; }},
       meta,
-      push: (f32) => {{ queue.push(f32); }},
+      push: (f32) => {{
+        queue.push(f32);
+        if (!meta.t_first_push_ms) {{
+          meta.t_first_push_ms = performance.now();
+          logUpdate();
+        }}
+        if (!playing && queue.length >= PRIME_CHUNKS) {{
+          // стартуем толькі калі ёсць мінімум 2 чанкі ў чарзе
+          window.__wa.start();
+        }}
+      }},
       start: async () => {{ try {{ await ctx.resume(); }} catch(e){{}} playing = true; logUpdate(); }},
       stop: () => {{ playing = false; logUpdate(); }},
-      reset: () => {{ playing = false; eos = false; queue = []; meta.t_first_push_ms = null; meta.t_first_audio_ms = null; logUpdate(); }},
+      reset: () => {{
+        playing = false; eos = false; queue = [];
+        primeCounter = 0;
+        meta.t_first_push_ms = null; meta.t_first_audio_ms = null;
+        logUpdate();
+      }},
       updateLog: logUpdate,
     }};
   }} else {{
     window.__wa.reset();
     window.__wa.meta.t_click_ms = performance.now();
   }}
-
-  window.__wa.start();
 }}
 """
 
@@ -663,16 +652,7 @@ with gr.Blocks() as demo:
     PUSH_JS = """
 (b64) => {
   if (!window.__wa || !b64) return;
-  const meta = window.__wa.meta || {};
-  if (b64 === "__STOP__") {
-    window.__wa.eos = true;
-    window.__wa.updateLog && window.__wa.updateLog();
-    return;
-  }
-  if (!meta.t_first_push_ms) {
-    meta.t_first_push_ms = performance.now();
-    window.__wa.updateLog && window.__wa.updateLog();
-  }
+  if (b64 === "__STOP__") { window.__wa.eos = true; window.__wa.updateLog && window.__wa.updateLog(); return; }
   const bin = atob(b64);
   const len = bin.length;
   const buf = new ArrayBuffer(len);
@@ -689,7 +669,7 @@ with gr.Blocks() as demo:
   try {
     if (js) {
       const obj = JSON.parse(js);
-      window.__wa.meta.server = obj;  // значэнні ў секундах
+      window.__wa.meta.server = obj;
       window.__wa.updateLog && window.__wa.updateLog();
     }
   } catch (e) {}
@@ -705,7 +685,6 @@ with gr.Blocks() as demo:
 }
 """
 
-    # кнопкі
     play_btn.click(fn=None, inputs=[], outputs=[], js=PLAY_JS)
     stop_btn.click(fn=None, inputs=[], outputs=[], js=STOP_JS)