Update app.py

app.py

@@ -1,12 +1,11 @@
 # -*- coding: utf-8 -*-
-
+import re
 import numpy as np
 import torch
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from transformers import PreTrainedTokenizerFast
-from qwen_tts import Qwen3TTSTokenizer
 from transformers import AutoModelForSeq2SeqLM, PreTrainedTokenizerFast
+from qwen_tts import Qwen3TTSTokenizer
+
 REPO = "telecomadm1145/Kiseki-TTS-1.2"
 CODEC_REPO = "Qwen/Qwen3-TTS-Tokenizer-12Hz"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -14,56 +13,218 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Loading models onto device={DEVICE} ...")
 tok = PreTrainedTokenizerFast.from_pretrained(REPO, trust_remote_code=True)
 codec_tokenizer = Qwen3TTSTokenizer.from_pretrained(CODEC_REPO, device_map=DEVICE)
-
 model = AutoModelForSeq2SeqLM.from_pretrained(REPO, trust_remote_code=True).eval().to(DEVICE)
+NUM_Q = model.config.num_q
+FPS = 12  # Qwen3-TTS-Tokenizer-12Hz
+
 
+# ---------------------------------------------------------------- 分层参数解析
+def parse_per_layer(spec, Q, default, cast=float):
+    """
+    支持:
+      ""            -> default
+      "0.7"         -> 全层 0.7
+      "0.9,0.7,0.6" -> 逐层(不足按最后一个补齐)
+      "default:0.7, 0:0.9, 1-3:0.8"
+    """
+    if spec is None or not str(spec).strip():
+        return default
+    s = str(spec).strip()
+    if ":" in s:
+        out = [cast(default)] * Q
+        for item in re.split(r"[,\s;]+", s):
+            if not item or ":" not in item:
+                continue
+            k, v = item.split(":", 1)
+            k, v = k.strip(), cast(float(v))
+            if k in ("default", "*", "all"):
+                out = [v] * Q
+            elif "-" in k:
+                a, b = k.split("-")
+                for i in range(int(a), min(int(b), Q - 1) + 1):
+                    out[i] = v
+            else:
+                i = int(k)
+                if 0 <= i < Q:
+                    out[i] = v
+        return out
+    vals = [cast(float(p)) for p in re.split(r"[,\s;]+", s) if p]
+    return vals if len(vals) > 1 else vals[0]
 
-def tts_generate(text, max_new_frames, temperature, top_k, repetition_penalty):
+
+# ---------------------------------------------------------------- 参考音频编码
+def encode_reference_audio(path):
+    """返回 (T, Q) 的 LongTensor。"""
+    if not path:
+        return None
+    res = None
+    for call in (lambda: codec_tokenizer.encode(path),
+                 lambda: codec_tokenizer.encode([path]),
+                 lambda: codec_tokenizer(path)):
+        try:
+            res = call()
+            break
+        except Exception:
+            continue
+    if res is None:
+        raise gr.Error("参考音频编码失败：codec tokenizer 不支持该输入形式。")
+
+    if isinstance(res, dict):
+        codes = res.get("audio_codes", res.get("codes", next(iter(res.values()))))
+    elif isinstance(res, (list, tuple)):
+        codes = res[0]
+    else:
+        codes = res
+    if isinstance(codes, (list, tuple)):
+        codes = codes[0]
+    codes = torch.as_tensor(codes)
+    while codes.dim() > 2:
+        codes = codes[0]
+    if codes.dim() != 2:
+        raise gr.Error(f"参考音频 codes 维度异常: {tuple(codes.shape)}")
+    # 统一成 (T, Q)
+    if codes.shape[-1] != NUM_Q and codes.shape[0] == NUM_Q:
+        codes = codes.transpose(0, 1)
+    if codes.shape[-1] != NUM_Q:
+        raise gr.Error(f"codebook 数不匹配: 模型 num_q={NUM_Q}, 音频 {tuple(codes.shape)}")
+    return codes.long()
+
+
+# ---------------------------------------------------------------- 主流程
+def tts_generate(text, ref_audio, ref_text, prompt_seconds,
+                 max_new_frames, min_new_frames,
+                 temperature, top_k, top_p, repetition_penalty, rep_window,
+                 temp_spec, topk_spec, topp_spec, rep_spec,
+                 cfg_scale, cfg_scale_q0, cfg_spec, cfg_rescale, neg_text,
+                 seed):
     if not text or not text.strip():
         raise gr.Error("Please enter some text to synthesize.")
 
