Update app.py

app.py

@@ -1,5 +1,5 @@
 # Калі запускаеце ў чыстым асяроддзі (раскаментуйце):
-# !pip install -q gradio spaces huggingface_hub torch scipy tqdm gitpython
+# !pip install -q gradio spaces huggingface_hub torch scipy gitpython
 
 import os
 import sys
@@ -8,7 +8,7 @@ import time
 import tempfile
 import subprocess
 import inspect
-from typing import Iterator, Iterable, Optional, Tuple, Any
+from typing import Iterator, Iterable, Optional, Tuple, Any, List
 
 import spaces
 import gradio as gr
@@ -76,18 +76,34 @@ tokenizer = VoiceBpeTokenizer(vocab_file=vocab_file)
 XTTS_MODEL.tokenizer = tokenizer
 
 # =========================================================
-# 4) «Як у прыкладзе»: патч Xtts.generate / sample_stream
+# 4) «Як у transformers-stream-generator»: патч Xtts.generate/sample_stream
 # =========================================================
 
 # Канстанты латэнтнасці/буферу
 MIN_BUFFER_S = 0.050   # ~50 ms цэлявы буфер для аўдыя
-FADE_MS      = 8e-3    # кароткі cross-fade паміж чанкамі
+FADE_S       = 0.008   # кароткі cross-fade паміж чанкамі
 TOKENS_PER_STEP = 4    # памер кроку «токенаў» у fallback (BPE/субсловы)
 
 def _seconds_to_samples(sec: float, sr: int) -> int:
     return max(1, int(sec * sr))
 
-def _crossfade_concat(a: np.ndarray, b: np.ndarray, sr: int, fade_ms: float) -> np.ndarray:
+def _to_np_audio(x) -> np.ndarray:
+    """Прывесці аўдыя да 1D np.float32 на CPU (падтрымка dict/torch/np)."""
+    if isinstance(x, dict) and "wav" in x:
+        x = x["wav"]
+    if isinstance(x, torch.Tensor):
+        if x.dtype != torch.float32:
+            x = x.float()
+        x = x.detach().cpu().contiguous().view(-1)
+        return x.numpy()
+    x = np.asarray(x)
+    if x.ndim > 1:
+        x = x.reshape(-1)
+    if x.dtype != np.float32:
+        x = x.astype(np.float32, copy=False)
+    return x
+
+def _crossfade_concat(a: np.ndarray, b: np.ndarray, sr: int, fade_s: float) -> np.ndarray:
     """Плыўна зліць два кавалкі без клікаў."""
     if a.size == 0:
         return b.astype(np.float32, copy=False)
