app.py

"""BlueMagpie-TTS 線上試用 Space (Gradio + ZeroGPU)。

提供三種試用情境，並一律套用模型官方建議的最佳生成參數：
  1. 指定語者：李宏毅 —— 用模型內附、已取得授權的語者向量控制音色
  2. 聲音複製       —— 上傳一段參考音檔，模仿該語者的音色
  3. 長文逐句串流   —— 把長文切句，合成一句播一句，做出串流效果（每句含自動重試）

最佳參數讀自模型發佈中繼資料 release_metadata.json 的 recommended_generation_defaults
（以 500 句難句 + ASR CER 調出，最佳 trial：cfg=2.8 / steps=9），確保 demo 永遠跟著模型走。
"""

import json
import os
import time

import numpy as np
import torch
import gradio as gr
from huggingface_hub import snapshot_download
from transformers import PreTrainedTokenizerFast

from bluemagpie import BlueMagpieModel

# --------------------------------------------------------------------------- #
# ZeroGPU：用 @spaces.GPU 動態取得 GPU；本機無 spaces 時退回 cuda/cpu 自動偵測。
# --------------------------------------------------------------------------- #
try:
    import spaces

    gpu = spaces.GPU(duration=120)
    DEVICE = "cuda"
except ImportError:
    def gpu(fn):
        return fn

    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

REPO_ID = "OpenFormosa/BlueMagpie-TTS"

# --------------------------------------------------------------------------- #
# 載入模型（啟動時做一次）
# --------------------------------------------------------------------------- #
print(f"[BlueMagpie] downloading model from {REPO_ID} ...")
MODEL_DIR = snapshot_download(REPO_ID)

tokenizer = PreTrainedTokenizerFast(
    tokenizer_file=os.path.join(MODEL_DIR, "tokenizer.json")
)
print(f"[BlueMagpie] loading model on device={DEVICE} ...")
model = BlueMagpieModel.from_local(
    MODEL_DIR, tokenizer=tokenizer, training=False, device=DEVICE
)
SR = model.sample_rate

# --------------------------------------------------------------------------- #
# 最佳生成參數：讀自模型發佈中繼資料；缺檔時退回官方建議值。
# --------------------------------------------------------------------------- #
_GEN_KEYS = {
    "cfg_value",
    "inference_timesteps",
    "max_len",
    "min_len",
    "retry_badcase",
    "retry_badcase_max_times",
    "retry_badcase_ratio_threshold",
}
_BEST_FALLBACK = {
    "cfg_value": 2.8,
    "inference_timesteps": 9,
    "max_len": 2000,
    "retry_badcase": True,
    "retry_badcase_max_times": 3,
    "retry_badcase_ratio_threshold": 6.0,
}
SPEAKER_ID = "hung_yi_lee"

try:
    with open(os.path.join(MODEL_DIR, "release_metadata.json"), encoding="utf-8") as f:
        _meta = json.load(f)
    _rec = _meta.get("recommended_generation_defaults", {}) or {}
    GEN_KWARGS = {k: v for k, v in _rec.items() if k in _GEN_KEYS} or dict(_BEST_FALLBACK)
    SPEAKER_ID = _rec.get("speaker_id", SPEAKER_ID)
except Exception as e:  # 任何讀取/解析問題都退回安全的最佳值
    print(f"[BlueMagpie] release_metadata.json unavailable ({e}); using fallback best params")
    GEN_KWARGS = dict(_BEST_FALLBACK)

print(f"[BlueMagpie] best generation params: {GEN_KWARGS}")

# 語者向量（已取得本人授權）
_centroids = torch.load(
    os.path.join(MODEL_DIR, "checkpoints", "hung_yi_lee_speaker_centroids.pt"),
    map_location="cpu",
    weights_only=True,
)
SPK_CENTROID = _centroids["centroids"][_centroids["speaker_ids"].index(SPEAKER_ID)]
print(f"[BlueMagpie] ready. sample_rate={SR}, speaker={SPEAKER_ID}")

# 註：repo 的 bluemagpie.serving 加速引擎（torch.compile / 連續批次）是為「常駐
# 專用 GPU、追求吞吐」設計。在 ZeroGPU（GPU 逐請求序列化掛載）上實測後，單一
# 請求延遲沒有改善（warm ~7.0s vs eager ~6.2s），反而帶來首呼叫 ~50s 編譯成本與
# CUDA graph 反覆重編譯的尖峰，故此 demo 維持 eager。詳見對話中的量測數據。


