Update app.py

app.py

@@ -3,18 +3,19 @@
 
 import os
 import sys
+import re
 import time
 import tempfile
 import subprocess
 import inspect
-import re
+from typing import Iterator, Iterable, Optional, Tuple, Any
 
 import spaces
 import gradio as gr
 import torch
+import numpy as np
 from huggingface_hub import hf_hub_download
 from scipy.io.wavfile import write
-import numpy as np
 
 # ---------------------------------------------------------
 # 1) Клануем і падключаем coqui-ai-TTS (fork з падтрымкай BE)
@@ -74,13 +75,13 @@ sampling_rate = int(XTTS_MODEL.config.audio["sample_rate"])
 tokenizer = VoiceBpeTokenizer(vocab_file=vocab_file)
 XTTS_MODEL.tokenizer = tokenizer
 
-# ---------------------------------------------------------
-# 4) Патокавая TTS па "токенах" з мінімальнай затрымкай
-#    (натыйўны стримінг -> fallback інкрементальны прэфікс)
-# ---------------------------------------------------------
+# =========================================================
+# 4) «Як у прыкладзе»: патч Xtts.generate / sample_stream
+# =========================================================
 
-MIN_BUFFER_MS = 0.05   # мэтавы мінімальны буфер ~50 ms
-FADE_MS = 0.008        # кароткі cross-fade паміж чанкамі
+# Канстанты латэнтнасці/буферу
+MIN_BUFFER_S = 0.050   # ~50 ms цэлявы буфер для аўдыя
+FADE_MS      = 8e-3    # кароткі cross-fade паміж чанкамі
 TOKENS_PER_STEP = 4    # памер кроку «токенаў» у fallback (BPE/субсловы)
 
 def _seconds_to_samples(sec: float, sr: int) -> int:
@@ -94,8 +95,7 @@ def _crossfade_concat(a: np.ndarray, b: np.ndarray, sr: int, fade_ms: float) ->
         return a.astype(np.float32, copy=False)
     a = a.astype(np.float32, copy=False)
     b = b.astype(np.float32, copy=False)
-    fade_n = _seconds_to_samples(fade_ms, sr)
-    fade_n = min(fade_n, a.size, b.size)
+    fade_n = min(_seconds_to_samples(fade_ms, sr), a.size, b.size)
     if fade_n <= 1:
         return np.concatenate([a, b], axis=0)
     fade_out = np.linspace(1.0, 0.0, fade_n, endpoint=True, dtype=np.float32)
@@ -105,143 +105,221 @@ def _crossfade_concat(a: np.ndarray, b: np.ndarray, sr: int, fade_ms: float) ->
     rest = b[fade_n:]
     return np.concatenate([head, tail, rest], axis=0)
 
-def _merge_chunks_with_crossfade(chunks: list[np.ndarray], sr: int) -> np.ndarray:
-    merged = np.zeros((0,), dtype=np.float32)
-    for c in chunks:
-        if c is None or np.asarray(c).size == 0:
-            continue
-        merged = _crossfade_concat(merged, np.asarray(c, dtype=np.float32), sr, FADE_MS)
-    return merged
-
-def _yield_buffered_chunks(chunks: list[np.ndarray], sr: int, target_ms: float):
-    """
-    Збіраем маленькія кавалкі пакуль не назапасім ~target_ms,
-    пасля чаго yield (sr, buffer) і спім роўна на працягласць buffer.
-    """
-    target_samples = _seconds_to_samples(target_ms, sr)
-    buf = np.zeros((0,), dtype=np.float32)
-    for c in chunks:
-        if c is None:
-            continue
-        c = np.asarray(c, dtype=np.float32)
-        if c.size == 0:
-            continue
-        if buf.size == 0:
-            buf = c
-        else:
-            buf = _crossfade_concat(buf, c, sr, FADE_MS)
-        if buf.size >= target_samples:
-            yield (sr, buf)
-            # даём плэеру «дагуляць» без накладання
-            time.sleep(buf.size / float(sr))
-            buf = np.zeros((0,), dtype=np.float32)
-    if buf.size:
-        yield (sr, buf)
-        time.sleep(buf.size / float(sr))
-
-def _bpe_prefixes(text: str, lang: str, step_tokens: int):
+def _bpe_prefixes(text: str, lang: str, step_tokens: int) -> Iterable[str]:
     """
-    Вяртае паступовыя прэфіксы тэксту па BPE/субсловах, калі атрымліваецца.
-    Інакш — fallback па «псэўда-токенах» (словы+прабелы/пунктуацыя).
+    Вяртае прэфіксы па BPE/субсловах; калі encode/decode недаступны — псэўда-токены (словы+прабелы).
     """