@@ -95,7 +111,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 = min(_seconds_to_samples(fade_ms, sr), a.size, b.size)
+    fade_n = min(_seconds_to_samples(fade_s, 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)
@@ -107,9 +123,8 @@ def _crossfade_concat(a: np.ndarray, b: np.ndarray, sr: int, fade_ms: float) ->
 
 def _bpe_prefixes(text: str, lang: str, step_tokens: int) -> Iterable[str]:
     """
-    Вяртае прэфіксы па BPE/субсловах; калі encode/decode недаступны — псэўда-токены (словы+прабелы).
+    Прэфіксы па BPE/субсловах; калі encode/decode недаступны — псэўда-токены (словы+прабелы).
     """
-    # 1) BPE праз VoiceBpeTokenizer, калі падтрымліваецца
     try:
         ids = tokenizer.encode(text, lang=lang)
         n = len(ids)
@@ -120,7 +135,6 @@ def _bpe_prefixes(text: str, lang: str, step_tokens: int) -> Iterable[str]:
         return
     except Exception:
         pass
-    # 2) Падстрахоўка: «словы+раздзяляльнікі»
     pseudo_tokens = re.findall(r"\S+|\s+", text)
     acc = ""
     for i in range(0, len(pseudo_tokens), step_tokens):
@@ -138,7 +152,7 @@ def _native_stream(
     **gen_kwargs,
 ) -> Iterator[np.ndarray]:
     """
-    Натыўны паток, калі ў форку ёсць model.inference_stream(...)-> iterator of PCM/ndarray.
+    Натыўны паток, калі ў форку ёсць model.inference_stream(...)-> iterator of PCM/ndarray/torch.Tensor.
     """
     sig = inspect.signature(model.inference_stream)
     call_kwargs = dict(
@@ -147,18 +161,15 @@ def _native_stream(
         gpt_cond_latent=gpt_cond_latent,
         speaker_embedding=speaker_embedding,
     )
-    # Перадаём тыповыя параметры генерацыі, калі яны ёсць у подпісе
     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:
         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)
+        yield _to_np_audio(out)
 
 def _fallback_incremental(
     model: Xtts,
@@ -175,7 +186,7 @@ def _fallback_incremental(
     emitted = 0
     for prefix in _bpe_prefixes(text, language, tokens_per_step):
         with torch.no_grad():
-            wav = model.inference(
+            out = model.inference(
                 text=prefix,
                 language=language,
                 gpt_cond_latent=gpt_cond_latent,
@@ -185,7 +196,8 @@ def _fallback_incremental(
                 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)
+            )
+        wav = _to_np_audio(out)
         new_part = wav[emitted:]
         emitted = wav.size
         if new_part.size:
@@ -193,8 +205,7 @@ def _fallback_incremental(
 
 class NewTTSGenerationMixin:
     """
-    «Як у transformers-stream-generator»: дадаём generate() і sample_stream()
-    у мадэль Xtts. return: або поўны wav (ndarray), або ітэратар чанкаў (ndarray).
+    «Як у transformers-stream-generator»: дадаём generate() і sample_stream() у Xtts.
     """
 
     @torch.inference_mode()
@@ -210,12 +221,7 @@ class NewTTSGenerationMixin:
         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,
@@ -228,9 +234,8 @@ class NewTTSGenerationMixin:
                 top_k=gen_kwargs.get("top_k", 10),
                 top_p=gen_kwargs.get("top_p", 0.3),
             )
-            return out["wav"].astype(np.float32)
+            return _to_np_audio(out)
 
-        # Стрымінгавы рэжым — як у прыкладзе: асобны генератар
         return self.sample_stream(
             text=text,
             language=language,
@@ -253,27 +258,20 @@ class NewTTSGenerationMixin:
         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
 
-
 def init_stream_support():
-    """Прапатчыць Xtts, дадаўшы generate/sample_stream (як у прыкладзе)."""
+    """Прапатчыць Xtts, дадаўшы generate/sample_stream."""
     Xtts.generate = NewTTSGenerationMixin.generate
     Xtts.sample_stream = NewTTSGenerationMixin.sample_stream
 
@@ -289,19 +287,18 @@ def _yield_buffered_chunks_for_gradio(
     target_s: float = MIN_BUFFER_S,
 ) -> Iterator[Tuple[int, np.ndarray]]:
     """
-    Назапашваем невялікі буфер (~50 ms), каб плэер Gradio паспеў «дагуляць»
-    і не накладваў наступны чанк.
+    Назапашваем невялікі буфер (~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)
+        c = _to_np_audio(c)
         if c.size == 0:
             continue
         if buf.size == 0:
             buf = c
         else:
-            buf = _crossfade_concat(buf, c, sr, FADE_MS)
+            buf = _crossfade_concat(buf, c, sr, FADE_S)
         if buf.size >= target_samples:
             yield (sr, buf)
             time.sleep(buf.size / float(sr))
@@ -311,14 +308,13 @@ def _yield_buffered_chunks_for_gradio(
         time.sleep(buf.size / float(sr))
 
 # ---------------------------------------------------------
-# 6) Асноўная функцыя TTS для Gradio (як у цябе, але праз model.generate do_stream)
+# 6) Асноўная функцыя TTS для Gradio (у стылі .generate(do_stream=True))
 # ---------------------------------------------------------
 @spaces.GPU(duration=60)
 def text_to_speech(belarusian_story, speaker_audio_file=None):
     """
     Streaming для gr.Audio:
-      - падобна да прыкладу з transformers-stream-generator: model.generate(..., do_stream=True)
-      - аддаём невялікія чанкі (sr, chunk) з мінімальнай затрымкай;
+      - model.generate(..., do_stream=True) -> чанкі (sr, chunk) з мінімальнай затрымкай;
       - у фінале — шлях да поўнага WAV.
     """
     if not belarusian_story or str(belarusian_story).strip() == "":
@@ -342,7 +338,6 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
     except Exception as e:
         raise gr.Error(f"Памылка пры атрыманні латэнтаў голасу: {e}")
 
-    # --- Генератар па аналагіі з .generate(... do_stream=True) ---
     generator = XTTS_MODEL.generate(
         text=str(belarusian_story).strip(),
         do_stream=True,
@@ -358,20 +353,18 @@ def text_to_speech(belarusian_story, speaker_audio_file=None):
         top_p=0.3,
     )
 
-    # Будзем назапашваць увесь аўдыё для фінальнага WAV
-    full_audio_chunks: list[np.ndarray] = []
+    full_audio_chunks: List[np.ndarray] = []
 
-    # Аддаём у Gradio дробныя порцыі з невялікім буферам і рэальным «сном»
     for sr, chunk in _yield_buffered_chunks_for_gradio(generator, sampling_rate, MIN_BUFFER_S):
         full_audio_chunks.append(chunk)
         yield (sr, chunk)
 
-    # Гатовы поўны 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)
+        full_audio = _crossfade_concat(full_audio, full_audio_chunks[i], sampling_rate, FADE_S)
 
     try:
         tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
@@ -426,7 +419,7 @@ analytics_script = """
 with gr.Blocks() as demo:
     gr.HTML(analytics_script)
     gr.Interface(
-        fn=text_to_speech,  # генератар
+        fn=text_to_speech,
         inputs=[
             gr.Textbox(lines=5, label="Тэкст на беларускай мове"),
             gr.Audio(
@@ -436,7 +429,7 @@ with gr.Blocks() as demo:
             ),
         ],
         outputs=gr.Audio(
-            type="filepath",   # падчас стриму — (sr, ndarray); у фінале — шлях
+            type="filepath",
             label="Згенераванае аўдыя (па токенах, мінімальная затрымка)",
             autoplay=True,
         ),