def _to_numpy(audio: torch.Tensor):
    return audio.squeeze().float().cpu().numpy()


def _pcm16(x: np.ndarray) -> np.ndarray:
    """float 波形 -> 16-bit PCM，供 Gradio 串流輸出（格式最穩定）。"""
    return (np.clip(x, -1.0, 1.0) * 32767.0).astype(np.int16)


# 長文逐句串流：以句末標點／換行切句。逐句串流直到接近 ZeroGPU 單次 GPU 時間配額
# 為止（不是固定句數），盡量把整段長文串完；另設一個防濫用的硬上限。
_SENT_ENDERS = "。！？；!?…\n"
_GPU_BUDGET_S = 100         # @spaces.GPU(duration=120)，留約 20s 緩衝給最後一句
_MAX_STREAM_SENTENCES = 80  # 防濫用硬上限（通常時間配額會先到）


def _split_sentences(text):
    out, cur = [], []
    for ch in text:
        cur.append(ch)
        if ch in _SENT_ENDERS:
            s = "".join(cur).strip()
            if s:
                out.append(s)
            cur = []
    s = "".join(cur).strip()
    if s:
        out.append(s)
    # 把落單的標點（如連續「！？」被拆出的單一字元）併回前一句
    merged = []
    for s in out:
        if merged and len(s) <= 1:
            merged[-1] += s
        else:
            merged.append(s)
    return merged


# --------------------------------------------------------------------------- #
# 推論：三種模式都套用 GEN_KWARGS（官方最佳參數）
# --------------------------------------------------------------------------- #
@gpu
def tts_speaker(text):
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    audio = model.generate(
        target_text=text,
        speaker_centroid=SPK_CENTROID,
        **GEN_KWARGS,
    )
    return (SR, _to_numpy(audio))


@gpu
def tts_clone(text, reference_audio):
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    if not reference_audio:
        raise gr.Error("請先上傳或錄製一段參考音檔。")
    audio = model.generate(
        target_text=text,
        reference_wav_path=reference_audio,  # gr.Audio(type="filepath") 回傳路徑
        **GEN_KWARGS,
    )
    return (SR, _to_numpy(audio))


@gpu
def tts_stream(text, reference_audio):
    """長文逐句串流：把長文切成句子，逐句合成、合成一句就播一句，做出串流效果。

    逐句串流直到接近單次 GPU 時間配額（約 100s）為止，盡量把整段長文串完，而非固定
    句數；真的超過才截斷並提示分批。每句都用一般 generate（含 retry_badcase）產生
    完整音檔再 yield；句與句之間可能有短暫間隔（ZeroGPU 約 0.44x 實時所致）。預設用
    李宏毅語者向量，上傳參考音檔則改為聲音複製。
    """
    text = (text or "").strip()
    if not text:
        raise gr.Error("請先輸入要合成的文字。")
    sentences = _split_sentences(text)
    if not sentences:
        raise gr.Error("找不到可合成的句子。")
    total = len(sentences)
    sentences = sentences[:_MAX_STREAM_SENTENCES]
    start = time.time()
    done = 0
    for sent in sentences:
        # 至少先產出第一句；之後一旦逼近 GPU 時間配額就停，避免配額被回收而報錯
        if done and time.time() - start > _GPU_BUDGET_S:
            break
        if reference_audio:
            audio = model.generate(
                target_text=sent, reference_wav_path=reference_audio, **GEN_KWARGS
            )
        else:
            audio = model.generate(
                target_text=sent, speaker_centroid=SPK_CENTROID, **GEN_KWARGS
            )
        yield (SR, _pcm16(_to_numpy(audio)))
        done += 1
    if done < total:
        gr.Warning(
            f"受單次 GPU 時間配額限制，本次串流前 {done} 句（共 {total} 句）；"
            "想念完整段可分批貼上。"
        )


# --------------------------------------------------------------------------- #
# 介面
# --------------------------------------------------------------------------- #
EXAMPLE_TEXTS = [
    "今天天氣真好，我們一起去散步吧。",
    "這是 AI TTS code switching 測試，混合中英文也沒問題。",
    "台灣藍鵲是一種叫聲響亮、辨識度很高的鳥。",
]

