app.py

"""
VieNeu-TTS v3 Turbo — Hugging Face ZeroGPU Space
================================================
Vietnamese text-to-speech, 48 kHz, with built-in voices + instant voice cloning.
Shows generation speed / RTF per request.
ZeroGPU notes:
  * The model is placed on ``cuda`` at module import (ZeroGPU runs a CUDA
    emulation outside ``@spaces.GPU`` so this is the recommended, fastest path).
  * Real GPU compute happens only inside the ``@spaces.GPU`` decorated functions.
"""
import os
import time
import logging

import numpy as np
import soundfile as sf
import gradio as gr
import spaces

# ── Logging request infer của người dùng (xem ở tab "Logs" của Space) ──────────
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
logger = logging.getLogger("vieneu.space")


def _log_infer(mode: str, text: str, *, voice=None, cloned: bool = False) -> None:
    """Log một request infer: chế độ, giọng/clone, độ dài và nội dung text.

    Dùng repr() để text gói gọn trên một dòng log.
    """
    src = "clone" if cloned else f"voice={voice}"
    logger.info("[%s] %s | chars=%d | text=%r", mode, src, len(text or ""), text)

from vieneu import Vieneu
from vieneu_utils.core_utils import join_audio_chunks, gaps_to_silence
from vieneu_utils.phonemize_text import (
    phonemize_text_with_emotions,
    normalize_to_chunks_v3_with_gaps,
)

# ── TẠM THỜI: vá lỗi ZeroGPU "No CUDA GPUs are available" lúc import ───────────
# Bản `spaces` mới không cho safetensors materialize tensor thẳng lên CUDA khi
# chạy ngoài @spaces.GPU. Ta ép load_model nạp weights lên CPU; thư viện
# (_apply_dtype) sau đó tự .to(cuda) — đường .to() này ZeroGPU defer được nên
# hợp lệ ở import time. (Bản vá vĩnh viễn nằm trong vieneu hub_load_v3_turbo.)
# Phải patch TRƯỚC khi gọi Vieneu(...).
import vieneu._v3_turbo_engine.hub_load_v3_turbo as _hub_load

_orig_load_model = _hub_load.load_model

def _load_model_cpu(model, filename, *args, **kwargs):
    kwargs.pop("device", None)          # bỏ device -> nạp lên CPU
    return _orig_load_model(model, filename, *args, **kwargs)

_hub_load.load_model = _load_model_cpu

# ── Load model once, on GPU (CUDA emulation makes this valid at startup) ───────
print("⏳ Loading VieNeu-TTS v3 Turbo (PyTorch / CUDA) ...")
tts = Vieneu(
    mode="v3turbo",
    device="cuda",       # ZeroGPU: keep weights on cuda from the start
    backend="pytorch",   # force the PyTorch engine (ONNX is the CPU-only path)
    hf_token=os.getenv("HF_TOKEN"),
)
print("✅ Model ready.")

SR = tts.sample_rate
PRESET_VOICES = tts.list_preset_voices()                  # [(label, voice_id), ...]
VOICE_CHOICES = [(label, vid) for label, vid in PRESET_VOICES]
VOICE_IDS = [vid for _, vid in PRESET_VOICES]
VOICE_SET = set(VOICE_IDS)
DEFAULT_VOICE = tts._default_voice or (VOICE_IDS[0] if VOICE_IDS else None)
# Default voice shown in the read-text dropdown.
DEFAULT_UI_VOICE = "Phạm Tuyên" if "Phạm Tuyên" in VOICE_SET else DEFAULT_VOICE

DEFAULT_TEXT = (
    "Xin chào mọi người! [hắng giọng] Như bạn đang nghe thấy đấy, tốc độ xử lý của mình "
    "cực kỳ nhanh và mượt mà, giúp phản hồi gần như ngay lập tức theo thời gian thực. "
    "Chính vì vậy, mình rất phù hợp để ứng dụng trực tiếp vào các hệ thống Chatbot thông minh, "
    "trợ lý ảo, hoặc làm tổng đài viên tự động cho các doanh nghiệp. Tiện lợi quá đúng không ạ? "
    "[cười] Hi vọng phiên bản nâng cấp v3 này sẽ mang lại trải nghiệm tuyệt vời cho dự án của bạn."
)


def _stats_md(elapsed: float, n_samples: int) -> str:
    """Generation-speed / real-time-factor report."""
    dur = n_samples / SR if SR else 0.0
    rtf = (elapsed / dur) if dur > 0 else 0.0
    speed = (dur / elapsed) if elapsed > 0 else 0.0
    return (
        f"⏱️ Thời gian sinh: **{elapsed:.2f}s**\n"
        f"🔊 Độ dài audio: **{dur:.2f}s**\n"
        f"⚡ RTF: **{rtf:.3f}** (×{speed:.1f} so với thời gian thực)"
    )