+    info = []
     with torch.no_grad():
-        ids = model.build_tts_input_ids(tok.encode(text), device=DEVICE)
+        # ---- 引导音频 ----
+        prompt_codes = None
+        if ref_audio:
+            prompt_codes = encode_reference_audio(ref_audio)
+            cap = int(float(prompt_seconds) * FPS)
+            if cap > 0 and prompt_codes.shape[0] > cap:
+                prompt_codes = prompt_codes[:cap]
+            prompt_codes = prompt_codes.to(DEVICE)
+            info.append(f"prompt frames = {prompt_codes.shape[0]} "
+                        f"({prompt_codes.shape[0] / FPS:.2f}s)")
+
+        # ---- encoder 输入 ----
+        prompt_text_ids = tok.encode(ref_text) if (ref_audio and ref_text and ref_text.strip()) else None
+        ids = model.build_tts_input_ids(tok.encode(text),
+                                        prompt_text_ids=prompt_text_ids,
+                                        device=DEVICE)
+
+        neg_ids = None
+        if neg_text and neg_text.strip():
+            neg_ids = model.build_tts_input_ids(tok.encode(neg_text), device=DEVICE)
+
+        # ---- 分层采样参数 ----
+        temps = parse_per_layer(temp_spec, NUM_Q, float(temperature))
+        topks = parse_per_layer(topk_spec, NUM_Q, int(top_k), cast=int)
+        topps = parse_per_layer(topp_spec, NUM_Q, float(top_p))
+        reps  = parse_per_layer(rep_spec,  NUM_Q, float(repetition_penalty))
+        cfgs  = parse_per_layer(cfg_spec,  NUM_Q, float(cfg_scale))
+
         out = model.generate_audio(
             ids,
+            prompt_codes=prompt_codes,
             max_new_frames=int(max_new_frames),
-            temperature=float(temperature),
-            top_k=int(top_k),
-            repetition_penalty=float(repetition_penalty),
+            min_new_frames=int(min_new_frames),
+            temperature=temps, top_k=topks, top_p=topps,
+            repetition_penalty=reps, repetition_window=int(rep_window),
+            cfg_scale=cfgs,
+            cfg_scale_q0=(None if cfg_scale_q0 <= 0 else float(cfg_scale_q0)),
+            cfg_rescale=float(cfg_rescale),
+            negative_input_ids=neg_ids,
+            seed=(None if int(seed) < 0 else int(seed)),
         )
+
         codes = out["audio_codes"][0][out["valid_mask"][0]]
+        if codes.shape[0] == 0:
+            raise gr.Error("模型立刻输出了 EOS，试着提高 min_new_frames 或降低 CFG。")
         wavs, sr = codec_tokenizer.decode({"audio_codes": codes})
 
     wav = wavs[0]
     if isinstance(wav, torch.Tensor):
         wav = wav.detach().cpu().numpy()
-    wav = np.asarray(wav, dtype=np.float32)
+    wav = np.asarray(wav, dtype=np.float32).squeeze()
 
-    return (sr, wav), "Generated."
+    info.append(f"generated frames = {codes.shape[0]} ({codes.shape[0] / FPS:.2f}s)")
+    info.append(f"temp={temps} | top_k={topks} | cfg={cfgs}")
+    return (sr, wav), " · ".join(info)
 
 
+# ---------------------------------------------------------------- UI
 with gr.Blocks(title="Kiseki-TTS") as demo:
     gr.Markdown(
         "# Kiseki-TTS\n"
-        f"Model: [`{REPO}`](https://huggingface.co/{REPO}) · Codec: `{CODEC_REPO}`"
+        f"Model: [`{REPO}`](https://huggingface.co/{REPO}) · Codec: `{CODEC_REPO}` · num_q={NUM_Q}"
     )
-
     with gr.Row():
         with gr.Column():
             text_in = gr.Textbox(label="Text", placeholder="うん、今日はいい天気です", lines=3)
