Skip to content

Add VoxCPM2 TTS model (2B params, 48kHz, voice cloning) - #175

Open
sb1992 wants to merge 8 commits into
Blaizzy:mainfrom
sb1992:voxcpm2-upstream
Open

Add VoxCPM2 TTS model (2B params, 48kHz, voice cloning)#175
sb1992 wants to merge 8 commits into
Blaizzy:mainfrom
sb1992:voxcpm2-upstream

Conversation

@sb1992

@sb1992 sb1992 commented May 1, 2026

Copy link
Copy Markdown

Swift port of VoxCPM2, based on the Python implementation in Blaizzy/mlx-audio#641.

Autoregressive MiniCPM backbone + CFM diffusion + AudioVAE decoder.
Supports zero-shot generation and reference-audio voice cloning across 30+ languages.

Supported model

Example

import MLXAudioTTS
import MLXAudioCore

let model = try await VoxCPM2Model.fromPretrained("mlx-community/VoxCPM2-4bit")

// Voice cloning
let (_, refAudio) = try loadAudioArray(from: referenceAudioURL)
let audio = try await model.generate(
    text: "Hello, this is VoxCPM2.",
    voice: nil, refAudio: refAudio, refText: nil, language: nil,
    generationParameters: GenerateParameters(maxTokens: 100, temperature: 1.0)
)

Smoke test results (M5 Max, macOS 26.3.2)

Test Audio Duration Peak Memory Time
voxcpm2Generate() 92,160 samples @ 48kHz 1.92s 3.72 GB 6.96s
voxcpm2GenerateStream() 69,120 samples @ 48kHz 1.44s 2.03s

Both produce valid mono Float32 WAV at 48kHz.

Known limitations

  • 4-bit stop predictor can be unreliable — callers should set maxTokens conservatively
  • Short Chinese prompts may have text-following issues (consistent with Python, see mlx-audio#641)
  • Streaming wraps the full generate() call rather than yielding incremental chunks
  • Voice design and continuation modes from the Python PR are not yet ported

Checklist

  • Conforms to SpeechGenerationModel
  • Factory registration in TTSModel.swift
  • Smoke tests for generate + stream
  • README.md with usage and limitations
  • swift build clean (zero warnings)
  • Package.swift README exclude

Shraey and others added 6 commits May 2, 2026 03:15
Autoregressive MiniCPM backbone + CFM diffusion + AudioVAE decoder.
Ported from mlx-audio Python voxcpm2 implementation with:
- Config-driven generation constants (inferenceTimesteps, cfgScale)
- Resilient Codable decoding with defaults for all config structs
- 48kHz→16kHz resampling for encoder input
- BOS/EOS stripping in tokenizer to match upstream Python behavior
- Proper weight sanitization (snake_case→camelCase, RoPE buffers)
- Smoke tests for generate and stream paths
- Qualify Tokenizer as Tokenizers.Tokenizer to resolve ambiguity
  with MLXLMCommon.Tokenizer (build blocker)
- Remove unused variable (MiniCPM.swift), change var→let where
  not mutated (VoxAudioVAE.swift, VoxCPM2Model.swift)
- Add README.md with usage examples, streaming behavior, output
  format, and known limitations
Matches the existing exclude pattern for all other model READMEs.
- Broaden eval() to include predFeat, breaking MLX lazy graph chains
  that accumulate across 2000 AR steps and cause OOM
- Add Task.checkCancellation() in generation loop for cooperative cancel
- Add reserveCapacity for predFeatSeq and KV cache arrays
- Cache SinusoidalPosEmb freqs in init (was recomputed ~20k times/gen)
- Make rope, specialToken immutable (let) where never mutated after init
- Make tokenizer a let property injected via init (eliminates data race)
- Store Task handle in generateStream and hook onTermination for cancellation
- Replace Array(repeating:) with MLX.broadcast in Euler solver (avoids CPU alloc per ODE step)
- Add unmapped weight key logging in sanitize() (excludes quantization artifacts)
- Promote NSRegularExpression to static let in VoxAudioVAE sanitize
@sb1992

sb1992 commented May 7, 2026

Copy link
Copy Markdown
Author

Pushed a small compatibility fix for the failed CI build.

The failure was from Xcode 26.3 rejecting Float * MLXArray in VoxCPMDiT.swift. I reproduced the failure locally with Xcode 26.3 (17C529), changed the scalar to MLXArray(scale), and verified xcodebuild build-for-testing now succeeds with the same Xcode version.

Shraey and others added 2 commits May 7, 2026 22:29
Resolve TTSModel.swift by keeping upstream's loader refactor and MOSS-TTS registration, then adding VoxCPM2 registration in the same load helper style.\n\nVerified with Xcode 26.3 build-for-testing.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants