update readme

Made-with: Cursor

.gitignore

@@ -0,0 +1,4 @@
+launch.json
+__pycache__
+voxcpm.egg-info
+.DS_Store

README.md

@@ -1,183 +1,282 @@
-## 🎙️ VoxCPM: Tokenizer-Free TTS for Context-Aware Speech Generation and True-to-Life Voice Cloning
-
-
-[![Project Page](https://img.shields.io/badge/Project%20Page-GitHub-blue)](https://github.com/OpenBMB/VoxCPM/) [![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-OpenBMB-yellow)](https://huggingface.co/openbmb/VoxCPM-0.5B) [![Live Playground](https://img.shields.io/badge/Live%20PlayGround-Demo-orange)](https://huggingface.co/spaces/OpenBMB/VoxCPM-Demo) [![Samples](https://img.shields.io/badge/Page-Samples-red)](https://thuhcsi.github.io/VoxCPM/)
+<h2 align="center">VoxCPM2: Tokenizer-Free TTS for Multilingual Speech Generation, Creative Voice Design, and True-to-Life Cloning</h2>
 
+<p align="center">
+  <a href="https://github.com/OpenBMB/VoxCPM/"><img src="https://img.shields.io/badge/Project%20Page-GitHub-blue" alt="Project Page"></a>
+  <a href="https://huggingface.co/spaces/OpenBMB/VoxCPM-Demo"><img src="https://img.shields.io/badge/Live%20Playground-Demo-orange" alt="Live Playground"></a>
+  <a href="https://voxcpm.readthedocs.io/en/latest/"><img src="https://img.shields.io/badge/Docs-ReadTheDocs-8CA1AF" alt="Documentation"></a>
+  <a href="https://huggingface.co/openbmb/VoxCPM2"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-VoxCPM2-yellow" alt="Hugging Face"></a>
+  <a href="https://modelscope.cn/models/OpenBMB/VoxCPM2"><img src="https://img.shields.io/badge/ModelScope-VoxCPM2-purple" alt="ModelScope"></a>
+  <a href="https://openbmb.github.io/voxcpm2-demopage/"><img src="https://img.shields.io/badge/DemoPage-Audio Samples-red"></a>
+  
+</p>
 
 <div align="center">
-  <img src="assets/voxcpm_logo.png" alt="VoxCPM Logo" width="40%">
+  <img src="assets/voxcpm_logo.png" alt="VoxCPM Logo" width="35%">
+  <br><br>
+  <a href="https://trendshift.io/repositories/17704" target="_blank"><img src="https://trendshift.io/api/badge/repositories/17704" alt="OpenBMB%2FVoxCPM | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
 </div>
 
-## News 
-* [2025.09.16] 🔥 🔥 🔥  We Open Source the VoxCPM-0.5B [weights](https://huggingface.co/openbmb/VoxCPM-0.5B)!
-* [2025.09.16] 🎉 🎉 🎉  We Provide the [Gradio PlayGround](https://huggingface.co/spaces/OpenBMB/VoxCPM-Demo) for VoxCPM-0.5B, try it now! 
+<br>
 
-## Overview
+<p align="center">
+  👋 Join our community for discussion and support!
+  <br>
+  <a href="./assets/feishu-group.png" style="display:inline-block;vertical-align:middle; margin-left: 10px;">
+    <img src="./assets/feishu-logo.png" width="16" height="16" style="vertical-align:middle;"> Feishu
+  </a>
+  &nbsp;|&nbsp;
+  <a href="https://discord.gg/KZUx7tVNwz" style="display:inline-block;vertical-align:middle;">
+    <img src="./assets/discord-logo.png" width="16" height="16" style="vertical-align:middle;"> Discord
+  </a>
+</p>
 
-VoxCPM is a novel tokenizer-free Text-to-Speech (TTS) system that redefines realism in speech synthesis. By modeling speech in a continuous space, it overcomes the limitations of discrete tokenization and enables two flagship capabilities: context-aware speech generation and true-to-life zero-shot voice cloning.
+VoxCPM is a **tokenizer-free** Text-to-Speech system that directly generates continuous speech representations via an end-to-end **diffusion autoregressive architecture**, bypassing discrete tokenization to achieve highly natural and expressive synthesis.
 
-Unlike mainstream approaches that convert speech to discrete tokens, VoxCPM uses an end-to-end diffusion autoregressive architecture that directly generates continuous speech representations from text. Built on [MiniCPM-4](https://huggingface.co/openbmb/MiniCPM4-0.5B) backbone, it achieves implicit semantic-acoustic decoupling through hierachical language modeling and FSQ constraints, greatly enhancing both expressiveness and generation stability.
+**VoxCPM2** is the latest major release — a **2B** parameter model trained on **over 2 million hours** of multilingual speech data, now supporting **30 languages**, **Voice Design**, **Controllable Voice Cloning**, and **48kHz** studio-quality audio output. Built on a [MiniCPM-4](https://github.com/OpenBMB/MiniCPM) backbone.
 
-<div align="center">
-  <img src="assets/voxcpm_model.png" alt="VoxCPM Model Architecture" width="90%">
-</div>
+### ✨ Highlights
 
+- 🌍 **30-Language Multilingual** — Input text in any of the 30 supported languages and synthesize directly, no language tag needed
+- 🎨 **Voice Design** — Create a brand-new voice from a natural-language description alone (gender, age, tone, emotion, pace …), no reference audio required
+- 🎛️ **Controllable Cloning** — Clone any voice from a short reference clip, with optional style guidance to steer emotion, pace, and expression while preserving the original timbre
+- 🎙️ **Ultimate Cloning** — Reproduce every vocal nuance: provide both reference audio and its transcript, and the model continues seamlessly from the reference, faithfully preserving every vocal detail — timbre, rhythm, emotion, and style (same as VoxCPM1.5)
+- 🔊 **48kHz High-Quality Audio** — Accepts 16kHz reference audio and directly outputs 48kHz studio-quality audio via AudioVAE V2's asymmetric encode/decode design, with built-in super-resolution — no external upsampler needed
+- 🧠 **Context-Aware Synthesis** — Automatically infers appropriate prosody and expressiveness from text content
+- ⚡ **Real-Time Streaming** — RTF as low as ~0.3 on NVIDIA RTX 4090, and ~0.13  accelerated by [Nano-VLLM](https://github.com/a710128/nanovllm-voxcpm)
+- 📜 **Fully Open-Source & Commercial-Ready** — Weights and code released under the [Apache-2.0](LICENSE) license, free for commercial use
 
-###  🚀 Key Features
-- **Context-Aware, Expressive Speech Generation** - VoxCPM comprehends text to infer and generate appropriate prosody, delivering speech with remarkable expressiveness and natural flow. It spontaneously adapts speaking style based on content, producing highly fitting vocal expression trained on a massive 1.8 million-hour bilingual corpus.
-- **True-to-Life Voice Cloning** - With only a short reference audio clip, VoxCPM performs accurate zero-shot voice cloning, capturing not only the speaker’s timbre but also fine-grained characteristics such as accent, emotional tone, rhythm, and pacing to create a faithful and natural replica.
-- **High-Efficiency Synthesis** - VoxCPM supports streaming synthesis with a Real-Time Factor (RTF) as low as 0.17 on a consumer-grade NVIDIA RTX 4090 GPU, making it possible for real-time applications.
+<details>
+<summary><b>🌍 Supported Languages (30)</b></summary>
+<br>
+Arabic, Burmese, Chinese, Danish, Dutch, English, Finnish, French, German, Greek, Hebrew, Hindi, Indonesian, Italian, Japanese, Khmer, Korean, Lao, Malay, Norwegian, Polish, Portuguese, Russian, Spanish, Swahili, Swedish, Tagalog, Thai, Turkish, Vietnamese
 
+Chinese Dialect: 四川话, 粤语, 吴语, 东北话, 河南话, 陕西话, 山东话, 天津话, 闽南话
+</details>
 
+### News
 
+* **[2026.04]** 🔥 We release **VoxCPM2** — 2B, 30 languages, Voice Design & Controllable Voice Cloning, 48kHz audio output! [Weights](https://huggingface.co/openbmb/VoxCPM2) | [Docs](https://voxcpm.readthedocs.io/en/latest/) | [Playground](https://huggingface.co/spaces/OpenBMB/VoxCPM-Demo)
+* **[2025.12]** 🎉 Open-source **VoxCPM1.5** [weights](https://huggingface.co/openbmb/VoxCPM1.5) with SFT & LoRA fine-tuning. (**🏆 #1 GitHub Trending**)
+* **[2025.09]** 🔥 Release VoxCPM [Technical Report](https://arxiv.org/abs/2509.24650).
+* **[2025.09]** 🎉 Open-source **VoxCPM-0.5B** [weights](https://huggingface.co/openbmb/VoxCPM-0.5B) (**🏆 #1 HuggingFace Trending**)
 
+---
 
+## Contents
 
+- [Quick Start](#-quick-start)
+  - [Installation](#installation)
+  - [Python API](#python-api)
+  - [CLI Usage](#cli-usage)
+  - [Web Demo](#web-demo)
+  - [Production Deployment](#-production-deployment-nano-vllm)
+- [Models & Versions](#-models--versions)
+- [Performance](#-performance)
+- [Fine-tuning](#%EF%B8%8F-fine-tuning)
+- [Documentation](#-documentation)
+- [Ecosystem & Community](#-ecosystem--community)
+- [Risks and Limitations](#%EF%B8%8F-risks-and-limitations)
+- [Citation](#-citation)
 
+---
 
+## 🚀 Quick Start
 
+### Installation
 
-##  Quick Start
-
-### 🔧 Install from PyPI
-``` sh
+```sh
 pip install voxcpm
 ```
-### 1.  Model Download (Optional)
-By default, when you first run the script, the model will be downloaded automatically, but you can also download the model in advance.
-- Download VoxCPM-0.5B
-    ```
-    from huggingface_hub import snapshot_download
-    snapshot_download("openbmb/VoxCPM-0.5B",local_files_only=local_files_only)
-    ```
-- Download ZipEnhancer and SenseVoice-Small. We use ZipEnhancer to enhance speech prompts and SenseVoice-Small for speech prompt ASR in the web demo.
-    ```
-    from modelscope import snapshot_download
-    snapshot_download('iic/speech_zipenhancer_ans_multiloss_16k_base')
-    snapshot_download('iic/SenseVoiceSmall')
-    ```
 
-### 2. Basic Usage
+> **Requirements:** Python ≥ 3.10, PyTorch ≥ 2.5.0, CUDA ≥ 12.0. See [Quick Start Docs](https://voxcpm.readthedocs.io/en/latest/quickstart.html) for details.
+
+### Python API
+
+#### 🗣️ Text-to-Speech
+
 ```python
-import soundfile as sf
 from voxcpm import VoxCPM
+import soundfile as sf
 
-model = VoxCPM.from_pretrained("openbmb/VoxCPM-0.5B")
-
-wav = model.generate(
-    text="VoxCPM is an innovative end-to-end TTS model from ModelBest, designed to generate highly expressive speech.",
-    prompt_wav_path=None,      # optional: path to a prompt speech for voice cloning
-    prompt_text=None,          # optional: reference text
-    cfg_value=2.0,             # LM guidance on LocDiT, higher for better adherence to the prompt, but maybe worse
-    inference_timesteps=10,   # LocDiT inference timesteps, higher for better result, lower for fast speed
-    normalize=True,           # enable external TN tool
-    denoise=True,             # enable external Denoise tool
-    retry_badcase=True,        # enable retrying mode for some bad cases (unstoppable)
-    retry_badcase_max_times=3,  # maximum retrying times
-    retry_badcase_ratio_threshold=6.0, # maximum length restriction for bad case detection (simple but effective), it could be adjusted for slow pace speech
+model = VoxCPM.from_pretrained(
+  "openbmb/VoxCPM2"
+  load_denoiser=False,
 )
 
-sf.write("output.wav", wav, 16000)
-print("saved: output.wav")
+wav = model.generate(
+    text="VoxCPM2 is the current recommended release for realistic multilingual speech synthesis.",
+    cfg_value=2.0,
+    inference_timesteps=10,
+)
+sf.write("demo.wav", wav, model.tts_model.sample_rate)
+print("saved: demo.wav")
 ```
 
-### 3. CLI Usage
+#### 🎨 Voice Design
 
-After installation, the entry point is `voxcpm` (or use `python -m voxcpm.cli`).
+Create a voice from a natural-language description — no reference audio needed. **Format:** put the description in parentheses at the start of `text`(e.g. `"(your voice description)The text to synthesize."`):
+
+```python
+wav = model.generate(
+    text="(A young woman, gentle and sweet voice)Hello, welcome to VoxCPM2!",
+    cfg_value=2.0,
+    inference_timesteps=10,
+)
+sf.write("voice_design.wav", wav, model.tts_model.sample_rate)
+```
+
+#### 🎛️ Controllable Voice Cloning
+
+Upload a reference audio. The model clones the timbre, and you can still use control instructions to adjust speed, emotion, or style.
+
+```python
+wav = model.generate(
+    text="This is a cloned voice generated by VoxCPM2.",
+    reference_wav_path="speaker.wav",
+)
+sf.write("clone.wav", wav, model.tts_model.sample_rate)
+
+wav = model.generate(
+    text="(slightly faster, cheerful tone)This is a cloned voice with style control.",
+    reference_wav_path="speaker.wav",
+    cfg_value=2.0,
+    inference_timesteps=10,
+)
+sf.write("controllable_clone.wav", wav, model.tts_model.sample_rate)
+```
+
+#### 🎙️ Ultimate Cloning
+
+Provide both the reference audio and its exact transcript for audio-continuation-based cloning with every vocal nuance reproduced. For maximum cloning similarity, pass the same reference clip to both `reference_wav_path` and `prompt_wav_path` as shown below:
+
+```python
+wav = model.generate(
+    text="This is an ultimate cloning demonstration using VoxCPM2.",
+    prompt_wav_path="speaker_reference.wav",
+    prompt_text="The transcript of the reference audio.",
+    reference_wav_path="speaker_reference.wav",
+)
+sf.write("hifi_clone.wav", wav, model.tts_model.sample_rate)
+```
+
+<details>
+<summary><b>🔄 Streaming API</b></summary>
+
+```python
+import numpy as np
+
+chunks = []
+for chunk in model.generate_streaming(
+    text="Streaming text to speech is easy with VoxCPM!",
+):
+    chunks.append(chunk)
+wav = np.concatenate(chunks)
+sf.write("streaming.wav", wav, model.tts_model.sample_rate)
+```
+</details>
+
+### CLI Usage
 
 ```bash
-# 1) Direct synthesis (single text)
-voxcpm --text "Hello VoxCPM" --output out.wav
+# Voice design (no reference audio needed)
+voxcpm design \
+  --text "VoxCPM2 brings studio-quality multilingual speech synthesis." \
+  --output out.wav
 
-# 2) Voice cloning (reference audio + transcript)
-voxcpm --text "Hello" \
+# Controllable voice cloning with style control
+voxcpm design \
+  --text "VoxCPM2 brings studio-quality multilingual speech synthesis." \
+  --control "Young female voice, warm and gentle, slightly smiling" \
+  --output out.wav
+
+# Voice cloning (reference audio)
+voxcpm clone \
+  --text "This is a voice cloning demo." \
+  --reference-audio path/to/voice.wav \
+  --output out.wav
+
+# Ultimate cloning (prompt audio + transcript)
+voxcpm clone \
+  --text "This is a voice cloning demo." \
   --prompt-audio path/to/voice.wav \
   --prompt-text "reference transcript" \
-  --output out.wav \
-  --denoise
+  --output out.wav
 
-# 3) Batch processing (one text per line)
-voxcpm --input examples/input.txt --output-dir outs
-# (optional) Batch + cloning
-voxcpm --input examples/input.txt --output-dir outs \
-  --prompt-audio path/to/voice.wav \
-  --prompt-text "reference transcript" \
-  --denoise
+# Batch processing
+voxcpm batch --input examples/input.txt --output-dir outs
 
-# 4) Inference parameters (quality/speed)
-voxcpm --text "..." --output out.wav \
-  --cfg-value 2.0 --inference-timesteps 10 --normalize
-
-# 5) Model loading
-# Prefer local path
-voxcpm --text "..." --output out.wav --model-path /path/to/VoxCPM_model_dir
-# Or from Hugging Face (auto download/cache)
-voxcpm --text "..." --output out.wav \
-  --hf-model-id openbmb/VoxCPM-0.5B --cache-dir ~/.cache/huggingface --local-files-only
-
-# 6) Denoiser control
-voxcpm --text "..." --output out.wav \
-  --no-denoiser --zipenhancer-path iic/speech_zipenhancer_ans_multiloss_16k_base
-
-# 7) Help
+# Help
 voxcpm --help
-python -m voxcpm.cli --help
 ```
 
-### 4. Start web demo
+### Web Demo
 
-You can start the UI interface by running `python app.py`, which allows you to perform Voice Cloning and Voice Creation.
+```bash
+python app.py   # then open http://localhost:7860
+```
 
+### 🚢 Production Deployment (Nano-vLLM)
 
+For high-throughput serving, use [**Nano-vLLM-VoxCPM**](https://github.com/a710128/nanovllm-voxcpm) — a dedicated inference engine built on Nano-vLLM with concurrent request support and an async API.
 
-## 👩‍🍳 A Voice Chef's Guide
-Welcome to the VoxCPM kitchen! Follow this recipe to cook up perfect generated speech. Let’s begin.
+```bash
+pip install nano-vllm-voxcpm
+```
 
----
-### 🥚 Step 1: Prepare Your Base Ingredients (Content)
+```python
+from nanovllm_voxcpm import VoxCPM
+import numpy as np, soundfile as sf
 
-First, choose how you’d like to input your text:.
-1. Regular Text (Classic Mode)
-- ✅ Keep "Text Normalization" ON. Type naturally (e.g., "Hello, world! 123"). The system will automatically process numbers, abbreviations, and punctuation using WeTextProcessing library.
-2. Phoneme Input (Native Mode)
-- ❌ Turn "Text Normalization" OFF. Enter phoneme text like {HH AH0 L OW1} (EN) or {ni3}{hao3} (ZH) for precise pronunciation  control. In this mode, VoxCPM also supports native understanding of other complex non-normalized text—try it out!
+server = VoxCPM.from_pretrained(model="/path/to/VoxCPM", devices=[0])
+chunks = list(server.generate(target_text="Hello from VoxCPM!"))
+sf.write("out.wav", np.concatenate(chunks), 48000)
+server.stop()
+```
 
----
-### 🍳 Step 2: Choose Your Flavor Profile (Voice Style) 
-
-This is the secret sauce that gives your audio its unique sound.
-1. Cooking with a Prompt Speech (Following a Famous Recipe)
-  - A prompt speech provides the desired acoustic characteristics for VoxCPM. The speaker's timbre, speaking style, and even the background sounds and ambiance will be replicated.
-  - For a Clean, Studio-Quality Voice:
-    - ✅ Enable "Prompt Speech Enhancement". This acts like a noise filter, removing background hiss and rumble to give you a pure, clean voice clone.
-2. Cooking au Naturel (Letting the Model Improvise)
-  - If no reference is provided, VoxCPM becomes a creative chef! It will infer a fitting speaking style based on the text itself, thanks to the text-smartness of its foundation model, MiniCPM-4.
-  - Pro Tip: Challenge VoxCPM with any text—poetry, song lyrics, dramatic monologues—it may deliver some interesting results!
-
----
-### 🧂 Step 3: The Final Seasoning (Fine-Tuning Your Results)
-You're ready to serve! But for master chefs who want to tweak the flavor, here are two key spices.
-- CFG Value (How Closely to Follow the Recipe)
-  - Default: A great starting point.
-  - Voice sounds strained or weird? Lower this value. It tells the model to be more relaxed and improvisational, great for expressive prompts.
-  - Need maximum clarity and adherence to the text? Raise it slightly to keep the model on a tighter leash.
-- Inference Timesteps (Simmering Time: Quality vs. Speed)
-  - Need a quick snack? Use a lower number. Perfect for fast drafts and experiments.
-  - Cooking a gourmet meal? Use a higher number. This lets the model "simmer" longer, refining the audio for superior detail and naturalness.
-
----
-Happy creating! 🎉 Start with the default settings and tweak from there to suit your project. The kitchen is yours!
+> **RTF as low as ~0.13 on NVIDIA RTX 4090** (vs ~0.3 with the standard PyTorch implementation), with support for batched concurrent requests and a FastAPI HTTP server. See the [Nano-vLLM-VoxCPM repo](https://github.com/a710128/nanovllm-voxcpm) for deployment details.
 
+> **Full parameter reference, multi-scenario examples, and voice cloning tips →** [Quick Start Guide](https://voxcpm.readthedocs.io/en/latest/quickstart.html) | [Usage Guide](https://voxcpm.readthedocs.io/en/latest/usage_guide.html) | [Cookbook](https://voxcpm.readthedocs.io/en/latest/cookbook.html)
 
 ---
 
+## 📦 Models & Versions
 
+| | **VoxCPM2** | **VoxCPM1.5** | **VoxCPM-0.5B** |
+|---|:---:|:---:|:---:|
+| **Status** | 🟢 Latest | Stable | Legacy |
+| **Backbone Parameters** | 2B | 0.6B | 0.5B |
+| **Audio Sample Rate** | 48kHz | 44.1kHz | 16kHz |
+| **LM Token Rate** | 6.25Hz | 6.25Hz | 12.5Hz |
+| **Languages** | 30 | 2 (zh, en) | 2 (zh, en) |
+| **Cloning Mode** | Isolated Reference & Continuation | Continuation only | Continuation only |
+| **Voice Design** | ✅ | — | — |
+| **Controllable Voice Cloning** | ✅ | — | — |
+| **SFT / LoRA** | ✅ | ✅ | ✅ |
+| **RTF (RTX 4090)** | ~0.30 | ~0.15 | ~0.17 |
+| **RTF in Nano-VLLM (RTX 4090)** | ~0.13 | ~0.08 | ~0.10 |
+| **VRAM** | ~8 GB | ~6 GB | ~5 GB |
+| **Weights** | [🤗 HF](https://huggingface.co/openbmb/VoxCPM2) / [MS](https://modelscope.cn/models/OpenBMB/VoxCPM2) | [🤗 HF](https://huggingface.co/openbmb/VoxCPM1.5) / [MS](https://modelscope.cn/models/OpenBMB/VoxCPM1.5) | [🤗 HF](https://huggingface.co/openbmb/VoxCPM-0.5B) / [MS](https://modelscope.cn/models/OpenBMB/VoxCPM-0.5B) |
+| **Technical Report** | Coming soon | — | [arXiv](https://arxiv.org/abs/2509.24650) [ICLR 2026](https://openreview.net/forum?id=h5KLpGoqzC) |
+| **Demo Page** | [Audio Samples](https://openbmb.github.io/voxcpm2-demopage) | — | [Audio Samples](https://openbmb.github.io/VoxCPM-demopage) |
 
-## 📊 Performance Highlights
+VoxCPM2 is built on a **tokenizer-free, diffusion autoregressive** paradigm. The model operates entirely in the latent space of **AudioVAE V2**, following a four-stage pipeline: **LocEnc → TSLM → RALM → LocDiT**, enabling rich expressiveness and 48kHz native audio output.
 
-VoxCPM achieves competitive results on public zero-shot TTS benchmarks:
+<div align="center">
+  <img src="assets/voxcpm_model.png" alt="VoxCPM2 Model Architecture" width="90%">
+</div>
 
-### Seed-TTS-eval Benchmark
+> For full architectural details, VoxCPM2-specific upgrades, and a model comparison table, see the [Architecture Design](https://voxcpm.readthedocs.io/en/latest/models/architecture.html).
+
+---
+
+## 📊 Performance
+
+VoxCPM2 achieves state-of-the-art or comparable results on public zero-shot and controllable TTS benchmarks. 
+
+### Seed-TTS-eval
+
+<details>
+<summary><b>Seed-TTS-eval WER(⬇)&SIM(⬆) Results (click to expand)</b></summary>
 
 | Model | Parameters | Open-Source | test-EN | | test-ZH | | test-Hard | |
 |------|------|------|:------------:|:--:|:------------:|:--:|:-------------:|:--:|
@@ -188,35 +287,127 @@ VoxCPM achieves competitive results on public zero-shot TTS benchmarks:
 | CosyVoice3 | 1.5B | ❌ | 2.22 | 72.0 | 1.12 | 78.1 | 5.83 | 75.8 |
 | Seed-TTS | - | ❌ | 2.25 | 76.2 | 1.12 | 79.6 | 7.59 | 77.6 |
 | MiniMax-Speech | - | ❌ | 1.65 | 69.2 | 0.83 | 78.3 | - | - |
-| CosyVoice | 0.3B | ✅ | 4.29 | 60.9 | 3.63 | 72.3 | 11.75 | 70.9 |
-| CosyVoice2 | 0.5B | ✅ | 3.09 | 65.9 | 1.38 | 75.7 | **6.83** | 72.4 |
 | F5-TTS | 0.3B | ✅ | 2.00 | 67.0 | 1.53 | 76.0 | 8.67 | 71.3 |
+| MaskGCT | 1B | ✅ | 2.62 | 71.7 | 2.27 | 77.4 | - | - |
+| CosyVoice | 0.3B | ✅ | 4.29 | 60.9 | 3.63 | 72.3 | 11.75 | 70.9 |
+| CosyVoice2 | 0.5B | ✅ | 3.09 | 65.9 | 1.38 | 75.7 | 6.83 | 72.4 |
 | SparkTTS | 0.5B | ✅ | 3.14 | 57.3 | 1.54 | 66.0 | - | - |
 | FireRedTTS | 0.5B | ✅ | 3.82 | 46.0 | 1.51 | 63.5 | 17.45 | 62.1 |
 | FireRedTTS-2 | 1.5B | ✅ | 1.95 | 66.5 | 1.14 | 73.6 | - | - |
-| Qwen2.5-Omni | 7B | ✅ | 2.72 | 63.2 | 1.70 | 75.2 | 7.97 | **74.7** |
-| OpenAudio-s1-mini | 0.5B | ✅ | 1.94 | 55.0 | 1.18 | 68.5 | - | - |
-| IndexTTS2 | 1.5B | ✅ | 2.23 | 70.6 | 1.03 | 76.5 | - | - |
+| Qwen2.5-Omni | 7B | ✅ | 2.72 | 63.2 | 1.70 | 75.2 | 7.97 | 74.7 |
+| Qwen3-Omni | 30B-A3B | ✅ | 1.39 | - | 1.07 | - | - | - |
+| OpenAudio-s1-mini | 0.5B | ✅ | 1.94 | 55.0 | 1.18 | 68.5 | 23.37 | 64.3 |
+| IndexTTS2 | 1.5B | ✅ | 2.23 | 70.6 | 1.03 | 76.5 | 7.12 | 75.5 |
 | VibeVoice | 1.5B | ✅ | 3.04 | 68.9 | 1.16 | 74.4 | - | - |
-| HiggsAudio-v2 | 3B | ✅ | 2.44 | 67.7 | 1.50 | 74.0 | - | - |
-| **VoxCPM** | 0.5B | ✅ | **1.85** | **72.9** | **0.93** | **77.2** | 8.87 | 73.0 |
+| HiggsAudio-v2 | 3B | ✅ | 2.44 | 67.7 | 1.50 | 74.0 | 55.07 | 65.6 |
+| VoxCPM-0.5B | 0.6B | ✅ | 1.85 | 72.9 | 0.93 | 77.2 | 8.87 | 73.0 |
+| VoxCPM1.5 | 0.8B | ✅ | 2.12 | 71.4 | 1.18 | 77.0 | 7.74 | 73.1 |
+| MOSS-TTS |  | ✅ | 1.85 | 73.4 | 1.20 | 78.8 | - | - |
+| Qwen3-TTS | 1.7B | ✅ | 1.23 | 71.7 | 1.22 | 77.0 | 6.76 | 74.8 |
+| FishAudio S2 | 4B | ✅ | 0.99 | - | 0.54 | - | 5.99 | - |
+| LongCat-Audio-DiT | 3.5B | ✅ | 1.50 | 78.6 | 1.09 | 81.8 | 6.04 | 79.7 |
+| **VoxCPM2** | 2B | ✅ | 1.84 | 75.3 | 0.97| 79.5| 8.13 | 75.3 |  
+</details>
 
 
-###  CV3-eval Benchmark
+### CV3-eval 
+<details>
+<summary><b>CV3-eval Multilingual WER/CER(⬇) Results (click to expand)</b></summary>
 
-| Model | zh | en | hard-zh | | | hard-en | | |
-|-------|:--:|:--:|:-------:|:--:|:--:|:-------:|:--:|:--:|
-| | CER/%⬇ | WER/%⬇ | CER/%⬇ | SIM/%⬆ | DNSMOS⬆ | WER/%⬇ | SIM/%⬆ | DNSMOS⬆ |
-| F5-TTS | 5.47 | 8.90 | - | - | - | - | - | - |
-| SparkTTS | 5.15 | 11.0 | - | - | - | - | - | - |
-| GPT-SoVits | 7.34 | 12.5 | - | - | - | - | - | - |
-| CosyVoice2 | 4.08 | 6.32 | 12.58 | 72.6 | 3.81 | 11.96 | 66.7 | 3.95 |
-| OpenAudio-s1-mini | 4.00 | 5.54 | 18.1 | 58.2 | 3.77 | 12.4 | 55.7 | 3.89 |
-| IndexTTS2 | 3.58 | 4.45 | 12.8 | 74.6 | 3.65 | - | - | - |
-| HiggsAudio-v2 | 9.54 | 7.89 | 41.0 | 60.2 | 3.39 | 10.3 | 61.8 | 3.68 |
-| CosyVoice3-0.5B | 3.89 | 5.24 | 14.15 | 78.6 | 3.75 | 9.04 | 75.9 | 3.92 |
-| CosyVoice3-1.5B | 3.91 | 4.99 | 9.77 | 78.5 | 3.79 | 10.55 | 76.1 | 3.95 |
-| **VoxCPM** | **3.40** | **4.04** | 12.9 | 66.1 | 3.59 | **7.89** | 64.3 | 3.74 |
+| Model | zh | en | hard-zh | hard-en | ja | ko | de | es | fr | it | ru |
+|-------|:--:|:--:|:--:|:--:|:--:|:--:|:--:|:--:|:--:|:--:|:--:|
+| CosyVoice2 | 4.08 | 6.32 | 12.58| 11.96| 9.13 | 19.7 |- | - | - | - | - |
+| CosyVoice3-1.5B | 3.91 | 4.99 | 9.77 | 10.55 | 7.57 | 5.69 | 6.43 | 4.47 | 11.8 | 10.5 | 6.64 |
+| Fish Audio S2 | 2.65 | 2.43 | 9.10 | 4.40 | 3.96 | 2.76 | 2.22 | 2.00 | 6.26 | 2.04 | 2.78 |
+| **VoxCPM2** | 3.65 | 5.00 | 8.55 | 8.48 | 5.96 | 5.69 | 4.77 | 3.80 | 9.85 | 4.25 | 5.21 |
+</details>
+
+### MiniMax-Multilingual-Test
+
+<details>
+<summary><b>Minimax-MLS-test WER(⬇) Results (click to expand)</b></summary>
+
+| Language | Minimax | ElevenLabs | Qwen3-TTS | FishAudio S2 | **VoxCPM2** |
+|----------|:-------:|:----------:|:--------------------:|:------------:|:-----------:|
+| Arabic | **1.665** | 1.666 | – | 3.500 | 13.046 |
+| Cantonese | 34.111 | 51.513 | – | **30.670** | 38.584 |
+| Chinese | 2.252 | 16.026 | 0.928 | **0.730** | 1.136 |
+| Czech | 3.875 | **2.108** | – | 2.840 | 24.132 |
+| Dutch | 1.143 | **0.803** | – | 0.990 | 0.913 |
+| English | 2.164 | 2.339 | **0.934** | 1.620 | 2.289 |
+| Finnish | 4.666 | 2.964 | – | 3.330 | **2.632** |
+| French | 4.099 | 5.216 | **2.858** | 3.050 | 4.534 |
+| German | 1.906 | 0.572 | 1.235 | **0.550** | 0.679 |
+| Greek | 2.016 | **0.991** | – | 5.740 | 2.844 |
+| Hindi | 6.962 | **5.827** | – | 14.640 | 19.699 |
+| Indonesian | 1.237 | **1.059** | – | 1.460 | 1.084 |
+| Italian | 1.543 | 1.743 | **0.948** | 1.270 | 1.563 |
+| Japanese | 3.519 | 10.646 | 3.823 | **2.760** | 4.628 |
+| Korean | 1.747 | 1.865 | 1.755 | **1.180** | 1.962 |
+| Polish | 1.415 | **0.766** | – | 1.260 | 1.141 |
+| Portuguese | 1.877 | 1.331 | 1.526 | **1.140** | 1.938 |
+| Romanian | 2.878 | **1.347** | – | 10.740 | 21.577 |
+| Russian | 4.281 | 3.878 | 3.212 | **2.400** | 3.634 |
+| Spanish | 1.029 | 1.084 | 1.126 | **0.910** | 1.438 |
+| Thai | 2.701 | 73.936 | – | 4.230 | 2.961 |
+| Turkish | 1.52 | 0.699 | – | 0.870 | 0.817 |
+| Ukrainian | 1.082 | **0.997** | – | 2.300 | 6.316 |
+| Vietnamese | **0.88** | 73.415 | – | 7.410 | 3.307 |
+
+</details>
+
+<details>
+<summary><b>Minimax-MLS-test SIM(⬆) Results (click to expand)</b></summary>
+
+| Language | Minimax | ElevenLabs | Qwen3-TTS | FishAudio S2 | **VoxCPM2** |
+|----------|:-------:|:----------:|:--------------------:|:------------:|:-----------:|
+| Arabic | 73.6 | 70.6 | – | 75.0 | **79.1** |
+| Cantonese | 77.8 | 67.0 | – | 80.5 | **83.5** |
+| Chinese | 78.0 | 67.7 | 79.9 | 81.6 | **82.5** |
+| Czech | 79.6 | 68.5 | – | **79.8** | 78.3 |
+| Dutch | 73.8 | 68.0 | – | 73.0 | **80.8** |
+| English | 75.6 | 61.3 | 77.5 | 79.7 | **85.4** |
+| Finnish | 83.5 | 75.9 | – | 81.9 | **89.0** |
+| French | 62.8 | 53.5 | 62.8 | 69.8 | **73.5** |
+| German | 73.3 | 61.4 | 77.5 | 76.7 | **80.3** |
+| Greek | 82.6 | 73.3 | – | 79.5 | **86.0** |
+| Hindi | 81.8 | 73.0 | – | 82.1 | **85.6** |
+| Indonesian | 72.9 | 66.0 | – | 76.3 | **80.0** |
+| Italian | 69.9 | 57.9 | 81.7 | 74.7 | **78.0** |
+| Japanese | 77.6 | 73.8 | 78.8 | 79.6 | **82.8** |
+| Korean | 77.6 | 70.0 | 79.9 | 81.7 | **83.3** |
+| Polish | 80.2 | 72.9 | – | 81.9 | **88.4** |
+| Portuguese | 80.5 | 71.1 | 81.7 | 78.1 | **83.7** |
+| Romanian | **80.9** | 69.9 | – | 73.3 | 79.7 |
+| Russian | 76.1 | 67.6 | 79.2 | 79.0 | **81.1** |
+| Spanish | 76.2 | 61.5 | 81.4 | 77.6 | **83.1** |
+| Thai | 80.0 | 58.8 | – | 78.6 | **84.0** |
+| Turkish | 77.9 | 59.6 | – | 83.5 | **87.1** |
+| Ukrainian | 73.0 | 64.7 | – | 74.7 | **79.8** |
+| Vietnamese | 74.3 | 36.9 | – | 74.0 | **80.6** |
+
+</details>
+
+### InstructTTSEval
+
+<details>
+<summary><b>Instruction-Guided Voice Design Results</b></summary>
+
+| Model | InstructTTSEval-ZH | | | InstructTTSEval-EN | | |
+|-------|:---:|:----:|:----:|:----:|:----:|:----:|
+| | APS⬆| DSD⬆ | RP⬆| APS⬆ | DSD⬆ | RP⬆ |
+| Hume | – | – | – | 83.0 | 75.3 | 54.3 |
+| VoxInstruct | 47.5 | 52.3 | 42.6 | 54.9 | 57.0 | 39.3 |
+| Parler-tts-mini | – | – | – | 63.4 | 48.7 | 28.6 |
+| Parler-tts-large | – | – | – | 60.0 | 45.9 | 31.2 |
+| PromptTTS | – | – | – | 64.3 | 47.2 | 31.4 |
+| PromptStyle | – | – | – | 57.4 | 46.4 | 30.9 |
+| VoiceSculptor | 75.7 | 64.7 | 61.5 | – | – | – |
+| Mimo-Audio-7B-Instruct | 75.7 | 74.3 | 61.5 | 80.6 | 77.6 | 59.5 |
+| Qwen3TTS-12Hz-1.7B-VD | **85.2** | **81.1** | **65.1** | 82.9 | 82.4 | 68.4 |
+| **VoxCPM2** | **85.2** | 71.5 | 60.8 | **84.2** | **83.2** | **71.4** |
+
+</details>
 
 
 
@@ -224,52 +415,112 @@ VoxCPM achieves competitive results on public zero-shot TTS benchmarks:
 
 
 
+---
 
+## ⚙️ Fine-tuning
 
+VoxCPM supports both **full fine-tuning (SFT)** and **LoRA fine-tuning**. With as little as **5–10 minutes** of audio, you can adapt to a specific speaker, language, or domain.
 
+```bash
+# LoRA fine-tuning (parameter-efficient, recommended)
+python scripts/train_voxcpm_finetune.py \
+    --config_path conf/voxcpm_v2/voxcpm_finetune_lora.yaml
 
+# Full fine-tuning
+python scripts/train_voxcpm_finetune.py \
+    --config_path conf/voxcpm_v2/voxcpm_finetune_all.yaml
 
-## ⚠️ Risks and limitations
-- General Model Behavior: While VoxCPM has been trained on a large-scale dataset, it may still produce outputs that are unexpected, biased, or contain artifacts.
-- Potential for Misuse of Voice Cloning: VoxCPM's powerful zero-shot voice cloning capability can generate highly realistic synthetic speech. This technology could be misused for creating convincing deepfakes for purposes of impersonation, fraud, or spreading disinformation. Users of this model must not use it to create content that infringes upon the rights of individuals. It is strictly forbidden to use VoxCPM for any illegal or unethical purposes. We strongly recommend that any publicly shared content generated with this model be clearly marked as AI-generated.
-- Current Technical Limitations: Although generally stable, the model may occasionally exhibit instability, especially with very long or expressive inputs. Furthermore, the current version offers limited direct control over specific speech attributes like emotion or speaking style.
-- Bilingual Model: VoxCPM is trained primarily on Chinese and English data. Performance on other languages is not guaranteed and may result in unpredictable or low-quality audio.
-- This model is released for research and development purposes only. We do not recommend its use in production or commercial applications without rigorous testing and safety evaluations. Please use VoxCPM responsibly.
+# WebUI for training & inference
+python lora_ft_webui.py   # then open http://localhost:7860
+```
 
+> **Full guide →** [Fine-tuning Guide](https://voxcpm.readthedocs.io/en/latest/finetuning/finetune.html) (data preparation, configuration, training, LoRA hot-swapping, FAQ)
 
+---
+
+## 📚 Documentation
+
+Full documentation: **[voxcpm.readthedocs.io](https://voxcpm.readthedocs.io/en/latest/)**
+
+| Topic | Link |
+|---|---|
+| Quick Start & Installation | [Quick Start](https://voxcpm.readthedocs.io/en/latest/quickstart.html) |
+| Usage Guide & Cookbook | [User Guide](https://voxcpm.readthedocs.io/en/latest/usage_guide.html) |
+| VoxCPM Series | [Models](https://voxcpm.readthedocs.io/en/latest/models/version_history.html) |
+| Fine-tuning (SFT & LoRA) | [Fine-tuning Guide](https://voxcpm.readthedocs.io/en/latest/finetuning/finetune.html) |
+| FAQ & Troubleshooting | [FAQ](https://voxcpm.readthedocs.io/en/latest/faq.html) |
+
+---
+
+## 🌟 Ecosystem & Community
+
+| Project | Description |
+|---|---|
+| [**Nano-vLLM**](https://github.com/a710128/nanovllm-voxcpm) | High-throughput and Fast GPU serving |
+| [**VoxCPM.cpp**](https://github.com/bluryar/VoxCPM.cpp) | GGML/GGUF: CPU, CUDA, Vulkan inference |
+| [**VoxCPM-ONNX**](https://github.com/bluryar/VoxCPM-ONNX) | ONNX export for CPU inference |
+| [**VoxCPMANE**](https://github.com/0seba/VoxCPMANE) | Apple Neural Engine backend |
+| [**voxcpm_rs**](https://github.com/madushan1000/voxcpm_rs) | Rust re-implementation |
+| [**ComfyUI-VoxCPM**](https://github.com/wildminder/ComfyUI-VoxCPM) | ComfyUI node-based workflows |
+| [**ComfyUI-VoxCPMTTS**](https://github.com/1038lab/ComfyUI-VoxCPMTTS) | ComfyUI TTS extension |
+| [**TTS WebUI**](https://github.com/rsxdalv/tts_webui_extension.vox_cpm) | Browser-based TTS extension |
+
+> See the full [Ecosystem](https://voxcpm.readthedocs.io/en/latest/) in the docs. Community projects are not officially maintained by OpenBMB. Built something cool? [Open an issue or PR](https://github.com/OpenBMB/VoxCPM/issues) to add it!
+
+---
+
+## ⚠️ Risks and Limitations
+
+- **Potential for Misuse:** VoxCPM's voice cloning can generate highly realistic synthetic speech. It is **strictly forbidden** to use VoxCPM for impersonation, fraud, or disinformation. We strongly recommend clearly marking any AI-generated content.
+- **Controllable Generation Stability:** Voice Design and Controllable Voice Cloning results can vary between runs — you may try to generate 1~3 times to obtain the desired voice or style. We are actively working on improving controllability consistency.
+- **Language Coverage:** VoxCPM2 officially supports 30 languages. For languages not on the list, you are welcome to test directly or try fine-tuning on your own data. We plan to expand language coverage in future releases.
+- **Usage:** This model is released under the Apache-2.0 license. For production deployments, we recommend conducting thorough testing and safety evaluation tailored to your use case.
+
+---
+
+## 📖 Citation
+
+If you find VoxCPM helpful, please consider citing our work and starring ⭐ the repository!
+
+```bib
+@article{voxcpm2_2026,
+  title   = {VoxCPM2: Tokenizer-Free TTS for Multilingual Speech Generation, Creative Voice Design, and True-to-Life Cloning},
+  author  = {VoxCPM Team},
+  journal = {GitHub},
+  year    = {2026},
+}
+
+@article{voxcpm2025,
+  title   = {VoxCPM: Tokenizer-Free TTS for Context-Aware Speech Generation
+             and True-to-Life Voice Cloning},
+  author  = {Zhou, Yixuan and Zeng, Guoyang and Liu, Xin and Li, Xiang and
+             Yu, Renjie and Wang, Ziyang and Ye, Runchuan and Sun, Weiyue and
+             Gui, Jiancheng and Li, Kehan and Wu, Zhiyong and Liu, Zhiyuan},
+  journal = {arXiv preprint arXiv:2509.24650},
+  year    = {2025},
+}
+```
 
 ## 📄 License
-The VoxCPM model weights and code are open-sourced under the [Apache-2.0](LICENSE) license.
+
+VoxCPM model weights and code are open-sourced under the [Apache-2.0](LICENSE) license.
 
 ## 🙏 Acknowledgments
 
-We extend our sincere gratitude to the following works and resources for their inspiration and contributions:
-
-- [DiTAR](https://arxiv.org/abs/2502.03930) for the diffusion autoregressive backbone used in speech generation
-- [MiniCPM-4](https://github.com/OpenBMB/MiniCPM) for serving as the language model foundation
-- [CosyVoice](https://github.com/FunAudioLLM/CosyVoice) for the implementation of Flow Matching-based LocDiT
-- [DAC](https://github.com/descriptinc/descript-audio-codec) for providing the Audio VAE backbone
+- [DiTAR](https://arxiv.org/abs/2502.03930) for the diffusion autoregressive backbone
+- [MiniCPM-4](https://github.com/OpenBMB/MiniCPM) for the language model foundation
+- [CosyVoice](https://github.com/FunAudioLLM/CosyVoice) for the Flow Matching-based LocDiT implementation
+- [DAC](https://github.com/descriptinc/descript-audio-codec) for the Audio VAE backbone
+- Our community users for trying VoxCPM, reporting issues, sharing ideas, and contributing—your support helps the project keep getting better
 
 ## Institutions
 
-This project is developed by the following institutions:
-- <img src="assets/modelbest_logo.png" width="28px"> [ModelBest](https://modelbest.cn/)
+<p>
+  <a href="https://modelbest.cn/"><img src="assets/modelbest_logo.png" width="28px"> ModelBest</a>
+  &nbsp;&nbsp;&nbsp;
+  <a href="https://github.com/thuhcsi"><img src="assets/thuhcsi_logo.png" width="28px"> THUHCSI</a>
+</p>
 
-- <img src="assets/thuhcsi_logo.png" width="28px"> [THUHCSI](https://github.com/thuhcsi)
+## ⭐ Star History
 
-
-
-
-## 📚 Citation
-
-If you find our model helpful, please consider citing our projects 📝 and staring us ⭐️！
-
-```bib
-@misc{voxcpm2025,
-  author       = {{Yixuan Zhou, Guoyang Zeng, Xin Liu, Xiang Li, Renjie Yu, Ziyang Wang, Runchuan Ye, Weiyue Sun, Jiancheng Gui, Kehan Li, Zhiyong Wu, Zhiyuan Liu}},
-  title        = {{VoxCPM}},
-  year         = {2025},
-  publish = {\url{https://github.com/OpenBMB/VoxCPM}},
-  note         = {GitHub repository}
-}
-```
+[![Star History Chart](https://api.star-history.com/svg?repos=OpenBMB/VoxCPM&type=Date)](https://star-history.com/#OpenBMB/VoxCPM&Date)

app.py

@@ -1,58 +1,258 @@
 import os
+import sys
+import logging
 import numpy as np
 import torch
-import gradio as gr  
-import spaces
+import gradio as gr
 from typing import Optional, Tuple
 from funasr import AutoModel
 from pathlib import Path
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 if os.environ.get("HF_REPO_ID", "").strip() == "":
-    os.environ["HF_REPO_ID"] = "openbmb/VoxCPM-0.5B"
+    os.environ["HF_REPO_ID"] = "openbmb/VoxCPM2"
 
 import voxcpm
 
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+logger = logging.getLogger(__name__)
+
+# ---------- Inline i18n (en + zh-CN only) ----------
+
+_USAGE_INSTRUCTIONS_EN = (
+    "**VoxCPM2 — Three Modes of Speech Generation:**\n\n"
+    "🎨 **Voice Design** — Create a brand-new voice  \n"
+    "No reference audio required. Describe the desired voice characteristics "
+    "(gender, age, tone, emotion, pace …) in **Control Instruction**, and VoxCPM2 "
+    "will craft a unique voice from your description alone.\n\n"
+    "🎛️ **Controllable Cloning** — Clone a voice with optional style guidance  \n"
+    "Upload a reference audio clip, then use **Control Instruction** to steer "
+    "emotion, speaking pace, and overall style while preserving the original timbre.\n\n"
+    "🎙️ **Ultimate Cloning** — Reproduce every vocal nuance through audio continuation  \n"
+    "Turn on **Ultimate Cloning Mode** and provide (or auto-transcribe) the reference audio's transcript. "
+    "The model treats the reference clip as a spoken prefix and seamlessly **continues** from it, faithfully preserving every vocal detail."
+    "Note: This mode will disable Control Instruction."
+)
+
+_EXAMPLES_FOOTER_EN = (
+    "---\n"
+    "**💡 Voice Description Examples:**  \n"
+    "Try the following Control Instructions to explore different voices:  \n\n"
+    "**Example 1 — Gentle & Melancholic Girl**  \n"
+    '`Control Instruction`: *"A young girl with a soft, sweet voice. '
+    'Speaks slowly with a melancholic, slightly tsundere tone."*  \n'
+    '`Target Text`: *"I never asked you to stay… It\'s not like I care or anything. '
+    'But… why does it still hurt so much now that you\'re gone?"*  \n\n'
+    "**Example 2 — Laid-Back Surfer Dude**  \n"
+    '`Control Instruction`: *"Relaxed young male voice, slightly nasal, '
+    'lazy drawl, very casual and chill."*  \n'
+    '`Target Text`: *"Dude, did you see that set? The waves out there are totally gnarly today. '
+    "Just catching barrels all morning — it's like, totally righteous, you know what I mean?\"*"
+)
+
+_USAGE_INSTRUCTIONS_ZH = (
+    "**VoxCPM2 — 三种语音生成方式：**\n\n"
+    "🎨 **声音设计（Voice Design）**  \n"
+    "无需参考音频。在 **Control Instruction** 中描述目标音色特征"
+    "（性别、年龄、语气、情绪、语速等），VoxCPM2 即可为你从零创造独一无二的声音。\n\n"
+    "🎛️ **可控克隆（Controllable Cloning）**  \n"
+    "上传参考音频，同时可选地使用 **Control Instruction** 来指定情绪、语速、风格等表达方式，"
+    "在保留原始音色的基础上灵活控制说话风格。\n\n"
+    "🎙️ **极致克隆（Ultimate Cloning）**  \n"
+    "开启 **极致克隆模式** 并提供参考音频的文字内容（可自动识别）。"
+    "模型会将参考音频视为已说出的前文，以**音频续写**的方式完整还原参考音频中的所有声音细节。"
+    "注意：该模式与可控克隆模式互斥，将禁用Control Instruction。\n\n"
+)
+
+_EXAMPLES_FOOTER_ZH = (
+    "---\n"
+    "**💡 声音描述示例（中英文均可）：**  \n\n"
+    "**示例 1 — 深宫太后**  \n"
+    '`Control Instruction`: *"中老年女性，声音低沉阴冷，语速缓慢而有力，'
+    '字字深思熟虑，带有深不可测的城府与威慑感。"*  \n'
+    '`Target Text`: *"哀家在这深宫待了四十年，什么风浪没见过？你以为瞒得过哀家？"*  \n\n'
+    "**示例 2 — 暴躁驾校教练**  \n"
+    '`Control Instruction`: *"暴躁的中年男声，语速快，充满无奈和愤怒"*  \n'
+    '`Target Text`: *"踩离合！踩刹车啊！你往哪儿开呢？前面是树你看不见吗？'
+    '我教了你八百遍了，打死方向盘！你是不是想把车给我开到沟里去？"*  \n\n'
+    "---\n"
+    "**🗣️ 方言生成指南：**  \n"
+    "要生成地道的方言语音，请在 **Target Text** 中直接使用方言词汇和句式，"
+    "并在 **Control Instruction** 中描述方言特征。  \n\n"
+    "**示例 — 广东话**  \n"
+    '`Control Instruction`: *"粤语，中年男性，语气平淡"*  \n'
+    '✅ 正确（粤语表达）：*"伙計，唔該一個A餐，凍奶茶少甜！"*  \n'
+    '❌ 错误（普通话原文）：*"伙计，麻烦来一个A餐，冻奶茶少甜！"*  \n\n'
+    "**示例 — 河南话**  \n"
+    '`Control Instruction`: *"河南话，接地气的大叔"*  \n'
+    '✅ 正确（河南话表达）：*"恁这是弄啥嘞？晌午吃啥饭？"*  \n'
+    '❌ 错误（普通话原文）：*"你这是在干什么呢？中午吃什么饭？"*  \n\n'
+    "🤖 **小技巧：** 不知道方言怎么写？可以用豆包、DeepSeek、Kimi 等 AI 助手"
+    "将普通话翻译为方言文本，再粘贴到 Target Text 中即可。  \n\n"
+)
+
+_I18N_TRANSLATIONS = {
+    "en": {
+        "reference_audio_label": "🎤 Reference Audio (optional — upload for cloning)",
+        "show_prompt_text_label": "🎙️ Ultimate Cloning Mode (transcript-guided cloning)",
+        "show_prompt_text_info": "Auto-transcribes reference audio for every vocal nuance reproduced. Control Instruction will be disabled when active.",
+        "prompt_text_label": "Transcript of Reference Audio (auto-filled via ASR, editable)",
+        "prompt_text_placeholder": "The transcript of your reference audio will appear here …",
+        "control_label": "🎛️ Control Instruction (optional — supports Chinese & English)",
+        "control_placeholder": "e.g. A warm young woman / 年轻女性，温柔甜美 / Excited and fast-paced",
+        "target_text_label": "✍️ Target Text — the content to speak",
+        "generate_btn": "🔊 Generate Speech",
+        "generated_audio_label": "Generated Audio",
+        "advanced_settings_title": "⚙️ Advanced Settings",
+        "ref_denoise_label": "Reference audio enhancement",
+        "ref_denoise_info": "Apply ZipEnhancer denoising to the reference audio before cloning",
+        "normalize_label": "Text normalization",
+        "normalize_info": "Normalize numbers, dates, and abbreviations via wetext",
+        "cfg_label": "CFG (guidance scale)",
+        "cfg_info": "Higher → closer to the prompt / reference; lower → more creative variation",
+        "dit_steps_label": "LocDiT flow-matching steps",
+        "dit_steps_info": "LocDiT flow-matching steps — more steps → maybe better audio quality, but slower",
+        "usage_instructions": _USAGE_INSTRUCTIONS_EN,
+        "examples_footer": _EXAMPLES_FOOTER_EN,
+    },
+    "zh-CN": {
+        "reference_audio_label": "🎤 参考音频（可选 — 上传后用于克隆）",
+        "show_prompt_text_label": "🎙️ 极致克隆模式（基于文本引导的极致克隆）",
+        "show_prompt_text_info": "自动识别参考音频文本，完整还原音色、节奏、情感等全部声音细节。开启后 Control Instruction 将暂时禁用",
+        "prompt_text_label": "参考音频内容文本（ASR 自动填充，可手动编辑）",
+        "prompt_text_placeholder": "参考音频的文字内容将自动识别并显示在此处 …",
+        "control_label": "🎛️ Control Instruction（可选 — 支持中英文描述）",
+        "control_placeholder": "如：年轻女性，温柔甜美 / A warm young woman / 暴躁老哥，语速飞快",
+        "target_text_label": "✍️ Target Text — 要合成的目标文本",
+        "generate_btn": "🔊 开始生成",
+        "generated_audio_label": "生成结果",
+        "advanced_settings_title": "⚙️ 高级设置",
+        "ref_denoise_label": "参考音频降噪增强",
+        "ref_denoise_info": "克隆前使用 ZipEnhancer 对参考音频进行降噪处理",
+        "normalize_label": "文本规范化",
+        "normalize_info": "自动规范化数字、日期及缩写（基于 wetext）",
+        "cfg_label": "CFG（引导强度）",
+        "cfg_info": "数值越高 → 越贴合提示/参考音色；数值越低 → 生成风格更自由",
+        "dit_steps_label": "LocDiT 流匹配迭代步数",
+        "dit_steps_info": "LocDiT 流匹配生成迭代步数 — 步数越多 → 可能生成更好的音频质量，但速度变慢",
+        "usage_instructions": _USAGE_INSTRUCTIONS_ZH,
+        "examples_footer": _EXAMPLES_FOOTER_ZH,
+    },
+    "zh-Hans": None,  # alias, filled below
+    "zh": None,       # alias, filled below
+}
+_I18N_TRANSLATIONS["zh-Hans"] = _I18N_TRANSLATIONS["zh-CN"]
+_I18N_TRANSLATIONS["zh"] = _I18N_TRANSLATIONS["zh-CN"]
+
+for _d in _I18N_TRANSLATIONS.values():
+    if _d is not None:
+        for _k, _v in _I18N_TRANSLATIONS["en"].items():
+            _d.setdefault(_k, _v)
+
+I18N = gr.I18n(**_I18N_TRANSLATIONS)
+
+DEFAULT_TARGET_TEXT = (
+    "VoxCPM2 is a creative multilingual TTS model from ModelBest, "
+    "designed to generate highly realistic speech."
+)
+
+_CUSTOM_CSS = """
+.logo-container {
+    text-align: center;
+    margin: 0.5rem 0 1rem 0;
+}
+.logo-container img {
+    height: 80px;
+    width: auto;
+    max-width: 200px;
+    display: inline-block;
+}
+
+/* Toggle switch style */
+.switch-toggle {
+    padding: 8px 12px;
+    border-radius: 8px;
+    background: var(--block-background-fill);
+}
+.switch-toggle input[type="checkbox"] {
+    appearance: none;
+    -webkit-appearance: none;
+    width: 44px;
+    height: 24px;
+    background: #ccc;
+    border-radius: 12px;
+    position: relative;
+    cursor: pointer;
+    transition: background 0.3s ease;
+    flex-shrink: 0;
+}
+.switch-toggle input[type="checkbox"]::after {
+    content: "";
+    position: absolute;
+    top: 2px;
+    left: 2px;
+    width: 20px;
+    height: 20px;
+    background: white;
+    border-radius: 50%;
+    transition: transform 0.3s ease;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.2);
+}
+.switch-toggle input[type="checkbox"]:checked {
+    background: var(--color-accent);
+}
+.switch-toggle input[type="checkbox"]:checked::after {
+    transform: translateX(20px);
+}
+"""
+
+_APP_THEME = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="gray",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"],
+)
+
+
+# ---------- Model ----------
 
 class VoxCPMDemo:
-    def __init__(self) -> None:
+    def __init__(self, model_dir: Optional[str] = None) -> None:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"🚀 Running on device: {self.device}")
+        logger.info(f"Running on device: {self.device}")
 
-        # ASR model for prompt text recognition
         self.asr_model_id = "iic/SenseVoiceSmall"
         self.asr_model: Optional[AutoModel] = AutoModel(
             model=self.asr_model_id,
             disable_update=True,
-            log_level='DEBUG',
+            log_level="DEBUG",
             device="cuda:0" if self.device == "cuda" else "cpu",
         )
 
-        # TTS model (lazy init)
         self.voxcpm_model: Optional[voxcpm.VoxCPM] = None
-        self.default_local_model_dir = "./models/VoxCPM-0.5B"
+        self.explicit_model_dir = model_dir
 
-    # ---------- Model helpers ----------
     def _resolve_model_dir(self) -> str:
-        """
-        Resolve model directory:
-        1) Use local checkpoint directory if exists
-        2) If HF_REPO_ID env is set, download into models/{repo}
-        3) Fallback to 'models'
-        """
-        if os.path.isdir(self.default_local_model_dir):
-            return self.default_local_model_dir
-
+        if self.explicit_model_dir and os.path.isdir(self.explicit_model_dir):
+            return self.explicit_model_dir
+        env_model_dir = os.environ.get("VOXCPM_MODEL_DIR", "").strip()
+        if env_model_dir and os.path.isdir(env_model_dir):
+            return env_model_dir
         repo_id = os.environ.get("HF_REPO_ID", "").strip()
         if len(repo_id) > 0:
             target_dir = os.path.join("models", repo_id.replace("/", "__"))
             if not os.path.isdir(target_dir):
                 try:
-                    from huggingface_hub import snapshot_download  # type: ignore
+                    from huggingface_hub import snapshot_download
                     os.makedirs(target_dir, exist_ok=True)
-                    print(f"Downloading model from HF repo '{repo_id}' to '{target_dir}' ...")
+                    logger.info(f"Downloading model from HF repo '{repo_id}' to '{target_dir}' ...")
                     snapshot_download(repo_id=repo_id, local_dir=target_dir, local_dir_use_symlinks=False)
                 except Exception as e:
-                    print(f"Warning: HF download failed: {e}. Falling back to 'data'.")
+                    logger.warning(f"HF download failed: {e}. Falling back to 'models'.")
                     return "models"
             return target_dir
         return "models"
@@ -60,220 +260,271 @@ class VoxCPMDemo:
     def get_or_load_voxcpm(self) -> voxcpm.VoxCPM:
         if self.voxcpm_model is not None:
             return self.voxcpm_model
-        print("Model not loaded, initializing...")
+        logger.info("Model not loaded, initializing...")
         model_dir = self._resolve_model_dir()
-        print(f"Using model dir: {model_dir}")
-        self.voxcpm_model = voxcpm.VoxCPM(voxcpm_model_path=model_dir)
-        print("Model loaded successfully.")
+        logger.info(f"Using model dir: {model_dir}")
+        self.voxcpm_model = voxcpm.VoxCPM(voxcpm_model_path=model_dir, optimize=True)
+        logger.info("Model loaded successfully.")
         return self.voxcpm_model
 
-    # ---------- Functional endpoints ----------
     def prompt_wav_recognition(self, prompt_wav: Optional[str]) -> str:
         if prompt_wav is None:
             return ""
         res = self.asr_model.generate(input=prompt_wav, language="auto", use_itn=True)
-        text = res[0]["text"].split('|>')[-1]
-        return text
+        return res[0]["text"].split("|>")[-1]
+
+    def _build_generate_kwargs(
+        self,
+        *,
+        final_text: str,
+        audio_path: Optional[str],
+        prompt_text_clean: Optional[str],
+        cfg_value_input: float,
+        do_normalize: bool,
+        denoise: bool,
+        inference_timesteps: int = 10,
+    ) -> dict:
+        generate_kwargs = dict(
+            text=final_text,
+            reference_wav_path=audio_path,
+            cfg_value=float(cfg_value_input),
+            inference_timesteps=inference_timesteps,
+            normalize=do_normalize,
+            denoise=denoise,
+        )
+        if prompt_text_clean and audio_path:
+            generate_kwargs["prompt_wav_path"] = audio_path
+            generate_kwargs["prompt_text"] = prompt_text_clean
+        return generate_kwargs
 
     def generate_tts_audio(
         self,
         text_input: str,
-        prompt_wav_path_input: Optional[str] = None,
-        prompt_text_input: Optional[str] = None,
+        control_instruction: str = "",
+        reference_wav_path_input: Optional[str] = None,
+        prompt_text: str = "",
         cfg_value_input: float = 2.0,
-        inference_timesteps_input: int = 10,
         do_normalize: bool = True,
         denoise: bool = True,
+        inference_timesteps: int = 10,
     ) -> Tuple[int, np.ndarray]:
-        """
-        Generate speech from text using VoxCPM; optional reference audio for voice style guidance.
-        Returns (sample_rate, waveform_numpy)
-        """
         current_model = self.get_or_load_voxcpm()
 
         text = (text_input or "").strip()
         if len(text) == 0:
             raise ValueError("Please input text to synthesize.")
 
-        prompt_wav_path = prompt_wav_path_input if prompt_wav_path_input else None
-        prompt_text = prompt_text_input if prompt_text_input else None
+        control = (control_instruction or "").strip()
+        final_text = f"({control}){text}" if control else text
 
-        print(f"Generating audio for text: '{text[:60]}...'")
-        wav = current_model.generate(
-            text=text,
-            prompt_text=prompt_text,
-            prompt_wav_path=prompt_wav_path,
-            cfg_value=float(cfg_value_input),
-            inference_timesteps=int(inference_timesteps_input),
-            normalize=do_normalize,
+        audio_path = reference_wav_path_input if reference_wav_path_input else None
+        prompt_text_clean = (prompt_text or "").strip() or None
+
+        if audio_path and prompt_text_clean:
+            logger.info(f"[Voice Cloning] prompt_wav + prompt_text + reference_wav")
+        elif audio_path:
+            logger.info(f"[Voice Control] reference_wav only")
+        else:
+            logger.info(f"[Voice Design] control: {control[:50] if control else 'None'}...")
+
+        logger.info(f"Generating audio for text: '{final_text[:80]}...'")
+        generate_kwargs = self._build_generate_kwargs(
+            final_text=final_text,
+            audio_path=audio_path,
+            prompt_text_clean=prompt_text_clean,
+            cfg_value_input=cfg_value_input,
+            do_normalize=do_normalize,
             denoise=denoise,
+            inference_timesteps=inference_timesteps,
         )
-        return (16000, wav)
+        wav = current_model.generate(**generate_kwargs)
+        return (current_model.tts_model.sample_rate, wav)
 
 
-# ---------- UI Builders ----------
+# ---------- UI ----------
 
 def create_demo_interface(demo: VoxCPMDemo):
-    """Build the Gradio UI for VoxCPM demo."""
-    # static assets (logo path)
-    gr.set_static_paths(paths=[Path.cwd().absolute()/"assets"])
+    gr.set_static_paths(paths=[Path.cwd().absolute() / "assets"])
 
-    with gr.Blocks(
-        theme=gr.themes.Soft(
-            primary_hue="blue",
-            secondary_hue="gray",
-            neutral_hue="slate",
-            font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"]
-        ),
-        css="""
-        .logo-container {
-            text-align: center;
-            margin: 0.5rem 0 1rem 0;
-        }
-        .logo-container img {
-            height: 80px;
-            width: auto;
-            max-width: 200px;
-            display: inline-block;
-        }
-        /* Bold accordion labels */
-        #acc_quick details > summary,
-        #acc_tips details > summary {
-            font-weight: 600 !important;
-            font-size: 1.1em !important;
-        }
-        /* Bold labels for specific checkboxes */
-        #chk_denoise label,
-        #chk_denoise span,
-        #chk_normalize label,
-        #chk_normalize span {
-            font-weight: 600;
-        }
-        """
-    ) as interface:
-        # Header logo
-        gr.HTML('<div class="logo-container"><img src="/gradio_api/file=assets/voxcpm_logo.png" alt="VoxCPM Logo"></div>')
+    def _generate(
+        text: str,
+        control_instruction: str,
+        ref_wav: Optional[str],
+        use_prompt_text: bool,
+        prompt_text_value: str,
+        cfg_value: float,
+        do_normalize: bool,
+        denoise: bool,
+        dit_steps: int,
+    ):
+        actual_prompt_text = prompt_text_value.strip() if use_prompt_text else ""
+        actual_control = "" if use_prompt_text else control_instruction
+        sr, wav_np = demo.generate_tts_audio(
+            text_input=text,
+            control_instruction=actual_control,
+            reference_wav_path_input=ref_wav,
+            prompt_text=actual_prompt_text,
+            cfg_value_input=cfg_value,
+            do_normalize=do_normalize,
+            denoise=denoise,
+            inference_timesteps=int(dit_steps),
+        )
+        return (sr, wav_np)
 
-        # Quick Start
-        with gr.Accordion("📋 Quick Start Guide ｜快速入门", open=False, elem_id="acc_quick"):
-            gr.Markdown("""
-            ### How to Use ｜使用说明
-            1. **(Optional) Provide a Voice Prompt** - Upload or record an audio clip to provide the desired voice characteristics for synthesis.  
-               **（可选）提供参考声音** - 上传或录制一段音频，为声音合成提供音色、语调和情感等个性化特征
-            2. **(Optional) Enter prompt text** - If you provided a voice prompt, enter the corresponding transcript here (auto-recognition available).  
-               **（可选项）输入参考文本** - 如果提供了参考语音，请输入其对应的文本内容（支持自动识别）。
-            3. **Enter target text** - Type the text you want the model to speak.  
-               **输入目标文本** - 输入您希望模型朗读的文字内容。
-            4. **Generate Speech** - Click the "Generate" button to create your audio.  
-               **生成语音** - 点击"生成"按钮，即可为您创造出音频。
-            """)
+    def _on_toggle_instant(checked):
+        """Instant UI toggle — no ASR, no blocking."""
+        if checked:
+            return (
+                gr.update(visible=True, value="", placeholder="Recognizing reference audio..."),
+                gr.update(visible=False),
+            )
+        return (
+            gr.update(visible=False),
+            gr.update(visible=True, interactive=True),
+        )
 
-        # Pro Tips
-        with gr.Accordion("💡 Pro Tips ｜使用建议", open=False, elem_id="acc_tips"):
-            gr.Markdown(f"""
-            ### Prompt Speech Enhancement｜参考语音降噪
-            - **Enable** to remove background noise for a clean, studio-like voice, with an external ZipEnhancer component.  
-              **启用**：通过 ZipEnhancer 组件消除背景噪音，获得更好的音质。
-            - **Disable** to preserve the original audio's background atmosphere.  
-              **禁用**：保留原始音频的背景环境声，如果想复刻相应声学环境。
+    def _run_asr_if_needed(checked, audio_path):
+        """Run ASR after the UI has updated. Only when toggled ON."""
+        if not checked or not audio_path:
+            return gr.update()
+        try:
+            logger.info("Running ASR on reference audio...")
+            asr_text = demo.prompt_wav_recognition(audio_path)
+            logger.info(f"ASR result: {asr_text[:60]}...")
+            return gr.update(value=asr_text)
+        except Exception as e:
+            logger.warning(f"ASR recognition failed: {e}")
+            return gr.update(value="")
 
-            ### Text Normalization｜文本正则化
-            - **Enable** to process general text with an external WeTextProcessing component.  
-              **启用**：使用 WeTextProcessing 组件，可处理常见文本。
-            - **Disable** to use VoxCPM's native text understanding ability. For example, it supports phonemes input ({HH AH0 L OW1}), try it!  
-              **禁用**：将使用 VoxCPM 内置的文本理解能力。如，支持音素输入（如 {da4}{jia1}好）和公式符号合成，尝试一下！
+    with gr.Blocks() as interface:
+        gr.HTML(
+            '<div class="logo-container">'
+            '<img src="/gradio_api/file=assets/voxcpm_logo.png" alt="VoxCPM Logo">'
+            "</div>"
+        )
 
-            ### CFG Value｜CFG 值
-            - **Lower CFG** if the voice prompt sounds strained or expressive.  
-              **调低**：如果提示语音听起来不自然或过于夸张。
-            - **Higher CFG** for better adherence to the prompt speech style or input text.  
-              **调高**：为更好地贴合提示音频的风格或输入文本。
+        gr.Markdown(I18N("usage_instructions"))
 
-            ### Inference Timesteps｜推理时间步
-            - **Lower** for faster synthesis speed.  
-              **调低**：合成速度更快。
-            - **Higher** for better synthesis quality.  
-              **调高**：合成质量更佳。
-
-            ### Long Text (e.g., >5 min speech)｜长文本 (如 >5分钟的合成语音)
-            While VoxCPM can handle long texts directly, we recommend using empty lines to break very long content into paragraphs; the model will then synthesize each paragraph individually.  
-            虽然 VoxCPM 支持直接生成长文本，但如果目标文本过长，我们建议使用换行符将内容分段；模型将对每个段落分别合成。
-            """)
-
-        # Main controls
         with gr.Row():
             with gr.Column():
-                prompt_wav = gr.Audio(
-                    sources=["upload", 'microphone'],
+                reference_wav = gr.Audio(
+                    sources=["upload", "microphone"],
                     type="filepath",
-                    label="Prompt Speech",
-                    value="./examples/example.wav",
+                    label=I18N("reference_audio_label"),
                 )
-                DoDenoisePromptAudio = gr.Checkbox(
+                show_prompt_text = gr.Checkbox(
                     value=False,
-                    label="Prompt Speech Enhancement",
-                    elem_id="chk_denoise",
-                    info="We use ZipEnhancer model to denoise the prompt audio."
+                    label=I18N("show_prompt_text_label"),
+                    info=I18N("show_prompt_text_info"),
+                    elem_classes=["switch-toggle"],
+                )
+                prompt_text = gr.Textbox(
+                    value="",
+                    label=I18N("prompt_text_label"),
+                    placeholder=I18N("prompt_text_placeholder"),
+                    lines=2,
+                    visible=False,
+                )
+                control_instruction = gr.Textbox(
+                    value="",
+                    label=I18N("control_label"),
+                    placeholder=I18N("control_placeholder"),
+                    lines=2,
+                )
+                text = gr.Textbox(
+                    value=DEFAULT_TARGET_TEXT,
+                    label=I18N("target_text_label"),
+                    lines=3,
                 )
-                with gr.Row():
-                    prompt_text = gr.Textbox(
-                        value="Just by listening a few minutes a day, you'll be able to eliminate negative thoughts by conditioning your mind to be more positive.",
-                        label="Prompt Text",
-                        placeholder="Please enter the prompt text. Automatic recognition is supported, and you can correct the results yourself..."
-                    )
-                run_btn = gr.Button("Generate Speech", variant="primary")
 
-            with gr.Column():
-                cfg_value = gr.Slider(
-                    minimum=1.0,
-                    maximum=3.0,
-                    value=2.0,
-                    step=0.1,
-                    label="CFG Value (Guidance Scale)",
-                    info="Higher values increase adherence to prompt, lower values allow more creativity"
-                )
-                inference_timesteps = gr.Slider(
-                    minimum=4,
-                    maximum=30,
-                    value=10,
-                    step=1,
-                    label="Inference Timesteps",
-                    info="Number of inference timesteps for generation (higher values may improve quality but slower)"
-                )
-                with gr.Row():
-                    text = gr.Textbox(
-                        value="VoxCPM is an innovative end-to-end TTS model from ModelBest, designed to generate highly realistic speech.",
-                        label="Target Text",
-                        info="Default processing splits text on \\n into paragraphs; each is synthesized as a chunk and then concatenated into the final audio."
+                with gr.Accordion(I18N("advanced_settings_title"), open=False):
+                    DoDenoisePromptAudio = gr.Checkbox(
+                        value=False,
+                        label=I18N("ref_denoise_label"),
+                        elem_classes=["switch-toggle"],
+                        info=I18N("ref_denoise_info"),
                     )
-                with gr.Row():
                     DoNormalizeText = gr.Checkbox(
                         value=False,
-                        label="Text Normalization",
-                        elem_id="chk_normalize",
-                        info="We use WeTextPorcessing library to normalize the input text."
+                        label=I18N("normalize_label"),
+                        elem_classes=["switch-toggle"],
+                        info=I18N("normalize_info"),
+                    )
+                    cfg_value = gr.Slider(
+                        minimum=1.0,
+                        maximum=3.0,
+                        value=2.0,
+                        step=0.1,
+                        label=I18N("cfg_label"),
+                        info=I18N("cfg_info"),
+                    )
+                    dit_steps = gr.Slider(
+                        minimum=1,
+                        maximum=50,
+                        value=10,
+                        step=1,
+                        label=I18N("dit_steps_label"),
+                        info=I18N("dit_steps_info"),
                     )
-                audio_output = gr.Audio(label="Output Audio")
 
-        # Wiring
+                run_btn = gr.Button(I18N("generate_btn"), variant="primary", size="lg")
+
+            with gr.Column():
+                audio_output = gr.Audio(label=I18N("generated_audio_label"))
+                gr.Markdown(I18N("examples_footer"))
+
+        show_prompt_text.change(
+            fn=_on_toggle_instant,
+            inputs=[show_prompt_text],
+            outputs=[prompt_text, control_instruction],
+        ).then(
+            fn=_run_asr_if_needed,
+            inputs=[show_prompt_text, reference_wav],
+            outputs=[prompt_text],
+        )
+
         run_btn.click(
-            fn=demo.generate_tts_audio,
-            inputs=[text, prompt_wav, prompt_text, cfg_value, inference_timesteps, DoNormalizeText, DoDenoisePromptAudio],
+            fn=_generate,
+            inputs=[
+                text,
+                control_instruction,
+                reference_wav,
+                show_prompt_text,
+                prompt_text,
+                cfg_value,
+                DoNormalizeText,
+                DoDenoisePromptAudio,
+                dit_steps,
+            ],
             outputs=[audio_output],
             show_progress=True,
             api_name="generate",
         )
-        prompt_wav.change(fn=demo.prompt_wav_recognition, inputs=[prompt_wav], outputs=[prompt_text])
 
     return interface
 
-
-def run_demo(server_name: str = "localhost", server_port: int = 7860, show_error: bool = True):
-    demo = VoxCPMDemo()
+def run_demo(
+    server_name: str = "0.0.0.0",
+    server_port: int = 8808,
+    show_error: bool = True,
+    model_dir: Optional[str] = None,
+):
+    demo = VoxCPMDemo(model_dir=model_dir)
     interface = create_demo_interface(demo)
-    # Recommended to enable queue on Spaces for better throughput
-    interface.queue(max_size=10).launch(server_name=server_name, server_port=server_port, show_error=show_error)
+    interface.queue(max_size=10, default_concurrency_limit=1).launch(
+        server_name=server_name,
+        server_port=server_port,
+        show_error=show_error,
+        i18n=I18N,
+        theme=_APP_THEME,
+        css=_CUSTOM_CSS,
+    )
 
 
 if __name__ == "__main__":
-    run_demo()
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=str, default=None, help="Path to VoxCPM2 checkpoint directory")
+    parser.add_argument("--port", type=int, default=8808, help="Server port")
+    args = parser.parse_args()
+    run_demo(model_dir=args.model_dir, server_port=args.port)

app_old.py

@@ -0,0 +1,280 @@
+import os
+import sys
+import numpy as np
+import torch
+import gradio as gr
+from typing import Optional, Tuple
+from funasr import AutoModel
+from pathlib import Path
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+if os.environ.get("HF_REPO_ID", "").strip() == "":
+    os.environ["HF_REPO_ID"] = "openbmb/VoxCPM1.5"
+
+import voxcpm
+
+
+class VoxCPMDemo:
+    def __init__(self) -> None:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"🚀 Running on device: {self.device}", file=sys.stderr)
+
+        # ASR model for prompt text recognition
+        self.asr_model_id = "iic/SenseVoiceSmall"
+        self.asr_model: Optional[AutoModel] = AutoModel(
+            model=self.asr_model_id,
+            disable_update=True,
+            log_level='DEBUG',
+            device="cuda:0" if self.device == "cuda" else "cpu",
+        )
+
+        # TTS model (lazy init)
+        self.voxcpm_model: Optional[voxcpm.VoxCPM] = None
+        self.default_local_model_dir = "./models/VoxCPM1.5"
+
+    # ---------- Model helpers ----------
+    def _resolve_model_dir(self) -> str:
+        """
+        Resolve model directory:
+        1) Use local checkpoint directory if exists
+        2) If HF_REPO_ID env is set, download into models/{repo}
+        3) Fallback to 'models'
+        """
+        if os.path.isdir(self.default_local_model_dir):
+            return self.default_local_model_dir
+
+        repo_id = os.environ.get("HF_REPO_ID", "").strip()
+        if len(repo_id) > 0:
+            target_dir = os.path.join("models", repo_id.replace("/", "__"))
+            if not os.path.isdir(target_dir):
+                try:
+                    from huggingface_hub import snapshot_download  # type: ignore
+                    os.makedirs(target_dir, exist_ok=True)
+                    print(f"Downloading model from HF repo '{repo_id}' to '{target_dir}' ...", file=sys.stderr)
+                    snapshot_download(repo_id=repo_id, local_dir=target_dir, local_dir_use_symlinks=False)
+                except Exception as e:
+                    print(f"Warning: HF download failed: {e}. Falling back to 'data'.", file=sys.stderr)
+                    return "models"
+            return target_dir
+        return "models"
+
+    def get_or_load_voxcpm(self) -> voxcpm.VoxCPM:
+        if self.voxcpm_model is not None:
+            return self.voxcpm_model
+        print("Model not loaded, initializing...", file=sys.stderr)
+        model_dir = self._resolve_model_dir()
+        print(f"Using model dir: {model_dir}", file=sys.stderr)
+        self.voxcpm_model = voxcpm.VoxCPM(voxcpm_model_path=model_dir)
+        print("Model loaded successfully.", file=sys.stderr)
+        return self.voxcpm_model
+
+    # ---------- Functional endpoints ----------
+    def prompt_wav_recognition(self, prompt_wav: Optional[str]) -> str:
+        if prompt_wav is None:
+            return ""
+        res = self.asr_model.generate(input=prompt_wav, language="auto", use_itn=True)
+        text = res[0]["text"].split('|>')[-1]
+        return text
+
+    def generate_tts_audio(
+        self,
+        text_input: str,
+        prompt_wav_path_input: Optional[str] = None,
+        prompt_text_input: Optional[str] = None,
+        cfg_value_input: float = 2.0,
+        inference_timesteps_input: int = 10,
+        do_normalize: bool = True,
+        denoise: bool = True,
+    ) -> Tuple[int, np.ndarray]:
+        """
+        Generate speech from text using VoxCPM; optional reference audio for voice style guidance.
+        Returns (sample_rate, waveform_numpy)
+        """
+        current_model = self.get_or_load_voxcpm()
+
+        text = (text_input or "").strip()
+        if len(text) == 0:
+            raise ValueError("Please input text to synthesize.")
+
+        prompt_wav_path = prompt_wav_path_input if prompt_wav_path_input else None
+        prompt_text = prompt_text_input if prompt_text_input else None
+
+        print(f"Generating audio for text: '{text[:60]}...'", file=sys.stderr)
+        wav = current_model.generate(
+            text=text,
+            prompt_text=prompt_text,
+            prompt_wav_path=prompt_wav_path,
+            cfg_value=float(cfg_value_input),
+            inference_timesteps=int(inference_timesteps_input),
+            normalize=do_normalize,
+            denoise=denoise,
+        )
+        return (current_model.tts_model.sample_rate, wav)
+
+
+# ---------- UI Builders ----------
+
+_APP_THEME = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="gray",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"],
+)
+
+_CUSTOM_CSS = """
+.logo-container {
+    text-align: center;
+    margin: 0.5rem 0 1rem 0;
+}
+.logo-container img {
+    height: 80px;
+    width: auto;
+    max-width: 200px;
+    display: inline-block;
+}
+/* Bold accordion labels */
+#acc_quick details > summary,
+#acc_tips details > summary {
+    font-weight: 600 !important;
+    font-size: 1.1em !important;
+}
+/* Bold labels for specific checkboxes */
+#chk_denoise label,
+#chk_denoise span,
+#chk_normalize label,
+#chk_normalize span {
+    font-weight: 600;
+}
+"""
+
+
+def create_demo_interface(demo: VoxCPMDemo):
+    """Build the Gradio UI for VoxCPM demo."""
+    gr.set_static_paths(paths=[Path.cwd().absolute()/"assets"])
+
+    with gr.Blocks() as interface:
+        # Header logo
+        gr.HTML('<div class="logo-container"><img src="/gradio_api/file=assets/voxcpm_logo.png" alt="VoxCPM Logo"></div>')
+
+        # Quick Start
+        with gr.Accordion("📋 Quick Start Guide ｜快速入门", open=False, elem_id="acc_quick"):
+            gr.Markdown("""
+            ### How to Use ｜使用说明
+            1. **(Optional) Provide a Voice Prompt** - Upload or record an audio clip to provide the desired voice characteristics for synthesis.  
+               **（可选）提供参考声音** - 上传或录制一段音频，为声音合成提供音色、语调和情感等个性化特征
+            2. **(Optional) Enter prompt text** - If you provided a voice prompt, enter the corresponding transcript here (auto-recognition available).  
+               **（可选项）输入参考文本** - 如果提供了参考语音，请输入其对应的文本内容（支持自动识别）。
+            3. **Enter target text** - Type the text you want the model to speak.  
+               **输入目标文本** - 输入您希望模型朗读的文字内容。
+            4. **Generate Speech** - Click the "Generate" button to create your audio.  
+               **生成语音** - 点击"生成"按钮，即可为您创造出音频。
+            """)
+
+        # Pro Tips
+        with gr.Accordion("💡 Pro Tips ｜使用建议", open=False, elem_id="acc_tips"):
+            gr.Markdown("""
+            ### Prompt Speech Enhancement｜参考语音降噪
+            - **Enable** to remove background noise for a clean voice, with an external ZipEnhancer component. However, this will limit the audio sampling rate to 16kHz, restricting the cloning quality ceiling.  
+              **启用**：通过 ZipEnhancer 组件消除背景噪音，但会将音频采样率限制在16kHz，限制克隆上限。
+            - **Disable** to preserve the original audio's all information, including background atmosphere, and support audio cloning up to 44.1kHz sampling rate.  
+              **禁用**：保留原始音频的全部信息，包括背景环境声，最高支持44.1kHz的音频复刻。
+
+            ### Text Normalization｜文本正则化
+            - **Enable** to process general text with an external WeTextProcessing component.  
+              **启用**：使用 WeTextProcessing 组件，可支持常见文本的正则化处理。
+            - **Disable** to use VoxCPM's native text understanding ability. For example, it supports phonemes input (For Chinese, phonemes are converted using pinyin, {ni3}{hao3}; For English, phonemes are converted using CMUDict, {HH AH0 L OW1}), try it!  
+              **禁用**：将使用 VoxCPM 内置的文本理解能力。如，支持音素输入（如中文转拼音：{ni3}{hao3}；英文转CMUDict：{HH AH0 L OW1}）和公式符号合成，尝试一下！
+
+            ### CFG Value｜CFG 值
+            - **Lower CFG** if the voice prompt sounds strained or expressive, or instability occurs with long text input.  
+              **调低**：如果提示语音听起来不自然或过于夸张，或者长文本输入出现稳定性问题。
+            - **Higher CFG** for better adherence to the prompt speech style or input text, or instability occurs with too short text input.
+              **调高**：为更好地贴合提示音频的风格或输入文本， 或者极短文本输入出现稳定性问题。
+
+            ### Inference Timesteps｜推理时间步
+            - **Lower** for faster synthesis speed.  
+              **调低**：合成速度更快。
+            - **Higher** for better synthesis quality.  
+              **调高**：合成质量更佳。
+            """)
+
+        # Main controls
+        with gr.Row():
+            with gr.Column():
+                prompt_wav = gr.Audio(
+                    sources=["upload", 'microphone'],
+                    type="filepath",
+                    label="Prompt Speech (Optional, or let VoxCPM improvise)",
+                    value="./examples/example.wav",
+                )
+                DoDenoisePromptAudio = gr.Checkbox(
+                    value=False,
+                    label="Prompt Speech Enhancement",
+                    elem_id="chk_denoise",
+                    info="We use ZipEnhancer model to denoise the prompt audio."
+                )
+                with gr.Row():
+                    prompt_text = gr.Textbox(
+                        value="Just by listening a few minutes a day, you'll be able to eliminate negative thoughts by conditioning your mind to be more positive.",
+                        label="Prompt Text",
+                        placeholder="Please enter the prompt text. Automatic recognition is supported, and you can correct the results yourself..."
+                    )
+                run_btn = gr.Button("Generate Speech", variant="primary")
+
+            with gr.Column():
+                cfg_value = gr.Slider(
+                    minimum=1.0,
+                    maximum=3.0,
+                    value=2.0,
+                    step=0.1,
+                    label="CFG Value (Guidance Scale)",
+                    info="Higher values increase adherence to prompt, lower values allow more creativity"
+                )
+                inference_timesteps = gr.Slider(
+                    minimum=4,
+                    maximum=30,
+                    value=10,
+                    step=1,
+                    label="Inference Timesteps",
+                    info="Number of inference timesteps for generation (higher values may improve quality but slower)"
+                )
+                with gr.Row():
+                    text = gr.Textbox(
+                        value="VoxCPM is an innovative end-to-end TTS model from ModelBest, designed to generate highly realistic speech.",
+                        label="Target Text",
+                    )
+                with gr.Row():
+                    DoNormalizeText = gr.Checkbox(
+                        value=False,
+                        label="Text Normalization",
+                        elem_id="chk_normalize",
+                        info="We use wetext library to normalize the input text."
+                    )
+                audio_output = gr.Audio(label="Output Audio")
+
+        # Wiring
+        run_btn.click(
+            fn=demo.generate_tts_audio,
+            inputs=[text, prompt_wav, prompt_text, cfg_value, inference_timesteps, DoNormalizeText, DoDenoisePromptAudio],
+            outputs=[audio_output],
+            show_progress=True,
+            api_name="generate",
+        )
+        prompt_wav.change(fn=demo.prompt_wav_recognition, inputs=[prompt_wav], outputs=[prompt_text])
+
+    return interface
+
+
+def run_demo(server_name: str = "localhost", server_port: int = 7860, show_error: bool = True):
+    demo = VoxCPMDemo()
+    interface = create_demo_interface(demo)
+    interface.queue(max_size=10, default_concurrency_limit=1).launch(
+        server_name=server_name,
+        server_port=server_port,
+        show_error=show_error,
+        theme=_APP_THEME,
+        css=_CUSTOM_CSS,
+    )
+
+
+if __name__ == "__main__":
+    run_demo()

conf/voxcpm_v1.5/voxcpm_finetune_all.yaml

@@ -0,0 +1,21 @@
+pretrained_path: /path/to/VoxCPM1.5/
+train_manifest: /path/to/train.jsonl
+val_manifest: null
+sample_rate: 44100
+batch_size: 16
+grad_accum_steps: 1  # Gradient accumulation steps, >1 can increase effective batch size without increasing memory
+num_workers: 2
+num_iters: 2000
+log_interval: 10
+valid_interval: 1000
+save_interval: 1000
+learning_rate: 0.00001 
+weight_decay: 0.01
+warmup_steps: 100
+max_steps: 2000
+max_batch_tokens: 8192  # Example: single batch can have at most 16k tokens, with batch_size=4, each sample can have at most 4096 tokens
+save_path: /path/to/checkpoints/finetune_all
+tensorboard: /path/to/logs/finetune_all
+lambdas:
+  loss/diff: 1.0
+  loss/stop: 1.0

conf/voxcpm_v1.5/voxcpm_finetune_lora.yaml

@@ -0,0 +1,36 @@
+pretrained_path: /path/to/VoxCPM1.5/
+train_manifest: /path/to/train.jsonl
+val_manifest: null
+sample_rate: 44100
+batch_size: 16
+grad_accum_steps: 1  # Gradient accumulation steps, >1 can increase effective batch size without increasing memory
+num_workers: 2
+num_iters: 2000
+log_interval: 10
+valid_interval: 1000
+save_interval: 1000
+learning_rate: 0.0001  
+weight_decay: 0.01
+warmup_steps: 100
+max_steps: 2000
+max_batch_tokens: 8192  # Example: single batch can have at most 16k tokens, with batch_size=4, each sample can have at most 4096 tokens
+save_path: /path/to/checkpoints/finetune_lora
+tensorboard: /path/to/logs/finetune_lora
+lambdas:
+  loss/diff: 1.0
+  loss/stop: 1.0
+
+# LoRA configuration
+lora:
+  enable_lm: true
+  enable_dit: true
+  enable_proj: false
+  r: 8               
+  alpha: 16          
+  dropout: 0.0
+
+# Distribution options (optional)
+# - If distribute=false (default): save pretrained_path as base_model in lora_config.json
+# - If distribute=true: save hf_model_id as base_model (hf_model_id is required)
+# hf_model_id: "openbmb/VoxCPM1.5"
+# distribute: true

conf/voxcpm_v1/voxcpm_finetune_all.yaml

@@ -0,0 +1,21 @@
+pretrained_path: /path/to/VoxCPM-0.5B/
+train_manifest: /path/to/train.jsonl
+val_manifest: null
+sample_rate: 16000
+batch_size: 16
+grad_accum_steps: 1  # Gradient accumulation steps, >1 can increase effective batch size without increasing memory
+num_workers: 2
+num_iters: 2000
+log_interval: 10
+valid_interval: 1000
+save_interval: 1000
+learning_rate: 0.00001 
+weight_decay: 0.01
+warmup_steps: 100
+max_steps: 2000
+max_batch_tokens: 8192  # Example: single batch can have at most 16k tokens, with batch_size=4, each sample can have at most 4096 tokens
+save_path: /path/to/checkpoints/finetune_all
+tensorboard: /path/to/logs/finetune_all
+lambdas:
+  loss/diff: 1.0
+  loss/stop: 1.0

conf/voxcpm_v1/voxcpm_finetune_lora.yaml

@@ -0,0 +1,36 @@
+pretrained_path: /path/to/VoxCPM-0.5B/
+train_manifest: /path/to/train.jsonl
+val_manifest: null
+sample_rate: 16000
+batch_size: 16
+grad_accum_steps: 1  # Gradient accumulation steps, >1 can increase effective batch size without increasing memory
+num_workers: 2
+num_iters: 2000
+log_interval: 10
+valid_interval: 1000
+save_interval: 1000
+learning_rate: 0.0001  
+weight_decay: 0.01
+warmup_steps: 100
+max_steps: 2000
+max_batch_tokens: 8192  # Example: single batch can have at most 16k tokens, with batch_size=4, each sample can have at most 4096 tokens
+save_path: /path/to/checkpoints/finetune_lora
+tensorboard: /path/to/logs/finetune_lora
+lambdas:
+  loss/diff: 1.0
+  loss/stop: 1.0
+
+# LoRA configuration
+lora:
+  enable_lm: true
+  enable_dit: true
+  enable_proj: false
+  r: 8             
+  alpha: 16           
+  dropout: 0.0
+
+# Distribution options (optional)
+# - If distribute=false (default): save pretrained_path as base_model in lora_config.json
+# - If distribute=true: save hf_model_id as base_model (hf_model_id is required)
+# hf_model_id: "openbmb/VoxCPM-0.5B"
+# distribute: true       

conf/voxcpm_v2/voxcpm_finetune_all.yaml

@@ -0,0 +1,21 @@
+pretrained_path: /path/to/VoxCPM2/
+train_manifest: /path/to/train.jsonl
+val_manifest: null
+sample_rate: 48000
+batch_size: 2
+grad_accum_steps: 8  # effective batch size = batch_size × grad_accum_steps = 16
+num_workers: 8
+num_iters: 1000
+log_interval: 10
+valid_interval: 500
+save_interval: 500
+learning_rate: 0.00001
+weight_decay: 0.01
+warmup_steps: 100
+max_steps: 1000
+max_batch_tokens: 8192
+save_path: /path/to/checkpoints/finetune_all
+tensorboard: /path/to/logs/finetune_all
+lambdas:
+  loss/diff: 1.0
+  loss/stop: 1.0

conf/voxcpm_v2/voxcpm_finetune_lora.yaml

@@ -0,0 +1,36 @@
+pretrained_path: /path/to/VoxCPM2/
+train_manifest: /path/to/train.jsonl
+val_manifest: null
+sample_rate: 48000
+batch_size: 2
+grad_accum_steps: 8  # effective batch size = batch_size × grad_accum_steps = 16
+num_workers: 8
+num_iters: 1000
+log_interval: 10
+valid_interval: 500
+save_interval: 500
+learning_rate: 0.0001
+weight_decay: 0.01
+warmup_steps: 100
+max_steps: 1000
+max_batch_tokens: 8192
+save_path: /path/to/checkpoints/finetune_lora
+tensorboard: /path/to/logs/finetune_lora
+lambdas:
+  loss/diff: 1.0
+  loss/stop: 1.0
+
+# LoRA configuration
+lora:
+  enable_lm: true
+  enable_dit: true
+  enable_proj: false
+  r: 32
+  alpha: 32
+  dropout: 0.0
+
+# Distribution options (optional)
+# - If distribute=false (default): save pretrained_path as base_model in lora_config.json
+# - If distribute=true: save hf_model_id as base_model (hf_model_id is required)
+# hf_model_id: "openbmb/VoxCPM2"
+# distribute: true

examples/train_data_example.jsonl

@@ -0,0 +1,6 @@
+{"audio": "examples/example.wav", "text": "This is an example audio transcript for training."}
+{"audio": "/absolute/path/to/audio1.wav", "text": "You can use absolute paths for audio files."}
+{"audio": "relative/path/to/audio2.wav", "text": "Or relative paths from the working directory."}
+{"audio": "data/audio3.wav", "text": "Each line is a JSON object with audio path and text.", "duration": 3.5}
+{"audio": "data/audio4.wav", "text": "Optional: add duration field to skip audio loading during filtering.", "duration": 2.8}
+{"audio": "data/audio5.wav", "text": "Optional: add dataset_id for multi-dataset training.", "dataset_id": 1}

pyproject.toml

@@ -20,28 +20,35 @@ classifiers = [
     "Intended Audience :: Developers",
     "Operating System :: OS Independent",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 dependencies = [
     "torch>=2.5.0",
     "torchaudio>=2.5.0",
+    "torchcodec",
     "transformers>=4.36.2",
     "einops",
-    "gradio",
+    "gradio>=6,<7",
     "inflect",
     "addict",
-    "WeTextProcessing",
-    "modelscope",
+    "wetext",
+    "modelscope>=1.22.0",
+    "datasets>=3,<4",
     "huggingface-hub",
     "pydantic",
     "tqdm",
+    "simplejson",
+    "sortedcontainers",
     "soundfile",
+    "librosa",
+    "matplotlib",
     "funasr",
-    "spaces"
+    "spaces",
+    "argbind",
+    "safetensors"
+
 ]
 
 [project.optional-dependencies]
@@ -50,7 +57,6 @@ dev = [
     "pytest-cov>=2.0",
     "black>=21.0",
     "flake8>=3.8",
-    "mypy>=0.800",
     "pre-commit>=2.0",
 ]
 
@@ -75,7 +81,7 @@ version_scheme = "post-release"
 
 [tool.black]
 line-length = 120
-target-version = ['py38']
+target-version = ['py310']
 include = '\.pyi?$'
 extend-exclude = '''
 /(
@@ -83,7 +89,6 @@ extend-exclude = '''
   \.eggs
   | \.git
   | \.hg
-  | \.mypy_cache
   | \.tox
   | \.venv
   | build

scripts/test_voxcpm_ft_infer.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""
+Full finetune inference script (no LoRA).
+
+Checkpoint directory contains complete model files (pytorch_model.bin, config.json, audiovae.pth, etc.),
+can be loaded directly via VoxCPM.
+
+Usage:
+
+    python scripts/test_voxcpm_ft_infer.py \
+        --ckpt_dir /path/to/checkpoints/step_0001000 \
+        --text "Hello, I am the finetuned VoxCPM." \
+        --output ft_test.wav
+
+With voice cloning:
+
+    python scripts/test_voxcpm_ft_infer.py \
+        --ckpt_dir /path/to/checkpoints/step_0001000 \
+        --text "Hello, this is voice cloning result." \
+        --prompt_audio path/to/ref.wav \
+        --prompt_text "Reference audio transcript" \
+        --output ft_clone.wav
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import soundfile as sf
+
+from voxcpm.core import VoxCPM
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("VoxCPM full-finetune inference test (no LoRA)")
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        required=True,
+        help="Checkpoint directory (contains pytorch_model.bin, config.json, audiovae.pth, etc.)",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Target text to synthesize",
+    )
+    parser.add_argument(
+        "--prompt_audio",
+        type=str,
+        default="",
+        help="Optional: reference audio path for voice cloning",
+    )
+    parser.add_argument(
+        "--prompt_text",
+        type=str,
+        default="",
+        help="Optional: transcript of reference audio",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="ft_test.wav",
+        help="Output wav file path",
+    )
+    parser.add_argument(
+        "--cfg_value",
+        type=float,
+        default=2.0,
+        help="CFG scale (default: 2.0)",
+    )
+    parser.add_argument(
+        "--inference_timesteps",
+        type=int,
+        default=10,
+        help="Diffusion inference steps (default: 10)",
+    )
+    parser.add_argument(
+        "--max_len",
+        type=int,
+        default=600,
+        help="Max generation steps",
+    )
+    parser.add_argument(
+        "--normalize",
+        action="store_true",
+        help="Enable text normalization",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # Load model from checkpoint directory (no denoiser)
+    print(f"[FT Inference] Loading model: {args.ckpt_dir}", file=sys.stderr)
+    model = VoxCPM.from_pretrained(
+        hf_model_id=args.ckpt_dir,
+        load_denoiser=False,
+        optimize=True,
+    )
+
+    # Run inference
+    prompt_wav_path = args.prompt_audio if args.prompt_audio else None
+    prompt_text = args.prompt_text if args.prompt_text else None
+
+    print(f"[FT Inference] Synthesizing: text='{args.text}'", file=sys.stderr)
+    if prompt_wav_path:
+        print(f"[FT Inference] Using reference audio: {prompt_wav_path}", file=sys.stderr)
+        print(f"[FT Inference] Reference text: {prompt_text}", file=sys.stderr)
+
+    audio_np = model.generate(
+        text=args.text,
+        prompt_wav_path=prompt_wav_path,
+        prompt_text=prompt_text,
+        cfg_value=args.cfg_value,
+        inference_timesteps=args.inference_timesteps,
+        max_len=args.max_len,
+        normalize=args.normalize,
+        denoise=False,
+    )
+
+    # Save audio
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    sf.write(str(out_path), audio_np, model.tts_model.sample_rate)
+
+    print(
+        f"[FT Inference] Saved to: {out_path}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
+        file=sys.stderr,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/test_voxcpm_lora_infer.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python3
+"""
+LoRA inference test script.
+
+Usage:
+
+    python scripts/test_voxcpm_lora_infer.py \
+        --lora_ckpt checkpoints/step_0002000 \
+        --text "Hello, this is LoRA finetuned result." \
+        --output lora_test.wav
+
+With voice cloning:
+
+    python scripts/test_voxcpm_lora_infer.py \
+        --lora_ckpt checkpoints/step_0002000 \
+        --text "This is voice cloning result." \
+        --prompt_audio path/to/ref.wav \
+        --prompt_text "Reference audio transcript" \
+        --output lora_clone.wav
+
+Note: The script reads base_model path and lora_config from lora_config.json
+      in the checkpoint directory (saved automatically during training).
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import soundfile as sf
+
+from voxcpm.core import VoxCPM
+from voxcpm.model.voxcpm import LoRAConfig
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("VoxCPM LoRA inference test")
+    parser.add_argument(
+        "--lora_ckpt",
+        type=str,
+        required=True,
+        help="LoRA checkpoint directory (contains lora_weights.safetensors and lora_config.json)",
+    )
+    parser.add_argument(
+        "--base_model",
+        type=str,
+        default="",
+        help="Optional: override base model path (default: read from lora_config.json)",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Target text to synthesize",
+    )
+    parser.add_argument(
+        "--prompt_audio",
+        type=str,
+        default="",
+        help="Optional: reference audio path for voice cloning",
+    )
+    parser.add_argument(
+        "--prompt_text",
+        type=str,
+        default="",
+        help="Optional: transcript of reference audio",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="lora_test.wav",
+        help="Output wav file path",
+    )
+    parser.add_argument(
+        "--cfg_value",
+        type=float,
+        default=2.0,
+        help="CFG scale (default: 2.0)",
+    )
+    parser.add_argument(
+        "--inference_timesteps",
+        type=int,
+        default=10,
+        help="Diffusion inference steps (default: 10)",
+    )
+    parser.add_argument(
+        "--max_len",
+        type=int,
+        default=600,
+        help="Max generation steps",
+    )
+    parser.add_argument(
+        "--normalize",
+        action="store_true",
+        help="Enable text normalization",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # 1. Check LoRA checkpoint directory
+    ckpt_dir = Path(args.lora_ckpt)
+    if not ckpt_dir.exists():
+        raise FileNotFoundError(f"LoRA checkpoint not found: {ckpt_dir}")
+
+    # 2. Load lora_config.json from checkpoint
+    lora_config_path = ckpt_dir / "lora_config.json"
+    if not lora_config_path.exists():
+        raise FileNotFoundError(
+            f"lora_config.json not found in {ckpt_dir}. "
+            "Make sure the checkpoint was saved with the updated training script."
+        )
+
+    with open(lora_config_path, "r", encoding="utf-8") as f:
+        lora_info = json.load(f)
+
+    # Get base model path (command line arg overrides config)
+    pretrained_path = args.base_model if args.base_model else lora_info.get("base_model")
+    if not pretrained_path:
+        raise ValueError("base_model not found in lora_config.json and --base_model not provided")
+
+    # Get LoRA config
+    lora_cfg_dict = lora_info.get("lora_config", {})
+    lora_cfg = LoRAConfig(**lora_cfg_dict) if lora_cfg_dict else None
+
+    print(f"Loaded config from: {lora_config_path}", file=sys.stderr)
+    print(f"  Base model: {pretrained_path}", file=sys.stderr)
+    print(
+        f"  LoRA config: r={lora_cfg.r}, alpha={lora_cfg.alpha}" if lora_cfg else "  LoRA config: None", file=sys.stderr
+    )
+
+    # 3. Load model with LoRA (no denoiser)
+    print(f"\n[1/2] Loading model with LoRA: {pretrained_path}", file=sys.stderr)
+    print(f"      LoRA weights: {ckpt_dir}", file=sys.stderr)
+    model = VoxCPM.from_pretrained(
+        hf_model_id=pretrained_path,
+        load_denoiser=False,
+        optimize=True,
+        lora_config=lora_cfg,
+        lora_weights_path=str(ckpt_dir),
+    )
+
+    # 4. Synthesize audio
+    prompt_wav_path = args.prompt_audio if args.prompt_audio else None
+    prompt_text = args.prompt_text if args.prompt_text else None
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+
+    print("\n[2/2] Starting synthesis tests...", file=sys.stderr)
+
+    # === Test 1: With LoRA ===
+    print("\n  [Test 1] Synthesize with LoRA...", file=sys.stderr)
+    audio_np = model.generate(
+        text=args.text,
+        prompt_wav_path=prompt_wav_path,
+        prompt_text=prompt_text,
+        cfg_value=args.cfg_value,
+        inference_timesteps=args.inference_timesteps,
+        max_len=args.max_len,
+        normalize=args.normalize,
+        denoise=False,
+    )
+    lora_output = out_path.with_stem(out_path.stem + "_with_lora")
+    sf.write(str(lora_output), audio_np, model.tts_model.sample_rate)
+    print(
+        f"           Saved: {lora_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
+        file=sys.stderr,
+    )
+
+    # === Test 2: Disable LoRA (via set_lora_enabled) ===
+    print("\n  [Test 2] Disable LoRA (set_lora_enabled=False)...", file=sys.stderr)
+    model.set_lora_enabled(False)
+    audio_np = model.generate(
+        text=args.text,
+        prompt_wav_path=prompt_wav_path,
+        prompt_text=prompt_text,
+        cfg_value=args.cfg_value,
+        inference_timesteps=args.inference_timesteps,
+        max_len=args.max_len,
+        normalize=args.normalize,
+        denoise=False,
+    )
+    disabled_output = out_path.with_stem(out_path.stem + "_lora_disabled")
+    sf.write(str(disabled_output), audio_np, model.tts_model.sample_rate)
+    print(
+        f"           Saved: {disabled_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
+        file=sys.stderr,
+    )
+
+    # === Test 3: Re-enable LoRA ===
+    print("\n  [Test 3] Re-enable LoRA (set_lora_enabled=True)...", file=sys.stderr)
+    model.set_lora_enabled(True)
+    audio_np = model.generate(
+        text=args.text,
+        prompt_wav_path=prompt_wav_path,
+        prompt_text=prompt_text,
+        cfg_value=args.cfg_value,
+        inference_timesteps=args.inference_timesteps,
+        max_len=args.max_len,
+        normalize=args.normalize,
+        denoise=False,
+    )
+    reenabled_output = out_path.with_stem(out_path.stem + "_lora_reenabled")
+    sf.write(str(reenabled_output), audio_np, model.tts_model.sample_rate)
+    print(
+        f"           Saved: {reenabled_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
+        file=sys.stderr,
+    )
+
+    # === Test 4: Unload LoRA (reset_lora_weights) ===
+    print("\n  [Test 4] Unload LoRA (unload_lora)...", file=sys.stderr)
+    model.unload_lora()
+    audio_np = model.generate(
+        text=args.text,
+        prompt_wav_path=prompt_wav_path,
+        prompt_text=prompt_text,
+        cfg_value=args.cfg_value,
+        inference_timesteps=args.inference_timesteps,
+        max_len=args.max_len,
+        normalize=args.normalize,
+        denoise=False,
+    )
+    reset_output = out_path.with_stem(out_path.stem + "_lora_reset")
+    sf.write(str(reset_output), audio_np, model.tts_model.sample_rate)
+    print(
+        f"           Saved: {reset_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
+        file=sys.stderr,
+    )
+
+    # === Test 5: Hot-reload LoRA (load_lora) ===
+    print("\n  [Test 5] Hot-reload LoRA (load_lora)...", file=sys.stderr)
+    loaded, skipped = model.load_lora(ckpt_dir)
+    print(f"           Reloaded {len(loaded)} parameters", file=sys.stderr)
+    audio_np = model.generate(
+        text=args.text,
+        prompt_wav_path=prompt_wav_path,
+        prompt_text=prompt_text,
+        cfg_value=args.cfg_value,
+        inference_timesteps=args.inference_timesteps,
+        max_len=args.max_len,
+        normalize=args.normalize,
+        denoise=False,
+    )
+    reload_output = out_path.with_stem(out_path.stem + "_lora_reloaded")
+    sf.write(str(reload_output), audio_np, model.tts_model.sample_rate)
+    print(
+        f"           Saved: {reload_output}, duration: {len(audio_np) / model.tts_model.sample_rate:.2f}s",
+        file=sys.stderr,
+    )
+
+    print("\n[Done] All tests completed!", file=sys.stderr)
+    print(f"  - with_lora:      {lora_output}", file=sys.stderr)
+    print(f"  - lora_disabled:  {disabled_output}", file=sys.stderr)
+    print(f"  - lora_reenabled: {reenabled_output}", file=sys.stderr)
+    print(f"  - lora_reset:     {reset_output}", file=sys.stderr)
+    print(f"  - lora_reloaded:  {reload_output}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_voxcpm_finetune.py

@@ -0,0 +1,817 @@
+#!/usr/bin/env python3
+
+import sys
+from pathlib import Path
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root / "src"))
+
+import contextlib
+from typing import Dict
+
+import argbind
+import torch
+from tensorboardX import SummaryWriter
+from torch.optim import AdamW
+from transformers import get_cosine_schedule_with_warmup
+import signal
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+try:
+    from safetensors.torch import save_file
+
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+    print("Warning: safetensors not available, will use pytorch format", file=sys.stderr)
+
+import json
+
+from voxcpm.model import VoxCPMModel, VoxCPM2Model
+from voxcpm.model.voxcpm import LoRAConfig
+from voxcpm.training import (
+    Accelerator,
+    BatchProcessor,
+    TrainingTracker,
+    build_dataloader,
+    load_audio_text_datasets,
+)
+
+
+@argbind.bind(without_prefix=True)
+def train(
+    pretrained_path: str,
+    train_manifest: str,
+    val_manifest: str = "",
+    sample_rate: int = 16_000,
+    batch_size: int = 1,
+    grad_accum_steps: int = 1,
+    num_workers: int = 2,
+    num_iters: int = 100_000,
+    log_interval: int = 100,
+    valid_interval: int = 1_000,
+    save_interval: int = 10_000,
+    learning_rate: float = 1e-4,
+    weight_decay: float = 1e-2,
+    warmup_steps: int = 1_000,
+    max_steps: int = 100_000,
+    max_batch_tokens: int = 0,
+    save_path: str = "checkpoints",
+    tensorboard: str = "",
+    lambdas: Dict[str, float] = {"loss/diff": 1.0, "loss/stop": 1.0},
+    lora: dict = None,
+    config_path: str = "",
+    # Distribution options (for LoRA checkpoints)
+    hf_model_id: str = "",  # HuggingFace model ID (e.g., "openbmb/VoxCPM1.5")
+    distribute: bool = False,  # If True, save hf_model_id as base_model; otherwise save pretrained_path
+):
+    _ = config_path
+
+    # Validate distribution options
+    if lora is not None and distribute and not hf_model_id:
+        raise ValueError("hf_model_id is required when distribute=True")
+
+    accelerator = Accelerator(amp=True)
+
+    save_dir = Path(save_path)
+    tb_dir = Path(tensorboard) if tensorboard else save_dir / "logs"
+
+    # Only create directories on rank 0 to avoid race conditions
+    if accelerator.rank == 0:
+        save_dir.mkdir(parents=True, exist_ok=True)
+        tb_dir.mkdir(parents=True, exist_ok=True)
+    accelerator.barrier()  # Wait for directory creation
+
+    writer = SummaryWriter(log_dir=str(tb_dir)) if accelerator.rank == 0 else None
+    tracker = TrainingTracker(writer=writer, log_file=str(save_dir / "train.log"), rank=accelerator.rank)
+
+    # Auto-detect model architecture from config.json
+    with open(os.path.join(pretrained_path, "config.json"), "r", encoding="utf-8") as _f:
+        _arch = json.load(_f).get("architecture", "voxcpm").lower()
+    _model_cls = VoxCPM2Model if _arch == "voxcpm2" else VoxCPMModel
+    if accelerator.rank == 0:
+        print(f"Detected architecture: {_arch} -> {_model_cls.__name__}", file=sys.stderr)
+    base_model = _model_cls.from_local(
+        pretrained_path, optimize=False, training=True, lora_config=LoRAConfig(**lora) if lora else None
+    )
+    tokenizer = base_model.text_tokenizer
+
+    train_ds, val_ds = load_audio_text_datasets(
+        train_manifest=train_manifest,
+        val_manifest=val_manifest,
+        sample_rate=sample_rate,
+    )
+
+    def tokenize(batch):
+        text_list = batch["text"]
+        text_ids = [tokenizer(text) for text in text_list]
+        return {"text_ids": text_ids}
+
+    train_ds = train_ds.map(tokenize, batched=True, remove_columns=["text"])
+    # Save original validation texts for audio generation display
+    val_texts = None
+    if val_ds is not None:
+        val_texts = list(val_ds["text"])  # Save original texts
+        val_ds = val_ds.map(tokenize, batched=True, remove_columns=["text"])
+
+    dataset_cnt = int(max(train_ds["dataset_id"])) + 1 if "dataset_id" in train_ds.column_names else 1
+    num_train_samples = len(train_ds)
+
+    # ------------------------------------------------------------------ #
+    # Optional: filter samples by estimated token count to avoid OOM
+    # Enabled when max_batch_tokens > 0:
+    #   max_sample_len = max_batch_tokens // batch_size
+    #   Samples exceeding this length will be dropped
+    # ------------------------------------------------------------------ #
+    if max_batch_tokens and max_batch_tokens > 0:
+        from voxcpm.training.data import compute_sample_lengths
+
+        audio_vae_fps = base_model.audio_vae.sample_rate / base_model.audio_vae.hop_length
+        est_lengths = compute_sample_lengths(
+            train_ds,
+            audio_vae_fps=audio_vae_fps,
+            patch_size=base_model.config.patch_size,
+        )
+        max_sample_len = max_batch_tokens // batch_size if batch_size > 0 else max(est_lengths)
+        keep_indices = [i for i, L in enumerate(est_lengths) if L <= max_sample_len]
+
+        if len(keep_indices) < len(train_ds) and accelerator.rank == 0:
+            tracker.print(
+                f"Filtering {len(train_ds) - len(keep_indices)} / {len(train_ds)} "
+                f"training samples longer than {max_sample_len} tokens "
+                f"(max_batch_tokens={max_batch_tokens})."
+            )
+        train_ds = train_ds.select(keep_indices)
+
+    train_loader = build_dataloader(
+        train_ds,
+        accelerator=accelerator,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        drop_last=True,
+    )
+    val_loader = (
+        build_dataloader(
+            val_ds,
+            accelerator=accelerator,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            drop_last=False,
+        )
+        if val_ds is not None
+        else None
+    )
+
+    batch_processor = BatchProcessor(
+        config=base_model.config,
+        audio_vae=base_model.audio_vae,
+        dataset_cnt=dataset_cnt,
+        device=accelerator.device,
+    )
+    # Save audio_vae for audio generation
+    audio_vae_for_gen = base_model.audio_vae
+    del base_model.audio_vae
+    model = accelerator.prepare_model(base_model)
+    unwrapped_model = accelerator.unwrap(model)
+    unwrapped_model.train()
+
+    # Only print param info on rank 0 to avoid cluttered output
+    if accelerator.rank == 0:
+        for name, param in model.named_parameters():
+            print(name, param.requires_grad, file=sys.stderr)
+
+    optimizer = AdamW(
+        (p for p in model.parameters() if p.requires_grad),
+        lr=learning_rate,
+        weight_decay=weight_decay,
+    )
+
+    # Cosine + warmup scheduler from transformers:
+    # - num_warmup_steps: warmup steps
+    # - num_training_steps: total training steps (outer step count)
+    total_training_steps = max_steps if max_steps > 0 else num_iters
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=warmup_steps,
+        num_training_steps=total_training_steps,
+    )
+
+    # All ranks load the same checkpoint to keep model and optimizer state in sync.
+    start_step = load_checkpoint(model, optimizer, scheduler, save_dir, rank=accelerator.rank)
+    accelerator.barrier()
+
+    if start_step > 0 and accelerator.rank == 0:
+        tracker.print(f"Resuming training from step {start_step}")
+
+    # Resume tracker for signal handler to read current step
+    resume = {"step": start_step}
+
+    # Register signal handler to save checkpoint on termination (SIGTERM/SIGINT)
+    def _signal_handler(
+        signum,
+        frame,
+        _model=model,
+        _optim=optimizer,
+        _sched=scheduler,
+        _save_dir=save_dir,
+        _pretrained=pretrained_path,
+        _hf_id=hf_model_id,
+        _dist=distribute,
+        _resume=resume,
+        _rank=accelerator.rank,
+    ):
+        try:
+            cur_step = int(_resume.get("step", start_step))
+        except Exception:
+            cur_step = start_step
+        if _rank == 0:
+            print(f"Signal {signum} received. Saving checkpoint at step {cur_step} ...", file=sys.stderr)
+            try:
+                save_checkpoint(_model, _optim, _sched, _save_dir, cur_step, _pretrained, _hf_id, _dist)
+                print("Checkpoint saved. Exiting.", file=sys.stderr)
+            except Exception as e:
+                print(f"Error saving checkpoint on signal: {e}", file=sys.stderr)
+        os._exit(0)
+
+    signal.signal(signal.SIGTERM, _signal_handler)
+    signal.signal(signal.SIGINT, _signal_handler)
+
+    # Manual epoch management instead of itertools.cycle to support DistributedSampler.set_epoch()
+    grad_accum_steps = max(int(grad_accum_steps), 1)
+    data_epoch = 0
+    train_iter = iter(train_loader)
+
+    def get_next_batch():
+        """Get next batch, handles epoch boundary and DistributedSampler."""
+        nonlocal train_iter, data_epoch
+        try:
+            return next(train_iter)
+        except StopIteration:
+            data_epoch += 1
+            # Key: set DistributedSampler epoch to ensure different data order each epoch
+            sampler = getattr(train_loader, "sampler", None)
+            if hasattr(sampler, "set_epoch"):
+                sampler.set_epoch(data_epoch)
+            train_iter = iter(train_loader)
+            return next(train_iter)
+
+    with tracker.live():
+        for step in range(start_step, num_iters):
+            # update resume step so signal handler can save current progress
+            resume["step"] = step
+            tracker.step = step
+            optimizer.zero_grad(set_to_none=True)
+
+            # Gradient accumulation: accumulate gradients over micro-batches before optimizer step
+            loss_dict = {}
+            for micro_step in range(grad_accum_steps):
+                batch = get_next_batch()
+                processed = batch_processor(batch)
+
+                # Only sync gradients on the last micro-batch
+                # Use no_sync() for intermediate steps to reduce communication overhead
+                is_last_micro_step = micro_step == grad_accum_steps - 1
+                sync_context = contextlib.nullcontext() if is_last_micro_step else accelerator.no_sync()
+
+                with sync_context:
+                    with accelerator.autocast(dtype=torch.bfloat16):
+                        outputs = model(
+                            processed["text_tokens"],
+                            processed["text_mask"],
+                            processed["audio_feats"],
+                            processed["audio_mask"],
+                            processed["loss_mask"],
+                            processed["position_ids"],
+                            processed["labels"],
+                            progress=step / max(1, num_iters),
+                        )
+
+                    total_loss = 0.0
+                    for key, value in outputs.items():
+                        if key.startswith("loss/"):
+                            weight = lambdas.get(key, 1.0)
+                            loss_value = value * weight / grad_accum_steps
+                            total_loss = total_loss + loss_value
+                            # Record raw loss from last micro-batch for logging
+                            loss_dict[key] = value.detach()
+
+                    # Accumulate gradients (normalized by grad_accum_steps)
+                    accelerator.backward(total_loss)
+
+            # After all micro-batches, do unscale / grad_norm / step
+            scaler = getattr(accelerator, "scaler", None)
+            if scaler is not None:
+                scaler.unscale_(optimizer)
+            # Use large max_norm to only compute grad_norm without actual clipping
+            grad_norm = torch.nn.utils.clip_grad_norm_(unwrapped_model.parameters(), max_norm=1e9)
+
+            accelerator.step(optimizer)
+            accelerator.update()
+            scheduler.step()
+
+            if step % log_interval == 0 or step == num_iters - 1:
+                loss_values = {k: v.item() if isinstance(v, torch.Tensor) else float(v) for k, v in loss_dict.items()}
+                loss_values["lr"] = float(optimizer.param_groups[0]["lr"])
+                # Account for all GPUs when converting steps to epochs.
+                epoch = (step * grad_accum_steps * batch_size * accelerator.world_size) / max(1, num_train_samples)
+                loss_values["epoch"] = float(epoch)
+                loss_values["grad_norm"] = float(grad_norm)
+                tracker.log_metrics(loss_values, split="train")
+
+            if val_loader is not None and (step % valid_interval == 0 or step == num_iters - 1):
+                validate(
+                    model,
+                    val_loader,
+                    batch_processor,
+                    accelerator,
+                    tracker,
+                    lambdas,
+                    writer=writer,
+                    step=step,
+                    val_ds=val_ds,
+                    audio_vae=audio_vae_for_gen,
+                    sample_rate=sample_rate,
+                    val_texts=val_texts,
+                    tokenizer=tokenizer,
+                    valid_interval=valid_interval,
+                )
+
+            if (step % save_interval == 0 or step == num_iters - 1) and accelerator.rank == 0:
+                save_checkpoint(model, optimizer, scheduler, save_dir, step, pretrained_path, hf_model_id, distribute)
+
+    if accelerator.rank == 0:
+        save_checkpoint(model, optimizer, scheduler, save_dir, num_iters, pretrained_path, hf_model_id, distribute)
+    if writer:
+        writer.close()
+
+
+def validate(
+    model,
+    val_loader,
+    batch_processor,
+    accelerator,
+    tracker,
+    lambdas,
+    writer=None,
+    step=0,
+    val_ds=None,
+    audio_vae=None,
+    sample_rate=22050,
+    val_texts=None,
+    tokenizer=None,
+    valid_interval=1000,
+):
+    """Validate and generate sample audio"""
+    import numpy as np  # noqa: F401
+    from collections import defaultdict
+
+    model.eval()
+    total_losses = []
+    sub_losses = defaultdict(list)  # Track individual sub-losses
+    num_batches = 0
+    max_val_batches = 10
+
+    with torch.no_grad():
+        for batch in val_loader:
+            if num_batches >= max_val_batches:
+                break
+            processed = batch_processor(batch)
+            with accelerator.autocast(dtype=torch.bfloat16):
+                outputs = model(
+                    processed["text_tokens"],
+                    processed["text_mask"],
+                    processed["audio_feats"],
+                    processed["audio_mask"],
+                    processed["loss_mask"],
+                    processed["position_ids"],
+                    processed["labels"],
+                    progress=0.0,
+                    sample_generate=False,
+                )
+            total = 0.0
+            for key, value in outputs.items():
+                if key.startswith("loss/"):
+                    weighted_loss = lambdas.get(key, 1.0) * value
+                    total += weighted_loss
+                    sub_losses[key].append(value.detach())
+            total_losses.append(total.detach())
+            num_batches += 1
+
+    if total_losses:
+        # Compute mean total loss
+        mean_total_loss = torch.stack(total_losses).mean()
+        accelerator.all_reduce(mean_total_loss)
+
+        # Compute mean of each sub-loss
+        val_metrics = {"loss/total": mean_total_loss.item()}
+        for key, values in sub_losses.items():
+            mean_sub_loss = torch.stack(values).mean()
+            accelerator.all_reduce(mean_sub_loss)
+            val_metrics[key] = mean_sub_loss.item()
+
+        tracker.log_metrics(val_metrics, split="val")
+
+    # Generate sample audio for TensorBoard display
+    if writer is not None and val_ds is not None and audio_vae is not None and accelerator.rank == 0:
+        try:
+            generate_sample_audio(
+                model,
+                val_ds,
+                audio_vae,
+                writer,
+                step,
+                accelerator,
+                sample_rate,
+                val_texts=val_texts,
+                tokenizer=tokenizer,
+                valid_interval=valid_interval,
+                tracker=tracker,
+            )
+        except Exception as e:
+            tracker.print(f"[Warning] Failed to generate sample audio: {e}")
+            import traceback
+            import io
+
+            buf = io.StringIO()
+            traceback.print_exc(file=buf)
+            tracker.print(buf.getvalue())
+    else:
+        # Log why audio generation was skipped
+        missing = []
+        if writer is None:
+            missing.append("writer")
+        if val_ds is None:
+            missing.append("val_ds")
+        if audio_vae is None:
+            missing.append("audio_vae")
+        if missing and accelerator.rank == 0:
+            tracker.print(f"[Warning] Skip audio generation: missing {', '.join(missing)}")
+
+    model.train()
+
+
+def compute_mel_spectrogram(audio_np, sample_rate, n_mels=128):
+    """Compute Mel Spectrogram (dB) using librosa"""
+    import numpy as np
+    import librosa
+
+    audio_np = audio_np.flatten().astype(np.float32)
+    mel = librosa.feature.melspectrogram(y=audio_np, sr=sample_rate, n_mels=n_mels, fmax=sample_rate // 2)
+    return librosa.power_to_db(mel, ref=np.max)
+
+
+def create_mel_figure(gen_audio_np, gen_mel, sample_rate, step=None, ref_audio_np=None, ref_mel=None):
+    """
+    Create mel spectrogram figure: show comparison if reference audio exists, otherwise show generated only
+    """
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import librosa.display
+
+    fmax = sample_rate // 2
+    step_str = f" @ Step {step}" if step is not None else ""
+
+    if ref_audio_np is not None and ref_mel is not None:
+        # Comparison mode: reference vs generated
+        fig, (ax_ref, ax_gen) = plt.subplots(2, 1, figsize=(12, 8))
+
+        img_ref = librosa.display.specshow(
+            ref_mel, sr=sample_rate, x_axis="time", y_axis="mel", fmax=fmax, cmap="viridis", ax=ax_ref
+        )
+        ax_ref.set_title(
+            f"Reference (GT) - {len(ref_audio_np)/sample_rate:.2f}s{step_str}",
+            fontsize=10,
+            fontweight="bold",
+            color="#28A745",
+        )
+        plt.colorbar(img_ref, ax=ax_ref, format="%+2.0f dB", pad=0.02)
+
+        img_gen = librosa.display.specshow(
+            gen_mel, sr=sample_rate, x_axis="time", y_axis="mel", fmax=fmax, cmap="viridis", ax=ax_gen
+        )
+        ax_gen.set_title(
+            f"Generated - {len(gen_audio_np)/sample_rate:.2f}s", fontsize=10, fontweight="bold", color="#DC3545"
+        )
+        plt.colorbar(img_gen, ax=ax_gen, format="%+2.0f dB", pad=0.02)
+    else:
+        # Single figure mode: show generated only
+        fig, ax = plt.subplots(figsize=(12, 4))
+        img = librosa.display.specshow(
+            gen_mel, sr=sample_rate, x_axis="time", y_axis="mel", fmax=fmax, cmap="viridis", ax=ax
+        )
+        ax.set_title(f"Generated - {len(gen_audio_np)/sample_rate:.2f}s{step_str}", fontsize=11, fontweight="bold")
+        plt.colorbar(img, ax=ax, format="%+2.0f dB", pad=0.02)
+
+    plt.tight_layout()
+    return fig
+
+
+def normalize_audio(audio_np):
+    """Normalize audio to [-0.9, 0.9]"""
+    import numpy as np
+
+    max_val = np.abs(audio_np).max()
+    return audio_np / max_val * 0.9 if max_val > 0 else audio_np
+
+
+def generate_sample_audio(
+    model,
+    val_ds,
+    audio_vae,
+    writer,
+    step,
+    accelerator,
+    sample_rate=22050,
+    val_texts=None,
+    tokenizer=None,
+    pretrained_path=None,
+    valid_interval=1000,
+    tracker=None,
+):
+    """Select 2 fixed validation samples, generate audio and log to TensorBoard"""
+    import numpy as np
+
+    log = tracker.print if tracker else print
+    num_samples = min(2, len(val_ds))
+    log(f"[Audio] Starting audio generation for {num_samples} samples at step {step}")
+
+    unwrapped_model = accelerator.unwrap(model)
+
+    for i in range(num_samples):
+        sample = val_ds[i]
+        text = val_texts[i] if val_texts and i < len(val_texts) else "Hello, this is a test."
+
+        # Load reference audio
+        ref_audio_np = None
+        try:
+            if "audio" in sample and isinstance(sample["audio"], dict) and "array" in sample["audio"]:
+                ref_audio_np = np.array(sample["audio"]["array"], dtype=np.float32)
+                ref_sr = sample["audio"].get("sampling_rate", sample_rate)
+                if ref_sr != sample_rate:
+                    import torchaudio.functional as F
+
+                    ref_audio_np = (
+                        F.resample(torch.from_numpy(ref_audio_np).unsqueeze(0), ref_sr, sample_rate).squeeze(0).numpy()
+                    )
+                log(f"[Audio] Loaded reference audio for sample {i}: duration={len(ref_audio_np)/sample_rate:.2f}s")
+        except Exception as e:
+            log(f"[Warning] Failed to load reference audio: {e}")
+
+        # Preserve the original mode so validation failures do not leak into training.
+        prev_training = unwrapped_model.training
+        try:
+            # Inference setup
+            unwrapped_model.eval()
+            # unwrapped_model.to(torch.bfloat16)
+            unwrapped_model.audio_vae = audio_vae.to(torch.float32)
+
+            log(f"[Audio] Generating sample {i} with text: '{text[:50]}...'")
+            autocast_ctx = (
+                torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+                if torch.cuda.is_available()
+                else contextlib.nullcontext()
+            )
+            with torch.no_grad():
+                with autocast_ctx:
+                    generated = unwrapped_model.generate(target_text=text, inference_timesteps=10, cfg_value=2.0)
+
+            # Restore training setup
+            # unwrapped_model.to(torch.float32)
+            # unwrapped_model.audio_vae = None
+
+            if generated is None or len(generated) == 0:
+                log(f"[Warning] Generated audio is empty for sample {i}")
+                continue
+
+            # Process generated audio
+            gen_audio_np = (
+                generated.cpu().float().numpy().flatten()
+                if isinstance(generated, torch.Tensor)
+                else np.array(generated, dtype=np.float32).flatten()
+            )
+            gen_audio_np = normalize_audio(gen_audio_np)
+
+            tag = f"val_sample_{i}"
+            writer.add_audio(f"{tag}/generated_audio", gen_audio_np, global_step=step, sample_rate=sample_rate)
+            log(f"[Audio] Generated audio for sample {i}: duration={len(gen_audio_np)/sample_rate:.2f}s")
+
+            # Log reference audio
+            if ref_audio_np is not None:
+                writer.add_audio(
+                    f"{tag}/reference_audio", normalize_audio(ref_audio_np), global_step=step, sample_rate=sample_rate
+                )
+
+            # Generate mel spectrogram figure
+            try:
+                mel_gen = compute_mel_spectrogram(gen_audio_np, sample_rate)
+                mel_ref = compute_mel_spectrogram(ref_audio_np, sample_rate) if ref_audio_np is not None else None
+                fig = create_mel_figure(gen_audio_np, mel_gen, sample_rate, step, ref_audio_np, mel_ref)
+                writer.add_figure(f"{tag}/mel_spectrogram", fig, global_step=step)
+                log(f"[Audio] Created mel spectrogram figure for sample {i}")
+            except Exception as e:
+                log(f"[Warning] Failed to create mel spectrogram: {e}")
+
+        except Exception as e:
+            log(f"[Warning] Failed to generate audio for sample {i}: {e}")
+            import traceback
+
+            traceback.print_exc()
+
+        finally:
+            # Always restore the training state, even if generation fails.
+            try:
+                # unwrapped_model.to(torch.float32)
+                unwrapped_model.audio_vae = None
+                if prev_training:
+                    unwrapped_model.train()
+                else:
+                    unwrapped_model.eval()
+            except Exception as e:
+                log(f"[Warning] Failed to restore model state: {e}")
+
+
+def load_checkpoint(model, optimizer, scheduler, save_dir: Path, rank: int = 0):
+    """
+    Load the latest checkpoint if it exists.
+    Called by all ranks so that distributed state stays aligned.
+    Returns the step number to resume from, or 0 if no checkpoint found.
+    """
+    latest_folder = save_dir / "latest"
+    if not latest_folder.exists():
+        return 0
+
+    unwrapped = model.module if hasattr(model, "module") else model
+    lora_cfg = unwrapped.lora_config
+
+    # Load model weights
+    if lora_cfg is not None:
+        # LoRA: load lora_weights
+        lora_weights_path = latest_folder / "lora_weights.safetensors"
+        if not lora_weights_path.exists():
+            lora_weights_path = latest_folder / "lora_weights.ckpt"
+
+        if lora_weights_path.exists():
+            if lora_weights_path.suffix == ".safetensors":
+                from safetensors.torch import load_file
+
+                state_dict = load_file(str(lora_weights_path))
+            else:
+                ckpt = torch.load(lora_weights_path, map_location="cpu")
+                state_dict = ckpt.get("state_dict", ckpt)
+
+            unwrapped.load_state_dict(state_dict, strict=False)
+            if rank == 0:
+                print(f"Loaded LoRA weights from {lora_weights_path}", file=sys.stderr)
+    else:
+        # Full finetune: load model.safetensors or pytorch_model.bin
+        model_path = latest_folder / "model.safetensors"
+        if not model_path.exists():
+            model_path = latest_folder / "pytorch_model.bin"
+
+        if model_path.exists():
+            if model_path.suffix == ".safetensors":
+                from safetensors.torch import load_file
+
+                state_dict = load_file(str(model_path))
+            else:
+                ckpt = torch.load(model_path, map_location="cpu")
+                state_dict = ckpt.get("state_dict", ckpt)
+
+            unwrapped.load_state_dict(state_dict, strict=False)
+            if rank == 0:
+                print(f"Loaded model weights from {model_path}", file=sys.stderr)
+
+    # Load optimizer state
+    optimizer_path = latest_folder / "optimizer.pth"
+    if optimizer_path.exists():
+        optimizer.load_state_dict(torch.load(optimizer_path, map_location="cpu"))
+        if rank == 0:
+            print(f"Loaded optimizer state from {optimizer_path}", file=sys.stderr)
+
+    # Load scheduler state
+    scheduler_path = latest_folder / "scheduler.pth"
+    if scheduler_path.exists():
+        scheduler.load_state_dict(torch.load(scheduler_path, map_location="cpu"))
+        if rank == 0:
+            print(f"Loaded scheduler state from {scheduler_path}", file=sys.stderr)
+
+    state_path = latest_folder / "training_state.json"
+    if state_path.exists():
+        with open(state_path, "r", encoding="utf-8") as f:
+            state = json.load(f)
+        resume_step = int(state.get("step", 0))
+        if rank == 0:
+            print(f"Resuming from step {resume_step}", file=sys.stderr)
+        return resume_step
+
+    # Fallback for older checkpoints without metadata.
+    step_folders = [d for d in save_dir.iterdir() if d.is_dir() and d.name.startswith("step_")]
+    if step_folders:
+        steps = [int(d.name.split("_")[1]) for d in step_folders]
+        resume_step = max(steps)
+        if rank == 0:
+            print(f"Resuming from step {resume_step}", file=sys.stderr)
+        return resume_step
+
+    return 0
+
+
+def save_checkpoint(
+    model,
+    optimizer,
+    scheduler,
+    save_dir: Path,
+    step: int,
+    pretrained_path: str = None,
+    hf_model_id: str = "",
+    distribute: bool = False,
+):
+    """
+    Save checkpoint with different strategies for full finetune vs LoRA:
+    - Full finetune: save non-vae weights to model.safetensors (or pytorch_model.bin if safetensors unavailable)
+    - LoRA: save only lora weights to lora_weights.safetensors (or lora_weights.ckpt if safetensors unavailable)
+    """
+    import shutil
+
+    save_dir.mkdir(parents=True, exist_ok=True)
+    tag = f"step_{step:07d}"
+    folder = save_dir / tag
+    folder.mkdir(parents=True, exist_ok=True)
+
+    unwrapped = model.module if hasattr(model, "module") else model
+    full_state = unwrapped.state_dict()
+    lora_cfg = unwrapped.lora_config
+
+    if lora_cfg is not None:
+        # LoRA finetune: save only lora_A/lora_B weights
+        state_dict = {k: v for k, v in full_state.items() if "lora_" in k}
+        if SAFETENSORS_AVAILABLE:
+            save_file(state_dict, folder / "lora_weights.safetensors")
+        else:
+            torch.save({"state_dict": state_dict}, folder / "lora_weights.ckpt")
+
+        # Save LoRA config and base model path to a separate JSON file
+        # If distribute=True, save hf_model_id; otherwise save local pretrained_path
+        base_model_to_save = hf_model_id if distribute else (str(pretrained_path) if pretrained_path else None)
+        lora_info = {
+            "base_model": base_model_to_save,
+            "lora_config": lora_cfg.model_dump() if hasattr(lora_cfg, "model_dump") else vars(lora_cfg),
+        }
+        with open(folder / "lora_config.json", "w", encoding="utf-8") as f:
+            json.dump(lora_info, f, indent=2, ensure_ascii=False)
+    else:
+        # Full finetune: save non-vae weights to model.safetensors
+        state_dict = {k: v for k, v in full_state.items() if not k.startswith("audio_vae.")}
+        if SAFETENSORS_AVAILABLE:
+            save_file(state_dict, folder / "model.safetensors")
+        else:
+            torch.save({"state_dict": state_dict}, folder / "pytorch_model.bin")
+
+        # Copy config files from pretrained path
+        if pretrained_path:
+            pretrained_dir = Path(pretrained_path)
+            files_to_copy = [
+                "config.json",
+                "audiovae.pth",
+                "audiovae.safetensors",
+                "tokenizer.json",
+                "special_tokens_map.json",
+                "tokenizer_config.json",
+            ]
+            for fname in files_to_copy:
+                src = pretrained_dir / fname
+                if src.exists():
+                    shutil.copy2(src, folder / fname)
+
+    torch.save(optimizer.state_dict(), folder / "optimizer.pth")
+    torch.save(scheduler.state_dict(), folder / "scheduler.pth")
+    with open(folder / "training_state.json", "w", encoding="utf-8") as f:
+        json.dump({"step": int(step)}, f)
+
+    # Update (or create) a `latest` folder by copying the most recent checkpoint
+    latest_link = save_dir / "latest"
+    try:
+        if latest_link.exists():
+            shutil.rmtree(latest_link)
+        shutil.copytree(folder, latest_link)
+    except Exception:
+        print(f"Warning: failed to update latest checkpoint at {latest_link}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    from voxcpm.training.config import load_yaml_config
+
+    args = argbind.parse_args()
+    config_file = args.get("config_path")
+    # If YAML config provided, use YAML args to call train
+    if config_file:
+        yaml_args = load_yaml_config(config_file)
+        train(**yaml_args)
+    else:
+        # Otherwise use command line args (parsed by argbind)
+        with argbind.scope(args):
+            train()

src/voxcpm/cli.py

@@ -2,290 +2,596 @@
 """
 VoxCPM Command Line Interface
 
-Unified CLI for voice cloning, direct TTS synthesis, and batch processing.
-
-Usage examples:
-    # Direct synthesis (single sample)
-    voxcpm --text "Hello world" --output output.wav
-
-    # Voice cloning (with reference audio and text)
-    voxcpm --text "Hello world" --prompt-audio voice.wav --prompt-text "reference text" --output output.wav --denoise
-
-    # Batch processing (each line in the file is one sample)
-    voxcpm --input texts.txt --output-dir ./outputs/
+VoxCPM2-first CLI for voice design, cloning, and batch processing.
 """
 
 import argparse
+import json
 import os
 import sys
 from pathlib import Path
-from typing import Optional, List
+
 import soundfile as sf
 
 from voxcpm.core import VoxCPM
 
 
+DEFAULT_HF_MODEL_ID = "openbmb/VoxCPM2"
+
+# -----------------------------
+# Validators
+# -----------------------------
+
+
 def validate_file_exists(file_path: str, file_type: str = "file") -> Path:
-    """Validate that a file exists."""
     path = Path(file_path)
     if not path.exists():
         raise FileNotFoundError(f"{file_type} '{file_path}' does not exist")
     return path
 
 
+def require_file_exists(file_path: str, parser, file_type: str = "file") -> Path:
+    try:
+        return validate_file_exists(file_path, file_type)
+    except FileNotFoundError as exc:
+        parser.error(str(exc))
+
+
 def validate_output_path(output_path: str) -> Path:
-    """Validate the output path and create parent directories if needed."""
     path = Path(output_path)
     path.parent.mkdir(parents=True, exist_ok=True)
     return path
 
 
+def validate_ranges(args, parser):
+    """Validate numeric argument ranges."""
+    if not (0.1 <= args.cfg_value <= 10.0):
+        parser.error("--cfg-value must be between 0.1 and 10.0 (recommended: 1.0–3.0)")
+
+    if not (1 <= args.inference_timesteps <= 100):
+        parser.error("--inference-timesteps must be between 1 and 100 (recommended: 4–30)")
+
+    if args.lora_r <= 0:
+        parser.error("--lora-r must be a positive integer")
+
+    if args.lora_alpha <= 0:
+        parser.error("--lora-alpha must be a positive integer")
+
+    if not (0.0 <= args.lora_dropout <= 1.0):
+        parser.error("--lora-dropout must be between 0.0 and 1.0")
+
+
+def warn_legacy_mode():
+    print(
+        "Warning: legacy root CLI arguments are deprecated. Prefer `voxcpm design|clone|batch ...`.",
+        file=sys.stderr,
+    )
+
+
+def build_final_text(text: str, control: str | None) -> str:
+    control = (control or "").strip()
+    return f"({control}){text}" if control else text
+
+
+def resolve_prompt_text(args, parser) -> str | None:
+    prompt_text = getattr(args, "prompt_text", None)
+    prompt_file = getattr(args, "prompt_file", None)
+
+    if prompt_text and prompt_file:
+        parser.error("Use either --prompt-text or --prompt-file, not both.")
+
+    if prompt_file:
+        prompt_path = require_file_exists(prompt_file, parser, "prompt text file")
+        return prompt_path.read_text(encoding="utf-8").strip()
+
+    if prompt_text:
+        return prompt_text.strip()
+
+    return None
+
+
+def detect_model_architecture(args) -> str | None:
+    model_location = getattr(args, "model_path", None) or getattr(
+        args, "hf_model_id", None
+    )
+    if not model_location:
+        return None
+
+    if os.path.isdir(model_location):
+        config_path = Path(model_location) / "config.json"
+        if not config_path.exists():
+            return None
+
+        with open(config_path, "r", encoding="utf-8") as f:
+            return json.load(f).get("architecture", "voxcpm").lower()
+
+    model_hint = str(model_location).lower()
+    if "voxcpm2" in model_hint:
+        return "voxcpm2"
+    if (
+        "voxcpm1.5" in model_hint
+        or "voxcpm-1.5" in model_hint
+        or "voxcpm_1.5" in model_hint
+    ):
+        return "voxcpm"
+
+    return None
+
+
+def validate_prompt_related_args(args, parser, prompt_text: str | None):
+    if prompt_text and not args.prompt_audio:
+        parser.error("--prompt-text/--prompt-file requires --prompt-audio.")
+
+    if args.prompt_audio and not prompt_text:
+        parser.error("--prompt-audio requires --prompt-text or --prompt-file.")
+
+    if args.control and prompt_text:
+        parser.error(
+            "--control cannot be used together with --prompt-text or --prompt-file."
+        )
+
+
+def validate_reference_support(args, parser):
+    if not getattr(args, "reference_audio", None):
+        return
+
+    arch = detect_model_architecture(args)
+    if arch == "voxcpm":
+        parser.error("--reference-audio is only supported with VoxCPM2 models.")
+
+
+def validate_design_args(args, parser):
+    prompt_text = resolve_prompt_text(args, parser)
+    if args.prompt_audio or args.reference_audio or prompt_text:
+        parser.error(
+            "`design` does not accept prompt/reference audio. Use `clone` instead."
+        )
+
+
+def validate_clone_args(args, parser):
+    prompt_text = resolve_prompt_text(args, parser)
+    validate_prompt_related_args(args, parser, prompt_text)
+    validate_reference_support(args, parser)
+
+    if not args.prompt_audio and not args.reference_audio:
+        parser.error(
+            "`clone` requires --reference-audio, or --prompt-audio with --prompt-text/--prompt-file."
+        )
+
+    return prompt_text
+
+
+def validate_batch_args(args, parser):
+    prompt_text = resolve_prompt_text(args, parser)
+    validate_prompt_related_args(args, parser, prompt_text)
+    validate_reference_support(args, parser)
+    return prompt_text
+
+
+# -----------------------------
+# Model loading
+# -----------------------------
+
+
 def load_model(args) -> VoxCPM:
-    """Load VoxCPM model.
+    print("Loading VoxCPM model...", file=sys.stderr)
 
-    Prefer --model-path if provided; otherwise use from_pretrained (Hub).
-    """
-    print("Loading VoxCPM model...")
-
-    # 兼容旧参数：ZIPENHANCER_MODEL_PATH 环境变量作为默认
     zipenhancer_path = getattr(args, "zipenhancer_path", None) or os.environ.get(
         "ZIPENHANCER_MODEL_PATH", None
     )
 
-    # Load from local path if provided
-    if getattr(args, "model_path", None):
+    # Build LoRA config if provided
+    lora_config = None
+    lora_weights_path = getattr(args, "lora_path", None)
+    if lora_weights_path:
+        from voxcpm.model.voxcpm import LoRAConfig
+
+        lora_config = LoRAConfig(
+            enable_lm=not args.lora_disable_lm,
+            enable_dit=not args.lora_disable_dit,
+            enable_proj=args.lora_enable_proj,
+            r=args.lora_r,
+            alpha=args.lora_alpha,
+            dropout=args.lora_dropout,
+        )
+
+        print(
+            f"LoRA config: r={lora_config.r}, alpha={lora_config.alpha}, "
+            f"lm={lora_config.enable_lm}, dit={lora_config.enable_dit}, proj={lora_config.enable_proj}",
+            file=sys.stderr,
+        )
+
+    # Load local model if specified
+    if args.model_path:
         try:
             model = VoxCPM(
                 voxcpm_model_path=args.model_path,
                 zipenhancer_model_path=zipenhancer_path,
-                enable_denoiser=not getattr(args, "no_denoiser", False),
+                enable_denoiser=not args.no_denoiser,
+                optimize=not args.no_optimize,
+                lora_config=lora_config,
+                lora_weights_path=lora_weights_path,
             )
-            print("Model loaded (local).")
+            print("Model loaded (local).", file=sys.stderr)
             return model
         except Exception as e:
-            print(f"Failed to load model (local): {e}")
+            print(f"Failed to load model (local): {e}", file=sys.stderr)
             sys.exit(1)
 
-    # Otherwise, try from_pretrained (Hub); exit on failure
+    # Load from Hugging Face Hub
     try:
         model = VoxCPM.from_pretrained(
-            hf_model_id=getattr(args, "hf_model_id", "openbmb/VoxCPM-0.5B"),
-            load_denoiser=not getattr(args, "no_denoiser", False),
+            hf_model_id=args.hf_model_id,
+            load_denoiser=not args.no_denoiser,
             zipenhancer_model_id=zipenhancer_path,
-            cache_dir=getattr(args, "cache_dir", None),
-            local_files_only=getattr(args, "local_files_only", False),
+            cache_dir=args.cache_dir,
+            local_files_only=args.local_files_only,
+            optimize=not args.no_optimize,
+            lora_config=lora_config,
+            lora_weights_path=lora_weights_path,
         )
-        print("Model loaded (from_pretrained).")
+        print("Model loaded (from_pretrained).", file=sys.stderr)
         return model
     except Exception as e:
-        print(f"Failed to load model (from_pretrained): {e}")
+        print(f"Failed to load model (from_pretrained): {e}", file=sys.stderr)
         sys.exit(1)
 
 
-def cmd_clone(args):
-    """Voice cloning command."""
-    # Validate inputs
-    if not args.text:
-        print("Error: Please provide text to synthesize (--text)")
-        sys.exit(1)
-    
-    if not args.prompt_audio:
-        print("Error: Voice cloning requires a reference audio (--prompt-audio)")
-        sys.exit(1)
-        
-    if not args.prompt_text:
-        print("Error: Voice cloning requires a reference text (--prompt-text)")
-        sys.exit(1)
-    
-    # Validate files
-    prompt_audio_path = validate_file_exists(args.prompt_audio, "reference audio file")
-    output_path = validate_output_path(args.output)
-    
-    # Load model
+# -----------------------------
+# Commands
+# -----------------------------
+
+
+def _run_single(args, parser, *, text: str, output: str, prompt_text: str | None):
+    output_path = validate_output_path(output)
+
+    if args.prompt_audio:
+        require_file_exists(args.prompt_audio, parser, "prompt audio file")
+    if args.reference_audio:
+        require_file_exists(args.reference_audio, parser, "reference audio file")
+
     model = load_model(args)
-    
-    # Generate audio
-    print(f"Synthesizing text: {args.text}")
-    print(f"Reference audio: {prompt_audio_path}")
-    print(f"Reference text: {args.prompt_text}")
-    
+
     audio_array = model.generate(
-        text=args.text,
-        prompt_wav_path=str(prompt_audio_path),
-        prompt_text=args.prompt_text,
+        text=text,
+        prompt_wav_path=args.prompt_audio,
+        prompt_text=prompt_text,
+        reference_wav_path=args.reference_audio,
         cfg_value=args.cfg_value,
         inference_timesteps=args.inference_timesteps,
         normalize=args.normalize,
         denoise=args.denoise
+        and (args.prompt_audio is not None or args.reference_audio is not None),
     )
-    
-    # Save audio
-    sf.write(str(output_path), audio_array, 16000)
-    print(f"Saved audio to: {output_path}")
-    
-    # Stats
-    duration = len(audio_array) / 16000
-    print(f"Duration: {duration:.2f}s")
+
+    sf.write(str(output_path), audio_array, model.tts_model.sample_rate)
+
+    duration = len(audio_array) / model.tts_model.sample_rate
+    print(f"Saved audio to: {output_path} ({duration:.2f}s)", file=sys.stderr)
 
 
-def cmd_synthesize(args):
-    """Direct TTS synthesis command."""
-    # Validate inputs
-    if not args.text:
-        print("Error: Please provide text to synthesize (--text)")
-        sys.exit(1)
-    # Validate output path
-    output_path = validate_output_path(args.output)
-    # Load model
-    model = load_model(args)
-    # Generate audio
-    print(f"Synthesizing text: {args.text}")
-    
-    audio_array = model.generate(
-        text=args.text,
-        prompt_wav_path=None,
-        prompt_text=None,
-        cfg_value=args.cfg_value,
-        inference_timesteps=args.inference_timesteps,
-        normalize=args.normalize,
-        denoise=False  # 无参考音频时不需要降噪
+def cmd_design(args, parser):
+    validate_design_args(args, parser)
+    final_text = build_final_text(args.text, args.control)
+    return _run_single(
+        args, parser, text=final_text, output=args.output, prompt_text=None
     )
-    
-    # Save audio
-    sf.write(str(output_path), audio_array, 16000)
-    print(f"Saved audio to: {output_path}")
-    
-    # Stats
-    duration = len(audio_array) / 16000
-    print(f"Duration: {duration:.2f}s")
 
 
-def cmd_batch(args):
-    """Batch synthesis command."""
-    # Validate input file
-    input_file = validate_file_exists(args.input, "input file")
+def cmd_clone(args, parser):
+    prompt_text = validate_clone_args(args, parser)
+    final_text = build_final_text(args.text, args.control)
+    return _run_single(
+        args, parser, text=final_text, output=args.output, prompt_text=prompt_text
+    )
+
+
+def cmd_batch(args, parser):
+    input_file = require_file_exists(args.input, parser, "input file")
     output_dir = Path(args.output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
-    
-    try:
-        with open(input_file, 'r', encoding='utf-8') as f:
-            texts = [line.strip() for line in f if line.strip()]
-    except Exception as e:
-        print(f"Failed to read input file: {e}")
-        sys.exit(1)
+
+    with open(input_file, "r", encoding="utf-8") as f:
+        texts = [line.strip() for line in f if line.strip()]
+
     if not texts:
-        print("Error: Input file is empty or contains no valid lines")
-        sys.exit(1)
-    print(f"Found {len(texts)} lines to process")
-    
+        sys.exit("Error: Input file is empty")
+
+    prompt_text = validate_batch_args(args, parser)
     model = load_model(args)
+
     prompt_audio_path = None
     if args.prompt_audio:
-        prompt_audio_path = str(validate_file_exists(args.prompt_audio, "reference audio file"))
-    
+        prompt_audio_path = str(
+            require_file_exists(args.prompt_audio, parser, "prompt audio file")
+        )
+
+    reference_audio_path = None
+    if args.reference_audio:
+        reference_audio_path = str(
+            require_file_exists(args.reference_audio, parser, "reference audio file")
+        )
+
     success_count = 0
+
     for i, text in enumerate(texts, 1):
-        print(f"\nProcessing {i}/{len(texts)}: {text[:50]}...")
-        
         try:
+            final_text = build_final_text(text, args.control)
             audio_array = model.generate(
-                text=text,
+                text=final_text,
                 prompt_wav_path=prompt_audio_path,
-                prompt_text=args.prompt_text,
+                prompt_text=prompt_text,
+                reference_wav_path=reference_audio_path,
                 cfg_value=args.cfg_value,
                 inference_timesteps=args.inference_timesteps,
                 normalize=args.normalize,
-                denoise=args.denoise and prompt_audio_path is not None
+                denoise=args.denoise
+                and (prompt_audio_path is not None or reference_audio_path is not None),
             )
-            output_file = output_dir / f"output_{i:03d}.wav"
-            sf.write(str(output_file), audio_array, 16000)
-            
-            duration = len(audio_array) / 16000
-            print(f"  Saved: {output_file} ({duration:.2f}s)")
-            success_count += 1
-            
-        except Exception as e:
-            print(f"  Failed: {e}")
-            continue
-    
-    print(f"\nBatch finished: {success_count}/{len(texts)} succeeded")
 
-def _build_unified_parser():
-    """Build unified argument parser (no subcommands, route by args)."""
+            output_file = output_dir / f"output_{i:03d}.wav"
+            sf.write(str(output_file), audio_array, model.tts_model.sample_rate)
+
+            duration = len(audio_array) / model.tts_model.sample_rate
+            print(f"Saved: {output_file} ({duration:.2f}s)", file=sys.stderr)
+            success_count += 1
+
+        except Exception as e:
+            print(f"Failed on line {i}: {e}", file=sys.stderr)
+
+    print(f"\nBatch finished: {success_count}/{len(texts)} succeeded", file=sys.stderr)
+
+
+# -----------------------------
+# Parser
+# -----------------------------
+
+
+def _add_common_generation_args(parser):
+    parser.add_argument("--text", "-t", help="Text to synthesize")
+    parser.add_argument(
+        "--control",
+        type=str,
+        help="Control instruction for VoxCPM2 voice design/cloning",
+    )
+    parser.add_argument(
+        "--cfg-value",
+        type=float,
+        default=2.0,
+        help="CFG guidance scale (float, recommended 1.0–3.0, default: 2.0)",
+    )
+    parser.add_argument(
+        "--inference-timesteps",
+        type=int,
+        default=10,
+        help="Inference steps (int, recommended 4–30, default: 10)",
+    )
+    parser.add_argument(
+        "--normalize", action="store_true", help="Enable text normalization"
+    )
+
+
+def _add_prompt_reference_args(parser):
+    parser.add_argument(
+        "--prompt-audio",
+        "-pa",
+        help="Prompt audio file path (continuation mode, requires --prompt-text or --prompt-file)",
+    )
+    parser.add_argument(
+        "--prompt-text", "-pt", help="Text corresponding to the prompt audio"
+    )
+    parser.add_argument(
+        "--prompt-file", type=str, help="Text file corresponding to the prompt audio"
+    )
+    parser.add_argument(
+        "--reference-audio",
+        "-ra",
+        help="Reference audio for voice cloning (VoxCPM2 only)",
+    )
+    parser.add_argument(
+        "--denoise",
+        action="store_true",
+        help="Enable prompt/reference speech enhancement",
+    )
+
+
+def _add_model_args(parser):
+    parser.add_argument("--model-path", type=str, help="Local VoxCPM model path")
+    parser.add_argument(
+        "--hf-model-id",
+        type=str,
+        default=DEFAULT_HF_MODEL_ID,
+        help=f"Hugging Face repo id (default: {DEFAULT_HF_MODEL_ID})",
+    )
+    parser.add_argument(
+        "--cache-dir", type=str, help="Cache directory for Hub downloads"
+    )
+    parser.add_argument(
+        "--local-files-only", action="store_true", help="Disable network access"
+    )
+    parser.add_argument(
+        "--no-denoiser", action="store_true", help="Disable denoiser model loading"
+    )
+    parser.add_argument(
+        "--no-optimize",
+        action="store_true",
+        help="Disable model optimization during loading",
+    )
+    parser.add_argument(
+        "--zipenhancer-path",
+        type=str,
+        help="ZipEnhancer model id or local path (or env ZIPENHANCER_MODEL_PATH)",
+    )
+
+
+def _add_lora_args(parser):
+    parser.add_argument("--lora-path", type=str, help="Path to LoRA weights")
+    parser.add_argument(
+        "--lora-r", type=int, default=32, help="LoRA rank (positive int, default: 32)"
+    )
+    parser.add_argument(
+        "--lora-alpha",
+        type=int,
+        default=16,
+        help="LoRA alpha (positive int, default: 16)",
+    )
+    parser.add_argument(
+        "--lora-dropout",
+        type=float,
+        default=0.0,
+        help="LoRA dropout rate (0.0–1.0, default: 0.0)",
+    )
+    parser.add_argument(
+        "--lora-disable-lm", action="store_true", help="Disable LoRA on LM layers"
+    )
+    parser.add_argument(
+        "--lora-disable-dit", action="store_true", help="Disable LoRA on DiT layers"
+    )
+    parser.add_argument(
+        "--lora-enable-proj",
+        action="store_true",
+        help="Enable LoRA on projection layers",
+    )
+
+
+def _build_parser():
     parser = argparse.ArgumentParser(
-        description="VoxCPM CLI (single parser) - voice cloning, direct TTS, and batch processing",
+        description="VoxCPM CLI - VoxCPM2-first voice design, cloning, and batch processing",
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog="""
 Examples:
-  # Direct synthesis (single sample)
-  voxcpm --text "Hello world" --output out.wav
-
-  # Voice cloning (reference audio + text)
-  voxcpm --text "Hello world" --prompt-audio voice.wav --prompt-text "reference text" --output out.wav --denoise
-
-  # Batch processing
-  voxcpm --input texts.txt --output-dir ./outs
-
-  # Select model (from Hub)
-  voxcpm --text "Hello" --output out.wav --hf-model-id openbmb/VoxCPM-0.5B
-        """
+  voxcpm design --text "Hello world" --output out.wav
+  voxcpm design --text "Hello world" --control "warm female voice" --output out.wav
+  voxcpm clone --text "Hello" --reference-audio ref.wav --output out.wav
+  voxcpm batch --input texts.txt --output-dir ./outs --reference-audio ref.wav
+        """,
     )
 
-    # Task selection (automatic routing by presence of args)
-    parser.add_argument("--input", "-i", help="Input text file (one line per sample)")
-    parser.add_argument("--output-dir", "-od", help="Output directory (for batch mode)")
-    parser.add_argument("--text", "-t", help="Text to synthesize (single-sample mode)")
-    parser.add_argument("--output", "-o", help="Output audio file path (single-sample mode)")
+    subparsers = parser.add_subparsers(dest="command")
 
-    # Prompt audio (for voice cloning)
-    parser.add_argument("--prompt-audio", "-pa", help="Reference audio file path")
-    parser.add_argument("--prompt-text", "-pt", help="Reference text corresponding to the audio")
-    parser.add_argument("--denoise", action="store_true", help="Enable prompt speech enhancement (denoising)")
+    design_parser = subparsers.add_parser(
+        "design", help="Generate speech with VoxCPM2-first voice design"
+    )
+    _add_common_generation_args(design_parser)
+    _add_prompt_reference_args(design_parser)
+    _add_model_args(design_parser)
+    _add_lora_args(design_parser)
+    design_parser.add_argument(
+        "--output", "-o", required=True, help="Output audio file path"
+    )
 
-    # Generation parameters
-    parser.add_argument("--cfg-value", type=float, default=2.0, help="CFG guidance scale (default: 2.0)")
-    parser.add_argument("--inference-timesteps", type=int, default=10, help="Inference steps (default: 10)")
-    parser.add_argument("--normalize", action="store_true", help="Enable text normalization")
+    clone_parser = subparsers.add_parser(
+        "clone", help="Clone a voice with reference/prompt audio"
+    )
+    _add_common_generation_args(clone_parser)
+    _add_prompt_reference_args(clone_parser)
+    _add_model_args(clone_parser)
+    _add_lora_args(clone_parser)
+    clone_parser.add_argument(
+        "--output", "-o", required=True, help="Output audio file path"
+    )
 
-    # Model loading parameters
-    parser.add_argument("--model-path", type=str, help="Local VoxCPM model path (overrides Hub download)")
-    parser.add_argument("--hf-model-id", type=str, default="openbmb/VoxCPM-0.5B", help="Hugging Face repo id (e.g., openbmb/VoxCPM-0.5B)")
-    parser.add_argument("--cache-dir", type=str, help="Cache directory for Hub downloads")
-    parser.add_argument("--local-files-only", action="store_true", help="Use only local files (no network)")
-    parser.add_argument("--no-denoiser", action="store_true", help="Disable denoiser model loading")
-    parser.add_argument("--zipenhancer-path", type=str, default="iic/speech_zipenhancer_ans_multiloss_16k_base", help="ZipEnhancer model id or local path (default reads from env)")
+    batch_parser = subparsers.add_parser(
+        "batch", help="Batch-generate one line per output file"
+    )
+    batch_parser.add_argument(
+        "--input", "-i", required=True, help="Input text file (one text per line)"
+    )
+    batch_parser.add_argument(
+        "--output-dir", "-od", required=True, help="Output directory"
+    )
+    batch_parser.add_argument(
+        "--control",
+        type=str,
+        help="Control instruction for VoxCPM2 voice design/cloning",
+    )
+    _add_prompt_reference_args(batch_parser)
+    batch_parser.add_argument(
+        "--cfg-value",
+        type=float,
+        default=2.0,
+        help="CFG guidance scale (float, recommended 1.0–3.0, default: 2.0)",
+    )
+    batch_parser.add_argument(
+        "--inference-timesteps",
+        type=int,
+        default=10,
+        help="Inference steps (int, recommended 4–30, default: 10)",
+    )
+    batch_parser.add_argument(
+        "--normalize", action="store_true", help="Enable text normalization"
+    )
+    _add_model_args(batch_parser)
+    _add_lora_args(batch_parser)
+
+    # Legacy root arguments
+    parser.add_argument("--input", "-i", help="Input text file (batch mode only)")
+    parser.add_argument(
+        "--output-dir", "-od", help="Output directory (batch mode only)"
+    )
+    _add_common_generation_args(parser)
+    parser.add_argument(
+        "--output", "-o", help="Output audio file path (single or clone mode)"
+    )
+    _add_prompt_reference_args(parser)
+    _add_model_args(parser)
+    _add_lora_args(parser)
 
     return parser
 
 
-def main():
-    """Unified CLI entrypoint: route by provided arguments."""
-    parser = _build_unified_parser()
-    args = parser.parse_args()
+def _dispatch_legacy(args, parser):
+    warn_legacy_mode()
+
+    if args.input and args.text:
+        parser.error(
+            "Use either batch mode (--input) or single mode (--text), not both."
+        )
 
-    # Routing: prefer batch → single (clone/direct)
     if args.input:
         if not args.output_dir:
-            print("Error: Batch mode requires --output-dir")
-            parser.print_help()
-            sys.exit(1)
-        return cmd_batch(args)
+            parser.error("Batch mode requires --output-dir")
+        return cmd_batch(args, parser)
 
-    # Single-sample mode
     if not args.text or not args.output:
-        print("Error: Single-sample mode requires --text and --output")
-        parser.print_help()
-        sys.exit(1)
+        parser.error("Single-sample legacy mode requires --text and --output")
 
-    # If prompt audio+text provided → voice cloning
-    if args.prompt_audio or args.prompt_text:
-        if not args.prompt_audio or not args.prompt_text:
-            print("Error: Voice cloning requires both --prompt-audio and --prompt-text")
-            sys.exit(1)
-        return cmd_clone(args)
+    if (
+        args.prompt_audio
+        or args.prompt_text
+        or args.prompt_file
+        or args.reference_audio
+    ):
+        return cmd_clone(args, parser)
 
-    # Otherwise → direct synthesis
-    return cmd_synthesize(args)
+    return cmd_design(args, parser)
+
+
+# -----------------------------
+# Entrypoint
+# -----------------------------
+
+
+def main():
+    parser = _build_parser()
+    args = parser.parse_args()
+
+    validate_ranges(args, parser)
+
+    if args.command == "design":
+        if not args.text:
+            parser.error("`design` requires --text")
+        return cmd_design(args, parser)
+
+    if args.command == "clone":
+        if not args.text or not args.output:
+            parser.error("`clone` requires --text and --output")
+        return cmd_clone(args, parser)
+
+    if args.command == "batch":
+        return cmd_batch(args, parser)
+
+    return _dispatch_legacy(args, parser)
 
 
 if __name__ == "__main__":

src/voxcpm/core.py

@@ -1,20 +1,25 @@
-import torch
-import torchaudio
 import os
+import sys
+import re
+import json
 import tempfile
-from modelscope.pipelines import pipeline
-from modelscope.utils.constant import Tasks
+import numpy as np
+from typing import Generator, Optional
 from huggingface_hub import snapshot_download
-from .model.voxcpm import VoxCPMModel
-from .utils.text_normalize import TextNormalizer
+from .model.voxcpm import VoxCPMModel, LoRAConfig
+from .model.voxcpm2 import VoxCPM2Model
 
 
 class VoxCPM:
-    def __init__(self,
-            voxcpm_model_path : str,
-            zipenhancer_model_path : str = "iic/speech_zipenhancer_ans_multiloss_16k_base",
-            enable_denoiser : bool = True,
-        ):
+    def __init__(
+        self,
+        voxcpm_model_path: str,
+        zipenhancer_model_path: str | None = "iic/speech_zipenhancer_ans_multiloss_16k_base",
+        enable_denoiser: bool = True,
+        optimize: bool = True,
+        lora_config: Optional[LoRAConfig] = None,
+        lora_weights_path: Optional[str] = None,
+    ):
         """Initialize VoxCPM TTS pipeline.
 
         Args:
@@ -24,39 +29,94 @@ class VoxCPM:
             zipenhancer_model_path: ModelScope acoustic noise suppression model
                 id or local path. If None, denoiser will not be initialized.
             enable_denoiser: Whether to initialize the denoiser pipeline.
+            optimize: Whether to optimize the model with torch.compile. True by default, but can be disabled for debugging.
+            lora_config: LoRA configuration for fine-tuning. If lora_weights_path is
+                provided without lora_config, a default config will be created.
+            lora_weights_path: Path to pre-trained LoRA weights (.pth file or directory
+                containing lora_weights.ckpt). If provided, LoRA weights will be loaded.
         """
-        print(f"voxcpm_model_path: {voxcpm_model_path}, zipenhancer_model_path: {zipenhancer_model_path}, enable_denoiser: {enable_denoiser}")
-        self.tts_model = VoxCPMModel.from_local(voxcpm_model_path)
-        self.text_normalizer = TextNormalizer()
+        print(
+            f"voxcpm_model_path: {voxcpm_model_path}, zipenhancer_model_path: {zipenhancer_model_path}, enable_denoiser: {enable_denoiser}",
+            file=sys.stderr,
+        )
+
+        # If lora_weights_path is provided but no lora_config, create a default one
+        if lora_weights_path is not None and lora_config is None:
+            lora_config = LoRAConfig(
+                enable_lm=True,
+                enable_dit=True,
+                enable_proj=False,
+            )
+            print(f"Auto-created default LoRAConfig for loading weights from: {lora_weights_path}", file=sys.stderr)
+
+        # Determine model type from config.json architecture field
+        config_path = os.path.join(voxcpm_model_path, "config.json")
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+        arch = config.get("architecture", "voxcpm").lower()
+
+        if arch == "voxcpm2":
+            self.tts_model = VoxCPM2Model.from_local(voxcpm_model_path, optimize=optimize, lora_config=lora_config)
+            print("Loaded VoxCPM2Model", file=sys.stderr)
+        elif arch == "voxcpm":
+            self.tts_model = VoxCPMModel.from_local(voxcpm_model_path, optimize=optimize, lora_config=lora_config)
+            print("Loaded VoxCPMModel", file=sys.stderr)
+        else:
+            raise ValueError(f"Unsupported architecture: {arch}")
+
+        # Load LoRA weights if path is provided
+        if lora_weights_path is not None:
+            print(f"Loading LoRA weights from: {lora_weights_path}", file=sys.stderr)
+            loaded_keys, skipped_keys = self.tts_model.load_lora_weights(lora_weights_path)
+            print(f"Loaded {len(loaded_keys)} LoRA parameters, skipped {len(skipped_keys)}", file=sys.stderr)
+
+        self.text_normalizer = None
+        self.denoiser = None
         if enable_denoiser and zipenhancer_model_path is not None:
-            self.denoiser = pipeline(
-                Tasks.acoustic_noise_suppression,
-                model=zipenhancer_model_path)
+            from .zipenhancer import ZipEnhancer
+
+            self.denoiser = ZipEnhancer(zipenhancer_model_path)
         else:
             self.denoiser = None
-        print("Warm up VoxCPMModel...")
-        self.tts_model.generate(
-            target_text="Hello, this is the first test sentence."
-        ) 
+        if optimize:
+            print("Warm up VoxCPMModel...", file=sys.stderr)
+            self.tts_model.generate(
+                target_text="Hello, this is the first test sentence.",
+                max_len=10,
+            )
 
     @classmethod
-    def from_pretrained(cls,
-            hf_model_id: str = "openbmb/VoxCPM-0.5B",
-            load_denoiser: bool = True,
-            zipenhancer_model_id: str = "iic/speech_zipenhancer_ans_multiloss_16k_base",
-            cache_dir: str = None,
-            local_files_only: bool = False,
-        ):
+    def from_pretrained(
+        cls,
+        hf_model_id: str = "openbmb/VoxCPM2",
+        load_denoiser: bool = True,
+        zipenhancer_model_id: str = "iic/speech_zipenhancer_ans_multiloss_16k_base",
+        cache_dir: str = None,
+        local_files_only: bool = False,
+        optimize: bool = True,
+        lora_config: Optional[LoRAConfig] = None,
+        lora_weights_path: Optional[str] = None,
+        **kwargs,
+    ):
         """Instantiate ``VoxCPM`` from a Hugging Face Hub snapshot.
 
         Args:
-            hf_model_id: Explicit Hugging Face repository id (e.g. "org/repo").
+            hf_model_id: Explicit Hugging Face repository id (e.g. "org/repo") or local path.
             load_denoiser: Whether to initialize the denoiser pipeline.
+            optimize: Whether to optimize the model with torch.compile. True by default, but can be disabled for debugging.
             zipenhancer_model_id: Denoiser model id or path for ModelScope
                 acoustic noise suppression.
             cache_dir: Custom cache directory for the snapshot.
             local_files_only: If True, only use local files and do not attempt
                 to download.
+            lora_config: LoRA configuration for fine-tuning. If lora_weights_path is
+                provided without lora_config, a default config will be created with
+                enable_lm=True and enable_dit=True.
+            lora_weights_path: Path to pre-trained LoRA weights (.pth file or directory
+                containing lora_weights.ckpt). If provided, LoRA weights will be loaded
+                after model initialization.
+        Kwargs:
+            Additional keyword arguments passed to the ``VoxCPM`` constructor.
 
         Returns:
             VoxCPM: Initialized instance whose ``voxcpm_model_path`` points to
@@ -67,115 +127,207 @@ class VoxCPM:
                 ``hf_model_id`` is provided.
         """
         repo_id = hf_model_id
-        if not repo_id or repo_id.strip() == "":
-            raise ValueError("You must provide a valid hf_model_id")
+        if not repo_id:
+            raise ValueError("You must provide hf_model_id")
 
-        local_path = snapshot_download(
-            repo_id=repo_id,
-            cache_dir=cache_dir,
-            local_files_only=local_files_only,
-        )
+        # Load from local path if provided
+        if os.path.isdir(repo_id):
+            local_path = repo_id
+        else:
+            # Otherwise, try from_pretrained (Hub); exit on failure
+            local_path = snapshot_download(
+                repo_id=repo_id,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+            )
 
         return cls(
             voxcpm_model_path=local_path,
             zipenhancer_model_path=zipenhancer_model_id if load_denoiser else None,
             enable_denoiser=load_denoiser,
+            optimize=optimize,
+            lora_config=lora_config,
+            lora_weights_path=lora_weights_path,
+            **kwargs,
         )
-        
-    def _normalize_loudness(self, wav_path: str):
-        audio, sr = torchaudio.load(wav_path)
-        loudness = torchaudio.functional.loudness(audio, sr)
-        normalized_audio = torchaudio.functional.gain(audio, -20-loudness)
-        torchaudio.save(wav_path, normalized_audio, sr)
 
-    def generate(self, 
-            text : str,
-            prompt_wav_path : str = None,
-            prompt_text : str = None,
-            cfg_value : float = 2.0,    
-            inference_timesteps : int = 10,
-            max_length : int = 4096,
-            normalize : bool = True,
-            denoise : bool = True,
-            retry_badcase : bool = True,
-            retry_badcase_max_times : int = 3,
-            retry_badcase_ratio_threshold : float = 6.0,
-        ):
+    def generate(self, *args, **kwargs) -> np.ndarray:
+        return next(self._generate(*args, streaming=False, **kwargs))
+
+    def generate_streaming(self, *args, **kwargs) -> Generator[np.ndarray, None, None]:
+        return self._generate(*args, streaming=True, **kwargs)
+
+    def _generate(
+        self,
+        text: str,
+        prompt_wav_path: str = None,
+        prompt_text: str = None,
+        reference_wav_path: str = None,
+        cfg_value: float = 2.0,
+        inference_timesteps: int = 10,
+        min_len: int = 2,
+        max_len: int = 4096,
+        normalize: bool = False,
+        denoise: bool = False,
+        retry_badcase: bool = True,
+        retry_badcase_max_times: int = 3,
+        retry_badcase_ratio_threshold: float = 6.0,
+        streaming: bool = False,
+    ) -> Generator[np.ndarray, None, None]:
         """Synthesize speech for the given text and return a single waveform.
 
-        This method optionally builds and reuses a prompt cache. If an external
-        prompt (``prompt_wav_path`` + ``prompt_text``) is provided, it will be
-        used for all sub-sentences. Otherwise, the prompt cache is built from
-        the first generated result and reused for the remaining text chunks.
-
         Args:
-            text: Input text. Can include newlines; each non-empty line is
-                treated as a sub-sentence.
-            prompt_wav_path: Path to a reference audio file for prompting.
+            text: Input text to synthesize.
+            prompt_wav_path: Path to prompt audio for continuation mode.
+                Must be paired with ``prompt_text``.
             prompt_text: Text content corresponding to the prompt audio.
+            reference_wav_path: Path to reference audio for voice cloning
+                (structurally isolated via ref_audio tokens). Can be used
+                alone or combined with ``prompt_wav_path`` + ``prompt_text``.
             cfg_value: Guidance scale for the generation model.
             inference_timesteps: Number of inference steps.
-            max_length: Maximum token length during generation.
+            min_len: Minimum audio length.
+            max_len: Maximum token length during generation.
             normalize: Whether to run text normalization before generation.
-            denoise: Whether to denoise the prompt audio if a denoiser is
-                available.
+            denoise: Whether to denoise the prompt/reference audio if a
+                denoiser is available.
             retry_badcase: Whether to retry badcase.
             retry_badcase_max_times: Maximum number of times to retry badcase.
             retry_badcase_ratio_threshold: Threshold for audio-to-text ratio.
+            streaming: Whether to return a generator of audio chunks.
         Returns:
-            numpy.ndarray: 1D waveform array (float32) on CPU.
+            Generator of numpy.ndarray: 1D waveform array (float32) on CPU.
+            Yields audio chunks for each generation step if ``streaming=True``,
+            otherwise yields a single array containing the final audio.
         """
-        texts = text.split("\n")
-        texts = [t.strip() for t in texts if t.strip()]
-        final_wav = []
-        temp_prompt_wav_path = None 
-        
+        if not text.strip() or not isinstance(text, str):
+            raise ValueError("target text must be a non-empty string")
+
+        if prompt_wav_path is not None:
+            if not os.path.exists(prompt_wav_path):
+                raise FileNotFoundError(f"prompt_wav_path does not exist: {prompt_wav_path}")
+
+        if reference_wav_path is not None:
+            if not os.path.exists(reference_wav_path):
+                raise FileNotFoundError(f"reference_wav_path does not exist: {reference_wav_path}")
+
+        if (prompt_wav_path is None) != (prompt_text is None):
+            raise ValueError("prompt_wav_path and prompt_text must both be provided or both be None")
+
+        is_v2 = isinstance(self.tts_model, VoxCPM2Model)
+        if reference_wav_path is not None and not is_v2:
+            raise ValueError("reference_wav_path is only supported with VoxCPM2 models")
+
+        text = text.replace("\n", " ")
+        text = re.sub(r"\s+", " ", text)
+        temp_files = []
+
         try:
-            if prompt_wav_path is not None and prompt_text is not None:
-                if denoise and self.denoiser is not None:
-                    with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp_file:
-                        temp_prompt_wav_path = tmp_file.name
-                    
-                    self.denoiser(prompt_wav_path, output_path=temp_prompt_wav_path)
-                    self._normalize_loudness(temp_prompt_wav_path)
-                    prompt_wav_path = temp_prompt_wav_path
-                fixed_prompt_cache = self.tts_model.build_prompt_cache(
-                    prompt_wav_path=prompt_wav_path,
-                    prompt_text=prompt_text
-                )
-            else:
-                fixed_prompt_cache = None  # will be built from the first inference
-            
-            for sub_text in texts:
-                if sub_text.strip() == "":
-                    continue
-                print("sub_text:", sub_text)
-                if normalize:
-                    sub_text = self.text_normalizer.normalize(sub_text)
-                wav, target_text_token, generated_audio_feat = self.tts_model.generate_with_prompt_cache(
-                                target_text=sub_text,
-                                prompt_cache=fixed_prompt_cache,
-                                min_len=2,
-                                max_len=max_length,
-                                inference_timesteps=inference_timesteps,
-                                cfg_value=cfg_value,
-                                retry_badcase=retry_badcase,
-                                retry_badcase_max_times=retry_badcase_max_times,
-                                retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
-                            )
-                if fixed_prompt_cache is None:
-                    fixed_prompt_cache = self.tts_model.merge_prompt_cache(
-                        original_cache=None,
-                        new_text_token=target_text_token,
-                        new_audio_feat=generated_audio_feat
+            actual_prompt_path = prompt_wav_path
+            actual_ref_path = reference_wav_path
+
+            if denoise and self.denoiser is not None:
+                if prompt_wav_path is not None:
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
+                        temp_files.append(tmp.name)
+                    self.denoiser.enhance(prompt_wav_path, output_path=temp_files[-1])
+                    actual_prompt_path = temp_files[-1]
+                if reference_wav_path is not None:
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
+                        temp_files.append(tmp.name)
+                    self.denoiser.enhance(reference_wav_path, output_path=temp_files[-1])
+                    actual_ref_path = temp_files[-1]
+
+            if actual_prompt_path is not None or actual_ref_path is not None:
+                if is_v2:
+                    fixed_prompt_cache = self.tts_model.build_prompt_cache(
+                        prompt_text=prompt_text,
+                        prompt_wav_path=actual_prompt_path,
+                        reference_wav_path=actual_ref_path,
                     )
-                final_wav.append(wav)
-        
-            return torch.cat(final_wav, dim=1).squeeze(0).cpu().numpy()
-        
+                else:
+                    fixed_prompt_cache = self.tts_model.build_prompt_cache(
+                        prompt_text=prompt_text,
+                        prompt_wav_path=actual_prompt_path,
+                    )
+            else:
+                fixed_prompt_cache = None
+
+            if normalize:
+                if self.text_normalizer is None:
+                    from .utils.text_normalize import TextNormalizer
+
+                    self.text_normalizer = TextNormalizer()
+                text = self.text_normalizer.normalize(text)
+
+            generate_result = self.tts_model._generate_with_prompt_cache(
+                target_text=text,
+                prompt_cache=fixed_prompt_cache,
+                min_len=min_len,
+                max_len=max_len,
+                inference_timesteps=inference_timesteps,
+                cfg_value=cfg_value,
+                retry_badcase=retry_badcase,
+                retry_badcase_max_times=retry_badcase_max_times,
+                retry_badcase_ratio_threshold=retry_badcase_ratio_threshold,
+                streaming=streaming,
+            )
+
+            for wav, _, _ in generate_result:
+                yield wav.squeeze(0).cpu().numpy()
+
         finally:
-            if temp_prompt_wav_path and os.path.exists(temp_prompt_wav_path):
-                try:
-                    os.unlink(temp_prompt_wav_path)
-                except OSError:
-                    pass  
+            for tmp_path in temp_files:
+                if tmp_path and os.path.exists(tmp_path):
+                    try:
+                        os.unlink(tmp_path)
+                    except OSError:
+                        pass
+
+    # ------------------------------------------------------------------ #
+    # LoRA Interface (delegated to VoxCPMModel)
+    # ------------------------------------------------------------------ #
+    def load_lora(self, lora_weights_path: str) -> tuple:
+        """Load LoRA weights from a checkpoint file.
+
+        Args:
+            lora_weights_path: Path to LoRA weights (.pth file or directory
+                containing lora_weights.ckpt).
+
+        Returns:
+            tuple: (loaded_keys, skipped_keys) - lists of loaded and skipped parameter names.
+
+        Raises:
+            RuntimeError: If model was not initialized with LoRA config.
+        """
+        if self.tts_model.lora_config is None:
+            raise RuntimeError(
+                "Cannot load LoRA weights: model was not initialized with LoRA config. "
+                "Please reinitialize with lora_config or lora_weights_path parameter."
+            )
+        return self.tts_model.load_lora_weights(lora_weights_path)
+
+    def unload_lora(self):
+        """Unload LoRA by resetting all LoRA weights to initial state (effectively disabling LoRA)."""
+        self.tts_model.reset_lora_weights()
+
+    def set_lora_enabled(self, enabled: bool):
+        """Enable or disable LoRA layers without unloading weights.
+
+        Args:
+            enabled: If True, LoRA layers are active; if False, only base model is used.
+        """
+        self.tts_model.set_lora_enabled(enabled)
+
+    def get_lora_state_dict(self) -> dict:
+        """Get current LoRA parameters state dict.
+
+        Returns:
+            dict: State dict containing all LoRA parameters (lora_A, lora_B).
+        """
+        return self.tts_model.get_lora_state_dict()
+
+    @property
+    def lora_enabled(self) -> bool:
+        """Check if LoRA is currently configured."""
+        return self.tts_model.lora_config is not None

src/voxcpm/model/__init__.py

@@ -1,3 +1,4 @@
 from .voxcpm import VoxCPMModel
+from .voxcpm2 import VoxCPM2Model
 
-__all__ = ["VoxCPMModel"]
+__all__ = ["VoxCPMModel", "VoxCPM2Model"]

src/voxcpm/model/utils.py

@@ -5,17 +5,17 @@ from transformers import PreTrainedTokenizer
 
 def mask_multichar_chinese_tokens(tokenizer: PreTrainedTokenizer):
     """Create a tokenizer wrapper that converts multi-character Chinese tokens to single characters.
-    
+
     This function creates a wrapper around the provided tokenizer that automatically
     splits multi-character Chinese tokens into individual characters. This is useful
     for ensuring consistent tokenization of Chinese text.
-    
+
     Args:
         tokenizer: The base tokenizer to wrap
-        
+
     Returns:
         A CharTokenizerWrapper instance that handles multi-character Chinese tokens
-        
+
     Example:
         >>> from transformers import LlamaTokenizerFast
         >>> tokenizer = LlamaTokenizerFast.from_pretrained("path/to/tokenizer")
@@ -24,20 +24,19 @@ def mask_multichar_chinese_tokens(tokenizer: PreTrainedTokenizer):
     """
     # Pre-compute multi-character tokens (length >= 2, pure Chinese characters)
     multichar_tokens = {
-        token for token in tokenizer.vocab.keys() 
-        if len(token) >= 2 and all("\u4e00" <= c <= "\u9fff" for c in token)
+        token for token in tokenizer.vocab.keys() if len(token) >= 2 and all("\u4e00" <= c <= "\u9fff" for c in token)
     }
 
     class CharTokenizerWrapper:
         """Wrapper class for tokenizers that handles multi-character Chinese tokens.
-        
+
         This wrapper automatically splits multi-character Chinese tokens into
         individual characters while preserving the original tokenizer's interface.
         """
-        
+
         def __init__(self, base_tokenizer: PreTrainedTokenizer) -> None:
             """Initialize the wrapper with a base tokenizer.
-            
+
             Args:
                 base_tokenizer: The tokenizer to wrap
             """
@@ -46,14 +45,14 @@ def mask_multichar_chinese_tokens(tokenizer: PreTrainedTokenizer):
 
         def tokenize(self, text: str, **kwargs) -> List[str]:
             """Tokenize text and split multi-character Chinese tokens into single characters.
-            
+
             Args:
                 text: Input text to tokenize
                 **kwargs: Additional arguments passed to the base tokenizer
-                
+
             Returns:
                 List of processed tokens with multi-character Chinese tokens split
-                
+
             Example:
                 >>> wrapper = CharTokenizerWrapper(tokenizer)
                 >>> tokens = wrapper.tokenize("你好世界")
@@ -61,10 +60,10 @@ def mask_multichar_chinese_tokens(tokenizer: PreTrainedTokenizer):
             """
             if not isinstance(text, str):
                 raise TypeError(f"Expected string input, got {type(text)}")
-                
+
             tokens = self.tokenizer.tokenize(text, **kwargs)
             processed = []
-            
+
             for token in tokens:
                 # Remove possible subword prefix
                 clean_token = token.replace("▁", "")
@@ -75,22 +74,22 @@ def mask_multichar_chinese_tokens(tokenizer: PreTrainedTokenizer):
                     processed.extend(chars)
                 else:
                     processed.append(token)
-                    
+
             return processed
 
         def __call__(self, text: str, **kwargs) -> List[int]:
             """Call the tokenizer and return token IDs.
-            
+
             This method provides the same interface as the original tokenizer
             but with multi-character Chinese token handling.
-            
+
             Args:
                 text: Input text to tokenize
                 **kwargs: Additional arguments passed to the base tokenizer
-                
+
             Returns:
                 List of token IDs
-                
+
             Raises:
                 TypeError: If input is not a string
                 ValueError: If tokenization fails

src/voxcpm/model/voxcpm.py

@@ -19,18 +19,28 @@ limitations under the License.
 """
 
 import os
-from typing import Dict, Optional, Tuple, Union
+import sys
+from typing import Tuple, Union, Generator, List, Optional
 
 import torch
 import torch.nn as nn
 import torchaudio
+import warnings
 from einops import rearrange
 from pydantic import BaseModel
+
+try:
+    from safetensors.torch import load_file
+
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
 from tqdm import tqdm
 from transformers import LlamaTokenizerFast
 
-from ..modules.audiovae import AudioVAE
+from ..modules.audiovae import AudioVAE, AudioVAEConfig
 from ..modules.layers import ScalarQuantizationLayer
+from ..modules.layers.lora import apply_lora_to_named_linear_modules
 from ..modules.locdit import CfmConfig, UnifiedCFM, VoxCPMLocDiT
 from ..modules.locenc import VoxCPMLocEnc
 from ..modules.minicpm4 import MiniCPM4Config, MiniCPMModel
@@ -65,10 +75,31 @@ class VoxCPMConfig(BaseModel):
 
     encoder_config: VoxCPMEncoderConfig
     dit_config: VoxCPMDitConfig
+    audio_vae_config: Optional[AudioVAEConfig] = None
 
     max_length: int = 4096
     device: str = "cuda"
     dtype: str = "bfloat16"
+    dit_mean_mode: bool = False
+
+
+class LoRAConfig(BaseModel):
+    enable_lm: bool = False  # Apply LoRA to base_lm + residual_lm
+    enable_dit: bool = False  # Apply LoRA to VoxCPMLocDiT
+    enable_proj: bool = False  # Apply LoRA to projection Linear layers
+
+    r: int = 8
+    alpha: int = 16
+    dropout: float = 0.0
+
+    # Target linear layer names for LM & DiT (matched by attribute name)
+    target_modules_lm: list[str] = ["q_proj", "v_proj", "k_proj", "o_proj"]
+    target_modules_dit: list[str] = ["q_proj", "v_proj", "k_proj", "o_proj"]
+    # Projection layer attribute names to find on VoxCPMModel
+    target_proj_modules: list[str] = ["enc_to_lm_proj", "lm_to_dit_proj", "res_to_dit_proj"]
+
+
+VoxCPMConfig.model_rebuild()
 
 
 class VoxCPMModel(nn.Module):
@@ -77,18 +108,24 @@ class VoxCPMModel(nn.Module):
         config: VoxCPMConfig,
         tokenizer: LlamaTokenizerFast,
         audio_vae: AudioVAE,
+        lora_config: LoRAConfig = None,
     ):
         super().__init__()
         self.config = config
+        self.lora_config = lora_config
         self.feat_dim = config.feat_dim
         self.patch_size = config.patch_size
         self.device = config.device
         if not torch.cuda.is_available():
-            self.device = "cpu"
+            if torch.backends.mps.is_available():
+                self.device = "mps"
+            else:
+                self.device = "cpu"
+        print(f"Running on device: {self.device}, dtype: {self.config.dtype}", file=sys.stderr)
 
         # Text-Semantic LM
         self.base_lm = MiniCPMModel(config.lm_config)
-        self.base_lm.setup_cache(1, config.max_length, self.device, get_dtype(config.dtype))
+        self.base_lm.setup_cache(1, config.max_length, self.device, get_dtype(self.config.dtype))
 
         self.text_tokenizer = mask_multichar_chinese_tokens(tokenizer)
         self.audio_start_token = 101
@@ -99,7 +136,7 @@ class VoxCPMModel(nn.Module):
         residual_lm_config.num_hidden_layers = config.residual_lm_num_layers
         residual_lm_config.vocab_size = 0
         self.residual_lm = MiniCPMModel(residual_lm_config)
-        self.residual_lm.setup_cache(1, config.max_length, self.device, get_dtype(config.dtype))
+        self.residual_lm.setup_cache(1, config.max_length, self.device, get_dtype(self.config.dtype))
 
         # Local Encoder
         encoder_config = config.lm_config.model_copy(deep=True)
@@ -123,15 +160,16 @@ class VoxCPMModel(nn.Module):
             in_channels=config.feat_dim,
             cfm_params=config.dit_config.cfm_config,
             estimator=VoxCPMLocDiT(decoder_config, in_channels=config.feat_dim),
+            mean_mode=config.dit_mean_mode,
         )
 
         # Projection layers
         self.fsq_layer = ScalarQuantizationLayer(
-            config.lm_config.hidden_size, 
-            config.lm_config.hidden_size, 
-            config.scalar_quantization_latent_dim, 
-            config.scalar_quantization_scale
-        ) 
+            config.lm_config.hidden_size,
+            config.lm_config.hidden_size,
+            config.scalar_quantization_latent_dim,
+            config.scalar_quantization_scale,
+        )
         self.enc_to_lm_proj = nn.Linear(config.encoder_config.hidden_dim, config.lm_config.hidden_size)
         self.lm_to_dit_proj = nn.Linear(config.lm_config.hidden_size, config.dit_config.hidden_dim)
         self.res_to_dit_proj = nn.Linear(config.lm_config.hidden_size, config.dit_config.hidden_dim)
@@ -140,29 +178,172 @@ class VoxCPMModel(nn.Module):
         self.stop_proj = nn.Linear(config.lm_config.hidden_size, config.lm_config.hidden_size)
         self.stop_actn = nn.SiLU()
         self.stop_head = nn.Linear(config.lm_config.hidden_size, 2, bias=False)
+        self.stop_loss = nn.CrossEntropyLoss(reduction="none")
 
         # Audio VAE
         self.audio_vae = audio_vae
         self.chunk_size = audio_vae.chunk_size
         self.sample_rate = audio_vae.sample_rate
 
-    
-    def optimize(self):
-        if self.device == "cuda":
+        if self.lora_config is not None:
+            self._apply_lora()
+
+    def _apply_lora(self):
+        """注入 LoRA 到 LM / DiT / 投影层"""
+        cfg = self.lora_config
+        lora_kwargs = dict(r=cfg.r, alpha=cfg.alpha, dropout=cfg.dropout)
+
+        # LM: base_lm + residual_lm
+        if cfg.enable_lm:
+            for lm in [self.base_lm, self.residual_lm]:
+                apply_lora_to_named_linear_modules(lm, target_submodule_names=cfg.target_modules_lm, **lora_kwargs)
+
+        # DiT: feat_decoder.estimator
+        if cfg.enable_dit:
+            apply_lora_to_named_linear_modules(
+                self.feat_decoder.estimator, target_submodule_names=cfg.target_modules_dit, **lora_kwargs
+            )
+
+        # 投影层
+        if cfg.enable_proj:
+            from ..modules.layers.lora import LoRALinear
+
+            for attr_name in cfg.target_proj_modules:
+                module = getattr(self, attr_name, None)
+                if isinstance(module, nn.Linear):
+                    setattr(self, attr_name, LoRALinear(base=module, **lora_kwargs))
+
+    def optimize(self, disable: bool = False):
+        if disable:
+            return self
+        try:
+            if self.device != "cuda":
+                raise ValueError("VoxCPMModel can only be optimized on CUDA device")
+            try:
+                import triton  # noqa: F401
+            except ImportError:
+                raise ValueError("triton is not installed")
             self.base_lm.forward_step = torch.compile(self.base_lm.forward_step, mode="reduce-overhead", fullgraph=True)
-            self.residual_lm.forward_step = torch.compile(self.residual_lm.forward_step, mode="reduce-overhead", fullgraph=True)
-            self.feat_encoder_step = torch.compile(self.feat_encoder, mode="reduce-overhead", fullgraph=True)
-            self.feat_decoder.estimator = torch.compile(self.feat_decoder.estimator, mode="reduce-overhead", fullgraph=True)
-        else:
-            self.base_lm.forward_step = self.base_lm.forward_step
-            self.residual_lm.forward_step = self.residual_lm.forward_step
-            self.feat_encoder_step = self.feat_encoder
-            self.feat_decoder.estimator = self.feat_decoder.estimator
+            self.residual_lm.forward_step = torch.compile(
+                self.residual_lm.forward_step, mode="reduce-overhead", fullgraph=True
+            )
+            self.feat_encoder = torch.compile(self.feat_encoder, mode="reduce-overhead", fullgraph=True)
+            self.feat_decoder.estimator = torch.compile(
+                self.feat_decoder.estimator, mode="reduce-overhead", fullgraph=True
+            )
+        except Exception as e:
+            print(f"Warning: torch.compile disabled - {e}", file=sys.stderr)
         return self
 
+    def forward(
+        self,
+        text_tokens: torch.Tensor,
+        text_mask: torch.Tensor,
+        audio_feats: torch.Tensor,
+        audio_mask: torch.Tensor,
+        loss_mask: torch.Tensor,
+        position_ids: torch.Tensor,
+        labels: torch.Tensor,
+        *,
+        progress: float = 0.0,
+        sample_generate: bool = False,
+    ):
+        del position_ids  # not used yet
+
+        text_tokens = text_tokens.to(self.device, dtype=torch.long)
+        text_mask = text_mask.to(self.device, dtype=self._dtype())
+        audio_feats = audio_feats.to(self.device, dtype=self._dtype())
+        audio_mask = audio_mask.to(self.device, dtype=self._dtype())
+        loss_mask = loss_mask.to(self.device, dtype=self._dtype())
+        labels = labels.to(self.device, dtype=torch.long)
+
+        B, T, P, D = audio_feats.shape
+        feat_embed = self.feat_encoder(audio_feats)
+        feat_embed = self.enc_to_lm_proj(feat_embed)
+
+        scale_emb = getattr(self.config.lm_config, "scale_emb", 1.0)
+        if not getattr(self.config.lm_config, "use_mup", False):
+            scale_emb = 1.0
+        text_embed = self.base_lm.embed_tokens(text_tokens) * scale_emb
+        combined_embed = text_mask.unsqueeze(-1) * text_embed + audio_mask.unsqueeze(-1) * feat_embed
+
+        enc_outputs, _ = self.base_lm(inputs_embeds=combined_embed, is_causal=True)
+        enc_outputs = enc_outputs.to(self._dtype())
+        enc_outputs = self.fsq_layer(enc_outputs) * audio_mask.unsqueeze(-1) + enc_outputs * text_mask.unsqueeze(-1)
+        lm_hidden = torch.cat((torch.zeros_like(enc_outputs[:, 0:1, :]), enc_outputs[:, :-1, :]), dim=1)
+
+        residual_inputs = enc_outputs + audio_mask.unsqueeze(-1) * feat_embed
+        residual_outputs, _ = self.residual_lm(inputs_embeds=residual_inputs, is_causal=True)
+        residual_outputs = residual_outputs.to(self._dtype())
+        residual_hidden = torch.cat(
+            (torch.zeros_like(residual_outputs[:, 0:1, :]), residual_outputs[:, :-1, :]),
+            dim=1,
+        )
+
+        dit_hidden = self.lm_to_dit_proj(lm_hidden) + self.res_to_dit_proj(residual_hidden)
+        dit_hidden = rearrange(dit_hidden, "b t c -> (b t) c")
+
+        # Keep diffusion inputs in the same dtype as the model (e.g., bfloat16)
+        target_dtype = self._dtype()
+
+        feat_gt = rearrange(audio_feats.to(target_dtype), "b t p d -> (b t) p d")
+        feat_cond = torch.cat(
+            (torch.zeros_like(audio_feats[:, 0:1, ...]), audio_feats[:, :-1, ...]),
+            dim=1,
+        )
+        feat_cond = rearrange(feat_cond.to(target_dtype), "b t p d -> (b t) p d")
+
+        loss_seq_mask = loss_mask.unsqueeze(-1).repeat(1, 1, self.patch_size)
+        loss_seq_mask = rearrange(loss_seq_mask, "b t p -> (b t) p 1").to(target_dtype)
+
+        diff_loss = self.feat_decoder.compute_loss(
+            feat_gt.transpose(1, 2).contiguous(),
+            dit_hidden,
+            cond=feat_cond.transpose(1, 2).contiguous(),
+            tgt_mask=loss_seq_mask.transpose(1, 2).contiguous(),
+            progress=progress,
+        )
+
+        stop_logits = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden)))
+        stop_losses = self.stop_loss(stop_logits.transpose(1, 2), labels)
+        denom = torch.clamp(loss_mask.sum(), min=1.0)
+        stop_loss = (stop_losses * loss_mask).sum() / denom
+
+        feat_pred = None
+        if sample_generate:
+            feat_cond_for_sample = feat_cond.transpose(1, 2).contiguous()
+            feat_pred_seq = self.feat_decoder(
+                mu=dit_hidden,
+                patch_size=self.patch_size,
+                cond=feat_cond_for_sample,
+                n_timesteps=(
+                    self.config.dit_config.cfm_config.inference_cfg_rate
+                    if hasattr(self.config.dit_config.cfm_config, "inference_cfg_rate")
+                    else 10
+                ),
+            )
+            feat_pred = rearrange(feat_pred_seq.transpose(1, 2), "(b t) d p -> b d (t p)", b=B, p=self.patch_size)
+
+        feat_gt_tensor = rearrange(feat_gt, "(b t) p d -> b d (t p)", b=B, p=self.patch_size)
+
+        return {
+            "loss/diff": diff_loss,
+            "loss/stop": stop_loss,
+            "feat_gt": feat_gt_tensor,
+            "feat_pred": feat_pred,
+        }
+
+    def _dtype(self):
+        return get_dtype(self.config.dtype)
+
+    def generate(self, *args, **kwargs) -> torch.Tensor:
+        return next(self._generate(*args, streaming=False, **kwargs))
+
+    def generate_streaming(self, *args, **kwargs) -> Generator[torch.Tensor, None, None]:
+        return self._generate(*args, streaming=True, **kwargs)
 
     @torch.inference_mode()
-    def generate(
+    def _generate(
         self,
         target_text: str,
         prompt_text: str = "",
@@ -173,8 +354,12 @@ class VoxCPMModel(nn.Module):
         cfg_value: float = 2.0,
         retry_badcase: bool = False,
         retry_badcase_max_times: int = 3,
-        retry_badcase_ratio_threshold: float = 6.0, # setting acceptable ratio of audio length to text length (for badcase detection)
-    ):
+        retry_badcase_ratio_threshold: float = 6.0,  # setting acceptable ratio of audio length to text length (for badcase detection)
+        streaming: bool = False,
+    ) -> Generator[torch.Tensor, None, None]:
+        if retry_badcase and streaming:
+            warnings.warn("Retry on bad cases is not supported in streaming mode, setting retry_badcase=False.")
+            retry_badcase = False
         if len(prompt_wav_path) == 0:
             text = target_text
             text_token = torch.LongTensor(self.text_tokenizer(text))
@@ -214,24 +399,24 @@ class VoxCPMModel(nn.Module):
             audio, sr = torchaudio.load(prompt_wav_path)
             if audio.size(0) > 1:
                 audio = audio.mean(dim=0, keepdim=True)
-                
+
             if sr != self.sample_rate:
                 audio = torchaudio.functional.resample(audio, sr, self.sample_rate)
 
             patch_len = self.patch_size * self.chunk_size
 
             if audio.size(1) % patch_len != 0:
-                audio = torch.nn.functional.pad(audio, (0, patch_len - audio.size(1) % patch_len))
+                # 左填充：在音频开头填充，保持有效音频数据在序列末尾
+                padding_size = patch_len - audio.size(1) % patch_len
+                audio = torch.nn.functional.pad(audio, (padding_size, 0))
 
             # (B, D, T)
             audio_feat = self.audio_vae.encode(audio.to(self.device), self.sample_rate).cpu()
-
             audio_feat = audio_feat.view(
                 self.audio_vae.latent_dim,
                 -1,
                 self.patch_size,
             ).permute(1, 2, 0)
-            audio_feat = audio_feat[:-1, ...] # trick: remove the last padding token
             audio_length = audio_feat.size(0)
             text_pad_token = torch.zeros(audio_length, dtype=torch.int32, device=text_token.device)
             text_token = torch.cat([text_token, text_pad_token])
@@ -250,34 +435,52 @@ class VoxCPMModel(nn.Module):
 
         text_token = text_token.unsqueeze(0).to(self.device)
         text_mask = text_mask.unsqueeze(0).to(self.device)
-        audio_feat = audio_feat.unsqueeze(0).to(self.device).to(torch.bfloat16)
+        audio_feat = audio_feat.unsqueeze(0).to(self.device).to(get_dtype(self.config.dtype))
         audio_mask = audio_mask.unsqueeze(0).to(self.device)
 
         target_text_length = len(self.text_tokenizer(target_text))
-        
+
         retry_badcase_times = 0
         while retry_badcase_times < retry_badcase_max_times:
-            latent_pred, pred_audio_feat = self.inference(
+            inference_result = self._inference(
                 text_token,
                 text_mask,
                 audio_feat,
                 audio_mask,
                 min_len=min_len,
-                max_len=int(target_text_length * retry_badcase_ratio_threshold + 10) if retry_badcase else max_len,
+                max_len=min(
+                    int(target_text_length * retry_badcase_ratio_threshold + 10), max_len
+                ),  # avoid too long audio
                 inference_timesteps=inference_timesteps,
                 cfg_value=cfg_value,
+                streaming=streaming,
             )
-            if retry_badcase:
-                if pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold:
-                    print(f"  Badcase detected, audio_text_ratio={pred_audio_feat.shape[0] / target_text_length}, retrying...")
-                    retry_badcase_times += 1
-                    continue
+            if streaming:
+                patch_len = self.patch_size * self.chunk_size
+                for latent_pred, _ in inference_result:
+                    decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
+                    decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu()
+                    yield decode_audio
+                break
+            else:
+                latent_pred, pred_audio_feat = next(inference_result)
+                if retry_badcase:
+                    if pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold:
+                        print(
+                            f"  Badcase detected, audio_text_ratio={pred_audio_feat.shape[0] / target_text_length}, retrying...",
+                            file=sys.stderr,
+                        )
+                        retry_badcase_times += 1
+                        continue
+                    else:
+                        break
                 else:
                     break
-            else:
-                break
-        return self.audio_vae.decode(latent_pred.to(torch.float32)).squeeze(1).cpu()
-    
+
+        if not streaming:
+            decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)).squeeze(1).cpu()
+            yield decode_audio
+
     @torch.inference_mode()
     def build_prompt_cache(
         self,
@@ -286,88 +489,97 @@ class VoxCPMModel(nn.Module):
     ):
         """
         Build prompt cache for subsequent fast generation.
-        
+
         Args:
             prompt_text: prompt text (required)
             prompt_wav_path: prompt audio path (required)
-            
+
         Returns:
-            prompt_cache: dict with text tokens and audio features
+            prompt_cache: dict with prompt_text (raw text) and audio features.
+                         Text tokenization will be done during generation for consistency.
         """
         if not prompt_text or not prompt_wav_path:
             raise ValueError("prompt_text and prompt_wav_path are required")
-        
-        # build text tokens
-        text_token = torch.LongTensor(self.text_tokenizer(prompt_text))
 
         # load audio
         audio, sr = torchaudio.load(prompt_wav_path)
         if audio.size(0) > 1:
             audio = audio.mean(dim=0, keepdim=True)
-            
+
         if sr != self.sample_rate:
             audio = torchaudio.functional.resample(audio, sr, self.sample_rate)
 
         patch_len = self.patch_size * self.chunk_size
 
         if audio.size(1) % patch_len != 0:
-            audio = torch.nn.functional.pad(audio, (0, patch_len - audio.size(1) % patch_len))
+            # Left padding: pad at the beginning of the audio to keep valid audio data at the end of the sequence
+            padding_size = patch_len - audio.size(1) % patch_len
+            audio = torch.nn.functional.pad(audio, (padding_size, 0))
 
         # extract audio features
-        audio_feat = self.audio_vae.encode(audio.cuda(), self.sample_rate).cpu()
+        audio_feat = self.audio_vae.encode(audio.to(self.device), self.sample_rate).cpu()
 
         audio_feat = audio_feat.view(
             self.audio_vae.latent_dim,
             -1,
             self.patch_size,
-        ).permute(1, 2, 0) # (D, T, P)
-        audio_feat = audio_feat[:-1, ...] # trick: remove the last padding token
-        # build prompt cache
+        ).permute(
+            1, 2, 0
+        )  # (D, T, P)
+        # build prompt cache - only save raw text and audio features
         prompt_cache = {
-            "text_token": text_token,
+            "prompt_text": prompt_text,
             "audio_feat": audio_feat,
         }
-        
+
         return prompt_cache
 
-    
     def merge_prompt_cache(
         self,
         original_cache: dict,
-        new_text_token: torch.Tensor,
+        new_text: str,
         new_audio_feat: torch.Tensor,
     ):
         """
         Merge original prompt cache with newly generated content to stabilize voice.
-        
+
         Args:
             original_cache: original prompt cache
-            new_text_token: newly generated text tokens
+            new_text: newly generated text
             new_audio_feat: newly generated audio features
-            
+
         Returns:
-            merged_cache: merged cache
+            merged_cache: merged cache with prompt_text and audio_feat
         """
         if original_cache is None:
             return {
-                "text_token": new_text_token,
+                "prompt_text": new_text,
                 "audio_feat": new_audio_feat,
             }
-        original_text_token = original_cache["text_token"]
+        original_prompt_text = original_cache["prompt_text"]
         original_audio_feat = original_cache["audio_feat"]
-        merged_text_token = torch.cat([original_text_token, new_text_token], dim=0)
+        # Merge text by concatenation
+        merged_prompt_text = original_prompt_text + new_text
         merged_audio_feat = torch.cat([original_audio_feat, new_audio_feat], dim=0)
 
         # build new cache
         merged_cache = {
-            "text_token": merged_text_token,
+            "prompt_text": merged_prompt_text,
             "audio_feat": merged_audio_feat,
         }
-        
+
         return merged_cache
-    
+
+    def generate_with_prompt_cache(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return next(self._generate_with_prompt_cache(*args, streaming=False, **kwargs))
+
+    def generate_with_prompt_cache_streaming(
+        self, *args, **kwargs
+    ) -> Generator[Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]], None, None]:
+        return self._generate_with_prompt_cache(*args, streaming=True, **kwargs)
+
     @torch.inference_mode()
-    def generate_with_prompt_cache(
+    def _generate_with_prompt_cache(
         self,
         target_text: str,
         prompt_cache: dict,
@@ -378,10 +590,12 @@ class VoxCPMModel(nn.Module):
         retry_badcase: bool = False,
         retry_badcase_max_times: int = 3,
         retry_badcase_ratio_threshold: float = 6.0,
-    ):
+        streaming: bool = False,
+        streaming_prefix_len: int = 3,
+    ) -> Generator[Tuple[torch.Tensor, torch.Tensor, Union[torch.Tensor, List[torch.Tensor]]], None, None]:
         """
         Generate audio using pre-built prompt cache.
-        
+
         Args:
             target_text: Text to convert to speech
             prompt_cache: Cache built by build_prompt_cache (can be None)
@@ -392,20 +606,28 @@ class VoxCPMModel(nn.Module):
             retry_badcase: Whether to retry on bad cases
             retry_badcase_max_times: Maximum retry attempts
             retry_badcase_ratio_threshold: Threshold for audio-to-text ratio
-            
+            streaming: Whether to return a generator of audio chunks
+            streaming_prefix_len: Number of prefix audio patches to use for streaming mode
+
         Returns:
-            tuple: (decoded audio tensor, new text tokens, new audio features)
+            Generator of Tuple containing:
+                - Decoded audio tensor for the current step if ``streaming=True``, else final decoded audio tensor
+                - Tensor of new text tokens
+                - New audio features up to the current step as a List if ``streaming=True``, else as a concatenated Tensor
         """
+        if retry_badcase and streaming:
+            warnings.warn("Retry on bad cases is not supported in streaming mode, setting retry_badcase=False.")
+            retry_badcase = False
         # get prompt from cache
         if prompt_cache is None:
-            prompt_text_token = torch.empty(0, dtype=torch.int32)
             prompt_audio_feat = torch.empty((0, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32)
+            text = target_text
         else:
-            prompt_text_token = prompt_cache["text_token"]
             prompt_audio_feat = prompt_cache["audio_feat"]
-        # build target text tokens
-        target_text_token = torch.LongTensor(self.text_tokenizer(target_text))
-        text_token = torch.cat([prompt_text_token, target_text_token], dim=0)
+            prompt_text = prompt_cache["prompt_text"]
+            text = prompt_text + target_text
+
+        text_token = torch.LongTensor(self.text_tokenizer(text))
         text_token = torch.cat(
             [
                 text_token,
@@ -418,6 +640,8 @@ class VoxCPMModel(nn.Module):
             dim=-1,
         )
 
+        target_text_token = torch.LongTensor(self.text_tokenizer(target_text))
+
         audio_length = prompt_audio_feat.size(0)
         text_length = text_token.shape[0]
         text_pad_token = torch.zeros(audio_length, dtype=torch.int32, device=text_token.device)
@@ -428,47 +652,74 @@ class VoxCPMModel(nn.Module):
         )
         text_token = torch.cat([text_token, text_pad_token])
         audio_feat = torch.cat([audio_pad_feat, prompt_audio_feat], dim=0)
-        text_mask = torch.cat([torch.ones(text_length), torch.zeros(audio_length)]).type(torch.int32).to(text_token.device)
-        audio_mask = torch.cat([torch.zeros(text_length), torch.ones(audio_length)]).type(torch.int32).to(text_token.device)
+        text_mask = (
+            torch.cat([torch.ones(text_length), torch.zeros(audio_length)]).type(torch.int32).to(text_token.device)
+        )
+        audio_mask = (
+            torch.cat([torch.zeros(text_length), torch.ones(audio_length)]).type(torch.int32).to(text_token.device)
+        )
 
         text_token = text_token.unsqueeze(0).to(self.device)
         text_mask = text_mask.unsqueeze(0).to(self.device)
-        audio_feat = audio_feat.unsqueeze(0).to(self.device).to(torch.bfloat16)
+        audio_feat = audio_feat.unsqueeze(0).to(self.device).to(get_dtype(self.config.dtype))
         audio_mask = audio_mask.unsqueeze(0).to(self.device)
-    
+
         # run inference
         target_text_length = len(self.text_tokenizer(target_text))
         retry_badcase_times = 0
         while retry_badcase_times < retry_badcase_max_times:
-            latent_pred, pred_audio_feat = self.inference(
+            inference_result = self._inference(
                 text_token,
                 text_mask,
                 audio_feat,
                 audio_mask,
                 min_len=min_len,
-                max_len=int(target_text_length * retry_badcase_ratio_threshold + 10) if retry_badcase else max_len,
+                max_len=min(
+                    int(target_text_length * retry_badcase_ratio_threshold + 10), max_len
+                ),  # avoid too long audio
                 inference_timesteps=inference_timesteps,
                 cfg_value=cfg_value,
+                streaming=streaming,
+                streaming_prefix_len=streaming_prefix_len,
             )
-            if retry_badcase:
-                if pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold:
-                    print(f"  Badcase detected, audio_text_ratio={pred_audio_feat.shape[0] / target_text_length}, retrying...")
-                    retry_badcase_times += 1
-                    continue
+            if streaming:
+                patch_len = self.patch_size * self.chunk_size
+                for latent_pred, pred_audio_feat in inference_result:
+                    decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
+                    decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu()
+                    yield (decode_audio, target_text_token, pred_audio_feat)
+                break
+            else:
+                latent_pred, pred_audio_feat = next(inference_result)
+                if retry_badcase:
+                    if pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold:
+                        print(
+                            f"  Badcase detected, audio_text_ratio={pred_audio_feat.shape[0] / target_text_length}, retrying...",
+                            file=sys.stderr,
+                        )
+                        retry_badcase_times += 1
+                        continue
+                    else:
+                        break
                 else:
                     break
+        if not streaming:
+            decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32))
+            patch_len = self.patch_size * self.chunk_size
+            if audio_mask.sum().item() > 0:
+                decode_audio = decode_audio[..., patch_len * (streaming_prefix_len - 1) :].squeeze(1).cpu()
             else:
-                break
-        decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)).squeeze(1).cpu()
-            
-        return (
-            decode_audio,
-            target_text_token,
-            pred_audio_feat
-        )
+                decode_audio = decode_audio[..., :].squeeze(1).cpu()
+            yield (decode_audio, target_text_token, pred_audio_feat)
+
+    def inference(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
+        return next(self._inference(*args, streaming=False, **kwargs))
+
+    def inference_streaming(self, *args, **kwargs) -> Generator[Tuple[torch.Tensor, List[torch.Tensor]], None, None]:
+        return self._inference(*args, streaming=True, **kwargs)
 
     @torch.inference_mode()
-    def inference(
+    def _inference(
         self,
         text: torch.Tensor,
         text_mask: torch.Tensor,
@@ -478,12 +729,14 @@ class VoxCPMModel(nn.Module):
         max_len: int = 2000,
         inference_timesteps: int = 10,
         cfg_value: float = 2.0,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        streaming: bool = False,
+        streaming_prefix_len: int = 3,
+    ) -> Generator[Tuple[torch.Tensor, Union[torch.Tensor, List[torch.Tensor]]], None, None]:
         """Core inference method for audio generation.
-        
+
         This is the main inference loop that generates audio features
         using the language model and diffusion transformer.
-        
+
         Args:
             text: Input text tokens
             text_mask: Mask for text tokens
@@ -493,22 +746,23 @@ class VoxCPMModel(nn.Module):
             max_len: Maximum generation length
             inference_timesteps: Number of diffusion steps
             cfg_value: Classifier-free guidance value
-            
+            streaming: Whether to yield each step latent feature or just the final result
+
         Returns:
-            Tuple containing:
-                - Predicted latent features
-                - Predicted audio feature sequence
+            Generator of Tuple containing:
+                - Predicted latent feature at the current step if ``streaming=True``, else final latent features
+                - Predicted audio feature sequence so far as a List if ``streaming=True``, else as a concatenated Tensor
         """
         B, T, P, D = feat.shape
 
         feat_embed = self.feat_encoder(feat)  # [b, t, h_feat]
         feat_embed = self.enc_to_lm_proj(feat_embed)
-        
+
         if self.config.lm_config.use_mup:
             scale_emb = self.config.lm_config.scale_emb
         else:
             scale_emb = 1.0
-       
+
         text_embed = self.base_lm.embed_tokens(text) * scale_emb
         combined_embed = text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed
 
@@ -516,16 +770,26 @@ class VoxCPMModel(nn.Module):
         pred_feat_seq = []  # b, t, p, d
         curr_embed = None
 
+        # Prepare prompt context patches for streaming mode
+        # When there's a prompt audio, use its last (streaming_prefix_len - 1) patches as initial context
+        prompt_context_patches = []
+        audio_patch_count = int(feat_mask.sum().item())
+        if audio_patch_count > 0:
+            context_len = min(streaming_prefix_len - 1, audio_patch_count)
+            # Take the last context_len patches from prompt audio as initial context
+            # Split into list of [b, 1, p, d] tensors to match pred_feat_seq format
+            prompt_context_patches = list(feat[:, -context_len:, :, :].split(1, dim=1))
+            pred_feat_seq = prompt_context_patches + pred_feat_seq
+
         enc_outputs, kv_cache_tuple = self.base_lm(
             inputs_embeds=combined_embed,
             is_causal=True,
         )
         self.base_lm.kv_cache.fill_caches(kv_cache_tuple)
-        
+
         enc_outputs = self.fsq_layer(enc_outputs) * feat_mask.unsqueeze(-1) + enc_outputs * text_mask.unsqueeze(-1)
         lm_hidden = enc_outputs[:, -1, :]
 
-         
         residual_enc_outputs, residual_kv_cache_tuple = self.residual_lm(
             inputs_embeds=enc_outputs + feat_mask.unsqueeze(-1) * feat_embed,
             is_causal=True,
@@ -533,7 +797,6 @@ class VoxCPMModel(nn.Module):
         self.residual_lm.kv_cache.fill_caches(residual_kv_cache_tuple)
         residual_hidden = residual_enc_outputs[:, -1, :]
 
-
         for i in tqdm(range(max_len)):
             dit_hidden_1 = self.lm_to_dit_proj(lm_hidden)  # [b, h_dit]
             dit_hidden_2 = self.res_to_dit_proj(residual_hidden)  # [b, h_dit]
@@ -548,58 +811,175 @@ class VoxCPMModel(nn.Module):
             ).transpose(
                 1, 2
             )  # [b, p, d]
-            
-            curr_embed = self.feat_encoder_step(pred_feat.unsqueeze(1))  # b, 1, c
+
+            curr_embed = self.feat_encoder(pred_feat.unsqueeze(1))  # b, 1, c
             curr_embed = self.enc_to_lm_proj(curr_embed)
-            
+
             pred_feat_seq.append(pred_feat.unsqueeze(1))  # b, 1, p, d
             prefix_feat_cond = pred_feat
-            
+
+            if streaming:
+                # return the last three predicted latent features to provide enough context for smooth decoding
+                pred_feat_chunk = torch.cat(pred_feat_seq[-streaming_prefix_len:], dim=1)
+                feat_pred = rearrange(pred_feat_chunk, "b t p d -> b d (t p)", b=B, p=self.patch_size)
+
+                yield feat_pred, pred_feat_seq
+
             stop_flag = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden))).argmax(dim=-1)[0].cpu().item()
             if i > min_len and stop_flag == 1:
                 break
-    
+
             lm_hidden = self.base_lm.forward_step(
                 curr_embed[:, 0, :], torch.tensor([self.base_lm.kv_cache.step()], device=curr_embed.device)
             ).clone()
-           
 
             lm_hidden = self.fsq_layer(lm_hidden)
             residual_hidden = self.residual_lm.forward_step(
-                lm_hidden + curr_embed[:, 0, :], torch.tensor([self.residual_lm.kv_cache.step()], device=curr_embed.device)
+                lm_hidden + curr_embed[:, 0, :],
+                torch.tensor([self.residual_lm.kv_cache.step()], device=curr_embed.device),
             ).clone()
-                
-        pred_feat_seq = torch.cat(pred_feat_seq, dim=1)  # b, t, p, d
 
-        feat_pred = rearrange(pred_feat_seq, "b t p d -> b d (t p)", b=B, p=self.patch_size)
-        feat_pred = feat_pred[..., 1:-1] # trick: remove the first and last token
-        return feat_pred, pred_feat_seq.squeeze(0).cpu()
+        if not streaming:
+            pred_feat_seq = torch.cat(pred_feat_seq, dim=1)  # b, t, p, d
+            feat_pred = rearrange(pred_feat_seq, "b t p d -> b d (t p)", b=B, p=self.patch_size)
+            yield feat_pred, pred_feat_seq.squeeze(0).cpu()
 
     @classmethod
-    def from_local(cls, path: str):
+    def from_local(cls, path: str, optimize: bool = True, training: bool = False, lora_config: LoRAConfig = None):
         config = VoxCPMConfig.model_validate_json(open(os.path.join(path, "config.json")).read())
-
         tokenizer = LlamaTokenizerFast.from_pretrained(path)
-
-        audio_vae = AudioVAE()
-        vae_state_dict = torch.load(
-            os.path.join(path, "audiovae.pth"),
-            map_location="cpu",
-            weights_only=True,
-        )["state_dict"]
-
-        model = cls(config, tokenizer, audio_vae)
-        lm_dtype = get_dtype(config.dtype)
-        model = model.to(lm_dtype)
+        audio_vae_config = getattr(config, "audio_vae_config", None)
+        audio_vae = AudioVAE(config=audio_vae_config) if audio_vae_config else AudioVAE()
+        # Try to load AudioVAE from safetensors first, fallback to pytorch
+        audiovae_safetensors_path = os.path.join(path, "audiovae.safetensors")
+        audiovae_pth_path = os.path.join(path, "audiovae.pth")
+        if os.path.exists(audiovae_safetensors_path) and SAFETENSORS_AVAILABLE:
+            print(f"Loading AudioVAE from safetensors: {audiovae_safetensors_path}", file=sys.stderr)
+            vae_state_dict = load_file(audiovae_safetensors_path, device="cpu")
+        elif os.path.exists(audiovae_pth_path):
+            print(f"Loading AudioVAE from pytorch: {audiovae_pth_path}", file=sys.stderr)
+            checkpoint = torch.load(
+                audiovae_pth_path,
+                map_location="cpu",
+                weights_only=True,
+            )
+            vae_state_dict = checkpoint.get("state_dict", checkpoint)
+        else:
+            raise FileNotFoundError(
+                f"AudioVAE checkpoint not found. Expected either {audiovae_safetensors_path} or {audiovae_pth_path}"
+            )
+        model = cls(config, tokenizer, audio_vae, lora_config)
+        if not training:
+            lm_dtype = get_dtype(model.config.dtype)
+            model = model.to(lm_dtype)
+        else:  # training mode
+            for name, param in model.named_parameters():
+                if "audio_vae" in name:  # freeze VAE weights
+                    param.requires_grad = False
+                    continue
+                if lora_config is not None:
+                    if "lora" not in name:  # freeze non-LoRA weights
+                        param.requires_grad = False
         model.audio_vae = model.audio_vae.to(torch.float32)
 
-        model_state_dict = torch.load(
-            os.path.join(path, "pytorch_model.bin"),
-            map_location="cpu",
-            weights_only=True,
-        )["state_dict"]
+        # Try to load from safetensors first, fallback to pytorch_model.bin
+        safetensors_path = os.path.join(path, "model.safetensors")
+        pytorch_model_path = os.path.join(path, "pytorch_model.bin")
+
+        if os.path.exists(safetensors_path) and SAFETENSORS_AVAILABLE:
+            print(f"Loading model from safetensors: {safetensors_path}", file=sys.stderr)
+            model_state_dict = load_file(safetensors_path)
+        elif os.path.exists(pytorch_model_path):
+            print(f"Loading model from pytorch_model.bin: {pytorch_model_path}", file=sys.stderr)
+            checkpoint = torch.load(
+                pytorch_model_path,
+                map_location="cpu",
+                weights_only=True,
+            )
+            model_state_dict = checkpoint.get("state_dict", checkpoint)
+        else:
+            raise FileNotFoundError(f"Model file not found. Expected either {safetensors_path} or {pytorch_model_path}")
 
         for kw, val in vae_state_dict.items():
             model_state_dict[f"audio_vae.{kw}"] = val
-        model.load_state_dict(model_state_dict, strict=True)
-        return model.to(model.device).eval().optimize()
+
+        # LoRALinear holds weight/bias directly, compatible with nn.Linear state_dict keys.
+        # Using strict=False since pretrained weights don't contain lora_A/lora_B.
+        model.load_state_dict(model_state_dict, strict=False)
+        if training:
+            return model
+        return model.to(model.device).eval().optimize(disable=not optimize)
+
+    # ------------------------------------------------------------------ #
+    # LoRA Weight Management
+    # ------------------------------------------------------------------ #
+    def _iter_lora_modules(self):
+        """Iterate over all LoRA modules."""
+        from ..modules.layers.lora import LoRALinear
+
+        for module in self.modules():
+            if isinstance(module, LoRALinear):
+                yield module
+
+    def load_lora_weights(self, lora_path: str, device: str = None):
+        """
+        Load LoRA weights from file, supports calling after torch.compile.
+        Uses named_parameters() to handle compile's _orig_mod wrapper.
+        Supports both safetensors and pytorch formats.
+
+        Args:
+            lora_path: Checkpoint path (directory or .safetensors/.ckpt file)
+            device: Target device, defaults to model's current device
+        Returns:
+            tuple: (loaded_keys, skipped_keys)
+        """
+        from pathlib import Path
+
+        device = device or self.device
+        lora_p = Path(lora_path)
+
+        # Try safetensors first, then fallback to .ckpt
+        if lora_p.is_dir():
+            safetensors_file = lora_p / "lora_weights.safetensors"
+            ckpt_file = lora_p / "lora_weights.ckpt"
+        else:
+            safetensors_file = lora_p if lora_p.suffix == ".safetensors" else None
+            ckpt_file = lora_p if lora_p.suffix in [".ckpt", ".pth"] else None
+
+        # Load from safetensors if available
+        if safetensors_file and safetensors_file.exists() and SAFETENSORS_AVAILABLE:
+            state_dict = load_file(str(safetensors_file), device=device)
+        elif ckpt_file and ckpt_file.exists():
+            ckpt = torch.load(ckpt_file, map_location=device, weights_only=False)
+            state_dict = ckpt.get("state_dict", ckpt)
+        else:
+            raise FileNotFoundError(f"LoRA checkpoint not found. Expected either {safetensors_file} or {ckpt_file}")
+
+        # Build param mapping (handle torch.compile's _orig_mod prefix)
+        model_params = dict(self.named_parameters())
+        key_mapping = {k.replace("._orig_mod.", "."): k for k in model_params if "._orig_mod." in k}
+
+        loaded_keys, skipped_keys = [], []
+        for key, value in state_dict.items():
+            target_key = key if key in model_params else key_mapping.get(key)
+            if target_key:
+                model_params[target_key].data.copy_(value.to(device))
+                loaded_keys.append(key)
+            else:
+                skipped_keys.append(key)
+
+        return loaded_keys, skipped_keys
+
+    def set_lora_enabled(self, enabled: bool):
+        """Enable/disable all LoRA layers."""
+        for module in self._iter_lora_modules():
+            module.set_enabled(enabled)
+
+    def reset_lora_weights(self):
+        """Reset all LoRA weights (A: kaiming, B: zeros), effectively unloading LoRA."""
+        for module in self._iter_lora_modules():
+            module.reset_lora_parameters()
+
+    def get_lora_state_dict(self) -> dict:
+        """Get all LoRA parameters (lora_A/lora_B)."""
+        return {name: param.data.clone() for name, param in self.named_parameters() if "lora_" in name}

src/voxcpm/modules/audiovae/__init__.py

@@ -1 +1,2 @@
-from .audio_vae import AudioVAE
+from .audio_vae import AudioVAE, AudioVAEConfig
+from .audio_vae_v2 import AudioVAE as AudioVAEV2, AudioVAEConfig as AudioVAEConfigV2

src/voxcpm/modules/audiovae/audio_vae.py

@@ -1,11 +1,12 @@
 import math
-from typing import List, Union
+from typing import List
 
 import numpy as np
 import torch
 from torch import nn
 import torch.nn.functional as F
 from torch.nn.utils import weight_norm
+from pydantic import BaseModel
 
 
 def WNConv1d(*args, **kwargs):
@@ -266,6 +267,17 @@ class CausalDecoder(nn.Module):
         return self.model(x)
 
 
+class AudioVAEConfig(BaseModel):
+    encoder_dim: int = 128
+    encoder_rates: List[int] = [2, 5, 8, 8]
+    latent_dim: int = 64
+    decoder_dim: int = 1536
+    decoder_rates: List[int] = [8, 8, 5, 2]
+    depthwise: bool = True
+    sample_rate: int = 16000
+    use_noise_block: bool = False
+
+
 class AudioVAE(nn.Module):
     """
     Args:
@@ -273,17 +285,23 @@ class AudioVAE(nn.Module):
 
     def __init__(
         self,
-        encoder_dim: int = 128,
-        encoder_rates: List[int] = [2, 5, 8, 8],
-        latent_dim: int = 64,
-        decoder_dim: int = 1536,
-        decoder_rates: List[int] = [8, 8, 5, 2],
-        depthwise: bool = True,
-        sample_rate: int = 16000,
-        use_noise_block: bool = False,
+        config: AudioVAEConfig = None,
     ):
+        # 如果没有传入config，使用默认配置
+        if config is None:
+            config = AudioVAEConfig()
+
         super().__init__()
 
+        encoder_dim = config.encoder_dim
+        encoder_rates = config.encoder_rates
+        latent_dim = config.latent_dim
+        decoder_dim = config.decoder_dim
+        decoder_rates = config.decoder_rates
+        depthwise = config.depthwise
+        sample_rate = config.sample_rate
+        use_noise_block = config.use_noise_block
+
         self.encoder_dim = encoder_dim
         self.encoder_rates = encoder_rates
         self.decoder_dim = decoder_dim

src/voxcpm/modules/audiovae/audio_vae_v2.py

@@ -0,0 +1,486 @@
+import math
+from typing import List, Optional
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.utils import weight_norm
+from pydantic import BaseModel
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class CausalConv1d(nn.Conv1d):
+    def __init__(self, *args, padding: int = 0, output_padding: int = 0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__padding = padding
+        self.__output_padding = output_padding
+
+    def forward(self, x):
+        x_pad = F.pad(x, (self.__padding * 2 - self.__output_padding, 0))
+        return super().forward(x_pad)
+
+
+class CausalTransposeConv1d(nn.ConvTranspose1d):
+    def __init__(self, *args, padding: int = 0, output_padding: int = 0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__padding = padding
+        self.__output_padding = output_padding
+
+    def forward(self, x):
+        return super().forward(x)[..., : -(self.__padding * 2 - self.__output_padding)]
+
+
+def WNCausalConv1d(*args, **kwargs):
+    return weight_norm(CausalConv1d(*args, **kwargs))
+
+
+def WNCausalTransposeConv1d(*args, **kwargs):
+    return weight_norm(CausalTransposeConv1d(*args, **kwargs))
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+
+
+class CausalResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1, kernel: int = 7, groups: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            WNCausalConv1d(
+                dim,
+                dim,
+                kernel_size=kernel,
+                dilation=dilation,
+                padding=pad,
+                groups=groups,
+            ),
+            Snake1d(dim),
+            WNCausalConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        assert pad == 0
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class CausalEncoderBlock(nn.Module):
+    def __init__(self, output_dim: int = 16, input_dim=None, stride: int = 1, groups=1):
+        super().__init__()
+        input_dim = input_dim or output_dim // 2
+        self.block = nn.Sequential(
+            CausalResidualUnit(input_dim, dilation=1, groups=groups),
+            CausalResidualUnit(input_dim, dilation=3, groups=groups),
+            CausalResidualUnit(input_dim, dilation=9, groups=groups),
+            Snake1d(input_dim),
+            WNCausalConv1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class CausalEncoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        latent_dim: int = 32,
+        strides: list = [2, 4, 8, 8],
+        depthwise: bool = False,
+    ):
+        super().__init__()
+        # Create first convolution
+        self.block = [WNCausalConv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in strides:
+            d_model *= 2
+            groups = d_model // 2 if depthwise else 1
+            self.block += [CausalEncoderBlock(output_dim=d_model, stride=stride, groups=groups)]
+
+        groups = d_model if depthwise else 1
+
+        # Create two convolution, for mu and logvar
+        self.fc_mu = WNCausalConv1d(d_model, latent_dim, kernel_size=3, padding=1)
+        self.fc_logvar = WNCausalConv1d(d_model, latent_dim, kernel_size=3, padding=1)
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+    def forward(self, x):
+        hidden_state = self.block(x)
+        return {
+            "hidden_state": hidden_state,
+            "mu": self.fc_mu(hidden_state),
+            "logvar": self.fc_logvar(hidden_state),
+        }
+
+
+class NoiseBlock(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.linear = WNCausalConv1d(dim, dim, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        B, C, T = x.shape
+        noise = torch.randn((B, 1, T), device=x.device, dtype=x.dtype)
+        h = self.linear(x)
+        n = noise * h
+        x = x + n
+        return x
+
+
+class CausalDecoderBlock(nn.Module):
+    def __init__(
+        self,
+        input_dim: int = 16,
+        output_dim: int = 8,
+        stride: int = 1,
+        groups=1,
+        use_noise_block: bool = False,
+    ):
+        super().__init__()
+        layers = [
+            Snake1d(input_dim),
+            WNCausalTransposeConv1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+        ]
+        if use_noise_block:
+            layers.append(NoiseBlock(output_dim))
+        layers.extend(
+            [
+                CausalResidualUnit(output_dim, dilation=1, groups=groups),
+                CausalResidualUnit(output_dim, dilation=3, groups=groups),
+                CausalResidualUnit(output_dim, dilation=9, groups=groups),
+            ]
+        )
+        self.block = nn.Sequential(*layers)
+        self.input_channels = input_dim
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class TransposeLastTwoDim(torch.nn.Module):
+    def forward(self, x):
+        return torch.transpose(x, -1, -2)
+
+
+class SampleRateConditionLayer(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        sr_bin_buckets: int = None,
+        cond_type: str = "scale_bias",
+        cond_dim: int = 128,
+        out_layer: bool = False,
+    ):
+        super().__init__()
+
+        self.cond_type, out_layer_in_dim = cond_type, input_dim
+
+        if cond_type == "scale_bias":
+            self.scale_embed = nn.Embedding(sr_bin_buckets, input_dim)
+            self.bias_embed = nn.Embedding(sr_bin_buckets, input_dim)
+            nn.init.ones_(self.scale_embed.weight)
+            nn.init.zeros_(self.bias_embed.weight)
+        elif cond_type == "scale_bias_init":
+            self.scale_embed = nn.Embedding(sr_bin_buckets, input_dim)
+            self.bias_embed = nn.Embedding(sr_bin_buckets, input_dim)
+            nn.init.normal_(self.scale_embed.weight, mean=1)
+            nn.init.normal_(self.bias_embed.weight)
+        elif cond_type == "add":
+            self.cond_embed = nn.Embedding(sr_bin_buckets, input_dim)
+            nn.init.normal_(self.cond_embed.weight)
+        elif cond_type == "concat":
+            self.cond_embed = nn.Embedding(sr_bin_buckets, cond_dim)
+            assert out_layer, "out_layer must be True for concat cond_type"
+            out_layer_in_dim = input_dim + cond_dim
+        else:
+            raise ValueError(f"Invalid cond_type: {cond_type}")
+
+        if out_layer:
+            self.out_layer = nn.Sequential(
+                Snake1d(out_layer_in_dim),
+                WNCausalConv1d(out_layer_in_dim, input_dim, kernel_size=1),
+            )
+        else:
+            self.out_layer = nn.Identity()
+
+    def forward(self, x, sr_cond):
+        if self.cond_type == "scale_bias" or self.cond_type == "scale_bias_init":
+            x = x * self.scale_embed(sr_cond).unsqueeze(-1) + self.bias_embed(sr_cond).unsqueeze(-1)
+        elif self.cond_type == "add":
+            x = x + self.cond_embed(sr_cond).unsqueeze(-1)
+        elif self.cond_type == "concat":
+            x = torch.cat([x, self.cond_embed(sr_cond).unsqueeze(-1).repeat(1, 1, x.shape[-1])], dim=1)
+
+        return self.out_layer(x)
+
+
+class CausalDecoder(nn.Module):
+    def __init__(
+        self,
+        input_channel,
+        channels,
+        rates,
+        depthwise: bool = False,
+        d_out: int = 1,
+        use_noise_block: bool = False,
+        sr_bin_boundaries: List[int] = None,
+        cond_type: str = "scale_bias",
+        cond_dim: int = 128,
+        cond_out_layer: bool = False,
+    ):
+        super().__init__()
+
+        # Add first conv layer
+        if depthwise:
+            layers = [
+                WNCausalConv1d(input_channel, input_channel, kernel_size=7, padding=3, groups=input_channel),
+                WNCausalConv1d(input_channel, channels, kernel_size=1),
+            ]
+        else:
+            layers = [WNCausalConv1d(input_channel, channels, kernel_size=7, padding=3)]
+
+        # Add upsampling + MRF blocks
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            groups = output_dim if depthwise else 1
+            layers += [
+                CausalDecoderBlock(
+                    input_dim,
+                    output_dim,
+                    stride,
+                    groups=groups,
+                    use_noise_block=use_noise_block,
+                )
+            ]
+
+        # Add final conv layer
+        layers += [
+            Snake1d(output_dim),
+            WNCausalConv1d(output_dim, d_out, kernel_size=7, padding=3),
+            nn.Tanh(),
+        ]
+
+        if sr_bin_boundaries is None:
+            self.model = nn.Sequential(*layers)
+            self.sr_bin_boundaries = None
+        else:
+            self.model = nn.ModuleList(layers)
+
+            self.register_buffer("sr_bin_boundaries", torch.tensor(sr_bin_boundaries, dtype=torch.int32))
+            self.sr_bin_buckets = len(sr_bin_boundaries) + 1
+
+            cond_layers = []
+            for layer in self.model:
+                if layer.__class__.__name__ == "CausalDecoderBlock":
+                    cond_layers.append(
+                        SampleRateConditionLayer(
+                            input_dim=layer.input_channels,
+                            sr_bin_buckets=self.sr_bin_buckets,
+                            cond_type=cond_type,
+                            cond_dim=cond_dim,
+                            out_layer=cond_out_layer,
+                        )
+                    )
+                else:
+                    cond_layers.append(None)
+            self.sr_cond_model = nn.ModuleList(cond_layers)
+
+    def get_sr_idx(self, sr):
+        return torch.bucketize(sr, self.sr_bin_boundaries)
+
+    def forward(self, x, sr_cond=None):
+        if self.sr_bin_boundaries is not None:
+            # assert sr_cond is not None
+            sr_cond = self.get_sr_idx(sr_cond)
+
+            for layer, sr_cond_layer in zip(self.model, self.sr_cond_model):
+                if sr_cond_layer is not None:
+                    x = sr_cond_layer(x, sr_cond)
+                x = layer(x)
+            return x
+        else:
+            return self.model(x)
+
+
+class AudioVAEConfig(BaseModel):
+    encoder_dim: int = 128
+    encoder_rates: List[int] = [2, 5, 8, 8]
+    latent_dim: int = 64
+    decoder_dim: int = 2048
+    decoder_rates: List[int] = [8, 6, 5, 2, 2, 2]
+    depthwise: bool = True
+    sample_rate: int = 16000
+    out_sample_rate: int = 48000
+    use_noise_block: bool = False
+    sr_bin_boundaries: Optional[List[int]] = [20000, 30000, 40000]
+    cond_type: str = "scale_bias"
+    cond_dim: int = 128
+    cond_out_layer: bool = False
+
+
+class AudioVAE(nn.Module):
+    """
+    Args:
+    """
+
+    def __init__(
+        self,
+        config: AudioVAEConfig = None,
+    ):
+        # 如果没有传入config，使用默认配置
+        if config is None:
+            config = AudioVAEConfig()
+
+        super().__init__()
+
+        encoder_dim = config.encoder_dim
+        encoder_rates = config.encoder_rates
+        latent_dim = config.latent_dim
+        decoder_dim = config.decoder_dim
+        decoder_rates = config.decoder_rates
+        depthwise = config.depthwise
+        sample_rate = config.sample_rate
+        out_sample_rate = config.out_sample_rate
+        use_noise_block = config.use_noise_block
+        sr_bin_boundaries = config.sr_bin_boundaries
+        cond_type = config.cond_type
+        cond_dim = config.cond_dim
+        cond_out_layer = config.cond_out_layer
+
+        self.encoder_dim = encoder_dim
+        self.encoder_rates = encoder_rates
+        self.decoder_dim = decoder_dim
+        self.decoder_rates = decoder_rates
+        self.depthwise = depthwise
+
+        self.use_noise_block = use_noise_block
+
+        if latent_dim is None:
+            latent_dim = encoder_dim * (2 ** len(encoder_rates))
+
+        self.latent_dim = latent_dim
+        self.hop_length = np.prod(encoder_rates)
+        self.encoder = CausalEncoder(
+            encoder_dim,
+            latent_dim,
+            encoder_rates,
+            depthwise=depthwise,
+        )
+
+        self.decoder = CausalDecoder(
+            latent_dim,
+            decoder_dim,
+            decoder_rates,
+            depthwise=depthwise,
+            use_noise_block=use_noise_block,
+            sr_bin_boundaries=sr_bin_boundaries,
+            cond_type=cond_type,
+            cond_dim=cond_dim,
+            cond_out_layer=cond_out_layer,
+        )
+        self.sample_rate = sample_rate
+        self.out_sample_rate = out_sample_rate
+        self.sr_bin_boundaries = sr_bin_boundaries
+        self.chunk_size = math.prod(encoder_rates)
+
+    def preprocess(self, audio_data, sample_rate):
+        if sample_rate is None:
+            sample_rate = self.sample_rate
+        assert sample_rate == self.sample_rate
+        pad_to = self.hop_length
+        length = audio_data.shape[-1]
+        right_pad = math.ceil(length / pad_to) * pad_to - length
+        audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+        return audio_data
+
+    def decode(self, z: torch.Tensor, sr_cond: torch.Tensor = None):
+        """Decode given latent codes and return audio data
+
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+            Quantized continuous representation of input
+        length : int, optional
+            Number of samples in output audio, by default None
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "audio" : Tensor[B x 1 x length]
+                Decoded audio data.
+        """
+        if self.sr_bin_boundaries is not None:
+            # use default output sample rate
+            if sr_cond is None:
+                sr_cond = torch.tensor([self.out_sample_rate], device=z.device, dtype=torch.int32)
+        return self.decoder(z, sr_cond)
+
+    def encode(self, audio_data: torch.Tensor, sample_rate: int):
+        """
+        Args:
+            audio_data: Tensor[B x 1 x T]
+            sample_rate: int
+        Returns:
+            z: Tensor[B x D x T]
+        """
+        if audio_data.ndim == 2:
+            audio_data = audio_data.unsqueeze(1)
+
+        audio_data = self.preprocess(audio_data, sample_rate)
+        return self.encoder(audio_data)["mu"]

src/voxcpm/modules/layers/__init__.py

@@ -1 +1 @@
-from .scalar_quantization_layer import ScalarQuantizationLayer
+from .scalar_quantization_layer import ScalarQuantizationLayer

src/voxcpm/modules/layers/lora.py

@@ -0,0 +1,130 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LoRALinear(nn.Module):
+    """
+    LoRA 线性层：直接持有 weight/bias，保持与 nn.Linear 相同的 state_dict key 结构。
+
+    state_dict 结构：
+        - weight: 原始权重（与 nn.Linear 一致）
+        - bias: 原始偏置（与 nn.Linear 一致）
+        - lora_A: LoRA 低秩矩阵 A
+        - lora_B: LoRA 低秩矩阵 B
+
+    这样设计的好处：加载预训练权重时无需做 key 转换。
+    """
+
+    def __init__(
+        self,
+        base: nn.Linear,
+        r: int,
+        alpha: float = 1.0,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        assert isinstance(base, nn.Linear), "LoRALinear only supports wrapping nn.Linear."
+
+        self.in_features = base.in_features
+        self.out_features = base.out_features
+        self.r = r
+        self.alpha = alpha
+        self._base_scaling = alpha / r if r > 0 else 0.0
+
+        # 使用 buffer 存储 scaling，这样修改值不会触发 torch.compile 重编译
+        # persistent=False 表示不保存到 state_dict，避免加载时 missing key
+        self.register_buffer("scaling", torch.tensor(self._base_scaling), persistent=False)
+
+        # 直接持有 weight 和 bias（从原始 Linear 转移过来）
+        self.weight = base.weight
+        self.bias = base.bias  # 可能是 None
+
+        # LoRA 参数
+        if r > 0:
+            self.lora_A = nn.Parameter(torch.zeros(r, self.in_features))
+            self.lora_B = nn.Parameter(torch.zeros(self.out_features, r))
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B)
+        else:
+            self.register_parameter("lora_A", None)
+            self.register_parameter("lora_B", None)
+
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 基础 Linear 计算
+        result = F.linear(x, self.weight, self.bias)
+        if self.r <= 0 or self.lora_A is None:
+            return result
+        # LoRA: result + dropout(x @ A^T @ B^T) * scaling
+        lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
+        return result + self.dropout(lora_out) * self.scaling
+
+    def reset_lora_parameters(self):
+        """重置 LoRA 参数到初始状态"""
+        if self.r > 0 and self.lora_A is not None:
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B)
+
+    def set_enabled(self, enabled: bool):
+        """启用/禁用 LoRA（通过 scaling 控制，兼容 torch.compile）"""
+        # 使用 fill_ 原地修改 buffer 值，不会触发重编译
+        self.scaling.fill_(self._base_scaling if enabled else 0.0)
+
+    @property
+    def enabled(self) -> bool:
+        return self.scaling.item() != 0.0
+
+
+def _get_parent_module(root: nn.Module, name: str) -> Optional[nn.Module]:
+    """
+    根据类似 'layers.0.self_attn.q_proj' 的全名，返回 parent module（即 q_proj 的上一级）。
+    """
+    parts = name.split(".")
+    if len(parts) == 1:
+        return root
+    parent = root
+    for p in parts[:-1]:
+        if not hasattr(parent, p):
+            return None
+        parent = getattr(parent, p)
+    return parent
+
+
+def apply_lora_to_named_linear_modules(
+    root: nn.Module,
+    *,
+    target_submodule_names: list[str],
+    r: int,
+    alpha: float,
+    dropout: float,
+) -> None:
+    """
+    在给定模块及其子模块中，对名字以 target_submodule_names 结尾的 Linear 层注入 LoRA。
+
+    例如 target_submodule_names=["q_proj", "v_proj"] 时，
+    会在所有名为 *.q_proj / *.v_proj 的 nn.Linear 上替换为 LoRALinear。
+    """
+    for full_name, module in list(root.named_modules()):
+        if not isinstance(module, nn.Linear):
+            continue
+        short_name = full_name.split(".")[-1]
+        if short_name not in target_submodule_names:
+            continue
+
+        parent = _get_parent_module(root, full_name)
+        if parent is None:
+            continue
+
+        # 用 LoRALinear 替换原始 Linear
+        lora_layer = LoRALinear(
+            base=module,
+            r=r,
+            alpha=alpha,
+            dropout=dropout,
+        )
+        setattr(parent, short_name, lora_layer)

src/voxcpm/modules/layers/scalar_quantization_layer.py

@@ -12,7 +12,7 @@ class ScalarQuantizationLayer(nn.Module):
 
         self.in_proj = nn.Linear(in_dim, latent_dim)
         self.out_proj = nn.Linear(latent_dim, out_dim)
-    
+
     def forward(self, hidden):
         hidden = self.in_proj(hidden)
         hidden = torch.tanh(hidden)
@@ -23,4 +23,4 @@ class ScalarQuantizationLayer(nn.Module):
         else:
             hidden = torch.round(hidden * self.scale) / self.scale
 
-        return self.out_proj(hidden)
+        return self.out_proj(hidden)

src/voxcpm/modules/locdit/__init__.py

@@ -1,2 +1,3 @@
 from .unified_cfm import UnifiedCFM, CfmConfig
 from .local_dit import VoxCPMLocDiT
+from .local_dit_v2 import VoxCPMLocDiT as VoxCPMLocDiTV2

src/voxcpm/modules/locdit/local_dit_v2.py

@@ -0,0 +1,116 @@
+import torch
+from ..minicpm4 import MiniCPMModel, MiniCPM4Config
+import torch.nn as nn
+import math
+
+
+class SinusoidalPosEmb(torch.nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        assert self.dim % 2 == 0, "SinusoidalPosEmb requires dim to be even"
+
+    def forward(self, x, scale=1000):
+        if x.ndim < 1:
+            x = x.unsqueeze(0)
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=x.dtype, device=device) * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        out_dim: int = None,
+    ):
+        super().__init__()
+
+        self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=True)
+        self.act = nn.SiLU()
+        if out_dim is not None:
+            time_embed_dim_out = out_dim
+        else:
+            time_embed_dim_out = time_embed_dim
+
+        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, bias=True)
+
+    def forward(self, sample):
+        sample = self.linear_1(sample)
+        sample = self.act(sample)
+        sample = self.linear_2(sample)
+        return sample
+
+
+class VoxCPMLocDiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+
+    def __init__(
+        self,
+        config: MiniCPM4Config,
+        in_channels: int = 64,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.config = config
+
+        self.in_proj = nn.Linear(in_channels, config.hidden_size, bias=True)
+        self.cond_proj = nn.Linear(in_channels, config.hidden_size, bias=True)
+        self.out_proj = nn.Linear(config.hidden_size, self.out_channels, bias=True)
+
+        self.time_embeddings = SinusoidalPosEmb(config.hidden_size)
+        self.time_mlp = TimestepEmbedding(
+            in_channels=config.hidden_size,
+            time_embed_dim=config.hidden_size,
+        )
+        self.delta_time_mlp = TimestepEmbedding(
+            in_channels=config.hidden_size,
+            time_embed_dim=config.hidden_size,
+        )
+
+        assert config.vocab_size == 0, "vocab_size must be 0 for local DiT"
+        self.decoder = MiniCPMModel(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mu: torch.Tensor,
+        t: torch.Tensor,
+        cond: torch.Tensor,
+        dt: torch.Tensor,
+    ):
+        """
+        Forward pass of DiT.
+        x: (N, C, T) tensor of inputs
+        mu: (N, C) tensor of hidden embedding
+        t: (N,) tensor of diffusion timesteps
+        cond: (N, C, T') tensor of prefix conditions
+        dt: (N,) used for mean velocity (may be supported in the future...)
+        """
+        x = self.in_proj(x.transpose(1, 2).contiguous())
+
+        cond = self.cond_proj(cond.transpose(1, 2).contiguous())
+        prefix = cond.size(1)
+
+        t = self.time_embeddings(t).to(x.dtype)
+        t = self.time_mlp(t)
+        dt = self.time_embeddings(dt).to(x.dtype)
+        dt = self.delta_time_mlp(dt)
+        t = t + dt
+
+        mu = mu.view(x.size(0), -1, x.size(-1))
+        x = torch.cat([mu, (t).unsqueeze(1), cond, x], dim=1)
+
+        hidden, _ = self.decoder(x, is_causal=False)
+        hidden = hidden[:, prefix + mu.size(1) + 1 :, :]
+        hidden = self.out_proj(hidden)
+
+        return hidden.transpose(1, 2).contiguous()

src/voxcpm/modules/locdit/unified_cfm.py

@@ -1,20 +1,29 @@
+from typing import Tuple
+
 import torch
-from typing import List
-from .local_dit import VoxCPMLocDiT
-import math
+import torch.nn.functional as F
+from torch.func import jvp
 from pydantic import BaseModel
 
+from .local_dit import VoxCPMLocDiT
+
 
 class CfmConfig(BaseModel):
-    sigma_min: float = 1e-06
+    sigma_min: float = 1e-6
     solver: str = "euler"
     t_scheduler: str = "log-norm"
+    training_cfg_rate: float = 0.1
+    inference_cfg_rate: float = 1.0
+    reg_loss_type: str = "l1"
+    ratio_r_neq_t_range: Tuple[float, float] = (0.25, 0.75)
+    noise_cond_prob_range: Tuple[float, float] = (0.0, 0.0)
+    noise_cond_scale: float = 0.0
 
 
 class UnifiedCFM(torch.nn.Module):
     def __init__(
         self,
-        in_channels,
+        in_channels: int,
         cfm_params: CfmConfig,
         estimator: VoxCPMLocDiT,
         mean_mode: bool = False,
@@ -23,12 +32,21 @@ class UnifiedCFM(torch.nn.Module):
         self.solver = cfm_params.solver
         self.sigma_min = cfm_params.sigma_min
         self.t_scheduler = cfm_params.t_scheduler
+        self.training_cfg_rate = cfm_params.training_cfg_rate
+        self.inference_cfg_rate = cfm_params.inference_cfg_rate
+        self.reg_loss_type = cfm_params.reg_loss_type
+        self.ratio_r_neq_t_range = cfm_params.ratio_r_neq_t_range
+        self.noise_cond_prob_range = cfm_params.noise_cond_prob_range
+        self.noise_cond_scale = cfm_params.noise_cond_scale
+
         self.in_channels = in_channels
         self.mean_mode = mean_mode
 
-        # Just change the architecture of the estimator here
         self.estimator = estimator
 
+    # ------------------------------------------------------------------ #
+    # Inference
+    # ------------------------------------------------------------------ #
     @torch.inference_mode()
     def forward(
         self,
@@ -38,36 +56,28 @@ class UnifiedCFM(torch.nn.Module):
         cond: torch.Tensor,
         temperature: float = 1.0,
         cfg_value: float = 1.0,
-        sway_sampling_coef: float = 1.0, 
+        sway_sampling_coef: float = 1.0,
         use_cfg_zero_star: bool = True,
     ):
-        """Forward diffusion
-
-        Args:
-            mu (torch.Tensor): output of encoder
-                shape: (batch_size, n_feats)
-            n_timesteps (int): number of diffusion steps
-            cond: Not used but kept for future purposes
-            temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
-
-        Returns:
-            sample: generated mel-spectrogram
-                shape: (batch_size, n_feats, mel_timesteps)
-        """
-        b, c = mu.shape
+        b, _ = mu.shape
         t = patch_size
         z = torch.randn((b, self.in_channels, t), device=mu.device, dtype=mu.dtype) * temperature
 
         t_span = torch.linspace(1, 0, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
-        # Sway sampling strategy
         t_span = t_span + sway_sampling_coef * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span)
 
-        return self.solve_euler(z, t_span=t_span, mu=mu, cond=cond, cfg_value=cfg_value, use_cfg_zero_star=use_cfg_zero_star)
+        return self.solve_euler(
+            x=z,
+            t_span=t_span,
+            mu=mu,
+            cond=cond,
+            cfg_value=cfg_value,
+            use_cfg_zero_star=use_cfg_zero_star,
+        )
 
-    def optimized_scale(self, positive_flat, negative_flat):
+    def optimized_scale(self, positive_flat: torch.Tensor, negative_flat: torch.Tensor):
         dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
-        squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
-        
+        squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
         st_star = dot_product / squared_norm
         return st_star
 
@@ -80,24 +90,13 @@ class UnifiedCFM(torch.nn.Module):
         cfg_value: float = 1.0,
         use_cfg_zero_star: bool = True,
     ):
-        """
-        Fixed euler solver for ODEs.
-        Args:
-            x (torch.Tensor): random noise
-            t_span (torch.Tensor): n_timesteps interpolated
-                shape: (n_timesteps + 1,)
-            mu (torch.Tensor): output of encoder
-                shape: (batch_size, n_feats)
-            cond: condition -- prefix prompt
-            cfg_value (float, optional): cfg value for guidance. Defaults to 1.0.
-        """
         t, _, dt = t_span[0], t_span[-1], t_span[0] - t_span[1]
 
         sol = []
         zero_init_steps = max(1, int(len(t_span) * 0.04))
         for step in range(1, len(t_span)):
             if use_cfg_zero_star and step <= zero_init_steps:
-                dphi_dt = 0.
+                dphi_dt = torch.zeros_like(x)
             else:
                 # Classifier-Free Guidance inference introduced in VoiceBox
                 b = x.size(0)
@@ -105,7 +104,7 @@ class UnifiedCFM(torch.nn.Module):
                 mu_in = torch.zeros([2 * b, mu.size(1)], device=x.device, dtype=x.dtype)
                 t_in = torch.zeros([2 * b], device=x.device, dtype=x.dtype)
                 dt_in = torch.zeros([2 * b], device=x.device, dtype=x.dtype)
-                cond_in = torch.zeros([2 * b, self.in_channels, x.size(2)], device=x.device, dtype=x.dtype)
+                cond_in = torch.zeros([2 * b, self.in_channels, cond.size(2)], device=x.device, dtype=x.dtype)
                 x_in[:b], x_in[b:] = x, x
                 mu_in[:b] = mu
                 t_in[:b], t_in[b:] = t.unsqueeze(0), t.unsqueeze(0)
@@ -117,7 +116,7 @@ class UnifiedCFM(torch.nn.Module):
 
                 dphi_dt = self.estimator(x_in, mu_in, t_in, cond_in, dt_in)
                 dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0)
-                
+
                 if use_cfg_zero_star:
                     positive_flat = dphi_dt.view(b, -1)
                     negative_flat = cfg_dphi_dt.view(b, -1)
@@ -125,7 +124,7 @@ class UnifiedCFM(torch.nn.Module):
                     st_star = st_star.view(b, *([1] * (len(dphi_dt.shape) - 1)))
                 else:
                     st_star = 1.0
-                
+
                 dphi_dt = cfg_dphi_dt * st_star + cfg_value * (dphi_dt - cfg_dphi_dt * st_star)
 
             x = x - dt * dphi_dt
@@ -135,3 +134,99 @@ class UnifiedCFM(torch.nn.Module):
                 dt = t - t_span[step + 1]
 
         return sol[-1]
+
+    # ------------------------------------------------------------------ #
+    # Training loss
+    # ------------------------------------------------------------------ #
+    def adaptive_loss_weighting(
+        self, losses: torch.Tensor, mask: torch.Tensor | None = None, p: float = 0.0, epsilon: float = 1e-3
+    ):
+        weights = 1.0 / ((losses + epsilon).pow(p))
+        if mask is not None:
+            weights = weights * mask
+        return weights.detach()
+
+    def sample_r_t(self, x: torch.Tensor, mu: float = -0.4, sigma: float = 1.0, ratio_r_neq_t: float = 0.0):
+        batch_size = x.shape[0]
+        if self.t_scheduler == "log-norm":
+            s_r = torch.randn(batch_size, device=x.device, dtype=x.dtype) * sigma + mu
+            s_t = torch.randn(batch_size, device=x.device, dtype=x.dtype) * sigma + mu
+            r = torch.sigmoid(s_r)
+            t = torch.sigmoid(s_t)
+        elif self.t_scheduler == "uniform":
+            r = torch.rand(batch_size, device=x.device, dtype=x.dtype)
+            t = torch.rand(batch_size, device=x.device, dtype=x.dtype)
+        else:
+            raise ValueError(f"Unsupported t_scheduler: {self.t_scheduler}")
+
+        mask = torch.rand(batch_size, device=x.device, dtype=x.dtype) < ratio_r_neq_t
+        r, t = torch.where(
+            mask,
+            torch.stack([torch.min(r, t), torch.max(r, t)], dim=0),
+            torch.stack([t, t], dim=0),
+        )
+
+        return r.squeeze(), t.squeeze()
+
+    def compute_loss(
+        self,
+        x1: torch.Tensor,
+        mu: torch.Tensor,
+        cond: torch.Tensor | None = None,
+        tgt_mask: torch.Tensor | None = None,
+        progress: float = 0.0,
+    ):
+        b, _, _ = x1.shape
+
+        if self.training_cfg_rate > 0:
+            cfg_mask = torch.rand(b, device=x1.device) > self.training_cfg_rate
+            mu = mu * cfg_mask.view(-1, 1)
+
+        if cond is None:
+            cond = torch.zeros_like(x1)
+
+        noisy_mask = torch.rand(b, device=x1.device) > (
+            1.0
+            - (
+                self.noise_cond_prob_range[0]
+                + progress * (self.noise_cond_prob_range[1] - self.noise_cond_prob_range[0])
+            )
+        )
+        cond = cond + noisy_mask.view(-1, 1, 1) * torch.randn_like(cond) * self.noise_cond_scale
+
+        ratio_r_neq_t = (
+            self.ratio_r_neq_t_range[0] + progress * (self.ratio_r_neq_t_range[1] - self.ratio_r_neq_t_range[0])
+            if self.mean_mode
+            else 0.0
+        )
+
+        r, t = self.sample_r_t(x1, ratio_r_neq_t=ratio_r_neq_t)
+        r_ = r.detach().clone()
+        t_ = t.detach().clone()
+        z = torch.randn_like(x1)
+        y = (1 - t_.view(-1, 1, 1)) * x1 + t_.view(-1, 1, 1) * z
+        v = z - x1
+
+        def model_fn(z_sample, r_sample, t_sample):
+            return self.estimator(z_sample, mu, t_sample, cond, dt=t_sample - r_sample)
+
+        if self.mean_mode:
+            v_r = torch.zeros_like(r)
+            v_t = torch.ones_like(t)
+            from torch.backends.cuda import sdp_kernel
+
+            with sdp_kernel(enable_flash=False, enable_mem_efficient=False):
+                u_pred, dudt = jvp(model_fn, (y, r, t), (v, v_r, v_t))
+            u_tgt = v - (t_ - r_).view(-1, 1, 1) * dudt
+        else:
+            u_pred = model_fn(y, r, t)
+            u_tgt = v
+
+        losses = F.mse_loss(u_pred, u_tgt.detach(), reduction="none").mean(dim=1)
+        if tgt_mask is not None:
+            weights = self.adaptive_loss_weighting(losses, tgt_mask.squeeze(1))
+            loss = (weights * losses).sum() / torch.sum(tgt_mask)
+        else:
+            loss = losses.mean()
+
+        return loss

src/voxcpm/modules/minicpm4/config.py

@@ -26,4 +26,5 @@ class MiniCPM4Config(BaseModel):
     dim_model_base: int
     scale_depth: float
     rope_theta: float
-    kv_channels: int = None
+    kv_channels: int = None
+    no_rope: bool = False

src/voxcpm/modules/minicpm4/model.py

@@ -64,10 +64,8 @@ class MiniCPMLongRoPE(nn.Module):
         self.long_factor = config.rope_scaling.long_factor
         self.original_max_position_embeddings = config.rope_scaling.original_max_position_embeddings
 
-        scale = (self.max_position_embeddings / self.original_max_position_embeddings)
-        self.scaling_factor = math.sqrt(
-            1 + math.log(scale) / math.log(self.original_max_position_embeddings)
-        )
+        scale = self.max_position_embeddings / self.original_max_position_embeddings
+        self.scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))
         inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
 
@@ -76,11 +74,7 @@ class MiniCPMLongRoPE(nn.Module):
         self.register_buffer("cos_cached", torch.empty(0), persistent=False)
         self.register_buffer("sin_cached", torch.empty(0), persistent=False)
 
-        self._set_cos_sin_cache(
-            seq_len=self.max_position_embeddings,
-            device=self.inv_freq.device,
-            dtype=torch.float32
-        )
+        self._set_cos_sin_cache(seq_len=self.max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32)
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
         """设置cos和sin缓存"""
@@ -93,8 +87,7 @@ class MiniCPMLongRoPE(nn.Module):
             ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=device)
 
         freqs = torch.mul(
-            torch.outer(t, 1.0 / ext_factors).to(device=device),
-            self.inv_freq.to(device=device).to(dtype)
+            torch.outer(t, 1.0 / ext_factors).to(device=device), self.inv_freq.to(device=device).to(dtype)
         )
 
         # 创建embeddings
@@ -123,7 +116,9 @@ class MiniCPMAttention(nn.Module):
         self.layer_idx = layer_idx
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
-        self.head_dim = config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
+        self.head_dim = (
+            config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
+        )
         self.num_key_value_heads = config.num_key_value_heads
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
         self.max_position_embeddings = config.max_position_embeddings
@@ -150,10 +145,15 @@ class MiniCPMAttention(nn.Module):
         key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
 
-        cos, sin = position_emb
-
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        if position_emb is not None:
+            cos, sin = position_emb
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
+        # ref: https://github.com/pytorch/pytorch/issues/163597
+        # there is a bug in MPS for non-contiguous tensors, so we need to make them contiguous
+        query_states = query_states.contiguous()
+        key_states = key_states.contiguous()
+        value_states = value_states.contiguous()
         attn_output = torch.nn.functional.scaled_dot_product_attention(
             query_states,
             key_states,
@@ -187,9 +187,9 @@ class MiniCPMAttention(nn.Module):
         key_states = key_states.view(bsz, 1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
         value_states = value_states.view(bsz, 1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
 
-        cos, sin = position_emb
-
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        if position_emb is not None:
+            cos, sin = position_emb
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         key_cache, value_cache = kv_cache
 
@@ -198,6 +198,11 @@ class MiniCPMAttention(nn.Module):
 
         attn_mask = torch.arange(key_cache.size(2), device=key_cache.device) <= position_id
 
+        # ref: https://github.com/pytorch/pytorch/issues/163597
+        # there is a bug in MPS for non-contiguous tensors, so we need to make them contiguous
+        query_states = query_states.contiguous()
+        key_cache = key_cache.contiguous()
+        value_cache = value_cache.contiguous()
         attn_output = torch.nn.functional.scaled_dot_product_attention(
             query_states,
             key_cache,
@@ -338,7 +343,10 @@ class MiniCPMModel(nn.Module):
         )
 
         self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rope_emb = MiniCPMLongRoPE(config)
+        if config.no_rope:
+            self.rope_emb = None
+        else:
+            self.rope_emb = MiniCPMLongRoPE(config)
 
         self.kv_cache = None
 
@@ -355,8 +363,11 @@ class MiniCPMModel(nn.Module):
             hidden_states: Tensor(batch_size, seq_length, hidden_size)
             next_decoder_cache: List[(batch_size, num_heads, seq_length, head_dim), (batch_size, num_heads, seq_length, head_dim)]
         """
-        position_ids = torch.arange(0, inputs_embeds.size(1), dtype=torch.long, device=inputs_embeds.device)
-        position_emb = self.rope_emb(position_ids)
+        if self.rope_emb is not None:
+            position_ids = torch.arange(0, inputs_embeds.size(1), dtype=torch.long, device=inputs_embeds.device)
+            position_emb = self.rope_emb(position_ids)
+        else:
+            position_emb = None
         hidden_states = inputs_embeds
 
         next_decoder_cache = []
@@ -385,7 +396,10 @@ class MiniCPMModel(nn.Module):
         """
         assert self.kv_cache is not None, "KV cache is not setup"
 
-        position_emb = self.rope_emb(position_id)
+        if self.rope_emb is not None:
+            position_emb = self.rope_emb(position_id)
+        else:
+            position_emb = None
         hidden_states = inputs_embeds
 
         for i, decoder_layer in enumerate(self.layers):
@@ -403,7 +417,11 @@ class MiniCPMModel(nn.Module):
         self.kv_cache = StaticKVCache(
             num_layers=self.config.num_hidden_layers,
             num_kv_heads=self.config.num_key_value_heads,
-            dim_kv_head=self.config.hidden_size // self.config.num_attention_heads if self.config.kv_channels is None else self.config.kv_channels,
+            dim_kv_head=(
+                self.config.hidden_size // self.config.num_attention_heads
+                if self.config.kv_channels is None
+                else self.config.kv_channels
+            ),
             batch_size=batch_size,
             device=device,
             dtype=dtype,

src/voxcpm/training/__init__.py

@@ -0,0 +1,27 @@
+"""
+Training utilities for VoxCPM fine-tuning.
+
+This package mirrors the training mechanics used in the minicpm-audio
+tooling while relying solely on local audio-text datasets managed via
+the HuggingFace ``datasets`` library.
+"""
+
+from .accelerator import Accelerator
+from .tracker import TrainingTracker
+from .data import (
+    load_audio_text_datasets,
+    HFVoxCPMDataset,
+    build_dataloader,
+    BatchProcessor,
+)
+from .state import TrainingState
+
+__all__ = [
+    "Accelerator",
+    "TrainingTracker",
+    "HFVoxCPMDataset",
+    "BatchProcessor",
+    "TrainingState",
+    "load_audio_text_datasets",
+    "build_dataloader",
+]

src/voxcpm/training/accelerator.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import contextlib
+import os
+import random
+import typing
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.utils.data
+from torch.nn.parallel import DistributedDataParallel
+
+
+class Accelerator:
+    """
+    Simplified accelerator that mirrors the behaviour of the minicpm-audio
+    training utilities. It initializes a distributed process group when
+    ``torchrun`` is used and exposes helpers for AMP, gradient scaling and
+    preparing models/dataloaders for DDP.
+    """
+
+    def __init__(self, amp: bool = False, seed: int = 42):
+        self.world_size = int(os.getenv("WORLD_SIZE", "1"))
+
+        if self.world_size > 1 and not dist.is_initialized():
+            dist.init_process_group("nccl", init_method="env://")
+
+        self.rank = dist.get_rank() if dist.is_initialized() else 0
+        self.local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+        self.amp = amp
+
+        # Set random seed to ensure model initialization consistency
+        self._set_seed(seed)
+
+        class DummyScaler:
+            def step(self, optimizer):
+                optimizer.step()
+
+            def scale(self, loss):
+                return loss
+
+            def unscale_(self, optimizer):
+                return optimizer
+
+            def update(self):
+                pass
+
+        self.scaler = torch.amp.GradScaler("cuda") if (amp and torch.cuda.is_available()) else DummyScaler()
+        self.device_ctx = torch.cuda.device(self.local_rank) if torch.cuda.is_available() else None
+        self._ddp_model = None  # For no_sync support
+
+    def _set_seed(self, seed: int):
+        """Set random seed to ensure model initialization consistency across multiple GPUs"""
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+
+    def __enter__(self):
+        if self.device_ctx is not None:
+            self.device_ctx.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.device_ctx is not None:
+            self.device_ctx.__exit__(exc_type, exc_value, traceback)
+
+    def barrier(self):
+        """Synchronize all processes"""
+        if dist.is_initialized():
+            dist.barrier()
+
+    def all_reduce(self, tensor: torch.Tensor, op=dist.ReduceOp.AVG):
+        """All-reduce tensor across processes"""
+        if dist.is_initialized():
+            dist.all_reduce(tensor, op=op)
+        return tensor
+
+    # ------------------------------------------------------------------ #
+    # Model helpers
+    # ------------------------------------------------------------------ #
+    def prepare_model(self, model: torch.nn.Module, **kwargs):
+        if hasattr(model, "device"):  # make sure the matrix will be moved to the correct device
+            model.device = self.device
+        model = model.to(self.device)
+        if self.world_size > 1:
+            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+            model = DistributedDataParallel(model, device_ids=[self.local_rank], **kwargs)
+            self._ddp_model = model  # Save DDP model reference for no_sync support
+        return model
+
+    @contextlib.contextmanager
+    def no_sync(self):
+        """
+        Context manager to skip gradient synchronization during gradient accumulation.
+        Only used outside the last micro-batch.
+        """
+        if self._ddp_model is not None:
+            with self._ddp_model.no_sync():
+                yield
+        else:
+            yield
+
+    @property
+    def device(self):
+        if torch.cuda.is_available():
+            return torch.device("cuda", self.local_rank)
+        if torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+
+    # ------------------------------------------------------------------ #
+    # AMP helpers
+    # ------------------------------------------------------------------ #
+    def autocast(self, *args, **kwargs):
+        return torch.amp.autocast("cuda", enabled=self.amp, *args, **kwargs)
+
+    def backward(self, loss: torch.Tensor):
+        self.scaler.scale(loss).backward()
+
+    def step(self, optimizer: torch.optim.Optimizer):
+        self.scaler.step(optimizer)
+
+    def update(self):
+        self.scaler.update()
+
+    # ------------------------------------------------------------------ #
+    # Data helpers
+    # ------------------------------------------------------------------ #
+    def prepare_dataloader(
+        self,
+        dataset: typing.Iterable,
+        *,
+        batch_size: int,
+        num_workers: int = 0,
+        shuffle: bool = True,
+        collate_fn=None,
+        drop_last: bool = False,
+    ) -> torch.utils.data.DataLoader:
+        if self.world_size > 1:
+            sampler = torch.utils.data.distributed.DistributedSampler(
+                dataset, num_replicas=self.world_size, rank=self.rank, shuffle=shuffle
+            )
+            shuffle = False
+        else:
+            sampler = None
+
+        return torch.utils.data.DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffle if sampler is None else False,
+            sampler=sampler,
+            num_workers=num_workers,
+            collate_fn=collate_fn,
+            drop_last=drop_last,
+            pin_memory=True,
+        )
+
+    @staticmethod
+    def unwrap(model: torch.nn.Module) -> torch.nn.Module:
+        return model.module if hasattr(model, "module") else model

src/voxcpm/training/config.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import argbind
+import yaml
+from pathlib import Path
+from typing import Dict, Any
+
+
+def load_yaml_config(path: str | Path) -> Dict[str, Any]:
+    """
+    Load a YAML configuration file into a dictionary suitable for argbind.
+    """
+    path = Path(path)
+    with path.open("r", encoding="utf-8") as f:
+        data = yaml.safe_load(f)
+    if not isinstance(data, dict):
+        raise ValueError(f"Configuration file {path} must contain a top-level mapping.")
+    return data
+
+
+def parse_args_with_config(config_path: str | Path | None = None):
+    """
+    Helper to unify CLI arguments and YAML configuration.
+
+    Usage mirrors minicpm-audio:
+        args = parse_args_with_config("conf/voxcpm/finetune.yml")
+        with argbind.scope(args):
+            ...
+    """
+    cli_args = argbind.parse_args()
+    if config_path is None:
+        return cli_args
+
+    yaml_args = load_yaml_config(config_path)
+    with argbind.scope(cli_args):
+        yaml_args = argbind.parse_args(yaml_args=yaml_args, argv=[])
+    cli_args.update(yaml_args)
+    return cli_args

src/voxcpm/training/data.py

@@ -0,0 +1,214 @@
+import math
+from typing import Dict, List, Optional, Tuple
+
+import argbind
+import torch
+from datasets import Audio, Dataset, DatasetDict, load_dataset
+from torch.utils.data import Dataset as TorchDataset
+
+from ..model.voxcpm import VoxCPMConfig
+from ..modules.audiovae import AudioVAE
+from .packers import AudioFeatureProcessingPacker
+
+DEFAULT_TEXT_COLUMN = "text"
+DEFAULT_AUDIO_COLUMN = "audio"
+DEFAULT_ID_COLUMN = "dataset_id"
+
+
+@argbind.bind()
+def load_audio_text_datasets(
+    train_manifest: str,
+    val_manifest: str = "",
+    text_column: str = DEFAULT_TEXT_COLUMN,
+    audio_column: str = DEFAULT_AUDIO_COLUMN,
+    dataset_id_column: str = DEFAULT_ID_COLUMN,
+    sample_rate: int = 16_000,
+    num_proc: int = 1,
+) -> Tuple[Dataset, Optional[Dataset]]:
+    data_files = {"train": train_manifest}
+    if val_manifest:
+        data_files["validation"] = val_manifest
+
+    dataset_dict: DatasetDict = load_dataset("json", data_files=data_files)
+
+    def prepare(ds: Dataset) -> Dataset:
+        if audio_column not in ds.column_names:
+            raise ValueError(f"Expected '{audio_column}' column in manifest.")
+        # We cast to Audio to ensure proper handling during training,
+        # but for length calculation we might need raw path or duration if available.
+        # HF datasets usually don't compute duration automatically for 'Audio' column.
+        ds = ds.cast_column(audio_column, Audio(sampling_rate=sample_rate))
+        if audio_column != DEFAULT_AUDIO_COLUMN:
+            ds = ds.rename_column(audio_column, DEFAULT_AUDIO_COLUMN)
+        if text_column != DEFAULT_TEXT_COLUMN:
+            ds = ds.rename_column(text_column, DEFAULT_TEXT_COLUMN)
+        if dataset_id_column and dataset_id_column in ds.column_names:
+            if dataset_id_column != DEFAULT_ID_COLUMN:
+                ds = ds.rename_column(dataset_id_column, DEFAULT_ID_COLUMN)
+        else:
+            ds = ds.add_column(DEFAULT_ID_COLUMN, [0] * len(ds))
+        return ds
+
+    train_ds = prepare(dataset_dict["train"])
+    val_ds = prepare(dataset_dict["validation"]) if "validation" in dataset_dict else None
+    return train_ds, val_ds
+
+
+def compute_sample_lengths(
+    ds: Dataset,
+    audio_vae_fps: int = 25,
+    patch_size: int = 1,
+) -> List[int]:
+    """
+    预估每个样本经过 packer 之后的大致序列长度（text+audio），用于过滤超长样本。
+
+    逻辑与 AudioFeatureProcessingPacker / AudioVAE 一致：
+    - 文本长度: len(text_ids)
+    - 音频长度:
+        duration(s) * audio_vae_fps -> 近似 VAE 帧数 t_vae
+        t_seq = ceil(t_vae / patch_size)
+    - 序列总长约为: text_len + t_seq + 2
+
+    Optimized: Use batch column access instead of iterating item by item.
+    """
+    # Batch access columns - much faster than per-item access
+    text_ids_list = ds["text_ids"]
+    text_lens = [len(t) for t in text_ids_list]
+
+    has_duration = "duration" in ds.column_names
+    if has_duration:
+        durations = ds["duration"]
+    else:
+        # Fallback: need to compute from audio (slow, but unavoidable without duration column)
+        durations = []
+        for i in range(len(ds)):
+            audio = ds[i][DEFAULT_AUDIO_COLUMN]
+            durations.append(len(audio["array"]) / float(audio["sampling_rate"]))
+
+    # Vectorized length computation
+    lengths = []
+    for text_len, duration in zip(text_lens, durations):
+        t_vae = math.ceil(float(duration) * audio_vae_fps)
+        t_seq = math.ceil(t_vae / patch_size)
+        total_len = text_len + t_seq + 2
+        lengths.append(total_len)
+
+    return lengths
+
+
+class HFVoxCPMDataset(TorchDataset):
+    """
+    Thin wrapper around a tokenized HuggingFace dataset that returns
+    PyTorch-friendly samples.
+    """
+
+    def __init__(self, dataset: Dataset):
+        self.dataset = dataset
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int):
+        item = self.dataset[idx]
+        audio = item[DEFAULT_AUDIO_COLUMN]
+        return {
+            "text_ids": item["text_ids"],
+            "audio_array": audio["array"],
+            "audio_sampling_rate": audio["sampling_rate"],
+            "dataset_id": item.get(DEFAULT_ID_COLUMN, 0),
+            "is_prompt": item.get("is_prompt", False),
+        }
+
+    @staticmethod
+    def pad_sequences(seqs: List[torch.Tensor], pad_value: float):
+        if not seqs:
+            return torch.empty(0)
+        max_len = max(seq.shape[0] for seq in seqs)
+        padded = []
+        for seq in seqs:
+            if seq.shape[0] < max_len:
+                pad_width = (0, max_len - seq.shape[0])
+                seq = torch.nn.functional.pad(seq, pad_width, value=pad_value)
+            padded.append(seq)
+        return torch.stack(padded)
+
+    @classmethod
+    def collate_fn(cls, batch: List[Dict]):
+        text_tensors = [torch.tensor(sample["text_ids"], dtype=torch.int32) for sample in batch]
+        audio_tensors = [torch.tensor(sample["audio_array"], dtype=torch.float32) for sample in batch]
+        dataset_ids = torch.tensor([sample["dataset_id"] for sample in batch], dtype=torch.int32)
+        is_prompts = [bool(sample.get("is_prompt", False)) for sample in batch]
+
+        text_padded = cls.pad_sequences(text_tensors, pad_value=-100)
+        audio_padded = cls.pad_sequences(audio_tensors, pad_value=-100.0)
+        task_ids = torch.ones(text_padded.size(0), dtype=torch.int32)
+
+        return {
+            "text_tokens": text_padded,
+            "audio_tokens": audio_padded,
+            "task_ids": task_ids,
+            "dataset_ids": dataset_ids,
+            "is_prompts": is_prompts,
+        }
+
+
+class BatchProcessor:
+    """
+    Wraps ``AudioFeatureProcessingPacker`` so the training loop can mirror
+    the minicpm-audio mechanics.
+    """
+
+    def __init__(
+        self,
+        *,
+        config: VoxCPMConfig,
+        audio_vae: AudioVAE,
+        dataset_cnt: int,
+        device: torch.device,
+    ):
+        self.device = device
+        self.dataset_cnt = dataset_cnt
+        self.audio_vae = audio_vae
+        self.audio_vae.to(device)
+        self.packer = AudioFeatureProcessingPacker(
+            dataset_cnt=dataset_cnt,
+            max_len=config.max_length,
+            patch_size=config.patch_size,
+            feat_dim=config.feat_dim,
+            audio_vae=self.audio_vae,
+        )
+
+    def __call__(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        audio_tokens = batch["audio_tokens"].to(self.device)
+        text_tokens = batch["text_tokens"].to(self.device)
+        task_ids = batch["task_ids"].to(self.device)
+        dataset_ids = batch["dataset_ids"].to(self.device)
+
+        packed = self.packer(
+            audio_tokens=audio_tokens,
+            text_tokens=text_tokens,
+            task_ids=task_ids,
+            dataset_ids=dataset_ids,
+            is_prompts=batch["is_prompts"],
+        )
+        return packed
+
+
+def build_dataloader(
+    hf_dataset: Dataset,
+    *,
+    accelerator,
+    batch_size: int,
+    num_workers: int,
+    drop_last: bool = False,
+) -> torch.utils.data.DataLoader:
+    torch_dataset = HFVoxCPMDataset(hf_dataset)
+    # Standard padding-based batching; Accelerator will attach DistributedSampler if needed.
+    return accelerator.prepare_dataloader(
+        torch_dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        shuffle=True,
+        collate_fn=HFVoxCPMDataset.collate_fn,
+        drop_last=drop_last,
+    )

src/voxcpm/training/packers.py

@@ -0,0 +1,296 @@
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+class AudioFeatureProcessingPacker:
+    """
+    Adapted from the minicpm-audio training utilities. It converts raw text and
+    audio tokens into the packed multimodal representation required by VoxCPM.
+    """
+
+    def __init__(self, dataset_cnt: int, max_len: int, patch_size: int, feat_dim: int, audio_vae: nn.Module):
+        self.audio_start_id = 101
+        self.audio_end_id = 102
+        # unused now
+        self.audio_prompt_start_id = 103
+        self.audio_prompt_end_id = 104
+        self.text_eos_token_id = 2
+
+        self.patch_size = patch_size
+        self.patch_len = audio_vae.hop_length * self.patch_size
+        self.feat_dim = feat_dim
+        self.dataset_cnt = max(dataset_cnt, 1)
+        self.max_len = max_len
+
+        self.audio_vae = audio_vae
+
+        self.process_functions = {"tts": self.process_tts_data}
+        self.task_id_map = {"tts": 1}
+        self.id_to_task = {idx: usage for usage, idx in self.task_id_map.items()}
+
+    # ------------------------------------------------------------------ #
+    # Helpers
+    # ------------------------------------------------------------------ #
+    @staticmethod
+    def _first_pad_position(tokens: torch.Tensor):
+        positions = (tokens == -100).nonzero(as_tuple=True)
+        if positions[0].numel() == 0:
+            return None
+        return int(positions[0][0])
+
+    def unpad_text_tokens(self, tokens: torch.Tensor):
+        pad_pos = self._first_pad_position(tokens)
+        return tokens if pad_pos is None else tokens[:pad_pos]
+
+    def unpad_audio_tokens(self, tokens: torch.Tensor):
+        pad_pos = self._first_pad_position(tokens)
+        return tokens if pad_pos is None else tokens[:pad_pos]
+
+    def encode_audio(self, wav: torch.Tensor):
+        """
+        Encode raw waveform into latent features using AudioVAE.
+
+        AudioVAE.encode expects shape [B, 1, T'] and returns [B, D, T].
+        We then transpose to [B, T, D] to match downstream expectations.
+        """
+        wav = wav.unsqueeze(0)  # [1, T]
+        wav = wav.unsqueeze(1)  # [1, 1, T]
+        wav_len = wav.size(-1)
+        if wav_len % self.patch_len != 0:
+            padding_size = self.patch_len - wav_len % self.patch_len
+            wav = torch.nn.functional.pad(wav, (0, padding_size))
+
+        with torch.no_grad():
+            z = self.audio_vae.encode(wav, self.audio_vae.sample_rate)  # [1, D, T']
+            feat = z.transpose(1, 2)  # [1, T', D]
+        return feat
+
+    # ------------------------------------------------------------------ #
+    # Main entry point
+    # ------------------------------------------------------------------ #
+    def __call__(
+        self,
+        audio_tokens: torch.Tensor,
+        text_tokens: torch.Tensor,
+        task_ids: torch.Tensor,
+        dataset_ids: torch.Tensor,
+        is_prompts: List[bool],
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Padding-based batching: each sample in the input batch is processed
+        independently and then padded to a common length (capped by ``max_len``).
+        The result tensors all have shape [B, T, ...].
+        """
+        device = audio_tokens.device
+        max_dataset_id = int(dataset_ids.max().item()) if dataset_ids.numel() > 0 else -1
+        dataset_cnt = max(self.dataset_cnt, max_dataset_id + 1)
+
+        text_tokens_list: List[torch.Tensor] = []
+        audio_feats_list: List[torch.Tensor] = []
+        text_mask_list: List[torch.Tensor] = []
+        audio_mask_list: List[torch.Tensor] = []
+        loss_mask_list: List[torch.Tensor] = []
+        labels_list: List[torch.Tensor] = []
+        audio_task_ids_list: List[torch.Tensor] = []
+        audio_dataset_ids_list: List[torch.Tensor] = []
+        lengths: List[int] = []
+
+        audio_duration_consumed = torch.zeros(dataset_cnt, dtype=torch.float32, device=device)
+        text_token_consumed = torch.zeros(dataset_cnt, dtype=torch.float32, device=device)
+
+        for audio_token, text_token, task_id, dataset_idx, is_prompt in zip(
+            audio_tokens, text_tokens, task_ids.tolist(), dataset_ids.tolist(), is_prompts
+        ):
+            unpad_audio_token = self.unpad_audio_tokens(audio_token).to(torch.float32)
+            unpad_text_token = self.unpad_text_tokens(text_token)
+            usage = self.id_to_task[task_id]
+
+            (
+                packed_text,
+                audio_feat,
+                text_mask,
+                audio_mask,
+                loss_mask,
+                labels,
+                audio_duration,
+                text_token_count,
+            ) = self.process_functions[usage](unpad_audio_token, unpad_text_token, is_prompt)
+
+            audio_duration_consumed[dataset_idx] += audio_duration
+            text_token_consumed[dataset_idx] += text_token_count
+
+            audio_task_id = torch.zeros_like(audio_mask)
+            audio_task_id[audio_mask == 1] = self.task_id_map[usage]
+
+            audio_dataset_id = torch.zeros_like(audio_mask)
+            audio_dataset_id[audio_mask == 1] = dataset_idx + 1
+
+            text_tokens_list.append(packed_text)
+            text_mask_list.append(text_mask)
+            audio_feats_list.append(audio_feat)
+            audio_mask_list.append(audio_mask)
+            loss_mask_list.append(loss_mask)
+            labels_list.append(labels)
+            audio_task_ids_list.append(audio_task_id)
+            audio_dataset_ids_list.append(audio_dataset_id)
+            lengths.append(packed_text.shape[0])
+
+        # Determine padded length per batch (cap by self.max_len)
+        if lengths:
+            max_len = min(self.max_len, max(lengths))
+        else:
+            max_len = self.max_len
+
+        def pad_1d(x: torch.Tensor, pad_value: int = 0) -> torch.Tensor:
+            if x.size(0) >= max_len:
+                return x[:max_len]
+            pad = torch.full((max_len - x.size(0),), pad_value, dtype=x.dtype, device=x.device)
+            return torch.cat([x, pad], dim=0)
+
+        def pad_3d(x: torch.Tensor) -> torch.Tensor:
+            # x: [T, P, D]
+            if x.size(0) >= max_len:
+                return x[:max_len]
+            pad = torch.zeros((max_len - x.size(0),) + x.shape[1:], dtype=x.dtype, device=x.device)
+            return torch.cat([x, pad], dim=0)
+
+        if lengths:
+            text_tokens_batch = torch.stack([pad_1d(t, pad_value=0) for t in text_tokens_list], dim=0)
+            text_mask_batch = torch.stack([pad_1d(m, pad_value=0) for m in text_mask_list], dim=0)
+            audio_feats_batch = torch.stack([pad_3d(f) for f in audio_feats_list], dim=0)
+            audio_mask_batch = torch.stack([pad_1d(m, pad_value=0) for m in audio_mask_list], dim=0)
+            loss_mask_batch = torch.stack([pad_1d(m, pad_value=0) for m in loss_mask_list], dim=0)
+            labels_batch = torch.stack([pad_1d(lbl, pad_value=0) for lbl in labels_list], dim=0)
+            audio_task_ids_batch = torch.stack([pad_1d(t, pad_value=0) for t in audio_task_ids_list], dim=0)
+            audio_dataset_ids_batch = torch.stack([pad_1d(d, pad_value=0) for d in audio_dataset_ids_list], dim=0)
+
+            # Position ids: [B, T], simple 0..L_i-1 then padded with 0
+            position_ids_list = []
+            for L in lengths:
+                L_clip = min(L, max_len)
+                pos = torch.arange(0, L_clip, device=device)
+                if L_clip < max_len:
+                    pad = torch.zeros(max_len - L_clip, dtype=pos.dtype, device=device)
+                    pos = torch.cat([pos, pad], dim=0)
+                position_ids_list.append(pos)
+            position_ids = torch.stack(position_ids_list, dim=0)
+        else:
+            # Empty batch fallback (shouldn't really happen)
+            text_tokens_batch = torch.zeros((0, self.max_len), dtype=torch.int32, device=device)
+            text_mask_batch = torch.zeros_like(text_tokens_batch)
+            audio_feats_batch = torch.zeros(
+                (0, self.max_len, self.patch_size, self.feat_dim), dtype=torch.float32, device=device
+            )
+            audio_mask_batch = torch.zeros_like(text_tokens_batch)
+            loss_mask_batch = torch.zeros_like(text_tokens_batch)
+            labels_batch = torch.zeros_like(text_tokens_batch)
+            audio_task_ids_batch = torch.zeros_like(text_tokens_batch)
+            audio_dataset_ids_batch = torch.zeros_like(text_tokens_batch)
+            position_ids = torch.zeros_like(text_tokens_batch)
+
+        audio_duration_consumed = audio_duration_consumed.to(torch.long)
+        text_token_consumed = text_token_consumed.to(torch.long)
+
+        return {
+            "text_tokens": text_tokens_batch,
+            "audio_feats": audio_feats_batch,
+            "text_mask": text_mask_batch,
+            "audio_mask": audio_mask_batch,
+            "loss_mask": loss_mask_batch,
+            "position_ids": position_ids,
+            "labels": labels_batch,
+            "audio_task_ids": audio_task_ids_batch,
+            "audio_dataset_ids": audio_dataset_ids_batch,
+            "audio_duration_consumed": audio_duration_consumed,
+            "text_token_consumed": text_token_consumed,
+        }
+
+    # ------------------------------------------------------------------ #
+    # Feature extraction helpers
+    # ------------------------------------------------------------------ #
+    def extract_audio_feats(self, audio_data: torch.Tensor):
+        audio_feats = self.encode_audio(audio_data)
+        if audio_feats.size(1) % self.patch_size != 0:
+            audio_feats_ = audio_feats.transpose(1, 2)
+            padding = nn.functional.pad(audio_feats_, (0, self.patch_size - audio_feats.size(1) % self.patch_size))
+            audio_feats = padding.transpose(1, 2)
+
+        audio_duration = audio_feats.size(1) / 25
+        audio_feats = rearrange(audio_feats, "b (t p) c -> b t p c", p=self.patch_size)
+        return audio_feats, audio_duration
+
+    def process_tts_data(self, audio_token: torch.Tensor, text_token: torch.Tensor, is_prompt: bool = False):
+        text_token_info = torch.cat(
+            [
+                text_token,
+                torch.tensor(
+                    [self.audio_prompt_start_id if is_prompt else self.audio_start_id],
+                    dtype=torch.int32,
+                    device=text_token.device,
+                ),
+            ],
+            dim=-1,
+        )
+        text_token_count = len(text_token)
+        text_length = text_token_info.shape[0]
+        audio_feat_info, audio_duration = self.extract_audio_feats(audio_token)
+        audio_feat_info = audio_feat_info.squeeze(0)
+        audio_length = audio_feat_info.shape[0]
+
+        text_pad_token = torch.zeros(audio_length, dtype=torch.int32, device=text_token.device)
+        text_token_info = torch.cat(
+            [
+                text_token_info,
+                text_pad_token,
+                torch.tensor(
+                    [self.audio_prompt_end_id if is_prompt else self.audio_end_id],
+                    dtype=torch.int32,
+                    device=text_token.device,
+                ),
+            ]
+        )
+        audio_pad_feat = torch.zeros(
+            (text_length, self.patch_size, audio_feat_info.size(-1)),
+            dtype=torch.float32,
+            device=text_token.device,
+        )
+        audio_feat_info = torch.cat([audio_pad_feat, audio_feat_info, audio_pad_feat[0:1, ...]], dim=0)
+
+        text_mask = (
+            torch.cat([torch.ones(text_length), torch.zeros(audio_length), torch.ones(1)])
+            .type(torch.int32)
+            .to(text_token.device)
+        )
+        audio_mask = (
+            torch.cat([torch.zeros(text_length), torch.ones(audio_length), torch.zeros(1)])
+            .type(torch.int32)
+            .to(text_token.device)
+        )
+        loss_mask = (
+            torch.cat(
+                [
+                    torch.zeros(text_length),
+                    torch.zeros(audio_length) if is_prompt else torch.ones(audio_length),
+                    torch.zeros(1),
+                ]
+            )
+            .type(torch.int32)
+            .to(text_token.device)
+        )
+
+        labels = torch.zeros(text_length + audio_length + 1).type(torch.int32).to(text_token.device)
+        labels[-2] = 1
+
+        return (
+            text_token_info,
+            audio_feat_info,
+            text_mask,
+            audio_mask,
+            loss_mask,
+            labels,
+            audio_duration,
+            text_token_count,
+        )

src/voxcpm/training/state.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class TrainingState:
+    """
+    Container that mirrors the object returned in the minicpm-audio training
+    loop. It holds persistent references to the model, optimizer, scheduler,
+    dataloaders and tracker.
+    """
+
+    generator: object
+    optimizer: object
+    scheduler: object
+    train_loader: object
+    val_loader: object
+    tracker: object
+    batch_processor: object

src/voxcpm/training/tracker.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import contextlib
+import sys
+import time
+from pathlib import Path
+from typing import Dict, Optional
+
+
+class TrainingTracker:
+    """
+    Lightweight tracker inspired by the minimcpm-audio training workflow.
+
+    It keeps track of the current global step, prints rank-aware messages,
+    optionally writes to TensorBoard via a provided writer, and stores progress
+    in a logfile for later inspection.
+    """
+
+    def __init__(
+        self,
+        *,
+        writer=None,
+        log_file: Optional[str] = None,
+        rank: int = 0,
+    ):
+        self.writer = writer
+        self.log_file = Path(log_file) if log_file else None
+        if self.log_file:
+            self.log_file.parent.mkdir(parents=True, exist_ok=True)
+        self.rank = rank
+        self.step = 0
+        # Record the time of the last log to calculate the interval
+        self._last_log_time: float | None = None
+
+    # ------------------------------------------------------------------ #
+    # Logging helpers
+    # ------------------------------------------------------------------ #
+    def print(self, message: str):
+        if self.rank == 0:
+            print(message, flush=True, file=sys.stderr)
+            if self.log_file:
+                with self.log_file.open("a", encoding="utf-8") as f:
+                    f.write(message + "\n")
+
+    def log_metrics(self, metrics: Dict[str, float], split: str):
+        if self.rank == 0:
+            now = time.time()
+            dt_str = ""
+            if self._last_log_time is not None:
+                dt = now - self._last_log_time
+                dt_str = f", log interval: {dt:.2f}s"
+            self._last_log_time = now
+
+            formatted = ", ".join(f"{k}: {v:.6f}" for k, v in metrics.items())
+            self.print(f"[{split}] step {self.step}: {formatted}{dt_str}")
+        if self.writer is not None:
+            for key, value in metrics.items():
+                if isinstance(value, (int, float)):
+                    self.writer.add_scalar(f"{split}/{key}", value, self.step)
+
+    def done(self, split: str, message: str):
+        self.print(f"[{split}] {message}")
+
+    # ------------------------------------------------------------------ #
+    # State dict
+    # ------------------------------------------------------------------ #
+    def state_dict(self):
+        return {"step": self.step}
+
+    def load_state_dict(self, state):
+        self.step = int(state.get("step", 0))
+
+    # ------------------------------------------------------------------ #
+    # Context manager compatibility (for parity with minicpm-audio code)
+    # ------------------------------------------------------------------ #
+    @contextlib.contextmanager
+    def live(self):
+        yield

src/voxcpm/utils/text_normalize.py

@@ -2,44 +2,11 @@
 import re
 import regex
 import inflect
-from functools import partial
-from tn.chinese.normalizer import Normalizer as ZhNormalizer
-from tn.english.normalizer import Normalizer as EnNormalizer
+from wetext import Normalizer
 
-def normal_cut_sentence(text):
-    # 先处理括号内的逗号，将其替换为特殊标记
-    text = re.sub(r'([（(][^）)]*)([,，])([^）)]*[）)])', r'\1&&&\3', text)
-    text = re.sub('([。！，？\?])([^’”])',r'\1\n\2',text)#普通断句符号且后面没有引号
-    text = re.sub('(\.{6})([^’”])',r'\1\n\2',text)#英文省略号且后面没有引号
-    text = re.sub('(\…{2})([^’”])',r'\1\n\2',text)#中文省略号且后面没有引号
-    text = re.sub('([. ,。！；？\?\.{6}\…{2}][’”])([^’”])',r'\1\n\2',text)#断句号+引号且后面没有引号
-    # 处理英文句子的分隔
-    text = re.sub(r'([.,!?])([^’”\'"])', r'\1\n\2', text)  # 句号、感叹号、问号后面没有引号
-    text = re.sub(r'([.!?][’”\'"])([^’”\'"])', r'\1\n\2', text)  # 句号、感叹号、问号加引号后面的部分
-    text = re.sub(r'([（(][^）)]*)(&&&)([^）)]*[）)])', r'\1，\3', text)
-    text = [t for t in text.split("\n") if t]
-    return text
+chinese_char_pattern = re.compile(r"[\u4e00-\u9fff]+")
 
 
-def cut_sentence_with_fix_length(text : str, length : int):
-    sentences = normal_cut_sentence(text)
-    cur_length = 0
-    res = ""
-    for sentence in sentences:
-        if not sentence:
-            continue
-        if cur_length > length or cur_length + len(sentence) > length:
-            yield res
-            res = ""
-            cur_length = 0
-        res += sentence
-        cur_length += len(sentence)
-    if res:
-        yield res
-
- 
-chinese_char_pattern = re.compile(r'[\u4e00-\u9fff]+')
-
 # whether contain chinese character
 def contains_chinese(text):
     return bool(chinese_char_pattern.search(text))
@@ -47,19 +14,19 @@ def contains_chinese(text):
 
 # replace special symbol
 def replace_corner_mark(text):
-    text = text.replace('²', '平方')
-    text = text.replace('³', '立方')
-    text = text.replace('√', '根号')
-    text = text.replace('≈', '约等于')
-    text = text.replace('<', '小于')
+    text = text.replace("²", "平方")
+    text = text.replace("³", "立方")
+    text = text.replace("√", "根号")
+    text = text.replace("≈", "约等于")
+    text = text.replace("<", "小于")
     return text
 
 
 # remove meaningless symbol
 def remove_bracket(text):
-    text = text.replace('（', ' ').replace('）', ' ')
-    text = text.replace('【', ' ').replace('】', ' ')
-    text = text.replace('`', '').replace('`', '')
+    text = text.replace("（", " ").replace("）", " ")
+    text = text.replace("【", " ").replace("】", " ")
+    text = text.replace("`", "").replace("`", "")
     text = text.replace("——", " ")
     return text
 
@@ -71,7 +38,7 @@ def spell_out_number(text: str, inflect_parser):
     for i, c in enumerate(text):
         if not c.isdigit():
             if st is not None:
-                num_str = inflect_parser.number_to_words(text[st: i])
+                num_str = inflect_parser.number_to_words(text[st:i])
                 new_text.append(num_str)
                 st = None
             new_text.append(c)
@@ -81,7 +48,7 @@ def spell_out_number(text: str, inflect_parser):
     if st is not None and st < len(text):
         num_str = inflect_parser.number_to_words(text[st:])
         new_text.append(num_str)
-    return ''.join(new_text)
+    return "".join(new_text)
 
 
 # split paragrah logic：
@@ -102,18 +69,18 @@ def split_paragraph(text: str, tokenize, lang="zh", token_max_n=80, token_min_n=
             return len(tokenize(_text)) < merge_len
 
     if lang == "zh":
-        pounc = ['。', '？', '！', '；', '：', '、', '.', '?', '!', ';']
+        pounc = ["。", "？", "！", "；", "：", "、", ".", "?", "!", ";"]
     else:
-        pounc = ['.', '?', '!', ';', ':']
+        pounc = [".", "?", "!", ";", ":"]
     if comma_split:
-        pounc.extend(['，', ','])
+        pounc.extend(["，", ","])
     st = 0
     utts = []
     for i, c in enumerate(text):
         if c in pounc:
-            if len(text[st: i]) > 0:
-                utts.append(text[st: i] + c)
-            if i + 1 < len(text) and text[i + 1] in ['"', '”']:
+            if len(text[st:i]) > 0:
+                utts.append(text[st:i] + c)
+            if i + 1 < len(text) and text[i + 1] in ['"', "”"]:
                 tmp = utts.pop(-1)
                 utts.append(tmp + text[i + 1])
                 st = i + 2
@@ -121,9 +88,9 @@ def split_paragraph(text: str, tokenize, lang="zh", token_max_n=80, token_min_n=
                 st = i + 1
     if len(utts) == 0:
         if lang == "zh":
-            utts.append(text + '。')
+            utts.append(text + "。")
         else:
-            utts.append(text + '.')
+            utts.append(text + ".")
     final_utts = []
     cur_utt = ""
     for utt in utts:
@@ -145,13 +112,13 @@ def replace_blank(text: str):
     out_str = []
     for i, c in enumerate(text):
         if c == " ":
-            if ((text[i + 1].isascii() and text[i + 1] != " ") and
-                    (text[i - 1].isascii() and text[i - 1] != " ")):
+            if (text[i + 1].isascii() and text[i + 1] != " ") and (text[i - 1].isascii() and text[i - 1] != " "):
                 out_str.append(c)
         else:
             out_str.append(c)
     return "".join(out_str)
 
+
 def clean_markdown(md_text: str) -> str:
     # 去除代码块 ``` ```（包括多行）
     md_text = re.sub(r"```.*?```", "", md_text, flags=re.DOTALL)
@@ -164,9 +131,9 @@ def clean_markdown(md_text: str) -> str:
 
     # 去除链接但保留文本 [text](url) -> text
     md_text = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", md_text)
-    
+
     # 替换无序列表符号
-    md_text = re.sub(r'^(\s*)-\s+', r'\1', md_text, flags=re.MULTILINE)
+    md_text = re.sub(r"^(\s*)-\s+", r"\1", md_text, flags=re.MULTILINE)
 
     # 去除HTML标签
     md_text = re.sub(r"<[^>]+>", "", md_text)
@@ -185,60 +152,37 @@ def clean_text(text):
     # 去除 Markdown 语法
     text = clean_markdown(text)
     # 匹配并移除表情符号
-    text = regex.compile(r'\p{Emoji_Presentation}|\p{Emoji}\uFE0F', flags=regex.UNICODE).sub("",text)
+    text = regex.compile(r"\p{Emoji_Presentation}|\p{Emoji}\uFE0F", flags=regex.UNICODE).sub("", text)
     # 去除换行符
     text = text.replace("\n", " ")
     text = text.replace("\t", " ")
-    text = text.replace('"', "\“")
+    text = text.replace("“", '"').replace("”", '"')
     return text
 
+
 class TextNormalizer:
     def __init__(self, tokenizer=None):
         self.tokenizer = tokenizer
-        self.zh_tn_model = ZhNormalizer(remove_erhua=False, full_to_half=False, remove_interjections=False, overwrite_cache=True)
-        self.en_tn_model = EnNormalizer()
+        self.zh_tn_model = Normalizer(lang="zh", operator="tn", remove_erhua=True)
+        self.en_tn_model = Normalizer(lang="en", operator="tn")
         self.inflect_parser = inflect.engine()
-    
+
     def normalize(self, text, split=False):
         # 去除 Markdown 语法，去除表情符号，去除换行符
         lang = "zh" if contains_chinese(text) else "en"
         text = clean_text(text)
         if lang == "zh":
-            text = text.replace("=", "等于") # 修复 ”550 + 320 等于 870 千卡。“ 被错误正则为 ”五百五十加三百二十等于八七十千卡.“
-            if re.search(r'([\d$%^*_+≥≤≠×÷?=])', text): # 避免 英文连字符被错误正则为减
-                text = re.sub(r'(?<=[a-zA-Z0-9])-(?=\d)', ' - ', text) # 修复 x-2 被正则为 x负2
-                text = self.zh_tn_model.normalize(text)
-            text = re.sub(r'(?<=[a-zA-Z0-9])-(?=\d)', ' - ', text) # 修复 x-2 被正则为 x负2
+            text = text.replace(
+                "=", "等于"
+            )  # 修复 ”550 + 320 等于 870 千卡。“ 被错误正则为 ”五百五十加三百二十等于八七十千卡.“
+            if re.search(r"([\d$%^*_+≥≤≠×÷?=])", text):  # 避免 英文连字符被错误正则为减
+                text = re.sub(r"(?<=[a-zA-Z0-9])-(?=\d)", " - ", text)  # 修复 x-2 被正则为 x负2
             text = self.zh_tn_model.normalize(text)
             text = replace_blank(text)
             text = replace_corner_mark(text)
             text = remove_bracket(text)
-            text = re.sub(r'[，,]+$', '。', text)
         else:
             text = self.en_tn_model.normalize(text)
             text = spell_out_number(text, self.inflect_parser)
         if split is False:
             return text
-        
-        
-if __name__ == "__main__":
-    text_normalizer = TextNormalizer()
-    text = r"""今天我们学习一元二次方程。一元二次方程的标准形式是：
-ax2+bx+c=0ax^2 + bx + c = 0ax2+bx+c=0 
-其中，aaa、bbb 和 ccc 是常数，xxx 是变量。这个方程的解可以通过求根公式来找到。
-一元二次方程的解法有几种：
-  - 因式分解法：通过将方程因式分解来求解。我们首先尝试将方程表达成两个括号的形式，解决方程的解。比如，方程x2−5x+6=0x^2 - 5x + 6 = 0x2−5x+6=0可以因式分解为(x−2)(x−3)=0(x - 2)(x - 3) = 0(x−2)(x−3)=0，因此根为2和3。
-  - 配方法：通过配方将方程转化为完全平方的形式，从而解出。我们通过加上或减去适当的常数来完成这一过程，使得方程可以直接写成一个完全平方的形式。
-  - 求根公式：我们可以使用求根公式直接求出方程的解。这个公式适用于所有的一元二次方程，即使我们无法通过因式分解或配方法来解决时，也能使用该公式。
-公式：x=−b±b2−4ac2ax = \frac{-b \pm \sqrt{b^2 - 4ac}}{2a}x=2a−b±b2−4ac​​这个公式可以帮助我们求解任何一元二次方程的根。
-对于一元二次方程，我们需要了解判别式。判别式的作用是帮助我们判断方程的解的个数和性质。判别式 Δ\DeltaΔ 由下式给出：Δ=b2−4ac\Delta = b^2 - 4acΔ=b2−4ac 根据判别式的值，我们可以知道：
-  - 如果 Δ>0\Delta > 0Δ>0，方程有两个不相等的实数解。这是因为判别式大于0时，根号内的值是正数，所以我们可以得到两个不同的解。
-  - 如果 Δ=0\Delta = 0Δ=0，方程有一个实数解。这是因为根号内的值为零，导致两个解相等，也就是说方程有一个解。
-  - 如果 Δ<0\Delta < 0Δ<0，方程没有实数解。这意味着根号内的值是负数，无法进行实数运算，因此方程没有实数解，可能有复数解。"""
-    texts = ["这是一个公式 (a+b)³=a³+3a²b+3ab²+b³ S=(a×b)÷2", "这样的发展为AI仅仅作为“工具”这一观点提出了新的挑战，", "550 + 320 = 870千卡。", "解一元二次方程：3x^2+x-2=0", "你好啊"]
-    texts = [text]
-    for text in texts:
-        text = text_normalizer.normalize(text)
-        print(text)
-        for t in cut_sentence_with_fix_length(text, 15):
-            print(t)

src/voxcpm/zipenhancer.py

@@ -0,0 +1,72 @@
+"""
+ZipEnhancer Module - Audio Denoising Enhancer
+
+Provides on-demand import ZipEnhancer functionality for audio denoising processing.
+Related dependencies are imported only when denoising functionality is needed.
+"""
+
+import os
+import tempfile
+from typing import Optional
+import torchaudio
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+
+
+class ZipEnhancer:
+    """ZipEnhancer Audio Denoising Enhancer"""
+
+    def __init__(self, model_path: str = "iic/speech_zipenhancer_ans_multiloss_16k_base"):
+        """
+        Initialize ZipEnhancer
+        Args:
+            model_path: ModelScope model path or local path
+        """
+        self.model_path = model_path
+        self._pipeline = pipeline(Tasks.acoustic_noise_suppression, model=self.model_path)
+
+    def _normalize_loudness(self, wav_path: str):
+        """
+        Audio loudness normalization
+
+        Args:
+            wav_path: Audio file path
+        """
+        audio, sr = torchaudio.load(wav_path)
+        loudness = torchaudio.functional.loudness(audio, sr)
+        normalized_audio = torchaudio.functional.gain(audio, -20 - loudness)
+        torchaudio.save(wav_path, normalized_audio, sr)
+
+    def enhance(self, input_path: str, output_path: Optional[str] = None, normalize_loudness: bool = True) -> str:
+        """
+        Audio denoising enhancement
+        Args:
+            input_path: Input audio file path
+            output_path: Output audio file path (optional, creates temp file by default)
+            normalize_loudness: Whether to perform loudness normalization
+        Returns:
+            str: Output audio file path
+        Raises:
+            RuntimeError: If pipeline is not initialized or processing fails
+        """
+        if not os.path.exists(input_path):
+            raise FileNotFoundError(f"Input audio file does not exist: {input_path}")
+        # Create temporary file if no output path is specified
+        if output_path is None:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
+                output_path = tmp_file.name
+        try:
+            # Perform denoising processing
+            self._pipeline(input_path, output_path=output_path)
+            # Loudness normalization
+            if normalize_loudness:
+                self._normalize_loudness(output_path)
+            return output_path
+        except Exception as e:
+            # Clean up possibly created temporary files
+            if output_path and os.path.exists(output_path):
+                try:
+                    os.unlink(output_path)
+                except OSError:
+                    pass
+            raise RuntimeError(f"Audio denoising processing failed: {e}")

tests/test_cli.py

@@ -0,0 +1,512 @@
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+ROOT = Path(__file__).resolve().parents[1]
+CLI_PATH = ROOT / "src" / "voxcpm" / "cli.py"
+V1_MODEL_PATH = ROOT / "models" / "openbmb__VoxCPM1.5"
+V2_MODEL_PATH = ROOT / "models" / "VoxCPM2-1B-newaudiovae-6hz-nope-sft"
+
+
+pkg = types.ModuleType("voxcpm")
+pkg.__path__ = [str(ROOT / "src" / "voxcpm")]
+sys.modules.setdefault("voxcpm", pkg)
+
+core_stub = types.ModuleType("voxcpm.core")
+
+
+class StubVoxCPM:
+    pass
+
+
+core_stub.VoxCPM = StubVoxCPM
+sys.modules["voxcpm.core"] = core_stub
+
+spec = importlib.util.spec_from_file_location("voxcpm.cli", CLI_PATH)
+cli = importlib.util.module_from_spec(spec)
+sys.modules["voxcpm.cli"] = cli
+assert spec.loader is not None
+spec.loader.exec_module(cli)
+
+
+class DummyTTSModel:
+    sample_rate = 16000
+
+
+class DummyModel:
+    def __init__(self):
+        self.tts_model = DummyTTSModel()
+        self.calls = []
+
+    def generate(self, **kwargs):
+        self.calls.append(kwargs)
+        return np.zeros(160, dtype=np.float32)
+
+
+def run_main(monkeypatch, argv):
+    monkeypatch.setattr(sys, "argv", ["voxcpm", *argv])
+    cli.main()
+
+
+def test_parser_defaults_to_voxcpm2():
+    parser = cli._build_parser()
+    args = parser.parse_args(["design", "--text", "hello", "--output", "out.wav"])
+    assert args.hf_model_id == "openbmb/VoxCPM2"
+    assert args.no_optimize is False
+
+
+def test_load_model_respects_no_optimize_for_local_model(monkeypatch):
+    calls = {}
+
+    class FakeVoxCPM:
+        def __init__(self, **kwargs):
+            calls["kwargs"] = kwargs
+            self.tts_model = DummyTTSModel()
+
+    monkeypatch.setattr(cli, "VoxCPM", FakeVoxCPM)
+    args = cli._build_parser().parse_args(
+        [
+            "design",
+            "--text",
+            "hello",
+            "--output",
+            "out.wav",
+            "--model-path",
+            str(V2_MODEL_PATH),
+            "--no-optimize",
+        ]
+    )
+
+    cli.load_model(args)
+
+    assert calls["kwargs"]["optimize"] is False
+
+
+def test_load_model_defaults_optimize_for_hf(monkeypatch):
+    calls = {}
+
+    class FakeVoxCPM:
+        @classmethod
+        def from_pretrained(cls, **kwargs):
+            calls["kwargs"] = kwargs
+            return DummyModel()
+
+    monkeypatch.setattr(cli, "VoxCPM", FakeVoxCPM)
+    args = cli._build_parser().parse_args(
+        [
+            "design",
+            "--text",
+            "hello",
+            "--output",
+            "out.wav",
+        ]
+    )
+
+    cli.load_model(args)
+
+    assert calls["kwargs"]["optimize"] is True
+
+
+def test_load_model_respects_no_optimize_for_hf(monkeypatch):
+    calls = {}
+
+    class FakeVoxCPM:
+        @classmethod
+        def from_pretrained(cls, **kwargs):
+            calls["kwargs"] = kwargs
+            return DummyModel()
+
+    monkeypatch.setattr(cli, "VoxCPM", FakeVoxCPM)
+    args = cli._build_parser().parse_args(
+        [
+            "design",
+            "--text",
+            "hello",
+            "--output",
+            "out.wav",
+            "--no-optimize",
+        ]
+    )
+
+    cli.load_model(args)
+
+    assert calls["kwargs"]["optimize"] is False
+
+
+def test_design_subcommand_applies_control(monkeypatch, tmp_path):
+    dummy_model = DummyModel()
+    monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
+    monkeypatch.setattr(cli.sf, "write", lambda *args, **kwargs: None)
+
+    run_main(
+        monkeypatch,
+        [
+            "design",
+            "--text",
+            "hello",
+            "--control",
+            "warm female voice",
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    assert dummy_model.calls[0]["text"] == "(warm female voice)hello"
+    assert dummy_model.calls[0]["prompt_wav_path"] is None
+    assert dummy_model.calls[0]["reference_wav_path"] is None
+
+
+def test_clone_subcommand_reads_prompt_file(monkeypatch, tmp_path):
+    dummy_model = DummyModel()
+    prompt_audio = tmp_path / "prompt.wav"
+    prompt_audio.write_bytes(b"RIFF")
+    prompt_file = tmp_path / "prompt.txt"
+    prompt_file.write_text("prompt transcript\n", encoding="utf-8")
+
+    monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
+    monkeypatch.setattr(cli.sf, "write", lambda *args, **kwargs: None)
+
+    run_main(
+        monkeypatch,
+        [
+            "clone",
+            "--text",
+            "hello",
+            "--prompt-audio",
+            str(prompt_audio),
+            "--prompt-file",
+            str(prompt_file),
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    assert dummy_model.calls[0]["prompt_wav_path"] == str(prompt_audio)
+    assert dummy_model.calls[0]["prompt_text"] == "prompt transcript"
+
+
+def test_clone_rejects_reference_audio_for_v1_local_model(monkeypatch, tmp_path):
+    reference_audio = tmp_path / "ref.wav"
+    reference_audio.write_bytes(b"RIFF")
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "clone",
+            "--text",
+            "hello",
+            "--reference-audio",
+            str(reference_audio),
+            "--model-path",
+            str(V1_MODEL_PATH),
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+
+def test_clone_rejects_reference_audio_for_v1_hf_model_id(monkeypatch, tmp_path):
+    reference_audio = tmp_path / "ref.wav"
+    reference_audio.write_bytes(b"RIFF")
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "clone",
+            "--text",
+            "hello",
+            "--reference-audio",
+            str(reference_audio),
+            "--hf-model-id",
+            "openbmb/VoxCPM1.5",
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+
+def test_legacy_root_args_still_work_and_warn(monkeypatch, tmp_path, capsys):
+    dummy_model = DummyModel()
+    monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
+    monkeypatch.setattr(cli.sf, "write", lambda *args, **kwargs: None)
+
+    run_main(
+        monkeypatch,
+        [
+            "--text",
+            "hello",
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    captured = capsys.readouterr()
+    assert "deprecated" in captured.err
+    assert dummy_model.calls[0]["text"] == "hello"
+
+
+def test_batch_subcommand_applies_control(monkeypatch, tmp_path):
+    dummy_model = DummyModel()
+    input_file = tmp_path / "texts.txt"
+    input_file.write_text("hello\nworld\n", encoding="utf-8")
+
+    monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
+    monkeypatch.setattr(cli.sf, "write", lambda *args, **kwargs: None)
+
+    run_main(
+        monkeypatch,
+        [
+            "batch",
+            "--input",
+            str(input_file),
+            "--output-dir",
+            str(tmp_path / "outs"),
+            "--control",
+            "calm narrator",
+        ],
+    )
+
+    assert [call["text"] for call in dummy_model.calls] == [
+        "(calm narrator)hello",
+        "(calm narrator)world",
+    ]
+
+
+def test_legacy_clone_with_prompt_file_still_works(monkeypatch, tmp_path, capsys):
+    dummy_model = DummyModel()
+    prompt_audio = tmp_path / "prompt.wav"
+    prompt_audio.write_bytes(b"RIFF")
+    prompt_file = tmp_path / "prompt.txt"
+    prompt_file.write_text("legacy transcript", encoding="utf-8")
+
+    monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
+    monkeypatch.setattr(cli.sf, "write", lambda *args, **kwargs: None)
+
+    run_main(
+        monkeypatch,
+        [
+            "--text",
+            "hello",
+            "--prompt-audio",
+            str(prompt_audio),
+            "--prompt-file",
+            str(prompt_file),
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    captured = capsys.readouterr()
+    assert "deprecated" in captured.err
+    assert dummy_model.calls[0]["prompt_text"] == "legacy transcript"
+
+
+def test_invalid_prompt_text_and_prompt_file_combination(monkeypatch, tmp_path, capsys):
+    prompt_audio = tmp_path / "prompt.wav"
+    prompt_audio.write_bytes(b"RIFF")
+    prompt_file = tmp_path / "prompt.txt"
+    prompt_file.write_text("transcript", encoding="utf-8")
+
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "clone",
+            "--text",
+            "hello",
+            "--prompt-audio",
+            str(prompt_audio),
+            "--prompt-text",
+            "inline transcript",
+            "--prompt-file",
+            str(prompt_file),
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+    assert "Use either --prompt-text or --prompt-file" in capsys.readouterr().err
+
+
+def test_missing_prompt_file_reports_parser_error(monkeypatch, tmp_path, capsys):
+    prompt_audio = tmp_path / "prompt.wav"
+    prompt_audio.write_bytes(b"RIFF")
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "clone",
+            "--text",
+            "hello",
+            "--prompt-audio",
+            str(prompt_audio),
+            "--prompt-file",
+            str(tmp_path / "missing.txt"),
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+    assert "prompt text file" in capsys.readouterr().err
+
+
+def test_design_rejects_prompt_audio_args(monkeypatch, tmp_path, capsys):
+    prompt_audio = tmp_path / "prompt.wav"
+    prompt_audio.write_bytes(b"RIFF")
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "design",
+            "--text",
+            "hello",
+            "--prompt-audio",
+            str(prompt_audio),
+            "--prompt-text",
+            "transcript",
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+    assert "does not accept prompt/reference audio" in capsys.readouterr().err
+
+
+def test_clone_rejects_prompt_audio_without_transcript(monkeypatch, tmp_path, capsys):
+    prompt_audio = tmp_path / "prompt.wav"
+    prompt_audio.write_bytes(b"RIFF")
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "clone",
+            "--text",
+            "hello",
+            "--prompt-audio",
+            str(prompt_audio),
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+    assert (
+        "--prompt-audio requires --prompt-text or --prompt-file"
+        in capsys.readouterr().err
+    )
+
+
+def test_clone_rejects_transcript_without_prompt_audio(monkeypatch, tmp_path, capsys):
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "clone",
+            "--text",
+            "hello",
+            "--prompt-text",
+            "transcript",
+            "--output",
+            str(tmp_path / "out.wav"),
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+    assert (
+        "--prompt-text/--prompt-file requires --prompt-audio" in capsys.readouterr().err
+    )
+
+
+def test_batch_rejects_control_with_prompt_transcript(monkeypatch, tmp_path, capsys):
+    input_file = tmp_path / "texts.txt"
+    input_file.write_text("hello\n", encoding="utf-8")
+    prompt_audio = tmp_path / "prompt.wav"
+    prompt_audio.write_bytes(b"RIFF")
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "voxcpm",
+            "batch",
+            "--input",
+            str(input_file),
+            "--output-dir",
+            str(tmp_path / "outs"),
+            "--control",
+            "calm narrator",
+            "--prompt-audio",
+            str(prompt_audio),
+            "--prompt-text",
+            "transcript",
+        ],
+    )
+
+    with pytest.raises(SystemExit):
+        cli.main()
+
+    assert "--control cannot be used together" in capsys.readouterr().err
+
+
+def test_detect_model_architecture_uses_local_configs():
+    parser = cli._build_parser()
+    v1_args = parser.parse_args(
+        [
+            "clone",
+            "--text",
+            "hello",
+            "--reference-audio",
+            "ref.wav",
+            "--model-path",
+            str(V1_MODEL_PATH),
+            "--output",
+            "out.wav",
+        ]
+    )
+    v2_args = parser.parse_args(
+        [
+            "clone",
+            "--text",
+            "hello",
+            "--reference-audio",
+            "ref.wav",
+            "--model-path",
+            str(V2_MODEL_PATH),
+            "--output",
+            "out.wav",
+        ]
+    )
+
+    assert cli.detect_model_architecture(v1_args) == "voxcpm"
+    assert cli.detect_model_architecture(v2_args) == "voxcpm2"