wehub/openbmb--voxcpm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:37:31 +08:00
commit d30c1ed902
76 changed files with 19614 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/publish-to-pypi.yml
+54
View File
@@ -0,0 +1,54 @@
name: Publish Python 🐍 distribution 📦 to PyPI
on:
release:
types: [created]
jobs:
build:
name: Build distribution 📦
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: "3.x"
- name: Install pypa/build
run: >-
python3 -m
pip install
build
--user
- name: Build a binary wheel and a source tarball
run: python3 -m build
- name: Store the distribution packages
uses: actions/upload-artifact@v4
with:
name: python-package-distributions
path: dist/
publish-to-pypi:
name: >-
Publish Python 🐍 distribution 📦 to PyPI
if: startsWith(github.ref, 'refs/tags/') # only publish to PyPI on tag pushes
needs:
- build
runs-on: ubuntu-latest
environment:
name: pypi
url: https://pypi.org/p/voxcpm
permissions:
id-token: write
steps:
- name: Download all the dists
uses: actions/download-artifact@v4
with:
name: python-package-distributions
path: dist/
- name: Publish distribution 📦 to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
.gitignore
+7
View File
@@ -0,0 +1,7 @@
launch.json
.venv/
__pycache__
voxcpm.egg-info
.DS_Store
./pretrained_models/
app_local.py
LICENSE
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright OpenBMB
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
README.md
+692
View File
@@ -0,0 +1,692 @@
<h2 align="center">VoxCPM2: Tokenizer-Free TTS for Multilingual Speech Generation, Creative Voice Design, and True-to-Life Cloning</h2>
<p align="center">
<b>English</b> | <a href="./README_zh.md">中文</a>
</p>
<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" alt="DemoPage"></a>
<a href="https://arxiv.org/abs/2606.06928"><img src="https://img.shields.io/badge/arXiv-VoxCPM2%20Technical%20Report-red" alt="VoxCPM2 Technical Report"></a>
</p>
<div align="center">
<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>
<br>
<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 **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.
**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.
### ✨ 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) or [vLLM-Omni](https://github.com/vllm-project/vllm-omni) — official vLLM omni-modal serving for VoxCPM2 with PagedAttention and an OpenAI-compatible API
- 📜 **Fully Open-Source & Commercial-Ready** — Weights and code released under the [Apache-2.0](LICENSE) license, free for commercial use
**🌍 Supported Languages (30)**
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: 四川话, 粤语, 吴语, 东北话, 河南话, 陕西话, 山东话, 天津话, 闽南话
### 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) | [Technical Report](https://arxiv.org/abs/2606.06928)
- **[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)
- [On-Device Inference (llama.cpp-omni)](#-on-device-inference-llamacpp-omni)
- [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
```sh
pip install voxcpm
```
> **Requirements:** Python ≥ 3.10 (<3.13), 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
from voxcpm import VoxCPM
import soundfile as sf
model = VoxCPM.from_pretrained(
"openbmb/VoxCPM2",
load_denoiser=False,
)
wav = model.generate(
text="VoxCPM2 is the current recommended release for realistic multilingual speech synthesis.",
cfg_value=2.0,
inference_timesteps=10,
seed=42,
)
sf.write("demo.wav", wav, model.tts_model.sample_rate)
print("saved: demo.wav")
```
If you prefer downloading from ModelScope first, you can use:
```bash
pip install modelscope
```
```python
from modelscope import snapshot_download
snapshot_download("OpenBMB/VoxCPM2", local_dir='./pretrained_models/VoxCPM2') # specify the local directory to save the model
from voxcpm import VoxCPM
import soundfile as sf
model = VoxCPM.from_pretrained("./pretrained_models/VoxCPM2", load_denoiser=False)
wav = model.generate(
text="VoxCPM2 is the current recommended release for realistic multilingual speech synthesis.",
cfg_value=2.0,
inference_timesteps=10,
seed=42,
)
sf.write("demo.wav", wav, model.tts_model.sample_rate)
```
#### 🎨 Voice Design
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,
seed=42,
)
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="path/to/voice.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="path/to/voice.wav",
cfg_value=2.0,
inference_timesteps=10,
seed=42,
)
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="path/to/voice.wav",
prompt_text="The transcript of the reference audio.",
reference_wav_path="path/to/voice.wav", # optional, for better simliarity
)
sf.write("hifi_clone.wav", wav, model.tts_model.sample_rate)
```
**🔄 Streaming API**
```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)
```
### CLI Usage
```bash
# Voice design (no reference audio needed)
voxcpm design \
--text "VoxCPM2 brings studio-quality multilingual speech synthesis." \
--output out.wav
# 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" \
--seed 42 \
--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" \
--reference-audio path/to/voice.wav \ # optional, for better simliarity
--output out.wav
# Batch processing
voxcpm batch --input examples/input.txt --output-dir outs
# Optional post-generation timestamps with stable-ts
pip install "voxcpm[timestamps]"
voxcpm design \
--text "VoxCPM2 brings studio-quality multilingual speech synthesis." \
--output out.wav \
--timestamps \
--timestamp-level word \
--timestamp-language en
# Character timestamps are best-effort and are derived from word alignment
voxcpm design \
--text "欢迎使用 VoxCPM2。" \
--output out.wav \
--timestamps \
--timestamp-level char \
--timestamp-language zh
# Help
voxcpm --help
```
### Web Demo
```bash
python app.py --port 8808 # then open in browser: http://localhost:8808
```
Use `--device` to choose the runtime device:
```bash
python app.py --device auto
```
Supported values are `auto`, `cpu`, `mps`, `cuda`, and `cuda:N`. On Apple Silicon Macs, `auto` uses MPS when available.
### 🚢 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.
```bash
pip install nano-vllm-voxcpm
```
```python
from nanovllm_voxcpm import VoxCPM
import numpy as np, soundfile as sf
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()
```
> **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.
### 🏭 Production Serving (vLLM-Omni)
For production multi-tenant deployments, use **[vLLM-Omni](https://github.com/vllm-project/vllm-omni)** — the official vLLM project's omni-modal extension with native **VoxCPM2** support. PagedAttention KV cache, continuous batching, and a drop-in **OpenAI-compatible** `/v1/audio/speech` endpoint.
```bash
# Install from source (latest main — vllm-omni is rapidly evolving)
uv pip install vllm==0.19.0 --torch-backend=auto
git clone https://github.com/vllm-project/vllm-omni.git && cd vllm-omni
uv pip install -e .
```
See the [vLLM-Omni installation guide](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/) for other platforms (ROCm, XPU, MUSA, NPU) and Docker images.
```bash
# Launch an OpenAI-compatible TTS server (--omni enables omni-modal serving)
vllm serve openbmb/VoxCPM2 --omni --port 8000
# Call it from any OpenAI client
curl http://localhost:8000/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"model":"openbmb/VoxCPM2","input":"Hello from VoxCPM2 on vLLM-Omni!","voice":"default"}' \
--output out.wav
```
> Built on the upstream vLLM scheduler, with batched concurrent requests, streaming chunk delivery, and multi-GPU deployment out of the box. See the [VoxCPM2 example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm2) for full deployment recipes.
### 📱 On-Device Inference (llama.cpp-omni)
For on-device / edge deployment without Python, use **[llama.cpp-omni](https://github.com/tc-mb/llama.cpp-omni)** — a high-performance C++ inference engine built on llama.cpp, with native VoxCPM2 GGUF support on **CPU / Metal / CUDA / Vulkan**.
**1. Download GGUF weights** from [HuggingFace](https://huggingface.co/DennisHuang648/VoxCPM2-GGUF) | [ModelScope](https://modelscope.cn/models/DennisHuang/VoxCPM2-GGUF) — you need one **BaseLM** (F16 or Q8_0) + the **Acoustic** file. Q8_0 halves the download with negligible quality loss.
**2. Build**
```bash
git clone https://github.com/tc-mb/llama.cpp-omni.git && cd llama.cpp-omni
cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build --target voxcpm2-cli -j
```
> CMake auto-detects Metal (macOS) or CUDA (Linux with NVIDIA GPU).
**3. Run**
```bash
# Basic TTS
./build/bin/voxcpm2-cli \
-t "Hello, this is VoxCPM2 running through llama.cpp-omni." \
-o output.wav VoxCPM2-BaseLM-Q8_0.gguf VoxCPM2-Acoustic-F16.gguf
# Voice cloning (reference audio)
./build/bin/voxcpm2-cli \
-t "Cloned voice." -r speaker.wav -o clone.wav \
VoxCPM2-BaseLM-Q8_0.gguf VoxCPM2-Acoustic-F16.gguf
# Ultimate cloning (reference audio + transcript)
./build/bin/voxcpm2-cli \
-t "Target text." --prompt-wav speaker.wav --prompt-text "transcript of speaker.wav" \
-o clone.wav VoxCPM2-BaseLM-Q8_0.gguf VoxCPM2-Acoustic-F16.gguf
```
> **RTF ~1.76 (Q8_0) on Apple M4 Pro / Metal.** Key flags: `--cfg` (guidance scale), `--timesteps` (CFM steps), `--seed`, `--temperature`, `--stream`. See the [llama.cpp-omni repo](https://github.com/tc-mb/llama.cpp-omni) and [GGUF weights page](https://huggingface.co/DennisHuang648/VoxCPM2-GGUF) for full 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** | [arXiv](https://arxiv.org/abs/2606.06928) | — | [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) |
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.
<div align="center">
<img src="assets/voxcpm2_model.png" alt="VoxCPM2 Model Architecture" width="90%">
</div>
> 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
**Seed-TTS-eval WER(⬇)&SIM(⬆) Results (click to expand)**
| Model | Parameters | Open-Source | test-EN | | test-ZH | | test-Hard | |
| ----------------- | ---------- | ----------- | ------- | ------ | ------- | ------ | --------- | ------ |
| | | | WER/%⬇ | SIM/%⬆ | CER/%⬇ | SIM/%⬆ | CER/%⬇ | SIM/%⬆ |
| MegaTTS3 | 0.5B | ❌ | 2.79 | 77.1 | 1.52 | 79.0 | - | - |
| DiTAR | 0.6B | ❌ | 1.69 | 73.5 | 1.02 | 75.3 | - | - |
| CosyVoice3 | 0.5B | ❌ | 2.02 | 71.8 | 1.16 | 78.0 | 6.08 | 75.8 |
| 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 | - | - |
| 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 |
| 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 | 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 |
### CV3-eval
**CV3-eval Multilingual WER/CER(⬇) Results (click to expand)**
| 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 |
### MiniMax-Multilingual-Test
**Minimax-MLS-test WER(⬇) Results (click to expand)**
| 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 |
**Minimax-MLS-test SIM(⬆) Results (click to expand)**
| 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** |
### Internal 30-Language ASR Benchmark
We additionally run an internal multilingual intelligibility benchmark with **30 languages × 500 samples**. ASR transcription is evaluated via **Gemini 3.1 Flash Lite API**.
**Internal 30-Language ASR Benchmark (click to expand)**
| Language | Metric | VoxCPM2 | Fish S2-Pro |
| ---------------------- | ------ | --------- | ----------- |
| ar (Arabic) | CER | 1.23% | 0.30% |
| da (Danish) | WER | 2.70% | 3.52% |
| de (German) | WER | 0.96% | 0.64% |
| el (Greek) | WER | 3.17% | 4.61% |
| en (English) | WER | 0.42% | 1.03% |
| es (Spanish) | WER | 1.33% | 0.64% |
| fi (Finnish) | WER | 2.24% | 2.80% |
| fr (French) | WER | 2.16% | 2.34% |
| he (Hebrew) | CER | 2.98% | 15.27% |
| hi (Hindi) | CER | 0.79% | 0.91% |
| id (Indonesian) | WER | 1.36% | 1.68% |
| it (Italian) | WER | 1.65% | 1.08% |
| ja (Japanese) | CER | 2.40% | 1.82% |
| km (Khmer) | CER | 2.05% | 75.15% |
| ko (Korean) | CER | 0.95% | 0.29% |
| lo (Lao) | CER | 1.90% | 87.40% |
| ms (Malay) | WER | 1.75% | 1.41% |
| my (Burmese) | CER | 1.42% | 85.27% |
| nl (Dutch) | WER | 1.25% | 1.68% |
| no (Norwegian) | WER | 2.49% | 3.76% |
| pl (Polish) | WER | 1.90% | 1.65% |
| pt (Portuguese) | WER | 1.48% | 1.49% |
| ru (Russian) | WER | 0.90% | 0.86% |
| sv (Swedish) | WER | 2.22% | 2.63% |
| sw (Swahili) | CER | 1.07% | 2.02% |
| th (Thai) | CER | 0.94% | 1.92% |
| tl (Tagalog) | WER | 2.63% | 4.00% |
| tr (Turkish) | WER | 1.65% | 1.65% |
| vi (Vietnamese) | WER | 1.56% | 5.56% |
| zh (Chinese) | CER | 0.92% | 1.02% |
| Average (30 languages) | | **1.68%** | - |
### InstructTTSEval
**Instruction-Guided Voice Design Results (click to expand)**
| 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** |
---
## ⚙️ 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
# 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 |
| **[vLLM-Omni](https://github.com/vllm-project/vllm-omni)** | Official vLLM omni-modal serving for VoxCPM2 — PagedAttention, OpenAI-compatible API |
| **[llama.cpp-omni](https://github.com/tc-mb/llama.cpp-omni)** | Full-duplex omni inference engine — VoxCPM2 GGUF on CPU / Metal / CUDA / Vulkan |
| **[VoxCPM.cpp](https://github.com/bluryar/VoxCPM.cpp)** | GGML/GGUF: CPU, CUDA, Vulkan inference |
| **[audio.cpp](https://github.com/0xShug0/audio.cpp)** | ggml-based unified C++ inference framework — CPU/CUDA/Vulkan/Metal, CLI & server, no Python |
| **[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_RH_VoxCPM](https://github.com/HM-RunningHub/ComfyUI_RH_VoxCPM)** | Feature-complete ComfyUI workflow for VoxCPM 2 with multi-speaker generation, LoRA, and auto-ASR |
| **[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{zhou2026voxcpm2,
title = {VoxCPM2 Technical Report},
author = {Zhou, Yixuan and Zeng, Guoyang and Liu, Xin and Li, Xiang and Yu, Renjie and Gui, Jiancheng and Wu, Jiaheng and Wang, Ziyang and Shen, Xudong and Ye, Runchuan and Zhang, Zhisheng and Zhou, Jiuyang and Bai, Bingsong and Sun, Weiyue and Deng, Mengyuan and Shi, Qundong and Wu, Zhiyong and Liu, Zhiyuan},
journal = {arXiv preprint arXiv:2606.06928},
year = {2026},
}
@article{zhou2025voxcpm,
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
VoxCPM model weights and code are open-sourced under the [Apache-2.0](LICENSE) license.
## 🙏 Acknowledgments
- [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
[ModelBest](https://modelbest.cn/)     [THUHCSI](https://github.com/thuhcsi)
## ⭐ Star History
[Star History Chart](https://star-history.com/#OpenBMB/VoxCPM&Date)
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`OpenBMB/VoxCPM`
- 原始仓库：https://github.com/OpenBMB/VoxCPM
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
README_zh.md
+685
View File
@@ -0,0 +1,685 @@
<h2 align="center">VoxCPM2：基于连续表征的多语言语音合成、创意音色设计与高保真声音克隆</h2>
<p align="center">
<a href="./README.md">English</a> | <b>中文</b>
</p>
<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/zh-cn/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" alt="DemoPage"></a>
</p>
<div align="center">
<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>
<br>
<p align="center">
👋 欢迎加入社区，参与讨论与交流！
<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;"> 飞书群
</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 是一个**无离散音频分词器**（Tokenizer-Free）的语音合成系统，通过端到端的**扩散自回归架构**直接生成连续语音表征，绕过对音频的离散编码步骤，实现高度自然且富有表现力的语音合成。
**VoxCPM2** 是最新的版本 — 基于 [MiniCPM-4](https://github.com/OpenBMB/MiniCPM) 基座构建，总计 **20亿** 参数，在超过 **200万小时** 的多语种音频数据上训练，支持 **30种全球语言+9种中文方言**、**音色设计**、**可控声音克隆**，原生输出 **48kHz** 高质量音频。
### ✨ 核心特性
- 🌍 **30种语言语音合成** — 直接输入原始文本即可合成（支持语言详见下文），无需额外语言标签
- 🎨 **音色设计** — 用自然语言描述（性别、年龄、音色、情绪、语速……）凭空创建全新音色，无需参考音频
- 🎛️ **可控声音克隆** — 从参考音频片段克隆任意声音，可叠加风格指令控制情绪、语速和表现力，同时保持原始音色
- 🎙️ **极致克隆** — 提供参考音频及其文本内容，模型接着参考音频进行无缝续写，从而精准还原声音细节特征（与 VoxCPM1.5 一致）
- 🔊 **48kHz 高质量音频** — 输入 16kHz 参考音频，通过 AudioVAE V2 的非对称编解码设计直接输出 48kHz 高质量音频，内置超分能力
- 🧠 **语境感知合成** — 根据文本内容自动推断合适的韵律和表现力
- ⚡ **实时流式合成** — 在 NVIDIA RTX 4090 上 RTF 低至 ~0.3，通过 [Nano-vLLM](https://github.com/a710128/nanovllm-voxcpm) 或 [vLLM-Omni](https://github.com/vllm-project/vllm-omni)（官方 vLLM 全模态服务，原生支持 VoxCPM2，提供 PagedAttention 与 OpenAI 兼容 API）加速后可达 ~0.13
- 📜 **完全开源，商用就绪** — 权重和代码基于 [Apache-2.0](LICENSE) 协议发布，免费商用
<summary><b>🌍 支持的语言（30种）</b></summary>
<br>
阿拉伯语、缅甸语、中文、丹麦语、荷兰语、英语、芬兰语、法语、德语、希腊语、希伯来语、印地语、印尼语、意大利语、日语、高棉语、韩语、老挝语、马来语、挪威语、波兰语、葡萄牙语、俄语、西班牙语、斯瓦希里语、瑞典语、菲律宾语、泰语、土耳其语、越南语
中国方言：四川话、粤语、吴语、东北话、河南话、陕西话、山东话、天津话、闽南话
### 最新动态
* **[2026.04]** 🔥 发布 **VoxCPM2** — 20亿参数，30种语言，音色设计与可控声音克隆，48kHz 音频输出！[模型权重](https://huggingface.co/openbmb/VoxCPM2) | [使用文档](https://voxcpm.readthedocs.io/zh-cn/latest/) | [在线体验](https://huggingface.co/spaces/OpenBMB/VoxCPM-Demo) | [官网体验](https://voxcpm.modelbest.cn/) (适用国内访问) | [技术报告](https://arxiv.org/abs/2606.06928)
* **[2025.12]** 🎉 开源 **VoxCPM1.5** [模型权重](https://huggingface.co/openbmb/VoxCPM1.5)，支持 SFT 和 LoRA 微调。(**🏆 GitHub Trending #1**)
* **[2025.09]** 🔥 发布 VoxCPM [技术报告](https://arxiv.org/abs/2509.24650)。
* **[2025.09]** 🎉 开源 **VoxCPM-0.5B** [模型权重](https://huggingface.co/openbmb/VoxCPM-0.5B) (**🏆 HuggingFace Trending #1**)
---
## 目录
- [快速开始](#-快速开始)
- [安装](#安装)
- [Python API](#python-api)
- [命令行使用](#命令行使用)
- [Web Demo](#web-demo)
- [生产部署](#-生产部署nano-vllm)
- [端侧推理（llama.cpp-omni）](#-端侧推理llamacpp-omni)
- [模型与版本](#-模型与版本)
- [性能评测](#-性能评测)
- [微调](#%EF%B8%8F-微调)
- [文档](#-文档)
- [生态与社区](#-生态与社区)
- [风险与局限性](#%EF%B8%8F-风险与局限性)
- [引用](#-引用)
---
## 🚀 快速开始
### 安装
```sh
pip install voxcpm
```
> **环境要求：** Python ≥ 3.10 (<3.13)，PyTorch ≥ 2.5.0，CUDA ≥ 12.0。详见 [快速开始文档](https://voxcpm.readthedocs.io/zh-cn/latest/quickstart.html)。
### Python API
#### 🗣️ 文本转语音
```python
from voxcpm import VoxCPM
import soundfile as sf
model = VoxCPM.from_pretrained(
"openbmb/VoxCPM2",
load_denoiser=False,
)
wav = model.generate(
text="VoxCPM2 是目前推荐使用的多语言语音合成版本。",
cfg_value=2.0,
inference_timesteps=10,
seed=42,
)
sf.write("demo.wav", wav, model.tts_model.sample_rate)
print("已保存: demo.wav")
```
如果你希望先从 ModelScope 下载模型到本地（适用于国内网络访问），可以使用：
```bash
pip install modelscope
```
```python
from modelscope import snapshot_download
snapshot_download("OpenBMB/VoxCPM2", local_dir='./pretrained_models/VoxCPM2') # 指定模型保存的本地路径
from voxcpm import VoxCPM
import soundfile as sf
model = VoxCPM.from_pretrained('./pretrained_models/VoxCPM2', load_denoiser=False)
wav = model.generate(
text="VoxCPM2 是目前推荐使用的多语言语音合成版本。",
cfg_value=2.0,
inference_timesteps=10,
seed=42,
)
sf.write("demo.wav", wav, model.tts_model.sample_rate)
```
#### 🎨 音色设计
用自然语言描述创建全新音色，无需参考音频。**格式：** 在 `text` 开头用括号写入音色描述（如 `"(音色描述)要合成的文本。"`）：
```python
wav = model.generate(
text="(年轻女性，声音温柔甜美)你好，欢迎使用VoxCPM2！",
cfg_value=2.0,
inference_timesteps=10,
seed=42,
)
sf.write("voice_design.wav", wav, model.tts_model.sample_rate)
```
#### 🎛️ 可控声音克隆
上传一段参考音频，模型克隆其音色，同时可以使用控制指令调节语速、情绪或风格。
```python
wav = model.generate(
text="这是VoxCPM2生成的克隆语音。",
reference_wav_path="path/to/voice.wav",
)
sf.write("clone.wav", wav, model.tts_model.sample_rate)
wav = model.generate(
text="(稍快一点，欢快的语气)这是带风格控制的克隆语音。",
reference_wav_path="path/to/voice.wav",
cfg_value=2.0,
inference_timesteps=10,
seed=42,
)
sf.write("controllable_clone.wav", wav, model.tts_model.sample_rate)
```
#### 🎙️ 极致克隆
提供参考音频及其精确文本转录，实现基于音频续写的高保真克隆。为获得最高克隆相似度，可将同一音频同时传给 `reference_wav_path` 和 `prompt_wav_path`：
```python
wav = model.generate(
text="这是使用VoxCPM2的极致克隆演示。",
prompt_wav_path="path/to/voice.wav",
prompt_text="参考音频的文本转录。",
reference_wav_path="path/to/voice.wav", # 可选，提升相似度
)
sf.write("hifi_clone.wav", wav, model.tts_model.sample_rate)
```
<details>
<summary><b>🔄 流式 API</b></summary>
```python
import numpy as np
chunks = []
for chunk in model.generate_streaming(
text="使用VoxCPM进行流式语音合成非常简单！",
):
chunks.append(chunk)
wav = np.concatenate(chunks)
sf.write("streaming.wav", wav, model.tts_model.sample_rate)
```
</details>
### 命令行使用
```bash
# 音色设计（无需参考音频）
voxcpm design \
--text "VoxCPM2带来全新语音合成体验。" \
--output out.wav
# 可控声音克隆（带风格控制）
voxcpm design \
--text "VoxCPM2带来全新语音合成体验。" \
--control "年轻女声，温暖温柔，略带微笑" \
--seed 42 \
--output out.wav
# 声音克隆（参考音频）
voxcpm clone \
--text "这是一个声音克隆的演示。" \
--reference-audio path/to/voice.wav \
--output out.wav
# 极致克隆（提示音频 + 转录文本）
voxcpm clone \
--text "这是一个声音克隆的演示。" \
--prompt-audio path/to/voice.wav \
--prompt-text "参考音频转录文本" \
--reference-audio path/to/voice.wav \
--output out.wav
# 批量处理
voxcpm batch --input examples/input.txt --output-dir outs
# 可选的生成后时间戳对齐（基于 stable-ts）
pip install "voxcpm[timestamps]"
voxcpm design \
--text "VoxCPM2带来全新语音合成体验。" \
--output out.wav \
--timestamps \
--timestamp-level word \
--timestamp-language zh
# 字级时间戳是 best-effort，会基于词级对齐结果拆分
voxcpm design \
--text "欢迎使用 VoxCPM2。" \
--output out.wav \
--timestamps \
--timestamp-level char \
--timestamp-language zh
# 帮助
voxcpm --help
```
### Web Demo
```bash
python app.py --port 8808 # 然后在浏览器打开 http://localhost:8808
```
使用 `--device` 选择运行设备：
```bash
python app.py --device auto
```
支持的取值包括 `auto`、`cpu`、`mps`、`cuda` 和 `cuda:N`。在 Apple Silicon Mac 上，`auto` 会在可用时使用 MPS。
### 🚢 生产部署（Nano-vLLM）
如需高吞吐量部署，使用 [**Nano-vLLM-VoxCPM**](https://github.com/a710128/nanovllm-voxcpm) — 基于 Nano-vLLM 构建的专用推理引擎，支持并发请求和异步 API。
```bash
pip install nano-vllm-voxcpm
```
```python
from nanovllm_voxcpm import VoxCPM
import numpy as np, soundfile as sf
server = VoxCPM.from_pretrained(model="/path/to/VoxCPM", devices=[0])
chunks = list(server.generate(target_text="你好，我来自VoxCPM！"))
sf.write("out.wav", np.concatenate(chunks), 48000)
server.stop()
```
> **在 NVIDIA RTX 4090 上 RTF 低至 ~0.13**（标准 PyTorch 实现约 ~0.3），支持批量并发请求和 FastAPI HTTP 服务。详见 [Nano-vLLM-VoxCPM 仓库](https://github.com/a710128/nanovllm-voxcpm)。
### 🏭 生产环境部署（vLLM-Omni）
如需生产级多租户部署，使用 [**vLLM-Omni**](https://github.com/vllm-project/vllm-omni) — 官方 vLLM 项目的全模态扩展，原生支持 **VoxCPM2**。具备 PagedAttention KV 缓存、连续批处理，以及与 OpenAI 完全兼容的 `/v1/audio/speech` 接口。
```bash
# 从源码安装（最新 main 分支 —— vllm-omni 正在快速迭代）
uv pip install vllm==0.19.0 --torch-backend=auto
git clone https://github.com/vllm-project/vllm-omni.git && cd vllm-omni
uv pip install -e .
```
其他平台（ROCm、XPU、MUSA、NPU）与 Docker 镜像请参考 [vLLM-Omni 安装文档](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)。
```bash
# 启动 OpenAI 兼容的 TTS 服务（--omni 启用全模态服务）
vllm serve openbmb/VoxCPM2 --omni --port 8000
# 任意 OpenAI 客户端均可调用
curl http://localhost:8000/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{"model":"openbmb/VoxCPM2","input":"你好，欢迎使用 VoxCPM2 on vLLM-Omni！","voice":"default"}' \
--output out.wav
```
> 基于上游 vLLM 调度器构建，开箱即用支持批量并发、流式分块输出和多 GPU 部署。完整示例见 [VoxCPM2 部署样例](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm2)。
### 📱 端侧推理（llama.cpp-omni）
如需在端侧/消费级硬件上无 Python 运行，使用 **[llama.cpp-omni](https://github.com/tc-mb/llama.cpp-omni)** — 基于 llama.cpp 的高性能 C++ 推理引擎，原生支持 VoxCPM2 GGUF，可在 **CPU / Metal / CUDA / Vulkan** 上运行。
**1. 下载 GGUF 权重**：从 [HF下载](https://huggingface.co/DennisHuang648/VoxCPM2-GGUF) | [ModelScope](https://modelscope.cn/models/DennisHuang/VoxCPM2-GGUF)，需要一个 **BaseLM**（F16 或 Q8_0）+ **Acoustic** 文件。Q8_0 体积减半，质量损失可忽略。
**2. 编译**
```bash
git clone https://github.com/tc-mb/llama.cpp-omni.git && cd llama.cpp-omni
cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build --target voxcpm2-cli -j
```
> CMake 会自动检测并启用 Metal（macOS）或 CUDA（Linux + NVIDIA GPU）。
**3. 运行**
```bash
# 基础 TTS
./build/bin/voxcpm2-cli \
-t "你好，我是通过 llama.cpp-omni 运行的 VoxCPM2。" \
-o output.wav VoxCPM2-BaseLM-Q8_0.gguf VoxCPM2-Acoustic-F16.gguf
# 声音克隆（参考音频）
./build/bin/voxcpm2-cli \
-t "克隆的声音。" -r speaker.wav -o clone.wav \
VoxCPM2-BaseLM-Q8_0.gguf VoxCPM2-Acoustic-F16.gguf
# 精准克隆（参考音频 + 转写文本）
./build/bin/voxcpm2-cli \
-t "目标文本。" --prompt-wav speaker.wav --prompt-text "参考音频的转写文本" \
-o clone.wav VoxCPM2-BaseLM-Q8_0.gguf VoxCPM2-Acoustic-F16.gguf
```
> **在 Apple M4 Pro / Metal 上 RTF ~1.76（Q8_0）。** 主要参数：`--cfg`（引导尺度）、`--timesteps`（CFM 步数）、`--seed`、`--temperature`、`--stream`。详见 [llama.cpp-omni 仓库](https://github.com/tc-mb/llama.cpp-omni) 和 [GGUF 权重页面](https://huggingface.co/DennisHuang648/VoxCPM2-GGUF)。
> **完整参数说明、多场景示例与声音克隆技巧 →** [快速开始指南](https://voxcpm.readthedocs.io/zh-cn/latest/quickstart.html) | [使用指南](https://voxcpm.readthedocs.io/zh-cn/latest/usage_guide.html) | [Cookbook](https://voxcpm.readthedocs.io/zh-cn/latest/cookbook.html)
---
## 📦 模型与版本
| | **VoxCPM2** | **VoxCPM1.5** | **VoxCPM-0.5B** |
|---|:---:|:---:|:---:|
| **状态** | 🟢 最新版本 | 稳定版 | 旧版 |
| **主模型参数量** | 2B | 0.6B | 0.5B |
| **音频采样率** | 48kHz | 44.1kHz | 16kHz |
| **LM处理码率** | 6.25Hz | 6.25Hz | 12.5Hz |
| **语言支持数量** | 30 | 2（中文、英文） | 2（中文、英文） |
| **克隆模式** | 隔离参考音频（无需文本） & 音频续写 | 仅音频续写 | 仅音频续写 |
| **音色设计** | ✅ | — | — |
| **可控声音克隆** | ✅ | — | — |
| **SFT / LoRA** | ✅ | ✅ | ✅ |
| **RTF (RTX 4090)** | ~0.30 | ~0.15 | ~0.17 |
| **RTF Nano-VLLM (RTX 4090)** | ~0.13 | ~0.08 | ~0.10 |
| **显存占用** | ~8 GB | ~6 GB | ~5 GB |
| **模型权重** | [🤗 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) |
| **技术报告** | [arXiv](https://arxiv.org/abs/2606.06928) | — | [arXiv](https://arxiv.org/abs/2509.24650) [ICLR 2026](https://openreview.net/forum?id=h5KLpGoqzC) |
| **Demo 页面** | [音频示例](https://openbmb.github.io/voxcpm2-demopage) | — | [音频示例](https://openbmb.github.io/VoxCPM-demopage) |
VoxCPM2 采用**连续音频表征、扩散自回归**范式，模型在 **AudioVAE** 的连续隐空间中通过四阶段处理：**LocEnc → TSLM → RALM → LocDiT**，实现丰富的表现力语音合成和 48kHz 原生音频输出。
<div align="center">
<img src="assets/voxcpm2_model.png" alt="VoxCPM2 模型架构" width="90%">
</div>
> 完整架构细节、VoxCPM2 升级内容和模型对比表见 [架构设计文档](https://voxcpm.readthedocs.io/zh-cn/latest/models/architecture.html)。
---
## 📊 性能评测
VoxCPM2 在公开的零样本和可控 TTS 基准测试中取得了 SOTA 或可比的结果。
### Seed-TTS-eval
<details>
<summary><b>Seed-TTS-eval WER(⬇)&SIM(⬆) 结果（点击展开）</b></summary>
| Model | Parameters | Open-Source | test-EN | | test-ZH | | test-Hard | |
|------|------|------|:------------:|:--:|:------------:|:--:|:-------------:|:--:|
| | | | WER/%⬇ | SIM/%⬆| CER/%⬇| SIM/%⬆ | CER/%⬇ | SIM/%⬆ |
| MegaTTS3 | 0.5B | ❌ | 2.79 | 77.1 | 1.52 | 79.0 | - | - |
| DiTAR | 0.6B | ❌ | 1.69 | 73.5 | 1.02 | 75.3 | - | - |
| CosyVoice3 | 0.5B | ❌ | 2.02 | 71.8 | 1.16 | 78.0 | 6.08 | 75.8 |
| 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 | - | - |
| 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 |
| 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 | 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
<details>
<summary><b>CV3-eval 多语言 WER/CER(⬇) 结果（点击展开）</b></summary>
| 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(⬇) 结果（点击展开）</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(⬆) 结果（点击展开）</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>
### Internal 30-Language ASR Benchmark
我们额外进行了内部多语言可懂度评测：**30 语种 × 500 样本**，ASR 转写评估使用 **Gemini 3.1 Flash Lite API**。
<details>
<summary><b>内部30语种评测集ASR结果（点击展开）</b></summary>
| 语言 | 指标 | VoxCPM2 | Fish S2-Pro |
|---|---:|---:|---:|
| ar (阿拉伯语) | CER | 1.23% | 0.30% |
| da (丹麦语) | WER | 2.70% | 3.52% |
| de (德语) | WER | 0.96% | 0.64% |
| el (希腊语) | WER | 3.17% | 4.61% |
| en (英语) | WER | 0.42% | 1.03% |
| es (西班牙语) | WER | 1.33% | 0.64% |
| fi (芬兰语) | WER | 2.24% | 2.80% |
| fr (法语) | WER | 2.16% | 2.34% |
| he (希伯来语) | CER | 2.98% | 15.27% |
| hi (印地语) | CER | 0.79% | 0.91% |
| id (印尼语) | WER | 1.36% | 1.68% |
| it (意大利语) | WER | 1.65% | 1.08% |
| ja (日语) | CER | 2.40% | 1.82% |
| km (高棉语) | CER | 2.05% | 75.15% |
| ko (韩语) | CER | 0.95% | 0.29% |
| lo (老挝语) | CER | 1.90% | 87.40% |
| ms (马来语) | WER | 1.75% | 1.41% |
| my (缅甸语) | CER | 1.42% | 85.27% |
| nl (荷兰语) | WER | 1.25% | 1.68% |
| no (挪威语) | WER | 2.49% | 3.76% |
| pl (波兰语) | WER | 1.90% | 1.65% |
| pt (葡萄牙语) | WER | 1.48% | 1.49% |
| ru (俄语) | WER | 0.90% | 0.86% |
| sv (瑞典语) | WER | 2.22% | 2.63% |
| sw (斯瓦希里语) | CER | 1.07% | 2.02% |
| th (泰语) | CER | 0.94% | 1.92% |
| tl (菲律宾语) | WER | 2.63% | 4.00% |
| tr (土耳其语) | WER | 1.65% | 1.65% |
| vi (越南语) | WER | 1.56% | 5.56% |
| zh (中文) | CER | 0.92% | 1.02% |
| 平均（30 语种） | | **1.68%** | - |
</details>
### InstructTTSEval
<details>
<summary><b>指令驱动音色设计结果 (点击展开)</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>
---
## ⚙️ 微调
VoxCPM 支持**全参数微调（SFT）** 和 **LoRA 微调**。仅需 **5-10分钟** 的音频数据，即可适配特定说话人、语言或领域。
```bash
# LoRA 微调（参数高效，推荐）
python scripts/train_voxcpm_finetune.py \
--config_path conf/voxcpm_v2/voxcpm_finetune_lora.yaml
# 全参数微调
python scripts/train_voxcpm_finetune.py \
--config_path conf/voxcpm_v2/voxcpm_finetune_all.yaml
# WebUI 训练与推理
python lora_ft_webui.py # 然后打开 http://localhost:7860
```
> **完整指南 →** [微调文档](https://voxcpm.readthedocs.io/zh-cn/latest/finetuning/finetune.html)（数据准备、配置、训练、LoRA 热切换、常见问题）
---
## 📚 文档
完整文档：**[voxcpm.readthedocs.io](https://voxcpm.readthedocs.io/zh-cn/latest/)**
| 主题 | 链接 |
|---|---|
| 快速开始与安装 | [快速开始](https://voxcpm.readthedocs.io/zh-cn/latest/quickstart.html) |
| 使用指南与 Cookbook | [使用指南](https://voxcpm.readthedocs.io/zh-cn/latest/usage_guide.html) |
| VoxCPM 系列模型 | [模型列表](https://voxcpm.readthedocs.io/zh-cn/latest/models/version_history.html) |
| 微调（SFT & LoRA） | [微调指南](https://voxcpm.readthedocs.io/zh-cn/latest/finetuning/finetune.html) |
| 常见问题 | [FAQ](https://voxcpm.readthedocs.io/zh-cn/latest/faq.html) |
---
## 🌟 生态与社区
| 项目 | 说明 |
|---|---|
| [**Nano-vLLM**](https://github.com/a710128/nanovllm-voxcpm) | 高吞吐快速 GPU 推理引擎 |
| [**vLLM-Omni**](https://github.com/vllm-project/vllm-omni) | 官方 vLLM 全模态服务（原生支持 VoxCPM2）— PagedAttention、OpenAI 兼容 API |
| [**llama.cpp-omni**](https://github.com/tc-mb/llama.cpp-omni) | 全双工全模态推理引擎 — VoxCPM2 GGUF，支持 CPU / Metal / CUDA / Vulkan |
| [**VoxCPM.cpp**](https://github.com/bluryar/VoxCPM.cpp) | GGML/GGUF：CPU、CUDA、Vulkan 推理 |
| [**audio.cpp**](https://github.com/0xShug0/audio.cpp) | 基于 ggml 的统一 C++ 推理框架 — CPU/CUDA/Vulkan/Metal，CLI 与服务端，无需 Python |
| [**VoxCPM-ONNX**](https://github.com/bluryar/VoxCPM-ONNX) | ONNX 导出，支持 CPU 推理 |
| [**VoxCPMANE**](https://github.com/0seba/VoxCPMANE) | Apple Neural Engine 后端 |
| [**voxcpm_rs**](https://github.com/madushan1000/voxcpm_rs) | Rust 重新实现 |
| [**ComfyUI-VoxCPM**](https://github.com/wildminder/ComfyUI-VoxCPM) | ComfyUI 节点工作流 |
| [**ComfyUI_RH_VoxCPM**](https://github.com/HM-RunningHub/ComfyUI_RH_VoxCPM) | 面向 VoxCPM 2 的功能更完整的 ComfyUI 工作流，支持多说话人、LoRA 和自动 ASR |
| [**ComfyUI-VoxCPMTTS**](https://github.com/1038lab/ComfyUI-VoxCPMTTS) | ComfyUI TTS 扩展 |
| [**TTS WebUI**](https://github.com/rsxdalv/tts_webui_extension.vox_cpm) | 浏览器端 TTS 扩展 |
> 完整生态见[文档](https://voxcpm.readthedocs.io/zh-cn/latest/)。社区项目非 OpenBMB 官方维护。做了什么有趣的东西？[提 Issue 或 PR](https://github.com/OpenBMB/VoxCPM/issues) 把它加进来！
---
## ⚠️ 风险与局限性
- **滥用风险：** VoxCPM 的声音克隆能力可生成高度逼真的合成语音。**严禁**将 VoxCPM 用于冒充他人、欺诈或虚假信息传播。我们强烈建议对所有 AI 生成的内容进行明确标注。
- **可控生成稳定性：** 音色设计和可控声音克隆的结果可能因生成次数而异 — 建议尝试生成 1~3 次以获得理想的音色或风格。我们正在积极提升可控性的一致性。
- **语言覆盖：** VoxCPM2 官方支持 30 种语言。对于未列入的语言，欢迎直接测试或使用自有数据进行微调。我们计划在未来版本中扩展语言覆盖。
- **使用说明：** 本模型基于 Apache-2.0 协议发布。用于生产部署时，我们建议针对具体场景进行充分的测试和安全评估。
---
## 📖 引用
如果 VoxCPM 对您有帮助，请考虑引用我们的工作并为仓库加星 ⭐！
```bib
@article{zhou2026voxcpm2,
title = {VoxCPM2 Technical Report},
author = {Zhou, Yixuan and Zeng, Guoyang and Liu, Xin and Li, Xiang and Yu, Renjie and Gui, Jiancheng and Wu, Jiaheng and Wang, Ziyang and Shen, Xudong and Ye, Runchuan and Zhang, Zhisheng and Zhou, Jiuyang and Bai, Bingsong and Sun, Weiyue and Deng, Mengyuan and Shi, Qundong and Wu, Zhiyong and Liu, Zhiyuan},
journal = {arXiv preprint arXiv:2606.06928},
year = {2026},
}
@article{zhou2025voxcpm,
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}
}
```
## 📄 许可证
VoxCPM 模型权重和代码基于 [Apache-2.0](LICENSE) 协议开源。
## 🙏 致谢
- [DiTAR](https://arxiv.org/abs/2502.03930) 扩散自回归骨干架构
- [MiniCPM-4](https://github.com/OpenBMB/MiniCPM) 语言模型基座
- [CosyVoice](https://github.com/FunAudioLLM/CosyVoice) 基于 Flow Matching 的 LocDiT 实现
- [DAC](https://github.com/descriptinc/descript-audio-codec) Audio VAE 骨干
- 感谢所有社区用户试用 VoxCPM、反馈问题、分享想法和贡献——你们的支持让项目持续进步
## 机构
<p>
<a href="https://modelbest.cn/"><img src="assets/modelbest_logo.png" width="28px"> 面壁智能</a>
&nbsp;&nbsp;&nbsp;
<a href="https://github.com/thuhcsi"><img src="assets/thuhcsi_logo.png" width="28px"> 清华大学人机交互实验室</a>
</p>
## ⭐ Star 历史
[![Star History Chart](https://api.star-history.com/svg?repos=OpenBMB/VoxCPM&type=Date)](https://star-history.com/#OpenBMB/VoxCPM&Date)
app.py
+606
View File
@@ -0,0 +1,606 @@
import os
import re
import sys
import logging
import random
import numpy as np
import gradio as gr
from typing import Optional, Tuple
from funasr import AutoModel
from pathlib import Path
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import voxcpm
from voxcpm.model.utils import resolve_runtime_device
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",
"seed_label": "Seed",
"seed_info": "Seed used for reproducible generation. Updated with the actual successful seed after generation.",
"random_seed_label": "Random Seed",
"random_seed_info": "Generate a new seed before each inference run.",
"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, model_id: str = "openbmb/VoxCPM2", device: str = "auto") -> None:
self.device = resolve_runtime_device(device, "cuda")
logger.info(f"Running VoxCPM on device: {self.device}")
self.optimize = self.device.startswith("cuda")
self.asr_model_id = "iic/SenseVoiceSmall"
self.asr_device = "cuda:0" if self.device.startswith("cuda") else "cpu"
self.asr_model: Optional[AutoModel] = None
self.voxcpm_model: Optional[voxcpm.VoxCPM] = None
self._model_id = model_id
def get_or_load_voxcpm(self) -> voxcpm.VoxCPM:
if self.voxcpm_model is not None:
return self.voxcpm_model
logger.info(f"Loading model: {self._model_id}")
self.voxcpm_model = voxcpm.VoxCPM.from_pretrained(
self._model_id,
optimize=self.optimize,
device=self.device,
)
logger.info("Model loaded successfully.")
return self.voxcpm_model
def get_or_load_asr_model(self) -> AutoModel:
if self.asr_model is not None:
return self.asr_model
logger.info(f"Loading ASR model: {self.asr_model_id} on device: {self.asr_device}")
self.asr_model = AutoModel(
model=self.asr_model_id,
disable_update=True,
log_level="DEBUG",
device=self.asr_device,
)
logger.info("ASR model loaded successfully.")
return self.asr_model
def prompt_wav_recognition(self, prompt_wav: Optional[str]) -> str:
if prompt_wav is None:
return ""
res = self.get_or_load_asr_model().generate(
input=prompt_wav,
language="auto",
use_itn=True,
)
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,
seed: Optional[int] = None,
) -> 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,
seed=seed,
)
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,
control_instruction: str = "",
reference_wav_path_input: Optional[str] = None,
prompt_text: str = "",
cfg_value_input: float = 2.0,
do_normalize: bool = True,
denoise: bool = True,
inference_timesteps: int = 10,
seed: Optional[int] = None,
) -> Tuple[int, np.ndarray, Optional[int]]:
current_model = self.get_or_load_voxcpm()
text = (text_input or "").strip()
if len(text) == 0:
raise ValueError("Please input text to synthesize.")
control = (control_instruction or "").strip()
# Strip any parentheses (half-width/full-width) from control text to avoid
# breaking the "(control)text" prompt format expected by the model.
control = re.sub(r"[()（）]", "", control).strip()
final_text = f"({control}){text}" if control else text
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,
seed=seed,
)
wav = current_model.generate(**generate_kwargs)
last_successful_seed = getattr(current_model.tts_model, "last_successful_seed", seed)
return (current_model.tts_model.sample_rate, wav, last_successful_seed)
# ---------- UI ----------
def create_demo_interface(demo: VoxCPMDemo):
gr.set_static_paths(paths=[Path.cwd().absolute() / "assets"])
def _coerce_seed(seed_value) -> Optional[int]:
if seed_value is None or seed_value == "":
return None
return int(seed_value)
def _prepare_seed(use_random_seed: bool, seed_value):
if use_random_seed:
return random.randint(0, 2**32 - 1)
return _coerce_seed(seed_value)
def _on_random_seed_toggle(checked):
return gr.update(interactive=not checked)
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,
seed_value,
):
actual_prompt_text = prompt_text_value.strip() if use_prompt_text else ""
actual_control = "" if use_prompt_text else control_instruction
seed = _coerce_seed(seed_value)
sr, wav_np, last_successful_seed = 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),
seed=seed,
)
return (sr, wav_np), last_successful_seed
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),
)
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="")
with gr.Blocks() as interface:
gr.HTML(
'<div class="logo-container">'
'<img src="/gradio_api/file=assets/voxcpm_logo.png" alt="VoxCPM Logo">'
"</div>"
)
gr.Markdown(I18N("usage_instructions"))
with gr.Row():
with gr.Column():
reference_wav = gr.Audio(
sources=["upload", "microphone"],
type="filepath",
label=I18N("reference_audio_label"),
)
show_prompt_text = gr.Checkbox(
value=False,
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.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"),
)
DoNormalizeText = gr.Checkbox(
value=False,
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"),
)
with gr.Row():
seed_value = gr.Number(
value=random.randint(0, 2**32 - 1),
precision=0,
label=I18N("seed_label"),
info=I18N("seed_info"),
interactive=False,
)
random_seed = gr.Checkbox(
value=True,
label=I18N("random_seed_label"),
elem_classes=["switch-toggle"],
info=I18N("random_seed_info"),
)
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],
)
random_seed.change(
fn=_on_random_seed_toggle,
inputs=[random_seed],
outputs=[seed_value],
)
run_btn.click(
fn=_prepare_seed,
inputs=[random_seed, seed_value],
outputs=[seed_value],
show_progress=False,
).then(
fn=_generate,
inputs=[
text,
control_instruction,
reference_wav,
show_prompt_text,
prompt_text,
cfg_value,
DoNormalizeText,
DoDenoisePromptAudio,
dit_steps,
seed_value,
],
outputs=[audio_output, seed_value],
show_progress=True,
api_name="generate",
)
return interface
def run_demo(
server_name: str = "0.0.0.0",
server_port: int = 8808,
show_error: bool = True,
model_id: str = "openbmb/VoxCPM2",
device: str = "auto",
):
demo = VoxCPMDemo(model_id=model_id, device=device)
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,
i18n=I18N,
theme=_APP_THEME,
css=_CUSTOM_CSS,
)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument(
"--model-id",
type=str,
default="openbmb/VoxCPM2",
help="Local path or HuggingFace repo ID (default: openbmb/VoxCPM2)",
)
parser.add_argument("--port", type=int, default=8808, help="Server port")
parser.add_argument(
"--host",
type=str,
default="0.0.0.0",
help="Bind address. Use 127.0.0.1 to restrict access to the local machine; "
"the default 0.0.0.0 exposes the unauthenticated UI/API to the network (default: 0.0.0.0)",
)
parser.add_argument(
"--device",
type=str,
default="auto",
help="Runtime device: auto, cpu, mps, cuda, or cuda:N (default: auto)",
)
args = parser.parse_args()
run_demo(
model_id=args.model_id,
server_name=args.host,
server_port=args.port,
device=args.device,
)
app_old.py
+292
View File
@@ -0,0 +1,292 @@
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()
assets/discord-logo.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 9.3 KiB

assets/feishu-group.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 210 KiB

assets/feishu-logo.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 512 B

assets/modelbest_logo.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 54 KiB

assets/thuhcsi_logo.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 9.9 KiB

assets/voxcpm2_model.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.1 MiB

assets/voxcpm_logo.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 24 KiB

assets/voxcpm_model.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 142 KiB

conf/voxcpm_v1.5/voxcpm_finetune_all.yaml
+21
View File
@@ -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
+36
View File
@@ -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
+21
View File
@@ -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
+36
View File
@@ -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
+23
View File
@@ -0,0 +1,23 @@
pretrained_path: /path/to/VoxCPM2/
train_manifest: /path/to/train.jsonl
val_manifest: null
sample_rate: 16000 # AudioVAE encoder input rate; must match audio_vae_config.sample_rate
out_sample_rate: 48000 # AudioVAE decoder output rate; used for TensorBoard audio logging
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
max_grad_norm: 1.0 # gradient clipping max norm; 0 = disabled
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
+38
View File
@@ -0,0 +1,38 @@
pretrained_path: /path/to/VoxCPM2/
train_manifest: /path/to/train.jsonl
val_manifest: null
sample_rate: 16000 # AudioVAE encoder input rate; must match audio_vae_config.sample_rate
out_sample_rate: 48000 # AudioVAE decoder output rate; used for TensorBoard audio logging
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
max_grad_norm: 1.0 # gradient clipping max norm; 0 = disabled
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/example.wav
BIN
View File
Binary file not shown.
examples/reference_speaker.wav
BIN
View File
Binary file not shown.
examples/train_data_example.jsonl
+6
View File
@@ -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}
lora_ft_webui.py
+1327
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+100
View File
@@ -0,0 +1,100 @@
[build-system]
requires = ["setuptools>=64", "wheel", "setuptools_scm>=8"]
build-backend = "setuptools.build_meta"
[project]
name = "voxcpm"
dynamic = ["version"]
description = "VoxCPM: Tokenizer-Free TTS for Context-Aware Speech Generation and True-to-Life Voice Cloning"
readme = "README.md"
license = "Apache-2.0"
authors = [
{name = "OpenBMB", email = "openbmb@gmail.com"}
]
maintainers = [
{name = "OpenBMB", email = "openbmb@gmail.com"}
]
keywords = ["voxcpm", "text-to-speech", "tts", "speech-synthesis", "voice-cloning", "ai", "deep-learning", "pytorch"]
classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Developers",
"Operating System :: OS Independent",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
]
requires-python = ">=3.10"
dependencies = [
"torch>=2.5.0",
"torchaudio>=2.5.0",
"torchcodec",
"transformers>=4.36.2",
"einops",
"gradio>=6,<7",
"inflect",
"addict",
"wetext",
"modelscope>=1.22.0",
"datasets>=3,<4",
"huggingface-hub",
"pydantic",
"tqdm",
"simplejson",
"sortedcontainers",
"soundfile",
"librosa",
"matplotlib",
"funasr",
"spaces",
"argbind",
"safetensors"
]
[project.optional-dependencies]
dev = [
"pytest>=6.0",
"pytest-cov>=2.0",
"black>=21.0",
"flake8>=3.8",
"pre-commit>=2.0",
]
timestamps = [
"stable-ts>=2.19.1",
]
[project.scripts]
voxcpm = "voxcpm.cli:main"
[project.urls]
Homepage = "https://github.com/OpenBMB/VoxCPM"
Repository = "https://github.com/OpenBMB/VoxCPM.git"
Documentation = "https://github.com/OpenBMB/VoxCPM#readme"
"Bug Tracker" = "https://github.com/OpenBMB/VoxCPM/issues"
[tool.setuptools.packages.find]
where = ["src"]
include = ["voxcpm*"]
[tool.setuptools.package-dir]
"" = "src"
[tool.setuptools_scm]
version_scheme = "post-release"
[tool.black]
line-length = 120
target-version = ['py310']
include = '\.pyi?$'
extend-exclude = '''
/(
# directories
\.eggs
| \.git
| \.hg
| \.tox
| \.venv
| build
| dist
)/
'''
scripts/test_pick_runtime_dtype.py
+109
View File
@@ -0,0 +1,109 @@
"""Unit checks for pick_runtime_dtype / get_dtype consistency.
Loads src/voxcpm/model/utils.py directly to avoid the heavy voxcpm package
init. Run with: `python scripts/test_pick_runtime_dtype.py`.
"""
import importlib.util
import os
import pathlib
import sys
REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
UTILS = str(REPO_ROOT / "src" / "voxcpm" / "model" / "utils.py")
spec = importlib.util.spec_from_file_location("voxcpm_utils", UTILS)
utils = importlib.util.module_from_spec(spec)
spec.loader.exec_module(utils)
_LOW_PRECISION_DTYPES = utils._LOW_PRECISION_DTYPES
_VALID_DTYPE_OVERRIDES = utils._VALID_DTYPE_OVERRIDES
get_dtype = utils.get_dtype
pick_runtime_dtype = utils.pick_runtime_dtype
def expect(actual, expected, label):
ok = actual == expected
mark = "OK " if ok else "FAIL"
print(f"[{mark}] {label}: got={actual!r} expected={expected!r}")
return ok
def expect_raises(fn, exc_type, label):
try:
fn()
except exc_type as e:
print(f"[OK ] {label}: raised {exc_type.__name__}: {e}")
return True
except Exception as e:
print(f"[FAIL] {label}: raised {type(e).__name__} not {exc_type.__name__}: {e}")
return False
print(f"[FAIL] {label}: no exception raised")
return False
results = []
print("=== override set sanity ===")
results.append(expect("half" not in _VALID_DTYPE_OVERRIDES, True, "half removed from _VALID_DTYPE_OVERRIDES"))
results.append(expect("half" not in _LOW_PRECISION_DTYPES, True, "half removed from _LOW_PRECISION_DTYPES"))
print("\n=== every accepted override parses through get_dtype ===")
for dt in sorted(_VALID_DTYPE_OVERRIDES):
try:
torch_dtype = get_dtype(dt)
print(f"[OK ] get_dtype({dt!r}) -> {torch_dtype}")
results.append(True)
except Exception as e:
print(f"[FAIL] get_dtype({dt!r}) raised: {e}")
results.append(False)
print("\n=== pick_runtime_dtype: non-mps is a no-op ===")
results.append(expect(pick_runtime_dtype("cuda", "bfloat16"), "bfloat16", "cuda/bf16 untouched"))
results.append(expect(pick_runtime_dtype("cpu", "float16"), "float16", "cpu/fp16 untouched"))
results.append(expect(pick_runtime_dtype("cuda", "float32"), "float32", "cuda/fp32 untouched"))
print("\n=== pick_runtime_dtype: mps forces fp32 for low-precision ===")
os.environ.pop("VOXCPM_MPS_DTYPE", None)
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "float32", "mps/bf16 -> fp32"))
results.append(expect(pick_runtime_dtype("mps", "bf16"), "float32", "mps/bf16-alias -> fp32"))
results.append(expect(pick_runtime_dtype("mps", "float16"), "float32", "mps/fp16 -> fp32"))
results.append(expect(pick_runtime_dtype("mps", "fp16"), "float32", "mps/fp16-alias -> fp32"))
results.append(expect(pick_runtime_dtype("mps", "float32"), "float32", "mps/fp32 stays"))
results.append(expect(pick_runtime_dtype("mps", "fp32"), "fp32", "mps/fp32-alias stays"))
print("\n=== pick_runtime_dtype: VOXCPM_MPS_DTYPE override ===")
os.environ["VOXCPM_MPS_DTYPE"] = "bfloat16"
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "bfloat16", "override bf16 honored"))
os.environ["VOXCPM_MPS_DTYPE"] = "FP16"
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "fp16", "override is case-insensitive"))
os.environ["VOXCPM_MPS_DTYPE"] = " float32 "
results.append(expect(pick_runtime_dtype("mps", "bfloat16"), "float32", "override is whitespace-trimmed"))
print("\n=== pick_runtime_dtype: 'half' is no longer a valid override ===")
os.environ["VOXCPM_MPS_DTYPE"] = "half"
results.append(
expect_raises(
lambda: pick_runtime_dtype("mps", "bfloat16"),
ValueError,
"override=half now rejected (was the bug)",
)
)
os.environ["VOXCPM_MPS_DTYPE"] = "garbage"
results.append(
expect_raises(
lambda: pick_runtime_dtype("mps", "bfloat16"),
ValueError,
"override=garbage still rejected",
)
)
os.environ.pop("VOXCPM_MPS_DTYPE", None)
print("\n=== summary ===")
passed = sum(results)
total = len(results)
print(f"{passed}/{total} passed")
sys.exit(0 if passed == total else 1)
scripts/test_voxcpm_ft_infer.py
+146
View File
@@ -0,0 +1,146 @@
#!/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" \
--seed 42 \
--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",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for generation (default: None)",
)
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)
if args.seed is not None:
print(f"[FT Inference] Using seed: {args.seed}", 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,
seed=args.seed,
)
# 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
+276
View File
@@ -0,0 +1,276 @@
#!/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" \
--seed 42 \
--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",
)
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for generation (default: None)",
)
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
)
if args.seed is not None:
print(f" Seed: {args.seed}", 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,
seed=args.seed,
)
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,
seed=args.seed,
)
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,
seed=args.seed,
)
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,
seed=args.seed,
)
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,
seed=args.seed,
)
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
+841
View File
@@ -0,0 +1,841 @@
#!/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 as LoRAConfigV1
from voxcpm.model.voxcpm2 import LoRAConfig as LoRAConfigV2
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,
out_sample_rate: int = 0, # AudioVAE decoder output rate; used for TensorBoard audio logging
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 = "",
max_grad_norm: float = 0.0, # gradient clipping; 0 = disabled (backward compat)
# 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
LoRAConfig = LoRAConfigV2 if _arch == "voxcpm2" else LoRAConfigV1
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
expected_sr = base_model.audio_vae.sample_rate
assert sample_rate == expected_sr, (
f"sample_rate mismatch: config says {sample_rate}, but the AudioVAE encoder expects {expected_sr}. "
f"Please set sample_rate: {expected_sr} in your training config. "
)
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 and output sample rate for audio generation.
# Prefer model's actual output rate; fall back to YAML out_sample_rate or encode rate.
audio_vae_for_gen = base_model.audio_vae
out_sr = base_model.sample_rate # decoder output rate (e.g. 48000 for V2)
if out_sr == 0 and out_sample_rate > 0:
out_sr = out_sample_rate
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)
effective_max_norm = max_grad_norm if max_grad_norm > 0 else 1e9
grad_norm = torch.nn.utils.clip_grad_norm_(unwrapped_model.parameters(), max_norm=effective_max_norm)
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,
out_sample_rate=out_sr,
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,
out_sample_rate=0,
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,
out_sample_rate=out_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,
out_sample_rate=0,
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)
# Determine the correct output sample rate for generated audio.
# out_sample_rate is the decoder output rate (e.g. 48kHz for V2);
# sample_rate is the encoder input rate (e.g. 16kHz for V2).
gen_sr = out_sample_rate if out_sample_rate > 0 else sample_rate
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, seed=42
)
# 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=gen_sr)
log(f"[Audio] Generated audio for sample {i}: duration={len(gen_audio_np)/gen_sr:.2f}s")
# Log reference audio (at encoder input rate, which is what val_ds provides)
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, gen_sr)
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, gen_sr, 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", weights_only=True)
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", weights_only=True)
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", weights_only=True))
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", weights_only=True))
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/__init__.py
+5
View File
@@ -0,0 +1,5 @@
from .core import VoxCPM
__all__ = [
"VoxCPM",
]
src/voxcpm/cli.py
+692
View File
@@ -0,0 +1,692 @@
#!/usr/bin/env python3
"""
VoxCPM Command Line Interface
VoxCPM2-first CLI for voice design, cloning, and batch processing.
"""
import argparse
import json
import os
import sys
from pathlib import Path
from voxcpm.timestamps import align_audio_file
DEFAULT_HF_MODEL_ID = "openbmb/VoxCPM2"
# -----------------------------
# Validators
# -----------------------------
def validate_file_exists(file_path: str, file_type: str = "file") -> Path:
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:
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):
from voxcpm.core import VoxCPM
print("Loading VoxCPM model...", file=sys.stderr)
zipenhancer_path = getattr(args, "zipenhancer_path", None) or os.environ.get("ZIPENHANCER_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 args.no_denoiser,
optimize=not args.no_optimize,
device=args.device,
lora_config=lora_config,
lora_weights_path=lora_weights_path,
)
print("Model loaded (local).", file=sys.stderr)
return model
except Exception as e:
print(f"Failed to load model (local): {e}", file=sys.stderr)
sys.exit(1)
# Load from Hugging Face Hub
try:
model = VoxCPM.from_pretrained(
hf_model_id=args.hf_model_id,
load_denoiser=not args.no_denoiser,
zipenhancer_model_id=zipenhancer_path,
cache_dir=args.cache_dir,
local_files_only=args.local_files_only,
optimize=not args.no_optimize,
device=args.device,
lora_config=lora_config,
lora_weights_path=lora_weights_path,
)
print("Model loaded (from_pretrained).", file=sys.stderr)
return model
except Exception as e:
print(f"Failed to load model (from_pretrained): {e}", file=sys.stderr)
sys.exit(1)
# -----------------------------
# 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)
audio_array = model.generate(
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),
seed=args.seed,
)
import soundfile as sf
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)
maybe_write_timestamps(
args,
text=text,
audio_path=output_path,
sample_rate=model.tts_model.sample_rate,
)
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)
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_validate(args, parser):
from voxcpm.training.validate import (
print_validation_report,
validate_manifest,
)
manifest = str(require_file_exists(args.manifest, parser, "manifest file"))
result = validate_manifest(
manifest_path=manifest,
sample_rate=args.sample_rate,
max_samples=args.max_samples,
verbose=args.verbose,
)
print_validation_report(result, manifest)
if not result.is_valid:
sys.exit(1)
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)
with open(input_file, "r", encoding="utf-8") as f:
texts = [line.strip() for line in f if line.strip()]
if not texts:
sys.exit("Error: Input file is empty")
prompt_text = validate_batch_args(args, parser)
model = load_model(args)
import soundfile as sf
prompt_audio_path = None
if args.prompt_audio:
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):
try:
final_text = build_final_text(text, args.control)
audio_array = model.generate(
text=final_text,
prompt_wav_path=prompt_audio_path,
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 or reference_audio_path is not None),
seed=args.seed,
)
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)
maybe_write_timestamps(
args,
text=final_text,
audio_path=output_file,
sample_rate=model.tts_model.sample_rate,
)
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)
def default_timestamp_path(audio_path: Path) -> Path:
return audio_path.with_suffix(".timestamps.json")
def maybe_write_timestamps(args, *, text: str, audio_path: Path, sample_rate: int) -> None:
if not getattr(args, "timestamps", False):
return
timestamp_output = getattr(args, "timestamp_output", None)
output_path = Path(timestamp_output) if timestamp_output else default_timestamp_path(audio_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
try:
result = align_audio_file(
audio_path=str(audio_path),
text=text,
sample_rate=sample_rate,
backend=args.timestamp_backend,
level=args.timestamp_level,
model_name=args.timestamp_model,
device=args.timestamp_device,
language=args.timestamp_language,
)
except Exception as exc:
if getattr(args, "timestamp_strict", False):
raise SystemExit(f"Timestamp alignment failed: {exc}") from exc
print(f"Warning: Timestamp alignment failed: {exc}", file=sys.stderr)
return
with open(output_path, "w", encoding="utf-8") as f:
json.dump(result, f, ensure_ascii=False, indent=2)
f.write("\n")
print(f"Saved timestamps to: {output_path}", 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")
parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for generation (default: None)",
)
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(
"--device",
type=str,
default="auto",
help="Runtime device: auto, cpu, mps, cuda, or cuda:N (default: auto)",
)
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 _add_timestamp_args(parser, *, include_output: bool = True):
parser.add_argument(
"--timestamps",
action="store_true",
help="Run post-generation timestamp alignment and write a JSON sidecar file",
)
if include_output:
parser.add_argument(
"--timestamp-output",
type=str,
help="Output timestamp JSON path (default: output audio path with .timestamps.json suffix)",
)
parser.add_argument(
"--timestamp-level",
choices=["segment", "word", "char"],
default="word",
help="Timestamp granularity (default: word; char is best-effort)",
)
parser.add_argument(
"--timestamp-backend",
choices=["stable-ts"],
default="stable-ts",
help="Timestamp alignment backend (default: stable-ts)",
)
parser.add_argument(
"--timestamp-model",
default="base",
help="stable-ts Whisper model name (default: base)",
)
parser.add_argument(
"--timestamp-language",
default=None,
help="Language hint for timestamp alignment, e.g. zh or en",
)
parser.add_argument(
"--timestamp-device",
default=None,
help="Device for timestamp alignment, e.g. cuda or cpu",
)
parser.add_argument(
"--timestamp-strict",
action="store_true",
help="Fail the command if timestamp alignment fails",
)
def _build_parser():
parser = argparse.ArgumentParser(
description="VoxCPM CLI - VoxCPM2-first voice design, cloning, and batch processing",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
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
""",
)
subparsers = parser.add_subparsers(dest="command")
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)
_add_timestamp_args(design_parser)
design_parser.add_argument("--output", "-o", required=True, help="Output audio file path")
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)
_add_timestamp_args(clone_parser)
clone_parser.add_argument("--output", "-o", required=True, help="Output audio file path")
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")
batch_parser.add_argument(
"--seed",
type=int,
default=None,
help="Random seed for generation (default: None)",
)
_add_model_args(batch_parser)
_add_lora_args(batch_parser)
_add_timestamp_args(batch_parser, include_output=False)
# Validate subcommand
validate_parser = subparsers.add_parser(
"validate",
help="Validate a training data manifest (JSONL) before fine-tuning",
)
validate_parser.add_argument("--manifest", "-m", required=True, help="Path to JSONL training manifest")
validate_parser.add_argument(
"--sample-rate",
type=int,
default=16_000,
help="Expected audio sample rate in Hz (default: 16000)",
)
validate_parser.add_argument(
"--max-samples",
type=int,
default=0,
help="Maximum number of samples to validate (0 = all, default: 0)",
)
validate_parser.add_argument("--verbose", "-v", action="store_true", help="Print per-sample progress")
# 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)
_add_timestamp_args(parser)
return parser
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.")
if args.input:
if not args.output_dir:
parser.error("Batch mode requires --output-dir")
return cmd_batch(args, parser)
if not args.text or not args.output:
parser.error("Single-sample legacy mode requires --text and --output")
if args.prompt_audio or args.prompt_text or args.prompt_file or args.reference_audio:
return cmd_clone(args, parser)
return cmd_design(args, parser)
# -----------------------------
# Entrypoint
# -----------------------------
def main():
parser = _build_parser()
args = parser.parse_args()
if args.command == "validate":
return cmd_validate(args, parser)
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__":
main()
src/voxcpm/core.py
+366
View File
@@ -0,0 +1,366 @@
import os
import sys
import re
import json
import tempfile
import numpy as np
from typing import Generator, Optional
from huggingface_hub import snapshot_download
from .model.voxcpm import VoxCPMModel, LoRAConfig
from .model.voxcpm2 import VoxCPM2Model
from .model.utils import next_and_close
class VoxCPM:
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,
device: str | None = None,
lora_config: Optional[LoRAConfig] = None,
lora_weights_path: Optional[str] = None,
):
"""Initialize VoxCPM TTS pipeline.
Args:
voxcpm_model_path: Local filesystem path to the VoxCPM model assets
(weights, configs, etc.). Typically the directory returned by
a prior download step.
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.
device: Runtime device. If set to ``None`` or ``"auto"``, VoxCPM
will choose automatically (preferring CUDA, then MPS, then CPU).
If set explicitly, that device is used or a clear error is raised.
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}",
file=sys.stderr,
)
# If lora_weights_path is provided but no lora_config, load the saved
# lora_config.json (so r/alpha match the checkpoint); else use a default.
if lora_weights_path is not None and lora_config is None:
cfg_path = os.path.join(lora_weights_path, "lora_config.json")
if os.path.isdir(lora_weights_path) and os.path.isfile(cfg_path):
with open(cfg_path, "r", encoding="utf-8") as f:
lora_config = LoRAConfig(**json.load(f)["lora_config"])
print(f"Loaded LoRAConfig from: {cfg_path}", file=sys.stderr)
else:
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,
device=device,
lora_config=lora_config,
)
print("Loaded VoxCPM2Model", file=sys.stderr)
elif arch == "voxcpm":
self.tts_model = VoxCPMModel.from_local(
voxcpm_model_path,
optimize=optimize,
device=device,
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:
from .zipenhancer import ZipEnhancer
self.denoiser = ZipEnhancer(zipenhancer_model_path)
else:
self.denoiser = None
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/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,
device: str | None = None,
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") 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.
device: Runtime device. Use ``None``/``"auto"`` for automatic
fallback, or an explicit value such as ``"cpu"``, ``"mps"``,
``"cuda"``, or ``"cuda:0"``.
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
the downloaded snapshot directory.
Raises:
ValueError: If neither a valid ``hf_model_id`` nor a resolvable
``hf_model_id`` is provided.
"""
repo_id = hf_model_id
if not repo_id:
raise ValueError("You must provide hf_model_id")
# 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,
device=device,
lora_config=lora_config,
lora_weights_path=lora_weights_path,
**kwargs,
)
def generate(self, *args, **kwargs) -> np.ndarray:
return next_and_close(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,
seed: Optional[int] = None,
) -> Generator[np.ndarray, None, None]:
"""Synthesize speech for the given text and return a single waveform.
Args:
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.
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/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.
seed: Optional random seed for reproducibility.
Returns:
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.
"""
if not isinstance(text, str) or not text.strip():
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:
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,
)
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,
seed=seed,
)
if streaming:
try:
for wav, _, _ in generate_result:
yield wav.squeeze(0).cpu().numpy()
finally:
generate_result.close()
else:
wav, _, _ = next_and_close(generate_result)
yield wav.squeeze(0).cpu().numpy()
finally:
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
+4
View File
@@ -0,0 +1,4 @@
from .voxcpm import VoxCPMModel
from .voxcpm2 import VoxCPM2Model
__all__ = ["VoxCPMModel", "VoxCPM2Model"]
src/voxcpm/model/utils.py
+242
View File
@@ -0,0 +1,242 @@
import os
from typing import List, Optional
import torch
from transformers import PreTrainedTokenizer
_LOW_PRECISION_DTYPES = {"bfloat16", "bf16", "float16", "fp16"}
_VALID_DTYPE_OVERRIDES = {
"bfloat16",
"bf16",
"float16",
"fp16",
"float32",
"fp32",
}
# Ref: https://github.com/OpenBMB/VoxCPM/issues/256#issuecomment-4235252732
# Explicitly close partially-consumed generators so inference_mode cleanup
# does not get deferred to Python's GC/finalizer path.
def next_and_close(gen):
try:
return next(gen)
finally:
gen.close()
def materialize_generation_seed(seed: Optional[int]) -> int:
"""Return a concrete seed for a generation request."""
if seed is not None:
return int(seed)
return int(torch.seed() & 0xFFFFFFFF)
def apply_generation_seed(seed: int) -> None:
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
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")
>>> wrapped_tokenizer = mask_multichar_chinese_tokens(tokenizer)
>>> tokens = wrapped_tokenizer("你好世界")
"""
# 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)
}
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
"""
self.tokenizer = base_tokenizer
self.multichar_tokens = multichar_tokens
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("你好世界")
>>> # Returns ["你", "好", "世", "界"] instead of ["你好", "世界"]
"""
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("▁", "")
if clean_token in self.multichar_tokens:
# Split multi-character token into single characters
chars = list(clean_token)
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
"""
try:
tokens = self.tokenize(text, **kwargs)
result = self.tokenizer.convert_tokens_to_ids(tokens)
return result
except Exception as e:
raise ValueError(f"Tokenization failed: {str(e)}") from e
return CharTokenizerWrapper(tokenizer)
def get_dtype(dtype: str):
if dtype == "bfloat16":
return torch.bfloat16
elif dtype == "bf16":
return torch.bfloat16
elif dtype == "float16":
return torch.float16
elif dtype == "fp16":
return torch.float16
elif dtype == "float32":
return torch.float32
elif dtype == "fp32":
return torch.float32
else:
raise ValueError(f"Unsupported dtype: {dtype}")
def _has_mps() -> bool:
return hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
def pick_runtime_dtype(device: str, configured_dtype: str) -> str:
"""Pick a safe runtime dtype for the resolved device.
On Apple Silicon (MPS), bfloat16/float16 produce enough numerical drift
in the diffusion AR loop that the output is glitched and the model's
badcase detector triggers infinite retries. float32 is the only stable
option today. CUDA and CPU keep whatever the checkpoint was trained with.
Users can override with ``VOXCPM_MPS_DTYPE`` (e.g. ``bfloat16``) when
they want to test future MPS improvements.
"""
if device != "mps":
return configured_dtype
override = os.environ.get("VOXCPM_MPS_DTYPE", "").strip().lower()
if override:
if override not in _VALID_DTYPE_OVERRIDES:
raise ValueError(f"VOXCPM_MPS_DTYPE='{override}' is not one of " f"{sorted(_VALID_DTYPE_OVERRIDES)}")
return override
if (configured_dtype or "").lower() in _LOW_PRECISION_DTYPES:
return "float32"
return configured_dtype
def auto_select_device(preferred_device: Optional[str] = "cuda") -> str:
"""
Choose a runtime device automatically.
Preference order:
- if the preferred device is available, use it
- otherwise fall back to CUDA -> MPS -> CPU
"""
preferred = (preferred_device or "cuda").strip().lower()
if preferred.startswith("cuda") and torch.cuda.is_available():
return preferred
if preferred == "mps" and _has_mps():
return "mps"
if preferred == "cpu":
return "cpu"
if torch.cuda.is_available():
return "cuda"
if _has_mps():
return "mps"
return "cpu"
def resolve_runtime_device(device: Optional[str], configured_device: str = "cuda") -> str:
"""
Resolve the actual runtime device.
Semantics:
- ``device`` is ``None`` or ``"auto"``: use automatic fallback selection
- otherwise: treat it as an explicit user choice and validate availability
"""
explicit = None if device is None else device.strip().lower()
if explicit is None or explicit == "auto":
return auto_select_device(configured_device)
if explicit.startswith("cuda"):
if not torch.cuda.is_available():
raise ValueError(
f"Requested device '{device}', but CUDA is not available. " "Use device='auto' for automatic fallback."
)
return explicit
if explicit == "mps":
if not _has_mps():
raise ValueError(
"Requested device 'mps', but MPS is not available. " "Use device='auto' for automatic fallback."
)
return "mps"
if explicit == "cpu":
return "cpu"
raise ValueError(
f"Unsupported device '{device}'. Supported values are 'auto', 'cpu', 'mps', "
"'cuda', or indexed CUDA devices like 'cuda:0'."
)
src/voxcpm/model/voxcpm.py
+1026
View File
File diff suppressed because it is too large Load Diff
src/voxcpm/model/voxcpm2.py
+1279
View File
File diff suppressed because it is too large Load Diff
src/voxcpm/modules/__init__.py
View File
src/voxcpm/modules/audiovae/__init__.py
+2
View File
@@ -0,0 +1,2 @@
from .audio_vae import AudioVAE, AudioVAEConfig
from .audio_vae_v2 import AudioVAE as AudioVAEV2, AudioVAEConfig as AudioVAEConfigV2
src/voxcpm/modules/audiovae/audio_vae.py
+377
View File
@@ -0,0 +1,377 @@
import math
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):
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, **kwargs):
super().__init__(*args, **kwargs)
self.__padding = padding
def forward(self, x):
x_pad = F.pad(x, (self.__padding * 2, 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),
),
)
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)
def forward(self, x):
return self.block(x)
class TransposeLastTwoDim(torch.nn.Module):
def forward(self, x):
return torch.transpose(x, -1, -2)
class CausalDecoder(nn.Module):
def __init__(
self,
input_channel,
channels,
rates,
depthwise: bool = False,
d_out: int = 1,
use_noise_block: 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(),
]
self.model = nn.Sequential(*layers)
def forward(self, x):
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:
"""
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
use_noise_block = config.use_noise_block
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,
)
self.sample_rate = sample_rate
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):
"""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.
"""
return self.decoder(z)
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/audiovae/audio_vae_v2.py
+579
View File
@@ -0,0 +1,579 @@
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)
self.decode_chunk_size = math.prod(decoder_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 streaming_decode(self):
"""Return a ``StreamingVAEDecoder`` context manager for stateful
chunk-by-chunk decoding. Each call to ``decode_chunk`` processes only
the new latent patch and carries causal-conv state internally, avoiding
the redundant overlap decode used previously.
Usage::
with vae.streaming_decode() as dec:
for patch in patches:
audio_chunk = dec.decode_chunk(patch)
"""
return StreamingVAEDecoder(self)
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"]
class StreamingVAEDecoder:
"""Stateful streaming wrapper for :class:`AudioVAE`.
Carries causal-convolution padding buffers between calls so that each
``decode_chunk`` processes only the new latent patch — no overlap needed.
"""
def __init__(self, vae: AudioVAE):
self._vae = vae
self._states: dict = {}
self._originals: list = []
# -- context manager --------------------------------------------------
def __enter__(self):
self._states.clear()
self._install()
return self
def __exit__(self, *exc):
self._restore()
self._states.clear()
# -- public API --------------------------------------------------------
def decode_chunk(self, z_chunk: torch.Tensor) -> torch.Tensor:
"""Decode a single latent chunk and return the audio waveform."""
return self._vae.decode(z_chunk)
# -- internals ---------------------------------------------------------
def _install(self):
for name, mod in self._vae.decoder.named_modules():
if isinstance(mod, CausalConv1d):
pad = mod._CausalConv1d__padding * 2 - mod._CausalConv1d__output_padding
if pad > 0:
self._patch_causal_conv(mod, pad)
elif isinstance(mod, CausalTransposeConv1d):
trim = mod._CausalTransposeConv1d__padding * 2 - mod._CausalTransposeConv1d__output_padding
ctx = (mod.kernel_size[0] - 1) // mod.stride[0]
if ctx > 0:
self._patch_transpose_conv(mod, ctx, trim)
def _patch_causal_conv(self, mod, pad_size):
states = self._states
key = id(mod)
orig = mod.forward
def fwd(x, _k=key, _p=pad_size, _m=mod):
x_pad = torch.cat([states[_k], x], dim=-1) if _k in states else F.pad(x, (_p, 0))
if x.shape[-1] >= _p:
states[_k] = x[:, :, -_p:].detach()
else:
prev = states.get(_k, torch.zeros(x.shape[0], x.shape[1], _p, device=x.device, dtype=x.dtype))
states[_k] = torch.cat([prev, x], dim=-1)[:, :, -_p:].detach()
return nn.Conv1d.forward(_m, x_pad)
mod.forward = fwd
self._originals.append((mod, orig))
def _patch_transpose_conv(self, mod, ctx, trim):
states = self._states
key = id(mod)
orig = mod.forward
def fwd(x, _k=key, _c=ctx, _t=trim, _m=mod):
x_full = torch.cat([states[_k], x], dim=-1) if _k in states else F.pad(x, (_c, 0))
states[_k] = x[:, :, -_c:].detach()
out = nn.ConvTranspose1d.forward(_m, x_full)
left = _c * _m.stride[0]
return out[..., left:-_t] if _t > 0 else out[..., left:]
mod.forward = fwd
self._originals.append((mod, orig))
def _restore(self):
for mod, orig in self._originals:
mod.forward = orig
self._originals.clear()
src/voxcpm/modules/layers/__init__.py
+1
View File
@@ -0,0 +1 @@
from .scalar_quantization_layer import ScalarQuantizationLayer
src/voxcpm/modules/layers/lora.py
+130
View File
@@ -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
+26
View File
@@ -0,0 +1,26 @@
import torch
import torch.nn as nn
class ScalarQuantizationLayer(nn.Module):
def __init__(self, in_dim, out_dim, latent_dim: int = 64, scale: int = 9):
super().__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.latent_dim = latent_dim
self.scale = scale
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)
if self.training:
quantized = torch.round(hidden * self.scale) / self.scale
hidden = hidden + (quantized - hidden).detach()
else:
hidden = torch.round(hidden * self.scale) / self.scale
return self.out_proj(hidden)
src/voxcpm/modules/locdit/__init__.py
+3
View File
@@ -0,0 +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.py
+114
View File
@@ -0,0 +1,114 @@
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
x = torch.cat([(mu + t).unsqueeze(1), cond, x], dim=1)
hidden, _ = self.decoder(x, is_causal=False)
hidden = hidden[:, prefix + 1 :, :]
hidden = self.out_proj(hidden)
return hidden.transpose(1, 2).contiguous()
src/voxcpm/modules/locdit/local_dit_v2.py
+116
View File
@@ -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
+232
View File
@@ -0,0 +1,232 @@
from typing import Tuple
import torch
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-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: int,
cfm_params: CfmConfig,
estimator: VoxCPMLocDiT,
mean_mode: bool = False,
):
super().__init__()
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
self.estimator = estimator
# ------------------------------------------------------------------ #
# Inference
# ------------------------------------------------------------------ #
@torch.inference_mode()
def forward(
self,
mu: torch.Tensor,
n_timesteps: int,
patch_size: int,
cond: torch.Tensor,
temperature: float = 1.0,
cfg_value: float = 1.0,
sway_sampling_coef: float = 1.0,
use_cfg_zero_star: bool = True,
):
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)
t_span = t_span + sway_sampling_coef * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span)
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: 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
st_star = dot_product / squared_norm
return st_star
def solve_euler(
self,
x: torch.Tensor,
t_span: torch.Tensor,
mu: torch.Tensor,
cond: torch.Tensor,
cfg_value: float = 1.0,
use_cfg_zero_star: bool = True,
):
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 = torch.zeros_like(x)
else:
# Classifier-Free Guidance inference introduced in VoiceBox
b = x.size(0)
x_in = torch.zeros([2 * b, self.in_channels, x.size(2)], device=x.device, dtype=x.dtype)
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, 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)
dt_in[:b], dt_in[b:] = dt.unsqueeze(0), dt.unsqueeze(0)
# not used now
if not self.mean_mode:
dt_in = torch.zeros_like(dt_in)
cond_in[:b], cond_in[b:] = cond, cond
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)
st_star = self.optimized_scale(positive_flat, negative_flat)
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
t = t - dt
sol.append(x)
if step < len(t_span) - 1:
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.clamp(torch.sum(tgt_mask), min=1.0)
else:
loss = losses.mean()
return loss
src/voxcpm/modules/locenc/__init__.py
+1
View File
@@ -0,0 +1 @@
from .local_encoder import VoxCPMLocEnc
src/voxcpm/modules/locenc/local_encoder.py
+30
View File
@@ -0,0 +1,30 @@
import torch
import torch.nn as nn
from ..minicpm4 import MiniCPMModel, MiniCPM4Config
from einops import rearrange
class VoxCPMLocEnc(nn.Module):
def __init__(self, config: MiniCPM4Config, input_dim: int = 64):
super().__init__()
self.config = config
self.special_token = nn.Parameter(torch.randn(1, 1, 1, config.hidden_size))
self.in_proj = nn.Linear(input_dim, config.hidden_size, bias=True)
assert config.vocab_size == 0, "vocab_size must be 0 for local encoder"
self.encoder = MiniCPMModel(config)
def forward(self, x):
"""
x: [B, T, P, D]
"""
B, T, P, D = x.shape
x = self.in_proj(x)
special_tokens = self.special_token.expand(B, T, 1, -1)
x = torch.cat([special_tokens, x], dim=2)
x = rearrange(x, "b t p c -> (b t) p c")
outputs, _ = self.encoder(x, is_causal=False)
cls_output = outputs[:, 0, :]
return rearrange(cls_output, "(b t) c -> b t c", b=B)
src/voxcpm/modules/minicpm4/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .config import MiniCPM4Config
from .model import MiniCPMModel
from .cache import StaticKVCache
src/voxcpm/modules/minicpm4/cache.py
+47
View File
@@ -0,0 +1,47 @@
from typing import List, Tuple
import torch
class StaticKVCache:
def __init__(
self,
num_layers: int,
num_kv_heads: int,
dim_kv_head: int,
batch_size: int,
device: torch.device,
dtype: torch.dtype,
max_length: int = 8192,
):
self.max_length = max_length
self.num_layers = num_layers
self.kv_cache = torch.zeros(
2,
num_layers,
batch_size,
num_kv_heads,
max_length,
dim_kv_head,
device=device,
dtype=dtype,
)
self.current_length = 0
def get_layer_cache(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
return self.kv_cache[0, layer_idx], self.kv_cache[1, layer_idx]
def step(self) -> int:
if self.current_length >= self.max_length:
raise ValueError("KV cache is full")
ret = self.current_length
self.current_length += 1
return ret
def fill_caches(self, kv_caches: List[Tuple[torch.Tensor, torch.Tensor]]):
self.current_length = kv_caches[0][0].size(2)
self.kv_cache.zero_()
for i in range(self.num_layers):
self.kv_cache[0, i, :, :, : self.current_length, :] = kv_caches[i][0]
self.kv_cache[1, i, :, :, : self.current_length, :] = kv_caches[i][1]
src/voxcpm/modules/minicpm4/config.py
+30
View File
@@ -0,0 +1,30 @@
from pydantic import BaseModel
from typing import List
class RopeScalingConfig(BaseModel):
type: str
long_factor: List[float]
short_factor: List[float]
original_max_position_embeddings: int
class MiniCPM4Config(BaseModel):
bos_token_id: int
eos_token_id: int
hidden_size: int
intermediate_size: int
max_position_embeddings: int
num_attention_heads: int
num_hidden_layers: int
num_key_value_heads: int
rms_norm_eps: float
rope_scaling: RopeScalingConfig
vocab_size: int
use_mup: bool = True
scale_emb: float
dim_model_base: int
scale_depth: float
rope_theta: float
kv_channels: int = None
no_rope: bool = False
src/voxcpm/modules/minicpm4/model.py
+431
View File
@@ -0,0 +1,431 @@
from .config import MiniCPM4Config
import torch
import torch.nn as nn
from typing import List, Tuple
import math
from .cache import StaticKVCache
def rms_layernorm(hidden: torch.Tensor, weight: torch.Tensor, eps: float):
old_dtype = hidden.dtype
variance = hidden.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
hidden = (hidden * torch.rsqrt(variance + eps)).to(old_dtype)
return hidden * weight
class MiniCPMRMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
"""
MiniCPMRMSNorm is equivalent to T5LayerNorm
"""
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
return rms_layernorm(hidden_states, self.weight, self.variance_epsilon)
def rotate_half(x):
"""Rotates half the hidden dims of the input."""
x1, x2 = x.chunk(2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
"""
Args:
q: Tensor(batch_size, num_heads, seq_len, head_dim)
k: Tensor(batch_size, num_key_value_heads, seq_len, head_dim)
cos: Tensor(seq_len, head_dim)
sin: Tensor(seq_len, head_dim)
Returns:
Tensor(batch_size, num_heads, seq_len, head_dim), Tensor(batch_size, num_key_value_heads, seq_len, head_dim)
"""
orig_dtype = q.dtype
q = q.to(torch.float32)
k = k.to(torch.float32)
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed.to(orig_dtype), k_embed.to(orig_dtype)
class MiniCPMLongRoPE(nn.Module):
"""MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
def __init__(self, config: MiniCPM4Config):
super().__init__()
self.config = config
self.dim = config.kv_channels if config.kv_channels else config.hidden_size // config.num_attention_heads
self.base = config.rope_theta
self.max_position_embeddings = config.max_position_embeddings
self.short_factor = config.rope_scaling.short_factor
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))
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
self.max_seq_len_cached = 0
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)
def _set_cos_sin_cache(self, seq_len, device, dtype):
"""设置cos和sin缓存"""
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
if seq_len > self.original_max_position_embeddings:
ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=device)
else:
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)
)
# 创建embeddings
emb = torch.cat((freqs, freqs), dim=-1)
self.cos_cached = emb.cos().to(dtype) * self.scaling_factor
self.sin_cached = emb.sin().to(dtype) * self.scaling_factor
def forward(self, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Args:
position_ids: Tensor(seq_len) 或 Tensor(batch_size, seq_len)
Returns:
Tensor(seq_len, head_dim), Tensor(seq_len, head_dim)
"""
cos = self.cos_cached[position_ids]
sin = self.sin_cached[position_ids]
return cos, sin
class MiniCPMAttention(nn.Module):
def __init__(self, config: MiniCPM4Config, layer_idx: int):
super().__init__()
self.config = config
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.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
self.rope_theta = 10000.0
self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
def forward(
self,
hidden_states: torch.Tensor,
position_emb: Tuple[torch.Tensor, torch.Tensor],
is_causal: bool,
) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
bsz, q_len, _ = hidden_states.size()
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
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)
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,
value_states,
is_causal=is_causal,
enable_gqa=True,
)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim)
attn_output = self.o_proj(attn_output)
past_key_value = (key_states, value_states)
return attn_output, past_key_value
def forward_step(
self,
hidden_states: torch.Tensor,
position_emb: Tuple[torch.Tensor, torch.Tensor],
position_id: int,
kv_cache: Tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
bsz, _ = hidden_states.size()
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
query_states = query_states.view(bsz, 1, self.num_heads, self.head_dim).transpose(1, 2)
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)
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
key_cache[:, :, position_id, :] = key_states
value_cache[:, :, position_id, :] = value_states
# Use an explicit broadcastable mask shape for SDPA. A 1D mask can
# trigger a CPU-side dimension bug in some PyTorch versions.
attn_mask = (torch.arange(key_cache.size(2), device=key_cache.device) <= position_id).view(1, 1, 1, -1)
# 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,
value_cache,
attn_mask=attn_mask,
enable_gqa=True,
)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, self.num_heads * self.head_dim)
attn_output = self.o_proj(attn_output)
return attn_output
class MiniCPMMLP(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
self.act_fn = nn.SiLU()
def forward(self, x):
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
class MiniCPMDecoderLayer(nn.Module):
def __init__(self, config: MiniCPM4Config, layer_idx: int):
super().__init__()
self.hidden_size = config.hidden_size
self.self_attn = MiniCPMAttention(config=config, layer_idx=layer_idx)
self.mlp = MiniCPMMLP(config)
self.input_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.scale_depth = config.scale_depth
self.num_hidden_layers = config.num_hidden_layers
self.use_mup = config.use_mup
def forward(
self,
hidden_states: torch.Tensor,
position_emb: Tuple[torch.Tensor, torch.Tensor],
is_causal: bool,
) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
"""
Args:
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
position_ids (`torch.LongTensor`): position ids of shape `(batch_size, seq_len)`
is_causal (`bool`): whether the attention mask is causal
"""
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
hidden_states, present_key_value = self.self_attn(
hidden_states=hidden_states,
position_emb=position_emb,
is_causal=is_causal,
)
if self.use_mup:
hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
else:
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
if self.use_mup:
hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
else:
hidden_states = residual + hidden_states
return hidden_states, present_key_value
def forward_step(
self,
hidden_states: torch.Tensor,
position_emb: Tuple[torch.Tensor, torch.Tensor],
position_id: torch.Tensor,
kv_cache: Tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
hidden_states = self.self_attn.forward_step(
hidden_states=hidden_states,
position_emb=position_emb,
position_id=position_id,
kv_cache=kv_cache,
)
if self.use_mup:
hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
else:
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
if self.use_mup:
hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
else:
hidden_states = residual + hidden_states
return hidden_states
class MiniCPMModel(nn.Module):
"""
Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MiniCPMDecoderLayer`]
Args:
config: MiniCPMConfig
"""
def __init__(self, config: MiniCPM4Config):
super().__init__()
self.vocab_size = config.vocab_size
self.config = config
if config.vocab_size > 0:
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
else:
self.embed_tokens = nn.Identity()
self.layers = nn.ModuleList(
[MiniCPMDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
)
self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
if config.no_rope:
self.rope_emb = None
else:
self.rope_emb = MiniCPMLongRoPE(config)
self.kv_cache = None
def forward(
self,
inputs_embeds: torch.Tensor,
is_causal: bool = True,
) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
"""
Args:
inputs_embeds: Tensor(batch_size, seq_length, hidden_size)
is_causal: bool, whether the attention mask is causal
Returns:
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)]
"""
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 = []
for decoder_layer in self.layers:
hidden_states, this_cache = decoder_layer(
hidden_states,
position_emb,
is_causal,
)
next_decoder_cache.append(this_cache)
hidden_states = self.norm(hidden_states)
return hidden_states, next_decoder_cache
def forward_step(
self,
inputs_embeds: torch.Tensor,
position_id: torch.Tensor,
) -> torch.Tensor:
"""
Args:
inputs_embeds: Tensor(batch_size, hidden_size)
Returns:
hidden_states: Tensor(batch_size, hidden_size)
"""
assert self.kv_cache is not None, "KV cache is not setup"
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):
hidden_states = decoder_layer.forward_step(
hidden_states,
position_emb,
position_id,
self.kv_cache.get_layer_cache(i),
)
hidden_states = self.norm(hidden_states)
return hidden_states
def setup_cache(self, batch_size: int, max_length: int, device, dtype: torch.dtype):
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
),
batch_size=batch_size,
device=device,
dtype=dtype,
max_length=max_length,
)
src/voxcpm/timestamps/__init__.py
+8
View File
@@ -0,0 +1,8 @@
from .base import TimestampItem, TimestampResult
from .postprocess import align_audio_file
__all__ = [
"TimestampItem",
"TimestampResult",
"align_audio_file",
]
src/voxcpm/timestamps/base.py
+33
View File
@@ -0,0 +1,33 @@
from __future__ import annotations
from dataclasses import asdict, dataclass
from typing import Literal
TimestampLevel = Literal["segment", "word", "char"]
@dataclass
class TimestampItem:
text: str
start: float
end: float
level: str
def to_dict(self) -> dict:
return asdict(self)
@dataclass
class TimestampResult:
audio_path: str
sample_rate: int | None
backend: str
level: str
text: str
items: list[TimestampItem]
warning: str | None = None
def to_dict(self) -> dict:
payload = asdict(self)
payload["items"] = [item.to_dict() for item in self.items]
return payload
src/voxcpm/timestamps/postprocess.py
+33
View File
@@ -0,0 +1,33 @@
from __future__ import annotations
from .base import TimestampLevel
def align_audio_file(
*,
audio_path: str,
text: str,
sample_rate: int | None = None,
backend: str = "stable-ts",
level: TimestampLevel = "word",
model_name: str = "base",
device: str | None = None,
language: str | None = None,
) -> dict:
if backend != "stable-ts":
raise ValueError(f"Unsupported timestamp backend: {backend}")
from .stable_ts import StableTSAligner
aligner = StableTSAligner(
model_name=model_name,
device=device,
language=language,
)
result = aligner.align(
audio_path=audio_path,
text=text,
sample_rate=sample_rate,
level=level,
)
return result.to_dict()
src/voxcpm/timestamps/stable_ts.py
+105
View File
@@ -0,0 +1,105 @@
from __future__ import annotations
from typing import Any
from .base import TimestampItem, TimestampLevel, TimestampResult
class StableTSAligner:
def __init__(
self,
model_name: str = "base",
device: str | None = None,
language: str | None = None,
) -> None:
try:
import stable_whisper
except ImportError as exc:
raise ImportError(
"stable-ts is required for timestamp alignment. " 'Install with: pip install "voxcpm[timestamps]"'
) from exc
self.model = stable_whisper.load_model(model_name, device=device)
self.model_name = model_name
self.device = device
self.language = language
def align(
self,
*,
audio_path: str,
text: str,
sample_rate: int | None = None,
level: TimestampLevel = "word",
) -> TimestampResult:
result = self.model.align(audio_path, text, language=self.language)
items = extract_timestamp_items(result, level)
return TimestampResult(
audio_path=audio_path,
sample_rate=sample_rate,
backend="stable-ts",
level=level,
text=text,
items=items,
)
def extract_timestamp_items(result: Any, level: TimestampLevel) -> list[TimestampItem]:
segments = _get_value(result, "segments", []) or []
if level == "segment":
return [
TimestampItem(
text=str(_get_value(segment, "text", "")).strip(),
start=float(_get_value(segment, "start", 0.0) or 0.0),
end=float(_get_value(segment, "end", 0.0) or 0.0),
level="segment",
)
for segment in segments
if str(_get_value(segment, "text", "")).strip()
]
words = []
for segment in segments:
for word in _get_value(segment, "words", []) or []:
text = str(_get_value(word, "word", _get_value(word, "text", ""))).strip()
if not text:
continue
words.append(
TimestampItem(
text=text,
start=float(_get_value(word, "start", 0.0) or 0.0),
end=float(_get_value(word, "end", 0.0) or 0.0),
level="word",
)
)
if level == "char":
return split_word_items_to_chars(words)
return words
def split_word_items_to_chars(words: list[TimestampItem]) -> list[TimestampItem]:
chars = []
for word in words:
text = word.text.strip()
if not text:
continue
duration = max(word.end - word.start, 0.0)
step = duration / len(text)
for idx, char in enumerate(text):
chars.append(
TimestampItem(
text=char,
start=word.start + idx * step,
end=word.start + (idx + 1) * step,
level="char",
)
)
return chars
def _get_value(obj: Any, key: str, default: Any = None) -> Any:
if isinstance(obj, dict):
return obj.get(key, default)
return getattr(obj, key, default)
src/voxcpm/training/__init__.py
+30
View File
@@ -0,0 +1,30 @@
"""
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
from .validate import validate_manifest, ValidationResult
__all__ = [
"Accelerator",
"TrainingTracker",
"HFVoxCPMDataset",
"BatchProcessor",
"TrainingState",
"load_audio_text_datasets",
"build_dataloader",
"validate_manifest",
"ValidationResult",
]
src/voxcpm/training/accelerator.py
+163
View File
@@ -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
+38
View File
@@ -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
+256
View File
@@ -0,0 +1,256 @@
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_REF_AUDIO_COLUMN = "ref_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,
ref_audio_column: str = DEFAULT_REF_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.")
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)
# ref_audio is optional — cast to Audio if the column exists
ref_col = ref_audio_column if ref_audio_column in ds.column_names else DEFAULT_REF_AUDIO_COLUMN
if ref_col in ds.column_names:
ds = ds.cast_column(ref_col, Audio(sampling_rate=sample_rate))
if ref_col != DEFAULT_REF_AUDIO_COLUMN:
ds = ds.rename_column(ref_col, DEFAULT_REF_AUDIO_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)
- 无 ref_audio: text_len + t_seq + 2
- 有 ref_audio: text_len + t_seq + ref_seq + 4
Optimized: Use batch column access instead of iterating item by item.
"""
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:
durations = []
for i in range(len(ds)):
audio = ds[i][DEFAULT_AUDIO_COLUMN]
durations.append(len(audio["array"]) / float(audio["sampling_rate"]))
has_ref_audio = DEFAULT_REF_AUDIO_COLUMN in ds.column_names
if has_ref_audio:
ref_duration_col = "ref_duration" if "ref_duration" in ds.column_names else None
lengths = []
for i, (text_len, duration) in enumerate(zip(text_lens, durations)):
t_vae = math.ceil(float(duration) * audio_vae_fps)
t_seq = math.ceil(t_vae / patch_size)
ref_seq = 0
if has_ref_audio:
# Estimate ref_audio length; ref_audio is None for samples without it
if ref_duration_col:
ref_dur = ds[i].get(ref_duration_col)
else:
ref_item = ds[i].get(DEFAULT_REF_AUDIO_COLUMN)
ref_dur = len(ref_item["array"]) / float(ref_item["sampling_rate"]) if ref_item else None
if ref_dur is not None and float(ref_dur) > 0:
ref_vae = math.ceil(float(ref_dur) * audio_vae_fps)
ref_seq = math.ceil(ref_vae / patch_size)
# +2 for 101/102; +2 more for 103/104 when ref_audio present
overhead = 4 if ref_seq > 0 else 2
total_len = text_len + t_seq + ref_seq + overhead
lengths.append(total_len)
return lengths
class HFVoxCPMDataset(TorchDataset):
"""
Thin wrapper around a tokenized HuggingFace dataset that returns
PyTorch-friendly samples.
"""
_SENTINEL = [-100.0]
def __init__(self, dataset: Dataset):
self.dataset = dataset
self.has_ref_audio = DEFAULT_REF_AUDIO_COLUMN in dataset.column_names
def __len__(self):
return len(self.dataset)
def __getitem__(self, idx: int):
item = self.dataset[idx]
audio = item[DEFAULT_AUDIO_COLUMN]
sample = {
"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),
}
if self.has_ref_audio:
ref = item.get(DEFAULT_REF_AUDIO_COLUMN)
sample["ref_audio_array"] = ref["array"] if ref else self._SENTINEL
return sample
@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)
result = {
"text_tokens": text_padded,
"audio_tokens": audio_padded,
"task_ids": task_ids,
"dataset_ids": dataset_ids,
"is_prompts": is_prompts,
}
if "ref_audio_array" in batch[0]:
ref_tensors = [torch.tensor(s["ref_audio_array"], dtype=torch.float32) for s in batch]
result["ref_audio_tokens"] = cls.pad_sequences(ref_tensors, pad_value=-100.0)
return result
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)
ref_audio_tokens = None
if "ref_audio_tokens" in batch:
ref_audio_tokens = batch["ref_audio_tokens"].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"],
ref_audio_tokens=ref_audio_tokens,
)
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
+441
View File
@@ -0,0 +1,441 @@
from typing import Dict, List, Optional
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
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],
ref_audio_tokens: Optional[torch.Tensor] = None,
) -> 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, ...].
If ``ref_audio_tokens`` is provided (same batch dim as ``audio_tokens``),
samples whose unpadded ref_audio length > 0 will be processed with the
reference-audio path (tokens 103/104 prepended, loss only on target audio).
"""
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)
ref_iter = ref_audio_tokens if ref_audio_tokens is not None else [None] * audio_tokens.size(0)
for audio_token, text_token, task_id, dataset_idx, is_prompt, ref_token in zip(
audio_tokens, text_tokens, task_ids.tolist(), dataset_ids.tolist(), is_prompts, ref_iter
):
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]
has_ref = False
if ref_token is not None:
unpad_ref_token = self.unpad_audio_tokens(ref_token).to(torch.float32)
if unpad_ref_token.numel() > 0:
has_ref = True
if has_ref:
(
packed_text,
audio_feat,
text_mask,
audio_mask,
loss_mask,
labels,
audio_duration,
text_token_count,
) = self.process_tts_data_with_ref(unpad_ref_token, unpad_audio_token, unpad_text_token)
else:
(
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,
)
def process_tts_data_with_ref(
self,
ref_audio_token: torch.Tensor,
target_audio_token: torch.Tensor,
text_token: torch.Tensor,
):
"""
Build a training sequence with reference audio prepended:
[103, ref_feats, 104, text, 101, target_feats, 102]
Loss is computed only on the target audio segment.
"""
device = text_token.device
txt_len = len(text_token)
ref_feats, ref_duration = self.extract_audio_feats(ref_audio_token)
ref_feats = ref_feats.squeeze(0) # [R, P, D]
ref_len = ref_feats.shape[0]
tgt_feats, tgt_duration = self.extract_audio_feats(target_audio_token)
tgt_feats = tgt_feats.squeeze(0) # [A, P, D]
tgt_len = tgt_feats.shape[0]
feat_shape = (self.patch_size, ref_feats.size(-1))
def _tok(ids):
return torch.tensor(ids, dtype=torch.int32, device=device)
# -- text token track --
# [103, 0×R, 104, text_ids, 101, 0×A, 102]
text_token_info = torch.cat(
[
_tok([self.audio_prompt_start_id]),
torch.zeros(ref_len, dtype=torch.int32, device=device),
_tok([self.audio_prompt_end_id]),
text_token,
_tok([self.audio_start_id]),
torch.zeros(tgt_len, dtype=torch.int32, device=device),
_tok([self.audio_end_id]),
]
)
# -- audio feature track --
zero_1 = torch.zeros((1,) + feat_shape, dtype=torch.float32, device=device)
zero_txt = torch.zeros((txt_len,) + feat_shape, dtype=torch.float32, device=device)
audio_feat_info = torch.cat(
[
zero_1,
ref_feats,
zero_1, # 103, ref, 104
zero_txt, # text
zero_1,
tgt_feats,
zero_1, # 101, target, 102
],
dim=0,
)
# -- masks --
text_mask = (
torch.cat(
[
torch.ones(1),
torch.zeros(ref_len),
torch.ones(1),
torch.ones(txt_len),
torch.ones(1),
torch.zeros(tgt_len),
torch.ones(1),
]
)
.to(torch.int32)
.to(device)
)
audio_mask = (
torch.cat(
[
torch.zeros(1),
torch.ones(ref_len),
torch.zeros(1),
torch.zeros(txt_len),
torch.zeros(1),
torch.ones(tgt_len),
torch.zeros(1),
]
)
.to(torch.int32)
.to(device)
)
loss_mask = (
torch.cat(
[
torch.zeros(1 + ref_len + 1), # ref part: no loss
torch.zeros(txt_len), # text: no loss
torch.zeros(1), # 101: no loss
torch.ones(tgt_len), # target audio: LOSS
torch.zeros(1), # 102: no loss
]
)
.to(torch.int32)
.to(device)
)
total_len = 1 + ref_len + 1 + txt_len + 1 + tgt_len + 1
labels = torch.zeros(total_len, dtype=torch.int32, device=device)
labels[-2] = 1 # stop label at last target audio position
return (
text_token_info,
audio_feat_info,
text_mask,
audio_mask,
loss_mask,
labels,
ref_duration + tgt_duration,
txt_len,
)
src/voxcpm/training/state.py
+20
View File
@@ -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
+78
View File
@@ -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/training/validate.py
+296
View File
@@ -0,0 +1,296 @@
"""
Pre-flight validation for VoxCPM training data manifests.
Validates JSONL manifest files before starting expensive fine-tuning jobs,
catching format issues, missing files, and data quality problems early.
"""
import json
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional
@dataclass
class ValidationResult:
"""Structured result of a manifest validation run."""
total_samples: int = 0
valid_samples: int = 0
errors: List[str] = field(default_factory=list)
warnings: List[str] = field(default_factory=list)
audio_durations: List[float] = field(default_factory=list)
text_lengths: List[int] = field(default_factory=list)
has_ref_audio: int = 0
@property
def is_valid(self) -> bool:
return len(self.errors) == 0 and self.valid_samples > 0
def _check_audio_file(audio_path: str, sample_rate: int) -> Optional[str]:
"""Check if an audio file exists, is readable, and matches expected sample rate.
Returns an error message, or None if the file is valid.
"""
if not os.path.isfile(audio_path):
return f"Audio file not found: {audio_path}"
try:
import soundfile as sf
info = sf.info(audio_path)
if info.frames == 0:
return f"Audio file is empty: {audio_path}"
if info.samplerate != sample_rate:
return f"Sample rate mismatch in {audio_path}: " f"expected {sample_rate} Hz, got {info.samplerate} Hz"
return None
except ImportError:
# soundfile not available; just check existence
return None
except Exception as e:
return f"Cannot read audio file {audio_path}: {e}"
def _get_audio_duration(audio_path: str) -> Optional[float]:
"""Get audio duration in seconds. Returns None if unavailable."""
try:
import soundfile as sf
info = sf.info(audio_path)
return info.duration
except Exception:
return None
def validate_manifest(
manifest_path: str,
sample_rate: int = 16_000,
max_samples: int = 0,
verbose: bool = False,
) -> ValidationResult:
"""Validate a JSONL training manifest file.
Checks:
1. File exists and is readable
2. Each line is valid JSON
3. Required columns present (text, audio)
4. Audio files exist and are readable
5. Text content is non-empty
6. Collects duration and text length statistics
7. Validates optional ref_audio column
Args:
manifest_path: Path to the JSONL manifest file.
sample_rate: Expected audio sample rate (for informational purposes).
max_samples: Maximum number of samples to validate (0 = all).
verbose: Print per-sample progress.
Returns:
ValidationResult with errors, warnings, and statistics.
"""
result = ValidationResult()
path = Path(manifest_path)
if not path.exists():
result.errors.append(f"Manifest file not found: {manifest_path}")
return result
if not path.is_file():
result.errors.append(f"Manifest path is not a file: {manifest_path}")
return result
manifest_dir = path.parent
try:
with open(path, "r", encoding="utf-8") as f:
lines = f.readlines()
except Exception as e:
result.errors.append(f"Cannot read manifest file: {e}")
return result
if not lines:
result.errors.append("Manifest file is empty")
return result
samples_to_check = len(lines)
if max_samples > 0:
samples_to_check = min(samples_to_check, max_samples)
missing_audio_count = 0
empty_text_count = 0
for i, line in enumerate(lines[:samples_to_check]):
line = line.strip()
if not line:
continue
result.total_samples += 1
# Check JSON validity
try:
entry = json.loads(line)
except json.JSONDecodeError as e:
result.errors.append(f"Line {i + 1}: Invalid JSON — {e}")
continue
if not isinstance(entry, dict):
result.errors.append(f"Line {i + 1}: Expected JSON object, got {type(entry).__name__}")
continue
# Check required columns
has_error = False
if "text" not in entry:
result.errors.append(f"Line {i + 1}: Missing required column 'text'")
has_error = True
if "audio" not in entry:
result.errors.append(f"Line {i + 1}: Missing required column 'audio'")
has_error = True
if has_error:
continue
# Validate text
text = entry["text"]
if not isinstance(text, str) or not text.strip():
empty_text_count += 1
if empty_text_count <= 5:
result.warnings.append(f"Line {i + 1}: Empty or non-string text")
else:
result.text_lengths.append(len(text))
# Validate audio path
audio_path = entry["audio"]
if isinstance(audio_path, dict):
# HuggingFace Audio format with {"path": ..., "array": ...}
audio_path = audio_path.get("path", "")
if isinstance(audio_path, str) and audio_path:
# Resolve relative paths against manifest directory
if not os.path.isabs(audio_path):
audio_path = str(manifest_dir / audio_path)
audio_error = _check_audio_file(audio_path, sample_rate)
if audio_error:
missing_audio_count += 1
if missing_audio_count <= 5:
result.errors.append(f"Line {i + 1}: {audio_error}")
has_error = True
else:
duration = _get_audio_duration(audio_path)
if duration is not None:
result.audio_durations.append(duration)
if duration < 0.3:
result.warnings.append(f"Line {i + 1}: Very short audio ({duration:.2f}s)")
elif duration > 30.0:
result.warnings.append(f"Line {i + 1}: Very long audio ({duration:.1f}s), may cause OOM")
else:
result.errors.append(f"Line {i + 1}: Invalid audio path")
has_error = True
# Validate optional ref_audio
if "ref_audio" in entry:
ref_path = entry["ref_audio"]
if isinstance(ref_path, dict):
ref_path = ref_path.get("path", "")
if isinstance(ref_path, str) and ref_path:
if not os.path.isabs(ref_path):
ref_path = str(manifest_dir / ref_path)
if os.path.isfile(ref_path):
result.has_ref_audio += 1
else:
result.warnings.append(f"Line {i + 1}: ref_audio file not found: {ref_path}")
if not has_error:
result.valid_samples += 1
if verbose and (i + 1) % 100 == 0:
print(f" Validated {i + 1}/{samples_to_check} samples...", file=sys.stderr)
# Summarize truncated errors
if missing_audio_count > 5:
result.errors.append(
f"... and {missing_audio_count - 5} more missing audio files " f"({missing_audio_count} total)"
)
if empty_text_count > 5:
result.warnings.append(f"... and {empty_text_count - 5} more empty text entries " f"({empty_text_count} total)")
return result
def print_validation_report(result: ValidationResult, manifest_path: str) -> None:
"""Print a human-readable validation report to stderr."""
print(f"\n{'=' * 60}", file=sys.stderr)
print(f" VoxCPM Training Data Validation Report", file=sys.stderr)
print(f"{'=' * 60}", file=sys.stderr)
print(f" Manifest : {manifest_path}", file=sys.stderr)
print(f" Samples : {result.valid_samples}/{result.total_samples} valid", file=sys.stderr)
if result.has_ref_audio > 0:
print(
f" Ref Audio: {result.has_ref_audio} samples with reference audio",
file=sys.stderr,
)
# Audio duration statistics
if result.audio_durations:
durations = sorted(result.audio_durations)
total_hrs = sum(durations) / 3600
print(f"\n Audio Duration Statistics:", file=sys.stderr)
print(f" Total : {total_hrs:.2f} hours", file=sys.stderr)
print(
f" Range : {durations[0]:.2f}s — {durations[-1]:.1f}s",
file=sys.stderr,
)
print(
f" Mean : {sum(durations) / len(durations):.2f}s",
file=sys.stderr,
)
median_idx = len(durations) // 2
print(f" Median : {durations[median_idx]:.2f}s", file=sys.stderr)
# Text length statistics
if result.text_lengths:
lengths = sorted(result.text_lengths)
print(f"\n Text Length Statistics (characters):", file=sys.stderr)
print(
f" Range : {lengths[0]} — {lengths[-1]}",
file=sys.stderr,
)
print(
f" Mean : {sum(lengths) / len(lengths):.0f}",
file=sys.stderr,
)
# Errors
if result.errors:
print(f"\n ERRORS ({len(result.errors)}):", file=sys.stderr)
for err in result.errors[:20]:
print(f" x {err}", file=sys.stderr)
if len(result.errors) > 20:
print(
f" ... ({len(result.errors) - 20} more errors omitted)",
file=sys.stderr,
)
# Warnings
if result.warnings:
print(f"\n WARNINGS ({len(result.warnings)}):", file=sys.stderr)
for warn in result.warnings[:10]:
print(f" ! {warn}", file=sys.stderr)
if len(result.warnings) > 10:
print(
f" ... ({len(result.warnings) - 10} more warnings omitted)",
file=sys.stderr,
)
# Summary
print(f"\n{'=' * 60}", file=sys.stderr)
if result.is_valid:
print(" PASSED: Manifest is valid for training.", file=sys.stderr)
else:
print(" FAILED: Fix errors above before starting training.", file=sys.stderr)
print(f"{'=' * 60}\n", file=sys.stderr)
src/voxcpm/utils/text_normalize.py
+188
View File
@@ -0,0 +1,188 @@
# some functions are copied from https://github.com/FunAudioLLM/CosyVoice/blob/main/cosyvoice/utils/frontend_utils.py
import re
import regex
import inflect
from wetext import Normalizer
chinese_char_pattern = re.compile(r"[\u4e00-\u9fff]+")
# whether contain chinese character
def contains_chinese(text):
return bool(chinese_char_pattern.search(text))
# replace special symbol
def replace_corner_mark(text):
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("——", " ")
return text
# spell Arabic numerals
def spell_out_number(text: str, inflect_parser):
new_text = []
st = None
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])
new_text.append(num_str)
st = None
new_text.append(c)
else:
if st is None:
st = i
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)
# split paragrah logic：
# 1. per sentence max len token_max_n, min len token_min_n, merge if last sentence len less than merge_len
# 2. cal sentence len according to lang
# 3. split sentence according to puncatation
def split_paragraph(text: str, tokenize, lang="zh", token_max_n=80, token_min_n=60, merge_len=20, comma_split=False):
def calc_utt_length(_text: str):
if lang == "zh":
return len(_text)
else:
return len(tokenize(_text))
def should_merge(_text: str):
if lang == "zh":
return len(_text) < merge_len
else:
return len(tokenize(_text)) < merge_len
if lang == "zh":
pounc = ["。", "？", "！", "；", "：", "、", ".", "?", "!", ";"]
else:
pounc = [".", "?", "!", ";", ":"]
if comma_split:
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 ['"', "”"]:
tmp = utts.pop(-1)
utts.append(tmp + text[i + 1])
st = i + 2
else:
st = i + 1
if len(utts) == 0:
if lang == "zh":
utts.append(text + "。")
else:
utts.append(text + ".")
final_utts = []
cur_utt = ""
for utt in utts:
if calc_utt_length(cur_utt + utt) > token_max_n and calc_utt_length(cur_utt) > token_min_n:
final_utts.append(cur_utt)
cur_utt = ""
cur_utt = cur_utt + utt
if len(cur_utt) > 0:
if should_merge(cur_utt) and len(final_utts) != 0:
final_utts[-1] = final_utts[-1] + cur_utt
else:
final_utts.append(cur_utt)
return final_utts
# remove blank between chinese character
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] != " "):
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)
# 去除内联代码 `code`
md_text = re.sub(r"`[^`]*`", "", md_text)
# 去除图片语法 ![alt](url)
md_text = re.sub(r"!\[[^\]]*\]\([^\)]+\)", "", md_text)
# 去除链接但保留文本 [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)
# 去除HTML标签
md_text = re.sub(r"<[^>]+>", "", md_text)
# 去除标题符号（#）
md_text = re.sub(r"^#{1,6}\s*", "", md_text, flags=re.MULTILINE)
# 去除多余空格和空行
md_text = re.sub(r"\n\s*\n", "\n", md_text) # 多余空行
md_text = md_text.strip()
return md_text
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 = text.replace("\n", " ")
text = text.replace("\t", " ")
text = text.replace("“", '"').replace("”", '"')
return text
class TextNormalizer:
def __init__(self, tokenizer=None):
self.tokenizer = tokenizer
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 = replace_blank(text)
text = replace_corner_mark(text)
text = remove_bracket(text)
else:
text = self.en_tn_model.normalize(text)
text = spell_out_number(text, self.inflect_parser)
if split is False:
return text
src/voxcpm/zipenhancer.py
+72
View File
@@ -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
+736
View File
@@ -0,0 +1,736 @@
from __future__ import annotations
import importlib.util
import json
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 patch_soundfile_write(monkeypatch):
soundfile_stub = types.SimpleNamespace(write=lambda *args, **kwargs: None)
monkeypatch.setitem(sys.modules, "soundfile", soundfile_stub)
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.device == "auto"
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(core_stub, "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"]["device"] == "auto"
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(core_stub, "VoxCPM", FakeVoxCPM)
args = cli._build_parser().parse_args(
[
"design",
"--text",
"hello",
"--output",
"out.wav",
]
)
cli.load_model(args)
assert calls["kwargs"]["device"] == "auto"
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(core_stub, "VoxCPM", FakeVoxCPM)
args = cli._build_parser().parse_args(
[
"design",
"--text",
"hello",
"--output",
"out.wav",
"--no-optimize",
]
)
cli.load_model(args)
assert calls["kwargs"]["device"] == "auto"
assert calls["kwargs"]["optimize"] is False
def test_load_model_passes_explicit_device_to_hf(monkeypatch):
calls = {}
class FakeVoxCPM:
@classmethod
def from_pretrained(cls, **kwargs):
calls["kwargs"] = kwargs
return DummyModel()
monkeypatch.setattr(core_stub, "VoxCPM", FakeVoxCPM)
args = cli._build_parser().parse_args(
[
"design",
"--text",
"hello",
"--output",
"out.wav",
"--device",
"mps",
]
)
cli.load_model(args)
assert calls["kwargs"]["device"] == "mps"
def test_design_subcommand_applies_control(monkeypatch, tmp_path):
dummy_model = DummyModel()
monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
patch_soundfile_write(monkeypatch)
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)
patch_soundfile_write(monkeypatch)
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)
patch_soundfile_write(monkeypatch)
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)
patch_soundfile_write(monkeypatch)
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_design_writes_timestamp_json_when_requested(monkeypatch, tmp_path):
dummy_model = DummyModel()
timestamp_calls = []
def fake_align_audio_file(**kwargs):
timestamp_calls.append(kwargs)
return {
"audio_path": kwargs["audio_path"],
"sample_rate": kwargs["sample_rate"],
"backend": kwargs["backend"],
"level": kwargs["level"],
"text": kwargs["text"],
"items": [{"text": "hello", "start": 0.0, "end": 0.5, "level": "word"}],
"warning": None,
}
monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
monkeypatch.setattr(cli, "align_audio_file", fake_align_audio_file)
patch_soundfile_write(monkeypatch)
output = tmp_path / "out.wav"
run_main(
monkeypatch,
[
"design",
"--text",
"hello",
"--output",
str(output),
"--timestamps",
"--timestamp-language",
"en",
],
)
timestamp_path = tmp_path / "out.timestamps.json"
payload = json.loads(timestamp_path.read_text(encoding="utf-8"))
assert payload["audio_path"] == str(output)
assert payload["text"] == "hello"
assert payload["items"][0]["text"] == "hello"
assert timestamp_calls[0]["language"] == "en"
assert timestamp_calls[0]["level"] == "word"
def test_timestamp_alignment_failure_warns_by_default(monkeypatch, tmp_path, capsys):
dummy_model = DummyModel()
def fake_align_audio_file(**kwargs):
raise RuntimeError("alignment failed")
monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
monkeypatch.setattr(cli, "align_audio_file", fake_align_audio_file)
patch_soundfile_write(monkeypatch)
run_main(
monkeypatch,
[
"design",
"--text",
"hello",
"--output",
str(tmp_path / "out.wav"),
"--timestamps",
],
)
assert "Timestamp alignment failed" in capsys.readouterr().err
def test_timestamp_alignment_failure_exits_in_strict_mode(monkeypatch, tmp_path):
dummy_model = DummyModel()
def fake_align_audio_file(**kwargs):
raise RuntimeError("alignment failed")
monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
monkeypatch.setattr(cli, "align_audio_file", fake_align_audio_file)
patch_soundfile_write(monkeypatch)
with pytest.raises(SystemExit):
run_main(
monkeypatch,
[
"design",
"--text",
"hello",
"--output",
str(tmp_path / "out.wav"),
"--timestamps",
"--timestamp-strict",
],
)
def test_batch_writes_one_timestamp_json_per_output(monkeypatch, tmp_path):
dummy_model = DummyModel()
def fake_align_audio_file(**kwargs):
return {
"audio_path": kwargs["audio_path"],
"sample_rate": kwargs["sample_rate"],
"backend": kwargs["backend"],
"level": kwargs["level"],
"text": kwargs["text"],
"items": [{"text": kwargs["text"], "start": 0.0, "end": 0.5, "level": "word"}],
"warning": None,
}
input_file = tmp_path / "texts.txt"
input_file.write_text("hello\nworld\n", encoding="utf-8")
output_dir = tmp_path / "outs"
monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
monkeypatch.setattr(cli, "align_audio_file", fake_align_audio_file)
patch_soundfile_write(monkeypatch)
run_main(
monkeypatch,
[
"batch",
"--input",
str(input_file),
"--output-dir",
str(output_dir),
"--timestamps",
],
)
first = json.loads((output_dir / "output_001.timestamps.json").read_text(encoding="utf-8"))
second = json.loads((output_dir / "output_002.timestamps.json").read_text(encoding="utf-8"))
assert first["text"] == "hello"
assert second["text"] == "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)
patch_soundfile_write(monkeypatch)
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"
def test_parser_accepts_seed():
parser = cli._build_parser()
# Default seed should be None
args = parser.parse_args(["design", "--text", "hello", "--output", "out.wav"])
assert args.seed is None
# Custom seed should be parsed as int
args = parser.parse_args(["design", "--text", "hello", "--output", "out.wav", "--seed", "42"])
assert args.seed == 42
def test_design_subcommand_passes_seed(monkeypatch, tmp_path):
dummy_model = DummyModel()
monkeypatch.setattr(cli, "load_model", lambda args: dummy_model)
patch_soundfile_write(monkeypatch)
run_main(
monkeypatch,
[
"design",
"--text",
"hello",
"--seed",
"123",
"--output",
str(tmp_path / "out.wav"),
],
)
assert dummy_model.calls[0]["seed"] == 123
def test_batch_subcommand_passes_seed(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)
patch_soundfile_write(monkeypatch)
run_main(
monkeypatch,
[
"batch",
"--input",
str(input_file),
"--output-dir",
str(tmp_path / "outs"),
"--seed",
"999",
],
)
assert len(dummy_model.calls) == 2
assert dummy_model.calls[0]["seed"] == 999
assert dummy_model.calls[1]["seed"] == 999
tests/test_lora_checkpoint_loading.py
+150
View File
@@ -0,0 +1,150 @@
from __future__ import annotations
import importlib.util
import sys
import types
from pathlib import Path
import pytest
import torch
ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
def _load_module(name: str, path: Path):
spec = importlib.util.spec_from_file_location(name, path)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
sys.modules[name] = module
spec.loader.exec_module(module)
return module
def bootstrap_repo_modules(monkeypatch):
for name, path in [
("voxcpm", SRC / "voxcpm"),
("voxcpm.model", SRC / "voxcpm" / "model"),
("voxcpm.modules", SRC / "voxcpm" / "modules"),
]:
pkg = types.ModuleType(name)
pkg.__path__ = [str(path)]
monkeypatch.setitem(sys.modules, name, pkg)
hh = types.ModuleType("huggingface_hub")
hh.snapshot_download = lambda *a, **k: "/tmp/fake"
monkeypatch.setitem(sys.modules, "huggingface_hub", hh)
pydantic = types.ModuleType("pydantic")
class BaseModel:
@classmethod
def model_rebuild(cls):
return None
@classmethod
def model_validate_json(cls, s):
return cls()
def model_dump(self):
return {}
pydantic.BaseModel = BaseModel
monkeypatch.setitem(sys.modules, "pydantic", pydantic)
torchaudio = types.ModuleType("torchaudio")
monkeypatch.setitem(sys.modules, "torchaudio", torchaudio)
librosa = types.ModuleType("librosa")
librosa.effects = types.SimpleNamespace(trim=lambda *a, **k: (None, (0, 0)))
monkeypatch.setitem(sys.modules, "librosa", librosa)
einops = types.ModuleType("einops")
einops.rearrange = lambda x, *a, **k: x
monkeypatch.setitem(sys.modules, "einops", einops)
tqdm_pkg = types.ModuleType("tqdm")
tqdm_pkg.__path__ = ["/nonexistent"]
tqdm_pkg.tqdm = lambda x, *a, **k: x
monkeypatch.setitem(sys.modules, "tqdm", tqdm_pkg)
tqdm_auto = types.ModuleType("tqdm.auto")
tqdm_auto.tqdm = lambda x, *a, **k: x
monkeypatch.setitem(sys.modules, "tqdm.auto", tqdm_auto)
transformers = types.ModuleType("transformers")
class LlamaTokenizerFast:
pass
class PreTrainedTokenizer:
pass
transformers.LlamaTokenizerFast = LlamaTokenizerFast
transformers.PreTrainedTokenizer = PreTrainedTokenizer
monkeypatch.setitem(sys.modules, "transformers", transformers)
internal_mods = {
"voxcpm.modules.audiovae": ["AudioVAE", "AudioVAEConfig", "AudioVAEV2", "AudioVAEConfigV2"],
"voxcpm.modules.layers": ["ScalarQuantizationLayer"],
"voxcpm.modules.locdit": ["CfmConfig", "UnifiedCFM", "VoxCPMLocDiT", "VoxCPMLocDiTV2"],
"voxcpm.modules.locenc": ["VoxCPMLocEnc"],
"voxcpm.modules.minicpm4": ["MiniCPM4Config", "MiniCPMModel"],
"voxcpm.modules.layers.lora": ["apply_lora_to_named_linear_modules", "LoRALinear"],
}
for modname, names in internal_mods.items():
module = types.ModuleType(modname)
for name in names:
if name == "apply_lora_to_named_linear_modules":
setattr(module, name, lambda *a, **k: None)
else:
setattr(module, name, type(name, (), {}))
monkeypatch.setitem(sys.modules, modname, module)
_load_module("voxcpm.model.utils", SRC / "voxcpm" / "model" / "utils.py")
voxcpm = _load_module("voxcpm.model.voxcpm", SRC / "voxcpm" / "model" / "voxcpm.py")
voxcpm2 = _load_module("voxcpm.model.voxcpm2", SRC / "voxcpm" / "model" / "voxcpm2.py")
return voxcpm.VoxCPMModel, voxcpm2.VoxCPM2Model
class DummyModel:
device = "cpu"
def named_parameters(self):
return []
@pytest.mark.parametrize("module_name", ["v1", "v2"])
def test_load_lora_weights_accepts_tensor_only_legacy_checkpoints(monkeypatch, tmp_path, module_name):
VoxCPMModel, VoxCPM2Model = bootstrap_repo_modules(monkeypatch)
cls = VoxCPMModel if module_name == "v1" else VoxCPM2Model
ckpt_path = tmp_path / "lora_weights.ckpt"
torch.save({"state_dict": {"fake": torch.zeros(1)}}, ckpt_path)
loaded, skipped = cls.load_lora_weights(DummyModel(), str(ckpt_path), device="cpu")
assert loaded == []
assert skipped == ["fake"]
@pytest.mark.parametrize("module_name", ["v1", "v2"])
def test_load_lora_weights_rejects_malicious_pickle_payloads(monkeypatch, tmp_path, module_name):
VoxCPMModel, VoxCPM2Model = bootstrap_repo_modules(monkeypatch)
cls = VoxCPMModel if module_name == "v1" else VoxCPM2Model
ckpt_path = tmp_path / "lora_weights.ckpt"
marker_path = tmp_path / f"{module_name}-marker.txt"
class Exploit:
def __reduce__(self):
import pathlib
return (pathlib.Path.write_text, (marker_path, f"{module_name} executed\n"))
torch.save({"state_dict": {"fake": torch.zeros(1)}, "boom": Exploit()}, ckpt_path)
with pytest.raises(Exception, match="Weights only load failed"):
cls.load_lora_weights(DummyModel(), str(ckpt_path), device="cpu")
assert not marker_path.exists()
tests/test_model_utils.py
+71
View File
@@ -0,0 +1,71 @@
from __future__ import annotations
import importlib.util
import sys
import types
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
UTILS_PATH = ROOT / "src" / "voxcpm" / "model" / "utils.py"
transformers_stub = types.ModuleType("transformers")
transformers_stub.PreTrainedTokenizer = object
sys.modules.setdefault("transformers", transformers_stub)
spec = importlib.util.spec_from_file_location("voxcpm.model.utils", UTILS_PATH)
utils = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(utils)
def test_resolve_runtime_device_auto_falls_back_to_cpu(monkeypatch):
monkeypatch.setattr(utils.torch.cuda, "is_available", lambda: False)
monkeypatch.setattr(utils, "_has_mps", lambda: False)
assert utils.resolve_runtime_device(None, "cuda") == "cpu"
def test_resolve_runtime_device_auto_uses_mps_when_available(monkeypatch):
monkeypatch.setattr(utils.torch.cuda, "is_available", lambda: False)
monkeypatch.setattr(utils, "_has_mps", lambda: True)
assert utils.resolve_runtime_device("auto", "cuda") == "mps"
def test_resolve_runtime_device_respects_explicit_cpu(monkeypatch):
monkeypatch.setattr(utils.torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(utils, "_has_mps", lambda: True)
assert utils.resolve_runtime_device("cpu", "cuda") == "cpu"
def test_resolve_runtime_device_rejects_unavailable_explicit_cuda(monkeypatch):
monkeypatch.setattr(utils.torch.cuda, "is_available", lambda: False)
monkeypatch.setattr(utils, "_has_mps", lambda: True)
with pytest.raises(ValueError, match="CUDA is not available"):
utils.resolve_runtime_device("cuda:0", "cuda")
def test_materialize_generation_seed_preserves_explicit_seed():
assert utils.materialize_generation_seed(42) == 42
def test_materialize_generation_seed_creates_concrete_seed_for_none(monkeypatch):
monkeypatch.setattr(utils.torch, "seed", lambda: 0x123456789)
assert utils.materialize_generation_seed(None) == 0x23456789
def test_apply_generation_seed_sets_cpu_and_cuda_rng(monkeypatch):
calls = []
monkeypatch.setattr(utils.torch, "manual_seed", lambda seed: calls.append(("cpu", seed)))
monkeypatch.setattr(utils.torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(utils.torch.cuda, "manual_seed_all", lambda seed: calls.append(("cuda", seed)))
utils.apply_generation_seed(123)
assert calls == [("cpu", 123), ("cuda", 123)]
tests/test_timestamps.py
+34
View File
@@ -0,0 +1,34 @@
from pathlib import Path
import sys
import types
ROOT = Path(__file__).resolve().parents[1]
pkg = types.ModuleType("voxcpm")
pkg.__path__ = [str(ROOT / "src" / "voxcpm")]
sys.modules.setdefault("voxcpm", pkg)
from voxcpm.timestamps.base import TimestampItem
from voxcpm.timestamps.stable_ts import split_word_items_to_chars
def test_split_word_items_to_chars_evenly_distributes_word_duration():
chars = split_word_items_to_chars([TimestampItem(text="欢迎", start=0.5, end=0.9, level="word")])
assert [item.text for item in chars] == ["欢", "迎"]
assert chars[0].start == 0.5
assert chars[0].end == 0.7
assert chars[1].start == 0.7
assert chars[1].end == 0.9
assert all(item.level == "char" for item in chars)
def test_split_word_items_to_chars_skips_empty_text():
chars = split_word_items_to_chars(
[
TimestampItem(text=" ", start=0.0, end=0.2, level="word"),
TimestampItem(text="你", start=0.2, end=0.4, level="word"),
]
)
assert len(chars) == 1
assert chars[0].text == "你"
tests/test_torch_load_safety.py
+85
View File
@@ -0,0 +1,85 @@
"""Regression guard: every ``torch.load`` call must set ``weights_only=True``.
VoxCPM deliberately loads checkpoints with ``weights_only=True`` so that a
crafted ``.ckpt``/``.pth``/``.bin`` file cannot execute arbitrary code via
pickle during unpickling (see
``tests/test_lora_checkpoint_loading.py::test_load_lora_weights_rejects_malicious_pickle_payloads``).
The fine-tuning resume path in ``scripts/train_voxcpm_finetune.py`` originally
called ``torch.load`` without that flag, leaving an arbitrary-code-execution
gap when resuming from an attacker-supplied checkpoint directory. This test
statically asserts the flag is present on every ``torch.load`` call across the
package and scripts so the gap cannot silently reappear.
"""
import ast
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[1]
# Directories whose Python files load checkpoints at runtime / on resume.
SCANNED_DIRS = [REPO_ROOT / "src", REPO_ROOT / "scripts", REPO_ROOT / "app.py", REPO_ROOT / "lora_ft_webui.py"]
def _python_files():
for entry in SCANNED_DIRS:
if entry.is_file() and entry.suffix == ".py":
yield entry
elif entry.is_dir():
yield from entry.rglob("*.py")
def _is_torch_load(node: ast.Call) -> bool:
func = node.func
# Matches ``torch.load(...)`` and ``load(...)`` aliased from torch.
if isinstance(func, ast.Attribute) and func.attr == "load":
return isinstance(func.value, ast.Name) and func.value.id == "torch"
return False
def _has_weights_only_true(node: ast.Call) -> bool:
for kw in node.keywords:
if kw.arg == "weights_only":
return isinstance(kw.value, ast.Constant) and kw.value.value is True
return False
def test_every_torch_load_sets_weights_only_true():
offenders = []
checked = 0
for path in _python_files():
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
for node in ast.walk(tree):
if isinstance(node, ast.Call) and _is_torch_load(node):
checked += 1
if not _has_weights_only_true(node):
offenders.append(f"{path.relative_to(REPO_ROOT)}:{node.lineno}")
assert checked > 0, "expected to find at least one torch.load call to verify"
assert not offenders, (
"torch.load without weights_only=True (pickle RCE risk):\n "
+ "\n ".join(offenders)
)
def test_torch_load_weights_only_blocks_malicious_pickle(tmp_path):
"""Behavioral check that weights_only=True actually rejects a code-exec payload."""
torch = pytest.importorskip("torch")
marker = tmp_path / "pwned.txt"
class Exploit:
def __reduce__(self):
import pathlib
return (pathlib.Path.write_text, (marker, "executed\n"))
ckpt = tmp_path / "optimizer.pth"
torch.save({"state_dict": {"w": torch.zeros(1)}, "boom": Exploit()}, ckpt)
with pytest.raises(Exception):
torch.load(ckpt, map_location="cpu", weights_only=True)
assert not marker.exists(), "malicious pickle executed despite weights_only=True"
tests/test_validate.py
+254
View File
@@ -0,0 +1,254 @@
"""Tests for the training data validation module."""
from __future__ import annotations
import json
import os
import sys
import tempfile
import types
from pathlib import Path
import pytest
ROOT = Path(__file__).resolve().parents[1]
# Stub voxcpm package so imports work without full dependencies
pkg = types.ModuleType("voxcpm")
pkg.__path__ = [str(ROOT / "src" / "voxcpm")]
sys.modules.setdefault("voxcpm", pkg)
training_pkg = types.ModuleType("voxcpm.training")
training_pkg.__path__ = [str(ROOT / "src" / "voxcpm" / "training")]
sys.modules.setdefault("voxcpm.training", training_pkg)
from voxcpm.training.validate import ValidationResult, validate_manifest
@pytest.fixture
def tmp_dir():
with tempfile.TemporaryDirectory() as d:
yield Path(d)
def _create_wav(path: Path, duration_s: float = 1.0, sr: int = 16000):
"""Create a minimal valid WAV file."""
try:
import soundfile as sf
import numpy as np
samples = int(duration_s * sr)
data = np.zeros(samples, dtype=np.float32)
sf.write(str(path), data, sr)
except ImportError:
# If soundfile is not available, create a minimal WAV header
import struct
samples = int(duration_s * sr)
data_size = samples * 2 # 16-bit PCM
with open(path, "wb") as f:
f.write(b"RIFF")
f.write(struct.pack("<I", 36 + data_size))
f.write(b"WAVEfmt ")
f.write(struct.pack("<IHHIIHH", 16, 1, 1, sr, sr * 2, 2, 16))
f.write(b"data")
f.write(struct.pack("<I", data_size))
f.write(b"\x00" * data_size)
def _write_manifest(path: Path, entries: list[dict]):
with open(path, "w", encoding="utf-8") as f:
for entry in entries:
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
class TestValidateManifest:
def test_valid_manifest(self, tmp_dir):
audio1 = tmp_dir / "audio1.wav"
audio2 = tmp_dir / "audio2.wav"
_create_wav(audio1, 2.0)
_create_wav(audio2, 3.0)
manifest = tmp_dir / "train.jsonl"
_write_manifest(
manifest,
[
{"text": "Hello world", "audio": str(audio1)},
{"text": "Goodbye world", "audio": str(audio2)},
],
)
result = validate_manifest(str(manifest))
assert result.total_samples == 2
assert result.valid_samples == 2
assert result.is_valid
assert len(result.errors) == 0
def test_missing_manifest(self):
result = validate_manifest("/nonexistent/path.jsonl")
assert not result.is_valid
assert any("not found" in e for e in result.errors)
def test_empty_manifest(self, tmp_dir):
manifest = tmp_dir / "empty.jsonl"
manifest.write_text("")
result = validate_manifest(str(manifest))
assert not result.is_valid
def test_invalid_json(self, tmp_dir):
manifest = tmp_dir / "bad.jsonl"
manifest.write_text("not json\n{bad json}\n")
result = validate_manifest(str(manifest))
assert len(result.errors) >= 2
assert any("Invalid JSON" in e for e in result.errors)
def test_missing_columns(self, tmp_dir):
manifest = tmp_dir / "missing.jsonl"
_write_manifest(
manifest,
[
{"text": "hello"}, # missing audio
{"audio": "test.wav"}, # missing text
],
)
result = validate_manifest(str(manifest))
assert len(result.errors) >= 2
assert any("'audio'" in e for e in result.errors)
assert any("'text'" in e for e in result.errors)
def test_missing_audio_file(self, tmp_dir):
manifest = tmp_dir / "missing_audio.jsonl"
_write_manifest(
manifest,
[{"text": "hello", "audio": "/nonexistent/audio.wav"}],
)
result = validate_manifest(str(manifest))
assert not result.is_valid
assert any("not found" in e for e in result.errors)
def test_empty_text_warning(self, tmp_dir):
audio = tmp_dir / "audio.wav"
_create_wav(audio)
manifest = tmp_dir / "empty_text.jsonl"
_write_manifest(
manifest,
[{"text": "", "audio": str(audio)}],
)
result = validate_manifest(str(manifest))
assert len(result.warnings) > 0
assert any("Empty" in w for w in result.warnings)
def test_relative_audio_path(self, tmp_dir):
audio = tmp_dir / "audio.wav"
_create_wav(audio)
manifest = tmp_dir / "rel.jsonl"
_write_manifest(
manifest,
[{"text": "hello", "audio": "audio.wav"}],
)
result = validate_manifest(str(manifest))
assert result.valid_samples == 1
assert result.is_valid
def test_max_samples_limit(self, tmp_dir):
audio = tmp_dir / "audio.wav"
_create_wav(audio)
manifest = tmp_dir / "many.jsonl"
_write_manifest(
manifest,
[{"text": f"sample {i}", "audio": str(audio)} for i in range(100)],
)
result = validate_manifest(str(manifest), max_samples=10)
assert result.total_samples == 10
def test_ref_audio_counted(self, tmp_dir):
audio = tmp_dir / "audio.wav"
ref = tmp_dir / "ref.wav"
_create_wav(audio)
_create_wav(ref)
manifest = tmp_dir / "ref.jsonl"
_write_manifest(
manifest,
[{"text": "hello", "audio": str(audio), "ref_audio": str(ref)}],
)
result = validate_manifest(str(manifest))
assert result.has_ref_audio == 1
def test_validation_result_properties(self):
r = ValidationResult(total_samples=5, valid_samples=5)
assert r.is_valid
r2 = ValidationResult(total_samples=5, valid_samples=5, errors=["err"])
assert not r2.is_valid
r3 = ValidationResult(total_samples=0, valid_samples=0)
assert not r3.is_valid
def test_invalid_audio_not_counted_as_valid(self, tmp_dir):
"""A row with a bad audio path must not increment valid_samples."""
manifest = tmp_dir / "bad_audio.jsonl"
_write_manifest(
manifest,
[{"text": "hello", "audio": "/nonexistent/audio.wav"}],
)
result = validate_manifest(str(manifest))
assert result.total_samples == 1
assert result.valid_samples == 0
assert not result.is_valid
assert any("not found" in e for e in result.errors)
def test_sample_rate_mismatch(self, tmp_dir):
"""A file with a different sample rate should be reported as an error."""
try:
import soundfile as sf
import numpy as np
except ImportError:
pytest.skip("soundfile not available")
audio = tmp_dir / "audio_8k.wav"
import numpy as np
samples = np.zeros(8000, dtype=np.float32)
sf.write(str(audio), samples, 8000)
manifest = tmp_dir / "sr_mismatch.jsonl"
_write_manifest(manifest, [{"text": "hello", "audio": str(audio)}])
result = validate_manifest(str(manifest), sample_rate=16000)
assert result.valid_samples == 0
assert not result.is_valid
assert any("Sample rate mismatch" in e or "sample rate" in e.lower() for e in result.errors)
def test_mixed_ref_audio_warns_for_each_missing(self, tmp_dir):
"""Missing ref_audio entries should each generate a warning independently."""
audio = tmp_dir / "audio.wav"
ref_good = tmp_dir / "ref_good.wav"
_create_wav(audio)
_create_wav(ref_good)
manifest = tmp_dir / "mixed_ref.jsonl"
_write_manifest(
manifest,
[
{"text": "row1", "audio": str(audio), "ref_audio": str(ref_good)},
{"text": "row2", "audio": str(audio), "ref_audio": "/nonexistent/ref.wav"},
],
)
result = validate_manifest(str(manifest))
assert result.has_ref_audio == 1
assert any("ref_audio file not found" in w for w in result.warnings)
def test_cli_validate_exit_code(self, tmp_dir):
"""validate subcommand must exit 1 on validation error (missing audio)."""
import subprocess
manifest = tmp_dir / "bad.jsonl"
_write_manifest(manifest, [{"text": "hi", "audio": "/nonexistent/x.wav"}])
proc = subprocess.run(
[sys.executable, "-m", "voxcpm.cli", "validate", "--manifest", str(manifest)],
capture_output=True,
text=True,
)
assert proc.returncode == 1, f"Expected exit 1, got {proc.returncode}"
assert "FAILED" in proc.stderr or "Audio file not found" in proc.stderr
uv.lock
Generated
+5263
View File
File diff suppressed because it is too large Load Diff