diff --git a/CMakeLists.txt b/CMakeLists.txt index 313f05b7e..1ffdb9418 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1469,6 +1469,24 @@ audiocpp_add_model(chatterbox engine::models::chatterbox::make_chatterbox_loader ) +audiocpp_add_model(chatterbox_turbo + SOURCES + src/community_models/chatterbox_turbo/t3_turbo_weights.cpp + src/community_models/chatterbox_turbo/t3_turbo_component.cpp + src/community_models/chatterbox_turbo/s3gen_turbo.cpp + src/community_models/chatterbox_turbo/text_tokenizer_turbo.cpp + src/community_models/chatterbox_turbo/assets.cpp + src/community_models/chatterbox_turbo/loader.cpp + src/community_models/chatterbox_turbo/session.cpp + src/community_models/chatterbox_turbo/tts.cpp + INCLUDES + engine/community_models/chatterbox_turbo/loader.h + LOADERS + engine::community_models::chatterbox_turbo::make_chatterbox_turbo_loader + DEPENDS + chatterbox +) + audiocpp_add_model(ace_step SOURCES src/models/ace_step/assets.cpp diff --git a/README.md b/README.md index 446d157fb..dab73be2b 100644 --- a/README.md +++ b/README.md @@ -145,6 +145,7 @@ Community model ports live under `community_models` to make the ownership bounda | Family | Task | Lang | Runtime | Contributor | What They Added | |---|---|---|---|---|---| | **audio8_asr** | ASR | en, zh, yue, ja, ko, fr, de | GGUF Q8, Safetensors | [@0xShug0](https://github.com/0xShug0) | [Audio8-ASR-0.1B](docs/community_models/audio8_asr.md) compact multilingual autoregressive ASR reusing the Qwen3-ASR encoder with an MLP-tower adapter and an 8-layer Qwen2-style decoder (CC-BY-NC, local conversion only) | +| **chatterbox_turbo** | TTS, Clone (testing) | en | GGUF 16/Q8 | [@pannagaps](https://github.com/pannagaps) | [Chatterbox Turbo](docs/community_models/chatterbox_turbo.md) distilled 350M GPT2 T3 backbone + 2-step meanflow S3Gen decoder; built-in voice only for now | | **f5_tts** | TTS, Clone | en, ar (Habibi) | GGUF | [@tareko](https://github.com/tareko) | [F5-TTS](docs/community_models/f5_tts.md) flow-matching DiT synthesis and voice cloning, with Habibi Arabic aliases `habibi`/`habibi_tts` | | **glm_tts** | TTS, Clone | zh, en | GGUF | Mirek [@mirek190](https://github.com/mirek190) | [GLM-TTS](docs/community_models/glm_tts.md) zero-shot synthesis and voice cloning support | | **granite5asr** | ASR | en | GGUF Q8 | [@ampersandru](https://github.com/ampersandru) | [IBM Granite Speech 5.0 470M TurboCTC](docs/community_models/granite5asr.md) ultra-fast Conformer-CTC ASR with Shaw relative positional embeddings and ByteLevel BPE | diff --git a/docs/community_models/chatterbox_turbo.md b/docs/community_models/chatterbox_turbo.md new file mode 100644 index 000000000..2ed43d89c --- /dev/null +++ b/docs/community_models/chatterbox_turbo.md @@ -0,0 +1,84 @@ +# Chatterbox Turbo (community model) + +[Chatterbox Turbo](https://huggingface.co/ResembleAI/chatterbox-turbo) is Resemble AI's +distilled 350M-parameter sibling of Chatterbox (see [the Chatterbox section in docs/tts.md](../tts.md#chatterbox)): a GPT2-style T3 +backbone (vs. the base model's 0.5B Llama-style backbone), a GPT2 BPE tokenizer with 19 built-in +emotion/style tags (`[laugh]`, `[sigh]`, ...), and a 2-step meanflow-distilled S3Gen decoder (vs. +the base model's 10-step CFG decoder) for substantially faster generation. English-only. + +**Status: testing.** The T3 backbone and the built-in default voice both load and generate +audio end to end. Custom voice cloning (a caller-supplied reference clip) is not implemented +yet — it depends on the checkpoint's speaker-encoder and S3-tokenizer sections, whose tensor +layout has not been validated. This family lives under `community_models` rather than the core +model tree because it does not yet have the CUDA/Vulkan/Metal runtime test coverage core models +carry. + +## Packaging: audio.cpp-native, self-contained GGUF + +Resemble AI has not published Chatterbox Turbo weights in a format audio.cpp can convert +directly. The only available conversion is a **third-party GGUF**, +[`cstr/chatterbox-turbo-GGUF`](https://huggingface.co/cstr/chatterbox-turbo-GGUF), published by +`cstr` for their own CrispASR project (MIT-relicensed) — not published by ResembleAI or +audio.cpp. It ships as two loose GGUF files (T3 and S3Gen) with a flat dot-separated tensor +namespace and abbreviated S3Gen tensor names that don't match this codebase's own naming. + +Rather than teach the runtime a compatibility layer for that third-party layout, +[`tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py`](../../tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py) +repacks it offline into one self-contained, audio.cpp-native GGUF: + +- T3 and built-in-conditional tensors are moved from the upstream flat `t3.`/`conds.` dot + namespace into this project's own `/`-delimited packed-GGUF namespace convention. +- S3Gen tensors are renamed back to the exact names base Chatterbox's own S3Gen flow/HiFT-vocoder + loader (`src/models/chatterbox/s3gen_flow.cpp`, + `src/framework/modules/vocoders/hift_vocoder.cpp`) already expects, so that loader runs + completely unmodified for Turbo — no tensor-name translation code exists in this family at + runtime. +- The GPT2 BPE tokenizer (vocab, merges, and the trailing emotion/style special tokens) is + extracted into plain `vocab.json`/`merges.txt`/`special_tokens.json` sidecar files instead of + being read from raw GGUF metadata at load time. +- The result is fed through this project's own `audiocpp_gguf` converter, which quantizes, + embeds the package spec (`model_specs/chatterbox_turbo.json`), and embeds the sidecar files — + producing one file that loads with nothing else needed, like every other GGUF family here. + +The `ve.*` (LSTM speaker-verification voice encoder) and `s3.se.*`/`s3.tok.*` (ResNet speaker +encoder / S3 speech tokenizer) sections of the upstream checkpoint are not repacked: nothing in +this codebase reads them yet (see Current limitations above). + +### Repacking it yourself + +```bash +# 1. Build the converter +cmake --build build/debug --parallel --target audiocpp_gguf + +# 2. Get the upstream third-party GGUF pair (~1 GB for Q8_0) +huggingface-cli download cstr/chatterbox-turbo-GGUF \ + chatterbox-turbo-t3-q8_0.gguf chatterbox-turbo-s3gen-q8_0.gguf \ + --local-dir /tmp/chatterbox-turbo-src + +# 3. Repack +pip install gguf numpy safetensors +python3 tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py \ + --t3-source /tmp/chatterbox-turbo-src/chatterbox-turbo-t3-q8_0.gguf \ + --s3gen-source /tmp/chatterbox-turbo-src/chatterbox-turbo-s3gen-q8_0.gguf \ + --output models/Chatterbox-Turbo-GGUF/chatterbox-turbo-q8_0.gguf \ + --type q8_0 --overwrite +``` + +Verify with `build/debug/bin/audiocpp_gguf --inspect models/Chatterbox-Turbo-GGUF/chatterbox-turbo-q8_0.gguf` (expect `embedded_sidecars=true`, `embedded_model_spec=true`, and `t3`/`conds`/`s3gen` namespaces). + +## Usage + +```bash +audiocpp_cli --task clon --family chatterbox_turbo \ + --model models/Chatterbox-Turbo-GGUF/chatterbox-turbo-q8_0.gguf \ + --backend cuda --text "Hello from Chatterbox Turbo." --out out.wav +``` + +See the [Chatterbox Turbo section in docs/tts.md](../tts.md#chatterbox-turbo) for the full option +table. + +## Checkpoints + +| Model | Source | License | +|---|---|---| +| Chatterbox Turbo (T3 + S3Gen) | `cstr/chatterbox-turbo-GGUF` (third-party repack of `ResembleAI/chatterbox-turbo`) | MIT | diff --git a/docs/community_models/models.md b/docs/community_models/models.md index fc6bc07cd..994556b52 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -17,6 +17,7 @@ Practical expectations: | Family | Task | Supported language(s) | Contributor | What They Added | |---|---|---|---|---| | **audio8_asr** | ASR | en, zh, yue, ja, ko, fr, de | [@0xShug0](https://github.com/0xShug0) | [Audio8-ASR-0.1B](audio8_asr.md) compact multilingual autoregressive ASR reusing the Qwen3-ASR encoder with an MLP-tower adapter and an 8-layer Qwen2-style decoder (CC-BY-NC, local conversion only) | +| **chatterbox_turbo** | TTS, voice cloning (testing) | en | [@pannagaps](https://github.com/pannagaps) | [Chatterbox Turbo](chatterbox_turbo.md) Resemble AI's distilled 350M GPT2 T3 backbone + 2-step meanflow S3Gen decoder for fast English TTS; built-in default voice only for now | | **echo_tts** | TTS, voice cloning | en | Tym [@5uck1ess](https://github.com/5uck1ess), [@dignome](https://github.com/dignome) | [Echo-TTS](echo_tts.md) 44.1 kHz zero-shot voice cloning: 2.8B diffusion transformer in 80-D PCA space, decoded by the Fish S1-DAC autoencoder. Byte-level text, no phonemiser, no reference transcript | | **f5_tts** | TTS, voice cloning | en, ar (Habibi) | Community | [F5-TTS](f5_tts.md) flow-matching DiT — M0 scaffolding, aliases `habibi`/`habibi_tts` | | **glm_tts** | TTS, voice cloning | zh, en | Mirek [@mirek190](https://github.com/mirek190) | [GLM-TTS](glm_tts.md) zero-shot synthesis and voice cloning support | diff --git a/docs/tts.md b/docs/tts.md index eba1c7794..28e0818ce 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -100,6 +100,53 @@ audiocpp_cli --task vc --family chatterbox --model models/chatterbox --backend c | `--max-tokens` | integer | `1000` | Maximum generated T3 tokens per chunk. | | `--do-sample` | `true`, `false` | `true` | Enable stochastic T3 sampling. | +## Chatterbox Turbo + +Chatterbox Turbo is a [community model](community_models/chatterbox_turbo.md): Resemble AI's +distilled 350M-parameter sibling of Chatterbox, with a GPT2-style T3 backbone (vs. the base +model's 0.5B Llama-style backbone), a GPT2 BPE tokenizer with 19 built-in emotion/style tags +(`[laugh]`, `[sigh]`, ...), and a 2-step meanflow-distilled S3Gen decoder (vs. the base model's +10-step CFG decoder) for substantially faster generation. It is English-only. + +`chatterbox_turbo` is a separate model family from `chatterbox` (not a variant selectable within +it): its T3 backbone and tokenizer differ from the base model's, and it reuses base Chatterbox's +own S3Gen/HiFT-vocoder loader code for the flow decoder and vocoder half. + +The package is one self-contained, audio.cpp-native GGUF produced by repacking Resemble AI's +weights (via the third-party `cstr/chatterbox-turbo-GGUF` conversion published for the CrispASR +project, MIT-relicensed) with +[`tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py`](../tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py) +— see that model's community doc for details. + +**Current limitations:** only the built-in default voice baked into the package is supported — +custom voice cloning (a caller-supplied reference clip) is not implemented yet, since it depends +on the checkpoint's speaker-encoder and S3-tokenizer sections, whose exact tensor layout hasn't +been validated. `--voice-ref` is rejected with an explicit error rather than silently ignored. + +| Field | Value | +|---|---| +| Family | `chatterbox_turbo` | +| Model directory | `Chatterbox-Turbo-GGUF/chatterbox-turbo-{q8_0,f16}.gguf` (single self-contained file) | +| Tasks | `clon` (built-in voice only) | +| Modes | `offline` | +| Languages | `en` | +| Voice input | Not yet supported — omit `--voice-ref` to use the built-in voice | +| Built-in voices | One, embedded in the package | + +```bash +audiocpp_cli --task clon --family chatterbox_turbo --model models/Chatterbox-Turbo-GGUF/chatterbox-turbo-q8_0.gguf --backend cuda --text "Hello from Chatterbox Turbo." --out out.wav +``` + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--temperature` | float | `0.8` | T3 sampling temperature. | +| `--top-p` | float | `0.95` | T3 nucleus sampling limit. | +| `top_k` (session option) | integer | `1000` | T3 top-k sampling limit. | +| `--repetition-penalty` | float | `1.2` | T3 repetition penalty. | +| `--max-tokens` | integer | `1000` | Maximum generated T3 tokens. | + +`--guidance-scale`/exaggeration/min_p have no effect on Turbo (it was distilled without CFG) and are accepted but ignored, matching upstream's own behavior. + ## Confucius4-TTS Confucius4-TTS is an experimental multilingual voice-cloning TTS model packaged as a standalone GGUF bundle. It supports offline generation and streaming text input, using reference speech, language-aware text normalization, T2S semantic generation, S2A flow matching, style encoding, semantic audio features, and BigVGAN vocoding. diff --git a/include/engine/community_models/chatterbox_turbo/assets.h b/include/engine/community_models/chatterbox_turbo/assets.h new file mode 100644 index 000000000..628e88174 --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/assets.h @@ -0,0 +1,20 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" + +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +struct ChatterboxTurboAssets { + engine::assets::ResourceBundle resources; + std::shared_ptr t3_turbo_weights; + std::shared_ptr builtin_conditionals_turbo; + std::shared_ptr s3gen_weights; +}; + +std::shared_ptr load_chatterbox_turbo_assets(const std::filesystem::path & model_path); + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/loader.h b/include/engine/community_models/chatterbox_turbo/loader.h new file mode 100644 index 000000000..97d843a0f --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/loader.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/community_models/chatterbox_turbo/assets.h" + +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +class ChatterboxTurboLoadedModel final : public runtime::ILoadedVoiceModel { +public: + ChatterboxTurboLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets); + + const runtime::ModelMetadata & metadata() const noexcept override; + const runtime::CapabilitySet & capabilities() const noexcept override; + std::unique_ptr create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const override; + +private: + runtime::ModelMetadata metadata_; + runtime::CapabilitySet capabilities_; + std::shared_ptr assets_; +}; + +std::unique_ptr load_chatterbox_turbo_model(const std::filesystem::path & model_root); +std::shared_ptr make_chatterbox_turbo_loader(); + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h b/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h new file mode 100644 index 000000000..a5e0213da --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h @@ -0,0 +1,49 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/vocoders/hift_vocoder.h" +#include "engine/models/chatterbox/components.h" +#include "engine/models/chatterbox/s3gen_flow.h" +#include "engine/models/chatterbox/s3gen_inference.h" + +#include +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +// Loads and runs Chatterbox Turbo's S3Gen half (speech tokens -> waveform) by delegating to the +// existing chatterbox family's flow encoder/decoder and HiFT vocoder code (architecturally +// identical apart from the meanflow decoder branch already added to +// engine::models::chatterbox::S3FlowDecoderWeights). The native, repacked GGUF (see +// tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py) stores S3Gen tensors +// under the same names those loaders already expect, so no name translation is needed here. +// +// MVP scope: only the built-in default voice baked into the T3 GGUF's `conds.gen.*` tensors is +// supported (no custom voice cloning yet) -- that path needs the turbo checkpoint's `s3.se` +// (ResNet-style speaker encoder, not CAMPPlus) and `s3.tok` (S3 speech tokenizer) sections, which +// are unverified and out of scope for this pass. See ChatterboxTurboAssets. +class ChatterboxTurboS3Gen { +public: + static std::shared_ptr load( + std::shared_ptr s3gen_source, + const engine::core::ExecutionContext & execution_context, + engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); + + // speech_tokens: T3 output (S3 codebook ids, < 6561; caller strips control tokens). + engine::models::chatterbox::S3GenInferenceOutputs synthesize( + const engine::models::chatterbox::EmbedReferenceOutputs & ref_dict, + const std::vector & speech_tokens, + uint64_t flow_seed, + uint64_t vocoder_seed) const; + +private: + std::shared_ptr encoder_weights_; + std::shared_ptr decoder_weights_; + std::shared_ptr vocoder_; + mutable engine::models::chatterbox::S3GenSessionCache cache_; + const engine::core::ExecutionContext * execution_context_ = nullptr; +}; + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/session.h b/include/engine/community_models/chatterbox_turbo/session.h new file mode 100644 index 000000000..60094400f --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/session.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/runtime/session_base.h" +#include "engine/community_models/chatterbox_turbo/assets.h" +#include "engine/community_models/chatterbox_turbo/tts.h" + +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +class ChatterboxTurboSession final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + ChatterboxTurboSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + ~ChatterboxTurboSession() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::unique_ptr component_; +}; + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/t3_turbo_component.h b/include/engine/community_models/chatterbox_turbo/t3_turbo_component.h new file mode 100644 index 000000000..01a4fd97d --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/t3_turbo_component.h @@ -0,0 +1,36 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/community_models/chatterbox_turbo/t3_turbo_types.h" + +#include + +namespace engine::community_models::chatterbox_turbo { + +std::shared_ptr load_t3_turbo_inference_weights( + const engine::assets::TensorSource & source, + const engine::core::ExecutionContext & execution_context, + engine::assets::TensorStorageType graph_weight_storage_type = engine::assets::TensorStorageType::Native, + bool load_reference_f32_graph_weights = true); + +class T3TurboInferenceComponent { +public: + explicit T3TurboInferenceComponent( + std::shared_ptr weights, + const engine::core::ExecutionContext & execution_context); + + T3TurboGenerateOutputs generate_speech_tokens(const T3TurboGenerateRequest & request) const; + void release_runtime_graphs() const; + void release_runtime_cache() const; + +private: + struct State; + + std::shared_ptr weights_; + const engine::core::ExecutionContext * execution_context_ = nullptr; + std::shared_ptr state_; +}; + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/t3_turbo_types.h b/include/engine/community_models/chatterbox_turbo/t3_turbo_types.h new file mode 100644 index 000000000..20401f0b1 --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/t3_turbo_types.h @@ -0,0 +1,97 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/core/module.h" + +#include +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +struct T3TurboGraphWeight { + std::vector values; + engine::core::TensorValue tensor; +}; + +// GPT2-style T3 backbone used by Chatterbox Turbo: LayerNorm (weight+bias), +// fused-QKV causal self attention, GELU(tanh) MLP with bias, absolute +// position embedding ("wpe") applied host-side before the transformer stack. +// No RoPE, no perceiver resampler, no emotion-exaggeration conditioning. +struct T3TurboInferenceWeights { + int64_t hidden_size = 1024; + int64_t speaker_embed_size = 256; + int64_t num_heads = 16; + int64_t mlp_intermediate_size = 4096; + int64_t text_vocab = 0; + int64_t speech_vocab = 0; + int64_t max_positions = 0; + + struct TransformerLayer { + std::vector ln1_weight; + std::vector ln1_bias; + engine::core::TensorValue ln1_weight_tensor; + engine::core::TensorValue ln1_bias_tensor; + T3TurboGraphWeight attn_qkv_weight; + std::vector attn_qkv_bias; + engine::core::TensorValue attn_qkv_bias_tensor; + T3TurboGraphWeight attn_output_weight; + std::vector attn_output_bias; + engine::core::TensorValue attn_output_bias_tensor; + std::vector ln2_weight; + std::vector ln2_bias; + engine::core::TensorValue ln2_weight_tensor; + engine::core::TensorValue ln2_bias_tensor; + T3TurboGraphWeight ffn_fc_weight; + std::vector ffn_fc_bias; + engine::core::TensorValue ffn_fc_bias_tensor; + T3TurboGraphWeight ffn_proj_weight; + std::vector ffn_proj_bias; + engine::core::TensorValue ffn_proj_bias_tensor; + }; + + engine::assets::TensorData spkr_enc_weight; + engine::assets::TensorData spkr_enc_bias; + engine::assets::TensorData text_embedding_weight; + engine::assets::TensorData speech_embedding_weight; + engine::assets::TensorData wpe_weight; + std::vector ln_f_weight; + std::vector ln_f_bias; + engine::core::TensorValue ln_f_weight_tensor; + engine::core::TensorValue ln_f_bias_tensor; + T3TurboGraphWeight text_head_weight; + T3TurboGraphWeight speech_head_weight; + std::vector speech_head_bias; + engine::core::TensorValue speech_head_bias_tensor; + std::vector layers; + const engine::core::ExecutionContext * execution_context = nullptr; + std::shared_ptr store; +}; + +struct T3TurboGenerateRequest { + std::vector speaker_embedding; + std::vector cond_prompt_speech_tokens; + std::vector text_tokens; + std::vector initial_speech_tokens; + int64_t max_new_tokens = 1000; + bool stop_on_eos = true; + bool do_sample = true; + float temperature = 0.8f; + float top_p = 0.95f; + int64_t top_k = 1000; + float repetition_penalty = 1.2f; + uint32_t seed = 0; +}; + +struct T3TurboGenerateOutputs { + std::vector predicted_tokens; + int64_t token_count = 0; + bool hit_eos = false; + double prefix_cache_build_ms = 0.0; + double prefill_runner_ms = 0.0; + double decode_runner_ms = 0.0; +}; + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/text_tokenizer_turbo.h b/include/engine/community_models/chatterbox_turbo/text_tokenizer_turbo.h new file mode 100644 index 000000000..c4eadb0ba --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/text_tokenizer_turbo.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include +#include +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +// Builds a GPT2-style BPE tokenizer from the vocab.json/merges.txt/special_tokens.json sidecar +// files embedded in the repacked native GGUF (see +// tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py, which splits the +// upstream tokenizer.ggml.tokens/merges GGUF metadata arrays into these plain files at conversion +// time -- the trailing `[laugh]`/`[sigh]`/... emotion tags are pulled out into +// special_tokens_path so they are matched as atomic tokens rather than shredded into byte-level +// BPE pieces). +std::shared_ptr load_chatterbox_turbo_tokenizer( + const std::filesystem::path & vocab_path, + const std::filesystem::path & merges_path, + const std::filesystem::path & special_tokens_path); + +// Direct port of tts_turbo.py::punc_norm: capitalizes the first letter, collapses whitespace, +// replaces punctuation characters uncommon in the training data, and ensures a trailing +// sentence-ending punctuation mark. +std::string chatterbox_turbo_punc_norm(const std::string & text); + +std::vector encode_chatterbox_turbo_text( + const engine::tokenizers::LlamaBpeTokenizer & tokenizer, + const std::string & text); + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/community_models/chatterbox_turbo/tts.h b/include/engine/community_models/chatterbox_turbo/tts.h new file mode 100644 index 000000000..2d74bb1f9 --- /dev/null +++ b/include/engine/community_models/chatterbox_turbo/tts.h @@ -0,0 +1,52 @@ +#pragma once + +#include "engine/framework/core/execution_context.h" +#include "engine/framework/tokenizers/llama_bpe.h" +#include "engine/models/chatterbox/s3gen_inference.h" +#include "engine/community_models/chatterbox_turbo/assets.h" +#include "engine/community_models/chatterbox_turbo/s3gen_turbo.h" +#include "engine/community_models/chatterbox_turbo/t3_turbo_component.h" + +#include +#include +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +struct ChatterboxTurboGenerateConfig { + float temperature = 0.8f; + float top_p = 0.95f; + int64_t top_k = 1000; + float repetition_penalty = 1.2f; + int64_t max_new_tokens = 1000; + uint32_t seed = 0; + // exaggeration/cfg_weight/min_p from the base Chatterbox request surface are accepted but + // ignored for Turbo (matches upstream tts_turbo.py's logger.warning(...ignored...)); no + // fields here since this MVP path only drives the built-in voice (see ChatterboxTurboAssets). +}; + +// Orchestrates the full Turbo TTS pipeline (T3 GPT2 backbone -> S3Gen meanflow decoder -> HiFT +// vocoder) using only the built-in default voice baked into the T3 GGUF's `conds.*` tensors — +// custom voice cloning is not implemented yet (see s3gen_turbo.h). +class ChatterboxTurboTtsComponent { +public: + ChatterboxTurboTtsComponent( + std::shared_ptr assets, + const engine::core::ExecutionContext & execution_context); + + engine::models::chatterbox::S3GenInferenceOutputs generate( + const std::string & text, + const ChatterboxTurboGenerateConfig & config) const; + +private: + std::shared_ptr assets_; + std::shared_ptr tokenizer_; + std::unique_ptr t3_; + std::shared_ptr s3gen_; + std::vector builtin_speaker_embedding_; + std::vector builtin_cond_prompt_speech_tokens_; + engine::models::chatterbox::EmbedReferenceOutputs builtin_ref_dict_; +}; + +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/models/chatterbox/s3gen_flow.h b/include/engine/models/chatterbox/s3gen_flow.h index 94b47178e..0f003051b 100644 --- a/include/engine/models/chatterbox/s3gen_flow.h +++ b/include/engine/models/chatterbox/s3gen_flow.h @@ -50,7 +50,8 @@ class S3FlowSessionCache { int64_t batch, int64_t frames, int64_t capacity_frames, - engine::core::BackendConfig backend); + engine::core::BackendConfig backend, + const std::vector * r); friend S3FlowCFMOutputs compute_s3_flow_cfm_euler( S3FlowSessionCache & cache, const S3FlowDecoderWeights & weights, @@ -67,6 +68,19 @@ class S3FlowSessionCache { bool cosine_schedule, engine::core::BackendConfig backend, S3FlowCFMTimingBreakdown * timing); + friend S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( + S3FlowSessionCache & cache, + const S3FlowDecoderWeights & weights, + const std::vector & noise, + const std::vector & mask, + const std::vector & mu, + const std::vector & spks, + const std::vector & cond, + int64_t batch, + int64_t frames, + int64_t capacity_frames, + int64_t num_steps, + engine::core::BackendConfig backend); }; struct S3FlowDecoderRunTiming { @@ -107,6 +121,12 @@ std::shared_ptr load_s3_flow_decoder_weights( const engine::assets::TensorSource & source, const engine::core::ExecutionContext & execution_context, engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); +// S3FlowDecoderWeights is intentionally opaque outside this family; this accessor lets callers +// (e.g. chatterbox_turbo) detect whether a loaded decoder is meanflow-distilled without needing +// the struct's internal layout. +bool s3_flow_decoder_is_meanflow(const S3FlowDecoderWeights & weights); +// r: meanflow end-time input (see S3FlowDecoderWeights::meanflow); nullptr for the base +// Chatterbox 10-step CFG decoder, which never reads it. S3FlowDecoderOutputs compute_s3_flow_decoder_forward( S3FlowSessionCache & cache, const S3FlowDecoderWeights & weights, @@ -119,7 +139,8 @@ S3FlowDecoderOutputs compute_s3_flow_decoder_forward( int64_t batch, int64_t frames, int64_t capacity_frames, - engine::core::BackendConfig backend = {}); + engine::core::BackendConfig backend = {}, + const std::vector * r = nullptr); S3FlowCFMOutputs compute_s3_flow_cfm_euler( S3FlowSessionCache & cache, const S3FlowDecoderWeights & weights, @@ -137,4 +158,20 @@ S3FlowCFMOutputs compute_s3_flow_cfm_euler( engine::core::BackendConfig backend = {}, S3FlowCFMTimingBreakdown * timing = nullptr); +// Non-CFG, non-batch-doubled Euler solve for meanflow-distilled decoders (Chatterbox Turbo); +// see S3FlowDecoderWeights::meanflow. num_steps defaults to 2 in upstream Python (tts_turbo.py). +S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( + S3FlowSessionCache & cache, + const S3FlowDecoderWeights & weights, + const std::vector & noise, + const std::vector & mask, + const std::vector & mu, + const std::vector & spks, + const std::vector & cond, + int64_t batch, + int64_t frames, + int64_t capacity_frames, + int64_t num_steps = 2, + engine::core::BackendConfig backend = {}); + } // namespace engine::models::chatterbox diff --git a/model_specs/chatterbox_turbo.json b/model_specs/chatterbox_turbo.json new file mode 100644 index 000000000..13cd446ab --- /dev/null +++ b/model_specs/chatterbox_turbo.json @@ -0,0 +1,100 @@ +{ + "family": "chatterbox_turbo", + "display_name": "Chatterbox Turbo", + "description": "Resemble AI's distilled 350M Chatterbox Turbo: GPT2 T3 backbone, GPT2 BPE tokenizer, and a 2-step meanflow S3Gen decoder for fast English TTS and zero-shot voice cloning.", + "category": "tts", + "status": "testing", + "tasks": [ + "tts", + "clone" + ], + "modes": [ + "offline" + ], + "languages": [ + "en" + ], + "capabilities": { + "clone": [ + "speaker_reference" + ] + }, + "runtime": { + "tags": [ + "gguf" + ] + }, + "ui": { + "recommended_package": "chatterbox_turbo_q8_0", + "tags": [ + "TTS", + "Clone", + "GGUF" + ], + "docs": [ + "docs/tts.md", + "docs/gguf.md" + ] + }, + "package_defaults": { + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/audio.cpp-gguf", + "revision": "main", + "gated": false + } + }, + "packages": [ + { + "id": "chatterbox_turbo_q8_0", + "display_name": "Chatterbox Turbo Q8_0 GGUF", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "Chatterbox-Turbo-GGUF", + "files": [ + "Chatterbox-Turbo-GGUF/chatterbox-turbo-q8_0.gguf" + ], + "strip_prefix": "Chatterbox-Turbo-GGUF" + }, + { + "id": "chatterbox_turbo_f16", + "display_name": "Chatterbox Turbo F16 GGUF", + "format": "gguf", + "precision": "f16", + "target_directory": "Chatterbox-Turbo-GGUF", + "files": [ + "Chatterbox-Turbo-GGUF/chatterbox-turbo-f16.gguf" + ], + "strip_prefix": "Chatterbox-Turbo-GGUF" + } + ], + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "tokenizer_vocab": "model:chatterbox_turbo_vocab.json", + "tokenizer_merges": "model:chatterbox_turbo_merges.txt", + "tokenizer_special_tokens": "model:chatterbox_turbo_special_tokens.json" + }, + "tensors": { + "t3_turbo_weights": { + "source": "weights:", + "prefix": "t3" + }, + "builtin_conditionals_turbo": { + "source": "weights:", + "prefix": "conds" + }, + "s3gen_weights": { + "source": "weights:", + "prefix": "s3gen" + } + } + } + ] +} diff --git a/src/community_models/chatterbox_turbo/assets.cpp b/src/community_models/chatterbox_turbo/assets.cpp new file mode 100644 index 000000000..4185efce5 --- /dev/null +++ b/src/community_models/chatterbox_turbo/assets.cpp @@ -0,0 +1,18 @@ +#include "engine/community_models/chatterbox_turbo/assets.h" + +#include "engine/framework/model_spec/package.h" + +namespace engine::community_models::chatterbox_turbo { + +std::shared_ptr load_chatterbox_turbo_assets(const std::filesystem::path & model_path) { + auto out = std::make_shared(); + out->resources = engine::model_spec::load_resource_bundle( + model_path, + engine::model_spec::default_spec_path("chatterbox_turbo")); + out->t3_turbo_weights = out->resources.open_tensor_source("t3_turbo_weights"); + out->builtin_conditionals_turbo = out->resources.open_tensor_source("builtin_conditionals_turbo"); + out->s3gen_weights = out->resources.open_tensor_source("s3gen_weights"); + return out; +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/components/t3_turbo_runtime.h b/src/community_models/chatterbox_turbo/components/t3_turbo_runtime.h new file mode 100644 index 000000000..27b263986 --- /dev/null +++ b/src/community_models/chatterbox_turbo/components/t3_turbo_runtime.h @@ -0,0 +1,811 @@ +#pragma once + +#include "engine/community_models/chatterbox_turbo/t3_turbo_component.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::chatterbox_turbo { +namespace { + +constexpr int32_t kTurboStartSpeechToken = 6561; +constexpr int32_t kTurboStopSpeechToken = 6562; +constexpr int32_t kTurboMaxSpeechToken = 6560; // tokens >= kTurboStartSpeechToken are control tokens + +std::vector gather_rows( + const std::vector & table, + int64_t rows, + int64_t cols, + const std::vector & indices) { + std::vector out(static_cast(indices.size() * static_cast(cols)), 0.0f); + for (size_t i = 0; i < indices.size(); ++i) { + const int32_t index = indices[i]; + if (index < 0 || index >= rows) { + throw std::runtime_error("chatterbox_turbo embedding index out of range"); + } + const float * src = table.data() + static_cast(static_cast(index) * cols); + float * dst = out.data() + static_cast(i * static_cast(cols)); + std::copy(src, src + cols, dst); + } + return out; +} + +void softmax_into(const std::vector & logits, std::vector & probs) { + float max_value = -std::numeric_limits::infinity(); + for (float value : logits) { + max_value = std::max(max_value, value); + } + probs.assign(logits.size(), 0.0f); + double sum = 0.0; + for (size_t i = 0; i < logits.size(); ++i) { + probs[i] = std::exp(logits[i] - max_value); + sum += static_cast(probs[i]); + } + for (float & value : probs) { + value = static_cast(static_cast(value) / sum); + } +} + +std::vector softmax(const std::vector & logits) { + std::vector probs; + softmax_into(logits, probs); + return probs; +} + +void apply_repetition_penalty_in_place( + std::vector & logits, + const std::vector & generated_ids, + float penalty, + std::vector & seen) { + if (penalty == 1.0f) { + return; + } + seen.assign(logits.size(), 0); + for (int32_t token : generated_ids) { + if (token < 0 || token >= static_cast(logits.size())) { + continue; + } + if (seen[static_cast(token)] != 0) { + continue; + } + seen[static_cast(token)] = 1; + float & value = logits[static_cast(token)]; + value = value < 0.0f ? value * penalty : value / penalty; + } +} + +void apply_top_k_in_place(std::vector & logits, int64_t top_k) { + if (top_k <= 0 || top_k >= static_cast(logits.size())) { + return; + } + std::vector sorted = logits; + std::nth_element(sorted.begin(), sorted.begin() + static_cast(top_k - 1), sorted.end(), std::greater()); + const float threshold = sorted[static_cast(top_k - 1)]; + for (float & value : logits) { + if (value < threshold) { + value = -std::numeric_limits::infinity(); + } + } +} + +void apply_top_p_in_place( + std::vector & logits, + float top_p, + std::vector & probs, + std::vector & order, + std::vector & remove) { + if (top_p >= 1.0f) { + return; + } + softmax_into(logits, probs); + order.resize(probs.size()); + for (size_t i = 0; i < order.size(); ++i) { + order[i] = i; + } + std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { + return probs[a] > probs[b]; + }); + remove.assign(probs.size(), 0); + float cumulative = 0.0f; + for (auto it = order.rbegin(); it != order.rend(); ++it) { + const size_t index = *it; + cumulative += probs[index]; + if (cumulative <= (1.0f - top_p)) { + remove[index] = 1; + } + } + if (!order.empty()) { + remove[order[0]] = 0; + } + for (size_t i = 0; i < logits.size(); ++i) { + if (remove[i] != 0) { + logits[i] = -std::numeric_limits::infinity(); + } + } +} + +class TurboMt19937 { +public: + explicit TurboMt19937(uint32_t seed) + : state_{}, left_(1), next_(0) { + state_[0] = seed; + for (size_t index = 1; index < state_.size(); ++index) { + state_[index] = static_cast( + 1812433253U * (state_[index - 1] ^ (state_[index - 1] >> 30U)) + static_cast(index)); + } + } + + uint32_t random() { + if (--left_ == 0) { + next_state(); + } + uint32_t value = state_[next_++]; + value ^= (value >> 11U); + value ^= (value << 7U) & 0x9d2c5680U; + value ^= (value << 15U) & 0xefc60000U; + value ^= (value >> 18U); + return value; + } + + uint64_t random64() { + const uint64_t high = static_cast(random()); + const uint64_t low = static_cast(random()); + return (high << 32U) | low; + } + +private: + static constexpr size_t kStateSize = 624; + static constexpr size_t kStateM = 397; + static constexpr uint32_t kMatrixA = 0x9908b0dfU; + static constexpr uint32_t kUpperMask = 0x80000000U; + static constexpr uint32_t kLowerMask = 0x7fffffffU; + + static uint32_t mix_bits(uint32_t first, uint32_t second) { + return (first & kUpperMask) | (second & kLowerMask); + } + static uint32_t twist(uint32_t first, uint32_t second) { + return (mix_bits(first, second) >> 1U) ^ ((second & 1U) ? kMatrixA : 0U); + } + void next_state() { + size_t offset = 0; + left_ = static_cast(kStateSize); + next_ = 0; + for (size_t count = 0; count < (kStateSize - kStateM); ++count, ++offset) { + state_[offset] = state_[offset + kStateM] ^ twist(state_[offset], state_[offset + 1]); + } + for (size_t count = 0; count < (kStateM - 1); ++count, ++offset) { + state_[offset] = state_[offset + kStateM - kStateSize] ^ twist(state_[offset], state_[offset + 1]); + } + state_[offset] = state_[offset + kStateM - kStateSize] ^ twist(state_[offset], state_[0]); + } + + std::array state_; + int left_; + size_t next_; +}; + +double uniform_double(TurboMt19937 & rng) { + constexpr uint64_t kMask = (static_cast(1) << std::numeric_limits::digits) - 1U; + constexpr double kDivisor = + 1.0 / static_cast(static_cast(1) << std::numeric_limits::digits); + return static_cast(rng.random64() & kMask) * kDivisor; +} + +int32_t sample_from_logits(const std::vector & logits, bool do_sample, TurboMt19937 & rng) { + if (!do_sample) { + return static_cast(std::distance(logits.begin(), std::max_element(logits.begin(), logits.end()))); + } + const auto probs = softmax(logits); + const double draw = uniform_double(rng); + double cumulative = 0.0; + for (size_t index = 0; index < probs.size(); ++index) { + cumulative += static_cast(probs[index]); + if (draw < cumulative) { + return static_cast(index); + } + } + return static_cast(probs.empty() ? 0 : (probs.size() - 1)); +} + +core::TensorValue contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + return core::ensure_backend_addressable_layout(ctx, input); +} + +core::TensorValue make_graph_param_tensor(const T3TurboGraphWeight & weight) { + return weight.tensor; +} + +// GPT2 pre-LN transformer block. hidden_size/num_heads/head_dim/mlp are read from the layer's +// tensor shapes at call sites; no RoPE, no position tensor (positions are summed into the +// input embeddings host-side before this graph runs, matching how GPT2's own wpe is applied +// once to combined inputs_embeds rather than per attention layer). +struct T3TurboLayerOutput { + core::TensorValue hidden; + core::TensorValue key; + core::TensorValue value; +}; + +core::TensorValue turbo_gpt_mlp( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t hidden_size, + int64_t intermediate_size, + const T3TurboInferenceWeights::TransformerLayer & layer) { + auto hidden = modules::LinearModule({hidden_size, intermediate_size, true}).build( + ctx, input, {make_graph_param_tensor(layer.ffn_fc_weight), layer.ffn_fc_bias_tensor}); + hidden = modules::GeluModule({modules::GeluApproximation::Tanh}).build(ctx, hidden); + return modules::LinearModule({intermediate_size, hidden_size, true}).build( + ctx, hidden, {make_graph_param_tensor(layer.ffn_proj_weight), layer.ffn_proj_bias_tensor}); +} + +core::TensorValue concat_along_axis( + core::ModuleBuildContext & ctx, + const core::TensorValue & lhs, + const core::TensorValue & rhs, + int logical_axis) { + auto output_shape = lhs.shape; + output_shape.dims[logical_axis] += rhs.shape.dims[logical_axis]; + return core::wrap_tensor( + ggml_concat(ctx.ggml, lhs.tensor, rhs.tensor, core::logical_axis_to_ggml_axis(lhs.shape.rank, logical_axis)), + output_shape, + lhs.type); +} + +// prefix_key/prefix_value (if present) hold the already-computed cond-prefix K/V for this layer, +// so the dynamic [text; speech] window here can attend over the full [cond; text; speech] +// history even though it is a fresh graph with no KV cache of its own. +T3TurboLayerOutput build_t3_turbo_layer_full( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const T3TurboInferenceWeights::TransformerLayer & layer, + int64_t hidden_size, + int64_t intermediate_size, + int64_t num_heads, + const std::optional & prefix_key = std::nullopt, + const std::optional & prefix_value = std::nullopt) { + const int64_t head_dim = hidden_size / num_heads; + auto normed = modules::LayerNormModule({hidden_size, 1.0e-5f, true, true}).build( + ctx, input, {layer.ln1_weight_tensor, layer.ln1_bias_tensor}); + auto qkv = modules::LinearModule({hidden_size, 3 * hidden_size, true}).build( + ctx, normed, {make_graph_param_tensor(layer.attn_qkv_weight), layer.attn_qkv_bias_tensor}); + auto q = modules::SliceModule({2, 0, hidden_size}).build(ctx, qkv); + auto k = modules::SliceModule({2, hidden_size, hidden_size}).build(ctx, qkv); + auto v = modules::SliceModule({2, 2 * hidden_size, hidden_size}).build(ctx, qkv); + + auto reshape_heads = [&](const core::TensorValue & value) { + return core::reshape_tensor( + ctx, contiguous(ctx, value), core::TensorShape::from_dims({value.shape.dims[0], value.shape.dims[1], num_heads, head_dim})); + }; + auto k_cache = reshape_heads(k); + auto v_cache = reshape_heads(v); + + auto all_k = prefix_key.has_value() ? concat_along_axis(ctx, contiguous(ctx, *prefix_key), k_cache, 1) : k_cache; + auto all_v = prefix_value.has_value() ? concat_along_axis(ctx, contiguous(ctx, *prefix_value), v_cache, 1) : v_cache; + + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(q)); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, all_k); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, all_v); + + auto scores = modules::MatMulModule{}.build(ctx, q_heads, modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k_heads)); + scores = core::wrap_tensor( + ggml_scale(ctx.ggml, scores.tensor, 1.0f / std::sqrt(static_cast(head_dim))), scores.shape, GGML_TYPE_F32); + const int n_past = prefix_key.has_value() ? static_cast(prefix_key->shape.dims[1]) : 0; + scores = core::wrap_tensor(ggml_diag_mask_inf(ctx.ggml, scores.tensor, n_past), scores.shape, GGML_TYPE_F32); + auto attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, contiguous(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); + auto context = modules::MatMulModule{}.build(ctx, attn, v_heads); + context = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context); + context = core::reshape_tensor(ctx, contiguous(ctx, context), input.shape); + context = modules::LinearModule({hidden_size, hidden_size, true}).build( + ctx, context, {make_graph_param_tensor(layer.attn_output_weight), layer.attn_output_bias_tensor}); + auto hidden = core::wrap_tensor( + ggml_add(ctx.ggml, contiguous(ctx, context).tensor, contiguous(ctx, input).tensor), input.shape, GGML_TYPE_F32); + + auto mlp_in = modules::LayerNormModule({hidden_size, 1.0e-5f, true, true}).build( + ctx, hidden, {layer.ln2_weight_tensor, layer.ln2_bias_tensor}); + auto mlp_out = turbo_gpt_mlp(ctx, mlp_in, hidden_size, intermediate_size, layer); + auto output = core::wrap_tensor( + ggml_add(ctx.ggml, contiguous(ctx, mlp_out).tensor, contiguous(ctx, hidden).tensor), hidden.shape, GGML_TYPE_F32); + return {output, k_cache, v_cache}; +} + +T3TurboLayerOutput build_t3_turbo_layer_cached( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const T3TurboInferenceWeights::TransformerLayer & layer, + int64_t hidden_size, + int64_t intermediate_size, + int64_t num_heads, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask) { + const int64_t head_dim = hidden_size / num_heads; + auto normed = modules::LayerNormModule({hidden_size, 1.0e-5f, true, true}).build( + ctx, input, {layer.ln1_weight_tensor, layer.ln1_bias_tensor}); + auto qkv = modules::LinearModule({hidden_size, 3 * hidden_size, true}).build( + ctx, normed, {make_graph_param_tensor(layer.attn_qkv_weight), layer.attn_qkv_bias_tensor}); + auto q = modules::SliceModule({2, 0, hidden_size}).build(ctx, qkv); + auto k = modules::SliceModule({2, hidden_size, hidden_size}).build(ctx, qkv); + auto v = modules::SliceModule({2, 2 * hidden_size, hidden_size}).build(ctx, qkv); + auto reshape_heads = [&](const core::TensorValue & value) { + return core::reshape_tensor( + ctx, contiguous(ctx, value), core::TensorShape::from_dims({value.shape.dims[0], value.shape.dims[1], num_heads, head_dim})); + }; + auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(q)); + k = reshape_heads(k); + v = reshape_heads(v); + + const modules::FastKVSetRowsModule set_rows; + auto updated_key = set_rows.build(ctx, cache_key, k, cache_slot); + auto updated_value = set_rows.build(ctx, cache_value, v, cache_slot); + + q_heads = contiguous(ctx, q_heads); + auto k_heads = contiguous(ctx, modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, updated_key)); + auto v_heads = contiguous(ctx, modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, updated_value)); + + auto * flash = ggml_flash_attn_ext( + ctx.ggml, + q_heads.tensor, + k_heads.tensor, + v_heads.tensor, + attention_mask.tensor, + 1.0f / std::sqrt(static_cast(head_dim)), + 0.0f, + 0.0f); + ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); + auto context = core::wrap_tensor( + flash, + core::TensorShape::from_dims({q_heads.shape.dims[0], q_heads.shape.dims[2], q_heads.shape.dims[1], head_dim}), + GGML_TYPE_F32); + context = contiguous(ctx, context); + context = core::reshape_tensor(ctx, context, input.shape); + context = modules::LinearModule({hidden_size, hidden_size, true}).build( + ctx, context, {make_graph_param_tensor(layer.attn_output_weight), layer.attn_output_bias_tensor}); + auto hidden = core::wrap_tensor( + ggml_add(ctx.ggml, contiguous(ctx, context).tensor, contiguous(ctx, input).tensor), input.shape, GGML_TYPE_F32); + + auto mlp_in = modules::LayerNormModule({hidden_size, 1.0e-5f, true, true}).build( + ctx, hidden, {layer.ln2_weight_tensor, layer.ln2_bias_tensor}); + auto mlp_out = turbo_gpt_mlp(ctx, mlp_in, hidden_size, intermediate_size, layer); + auto output = core::wrap_tensor( + ggml_add(ctx.ggml, contiguous(ctx, mlp_out).tensor, contiguous(ctx, hidden).tensor), hidden.shape, GGML_TYPE_F32); + return {output, k, v}; +} + +inline bool same_backend(const engine::core::BackendConfig & lhs, const engine::core::BackendConfig & rhs) { + return lhs.type == rhs.type && lhs.device == rhs.device && lhs.threads == rhs.threads; +} + +class T3TurboBackendOwner { +public: + T3TurboBackendOwner(const T3TurboInferenceWeights & weights, const engine::core::BackendConfig & config) + : config_(config), execution_context_(weights.execution_context) { + if (!execution_context_) { + throw std::runtime_error("T3 Turbo backend owner requires loaded backend weights"); + } + } + ggml_backend_t backend() const noexcept { return execution_context_->backend(); } + const engine::core::ExecutionContext & execution_context() const noexcept { return *execution_context_; } + const engine::core::BackendConfig & config() const noexcept { return config_; } + +private: + engine::core::BackendConfig config_; + const engine::core::ExecutionContext * execution_context_ = nullptr; +}; + +struct T3TurboLayerCacheState { + std::vector key; + std::vector value; +}; + +struct T3TurboCacheState { + int64_t hidden_size = 0; + int64_t num_heads = 0; + int64_t head_dim = 0; + int64_t steps = 0; + std::vector layers; +}; + +// Batch is always 1: Chatterbox Turbo never uses classifier-free guidance, so there is no +// unconditional duplicate branch the way base Chatterbox's T3 decode runners need. +class T3TurboDecodeBackendRunner { +public: + T3TurboDecodeBackendRunner( + const T3TurboInferenceWeights & weights, + int64_t cache_steps, + std::shared_ptr owner) + : owner_(std::move(owner)), + backend_config_(owner_->config()), + cache_steps_(cache_steps), + hidden_size_(weights.hidden_size), + num_heads_(weights.num_heads), + head_dim_(weights.hidden_size / weights.num_heads) { + if (cache_steps_ <= 0) { + throw std::runtime_error("T3TurboDecodeBackendRunner requires positive cache_steps"); + } + ggml_init_params params = {}; + params.mem_size = 256ull * 1024ull * 1024ull; + params.no_alloc = true; + ggml_ = ggml_init(params); + if (ggml_ == nullptr) { + throw std::runtime_error("failed to initialize ggml context for T3 Turbo decode runner"); + } + core::ModuleBuildContext ctx = {}; + ctx.ggml = ggml_; + ctx.module_instance_name = "t3_turbo_decode_runner"; + try { + input_hidden_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, hidden_size_})); + attention_mask_ = core::make_tensor(ctx, GGML_TYPE_F16, core::TensorShape::from_dims({1, 1, 1, cache_steps_})); + cache_slot_ = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1})); + + key_cache_tensors_.reserve(weights.layers.size()); + value_cache_tensors_.reserve(weights.layers.size()); + for (size_t layer_index = 0; layer_index < weights.layers.size(); ++layer_index) { + key_cache_tensors_.push_back(core::make_tensor( + ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, cache_steps_, num_heads_, head_dim_}))); + value_cache_tensors_.push_back(core::make_tensor( + ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, cache_steps_, num_heads_, head_dim_}))); + } + + auto hidden = input_hidden_; + current_keys_.reserve(weights.layers.size()); + current_values_.reserve(weights.layers.size()); + current_key_scratch_.resize(weights.layers.size()); + current_value_scratch_.resize(weights.layers.size()); + for (size_t layer_index = 0; layer_index < weights.layers.size(); ++layer_index) { + auto out = build_t3_turbo_layer_cached( + ctx, + hidden, + weights.layers[layer_index], + hidden_size_, + weights.mlp_intermediate_size, + num_heads_, + key_cache_tensors_[layer_index], + value_cache_tensors_[layer_index], + cache_slot_, + attention_mask_); + hidden = out.hidden; + current_keys_.push_back(out.key); + current_values_.push_back(out.value); + } + hidden_out_ = modules::LayerNormModule({hidden_size_, 1.0e-5f, true, true}).build( + ctx, hidden, {weights.ln_f_weight_tensor, weights.ln_f_bias_tensor}); + logits_out_ = modules::LinearModule({hidden_size_, weights.speech_vocab, true}).build( + ctx, hidden_out_, {make_graph_param_tensor(weights.speech_head_weight), weights.speech_head_bias_tensor}); + + graph_ = ggml_new_graph_custom(ggml_, 65536, false); + ggml_build_forward_expand(graph_, logits_out_.tensor); + for (const auto & key : current_keys_) { + ggml_build_forward_expand(graph_, key.tensor); + } + for (const auto & value : current_values_) { + ggml_build_forward_expand(graph_, value.tensor); + } + buffer_ = ggml_backend_alloc_ctx_tensors(ggml_, owner_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate backend tensors for T3 Turbo decode runner"); + } + engine::core::prepare_host_graph_plan(owner_->execution_context(), graph_, cpu_plan_); + } catch (...) { + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + buffer_ = nullptr; + } + if (ggml_ != nullptr) { + ggml_free(ggml_); + ggml_ = nullptr; + } + throw; + } + cache_state_.hidden_size = hidden_size_; + cache_state_.num_heads = num_heads_; + cache_state_.head_dim = head_dim_; + cache_state_.layers.resize(weights.layers.size()); + attention_mask_scratch_.resize(static_cast(cache_steps_), 0.0f); + import_state(cache_state_); + } + + ~T3TurboDecodeBackendRunner() { + if (owner_ != nullptr && graph_ != nullptr) { + engine::core::release_backend_graph_resources(owner_->backend(), graph_); + } + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + if (ggml_ != nullptr) { + ggml_free(ggml_); + } + } + + bool matches(int64_t cache_steps, const engine::core::BackendConfig & backend_config) const { + return cache_steps_ == cache_steps && same_backend(backend_config_, backend_config); + } + + void import_state(const T3TurboCacheState & state) { + if (state.layers.size() != cache_state_.layers.size()) { + throw std::runtime_error("T3TurboDecodeBackendRunner state layer count mismatch"); + } + if (state.steps > cache_steps_) { + throw std::runtime_error("T3TurboDecodeBackendRunner state exceeds cache capacity"); + } + cache_state_ = state; + cache_state_.hidden_size = hidden_size_; + cache_state_.num_heads = num_heads_; + cache_state_.head_dim = head_dim_; + for (size_t layer_index = 0; layer_index < cache_state_.layers.size(); ++layer_index) { + const auto & layer = cache_state_.layers[layer_index]; + if (cache_state_.steps == 0 || layer.key.empty() || layer.value.empty()) { + continue; + } + core::write_tensor_f32_slice(key_cache_tensors_[layer_index], 0, layer.key.data(), layer.key.size()); + core::write_tensor_f32_slice(value_cache_tensors_[layer_index], 0, layer.value.data(), layer.value.size()); + } + } + + T3TurboCacheState export_state() const { return cache_state_; } + + T3TurboCacheState export_state_from_device() const { + T3TurboCacheState state = cache_state_; + state.layers.assign(cache_state_.layers.size(), {}); + if (cache_state_.steps <= 0) { + return state; + } + for (size_t layer_index = 0; layer_index < key_cache_tensors_.size(); ++layer_index) { + state.layers[layer_index].key = core::read_tensor_f32(key_cache_tensors_[layer_index].tensor); + state.layers[layer_index].value = core::read_tensor_f32(value_cache_tensors_[layer_index].tensor); + } + return state; + } + + int64_t valid_steps() const noexcept { return cache_state_.steps; } + int64_t cache_capacity_steps() const noexcept { return cache_steps_; } + + void set_capture_cache_state(bool capture) noexcept { capture_cache_state_ = capture; } + + std::vector step(const std::vector & input_hidden, int64_t /*position*/) { + if (static_cast(input_hidden.size()) != hidden_size_) { + throw std::runtime_error("T3TurboDecodeBackendRunner step input size mismatch"); + } + const int64_t masked_prefix_begin = std::clamp(cache_state_.steps + 1, 0, cache_steps_); + attention_mask_scratch_.assign(static_cast(cache_steps_), 0.0f); + for (int64_t step = masked_prefix_begin; step < cache_steps_; ++step) { + attention_mask_scratch_[static_cast(step)] = -10000.0f; + } + core::write_tensor_f32(input_hidden_, input_hidden); + core::write_tensor_i32(cache_slot_, std::vector{static_cast(cache_state_.steps)}); + core::write_tensor_f16(attention_mask_, attention_mask_scratch_); + + const ggml_status status = engine::core::compute_graph(owner_->execution_context(), graph_, cpu_plan_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("ggml compute failed for T3 Turbo decode runner"); + } + if (capture_cache_state_) { + for (size_t layer_index = 0; layer_index < current_keys_.size(); ++layer_index) { + core::read_tensor_f32_into(current_keys_[layer_index].tensor, current_key_scratch_[layer_index]); + core::read_tensor_f32_into(current_values_[layer_index].tensor, current_value_scratch_[layer_index]); + auto & layer = cache_state_.layers[layer_index]; + layer.key.insert(layer.key.end(), current_key_scratch_[layer_index].begin(), current_key_scratch_[layer_index].end()); + layer.value.insert(layer.value.end(), current_value_scratch_[layer_index].begin(), current_value_scratch_[layer_index].end()); + } + } + cache_state_.steps += 1; + return core::read_tensor_f32(logits_out_.tensor); + } + +private: + std::shared_ptr owner_; + engine::core::BackendConfig backend_config_; + ggml_context * ggml_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; + ggml_cgraph * graph_ = nullptr; + engine::core::HostGraphPlan cpu_plan_; + int64_t cache_steps_ = 0; + int64_t hidden_size_ = 0; + int64_t num_heads_ = 0; + int64_t head_dim_ = 0; + core::TensorValue input_hidden_; + core::TensorValue cache_slot_; + core::TensorValue attention_mask_; + core::TensorValue hidden_out_; + core::TensorValue logits_out_; + std::vector key_cache_tensors_; + std::vector value_cache_tensors_; + std::vector current_keys_; + std::vector current_values_; + T3TurboCacheState cache_state_; + bool capture_cache_state_ = false; + std::vector attention_mask_scratch_; + std::vector> current_key_scratch_; + std::vector> current_value_scratch_; +}; + +struct T3TurboPrefillOutput { + std::vector logits; + T3TurboCacheState cache; +}; + +class T3TurboPrefillBackendRunner { +public: + T3TurboPrefillBackendRunner( + const T3TurboInferenceWeights & weights, + int64_t prefix_steps, + int64_t seq_len, + std::shared_ptr owner) + : owner_(std::move(owner)), + backend_config_(owner_->config()), + prefix_steps_(prefix_steps), + seq_len_(seq_len), + hidden_size_(weights.hidden_size), + speech_vocab_(weights.speech_vocab), + num_heads_(weights.num_heads), + head_dim_(weights.hidden_size / weights.num_heads) { + if (seq_len_ <= 0 || prefix_steps_ < 0) { + throw std::runtime_error("T3 Turbo prefill runner requires positive seq_len and non-negative prefix"); + } + ggml_init_params params = {}; + params.mem_size = 768ull * 1024ull * 1024ull; + params.no_alloc = true; + ggml_ = ggml_init(params); + if (ggml_ == nullptr) { + throw std::runtime_error("failed to initialize ggml context for T3 Turbo prefill runner"); + } + core::ModuleBuildContext ctx = {}; + ctx.ggml = ggml_; + ctx.module_instance_name = "t3_turbo_prefill_runner"; + try { + input_hidden_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, seq_len_, hidden_size_})); + if (prefix_steps_ > 0) { + prefix_key_tensors_.reserve(weights.layers.size()); + prefix_value_tensors_.reserve(weights.layers.size()); + for (size_t layer_index = 0; layer_index < weights.layers.size(); ++layer_index) { + prefix_key_tensors_.push_back(core::make_tensor( + ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, prefix_steps_, num_heads_, head_dim_}))); + prefix_value_tensors_.push_back(core::make_tensor( + ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, prefix_steps_, num_heads_, head_dim_}))); + } + } + auto hidden = input_hidden_; + current_keys_.reserve(weights.layers.size()); + current_values_.reserve(weights.layers.size()); + for (size_t layer_index = 0; layer_index < weights.layers.size(); ++layer_index) { + const std::optional prefix_key = + prefix_steps_ > 0 ? std::optional(prefix_key_tensors_[layer_index]) : std::nullopt; + const std::optional prefix_value = + prefix_steps_ > 0 ? std::optional(prefix_value_tensors_[layer_index]) : std::nullopt; + auto out = build_t3_turbo_layer_full( + ctx, + hidden, + weights.layers[layer_index], + hidden_size_, + weights.mlp_intermediate_size, + num_heads_, + prefix_key, + prefix_value); + hidden = out.hidden; + current_keys_.push_back(out.key); + current_values_.push_back(out.value); + } + hidden_out_ = modules::LayerNormModule({hidden_size_, 1.0e-5f, true, true}).build( + ctx, hidden, {weights.ln_f_weight_tensor, weights.ln_f_bias_tensor}); + logits_out_ = modules::LinearModule({hidden_size_, speech_vocab_, true}).build( + ctx, hidden_out_, {make_graph_param_tensor(weights.speech_head_weight), weights.speech_head_bias_tensor}); + + graph_ = ggml_new_graph_custom(ggml_, 131072, false); + ggml_build_forward_expand(graph_, logits_out_.tensor); + for (const auto & key : current_keys_) { + ggml_build_forward_expand(graph_, key.tensor); + } + for (const auto & value : current_values_) { + ggml_build_forward_expand(graph_, value.tensor); + } + buffer_ = ggml_backend_alloc_ctx_tensors(ggml_, owner_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate backend tensors for T3 Turbo prefill runner"); + } + engine::core::prepare_host_graph_plan(owner_->execution_context(), graph_, cpu_plan_); + } catch (...) { + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + buffer_ = nullptr; + } + if (ggml_ != nullptr) { + ggml_free(ggml_); + ggml_ = nullptr; + } + throw; + } + } + + ~T3TurboPrefillBackendRunner() { + if (owner_ != nullptr && graph_ != nullptr) { + engine::core::release_backend_graph_resources(owner_->backend(), graph_); + } + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + if (ggml_ != nullptr) { + ggml_free(ggml_); + } + } + + bool matches(int64_t prefix_steps, int64_t seq_len, const engine::core::BackendConfig & backend_config) const { + return prefix_steps_ == prefix_steps && seq_len_ == seq_len && same_backend(backend_config_, backend_config); + } + + T3TurboPrefillOutput run(const std::vector & input_hidden, const T3TurboCacheState & prefix_state) { + if (static_cast(input_hidden.size()) != seq_len_ * hidden_size_) { + throw std::runtime_error("T3 Turbo prefill runner input size mismatch"); + } + if (prefix_state.steps != prefix_steps_ || prefix_state.layers.size() != current_keys_.size()) { + throw std::runtime_error("T3 Turbo prefill runner prefix cache shape mismatch"); + } + core::write_tensor_f32(input_hidden_, input_hidden); + for (size_t layer_index = 0; layer_index < prefix_key_tensors_.size(); ++layer_index) { + core::write_tensor_f32(prefix_key_tensors_[layer_index], prefix_state.layers[layer_index].key); + core::write_tensor_f32(prefix_value_tensors_[layer_index], prefix_state.layers[layer_index].value); + } + const ggml_status status = engine::core::compute_graph(owner_->execution_context(), graph_, cpu_plan_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("ggml compute failed for T3 Turbo prefill runner"); + } + T3TurboPrefillOutput out; + const auto logits = core::read_tensor_f32(logits_out_.tensor); + out.logits.assign(logits.end() - speech_vocab_, logits.end()); + out.cache.hidden_size = hidden_size_; + out.cache.num_heads = num_heads_; + out.cache.head_dim = head_dim_; + out.cache.steps = prefix_steps_ + seq_len_; + out.cache.layers.resize(current_keys_.size()); + for (size_t layer_index = 0; layer_index < current_keys_.size(); ++layer_index) { + auto & layer = out.cache.layers[layer_index]; + if (prefix_steps_ > 0) { + layer.key = prefix_state.layers[layer_index].key; + layer.value = prefix_state.layers[layer_index].value; + } + const auto dyn_key = core::read_tensor_f32(current_keys_[layer_index].tensor); + const auto dyn_value = core::read_tensor_f32(current_values_[layer_index].tensor); + layer.key.insert(layer.key.end(), dyn_key.begin(), dyn_key.end()); + layer.value.insert(layer.value.end(), dyn_value.begin(), dyn_value.end()); + } + return out; + } + +private: + std::shared_ptr owner_; + engine::core::BackendConfig backend_config_; + int64_t prefix_steps_ = 0; + int64_t seq_len_ = 0; + int64_t hidden_size_ = 0; + int64_t speech_vocab_ = 0; + int64_t num_heads_ = 0; + int64_t head_dim_ = 0; + ggml_context * ggml_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; + ggml_cgraph * graph_ = nullptr; + engine::core::HostGraphPlan cpu_plan_; + core::TensorValue input_hidden_; + core::TensorValue hidden_out_; + core::TensorValue logits_out_; + std::vector prefix_key_tensors_; + std::vector prefix_value_tensors_; + std::vector current_keys_; + std::vector current_values_; +}; + +} // namespace +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/loader.cpp b/src/community_models/chatterbox_turbo/loader.cpp new file mode 100644 index 000000000..b5fbe8525 --- /dev/null +++ b/src/community_models/chatterbox_turbo/loader.cpp @@ -0,0 +1,128 @@ +#include "engine/community_models/chatterbox_turbo/loader.h" + +#include "engine/community_models/chatterbox_turbo/session.h" + +#include "engine/framework/model_spec/package.h" + +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +namespace { + +runtime::CapabilitySet capabilities(const ChatterboxTurboAssets &) { + runtime::CapabilitySet out; + out.supported_tasks.push_back({runtime::VoiceTaskKind::VoiceCloning, {runtime::RunMode::Offline}}); + out.languages = {"en"}; + // Custom voice cloning (a caller-supplied speaker reference) is not implemented yet -- only + // the built-in default voice baked into the GGUF's `conds.*` tensors is supported. Advertise + // this honestly rather than claiming speaker-reference support the loader would reject. + out.supports_speaker_reference = false; + out.supports_style_condition = false; + return out; +} + +runtime::ModelMetadata metadata(const ChatterboxTurboAssets & assets) { + runtime::ModelMetadata out; + out.family = "chatterbox_turbo"; + out.variant = assets.resources.model_root().filename().string(); + out.description = + "Chatterbox Turbo (distilled GPT2 T3 backbone + meanflow S3Gen decoder) loaded from local assets. " + "Built-in default voice only; custom voice cloning is not yet implemented."; + return out; +} + +class ChatterboxTurboLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "chatterbox_turbo"; + } + + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks.push_back({runtime::VoiceTaskKind::VoiceCloning, {runtime::RunMode::Offline}}); + out.languages = {"en"}; + out.supports_speaker_reference = false; + out.supports_style_condition = false; + return out; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + try { + (void) engine::model_spec::load_resource_bundle( + request.model_path, + engine::model_spec::default_spec_path(family())); + return !request.family_hint.has_value() || *request.family_hint == family(); + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto assets = load_chatterbox_turbo_assets(request.model_path); + runtime::ModelInspection inspection; + inspection.model_root = assets->resources.model_root(); + inspection.metadata = metadata(*assets); + inspection.capabilities = capabilities(*assets); + const auto spec_path = engine::model_spec::default_spec_path(family()); + inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( + request.model_path, + spec_path, + engine::model_spec::ResourceKind::Files); + inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( + request.model_path, + spec_path, + engine::model_spec::ResourceKind::Tensors); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_chatterbox_turbo_model(request.model_path); + } +}; + +} // namespace + +ChatterboxTurboLoadedModel::ChatterboxTurboLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets) + : metadata_(std::move(metadata)), + capabilities_(std::move(capabilities)), + assets_(std::move(assets)) { + if (assets_ == nullptr) { + throw std::runtime_error("Chatterbox Turbo loaded model requires assets"); + } +} + +const runtime::ModelMetadata & ChatterboxTurboLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & ChatterboxTurboLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr ChatterboxTurboLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.task != runtime::VoiceTaskKind::VoiceCloning) { + throw std::runtime_error("Chatterbox Turbo supports VoiceCloning only"); + } + if (task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Chatterbox Turbo only supports offline mode"); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_chatterbox_turbo_model(const std::filesystem::path & model_root) { + auto assets = load_chatterbox_turbo_assets(model_root); + return std::make_unique(metadata(*assets), capabilities(*assets), std::move(assets)); +} + +std::shared_ptr make_chatterbox_turbo_loader() { + return std::make_shared(); +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/s3gen_turbo.cpp b/src/community_models/chatterbox_turbo/s3gen_turbo.cpp new file mode 100644 index 000000000..a9ef66071 --- /dev/null +++ b/src/community_models/chatterbox_turbo/s3gen_turbo.cpp @@ -0,0 +1,103 @@ +#include "engine/community_models/chatterbox_turbo/s3gen_turbo.h" + +#include + +namespace engine::community_models::chatterbox_turbo { + +namespace { + +engine::modules::HiftVocoderConfig make_turbo_hift_config(engine::assets::TensorStorageType weight_storage_type) { + // Matches the S3Gen GGUF's `chatterbox.s3gen.*` metadata and base Chatterbox's own HiFT + // config (src/models/chatterbox/hift_vocoder_impl.cpp) -- the vocoder architecture is + // unchanged by the meanflow distillation, only the tensor names differ. The repacked native + // GGUF folds torch weight-norm parametrization into a plain "weight" tensor at conversion + // time (no "parametrizations.weight.original0/1" keys), so this uses the Canonical (plain) + // layout rather than TorchParametrizedWeightNorm. + engine::modules::HiftVocoderConfig config; + config.in_channels = 80; + config.base_channels = 512; + config.nb_harmonics = 8; + config.sampling_rate = 24000; + config.nsf_alpha = 0.1F; + config.nsf_sigma = 0.003F; + config.nsf_voiced_threshold = 10.0F; + config.upsample_rates = {8, 5, 3}; + config.upsample_kernel_sizes = {16, 11, 7}; + config.istft_n_fft = 16; + config.istft_hop = 4; + config.resblock_kernel_sizes = {3, 7, 11}; + config.resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; + config.source_resblock_kernel_sizes = {7, 7, 11}; + config.source_resblock_dilation_sizes = {{1, 3, 5}, {1, 3, 5}, {1, 3, 5}}; + config.lrelu_slope = 0.1F; + config.audio_limit = 0.99F; + config.f0_num_class = 1; + config.f0_in_channels = 80; + config.f0_cond_channels = 512; + config.weight_storage_type = weight_storage_type; + // The repacked GGUF stores every vocoder tensor under the top-level "v." namespace (matching + // base Chatterbox's own "mel2wav."-style tensor_prefix convention, just with a different + // literal prefix); HiftVocoderComponent applies config.tensor_prefix + name itself, so no + // separate translation layer is needed here. + config.tensor_prefix = "v."; + config.weight_layout = engine::modules::HiftVocoderWeightLayout::Canonical; + return config; +} + +} // namespace + +std::shared_ptr ChatterboxTurboS3Gen::load( + std::shared_ptr s3gen_source, + const engine::core::ExecutionContext & execution_context, + engine::assets::TensorStorageType weight_storage_type) { + auto out = std::make_shared(); + out->execution_context_ = &execution_context; + out->encoder_weights_ = + engine::models::chatterbox::load_s3_flow_encoder_weights(*s3gen_source, execution_context, weight_storage_type); + out->decoder_weights_ = + engine::models::chatterbox::load_s3_flow_decoder_weights(*s3gen_source, execution_context, weight_storage_type); + if (!engine::models::chatterbox::s3_flow_decoder_is_meanflow(*out->decoder_weights_)) { + throw std::runtime_error( + "Chatterbox Turbo S3Gen weights are missing the meanflow time_embed_mixer tensor (flow.decoder.estimator.time_embed_mixer)"); + } + out->vocoder_ = std::make_shared( + engine::modules::HiftVocoderComponent::load_from_tensor_source( + s3gen_source, execution_context.config(), make_turbo_hift_config(weight_storage_type))); + return out; +} + +engine::models::chatterbox::S3GenInferenceOutputs ChatterboxTurboS3Gen::synthesize( + const engine::models::chatterbox::EmbedReferenceOutputs & ref_dict, + const std::vector & speech_tokens, + uint64_t flow_seed, + uint64_t vocoder_seed) const { + // n_cfm_timesteps=2 matches tts_turbo.py's ChatterboxTurboTTS.generate default; cfg_rate and + // cosine_schedule are unused on the meanflow path (see s3gen_inference.cpp's + // decoder_weights.meanflow branch). + const auto mel = engine::models::chatterbox::compute_s3_token2mel_inference( + cache_, + *encoder_weights_, + *decoder_weights_, + ref_dict, + speech_tokens, + static_cast(speech_tokens.size()), + /*num_steps=*/2, + /*cfg_rate=*/0.0f, + /*cosine_schedule=*/false, + /*full_noise=*/{}, + flow_seed, + execution_context_->config(), + /*timing=*/nullptr); + + const auto voc = vocoder_->synthesize(mel.mel, mel.frames, vocoder_seed); + + engine::models::chatterbox::S3GenInferenceOutputs outputs; + outputs.waveform = voc.waveform; + outputs.samples = voc.samples; + outputs.mel = mel.mel; + outputs.mel_channels = mel.channels; + outputs.mel_frames = mel.frames; + return outputs; +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/session.cpp b/src/community_models/chatterbox_turbo/session.cpp new file mode 100644 index 000000000..19790b885 --- /dev/null +++ b/src/community_models/chatterbox_turbo/session.cpp @@ -0,0 +1,112 @@ +#include "engine/community_models/chatterbox_turbo/session.h" + +#include "engine/framework/runtime/options.h" +#include "engine/framework/text/chunking.h" + +#include + +namespace engine::community_models::chatterbox_turbo { + +namespace { + +// Chatterbox Turbo has no CFG-based prefix cache to amortize across chunks (unlike base +// Chatterbox), but chunking still matters here: the S3Gen flow encoder's attention buffer scales +// with token count, and without a bound a long paragraph can exceed VRAM on small GPUs even +// though the model weights themselves fit comfortably. Matches base Chatterbox's default +// (src/models/chatterbox/session.cpp). +constexpr int64_t kDefaultTextChunkSize = 128; + +ChatterboxTurboGenerateConfig make_generate_config(const std::unordered_map & options) { + ChatterboxTurboGenerateConfig config; + if (const auto value = runtime::parse_float_option(options, {"temperature"})) { + config.temperature = *value; + } + if (const auto value = runtime::parse_float_option(options, {"top_p"})) { + config.top_p = *value; + } + if (const auto value = runtime::parse_float_option(options, {"repetition_penalty"})) { + config.repetition_penalty = *value; + } + if (const auto value = runtime::parse_u32_option(options, {"top_k"})) { + config.top_k = *value; + } + if (const auto value = runtime::parse_u32_option(options, {"max_new_tokens"})) { + config.max_new_tokens = *value; + } + if (const auto value = runtime::parse_u32_option(options, {"seed"})) { + config.seed = *value; + } + // exaggeration/cfg_weight/min_p are accepted elsewhere in the request surface but have no + // effect on Turbo, matching upstream tts_turbo.py's logger.warning(...ignored...). + return config; +} + +} // namespace + +ChatterboxTurboSession::ChatterboxTurboSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : RuntimeSessionBase(options), task_(std::move(task)), assets_(std::move(assets)) { + if (!assets_) { + throw std::runtime_error("Chatterbox Turbo session requires assets"); + } + if (task_.task != runtime::VoiceTaskKind::VoiceCloning) { + throw std::runtime_error("Chatterbox Turbo session supports --task clone only"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Chatterbox Turbo session only supports offline mode"); + } +} + +ChatterboxTurboSession::~ChatterboxTurboSession() = default; + +std::string ChatterboxTurboSession::family() const { + return "chatterbox_turbo"; +} + +runtime::VoiceTaskKind ChatterboxTurboSession::task_kind() const { + return task_.task; +} + +runtime::RunMode ChatterboxTurboSession::run_mode() const { + return task_.mode; +} + +void ChatterboxTurboSession::prepare(const runtime::SessionPreparationRequest & request) { + if (!request.text.has_value() || request.text->text.empty()) { + throw std::runtime_error("Chatterbox Turbo prepare requires text input"); + } + if (request.voice.has_value() && request.voice->speaker.has_value() && request.voice->speaker->audio.has_value()) { + throw std::runtime_error( + "Chatterbox Turbo does not yet support custom voice cloning (only the built-in default voice) " + "-- omit the speaker reference audio to use the built-in voice"); + } + if (!component_) { + component_ = std::make_unique(assets_, execution_context()); + } + mark_prepared(); +} + +runtime::TaskResult ChatterboxTurboSession::run(const runtime::TaskRequest & request) { + require_prepared("Chatterbox Turbo run"); + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("Chatterbox Turbo run requires text input"); + } + const auto config = make_generate_config(request.options); + const int64_t text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); + const auto chunk_requests = runtime::chunk_text_request(request, text_chunk_size); + + runtime::AudioBuffer merged_audio; + for (const auto & chunk_request : chunk_requests) { + const auto outputs = component_->generate(chunk_request.text_input->text, config); + runtime::append_audio_buffer(merged_audio, runtime::AudioBuffer{24000, 1, outputs.waveform}); + } + + runtime::TaskResult result; + result.audio_output = std::move(merged_audio); + return result; +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/t3_turbo_component.cpp b/src/community_models/chatterbox_turbo/t3_turbo_component.cpp new file mode 100644 index 000000000..c25f880d2 --- /dev/null +++ b/src/community_models/chatterbox_turbo/t3_turbo_component.cpp @@ -0,0 +1,295 @@ +#include "components/t3_turbo_runtime.h" + +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +struct T3TurboInferenceComponent::State { + void release_runtime_graphs() { + std::lock_guard lock(mutex); + runner.reset(); + prefill_runner.reset(); + } + void release_runtime_cache() { + std::lock_guard lock(mutex); + owner.reset(); + runner.reset(); + prefill_runner.reset(); + prefix_cache.reset(); + } + + std::shared_ptr get_owner( + const engine::core::BackendConfig & backend, + const T3TurboInferenceWeights & weights) { + if (!owner || !same_backend(owner->config(), backend)) { + owner = std::make_shared(weights, backend); + runner.reset(); + prefill_runner.reset(); + prefix_cache.reset(); + } + return owner; + } + + struct PrefixCache { + std::vector speaker_embedding; + std::vector cond_prompt_speech_tokens; + int64_t cond_length = 0; + T3TurboCacheState cache; + }; + + std::shared_ptr get_runner( + int64_t cache_steps, + const engine::core::BackendConfig & backend, + const T3TurboInferenceWeights & weights) { + if (runner && runner->matches(cache_steps, backend)) { + return runner; + } + auto held_owner = get_owner(backend, weights); + runner.reset(); + runner = std::make_shared(weights, cache_steps, std::move(held_owner)); + return runner; + } + + std::shared_ptr get_prefill_runner( + int64_t prefix_steps, + int64_t seq_len, + const engine::core::BackendConfig & backend, + const T3TurboInferenceWeights & weights) { + if (prefill_runner && prefill_runner->matches(prefix_steps, seq_len, backend)) { + return prefill_runner; + } + auto held_owner = get_owner(backend, weights); + prefill_runner.reset(); + prefill_runner = std::make_shared(weights, prefix_steps, seq_len, std::move(held_owner)); + return prefill_runner; + } + + mutable std::mutex mutex; + std::shared_ptr owner; + std::optional prefix_cache; + std::shared_ptr prefill_runner; + std::shared_ptr runner; +}; + +T3TurboInferenceComponent::T3TurboInferenceComponent( + std::shared_ptr weights, + const engine::core::ExecutionContext & execution_context) + : weights_(std::move(weights)), execution_context_(&execution_context), state_(std::make_shared()) { + if (!weights_) { + throw std::runtime_error("T3TurboInferenceComponent requires weights"); + } +} + +namespace { + +// cond_embeds = [ spkr_proj(1 token), speech_emb(cond_prompt_speech_tokens) (speech_cond_prompt_len tokens) ] +// No perceiver resampler, no emotion conditioning, no per-segment position embeddings: Chatterbox +// Turbo's GPT2 backbone applies one absolute "wpe" position embedding across the whole combined +// [cond; text; speech] sequence, added host-side below (see t3.py prepare_input_embeds with +// hp.input_pos_emb=None: HF's GPT2Model applies its own wpe once over inputs_embeds). +std::vector build_cond_embeddings( + const T3TurboInferenceWeights & weights, + const std::vector & speaker_embedding, + const std::vector & cond_prompt_speech_tokens) { + const int64_t hidden = weights.hidden_size; + const auto spkr_enc_weight = engine::assets::tensor_data_to_f32("cond.spkr_enc.weight", weights.spkr_enc_weight); + const auto spkr_enc_bias = engine::assets::tensor_data_to_f32("cond.spkr_enc.bias", weights.spkr_enc_bias); + const auto speech_embedding_weight = + engine::assets::tensor_data_to_f32("speech_emb.weight", weights.speech_embedding_weight); + + std::vector spkr_proj(static_cast(hidden), 0.0f); + for (int64_t out_index = 0; out_index < hidden; ++out_index) { + double value = static_cast(spkr_enc_bias[static_cast(out_index)]); + const float * w = spkr_enc_weight.data() + static_cast(out_index * weights.speaker_embed_size); + for (int64_t in_index = 0; in_index < weights.speaker_embed_size; ++in_index) { + value += static_cast(w[in_index]) * static_cast(speaker_embedding[static_cast(in_index)]); + } + spkr_proj[static_cast(out_index)] = static_cast(value); + } + + const auto cond_prompt_emb = gather_rows(speech_embedding_weight, weights.speech_vocab, hidden, cond_prompt_speech_tokens); + + std::vector cond_embeddings; + cond_embeddings.reserve(spkr_proj.size() + cond_prompt_emb.size()); + cond_embeddings.insert(cond_embeddings.end(), spkr_proj.begin(), spkr_proj.end()); + cond_embeddings.insert(cond_embeddings.end(), cond_prompt_emb.begin(), cond_prompt_emb.end()); + return cond_embeddings; +} + +void add_wpe_in_place( + std::vector & embeddings, + const std::vector & wpe_weight, + int64_t hidden, + int64_t start_position) { + const int64_t seq_len = static_cast(embeddings.size()) / hidden; + for (int64_t t = 0; t < seq_len; ++t) { + const float * pos_row = wpe_weight.data() + static_cast((start_position + t) * hidden); + float * dst = embeddings.data() + static_cast(t * hidden); + for (int64_t i = 0; i < hidden; ++i) { + dst[i] += pos_row[i]; + } + } +} + +} // namespace + +T3TurboGenerateOutputs T3TurboInferenceComponent::generate_speech_tokens(const T3TurboGenerateRequest & request) const { + const auto & backend_config = execution_context_->config(); + const int64_t hidden_size = weights_->hidden_size; + const int64_t speech_vocab = weights_->speech_vocab; + TurboMt19937 rng(request.seed == 0 ? 0x2A2A2A2AU : request.seed); + + const auto wpe_weight = engine::assets::tensor_data_to_f32("wpe.weight", weights_->wpe_weight); + const auto text_embedding_weight = + engine::assets::tensor_data_to_f32("text_emb.weight", weights_->text_embedding_weight); + const auto speech_embedding_weight = + engine::assets::tensor_data_to_f32("speech_emb.weight", weights_->speech_embedding_weight); + + std::vector generated_ids = request.initial_speech_tokens; + if (generated_ids.empty()) { + generated_ids.push_back(kTurboStartSpeechToken); + } + + T3TurboGenerateOutputs outputs; + outputs.predicted_tokens.reserve(static_cast(request.max_new_tokens)); + + State::PrefixCache prefix_cache; + double prefix_cache_build_ms = 0.0; + { + std::lock_guard lock(state_->mutex); + const bool prefix_cache_hit = state_->prefix_cache.has_value() && + state_->prefix_cache->speaker_embedding == request.speaker_embedding && + state_->prefix_cache->cond_prompt_speech_tokens == request.cond_prompt_speech_tokens; + if (!prefix_cache_hit) { + const auto started = std::chrono::steady_clock::now(); + auto cond_embeddings = + build_cond_embeddings(*weights_, request.speaker_embedding, request.cond_prompt_speech_tokens); + const int64_t cond_length = static_cast(cond_embeddings.size()) / hidden_size; + add_wpe_in_place(cond_embeddings, wpe_weight, hidden_size, 0); + + auto runner = state_->get_runner(cond_length, backend_config, *weights_); + T3TurboCacheState initial_cache; + initial_cache.hidden_size = hidden_size; + initial_cache.num_heads = weights_->num_heads; + initial_cache.head_dim = hidden_size / weights_->num_heads; + initial_cache.layers = std::vector(weights_->layers.size()); + runner->import_state(initial_cache); + runner->set_capture_cache_state(true); + for (int64_t pos = 0; pos < cond_length; ++pos) { + std::vector step( + cond_embeddings.begin() + static_cast(pos * hidden_size), + cond_embeddings.begin() + static_cast((pos + 1) * hidden_size)); + runner->step(step, pos); + } + runner->set_capture_cache_state(false); + + State::PrefixCache new_cache; + new_cache.speaker_embedding = request.speaker_embedding; + new_cache.cond_prompt_speech_tokens = request.cond_prompt_speech_tokens; + new_cache.cond_length = cond_length; + new_cache.cache = runner->export_state(); + state_->prefix_cache = std::move(new_cache); + prefix_cache_build_ms = std::chrono::duration( + std::chrono::steady_clock::now() - started).count(); + } + prefix_cache = *state_->prefix_cache; + } + + const int64_t text_length = static_cast(request.text_tokens.size()); + const int64_t speech_length = static_cast(generated_ids.size()); + const int64_t dynamic_length = text_length + speech_length; + + std::vector dynamic_embeddings(static_cast(dynamic_length * hidden_size), 0.0f); + if (text_length > 0) { + const auto text_emb = gather_rows(text_embedding_weight, weights_->text_vocab, hidden_size, request.text_tokens); + std::copy(text_emb.begin(), text_emb.end(), dynamic_embeddings.begin()); + } + { + const auto speech_emb = gather_rows(speech_embedding_weight, weights_->speech_vocab, hidden_size, generated_ids); + std::copy(speech_emb.begin(), speech_emb.end(), dynamic_embeddings.begin() + static_cast(text_length * hidden_size)); + } + add_wpe_in_place(dynamic_embeddings, wpe_weight, hidden_size, prefix_cache.cond_length); + + auto step_started = std::chrono::steady_clock::now(); + T3TurboPrefillOutput prefill_output; + { + std::shared_ptr prefill_runner; + { + std::lock_guard lock(state_->mutex); + prefill_runner = state_->get_prefill_runner(prefix_cache.cond_length, dynamic_length, backend_config, *weights_); + } + prefill_output = prefill_runner->run(dynamic_embeddings, prefix_cache.cache); + } + outputs.prefill_runner_ms = std::chrono::duration( + std::chrono::steady_clock::now() - step_started).count(); + + std::vector last_logits = std::move(prefill_output.logits); + const int64_t prefill_cache_steps = prefix_cache.cond_length + dynamic_length; + { + std::lock_guard lock(state_->mutex); + auto runner = state_->get_runner(prefill_cache_steps + request.max_new_tokens, backend_config, *weights_); + runner->import_state(prefill_output.cache); + } + + std::vector logits(static_cast(speech_vocab), 0.0f); + std::vector sampling_probs; + std::vector sampling_order; + std::vector sampling_mask; + for (int64_t step = 0; step < request.max_new_tokens; ++step) { + std::copy_n(last_logits.data(), speech_vocab, logits.data()); + apply_repetition_penalty_in_place(logits, generated_ids, request.repetition_penalty, sampling_mask); + if (request.do_sample) { + if (request.temperature != 1.0f && request.temperature > 0.0f) { + for (float & value : logits) { + value /= request.temperature; + } + } + apply_top_k_in_place(logits, request.top_k); + apply_top_p_in_place(logits, request.top_p, sampling_probs, sampling_order, sampling_mask); + } + const int32_t next_token = sample_from_logits(logits, request.do_sample, rng); + generated_ids.push_back(next_token); + + if (request.stop_on_eos && next_token == kTurboStopSpeechToken) { + outputs.hit_eos = true; + break; + } + + std::vector next_embed(speech_embedding_weight.begin() + static_cast(next_token * hidden_size), + speech_embedding_weight.begin() + static_cast((next_token + 1) * hidden_size)); + const float * pos_row = wpe_weight.data() + + static_cast((prefill_cache_steps + step) * hidden_size); + for (int64_t i = 0; i < hidden_size; ++i) { + next_embed[static_cast(i)] += pos_row[i]; + } + + step_started = std::chrono::steady_clock::now(); + std::shared_ptr runner; + { + std::lock_guard lock(state_->mutex); + runner = state_->runner; + } + last_logits = runner->step(next_embed, prefill_cache_steps + step); + outputs.decode_runner_ms += std::chrono::duration( + std::chrono::steady_clock::now() - step_started).count(); + } + + const size_t initial_count = request.initial_speech_tokens.empty() ? 1 : request.initial_speech_tokens.size(); + if (generated_ids.size() > initial_count) { + outputs.predicted_tokens.assign(generated_ids.begin() + static_cast(initial_count), generated_ids.end()); + } + outputs.token_count = static_cast(outputs.predicted_tokens.size()); + outputs.prefix_cache_build_ms = prefix_cache_build_ms; + return outputs; +} + +void T3TurboInferenceComponent::release_runtime_graphs() const { + state_->release_runtime_graphs(); +} + +void T3TurboInferenceComponent::release_runtime_cache() const { + state_->release_runtime_cache(); +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/t3_turbo_weights.cpp b/src/community_models/chatterbox_turbo/t3_turbo_weights.cpp new file mode 100644 index 000000000..0b46085fc --- /dev/null +++ b/src/community_models/chatterbox_turbo/t3_turbo_weights.cpp @@ -0,0 +1,182 @@ +#include "engine/community_models/chatterbox_turbo/t3_turbo_component.h" + +#include + +namespace engine::community_models::chatterbox_turbo { + +namespace { + +T3TurboGraphWeight load_graph_weight( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & name, + engine::assets::TensorStorageType storage_type, + const std::vector & shape, + bool load_reference_f32_values) { + T3TurboGraphWeight weight; + if (load_reference_f32_values) { + weight.values = source.require_f32(name, shape); + } + weight.tensor = store.load_tensor(source, name, storage_type, shape); + return weight; +} + +engine::assets::TensorStorageType storage_for_shape( + engine::assets::TensorStorageType requested, + const std::vector & shape) { + if (requested == engine::assets::TensorStorageType::Native) { + return requested; + } + const ggml_type type = engine::assets::ggml_type_for_tensor_storage(requested); + if (!ggml_is_quantized(type)) { + return requested; + } + if (shape.size() < 2 || shape.back() % ggml_blck_size(type) != 0) { + return engine::assets::TensorStorageType::F32; + } + return requested; +} + +engine::assets::TensorData require_tensor_data( + const engine::assets::TensorSource & source, + const std::string & name, + engine::assets::TensorStorageType requested, + const std::vector & shape) { + return source.require_tensor(name, storage_for_shape(requested, shape), shape); +} + +int64_t count_turbo_layers(const engine::assets::TensorSource & source) { + int64_t count = 0; + while (source.has_tensor("blk." + std::to_string(count) + ".attn_norm.weight")) { + ++count; + } + if (count == 0) { + throw std::runtime_error("Chatterbox Turbo T3 weights contain no transformer layers"); + } + return count; +} + +} // namespace + +std::shared_ptr load_t3_turbo_inference_weights( + const engine::assets::TensorSource & source, + const engine::core::ExecutionContext & execution_context, + engine::assets::TensorStorageType graph_weight_storage_type, + bool load_reference_f32_graph_weights) { + auto weights = std::make_shared(); + weights->execution_context = &execution_context; + weights->store = std::make_shared( + execution_context.backend(), + execution_context.backend_type(), + "chatterbox_turbo.t3.weights", + 4096ull * 1024ull * 1024ull); + + const auto text_emb_info = source.require_metadata("text_emb.weight"); + const auto speech_emb_info = source.require_metadata("speech_emb.weight"); + const auto wpe_info = source.require_metadata("wpe.weight"); + + weights->hidden_size = text_emb_info.shape.at(1); + weights->text_vocab = text_emb_info.shape.at(0); + weights->speech_vocab = speech_emb_info.shape.at(0); + weights->max_positions = wpe_info.shape.at(0); + weights->speaker_embed_size = source.require_metadata("cond.spkr_enc.weight").shape.at(1); + weights->num_heads = 16; + weights->mlp_intermediate_size = source.require_metadata("blk.0.ffn_fc.weight").shape.at(0); + + const int64_t hidden = weights->hidden_size; + + weights->spkr_enc_weight = require_tensor_data( + source, "cond.spkr_enc.weight", graph_weight_storage_type, {hidden, weights->speaker_embed_size}); + weights->spkr_enc_bias = source.require_tensor("cond.spkr_enc.bias", engine::assets::TensorStorageType::F32, {hidden}); + + weights->text_embedding_weight = + require_tensor_data(source, "text_emb.weight", graph_weight_storage_type, text_emb_info.shape); + weights->speech_embedding_weight = + require_tensor_data(source, "speech_emb.weight", graph_weight_storage_type, speech_emb_info.shape); + weights->wpe_weight = require_tensor_data(source, "wpe.weight", graph_weight_storage_type, wpe_info.shape); + + weights->ln_f_weight = source.require_f32("output_norm.weight", {hidden}); + weights->ln_f_weight_tensor = weights->store->load_f32_tensor(source, "output_norm.weight", {hidden}); + weights->ln_f_bias = source.require_f32("output_norm.bias", {hidden}); + weights->ln_f_bias_tensor = weights->store->load_f32_tensor(source, "output_norm.bias", {hidden}); + + weights->text_head_weight = load_graph_weight( + *weights->store, + source, + "text_head.weight", + graph_weight_storage_type, + source.require_metadata("text_head.weight").shape, + load_reference_f32_graph_weights); + weights->speech_head_weight = load_graph_weight( + *weights->store, + source, + "speech_head.weight", + graph_weight_storage_type, + source.require_metadata("speech_head.weight").shape, + load_reference_f32_graph_weights); + weights->speech_head_bias = source.require_f32("speech_head.bias", {weights->speech_vocab}); + weights->speech_head_bias_tensor = weights->store->load_f32_tensor(source, "speech_head.bias", {weights->speech_vocab}); + + const int64_t num_layers = count_turbo_layers(source); + const int64_t intermediate = weights->mlp_intermediate_size; + weights->layers.resize(static_cast(num_layers)); + for (int64_t layer = 0; layer < num_layers; ++layer) { + const std::string prefix = "blk." + std::to_string(layer); + auto & out = weights->layers[static_cast(layer)]; + + out.ln1_weight = source.require_f32(prefix + ".attn_norm.weight", {hidden}); + out.ln1_weight_tensor = weights->store->load_f32_tensor(source, prefix + ".attn_norm.weight", {hidden}); + out.ln1_bias = source.require_f32(prefix + ".attn_norm.bias", {hidden}); + out.ln1_bias_tensor = weights->store->load_f32_tensor(source, prefix + ".attn_norm.bias", {hidden}); + + out.attn_qkv_weight = load_graph_weight( + *weights->store, + source, + prefix + ".attn_qkv.weight", + graph_weight_storage_type, + {3 * hidden, hidden}, + load_reference_f32_graph_weights); + out.attn_qkv_bias = source.require_f32(prefix + ".attn_qkv.bias", {3 * hidden}); + out.attn_qkv_bias_tensor = weights->store->load_f32_tensor(source, prefix + ".attn_qkv.bias", {3 * hidden}); + + out.attn_output_weight = load_graph_weight( + *weights->store, + source, + prefix + ".attn_output.weight", + graph_weight_storage_type, + {hidden, hidden}, + load_reference_f32_graph_weights); + out.attn_output_bias = source.require_f32(prefix + ".attn_output.bias", {hidden}); + out.attn_output_bias_tensor = weights->store->load_f32_tensor(source, prefix + ".attn_output.bias", {hidden}); + + out.ln2_weight = source.require_f32(prefix + ".ffn_norm.weight", {hidden}); + out.ln2_weight_tensor = weights->store->load_f32_tensor(source, prefix + ".ffn_norm.weight", {hidden}); + out.ln2_bias = source.require_f32(prefix + ".ffn_norm.bias", {hidden}); + out.ln2_bias_tensor = weights->store->load_f32_tensor(source, prefix + ".ffn_norm.bias", {hidden}); + + out.ffn_fc_weight = load_graph_weight( + *weights->store, + source, + prefix + ".ffn_fc.weight", + graph_weight_storage_type, + {intermediate, hidden}, + load_reference_f32_graph_weights); + out.ffn_fc_bias = source.require_f32(prefix + ".ffn_fc.bias", {intermediate}); + out.ffn_fc_bias_tensor = weights->store->load_f32_tensor(source, prefix + ".ffn_fc.bias", {intermediate}); + + out.ffn_proj_weight = load_graph_weight( + *weights->store, + source, + prefix + ".ffn_proj.weight", + graph_weight_storage_type, + {hidden, intermediate}, + load_reference_f32_graph_weights); + out.ffn_proj_bias = source.require_f32(prefix + ".ffn_proj.bias", {hidden}); + out.ffn_proj_bias_tensor = weights->store->load_f32_tensor(source, prefix + ".ffn_proj.bias", {hidden}); + } + + weights->store->upload(); + return weights; +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/text_tokenizer_turbo.cpp b/src/community_models/chatterbox_turbo/text_tokenizer_turbo.cpp new file mode 100644 index 000000000..d485fd8eb --- /dev/null +++ b/src/community_models/chatterbox_turbo/text_tokenizer_turbo.cpp @@ -0,0 +1,96 @@ +#include "engine/community_models/chatterbox_turbo/text_tokenizer_turbo.h" + +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +std::shared_ptr load_chatterbox_turbo_tokenizer( + const std::filesystem::path & vocab_path, + const std::filesystem::path & merges_path, + const std::filesystem::path & special_tokens_path) { + std::vector special_tokens; + const auto special_tokens_json = engine::io::json::parse_file(special_tokens_path); + for (const auto & entry : special_tokens_json.as_array()) { + special_tokens.emplace_back( + engine::io::json::require_string(entry, "token"), + static_cast(engine::io::json::require_i64(entry, "id"))); + } + + engine::tokenizers::LlamaBpeTokenizerSpec spec( + vocab_path, + merges_path, + /*tokenizer_config_path=*/{}, + /*tokenizer_json_path=*/std::nullopt, + engine::tokenizers::LlamaBpePreTokenizer::Gpt2, + std::move(special_tokens)); + return engine::tokenizers::load_llama_bpe_tokenizer(spec); +} + +std::string chatterbox_turbo_punc_norm(const std::string & text) { + if (text.empty()) { + return "You need to add some text for me to talk."; + } + std::string normalized = text; + if (std::islower(static_cast(normalized.front()))) { + normalized.front() = static_cast(std::toupper(static_cast(normalized.front()))); + } + + // Collapse runs of whitespace to single spaces (Python's " ".join(text.split())). + { + std::string collapsed; + collapsed.reserve(normalized.size()); + bool in_space = false; + for (char c : normalized) { + const bool is_space = std::isspace(static_cast(c)) != 0; + if (is_space) { + in_space = true; + continue; + } + if (in_space && !collapsed.empty()) { + collapsed += ' '; + } + in_space = false; + collapsed += c; + } + normalized = std::move(collapsed); + } + + static const std::vector> kReplacements = { + {"\xE2\x80\xA6", ", "}, // U+2026 HORIZONTAL ELLIPSIS + {":", ","}, + {"\xE2\x80\x94", "-"}, // U+2014 EM DASH + {"\xE2\x80\x93", "-"}, // U+2013 EN DASH + {" ,", ","}, + {"\xE2\x80\x9C", "\""}, // U+201C LEFT DOUBLE QUOTATION MARK + {"\xE2\x80\x9D", "\""}, // U+201D RIGHT DOUBLE QUOTATION MARK + {"\xE2\x80\x98", "'"}, // U+2018 LEFT SINGLE QUOTATION MARK + {"\xE2\x80\x99", "'"}, // U+2019 RIGHT SINGLE QUOTATION MARK + }; + for (const auto & [from, to] : kReplacements) { + size_t pos = 0; + while ((pos = normalized.find(from, pos)) != std::string::npos) { + normalized.replace(pos, from.size(), to); + pos += to.size(); + } + } + + while (!normalized.empty() && normalized.back() == ' ') { + normalized.pop_back(); + } + static const std::unordered_set kSentenceEnders = {'.', '!', '?', '-', ','}; + if (normalized.empty() || kSentenceEnders.find(normalized.back()) == kSentenceEnders.end()) { + normalized += '.'; + } + return normalized; +} + +std::vector encode_chatterbox_turbo_text( + const engine::tokenizers::LlamaBpeTokenizer & tokenizer, + const std::string & text) { + return tokenizer.encode(chatterbox_turbo_punc_norm(text), /*parse_special=*/true); +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/community_models/chatterbox_turbo/tts.cpp b/src/community_models/chatterbox_turbo/tts.cpp new file mode 100644 index 000000000..90f41a1e3 --- /dev/null +++ b/src/community_models/chatterbox_turbo/tts.cpp @@ -0,0 +1,97 @@ +#include "engine/community_models/chatterbox_turbo/tts.h" + +#include "engine/community_models/chatterbox_turbo/text_tokenizer_turbo.h" + +#include +#include + +namespace engine::community_models::chatterbox_turbo { + +namespace { + +// Chatterbox Turbo's speech-token control ids (tts_turbo.py / s3gen/const.py). Not exposed via +// GGUF metadata, so mirrored here from the upstream Python reference. +constexpr int32_t kS3GenSilenceToken = 4299; +constexpr int32_t kOovTokenThreshold = 6561; // speech_tokens = speech_tokens[speech_tokens < 6561] + +std::vector require_i32_array( + const engine::assets::TensorSource & source, + const std::string & name, + int64_t expected_count) { + const auto raw = source.require_tensor_data(name); + if (raw.metadata.dtype != "i32") { + throw std::runtime_error("Chatterbox Turbo expected an i32 tensor for " + name + ", got " + raw.metadata.dtype); + } + if (raw.bytes.size() != static_cast(expected_count) * sizeof(int32_t)) { + throw std::runtime_error("Chatterbox Turbo tensor size mismatch for " + name); + } + std::vector out(static_cast(expected_count)); + std::memcpy(out.data(), raw.bytes.data(), raw.bytes.size()); + return out; +} + +} // namespace + +ChatterboxTurboTtsComponent::ChatterboxTurboTtsComponent( + std::shared_ptr assets, + const engine::core::ExecutionContext & execution_context) + : assets_(std::move(assets)) { + tokenizer_ = load_chatterbox_turbo_tokenizer( + assets_->resources.require_file("tokenizer_vocab"), + assets_->resources.require_file("tokenizer_merges"), + assets_->resources.require_file("tokenizer_special_tokens")); + + auto t3_weights = load_t3_turbo_inference_weights(*assets_->t3_turbo_weights, execution_context); + t3_ = std::make_unique(t3_weights, execution_context); + + s3gen_ = ChatterboxTurboS3Gen::load(assets_->s3gen_weights, execution_context); + + const auto & conds = *assets_->builtin_conditionals_turbo; + builtin_speaker_embedding_ = conds.require_f32("t3.speaker_emb", {t3_weights->speaker_embed_size}); + builtin_cond_prompt_speech_tokens_ = require_i32_array( + conds, "t3.speech_prompt_tokens", static_cast(conds.require_metadata("t3.speech_prompt_tokens").shape.at(0))); + + const auto prompt_token_shape = conds.require_metadata("gen.prompt_token").shape; + const auto prompt_feat_shape = conds.require_metadata("gen.prompt_feat").shape; // [frames, mel_dim] + const auto embedding_shape = conds.require_metadata("gen.embedding").shape; + + builtin_ref_dict_.prompt_tokens = require_i32_array(conds, "gen.prompt_token", prompt_token_shape.at(0)); + builtin_ref_dict_.prompt_token_count = prompt_token_shape.at(0); + builtin_ref_dict_.prompt_feat = conds.require_f32("gen.prompt_feat", prompt_feat_shape); + builtin_ref_dict_.prompt_feat_frames = prompt_feat_shape.at(0); + builtin_ref_dict_.prompt_feat_dims = prompt_feat_shape.at(1); + builtin_ref_dict_.embedding = conds.require_f32("gen.embedding", embedding_shape); + builtin_ref_dict_.embedding_size = embedding_shape.at(0); +} + +engine::models::chatterbox::S3GenInferenceOutputs ChatterboxTurboTtsComponent::generate( + const std::string & text, + const ChatterboxTurboGenerateConfig & config) const { + T3TurboGenerateRequest request; + request.speaker_embedding = builtin_speaker_embedding_; + request.cond_prompt_speech_tokens = builtin_cond_prompt_speech_tokens_; + request.text_tokens = encode_chatterbox_turbo_text(*tokenizer_, text); + request.max_new_tokens = config.max_new_tokens; + request.temperature = config.temperature; + request.top_p = config.top_p; + request.top_k = config.top_k; + request.repetition_penalty = config.repetition_penalty; + request.seed = config.seed; + + const auto t3_outputs = t3_->generate_speech_tokens(request); + + std::vector speech_tokens; + speech_tokens.reserve(t3_outputs.predicted_tokens.size() + 3); + for (int32_t token : t3_outputs.predicted_tokens) { + if (token < kOovTokenThreshold) { + speech_tokens.push_back(token); + } + } + speech_tokens.push_back(kS3GenSilenceToken); + speech_tokens.push_back(kS3GenSilenceToken); + speech_tokens.push_back(kS3GenSilenceToken); + + return s3gen_->synthesize(builtin_ref_dict_, speech_tokens, config.seed, config.seed); +} + +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/models/chatterbox/components/s3gen_weights.h b/src/models/chatterbox/components/s3gen_weights.h index ede7aa5f5..82402fd78 100644 --- a/src/models/chatterbox/components/s3gen_weights.h +++ b/src/models/chatterbox/components/s3gen_weights.h @@ -261,6 +261,12 @@ struct S3FlowDecoderWeights { std::vector up_blocks; CausalBlockWeights final_block; Conv1dWeights final_proj; + // Meanflow-distilled decoders (Chatterbox Turbo) mix a second "r" (end-time) sinusoidal + // embedding into the time embedding via this diagonal-init, no-bias linear layer before + // feeding the UNet1D estimator; see decoder.py::get_intmeanflow_time_mixer upstream. Unset + // (weight_tensor.tensor == nullptr) for the base Chatterbox 10-step CFG decoder. + bool meanflow = false; + LinearWeights time_embed_mixer; const engine::core::ExecutionContext * execution_context = nullptr; std::shared_ptr store; }; diff --git a/src/models/chatterbox/s3gen_flow.cpp b/src/models/chatterbox/s3gen_flow.cpp index deb9e5fc5..e80912992 100644 --- a/src/models/chatterbox/s3gen_flow.cpp +++ b/src/models/chatterbox/s3gen_flow.cpp @@ -1497,6 +1497,19 @@ class FlowDecoderBackendRunner { auto time_hidden = linear_lastdim(ctx, time_in_, weights.time_mlp_1, writer); time_hidden = engine::core::wrap_tensor(ggml_silu(ctx.ggml, time_hidden.tensor), time_hidden.shape, GGML_TYPE_F32); time_hidden = linear_lastdim(ctx, time_hidden, weights.time_mlp_2, writer); + if (weights.meanflow) { + r_in_ = engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims({batch_, 320})); + ggml_set_input(r_in_.tensor); + ggml_set_output(r_in_.tensor); + auto r_hidden = linear_lastdim(ctx, r_in_, weights.time_mlp_1, writer); + r_hidden = engine::core::wrap_tensor(ggml_silu(ctx.ggml, r_hidden.tensor), r_hidden.shape, GGML_TYPE_F32); + r_hidden = linear_lastdim(ctx, r_hidden, weights.time_mlp_2, writer); + auto concat_shape = time_hidden.shape; + concat_shape.dims[concat_shape.rank - 1] += r_hidden.shape.dims[r_hidden.shape.rank - 1]; + auto concat_te = engine::core::wrap_tensor( + ggml_concat(ctx.ggml, time_hidden.tensor, r_hidden.tensor, 0), concat_shape, GGML_TYPE_F32); + time_hidden = linear_lastdim(ctx, concat_te, weights.time_embed_mixer, writer); + } auto hidden_bct = cat_channels_bct(ctx, x_in_, mu_in_); auto repeat_spks = repeat_spks_bct(ctx, spks_in_, frames); @@ -1567,7 +1580,8 @@ class FlowDecoderBackendRunner { const std::vector & x, const std::vector & mask, const std::vector & t, - S3FlowDecoderRunTiming * timing = nullptr) { + S3FlowDecoderRunTiming * timing = nullptr, + const std::vector * r = nullptr) { constexpr int64_t mel_channels = 80; constexpr int64_t time_dim = 320; if (timing != nullptr) { @@ -1587,6 +1601,12 @@ class FlowDecoderBackendRunner { } started = Clock::now(); engine::core::write_tensor_f32(time_in_, time_emb_values); + if (r != nullptr) { + if (r_in_.tensor == nullptr) { + throw std::runtime_error("S3 flow decoder runner was not built for meanflow but received r"); + } + engine::core::write_tensor_f32(r_in_, decoder_sinusoidal_pos_emb(*r, batch_, time_dim)); + } if (timing != nullptr) { timing->input_write_ms += engine::debug::elapsed_ms(started); } @@ -1631,6 +1651,7 @@ class FlowDecoderBackendRunner { engine::core::TensorValue x_in_; engine::core::TensorValue mu_in_; engine::core::TensorValue time_in_; + engine::core::TensorValue r_in_; engine::core::TensorValue spks_in_; engine::core::TensorValue cond_in_; engine::core::TensorValue attention_mask_; @@ -1740,10 +1761,14 @@ std::shared_ptr load_s3_flow_encoder_weights( execution_context.backend_type(), "chatterbox.s3_flow_encoder.weights", 1024ull * 1024ull * 1024ull); + // ggml_get_rows (used by S3TokenEmbeddingGraph for this lookup table) only supports F32/F16 + // and legacy quant types on the CUDA backend, not K-quants -- pin this one small table + // (~6.7 MB in F16) to F16 regardless of the requested weight_storage_type so a Q4_K/Q5_K/... + // package (e.g. Chatterbox Turbo's chatterbox-turbo-s3gen-q4_k.gguf) doesn't crash on it. weights->input_embedding_tensor = weights->store->load_tensor( source, "flow.input_embedding.weight", - weight_storage_type, + engine::assets::TensorStorageType::F16, {6561, 512}); weights->speaker_affine = load_flow_linear(*weights->store, source, "flow.spk_embed_affine_layer", 80, 192, true, weight_storage_type); weights->encoder_proj = load_flow_linear(*weights->store, source, "flow.encoder_proj", 80, 512, true, weight_storage_type); @@ -1867,10 +1892,21 @@ std::shared_ptr load_s3_flow_decoder_weights( weights->final_block = load_causal_block("flow.decoder.estimator.final_block", 256, 256); weights->final_proj = load_decoder_conv1d(*weights->store, source, "flow.decoder.estimator.final_proj", 80, 256, 1, 1, true, weight_storage_type); + + if (source.has_tensor("flow.decoder.estimator.time_embed_mixer.weight")) { + weights->meanflow = true; + weights->time_embed_mixer = load_decoder_linear( + *weights->store, source, "flow.decoder.estimator.time_embed_mixer", 1024, 2048, false, weight_storage_type); + } + weights->store->upload(); return weights; } +bool s3_flow_decoder_is_meanflow(const S3FlowDecoderWeights & weights) { + return weights.meanflow; +} + S3FlowDecoderOutputs compute_s3_flow_decoder_forward( S3FlowSessionCache & cache, const S3FlowDecoderWeights & weights, @@ -1883,11 +1919,12 @@ S3FlowDecoderOutputs compute_s3_flow_decoder_forward( int64_t batch, int64_t frames, int64_t capacity_frames, - engine::core::BackendConfig backend) { + engine::core::BackendConfig backend, + const std::vector * r) { const int64_t valid_frames = valid_frames_from_mask(mask, batch, frames); auto & runner = cache.state_->decoder_runner_for_capacity(weights, batch, capacity_frames, backend); runner.set_conditioning(weights, mu, spks, cond, frames, backend); - auto outputs = runner.run(x, mask, t); + auto outputs = runner.run(x, mask, t, nullptr, r); outputs.frames = valid_frames; return outputs; } @@ -2003,4 +2040,55 @@ S3FlowCFMOutputs compute_s3_flow_cfm_euler( return outputs; } +// Meanflow-distilled decoders (Chatterbox Turbo) were trained with CFG already baked in, so +// inference is a plain (non-CFG, non-batch-doubled) Euler solve driven by two time inputs per +// step (t, r); mirrors ConditionalCFM.basic_euler in the upstream Python reference. +S3FlowCFMOutputs compute_s3_flow_cfm_meanflow( + S3FlowSessionCache & cache, + const S3FlowDecoderWeights & weights, + const std::vector & noise, + const std::vector & mask, + const std::vector & mu, + const std::vector & spks, + const std::vector & cond, + int64_t batch, + int64_t frames, + int64_t capacity_frames, + int64_t num_steps, + engine::core::BackendConfig backend) { + constexpr int64_t mel_channels = 80; + if (!weights.meanflow) { + throw std::runtime_error("compute_s3_flow_cfm_meanflow requires meanflow-trained S3FlowDecoderWeights"); + } + std::vector x = noise; + std::vector t_span(static_cast(num_steps + 1), 0.0f); + for (int64_t i = 0; i <= num_steps; ++i) { + t_span[static_cast(i)] = static_cast(i) / static_cast(num_steps); + } + const int64_t valid_frames = valid_frames_from_mask(mask, batch, frames); + + auto & runner = cache.state_->decoder_runner_for_capacity(weights, batch, capacity_frames, backend); + runner.set_conditioning(weights, mu, spks, cond, frames, backend); + std::vector t_vec(static_cast(batch), 0.0f); + std::vector r_vec(static_cast(batch), 0.0f); + for (int64_t step = 0; step < num_steps; ++step) { + const float t0 = t_span[static_cast(step)]; + const float t1 = t_span[static_cast(step + 1)]; + const float dt = t1 - t0; + std::fill(t_vec.begin(), t_vec.end(), t0); + std::fill(r_vec.begin(), r_vec.end(), t1); + const auto outputs = runner.run(x, mask, t_vec, nullptr, &r_vec); + for (size_t i = 0; i < x.size(); ++i) { + x[i] += dt * outputs.mel[i]; + } + } + + S3FlowCFMOutputs outputs; + outputs.mel = std::move(x); + outputs.channels = mel_channels; + outputs.frames = valid_frames; + outputs.storage_frames = frames; + return outputs; +} + } // namespace engine::models::chatterbox diff --git a/src/models/chatterbox/s3gen_inference.cpp b/src/models/chatterbox/s3gen_inference.cpp index 3b6a1677a..9226ce0dc 100644 --- a/src/models/chatterbox/s3gen_inference.cpp +++ b/src/models/chatterbox/s3gen_inference.cpp @@ -539,22 +539,40 @@ S3Token2MelOutputs compute_s3_token2mel_inference( } cache.state_->release_pre_cfm_graphs(); const auto cfm_started = std::chrono::steady_clock::now(); - const auto mel_full = compute_s3_flow_cfm_euler( - cache.state_->flow_cache, - decoder_weights, - noise, - mask, - mu, - prepared.speaker, - cond, - batch, - total_frames, - frame_capacity, - num_steps, - cfg_rate, - cosine_schedule, - backend, - timing == nullptr ? nullptr : &timing->token2mel_cfm); + // Meanflow-distilled decoders (Chatterbox Turbo) use a plain, non-CFG Euler solve driven by + // two time inputs per step; see S3FlowDecoderWeights::meanflow and + // compute_s3_flow_cfm_meanflow. cfg_rate/cosine_schedule are meaningless for that path and + // are ignored, matching upstream's ConditionalCFM.forward branching on `meanflow`. + const auto mel_full = decoder_weights.meanflow + ? compute_s3_flow_cfm_meanflow( + cache.state_->flow_cache, + decoder_weights, + noise, + mask, + mu, + prepared.speaker, + cond, + batch, + total_frames, + frame_capacity, + num_steps, + backend) + : compute_s3_flow_cfm_euler( + cache.state_->flow_cache, + decoder_weights, + noise, + mask, + mu, + prepared.speaker, + cond, + batch, + total_frames, + frame_capacity, + num_steps, + cfg_rate, + cosine_schedule, + backend, + timing == nullptr ? nullptr : &timing->token2mel_cfm); if (timing != nullptr) { timing->token2mel_cfm_ms = std::chrono::duration(std::chrono::steady_clock::now() - cfm_started).count(); diff --git a/tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py b/tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py new file mode 100644 index 000000000..50da28bf2 --- /dev/null +++ b/tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +"""Repack the third-party `cstr/chatterbox-turbo-GGUF` pair into one audio.cpp-native, +self-contained GGUF. + +The upstream package (published by `cstr` for the CrispASR project, MIT-relicensed; not +published by ResembleAI or audio.cpp) ships as two loose GGUF files with a flat dot-separated +tensor namespace and abbreviated S3Gen tensor names. This script repacks both into audio.cpp's +own "/"-delimited packed-GGUF namespace convention, renames S3Gen tensors back to the exact names +base Chatterbox's own S3Gen/HiFT-vocoder loader code already expects (so that loader runs +unmodified, no runtime name-translation bridge needed), and produces the tokenizer as plain +vocab/merges/special-token sidecar files instead of runtime GGUF-metadata reads. The actual GGUF +writing (quantization, embedded package-spec metadata, embedded sidecars) is done by this +project's own `audiocpp_gguf` converter -- this script only stages inputs for it, exactly like +tools/community_models/convert_voxcpm1.py does for VoxCPM1. + +Only the T3 GPT2 backbone, the built-in default-voice conditionals, and the S3Gen flow/HiFT +vocoder are repacked. The upstream GGUF's `ve.*` (LSTM speaker-verification voice encoder) and +`s3.se.*`/`s3.tok.*` (ResNet speaker encoder / S3 speech tokenizer) sections are not staged: they +are not read by any loader in this codebase yet (custom voice cloning -- a caller-supplied +reference clip -- is not implemented; only the built-in baked-in voice is), and their tensor +layout has not been validated. +""" +import argparse +import json +import subprocess +from pathlib import Path + +import numpy as np +from gguf import GGUFReader +from gguf.quants import dequantize +from safetensors.numpy import save_file + + +# --------------------------------------------------------------------------------------------- +# S3Gen tensor-name mapping. This is the exact inverse of the rule tables that used to live in +# src/community_models/chatterbox_turbo/components/s3gen_name_bridge.cpp (deleted -- the repacked +# GGUF no longer needs a runtime bridge). Keep this in sync by hand if that history is ever +# revisited; the rules are transcribed here, not derived programmatically, to keep this script +# self-contained. +# --------------------------------------------------------------------------------------------- + +FLOW_PREFIX_RULES = [ + ("flow.encoder.encoders.", "fe.enc."), + ("flow.encoder.up_encoders.", "fe.ue."), + ("flow.encoder.up_layer.", "fe.ul."), + ("flow.encoder.up_embed.", "fe.uemb."), + ("flow.encoder.pre_lookahead_layer.", "fe.pla."), + ("flow.encoder.embed.", "fe.embed."), + ("flow.encoder.after_norm", "fe.an"), + ("flow.decoder.estimator.down_blocks.", "fd.db."), + ("flow.decoder.estimator.mid_blocks.", "fd.mb."), + ("flow.decoder.estimator.up_blocks.", "fd.ub."), + ("flow.decoder.estimator.final_block", "fd.fb"), + ("flow.decoder.estimator.final_proj", "fd.fp"), + ("flow.decoder.estimator.time_mlp.", "fd.tm."), + ("flow.decoder.estimator.time_embed_mixer", "fd.tmx"), + # flow.encoder_proj / flow.spk_embed_affine_layer / flow.input_embedding: unchanged. +] + +SEGMENT_RULES = [ + (".self_attn.linear_q", ".sa.lq"), + (".self_attn.linear_k", ".sa.lk"), + (".self_attn.linear_v", ".sa.lv"), + (".self_attn.linear_out", ".sa.lo"), + (".self_attn.linear_pos", ".sa.lp"), + (".self_attn.pos_bias_u", ".sa.pbu"), + (".self_attn.pos_bias_v", ".sa.pbv"), + (".norm_mha", ".nmha"), + (".norm_ff", ".nff"), + (".feed_forward.w_1", ".ff.w_1"), + (".feed_forward.w_2", ".ff.w_2"), + (".block1.block.0", ".b1.0"), + (".block1.block.2", ".b1.2"), + (".block2.block.0", ".b2.0"), + (".block2.block.2", ".b2.2"), +] + +TRANSFORMER_BLOCK_RULES = [ + (".res_conv", ".rc"), + (".attn1.to_q", ".attn1.q"), + (".attn1.to_k", ".attn1.k"), + (".attn1.to_v", ".attn1.v"), + (".attn1.to_out.0", ".attn1.o"), + (".ff.net.0.proj", ".ff.up"), +] + +VOCODER_PREFIX_RULES = [ + ("conv_pre", "cpre"), + ("conv_post", "cpost"), + ("resblocks.", "rb."), + ("source_downs.", "sd."), + ("source_resblocks.", "srb."), + ("f0_predictor.condnet.", "f0.cn."), +] + + +def reverse_translate_flow_name(name): + """abbrev (post 's3.'-strip, no leading 'v.') -> canonical base-chatterbox flow name.""" + translated = name + # Forward order was: prefix, then segment rules, then ".ff.net.2"->".ff.down", then + # transformer-block rules. Invert in the opposite order. + for canonical, abbrev in TRANSFORMER_BLOCK_RULES: + translated = translated.replace(abbrev, canonical) + translated = translated.replace(".ff.down", ".ff.net.2") + for canonical, abbrev in SEGMENT_RULES: + translated = translated.replace(abbrev, canonical) + for canonical, abbrev in FLOW_PREFIX_RULES: + if translated.startswith(abbrev): + translated = canonical + translated[len(abbrev):] + break + return translated + + +def reverse_translate_vocoder_name(name): + """abbrev (post 's3.v.'-strip) -> canonical HiFT vocoder tensor name (without the 'v.' the + C++ loader's HiftVocoderConfig.tensor_prefix adds back on).""" + translated = name + if translated == "f0.cls" or translated.startswith("f0.cls."): + translated = "f0_predictor.classifier" + translated[len("f0.cls"):] + elif translated == "ms.ll" or translated.startswith("ms.ll."): + translated = "m_source.l_linear" + translated[len("ms.ll"):] + else: + translated = translated.replace(".c1.", ".convs1.") + translated = translated.replace(".c2.", ".convs2.") + translated = translated.replace(".a1.", ".activations1.") + translated = translated.replace(".a2.", ".activations2.") + for canonical, abbrev in VOCODER_PREFIX_RULES: + if translated.startswith(abbrev): + translated = canonical + translated[len(abbrev):] + break + return translated + + +def reverse_translate_s3gen_name(name_without_s3_prefix): + if name_without_s3_prefix.startswith("v."): + return "v." + reverse_translate_vocoder_name(name_without_s3_prefix[len("v."):]) + return reverse_translate_flow_name(name_without_s3_prefix) + + +# --------------------------------------------------------------------------------------------- + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Repack the cstr/chatterbox-turbo-GGUF pair into an audio.cpp-native GGUF." + ) + parser.add_argument( + "--t3-source", + type=Path, + required=True, + help="Path to the upstream chatterbox-turbo-t3-*.gguf file.", + ) + parser.add_argument( + "--s3gen-source", + type=Path, + required=True, + help="Path to the upstream chatterbox-turbo-s3gen-*.gguf file.", + ) + parser.add_argument( + "--output", + type=Path, + required=True, + help="Output native GGUF path.", + ) + parser.add_argument( + "--converter", + type=Path, + default=Path("build/debug/bin/audiocpp_gguf"), + help="Path to the audiocpp_gguf converter.", + ) + parser.add_argument( + "--model-spec", + type=Path, + default=Path("model_specs/chatterbox_turbo.json"), + help="Chatterbox Turbo model spec JSON.", + ) + parser.add_argument( + "--work-dir", + type=Path, + default=Path("build/chatterbox_turbo_native_gguf"), + help="Directory for staged safetensors and sidecars.", + ) + parser.add_argument("--type", default="q8_0", help="Main model conversion type.") + parser.add_argument("--overwrite", action="store_true", help="Overwrite output GGUF.") + return parser.parse_args() + + +def require_file(path): + if not path.is_file(): + raise FileNotFoundError(path) + return path + + +def read_tensor_float(reader_tensor): + return dequantize(reader_tensor.data, reader_tensor.tensor_type).astype(np.float32) + + +def stage_t3(t3_reader, staging): + t3_tensors = {} + conds_tensors = {} + for tensor in t3_reader.tensors: + name = tensor.name + if name.startswith("t3."): + key = name[len("t3."):] + bucket = t3_tensors + elif name.startswith("conds."): + key = name[len("conds."):] + bucket = conds_tensors + elif name.startswith("ve."): + continue # LSTM speaker-verification encoder: unused, out of scope. + else: + raise KeyError(f"unrecognized top-level T3 GGUF tensor namespace: {name}") + + if key in ("gen.prompt_token", "t3.speech_prompt_tokens"): + # i32 speech-token id arrays: keep integer dtype, do not dequantize. tensor.data is + # already a properly shaped (not reversed) numpy view for GGUFReader tensors. + bucket[key] = np.array(tensor.data, dtype=np.int32) + else: + bucket[key] = read_tensor_float(tensor) + + save_file(t3_tensors, staging / "t3.safetensors") + save_file(conds_tensors, staging / "conds.safetensors") + return t3_tensors, conds_tensors + + +def stage_s3gen(s3gen_reader, staging): + s3gen_tensors = {} + for tensor in s3gen_reader.tensors: + name = tensor.name + if not name.startswith("s3."): + raise KeyError(f"unrecognized top-level S3Gen GGUF tensor namespace: {name}") + suffix = name[len("s3."):] + if suffix.startswith("se.") or suffix.startswith("tok."): + continue # ResNet speaker encoder / S3 speech tokenizer: unused, out of scope. + + mapped = reverse_translate_s3gen_name(suffix) + if mapped in s3gen_tensors: + raise RuntimeError(f"duplicate mapped S3Gen tensor: {mapped} (from {name})") + value = read_tensor_float(tensor) + + # Known shape fix: every nn.Linear(in, 1).weight in this checkpoint (out_features=1) is + # squeezed to 1-D [in] in the upstream GGUF; base Chatterbox's HiFT vocoder loader + # (load_linear in src/framework/modules/vocoders/hift_vocoder.cpp) does a strict + # {out_features, in_features} shape check, so restore the canonical 2-D [1, in] shape. + if mapped in ("v.f0_predictor.classifier.weight", "v.m_source.l_linear.weight") and value.ndim == 1: + value = value.reshape(1, -1) + + s3gen_tensors[mapped] = value + + save_file(s3gen_tensors, staging / "s3gen.safetensors") + return s3gen_tensors + + +def looks_like_bracket_tag(token): + return len(token) >= 3 and token.startswith("[") and token.endswith("]") + + +def stage_tokenizer(t3_reader, staging): + tokens = t3_reader.get_field("tokenizer.ggml.tokens").contents() + merges = t3_reader.get_field("tokenizer.ggml.merges").contents() + if not tokens: + raise RuntimeError("T3 GGUF has an empty tokenizer.ggml.tokens array") + + # Exact port of the trailing-bracket-tag scan that used to run at C++ load time in + # text_tokenizer_turbo.cpp: the 19 emotion/style control tags ([laugh], [sigh], ...) sit at + # the tail of the vocab and were never trained into the BPE merge table, so they must be + # registered as atomic special tokens rather than plain vocab entries. + special_tokens_start = len(tokens) + for token_id in range(len(tokens) - 1, -1, -1): + if not looks_like_bracket_tag(tokens[token_id]): + break + special_tokens_start = token_id + + vocab = {token: token_id for token_id, token in enumerate(tokens) if token_id < special_tokens_start} + special_tokens = [ + {"token": tokens[token_id], "id": token_id} + for token_id in range(special_tokens_start, len(tokens)) + ] + + with open(staging / "chatterbox_turbo_vocab.json", "w", encoding="utf-8") as handle: + json.dump(vocab, handle, ensure_ascii=False) + with open(staging / "chatterbox_turbo_merges.txt", "w", encoding="utf-8") as handle: + for merge in merges: + handle.write(merge + "\n") + with open(staging / "chatterbox_turbo_special_tokens.json", "w", encoding="utf-8") as handle: + json.dump(special_tokens, handle, ensure_ascii=False, indent=2) + + +def run_converter(args, staging): + command = [ + str(args.converter), + "--input", f"t3={staging / 't3.safetensors'}", + "--input", f"conds={staging / 'conds.safetensors'}", + "--input", f"s3gen={staging / 's3gen.safetensors'}", + "--root", str(staging), + "--output", str(args.output), + "--type", args.type, + "--keep-type", "conds/gen.prompt_token=orig", + "--keep-type", "conds/t3.speech_prompt_tokens=orig", + "--family", "chatterbox_turbo", + "--model-spec", str(args.model_spec), + ] + if args.overwrite: + command.append("--overwrite") + subprocess.run(command, check=True) + + +def main(): + args = parse_args() + require_file(args.t3_source) + require_file(args.s3gen_source) + require_file(args.converter) + require_file(args.model_spec) + + staging = args.work_dir / "staging" + staging.mkdir(parents=True, exist_ok=True) + args.output.parent.mkdir(parents=True, exist_ok=True) + + t3_reader = GGUFReader(str(args.t3_source)) + s3gen_reader = GGUFReader(str(args.s3gen_source)) + + stage_t3(t3_reader, staging) + stage_s3gen(s3gen_reader, staging) + stage_tokenizer(t3_reader, staging) + + run_converter(args, staging) + print(f"wrote {args.output}") + + +if __name__ == "__main__": + main()