STREAM_EXAMPLES = [
    "台灣藍鵲是台灣特有的鳥類，羽色以亮藍為主，尾羽修長。牠們常成群在樹林間活動，"
    "叫聲響亮而容易辨認。今天天氣真好，我們一起出去走走，順便看看這些美麗的鳥吧！",
    "歡迎使用這個示範。這是一段比較長的文字，會被切成好幾句。"
    "系統會合成一句、播放一句，做出串流的效果。希望你會喜歡這個聲音！",
]

_PARAMS_NOTE = (
    f"已套用模型官方建議最佳參數："
    f"`cfg_value={GEN_KWARGS.get('cfg_value')}`、"
    f"`inference_timesteps={GEN_KWARGS.get('inference_timesteps')}`、"
    f"`retry_badcase={'on' if GEN_KWARGS.get('retry_badcase') else 'off'}`。"
)

HEADER = f"""
# 🐦‍⬛ BlueMagpie-TTS · 文字轉語音線上試用

**OpenFormosa Blue Magpie TTS** —— 支援**台灣華語**與**中英混合**的文字轉語音模型，輸出 48 kHz 單聲道語音。
模型：[OpenFormosa/BlueMagpie-TTS](https://huggingface.co/OpenFormosa/BlueMagpie-TTS) ·
程式碼：[GitHub](https://github.com/OpenFormosa/BlueMagpie-TTS)

> {_PARAMS_NOTE}
"""

with gr.Blocks(title="BlueMagpie-TTS Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(HEADER)

    with gr.Tab("指定語者：李宏毅"):
        gr.Markdown(
            "使用模型內附、**已取得李宏毅老師本人授權**的語者向量來控制音色"
            "（官方最佳參數正是針對此語者調校）。"
        )
        with gr.Row():
            with gr.Column():
                s_text = gr.Textbox(
                    label="要合成的文字", lines=3,
                    placeholder="輸入中文或中英混合的文字…",
                )
                s_btn = gr.Button("以李宏毅音色合成", variant="primary")
            with gr.Column():
                s_out = gr.Audio(label="合成結果", type="numpy")
        gr.Examples(EXAMPLE_TEXTS, inputs=s_text, label="範例文字")
        s_btn.click(tts_speaker, s_text, s_out)

    with gr.Tab("聲音複製"):
        gr.Markdown(
            "上傳一段參考音檔，輸出會模仿該語者的音色。\n\n"
            "⚠️ **請只使用你已取得授權的聲音**，請勿在未經本人同意下複製他人聲音。"
        )
        with gr.Row():
            with gr.Column():
                c_text = gr.Textbox(
                    label="要合成的文字", lines=3,
                    placeholder="輸入中文或中英混合的文字…",
                )
                c_ref = gr.Audio(
                    label="參考音檔（上傳或錄製）", type="filepath",
                    sources=["upload", "microphone"],
                )
                c_btn = gr.Button("複製音色合成", variant="primary")
            with gr.Column():
                c_out = gr.Audio(label="合成結果", type="numpy")
        c_btn.click(tts_clone, [c_text, c_ref], c_out)

    with gr.Tab("長文逐句串流"):
        gr.Markdown(
            "**長文逐句串流**：把長文切成句子，合成一句就播一句，做出串流（邊聽邊等）效果。"
            "第一句很快就能聽到，後面邊播邊合成。\n\n"
            "預設用李宏毅語者向量；若上傳參考音檔則改為聲音複製（請只使用已取得授權的聲音）。\n"
            "> ZeroGPU 約 0.44x 實時，句與句之間可能有短暫間隔；每句都套用最佳參數（含自動重試）。"
            "單次串流以 ZeroGPU 時間配額為限（約可串數十句），超長會自動截斷並提示分批。"
        )
        with gr.Row():
            with gr.Column():
                t_text = gr.Textbox(
                    label="要合成的長文", lines=5,
                    placeholder="貼上一段多句的文字，會逐句合成、逐句播放…",
                )
                t_ref = gr.Audio(
                    label="參考音檔（選填；留空則用李宏毅音色）", type="filepath",
                    sources=["upload", "microphone"],
                )
                t_btn = gr.Button("逐句串流合成", variant="primary")
            with gr.Column():
                t_out = gr.Audio(label="逐句串流播放", streaming=True, autoplay=True)
        gr.Examples(STREAM_EXAMPLES, inputs=t_text, label="長文範例")
        t_btn.click(tts_stream, [t_text, t_ref], t_out)

    gr.Markdown(
        "---\n"
        "合成的語音僅供研究與評估展示用途，輸出可能不完美；正式使用前請人工檢視。"
        " · Apache-2.0"
    )


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