def _gpu_duration(text, *args, **kwargs):
    """Dynamic ZeroGPU budget: scale with text length, capped at 3 minutes."""
    n = len(text or "")
    return int(min(180, 30 + n // 8))


# ── Batched serving engine (static batching) ──────────────────────────────────
# v3 Turbo ships a static-batching runtime (``vieneu.v3_turbo_serve``) that advances
# many chunks through each forward step together — a big GPU-throughput win. We split
# the text into chunks and run them in groups of ``_BATCH_SIZE``, then join with the
# same boundary-aware pauses ``tts.infer`` uses. The engine needs CUDA (ZeroGPU makes
# it real inside ``@spaces.GPU``); it is built lazily and cached on first use. On CPU
# there is no batched engine, so we fall back to the sequential ``tts.infer`` path.
_BATCH_SIZE = 32


def _engine_on_cuda() -> bool:
    dev = getattr(getattr(tts, "engine", None), "device", None)
    return dev is not None and getattr(dev, "type", None) == "cuda"


def _get_batch_engine():
    eng = getattr(tts, "_v3_batch_engine", None)
    if eng is None:
        from vieneu.v3_turbo_serve import V3TurboBatchEngine
        eng = V3TurboBatchEngine(tts.engine)
        tts._v3_batch_engine = eng
    return eng


def _resolve_voice(voice, ref_audio):
    """Resolve to ``(speaker_emb, ref_codes)``. A cloned ``ref_audio`` wins over the preset."""
    return tts._resolve_ref(voice, ref_audio or None, True, True)


def _make_reqs(chunks, speaker_emb, ref_codes):
    """Build one batched request per (already-split) text chunk."""
    return [{"phonemes": phonemize_text_with_emotions(c), "speaker_emb": speaker_emb,
             "ref_codes": ref_codes, "style": "tu_nhien", "use_ref_codes": True}
            for c in chunks]


def _run_batched(reqs, *, temperature, top_k, top_p, repetition_penalty, max_new_frames):
    """Generate a waveform per request via the batched engine, in groups of _BATCH_SIZE."""
    eng = _get_batch_engine()
    wavs = []
    for i in range(0, len(reqs), _BATCH_SIZE):
        wavs.extend(eng.generate_batch(
            reqs[i:i + _BATCH_SIZE],
            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 wavs


def _check_ref_len(path):
    """Warn when the reference clip is longer than the ideal 3–5 s window."""
    if not path:
        return gr.update(visible=False)
    try:
        dur = sf.info(path).duration
    except Exception:
        return gr.update(visible=False)
    if dur > 5.5:
        return gr.update(
            visible=True,
            value=(
                f"⚠️ Audio mẫu đang dài **{dur:.1f} giây**. Hãy cắt còn "
                f"**3–5 giây** (một câu nói rõ, ít ồn) để clone giọng tốt nhất — "
                f"audio quá dài thường cho kết quả kém hơn."
            ),
        )
    return gr.update(visible=False)


# ── Single-speaker synthesis ──────────────────────────────────────────────────
@spaces.GPU(duration=_gpu_duration)
def 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.")
    _log_infer("full", text, voice=voice, cloned=bool(ref_audio))

    t0 = time.perf_counter()
    if _engine_on_cuda():
        # Batched path: split into chunks and run them through the static-batching
        # engine so multiple chunks share each forward step. A cloned ``ref_audio``
        # takes precedence over the preset ``voice`` (resolved inside ``_resolve_voice``).
        speaker_emb, ref_codes = _resolve_voice(voice, ref_audio)
        chunks, gaps = normalize_to_chunks_v3_with_gaps(text, max_chars=int(max_chars))
        if not chunks:
            raise gr.Error("Vui lòng nhập văn bản cần đọc.")
        reqs = _make_reqs(chunks, speaker_emb, ref_codes)
        wavs = _run_batched(
            reqs, temperature=temperature, top_k=top_k, top_p=top_p,
            repetition_penalty=repetition_penalty, max_new_frames=max_new_frames,
        )
        # Same boundary-aware pauses as ``tts.infer``, then watermark once over the whole clip.
        wav = join_audio_chunks(wavs, sr=SR, silence_ps=gaps_to_silence(gaps))
        if wav is not None and len(wav):
            wav = tts._apply_watermark(wav)
    else:
        # CPU / no batched engine → sequential engine via the public API.
        kwargs = dict(
            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), max_chars=int(max_chars),
        )
        if ref_audio:
            wav = tts.infer(text, ref_audio=ref_audio, **kwargs)
        else:
            wav = tts.infer(text, voice=voice, **kwargs)
    elapsed = time.perf_counter() - t0

    wav = np.asarray(wav, dtype=np.float32)
    return (SR, wav), _stats_md(elapsed, len(wav))


# ── UI ────────────────────────────────────────────────────────────────────────
HEADER = """
<div style="text-align:center;padding:22px;border-radius:14px;
     background:linear-gradient(135deg,#0f172a 0%,#1e293b 100%);color:#fff;margin-bottom:18px;">
  <div style="font-size:2.1rem;font-weight:800;">🦜 VieNeu-TTS <span
       style="background:-webkit-linear-gradient(45deg,#60A5FA,#22D3EE);
       -webkit-background-clip:text;-webkit-text-fill-color:transparent;">v3 Turbo</span></div>
  <div style="margin-top:8px;font-size:0.95rem;">
     <a href="https://huggingface.co/pnnbao-ump/VieNeu-TTS-v3-Turbo" target="_blank"
        style="color:#60A5FA;text-decoration:none;font-weight:600;">🤗 Model card</a>
     &nbsp;·&nbsp;
     <a href="https://github.com/pnnbao97/VieNeu-TTS" target="_blank"
        style="color:#60A5FA;text-decoration:none;font-weight:600;">💻 GitHub repo</a>
  </div>
  <div style="opacity:.85;margin-top:6px;">
     Text-to-Speech tiếng Việt · 48&nbsp;kHz · giọng dựng sẵn + nhân bản giọng
  </div>
</div>
"""

theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="cyan", neutral_hue="slate")

with gr.Blocks(theme=theme, title="VieNeu-TTS v3 Turbo") as demo:
    gr.HTML(HEADER)

    gr.Markdown(
        "Chèn tag cảm xúc trong văn bản (thử nghiệm): `[cười]`, `[thở dài]`, `[hắng giọng]`."
    )
    with gr.Row():
        with gr.Column(scale=3):
            text_in = gr.Textbox(
                label="Văn bản", value=DEFAULT_TEXT, lines=8,
                placeholder="Nhập văn bản tiếng Việt...",
            )
            with gr.Tabs():
                with gr.Tab("Giọng dựng sẵn"):
                    voice_in = gr.Dropdown(
                        label="Chọn giọng", choices=VOICE_CHOICES, value=DEFAULT_UI_VOICE,
                    )
                with gr.Tab("Nhân bản giọng"):
                    gr.Markdown(
                        "### ⏱️ Audio mẫu nên dài **3–5 giây**\n"
                        "Dùng **một câu nói rõ ràng, ít tiếng ồn**. "
                        "**Đừng** tải lên file dài (cả đoạn/cả bài) — audio càng dài "
                        "clone càng dễ sai giọng và méo tiếng."
                    )
                    ref_audio_in = gr.Audio(
                        label="Audio mẫu (3–5 giây)", type="filepath",
                        sources=["upload", "microphone"],
                    )
                    ref_warn = gr.Markdown(visible=False)
                    gr.Markdown(
                        "_Có audio mẫu ở đây sẽ **ghi đè** giọng dựng sẵn. "
                        "Xoá audio để quay lại giọng dựng sẵn._"
                    )
                    ref_audio_in.change(_check_ref_len, ref_audio_in, ref_warn)
            with gr.Accordion("Tuỳ chọn nâng cao", open=False):
                with gr.Row():
                    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=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)")
                    max_chars_in = gr.Slider(64, 400, value=256, step=8, label="max_chars (cắt đoạn)")
        with gr.Column(scale=2):
            audio_out = gr.Audio(label="Kết quả", type="numpy")
            stats_out = gr.Markdown(label="Tốc độ sinh / RTF")
            run_btn = gr.Button("Tạo giọng nói", variant="primary")

    run_btn.click(
        fn=synthesize,
        inputs=[text_in, voice_in, ref_audio_in, temperature_in, top_k_in,
                top_p_in, rep_pen_in, max_frames_in, max_chars_in],
        outputs=[audio_out, stats_out],
    )
    gr.Examples(
        examples=[
            [DEFAULT_TEXT, DEFAULT_UI_VOICE],
            ["Hello everyone! [hắng giọng] As you can hear, my processing speed is incredibly fast and smooth, allowing for near-instant responses in real time. Because of this, I’m a perfect fit for direct integration into smart chatbots, virtual assistants, or automated call centers for businesses. Pretty convenient, right? [cười] I hope this version three upgrade brings an amazing experience to your project.", DEFAULT_UI_VOICE],
        ],
        inputs=[text_in, voice_in],
    )

    gr.Markdown(
        "Model: [pnnbao-ump/VieNeu-TTS-v3-Turbo]"
        "(https://huggingface.co/pnnbao-ump/VieNeu-TTS-v3-Turbo) · "
        "Code: [github.com/pnnbao97/VieNeu-TTS](https://github.com/pnnbao97/VieNeu-TTS)"
    )


if __name__ == "__main__":
    demo.queue().launch()