-    # 1) Спроба праз VoiceBpeTokenizer (калі ёсць encode/decode)
+    # 1) BPE праз VoiceBpeTokenizer, калі падтрымліваецца
     try:
-        # у вашым форку можа быць encode(text, lang=...), decode(ids, lang=...)
         ids = tokenizer.encode(text, lang=lang)
         n = len(ids)
         for k in range(step_tokens, n + 1, step_tokens):
-            prefix = tokenizer.decode(ids[:k], lang=lang)
-            yield prefix
+            yield tokenizer.decode(ids[:k], lang=lang)
         if n % step_tokens != 0:
             yield tokenizer.decode(ids, lang=lang)
         return
     except Exception:
         pass
-
-    # 2) Падстрахоўка: разбі��ь на «словы+знакі»
+    # 2) Падстрахоўка: «словы+раздзяляльнікі»
     pseudo_tokens = re.findall(r"\S+|\s+", text)
-    buff = ""
+    acc = ""
     for i in range(0, len(pseudo_tokens), step_tokens):
-        buff = "".join(pseudo_tokens[: i + step_tokens])
-        yield buff
-    if buff.strip() != text.strip():
+        acc = "".join(pseudo_tokens[: i + step_tokens])
+        yield acc
+    if acc.strip() != text.strip():
         yield text
 
-def _stream_from_model_native(text: str, gpt_cond_latent, speaker_embedding, sr: int, lang: str):
+def _native_stream(
+    model: Xtts,
+    text: str,
+    language: str,
+    gpt_cond_latent: Any,
+    speaker_embedding: Any,
+    **gen_kwargs,
+) -> Iterator[np.ndarray]:
     """
-    Калі ў форку ёсць натыўны струмень (inference_stream) — карыстаемся ім.
-    Павінен yield'іць PCM фрагменты па меры дэкавання.
+    Натыўны паток, калі ў форку ёсць model.inference_stream(...)-> iterator of PCM/ndarray.
     """
