First commit

.gitattributes

@@ -33,3 +33,14 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/moss_prompts/en_4_4p5s.wav filter=lfs diff=lfs merge=lfs -text
+assets/moss_prompts/zh_1_4p5s.wav filter=lfs diff=lfs merge=lfs -text
+models/axmodels/decoder_part0.axmodel filter=lfs diff=lfs merge=lfs -text
+models/axmodels/decoder_part1.axmodel filter=lfs diff=lfs merge=lfs -text
+models/axmodels/decoder_part2.axmodel filter=lfs diff=lfs merge=lfs -text
+models/axmodels/decoder_part3.axmodel filter=lfs diff=lfs merge=lfs -text
+models/axmodels/encoder.axmodel filter=lfs diff=lfs merge=lfs -text
+outputs/en_long_paragraph.wav filter=lfs diff=lfs merge=lfs -text
+outputs/en_sentence.wav filter=lfs diff=lfs merge=lfs -text
+outputs/zh_long_paragraph.wav filter=lfs diff=lfs merge=lfs -text
+outputs/zh_sentence.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,168 @@
----
-license: mit
----
+# ZipVoice.AXERA
+
+ZipVoice 板端推理 demo。
+
+## 功能
+
+- 支持中文和英文语音生成。
+- 支持语音克隆。
+
+## 模型转换
+
+For those interested in model conversion, refer to:
+
+- [Pulsar2 Docs — How to Convert ONNX to axmodel](https://pulsar2-docs.readthedocs.io/en/latest/pulsar2/introduction.html)
+- [export / quantization Project](https://github.com/AXERA-TECH/ZipVoice.AXERA)
+
+## 支持平台
+
+- AX650
+  - [M4N-Dock（爱芯派Pro）](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
+  - [M.2 Accelerator Card](https://docs.m5stack.com/zh_CN/ai_hardware/LLM-8850_Card)
+
+## 目录结构
+
+```text
+ZipVoice.AXERA/
+├── assets/
+│   ├── moss_prompts/
+│   └── paragraphs/
+├── configs/
+├── models/
+│   └── axmodels/
+├── resources/
+│   ├── vocos-mel-24khz/
+│   └── zipvoice_hf/
+├── scripts/
+├── infer_zipvoice_axera.py
+├── requirements.txt
+└── README.md
+```
+
+## 环境
+
+### 安装 pyaxengine
+
+[pyaxengine Releases](https://github.com/AXERA-TECH/pyaxengine/releases/latest) 下载对应版本的 `.whl` 文件，然后安装：
+
+```bash
+pip3 install axengine-x.x.x-py3-none-any.whl
+```
+
+### Python 依赖
+
+```bash
+pip3 install -r requirements.txt
+```
+
+## 推理命令
+
+进入目录：
+
+```bash
+cd /root/huyuan/workspace/TTS/ZipVoice.AXERA
+```
+
+### 句子
+
+中文：
+```bash
+python3 infer_zipvoice_axera.py \
+  --text "今天午后天气很好，我打开窗户，听见远处有人聊天，水杯也轻轻晃了一下。" \
+  --prompt-text "不管怎么样我和汤姆还是要感谢贝尔卡金的援手" \
+  --prompt-wav assets/moss_prompts/zh_1_4p5s.wav \
+  --output-wav outputs/zh_sentence.wav \
+  --speed 1.0 \
+  --num-step 10 \
+  --seed 42
+```
+
+推理结果：
+```text
+推理耗时: 5.781s
+生成语音耗时: 6.411s
+RTF: 0.9018
+```
+
+音频：[outputs/zh_sentence.wav](outputs/zh_sentence.wav)
+
+英文：
+
+```bash
+python3 infer_zipvoice_axera.py \
+  --text "This morning, a small train left the station, carrying sleepy passengers toward a bright coastal town." \
+  --prompt-text "This is almost twice the current industry production level per train." \
+  --prompt-wav assets/moss_prompts/en_4_4p5s.wav \
+  --output-wav outputs/en_sentence.wav \
+  --speed 1.0 \
+  --num-step 10 \
+  --seed 42
+```
+
+推理结果：
+```text
+推理耗时: 5.711s
+生成语音耗时: 6.411s
+RTF: 0.8909
+```
+
+音频：[outputs/en_sentence.wav](outputs/en_sentence.wav)
+
+
+### 段落
+
+中文：
+```bash
+python3 infer_zipvoice_axera.py \
+  --text-file assets/paragraphs/zh_ginkgo.txt \
+  --prompt-text "不管怎么样我和汤姆还是要感谢贝尔卡金的援手" \
+  --prompt-wav assets/moss_prompts/zh_1_4p5s.wav \
+  --output-wav outputs/zh_long_paragraph.wav \
+  --speed 1.0 \
+  --num-step 10 \
+  --seed 42
+```
+
+推理结果：
+```text
+推理耗时: 40.292s
+生成语音耗时: 44.744s
+RTF: 0.9005
+```
+
+音频：[outputs/zh_long_paragraph.wav](outputs/zh_long_paragraph.wav)
+
+英文：
+
+```bash
+python3 infer_zipvoice_axera.py \
+  --text-file assets/paragraphs/en_scavenger.txt \
+  --prompt-text "This is almost twice the current industry production level per train." \
+  --prompt-wav assets/moss_prompts/en_4_4p5s.wav \
+  --output-wav outputs/en_long_paragraph.wav \
+  --speed 1.0 \
+  --num-step 10 \
+  --seed 42
+```
+
+推理结果：
+```
+推理耗时: 62.161s
+生成语音耗时: 64.749s
+RTF: 0.9600
+```
+
+音频：[outputs/en_long_paragraph.wav](outputs/en_long_paragraph.wav)
+
+## 参数说明
+
+- `--prompt-wav`：参考音频，用于控制音色。
+- `--prompt-text`：参考音频对应文本，必须尽量和 `prompt-wav` 内容一致，否则 tokenizer 长度和音频帧数不匹配，会影响语速和效果。
+- `--speed`：控制生成帧数，值越大音频越短，语速越快；值越小音频越长，语速越慢。
+- `--num-step`：采样步数。步数越大通常效果更稳，推理耗时线性增加(当 num-step=8 时，以上推理的 rtf=0.75 左右)。
+- `--max-feat-len`：decoder 固定 feature 长度，当前为 1024。
+
+
+## 参考
+
+- [ZipVoice](https://github.com/k2-fsa/ZipVoice)

assets/moss_prompts/en_4_4p5s.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa360ae39ff73bfee70f10ba997ff5ae5c15c1dc25bb382510f6e9654f0f11b8
+size 216078

assets/moss_prompts/zh_1_4p5s.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f377c00e0cce4c1ca699bd0e53b553151faeddcb1755a5e9604cdbbeb67e2eb
+size 216078

assets/paragraphs/en_scavenger.txt

@@ -0,0 +1 @@
+At dusk, I saw an old scavenger digging through a bin on a quiet street corner. His clothes were torn, his back bent from years of labor. I was about to walk past him, as most people did, when something made me stop. From the pile of waste, he pulled out a thin, battered book. Curious, I stepped closer. It was a collection of poems, its pages yellowed and its cover stained. Instead of tossing it aside, the old man wiped the dirt off with his sleeve. Then, in a soft, trembling voice, he began to read: "Beside the eastern fence, I pick chrysanthemums; in quiet leisure, I see the southern hills." The words floated in the evening air. At that moment, the setting sun touched his silver hair, and his eyes, bright and clear, seemed to belong to someone else entirely. He was no longer a scavenger standing by garbage; he was a traveler in a world of beauty and peace. After a long while, he closed the book, pressed it gently to his chest, and slipped it inside his worn jacket. Then he went back to searching for plastic bottles, his movements slow and steady. I stood there, speechless.

assets/paragraphs/zh_ginkgo.txt

@@ -0,0 +1 @@
+窗外的银杏叶又黄了。我坐在老家的院子里，看那一片片扇形的叶子在风中轻轻摇摆，像是无数小扇子，扇走了夏天最后的暑气，也扇来了秋天的凉意。阳光透过枝叶的缝隙洒下来，在地上投下斑驳的光影，那光影随着风不停地晃动，像一群调皮的孩子在捉迷藏。说起来，这棵银杏树是我爷爷年轻时种下的，算算已有十多年了。爷爷去年秋天走了，走的时候，银杏叶正黄得灿烂。我记得他临终前说的最后一句话是：“叶子落了，来年还会再长。”当时我不太明白这句话的意思，现在看着满树金黄，忽然好像懂了些什么。

configs/decoder4_split_manifest.json

@@ -0,0 +1,74 @@
+{
+    "version": 2,
+    "source_summary": "/data/huyuan/workspace/LLM/TTS/ZIPVOICE_RICH_CHECK0_REQUANT_SUMMARY.md",
+    "source_static_onnx_manifest": "/data/huyuan/workspace/LLM/TTS/ZipVoice.AXERA/models/static_mask_1024_cfgfix_nolog_encoder_attn_loop_sigmoid_bump_decoder4/decoder4_split_manifest.json",
+    "source_encoder_axmodel": "/data/huyuan/workspace/LLM/TTS/ZipVoice.AXERA/model_convert/axmodels_nolog_encoder_core_rich_check0/encoder_core_nolog_rich_check0.axmodel",
+    "source_decoder_axmodel_dir": "/data/huyuan/workspace/LLM/TTS/ZipVoice.AXERA/model_convert/axmodels_attn_loop_sigmoid_bump_decoder4_rich_check0",
+    "encoder": {
+        "name": "encoder_core",
+        "file": "encoder.axmodel",
+        "inputs": [
+            "cat_tokens"
+        ],
+        "outputs": [
+            "encoded"
+        ]
+    },
+    "decoder_parts": [
+        {
+            "name": "fm_decoder_part0",
+            "file": "decoder_part0.axmodel",
+            "inputs": [
+                "t",
+                "x",
+                "text_condition",
+                "speech_condition",
+                "guidance_scale",
+                "padding_mask"
+            ],
+            "outputs": [
+                "decoder_hidden_p0",
+                "time_hidden",
+                "padding_mask2",
+                "cfg_scale"
+            ]
+        },
+        {
+            "name": "fm_decoder_part1",
+            "file": "decoder_part1.axmodel",
+            "inputs": [
+                "decoder_hidden_p0",
+                "time_hidden",
+                "padding_mask2"
+            ],
+            "outputs": [
+                "decoder_hidden_p1"
+            ]
+        },
+        {
+            "name": "fm_decoder_part2",
+            "file": "decoder_part2.axmodel",
+            "inputs": [
+                "decoder_hidden_p1",
+                "time_hidden",
+                "padding_mask2"
+            ],
+            "outputs": [
+                "decoder_hidden_p2"
+            ]
+        },
+        {
+            "name": "fm_decoder_part3",
+            "file": "decoder_part3.axmodel",
+            "inputs": [
+                "decoder_hidden_p2",
+                "time_hidden",
+                "padding_mask2",
+                "cfg_scale"
+            ],
+            "outputs": [
+                "v"
+            ]
+        }
+    ]
+}

configs/runtime_config.json

@@ -0,0 +1,12 @@
+{
+    "model_type": "zipvoice_decoder4",
+    "max_tokens": 384,
+    "max_feat_len": 1024,
+    "feat_dim": 100,
+    "sampling_rate": 24000,
+    "hop_length": 256,
+    "num_step": 16,
+    "t_shift": 0.5,
+    "guidance_scale": 1.0,
+    "speed": 1.0
+}

infer_zipvoice_axera.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import argparse
+import logging
+import time
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+
+from scripts.common_infer import extract_prompt_features, load_tokenizer
+from scripts.common_infer import load_vocoder, vocoder_decode_loaded
+from scripts.text_processing import build_cat_tokens, build_segments
+from scripts.text_processing import load_text, normalize_punctuation, setup_jieba_cache
+from scripts.zipvoice_decoder4_runtime import Decoder4ZipVoiceBoardRuntime
+
+
+def parse_args() -> argparse.Namespace:
+    repo_dir = Path(__file__).resolve().parent
+    p = argparse.ArgumentParser(
+        description="ZipVoice encoder + decoder4 inference on AX650",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument("--config-dir", default=str(repo_dir / "configs"))
+    p.add_argument("--model-dir", default=str(repo_dir / "models" / "axmodels"))
+    p.add_argument("--text", default=None, help="Text to synthesize")
+    p.add_argument("--text-file", default=None, help="UTF-8 text file containing text")
+    p.add_argument("--prompt-text", required=True)
+    p.add_argument("--prompt-wav", required=True)
+    p.add_argument("--output-wav", default="outputs/zipvoice_axera.wav")
+    p.add_argument("--num-step", type=int, default=12)
+    p.add_argument("--guidance-scale", type=float, default=1.0)
+    p.add_argument("--speed", type=float, default=0.85)
+    p.add_argument("--t-shift", type=float, default=0.5)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--max-tokens", type=int, default=384)
+    p.add_argument("--max-feat-len", type=int, default=1024)
+    p.add_argument("--feat-scale", type=float, default=0.1)
+    p.add_argument("--target-rms", type=float, default=0.1)
+    p.add_argument("--min-generated-frames", type=int, default=360)
+    p.add_argument("--max-generated-frames", type=int, default=620)
+    p.add_argument(
+        "--max-raw-feat-ratio",
+        type=float,
+        default=1.2,
+        help="Allow raw features_len up to this ratio of max_feat_len before splitting",
+    )
+    p.add_argument("--silence-ms", type=int, default=140)
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+    repo_dir = Path(__file__).resolve().parent
+    setup_jieba_cache(repo_dir)
+    output_wav = Path(args.output_wav)
+    output_wav.parent.mkdir(parents=True, exist_ok=True)
+
+    t_target_text_start = time.perf_counter()
+    original_text = load_text(args)
+    text = normalize_punctuation(original_text)
+    target_text_preprocess_sec = time.perf_counter() - t_target_text_start
+    prompt_text = normalize_punctuation(args.prompt_text)
+    if text != original_text:
+        logging.debug("Normalized text punctuation: %s", text)
+    if prompt_text != args.prompt_text:
+        logging.debug("Normalized prompt punctuation: %s", prompt_text)
+
+    logging.debug("Loading tokenizer...")
+    tokenizer = load_tokenizer(repo_dir)
+    prompt_tokens = tokenizer.texts_to_token_ids([prompt_text])[0]
+    max_text_tokens = args.max_tokens - len(prompt_tokens) - 1
+    if max_text_tokens <= 0:
+        raise ValueError("prompt_tokens leaves no room for text")
+
+    logging.debug("Extracting prompt features from %s...", args.prompt_wav)
+    prompt_features, prompt_rms = extract_prompt_features(
+        args.prompt_wav,
+        repo_dir=repo_dir,
+        sampling_rate=24000,
+        feat_scale=args.feat_scale,
+        target_rms=args.target_rms,
+    )
+    prompt_frames = int(prompt_features.shape[1])
+    logging.debug(
+        "prompt_tokens=%d, prompt_frames=%d",
+        len(prompt_tokens),
+        prompt_frames,
+    )
+
+    t_segment_planning_start = time.perf_counter()
+    segments = build_segments(
+        tokenizer=tokenizer,
+        text=text,
+        prompt_frames=prompt_frames,
+        prompt_tokens_len=len(prompt_tokens),
+        speed=args.speed,
+        max_feat_len=args.max_feat_len,
+        max_text_tokens=max_text_tokens,
+        min_generated_frames=args.min_generated_frames,
+        max_generated_frames=args.max_generated_frames,
+        max_raw_feat_ratio=args.max_raw_feat_ratio,
+    )
+    segment_planning_sec = time.perf_counter() - t_segment_planning_start
+    logging.debug("Built %d segments", len(segments))
+    for index, segment in enumerate(segments, start=1):
+        logging.debug(
+            "segment %02d: text_tokens=%d generated_frames=%d features_len=%d clamped=%s text=%s",
+            index,
+            segment["text_tokens"],
+            segment["generated_frames"],
+            segment["features_len"],
+            segment["clamped"],
+            segment["text"],
+        )
+
+    logging.info("正在加载模型...")
+    runtime = Decoder4ZipVoiceBoardRuntime(
+        config_dir=args.config_dir,
+        models_dir=args.model_dir,
+        max_feat_len=args.max_feat_len,
+        max_tokens=args.max_tokens,
+        num_step=args.num_step,
+        t_shift=args.t_shift,
+    )
+    logging.debug("Loading vocoder...")
+    vocoder = load_vocoder(repo_dir)
+    logging.info("模型加载完成")
+
+    audios: list[np.ndarray] = []
+    silence = np.zeros(int(24000 * args.silence_ms / 1000), dtype=np.float32)
+    total_segment_sec = 0.0
+
+    for index, segment in enumerate(segments, start=1):
+        logging.info("开始推理第 %d/%d 句", index, len(segments))
+        t_segment_start = time.perf_counter()
+        seg_text = str(segment["text"])
+        text_tokens = tokenizer.texts_to_token_ids([seg_text])[0]
+        cat_tokens = build_cat_tokens(tokenizer, prompt_tokens, text_tokens, args.max_tokens)
+
+        pred_features, timing = runtime.sample(
+            cat_tokens=cat_tokens,
+            prompt_tokens_len=len(prompt_tokens),
+            text_tokens_len=len(text_tokens),
+            prompt_features=prompt_features,
+            prompt_features_len=prompt_frames,
+            speed=args.speed,
+            guidance_scale=args.guidance_scale,
+            seed=args.seed + index - 1,
+        )
+
+        audio = vocoder_decode_loaded(
+            vocoder,
+            pred_features,
+            feat_scale=args.feat_scale,
+            target_rms=args.target_rms,
+            prompt_rms=prompt_rms,
+        )
+        segment_sec = time.perf_counter() - t_segment_start
+        audio_sec = len(audio) / 24000
+
+        if audios:
+            audios.append(silence)
+        audios.append(audio.astype(np.float32, copy=False))
+
+        total_segment_sec += segment_sec
+        logging.debug(
+            "segment %02d done: audio=%.2fs runtime=%.2fs",
+            index,
+            audio_sec,
+            segment_sec,
+        )
+
+    final_audio = np.concatenate(audios) if audios else np.zeros(0, dtype=np.float32)
+    sf.write(str(output_wav), final_audio, samplerate=24000)
+    final_audio_sec = len(final_audio) / 24000
+    runtime_sec_for_rtf = target_text_preprocess_sec + segment_planning_sec + total_segment_sec
+    rtf = (
+        runtime_sec_for_rtf / final_audio_sec if final_audio_sec > 0 else float("inf")
+    )
+
+    logging.info("推理耗时: %.3fs", runtime_sec_for_rtf)
+    logging.info("生成语音耗时: %.3fs", final_audio_sec)
+    logging.info("RTF: %.4f", rtf)
+
+
+if __name__ == "__main__":
+    main()

models/axmodels/decoder_part0.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c174cd5458a1490936055a5060b00f9198431b5b97a1cd64176fcb67378e548b
+size 46276224

models/axmodels/decoder_part1.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94d2c75c7b2199f829fe0a9a44762ac8bac194e3e5ae589826a69671f3031442
+size 35358649

models/axmodels/decoder_part2.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9eb7ceec3baaef0b4e56513a791ffedaa3fbea63aeac41fcdd96f6850bf0fb8
+size 38312969

models/axmodels/decoder_part3.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2cb91ec05e2d4cb69de73b4e258ce8eb85ce7d507fa795029bb071a336c72821
+size 47283806

models/axmodels/encoder.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d6d7564e79edb513c0d456663c27a7606d272fa8f472b58364a59d9682a3ee3
+size 14093939

outputs/en_long_paragraph.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d59398ed12e49d51011777d80a0dbbf8b3c6322be7ec4f12338e8ca2740ddc2
+size 3108012

outputs/en_sentence.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc94c226559874053edde5331913364b6e2cd01ef7dde58e3c8df7735296e377
+size 307756

outputs/zh_long_paragraph.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6805938a16e5b4a257444c2b33897ffba0d3a5cdbdb7a2de9dd93dd629273947
+size 2147756

outputs/zh_sentence.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2997789e759fd6852ea30c73226d04a09369dabb777244c75d641798b09f1dab
+size 307756

requirements.txt

@@ -0,0 +1,9 @@
+numpy>=1.24.0
+soundfile
+torch
+torchaudio
+vocos
+jieba
+pypinyin
+piper_phonemize
+

resources/vocos-mel-24khz/config.yaml

@@ -0,0 +1,24 @@
+feature_extractor:
+  class_path: vocos.feature_extractors.MelSpectrogramFeatures
+  init_args:
+    sample_rate: 24000
+    n_fft: 1024
+    hop_length: 256
+    n_mels: 100
+    padding: center
+
+backbone:
+  class_path: vocos.models.VocosBackbone
+  init_args:
+    input_channels: 100
+    dim: 512
+    intermediate_dim: 1536
+    num_layers: 8
+
+head:
+  class_path: vocos.heads.ISTFTHead
+  init_args:
+    dim: 512
+    n_fft: 1024
+    hop_length: 256
+    padding: center

resources/vocos-mel-24khz/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97ec976ad1fd67a33ab2682d29c0ac7df85234fae875aefcc5fb215681a91b2a
+size 54365991

resources/zipvoice_hf/zipvoice/model.json

@@ -0,0 +1,26 @@
+{
+  "model" : {
+    "fm_decoder_downsampling_factor" : [1,2,4,2,1],
+    "fm_decoder_num_layers" : [2,2,4,4,4],
+    "fm_decoder_cnn_module_kernel" : [31,15,7,15,31],
+    "fm_decoder_feedforward_dim" : 1536,
+    "fm_decoder_num_heads" : 4,
+    "fm_decoder_dim" : 512,
+    "text_encoder_num_layers" : 4,
+    "text_encoder_feedforward_dim" : 512,
+    "text_encoder_cnn_module_kernel" : 9,
+    "text_encoder_num_heads" : 4,
+    "text_encoder_dim" : 192,
+    "query_head_dim" : 32,
+    "value_head_dim" : 12,
+    "pos_head_dim" : 4,
+    "pos_dim" : 48,
+    "time_embed_dim" : 192,
+    "text_embed_dim" : 192,
+    "feat_dim": 100
+  },
+  "feature" : {
+    "sampling_rate": 24000,
+    "type": "vocos"
+  }
+}

resources/zipvoice_hf/zipvoice/tokens.txt

@@ -0,0 +1,360 @@
+_	0
+^	1
+$	2
+ 	3
+!	4
+'	5
+(	6
+)	7
+,	8
+-	9
+.	10
+:	11
+;	12
+?	13
+a	14
+b	15
+c	16
+d	17
+e	18
+f	19
+h	20
+i	21
+j	22
+k	23
+l	24
+m	25
+n	26
+o	27
+p	28
+q	29
+r	30
+s	31
+t	32
+u	33
+v	34
+w	35
+x	36
+y	37
+z	38
+æ	39
+ç	40
+ð	41
+ø	42
+ħ	43
+ŋ	44
+œ	45
+ǀ	46
+ǁ	47
+ǂ	48
+ǃ	49
+ɐ	50
+ɑ	51
+ɒ	52
+ɓ	53
+ɔ	54
+ɕ	55
+ɖ	56
+ɗ	57
+ɘ	58
+ə	59
+ɚ	60
+ɛ	61
+ɜ	62
+ɞ	63
+ɟ	64
+ɠ	65
+ɡ	66
+ɢ	67
+ɣ	68
+ɤ	69
+ɥ	70
+ɦ	71
+ɧ	72
+ɨ	73
+ɪ	74
+ɫ	75
+ɬ	76
+ɭ	77
+ɮ	78
+ɯ	79
+ɰ	80
+ɱ	81
+ɲ	82
+ɳ	83
+ɴ	84
+ɵ	85
+ɶ	86
+ɸ	87
+ɹ	88
+ɺ	89
+ɻ	90
+ɽ	91
+ɾ	92
+ʀ	93
+ʁ	94
+ʂ	95
+ʃ	96
+ʄ	97
+ʈ	98
+ʉ	99
+ʊ	100
+ʋ	101
+ʌ	102
+ʍ	103
+ʎ	104
+ʏ	105
+ʐ	106
+ʑ	107
+ʒ	108
+ʔ	109
+ʕ	110
+ʘ	111
+ʙ	112
+ʛ	113
+ʜ	114
+ʝ	115
+ʟ	116
+ʡ	117
+ʢ	118
+ʲ	119
+ˈ	120
+ˌ	121
+ː	122
+ˑ	123
+˞	124
+β	125
+θ	126
+χ	127
+ᵻ	128
+ⱱ	129
+0	130
+1	131
+2	132
+3	133
+4	134
+5	135
+6	136
+7	137
+8	138
+9	139
+̧	140
+̃	141
+̪	142
+̯	143
+̩	144
+ʰ	145
+ˤ	146
+ε	147
+↓	148
+#	149
+"	150
+↑	151
+̺	152
+̻	153
+g	154
+ʦ	155
+X	156
+̝	157
+̊	158
+a1	159
+a2	160
+a3	161
+a4	162
+a5	163
+ai1	164
+ai2	165
+ai3	166
+ai4	167
+ai5	168
+an1	169
+an2	170
+an3	171
+an4	172
+an5	173
+ang1	174
+ang2	175
+ang3	176
+ang4	177
+ang5	178
+ao1	179
+ao2	180
+ao3	181
+ao4	182
+ao5	183
+b0	184
+c0	185
+ch0	186
+d0	187
+e1	188
+e2	189
+e3	190
+e4	191
+e5	192
+ei1	193
+ei2	194
+ei3	195
+ei4	196
+ei5	197
+en1	198
+en2	199
+en3	200
+en4	201
+en5	202
+eng1	203
+eng2	204
+eng3	205
+eng4	206
+eng5	207
+er2	208
+er3	209
+er4	210
+er5	211
+f0	212
+g0	213
+g2	214
+g3	215
+g4	216
+g5	217
+h0	218
+i1	219
+i2	220
+i3	221
+i4	222
+i5	223
+ia1	224
+ia2	225
+ia3	226
+ia4	227
+ia5	228
+ian1	229
+ian2	230
+ian3	231
+ian4	232
+ian5	233
+iang1	234
+iang2	235
+iang3	236
+iang4	237
+iang5	238
+iao1	239
+iao2	240
+iao3	241
+iao4	242
+iao5	243
+ie1	244
+ie2	245
+ie3	246
+ie4	247
+ie5	248
+in1	249
+in2	250
+in3	251
+in4	252
+in5	253
+ing1	254
+ing2	255
+ing3	256
+ing4	257
+ing5	258
+iong1	259
+iong2	260
+iong3	261
+iong4	262
+iu1	263
+iu2	264
+iu3	265
+iu4	266
+iu5	267
+j0	268
+k0	269
+l0	270
+m0	271
+m1	272
+m2	273
+m4	274
+m5	275
+n0	276
+n2	277
+n3	278
+n4	279
+n5	280
+ng5	281
+o1	282
+o2	283
+o3	284
+o4	285
+o5	286
+ong1	287
+ong2	288
+ong3	289
+ong4	290
+ong5	291
+ou1	292
+ou2	293
+ou3	294
+ou4	295
+ou5	296
+p0	297
+q0	298
+r0	299
+s0	300
+sh0	301
+t0	302
+u1	303
+u2	304
+u3	305
+u4	306
+u5	307
+ua1	308
+ua2	309
+ua3	310
+ua4	311
+uai1	312
+uai2	313
+uai3	314
+uai4	315
+uai5	316
+uan1	317
+uan2	318
+uan3	319
+uan4	320
+uan5	321
+uang1	322
+uang2	323
+uang3	324
+uang4	325
+uang5	326
+ue1	327
+ue2	328
+ue3	329
+ue4	330
+ui1	331
+ui2	332
+ui3	333
+ui4	334
+ui5	335
+un1	336
+un2	337
+un3	338
+un4	339
+un5	340
+uo1	341
+uo2	342
+uo3	343
+uo4	344
+uo5	345
+v2	346
+v3	347
+v4	348
+ve3	349
+ve4	350
+w0	351
+x0	352
+y0	353
+z0	354
+zh0	355
+ê1	356
+ê2	357
+ê3	358
+ê4	359

resources/zipvoice_hf/zipvoice/zipvoice_base.json

@@ -0,0 +1,26 @@
+{
+  "model" : {
+    "fm_decoder_downsampling_factor" : [1,2,4,2,1],
+    "fm_decoder_num_layers" : [2,2,4,4,4],
+    "fm_decoder_cnn_module_kernel" : [31,15,7,15,31],
+    "fm_decoder_feedforward_dim" : 1536,
+    "fm_decoder_num_heads" : 4,
+    "fm_decoder_dim" : 512,
+    "text_encoder_num_layers" : 4,
+    "text_encoder_feedforward_dim" : 512,
+    "text_encoder_cnn_module_kernel" : 9,
+    "text_encoder_num_heads" : 4,
+    "text_encoder_dim" : 192,
+    "query_head_dim" : 32,
+    "value_head_dim" : 12,
+    "pos_head_dim" : 4,
+    "pos_dim" : 48,
+    "time_embed_dim" : 192,
+    "text_embed_dim" : 192,
+    "feat_dim": 100
+  },
+  "feature" : {
+    "sampling_rate": 24000,
+    "type": "vocos"
+  }
+}

scripts/__init__.py

@@ -0,0 +1 @@
+# ZipVoice board inference scripts

scripts/common_infer.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+
+def load_tokenizer(repo_dir: str | Path):
+    from scripts.local_tokenizer import LocalEmiliaTokenizer
+
+    token_file = Path(repo_dir) / "resources" / "zipvoice_hf" / "zipvoice" / "tokens.txt"
+    if not token_file.is_file():
+        raise FileNotFoundError(f"tokens.txt not found: {token_file}")
+    return LocalEmiliaTokenizer(token_file=str(token_file))
+
+
+def extract_prompt_features(
+    prompt_wav: str | Path,
+    repo_dir: str | Path,
+    sampling_rate: int,
+    feat_scale: float,
+    target_rms: float,
+):
+    from scripts.local_audio import LocalVocosFbank, load_prompt_wav, rms_norm
+    import torch
+
+    wav = load_prompt_wav(prompt_wav, sampling_rate=sampling_rate)
+    wav, prompt_rms = rms_norm(wav, target_rms)
+    extractor = LocalVocosFbank()
+    features = extractor.extract(wav, sampling_rate=sampling_rate)
+    if not isinstance(features, torch.Tensor):
+        features = torch.from_numpy(features)
+    features = features.unsqueeze(0) * feat_scale
+    return features.cpu().numpy().astype(np.float32), float(prompt_rms)
+
+
+def load_vocoder(repo_dir: str | Path):
+    from scripts.local_audio import load_local_vocos
+    import torch
+
+    vocoder_dir = Path(repo_dir) / "resources" / "vocos-mel-24khz"
+    if not (vocoder_dir / "config.yaml").is_file() or not (
+        vocoder_dir / "pytorch_model.bin"
+    ).is_file():
+        raise FileNotFoundError(f"Local Vocos files not found in {vocoder_dir}")
+    vocoder = load_local_vocos(vocoder_dir)
+    vocoder = vocoder.to(torch.device("cpu"))
+    vocoder.eval()
+    return vocoder
+
+
+def vocoder_decode_loaded(
+    vocoder,
+    features: np.ndarray,
+    feat_scale: float,
+    target_rms: float,
+    prompt_rms: float,
+) -> np.ndarray:
+    from scripts.local_audio import rms_norm
+    import torch
+
+    feat_tensor = torch.from_numpy(features).float().permute(0, 2, 1) / feat_scale
+    with torch.no_grad():
+        wav = vocoder.decode(feat_tensor).squeeze(1).clamp(-1, 1)
+    wav = rms_norm(wav, target_rms)[0]
+    if prompt_rms < target_rms:
+        wav = wav * prompt_rms / target_rms
+    return wav.squeeze().cpu().numpy()

scripts/local_audio.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torchaudio
+from vocos import Vocos
+
+
+@dataclass
+class VocosFbankConfig:
+    sampling_rate: int = 24000
+    n_mels: int = 100
+    n_fft: int = 1024
+    hop_length: int = 256
+
+
+def compute_num_frames(num_samples: int, hop_length: int) -> int:
+    return int((num_samples + hop_length // 2) // hop_length)
+
+
+class LocalVocosFbank:
+    def __init__(self) -> None:
+        self.config = VocosFbankConfig()
+        self.fbank = torchaudio.transforms.MelSpectrogram(
+            sample_rate=self.config.sampling_rate,
+            n_fft=self.config.n_fft,
+            hop_length=self.config.hop_length,
+            n_mels=self.config.n_mels,
+            center=True,
+            power=1,
+        )
+
+    def extract(self, samples: torch.Tensor, sampling_rate: int) -> torch.Tensor:
+        if sampling_rate != self.config.sampling_rate:
+            raise ValueError(
+                f"Mismatched sampling rate: expected {self.config.sampling_rate}, got {sampling_rate}"
+            )
+        if samples.ndim == 1:
+            samples = samples.unsqueeze(0)
+        if samples.ndim != 2:
+            raise ValueError(f"Expected waveform shape [C, T], got {tuple(samples.shape)}")
+        if samples.shape[0] == 2:
+            samples = samples.mean(dim=0, keepdim=True)
+
+        mel = self.fbank(samples).clamp(min=1e-7).log()
+        mel = mel.reshape(-1, mel.shape[-1]).t()
+        num_frames = compute_num_frames(samples.shape[1], self.config.hop_length)
+
+        if mel.shape[0] > num_frames:
+            mel = mel[:num_frames]
+        elif mel.shape[0] < num_frames:
+            mel = torch.nn.functional.pad(
+                mel.unsqueeze(0),
+                (0, 0, 0, num_frames - mel.shape[0]),
+                mode="replicate",
+            ).squeeze(0)
+        return mel
+
+
+def load_prompt_wav(prompt_wav: str | Path, sampling_rate: int) -> torch.Tensor:
+    wav, sr = torchaudio.load(str(prompt_wav))
+    if sr != sampling_rate:
+        wav = torchaudio.transforms.Resample(orig_freq=sr, new_freq=sampling_rate)(wav)
+    return wav
+
+
+def rms_norm(wav: torch.Tensor, target_rms: float):
+    wav_rms = torch.sqrt(torch.mean(torch.square(wav)))
+    if wav_rms < target_rms:
+        wav = wav * target_rms / wav_rms
+    return wav, wav_rms
+
+
+def load_local_vocos(vocoder_dir: str | Path):
+    vocoder_dir = Path(vocoder_dir)
+    vocoder = Vocos.from_hparams(str(vocoder_dir / "config.yaml"))
+    state_dict = torch.load(
+        str(vocoder_dir / "pytorch_model.bin"),
+        weights_only=True,
+        map_location="cpu",
+    )
+    vocoder.load_state_dict(state_dict)
+    return vocoder

scripts/local_tokenizer.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import logging
+import re
+from functools import reduce
+from pathlib import Path
+from typing import Dict, List
+
+import jieba
+from pypinyin import Style, lazy_pinyin
+from pypinyin.contrib.tone_convert import to_finals_tone3, to_initials
+
+try:
+    from piper_phonemize import phonemize_espeak
+except Exception as ex:  # pragma: no cover - board dependency check
+    raise RuntimeError(
+        f"{ex}\nPlease install piper_phonemize for English tokenization."
+    )
+
+jieba.default_logger.setLevel(logging.INFO)
+
+
+class EnglishTextNormalizer:
+    def normalize(self, text: str) -> str:
+        return text
+
+
+class ChineseTextNormalizer:
+    def normalize(self, text: str) -> str:
+        try:
+            import cn2an
+
+            return cn2an.transform(text, "an2cn")
+        except Exception:
+            return text
+
+
+class LocalEmiliaTokenizer:
+    """Small board-side Emilia tokenizer without lhotse/CutSet dependencies."""
+
+    def __init__(self, token_file: str | Path):
+        self.english_normalizer = EnglishTextNormalizer()
+        self.chinese_normalizer = ChineseTextNormalizer()
+        self.token2id: Dict[str, int] = {}
+        with open(token_file, "r", encoding="utf-8") as f:
+            for line in f:
+                token, token_id = line.rstrip().split("\t")
+                self.token2id[token] = int(token_id)
+        self.pad_id = self.token2id["_"]
+        self.vocab_size = len(self.token2id)
+
+    def texts_to_token_ids(self, texts: List[str]) -> List[List[int]]:
+        return self.tokens_to_token_ids(self.texts_to_tokens(texts))
+
+    def texts_to_tokens(self, texts: List[str]) -> List[List[str]]:
+        phoneme_list = []
+        for text in texts:
+            text = self.map_punctuations(text)
+            segments = self.get_segment(text)
+            all_phoneme = []
+            for seg_text, seg_type in segments:
+                if seg_type == "zh":
+                    all_phoneme += self.tokenize_zh(seg_text)
+                elif seg_type == "en":
+                    all_phoneme += self.tokenize_en(seg_text)
+                elif seg_type == "pinyin":
+                    all_phoneme += self.tokenize_pinyin(seg_text)
+                elif seg_type == "tag":
+                    all_phoneme.append(seg_text)
+                else:
+                    logging.debug("Skipping unknown language segment: %r", (seg_text, seg_type))
+            phoneme_list.append(all_phoneme)
+        return phoneme_list
+
+    def tokens_to_token_ids(self, tokens_list: List[List[str]]) -> List[List[int]]:
+        token_ids_list = []
+        for tokens in tokens_list:
+            token_ids = []
+            for token in tokens:
+                if token not in self.token2id:
+                    logging.debug("Skip OOV token %s", token)
+                    continue
+                token_ids.append(self.token2id[token])
+            token_ids_list.append(token_ids)
+        return token_ids_list
+
+    def tokenize_zh(self, text: str) -> List[str]:
+        try:
+            text = self.chinese_normalizer.normalize(text)
+            segs = list(jieba.cut(text))
+            full = lazy_pinyin(
+                segs,
+                style=Style.TONE3,
+                tone_sandhi=True,
+                neutral_tone_with_five=True,
+            )
+            phones = []
+            for item in full:
+                if not (item[0:-1].isalpha() and item[-1] in ("1", "2", "3", "4", "5")):
+                    phones.append(item)
+                else:
+                    phones.extend(self.separate_pinyin(item))
+            return phones
+        except Exception as ex:
+            logging.debug("Tokenization of Chinese text failed: %s", ex)
+            return []
+
+    def tokenize_en(self, text: str) -> List[str]:
+        try:
+            text = self.english_normalizer.normalize(text)
+            tokens = phonemize_espeak(text, "en-us")
+            return reduce(lambda x, y: x + y, tokens)
+        except Exception as ex:
+            logging.debug("Tokenization of English text failed: %s", ex)
+            return []
+
+    def tokenize_pinyin(self, text: str) -> List[str]:
+        try:
+            text = text.lstrip("<").rstrip(">")
+            if not (text[0:-1].isalpha() and text[-1] in ("1", "2", "3", "4", "5")):
+                logging.debug("Invalid pinyin token: %s", text)
+                return []
+            return self.separate_pinyin(text)
+        except Exception as ex:
+            logging.debug("Tokenize pinyin failed: %s", ex)
+            return []
+
+    @staticmethod
+    def separate_pinyin(text: str) -> List[str]:
+        pinyins = []
+        initial = to_initials(text, strict=False)
+        final = to_finals_tone3(text, strict=False, neutral_tone_with_five=True)
+        if initial:
+            pinyins.append(initial + "0")
+        if final:
+            pinyins.append(final)
+        return pinyins
+
+    @staticmethod
+    def map_punctuations(text: str) -> str:
+        replacements = {
+            "，": ",",
+            "。": ".",
+            "！": "!",
+            "？": "?",
+            "；": ";",
+            "：": ":",
+            "、": ",",
+            "‘": "'",
+            "“": '"',
+            "”": '"',
+            "’": "'",
+            "⋯": "…",
+            "···": "…",
+            "・・・": "…",
+            "...": "…",
+        }
+        for src, dst in replacements.items():
+            text = text.replace(src, dst)
+        return text
+
+    def get_segment(self, text: str):
+        segments = []
+        types = []
+        temp_seg = ""
+        temp_lang = ""
+        parts = re.compile(r"[<[].*?[>\]]|.").findall(text)
+        for part in parts:
+            if self.is_chinese(part) or self.is_pinyin(part):
+                types.append("zh")
+            elif self.is_alphabet(part):
+                types.append("en")
+            else:
+                types.append("other")
+
+        for index, part_type in enumerate(types):
+            if index == 0:
+                temp_seg += parts[index]
+                temp_lang = part_type
+            elif temp_lang == "other":
+                temp_seg += parts[index]
+                temp_lang = part_type
+            elif part_type in [temp_lang, "other"]:
+                temp_seg += parts[index]
+            else:
+                segments.append((temp_seg, temp_lang))
+                temp_seg = parts[index]
+                temp_lang = part_type
+        segments.append((temp_seg, temp_lang))
+        return self.split_segments(segments)
+
+    @staticmethod
+    def split_segments(segments):
+        result = []
+        for temp_seg, temp_lang in segments:
+            parts = re.split(r"([<[].*?[>\]])", temp_seg)
+            for part in parts:
+                if not part:
+                    continue
+                if part.startswith("<") and part.endswith(">"):
+                    result.append((part, "pinyin"))
+                elif part.startswith("[") and part.endswith("]"):
+                    result.append((part, "tag"))
+                else:
+                    result.append((part, temp_lang))
+        return result
+
+    @staticmethod
+    def is_chinese(char: str) -> bool:
+        return "\u4e00" <= char <= "\u9fa5"
+
+    @staticmethod
+    def is_alphabet(char: str) -> bool:
+        return ("A" <= char <= "Z") or ("a" <= char <= "z")
+
+    @staticmethod
+    def is_pinyin(part: str) -> bool:
+        return part.startswith("<") and part.endswith(">")

scripts/text_processing.py

@@ -0,0 +1,390 @@
+from __future__ import annotations
+
+import argparse
+import re
+import unicodedata
+from pathlib import Path
+
+import numpy as np
+
+
+SENTENCE_PUNCT = "。！？.!?｡．"
+ELLIPSIS_PUNCT = {"…"}
+QUOTE_PUNCT = set("\"'“”‘’「」『』")
+
+
+def is_text_punctuation(ch: str) -> bool:
+    return ch in ELLIPSIS_PUNCT or unicodedata.category(ch).startswith("P")
+
+
+def comma_for_context(left: str, right: str) -> str:
+    context = left + right
+    return "," if re.search(r"[A-Za-z0-9]", context) else "，"
+
+
+def setup_jieba_cache(repo_dir: Path) -> None:
+    import jieba
+
+    cache_dir = repo_dir / ".work_tmp" / "jieba"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    jieba.dt.tmp_dir = str(cache_dir)
+    jieba.dt.cache_file = "jieba.cache"
+
+
+def load_text(args: argparse.Namespace) -> str:
+    if args.text_file:
+        text = Path(args.text_file).read_text(encoding="utf-8")
+    elif args.text:
+        text = args.text
+    else:
+        raise ValueError("Either --text or --text-file is required")
+    text = re.sub(r"[ \t\r\f\v]+", " ", text)
+    text = re.sub(r"\n+", " ", text)
+    return text.strip()
+
+
+def normalize_punctuation(text: str) -> str:
+    text = re.sub(r"\s+", " ", text.strip())
+    text = text.replace("……", "。").replace("...", ".")
+    chars: list[str] = []
+    i = 0
+    while i < len(text):
+        ch = text[i]
+        if is_text_punctuation(ch):
+            j = i
+            punct_run = []
+            while j < len(text) and is_text_punctuation(text[j]):
+                punct_run.append(text[j])
+                j += 1
+            if all(p in QUOTE_PUNCT for p in punct_run):
+                i = j
+                continue
+            sentence_marks = [p for p in punct_run if p in SENTENCE_PUNCT]
+            left = chars[-1] if chars else ""
+            right = text[j] if j < len(text) else ""
+            mark = sentence_marks[-1] if sentence_marks else comma_for_context(left, right)
+            if chars and is_text_punctuation(chars[-1]):
+                chars[-1] = mark
+            else:
+                chars.append(mark)
+            i = j
+            continue
+        chars.append(ch)
+        i += 1
+
+    normalized = "".join(chars)
+    normalized = re.sub(r"[ \t]+([，。！？,.!?．｡])", r"\1", normalized)
+    normalized = re.sub(r"([，。！？])\s+", r"\1", normalized)
+    normalized = re.sub(r"([,.!?．｡])(?=[A-Za-z0-9])", r"\1 ", normalized)
+    normalized = re.sub(r"\s{2,}", " ", normalized)
+    normalized = re.sub(r"([，,])\s*[，,]+", r"\1", normalized)
+    normalized = re.sub(r"[，,]+([。！？.!?．｡])", r"\1", normalized)
+    normalized = re.sub(r"([。！？.!?．｡])([，,]+)", r"\1", normalized)
+    return normalized.strip()
+
+
+def join_units(left: str, right: str) -> str:
+    if not left:
+        return right.strip()
+    right = right.strip()
+    if not right:
+        return left.strip()
+    if re.search(r"[\u4e00-\u9fff]$", left) or re.match(r"^[\u4e00-\u9fff]", right):
+        return left.rstrip() + right
+    return left.rstrip() + " " + right
+
+
+def split_units(text: str) -> list[str]:
+    text = text.strip()
+    if not text:
+        return []
+    pattern = r"[^。！？!?；;，,、：:\n]+[。！？!?；;，,、：:]?"
+    units = [m.group(0).strip() for m in re.finditer(pattern, text) if m.group(0).strip()]
+    return units or [text]
+
+
+def estimate_lengths(
+    prompt_frames: int,
+    prompt_tokens: int,
+    text_tokens: int,
+    speed: float,
+    max_feat_len: int,
+) -> dict[str, int | bool]:
+    raw_features_len = int(
+        np.ceil(prompt_frames / prompt_tokens * (prompt_tokens + text_tokens) / speed)
+    )
+    features_len = min(raw_features_len, max_feat_len)
+    generated_frames = features_len - prompt_frames
+    if generated_frames <= 0:
+        generated_frames = features_len
+    return {
+        "raw_features_len": raw_features_len,
+        "features_len": features_len,
+        "generated_frames": generated_frames,
+        "clamped": raw_features_len > max_feat_len,
+    }
+
+
+def token_count(tokenizer, text: str) -> int:
+    return len(tokenizer.texts_to_token_ids([text])[0])
+
+
+def split_long_unit(tokenizer, unit: str, max_text_tokens: int) -> list[str]:
+    if token_count(tokenizer, unit) <= max_text_tokens:
+        return [unit]
+
+    if " " in unit:
+        pieces = unit.split()
+        chunks: list[str] = []
+        current = ""
+        for piece in pieces:
+            candidate = join_units(current, piece)
+            if current and token_count(tokenizer, candidate) > max_text_tokens:
+                chunks.append(current)
+                current = piece
+            else:
+                current = candidate
+        if current:
+            chunks.append(current)
+        return chunks
+
+    chunks = []
+    current = ""
+    for char in unit:
+        candidate = current + char
+        if current and token_count(tokenizer, candidate) > max_text_tokens:
+            chunks.append(current)
+            current = char
+        else:
+            current = candidate
+    if current:
+        chunks.append(current)
+    return chunks
+
+
+def build_segments(
+    tokenizer,
+    text: str,
+    prompt_frames: int,
+    prompt_tokens_len: int,
+    speed: float,
+    max_feat_len: int,
+    max_text_tokens: int,
+    min_generated_frames: int,
+    max_generated_frames: int,
+    max_raw_feat_ratio: float,
+) -> list[dict[str, object]]:
+    raw_units = split_units(text)
+    units: list[str] = []
+    for unit in raw_units:
+        units.extend(split_long_unit(tokenizer, unit, max_text_tokens))
+
+    segments: list[dict[str, object]] = []
+    current = ""
+    for unit in units:
+        candidate = join_units(current, unit)
+        cand_tokens = token_count(tokenizer, candidate)
+        cand_est = estimate_lengths(
+            prompt_frames,
+            prompt_tokens_len,
+            cand_tokens,
+            speed,
+            max_feat_len,
+        )
+        raw_too_long = int(cand_est["raw_features_len"]) > int(max_feat_len * max_raw_feat_ratio)
+        too_long = (
+            cand_tokens > max_text_tokens
+            or int(cand_est["generated_frames"]) > max_generated_frames
+            or raw_too_long
+        )
+        if current and too_long:
+            cur_tokens = token_count(tokenizer, current)
+            cur_est = estimate_lengths(
+                prompt_frames,
+                prompt_tokens_len,
+                cur_tokens,
+                speed,
+                max_feat_len,
+            )
+            segments.append({"text": current, "text_tokens": cur_tokens, **cur_est})
+            current = unit
+        else:
+            current = candidate
+
+    if current:
+        cur_tokens = token_count(tokenizer, current)
+        cur_est = estimate_lengths(
+            prompt_frames,
+            prompt_tokens_len,
+            cur_tokens,
+            speed,
+            max_feat_len,
+        )
+        segments.append({"text": current, "text_tokens": cur_tokens, **cur_est})
+
+    if len(segments) >= 2 and int(segments[-1]["generated_frames"]) < min_generated_frames:
+        merged_text = join_units(str(segments[-2]["text"]), str(segments[-1]["text"]))
+        merged_tokens = token_count(tokenizer, merged_text)
+        merged_est = estimate_lengths(
+            prompt_frames,
+            prompt_tokens_len,
+            merged_tokens,
+            speed,
+            max_feat_len,
+        )
+        raw_ok = int(merged_est["raw_features_len"]) <= int(max_feat_len * max_raw_feat_ratio)
+        if (
+            merged_tokens <= max_text_tokens
+            and int(merged_est["generated_frames"]) <= max_generated_frames
+            and raw_ok
+        ):
+            segments[-2] = {"text": merged_text, "text_tokens": merged_tokens, **merged_est}
+            segments.pop()
+        else:
+            rebalanced = rebalance_short_tail(
+                tokenizer,
+                str(segments[-2]["text"]),
+                str(segments[-1]["text"]),
+                prompt_frames,
+                prompt_tokens_len,
+                speed,
+                max_feat_len,
+                max_text_tokens,
+                min_generated_frames,
+                max_generated_frames,
+                max_raw_feat_ratio,
+            )
+            if rebalanced is not None:
+                segments[-2], segments[-1] = rebalanced
+
+    return segments
+
+
+def split_rebalance_pieces(text: str) -> list[str]:
+    units = split_units(text)
+    if len(units) > 1:
+        return units
+    if " " in text:
+        return text.split()
+    return list(text)
+
+
+def segment_record(
+    tokenizer,
+    text: str,
+    prompt_frames: int,
+    prompt_tokens_len: int,
+    speed: float,
+    max_feat_len: int,
+) -> dict[str, object]:
+    text_tokens = token_count(tokenizer, text)
+    est = estimate_lengths(
+        prompt_frames,
+        prompt_tokens_len,
+        text_tokens,
+        speed,
+        max_feat_len,
+    )
+    return {"text": text, "text_tokens": text_tokens, **est}
+
+
+def segment_within_limits(
+    record: dict[str, object],
+    max_text_tokens: int,
+    max_generated_frames: int,
+    max_raw_feat: int,
+) -> bool:
+    return (
+        int(record["text_tokens"]) <= max_text_tokens
+        and int(record["generated_frames"]) <= max_generated_frames
+        and int(record["raw_features_len"]) <= max_raw_feat
+    )
+
+
+def rebalance_short_tail(
+    tokenizer,
+    prev_text: str,
+    tail_text: str,
+    prompt_frames: int,
+    prompt_tokens_len: int,
+    speed: float,
+    max_feat_len: int,
+    max_text_tokens: int,
+    min_generated_frames: int,
+    max_generated_frames: int,
+    max_raw_feat_ratio: float,
+) -> tuple[dict[str, object], dict[str, object]] | None:
+    max_raw_feat = int(max_feat_len * max_raw_feat_ratio)
+    best: tuple[dict[str, object], dict[str, object]] | None = None
+    best_tail_frames = -1
+
+    piece_sets = [split_rebalance_pieces(prev_text)]
+    if " " in prev_text:
+        word_pieces = prev_text.split()
+        if len(word_pieces) > 1 and word_pieces != piece_sets[0]:
+            piece_sets.append(word_pieces)
+
+    for pieces in piece_sets:
+        if len(pieces) < 2:
+            continue
+        for split_at in range(len(pieces) - 1, 0, -1):
+            new_prev = (
+                join_units("", " ".join(pieces[:split_at]))
+                if " " in prev_text
+                else "".join(pieces[:split_at])
+            )
+            moved_tail = (
+                join_units("", " ".join(pieces[split_at:]))
+                if " " in prev_text
+                else "".join(pieces[split_at:])
+            )
+            new_tail = join_units(moved_tail, tail_text)
+            prev_record = segment_record(
+                tokenizer,
+                new_prev,
+                prompt_frames,
+                prompt_tokens_len,
+                speed,
+                max_feat_len,
+            )
+            tail_record = segment_record(
+                tokenizer,
+                new_tail,
+                prompt_frames,
+                prompt_tokens_len,
+                speed,
+                max_feat_len,
+            )
+            if not segment_within_limits(
+                tail_record,
+                max_text_tokens,
+                max_generated_frames,
+                max_raw_feat,
+            ):
+                continue
+            if not segment_within_limits(
+                prev_record,
+                max_text_tokens,
+                max_generated_frames,
+                max_raw_feat,
+            ):
+                continue
+            prev_frames = int(prev_record["generated_frames"])
+            tail_frames = int(tail_record["generated_frames"])
+            if prev_frames < min_generated_frames:
+                continue
+            if tail_frames > best_tail_frames:
+                best = (prev_record, tail_record)
+                best_tail_frames = tail_frames
+            if tail_frames >= min_generated_frames:
+                return best
+
+    return best
+
+
+def build_cat_tokens(tokenizer, prompt_tokens: list[int], text_tokens: list[int], max_tokens: int) -> np.ndarray:
+    pad_id = tokenizer.pad_id
+    cat = prompt_tokens + text_tokens + [pad_id]
+    cat_tokens = np.full((1, max_tokens), pad_id, dtype=np.int64)
+    cat_tokens[0, : len(cat)] = np.array(cat, dtype=np.int64)
+    return cat_tokens

scripts/zipvoice_decoder4_runtime.py

@@ -0,0 +1,364 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from pathlib import Path
+from typing import Any, Dict, Iterable, List
+
+import numpy as np
+
+from scripts.zipvoice_runtime import AxeSession
+
+
+class Decoder4ZipVoiceBoardRuntime:
+    """Runs encoder_core_nolog.axmodel and fm_decoder_part0..part3.axmodel."""
+
+    def __init__(
+        self,
+        config_dir: str | Path,
+        models_dir: str | Path,
+        max_feat_len: int = 1024,
+        max_tokens: int = 384,
+        num_step: int = 16,
+        t_shift: float = 0.5,
+    ) -> None:
+        self.config_dir = Path(config_dir)
+        self.models_dir = Path(models_dir)
+        self.max_feat_len = int(max_feat_len)
+        self.max_tokens = int(max_tokens)
+        self.num_step = int(num_step)
+        self.t_shift = float(t_shift)
+
+        self._load_config()
+        self._load_manifest()
+        self.sessions: Dict[str, AxeSession] = {}
+        self._load_models()
+        self._load_decoder_input_metadata()
+
+    def _load_config(self) -> None:
+        config_path = self.config_dir / "runtime_config.json"
+        config = json.loads(config_path.read_text()) if config_path.exists() else {}
+        self.feat_dim = int(config.get("feat_dim", 100))
+        self.sampling_rate = int(config.get("sampling_rate", 24000))
+        self.hop_length = int(config.get("hop_length", 256))
+        logging.debug(
+            "Decoder4 runtime: max_tokens=%d, max_feat_len=%d, feat_dim=%d, num_step=%d",
+            self.max_tokens,
+            self.max_feat_len,
+            self.feat_dim,
+            self.num_step,
+        )
+
+    def _load_manifest(self) -> None:
+        manifest_path = self.config_dir / "decoder4_split_manifest.json"
+        if not manifest_path.exists():
+            raise FileNotFoundError(f"decoder4_split_manifest.json not found: {manifest_path}")
+        self.manifest = json.loads(manifest_path.read_text())
+        self.encoder_info = self.manifest["encoder"]
+        self.decoder_parts = self.manifest["decoder_parts"]
+
+    def _load_models(self) -> None:
+        model_infos = [self.encoder_info, *self.decoder_parts]
+        for info in model_infos:
+            name = info["name"]
+            path = self.models_dir / info["file"]
+            logging.debug("Loading %s from %s", name, path)
+            self.sessions[name] = AxeSession(path)
+        logging.debug("Loaded encoder + %d decoder4 parts", len(self.decoder_parts))
+
+    def _load_decoder_input_metadata(self) -> None:
+        part0 = self.decoder_parts[0]
+        sess = self.sessions[part0["name"]]
+        input_names = sess.input_names
+        self.decoder_has_padding_mask = "padding_mask" in input_names
+        self.decoder_seq_len = self.max_feat_len
+        if "x" in input_names:
+            index = input_names.index("x")
+            input_info = sess._inputs[index] if index < len(sess._inputs) else None
+            shape = getattr(input_info, "shape", None) if input_info is not None else None
+            if shape is not None and len(shape) >= 2 and isinstance(shape[1], (int, np.integer)):
+                self.decoder_seq_len = int(shape[1])
+
+        if self.decoder_seq_len != self.max_feat_len:
+            logging.debug(
+                "decoder x seq_len=%d differs from configured max_feat_len=%d; "
+                "using model seq_len for decoder feeds",
+                self.decoder_seq_len,
+                self.max_feat_len,
+            )
+        logging.debug(
+            "Decoder4 model metadata: seq_len=%d, has_padding_mask=%s",
+            self.decoder_seq_len,
+            self.decoder_has_padding_mask,
+        )
+
+    @staticmethod
+    def _coerce_input_dtype(value: np.ndarray, input_info: Any | None) -> np.ndarray:
+        if input_info is None:
+            return value
+
+        expected_dtype = getattr(input_info, "dtype", None)
+        if expected_dtype is None:
+            expected_dtype = getattr(input_info, "type", None)
+        if expected_dtype is None:
+            return value
+
+        dtype_text = str(expected_dtype).lower()
+        if "float32" in dtype_text:
+            return np.ascontiguousarray(value, dtype=np.float32)
+        if "int32" in dtype_text:
+            return np.ascontiguousarray(value, dtype=np.int32)
+        if "int64" in dtype_text:
+            return np.ascontiguousarray(value, dtype=np.int64)
+        if "uint8" in dtype_text:
+            return np.ascontiguousarray(value, dtype=np.uint8)
+        if "bool" in dtype_text:
+            return np.ascontiguousarray(value, dtype=np.bool_)
+        return np.ascontiguousarray(value)
+
+    def _run_model(
+        self,
+        name: str,
+        expected_inputs: Iterable[str],
+        expected_outputs: Iterable[str],
+        values: Dict[str, np.ndarray],
+    ) -> Dict[str, np.ndarray]:
+        sess = self.sessions[name]
+        expected_inputs = list(expected_inputs)
+        expected_outputs = list(expected_outputs)
+
+        feed: Dict[str, np.ndarray] = {}
+        for index, actual_name in enumerate(sess.input_names):
+            input_info = sess._inputs[index] if index < len(sess._inputs) else None
+            if actual_name in values:
+                feed[actual_name] = self._coerce_input_dtype(values[actual_name], input_info)
+                continue
+            if index < len(expected_inputs) and expected_inputs[index] in values:
+                feed[actual_name] = self._coerce_input_dtype(
+                    values[expected_inputs[index]], input_info
+                )
+                continue
+            expected = expected_inputs[index] if index < len(expected_inputs) else None
+            raise KeyError(
+                f"Missing input for {name}: actual={actual_name!r}, expected={expected!r}"
+            )
+
+        raw_outputs = sess.run(feed)
+        mapped: Dict[str, np.ndarray] = {}
+        for index, expected_name in enumerate(expected_outputs):
+            if expected_name in raw_outputs:
+                mapped[expected_name] = raw_outputs[expected_name]
+                continue
+            if index < len(sess.output_names) and sess.output_names[index] in raw_outputs:
+                mapped[expected_name] = raw_outputs[sess.output_names[index]]
+                continue
+            raise KeyError(f"Missing output for {name}: {expected_name!r}")
+        return mapped
+
+    def run_encoder(self, cat_tokens: np.ndarray) -> np.ndarray:
+        # Pulsar2/AXEngine exposes the quantized encoder token input as int32
+        # even though the reference ONNX path uses int64 token IDs.
+        cat_tokens = np.asarray(cat_tokens, dtype=np.int32)
+        outputs = self._run_model(
+            self.encoder_info["name"],
+            self.encoder_info["inputs"],
+            self.encoder_info["outputs"],
+            {"cat_tokens": cat_tokens},
+        )
+        return outputs[self.encoder_info["outputs"][0]].astype(np.float32)
+
+    def run_decoder(
+        self,
+        t: np.ndarray,
+        x: np.ndarray,
+        text_condition: np.ndarray,
+        speech_condition: np.ndarray,
+        guidance_scale: np.ndarray,
+        padding_mask: np.ndarray | None = None,
+    ) -> np.ndarray:
+        seq_len = x.shape[1]
+        values: Dict[str, np.ndarray] = {
+            "t": np.asarray(t, dtype=np.float32).reshape(1),
+            "x": x.astype(np.float32),
+            "text_condition": text_condition.astype(np.float32),
+            "speech_condition": speech_condition.astype(np.float32),
+            "guidance_scale": np.asarray(guidance_scale, dtype=np.float32).reshape(1),
+            "padding_mask": padding_mask.astype(np.bool_)
+            if padding_mask is not None
+            else np.zeros((1, seq_len), dtype=np.bool_),
+        }
+
+        for part in self.decoder_parts:
+            outputs = self._run_model(
+                part["name"],
+                part["inputs"],
+                part["outputs"],
+                values,
+            )
+            values.update(outputs)
+
+        final_output = self.decoder_parts[-1]["outputs"][0]
+        return values[final_output].astype(np.float32)
+
+    def duration_expand(
+        self,
+        encoded: np.ndarray,
+        prompt_tokens_len: int,
+        text_tokens_len: int,
+        prompt_features_len: int,
+        speed: float,
+    ) -> tuple[np.ndarray, int]:
+        total_tokens_len = prompt_tokens_len + text_tokens_len
+        features_len = int(
+            np.ceil(prompt_features_len / prompt_tokens_len * total_tokens_len / speed)
+        )
+        if features_len > self.max_feat_len:
+            logging.debug(
+                "features_len=%d > max_feat_len=%d, clamping",
+                features_len,
+                self.max_feat_len,
+            )
+            features_len = self.max_feat_len
+
+        token_dur = features_len // total_tokens_len
+        embed_no_pad = encoded[0, :total_tokens_len, :]
+        text_condition = np.repeat(embed_no_pad, token_dur, axis=0)
+
+        residual = features_len - text_condition.shape[0]
+        if residual > 0:
+            last_embed = encoded[0, total_tokens_len : total_tokens_len + 1, :]
+            text_condition = np.concatenate(
+                [text_condition, np.repeat(last_embed, residual, axis=0)],
+                axis=0,
+            )
+
+        text_condition = text_condition[:features_len, :]
+        return text_condition[np.newaxis, :, :].astype(np.float32), features_len
+
+    def _get_time_steps(self) -> np.ndarray:
+        t = np.linspace(0.0, 1.0, self.num_step + 1, dtype=np.float32)
+        ts = self.t_shift
+        return ts * t / (1.0 + (ts - 1.0) * t)
+
+    def sample(
+        self,
+        cat_tokens: np.ndarray,
+        prompt_tokens_len: int,
+        text_tokens_len: int,
+        prompt_features: np.ndarray,
+        prompt_features_len: int,
+        speed: float = 1.0,
+        guidance_scale: float = 1.0,
+        seed: int = 666,
+    ) -> tuple[np.ndarray, Dict[str, Any]]:
+        logging.debug(
+            "sample: prompt_tokens=%d, text_tokens=%d, prompt_frames=%d, "
+            "speed=%.2f, guidance_scale=%.2f, seed=%d",
+            prompt_tokens_len,
+            text_tokens_len,
+            prompt_features_len,
+            speed,
+            guidance_scale,
+            seed,
+        )
+        t_total_start = time.perf_counter()
+
+        t_start = time.perf_counter()
+        encoded = self.run_encoder(cat_tokens)
+        t_enc = time.perf_counter() - t_start
+        logging.debug("  encoder: %.3f s (output shape=%s)", t_enc, encoded.shape)
+
+        t_start = time.perf_counter()
+        text_condition, features_len = self.duration_expand(
+            encoded,
+            prompt_tokens_len,
+            text_tokens_len,
+            prompt_features_len,
+            speed,
+        )
+        t_dur = time.perf_counter() - t_start
+        logging.debug("  duration_expand: %.3f s (features_len=%d)", t_dur, features_len)
+
+        seq_len = self.decoder_seq_len or self.max_feat_len
+        if features_len > seq_len:
+            raise ValueError(
+                f"features_len={features_len} exceeds decoder sequence length {seq_len}"
+            )
+        if (
+            self.decoder_seq_len is not None
+            and not self.decoder_has_padding_mask
+            and features_len != seq_len
+        ):
+            raise ValueError(
+                "Fixed no-mask decoder requires exact feature length: "
+                f"features_len={features_len}, decoder_seq_len={seq_len}"
+            )
+        if prompt_features.shape[1] > seq_len:
+            raise ValueError(
+                f"prompt feature length {prompt_features.shape[1]} exceeds "
+                f"decoder sequence length {seq_len}"
+            )
+
+        text_cond_padded = np.zeros((1, seq_len, self.feat_dim), dtype=np.float32)
+        text_cond_padded[0, :features_len] = text_condition[0, :features_len]
+
+        speech_cond_padded = np.zeros((1, seq_len, self.feat_dim), dtype=np.float32)
+        prompt_actual_len = prompt_features.shape[1]
+        speech_cond_padded[0, :prompt_actual_len] = prompt_features[0].astype(np.float32)
+
+        padding_mask = np.zeros((1, seq_len), dtype=np.bool_)
+        padding_mask[:, features_len:] = True
+
+        rng = np.random.RandomState(seed)
+        x = rng.randn(1, seq_len, self.feat_dim).astype(np.float32)
+        x[:, features_len:, :] = 0.0
+
+        timesteps = self._get_time_steps()
+        gs = np.array([guidance_scale], dtype=np.float32)
+
+        t_dec_total = 0.0
+        for step in range(self.num_step):
+            t_val = np.array([float(timesteps[step])], dtype=np.float32)
+            t_start = time.perf_counter()
+            v = self.run_decoder(
+                t_val,
+                x,
+                text_cond_padded,
+                speech_cond_padded,
+                gs,
+                padding_mask,
+            )
+            t_dec_total += time.perf_counter() - t_start
+            dt = float(timesteps[step + 1] - timesteps[step])
+            x = (x + v * dt).astype(np.float32)
+            x[:, features_len:, :] = 0.0
+
+        logging.debug(
+            "  %s (NPU x%d): %.3f s total (avg %.3f ms/step)",
+            getattr(self, "decoder_label", "decoder4"),
+            self.num_step,
+            t_dec_total,
+            t_dec_total / self.num_step * 1000,
+        )
+
+        generated_frames = features_len - prompt_features_len
+        if generated_frames <= 0:
+            generated_frames = features_len
+            pred_features = x[0, :features_len, :]
+        else:
+            pred_features = x[0, prompt_features_len:features_len, :]
+
+        t_total = time.perf_counter() - t_total_start
+        timing = {
+            "encoder_time_sec": round(t_enc, 3),
+            "duration_expand_time_sec": round(t_dur, 3),
+            "decoder_time_sec": round(t_dec_total, 3),
+            "total_time_sec": round(t_total, 3),
+            "generated_frames": int(generated_frames),
+            "features_len": int(features_len),
+        }
+        logging.debug("  total: %.3f s", t_total)
+        return pred_features[np.newaxis, :, :].astype(np.float32), timing

scripts/zipvoice_runtime.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+
+from contextlib import contextmanager
+
+
+@contextmanager
+def _suppress_native_stdio():
+    if os.environ.get("ZIPVOICE_AXERA_VERBOSE"):
+        yield
+        return
+
+    stdout_fd = os.dup(1)
+    stderr_fd = os.dup(2)
+    try:
+        with open(os.devnull, "w") as devnull:
+            os.dup2(devnull.fileno(), 1)
+            os.dup2(devnull.fileno(), 2)
+            yield
+    finally:
+        os.dup2(stdout_fd, 1)
+        os.dup2(stderr_fd, 2)
+        os.close(stdout_fd)
+        os.close(stderr_fd)
+
+try:
+    with _suppress_native_stdio():
+        import axengine
+    HAS_AXENGINE = True
+except ImportError:
+    HAS_AXENGINE = False
+
+
+class AxeSession:
+    """Thin wrapper around axengine.InferenceSession for one axmodel."""
+
+    def __init__(self, model_path: str | Path):
+        if not HAS_AXENGINE:
+            raise RuntimeError("axengine not available; run this script on the AX650 board")
+        self.path = Path(model_path)
+        if not self.path.exists():
+            raise FileNotFoundError(f"Model not found: {self.path}")
+        with _suppress_native_stdio():
+            self._session = axengine.InferenceSession(str(self.path))
+        self._inputs = self._session.get_inputs()
+        self._outputs = self._session.get_outputs()
+
+    @property
+    def input_names(self) -> List[str]:
+        return [item.name for item in self._inputs]
+
+    @property
+    def output_names(self) -> List[str]:
+        return [item.name for item in self._outputs]
+
+    def run(self, feed_dict: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        outputs = self._session.run(None, feed_dict)
+        if isinstance(outputs, dict):
+            return outputs
+        if isinstance(outputs, (list, tuple)):
+            return {name: value for name, value in zip(self.output_names, outputs)}
+        if len(self.output_names) == 1:
+            return {self.output_names[0]: outputs}
+        raise RuntimeError(f"Unexpected AXEngine output type from {self.path.name}: {type(outputs)}")