From 34f06fcd9d8ee56be3c5dc1a9351493d00152953 Mon Sep 17 00:00:00 2001 From: pannagaps Date: Wed, 2 Sep 2026 09:58:29 +0530 Subject: [PATCH 1/2] feat: add Chatterbox Turbo TTS model family Adds chatterbox_turbo as a new model family: Resemble AI's distilled 350M-parameter Chatterbox variant with a GPT2-style T3 backbone and a 2-step meanflow-distilled S3Gen decoder, loaded from cstr's third-party GGUF conversion (built-in voice only; custom voice cloning is not yet implemented). Extends the shared chatterbox S3Gen flow decoder to support meanflow distillation (a second "r" end-time embedding and a non-CFG Euler solve via compute_s3_flow_cfm_meanflow), gated behind the presence of a time_embed_mixer tensor so the base Chatterbox 10-step CFG decoder is unaffected. Also pins the S3 flow encoder's token embedding table to F16 regardless of requested storage type, since ggml_get_rows on CUDA doesn't support K-quant lookups for that table. Co-Authored-By: Claude Sonnet 5 --- CMakeLists.txt | 20 + README.md | 1 + docs/tts.md | 34 + include/engine/models/chatterbox/s3gen_flow.h | 41 +- .../engine/models/chatterbox_turbo/assets.h | 26 + .../engine/models/chatterbox_turbo/loader.h | 33 + .../models/chatterbox_turbo/s3gen_turbo.h | 55 ++ .../engine/models/chatterbox_turbo/session.h | 34 + .../chatterbox_turbo/t3_turbo_component.h | 36 + .../models/chatterbox_turbo/t3_turbo_types.h | 97 +++ .../chatterbox_turbo/text_tokenizer_turbo.h | 31 + include/engine/models/chatterbox_turbo/tts.h | 53 ++ model_specs/chatterbox_turbo.json | 86 ++ .../chatterbox/components/s3gen_weights.h | 6 + src/models/chatterbox/s3gen_flow.cpp | 96 ++- src/models/chatterbox/s3gen_inference.cpp | 50 +- src/models/chatterbox_turbo/assets.cpp | 29 + .../components/dot_prefix_source.cpp | 56 ++ .../components/dot_prefix_source.h | 26 + .../components/s3gen_name_bridge.cpp | 207 +++++ .../components/s3gen_name_bridge.h | 61 ++ .../components/t3_turbo_runtime.h | 811 ++++++++++++++++++ src/models/chatterbox_turbo/loader.cpp | 128 +++ src/models/chatterbox_turbo/s3gen_turbo.cpp | 122 +++ src/models/chatterbox_turbo/session.cpp | 113 +++ .../chatterbox_turbo/t3_turbo_component.cpp | 295 +++++++ .../chatterbox_turbo/t3_turbo_weights.cpp | 182 ++++ .../chatterbox_turbo/text_tokenizer_turbo.cpp | 212 +++++ src/models/chatterbox_turbo/tts.cpp | 96 +++ 29 files changed, 3015 insertions(+), 22 deletions(-) create mode 100644 include/engine/models/chatterbox_turbo/assets.h create mode 100644 include/engine/models/chatterbox_turbo/loader.h create mode 100644 include/engine/models/chatterbox_turbo/s3gen_turbo.h create mode 100644 include/engine/models/chatterbox_turbo/session.h create mode 100644 include/engine/models/chatterbox_turbo/t3_turbo_component.h create mode 100644 include/engine/models/chatterbox_turbo/t3_turbo_types.h create mode 100644 include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h create mode 100644 include/engine/models/chatterbox_turbo/tts.h create mode 100644 model_specs/chatterbox_turbo.json create mode 100644 src/models/chatterbox_turbo/assets.cpp create mode 100644 src/models/chatterbox_turbo/components/dot_prefix_source.cpp create mode 100644 src/models/chatterbox_turbo/components/dot_prefix_source.h create mode 100644 src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp create mode 100644 src/models/chatterbox_turbo/components/s3gen_name_bridge.h create mode 100644 src/models/chatterbox_turbo/components/t3_turbo_runtime.h create mode 100644 src/models/chatterbox_turbo/loader.cpp create mode 100644 src/models/chatterbox_turbo/s3gen_turbo.cpp create mode 100644 src/models/chatterbox_turbo/session.cpp create mode 100644 src/models/chatterbox_turbo/t3_turbo_component.cpp create mode 100644 src/models/chatterbox_turbo/t3_turbo_weights.cpp create mode 100644 src/models/chatterbox_turbo/text_tokenizer_turbo.cpp create mode 100644 src/models/chatterbox_turbo/tts.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 6f6bfee09..7c45d23fe 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1442,6 +1442,26 @@ audiocpp_add_model(chatterbox engine::models::chatterbox::make_chatterbox_loader ) +audiocpp_add_model(chatterbox_turbo + SOURCES + src/models/chatterbox_turbo/t3_turbo_weights.cpp + src/models/chatterbox_turbo/t3_turbo_component.cpp + src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp + src/models/chatterbox_turbo/components/dot_prefix_source.cpp + src/models/chatterbox_turbo/s3gen_turbo.cpp + src/models/chatterbox_turbo/text_tokenizer_turbo.cpp + src/models/chatterbox_turbo/assets.cpp + src/models/chatterbox_turbo/loader.cpp + src/models/chatterbox_turbo/session.cpp + src/models/chatterbox_turbo/tts.cpp + INCLUDES + engine/models/chatterbox_turbo/loader.h + LOADERS + engine::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 c24e0951a..8a5e98f29 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,7 @@ Runtime tags summarize the supported loading paths. GGUF package precision varie | Family | Task | Lang | Variants | Runtime | |---|---|---|---|---| | **chatterbox** | TTS, Clone, VC| ar, da, de, el, en, es, fi, fr, hi, it, ko, ms, nl, no, pl, pt, sv, sw, tr | Chatterbox with 0.5B backbone | GGUF 16/Q8 | +| **chatterbox_turbo** | TTS, Clone (testing) | en | Chatterbox Turbo, distilled 350M GPT2 T3 backbone + 2-step meanflow S3Gen decoder; built-in voice only for now, see [docs/tts.md](docs/tts.md) | GGUF 16/Q8 | | **confucius4_tts** | Clone | zh, en, ja, ko, de, fr, es, id, it, th, pt, ru, ms, vi | Confucius4-TTS multilingual voice cloning | GGUF F32, Stream | | **dots_tts** | TTS, Clone, Edit, Ctrl | multilingual | DotTTS SOAR, MeanFlow, and Edit | GGUF 16/Q8, Stream | | **dramabox** | TTS, Clone | en | DramaBox expressive TTS and voice cloning | GGUF Q8 | diff --git a/docs/tts.md b/docs/tts.md index eba1c7794..4441bbb78 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -100,6 +100,40 @@ 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 Resemble AI's distilled 350M-parameter sibling of 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. It is English-only. + +`chatterbox_turbo` is a separate model family from `chatterbox` (not a variant selectable within it) because Turbo's model-loader resolution can only pick one GGUF source per family, and pointing this loader at a directory of loose GGUF files needs different handling than the base model's single-file packages — see the package notes below. + +The GGUF weights (`cstr/chatterbox-turbo-GGUF`) are a **third-party conversion**, published by `cstr` for their own `CrispASR` project (MIT-relicensed) — not published by ResembleAI or audio.cpp. The package ships as two loose GGUF files (T3 and S3Gen) in one directory; because this framework's model-directory resolution expects exactly one GGUF file per directory, point `--model` at the **T3 GGUF file directly**, not the containing folder (the loader finds the sibling S3Gen file itself by filename convention). + +**Current limitations:** only the built-in default voice baked into the T3 GGUF is supported — custom voice cloning (a caller-supplied reference clip) is not implemented yet, since it depends on the turbo GGUF'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 | Point `--model` at `Chatterbox-Turbo-GGUF/chatterbox-turbo-t3-{q8_0,f16}.gguf` directly | +| 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 T3 GGUF | + +```bash +audiocpp_cli --task clon --family chatterbox_turbo --model models/Chatterbox-Turbo-GGUF/chatterbox-turbo-t3-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/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/include/engine/models/chatterbox_turbo/assets.h b/include/engine/models/chatterbox_turbo/assets.h new file mode 100644 index 000000000..2381fc9ff --- /dev/null +++ b/include/engine/models/chatterbox_turbo/assets.h @@ -0,0 +1,26 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" + +#include +#include + +namespace engine::models::chatterbox_turbo { + +struct ChatterboxTurboAssets { + engine::assets::ResourceBundle resources; + std::shared_ptr t3_turbo_weights; + std::shared_ptr builtin_conditionals_turbo; + // Raw path to the T3 GGUF (not just a TensorSource) so the tokenizer can read the embedded + // `tokenizer.ggml.tokens`/`tokenizer.ggml.merges` metadata arrays directly, and so the + // sibling S3Gen GGUF's path can be derived (see s3gen_turbo.cpp: replaces "-t3-" with + // "-s3gen-" in the filename — the two GGUF files can't both live loosely in one directory + // under this framework's package loader, so only the T3 file is referenced by model_specs; + // the S3Gen file is located and opened directly here). + std::filesystem::path t3_gguf_path; +}; + +std::shared_ptr load_chatterbox_turbo_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/loader.h b/include/engine/models/chatterbox_turbo/loader.h new file mode 100644 index 000000000..cf7112155 --- /dev/null +++ b/include/engine/models/chatterbox_turbo/loader.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/models/chatterbox_turbo/assets.h" + +#include +#include + +namespace engine::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::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/s3gen_turbo.h b/include/engine/models/chatterbox_turbo/s3gen_turbo.h new file mode 100644 index 000000000..d3d85bba2 --- /dev/null +++ b/include/engine/models/chatterbox_turbo/s3gen_turbo.h @@ -0,0 +1,55 @@ +#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 +#include + +namespace engine::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), fed through a tensor-name bridge +// (components/s3gen_name_bridge.h) that translates the turbo GGUF's abbreviated names. +// +// 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 GGUF'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( + const std::filesystem::path & s3gen_gguf_path, + 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; +}; + +// Derives the sibling S3Gen GGUF path from the T3 GGUF path installed by the +// chatterbox_turbo_{q8_0,f16} packages (e.g. ".../chatterbox-turbo-t3-q8_0.gguf" -> +// ".../chatterbox-turbo-s3gen-q8_0.gguf"). Throws if the T3 filename doesn't match the expected +// "-t3-" naming convention. +std::filesystem::path derive_turbo_s3gen_gguf_path(const std::filesystem::path & t3_gguf_path); + +} // namespace engine::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/session.h b/include/engine/models/chatterbox_turbo/session.h new file mode 100644 index 000000000..6026284c6 --- /dev/null +++ b/include/engine/models/chatterbox_turbo/session.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/runtime/session_base.h" +#include "engine/models/chatterbox_turbo/assets.h" +#include "engine/models/chatterbox_turbo/tts.h" + +#include +#include + +namespace engine::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::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/t3_turbo_component.h b/include/engine/models/chatterbox_turbo/t3_turbo_component.h new file mode 100644 index 000000000..a6a1fd657 --- /dev/null +++ b/include/engine/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/models/chatterbox_turbo/t3_turbo_types.h" + +#include + +namespace engine::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::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/t3_turbo_types.h b/include/engine/models/chatterbox_turbo/t3_turbo_types.h new file mode 100644 index 000000000..99008c096 --- /dev/null +++ b/include/engine/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::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::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h b/include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h new file mode 100644 index 000000000..5a503329d --- /dev/null +++ b/include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h @@ -0,0 +1,31 @@ +#pragma once + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include +#include +#include +#include + +namespace engine::models::chatterbox_turbo { + +// Builds a GPT2-style BPE tokenizer from the vocab/merges arrays embedded as GGUF metadata +// (`tokenizer.ggml.tokens`/`tokenizer.ggml.merges`) in Chatterbox Turbo's T3 GGUF — there is no +// separate tokenizer.json file in this package. Materializes a vocab.json + merges.txt pair +// under a scratch directory and feeds those to the existing LlamaBpeTokenizer loader (Gpt2 +// pre-tokenizer mode) rather than teaching that shared loader a new in-memory-array input path. +std::shared_ptr load_chatterbox_turbo_tokenizer( + const std::filesystem::path & t3_gguf_path, + const std::filesystem::path & scratch_dir); + +// 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::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/tts.h b/include/engine/models/chatterbox_turbo/tts.h new file mode 100644 index 000000000..935b710ed --- /dev/null +++ b/include/engine/models/chatterbox_turbo/tts.h @@ -0,0 +1,53 @@ +#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/models/chatterbox_turbo/assets.h" +#include "engine/models/chatterbox_turbo/s3gen_turbo.h" +#include "engine/models/chatterbox_turbo/t3_turbo_component.h" + +#include +#include +#include +#include + +namespace engine::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, + const std::filesystem::path & scratch_dir); + + 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::models::chatterbox_turbo diff --git a/model_specs/chatterbox_turbo.json b/model_specs/chatterbox_turbo.json new file mode 100644 index 000000000..82e72523d --- /dev/null +++ b/model_specs/chatterbox_turbo.json @@ -0,0 +1,86 @@ +{ + "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": "cstr/chatterbox-turbo-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-t3-q8_0.gguf", + "chatterbox-turbo-s3gen-q8_0.gguf" + ], + "description": "Third-party GGUF conversion of ResembleAI/chatterbox-turbo, published by 'cstr' for the CrispASR project (MIT-relicensed); not published by ResembleAI or audio.cpp. Two loose GGUF files (T3 + S3Gen) live side by side in this package's directory, so point --model at the T3 file directly (chatterbox-turbo-t3-q8_0.gguf), not the containing folder, to avoid the ambiguous-directory-GGUF check." + }, + { + "id": "chatterbox_turbo_f16", + "display_name": "Chatterbox Turbo F16 GGUF", + "format": "gguf", + "precision": "f16", + "target_directory": "Chatterbox-Turbo-GGUF", + "files": [ + "chatterbox-turbo-t3-f16.gguf", + "chatterbox-turbo-s3gen-f16.gguf" + ], + "description": "Third-party GGUF conversion of ResembleAI/chatterbox-turbo, published by 'cstr' for the CrispASR project (MIT-relicensed); not published by ResembleAI or audio.cpp. Two loose GGUF files (T3 + S3Gen) live side by side in this package's directory, so point --model at the T3 file directly (chatterbox-turbo-f16.gguf), not the containing folder, to avoid the ambiguous-directory-GGUF check." + } + ], + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "t3_weights": "$gguf" + }, + "files": { + "t3_gguf_path": "t3_weights:" + } + } + ] +} 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/src/models/chatterbox_turbo/assets.cpp b/src/models/chatterbox_turbo/assets.cpp new file mode 100644 index 000000000..503f23595 --- /dev/null +++ b/src/models/chatterbox_turbo/assets.cpp @@ -0,0 +1,29 @@ +#include "engine/models/chatterbox_turbo/assets.h" + +#include "components/dot_prefix_source.h" + +#include "engine/framework/model_spec/package.h" + +#include + +namespace engine::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_gguf_path = out->resources.require_file("t3_gguf_path"); + if (out->t3_gguf_path.extension() != ".gguf") { + throw std::runtime_error("Chatterbox Turbo requires a GGUF model path (point --model at the T3 GGUF file directly)"); + } + // The turbo GGUF's tensor names are flat-dotted ("t3.blk.0...", "conds.gen...") with no "/" + // namespace separator, so the framework's own prefix-stripping (which expects "/") can't be + // used here -- see components/dot_prefix_source.h. + const auto t3_raw = engine::assets::open_tensor_source(out->t3_gguf_path); + out->t3_turbo_weights = make_dot_prefixed_tensor_source(t3_raw, "t3"); + out->builtin_conditionals_turbo = make_dot_prefixed_tensor_source(t3_raw, "conds"); + return out; +} + +} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/dot_prefix_source.cpp b/src/models/chatterbox_turbo/components/dot_prefix_source.cpp new file mode 100644 index 000000000..258f64c02 --- /dev/null +++ b/src/models/chatterbox_turbo/components/dot_prefix_source.cpp @@ -0,0 +1,56 @@ +#include "dot_prefix_source.h" + +#include + +namespace engine::models::chatterbox_turbo { + +namespace { + +class DotPrefixedTensorSource final : public engine::assets::TensorSource { +public: + DotPrefixedTensorSource(std::shared_ptr delegate, std::string prefix) + : delegate_(std::move(delegate)), prefix_(std::move(prefix) + ".") {} + + const std::filesystem::path & source_path() const noexcept override { + return delegate_->source_path(); + } + bool has_tensor(std::string_view name) const noexcept override { + return delegate_->has_tensor(prefix_ + std::string(name)); + } + engine::assets::TensorMetadata require_metadata(std::string_view name) const override { + return delegate_->require_metadata(prefix_ + std::string(name)); + } + std::vector tensors() const override { + return delegate_->tensors(); + } + engine::assets::RawTensorData require_tensor_data(std::string_view name) const override { + return delegate_->require_tensor_data(prefix_ + std::string(name)); + } + std::vector require_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + return delegate_->require_f32(prefix_ + std::string(name), expected_shape); + } + std::optional> optional_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + return delegate_->optional_f32(prefix_ + std::string(name), expected_shape); + } + int64_t require_i64_scalar(std::string_view name) const override { + return delegate_->require_i64_scalar(prefix_ + std::string(name)); + } + +private: + std::shared_ptr delegate_; + std::string prefix_; +}; + +} // namespace + +std::shared_ptr make_dot_prefixed_tensor_source( + std::shared_ptr source, + std::string prefix) { + return std::make_shared(std::move(source), std::move(prefix)); +} + +} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/dot_prefix_source.h b/src/models/chatterbox_turbo/components/dot_prefix_source.h new file mode 100644 index 000000000..b0858ab84 --- /dev/null +++ b/src/models/chatterbox_turbo/components/dot_prefix_source.h @@ -0,0 +1,26 @@ +#pragma once + +// The turbo GGUF (cstr/chatterbox-turbo-GGUF) names every tensor with a flat dot-separated path +// (e.g. "t3.blk.0.attn_norm.weight", "conds.gen.embedding", "s3.flow.encoder_proj") with no "/" +// namespace separator anywhere. The framework's own prefix-stripping helpers +// (engine::assets::open_tensor_source(path, prefix) / make_prefixed_tensor_source) expect a +// "/"-delimited packed-GGUF namespace instead (matching how this project's own GGUF converter +// packs multiple checkpoints into one file, e.g. "t3_english/tfmr.layers.0...") and throw +// "packed GGUF namespace does not exist" against a dot-separated name. This is therefore a +// separate, dot-based sibling used only for chatterbox_turbo's third-party GGUF. + +#include "engine/framework/assets/tensor_source.h" + +#include +#include + +namespace engine::models::chatterbox_turbo { + +// Strips a literal "." from every queried name before delegating to `source` (e.g. +// prefix="t3" turns a query for "blk.0.attn_norm.weight" into "t3.blk.0.attn_norm.weight" on the +// underlying source). +std::shared_ptr make_dot_prefixed_tensor_source( + std::shared_ptr source, + std::string prefix); + +} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp b/src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp new file mode 100644 index 000000000..c9b513134 --- /dev/null +++ b/src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp @@ -0,0 +1,207 @@ +#include "s3gen_name_bridge.h" + +#include +#include +#include +#include + +namespace engine::models::chatterbox_turbo { +namespace { + +using Rule = std::pair; + +// Longest/most-specific prefixes first so a shorter rule never shadows a longer one +// (e.g. "flow.encoder_proj" must not be caught by a "flow.encoder." rule). +constexpr std::array kFlowPrefixRules = {{ + {"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. +}}; + +// Applied to whatever remains after a prefix rule above consumed the "flow.decoder.estimator.*" +// or "flow.encoder.{encoders,up_encoders}.N." portion of the name. +constexpr std::array kSegmentRules = {{ + {".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"}, +}}; + +constexpr std::array kTransformerBlockRules = {{ + {".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"}, +}}; + +constexpr std::array kVocoderPrefixRules = {{ + {"conv_pre", "cpre"}, + {"conv_post", "cpost"}, + {"resblocks.", "rb."}, + {"source_downs.", "sd."}, + {"source_resblocks.", "srb."}, + {"f0_predictor.condnet.", "f0.cn."}, +}}; + +constexpr std::array kVocoderExactRules = {{ + {"f0_predictor.classifier", "f0.cls"}, + {"m_source.l_linear", "ms.ll"}, + {".convs1.", ".c1."}, + {".convs2.", ".c2."}, + {".activations1.", ".a1."}, + {".activations2.", ".a2."}, +}}; + +bool starts_with(std::string_view text, std::string_view prefix) { + return text.size() >= prefix.size() && text.compare(0, prefix.size(), prefix) == 0; +} + +std::string replace_all(std::string text, std::string_view search, std::string_view replace) { + if (search.empty()) { + return text; + } + size_t pos = 0; + while ((pos = text.find(search, pos)) != std::string::npos) { + text.replace(pos, search.size(), replace); + pos += replace.size(); + } + return text; +} + +std::string translate_flow_name(std::string_view name) { + std::string translated(name); + for (const auto & [from, to] : kFlowPrefixRules) { + if (starts_with(translated, from)) { + translated = std::string(to) + translated.substr(from.size()); + break; + } + } + for (const auto & [from, to] : kSegmentRules) { + translated = replace_all(std::move(translated), from, to); + } + // ".ff.net.2" would also match the shorter ".ff.net.0.proj" search text's neighbourhood, + // so handle it (and the rest of the transformer-block-only renames) after the segment pass. + translated = replace_all(std::move(translated), ".ff.net.2", ".ff.down"); + for (const auto & [from, to] : kTransformerBlockRules) { + translated = replace_all(std::move(translated), from, to); + } + return translated; +} + +std::string translate_vocoder_name(std::string_view name) { + std::string translated(name); + for (const auto & [from, to] : kVocoderPrefixRules) { + if (starts_with(translated, from)) { + translated = std::string(to) + translated.substr(from.size()); + break; + } + } + for (const auto & [from, to] : kVocoderExactRules) { + if (translated == from || starts_with(translated, std::string(from) + ".")) { + translated = std::string(to) + translated.substr(from.size()); + break; + } + translated = replace_all(translated, from, to); + } + // HiftVocoderConfig.tensor_prefix is left empty (see s3gen_turbo.cpp) so this translator + // sees HiftVocoderComponent's plain requested names; the turbo GGUF's vocoder tensors all + // live under the top-level "v." namespace, which is added here rather than via tensor_prefix. + return "v." + translated; +} + +class TranslatingTensorSource final : public engine::assets::TensorSource { +public: + TranslatingTensorSource( + std::shared_ptr delegate, + std::string (*translate)(std::string_view)) + : delegate_(std::move(delegate)), translate_(translate) {} + + const std::filesystem::path & source_path() const noexcept override { + return delegate_->source_path(); + } + bool has_tensor(std::string_view name) const noexcept override { + return delegate_->has_tensor(translate_(name)); + } + engine::assets::TensorMetadata require_metadata(std::string_view name) const override { + return delegate_->require_metadata(translate_(name)); + } + std::vector tensors() const override { + return delegate_->tensors(); + } + engine::assets::RawTensorData require_tensor_data(std::string_view name) const override { + return delegate_->require_tensor_data(translate_(name)); + } + // Some turbo GGUF tensors squeeze a leading size-1 dimension the base loader code still + // expects explicitly (e.g. the vocoder's f0 classifier: nn.Linear(512, 1).weight is stored + // as ne=[512] rather than [1, 512]). Rather than special-case those, skip the delegate's + // strict per-dimension shape check here and only verify the element count still matches. + std::vector require_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + auto values = delegate_->require_f32(translate_(name), std::nullopt); + if (expected_shape.has_value()) { + int64_t expected_count = 1; + for (int64_t dim : *expected_shape) { + expected_count *= dim; + } + if (static_cast(values.size()) != expected_count) { + throw std::runtime_error("tensor element count mismatch for " + std::string(name)); + } + } + return values; + } + std::optional> optional_f32( + std::string_view name, + const std::optional> & expected_shape) const override { + if (!has_tensor(name)) { + return std::nullopt; + } + return require_f32(name, expected_shape); + } + int64_t require_i64_scalar(std::string_view name) const override { + return delegate_->require_i64_scalar(translate_(name)); + } + +private: + std::shared_ptr delegate_; + std::string (*translate_)(std::string_view); +}; + +} // namespace + +std::shared_ptr make_turbo_s3gen_flow_bridge( + std::shared_ptr turbo_s3_source) { + return std::make_shared(std::move(turbo_s3_source), &translate_flow_name); +} + +std::shared_ptr make_turbo_vocoder_bridge( + std::shared_ptr turbo_s3_source) { + return std::make_shared(std::move(turbo_s3_source), &translate_vocoder_name); +} + +} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/s3gen_name_bridge.h b/src/models/chatterbox_turbo/components/s3gen_name_bridge.h new file mode 100644 index 000000000..9ae2f1c1f --- /dev/null +++ b/src/models/chatterbox_turbo/components/s3gen_name_bridge.h @@ -0,0 +1,61 @@ +#pragma once + +// Chatterbox Turbo's S3Gen weights ship in a third-party GGUF (cstr/chatterbox-turbo-GGUF) +// whose tensor names are a consistent, decodable abbreviation of the names the existing +// chatterbox S3Gen loader code (src/models/chatterbox/s3gen_flow.cpp, +// src/framework/modules/vocoders/hift_vocoder.cpp) already expects and knows how to build +// ggml graphs for. Rather than reimplement the flow encoder/decoder/vocoder, this header +// name-translates so the *existing* loaders can run unmodified against the turbo GGUF. +// +// Mapping verified against the actual GGUF header of chatterbox-turbo-s3gen-q8_0.gguf +// (parsed offline with a throwaway GGUF reader), matched segment-by-segment against the +// literal prefixes s3gen_flow.cpp/hift_vocoder.cpp construct: +// +// flow.encoder.encoders.N -> fe.enc.N +// flow.encoder.up_encoders.N -> fe.ue.N +// 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.encoder_proj / spk_embed_affine_layer / input_embedding (unchanged, top-level "flow.") +// (within a conformer layer) norm_mha -> nmha, self_attn.linear_{q,k,v,out,pos} -> sa.l{q,k,v,o,p}, +// self_attn.pos_bias_{u,v} -> sa.pb{u,v}, norm_ff -> nff, feed_forward.w_1/2 -> ff.w_1/2 +// 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 (unchanged linear_1/linear_2 suffix) +// flow.decoder.estimator.time_embed_mixer -> fd.tmx (turbo-only meanflow tensor, no base equivalent) +// (within a resnet block) blockN.block.0/2 -> bN.0/2, res_conv -> rc, mlp.1 unchanged +// (within a transformer block) norm1/norm3 unchanged, attn1.to_q/k/v -> attn1.q/k/v, +// attn1.to_out.0 -> attn1.o, ff.net.0.proj -> ff.up, ff.net.2 -> ff.down +// (vocoder, HiftVocoderConfig.tensor_prefix="v.") conv_pre -> cpre, conv_post -> cpost, +// ups.N unchanged, resblocks.N -> rb.N, source_downs.N -> sd.N, source_resblocks.N -> srb.N, +// f0_predictor.condnet.N -> f0.cn.N, f0_predictor.classifier -> f0.cls, m_source.l_linear -> ms.ll +// (within a resblock) convs1/2 -> c1/2, activations1/2 -> a1/2 +// +// This is deliberately additive/self-contained to chatterbox_turbo: it does not modify the +// base chatterbox family's own tensor naming or loading code. + +#include "engine/framework/assets/tensor_source.h" + +#include +#include + +namespace engine::models::chatterbox_turbo { + +// Wraps an already s3-prefix-stripped TensorSource (see open_tensor_source(path, "s3")) and +// rewrites incoming chatterbox-style names to the turbo GGUF's abbreviated names before +// delegating every lookup. +std::shared_ptr make_turbo_s3gen_flow_bridge( + std::shared_ptr turbo_s3_source); + +// Same idea for the HiFT vocoder; construct the shared HiftVocoderComponent with +// tensor_prefix = "v." against this bridge and weight_layout = Plain (the turbo GGUF already +// folds torch weight-norm parametrization into a plain "weight" tensor at conversion time). +std::shared_ptr make_turbo_vocoder_bridge( + std::shared_ptr turbo_s3_source); + +} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/t3_turbo_runtime.h b/src/models/chatterbox_turbo/components/t3_turbo_runtime.h new file mode 100644 index 000000000..006563f40 --- /dev/null +++ b/src/models/chatterbox_turbo/components/t3_turbo_runtime.h @@ -0,0 +1,811 @@ +#pragma once + +#include "engine/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::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::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/loader.cpp b/src/models/chatterbox_turbo/loader.cpp new file mode 100644 index 000000000..1ee809624 --- /dev/null +++ b/src/models/chatterbox_turbo/loader.cpp @@ -0,0 +1,128 @@ +#include "engine/models/chatterbox_turbo/loader.h" + +#include "engine/models/chatterbox_turbo/session.h" + +#include "engine/framework/model_spec/package.h" + +#include +#include + +namespace engine::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::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/s3gen_turbo.cpp b/src/models/chatterbox_turbo/s3gen_turbo.cpp new file mode 100644 index 000000000..23c8b35ba --- /dev/null +++ b/src/models/chatterbox_turbo/s3gen_turbo.cpp @@ -0,0 +1,122 @@ +#include "engine/models/chatterbox_turbo/s3gen_turbo.h" + +#include "components/dot_prefix_source.h" +#include "components/s3gen_name_bridge.h" + +#include + +namespace engine::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 (see + // s3gen_name_bridge.h). The turbo 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 "v." vocoder namespace is applied inside translate_vocoder_name (see + // s3gen_name_bridge.cpp) rather than here, since that translation runs on top of the raw + // dot-prefixed source and needs to see the plain (unprefixed) name HiftVocoderComponent asks + // for in order to match its own prefix rules. + config.tensor_prefix = ""; + config.weight_layout = engine::modules::HiftVocoderWeightLayout::Canonical; + return config; +} + +} // namespace + +std::filesystem::path derive_turbo_s3gen_gguf_path(const std::filesystem::path & t3_gguf_path) { + const std::string filename = t3_gguf_path.filename().string(); + const std::string marker = "-t3-"; + const auto pos = filename.find(marker); + if (pos == std::string::npos) { + throw std::runtime_error( + "Chatterbox Turbo T3 GGUF filename does not match the expected '...-t3-.gguf' naming: " + filename); + } + const std::string s3gen_filename = filename.substr(0, pos) + "-s3gen-" + filename.substr(pos + marker.size()); + return t3_gguf_path.parent_path() / s3gen_filename; +} + +std::shared_ptr ChatterboxTurboS3Gen::load( + const std::filesystem::path & s3gen_gguf_path, + const engine::core::ExecutionContext & execution_context, + engine::assets::TensorStorageType weight_storage_type) { + auto raw_source = make_dot_prefixed_tensor_source(engine::assets::open_tensor_source(s3gen_gguf_path), "s3"); + auto flow_source = make_turbo_s3gen_flow_bridge(raw_source); + auto vocoder_source = make_turbo_vocoder_bridge(raw_source); + + auto out = std::make_shared(); + out->execution_context_ = &execution_context; + out->encoder_weights_ = + engine::models::chatterbox::load_s3_flow_encoder_weights(*flow_source, execution_context, weight_storage_type); + out->decoder_weights_ = + engine::models::chatterbox::load_s3_flow_decoder_weights(*flow_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 GGUF is 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( + vocoder_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::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/session.cpp b/src/models/chatterbox_turbo/session.cpp new file mode 100644 index 000000000..5670fd176 --- /dev/null +++ b/src/models/chatterbox_turbo/session.cpp @@ -0,0 +1,113 @@ +#include "engine/models/chatterbox_turbo/session.h" + +#include "engine/framework/runtime/options.h" +#include "engine/framework/text/chunking.h" + +#include + +namespace engine::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(), assets_->resources.model_root() / ".chatterbox_turbo_tokenizer_cache"); + } + 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::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/t3_turbo_component.cpp b/src/models/chatterbox_turbo/t3_turbo_component.cpp new file mode 100644 index 000000000..3a50c5124 --- /dev/null +++ b/src/models/chatterbox_turbo/t3_turbo_component.cpp @@ -0,0 +1,295 @@ +#include "components/t3_turbo_runtime.h" + +#include +#include + +namespace engine::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::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/t3_turbo_weights.cpp b/src/models/chatterbox_turbo/t3_turbo_weights.cpp new file mode 100644 index 000000000..05a901259 --- /dev/null +++ b/src/models/chatterbox_turbo/t3_turbo_weights.cpp @@ -0,0 +1,182 @@ +#include "engine/models/chatterbox_turbo/t3_turbo_component.h" + +#include + +namespace engine::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::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/text_tokenizer_turbo.cpp b/src/models/chatterbox_turbo/text_tokenizer_turbo.cpp new file mode 100644 index 000000000..0aec66d90 --- /dev/null +++ b/src/models/chatterbox_turbo/text_tokenizer_turbo.cpp @@ -0,0 +1,212 @@ +#include "engine/models/chatterbox_turbo/text_tokenizer_turbo.h" + +#include +#include + +#include +#include +#include +#include + +namespace engine::models::chatterbox_turbo { + +namespace { + +struct GgufContextDeleter { + void operator()(gguf_context * ctx) const noexcept { + if (ctx != nullptr) { + gguf_free(ctx); + } + } +}; + +std::vector read_gguf_string_array(gguf_context * ctx, const char * key) { + const int64_t key_id = gguf_find_key(ctx, key); + if (key_id < 0) { + throw std::runtime_error(std::string("Chatterbox Turbo GGUF is missing metadata key: ") + key); + } + if (gguf_get_arr_type(ctx, key_id) != GGUF_TYPE_STRING) { + throw std::runtime_error(std::string("Chatterbox Turbo GGUF metadata key is not a string array: ") + key); + } + const size_t count = gguf_get_arr_n(ctx, key_id); + std::vector out; + out.reserve(count); + for (size_t i = 0; i < count; ++i) { + out.emplace_back(gguf_get_arr_str(ctx, key_id, i)); + } + return out; +} + +std::string json_escape(const std::string & value) { + std::string out; + out.reserve(value.size() + 8); + for (unsigned char c : value) { + switch (c) { + case '"': out += "\\\""; break; + case '\\': out += "\\\\"; break; + case '\n': out += "\\n"; break; + case '\r': out += "\\r"; break; + case '\t': out += "\\t"; break; + default: + if (c < 0x20) { + char buf[8]; + std::snprintf(buf, sizeof(buf), "\\u%04x", c); + out += buf; + } else { + out += static_cast(c); + } + } + } + return out; +} + +bool looks_like_bracket_tag(const std::string & token) { + return token.size() >= 3 && token.front() == '[' && token.back() == ']'; +} + +} // namespace + +std::shared_ptr load_chatterbox_turbo_tokenizer( + const std::filesystem::path & t3_gguf_path, + const std::filesystem::path & scratch_dir) { + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = nullptr; + std::unique_ptr ctx(gguf_init_from_file(t3_gguf_path.string().c_str(), params)); + if (!ctx) { + throw std::runtime_error("failed to open Chatterbox Turbo T3 GGUF for tokenizer metadata: " + t3_gguf_path.string()); + } + + const auto tokens = read_gguf_string_array(ctx.get(), "tokenizer.ggml.tokens"); + const auto merges = read_gguf_string_array(ctx.get(), "tokenizer.ggml.merges"); + if (tokens.empty()) { + throw std::runtime_error("Chatterbox Turbo T3 GGUF has an empty tokenizer.ggml.tokens array"); + } + + std::filesystem::create_directories(scratch_dir); + const auto vocab_path = scratch_dir / "chatterbox_turbo_vocab.json"; + const auto merges_path = scratch_dir / "chatterbox_turbo_merges.txt"; + + // The 19 emotion/style control tags ([laugh], [sigh], ...) sit at the tail of the vocab and + // cannot be reached through byte-pair merges (they were never trained into the merge table), + // so they must be registered as atomic special tokens the pre-tokenizer matches directly + // (LlamaBpeTokenizerSpec::additional_special_tokens) rather than as plain vocab entries. + // Crucially they must NOT also appear in vocab.json below: add_runtime_special_tokens only + // sets a token's "special" attribute when it is inserting a genuinely new id -- if the id + // already exists (which it would, from vocab.json) it just verifies the id and continues, + // silently leaving the token without its special/atomic-match flag, so it gets shredded into + // ordinary byte-level BPE pieces instead of matched as one unit. + std::vector special_tokens; + int64_t special_tokens_start = static_cast(tokens.size()); + for (int64_t id = static_cast(tokens.size()) - 1; id >= 0; --id) { + if (!looks_like_bracket_tag(tokens[static_cast(id)])) { + break; + } + special_tokens.emplace_back(tokens[static_cast(id)], static_cast(id)); + special_tokens_start = id; + } + + { + std::ofstream out(vocab_path, std::ios::trunc); + if (!out) { + throw std::runtime_error("failed to write Chatterbox Turbo tokenizer vocab: " + vocab_path.string()); + } + out << "{"; + bool first = true; + for (size_t i = 0; i < tokens.size(); ++i) { + if (static_cast(i) >= special_tokens_start) { + continue; + } + if (!first) { + out << ","; + } + first = false; + out << "\"" << json_escape(tokens[i]) << "\":" << i; + } + out << "}"; + } + { + std::ofstream out(merges_path, std::ios::trunc); + if (!out) { + throw std::runtime_error("failed to write Chatterbox Turbo tokenizer merges: " + merges_path.string()); + } + for (const auto & merge : merges) { + out << merge << "\n"; + } + } + + 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::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/tts.cpp b/src/models/chatterbox_turbo/tts.cpp new file mode 100644 index 000000000..c0778bf69 --- /dev/null +++ b/src/models/chatterbox_turbo/tts.cpp @@ -0,0 +1,96 @@ +#include "engine/models/chatterbox_turbo/tts.h" + +#include "engine/models/chatterbox_turbo/text_tokenizer_turbo.h" + +#include +#include + +namespace engine::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, + const std::filesystem::path & scratch_dir) + : assets_(std::move(assets)) { + tokenizer_ = load_chatterbox_turbo_tokenizer(assets_->t3_gguf_path, scratch_dir); + + auto t3_weights = load_t3_turbo_inference_weights(*assets_->t3_turbo_weights, execution_context); + t3_ = std::make_unique(t3_weights, execution_context); + + const auto s3gen_path = derive_turbo_s3gen_gguf_path(assets_->t3_gguf_path); + s3gen_ = ChatterboxTurboS3Gen::load(s3gen_path, 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::models::chatterbox_turbo From 65497d5482a098fd836dfa87b8030723afbf5f24 Mon Sep 17 00:00:00 2001 From: pannagaps Date: Thu, 3 Sep 2026 11:35:11 +0530 Subject: [PATCH 2/2] Refactor Chatterbox Turbo model components and repack GGUF - Removed obsolete assets and dot prefix source files from chatterbox_turbo model. - Deleted S3Gen name bridge and tokenizer implementation, simplifying the model structure. - Introduced a new Python script to repack the third-party chatterbox-turbo GGUF files into a format compatible with audio.cpp. - The repacking script includes tensor name translation to match the base Chatterbox's expectations, ensuring seamless integration. - Updated tokenizer handling to create separate vocab, merges, and special tokens files instead of relying on runtime GGUF metadata. --- CMakeLists.txt | 26 +- README.md | 2 +- docs/community_models/chatterbox_turbo.md | 84 +++++ docs/community_models/models.md | 1 + docs/tts.md | 33 +- .../chatterbox_turbo/assets.h | 20 ++ .../chatterbox_turbo/loader.h | 6 +- .../chatterbox_turbo/s3gen_turbo.h | 24 +- .../chatterbox_turbo/session.h | 8 +- .../chatterbox_turbo/t3_turbo_component.h | 6 +- .../chatterbox_turbo/t3_turbo_types.h | 4 +- .../chatterbox_turbo/text_tokenizer_turbo.h | 34 ++ .../chatterbox_turbo/tts.h | 13 +- .../engine/models/chatterbox_turbo/assets.h | 26 -- .../chatterbox_turbo/text_tokenizer_turbo.h | 31 -- model_specs/chatterbox_turbo.json | 32 +- .../chatterbox_turbo/assets.cpp | 18 + .../components/t3_turbo_runtime.h | 6 +- .../chatterbox_turbo/loader.cpp | 8 +- .../chatterbox_turbo/s3gen_turbo.cpp | 55 +-- .../chatterbox_turbo/session.cpp | 9 +- .../chatterbox_turbo/t3_turbo_component.cpp | 4 +- .../chatterbox_turbo/t3_turbo_weights.cpp | 6 +- .../chatterbox_turbo/text_tokenizer_turbo.cpp | 96 +++++ .../chatterbox_turbo/tts.cpp | 19 +- src/models/chatterbox_turbo/assets.cpp | 29 -- .../components/dot_prefix_source.cpp | 56 --- .../components/dot_prefix_source.h | 26 -- .../components/s3gen_name_bridge.cpp | 207 ----------- .../components/s3gen_name_bridge.h | 61 ---- .../chatterbox_turbo/text_tokenizer_turbo.cpp | 212 ----------- .../repack_chatterbox_turbo_gguf.py | 331 ++++++++++++++++++ 32 files changed, 714 insertions(+), 779 deletions(-) create mode 100644 docs/community_models/chatterbox_turbo.md create mode 100644 include/engine/community_models/chatterbox_turbo/assets.h rename include/engine/{models => community_models}/chatterbox_turbo/loader.h (85%) rename include/engine/{models => community_models}/chatterbox_turbo/s3gen_turbo.h (66%) rename include/engine/{models => community_models}/chatterbox_turbo/session.h (79%) rename include/engine/{models => community_models}/chatterbox_turbo/t3_turbo_component.h (86%) rename include/engine/{models => community_models}/chatterbox_turbo/t3_turbo_types.h (96%) create mode 100644 include/engine/community_models/chatterbox_turbo/text_tokenizer_turbo.h rename include/engine/{models => community_models}/chatterbox_turbo/tts.h (85%) delete mode 100644 include/engine/models/chatterbox_turbo/assets.h delete mode 100644 include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h create mode 100644 src/community_models/chatterbox_turbo/assets.cpp rename src/{models => community_models}/chatterbox_turbo/components/t3_turbo_runtime.h (99%) rename src/{models => community_models}/chatterbox_turbo/loader.cpp (95%) rename src/{models => community_models}/chatterbox_turbo/s3gen_turbo.cpp (60%) rename src/{models => community_models}/chatterbox_turbo/session.cpp (94%) rename src/{models => community_models}/chatterbox_turbo/t3_turbo_component.cpp (99%) rename src/{models => community_models}/chatterbox_turbo/t3_turbo_weights.cpp (97%) create mode 100644 src/community_models/chatterbox_turbo/text_tokenizer_turbo.cpp rename src/{models => community_models}/chatterbox_turbo/tts.cpp (85%) delete mode 100644 src/models/chatterbox_turbo/assets.cpp delete mode 100644 src/models/chatterbox_turbo/components/dot_prefix_source.cpp delete mode 100644 src/models/chatterbox_turbo/components/dot_prefix_source.h delete mode 100644 src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp delete mode 100644 src/models/chatterbox_turbo/components/s3gen_name_bridge.h delete mode 100644 src/models/chatterbox_turbo/text_tokenizer_turbo.cpp create mode 100644 tools/community_models/chatterbox_turbo/repack_chatterbox_turbo_gguf.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 0d632bd34..d15c475d9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1444,20 +1444,18 @@ audiocpp_add_model(chatterbox audiocpp_add_model(chatterbox_turbo SOURCES - src/models/chatterbox_turbo/t3_turbo_weights.cpp - src/models/chatterbox_turbo/t3_turbo_component.cpp - src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp - src/models/chatterbox_turbo/components/dot_prefix_source.cpp - src/models/chatterbox_turbo/s3gen_turbo.cpp - src/models/chatterbox_turbo/text_tokenizer_turbo.cpp - src/models/chatterbox_turbo/assets.cpp - src/models/chatterbox_turbo/loader.cpp - src/models/chatterbox_turbo/session.cpp - src/models/chatterbox_turbo/tts.cpp - INCLUDES - engine/models/chatterbox_turbo/loader.h - LOADERS - engine::models::chatterbox_turbo::make_chatterbox_turbo_loader + 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 ) diff --git a/README.md b/README.md index 8a5e98f29..dad65ee0a 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,6 @@ Runtime tags summarize the supported loading paths. GGUF package precision varie | Family | Task | Lang | Variants | Runtime | |---|---|---|---|---| | **chatterbox** | TTS, Clone, VC| ar, da, de, el, en, es, fi, fr, hi, it, ko, ms, nl, no, pl, pt, sv, sw, tr | Chatterbox with 0.5B backbone | GGUF 16/Q8 | -| **chatterbox_turbo** | TTS, Clone (testing) | en | Chatterbox Turbo, distilled 350M GPT2 T3 backbone + 2-step meanflow S3Gen decoder; built-in voice only for now, see [docs/tts.md](docs/tts.md) | GGUF 16/Q8 | | **confucius4_tts** | Clone | zh, en, ja, ko, de, fr, es, id, it, th, pt, ru, ms, vi | Confucius4-TTS multilingual voice cloning | GGUF F32, Stream | | **dots_tts** | TTS, Clone, Edit, Ctrl | multilingual | DotTTS SOAR, MeanFlow, and Edit | GGUF 16/Q8, Stream | | **dramabox** | TTS, Clone | en | DramaBox expressive TTS and voice cloning | GGUF Q8 | @@ -146,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 4441bbb78..28e0818ce 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -102,26 +102,39 @@ audiocpp_cli --task vc --family chatterbox --model models/chatterbox --backend c ## Chatterbox Turbo -Chatterbox Turbo is Resemble AI's distilled 350M-parameter sibling of 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. It is English-only. - -`chatterbox_turbo` is a separate model family from `chatterbox` (not a variant selectable within it) because Turbo's model-loader resolution can only pick one GGUF source per family, and pointing this loader at a directory of loose GGUF files needs different handling than the base model's single-file packages — see the package notes below. - -The GGUF weights (`cstr/chatterbox-turbo-GGUF`) are a **third-party conversion**, published by `cstr` for their own `CrispASR` project (MIT-relicensed) — not published by ResembleAI or audio.cpp. The package ships as two loose GGUF files (T3 and S3Gen) in one directory; because this framework's model-directory resolution expects exactly one GGUF file per directory, point `--model` at the **T3 GGUF file directly**, not the containing folder (the loader finds the sibling S3Gen file itself by filename convention). - -**Current limitations:** only the built-in default voice baked into the T3 GGUF is supported — custom voice cloning (a caller-supplied reference clip) is not implemented yet, since it depends on the turbo GGUF'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. +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 | Point `--model` at `Chatterbox-Turbo-GGUF/chatterbox-turbo-t3-{q8_0,f16}.gguf` directly | +| 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 T3 GGUF | +| Built-in voices | One, embedded in the package | ```bash -audiocpp_cli --task clon --family chatterbox_turbo --model models/Chatterbox-Turbo-GGUF/chatterbox-turbo-t3-q8_0.gguf --backend cuda --text "Hello from Chatterbox Turbo." --out out.wav +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 | 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/models/chatterbox_turbo/loader.h b/include/engine/community_models/chatterbox_turbo/loader.h similarity index 85% rename from include/engine/models/chatterbox_turbo/loader.h rename to include/engine/community_models/chatterbox_turbo/loader.h index cf7112155..97d843a0f 100644 --- a/include/engine/models/chatterbox_turbo/loader.h +++ b/include/engine/community_models/chatterbox_turbo/loader.h @@ -1,12 +1,12 @@ #pragma once #include "engine/framework/runtime/model.h" -#include "engine/models/chatterbox_turbo/assets.h" +#include "engine/community_models/chatterbox_turbo/assets.h" #include #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { class ChatterboxTurboLoadedModel final : public runtime::ILoadedVoiceModel { public: @@ -30,4 +30,4 @@ class ChatterboxTurboLoadedModel final : public runtime::ILoadedVoiceModel { std::unique_ptr load_chatterbox_turbo_model(const std::filesystem::path & model_root); std::shared_ptr make_chatterbox_turbo_loader(); -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/s3gen_turbo.h b/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h similarity index 66% rename from include/engine/models/chatterbox_turbo/s3gen_turbo.h rename to include/engine/community_models/chatterbox_turbo/s3gen_turbo.h index d3d85bba2..a5e0213da 100644 --- a/include/engine/models/chatterbox_turbo/s3gen_turbo.h +++ b/include/engine/community_models/chatterbox_turbo/s3gen_turbo.h @@ -8,26 +8,26 @@ #include "engine/models/chatterbox/s3gen_inference.h" #include -#include #include #include -namespace engine::models::chatterbox_turbo { +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), fed through a tensor-name bridge -// (components/s3gen_name_bridge.h) that translates the turbo GGUF's abbreviated names. +// 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 GGUF'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. +// 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( - const std::filesystem::path & s3gen_gguf_path, + std::shared_ptr s3gen_source, const engine::core::ExecutionContext & execution_context, engine::assets::TensorStorageType weight_storage_type = engine::assets::TensorStorageType::Native); @@ -46,10 +46,4 @@ class ChatterboxTurboS3Gen { const engine::core::ExecutionContext * execution_context_ = nullptr; }; -// Derives the sibling S3Gen GGUF path from the T3 GGUF path installed by the -// chatterbox_turbo_{q8_0,f16} packages (e.g. ".../chatterbox-turbo-t3-q8_0.gguf" -> -// ".../chatterbox-turbo-s3gen-q8_0.gguf"). Throws if the T3 filename doesn't match the expected -// "-t3-" naming convention. -std::filesystem::path derive_turbo_s3gen_gguf_path(const std::filesystem::path & t3_gguf_path); - -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/session.h b/include/engine/community_models/chatterbox_turbo/session.h similarity index 79% rename from include/engine/models/chatterbox_turbo/session.h rename to include/engine/community_models/chatterbox_turbo/session.h index 6026284c6..60094400f 100644 --- a/include/engine/models/chatterbox_turbo/session.h +++ b/include/engine/community_models/chatterbox_turbo/session.h @@ -1,13 +1,13 @@ #pragma once #include "engine/framework/runtime/session_base.h" -#include "engine/models/chatterbox_turbo/assets.h" -#include "engine/models/chatterbox_turbo/tts.h" +#include "engine/community_models/chatterbox_turbo/assets.h" +#include "engine/community_models/chatterbox_turbo/tts.h" #include #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { class ChatterboxTurboSession final : public runtime::RuntimeSessionBase @@ -31,4 +31,4 @@ class ChatterboxTurboSession final std::unique_ptr component_; }; -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/t3_turbo_component.h b/include/engine/community_models/chatterbox_turbo/t3_turbo_component.h similarity index 86% rename from include/engine/models/chatterbox_turbo/t3_turbo_component.h rename to include/engine/community_models/chatterbox_turbo/t3_turbo_component.h index a6a1fd657..01a4fd97d 100644 --- a/include/engine/models/chatterbox_turbo/t3_turbo_component.h +++ b/include/engine/community_models/chatterbox_turbo/t3_turbo_component.h @@ -3,11 +3,11 @@ #include "engine/framework/assets/tensor_source.h" #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" -#include "engine/models/chatterbox_turbo/t3_turbo_types.h" +#include "engine/community_models/chatterbox_turbo/t3_turbo_types.h" #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { std::shared_ptr load_t3_turbo_inference_weights( const engine::assets::TensorSource & source, @@ -33,4 +33,4 @@ class T3TurboInferenceComponent { std::shared_ptr state_; }; -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/t3_turbo_types.h b/include/engine/community_models/chatterbox_turbo/t3_turbo_types.h similarity index 96% rename from include/engine/models/chatterbox_turbo/t3_turbo_types.h rename to include/engine/community_models/chatterbox_turbo/t3_turbo_types.h index 99008c096..20401f0b1 100644 --- a/include/engine/models/chatterbox_turbo/t3_turbo_types.h +++ b/include/engine/community_models/chatterbox_turbo/t3_turbo_types.h @@ -9,7 +9,7 @@ #include #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { struct T3TurboGraphWeight { std::vector values; @@ -94,4 +94,4 @@ struct T3TurboGenerateOutputs { double decode_runner_ms = 0.0; }; -} // namespace engine::models::chatterbox_turbo +} // 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/models/chatterbox_turbo/tts.h b/include/engine/community_models/chatterbox_turbo/tts.h similarity index 85% rename from include/engine/models/chatterbox_turbo/tts.h rename to include/engine/community_models/chatterbox_turbo/tts.h index 935b710ed..2d74bb1f9 100644 --- a/include/engine/models/chatterbox_turbo/tts.h +++ b/include/engine/community_models/chatterbox_turbo/tts.h @@ -3,16 +3,16 @@ #include "engine/framework/core/execution_context.h" #include "engine/framework/tokenizers/llama_bpe.h" #include "engine/models/chatterbox/s3gen_inference.h" -#include "engine/models/chatterbox_turbo/assets.h" -#include "engine/models/chatterbox_turbo/s3gen_turbo.h" -#include "engine/models/chatterbox_turbo/t3_turbo_component.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::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { struct ChatterboxTurboGenerateConfig { float temperature = 0.8f; @@ -33,8 +33,7 @@ class ChatterboxTurboTtsComponent { public: ChatterboxTurboTtsComponent( std::shared_ptr assets, - const engine::core::ExecutionContext & execution_context, - const std::filesystem::path & scratch_dir); + const engine::core::ExecutionContext & execution_context); engine::models::chatterbox::S3GenInferenceOutputs generate( const std::string & text, @@ -50,4 +49,4 @@ class ChatterboxTurboTtsComponent { engine::models::chatterbox::EmbedReferenceOutputs builtin_ref_dict_; }; -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/assets.h b/include/engine/models/chatterbox_turbo/assets.h deleted file mode 100644 index 2381fc9ff..000000000 --- a/include/engine/models/chatterbox_turbo/assets.h +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include "engine/framework/assets/resource_bundle.h" -#include "engine/framework/assets/tensor_source.h" - -#include -#include - -namespace engine::models::chatterbox_turbo { - -struct ChatterboxTurboAssets { - engine::assets::ResourceBundle resources; - std::shared_ptr t3_turbo_weights; - std::shared_ptr builtin_conditionals_turbo; - // Raw path to the T3 GGUF (not just a TensorSource) so the tokenizer can read the embedded - // `tokenizer.ggml.tokens`/`tokenizer.ggml.merges` metadata arrays directly, and so the - // sibling S3Gen GGUF's path can be derived (see s3gen_turbo.cpp: replaces "-t3-" with - // "-s3gen-" in the filename — the two GGUF files can't both live loosely in one directory - // under this framework's package loader, so only the T3 file is referenced by model_specs; - // the S3Gen file is located and opened directly here). - std::filesystem::path t3_gguf_path; -}; - -std::shared_ptr load_chatterbox_turbo_assets(const std::filesystem::path & model_path); - -} // namespace engine::models::chatterbox_turbo diff --git a/include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h b/include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h deleted file mode 100644 index 5a503329d..000000000 --- a/include/engine/models/chatterbox_turbo/text_tokenizer_turbo.h +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include "engine/framework/tokenizers/llama_bpe.h" - -#include -#include -#include -#include -#include - -namespace engine::models::chatterbox_turbo { - -// Builds a GPT2-style BPE tokenizer from the vocab/merges arrays embedded as GGUF metadata -// (`tokenizer.ggml.tokens`/`tokenizer.ggml.merges`) in Chatterbox Turbo's T3 GGUF — there is no -// separate tokenizer.json file in this package. Materializes a vocab.json + merges.txt pair -// under a scratch directory and feeds those to the existing LlamaBpeTokenizer loader (Gpt2 -// pre-tokenizer mode) rather than teaching that shared loader a new in-memory-array input path. -std::shared_ptr load_chatterbox_turbo_tokenizer( - const std::filesystem::path & t3_gguf_path, - const std::filesystem::path & scratch_dir); - -// 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::models::chatterbox_turbo diff --git a/model_specs/chatterbox_turbo.json b/model_specs/chatterbox_turbo.json index 82e72523d..13cd446ab 100644 --- a/model_specs/chatterbox_turbo.json +++ b/model_specs/chatterbox_turbo.json @@ -39,7 +39,7 @@ "package_defaults": { "download": { "kind": "huggingface_snapshot", - "repo": "cstr/chatterbox-turbo-GGUF", + "repo": "audio-cpp/audio.cpp-gguf", "revision": "main", "gated": false } @@ -53,10 +53,9 @@ "precision": "q8_0", "target_directory": "Chatterbox-Turbo-GGUF", "files": [ - "chatterbox-turbo-t3-q8_0.gguf", - "chatterbox-turbo-s3gen-q8_0.gguf" + "Chatterbox-Turbo-GGUF/chatterbox-turbo-q8_0.gguf" ], - "description": "Third-party GGUF conversion of ResembleAI/chatterbox-turbo, published by 'cstr' for the CrispASR project (MIT-relicensed); not published by ResembleAI or audio.cpp. Two loose GGUF files (T3 + S3Gen) live side by side in this package's directory, so point --model at the T3 file directly (chatterbox-turbo-t3-q8_0.gguf), not the containing folder, to avoid the ambiguous-directory-GGUF check." + "strip_prefix": "Chatterbox-Turbo-GGUF" }, { "id": "chatterbox_turbo_f16", @@ -65,10 +64,9 @@ "precision": "f16", "target_directory": "Chatterbox-Turbo-GGUF", "files": [ - "chatterbox-turbo-t3-f16.gguf", - "chatterbox-turbo-s3gen-f16.gguf" + "Chatterbox-Turbo-GGUF/chatterbox-turbo-f16.gguf" ], - "description": "Third-party GGUF conversion of ResembleAI/chatterbox-turbo, published by 'cstr' for the CrispASR project (MIT-relicensed); not published by ResembleAI or audio.cpp. Two loose GGUF files (T3 + S3Gen) live side by side in this package's directory, so point --model at the T3 file directly (chatterbox-turbo-f16.gguf), not the containing folder, to avoid the ambiguous-directory-GGUF check." + "strip_prefix": "Chatterbox-Turbo-GGUF" } ], "sources": [ @@ -76,10 +74,26 @@ "format": "gguf", "roots": { "model": ".", - "t3_weights": "$gguf" + "weights": "$gguf" }, "files": { - "t3_gguf_path": "t3_weights:" + "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/models/chatterbox_turbo/components/t3_turbo_runtime.h b/src/community_models/chatterbox_turbo/components/t3_turbo_runtime.h similarity index 99% rename from src/models/chatterbox_turbo/components/t3_turbo_runtime.h rename to src/community_models/chatterbox_turbo/components/t3_turbo_runtime.h index 006563f40..27b263986 100644 --- a/src/models/chatterbox_turbo/components/t3_turbo_runtime.h +++ b/src/community_models/chatterbox_turbo/components/t3_turbo_runtime.h @@ -1,6 +1,6 @@ #pragma once -#include "engine/models/chatterbox_turbo/t3_turbo_component.h" +#include "engine/community_models/chatterbox_turbo/t3_turbo_component.h" #include "engine/framework/core/backend.h" #include "engine/framework/core/execution_context.h" @@ -20,7 +20,7 @@ #include #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { namespace { constexpr int32_t kTurboStartSpeechToken = 6561; @@ -808,4 +808,4 @@ class T3TurboPrefillBackendRunner { }; } // namespace -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/loader.cpp b/src/community_models/chatterbox_turbo/loader.cpp similarity index 95% rename from src/models/chatterbox_turbo/loader.cpp rename to src/community_models/chatterbox_turbo/loader.cpp index 1ee809624..b5fbe8525 100644 --- a/src/models/chatterbox_turbo/loader.cpp +++ b/src/community_models/chatterbox_turbo/loader.cpp @@ -1,13 +1,13 @@ -#include "engine/models/chatterbox_turbo/loader.h" +#include "engine/community_models/chatterbox_turbo/loader.h" -#include "engine/models/chatterbox_turbo/session.h" +#include "engine/community_models/chatterbox_turbo/session.h" #include "engine/framework/model_spec/package.h" #include #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { namespace { @@ -125,4 +125,4 @@ std::shared_ptr make_chatterbox_turbo_loader() { return std::make_shared(); } -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/s3gen_turbo.cpp b/src/community_models/chatterbox_turbo/s3gen_turbo.cpp similarity index 60% rename from src/models/chatterbox_turbo/s3gen_turbo.cpp rename to src/community_models/chatterbox_turbo/s3gen_turbo.cpp index 23c8b35ba..a9ef66071 100644 --- a/src/models/chatterbox_turbo/s3gen_turbo.cpp +++ b/src/community_models/chatterbox_turbo/s3gen_turbo.cpp @@ -1,21 +1,18 @@ -#include "engine/models/chatterbox_turbo/s3gen_turbo.h" - -#include "components/dot_prefix_source.h" -#include "components/s3gen_name_bridge.h" +#include "engine/community_models/chatterbox_turbo/s3gen_turbo.h" #include -namespace engine::models::chatterbox_turbo { +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 (see - // s3gen_name_bridge.h). The turbo 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. + // 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; @@ -38,50 +35,34 @@ engine::modules::HiftVocoderConfig make_turbo_hift_config(engine::assets::Tensor config.f0_in_channels = 80; config.f0_cond_channels = 512; config.weight_storage_type = weight_storage_type; - // The "v." vocoder namespace is applied inside translate_vocoder_name (see - // s3gen_name_bridge.cpp) rather than here, since that translation runs on top of the raw - // dot-prefixed source and needs to see the plain (unprefixed) name HiftVocoderComponent asks - // for in order to match its own prefix rules. - config.tensor_prefix = ""; + // 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::filesystem::path derive_turbo_s3gen_gguf_path(const std::filesystem::path & t3_gguf_path) { - const std::string filename = t3_gguf_path.filename().string(); - const std::string marker = "-t3-"; - const auto pos = filename.find(marker); - if (pos == std::string::npos) { - throw std::runtime_error( - "Chatterbox Turbo T3 GGUF filename does not match the expected '...-t3-.gguf' naming: " + filename); - } - const std::string s3gen_filename = filename.substr(0, pos) + "-s3gen-" + filename.substr(pos + marker.size()); - return t3_gguf_path.parent_path() / s3gen_filename; -} - std::shared_ptr ChatterboxTurboS3Gen::load( - const std::filesystem::path & s3gen_gguf_path, + std::shared_ptr s3gen_source, const engine::core::ExecutionContext & execution_context, engine::assets::TensorStorageType weight_storage_type) { - auto raw_source = make_dot_prefixed_tensor_source(engine::assets::open_tensor_source(s3gen_gguf_path), "s3"); - auto flow_source = make_turbo_s3gen_flow_bridge(raw_source); - auto vocoder_source = make_turbo_vocoder_bridge(raw_source); - auto out = std::make_shared(); out->execution_context_ = &execution_context; out->encoder_weights_ = - engine::models::chatterbox::load_s3_flow_encoder_weights(*flow_source, execution_context, weight_storage_type); + 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(*flow_source, execution_context, weight_storage_type); + 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 GGUF is missing the meanflow time_embed_mixer tensor (flow.decoder.estimator.time_embed_mixer)"); + "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( - vocoder_source, execution_context.config(), make_turbo_hift_config(weight_storage_type))); + s3gen_source, execution_context.config(), make_turbo_hift_config(weight_storage_type))); return out; } @@ -119,4 +100,4 @@ engine::models::chatterbox::S3GenInferenceOutputs ChatterboxTurboS3Gen::synthesi return outputs; } -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/session.cpp b/src/community_models/chatterbox_turbo/session.cpp similarity index 94% rename from src/models/chatterbox_turbo/session.cpp rename to src/community_models/chatterbox_turbo/session.cpp index 5670fd176..19790b885 100644 --- a/src/models/chatterbox_turbo/session.cpp +++ b/src/community_models/chatterbox_turbo/session.cpp @@ -1,11 +1,11 @@ -#include "engine/models/chatterbox_turbo/session.h" +#include "engine/community_models/chatterbox_turbo/session.h" #include "engine/framework/runtime/options.h" #include "engine/framework/text/chunking.h" #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { namespace { @@ -83,8 +83,7 @@ void ChatterboxTurboSession::prepare(const runtime::SessionPreparationRequest & "-- omit the speaker reference audio to use the built-in voice"); } if (!component_) { - component_ = std::make_unique( - assets_, execution_context(), assets_->resources.model_root() / ".chatterbox_turbo_tokenizer_cache"); + component_ = std::make_unique(assets_, execution_context()); } mark_prepared(); } @@ -110,4 +109,4 @@ runtime::TaskResult ChatterboxTurboSession::run(const runtime::TaskRequest & req return result; } -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/t3_turbo_component.cpp b/src/community_models/chatterbox_turbo/t3_turbo_component.cpp similarity index 99% rename from src/models/chatterbox_turbo/t3_turbo_component.cpp rename to src/community_models/chatterbox_turbo/t3_turbo_component.cpp index 3a50c5124..c25f880d2 100644 --- a/src/models/chatterbox_turbo/t3_turbo_component.cpp +++ b/src/community_models/chatterbox_turbo/t3_turbo_component.cpp @@ -3,7 +3,7 @@ #include #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { struct T3TurboInferenceComponent::State { void release_runtime_graphs() { @@ -292,4 +292,4 @@ void T3TurboInferenceComponent::release_runtime_cache() const { state_->release_runtime_cache(); } -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/t3_turbo_weights.cpp b/src/community_models/chatterbox_turbo/t3_turbo_weights.cpp similarity index 97% rename from src/models/chatterbox_turbo/t3_turbo_weights.cpp rename to src/community_models/chatterbox_turbo/t3_turbo_weights.cpp index 05a901259..0b46085fc 100644 --- a/src/models/chatterbox_turbo/t3_turbo_weights.cpp +++ b/src/community_models/chatterbox_turbo/t3_turbo_weights.cpp @@ -1,8 +1,8 @@ -#include "engine/models/chatterbox_turbo/t3_turbo_component.h" +#include "engine/community_models/chatterbox_turbo/t3_turbo_component.h" #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { namespace { @@ -179,4 +179,4 @@ std::shared_ptr load_t3_turbo_inference_weights( return weights; } -} // namespace engine::models::chatterbox_turbo +} // 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/models/chatterbox_turbo/tts.cpp b/src/community_models/chatterbox_turbo/tts.cpp similarity index 85% rename from src/models/chatterbox_turbo/tts.cpp rename to src/community_models/chatterbox_turbo/tts.cpp index c0778bf69..90f41a1e3 100644 --- a/src/models/chatterbox_turbo/tts.cpp +++ b/src/community_models/chatterbox_turbo/tts.cpp @@ -1,11 +1,11 @@ -#include "engine/models/chatterbox_turbo/tts.h" +#include "engine/community_models/chatterbox_turbo/tts.h" -#include "engine/models/chatterbox_turbo/text_tokenizer_turbo.h" +#include "engine/community_models/chatterbox_turbo/text_tokenizer_turbo.h" #include #include -namespace engine::models::chatterbox_turbo { +namespace engine::community_models::chatterbox_turbo { namespace { @@ -34,16 +34,17 @@ std::vector require_i32_array( ChatterboxTurboTtsComponent::ChatterboxTurboTtsComponent( std::shared_ptr assets, - const engine::core::ExecutionContext & execution_context, - const std::filesystem::path & scratch_dir) + const engine::core::ExecutionContext & execution_context) : assets_(std::move(assets)) { - tokenizer_ = load_chatterbox_turbo_tokenizer(assets_->t3_gguf_path, scratch_dir); + 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); - const auto s3gen_path = derive_turbo_s3gen_gguf_path(assets_->t3_gguf_path); - s3gen_ = ChatterboxTurboS3Gen::load(s3gen_path, 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}); @@ -93,4 +94,4 @@ engine::models::chatterbox::S3GenInferenceOutputs ChatterboxTurboTtsComponent::g return s3gen_->synthesize(builtin_ref_dict_, speech_tokens, config.seed, config.seed); } -} // namespace engine::models::chatterbox_turbo +} // namespace engine::community_models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/assets.cpp b/src/models/chatterbox_turbo/assets.cpp deleted file mode 100644 index 503f23595..000000000 --- a/src/models/chatterbox_turbo/assets.cpp +++ /dev/null @@ -1,29 +0,0 @@ -#include "engine/models/chatterbox_turbo/assets.h" - -#include "components/dot_prefix_source.h" - -#include "engine/framework/model_spec/package.h" - -#include - -namespace engine::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_gguf_path = out->resources.require_file("t3_gguf_path"); - if (out->t3_gguf_path.extension() != ".gguf") { - throw std::runtime_error("Chatterbox Turbo requires a GGUF model path (point --model at the T3 GGUF file directly)"); - } - // The turbo GGUF's tensor names are flat-dotted ("t3.blk.0...", "conds.gen...") with no "/" - // namespace separator, so the framework's own prefix-stripping (which expects "/") can't be - // used here -- see components/dot_prefix_source.h. - const auto t3_raw = engine::assets::open_tensor_source(out->t3_gguf_path); - out->t3_turbo_weights = make_dot_prefixed_tensor_source(t3_raw, "t3"); - out->builtin_conditionals_turbo = make_dot_prefixed_tensor_source(t3_raw, "conds"); - return out; -} - -} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/dot_prefix_source.cpp b/src/models/chatterbox_turbo/components/dot_prefix_source.cpp deleted file mode 100644 index 258f64c02..000000000 --- a/src/models/chatterbox_turbo/components/dot_prefix_source.cpp +++ /dev/null @@ -1,56 +0,0 @@ -#include "dot_prefix_source.h" - -#include - -namespace engine::models::chatterbox_turbo { - -namespace { - -class DotPrefixedTensorSource final : public engine::assets::TensorSource { -public: - DotPrefixedTensorSource(std::shared_ptr delegate, std::string prefix) - : delegate_(std::move(delegate)), prefix_(std::move(prefix) + ".") {} - - const std::filesystem::path & source_path() const noexcept override { - return delegate_->source_path(); - } - bool has_tensor(std::string_view name) const noexcept override { - return delegate_->has_tensor(prefix_ + std::string(name)); - } - engine::assets::TensorMetadata require_metadata(std::string_view name) const override { - return delegate_->require_metadata(prefix_ + std::string(name)); - } - std::vector tensors() const override { - return delegate_->tensors(); - } - engine::assets::RawTensorData require_tensor_data(std::string_view name) const override { - return delegate_->require_tensor_data(prefix_ + std::string(name)); - } - std::vector require_f32( - std::string_view name, - const std::optional> & expected_shape) const override { - return delegate_->require_f32(prefix_ + std::string(name), expected_shape); - } - std::optional> optional_f32( - std::string_view name, - const std::optional> & expected_shape) const override { - return delegate_->optional_f32(prefix_ + std::string(name), expected_shape); - } - int64_t require_i64_scalar(std::string_view name) const override { - return delegate_->require_i64_scalar(prefix_ + std::string(name)); - } - -private: - std::shared_ptr delegate_; - std::string prefix_; -}; - -} // namespace - -std::shared_ptr make_dot_prefixed_tensor_source( - std::shared_ptr source, - std::string prefix) { - return std::make_shared(std::move(source), std::move(prefix)); -} - -} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/dot_prefix_source.h b/src/models/chatterbox_turbo/components/dot_prefix_source.h deleted file mode 100644 index b0858ab84..000000000 --- a/src/models/chatterbox_turbo/components/dot_prefix_source.h +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -// The turbo GGUF (cstr/chatterbox-turbo-GGUF) names every tensor with a flat dot-separated path -// (e.g. "t3.blk.0.attn_norm.weight", "conds.gen.embedding", "s3.flow.encoder_proj") with no "/" -// namespace separator anywhere. The framework's own prefix-stripping helpers -// (engine::assets::open_tensor_source(path, prefix) / make_prefixed_tensor_source) expect a -// "/"-delimited packed-GGUF namespace instead (matching how this project's own GGUF converter -// packs multiple checkpoints into one file, e.g. "t3_english/tfmr.layers.0...") and throw -// "packed GGUF namespace does not exist" against a dot-separated name. This is therefore a -// separate, dot-based sibling used only for chatterbox_turbo's third-party GGUF. - -#include "engine/framework/assets/tensor_source.h" - -#include -#include - -namespace engine::models::chatterbox_turbo { - -// Strips a literal "." from every queried name before delegating to `source` (e.g. -// prefix="t3" turns a query for "blk.0.attn_norm.weight" into "t3.blk.0.attn_norm.weight" on the -// underlying source). -std::shared_ptr make_dot_prefixed_tensor_source( - std::shared_ptr source, - std::string prefix); - -} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp b/src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp deleted file mode 100644 index c9b513134..000000000 --- a/src/models/chatterbox_turbo/components/s3gen_name_bridge.cpp +++ /dev/null @@ -1,207 +0,0 @@ -#include "s3gen_name_bridge.h" - -#include -#include -#include -#include - -namespace engine::models::chatterbox_turbo { -namespace { - -using Rule = std::pair; - -// Longest/most-specific prefixes first so a shorter rule never shadows a longer one -// (e.g. "flow.encoder_proj" must not be caught by a "flow.encoder." rule). -constexpr std::array kFlowPrefixRules = {{ - {"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. -}}; - -// Applied to whatever remains after a prefix rule above consumed the "flow.decoder.estimator.*" -// or "flow.encoder.{encoders,up_encoders}.N." portion of the name. -constexpr std::array kSegmentRules = {{ - {".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"}, -}}; - -constexpr std::array kTransformerBlockRules = {{ - {".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"}, -}}; - -constexpr std::array kVocoderPrefixRules = {{ - {"conv_pre", "cpre"}, - {"conv_post", "cpost"}, - {"resblocks.", "rb."}, - {"source_downs.", "sd."}, - {"source_resblocks.", "srb."}, - {"f0_predictor.condnet.", "f0.cn."}, -}}; - -constexpr std::array kVocoderExactRules = {{ - {"f0_predictor.classifier", "f0.cls"}, - {"m_source.l_linear", "ms.ll"}, - {".convs1.", ".c1."}, - {".convs2.", ".c2."}, - {".activations1.", ".a1."}, - {".activations2.", ".a2."}, -}}; - -bool starts_with(std::string_view text, std::string_view prefix) { - return text.size() >= prefix.size() && text.compare(0, prefix.size(), prefix) == 0; -} - -std::string replace_all(std::string text, std::string_view search, std::string_view replace) { - if (search.empty()) { - return text; - } - size_t pos = 0; - while ((pos = text.find(search, pos)) != std::string::npos) { - text.replace(pos, search.size(), replace); - pos += replace.size(); - } - return text; -} - -std::string translate_flow_name(std::string_view name) { - std::string translated(name); - for (const auto & [from, to] : kFlowPrefixRules) { - if (starts_with(translated, from)) { - translated = std::string(to) + translated.substr(from.size()); - break; - } - } - for (const auto & [from, to] : kSegmentRules) { - translated = replace_all(std::move(translated), from, to); - } - // ".ff.net.2" would also match the shorter ".ff.net.0.proj" search text's neighbourhood, - // so handle it (and the rest of the transformer-block-only renames) after the segment pass. - translated = replace_all(std::move(translated), ".ff.net.2", ".ff.down"); - for (const auto & [from, to] : kTransformerBlockRules) { - translated = replace_all(std::move(translated), from, to); - } - return translated; -} - -std::string translate_vocoder_name(std::string_view name) { - std::string translated(name); - for (const auto & [from, to] : kVocoderPrefixRules) { - if (starts_with(translated, from)) { - translated = std::string(to) + translated.substr(from.size()); - break; - } - } - for (const auto & [from, to] : kVocoderExactRules) { - if (translated == from || starts_with(translated, std::string(from) + ".")) { - translated = std::string(to) + translated.substr(from.size()); - break; - } - translated = replace_all(translated, from, to); - } - // HiftVocoderConfig.tensor_prefix is left empty (see s3gen_turbo.cpp) so this translator - // sees HiftVocoderComponent's plain requested names; the turbo GGUF's vocoder tensors all - // live under the top-level "v." namespace, which is added here rather than via tensor_prefix. - return "v." + translated; -} - -class TranslatingTensorSource final : public engine::assets::TensorSource { -public: - TranslatingTensorSource( - std::shared_ptr delegate, - std::string (*translate)(std::string_view)) - : delegate_(std::move(delegate)), translate_(translate) {} - - const std::filesystem::path & source_path() const noexcept override { - return delegate_->source_path(); - } - bool has_tensor(std::string_view name) const noexcept override { - return delegate_->has_tensor(translate_(name)); - } - engine::assets::TensorMetadata require_metadata(std::string_view name) const override { - return delegate_->require_metadata(translate_(name)); - } - std::vector tensors() const override { - return delegate_->tensors(); - } - engine::assets::RawTensorData require_tensor_data(std::string_view name) const override { - return delegate_->require_tensor_data(translate_(name)); - } - // Some turbo GGUF tensors squeeze a leading size-1 dimension the base loader code still - // expects explicitly (e.g. the vocoder's f0 classifier: nn.Linear(512, 1).weight is stored - // as ne=[512] rather than [1, 512]). Rather than special-case those, skip the delegate's - // strict per-dimension shape check here and only verify the element count still matches. - std::vector require_f32( - std::string_view name, - const std::optional> & expected_shape) const override { - auto values = delegate_->require_f32(translate_(name), std::nullopt); - if (expected_shape.has_value()) { - int64_t expected_count = 1; - for (int64_t dim : *expected_shape) { - expected_count *= dim; - } - if (static_cast(values.size()) != expected_count) { - throw std::runtime_error("tensor element count mismatch for " + std::string(name)); - } - } - return values; - } - std::optional> optional_f32( - std::string_view name, - const std::optional> & expected_shape) const override { - if (!has_tensor(name)) { - return std::nullopt; - } - return require_f32(name, expected_shape); - } - int64_t require_i64_scalar(std::string_view name) const override { - return delegate_->require_i64_scalar(translate_(name)); - } - -private: - std::shared_ptr delegate_; - std::string (*translate_)(std::string_view); -}; - -} // namespace - -std::shared_ptr make_turbo_s3gen_flow_bridge( - std::shared_ptr turbo_s3_source) { - return std::make_shared(std::move(turbo_s3_source), &translate_flow_name); -} - -std::shared_ptr make_turbo_vocoder_bridge( - std::shared_ptr turbo_s3_source) { - return std::make_shared(std::move(turbo_s3_source), &translate_vocoder_name); -} - -} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/components/s3gen_name_bridge.h b/src/models/chatterbox_turbo/components/s3gen_name_bridge.h deleted file mode 100644 index 9ae2f1c1f..000000000 --- a/src/models/chatterbox_turbo/components/s3gen_name_bridge.h +++ /dev/null @@ -1,61 +0,0 @@ -#pragma once - -// Chatterbox Turbo's S3Gen weights ship in a third-party GGUF (cstr/chatterbox-turbo-GGUF) -// whose tensor names are a consistent, decodable abbreviation of the names the existing -// chatterbox S3Gen loader code (src/models/chatterbox/s3gen_flow.cpp, -// src/framework/modules/vocoders/hift_vocoder.cpp) already expects and knows how to build -// ggml graphs for. Rather than reimplement the flow encoder/decoder/vocoder, this header -// name-translates so the *existing* loaders can run unmodified against the turbo GGUF. -// -// Mapping verified against the actual GGUF header of chatterbox-turbo-s3gen-q8_0.gguf -// (parsed offline with a throwaway GGUF reader), matched segment-by-segment against the -// literal prefixes s3gen_flow.cpp/hift_vocoder.cpp construct: -// -// flow.encoder.encoders.N -> fe.enc.N -// flow.encoder.up_encoders.N -> fe.ue.N -// 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.encoder_proj / spk_embed_affine_layer / input_embedding (unchanged, top-level "flow.") -// (within a conformer layer) norm_mha -> nmha, self_attn.linear_{q,k,v,out,pos} -> sa.l{q,k,v,o,p}, -// self_attn.pos_bias_{u,v} -> sa.pb{u,v}, norm_ff -> nff, feed_forward.w_1/2 -> ff.w_1/2 -// 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 (unchanged linear_1/linear_2 suffix) -// flow.decoder.estimator.time_embed_mixer -> fd.tmx (turbo-only meanflow tensor, no base equivalent) -// (within a resnet block) blockN.block.0/2 -> bN.0/2, res_conv -> rc, mlp.1 unchanged -// (within a transformer block) norm1/norm3 unchanged, attn1.to_q/k/v -> attn1.q/k/v, -// attn1.to_out.0 -> attn1.o, ff.net.0.proj -> ff.up, ff.net.2 -> ff.down -// (vocoder, HiftVocoderConfig.tensor_prefix="v.") conv_pre -> cpre, conv_post -> cpost, -// ups.N unchanged, resblocks.N -> rb.N, source_downs.N -> sd.N, source_resblocks.N -> srb.N, -// f0_predictor.condnet.N -> f0.cn.N, f0_predictor.classifier -> f0.cls, m_source.l_linear -> ms.ll -// (within a resblock) convs1/2 -> c1/2, activations1/2 -> a1/2 -// -// This is deliberately additive/self-contained to chatterbox_turbo: it does not modify the -// base chatterbox family's own tensor naming or loading code. - -#include "engine/framework/assets/tensor_source.h" - -#include -#include - -namespace engine::models::chatterbox_turbo { - -// Wraps an already s3-prefix-stripped TensorSource (see open_tensor_source(path, "s3")) and -// rewrites incoming chatterbox-style names to the turbo GGUF's abbreviated names before -// delegating every lookup. -std::shared_ptr make_turbo_s3gen_flow_bridge( - std::shared_ptr turbo_s3_source); - -// Same idea for the HiFT vocoder; construct the shared HiftVocoderComponent with -// tensor_prefix = "v." against this bridge and weight_layout = Plain (the turbo GGUF already -// folds torch weight-norm parametrization into a plain "weight" tensor at conversion time). -std::shared_ptr make_turbo_vocoder_bridge( - std::shared_ptr turbo_s3_source); - -} // namespace engine::models::chatterbox_turbo diff --git a/src/models/chatterbox_turbo/text_tokenizer_turbo.cpp b/src/models/chatterbox_turbo/text_tokenizer_turbo.cpp deleted file mode 100644 index 0aec66d90..000000000 --- a/src/models/chatterbox_turbo/text_tokenizer_turbo.cpp +++ /dev/null @@ -1,212 +0,0 @@ -#include "engine/models/chatterbox_turbo/text_tokenizer_turbo.h" - -#include -#include - -#include -#include -#include -#include - -namespace engine::models::chatterbox_turbo { - -namespace { - -struct GgufContextDeleter { - void operator()(gguf_context * ctx) const noexcept { - if (ctx != nullptr) { - gguf_free(ctx); - } - } -}; - -std::vector read_gguf_string_array(gguf_context * ctx, const char * key) { - const int64_t key_id = gguf_find_key(ctx, key); - if (key_id < 0) { - throw std::runtime_error(std::string("Chatterbox Turbo GGUF is missing metadata key: ") + key); - } - if (gguf_get_arr_type(ctx, key_id) != GGUF_TYPE_STRING) { - throw std::runtime_error(std::string("Chatterbox Turbo GGUF metadata key is not a string array: ") + key); - } - const size_t count = gguf_get_arr_n(ctx, key_id); - std::vector out; - out.reserve(count); - for (size_t i = 0; i < count; ++i) { - out.emplace_back(gguf_get_arr_str(ctx, key_id, i)); - } - return out; -} - -std::string json_escape(const std::string & value) { - std::string out; - out.reserve(value.size() + 8); - for (unsigned char c : value) { - switch (c) { - case '"': out += "\\\""; break; - case '\\': out += "\\\\"; break; - case '\n': out += "\\n"; break; - case '\r': out += "\\r"; break; - case '\t': out += "\\t"; break; - default: - if (c < 0x20) { - char buf[8]; - std::snprintf(buf, sizeof(buf), "\\u%04x", c); - out += buf; - } else { - out += static_cast(c); - } - } - } - return out; -} - -bool looks_like_bracket_tag(const std::string & token) { - return token.size() >= 3 && token.front() == '[' && token.back() == ']'; -} - -} // namespace - -std::shared_ptr load_chatterbox_turbo_tokenizer( - const std::filesystem::path & t3_gguf_path, - const std::filesystem::path & scratch_dir) { - gguf_init_params params{}; - params.no_alloc = true; - params.ctx = nullptr; - std::unique_ptr ctx(gguf_init_from_file(t3_gguf_path.string().c_str(), params)); - if (!ctx) { - throw std::runtime_error("failed to open Chatterbox Turbo T3 GGUF for tokenizer metadata: " + t3_gguf_path.string()); - } - - const auto tokens = read_gguf_string_array(ctx.get(), "tokenizer.ggml.tokens"); - const auto merges = read_gguf_string_array(ctx.get(), "tokenizer.ggml.merges"); - if (tokens.empty()) { - throw std::runtime_error("Chatterbox Turbo T3 GGUF has an empty tokenizer.ggml.tokens array"); - } - - std::filesystem::create_directories(scratch_dir); - const auto vocab_path = scratch_dir / "chatterbox_turbo_vocab.json"; - const auto merges_path = scratch_dir / "chatterbox_turbo_merges.txt"; - - // The 19 emotion/style control tags ([laugh], [sigh], ...) sit at the tail of the vocab and - // cannot be reached through byte-pair merges (they were never trained into the merge table), - // so they must be registered as atomic special tokens the pre-tokenizer matches directly - // (LlamaBpeTokenizerSpec::additional_special_tokens) rather than as plain vocab entries. - // Crucially they must NOT also appear in vocab.json below: add_runtime_special_tokens only - // sets a token's "special" attribute when it is inserting a genuinely new id -- if the id - // already exists (which it would, from vocab.json) it just verifies the id and continues, - // silently leaving the token without its special/atomic-match flag, so it gets shredded into - // ordinary byte-level BPE pieces instead of matched as one unit. - std::vector special_tokens; - int64_t special_tokens_start = static_cast(tokens.size()); - for (int64_t id = static_cast(tokens.size()) - 1; id >= 0; --id) { - if (!looks_like_bracket_tag(tokens[static_cast(id)])) { - break; - } - special_tokens.emplace_back(tokens[static_cast(id)], static_cast(id)); - special_tokens_start = id; - } - - { - std::ofstream out(vocab_path, std::ios::trunc); - if (!out) { - throw std::runtime_error("failed to write Chatterbox Turbo tokenizer vocab: " + vocab_path.string()); - } - out << "{"; - bool first = true; - for (size_t i = 0; i < tokens.size(); ++i) { - if (static_cast(i) >= special_tokens_start) { - continue; - } - if (!first) { - out << ","; - } - first = false; - out << "\"" << json_escape(tokens[i]) << "\":" << i; - } - out << "}"; - } - { - std::ofstream out(merges_path, std::ios::trunc); - if (!out) { - throw std::runtime_error("failed to write Chatterbox Turbo tokenizer merges: " + merges_path.string()); - } - for (const auto & merge : merges) { - out << merge << "\n"; - } - } - - 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::models::chatterbox_turbo 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()