-    # Сфармуем kwargs у залежнасці ад подпісу функцыі
-    common_kwargs = dict(
+    sig = inspect.signature(model.inference_stream)
+    call_kwargs = dict(
         text=text,
-        language=lang,
+        language=language,
         gpt_cond_latent=gpt_cond_latent,
         speaker_embedding=speaker_embedding,
-        temperature=0.1,
-        length_penalty=1.0,
-        repetition_penalty=10.0,
-        top_k=10,
-        top_p=0.3,
     )
-    sig = inspect.signature(XTTS_MODEL.inference_stream)
+    # Перадаём тыповыя параметры генерацыі, калі яны ёсць у подпісе
+    for k in ("temperature", "length_penalty", "repetition_penalty", "top_k", "top_p"):
+        if k in gen_kwargs and k in sig.parameters:
+            call_kwargs[k] = gen_kwargs[k]
+    # Памер стрим-чанка (секунды), калі ёсць у подпісе
     if "stream_chunk_size_s" in sig.parameters:
-        common_kwargs["stream_chunk_size_s"] = MIN_BUFFER_MS
-
-    gen = XTTS_MODEL.inference_stream(**common_kwargs)
-    raw_chunks = []
-    for out in gen:
-        cur = out["wav"] if isinstance(out, dict) and "wav" in out else np.asarray(out, dtype=np.float32)
-        cur = cur.astype(np.float32, copy=False)
-        raw_chunks.append(cur)
-        # выдаём дробнымі порцыямі з невялікім буферам
-        yield from _yield_buffered_chunks([cur], sr, MIN_BUFFER_MS)
-
-    # Сабраць «хвост» у адзін WAV з лёгкім cross-fade
-    if raw_chunks:
-        final_full = _merge_chunks_with_crossfade(raw_chunks, sr)
-        yield ("__FINAL__", final_full)
-
-def _stream_fallback_incremental(text: str, gpt_cond_latent, speaker_embedding, sr: int, lang: str):
+        call_kwargs["stream_chunk_size_s"] = float(gen_kwargs.get("min_buffer_s", MIN_BUFFER_S))
+
+    generator = model.inference_stream(**call_kwargs)
+    for out in generator:
+        arr = out["wav"] if isinstance(out, dict) and "wav" in out else np.asarray(out, dtype=np.float32)
+        yield arr.astype(np.float32, copy=False)
+
+def _fallback_incremental(
+    model: Xtts,
+    text: str,
+    language: str,
+    gpt_cond_latent: Any,
+    speaker_embedding: Any,
+    tokens_per_step: int,
+    **gen_kwargs,
+) -> Iterator[np.ndarray]:
     """
-    Fallback: павялічваем прэфікс тэксту па токенах і кожны раз
-    генеруем гукавыя дадаткі (толькі «хвост» новай версіі).
+    Fallback: павялічваем прэфікс па токенах і вяртаем ТОЛЬКІ «новую» частку гуку.
     """
     emitted = 0
-    last_full = np.zeros((0,), dtype=np.float32)
-
-    for prefix in _bpe_prefixes(text, lang, TOKENS_PER_STEP):
+    for prefix in _bpe_prefixes(text, language, tokens_per_step):
         with torch.no_grad():
-            wav = XTTS_MODEL.inference(
+            wav = model.inference(
                 text=prefix,
-                language=lang,
+                language=language,
                 gpt_cond_latent=gpt_cond_latent,
                 speaker_embedding=speaker_embedding,
-                temperature=0.1,
-                length_penalty=1.0,
-                repetition_penalty=10.0,
-                top_k=10,
-                top_p=0.3,
+                temperature=gen_kwargs.get("temperature", 0.1),
+                length_penalty=gen_kwargs.get("length_penalty", 1.0),
+                repetition_penalty=gen_kwargs.get("repetition_penalty", 10.0),
+                top_k=gen_kwargs.get("top_k", 10),
+                top_p=gen_kwargs.get("top_p", 0.3),
             )["wav"].astype(np.float32)
-
-        # бярэм толькі новую частку адносна ўжо аддадзенага
         new_part = wav[emitted:]
-        if new_part.size > 0:
-            yield from _yield_buffered_chunks([new_part], sr, MIN_BUFFER_MS)
-            emitted = wav.size
-            last_full = wav
+        emitted = wav.size
+        if new_part.size:
+            yield new_part
+
+class NewTTSGenerationMixin:
+    """
+    «Як у transformers-stream-generator»: дадаём generate() і sample_stream()
+    у мадэль Xtts. return: або поўны wav (ndarray), або ітэратар чанкаў (ndarray).
+    """
+
+    @torch.inference_mode()
+    def generate(
+        self: Xtts,
+        text: Optional[str] = None,
+        *,
+        do_stream: bool = False,
+        language: str = "be",
+        gpt_cond_latent: Any = None,
+        speaker_embedding: Any = None,
+        min_buffer_s: float = MIN_BUFFER_S,
+        tokens_per_step: int = TOKENS_PER_STEP,
+        **gen_kwargs,
+    ):
+        """
+        Калі do_stream=False -> вяртае поўны wav (ndarray).
+        Калі do_stream=True  -> вяртае генератар чанкаў wav (Iterator[np.ndarray]).
+        """
+        assert isinstance(text, str) and text.strip(), "text is required"
+        # Блакіруючы рэжым — адным махам
+        if not do_stream:
+            out = self.inference(
+                text=text,
+                language=language,
+                gpt_cond_latent=gpt_cond_latent,
+                speaker_embedding=speaker_embedding,
+                temperature=gen_kwargs.get("temperature", 0.1),
+                length_penalty=gen_kwargs.get("length_penalty", 1.0),
+                repetition_penalty=gen_kwargs.get("repetition_penalty", 10.0),
+                top_k=gen_kwargs.get("top_k", 10),
+                top_p=gen_kwargs.get("top_p", 0.3),
+            )
+            return out["wav"].astype(np.float32)
+
+        # Стрымінгавы рэжым — як у прыкладзе: асобны генератар
+        return self.sample_stream(
+            text=text,
+            language=language,
+            gpt_cond_latent=gpt_cond_latent,
+            speaker_embedding=speaker_embedding,
+            min_buffer_s=min_buffer_s,
+            tokens_per_step=tokens_per_step,
+            **gen_kwargs,
+        )
+
+    @torch.inference_mode()
+    def sample_stream(
+        self: Xtts,
+        *,
+        text: str,
+        language: str,
+        gpt_cond_latent: Any,
+        speaker_embedding: Any,
+        min_buffer_s: float = MIN_BUFFER_S,
+        tokens_per_step: int = TOKENS_PER_STEP,
+        **gen_kwargs,
+    ) -> Iterator[np.ndarray]:
+        """
+        Вяртае генератар чанкаў wav. Стараемся даваць маленькія кавалкі як мага часцей.
+        """
+        # 1) Калі ёсць натыўны паток — проста перасылаем яго
+        if hasattr(self, "inference_stream"):
+            for chunk in _native_stream(
+                self, text, language, gpt_cond_latent, speaker_embedding, min_buffer_s=min_buffer_s, **gen_kwargs
+            ):
+                # тут мы не чакаем — верхні слой сам злімітуе плынь буферам
+                yield chunk
+            return
+
+        # 2) Інакш — інкрементальны fallback па токенах
+        for chunk in _fallback_incremental(
+            self, text, language, gpt_cond_latent, speaker_embedding, tokens_per_step, **gen_kwargs
+        ):
+            yield chunk
+
 
-    if last_full.size:
-        yield ("__FINAL__", last_full)
+def init_stream_support():
+    """Прапатчыць Xtts, дадаўшы generate/sample_stream (як у прыкладзе)."""
+    Xtts.generate = NewTTSGenerationMixin.generate
+    Xtts.sample_stream = NewTTSGenerationMixin.sample_stream
+
+# Актывуем стрим-падтрымку
+init_stream_support()
+
+# ---------------------------------------------------------
+# 5) Службовыя функцыі для Gradio (буферы, cross-fade, затрымкі)
+# ---------------------------------------------------------
+def _yield_buffered_chunks_for_gradio(
+    chunks: Iterable[np.ndarray],
+    sr: int,
+    target_s: float = MIN_BUFFER_S,
+) -> Iterator[Tuple[int, np.ndarray]]:
+    """
+    Назапашваем невялікі буфер (~50 ms), каб плэер Gradio паспеў «дагуляць»
+    і не накладваў наступны чанк.
+    """
+    target_samples = _seconds_to_samples(target_s, sr)
+    buf = np.zeros((0,), dtype=np.float32)
+    for c in chunks:
+        c = np.asarray(c, dtype=np.float32)
+        if c.size == 0:
+            continue
+        if buf.size == 0:
+            buf = c
+        else:
+            buf = _crossfade_concat(buf, c, sr, FADE_MS)
+        if buf.size >= target_samples:
+            yield (sr, buf)
+            time.sleep(buf.size / float(sr))
+            buf = np.zeros((0,), dtype=np.float32)
+    if buf.size:
+        yield (sr, buf)
+        time.sleep(buf.size / float(sr))
 
+# ---------------------------------------------------------
+# 6) Асноўная функцыя TTS для Gradio (як у цябе, але праз model.generate do_stream)
+# ---------------------------------------------------------
 @spaces.GPU(duration=60)
 def text_to_speech(belarusian_story, speaker_audio_file=None):
     """
-    Патокавы вывад з мінімальнай затрымкай.
-    - Крок 1: латэнты голасу.
-    - Крок 2: спроба натыўнага streaming з мадэлі.
-    - Крок 3: fallback — інкрементальны прэфікс (токен-крокі).
-    Выхад для gr.Audio: шмат (sr, chunk) + у фінале шлях да поўнага WAV.
+    Streaming для gr.Audio:
+      - падобна да прыкладу з transformers-stream-generator: model.generate(..., do_stream=True)
+      - аддаём невялікія чанкі (sr, chunk) з мінімальнай затрымкай;
+      - у фінале — шлях да поўнага WAV.
     """
     if not belarusian_story or str(belarusian_story).strip() == "":
         raise gr.Error("Увядзі хоць нейкі тэкст 🙂")
@@ -264,44 +342,46 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
     except Exception as e:
         raise gr.Error(f"Памылка пры атрыманні латэнтаў голасу: {e}")
 
-    lang = "be"
-    full_audio = None
+    # --- Генератар па аналагіі з .generate(... do_stream=True) ---
+    generator = XTTS_MODEL.generate(
+        text=str(belarusian_story).strip(),
+        do_stream=True,
+        language="be",
+        gpt_cond_latent=gpt_cond_latent,
+        speaker_embedding=speaker_embedding,
+        min_buffer_s=MIN_BUFFER_S,
+        tokens_per_step=TOKENS_PER_STEP,
+        temperature=0.1,
+        length_penalty=1.0,
+        repetition_penalty=10.0,
+        top_k=10,
+        top_p=0.3,
+    )
+
+    # Будзем назапашваць увесь аўдыё для фінальнага WAV
+    full_audio_chunks: list[np.ndarray] = []
 
-    # 1) Спачатку — натыўны паток, калі ёсць
-    try:
-        if hasattr(XTTS_MODEL, "inference_stream"):
-            for out in _stream_from_model_native(
-                belarusian_story, gpt_cond_latent, speaker_embedding, sampling_rate, lang
-            ):
-                if isinstance(out, tuple) and out and out[0] == "__FINAL__":
-                    full_audio = out[1]
-                else:
-                    yield out
-        else:
-            raise AttributeError("No native inference_stream in this build.")
-    except Exception:
-        # 2) fallback — інкрементальны прэфікс (токен-крокі)
-        for out in _stream_fallback_incremental(
-            belarusian_story, gpt_cond_latent, speaker_embedding, sampling_rate, lang
-        ):
-            if isinstance(out, tuple) and out and out[0] == "__FINAL__":
-                full_audio = out[1]
-            else:
-                yield out
+    # Аддаём у Gradio дробныя порцыі з невялікім буферам і рэальным «сном»
+    for sr, chunk in _yield_buffered_chunks_for_gradio(generator, sampling_rate, MIN_BUFFER_S):
+        full_audio_chunks.append(chunk)
+        yield (sr, chunk)
 
-    if full_audio is None:
+    # Гатовы поўны WAV
+    if not full_audio_chunks:
         raise gr.Error("Нічога не згенеравана. Праверце ўваходныя даныя або лагі.")
+    full_audio = full_audio_chunks[0]
+    for i in range(1, len(full_audio_chunks)):
+        full_audio = _crossfade_concat(full_audio, full_audio_chunks[i], sampling_rate, FADE_MS)
 
-    # Фінальны WAV у temp-файл
     try:
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
-        write(temp_file.name, sampling_rate, full_audio.astype(np.float32))
-        yield temp_file.name
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        write(tmp.name, sampling_rate, full_audio.astype(np.float32))
+        yield tmp.name
     except Exception as e:
         raise gr.Error(f"Памылка пры запісе фінальнага WAV: {e}")
 
 # ---------------------------------------------------------
-# 5) Прыклады (тэкст + файл голасу)
+# 7) Прыклады (тэкст + файл голасу)
 # ---------------------------------------------------------
 examples = [
     [
@@ -341,7 +421,7 @@ analytics_script = """
 """
 
 # ---------------------------------------------------------
-# 6) Gradio UI (аўтапрайграванне, мінімальная затрымка)
+# 8) Gradio UI (аўтапрайграванне)
 # ---------------------------------------------------------
 with gr.Blocks() as demo:
     gr.HTML(analytics_script)
@@ -356,13 +436,14 @@ with gr.Blocks() as demo:
             ),
         ],
         outputs=gr.Audio(
-            type="filepath",   # прымае (sr, ndarray) падчас стриму і фінальны шлях у фінале
+            type="filepath",   # падчас стриму — (sr, ndarray); у фінале — шлях
             label="Згенераванае аўдыя (па токенах, мінімальная затрымка)",
             autoplay=True,
         ),
-        title="Belarusian TTS — Token Streaming (мінімальная затрымка)",
+        title="Belarusian TTS — Token Streaming (як у transformers-stream-generator)",
         description="""
-        <p>Вывод гуку <b>па токенах</b> з буферам ~50&nbsp;мс. Калі мадэль падтрымлівае <code>inference_stream</code> — выкарыстоўваецца ён; інакш працуе інкрементальны fallback па BPE.</p>
+        <p>Мадэль <code>Xtts</code> мае метады <code>generate()</code> і <code>sample_stream()</code>, як у прыкладзе.
+        Калі даступны <code>inference_stream</code>, выкарыстоўваем яго; інакш — інкрементальна па «токенах» з ~50&nbsp;мс буферам.</p>
         """,
         examples=examples,
         cache_examples=False,