True frame-level streaming via engine.infer_stream (@spaces.GPU generator)

app.py

@@ -19,7 +19,6 @@ import spaces
 
 from vieneu import Vieneu
 from vieneu_utils.core_utils import split_text_into_chunks
-from vieneu_utils.phonemize_text import phonemize_text_with_emotions
 
 # ── Load model once, on GPU (CUDA emulation makes this valid at startup) ───────
 print("⏳ Loading VieNeu-TTS v3 Turbo (PyTorch / CUDA) ...")
@@ -178,54 +177,42 @@ def synthesize_conversation(script, silence_s, temperature, max_new_frames):
     return (SR, full), info
 
 
-# ── Streaming (Kokoro pattern: CPU generator drives a streaming gr.Audio, and
-#    each text chunk's GPU work is a short, separate @spaces.GPU call so ZeroGPU
-#    can allocate/release per chunk). Latency to first audio ≈ first chunk time. ─
-@spaces.GPU(duration=30)
-def _encode_ref_gpu(ref_audio):
-    """Encode a reference clip to MOSS codes on the GPU (for streamed cloning)."""
-    return np.asarray(tts.encode_reference(ref_audio))
-
-
+# ── True frame-level streaming ────────────────────────────────────────────────
+# The engine's ``infer_stream`` yields the waveform in small ~chunk_frames pieces
+# (with an adaptive lead-in) as the model generates — that's what makes playback
+# smooth, not chunky. On ZeroGPU the whole thing must hold the GPU, so this is a
+# single ``@spaces.GPU`` *generator*: it keeps the GPU for the session and pushes
+# each small piece out as it is produced. Long text is split by ``max_chars`` and
+# each chunk is streamed in turn (so arbitrarily long input still works).
 @spaces.GPU(duration=_gpu_duration)
-def _infer_chunk_gpu(phonemes, ref_codes, voice_token_id, temperature, top_k,
-                     top_p, repetition_penalty, max_new_frames):
-    wav = tts.engine.infer(
-        text="", phonemes=phonemes, ref_codes=ref_codes,
-        emotion="natural", voice_token_id=voice_token_id,
-        temperature=float(temperature), top_k=int(top_k), top_p=float(top_p),
-        repetition_penalty=float(repetition_penalty), max_new_frames=int(max_new_frames),
-    )
-    return np.asarray(wav, dtype=np.float32)
-
-
 def stream_synthesize(text, voice, ref_audio, temperature, top_k, top_p,
                       repetition_penalty, max_new_frames, max_chars):
     text = (text or "").strip()
     if not text:
         raise gr.Error("Vui lòng nhập văn bản cần đọc.")
 
-    # Resolve the voice once: preset → reserved token + cached codes (CPU only);
-    # cloning → encode the reference on the GPU a single time up front.
+    # Preset → reserved token + cached codes; cloning → encode the reference once.
     if ref_audio:
-        ref_codes, voice_token_id = _encode_ref_gpu(ref_audio), None
+        ref_codes, voice_token_id = tts.encode_reference(ref_audio), None
     else:
         ref_codes, voice_token_id = tts._resolve_v3_ref(voice, None, None)
 
-    # Smaller chunks = lower latency to first audio.
-    chunks = split_text_into_chunks(text, max_chars=min(int(max_chars), 140))
-
     first = True
-    for chunk in chunks:
-        phonemes = phonemize_text_with_emotions(chunk)
-        wav = _infer_chunk_gpu(phonemes, ref_codes, voice_token_id, temperature,
-                               top_k, top_p, repetition_penalty, max_new_frames)
-        yield (SR, wav)
-        if first:
-            # Workaround for the Gradio streaming quirk where the first chunk can
-            # be dropped; a 1-sample silence flushes the player (same as Kokoro).
-            first = False
-            yield (SR, np.zeros(1, dtype=np.float32))
+    for chunk in split_text_into_chunks(text, max_chars=int(max_chars)):
+        for wav in tts.engine.infer_stream(
+            text=chunk, ref_codes=ref_codes, voice_token_id=voice_token_id,
+            emotion="natural", temperature=float(temperature), top_k=int(top_k),
+            top_p=float(top_p), repetition_penalty=float(repetition_penalty),
+            max_new_frames=int(max_new_frames), chunk_frames=25,
+        ):
+            wav = np.asarray(wav, dtype=np.float32)
+            if wav.size == 0:
+                continue
+            yield (SR, wav)
+            if first:
+                # Flush the first piece (Gradio streaming sometimes drops it).
+                first = False
+                yield (SR, np.zeros(1, dtype=np.float32))
 
 
 # ── UI ────────────────────────────────────────────────────────────────────────
@@ -294,7 +281,7 @@ with gr.Blocks(theme=theme, title="VieNeu-TTS v3 Turbo") as demo:
                             temperature_in = gr.Slider(0.1, 1.5, value=0.8, step=0.05, label="temperature")
                             top_p_in = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="top_p")
                         with gr.Row():
-                            top_k_in = gr.Slider(1, 100, value=50, step=1, label="top_k")
+                            top_k_in = gr.Slider(1, 100, value=25, step=1, label="top_k")
                             rep_pen_in = gr.Slider(1.0, 2.0, value=1.2, step=0.05, label="repetition_penalty")
                         with gr.Row():
                             max_frames_in = gr.Slider(50, 1200, value=300, step=10, label="max_new_frames (mỗi đoạn)")
@@ -313,8 +300,9 @@ with gr.Blocks(theme=theme, title="VieNeu-TTS v3 Turbo") as demo:
                                 stream_btn = gr.Button("⚡ Phát trực tiếp", variant="primary")
                                 stop_btn = gr.Button("⏹ Dừng", variant="stop")
                             gr.Markdown(
-                                "_Phát dần theo từng câu (độ trễ thấp). Lần bấm **Phát trực tiếp** "
-                                "đầu tiên đôi khi Gradio nuốt mất chunk đầu — bấm lại là được._"
+                                "_Phát theo thời gian thực, mượt từng đoạn ngắn ngay khi model sinh. "
+                                "Lần bấm **Phát trực tiếp** đầu tiên đôi khi Gradio nuốt mất đoạn đầu "
+                                "— bấm lại là được._"
                             )
 
             run_btn.click(