-            with gr.Accordion("Generation settings", open=False):
-                max_frames = gr.Slider(50, 1000, value=250, step=10, label="max_new_frames")
-                temperature = gr.Slider(0.1, 1.5, value=0.5, step=0.05, label="temperature")
-                top_k = gr.Slider(1, 200, value=50, step=1, label="top_k")
-                repetition_penalty = gr.Slider(1.0, 2.0, value=1.0, step=0.05, label="repetition_penalty")
+
+            with gr.Accordion("🎤 引导音频 (voice prompt / 音色克隆)", open=True):
+                ref_audio = gr.Audio(label="Reference audio", type="filepath")
+                ref_text = gr.Textbox(
+                    label="Reference transcript（强烈建议填写，与参考音频内容一致）",
+                    lines=2)
+                prompt_seconds = gr.Slider(0, 30, value=10, step=0.5,
+                                           label="prompt 最长秒数 (0=不限)")
+
+            with gr.Accordion("🎚 CFG (classifier-free guidance)", open=False):
+                cfg_scale = gr.Slider(1.0, 5.0, value=1.0, step=0.05,
+                                      label="cfg_scale（1.0=关闭；建议 1.3~2.0）")
+                cfg_scale_q0 = gr.Slider(0.0, 5.0, value=0.0, step=0.05,
+                                         label="cfg_scale_q0（0=沿用上面）")
+                cfg_spec = gr.Textbox(label="逐层 cfg_scale（覆盖上面，如 `0:2.0, 1-15:1.2`）")
+                cfg_rescale = gr.Slider(0.0, 1.0, value=0.0, step=0.05,
+                                        label="cfg_rescale（std 回缩，抑制过锐）")
+                neg_text = gr.Textbox(label="Negative text（留空=丢弃全部文本做无条件分支）")
+
+            with gr.Accordion("⚙️ 采样参数（全局）", open=False):
+                max_frames = gr.Slider(50, 2000, value=250, step=10, label="max_new_frames")
+                min_frames = gr.Slider(0, 500, value=0, step=5, label="min_new_frames（抑制早停）")
+                temperature = gr.Slider(0.0, 1.5, value=0.5, step=0.05, label="temperature")
+                top_k = gr.Slider(0, 2048, value=50, step=1, label="top_k (0=off)")
+                top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="top_p")
+                repetition_penalty = gr.Slider(1.0, 2.0, value=1.0, step=0.05,
+                                               label="repetition_penalty")
+                rep_window = gr.Slider(0, 200, value=0, step=5,
+                                       label="repetition_window（0=全历史）")
+                seed = gr.Number(value=-1, precision=0, label="seed (-1 = random)")
+
+            with gr.Accordion("🧬 分层采样（按 RVQ 层覆盖全局值）", open=False):
+                gr.Markdown(
+                    "格式：`0.7` / `0.9,0.7,0.6,...`（不足按最后一个补齐）/ "
+                    "`default:0.7, 0:0.9, 1-3:0.8`。留空表示用全局值。\n"
+                    "经验：q0 决定内容与韵律（温度略高、top_k 大），q1+ 只是残差细节（低温甚至 0=贪心）。"
+                )
+                temp_spec = gr.Textbox(label="temperature / layer",
+                                       placeholder="default:0.4, 0:0.8")
+                topk_spec = gr.Textbox(label="top_k / layer", placeholder="default:20, 0:60")
+                topp_spec = gr.Textbox(label="top_p / layer", placeholder="default:0.9, 0:0.95")
+                rep_spec = gr.Textbox(label="repetition_penalty / layer", placeholder="0:1.15")
+
             tts_btn = gr.Button("Generate", variant="primary")
+
         with gr.Column():
             audio_out = gr.Audio(label="Output audio", type="numpy")
             tts_status = gr.Markdown()
 
     tts_btn.click(
         tts_generate,
-        inputs=[text_in, max_frames, temperature, top_k, repetition_penalty],
+        inputs=[text_in, ref_audio, ref_text, prompt_seconds,
+                max_frames, min_frames,
+                temperature, top_k, top_p, repetition_penalty, rep_window,
+                temp_spec, topk_spec, topp_spec, rep_spec,
+                cfg_scale, cfg_scale_q0, cfg_spec, cfg_rescale, neg_text,
+                seed],
         outputs=[audio_out, tts_status],
     )