From 73d707af574d09dc43ec57d32450285744e1851b Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Tue, 11 Aug 2026 16:10:38 +0800 Subject: [PATCH 1/8] index_tts2_5: add IndexTTS-2.5 multilingual TTS family Add support for IndexTTS-2.5 (GPT + DiT CFM + BigVGAN, zh/en/ja/es/ar) as a new model family alongside index_tts2: - tiktoken BPE tokenizer (multilingual_zh_ja_yue_char_del, 60509 vocab) on top of the vendored llama bpe-core, with <|lang|> prefixes, pronunciation annotations (), and language-aware casing - GPT speaker conditioning via CAMPPlus linear projection (3-token prefix) replacing the conformer+perceiver path, plus lang_embedding - semantic codec: full EnhancedCodec decode chain (decoder backbone + 2x nearest upsample + up conv); prompt_condition now consumes raw normalized w2v-bert semantics as upstream 2.5 does - s2mel/bucket flow drops the gpt-latent second pass (use_gpt_latent defaults to false upstream) - model spec, GGUF converter inputs, warm bench tests, CLI longform cases, WebUI catalog entries, and docs/tts.md section Verified against the official Python reference on CUDA: greedy zh output matches upstream token-for-prefix (99/143 codes, f16 drift after) with identical 5.70s duration; ASR transcripts match for zh/en/ja; q8_0 and native safetensors layouts both pass. --- CMakeLists.txt | 22 + README.md | 1 + docs/tts.md | 67 + include/engine/models/index_tts2_5/assets.h | 29 + .../models/index_tts2_5/audio_features.h | 54 + include/engine/models/index_tts2_5/gpt.h | 186 ++ include/engine/models/index_tts2_5/loader.h | 33 + .../engine/models/index_tts2_5/qwen_emotion.h | 82 + include/engine/models/index_tts2_5/request.h | 14 + include/engine/models/index_tts2_5/s2mel.h | 151 ++ .../models/index_tts2_5/semantic_codec.h | 94 + .../models/index_tts2_5/semantic_encoder.h | 100 + include/engine/models/index_tts2_5/session.h | 122 + .../models/index_tts2_5/style_encoder.h | 34 + .../models/index_tts2_5/tokenizer_text.h | 58 + include/engine/models/index_tts2_5/types.h | 153 ++ include/engine/models/index_tts2_5/vocoder.h | 34 + model_specs/index_tts2_5.json | 186 ++ src/models/index_tts2_5/assets.cpp | 259 ++ src/models/index_tts2_5/audio_features.cpp | 547 ++++ src/models/index_tts2_5/gpt.cpp | 2217 +++++++++++++++++ src/models/index_tts2_5/loader.cpp | 157 ++ src/models/index_tts2_5/qwen_emotion.cpp | 794 ++++++ src/models/index_tts2_5/request.cpp | 192 ++ src/models/index_tts2_5/s2mel.cpp | 1271 ++++++++++ src/models/index_tts2_5/semantic_codec.cpp | 723 ++++++ src/models/index_tts2_5/semantic_encoder.cpp | 622 +++++ src/models/index_tts2_5/session.cpp | 838 +++++++ src/models/index_tts2_5/style_encoder.cpp | 38 + src/models/index_tts2_5/tokenizer_text.cpp | 691 +++++ src/models/index_tts2_5/vocoder.cpp | 69 + .../index_tts2_5/index_tts2_5_warm_bench.cpp | 319 +++ .../index_tts2_5_warm_bench_cases.json | 131 + ...audiocpp_cli_longform_tts_clone_cases.json | 50 + webui/configs/model_params.json | 9 + webui/configs/models_catalog.json | 3 + webui/configs/required_files.json | 9 + 37 files changed, 10359 insertions(+) create mode 100644 include/engine/models/index_tts2_5/assets.h create mode 100644 include/engine/models/index_tts2_5/audio_features.h create mode 100644 include/engine/models/index_tts2_5/gpt.h create mode 100644 include/engine/models/index_tts2_5/loader.h create mode 100644 include/engine/models/index_tts2_5/qwen_emotion.h create mode 100644 include/engine/models/index_tts2_5/request.h create mode 100644 include/engine/models/index_tts2_5/s2mel.h create mode 100644 include/engine/models/index_tts2_5/semantic_codec.h create mode 100644 include/engine/models/index_tts2_5/semantic_encoder.h create mode 100644 include/engine/models/index_tts2_5/session.h create mode 100644 include/engine/models/index_tts2_5/style_encoder.h create mode 100644 include/engine/models/index_tts2_5/tokenizer_text.h create mode 100644 include/engine/models/index_tts2_5/types.h create mode 100644 include/engine/models/index_tts2_5/vocoder.h create mode 100644 model_specs/index_tts2_5.json create mode 100644 src/models/index_tts2_5/assets.cpp create mode 100644 src/models/index_tts2_5/audio_features.cpp create mode 100644 src/models/index_tts2_5/gpt.cpp create mode 100644 src/models/index_tts2_5/loader.cpp create mode 100644 src/models/index_tts2_5/qwen_emotion.cpp create mode 100644 src/models/index_tts2_5/request.cpp create mode 100644 src/models/index_tts2_5/s2mel.cpp create mode 100644 src/models/index_tts2_5/semantic_codec.cpp create mode 100644 src/models/index_tts2_5/semantic_encoder.cpp create mode 100644 src/models/index_tts2_5/session.cpp create mode 100644 src/models/index_tts2_5/style_encoder.cpp create mode 100644 src/models/index_tts2_5/tokenizer_text.cpp create mode 100644 src/models/index_tts2_5/vocoder.cpp create mode 100644 tests/index_tts2_5/index_tts2_5_warm_bench.cpp create mode 100644 tests/index_tts2_5/index_tts2_5_warm_bench_cases.json diff --git a/CMakeLists.txt b/CMakeLists.txt index 664cde5c..aa6da691 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -812,6 +812,27 @@ audiocpp_add_model(index_tts2 engine::models::index_tts2::make_index_tts2_loader ) +audiocpp_add_model(index_tts2_5 + SOURCES + src/models/index_tts2_5/assets.cpp + src/models/index_tts2_5/audio_features.cpp + src/models/index_tts2_5/gpt.cpp + src/models/index_tts2_5/loader.cpp + src/models/index_tts2_5/qwen_emotion.cpp + src/models/index_tts2_5/request.cpp + src/models/index_tts2_5/s2mel.cpp + src/models/index_tts2_5/semantic_codec.cpp + src/models/index_tts2_5/semantic_encoder.cpp + src/models/index_tts2_5/session.cpp + src/models/index_tts2_5/style_encoder.cpp + src/models/index_tts2_5/tokenizer_text.cpp + src/models/index_tts2_5/vocoder.cpp + INCLUDES + engine/models/index_tts2_5/loader.h + LOADERS + engine::models::index_tts2_5::make_index_tts2_5_loader +) + audiocpp_add_model(nemotron_asr SOURCES src/models/nemotron_asr/assets.cpp @@ -1357,6 +1378,7 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(higgs_audio_tts_warm_bench tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp) add_engine_warmbench(hviske_asr_warm_bench tests/hviske_asr/hviske_asr_warm_bench.cpp) add_engine_warmbench(index_tts2_warm_bench tests/index_tts2/index_tts2_warm_bench.cpp) + add_engine_warmbench(index_tts2_5_warm_bench tests/index_tts2_5/index_tts2_5_warm_bench.cpp) add_engine_warmbench(irodori_tts_warm_bench tests/irodori_tts/irodori_tts_warm_bench.cpp) add_engine_warmbench(marblenet_vad_warm_bench tests/marblenet_vad/marblenet_vad_warm_bench.cpp) add_engine_warmbench(miocodec_warm_bench tests/miocodec/miocodec_warm_bench.cpp) diff --git a/README.md b/README.md index 6fb0615d..39a85c51 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,7 @@ Runtime tags: safetensors is the default model loading path. `GGUF 16/Q8/Q4` mea | **voxtral_realtime** | ASR | auto | Voxtral-Mini-4B-Realtime-2602 | GGUF 16/Q8/Q4, Stream | | **voxcpm2** | TTS, Clone, Design, Ctrl | ar, da, de, el, en, es, fi, fr, he, hi, id, it, ja, km, ko, lo, ms, my, nl, no, pl, pt, ru, sv, sw, th, tl, tr, vi, zh | VoxCPM2-2B, 48 kHz | GGUF 16/Q8, Stream | | **index_tts2** | TTS, Clone, Ctrl | zh, en | IndexTTS-2 | GGUF 16/Q8 | +| **index_tts2_5** | TTS, Clone, Ctrl | zh, en, ja, es, ar | IndexTTS-2.5 | GGUF 16/Q8 | | **irodori_tts** | TTS, Clone, Design, Ctrl | ja | Irodori-TTS-v4-Small, Irodori-TTS-500M-v3, Irodori-TTS-600M-v3-VoiceDesign | GGUF 16/Q8 | | **moss_tts_nano** | TTS, Clone | auto | MOSS-TTS-Nano-100M | GGUF 16/Q8 | | **moss_tts_local** | TTS, Clone, Ctrl | auto, optional language hint | MOSS-TTS-Local-Transformer-v1.5 | GGUF 16/Q8 | diff --git a/docs/tts.md b/docs/tts.md index a7f2ee59..470d26c1 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -15,6 +15,7 @@ | Higgs Audio v3 TTS | `higgs_audio_tts` | `tts` | [Higgs Audio v3 TTS](#higgs-audio-v3-tts) | | Fish Audio S2 Pro | `fish_audio` | `tts` | [Fish Audio S2 Pro](#fish-audio-s2-pro) | | IndexTTS2 | `index_tts2` | `tts` | [IndexTTS2](#indextts2) | +| IndexTTS2.5 | `index_tts2_5` | `tts` | [IndexTTS2.5](#indextts25) | | Irodori-TTS | `irodori_tts` | `tts`, `vdes` | [Irodori-TTS](#irodori-tts) | | GLM-TTS | `glm_tts` | `tts`, `clon` | [GLM-TTS](#glm-tts) | | Inflect Micro v2 | `inflect_v2` | `tts` | [Inflect v2](#inflect-v2) | @@ -546,6 +547,72 @@ audiocpp_cli --task tts --family index_tts2 --model /path/to/IndexTTS-2 --backen | `--session-option index_tts2.emotion_text_max_new_tokens=` | tokens | `256` | Maximum generated tokens for emotion-text classification. | | `--session-option index_tts2.weight_context_mb=` | MB | `32` | Shared ggml weight metadata context size. | +## IndexTTS2.5 + +IndexTTS2.5 is IndexTeam/bilibili's multilingual zero-shot TTS model (released 2026-07): a 0.8B GPT (autoregressive) + DiT CFM + BigVGAN stack that keeps IndexTTS2's timbre-emotion decoupling and adds Japanese, Spanish, and Arabic on top of Chinese and English. It requires a speaker reference through the framework `--voice-ref` path. Inline `<文字|发音>` pronunciation overrides (pinyin, CMU phonemes, or kana) are supported. Upstream weights live at [IndexTeam/IndexTTS-2.5](https://huggingface.co/IndexTeam/IndexTTS-2.5); the reference implementation is [index-tts/index-tts](https://github.com/index-tts/index-tts) branch `indextts-2.5`. + +| Field | Value | +|---|---| +| Family | `index_tts2_5` | +| Model directory | `models/IndexTTS2.5-GGUF` (default GGUF package `index_tts2_5_q8_0`; `index_tts2_5_f16` and `index_tts2_5_orig` also available) | +| Task | `tts`, `clon` | +| Modes | `offline` | +| Languages | `zh`, `en`, `ja`, `es`, `ar` | +| Voice input | Required reference WAV through `--voice-ref` | +| Built-in voices | Not exposed | + +Voice clone: + +```bash +audiocpp_cli --task clon --family index_tts2_5 --model /path/to/IndexTTS2.5-GGUF --backend cuda --text "Hello from IndexTTS2.5." --voice-ref /path/to/reference.wav --out out.wav +``` + +Emotion text: + +```bash +audiocpp_cli --task tts --family index_tts2_5 --model /path/to/IndexTTS2.5-GGUF --backend cuda --text "今天的演示会更有情绪。" --voice-ref /path/to/reference.wav --emotion "你吓死我了！你是鬼吗？" --request-option emotion_alpha=0.6 --out out.wav +``` + +The `lang` request option selects the text language (`auto`, `zh`, `en`, `ja`, `es`, `ar`, or any tokenizer language code). The default `auto` picks `zh` when the text contains Han characters and `en` otherwise, so mixed Japanese/Spanish/Arabic text should set `--request-option lang=ja|es|ar` explicitly. + +Emotion conditioning supports all three IndexTTS2 paths: an emotion reference WAV through `--audio`, an explicit `emotion_vector`, and Qwen-based emotion-text classification through `--emotion` / `use_emotion_text`. Known limitation: the NeMo text normalizers for Japanese and Spanish are not ported, so ja/es input text is passed through without upstream-style normalization. + +License: IndexTTS-2.5 weights are distributed under the bilibili Model Use License, which is not OSI-approved. It requires separate commercial authorization when monthly active users exceed 100 million or annual revenue exceeds 1 billion RMB, and it forbids using model outputs to improve other AI models. Check the upstream repository for the full terms before redistribution or commercial use. + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--voice-ref` | WAV path | required | Reference speaker audio. | +| `--request-option lang=` | `auto`, `zh`, `en`, `ja`, `es`, `ar`, ... | `auto` | Text language hint; `auto` infers `zh` when the text contains Han characters, otherwise `en`. | +| `--emotion` | text | not set | Emotion-text conditioning through the framework style field. | +| `--request-option emotion_alpha=` | float in `[0, 1]` | `1.0` | Blend strength for explicit emotion conditioning. | +| `--request-option emotion_vector=` | 8 floats | not set | Explicit emotion vector. | +| `--request-option use_emotion_text=true|false` | bool | `false` | Infer emotion from text. | +| `--request-option use_random_emotion=true|false` | bool | `false` | Use random emotion weights in the emotion mixer. | +| `--request-option interval_silence_ms=` | milliseconds | `200` | Silence inserted between generated text chunks. | +| `--text-chunk-size` | characters | not set | Optional framework outer text chunk size. When omitted, IndexTTS2.5 keeps its internal tokenizer segmentation. | +| `--text-chunk-mode` | `default`, `tag_aware`, `japanese`, `endline` | `default` | Framework chunking mode used only when `--text-chunk-size` is set. | +| `--max-tokens` | integer | `1500` | Maximum generated GPT mel tokens. | +| `--temperature` | float | `0.8` | GPT sampling temperature. | +| `--top-p` | float | `0.8` | GPT nucleus sampling limit. | +| `--top-k` | integer | `30` | GPT top-k sampling limit. | +| `--repetition-penalty` | float | `10.0` | GPT repetition penalty. | +| `--do-sample` | `true`, `false` | `true` | Enable stochastic GPT sampling. | +| `--request-option length_penalty=` | float | `0.0` | GPT beam-search length penalty. | +| `--request-option num_beams=` | integer | `3` | GPT beam count. | +| `--session-option index_tts2_5.mem_saver=true|false` | bool | `false` | Release staged reference and conditioning graphs after request phases. | +| `--session-option index_tts2_5.weight_type=native|f32|f16|bf16|q8_0` | enum | `native` | Matmul weight storage type. | +| `--session-option index_tts2_5.conv_weight_type=native|f32|f16` | enum | `native` | Convolution weight storage type. | +| `--session-option index_tts2_5.speaker_cache_slots=` | integer slots | `1` | Prepared speaker-reference cache slots; set `0` to disable reuse. | +| `--session-option index_tts2_5.emotion_cache_slots=` | integer slots | `1` | Prepared emotion-reference cache slots; set `0` to disable reuse. | +| `--session-option index_tts2_5.emotion_text_cache_slots=` | integer slots | `1` | Emotion-text weight cache slots; set `0` to disable reuse. | +| `--session-option index_tts2_5.gpt_graph_arena_mb=` | MB | model default | GPT graph arena size. | +| `--session-option index_tts2_5.s2mel_graph_arena_mb=` | MB | model default | S2Mel graph arena size. | +| `--session-option index_tts2_5.reference_graph_arena_mb=` | MB | model default | Reference encoder and codec graph arena size. | +| `--session-option index_tts2_5.emotion_text_prefill_graph_arena_mb=` | MB | model default | Emotion-text prefill graph arena size. | +| `--session-option index_tts2_5.emotion_text_decode_graph_arena_mb=` | MB | model default | Emotion-text cached-step graph arena size. | +| `--session-option index_tts2_5.emotion_text_max_new_tokens=` | tokens | `256` | Maximum generated tokens for emotion-text classification. | +| `--session-option index_tts2_5.weight_context_mb=` | MB | `32` | Shared ggml weight metadata context size. | + ## Irodori-TTS Irodori-TTS is Japanese TTS under `--family irodori_tts`. v4 Small is the preferred GGUF-first package and supports no-reference speech, reference-conditioned speech, and caption-based voice design in one checkpoint. The older 500M v3 and 600M v3 VoiceDesign packages remain supported for existing users. See [Irodori-TTS](models/irodori_tts.md) for v3/v4 differences, GGUF variants, options, and compatibility aliases. diff --git a/include/engine/models/index_tts2_5/assets.h b/include/engine/models/index_tts2_5/assets.h new file mode 100644 index 00000000..6b0fc014 --- /dev/null +++ b/include/engine/models/index_tts2_5/assets.h @@ -0,0 +1,29 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/models/index_tts2_5/types.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25Assets { + assets::ResourceBundle resources; + IndexTTS25Config config; + std::shared_ptr gpt_weights; + std::shared_ptr s2mel_weights; + std::shared_ptr speaker_matrix; + std::shared_ptr emotion_matrix; + std::shared_ptr wav2vec2bert_stats; + std::shared_ptr wav2vec2bert_weights; + std::shared_ptr semantic_codec_weights; + std::shared_ptr campplus_weights; + std::shared_ptr bigvgan_weights; + std::shared_ptr qwen_emotion_weights; +}; + +std::shared_ptr load_index_tts2_5_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/audio_features.h b/include/engine/models/index_tts2_5/audio_features.h new file mode 100644 index 00000000..3022fdb3 --- /dev/null +++ b/include/engine/models/index_tts2_5/audio_features.h @@ -0,0 +1,54 @@ +#pragma once + +#include "engine/models/index_tts2_5/types.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25MelOutput { + std::vector values; + int64_t channels = 0; + int64_t frames = 0; +}; + +struct IndexTTS25FbankOutput { + std::vector values; + int64_t frames = 0; + int64_t dims = 0; +}; + +struct IndexTTS25SemanticFeatureOutput { + std::vector values; + std::vector attention_mask; + int64_t frames = 0; + int64_t dims = 0; +}; + +struct IndexTTS25PreparedReferenceAudio { + std::vector waveform_16k; + std::vector waveform_22k; + IndexTTS25MelOutput mel; + IndexTTS25FbankOutput campplus_fbank; + IndexTTS25SemanticFeatureOutput semantic_features; +}; + +IndexTTS25PreparedReferenceAudio prepare_index_tts2_5_reference_audio( + const std::vector & samples, + int sample_rate, + int channels, + const IndexTTS25S2MelConfig & mel_config, + size_t threads, + bool speaker_load_semantic = true); + +IndexTTS25MelOutput compute_index_tts2_5_mel_spectrogram( + const std::vector & waveform, + const IndexTTS25S2MelConfig & config, + size_t threads); + +IndexTTS25FbankOutput compute_index_tts2_5_campplus_fbank_16k(const std::vector & waveform_16k); + +IndexTTS25SemanticFeatureOutput compute_index_tts2_5_semantic_features_16k(const std::vector & waveform_16k); + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/gpt.h b/include/engine/models/index_tts2_5/gpt.h new file mode 100644 index 00000000..a8f9b93e --- /dev/null +++ b/include/engine/models/index_tts2_5/gpt.h @@ -0,0 +1,186 @@ +#pragma once + +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/attention/types.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/models/index_tts2_5/assets.h" + +#include "ggml-backend.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25GptConditionSubsamplingWeights { + engine::modules::Conv2dWeights conv; + engine::modules::LinearWeights out; + engine::core::TensorValue pos_enc; +}; + +struct IndexTTS25GptConditionLayerWeights { + engine::modules::NormWeights norm_ff; + engine::modules::NormWeights norm_mha; + engine::modules::NormWeights norm_conv; + engine::modules::NormWeights norm_final; + engine::modules::LinearWeights feed_forward_in; + engine::modules::LinearWeights feed_forward_out; + engine::modules::RelativeAttentionWeights self_attn; + engine::modules::Conv1dWeights conv_pointwise_in; + engine::modules::DepthwiseConv1dWeights conv_depthwise; + engine::modules::NormWeights conv_norm; + engine::modules::Conv1dWeights conv_pointwise_out; +}; + +struct IndexTTS25GptConditionEncoderWeights { + IndexTTS25GptConditionSubsamplingWeights subsampling; + std::vector layers; + engine::modules::NormWeights after_norm; +}; + +struct IndexTTS25PerceiverAttentionWeights { + engine::modules::LinearWeights q; + engine::modules::LinearWeights kv; + engine::modules::LinearWeights out; +}; + +struct IndexTTS25PerceiverFeedForwardWeights { + engine::modules::LinearWeights in; + engine::modules::LinearWeights out; +}; + +struct IndexTTS25PerceiverLayerWeights { + IndexTTS25PerceiverAttentionWeights attention; + IndexTTS25PerceiverFeedForwardWeights feed_forward; +}; + +struct IndexTTS25PerceiverWeights { + engine::core::TensorValue latents; + engine::modules::LinearWeights project_context; + std::vector layers; + engine::core::TensorValue norm_gamma; +}; + +struct IndexTTS25Gpt2LayerWeights { + engine::modules::NormWeights attn_norm; + engine::modules::LinearWeights qkv; + engine::modules::LinearWeights attn_out; + engine::modules::NormWeights mlp_norm; + engine::modules::LinearWeights mlp_in; + engine::modules::LinearWeights mlp_out; +}; + +struct IndexTTS25GptWeights { + std::shared_ptr store; + IndexTTS25GptConditionEncoderWeights emotion_conditioner; + IndexTTS25PerceiverWeights emotion_perceiver; + engine::modules::LinearWeights spk_emb_proj; + engine::core::TensorValue lang_embedding; + engine::core::TensorValue text_embedding; + engine::core::TensorValue mel_embedding; + engine::core::TensorValue text_pos_embedding; + engine::core::TensorValue mel_pos_embedding; + engine::modules::LinearWeights emotion_vec_projection; + engine::modules::LinearWeights emotion_layer; + std::vector gpt_layers; + engine::modules::NormWeights gpt_final_norm; + engine::modules::NormWeights final_norm; + engine::modules::LinearWeights mel_head; + engine::modules::LinearWeights text_head; +}; + +struct IndexTTS25GptLatent { + std::vector values; + int64_t frames = 0; + int64_t dims = 0; +}; + +struct IndexTTS25GptGeneration { + std::vector codes; + uint64_t rng_offset_blocks = 0; +}; + +struct IndexTTS25GptGenerationRequest { + std::vector text_tokens; + // 192-dim CAMPPlus speaker embedding, projected by spk_emb_proj inside the + // prefill graph (spk_cond_mode="campplus" in the official model_v2.py). + std::vector speaker_style; + // Row of the GPT lang_embedding table added to every text embedding. + int32_t lang_id = 0; + std::vector emotion_semantic; + int64_t emotion_frames = 0; + std::vector emotion_vector; + float top_p = 0.8F; + int top_k = 30; + float temperature = 0.8F; + float repetition_penalty = 10.0F; + bool do_sample = true; + float length_penalty = 0.0F; + int num_beams = 3; + int max_mel_tokens = 1500; + uint32_t seed = 0; +}; + +// Mirrors the valid_mask filtering in the official prepare_gpt_inputs: any +// start/stop text tokens in the segment (including the trailing pad appended by +// the tokenizer) are dropped before the start/stop pair is re-added around it. +std::vector align_index_tts2_5_gpt_text_tokens(const std::vector & text_tokens); + +std::shared_ptr load_index_tts2_5_gpt_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes); + +class IndexTTS25GptRuntime { +public: + IndexTTS25GptRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type); + ~IndexTTS25GptRuntime(); + + IndexTTS25GptRuntime(const IndexTTS25GptRuntime &) = delete; + IndexTTS25GptRuntime & operator=(const IndexTTS25GptRuntime &) = delete; + + void prepare_emotion_conditioning(int64_t frames); + void prepare_generation(int64_t text_tokens, int64_t max_mel_tokens, int64_t num_beams); + IndexTTS25GptLatent emotion_conditioning(const std::vector & semantic_btc, int64_t frames); + std::vector project_emotion_vector(const IndexTTS25GptLatent & emotion_conditioning); + std::vector merge_emotion_vector( + const std::vector & speaker_semantic, + int64_t speaker_frames, + const std::vector & emotion_semantic, + int64_t emotion_frames, + float alpha); + IndexTTS25GptGeneration generate_speech(const IndexTTS25GptGenerationRequest & request); + void release_conditioning_graphs(); + void release_generation_graphs(); + +private: + class ConditioningGraph; + class EmotionVectorGraph; + class PrefillGraph; + class DecodeGraph; + + std::shared_ptr assets_; + engine::core::ExecutionContext * execution_ = nullptr; + size_t graph_arena_bytes_ = 0; + std::shared_ptr weights_; + std::unique_ptr emotion_conditioning_graph_; + std::unique_ptr emotion_vector_graph_; + std::unique_ptr prefill_graph_; + std::unique_ptr decode_graph_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/loader.h b/include/engine/models/index_tts2_5/loader.h new file mode 100644 index 00000000..e063f716 --- /dev/null +++ b/include/engine/models/index_tts2_5/loader.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/models/index_tts2_5/assets.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +class IndexTTS25LoadedModel final : public runtime::ILoadedVoiceModel { +public: + IndexTTS25LoadedModel( + 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_index_tts2_5_model(const std::filesystem::path & model_path); +std::shared_ptr make_index_tts2_5_loader(); + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/qwen_emotion.h b/include/engine/models/index_tts2_5/qwen_emotion.h new file mode 100644 index 00000000..211a6cbb --- /dev/null +++ b/include/engine/models/index_tts2_5/qwen_emotion.h @@ -0,0 +1,82 @@ +#pragma once + +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/transformers/qwen_decoder.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/tokenizers/llama_bpe.h" +#include "engine/models/index_tts2_5/assets.h" + +#include "ggml-backend.h" + +#include +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25QwenEmotionWeights { + std::shared_ptr store; + engine::core::TensorValue token_embedding; + engine::modules::QwenDecoderStackWeights decoder; + engine::modules::NormWeights final_norm; +}; + +struct IndexTTS25EmotionVector { + std::vector values; +}; + +std::shared_ptr load_index_tts2_5_qwen_emotion_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType storage_type, + size_t weight_context_bytes); + +class IndexTTS25QwenEmotionTokenizer { +public: + explicit IndexTTS25QwenEmotionTokenizer(std::shared_ptr assets); + + std::vector encode_chat_prompt(const std::string & text) const; + std::string decode(const std::vector & token_ids, bool skip_special_tokens) const; + int32_t eos_token_id() const noexcept; + int32_t think_end_token_id() const noexcept; + +private: + std::shared_ptr tokenizer_; + int32_t eos_token_id_ = 151643; + int32_t think_end_token_id_ = 151668; +}; + +class IndexTTS25QwenEmotionRuntime { +public: + IndexTTS25QwenEmotionRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type); + ~IndexTTS25QwenEmotionRuntime(); + + IndexTTS25QwenEmotionRuntime(const IndexTTS25QwenEmotionRuntime &) = delete; + IndexTTS25QwenEmotionRuntime & operator=(const IndexTTS25QwenEmotionRuntime &) = delete; + + IndexTTS25EmotionVector infer(const std::string & text, int64_t max_new_tokens = 256); + void release_graphs(); + +private: + class PrefillGraph; + class DecodeGraph; + + std::shared_ptr assets_; + engine::core::ExecutionContext * execution_ = nullptr; + size_t prefill_graph_arena_bytes_ = 0; + size_t decode_graph_arena_bytes_ = 0; + std::shared_ptr weights_; + IndexTTS25QwenEmotionTokenizer tokenizer_; + std::unique_ptr prefill_graph_; + std::unique_ptr decode_graph_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/request.h b/include/engine/models/index_tts2_5/request.h new file mode 100644 index 00000000..474e30b0 --- /dev/null +++ b/include/engine/models/index_tts2_5/request.h @@ -0,0 +1,14 @@ +#pragma once + +#include "engine/framework/runtime/session.h" +#include "engine/models/index_tts2_5/types.h" + +namespace engine::models::index_tts2_5 { + +// Normalizes the "lang" request option: trims, lowercases, and maps "auto" to +// an empty string (tokenizer-side language inference). +std::string normalize_index_tts2_5_lang(const std::string & value); + +IndexTTS25Request parse_index_tts2_5_request(const runtime::TaskRequest & request); + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/s2mel.h b/include/engine/models/index_tts2_5/s2mel.h new file mode 100644 index 00000000..6f015502 --- /dev/null +++ b/include/engine/models/index_tts2_5/s2mel.h @@ -0,0 +1,151 @@ +#pragma once + +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/models/index_tts2_5/assets.h" + +#include "ggml-backend.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25LengthRegulatorWeights { + engine::modules::LinearWeights content_projection; + std::vector convs; + std::vector norms; + engine::modules::Conv1dWeights output; +}; + +struct IndexTTS25S2MelGptLayerWeights { + engine::modules::LinearWeights linear0; + engine::modules::LinearWeights linear1; + engine::modules::LinearWeights linear2; +}; + +struct IndexTTS25AdaLayerNormWeights { + engine::core::TensorValue norm_weight; + engine::modules::LinearWeights project; +}; + +struct IndexTTS25DitLayerWeights { + IndexTTS25AdaLayerNormWeights attention_norm; + engine::modules::LinearWeights qkv; + engine::modules::LinearWeights attention_out; + IndexTTS25AdaLayerNormWeights ffn_norm; + engine::modules::LinearWeights ffn_w1; + engine::modules::LinearWeights ffn_w2; + engine::modules::LinearWeights ffn_w3; + engine::modules::LinearWeights skip_in; +}; + +struct IndexTTS25WaveNetLayerWeights { + engine::modules::Conv1dWeights in_layer; + engine::modules::Conv1dWeights res_skip_layer; +}; + +struct IndexTTS25S2MelCfmWeights { + engine::modules::LinearWeights x_embedder; + engine::modules::LinearWeights cond_projection; + engine::modules::LinearWeights cond_x_merge; + engine::modules::LinearWeights skip_linear; + engine::modules::LinearWeights time_mlp0; + engine::modules::LinearWeights time_mlp2; + engine::core::TensorValue time_freqs; + engine::modules::LinearWeights time2_mlp0; + engine::modules::LinearWeights time2_mlp2; + engine::core::TensorValue time2_freqs; + std::vector dit_layers; + IndexTTS25AdaLayerNormWeights dit_norm; + engine::modules::LinearWeights conv1; + engine::modules::LinearWeights res_projection; + engine::modules::Conv1dWeights wavenet_cond; + std::vector wavenet_layers; + engine::modules::LinearWeights final_modulation; + engine::modules::LinearWeights final_linear; + engine::modules::Conv1dWeights conv2; +}; + +struct IndexTTS25S2MelWeights { + std::shared_ptr store; + IndexTTS25S2MelGptLayerWeights gpt_layer; + IndexTTS25LengthRegulatorWeights length_regulator; + IndexTTS25S2MelCfmWeights cfm; +}; + +struct IndexTTS25S2MelSequence { + std::vector values; + int64_t frames = 0; + int64_t dims = 0; +}; + +struct IndexTTS25S2MelMel { + std::vector values; + int64_t frames = 0; + int64_t channels = 80; +}; + +std::shared_ptr load_index_tts2_5_s2mel_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes); + +class IndexTTS25S2MelRuntime { +public: + IndexTTS25S2MelRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type); + ~IndexTTS25S2MelRuntime(); + + IndexTTS25S2MelRuntime(const IndexTTS25S2MelRuntime &) = delete; + IndexTTS25S2MelRuntime & operator=(const IndexTTS25S2MelRuntime &) = delete; + + void prepare_gpt_layer(int64_t frames); + void prepare_length_regulator(int64_t input_frames, int64_t output_frames); + void prepare_cfm(int64_t total_frames, bool use_cfg); + void release_pre_cfm_graphs(); + void release_cfm_graph(); + + IndexTTS25S2MelSequence project_gpt_latent(const std::vector & latent, int64_t frames); + IndexTTS25S2MelSequence regulate_length( + const std::vector & content, + int64_t input_frames, + int64_t output_frames); + IndexTTS25S2MelMel infer_mel( + const std::vector & condition, + int64_t total_frames, + const std::vector & reference_mel, + int64_t reference_frames, + const std::vector & style, + int64_t diffusion_steps, + float cfg_rate, + uint32_t seed, + uint64_t rng_offset_blocks); + +private: + class GptLayerGraph; + class LengthRegulatorGraph; + class CfmGraph; + + std::shared_ptr assets_; + engine::core::ExecutionContext * execution_ = nullptr; + size_t graph_arena_bytes_ = 0; + std::shared_ptr weights_; + std::unique_ptr gpt_layer_graph_; + std::unique_ptr length_regulator_graph_; + std::unique_ptr cfm_graph_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/semantic_codec.h b/include/engine/models/index_tts2_5/semantic_codec.h new file mode 100644 index 00000000..2762d7f7 --- /dev/null +++ b/include/engine/models/index_tts2_5/semantic_codec.h @@ -0,0 +1,94 @@ +#pragma once + +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/models/index_tts2_5/assets.h" +#include "engine/models/index_tts2_5/semantic_encoder.h" + +#include "ggml-backend.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25VocosConvNeXtBlockWeights { + engine::modules::DepthwiseConv1dWeights depthwise; + engine::modules::NormWeights norm; + engine::modules::LinearWeights pointwise_in; + engine::modules::LinearWeights pointwise_out; + engine::core::TensorValue gamma; +}; + +struct IndexTTS25VocosBackboneWeights { + engine::modules::Conv1dWeights embed; + engine::modules::NormWeights norm; + std::vector blocks; + engine::modules::NormWeights final_norm; +}; + +struct IndexTTS25SemanticCodecWeights { + std::shared_ptr store; + IndexTTS25VocosBackboneWeights encoder_backbone; + engine::modules::LinearWeights encoder_projection; + engine::modules::Conv1dWeights quantizer_in; + engine::core::TensorValue codebook; + engine::core::TensorValue normalized_codebook; + engine::modules::Conv1dWeights quantizer_out; + IndexTTS25VocosBackboneWeights decoder_backbone; + engine::modules::LinearWeights decoder_projection; + engine::modules::Conv1dWeights up; +}; + +struct IndexTTS25SemanticCodecOutput { + std::vector codes; + std::vector embedding_channel_first; + int64_t frames = 0; + int64_t dims = 0; +}; + +std::shared_ptr load_index_tts2_5_semantic_codec_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes); + +class IndexTTS25SemanticCodecRuntime { +public: + IndexTTS25SemanticCodecRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type); + ~IndexTTS25SemanticCodecRuntime(); + + IndexTTS25SemanticCodecRuntime(const IndexTTS25SemanticCodecRuntime &) = delete; + IndexTTS25SemanticCodecRuntime & operator=(const IndexTTS25SemanticCodecRuntime &) = delete; + + void prepare_quantize(int64_t frames); + void prepare_codes(int64_t frames); + IndexTTS25SemanticCodecOutput quantize(const IndexTTS25SemanticEmbedding & semantic); + IndexTTS25SemanticCodecOutput codes_to_embedding(const std::vector & codes, int64_t frames); + void release_graphs(); + +private: + class QuantizeGraph; + class CodesGraph; + + std::shared_ptr assets_; + engine::core::ExecutionContext * execution_ = nullptr; + size_t graph_arena_bytes_ = 0; + std::shared_ptr weights_; + std::unique_ptr quantize_graph_; + std::unique_ptr codes_graph_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/semantic_encoder.h b/include/engine/models/index_tts2_5/semantic_encoder.h new file mode 100644 index 00000000..5ffb21d9 --- /dev/null +++ b/include/engine/models/index_tts2_5/semantic_encoder.h @@ -0,0 +1,100 @@ +#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/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/models/index_tts2_5/assets.h" +#include "engine/models/index_tts2_5/audio_features.h" + +#include "ggml-backend.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25Wav2Vec2BertAttentionWeights { + engine::modules::LinearWeights q; + engine::modules::LinearWeights k; + engine::modules::LinearWeights v; + engine::modules::LinearWeights out; + engine::core::TensorValue distance_embedding; +}; + +struct IndexTTS25Wav2Vec2BertConvWeights { + engine::modules::NormWeights layer_norm; + engine::modules::Conv1dWeights pointwise_in; + engine::modules::DepthwiseConv1dWeights depthwise; + engine::modules::NormWeights depthwise_layer_norm; + engine::modules::Conv1dWeights pointwise_out; +}; + +struct IndexTTS25Wav2Vec2BertLayerWeights { + engine::modules::NormWeights ffn1_norm; + engine::modules::LinearWeights ffn1_in; + engine::modules::LinearWeights ffn1_out; + engine::modules::NormWeights self_attn_norm; + IndexTTS25Wav2Vec2BertAttentionWeights self_attn; + IndexTTS25Wav2Vec2BertConvWeights conv; + engine::modules::NormWeights ffn2_norm; + engine::modules::LinearWeights ffn2_in; + engine::modules::LinearWeights ffn2_out; + engine::modules::NormWeights final_norm; +}; + +struct IndexTTS25Wav2Vec2BertWeights { + std::shared_ptr store; + engine::modules::NormWeights feature_norm; + engine::modules::LinearWeights feature_projection; + std::vector layers; + engine::core::TensorValue semantic_mean; + engine::core::TensorValue semantic_std; +}; + +struct IndexTTS25SemanticEmbedding { + std::vector values; + int64_t frames = 0; + int64_t dims = 0; +}; + +std::shared_ptr load_index_tts2_5_wav2vec2bert_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes); + +class IndexTTS25Wav2Vec2BertRuntime { +public: + IndexTTS25Wav2Vec2BertRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type); + ~IndexTTS25Wav2Vec2BertRuntime(); + + IndexTTS25Wav2Vec2BertRuntime(const IndexTTS25Wav2Vec2BertRuntime &) = delete; + IndexTTS25Wav2Vec2BertRuntime & operator=(const IndexTTS25Wav2Vec2BertRuntime &) = delete; + + void prepare(int64_t frames); + IndexTTS25SemanticEmbedding encode(const IndexTTS25SemanticFeatureOutput & features); + void release_graph(); + +private: + class Graph; + + std::shared_ptr assets_; + engine::core::ExecutionContext * execution_ = nullptr; + size_t graph_arena_bytes_ = 0; + std::shared_ptr weights_; + std::unique_ptr graph_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/session.h b/include/engine/models/index_tts2_5/session.h new file mode 100644 index 00000000..65ab86ca --- /dev/null +++ b/include/engine/models/index_tts2_5/session.h @@ -0,0 +1,122 @@ +#pragma once + +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/index_tts2_5/assets.h" +#include "engine/models/index_tts2_5/audio_features.h" +#include "engine/models/index_tts2_5/gpt.h" +#include "engine/models/index_tts2_5/qwen_emotion.h" +#include "engine/models/index_tts2_5/request.h" +#include "engine/models/index_tts2_5/s2mel.h" +#include "engine/models/index_tts2_5/semantic_codec.h" +#include "engine/models/index_tts2_5/semantic_encoder.h" +#include "engine/models/index_tts2_5/style_encoder.h" +#include "engine/models/index_tts2_5/tokenizer_text.h" +#include "engine/models/index_tts2_5/vocoder.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25AudioIdentity { + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; +}; + +class IndexTTS25Session final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + IndexTTS25Session( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + + 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: + struct SpeakerState { + IndexTTS25AudioIdentity identity; + IndexTTS25SemanticEmbedding semantic; + IndexTTS25MelOutput reference_mel; + IndexTTS25StyleEmbedding style; + IndexTTS25S2MelSequence prompt_condition; + }; + + struct EmotionState { + IndexTTS25AudioIdentity identity; + IndexTTS25SemanticEmbedding semantic; + }; + + struct AudioIdentityEqual { + bool operator()( + const IndexTTS25AudioIdentity & lhs, + const IndexTTS25AudioIdentity & rhs) const; + }; + + const SpeakerState & resolve_speaker_state(const runtime::AudioBuffer & audio); + const EmotionState & resolve_emotion_state(const runtime::AudioBuffer & audio); + std::vector resolve_emotion_vector( + const IndexTTS25Request & request, + const SpeakerState & speaker, + const EmotionState & emotion); + runtime::AudioBuffer synthesize_segment( + const std::vector & text_tokens, + int32_t lang_id, + size_t segment_index, + const std::string & dump_dir, + const SpeakerState & speaker, + const EmotionState & emotion, + const std::vector & emotion_vector, + const IndexTTS25GenerationOptions & options, + uint32_t segment_seed); + + std::vector explicit_emotion_matrix_vector( + const std::vector & emotion_weights, + const IndexTTS25StyleEmbedding & style, + bool use_random, + uint32_t seed) const; + + runtime::TaskSpec task_; + std::shared_ptr assets_; + size_t gpt_graph_arena_bytes_ = 2048ull * 1024ull * 1024ull; + size_t s2mel_graph_arena_bytes_ = 2048ull * 1024ull * 1024ull; + size_t reference_graph_arena_bytes_ = 512ull * 1024ull * 1024ull; + size_t emotion_text_prefill_graph_arena_bytes_ = 2048ull * 1024ull * 1024ull; + size_t emotion_text_decode_graph_arena_bytes_ = 512ull * 1024ull * 1024ull; + size_t weight_context_bytes_ = 32ull * 1024ull * 1024ull; + int64_t emotion_text_max_new_tokens_ = 256; + engine::assets::TensorStorageType matmul_weight_storage_type_ = engine::assets::TensorStorageType::Native; + engine::assets::TensorStorageType conv_weight_storage_type_ = engine::assets::TensorStorageType::Native; + bool mem_saver_ = false; + + IndexTTS25TextTokenizer tokenizer_; + std::unique_ptr semantic_encoder_; + std::unique_ptr semantic_codec_; + std::unique_ptr style_encoder_; + std::unique_ptr gpt_; + std::unique_ptr s2mel_; + std::unique_ptr vocoder_; + std::unique_ptr qwen_emotion_; + + std::vector speaker_matrix_; + std::vector emotion_matrix_; + runtime::CacheSlots speaker_cache_; + runtime::CacheSlots emotion_cache_; + runtime::CacheSlots> emotion_text_weights_cache_; + std::optional uncached_speaker_state_; + std::optional uncached_emotion_state_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/style_encoder.h b/include/engine/models/index_tts2_5/style_encoder.h new file mode 100644 index 00000000..9ed6bc0f --- /dev/null +++ b/include/engine/models/index_tts2_5/style_encoder.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/modules/speech_encoders/campplus_encoder.h" +#include "engine/models/index_tts2_5/assets.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25StyleEmbedding { + std::vector values; + int64_t dims = 0; +}; + +class IndexTTS25StyleEncoder { +public: + IndexTTS25StyleEncoder( + std::shared_ptr assets, + core::BackendConfig backend, + engine::assets::TensorStorageType weight_storage_type); + + IndexTTS25StyleEmbedding embed_fbank( + const std::vector & features, + int64_t frames, + int64_t dims) const; + void release_graph(); + +private: + std::shared_ptr assets_; + engine::modules::CampplusEncoderComponent component_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/tokenizer_text.h b/include/engine/models/index_tts2_5/tokenizer_text.h new file mode 100644 index 00000000..65258fd2 --- /dev/null +++ b/include/engine/models/index_tts2_5/tokenizer_text.h @@ -0,0 +1,58 @@ +#pragma once + +#include "engine/models/index_tts2_5/assets.h" + +#include +#include +#include +#include + +namespace llama_tokenizer_vendor { +struct BpeVocabulary; +} // namespace llama_tokenizer_vendor + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25TextEncoding { + std::string lang; + std::string normalized_text; + std::vector segments; + std::vector> segment_token_ids; +}; + +// Whisper-style tiktoken BPE text tokenizer for IndexTTS-2.5 (vocab size 60509). +// Replaces the SentencePiece tokenizer used by index_tts2. +class IndexTTS25TextTokenizer { +public: + explicit IndexTTS25TextTokenizer(std::shared_ptr assets); + + std::string normalize_english(const std::string & text) const; + std::string normalize_chinese(const std::string & text) const; + + // Raw tiktoken encode with allowed_special="all"; does not apply any text + // normalization. Special tokens present in the text are recognized directly. + std::vector encode(const std::string & text) const; + + // Returns the id of an exact token text (e.g. "<|zh|>"), or -1 when unknown. + int32_t special_token_id(const std::string & token_text) const; + + // Maps a language code to the GPT lang_embedding row, following the + // LANGUAGES order of indextts/utils/tokenizer.py (en=0, zh=1, ...). + // Unknown codes map to "common". + static int32_t lang_to_id(const std::string & lang); + + // Full inference pipeline: normalize -> case rules -> pronunciation + // annotations -> special-token name uppercasing -> segment by token budget. + // Each segment is encoded as encode("<|{lang}|> " + segment) plus a trailing + // pad token id 1. When lang is empty, it is inferred (Han -> zh, else en). + IndexTTS25TextEncoding encode_for_inference( + const std::string & text, + int max_text_tokens_per_segment, + const std::string & lang = "") const; + +private: + std::shared_ptr assets_; + std::shared_ptr vocab_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/types.h b/include/engine/models/index_tts2_5/types.h new file mode 100644 index 00000000..07f105b9 --- /dev/null +++ b/include/engine/models/index_tts2_5/types.h @@ -0,0 +1,153 @@ +#pragma once + +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { + +// Rows in the GPT lang_embedding table. indextts/utils/tokenizer.py defines +// 106 language codes (including "common"); the checkpoint table has one extra +// unused row. +constexpr int64_t kIndexTTS25LangEmbeddingRows = 107; + +struct IndexTTS25GptConfig { + int64_t model_dim = 1280; + int64_t max_mel_tokens = 1815; + int64_t max_text_tokens = 600; + int64_t heads = 20; + bool use_mel_codes_as_input = true; + int64_t mel_length_compression = 1024; + int64_t layers = 24; + int64_t number_text_tokens = 12000; + int64_t number_mel_codes = 8194; + int64_t start_mel_token = 8192; + int64_t stop_mel_token = 8193; + int64_t start_text_token = 0; + int64_t stop_text_token = 1; + bool train_solo_embeddings = false; + std::string condition_type = "conformer_perceiver"; + int64_t condition_output_size = 512; + int64_t condition_linear_units = 2048; + int64_t condition_attention_heads = 8; + int64_t condition_num_blocks = 6; + std::string condition_input_layer = "conv2d2"; + int64_t condition_perceiver_mult = 2; + int64_t emo_condition_output_size = 512; + int64_t emo_condition_linear_units = 1024; + int64_t emo_condition_attention_heads = 4; + int64_t emo_condition_num_blocks = 4; + std::string emo_condition_input_layer = "conv2d2"; + int64_t emo_condition_perceiver_mult = 2; +}; + +struct IndexTTS25SemanticCodecConfig { + int64_t codebook_size = 8192; + int64_t hidden_size = 1024; + int64_t codebook_dim = 8; + int64_t vocos_dim = 384; + int64_t vocos_intermediate_dim = 2048; + int64_t vocos_num_layers = 12; +}; + +struct IndexTTS25S2MelConfig { + int sample_rate = 22050; + int64_t n_fft = 1024; + int64_t win_length = 1024; + int64_t hop_length = 256; + int64_t n_mels = 80; + float fmin = 0.0F; + std::optional fmax = std::nullopt; + std::string dit_type = "DiT"; + std::string reg_loss_type = "l1"; + int64_t style_dim = 192; + int64_t length_regulator_channels = 512; + bool length_regulator_is_discrete = false; + int64_t length_regulator_in_channels = 1024; + int64_t length_regulator_content_codebook_size = 2048; + std::vector length_regulator_sampling_ratios; + bool length_regulator_vector_quantize = false; + int64_t length_regulator_n_codebooks = 1; + float length_regulator_quantizer_dropout = 0.0F; + bool length_regulator_f0_condition = false; + int64_t length_regulator_n_f0_bins = 512; + int64_t dit_hidden_dim = 512; + int64_t dit_num_heads = 8; + int64_t dit_depth = 13; + float dit_class_dropout_prob = 0.1F; + int64_t dit_block_size = 8192; + int64_t dit_in_channels = 80; + bool dit_style_condition = true; + std::string dit_final_layer_type = "wavenet"; + std::string dit_target = "mel"; + int64_t dit_content_dim = 512; + int64_t dit_content_codebook_size = 1024; + std::string dit_content_type = "discrete"; + bool dit_f0_condition = false; + int64_t dit_n_f0_bins = 512; + int64_t dit_content_codebooks = 1; + bool dit_is_causal = false; + bool dit_long_skip_connection = true; + bool dit_zero_prompt_speech_token = false; + bool dit_time_as_token = false; + bool dit_style_as_token = false; + bool dit_uvit_skip_connection = true; + bool dit_add_resblock_in_transformer = false; + int64_t wavenet_hidden_dim = 512; + int64_t wavenet_num_layers = 8; + int64_t wavenet_kernel_size = 5; + int64_t wavenet_dilation_rate = 1; + float wavenet_dropout = 0.2F; + bool wavenet_style_condition = true; +}; + +struct IndexTTS25Config { + std::string version = "2.0"; + int dataset_sample_rate = 24000; + bool dataset_squeeze = false; + int dataset_mel_sample_rate = 24000; + int64_t dataset_mel_n_fft = 1024; + int64_t dataset_mel_hop_length = 256; + int64_t dataset_mel_win_length = 1024; + int64_t dataset_mel_n_mels = 100; + float dataset_mel_fmin = 0.0F; + bool dataset_mel_normalize = false; + IndexTTS25GptConfig gpt; + IndexTTS25SemanticCodecConfig semantic_codec; + IndexTTS25S2MelConfig s2mel; + std::vector emo_num; +}; + +struct IndexTTS25GenerationOptions { + bool do_sample = true; + float top_p = 0.8F; + int top_k = 30; + float temperature = 0.8F; + float length_penalty = 0.0F; + int num_beams = 3; + float repetition_penalty = 10.0F; + int max_mel_tokens = 1500; + uint32_t seed = 0; +}; + +struct IndexTTS25Request { + std::string text; + std::optional speaker_audio = std::nullopt; + std::optional emotion_audio = std::nullopt; + // Text language hint; empty means auto (zh when the text contains Han + // characters, otherwise en). + std::string lang; + float emotion_alpha = 1.0F; + std::optional> emotion_vector = std::nullopt; + bool use_emotion_text = false; + std::optional emotion_text = std::nullopt; + bool use_random_emotion = false; + int interval_silence_ms = 200; + int max_text_tokens_per_segment = 120; + IndexTTS25GenerationOptions generation; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/vocoder.h b/include/engine/models/index_tts2_5/vocoder.h new file mode 100644 index 00000000..f2ab1060 --- /dev/null +++ b/include/engine/models/index_tts2_5/vocoder.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/modules/vocoders/bigvgan_vocoder.h" +#include "engine/models/index_tts2_5/assets.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +struct IndexTTS25VocoderOutput { + std::vector waveform; + int64_t samples = 0; + int sample_rate = 0; +}; + +class IndexTTS25BigVganVocoder { +public: + IndexTTS25BigVganVocoder( + std::shared_ptr assets, + core::BackendConfig backend, + engine::assets::TensorStorageType weight_storage_type); + + IndexTTS25VocoderOutput synthesize( + const std::vector & mel, + int64_t frames) const; + void release_runtime_graph(); + +private: + std::shared_ptr assets_; + engine::modules::BigVganVocoderComponent component_; +}; + +} // namespace engine::models::index_tts2_5 diff --git a/model_specs/index_tts2_5.json b/model_specs/index_tts2_5.json new file mode 100644 index 00000000..5b4d8661 --- /dev/null +++ b/model_specs/index_tts2_5.json @@ -0,0 +1,186 @@ +{ + "family": "index_tts2_5", + "display_name": "IndexTTS2.5", + "description": "Multilingual zero-shot TTS system for Chinese, English, Japanese, Spanish and Arabic speech synthesis with voice cloning, timbre-emotion decoupling, text or audio emotion control, and explicit duration control.", + "category": "tts", + "status": "supported", + "tasks": [ + "tts", + "clone" + ], + "modes": [ + "offline" + ], + "languages": [ + "zh", + "en", + "ja", + "es", + "ar" + ], + "capabilities": { + "tts": [ + "emotion_control" + ], + "clone": [ + "speaker_reference", + "emotion_control" + ] + }, + "runtime": { + "tags": [ + "gguf" + ] + }, + "ui": { + "recommended_package": "index_tts2_5_q8_0", + "tags": [ + "TTS", + "Clone", + "GGUF" + ], + "docs": [ + "docs/tts.md", + "docs/gguf.md" + ] + }, + "package_defaults": { + "download": { + "kind": "huggingface_snapshot", + "repo": "audio-cpp/audio.cpp-gguf", + "revision": "main", + "gated": false + } + }, + "packages": [ + { + "id": "index_tts2_5_q8_0", + "display_name": "IndexTTS2.5 Q8_0 GGUF", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "IndexTTS2.5-GGUF", + "files": [ + "IndexTTS2.5-GGUF/index-tts2_5-q8_0.gguf" + ], + "strip_prefix": "IndexTTS2.5-GGUF" + }, + { + "id": "index_tts2_5_f16", + "display_name": "IndexTTS2.5 F16 GGUF", + "format": "gguf", + "precision": "f16", + "target_directory": "IndexTTS2.5-GGUF", + "files": [ + "IndexTTS2.5-GGUF/index-tts2_5-f16.gguf" + ], + "strip_prefix": "IndexTTS2.5-GGUF" + }, + { + "id": "index_tts2_5_orig", + "display_name": "IndexTTS2.5 Original-Dtype GGUF", + "format": "gguf", + "precision": "orig", + "target_directory": "IndexTTS2.5-GGUF", + "files": [ + "IndexTTS2.5-GGUF/index-tts2_5-orig.gguf" + ], + "strip_prefix": "IndexTTS2.5-GGUF" + } + ], + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.yaml", + "tiktoken": "model:multilingual_zh_ja_yue_char_del.tiktoken", + "wav2vec2bert_config": "model:w2v-bert-2.0/config.json", + "wav2vec2bert_preprocessor_config": "model:w2v-bert-2.0/preprocessor_config.json", + "bigvgan_config": "model:bigvgan/config.json", + "qwen_emotion_config": "model:qwen0.6bemo4-merge/config.json", + "qwen_emotion_generation_config": "model:qwen0.6bemo4-merge/generation_config.json", + "qwen_emotion_tokenizer": "model:qwen0.6bemo4-merge/tokenizer.json", + "qwen_emotion_tokenizer_config": "model:qwen0.6bemo4-merge/tokenizer_config.json", + "qwen_emotion_vocab": "model:qwen0.6bemo4-merge/vocab.json", + "qwen_emotion_merges": "model:qwen0.6bemo4-merge/merges.txt" + }, + "tensors": { + "gpt": { + "source": "weights:", + "prefix": "gpt" + }, + "s2mel": { + "source": "weights:", + "prefix": "s2mel" + }, + "speaker_matrix": { + "source": "weights:", + "prefix": "speaker_matrix" + }, + "emotion_matrix": { + "source": "weights:", + "prefix": "emotion_matrix" + }, + "wav2vec2bert_stats": { + "source": "weights:", + "prefix": "wav2vec2bert_stats" + }, + "wav2vec2bert": { + "source": "weights:", + "prefix": "wav2vec2bert" + }, + "semantic_codec": { + "source": "weights:", + "prefix": "semantic_codec" + }, + "campplus": { + "source": "weights:", + "prefix": "campplus" + }, + "bigvgan": { + "source": "weights:", + "prefix": "bigvgan" + }, + "qwen_emotion": { + "source": "weights:", + "prefix": "qwen_emotion" + } + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.yaml", + "tiktoken": "model:multilingual_zh_ja_yue_char_del.tiktoken", + "wav2vec2bert_config": "model:w2v-bert-2.0/config.json", + "wav2vec2bert_preprocessor_config": "model:w2v-bert-2.0/preprocessor_config.json", + "bigvgan_config": "model:bigvgan/config.json", + "qwen_emotion_config": "model:qwen0.6bemo4-merge/config.json", + "qwen_emotion_generation_config": "model:qwen0.6bemo4-merge/generation_config.json", + "qwen_emotion_tokenizer": "model:qwen0.6bemo4-merge/tokenizer.json", + "qwen_emotion_tokenizer_config": "model:qwen0.6bemo4-merge/tokenizer_config.json", + "qwen_emotion_vocab": "model:qwen0.6bemo4-merge/vocab.json", + "qwen_emotion_merges": "model:qwen0.6bemo4-merge/merges.txt" + }, + "tensors": { + "gpt": "model:gpt.safetensors", + "s2mel": "model:s2mel.safetensors", + "speaker_matrix": "model:feat1.safetensors", + "emotion_matrix": "model:feat2.safetensors", + "wav2vec2bert_stats": "model:wav2vec2bert_stats.safetensors", + "wav2vec2bert": "model:w2v-bert-2.0/model.safetensors", + "semantic_codec": "model:semantic_codec_model.safetensors", + "campplus": "model:campplus.safetensors", + "bigvgan": "model:bigvgan/model.safetensors", + "qwen_emotion": "model:qwen0.6bemo4-merge/model.safetensors" + } + } + ] +} diff --git a/src/models/index_tts2_5/assets.cpp b/src/models/index_tts2_5/assets.cpp new file mode 100644 index 00000000..cb3d0b11 --- /dev/null +++ b/src/models/index_tts2_5/assets.cpp @@ -0,0 +1,259 @@ +#include "engine/models/index_tts2_5/assets.h" + +#include "engine/framework/model_spec/package.h" +#include "engine/framework/io/config.h" +#include "engine/framework/io/json.h" +#include "engine/framework/io/yaml.h" + +#include + +namespace engine::models::index_tts2_5 { +namespace { + +namespace json = engine::io::json; +namespace yaml = engine::io::yaml; + +IndexTTS25Config parse_config(const assets::ResourceBundle & resources) { + const auto document = resources.parse_flattened_yaml("config"); + IndexTTS25Config config; + config.version = yaml::optional_string(document, "version", config.version); + // The official IndexTTS-2.5 config_v2_5.yaml has no dataset section; these values + // are parsed for compatibility but not used at inference time. + if (const auto value = yaml::optional_int(document, "dataset.sample_rate")) { + config.dataset_sample_rate = *value; + } + config.dataset_squeeze = yaml::optional_bool(document, "dataset.squeeze", config.dataset_squeeze); + if (const auto value = yaml::optional_int(document, "dataset.mel.sample_rate")) { + config.dataset_mel_sample_rate = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.n_fft")) { + config.dataset_mel_n_fft = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.hop_length")) { + config.dataset_mel_hop_length = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.win_length")) { + config.dataset_mel_win_length = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.n_mels")) { + config.dataset_mel_n_mels = *value; + } + config.dataset_mel_fmin = yaml::optional_f32(document, "dataset.mel.mel_fmin", config.dataset_mel_fmin); + config.dataset_mel_normalize = yaml::optional_bool(document, "dataset.mel.normalize", config.dataset_mel_normalize); + + config.gpt.model_dim = yaml::require_i64(document, "gpt.model_dim"); + config.gpt.max_mel_tokens = yaml::require_i64(document, "gpt.max_mel_tokens"); + config.gpt.max_text_tokens = yaml::require_i64(document, "gpt.max_text_tokens"); + config.gpt.heads = yaml::require_i64(document, "gpt.heads"); + config.gpt.use_mel_codes_as_input = yaml::optional_bool(document, "gpt.use_mel_codes_as_input", config.gpt.use_mel_codes_as_input); + config.gpt.mel_length_compression = yaml::require_i64(document, "gpt.mel_length_compression"); + config.gpt.layers = yaml::require_i64(document, "gpt.layers"); + config.gpt.number_text_tokens = yaml::require_i64(document, "gpt.number_text_tokens"); + config.gpt.number_mel_codes = yaml::require_i64(document, "gpt.number_mel_codes"); + config.gpt.start_mel_token = yaml::require_i64(document, "gpt.start_mel_token"); + config.gpt.stop_mel_token = yaml::require_i64(document, "gpt.stop_mel_token"); + config.gpt.start_text_token = yaml::require_i64(document, "gpt.start_text_token"); + config.gpt.stop_text_token = yaml::require_i64(document, "gpt.stop_text_token"); + config.gpt.train_solo_embeddings = yaml::optional_bool(document, "gpt.train_solo_embeddings", config.gpt.train_solo_embeddings); + config.gpt.condition_type = yaml::require_string(document, "gpt.condition_type"); + config.gpt.condition_output_size = yaml::require_i64(document, "gpt.condition_module.output_size"); + config.gpt.condition_linear_units = yaml::require_i64(document, "gpt.condition_module.linear_units"); + config.gpt.condition_attention_heads = yaml::require_i64(document, "gpt.condition_module.attention_heads"); + config.gpt.condition_num_blocks = yaml::require_i64(document, "gpt.condition_module.num_blocks"); + config.gpt.condition_input_layer = yaml::require_string(document, "gpt.condition_module.input_layer"); + config.gpt.condition_perceiver_mult = yaml::require_i64(document, "gpt.condition_module.perceiver_mult"); + config.gpt.emo_condition_output_size = yaml::require_i64(document, "gpt.emo_condition_module.output_size"); + config.gpt.emo_condition_linear_units = yaml::require_i64(document, "gpt.emo_condition_module.linear_units"); + config.gpt.emo_condition_attention_heads = yaml::require_i64(document, "gpt.emo_condition_module.attention_heads"); + config.gpt.emo_condition_num_blocks = yaml::require_i64(document, "gpt.emo_condition_module.num_blocks"); + config.gpt.emo_condition_input_layer = yaml::require_string(document, "gpt.emo_condition_module.input_layer"); + config.gpt.emo_condition_perceiver_mult = yaml::require_i64(document, "gpt.emo_condition_module.perceiver_mult"); + + config.semantic_codec.codebook_size = yaml::require_i64(document, "semantic_codec.codebook_size"); + config.semantic_codec.hidden_size = yaml::require_i64(document, "semantic_codec.hidden_size"); + config.semantic_codec.codebook_dim = yaml::require_i64(document, "semantic_codec.codebook_dim"); + config.semantic_codec.vocos_dim = yaml::require_i64(document, "semantic_codec.vocos_dim"); + config.semantic_codec.vocos_intermediate_dim = yaml::require_i64(document, "semantic_codec.vocos_intermediate_dim"); + config.semantic_codec.vocos_num_layers = yaml::require_i64(document, "semantic_codec.vocos_num_layers"); + + config.s2mel.sample_rate = static_cast(yaml::require_i64(document, "s2mel.preprocess_params.sr")); + config.s2mel.n_fft = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.n_fft"); + config.s2mel.win_length = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.win_length"); + config.s2mel.hop_length = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.hop_length"); + config.s2mel.n_mels = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.n_mels"); + config.s2mel.fmin = yaml::optional_f32(document, "s2mel.preprocess_params.spect_params.fmin", config.s2mel.fmin); + config.s2mel.fmax = yaml::optional_nullable_f32(document, "s2mel.preprocess_params.spect_params.fmax"); + config.s2mel.dit_type = yaml::require_string(document, "s2mel.dit_type"); + config.s2mel.reg_loss_type = yaml::require_string(document, "s2mel.reg_loss_type"); + config.s2mel.style_dim = yaml::require_i64(document, "s2mel.style_encoder.dim"); + config.s2mel.length_regulator_channels = yaml::require_i64(document, "s2mel.length_regulator.channels"); + config.s2mel.length_regulator_is_discrete = yaml::optional_bool(document, "s2mel.length_regulator.is_discrete", config.s2mel.length_regulator_is_discrete); + config.s2mel.length_regulator_in_channels = yaml::require_i64(document, "s2mel.length_regulator.in_channels"); + config.s2mel.length_regulator_content_codebook_size = yaml::require_i64(document, "s2mel.length_regulator.content_codebook_size"); + config.s2mel.length_regulator_sampling_ratios = yaml::require_list_i64(document, "s2mel.length_regulator.sampling_ratios"); + config.s2mel.length_regulator_vector_quantize = yaml::optional_bool(document, "s2mel.length_regulator.vector_quantize", config.s2mel.length_regulator_vector_quantize); + config.s2mel.length_regulator_n_codebooks = yaml::require_i64(document, "s2mel.length_regulator.n_codebooks"); + config.s2mel.length_regulator_quantizer_dropout = yaml::optional_f32(document, "s2mel.length_regulator.quantizer_dropout", config.s2mel.length_regulator_quantizer_dropout); + config.s2mel.length_regulator_f0_condition = yaml::optional_bool(document, "s2mel.length_regulator.f0_condition", config.s2mel.length_regulator_f0_condition); + config.s2mel.length_regulator_n_f0_bins = yaml::require_i64(document, "s2mel.length_regulator.n_f0_bins"); + config.s2mel.dit_hidden_dim = yaml::require_i64(document, "s2mel.DiT.hidden_dim"); + config.s2mel.dit_num_heads = yaml::require_i64(document, "s2mel.DiT.num_heads"); + config.s2mel.dit_depth = yaml::require_i64(document, "s2mel.DiT.depth"); + config.s2mel.dit_class_dropout_prob = yaml::optional_f32(document, "s2mel.DiT.class_dropout_prob", config.s2mel.dit_class_dropout_prob); + config.s2mel.dit_block_size = yaml::require_i64(document, "s2mel.DiT.block_size"); + config.s2mel.dit_in_channels = yaml::require_i64(document, "s2mel.DiT.in_channels"); + config.s2mel.dit_style_condition = yaml::optional_bool(document, "s2mel.DiT.style_condition", config.s2mel.dit_style_condition); + config.s2mel.dit_final_layer_type = yaml::require_string(document, "s2mel.DiT.final_layer_type"); + config.s2mel.dit_target = yaml::require_string(document, "s2mel.DiT.target"); + config.s2mel.dit_content_dim = yaml::require_i64(document, "s2mel.DiT.content_dim"); + config.s2mel.dit_content_codebook_size = yaml::require_i64(document, "s2mel.DiT.content_codebook_size"); + config.s2mel.dit_content_type = yaml::require_string(document, "s2mel.DiT.content_type"); + config.s2mel.dit_f0_condition = yaml::optional_bool(document, "s2mel.DiT.f0_condition", config.s2mel.dit_f0_condition); + config.s2mel.dit_n_f0_bins = yaml::require_i64(document, "s2mel.DiT.n_f0_bins"); + config.s2mel.dit_content_codebooks = yaml::require_i64(document, "s2mel.DiT.content_codebooks"); + config.s2mel.dit_is_causal = yaml::optional_bool(document, "s2mel.DiT.is_causal", config.s2mel.dit_is_causal); + config.s2mel.dit_long_skip_connection = yaml::optional_bool(document, "s2mel.DiT.long_skip_connection", config.s2mel.dit_long_skip_connection); + config.s2mel.dit_zero_prompt_speech_token = yaml::optional_bool(document, "s2mel.DiT.zero_prompt_speech_token", config.s2mel.dit_zero_prompt_speech_token); + config.s2mel.dit_time_as_token = yaml::optional_bool(document, "s2mel.DiT.time_as_token", config.s2mel.dit_time_as_token); + config.s2mel.dit_style_as_token = yaml::optional_bool(document, "s2mel.DiT.style_as_token", config.s2mel.dit_style_as_token); + config.s2mel.dit_uvit_skip_connection = yaml::optional_bool(document, "s2mel.DiT.uvit_skip_connection", config.s2mel.dit_uvit_skip_connection); + config.s2mel.dit_add_resblock_in_transformer = yaml::optional_bool(document, "s2mel.DiT.add_resblock_in_transformer", config.s2mel.dit_add_resblock_in_transformer); + config.s2mel.wavenet_hidden_dim = yaml::require_i64(document, "s2mel.wavenet.hidden_dim"); + config.s2mel.wavenet_num_layers = yaml::require_i64(document, "s2mel.wavenet.num_layers"); + config.s2mel.wavenet_kernel_size = yaml::require_i64(document, "s2mel.wavenet.kernel_size"); + config.s2mel.wavenet_dilation_rate = yaml::require_i64(document, "s2mel.wavenet.dilation_rate"); + config.s2mel.wavenet_dropout = yaml::optional_f32(document, "s2mel.wavenet.p_dropout", config.s2mel.wavenet_dropout); + config.s2mel.wavenet_style_condition = yaml::optional_bool(document, "s2mel.wavenet.style_condition", config.s2mel.wavenet_style_condition); + + config.emo_num = yaml::require_list_i64(document, "emo_num"); + return config; +} + +void validate_qwen_emotion_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("qwen_emotion_config"); + if (json::optional_string(root, "model_type", "") != "qwen3") { + throw std::runtime_error("IndexTTS2.5 Qwen emotion model must have model_type=qwen3"); + } + if (json::optional_i64(root, "hidden_size", 0) != 1024 || + json::optional_i64(root, "num_hidden_layers", 0) != 28 || + json::optional_i64(root, "num_attention_heads", 0) != 16) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion model config does not match expected 0.6B architecture"); + } +} + +void validate_config(const IndexTTS25Config & config, const assets::ResourceBundle & resources) { + engine::io::require_positive(config.dataset_sample_rate, "dataset.sample_rate"); + engine::io::require_positive(config.dataset_mel_sample_rate, "dataset.mel.sample_rate"); + engine::io::require_positive(config.dataset_mel_n_fft, "dataset.mel.n_fft"); + engine::io::require_positive(config.gpt.model_dim, "gpt.model_dim"); + engine::io::require_positive(config.gpt.layers, "gpt.layers"); + engine::io::require_divisible(config.gpt.model_dim, config.gpt.heads, "gpt.model_dim / gpt.heads"); + engine::io::require_positive(config.semantic_codec.codebook_size, "semantic_codec.codebook_size"); + engine::io::require_positive(config.s2mel.sample_rate, "s2mel.sample_rate"); + engine::io::require_positive(config.s2mel.n_mels, "s2mel.n_mels"); + engine::io::require_positive(config.s2mel.dit_hidden_dim, "s2mel.DiT.hidden_dim"); + engine::io::require_divisible(config.s2mel.dit_hidden_dim, config.s2mel.dit_num_heads, "s2mel.DiT.hidden_dim / num_heads"); + engine::io::require_nonnegative(config.s2mel.length_regulator_quantizer_dropout, "length_regulator.quantizer_dropout"); + engine::io::require_positive(config.s2mel.wavenet_dropout + 1.0F, "wavenet.p_dropout"); + if (config.emo_num.empty()) { + throw std::runtime_error("IndexTTS2.5 config emo_num must not be empty"); + } + validate_qwen_emotion_config(resources); +} + +void validate_gpt_weights(const IndexTTS25Config & config, const assets::TensorSource & source) { + assets::require_tensor_shape(source, "text_embedding.weight", {config.gpt.number_text_tokens + 1, config.gpt.model_dim}); + assets::require_tensor_shape(source, "mel_embedding.weight", {config.gpt.number_mel_codes, config.gpt.model_dim}); + assets::require_tensor_shape(source, "gpt.h.0.attn.c_attn.weight", {config.gpt.model_dim, config.gpt.model_dim * 3}); + assets::require_tensor_shape(source, "gpt.h.0.attn.c_proj.weight", {config.gpt.model_dim, config.gpt.model_dim}); + assets::require_tensor_shape(source, "gpt.h.0.mlp.c_fc.weight", {config.gpt.model_dim, config.gpt.model_dim * 4}); + assets::require_tensor_shape(source, "gpt.h.0.mlp.c_proj.weight", {config.gpt.model_dim * 4, config.gpt.model_dim}); + assets::require_tensor_shape(source, "spk_emb_proj.weight", {config.gpt.model_dim, config.s2mel.style_dim}); + assets::require_tensor_shape(source, "lang_embedding.weight", {kIndexTTS25LangEmbeddingRows, config.gpt.model_dim}); + assets::require_tensor_shape(source, "emo_conditioning_encoder.after_norm.weight", {config.gpt.emo_condition_output_size}); +} + +void validate_s2mel_weights(const IndexTTS25Config & config, const assets::TensorSource & source) { + assets::require_tensor_shape(source, "gpt_layer.0.weight", {256, config.gpt.model_dim}); + assets::require_tensor_shape(source, "gpt_layer.2.weight", {config.s2mel.length_regulator_in_channels, 128}); + assets::require_tensor_shape(source, "length_regulator.model.0.weight", {config.s2mel.length_regulator_channels, config.s2mel.length_regulator_channels, 3}); + assets::require_tensor_shape(source, "cfm.estimator.x_embedder.weight_v", {config.s2mel.dit_hidden_dim, config.s2mel.dit_in_channels}); + assets::require_tensor_shape(source, "cfm.estimator.transformer.layers.0.attention.wqkv.weight", {config.s2mel.dit_hidden_dim * 3, config.s2mel.dit_hidden_dim}); + assets::require_tensor_shape(source, "cfm.estimator.final_layer.adaLN_modulation.1.weight", {config.s2mel.dit_hidden_dim * 2, config.s2mel.dit_hidden_dim}); +} + +void validate_matrix_weights( + const IndexTTS25Config & config, + const assets::TensorSource & speaker_matrix, + const assets::TensorSource & emotion_matrix) { + int64_t total = 0; + for (const int64_t count : config.emo_num) { + engine::io::require_positive(count, "emo_num item"); + total += count; + } + assets::require_tensor_shape(speaker_matrix, "tensor", {total, config.s2mel.style_dim}); + assets::require_tensor_shape(emotion_matrix, "tensor", {total, config.gpt.model_dim}); +} + +void validate_w2v_stats(const IndexTTS25Config & config, const assets::TensorSource & source) { + assets::require_tensor_shape(source, "mean", {config.semantic_codec.hidden_size}); + assets::require_tensor_shape(source, "var", {config.semantic_codec.hidden_size}); +} + +void validate_w2v_weights(const assets::TensorSource & source) { + assets::require_tensor_shape(source, "feature_projection.projection.weight", {1024, 160}); + assets::require_tensor_shape(source, "encoder.layers.0.self_attn.linear_k.weight", {1024, 1024}); + assets::require_tensor_shape(source, "encoder.layers.0.conv_module.depthwise_conv.weight", {1024, 1, 31}); +} + +void validate_semantic_codec_weights(const IndexTTS25Config & config, const assets::TensorSource & source) { + assets::require_tensor_shape(source, "quantizer.quantizers.0.codebook.weight", {config.semantic_codec.codebook_size, config.semantic_codec.codebook_dim}); + assets::require_tensor_shape(source, "encoder.1.weight", {config.semantic_codec.hidden_size, config.semantic_codec.vocos_dim}); + assets::require_tensor_shape(source, "decoder.1.weight", {config.semantic_codec.hidden_size, config.semantic_codec.vocos_dim}); + assets::require_tensor_shape(source, "up.weight", {config.semantic_codec.hidden_size, config.semantic_codec.hidden_size, 3}); +} + +void validate_qwen_weights(const assets::TensorSource & source) { + assets::require_tensor_shape(source, "model.embed_tokens.weight", {151936, 1024}); + assets::require_tensor_shape(source, "model.layers.0.self_attn.q_proj.weight", {2048, 1024}); + assets::require_tensor_shape(source, "model.layers.0.self_attn.k_proj.weight", {1024, 1024}); + assets::require_tensor_shape(source, "model.layers.0.mlp.gate_proj.weight", {3072, 1024}); + assets::require_tensor_shape(source, "model.norm.weight", {1024}); +} + +void validate_weight_anchors(const IndexTTS25Assets & assets) { + validate_gpt_weights(assets.config, *assets.gpt_weights); + validate_s2mel_weights(assets.config, *assets.s2mel_weights); + validate_matrix_weights(assets.config, *assets.speaker_matrix, *assets.emotion_matrix); + validate_w2v_stats(assets.config, *assets.wav2vec2bert_stats); + validate_w2v_weights(*assets.wav2vec2bert_weights); + validate_semantic_codec_weights(assets.config, *assets.semantic_codec_weights); + validate_qwen_weights(*assets.qwen_emotion_weights); +} + +} // namespace + +std::shared_ptr load_index_tts2_5_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = engine::model_spec::load_resource_bundle( + model_path, + engine::model_spec::default_spec_path("index_tts2_5")); + assets->config = parse_config(assets->resources); + validate_config(assets->config, assets->resources); + + assets->gpt_weights = assets->resources.open_tensor_source("gpt"); + assets->s2mel_weights = assets->resources.open_tensor_source("s2mel"); + assets->speaker_matrix = assets->resources.open_tensor_source("speaker_matrix"); + assets->emotion_matrix = assets->resources.open_tensor_source("emotion_matrix"); + assets->wav2vec2bert_stats = assets->resources.open_tensor_source("wav2vec2bert_stats"); + assets->wav2vec2bert_weights = assets->resources.open_tensor_source("wav2vec2bert"); + assets->semantic_codec_weights = assets->resources.open_tensor_source("semantic_codec"); + assets->campplus_weights = assets->resources.open_tensor_source("campplus"); + assets->bigvgan_weights = assets->resources.open_tensor_source("bigvgan"); + assets->qwen_emotion_weights = assets->resources.open_tensor_source("qwen_emotion"); + + validate_weight_anchors(*assets); + return assets; +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/audio_features.cpp b/src/models/index_tts2_5/audio_features.cpp new file mode 100644 index 00000000..2e585500 --- /dev/null +++ b/src/models/index_tts2_5/audio_features.cpp @@ -0,0 +1,547 @@ +#include "engine/models/index_tts2_5/audio_features.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/dsp.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/audio/waveform_ops.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +struct MelFilterbankKey { + int64_t sample_rate = 0; + int64_t n_fft = 0; + int64_t num_mels = 0; + float fmin = 0.0F; + float fmax = 0.0F; + + bool operator==(const MelFilterbankKey & other) const noexcept { + return sample_rate == other.sample_rate && n_fft == other.n_fft && num_mels == other.num_mels && + fmin == other.fmin && fmax == other.fmax; + } +}; + +struct MelFilterbankKeyHash { + size_t operator()(const MelFilterbankKey & key) const noexcept { + size_t seed = std::hash{}(key.sample_rate); + seed ^= std::hash{}(key.n_fft) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.num_mels) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.fmin) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.fmax) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + return seed; + } +}; + +struct KaldiFilterbankKey { + int64_t sample_rate = 0; + int64_t padded_window_size = 0; + int64_t num_mels = 0; + float low_freq = 0.0F; + float high_freq = 0.0F; + + bool operator==(const KaldiFilterbankKey & other) const noexcept { + return sample_rate == other.sample_rate && + padded_window_size == other.padded_window_size && + num_mels == other.num_mels && + low_freq == other.low_freq && + high_freq == other.high_freq; + } +}; + +struct KaldiFilterbankKeyHash { + size_t operator()(const KaldiFilterbankKey & key) const noexcept { + size_t seed = 0; + seed ^= std::hash{}(key.sample_rate) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.padded_window_size) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.num_mels) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.low_freq) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.high_freq) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + return seed; + } +}; + +std::vector make_povey_window(int64_t window_size) { + std::vector window(static_cast(window_size), 0.0F); + constexpr float kPi = 3.14159265358979323846F; + for (int64_t i = 0; i < window_size; ++i) { + const float hann = + 0.5F - 0.5F * std::cos(2.0F * kPi * static_cast(i) / static_cast(window_size - 1)); + window[static_cast(i)] = std::pow(hann, 0.85F); + } + return window; +} + +std::vector make_kaldi_mel_filterbank( + int64_t sample_rate, + int64_t n_fft, + int64_t n_mels, + float low_freq, + float high_freq) { + const int64_t num_fft_bins = n_fft / 2 + 1; + const float nyquist = 0.5F * static_cast(sample_rate); + if (high_freq <= 0.0F) { + high_freq += nyquist; + } + const float fft_bin_width = static_cast(sample_rate) / static_cast(n_fft); + const float mel_low = 1127.0F * std::log(1.0F + low_freq / 700.0F); + const float mel_high = 1127.0F * std::log(1.0F + high_freq / 700.0F); + const float mel_delta = (mel_high - mel_low) / static_cast(n_mels + 1); + + std::vector filterbank(static_cast(n_mels * num_fft_bins), 0.0F); + for (int64_t mel_bin = 0; mel_bin < n_mels; ++mel_bin) { + const float left_mel = mel_low + static_cast(mel_bin) * mel_delta; + const float center_mel = mel_low + static_cast(mel_bin + 1) * mel_delta; + const float right_mel = mel_low + static_cast(mel_bin + 2) * mel_delta; + for (int64_t fft_bin = 0; fft_bin < num_fft_bins; ++fft_bin) { + const float freq = fft_bin_width * static_cast(fft_bin); + const float mel = 1127.0F * std::log(1.0F + freq / 700.0F); + const float up_slope = (mel - left_mel) / std::max(center_mel - left_mel, 1.0e-12F); + const float down_slope = (right_mel - mel) / std::max(right_mel - center_mel, 1.0e-12F); + filterbank[static_cast(mel_bin * num_fft_bins + fft_bin)] = + std::max(0.0F, std::min(up_slope, down_slope)); + } + } + return filterbank; +} + +const std::vector & cached_mel_filterbank(const IndexTTS25S2MelConfig & config) { + static std::mutex mutex; + static std::unordered_map, MelFilterbankKeyHash> cache; + const float fmax = config.fmax.value_or(static_cast(config.sample_rate) / 2.0F); + const MelFilterbankKey key{config.sample_rate, config.n_fft, config.n_mels, config.fmin, fmax}; + std::lock_guard lock(mutex); + const auto it = cache.find(key); + if (it != cache.end()) { + return it->second; + } + const auto filterbank = engine::audio::MelFilterbank().build({ + config.sample_rate, + config.n_fft, + config.n_mels, + config.fmin, + fmax, + true, + }); + return cache.emplace( + key, + filterbank.values).first->second; +} + +const std::vector & cached_povey_window(int64_t window_size) { + static std::mutex mutex; + static std::unordered_map> cache; + std::lock_guard lock(mutex); + const auto it = cache.find(window_size); + if (it != cache.end()) { + return it->second; + } + return cache.emplace(window_size, make_povey_window(window_size)).first->second; +} + +const std::vector & cached_kaldi_mel_filterbank( + int64_t sample_rate, + int64_t padded_window_size, + int64_t num_mels, + float low_freq, + float high_freq) { + static std::mutex mutex; + static std::unordered_map, KaldiFilterbankKeyHash> cache; + const KaldiFilterbankKey key{sample_rate, padded_window_size, num_mels, low_freq, high_freq}; + std::lock_guard lock(mutex); + const auto it = cache.find(key); + if (it != cache.end()) { + return it->second; + } + return cache.emplace( + key, + make_kaldi_mel_filterbank(sample_rate, padded_window_size, num_mels, low_freq, high_freq)).first->second; +} + +struct RealDftTables { + std::vector cos; + std::vector sin; +}; + +const RealDftTables & cached_real_dft_tables_512() { + static const RealDftTables tables = [] { + constexpr int64_t kFft = 512; + constexpr int64_t kFreqBins = kFft / 2 + 1; + constexpr double kPi = 3.14159265358979323846264338327950288; + RealDftTables out; + out.cos.resize(static_cast(kFreqBins * kFft)); + out.sin.resize(static_cast(kFreqBins * kFft)); + for (int64_t freq = 0; freq < kFreqBins; ++freq) { + for (int64_t n = 0; n < kFft; ++n) { + const double angle = 2.0 * kPi * static_cast(freq * n) / static_cast(kFft); + out.cos[static_cast(freq * kFft + n)] = std::cos(angle); + out.sin[static_cast(freq * kFft + n)] = std::sin(angle); + } + } + return out; + }(); + return tables; +} + +std::vector require_mono_samples(const std::vector & samples, int channels) { + if (channels <= 0) { + throw std::runtime_error("IndexTTS2.5 audio channel count must be positive"); + } + if (samples.empty()) { + throw std::runtime_error("IndexTTS2.5 audio must not be empty"); + } + if (samples.size() % static_cast(channels) != 0) { + throw std::runtime_error("IndexTTS2.5 audio sample count must be divisible by channel count"); + } + if (channels == 1) { + return samples; + } + return engine::audio::mixdown_interleaved_to_mono_average(samples, channels); +} + +std::vector resample_mono(const std::vector & input, int input_sample_rate, int output_sample_rate) { + if (input_sample_rate <= 0 || output_sample_rate <= 0) { + throw std::runtime_error("IndexTTS2.5 resampling requires positive sample rates"); + } + if (input_sample_rate == output_sample_rate || input.empty()) { + return input; + } + engine::audio::TorchaudioSincHannResampleOptions options; + options.kernel_mode = engine::audio::TorchaudioSincHannKernelMode::Float32ComputationStoredAsFloat32; + options.accumulation = engine::audio::TorchaudioSincHannAccumulation::Float32; + return engine::audio::resample_mono_torchaudio_sinc_hann(input, input_sample_rate, output_sample_rate, options); +} + +std::vector resample_mono_librosa(const std::vector & input, int input_sample_rate, int output_sample_rate) { + if (input_sample_rate <= 0 || output_sample_rate <= 0) { + throw std::runtime_error("IndexTTS2.5 librosa-style resampling requires positive sample rates"); + } + if (input_sample_rate == output_sample_rate || input.empty()) { + return input; + } + engine::audio::SoxrResampleOptions options; + options.profile = engine::audio::SoxrResampleProfile::ExplicitFloat32Runtime; + options.output_length_policy = engine::audio::SoxrOutputLengthPolicy::ExactExpected; + options.require_full_input = true; + if (auto output = engine::audio::try_resample_mono_soxr(input, input_sample_rate, output_sample_rate, options)) { + return *output; + } + return engine::audio::resample_mono_torchaudio_sinc_hann(input, input_sample_rate, output_sample_rate); +} + +} // namespace + +IndexTTS25MelOutput compute_index_tts2_5_mel_spectrogram( + const std::vector & waveform, + const IndexTTS25S2MelConfig & config, + size_t threads) { + if (config.sample_rate <= 0 || config.n_fft <= 0 || config.win_length <= 0 || + config.hop_length <= 0 || config.n_mels <= 0) { + throw std::runtime_error("IndexTTS2.5 mel spectrogram config is invalid"); + } + if (waveform.empty()) { + throw std::runtime_error("IndexTTS2.5 mel spectrogram requires non-empty waveform"); + } + + const int64_t pad = (config.n_fft - config.hop_length) / 2; + const auto padded = engine::audio::reflect_pad_samples(waveform, pad, pad); + const engine::audio::STFTConfig stft_config{ + config.n_fft, + config.hop_length, + config.win_length, + false, + engine::audio::STFTPadMode::Reflect, + engine::audio::STFTFamily::Kokoro, + }; + const auto & window = engine::audio::get_cached_stft_window(stft_config); + auto magnitude = engine::audio::STFT().compute_magnitude( + padded, + window, + 1, + static_cast(padded.size()), + stft_config, + threads); + for (float & value : magnitude.values) { + value = std::sqrt(value * value + 1.0e-9F); + } + + const auto & filterbank = cached_mel_filterbank(config); + const int64_t freq_bins = magnitude.shape[1]; + const int64_t frames = magnitude.shape[2]; + if (static_cast(filterbank.size()) != config.n_mels * freq_bins) { + throw std::runtime_error("IndexTTS2.5 mel filterbank shape mismatch"); + } + + IndexTTS25MelOutput output; + output.channels = config.n_mels; + output.frames = frames; + output.values.assign(static_cast(config.n_mels * frames), 0.0F); +#ifdef _OPENMP +#pragma omp parallel for collapse(2) if(config.n_mels * frames >= 4096) +#endif + for (int64_t mel = 0; mel < config.n_mels; ++mel) { + for (int64_t frame = 0; frame < frames; ++frame) { + float sum = 0.0F; + for (int64_t freq = 0; freq < freq_bins; ++freq) { + sum += filterbank[static_cast(mel * freq_bins + freq)] * + magnitude.values[static_cast(freq * frames + frame)]; + } + output.values[static_cast(mel * frames + frame)] = std::log(std::max(sum, 1.0e-5F)); + } + } + return output; +} + +IndexTTS25FbankOutput compute_index_tts2_5_campplus_fbank_16k(const std::vector & waveform_16k) { + constexpr int64_t kSampleRate = 16000; + constexpr int64_t kWindowSize = 400; + constexpr int64_t kWindowShift = 160; + constexpr int64_t kPaddedWindowSize = 512; + constexpr int64_t kNumMels = 80; + constexpr float kLowFreq = 20.0F; + constexpr float kHighFreq = 0.0F; + constexpr float kPreemphasis = 0.97F; + constexpr float kEpsilon = std::numeric_limits::epsilon(); + + if (static_cast(waveform_16k.size()) < kWindowSize) { + throw std::runtime_error("IndexTTS2.5 CAMPPlus fbank requires at least one 25 ms frame"); + } + + const int64_t frames = 1 + (static_cast(waveform_16k.size()) - kWindowSize) / kWindowShift; + const auto & window = cached_povey_window(kWindowSize); + const auto & mel_filterbank = cached_kaldi_mel_filterbank( + kSampleRate, + kPaddedWindowSize, + kNumMels, + kLowFreq, + kHighFreq); + + std::vector frame(static_cast(kWindowSize), 0.0F); + std::vector stft_batch(static_cast(frames * kPaddedWindowSize), 0.0F); + for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { + const int64_t start = frame_index * kWindowShift; + float mean = 0.0F; + for (int64_t i = 0; i < kWindowSize; ++i) { + const float sample = waveform_16k[static_cast(start + i)]; + frame[static_cast(i)] = sample; + mean += sample; + } + mean /= static_cast(kWindowSize); + for (int64_t i = 0; i < kWindowSize; ++i) { + frame[static_cast(i)] -= mean; + } + for (int64_t i = kWindowSize - 1; i > 0; --i) { + frame[static_cast(i)] -= kPreemphasis * frame[static_cast(i - 1)]; + } + frame[0] -= kPreemphasis * frame[0]; + for (int64_t i = 0; i < kWindowSize; ++i) { + stft_batch[static_cast(frame_index * kPaddedWindowSize + i)] = + frame[static_cast(i)] * window[static_cast(i)]; + } + } + + std::vector stft_window(static_cast(kPaddedWindowSize), 1.0F); + const engine::audio::STFTConfig stft_config{ + kPaddedWindowSize, + kPaddedWindowSize, + kPaddedWindowSize, + false, + engine::audio::STFTPadMode::Constant, + engine::audio::STFTFamily::Default, + }; + const auto magnitude = engine::audio::STFT().compute_magnitude( + stft_batch, + stft_window, + frames, + kPaddedWindowSize, + stft_config); + + const int64_t freq_bins = (kPaddedWindowSize / 2) + 1; + IndexTTS25FbankOutput output; + output.frames = frames; + output.dims = kNumMels; + output.values.assign(static_cast(frames * kNumMels), 0.0F); + for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { + for (int64_t mel_bin = 0; mel_bin < kNumMels; ++mel_bin) { + float energy = 0.0F; + for (int64_t freq = 0; freq < freq_bins; ++freq) { + const float mag = magnitude.values[static_cast(frame_index * freq_bins + freq)]; + energy += (mag * mag) * mel_filterbank[static_cast(mel_bin * freq_bins + freq)]; + } + output.values[static_cast(frame_index * kNumMels + mel_bin)] = + std::log(std::max(energy, kEpsilon)); + } + } + + for (int64_t mel_bin = 0; mel_bin < kNumMels; ++mel_bin) { + float mean = 0.0F; + for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { + mean += output.values[static_cast(frame_index * kNumMels + mel_bin)]; + } + mean /= static_cast(frames); + for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { + output.values[static_cast(frame_index * kNumMels + mel_bin)] -= mean; + } + } + return output; +} + +IndexTTS25SemanticFeatureOutput compute_index_tts2_5_semantic_features_16k(const std::vector & waveform_16k) { + constexpr int64_t kSampleRate = 16000; + constexpr int64_t kWindowSize = 400; + constexpr int64_t kWindowShift = 160; + constexpr int64_t kFftSize = 512; + constexpr int64_t kFreqBins = kFftSize / 2 + 1; + constexpr int64_t kNumMels = 80; + constexpr float kLowFreq = 20.0F; + constexpr float kHighFreq = 8000.0F; + constexpr float kPreemphasis = 0.97F; + constexpr double kInputScale = 32768.0; + constexpr double kMelFloor = 1.192092955078125e-07; + + if (static_cast(waveform_16k.size()) < kWindowSize) { + throw std::runtime_error("IndexTTS2.5 semantic fbank requires at least one 25 ms frame"); + } + + const int64_t frames = 1 + (static_cast(waveform_16k.size()) - kWindowSize) / kWindowShift; + const auto & window = cached_povey_window(kWindowSize); + const auto & mel_filterbank = cached_kaldi_mel_filterbank( + kSampleRate, + kFftSize, + kNumMels, + kLowFreq, + kHighFreq); + const auto & dft = cached_real_dft_tables_512(); + + IndexTTS25FbankOutput fbank; + fbank.frames = frames; + fbank.dims = kNumMels; + fbank.values.assign(static_cast(frames * kNumMels), 0.0F); + +#ifdef _OPENMP +#pragma omp parallel for if(frames >= 8) +#endif + for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { + const int64_t start = frame_index * kWindowShift; + double buffer[kFftSize] = {}; + double mean = 0.0; + for (int64_t i = 0; i < kWindowSize; ++i) { + const double sample = static_cast(waveform_16k[static_cast(start + i)]) * kInputScale; + buffer[i] = sample; + mean += sample; + } + mean /= static_cast(kWindowSize); + for (int64_t i = 0; i < kWindowSize; ++i) { + buffer[i] -= mean; + } + for (int64_t i = kWindowSize - 1; i > 0; --i) { + buffer[i] -= kPreemphasis * buffer[i - 1]; + } + buffer[0] *= (1.0 - kPreemphasis); + for (int64_t i = 0; i < kWindowSize; ++i) { + buffer[i] *= static_cast(window[static_cast(i)]); + } + + double power[kFreqBins] = {}; + for (int64_t freq = 0; freq < kFreqBins; ++freq) { + double re = 0.0; + double im = 0.0; + const size_t table_offset = static_cast(freq * kFftSize); + for (int64_t n = 0; n < kFftSize; ++n) { + const double sample = buffer[n]; + re += sample * dft.cos[table_offset + static_cast(n)]; + im -= sample * dft.sin[table_offset + static_cast(n)]; + } + power[freq] = re * re + im * im; + } + + for (int64_t mel_bin = 0; mel_bin < kNumMels; ++mel_bin) { + double energy = 0.0; + for (int64_t freq = 0; freq < kFreqBins; ++freq) { + energy += static_cast(mel_filterbank[static_cast(mel_bin * kFreqBins + freq)]) * + power[freq]; + } + fbank.values[static_cast(frame_index * kNumMels + mel_bin)] = + static_cast(std::log(std::max(energy, kMelFloor))); + } + } + + for (int64_t mel_bin = 0; mel_bin < fbank.dims; ++mel_bin) { + double mean = 0.0; + for (int64_t frame = 0; frame < fbank.frames; ++frame) { + mean += static_cast(fbank.values[static_cast(frame * fbank.dims + mel_bin)]); + } + mean /= static_cast(fbank.frames); + + double variance = 0.0; + for (int64_t frame = 0; frame < fbank.frames; ++frame) { + const double diff = + static_cast(fbank.values[static_cast(frame * fbank.dims + mel_bin)]) - mean; + variance += diff * diff; + } + variance = fbank.frames > 1 ? variance / static_cast(fbank.frames - 1) : 0.0; + const float scale = static_cast(1.0 / std::sqrt(variance + 1.0e-7)); + for (int64_t frame = 0; frame < fbank.frames; ++frame) { + float & value = fbank.values[static_cast(frame * fbank.dims + mel_bin)]; + value = (value - static_cast(mean)) * scale; + } + } + + const int64_t padded_frames = fbank.frames + (fbank.frames % 2); + std::vector padded(static_cast(padded_frames * fbank.dims), 1.0F); + std::copy(fbank.values.begin(), fbank.values.end(), padded.begin()); + + IndexTTS25SemanticFeatureOutput output; + output.frames = padded_frames / 2; + output.dims = fbank.dims * 2; + output.values.assign(static_cast(output.frames * output.dims), 0.0F); + output.attention_mask.assign(static_cast(output.frames), 0); + for (int64_t pair = 0; pair < output.frames; ++pair) { + const int64_t first_frame = pair * 2; + const int64_t second_frame = first_frame + 1; + std::copy_n( + padded.data() + static_cast(first_frame * fbank.dims), + static_cast(fbank.dims), + output.values.data() + static_cast(pair * output.dims)); + std::copy_n( + padded.data() + static_cast(second_frame * fbank.dims), + static_cast(fbank.dims), + output.values.data() + static_cast(pair * output.dims + fbank.dims)); + output.attention_mask[static_cast(pair)] = second_frame < fbank.frames ? 1 : 0; + } + return output; +} + +IndexTTS25PreparedReferenceAudio prepare_index_tts2_5_reference_audio( + const std::vector & samples, + int sample_rate, + int channels, + const IndexTTS25S2MelConfig & mel_config, + size_t threads, + bool speaker_load_semantic) { + constexpr int64_t kMaxReferenceSeconds = 15; + auto mono = require_mono_samples(samples, channels); + const int64_t max_input_samples = static_cast(sample_rate) * kMaxReferenceSeconds; + if (max_input_samples > 0 && static_cast(mono.size()) > max_input_samples) { + engine::audio::truncate_samples_to_count(mono, static_cast(max_input_samples)); + } + + IndexTTS25PreparedReferenceAudio output; + output.waveform_22k = resample_mono_librosa(mono, sample_rate, mel_config.sample_rate); + output.waveform_16k = speaker_load_semantic + ? resample_mono(output.waveform_22k, mel_config.sample_rate, 16000) + : resample_mono_librosa(mono, sample_rate, 16000); + output.mel = compute_index_tts2_5_mel_spectrogram(output.waveform_22k, mel_config, threads); + output.campplus_fbank = compute_index_tts2_5_campplus_fbank_16k(output.waveform_16k); + output.semantic_features = compute_index_tts2_5_semantic_features_16k(output.waveform_16k); + return output; +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/gpt.cpp b/src/models/index_tts2_5/gpt.cpp new file mode 100644 index 00000000..1637a8ec --- /dev/null +++ b/src/models/index_tts2_5/gpt.cpp @@ -0,0 +1,2217 @@ +#include "engine/models/index_tts2_5/gpt.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/relative_attention.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" +#include "engine/framework/modules/optimizations/fast_projection_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/torch_random.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kSemanticHidden = 1024; +constexpr int64_t kModelDim = 1280; +constexpr int64_t kConditionDim = 512; +constexpr int64_t kEmotionConditionLayers = 4; +constexpr int64_t kGptLayers = 24; +constexpr int64_t kGptMlpDim = 5120; +constexpr int64_t kTextTokens = 60510; +constexpr int64_t kMelCodes = 8194; +constexpr int64_t kMelPositions = 1818; +constexpr int64_t kTextPositions = 602; +constexpr int64_t kConditionPosFrames = 5000; +constexpr int64_t kConditionConvKernel = 15; +constexpr int64_t kGptHeads = 20; +constexpr int64_t kGptHeadDim = kModelDim / kGptHeads; +// spk_cond_mode="campplus": the projected 192-dim CAMPPlus embedding forms a +// single speaker token, followed by two all-zero tokens. +constexpr int64_t kCampplusStyleDim = 192; +constexpr int64_t kConditionTokens = 3; +constexpr int32_t kStartTextToken = 0; +constexpr int32_t kStopTextToken = 1; +constexpr int32_t kStartMelToken = 8192; +constexpr int32_t kStopMelToken = 8193; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +core::TensorValue div(core::ModuleBuildContext & ctx, const core::TensorValue & lhs, const core::TensorValue & rhs) { + core::validate_shape(rhs, lhs.shape, "Div rhs"); + return core::wrap_tensor(ggml_div(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32); +} + +core::TensorValue scale(core::ModuleBuildContext & ctx, const core::TensorValue & input, float value) { + return core::wrap_tensor(ggml_scale(ctx.ggml, input.tensor, value), input.shape, GGML_TYPE_F32); +} + +core::TensorValue transpose_btc_bct(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + return modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, input); +} + +core::TensorValue build_biased_gpt_projection( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t in_features, + int64_t out_features, + const modules::LinearWeights & weights) { + if (ctx.backend_type != core::BackendType::Cuda) { + return modules::LinearModule({in_features, out_features, true, GGML_PREC_F32}).build(ctx, input, weights); + } + if (out_features % 4 != 0) { + throw std::runtime_error("IndexTTS2.5 GPT fast projection requires output features divisible by 4"); + } + auto projected = modules::FastPackedProjection4Module({in_features, out_features, GGML_PREC_F32}) + .build(ctx, input, {weights.weight, std::nullopt}); + if (!weights.bias.has_value()) { + throw std::runtime_error("IndexTTS2.5 GPT linear bias is missing"); + } + const auto matrix_shape = core::TensorShape::from_dims({projected.shape.prefix_elements(), out_features}); + auto matrix = core::reshape_tensor(ctx, projected, matrix_shape); + matrix = core::wrap_tensor(ggml_add(ctx.ggml, matrix.tensor, weights.bias->tensor), matrix_shape, GGML_TYPE_F32); + return core::reshape_tensor(ctx, matrix, input.shape.with_last_dim(out_features)); +} + +core::TensorValue repeat_bias( + core::ModuleBuildContext & ctx, + const core::TensorValue & bias, + const core::TensorValue & like, + int64_t heads, + int64_t dim) { + auto view = core::reshape_tensor(ctx, bias, core::TensorShape::from_dims({1, heads, 1, dim})); + return modules::RepeatModule({like.shape}).build(ctx, view); +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + return core::reshape_tensor(ctx, contiguous, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); +} + +core::TensorValue gelu_geglu( + core::ModuleBuildContext & ctx, + const core::TensorValue & input) { + const int64_t half = input.shape.last_dim() / 2; + auto x = modules::SliceModule({static_cast(input.shape.rank - 1), 0, half}).build(ctx, input); + auto gate = modules::SliceModule({static_cast(input.shape.rank - 1), half, half}).build(ctx, input); + gate = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, gate); + return modules::MulModule{}.build(ctx, x, gate); +} + +core::TensorValue glu_axis( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int axis) { + if (axis < 0 || axis >= static_cast(input.shape.rank) || input.shape.dims[static_cast(axis)] % 2 != 0) { + throw std::runtime_error("IndexTTS2.5 GLU axis shape mismatch"); + } + const int64_t half = input.shape.dims[static_cast(axis)] / 2; + auto value = modules::SliceModule({axis, 0, half}).build(ctx, input); + auto gate = modules::SliceModule({axis, half, half}).build(ctx, input); + gate = modules::SigmoidModule{}.build(ctx, gate); + return modules::MulModule{}.build(ctx, value, gate); +} + +core::TensorValue rms_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & gamma) { + const auto squared = core::wrap_tensor(ggml_sqr(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); + auto sum = modules::ReduceSumModule({static_cast(input.shape.rank - 1)}).build(ctx, squared); + sum = core::wrap_tensor(ggml_sqrt(ctx.ggml, sum.tensor), sum.shape, GGML_TYPE_F32); + auto normed = div(ctx, input, modules::RepeatModule({input.shape}).build(ctx, sum)); + normed = scale(ctx, normed, std::sqrt(static_cast(input.shape.last_dim()))); + auto gamma_view = core::reshape_tensor(ctx, gamma, core::TensorShape::from_dims({1, 1, gamma.shape.dims[0]})); + return modules::MulModule{}.build(ctx, normed, modules::RepeatModule({input.shape}).build(ctx, gamma_view)); +} + +core::TensorValue condition_subsample( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25GptConditionEncoderWeights & weights) { + const int64_t frames = input.shape.dims[1]; + const int64_t frames_after = (frames - 3) / 2 + 1; + auto x = core::reshape_tensor(ctx, input, core::TensorShape::from_dims({1, 1, frames, kSemanticHidden})); + x = modules::Conv2dModule({1, kConditionDim, 3, 3, 2, 2, 0, 0, 1, 1, true}).build(ctx, x, weights.subsampling.conv); + x = modules::ReluModule{}.build(ctx, x); + x = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); + x = core::wrap_tensor(ggml_cont(ctx.ggml, x.tensor), x.shape, x.type); + x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, frames_after, kConditionDim * ((kSemanticHidden - 1) / 2)})); + x = modules::LinearModule({kConditionDim * ((kSemanticHidden - 1) / 2), kConditionDim, true}) + .build(ctx, x, weights.subsampling.out); + return scale(ctx, x, std::sqrt(static_cast(kConditionDim))); +} + +core::TensorValue condition_rel_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & pos_emb, + const IndexTTS25GptConditionLayerWeights & weights, + int64_t heads) { + const int64_t dim = kConditionDim / heads; + auto q = modules::LinearModule({kConditionDim, kConditionDim, true}) + .build(ctx, input, {weights.self_attn.attention.q_weight, weights.self_attn.attention.q_bias}); + auto k = modules::LinearModule({kConditionDim, kConditionDim, true}) + .build(ctx, input, {weights.self_attn.attention.k_weight, weights.self_attn.attention.k_bias}); + auto v = modules::LinearModule({kConditionDim, kConditionDim, true}) + .build(ctx, input, {weights.self_attn.attention.v_weight, weights.self_attn.attention.v_bias}); + auto p = modules::LinearModule({kConditionDim, kConditionDim, false}) + .build(ctx, pos_emb, {weights.self_attn.pos_weight, std::nullopt}); + + q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, heads, dim)); + k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, k, heads, dim)); + v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, v, heads, dim)); + p = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, p, heads, dim)); + + const auto q_u = modules::AddModule{}.build(ctx, q, repeat_bias(ctx, weights.self_attn.pos_bias_u, q, heads, dim)); + const auto q_v = modules::AddModule{}.build(ctx, q, repeat_bias(ctx, weights.self_attn.pos_bias_v, q, heads, dim)); + const auto k_t = modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k); + const auto p_t = modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, p); + auto scores = modules::AddModule{}.build(ctx, modules::MatMulModule{}.build(ctx, q_u, k_t), modules::MatMulModule{}.build(ctx, q_v, p_t)); + scores = scale(ctx, scores, 1.0F / std::sqrt(static_cast(dim))); + auto attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, core::ensure_backend_addressable_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); + auto context = modules::MatMulModule{}.build(ctx, attn, v); + context = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context); + context = core::ensure_backend_addressable_layout(ctx, context); + context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], kConditionDim})); + return modules::LinearModule({kConditionDim, kConditionDim, true}) + .build(ctx, context, {weights.self_attn.attention.out_weight, weights.self_attn.attention.out_bias}); +} + +core::TensorValue condition_conv_module( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25GptConditionLayerWeights & weights) { + auto x = transpose_btc_bct(ctx, input); + x = modules::Conv1dModule({kConditionDim, 2 * kConditionDim, 1, 1, 0, 1, true}).build(ctx, x, weights.conv_pointwise_in); + x = glu_axis(ctx, x, 1); + x = modules::DepthwiseConv1dModule({kConditionDim, kConditionConvKernel, 1, 7, 1, true}).build(ctx, x, weights.conv_depthwise); + x = transpose_btc_bct(ctx, x); + x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.conv_norm); + x = modules::SiluModule{}.build(ctx, x); + x = transpose_btc_bct(ctx, x); + x = modules::Conv1dModule({kConditionDim, kConditionDim, 1, 1, 0, 1, true}).build(ctx, x, weights.conv_pointwise_out); + return transpose_btc_bct(ctx, x); +} + +core::TensorValue condition_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & pos_emb, + const IndexTTS25GptConditionLayerWeights & weights, + int64_t heads) { + auto x = input; + auto y = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_mha); + y = condition_rel_attention(ctx, y, pos_emb, weights, heads); + x = modules::AddModule{}.build(ctx, x, y); + + y = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_conv); + y = condition_conv_module(ctx, y, weights); + x = modules::AddModule{}.build(ctx, x, y); + + y = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_ff); + y = modules::LinearModule({kConditionDim, weights.feed_forward_in.weight.shape.dims[0], true}).build(ctx, y, weights.feed_forward_in); + y = modules::SiluModule{}.build(ctx, y); + y = modules::LinearModule({weights.feed_forward_in.weight.shape.dims[0], kConditionDim, true}).build(ctx, y, weights.feed_forward_out); + x = modules::AddModule{}.build(ctx, x, y); + return modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_final); +} + +core::TensorValue perceiver_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & latents, + const core::TensorValue & context, + const IndexTTS25PerceiverAttentionWeights & weights, + int64_t dim, + int64_t heads, + int64_t inner) { + const int64_t head_dim = inner / heads; + const auto full_context = modules::ConcatModule({1}).build(ctx, latents, context); + auto q = modules::LinearModule({dim, inner, false}).build(ctx, latents, weights.q); + auto kv = modules::LinearModule({dim, 2 * inner, false}).build(ctx, full_context, weights.kv); + auto k = modules::SliceModule({2, 0, inner}).build(ctx, kv); + auto v = modules::SliceModule({2, inner, inner}).build(ctx, kv); + q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, heads, head_dim)); + k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, k, heads, head_dim)); + v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, v, heads, head_dim)); + auto scores = modules::MatMulModule{}.build(ctx, q, modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k)); + scores = scale(ctx, scores, 1.0F / std::sqrt(static_cast(head_dim))); + auto attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, core::ensure_backend_addressable_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); + auto output = modules::MatMulModule{}.build(ctx, attn, v); + output = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, output); + output = core::ensure_backend_addressable_layout(ctx, output); + output = core::reshape_tensor(ctx, output, core::TensorShape::from_dims({latents.shape.dims[0], latents.shape.dims[1], inner})); + return modules::LinearModule({inner, dim, false}).build(ctx, output, weights.out); +} + +core::TensorValue perceiver_ff( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25PerceiverFeedForwardWeights & weights, + int64_t dim, + int64_t ff_in) { + auto hidden = modules::LinearModule({dim, ff_in, true}).build(ctx, input, weights.in); + hidden = gelu_geglu(ctx, hidden); + return modules::LinearModule({ff_in / 2, dim, true}).build(ctx, hidden, weights.out); +} + +core::TensorValue perceiver( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25PerceiverWeights & weights, + int64_t latents_count, + int64_t dim, + int64_t heads, + int64_t inner, + int64_t ff_in) { + auto context = modules::LinearModule({kConditionDim, dim, true}).build(ctx, input, weights.project_context); + auto latents = modules::RepeatModule({core::TensorShape::from_dims({1, latents_count, dim})}) + .build(ctx, core::reshape_tensor(ctx, weights.latents, core::TensorShape::from_dims({1, latents_count, dim}))); + for (const auto & layer : weights.layers) { + latents = modules::AddModule{}.build(ctx, latents, perceiver_attention(ctx, latents, context, layer.attention, dim, heads, inner)); + latents = modules::AddModule{}.build(ctx, latents, perceiver_ff(ctx, latents, layer.feed_forward, dim, ff_in)); + } + return rms_norm(ctx, latents, weights.norm_gamma); +} + +core::TensorValue condition_encoder( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25GptConditionEncoderWeights & encoder, + const IndexTTS25PerceiverWeights & perceiver_weights, + int64_t encoder_heads, + int64_t perceiver_latents, + int64_t perceiver_dim, + int64_t perceiver_heads, + int64_t perceiver_inner, + int64_t perceiver_ff_in) { + auto x = condition_subsample(ctx, input, encoder); + const int64_t frames = x.shape.dims[1]; + auto pos_emb = modules::SliceModule({1, 0, frames}).build(ctx, encoder.subsampling.pos_enc); + for (const auto & layer : encoder.layers) { + x = condition_layer(ctx, x, pos_emb, layer, encoder_heads); + } + x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, encoder.after_norm); + return perceiver(ctx, x, perceiver_weights, perceiver_latents, perceiver_dim, perceiver_heads, perceiver_inner, perceiver_ff_in); +} + +engine::modules::LinearWeights load_hf_conv1d_linear( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type, + int64_t in_features, + int64_t out_features, + bool use_bias) { + engine::modules::LinearWeights weights; + const auto source_weight = source.require_f32(prefix + ".weight", {in_features, out_features}); + std::vector transposed(static_cast(out_features * in_features)); + for (int64_t in = 0; in < in_features; ++in) { + for (int64_t out = 0; out < out_features; ++out) { + transposed[static_cast(out * in_features + in)] = + source_weight[static_cast(in * out_features + out)]; + } + } + weights.weight = store.make_from_f32( + engine::core::TensorShape::from_dims({out_features, in_features}), + storage_type, + std::move(transposed)); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_features}); + } + return weights; +} + +engine::modules::LinearWeights load_biasless_linear( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type, + int64_t out_features, + int64_t in_features) { + return binding::linear_from_source(store, source, prefix, storage_type, out_features, in_features, false); +} + +engine::modules::RelativeAttentionWeights load_condition_relative_attention( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type, + int64_t heads) { + engine::modules::RelativeAttentionWeights weights; + weights.attention.q_weight = store.load_tensor(source, prefix + ".linear_q.weight", storage_type, {kConditionDim, kConditionDim}); + weights.attention.q_bias = store.load_f32_tensor(source, prefix + ".linear_q.bias", {kConditionDim}); + weights.attention.k_weight = store.load_tensor(source, prefix + ".linear_k.weight", storage_type, {kConditionDim, kConditionDim}); + weights.attention.k_bias = store.load_f32_tensor(source, prefix + ".linear_k.bias", {kConditionDim}); + weights.attention.v_weight = store.load_tensor(source, prefix + ".linear_v.weight", storage_type, {kConditionDim, kConditionDim}); + weights.attention.v_bias = store.load_f32_tensor(source, prefix + ".linear_v.bias", {kConditionDim}); + weights.attention.out_weight = store.load_tensor(source, prefix + ".linear_out.weight", storage_type, {kConditionDim, kConditionDim}); + weights.attention.out_bias = store.load_f32_tensor(source, prefix + ".linear_out.bias", {kConditionDim}); + weights.pos_weight = store.load_tensor(source, prefix + ".linear_pos.weight", storage_type, {kConditionDim, kConditionDim}); + weights.pos_bias_u = store.load_f32_tensor(source, prefix + ".pos_bias_u", {heads, kConditionDim / heads}); + weights.pos_bias_v = store.load_f32_tensor(source, prefix + ".pos_bias_v", {heads, kConditionDim / heads}); + return weights; +} + +IndexTTS25GptConditionLayerWeights load_condition_layer( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t linear_units, + int64_t heads, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) { + IndexTTS25GptConditionLayerWeights layer; + layer.norm_ff = binding::norm_from_source(store, source, prefix + ".norm_ff", kConditionDim); + layer.norm_mha = binding::norm_from_source(store, source, prefix + ".norm_mha", kConditionDim); + layer.norm_conv = binding::norm_from_source(store, source, prefix + ".norm_conv", kConditionDim); + layer.norm_final = binding::norm_from_source(store, source, prefix + ".norm_final", kConditionDim); + layer.feed_forward_in = binding::linear_from_source( + store, + source, + prefix + ".feed_forward.w_1", + matmul_storage_type, + linear_units, + kConditionDim, + true); + layer.feed_forward_out = binding::linear_from_source( + store, + source, + prefix + ".feed_forward.w_2", + matmul_storage_type, + kConditionDim, + linear_units, + true); + layer.self_attn = load_condition_relative_attention(store, source, prefix + ".self_attn", matmul_storage_type, heads); + layer.conv_pointwise_in = binding::conv1d_from_source( + store, + source, + prefix + ".conv_module.pointwise_conv1", + conv_storage_type, + 2 * kConditionDim, + kConditionDim, + 1, + true); + layer.conv_depthwise = binding::depthwise_conv1d_from_source( + store, + source, + prefix + ".conv_module.depthwise_conv", + conv_storage_type, + kConditionDim, + kConditionConvKernel, + true); + layer.conv_norm = binding::norm_from_source(store, source, prefix + ".conv_module.norm", kConditionDim); + layer.conv_pointwise_out = binding::conv1d_from_source( + store, + source, + prefix + ".conv_module.pointwise_conv2", + conv_storage_type, + kConditionDim, + kConditionDim, + 1, + true); + return layer; +} + +IndexTTS25GptConditionEncoderWeights load_condition_encoder( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t layers, + int64_t linear_units, + int64_t heads, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) { + IndexTTS25GptConditionEncoderWeights encoder; + encoder.subsampling.conv = binding::conv2d_from_source( + store, + source, + prefix + ".embed.conv.0", + conv_storage_type, + kConditionDim, + 1, + 3, + 3, + true); + encoder.subsampling.out = binding::linear_from_source( + store, + source, + prefix + ".embed.out.0", + matmul_storage_type, + kConditionDim, + kConditionDim * ((kSemanticHidden - 1) / 2), + true); + encoder.subsampling.pos_enc = store.load_f32_tensor( + source, + prefix + ".embed.pos_enc.pe", + {1, kConditionPosFrames, kConditionDim}); + encoder.layers.reserve(static_cast(layers)); + for (int64_t i = 0; i < layers; ++i) { + encoder.layers.push_back(load_condition_layer( + store, + source, + prefix + ".encoders." + std::to_string(i), + linear_units, + heads, + matmul_storage_type, + conv_storage_type)); + } + encoder.after_norm = binding::norm_from_source(store, source, prefix + ".after_norm", kConditionDim); + return encoder; +} + +IndexTTS25PerceiverLayerWeights load_perceiver_layer( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t dim, + int64_t inner, + int64_t ff_in, + engine::assets::TensorStorageType storage_type) { + IndexTTS25PerceiverLayerWeights layer; + layer.attention.q = load_biasless_linear(store, source, prefix + ".0.to_q", storage_type, inner, dim); + layer.attention.kv = load_biasless_linear(store, source, prefix + ".0.to_kv", storage_type, inner * 2, dim); + layer.attention.out = load_biasless_linear(store, source, prefix + ".0.to_out", storage_type, dim, inner); + layer.feed_forward.in = binding::linear_from_source(store, source, prefix + ".1.0", storage_type, ff_in, dim, true); + layer.feed_forward.out = binding::linear_from_source(store, source, prefix + ".1.2", storage_type, dim, ff_in / 2, true); + return layer; +} + +IndexTTS25PerceiverWeights load_perceiver( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t latents, + int64_t dim, + int64_t context_dim, + int64_t inner, + int64_t ff_in, + engine::assets::TensorStorageType storage_type) { + IndexTTS25PerceiverWeights weights; + weights.latents = store.load_f32_tensor(source, prefix + ".latents", {latents, dim}); + weights.project_context = binding::linear_from_source( + store, + source, + prefix + ".proj_context", + storage_type, + dim, + context_dim, + true); + weights.layers.reserve(2); + for (int64_t i = 0; i < 2; ++i) { + weights.layers.push_back(load_perceiver_layer( + store, + source, + prefix + ".layers." + std::to_string(i), + dim, + inner, + ff_in, + storage_type)); + } + weights.norm_gamma = store.load_f32_tensor(source, prefix + ".norm.gamma", {dim}); + return weights; +} + +IndexTTS25Gpt2LayerWeights load_gpt2_layer( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + int64_t layer_index, + engine::assets::TensorStorageType storage_type) { + const std::string prefix = "gpt.h." + std::to_string(layer_index); + IndexTTS25Gpt2LayerWeights layer; + layer.attn_norm = binding::norm_from_source(store, source, prefix + ".ln_1", kModelDim); + layer.qkv = load_hf_conv1d_linear(store, source, prefix + ".attn.c_attn", storage_type, kModelDim, 3 * kModelDim, true); + layer.attn_out = load_hf_conv1d_linear(store, source, prefix + ".attn.c_proj", storage_type, kModelDim, kModelDim, true); + layer.mlp_norm = binding::norm_from_source(store, source, prefix + ".ln_2", kModelDim); + layer.mlp_in = load_hf_conv1d_linear(store, source, prefix + ".mlp.c_fc", storage_type, kModelDim, kGptMlpDim, true); + layer.mlp_out = load_hf_conv1d_linear(store, source, prefix + ".mlp.c_proj", storage_type, kGptMlpDim, kModelDim, true); + return layer; +} + +struct Gpt2LayerOutput { + core::TensorValue output; + core::TensorValue key; + core::TensorValue value; +}; + +core::TensorValue gpt_attention_from_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + const std::optional & attention_mask) { + if (attention_mask.has_value()) { + auto q_contiguous = core::ensure_backend_addressable_layout(ctx, q_heads); + auto * flash = ggml_flash_attn_ext( + ctx.ggml, + q_contiguous.tensor, + k_heads.tensor, + v_heads.tensor, + attention_mask->tensor, + 1.0F / std::sqrt(static_cast(kGptHeadDim)), + 0.0F, + 0.0F); + ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); + return core::wrap_tensor( + flash, + core::TensorShape::from_dims({q_contiguous.shape.dims[0], q_contiguous.shape.dims[2], q_contiguous.shape.dims[1], kGptHeadDim}), + GGML_TYPE_F32); + } + 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(kGptHeadDim))), + scores.shape, + GGML_TYPE_F32); + scores = core::wrap_tensor(ggml_diag_mask_inf(ctx.ggml, scores.tensor, 0), scores.shape, GGML_TYPE_F32); + scores = core::wrap_tensor(ggml_soft_max(ctx.ggml, core::ensure_backend_addressable_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); + return modules::MatMulModule{}.build(ctx, scores, v_heads); +} + +core::TensorValue gpt_mlp( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25Gpt2LayerWeights & weights) { + auto hidden = build_biased_gpt_projection(ctx, input, kModelDim, kGptMlpDim, weights.mlp_in); + hidden = modules::GeluModule({modules::GeluApproximation::Tanh}).build(ctx, hidden); + return build_biased_gpt_projection(ctx, hidden, kGptMlpDim, kModelDim, weights.mlp_out); +} + +Gpt2LayerOutput gpt2_layer_full( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25Gpt2LayerWeights & weights, + const std::optional & attention_mask = std::nullopt) { + auto normed = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, input, weights.attn_norm); + auto qkv = build_biased_gpt_projection(ctx, normed, kModelDim, 3 * kModelDim, weights.qkv); + auto q = modules::SliceModule({2, 0, kModelDim}).build(ctx, qkv); + auto k = modules::SliceModule({2, kModelDim, kModelDim}).build(ctx, qkv); + auto v = modules::SliceModule({2, 2 * kModelDim, kModelDim}).build(ctx, qkv); + auto k_cache = reshape_heads(ctx, k, kGptHeads, kGptHeadDim); + auto v_cache = reshape_heads(ctx, v, kGptHeads, kGptHeadDim); + q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, kGptHeads, kGptHeadDim)); + auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k_cache); + auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v_cache); + auto context = gpt_attention_from_heads(ctx, q, k_heads, v_heads, attention_mask); + context = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context); + context = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, context), input.shape); + auto x = modules::AddModule{}.build(ctx, input, build_biased_gpt_projection(ctx, context, kModelDim, kModelDim, weights.attn_out)); + auto mlp_in = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights.mlp_norm); + return {modules::AddModule{}.build(ctx, x, gpt_mlp(ctx, mlp_in, weights)), k_cache, v_cache}; +} + +Gpt2LayerOutput gpt2_layer_cached_tail( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25Gpt2LayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slots, + const core::TensorValue & attention_mask) { + if (cache_key.shape.dims[0] != input.shape.dims[0] || + cache_value.shape.dims[0] != input.shape.dims[0] || + cache_key.shape.dims[1] != cache_value.shape.dims[1] || + cache_key.shape.dims[2] != cache_value.shape.dims[2] || + cache_key.shape.dims[3] != cache_value.shape.dims[3] || + cache_slots.shape.dims[0] != input.shape.dims[0]) { + throw std::runtime_error("IndexTTS2.5 GPT cached layer batch cache shape mismatch"); + } + auto normed = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, input, weights.attn_norm); + auto qkv = build_biased_gpt_projection(ctx, normed, kModelDim, 3 * kModelDim, weights.qkv); + auto q = modules::SliceModule({2, 0, kModelDim}).build(ctx, qkv); + auto k = modules::SliceModule({2, kModelDim, kModelDim}).build(ctx, qkv); + auto v = modules::SliceModule({2, 2 * kModelDim, kModelDim}).build(ctx, qkv); + q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, kGptHeads, kGptHeadDim)); + k = reshape_heads(ctx, k, kGptHeads, kGptHeadDim); + v = reshape_heads(ctx, v, kGptHeads, kGptHeadDim); + + const modules::FastKVSetRowsModule set_rows; + auto updated_key = set_rows.build(ctx, cache_key, k, cache_slots); + auto updated_value = set_rows.build(ctx, cache_value, v, cache_slots); + + auto context = gpt_attention_from_heads( + ctx, + q, + modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, updated_key), + modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, updated_value), + attention_mask); + context = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, context), input.shape); + auto x = modules::AddModule{}.build(ctx, input, build_biased_gpt_projection(ctx, context, kModelDim, kModelDim, weights.attn_out)); + auto mlp_in = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights.mlp_norm); + return {modules::AddModule{}.build(ctx, x, gpt_mlp(ctx, mlp_in, weights)), k, v}; +} + +struct TopPItem { + size_t index = 0; + float score = 0.0F; + float weight = 0.0F; +}; + +struct SampleScore { + size_t flat_index = 0; + float score = 0.0F; +}; + +struct RankedSample { + size_t score_index = 0; + size_t flat_index = 0; + double rank = 0.0; +}; + +struct IndexTTS25SamplerWorkspace { + std::vector scores; + std::vector top_k_heap; + std::vector top_p_items; + std::vector seen_tokens; + uint32_t seen_generation = 1; + std::vector finite_score_indices; + std::vector sample_scores; + std::vector ranked_samples; + std::vector selected_scores; +}; + +void apply_repetition_penalty( + std::vector & logits, + const std::vector & codes, + float penalty, + IndexTTS25SamplerWorkspace & workspace) { + if (penalty == 1.0F) { + return; + } + if (!(penalty > 0.0F)) { + throw std::runtime_error("IndexTTS2.5 GPT repetition_penalty must be positive"); + } + if (workspace.seen_tokens.size() != logits.size()) { + workspace.seen_tokens.assign(logits.size(), 0); + workspace.seen_generation = 1; + } else if (workspace.seen_generation == 0) { + std::fill(workspace.seen_tokens.begin(), workspace.seen_tokens.end(), 0); + workspace.seen_generation = 1; + } + const uint32_t generation = workspace.seen_generation++; + const auto apply_token = [&](int32_t token) { + if (token < 0 || static_cast(token) >= logits.size() || workspace.seen_tokens[static_cast(token)] == generation) { + return; + } + workspace.seen_tokens[static_cast(token)] = generation; + float & value = logits[static_cast(token)]; + value = value < 0.0F ? value * penalty : value / penalty; + }; + apply_token(kStopTextToken); + apply_token(kStartMelToken); + for (const int32_t token : codes) { + apply_token(token); + } +} + +float kth_largest_threshold(const std::vector & scores, size_t keep_count, std::vector & heap) { + heap.clear(); + heap.reserve(keep_count); + const auto greater = std::greater{}; + for (const float score : scores) { + if (heap.size() < keep_count) { + heap.push_back(score); + std::push_heap(heap.begin(), heap.end(), greater); + } else if (score > heap.front()) { + std::pop_heap(heap.begin(), heap.end(), greater); + heap.back() = score; + std::push_heap(heap.begin(), heap.end(), greater); + } + } + return heap.front(); +} + +void index_tts2_5_log_probs( + const std::vector & logits, + const std::vector & codes, + float repetition_penalty, + int top_k, + float top_p, + float temperature, + IndexTTS25SamplerWorkspace & workspace) { + if (!(temperature > 0.0F)) { + throw std::runtime_error("IndexTTS2.5 GPT temperature must be positive"); + } + float max_logit = -std::numeric_limits::infinity(); + for (float logit : logits) { + max_logit = std::max(max_logit, logit); + } + float total = 0.0F; + for (float logit : logits) { + total += std::exp(logit - max_logit); + } + if (!(total > 0.0F)) { + throw std::runtime_error("IndexTTS2.5 GPT sampler invalid logit mass"); + } + const float log_total = std::log(total); + auto & scores = workspace.scores; + scores.resize(logits.size()); + for (size_t i = 0; i < logits.size(); ++i) { + scores[i] = logits[i] - max_logit - log_total; + } + apply_repetition_penalty(scores, codes, repetition_penalty, workspace); + for (float & score : scores) { + score /= temperature; + } + const size_t min_tokens_to_keep = 2; + auto & finite_indices = workspace.finite_score_indices; + finite_indices.clear(); + if (top_k > 0 && static_cast(top_k) < scores.size()) { + const size_t keep_count = std::max(static_cast(top_k), min_tokens_to_keep); + const float threshold = kth_largest_threshold(scores, keep_count, workspace.top_k_heap); + float max_score = -std::numeric_limits::infinity(); + for (size_t i = 0; i < scores.size(); ++i) { + if (scores[i] < threshold) { + scores[i] = -std::numeric_limits::infinity(); + } else { + finite_indices.push_back(i); + max_score = std::max(max_score, scores[i]); + } + } + if (!std::isfinite(max_score)) { + throw std::runtime_error("IndexTTS2.5 GPT sampler has no finite score"); + } + } else { + finite_indices.reserve(scores.size()); + float max_score = -std::numeric_limits::infinity(); + for (size_t i = 0; i < scores.size(); ++i) { + if (std::isfinite(scores[i])) { + finite_indices.push_back(i); + max_score = std::max(max_score, scores[i]); + } + } + if (!std::isfinite(max_score)) { + throw std::runtime_error("IndexTTS2.5 GPT sampler has no finite score"); + } + } + if (top_p > 0.0F && top_p < 1.0F) { + auto & sorted = workspace.top_p_items; + sorted.clear(); + sorted.reserve(finite_indices.size()); + float total = 0.0F; + float max_score = -std::numeric_limits::infinity(); + for (const size_t i : finite_indices) { + max_score = std::max(max_score, scores[i]); + } + for (const size_t i : finite_indices) { + const float weight = std::exp(scores[i] - max_score); + sorted.push_back({i, scores[i], weight}); + total += weight; + } + if (!(total > 0.0F)) { + throw std::runtime_error("IndexTTS2.5 GPT sampler invalid top-p mass"); + } + std::sort(sorted.begin(), sorted.end(), [](const TopPItem & lhs, const TopPItem & rhs) { + if (lhs.score == rhs.score) { + return lhs.index < rhs.index; + } + return lhs.score < rhs.score; + }); + float cumulative = 0.0F; + const float remove_mass = 1.0F - top_p; + const size_t keep_from = sorted.size() > min_tokens_to_keep ? sorted.size() - min_tokens_to_keep : 0; + finite_indices.clear(); + for (size_t i = 0; i < sorted.size(); ++i) { + cumulative += sorted[i].weight / total; + if (i < keep_from && cumulative <= remove_mass) { + scores[sorted[i].index] = -std::numeric_limits::infinity(); + } else { + finite_indices.push_back(sorted[i].index); + } + } + } +} + +void sample_index_tts2_5_indices( + const std::vector & scores, + size_t total_score_count, + size_t count, + uint64_t seed, + uint64_t step, + const engine::sampling::TorchCudaSamplingPolicy & policy, + std::vector & ranked, + std::vector & selected) { + float max_score = -std::numeric_limits::infinity(); + for (const auto & score : scores) { + max_score = std::max(max_score, score.score); + } + if (!std::isfinite(max_score)) { + throw std::runtime_error("IndexTTS2.5 GPT sampler has no finite beam score"); + } + ranked.clear(); + ranked.reserve(scores.size()); + for (size_t i = 0; i < scores.size(); ++i) { + const auto & score = scores[i]; + const float probability = std::exp(score.score - max_score); + const float exponential = engine::sampling::torch_cuda_tensor_iterator_exponential_element( + seed, + static_cast(total_score_count), + static_cast(score.flat_index), + step, + policy.multiprocessor_count, + policy.max_threads_per_multiprocessor); + ranked.push_back({i, score.flat_index, static_cast(probability) / static_cast(exponential)}); + } + if (ranked.empty()) { + throw std::runtime_error("IndexTTS2.5 GPT sampler failed to select beam candidates"); + } + const size_t keep = std::min(count, ranked.size()); + std::partial_sort( + ranked.begin(), + ranked.begin() + static_cast(keep), + ranked.end(), + [](const RankedSample & lhs, const RankedSample & rhs) { + if (lhs.rank == rhs.rank) { + return lhs.flat_index < rhs.flat_index; + } + return lhs.rank > rhs.rank; + }); + selected.clear(); + selected.reserve(keep); + for (size_t i = 0; i < keep; ++i) { + selected.push_back(ranked[i].score_index); + } +} + +} // namespace + +std::shared_ptr load_index_tts2_5_gpt_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes) { + if (assets.gpt_weights == nullptr) { + throw std::runtime_error("IndexTTS2.5 GPT requires tensor source"); + } + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "index_tts2_5.gpt.weights", + weight_context_bytes); + + const auto & source = *assets.gpt_weights; + weights->emotion_conditioner = load_condition_encoder( + *weights->store, + source, + "emo_conditioning_encoder", + kEmotionConditionLayers, + 1024, + 4, + matmul_storage_type, + conv_storage_type); + weights->emotion_perceiver = load_perceiver( + *weights->store, + source, + "emo_perceiver_encoder", + 1, + kSemanticHidden, + kConditionDim, + 256, + 2730, + matmul_storage_type); + weights->spk_emb_proj = binding::linear_from_source( + *weights->store, + source, + "spk_emb_proj", + matmul_storage_type, + kModelDim, + kCampplusStyleDim, + true); + weights->lang_embedding = weights->store->load_tensor( + source, + "lang_embedding.weight", + matmul_storage_type, + {kIndexTTS25LangEmbeddingRows, kModelDim}); + weights->text_embedding = weights->store->load_tensor( + source, + "text_embedding.weight", + matmul_storage_type, + {kTextTokens, kModelDim}); + weights->mel_embedding = weights->store->load_tensor( + source, + "mel_embedding.weight", + matmul_storage_type, + {kMelCodes, kModelDim}); + weights->text_pos_embedding = weights->store->load_f32_tensor( + source, + "text_pos_embedding.emb.weight", + {kTextPositions, kModelDim}); + weights->mel_pos_embedding = weights->store->load_f32_tensor( + source, + "mel_pos_embedding.emb.weight", + {kMelPositions, kModelDim}); + weights->emotion_vec_projection = binding::linear_from_source( + *weights->store, + source, + "emovec_layer", + matmul_storage_type, + kModelDim, + kSemanticHidden, + true); + weights->emotion_layer = binding::linear_from_source( + *weights->store, + source, + "emo_layer", + matmul_storage_type, + kModelDim, + kModelDim, + true); + weights->gpt_layers.reserve(static_cast(kGptLayers)); + for (int64_t i = 0; i < kGptLayers; ++i) { + weights->gpt_layers.push_back(load_gpt2_layer(*weights->store, source, i, matmul_storage_type)); + } + weights->gpt_final_norm = binding::norm_from_source(*weights->store, source, "gpt.ln_f", kModelDim); + weights->final_norm = binding::norm_from_source(*weights->store, source, "final_norm", kModelDim); + weights->mel_head = binding::linear_from_source( + *weights->store, + source, + "mel_head", + matmul_storage_type, + kMelCodes, + kModelDim, + true); + weights->text_head = binding::linear_from_source( + *weights->store, + source, + "text_head", + matmul_storage_type, + kTextTokens, + kModelDim, + true); + + weights->store->upload(); + assets.gpt_weights->release_storage(); + return weights; +} + +class IndexTTS25GptRuntime::ConditioningGraph { +public: + ConditioningGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t frames, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + frames_(frames) { + if (frames_ <= 0) { + throw std::runtime_error("IndexTTS2.5 GPT conditioning graph requires positive frame count"); + } + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 GPT conditioning graph requires weights"); + } + + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT conditioning graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT conditioning input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.emo_condition", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.gpt.emo_condition.inputs", + execution_.backend_type()}; + semantic_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kSemanticHidden})).tensor; + ggml_set_input(semantic_); + auto input = core::wrap_tensor(semantic_, core::TensorShape::from_dims({1, frames_, kSemanticHidden}), GGML_TYPE_F32); + auto out = condition_encoder( + ctx, + input, + weights_->emotion_conditioner, + weights_->emotion_perceiver, + 4, + 1, + kSemanticHidden, + 4, + 256, + 2730); + output_ = core::ensure_backend_addressable_layout(ctx, out).tensor; + output_frames_ = out.shape.dims[1]; + output_dims_ = out.shape.dims[2]; + ggml_set_output(output_); + + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(65536, frames_ * 4096 + 8192)), false); + ggml_build_forward_expand(graph_, output_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT conditioning input buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT conditioning graph"); + } + debug::timing_log_scalar( + "index_tts2_5.gpt.emo_condition.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.gpt.emo_condition.frames", frames_); + } + + ~ConditioningGraph() { + clear_graph(); + } + + int64_t frames() const noexcept { + return frames_; + } + + IndexTTS25GptLatent run(const std::vector & semantic_btc) { + if (static_cast(semantic_btc.size()) != frames_ * kSemanticHidden) { + throw std::runtime_error("IndexTTS2.5 GPT conditioning input value count mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(semantic_, semantic_btc.data(), 0, semantic_btc.size() * sizeof(float)); + debug::timing_log_scalar( + "index_tts2_5.gpt.emo_condition.input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar( + "index_tts2_5.gpt.emo_condition.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 GPT conditioning graph compute failed"); + } + + IndexTTS25GptLatent output; + output.frames = output_frames_; + output.dims = output_dims_; + output.values.resize(static_cast(output.frames * output.dims)); + timing_start = Clock::now(); + ggml_backend_tensor_get(output_, output.values.data(), 0, output.values.size() * sizeof(float)); + debug::timing_log_scalar( + "index_tts2_5.gpt.emo_condition.output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return output; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t frames_ = 0; + int64_t output_frames_ = 0; + int64_t output_dims_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * semantic_ = nullptr; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; + }; + +class IndexTTS25GptRuntime::EmotionVectorGraph { +public: + EmotionVectorGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)) { + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 GPT emotion vector graph requires weights"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT emotion vector graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT emotion vector input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.emotion_vector", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.gpt.emotion_vector.inputs", + execution_.backend_type()}; + input_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kSemanticHidden})).tensor; + ggml_set_input(input_); + auto x = core::wrap_tensor(input_, core::TensorShape::from_dims({1, kSemanticHidden}), GGML_TYPE_F32); + x = modules::LinearModule({kSemanticHidden, kModelDim, true, GGML_PREC_F32}).build(ctx, x, weights_->emotion_vec_projection); + x = modules::LinearModule({kModelDim, kModelDim, true, GGML_PREC_F32}).build(ctx, x, weights_->emotion_layer); + output_ = core::ensure_backend_addressable_layout(ctx, x).tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 8192, false); + ggml_build_forward_expand(graph_, output_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT emotion vector input buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT emotion vector graph"); + } + debug::timing_log_scalar("index_tts2_5.gpt.emotion_vector.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~EmotionVectorGraph() { + clear_graph(); + } + + std::vector run(const IndexTTS25GptLatent & emotion_conditioning) { + if (emotion_conditioning.frames != 1 || + emotion_conditioning.dims != kSemanticHidden || + static_cast(emotion_conditioning.values.size()) != kSemanticHidden) { + throw std::runtime_error("IndexTTS2.5 GPT emotion vector input shape mismatch"); + } + ggml_backend_tensor_set(input_, emotion_conditioning.values.data(), 0, emotion_conditioning.values.size() * sizeof(float)); + core::set_backend_threads(execution_.backend(), execution_.config().threads); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 GPT emotion vector graph compute failed"); + } + std::vector out(static_cast(kModelDim)); + ggml_backend_tensor_get(output_, out.data(), 0, out.size() * sizeof(float)); + return out; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; +}; + +struct GptPrefillOutput { + std::vector logits; + std::vector latent; + runtime::TransformerKVState kv_state; +}; + +class IndexTTS25GptRuntime::PrefillGraph { +public: + PrefillGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t text_tokens, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + text_tokens_(text_tokens), + text_steps_(text_tokens + 2), + prompt_steps_(kConditionTokens + text_tokens + 3) { + if (weights_ == nullptr || text_tokens_ < 0) { + throw std::runtime_error("IndexTTS2.5 GPT prefill graph requires weights and non-negative text tokens"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT prefill graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT prefill input context"); + } + ggml_init_params output_params{16ull * 1024ull * 1024ull, nullptr, true}; + output_ctx_.reset(ggml_init(output_params)); + if (output_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT prefill output context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.prefill", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.gpt.prefill.inputs", + execution_.backend_type()}; + core::ModuleBuildContext output_ctx{ + output_ctx_.get(), + "index_tts2_5.gpt.prefill.outputs", + execution_.backend_type()}; + style_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kCampplusStyleDim})).tensor; + emo_vec_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kModelDim})).tensor; + lang_id_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 1); + text_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, text_steps_); + start_mel_id_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 1); + ggml_set_input(style_); + ggml_set_input(emo_vec_); + ggml_set_input(lang_id_); + ggml_set_input(text_ids_); + ggml_set_input(start_mel_id_); + // campplus conditioning prefix (model_v2.py inference_speech): + // conds = [spk_emb_proj(style) + emo_vec, zeros, zeros]. + auto style = core::wrap_tensor(style_, core::TensorShape::from_dims({1, kCampplusStyleDim}), GGML_TYPE_F32); + auto speaker_token = build_biased_gpt_projection(ctx, style, kCampplusStyleDim, kModelDim, weights_->spk_emb_proj); + speaker_token = core::reshape_tensor(ctx, speaker_token, core::TensorShape::from_dims({1, 1, kModelDim})); + auto emo_vec = core::wrap_tensor(emo_vec_, core::TensorShape::from_dims({1, kModelDim}), GGML_TYPE_F32); + emo_vec = core::reshape_tensor(ctx, emo_vec, core::TensorShape::from_dims({1, 1, kModelDim})); + auto conds = modules::AddModule{}.build(ctx, speaker_token, emo_vec); + auto zero_token = modules::RepeatModule({core::TensorShape::from_dims({1, kConditionTokens - 1, kModelDim})}) + .build(ctx, scale(ctx, conds, 0.0F)); + conds = modules::ConcatModule({1}).build(ctx, conds, zero_token); + auto text_ids = core::wrap_tensor(text_ids_, core::TensorShape::from_dims({text_steps_}), GGML_TYPE_I32); + auto text = modules::EmbeddingModule({kTextTokens, kModelDim}).build(ctx, text_ids, weights_->text_embedding); + auto text_pos = modules::SliceModule({0, 0, text_steps_}).build(ctx, weights_->text_pos_embedding); + text = modules::AddModule{}.build(ctx, text, text_pos); + auto lang_id = core::wrap_tensor(lang_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto lang = modules::EmbeddingModule({kIndexTTS25LangEmbeddingRows, kModelDim}).build(ctx, lang_id, weights_->lang_embedding); + text = modules::AddModule{}.build(ctx, text, modules::RepeatModule({text.shape}).build(ctx, lang)); + text = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, text), core::TensorShape::from_dims({1, text_steps_, kModelDim})); + auto mel_id = core::wrap_tensor(start_mel_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto mel = modules::EmbeddingModule({kMelCodes, kModelDim}).build(ctx, mel_id, weights_->mel_embedding); + auto mel_pos = modules::SliceModule({0, 0, 1}).build(ctx, weights_->mel_pos_embedding); + mel = modules::AddModule{}.build(ctx, mel, mel_pos); + mel = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, mel), core::TensorShape::from_dims({1, 1, kModelDim})); + auto x = modules::ConcatModule({1}).build(ctx, conds, text); + x = modules::ConcatModule({1}).build(ctx, x, mel); + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(65536, prompt_steps_ * 8192)), false); + for (const auto & layer : weights_->gpt_layers) { + auto out = gpt2_layer_full(ctx, x, layer); + x = out.output; + auto key = core::ensure_backend_addressable_layout(ctx, out.key); + auto value = core::ensure_backend_addressable_layout(ctx, out.value); + auto * key_output = core::make_tensor(output_ctx, GGML_TYPE_F32, key.shape).tensor; + auto * value_output = core::make_tensor(output_ctx, GGML_TYPE_F32, value.shape).tensor; + keys_.push_back(key_output); + values_.push_back(value_output); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), key.tensor, key_output)); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), value.tensor, value_output)); + } + x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->gpt_final_norm); + x = modules::SliceModule({1, prompt_steps_ - 1, 1}).build(ctx, x); + x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->final_norm); + auto latent = core::ensure_backend_addressable_layout(ctx, x); + auto logits = core::ensure_backend_addressable_layout( + ctx, + modules::LinearModule({kModelDim, kMelCodes, true, GGML_PREC_F32}).build(ctx, x, weights_->mel_head)); + latent_ = core::make_tensor(output_ctx, GGML_TYPE_F32, latent.shape).tensor; + logits_ = core::make_tensor(output_ctx, GGML_TYPE_F32, logits.shape).tensor; + ggml_set_output(latent_); + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), latent.tensor, latent_)); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), logits.tensor, logits_)); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT prefill input buffer"); + } + output_buffer_ = ggml_backend_alloc_ctx_tensors(output_ctx_.get(), execution_.backend()); + if (output_buffer_ == nullptr) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT prefill output buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT prefill graph"); + } + const int32_t start_mel = kStartMelToken; + ggml_backend_tensor_set(start_mel_id_, &start_mel, 0, sizeof(int32_t)); + debug::timing_log_scalar("index_tts2_5.gpt.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.gpt.prefill.prompt_steps", prompt_steps_); + } + + ~PrefillGraph() { + clear_graph(); + } + + bool matches(int64_t text_tokens) const noexcept { + return text_tokens_ == text_tokens; + } + + int64_t prompt_steps() const noexcept { + return prompt_steps_; + } + + GptPrefillOutput run( + const std::vector & speaker_style, + const std::vector & emotion_vector, + int32_t lang_id, + const std::vector & text_tokens) { + if (static_cast(speaker_style.size()) != kCampplusStyleDim || + static_cast(emotion_vector.size()) != kModelDim || + static_cast(text_tokens.size()) != text_tokens_) { + throw std::runtime_error("IndexTTS2.5 GPT prefill input shape mismatch"); + } + if (lang_id < 0 || lang_id >= kIndexTTS25LangEmbeddingRows) { + throw std::runtime_error("IndexTTS2.5 GPT prefill lang id is out of range"); + } + std::vector ids; + ids.reserve(static_cast(text_steps_)); + ids.push_back(kStartTextToken); + ids.insert(ids.end(), text_tokens.begin(), text_tokens.end()); + ids.push_back(kStopTextToken); + auto timing_start = Clock::now(); + ggml_backend_tensor_set(style_, speaker_style.data(), 0, speaker_style.size() * sizeof(float)); + ggml_backend_tensor_set(emo_vec_, emotion_vector.data(), 0, emotion_vector.size() * sizeof(float)); + ggml_backend_tensor_set(lang_id_, &lang_id, 0, sizeof(int32_t)); + ggml_backend_tensor_set(text_ids_, ids.data(), 0, ids.size() * sizeof(int32_t)); + debug::timing_log_scalar("index_tts2_5.gpt.prefill.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar("index_tts2_5.gpt.prefill.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 GPT prefill graph compute failed"); + } + GptPrefillOutput out; + out.logits.resize(static_cast(kMelCodes)); + out.latent.resize(static_cast(kModelDim)); + ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + ggml_backend_tensor_get(latent_, out.latent.data(), 0, out.latent.size() * sizeof(float)); + out.kv_state.current_end = prompt_steps_; + out.kv_state.layers.resize(keys_.size()); + const size_t layer_values = static_cast(prompt_steps_ * kGptHeads * kGptHeadDim); + for (size_t layer = 0; layer < keys_.size(); ++layer) { + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps_; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); + ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); + } + return out; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + if (output_buffer_ != nullptr) { + ggml_backend_buffer_free(output_buffer_); + output_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t text_tokens_ = 0; + int64_t text_steps_ = 0; + int64_t prompt_steps_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr output_ctx_; + std::unique_ptr ctx_; + ggml_tensor * style_ = nullptr; + ggml_tensor * emo_vec_ = nullptr; + ggml_tensor * lang_id_ = nullptr; + ggml_tensor * text_ids_ = nullptr; + ggml_tensor * start_mel_id_ = nullptr; + ggml_tensor * latent_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector keys_; + std::vector values_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; + ggml_backend_buffer_t output_buffer_ = nullptr; +}; + +class IndexTTS25GptRuntime::DecodeGraph { +public: + struct StepOutput { + std::vector logits; + }; + + struct BatchOutput { + std::vector steps; + }; + + DecodeGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t cache_steps, + int64_t beam_count, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + cache_steps_(cache_steps), + beam_count_(beam_count), + beam_slots_(2 * beam_count) { + if (weights_ == nullptr || cache_steps_ <= 0 || beam_count_ <= 0) { + throw std::runtime_error("IndexTTS2.5 GPT decode graph requires weights and cache steps"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT decode graph context"); + } + ggml_init_params state_params{256ull * 1024ull * 1024ull, nullptr, true}; + state_ctx_.reset(ggml_init(state_params)); + if (state_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 GPT decode state context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.decode", execution_.backend_type()}; + core::ModuleBuildContext state_ctx{ + state_ctx_.get(), + "index_tts2_5.gpt.decode.state", + execution_.backend_type()}; + token_ids_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, beam_count_); + mel_positions_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, beam_count_); + cache_slots_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, beam_count_); + attention_mask_ = ggml_new_tensor_4d(state_ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, beam_count_); + ggml_set_input(token_ids_); + ggml_set_input(mel_positions_); + ggml_set_input(cache_slots_); + ggml_set_input(attention_mask_); + for (int64_t bank = 0; bank < 2; ++bank) { + auto & keys = bank_keys_[static_cast(bank)]; + auto & values = bank_values_[static_cast(bank)]; + keys.reserve(weights_->gpt_layers.size()); + values.reserve(weights_->gpt_layers.size()); + for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { + keys.push_back(core::make_tensor( + state_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({beam_count_, cache_steps_, kGptHeads, kGptHeadDim}))); + values.push_back(core::make_tensor( + state_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({beam_count_, cache_steps_, kGptHeads, kGptHeadDim}))); + } + } + build_prefix_views(); + build_bank_graph(ctx, 0); + build_bank_graph(ctx, 1); + state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), execution_.backend()); + if (state_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT decode state buffer"); + } + debug::trace_log_scalar( + "index_tts2_5.gpt.decode.state_buffer_mib", + static_cast(ggml_backend_buffer_get_size(state_buffer_)) / (1024.0 * 1024.0)); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, bank_graphs_[0].graph) || + !ggml_gallocr_reserve(gallocr_, bank_graphs_[1].graph) || + !ggml_gallocr_alloc_graph(gallocr_, bank_graphs_[0].graph)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT decode graph"); + } + attention_mask_values_.assign(static_cast(beam_count_ * cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + token_values_.assign(static_cast(beam_count_), 0); + position_values_.assign(static_cast(beam_count_), 0); + cache_slot_values_.assign(static_cast(beam_count_), 0); + debug::timing_log_scalar("index_tts2_5.gpt.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.gpt.decode.cache_steps", cache_steps_); + debug::trace_log_scalar("index_tts2_5.gpt.decode.beam_batch", beam_count_); + } + + ~DecodeGraph() { + clear_graph(); + } + + void clear_graph() { + for (auto & graph : bank_graphs_) { + if (graph.graph != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph.graph); + graph.graph = nullptr; + } + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (state_buffer_ != nullptr) { + ggml_backend_buffer_free(state_buffer_); + state_buffer_ = nullptr; + } + } + + bool can_run(int64_t required_steps, int64_t required_beam_slots) const noexcept { + return cache_steps_ >= required_steps && beam_slots_ >= required_beam_slots && beam_count_ * 2 == required_beam_slots; + } + + void initialize_beam_slot(int64_t slot, const runtime::TransformerKVState & state) { + if (slot < 0 || slot >= beam_slots_) { + throw std::runtime_error("IndexTTS2.5 GPT beam slot is out of range"); + } + if (state.layers.size() != weights_->gpt_layers.size()) { + throw std::runtime_error("IndexTTS2.5 GPT beam state layer count mismatch"); + } + for (size_t layer = 0; layer < state.layers.size(); ++layer) { + const auto & layer_state = state.layers[layer]; + if (!state.layers.empty() && layer_state.valid_steps != state.layers.front().valid_steps) { + throw std::runtime_error("IndexTTS2.5 GPT beam state valid step mismatch"); + } + ggml_backend_tensor_set( + beam_key_prefix_views_[static_cast(slot)][static_cast(layer_state.valid_steps)][layer], + layer_state.key.data(), + 0, + layer_state.key.size() * sizeof(float)); + ggml_backend_tensor_set( + beam_value_prefix_views_[static_cast(slot)][static_cast(layer_state.valid_steps)][layer], + layer_state.value.data(), + 0, + layer_state.value.size() * sizeof(float)); + } + } + + BatchOutput run_batch_from_beams( + const std::vector & parent_slots, + const std::vector & child_slots, + int64_t valid_steps, + const std::vector & tokens, + int32_t mel_position) { + const size_t active = parent_slots.size(); + if (active == 0 || child_slots.size() != active || tokens.size() != active) { + throw std::runtime_error("IndexTTS2.5 GPT batched decode input shape mismatch"); + } + if (active > static_cast(beam_count_)) { + throw std::runtime_error("IndexTTS2.5 GPT batched decode exceeds beam batch"); + } + if (valid_steps >= cache_steps_) { + throw std::runtime_error("IndexTTS2.5 GPT decode cache exhausted"); + } + const int64_t child_bank = child_slots.front() / beam_count_; + if (child_bank < 0 || child_bank > 1) { + throw std::runtime_error("IndexTTS2.5 GPT child beam bank is out of range"); + } + for (size_t row = 0; row < active; ++row) { + if (parent_slots[row] < 0 || parent_slots[row] >= beam_slots_ || + child_slots[row] < 0 || child_slots[row] >= beam_slots_ || + child_slots[row] / beam_count_ != child_bank || + child_slots[row] % beam_count_ != static_cast(row)) { + throw std::runtime_error("IndexTTS2.5 GPT beam slot layout mismatch"); + } + for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { + if (parent_slots[row] != child_slots[row]) { + copy_beam_prefix(parent_slots[row], child_slots[row], valid_steps, layer); + } + } + token_values_[row] = tokens[row]; + position_values_[row] = mel_position; + } + for (size_t row = active; row < static_cast(beam_count_); ++row) { + const int64_t child_slot = child_bank * beam_count_ + static_cast(row); + for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { + copy_beam_prefix(child_slots.front(), child_slot, valid_steps, layer); + } + token_values_[row] = tokens.front(); + position_values_[row] = mel_position; + } + + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + std::fill(attention_mask_values_.begin(), attention_mask_values_.end(), masked); + for (int64_t row = 0; row < beam_count_; ++row) { + auto * row_values = attention_mask_values_.data() + static_cast(row * cache_steps_); + for (int64_t step = 0; step <= valid_steps; ++step) { + row_values[static_cast(step)] = visible; + } + cache_slot_values_[static_cast(row)] = static_cast(row * cache_steps_ + valid_steps); + } + ggml_backend_tensor_set(token_ids_, token_values_.data(), 0, token_values_.size() * sizeof(int32_t)); + ggml_backend_tensor_set(mel_positions_, position_values_.data(), 0, position_values_.size() * sizeof(int32_t)); + ggml_backend_tensor_set(cache_slots_, cache_slot_values_.data(), 0, cache_slot_values_.size() * sizeof(int32_t)); + ggml_backend_tensor_set(attention_mask_, attention_mask_values_.data(), 0, attention_mask_values_.size() * sizeof(ggml_fp16_t)); + + auto & graph = bank_graphs_[static_cast(child_bank)]; + if (!ggml_gallocr_alloc_graph(gallocr_, graph.graph)) { + throw std::runtime_error("failed to allocate IndexTTS2.5 GPT decode bank graph"); + } + core::set_backend_threads(execution_.backend(), execution_.config().threads); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph.graph); + ggml_backend_synchronize(execution_.backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 GPT decode graph compute failed"); + } + + std::vector logits(static_cast(beam_count_ * kMelCodes)); + ggml_backend_tensor_get(graph.logits, logits.data(), 0, logits.size() * sizeof(float)); + BatchOutput out; + out.steps.reserve(active); + for (size_t row = 0; row < active; ++row) { + StepOutput step; + step.logits.assign( + logits.begin() + static_cast(row * static_cast(kMelCodes)), + logits.begin() + static_cast((row + 1) * static_cast(kMelCodes))); + out.steps.push_back(std::move(step)); + } + return out; + } + +private: + struct BankGraph { + ggml_cgraph * graph = nullptr; + ggml_tensor * logits = nullptr; + }; + + void build_prefix_views() { + beam_key_prefix_views_.assign(static_cast(beam_slots_), {}); + beam_value_prefix_views_.assign(static_cast(beam_slots_), {}); + for (int64_t slot = 0; slot < beam_slots_; ++slot) { + const int64_t bank = slot / beam_count_; + const int64_t row = slot % beam_count_; + auto & key_steps = beam_key_prefix_views_[static_cast(slot)]; + auto & value_steps = beam_value_prefix_views_[static_cast(slot)]; + key_steps.assign(static_cast(cache_steps_ + 1), {}); + value_steps.assign(static_cast(cache_steps_ + 1), {}); + for (int64_t steps = 1; steps <= cache_steps_; ++steps) { + auto & key_layers = key_steps[static_cast(steps)]; + auto & value_layers = value_steps[static_cast(steps)]; + key_layers.reserve(weights_->gpt_layers.size()); + value_layers.reserve(weights_->gpt_layers.size()); + const int64_t elems = steps * kGptHeads * kGptHeadDim; + const size_t byte_offset = static_cast(row * cache_steps_ * kGptHeads * kGptHeadDim) * sizeof(float); + for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { + key_layers.push_back(ggml_view_1d(state_ctx_.get(), bank_keys_[static_cast(bank)][layer].tensor, elems, byte_offset)); + value_layers.push_back(ggml_view_1d(state_ctx_.get(), bank_values_[static_cast(bank)][layer].tensor, elems, byte_offset)); + } + } + } + } + + void copy_beam_prefix(int64_t parent_slot, int64_t child_slot, int64_t valid_steps, size_t layer) { + if (valid_steps <= 0 || parent_slot == child_slot) { + return; + } + if (valid_steps > cache_steps_) { + throw std::runtime_error("IndexTTS2.5 GPT beam prefix copy exceeds cache capacity"); + } + ggml_backend_tensor_copy( + beam_key_prefix_views_[static_cast(parent_slot)][static_cast(valid_steps)][layer], + beam_key_prefix_views_[static_cast(child_slot)][static_cast(valid_steps)][layer]); + ggml_backend_tensor_copy( + beam_value_prefix_views_[static_cast(parent_slot)][static_cast(valid_steps)][layer], + beam_value_prefix_views_[static_cast(child_slot)][static_cast(valid_steps)][layer]); + } + + void build_bank_graph(core::ModuleBuildContext & ctx, int64_t bank) { + BankGraph graph; + graph.graph = ggml_new_graph_custom(ctx_.get(), 65536, false); + auto token = core::wrap_tensor(token_ids_, core::TensorShape::from_dims({beam_count_}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({kMelCodes, kModelDim}).build(ctx, token, weights_->mel_embedding); + auto pos = core::wrap_tensor(mel_positions_, core::TensorShape::from_dims({beam_count_}), GGML_TYPE_I32); + auto pos_emb = modules::EmbeddingModule({kMelPositions, kModelDim}).build(ctx, pos, weights_->mel_pos_embedding); + x = modules::AddModule{}.build(ctx, x, pos_emb); + x = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, x), core::TensorShape::from_dims({beam_count_, 1, kModelDim})); + auto mask = core::wrap_tensor(attention_mask_, core::TensorShape::from_dims({beam_count_, 1, 1, cache_steps_}), GGML_TYPE_F16); + auto cache_slots = core::wrap_tensor(cache_slots_, core::TensorShape::from_dims({beam_count_}), GGML_TYPE_I32); + for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { + auto out = gpt2_layer_cached_tail( + ctx, + x, + weights_->gpt_layers[layer], + bank_keys_[static_cast(bank)][layer], + bank_values_[static_cast(bank)][layer], + cache_slots, + mask); + x = out.output; + } + x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->gpt_final_norm); + x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->final_norm); + graph.logits = modules::LinearModule({kModelDim, kMelCodes, true, GGML_PREC_F32}).build(ctx, x, weights_->mel_head).tensor; + ggml_set_output(graph.logits); + ggml_build_forward_expand(graph.graph, graph.logits); + bank_graphs_[static_cast(bank)] = graph; + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t cache_steps_ = 0; + int64_t beam_count_ = 0; + int64_t beam_slots_ = 0; + std::unique_ptr state_ctx_; + std::unique_ptr ctx_; + ggml_tensor * token_ids_ = nullptr; + ggml_tensor * mel_positions_ = nullptr; + ggml_tensor * cache_slots_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + std::array bank_graphs_; + std::array, 2> bank_keys_; + std::array, 2> bank_values_; + std::vector>> beam_key_prefix_views_; + std::vector>> beam_value_prefix_views_; + std::vector attention_mask_values_; + std::vector token_values_; + std::vector position_values_; + std::vector cache_slot_values_; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t state_buffer_ = nullptr; +}; + +IndexTTS25GptRuntime::IndexTTS25GptRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) + : assets_(std::move(assets)), + execution_(&execution), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 GPT runtime requires assets"); + } + if (graph_arena_bytes_ == 0) { + throw std::runtime_error("IndexTTS2.5 GPT graph arena must be non-zero"); + } + weights_ = load_index_tts2_5_gpt_weights( + *assets_, + execution.backend(), + execution.backend_type(), + matmul_storage_type, + conv_storage_type, + weight_context_bytes); +} + +IndexTTS25GptRuntime::~IndexTTS25GptRuntime() = default; + +void IndexTTS25GptRuntime::prepare_emotion_conditioning(int64_t frames) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 GPT runtime execution context is missing"); + } + if (frames <= 0) { + throw std::runtime_error("IndexTTS2.5 GPT emotion conditioning prepare requires positive frames"); + } + if (emotion_conditioning_graph_ != nullptr && emotion_conditioning_graph_->frames() == frames) { + return; + } + emotion_conditioning_graph_.reset(); + emotion_conditioning_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); +} + +IndexTTS25GptLatent IndexTTS25GptRuntime::emotion_conditioning(const std::vector & semantic_btc, int64_t frames) { + if (emotion_conditioning_graph_ == nullptr || emotion_conditioning_graph_->frames() != frames) { + throw std::runtime_error("IndexTTS2.5 GPT emotion conditioning graph was not prepared for this reference length"); + } + return emotion_conditioning_graph_->run(semantic_btc); +} + +void IndexTTS25GptRuntime::prepare_generation(int64_t text_tokens, int64_t max_mel_tokens, int64_t num_beams) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 GPT runtime execution context is missing"); + } + if (text_tokens < 0 || max_mel_tokens <= 0) { + throw std::runtime_error("IndexTTS2.5 GPT generation prepare requires non-negative text tokens and positive mel tokens"); + } + if (num_beams != 1) { + debug::trace_log_scalar("index_tts2_5.gpt.generation.num_beams", num_beams); + } + if (prefill_graph_ == nullptr || !prefill_graph_->matches(text_tokens)) { + prefill_graph_.reset(); + prefill_graph_ = std::make_unique(*execution_, weights_, text_tokens, graph_arena_bytes_); + } + const int64_t required_cache_steps = prefill_graph_->prompt_steps() + max_mel_tokens + 1; + const int64_t required_beam_slots = 2 * std::max(1, num_beams); + if (decode_graph_ == nullptr || !decode_graph_->can_run(required_cache_steps, required_beam_slots)) { + decode_graph_.reset(); + decode_graph_ = std::make_unique( + *execution_, + weights_, + required_cache_steps, + std::max(1, num_beams), + graph_arena_bytes_); + } +} + +std::vector IndexTTS25GptRuntime::project_emotion_vector(const IndexTTS25GptLatent & emotion_conditioning) { + if (emotion_vector_graph_ == nullptr) { + emotion_vector_graph_ = std::make_unique(*execution_, weights_, graph_arena_bytes_); + } + return emotion_vector_graph_->run(emotion_conditioning); +} + +std::vector IndexTTS25GptRuntime::merge_emotion_vector( + const std::vector & speaker_semantic, + int64_t speaker_frames, + const std::vector & emotion_semantic, + int64_t emotion_frames, + float alpha) { + prepare_emotion_conditioning(speaker_frames); + const auto base_condition = emotion_conditioning(speaker_semantic, speaker_frames); + prepare_emotion_conditioning(emotion_frames); + const auto emotion_condition = emotion_conditioning(emotion_semantic, emotion_frames); + auto base = project_emotion_vector(base_condition); + const auto emotion = project_emotion_vector(emotion_condition); + for (size_t i = 0; i < base.size(); ++i) { + base[i] = base[i] + alpha * (emotion[i] - base[i]); + } + return base; +} + +IndexTTS25GptGeneration IndexTTS25GptRuntime::generate_speech(const IndexTTS25GptGenerationRequest & request) { + if (request.text_tokens.empty()) { + throw std::runtime_error("IndexTTS2.5 GPT generation requires text tokens"); + } + const auto text_tokens = align_index_tts2_5_gpt_text_tokens(request.text_tokens); + if (static_cast(request.speaker_style.size()) != kCampplusStyleDim) { + throw std::runtime_error("IndexTTS2.5 GPT generation speaker style shape mismatch"); + } + if (request.lang_id < 0 || request.lang_id >= kIndexTTS25LangEmbeddingRows) { + throw std::runtime_error("IndexTTS2.5 GPT generation lang id is out of range"); + } + std::vector emotion_vector = request.emotion_vector; + if (emotion_vector.empty()) { + prepare_emotion_conditioning(request.emotion_frames); + emotion_vector = project_emotion_vector(emotion_conditioning(request.emotion_semantic, request.emotion_frames)); + } + if (static_cast(emotion_vector.size()) != kModelDim) { + throw std::runtime_error("IndexTTS2.5 GPT generation emotion vector shape mismatch"); + } + prepare_generation( + static_cast(text_tokens.size()), + request.max_mel_tokens, + request.num_beams); + auto prefill = prefill_graph_->run(request.speaker_style, emotion_vector, request.lang_id, text_tokens); + const auto sampling_policy = engine::sampling::resolve_torch_cuda_sampling_policy( + execution_->backend_type(), + execution_->config().device, + "index_tts2_5.gpt.cuda_sampling_policy", + "IndexTTS2.5", + engine::sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault); + + struct Beam { + std::vector codes; + std::vector logits; + int64_t slot = 0; + int64_t valid_steps = 0; + int64_t current_end = 0; + float score = 0.0F; + bool finished = false; + }; + auto normalized_score = [&](const Beam & beam) { + if (request.length_penalty == 0.0F) { + return beam.score; + } + const float length = static_cast(std::max(1, beam.codes.size())); + return beam.score / std::pow(length, request.length_penalty); + }; + struct BeamCandidate { + size_t parent = 0; + int32_t token = 0; + float score = 0.0F; + bool finished = false; + }; + const int beam_count = std::max(1, request.num_beams); + const int64_t prefill_valid_steps = prefill.kv_state.layers.empty() ? 0 : prefill.kv_state.layers.front().valid_steps; + std::vector beams; + beams.reserve(static_cast(beam_count)); + for (int beam = 0; beam < beam_count; ++beam) { + decode_graph_->initialize_beam_slot(beam, prefill.kv_state); + Beam initial; + initial.logits = prefill.logits; + initial.slot = beam; + initial.valid_steps = prefill_valid_steps; + initial.current_end = prefill.kv_state.current_end; + initial.score = beam == 0 ? 0.0F : -1.0e9F; + beams.push_back(std::move(initial)); + } + + std::vector completed; + auto add_completed = [&](Beam beam) { + completed.push_back(std::move(beam)); + if (static_cast(completed.size()) > beam_count) { + const auto worst = std::min_element(completed.begin(), completed.end(), [&](const Beam & lhs, const Beam & rhs) { + return normalized_score(lhs) < normalized_score(rhs); + }); + completed.erase(worst); + } + }; + auto completed_worst_score = [&]() { + if (completed.empty()) { + return -std::numeric_limits::infinity(); + } + const auto worst = std::min_element(completed.begin(), completed.end(), [&](const Beam & lhs, const Beam & rhs) { + return normalized_score(lhs) < normalized_score(rhs); + }); + return normalized_score(*worst); + }; + auto should_stop = [&](float best_sum_logprobs, int64_t cur_generated_len) { + if (static_cast(completed.size()) < beam_count) { + return false; + } + const float length = static_cast(std::max(1, cur_generated_len)); + const float highest_attainable = request.length_penalty == 0.0F + ? best_sum_logprobs + : best_sum_logprobs / std::pow(length, request.length_penalty); + return completed_worst_score() >= highest_attainable; + }; + bool first_decode_timing_logged = false; + int active_bank = 0; + bool beam_search_done = false; + uint64_t sample_call_index = 0; + uint64_t rng_offset_blocks = 0; + double sampling_ms = 0.0; + double decode_run_ms = 0.0; + IndexTTS25SamplerWorkspace sampler_workspace; + std::vector candidates; + std::vector next_beams; + std::vector parent_slots; + std::vector child_slots; + std::vector next_tokens; + candidates.reserve(static_cast(2 * beam_count)); + next_beams.reserve(static_cast(beam_count)); + parent_slots.reserve(static_cast(beam_count)); + child_slots.reserve(static_cast(beam_count)); + next_tokens.reserve(static_cast(beam_count)); + for (int step = 0; step < request.max_mel_tokens && !beams.empty(); ++step) { + const auto sampling_start = Clock::now(); + candidates.clear(); + sampler_workspace.sample_scores.clear(); + sampler_workspace.sample_scores.reserve(beams.size() * static_cast(std::max(request.top_k, 1))); + for (size_t beam_index = 0; beam_index < beams.size(); ++beam_index) { + index_tts2_5_log_probs( + beams[beam_index].logits, + beams[beam_index].codes, + request.repetition_penalty, + request.top_k, + request.top_p, + request.temperature, + sampler_workspace); + const size_t beam_offset = beam_index * static_cast(kMelCodes); + for (const size_t token : sampler_workspace.finite_score_indices) { + const float log_prob = sampler_workspace.scores[token]; + sampler_workspace.sample_scores.push_back({beam_offset + token, beams[beam_index].score + log_prob}); + } + } + const size_t flat_score_count = beams.size() * static_cast(kMelCodes); + const size_t keep = std::min(static_cast(2 * beam_count), flat_score_count); + if (request.do_sample) { + rng_offset_blocks += engine::sampling::torch_cuda_tensor_iterator_offset_blocks( + static_cast(flat_score_count), + sampling_policy); + sample_index_tts2_5_indices( + sampler_workspace.sample_scores, + flat_score_count, + keep, + request.seed, + sample_call_index++, + sampling_policy, + sampler_workspace.ranked_samples, + sampler_workspace.selected_scores); + std::sort(sampler_workspace.selected_scores.begin(), sampler_workspace.selected_scores.end(), [&](size_t lhs, size_t rhs) { + const auto & lhs_score = sampler_workspace.sample_scores[lhs]; + const auto & rhs_score = sampler_workspace.sample_scores[rhs]; + if (lhs_score.score == rhs_score.score) { + return lhs_score.flat_index < rhs_score.flat_index; + } + return lhs_score.score > rhs_score.score; + }); + } else { + sampler_workspace.selected_scores.resize(sampler_workspace.sample_scores.size()); + std::iota(sampler_workspace.selected_scores.begin(), sampler_workspace.selected_scores.end(), 0); + const size_t finite_keep = std::min(keep, sampler_workspace.selected_scores.size()); + std::partial_sort( + sampler_workspace.selected_scores.begin(), + sampler_workspace.selected_scores.begin() + static_cast(finite_keep), + sampler_workspace.selected_scores.end(), + [&](size_t lhs, size_t rhs) { + const auto & lhs_score = sampler_workspace.sample_scores[lhs]; + const auto & rhs_score = sampler_workspace.sample_scores[rhs]; + if (lhs_score.score == rhs_score.score) { + return lhs_score.flat_index < rhs_score.flat_index; + } + return lhs_score.score > rhs_score.score; + }); + sampler_workspace.selected_scores.resize(finite_keep); + } + for (size_t rank = 0; rank < sampler_workspace.selected_scores.size(); ++rank) { + const auto & selected = sampler_workspace.sample_scores[sampler_workspace.selected_scores[rank]]; + const size_t flat_index = selected.flat_index; + const size_t parent = flat_index / static_cast(kMelCodes); + const auto token = static_cast(flat_index % static_cast(kMelCodes)); + candidates.push_back({ + parent, + token, + selected.score, + token == kStopMelToken}); + } + sampling_ms += engine::debug::elapsed_ms(sampling_start, Clock::now()); + next_beams.clear(); + const int next_bank = 1 - active_bank; + parent_slots.clear(); + child_slots.clear(); + next_tokens.clear(); + for (size_t rank = 0; rank < candidates.size(); ++rank) { + const auto & candidate = candidates[rank]; + const Beam & parent = beams[candidate.parent]; + Beam next; + next.codes = parent.codes; + next.score = candidate.score; + if (candidate.finished) { + if (static_cast(rank) < beam_count) { + next.finished = true; + add_completed(std::move(next)); + } + continue; + } + next.codes.push_back(candidate.token); + next.slot = static_cast(next_bank * beam_count + static_cast(next_beams.size())); + next.valid_steps = parent.valid_steps + 1; + next.current_end = parent.current_end + 1; + parent_slots.push_back(parent.slot); + child_slots.push_back(next.slot); + next_tokens.push_back(candidate.token); + next_beams.push_back(std::move(next)); + if (static_cast(next_beams.size()) == beam_count) { + break; + } + } + if (!next_beams.empty()) { + const auto run_start = Clock::now(); + const int64_t parent_valid_steps = next_beams.front().valid_steps - 1; + const auto batch_out = decode_graph_->run_batch_from_beams( + parent_slots, + child_slots, + parent_valid_steps, + next_tokens, + static_cast(next_beams.front().codes.size() + 1)); + const auto run_ms = engine::debug::elapsed_ms(run_start, Clock::now()); + decode_run_ms += run_ms; + if (batch_out.steps.size() != next_beams.size()) { + throw std::runtime_error("IndexTTS2.5 GPT batched decode output size mismatch"); + } + for (size_t beam = 0; beam < next_beams.size(); ++beam) { + next_beams[beam].logits = batch_out.steps[beam].logits; + } + if (!first_decode_timing_logged) { + debug::timing_log_scalar("index_tts2_5.gpt.decode.first_run_ms", run_ms); + first_decode_timing_logged = true; + } + } + if (!candidates.empty() && should_stop(candidates.front().score, static_cast(step + 1))) { + beam_search_done = true; + break; + } + beams.swap(next_beams); + active_bank = next_bank; + } + debug::timing_log_scalar("index_tts2_5.gpt.sampling_ms", sampling_ms); + debug::timing_log_scalar("index_tts2_5.gpt.decode.run_ms", decode_run_ms); + if (!beam_search_done) { + for (auto & beam : beams) { + add_completed(std::move(beam)); + } + } + if (completed.empty()) { + throw std::runtime_error("IndexTTS2.5 GPT generation produced no beam candidates"); + } + const auto best = std::max_element(completed.begin(), completed.end(), [&](const Beam & lhs, const Beam & rhs) { + return normalized_score(lhs) < normalized_score(rhs); + }); + IndexTTS25GptGeneration out; + out.codes = best->codes; + out.rng_offset_blocks = rng_offset_blocks; + bool stop_seen = false; + for (size_t i = 0; i < out.codes.size(); ++i) { + if (out.codes[i] == kStopMelToken) { + stop_seen = true; + } + } + debug::trace_log_scalar("index_tts2_5.gpt.generated_code_count", static_cast(out.codes.size())); + debug::trace_log_scalar("index_tts2_5.gpt.generated_stop_seen", stop_seen); + return out; +} + +void IndexTTS25GptRuntime::release_conditioning_graphs() { + emotion_conditioning_graph_.reset(); + emotion_vector_graph_.reset(); +} + +void IndexTTS25GptRuntime::release_generation_graphs() { + prefill_graph_.reset(); + decode_graph_.reset(); +} + +std::vector align_index_tts2_5_gpt_text_tokens(const std::vector & text_tokens) { + std::vector out; + out.reserve(text_tokens.size()); + for (const int32_t token : text_tokens) { + if (token == kStartTextToken || token == kStopTextToken) { + continue; + } + out.push_back(token); + } + return out; +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/loader.cpp b/src/models/index_tts2_5/loader.cpp new file mode 100644 index 00000000..dd2ab65c --- /dev/null +++ b/src/models/index_tts2_5/loader.cpp @@ -0,0 +1,157 @@ +#include "engine/models/index_tts2_5/loader.h" + +#include "engine/framework/model_spec/package.h" +#include "engine/models/index_tts2_5/session.h" + +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +runtime::ModelMetadata metadata(const IndexTTS25Assets & assets) { + runtime::ModelMetadata out; + out.family = "index_tts2_5"; + out.variant = assets.config.version; + out.description = "IndexTTS2.5 loaded from local extracted assets."; + return out; +} + +runtime::CapabilitySet capabilities(const IndexTTS25Assets &) { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + {runtime::VoiceTaskKind::VoiceCloning, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + out.supports_style_condition = true; + out.languages = {"Chinese", "English", "Japanese", "Spanish", "Arabic"}; + return out; +} + +runtime::ModelCliInterface cli(const IndexTTS25Assets &) { + runtime::ModelCliInterface out; + out.request_options = { + {"lang", "auto|zh|en|ja|es|ar|...", "Text language hint; auto infers zh when the text contains Han characters, otherwise en."}, + {"emotion_alpha", "float", "Blend strength for explicit emotion conditioning."}, + {"emotion_vector", "float[,float...]", "Eight-value explicit emotion vector."}, + {"use_emotion_text", "bool", "Infer emotion from text instead of reference audio."}, + {"emotion_text", "text", "Text used when emotion-text conditioning is enabled."}, + {"use_random_emotion", "bool", "Use random emotion weights in the emotion mixer."}, + {"interval_silence_ms", "n", "Silence inserted between generated text chunks."}, + {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode used when text_chunk_size is set."}, + {"length_penalty", "float", "GPT beam-search length penalty."}, + {"num_beams", "n", "GPT beam count."}, + }; + out.session_options = { + {"index_tts2_5.weight_type", "native|f32|f16|bf16|q8_0", "Matmul weight storage type."}, + {"index_tts2_5.conv_weight_type", "native|f32|f16", "Convolution weight storage type."}, + {"index_tts2_5.gpt_graph_arena_mb", "n", "GPT graph arena size."}, + {"index_tts2_5.s2mel_graph_arena_mb", "n", "S2Mel graph arena size."}, + {"index_tts2_5.reference_graph_arena_mb", "n", "Reference encoder and codec graph arena size."}, + {"index_tts2_5.emotion_text_prefill_graph_arena_mb", "n", "Emotion-text prefill graph arena size."}, + {"index_tts2_5.emotion_text_decode_graph_arena_mb", "n", "Emotion-text cached-step graph arena size."}, + {"index_tts2_5.emotion_text_max_new_tokens", "n", "Maximum generated tokens for emotion-text classification; default 256."}, + {"index_tts2_5.weight_context_mb", "n", "Shared weight context size."}, + {"index_tts2_5.mem_saver", "true|false", "Release staged reference and conditioning graphs after request phases; default false."}, + {"index_tts2_5.speaker_cache_slots", "n", "Prepared speaker-reference cache slots; default 1."}, + {"index_tts2_5.emotion_cache_slots", "n", "Prepared emotion-reference cache slots; default 1."}, + {"index_tts2_5.emotion_text_cache_slots", "n", "Emotion-text weight cache slots; default 1."}, + }; + return out; +} + +class IndexTTS25Loader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "index_tts2_5"; + } + + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + {runtime::VoiceTaskKind::VoiceCloning, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + out.supports_style_condition = true; + return out; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + try { + const auto package_spec = engine::model_spec::default_spec_path(family()); + (void) engine::model_spec::load_resource_bundle( + request.model_path, + package_spec); + 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_index_tts2_5_assets(request.model_path); + runtime::ModelInspection inspection; + inspection.model_root = assets->resources.model_root(); + inspection.metadata = metadata(*assets); + inspection.capabilities = capabilities(*assets); + inspection.cli = cli(*assets); + const auto package_spec = engine::model_spec::default_spec_path(family()); + inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::model_spec::ResourceKind::Files); + inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::model_spec::ResourceKind::Tensors); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_index_tts2_5_model(request.model_path); + } +}; + +} // namespace + +IndexTTS25LoadedModel::IndexTTS25LoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets) + : metadata_(std::move(metadata)), + capabilities_(std::move(capabilities)), + assets_(std::move(assets)) {} + +const runtime::ModelMetadata & IndexTTS25LoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & IndexTTS25LoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr IndexTTS25LoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.mode != runtime::RunMode::Offline || + (task.task != runtime::VoiceTaskKind::Tts && task.task != runtime::VoiceTaskKind::VoiceCloning)) { + throw std::runtime_error("IndexTTS2.5 only supports offline TTS and voice-cloning sessions"); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_index_tts2_5_model(const std::filesystem::path & model_path) { + auto assets = load_index_tts2_5_assets(model_path); + return std::make_unique( + metadata(*assets), + capabilities(*assets), + std::move(assets)); +} + +std::shared_ptr make_index_tts2_5_loader() { + return std::make_shared(); +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/qwen_emotion.cpp b/src/models/index_tts2_5/qwen_emotion.cpp new file mode 100644 index 00000000..77955d47 --- /dev/null +++ b/src/models/index_tts2_5/qwen_emotion.cpp @@ -0,0 +1,794 @@ +#include "engine/models/index_tts2_5/qwen_emotion.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/runtime/kv_cache.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kHidden = 1024; +constexpr int64_t kIntermediate = 3072; +constexpr int64_t kLayers = 28; +constexpr int64_t kAttentionHeads = 16; +constexpr int64_t kKvHeads = 8; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kVocab = 151936; +constexpr float kRmsEps = 1.0e-6F; +constexpr float kRopeTheta = 1000000.0F; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +modules::QwenDecoderStackConfig qwen_config() { + modules::QwenDecoderStackConfig config; + config.hidden_size = kHidden; + config.num_attention_heads = kAttentionHeads; + config.num_key_value_heads = kKvHeads; + config.head_dim = kHeadDim; + config.intermediate_size = kIntermediate; + config.layers = kLayers; + config.rms_norm_eps = kRmsEps; + config.rope_theta = kRopeTheta; + config.attention_precision = GGML_PREC_F32; + config.projection_precision = GGML_PREC_DEFAULT; + return config; +} + +std::vector causal_mask(int64_t steps) { + std::vector mask(static_cast(steps * steps), 0.0F); + for (int64_t q = 0; q < steps; ++q) { + for (int64_t k = q + 1; k < steps; ++k) { + mask[static_cast(q * steps + k)] = -std::numeric_limits::infinity(); + } + } + return mask; +} + +float clamp_emotion(float value) { + return std::clamp(value, 0.0F, 1.2F); +} + +float parse_named_score(const std::string & content, const std::string & key) { + const std::regex pattern("\"?" + key + "\"?\\s*:\\s*([0-9]+(?:\\.[0-9]+)?)"); + std::smatch match; + if (!std::regex_search(content, match, pattern)) { + return 0.0F; + } + return clamp_emotion(std::stof(match[1].str())); +} + +IndexTTS25EmotionVector convert_emotion_json(const std::string & content, const std::string & source_text) { + IndexTTS25EmotionVector out; + out.values = { + parse_named_score(content, "高兴"), + parse_named_score(content, "愤怒"), + parse_named_score(content, "悲伤"), + parse_named_score(content, "恐惧"), + parse_named_score(content, "反感"), + parse_named_score(content, "低落"), + parse_named_score(content, "惊讶"), + parse_named_score(content, "自然"), + }; + std::string lower = source_text; + std::transform(lower.begin(), lower.end(), lower.begin(), [](unsigned char ch) { return static_cast(std::tolower(ch)); }); + if (lower.find("低落") != std::string::npos || + lower.find("melancholy") != std::string::npos || + lower.find("melancholic") != std::string::npos || + lower.find("depression") != std::string::npos || + lower.find("depressed") != std::string::npos || + lower.find("gloomy") != std::string::npos) { + std::swap(out.values[2], out.values[5]); + } + const bool all_zero = std::all_of(out.values.begin(), out.values.end(), [](float value) { return value <= 0.0F; }); + if (all_zero) { + out.values[7] = 1.0F; + } + return out; +} + +engine::modules::QwenDecoderLayerWeights load_qwen_layer( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + int64_t layer_index, + engine::assets::TensorStorageType storage_type) { + const std::string prefix = "model.layers." + std::to_string(layer_index); + engine::modules::QwenDecoderLayerWeights layer; + layer.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", kHidden); + layer.self_attention.q_weight = store.load_tensor( + source, + prefix + ".self_attn.q_proj.weight", + storage_type, + {kAttentionHeads * kHeadDim, kHidden}); + layer.self_attention.k_weight = store.load_tensor( + source, + prefix + ".self_attn.k_proj.weight", + storage_type, + {kKvHeads * kHeadDim, kHidden}); + layer.self_attention.v_weight = store.load_tensor( + source, + prefix + ".self_attn.v_proj.weight", + storage_type, + {kKvHeads * kHeadDim, kHidden}); + layer.self_attention.out_weight = store.load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + storage_type, + {kHidden, kAttentionHeads * kHeadDim}); + layer.q_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.q_norm", kHeadDim); + layer.k_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.k_norm", kHeadDim); + layer.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", kHidden); + layer.mlp.gate_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.gate_proj", + storage_type, + kIntermediate, + kHidden, + false); + layer.mlp.up_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.up_proj", + storage_type, + kIntermediate, + kHidden, + false); + layer.mlp.down_proj = binding::linear_from_source( + store, + source, + prefix + ".mlp.down_proj", + storage_type, + kHidden, + kIntermediate, + false); + return layer; +} + +} // namespace + +std::shared_ptr load_index_tts2_5_qwen_emotion_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType storage_type, + size_t weight_context_bytes) { + if (assets.qwen_emotion_weights == nullptr) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion requires tensor source"); + } + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "index_tts2_5.qwen_emotion.weights", + weight_context_bytes); + + const auto & source = *assets.qwen_emotion_weights; + weights->token_embedding = weights->store->load_tensor( + source, + "model.embed_tokens.weight", + storage_type, + {kVocab, kHidden}); + weights->decoder.layers.reserve(static_cast(kLayers)); + for (int64_t layer = 0; layer < kLayers; ++layer) { + weights->decoder.layers.push_back(load_qwen_layer(*weights->store, source, layer, storage_type)); + } + weights->final_norm = binding::norm_weight_from_source(*weights->store, source, "model.norm", kHidden); + weights->store->upload(); + assets.qwen_emotion_weights->release_storage(); + return weights; +} + +IndexTTS25QwenEmotionTokenizer::IndexTTS25QwenEmotionTokenizer(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion tokenizer requires assets"); + } + engine::tokenizers::LlamaBpeTokenizerSpec spec; + spec.vocab_path = assets->resources.require_file("qwen_emotion_vocab"); + spec.merges_path = assets->resources.require_file("qwen_emotion_merges"); + spec.tokenizer_config_path = assets->resources.require_file("qwen_emotion_tokenizer_config"); + spec.tokenizer_json_path = assets->resources.require_file("qwen_emotion_tokenizer"); + spec.pre_type = engine::tokenizers::LlamaBpePreTokenizer::Qwen2; + tokenizer_ = engine::tokenizers::load_llama_bpe_tokenizer(spec); + if (const auto id = tokenizer_->find_token_id("<|endoftext|>"); id.has_value()) { + eos_token_id_ = *id; + } + if (const auto id = tokenizer_->find_token_id(""); id.has_value()) { + think_end_token_id_ = *id; + } +} + +std::vector IndexTTS25QwenEmotionTokenizer::encode_chat_prompt(const std::string & text) const { + std::vector ids = tokenizer_->encode("System: 文本情感分类", true); + ids.push_back(eos_token_id_); + std::string user_text = "\nHuman: " + text; + size_t trailing_spaces = 0; + while (trailing_spaces < user_text.size() && user_text[user_text.size() - trailing_spaces - 1] == ' ') { + ++trailing_spaces; + } + if (trailing_spaces > 0) { + user_text.resize(user_text.size() - trailing_spaces); + } + auto user = tokenizer_->encode(user_text, true); + ids.insert(ids.end(), user.begin(), user.end()); + if (trailing_spaces > 0) { + const auto space_token = tokenizer_->find_token_id("Ġ"); + if (!space_token.has_value()) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion tokenizer missing space token"); + } + ids.insert(ids.end(), trailing_spaces, *space_token); + } + ids.push_back(eos_token_id_); + auto assistant = tokenizer_->encode("\nAssistant:", true); + ids.insert(ids.end(), assistant.begin(), assistant.end()); + return ids; +} + +std::string IndexTTS25QwenEmotionTokenizer::decode(const std::vector & token_ids, bool skip_special_tokens) const { + return tokenizer_->decode(token_ids, skip_special_tokens); +} + +int32_t IndexTTS25QwenEmotionTokenizer::eos_token_id() const noexcept { + return eos_token_id_; +} + +int32_t IndexTTS25QwenEmotionTokenizer::think_end_token_id() const noexcept { + return think_end_token_id_; +} + +class IndexTTS25QwenEmotionRuntime::PrefillGraph { +public: + PrefillGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t prompt_steps, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + prompt_steps_(prompt_steps) { + if (prompt_steps_ <= 0) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion prefill graph requires prompt tokens"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion prefill graph context"); + } + ggml_init_params input_params{64ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion prefill input context"); + } + ggml_init_params output_params{16ull * 1024ull * 1024ull, nullptr, true}; + output_ctx_.reset(ggml_init(output_params)); + if (output_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion prefill output context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.qwen_emotion.prefill", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.qwen_emotion.prefill.inputs", + execution_.backend_type()}; + core::ModuleBuildContext output_ctx{ + output_ctx_.get(), + "index_tts2_5.qwen_emotion.prefill.outputs", + execution_.backend_type()}; + token_ids_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, prompt_steps_})).tensor; + positions_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({prompt_steps_})).tensor; + mask_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_steps_})).tensor; + ggml_set_input(token_ids_); + ggml_set_input(positions_); + ggml_set_input(mask_); + auto x = modules::EmbeddingModule({kVocab, kHidden}).build( + ctx, + core::wrap_tensor(token_ids_, core::TensorShape::from_dims({1, prompt_steps_}), GGML_TYPE_I32), + weights_->token_embedding); + auto outputs = modules::QwenDecoderStackModule(qwen_config()).build( + ctx, + x, + core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32), + weights_->decoder, + std::nullopt, + core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_steps_}), GGML_TYPE_F32)); + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(131072, prompt_steps_ * 4096)), false); + for (const auto & layer : outputs.state.layers) { + auto key = core::ensure_backend_addressable_layout(ctx, *layer.key); + auto value = core::ensure_backend_addressable_layout(ctx, *layer.value); + auto * key_output = core::make_tensor(output_ctx, GGML_TYPE_F32, key.shape).tensor; + auto * value_output = core::make_tensor(output_ctx, GGML_TYPE_F32, value.shape).tensor; + keys_.push_back(key_output); + values_.push_back(value_output); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), key.tensor, key_output)); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), value.tensor, value_output)); + } + auto last = modules::SliceModule({1, prompt_steps_ - 1, 1}).build(ctx, outputs.output); + last = modules::RMSNormModule({kHidden, kRmsEps, true, false}).build(ctx, last, weights_->final_norm); + auto logits = modules::LinearModule({kHidden, kVocab, false}).build(ctx, last, {weights_->token_embedding, std::nullopt}); + auto flat_logits = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, logits), + core::TensorShape::from_dims({1, kVocab})); + auto * next_token_source = ggml_argmax(ctx.ggml, flat_logits.tensor); + next_token_ = core::make_tensor(output_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1})).tensor; + ggml_set_output(next_token_); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), next_token_source, next_token_)); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion prefill input buffer"); + } + output_buffer_ = ggml_backend_alloc_ctx_tensors(output_ctx_.get(), execution_.backend()); + if (output_buffer_ == nullptr) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion prefill output buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion prefill graph"); + } + std::vector positions(static_cast(prompt_steps_)); + for (int64_t i = 0; i < prompt_steps_; ++i) { + positions[static_cast(i)] = static_cast(i); + } + core::write_tensor_i32(core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32), positions); + const auto mask = causal_mask(prompt_steps_); + core::write_tensor_f32(core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_steps_}), GGML_TYPE_F32), mask); + debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.prefill.prompt_tokens", prompt_steps_); + } + + ~PrefillGraph() { + clear_graph(); + } + + bool matches(const IndexTTS25QwenEmotionWeights & weights, ggml_backend_t backend, int64_t prompt_steps) const noexcept { + return weights_.get() == &weights && execution_.backend() == backend && prompt_steps_ == prompt_steps; + } + + int32_t run(const std::vector & prompt_ids) { + if (static_cast(prompt_ids.size()) != prompt_steps_) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion prompt length mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(token_ids_, prompt_ids.data(), 0, prompt_ids.size() * sizeof(int32_t)); + debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_, nullptr, "IndexTTS2.5 Qwen emotion prefill"); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion prefill graph compute failed"); + } + timing_start = Clock::now(); + const auto next_token = core::read_tensor_i32(next_token_); + debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (next_token.size() != 1) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion prefill argmax output shape mismatch"); + } + return next_token.front(); + } + + int64_t prompt_steps() const noexcept { + return prompt_steps_; + } + + size_t layer_count() const noexcept { + return keys_.size(); + } + + void copy_layer_state_to(size_t layer, ggml_tensor * key_destination, ggml_tensor * value_destination) const { + const int64_t values = prompt_steps_ * kKvHeads * kHeadDim; + std::vector key(static_cast(values)); + std::vector value(static_cast(values)); + ggml_backend_tensor_get(keys_.at(layer), key.data(), 0, key.size() * sizeof(float)); + ggml_backend_tensor_get(values_.at(layer), value.data(), 0, value.size() * sizeof(float)); + ggml_backend_tensor_set(key_destination, key.data(), 0, key.size() * sizeof(float)); + ggml_backend_tensor_set(value_destination, value.data(), 0, value.size() * sizeof(float)); + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + if (output_buffer_ != nullptr) { + ggml_backend_buffer_free(output_buffer_); + output_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t prompt_steps_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr output_ctx_; + std::unique_ptr ctx_; + ggml_tensor * token_ids_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * next_token_ = nullptr; + std::vector keys_; + std::vector values_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; + ggml_backend_buffer_t output_buffer_ = nullptr; +}; + +class IndexTTS25QwenEmotionRuntime::DecodeGraph { +public: + DecodeGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t cache_steps, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + cache_steps_(cache_steps) { + if (cache_steps_ <= 0) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion decode graph requires cache capacity"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion decode graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion decode input context"); + } + ggml_init_params state_params{128ull * 1024ull * 1024ull, nullptr, true}; + state_ctx_.reset(ggml_init(state_params)); + if (state_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion decode state context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.qwen_emotion.decode", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.qwen_emotion.decode.inputs", + execution_.backend_type()}; + core::ModuleBuildContext state_ctx{ + state_ctx_.get(), + "index_tts2_5.qwen_emotion.decode.state", + execution_.backend_type()}; + token_id_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, 1})).tensor; + position_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1})).tensor; + mask_ = core::make_tensor(input_ctx, GGML_TYPE_F16, core::TensorShape::from_dims({1, 1, 1, cache_steps_ + 1})).tensor; + ggml_set_input(token_id_); + ggml_set_input(position_); + ggml_set_input(mask_); + graph_ = ggml_new_graph_custom(ctx_.get(), 131072, false); + std::vector cache_keys; + std::vector cache_values; + auto x = modules::EmbeddingModule({kVocab, kHidden}).build( + ctx, + core::wrap_tensor(token_id_, core::TensorShape::from_dims({1, 1}), GGML_TYPE_I32), + weights_->token_embedding); + const auto cfg = qwen_config(); + const modules::QwenDecoderLayerModule layer_module(modules::qwen_decoder_layer_config_from_stack(cfg)); + const auto mask = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_ + 1}), GGML_TYPE_F16); + for (const auto & layer : weights_->decoder.layers) { + cache_keys.push_back(core::make_tensor( + state_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps_ + 1, kKvHeads, kHeadDim}))); + cache_values.push_back(core::make_tensor( + state_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps_ + 1, kKvHeads, kHeadDim}))); + auto out = layer_module.build_with_static_cache_tail( + ctx, + graph_, + x, + core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32), + layer, + cache_keys.back(), + cache_values.back(), + std::nullopt, + mask); + x = out.output; + } + kv_cache_ = engine::runtime::TransformerKVCache(cache_steps_ + 1, kKvHeads * kHeadDim, std::move(cache_keys), std::move(cache_values)); + build_transfer_views(); + x = modules::RMSNormModule({kHidden, kRmsEps, true, false}).build(ctx, x, weights_->final_norm); + auto logits = modules::LinearModule({kHidden, kVocab, false}).build(ctx, x, {weights_->token_embedding, std::nullopt}); + auto flat_logits = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, logits), + core::TensorShape::from_dims({1, kVocab})); + next_token_ = ggml_argmax(ctx.ggml, flat_logits.tensor); + ggml_set_output(next_token_); + ggml_build_forward_expand(graph_, next_token_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion decode input buffer"); + } + state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), execution_.backend()); + if (state_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion decode state buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion decode graph"); + } + mask_values_.assign(static_cast(cache_steps_ + 1), ggml_fp32_to_fp16(-std::numeric_limits::infinity())); + debug::timing_log_scalar("index_tts2_5.qwen_emotion.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.decode.cache_steps", cache_steps_); + } + + ~DecodeGraph() { + clear_graph(); + } + + bool can_run(const IndexTTS25QwenEmotionWeights & weights, ggml_backend_t backend, int64_t required_steps) const noexcept { + return weights_.get() == &weights && execution_.backend() == backend && cache_steps_ >= required_steps; + } + + void import_state(const PrefillGraph & prefill) { + if (prefill.layer_count() != key_sources_.size() || prefill.prompt_steps() > cache_steps_) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion decode prefill state shape mismatch"); + } + kv_cache_.retain_prefix(0); + const size_t prefix = static_cast(prefill.prompt_steps()); + for (size_t layer = 0; layer < key_sources_.size(); ++layer) { + prefill.copy_layer_state_to(layer, key_prefix_destinations_[prefix][layer], value_prefix_destinations_[prefix][layer]); + } + kv_cache_.advance_after_direct_append(prefill.prompt_steps()); + } + + int32_t run_step(int32_t token) { + if (kv_cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion decode cache exhausted"); + } + ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t)); + const int32_t position = static_cast(kv_cache_.current_end()); + ggml_backend_tensor_set(position_, &position, 0, sizeof(int32_t)); + std::fill(mask_values_.begin(), mask_values_.end(), ggml_fp32_to_fp16(-std::numeric_limits::infinity())); + for (int64_t i = 0; i < kv_cache_.valid_steps(); ++i) { + mask_values_[static_cast(i)] = ggml_fp32_to_fp16(0.0F); + } + mask_values_[static_cast(cache_steps_)] = ggml_fp32_to_fp16(0.0F); + ggml_backend_tensor_set(mask_, mask_values_.data(), 0, mask_values_.size() * sizeof(ggml_fp16_t)); + core::set_backend_threads(execution_.backend(), execution_.config().threads); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_, nullptr, "IndexTTS2.5 Qwen emotion decode"); + ggml_backend_synchronize(execution_.backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion decode graph compute failed"); + } + const auto next_token = core::read_tensor_i32(next_token_); + if (next_token.size() != 1) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion decode argmax output shape mismatch"); + } + const size_t dst_slot = static_cast(kv_cache_.valid_steps()); + for (size_t layer = 0; layer < key_sources_.size(); ++layer) { + ggml_backend_tensor_copy(key_sources_[layer], key_destinations_[dst_slot][layer]); + ggml_backend_tensor_copy(value_sources_[layer], value_destinations_[dst_slot][layer]); + } + kv_cache_.advance_after_direct_append(1); + return next_token.front(); + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (state_buffer_ != nullptr) { + ggml_backend_buffer_free(state_buffer_); + state_buffer_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + void build_transfer_views() { + const int64_t step_elems = kKvHeads * kHeadDim; + const size_t scratch_offset = static_cast(cache_steps_ * step_elems) * sizeof(float); + key_sources_.clear(); + value_sources_.clear(); + key_sources_.reserve(weights_->decoder.layers.size()); + value_sources_.reserve(weights_->decoder.layers.size()); + for (size_t layer = 0; layer < weights_->decoder.layers.size(); ++layer) { + key_sources_.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.key_tensor(layer).tensor, step_elems, scratch_offset)); + value_sources_.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.value_tensor(layer).tensor, step_elems, scratch_offset)); + } + key_destinations_.assign(static_cast(cache_steps_), {}); + value_destinations_.assign(static_cast(cache_steps_), {}); + key_prefix_destinations_.assign(static_cast(cache_steps_ + 1), {}); + value_prefix_destinations_.assign(static_cast(cache_steps_ + 1), {}); + for (int64_t slot = 0; slot < cache_steps_; ++slot) { + const size_t byte_offset = static_cast(slot * step_elems) * sizeof(float); + auto & key_slot = key_destinations_[static_cast(slot)]; + auto & value_slot = value_destinations_[static_cast(slot)]; + key_slot.reserve(key_sources_.size()); + value_slot.reserve(value_sources_.size()); + for (size_t layer = 0; layer < key_sources_.size(); ++layer) { + key_slot.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.key_tensor(layer).tensor, step_elems, byte_offset)); + value_slot.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.value_tensor(layer).tensor, step_elems, byte_offset)); + } + } + for (int64_t prefix = 1; prefix <= cache_steps_; ++prefix) { + auto & key_prefix = key_prefix_destinations_[static_cast(prefix)]; + auto & value_prefix = value_prefix_destinations_[static_cast(prefix)]; + key_prefix.reserve(key_sources_.size()); + value_prefix.reserve(value_sources_.size()); + for (size_t layer = 0; layer < key_sources_.size(); ++layer) { + key_prefix.push_back(ggml_view_1d( + state_ctx_.get(), + kv_cache_.key_tensor(layer).tensor, + prefix * step_elems, + 0)); + value_prefix.push_back(ggml_view_1d( + state_ctx_.get(), + kv_cache_.value_tensor(layer).tensor, + prefix * step_elems, + 0)); + } + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t cache_steps_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr state_ctx_; + std::unique_ptr ctx_; + ggml_tensor * token_id_ = nullptr; + ggml_tensor * position_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * next_token_ = nullptr; + std::vector key_sources_; + std::vector value_sources_; + std::vector> key_destinations_; + std::vector> value_destinations_; + std::vector> key_prefix_destinations_; + std::vector> value_prefix_destinations_; + std::vector mask_values_; + engine::runtime::TransformerKVCache kv_cache_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; + ggml_backend_buffer_t state_buffer_ = nullptr; +}; + +IndexTTS25QwenEmotionRuntime::IndexTTS25QwenEmotionRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType storage_type) + : assets_(std::move(assets)), + execution_(&execution), + prefill_graph_arena_bytes_(prefill_graph_arena_bytes), + decode_graph_arena_bytes_(decode_graph_arena_bytes), + tokenizer_(assets_) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion runtime requires assets"); + } + if (prefill_graph_arena_bytes_ == 0 || decode_graph_arena_bytes_ == 0) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion graph arenas must be non-zero"); + } + weights_ = load_index_tts2_5_qwen_emotion_weights( + *assets_, + execution.backend(), + execution.backend_type(), + storage_type, + weight_context_bytes); +} + +IndexTTS25QwenEmotionRuntime::~IndexTTS25QwenEmotionRuntime() = default; + +IndexTTS25EmotionVector IndexTTS25QwenEmotionRuntime::infer(const std::string & text, int64_t max_new_tokens) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion runtime execution context is missing"); + } + if (max_new_tokens <= 0) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion max_new_tokens must be positive"); + } + const auto prompt_ids = tokenizer_.encode_chat_prompt(text); + const int64_t prompt_steps = static_cast(prompt_ids.size()); + if (prefill_graph_ == nullptr || !prefill_graph_->matches(*weights_, execution_->backend(), prompt_steps)) { + prefill_graph_.reset(); + prefill_graph_ = std::make_unique(*execution_, weights_, prompt_steps, prefill_graph_arena_bytes_); + } + const int64_t required_cache_steps = prompt_steps + max_new_tokens; + if (decode_graph_ == nullptr || !decode_graph_->can_run(*weights_, execution_->backend(), required_cache_steps)) { + decode_graph_.reset(); + decode_graph_ = std::make_unique(*execution_, weights_, required_cache_steps, decode_graph_arena_bytes_); + } + const int32_t prefill_token = prefill_graph_->run(prompt_ids); + decode_graph_->import_state(*prefill_graph_); + + std::vector generated; + generated.reserve(static_cast(max_new_tokens)); + int32_t token = prefill_token; + double decode_run_ms = 0.0; + bool saw_eos = false; + for (int64_t step = 0; step < max_new_tokens; ++step) { + if (token == tokenizer_.eos_token_id()) { + saw_eos = true; + break; + } + generated.push_back(token); + if (step + 1 >= max_new_tokens) { + break; + } + const auto decode_start = Clock::now(); + token = decode_graph_->run_step(token); + decode_run_ms += engine::debug::elapsed_ms(decode_start, Clock::now()); + } + if (!saw_eos && static_cast(generated.size()) >= max_new_tokens) { + throw std::runtime_error("IndexTTS2.5 Qwen emotion decode reached max_new_tokens before EOS"); + } + + size_t start = 0; + for (size_t i = generated.size(); i > 0; --i) { + if (generated[i - 1] == tokenizer_.think_end_token_id()) { + start = i; + break; + } + } + const std::vector answer(generated.begin() + static_cast(start), generated.end()); + const std::string content = tokenizer_.decode(answer, true); + debug::timing_log_scalar("index_tts2_5.qwen_emotion.decode.run_ms", decode_run_ms); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.generated_tokens", generated.size()); + return convert_emotion_json(content, text); +} + +void IndexTTS25QwenEmotionRuntime::release_graphs() { + prefill_graph_.reset(); + decode_graph_.reset(); +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/request.cpp b/src/models/index_tts2_5/request.cpp new file mode 100644 index 00000000..9c792973 --- /dev/null +++ b/src/models/index_tts2_5/request.cpp @@ -0,0 +1,192 @@ +#include "engine/models/index_tts2_5/request.h" + +#include "engine/framework/io/text.h" +#include "engine/framework/runtime/options.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +std::vector parse_emotion_vector(const std::string & value) { + std::string normalized = engine::io::trim_ascii_whitespace(value); + if (!normalized.empty() && normalized.front() == '[') { + normalized.erase(normalized.begin()); + } + if (!normalized.empty() && normalized.back() == ']') { + normalized.pop_back(); + } + std::vector values; + std::stringstream stream(normalized); + std::string item; + while (std::getline(stream, item, ',')) { + item = engine::io::trim_ascii_whitespace(item); + if (item.empty()) { + throw std::runtime_error("IndexTTS2.5 emotion_vector contains an empty item"); + } + size_t parsed = 0; + const float parsed_value = std::stof(item, &parsed); + if (parsed != item.size() || !std::isfinite(parsed_value)) { + throw std::runtime_error("IndexTTS2.5 emotion_vector must contain finite floats"); + } + values.push_back(parsed_value); + } + if (values.size() != 8) { + throw std::runtime_error("IndexTTS2.5 emotion_vector must contain exactly 8 values"); + } + return values; +} + +const runtime::AudioBuffer * speaker_audio_from_request(const runtime::TaskRequest & request) { + if (request.voice.has_value() && + request.voice->speaker.has_value() && + request.voice->speaker->audio.has_value()) { + return &*request.voice->speaker->audio; + } + return nullptr; +} + +void require_valid_audio(const runtime::AudioBuffer & audio, const char * label) { + if (audio.sample_rate <= 0 || audio.channels <= 0 || audio.samples.empty()) { + throw std::runtime_error(std::string("IndexTTS2.5 ") + label + " audio is empty or invalid"); + } +} + +} // namespace + +std::string normalize_index_tts2_5_lang(const std::string & value) { + std::string lang = engine::io::trim_ascii_whitespace(value); + std::transform(lang.begin(), lang.end(), lang.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + if (lang == "auto") { + lang.clear(); + } + return lang; +} + +IndexTTS25Request parse_index_tts2_5_request(const runtime::TaskRequest & request) { + IndexTTS25Request out; + if (request.text_input.has_value()) { + out.text = engine::io::trim_ascii_whitespace(request.text_input->text); + } else if (const auto value = runtime::find_option(request.options, {"text", "prompt"})) { + out.text = engine::io::trim_ascii_whitespace(*value); + } + if (out.text.empty()) { + throw std::runtime_error("IndexTTS2.5 request requires text_input or text option"); + } + + if (const auto * speaker = speaker_audio_from_request(request)) { + require_valid_audio(*speaker, "speaker reference"); + out.speaker_audio = *speaker; + } else { + throw std::runtime_error("IndexTTS2.5 request requires --voice-ref or voice.speaker.audio"); + } + + if (const auto value = runtime::find_option(request.options, {"lang"})) { + out.lang = normalize_index_tts2_5_lang(*value); + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"emotion_alpha"})) { + if (*value < 0.0F || *value > 1.0F) { + throw std::runtime_error("IndexTTS2.5 emotion_alpha must be in [0, 1]"); + } + out.emotion_alpha = *value; + } + if (const auto value = runtime::find_option(request.options, {"emotion_vector"})) { + out.emotion_vector = parse_emotion_vector(*value); + } + const auto use_emotion_text_value = runtime::find_option(request.options, {"use_emotion_text"}); + if (use_emotion_text_value.has_value()) { + out.use_emotion_text = runtime::parse_bool_option(*use_emotion_text_value, "use_emotion_text"); + } + if (const auto value = runtime::find_option(request.options, {"emotion_text"})) { + if (!engine::io::trim_ascii_whitespace(*value).empty()) { + out.emotion_text = *value; + } + } + if (request.voice.has_value() && + request.voice->style.has_value() && + request.voice->style->emotion.has_value()) { + const auto & text = *request.voice->style->emotion; + if (!engine::io::trim_ascii_whitespace(text).empty()) { + if (use_emotion_text_value.has_value() && !out.use_emotion_text) { + throw std::runtime_error("IndexTTS2.5 --emotion conflicts with use_emotion_text=false"); + } + if (out.emotion_text.has_value() && *out.emotion_text != text) { + throw std::runtime_error("IndexTTS2.5 --emotion conflicts with emotion_text"); + } + out.use_emotion_text = true; + out.emotion_text = text; + } + } + if (const auto value = runtime::find_option(request.options, {"use_random_emotion"})) { + out.use_random_emotion = runtime::parse_bool_option(*value, "use_random_emotion"); + } + if (const auto value = runtime::parse_int_option(request.options, {"interval_silence_ms"})) { + if (*value < 0) { + throw std::runtime_error("IndexTTS2.5 interval_silence_ms must be non-negative"); + } + out.interval_silence_ms = *value; + } + if (request.audio_input.has_value()) { + require_valid_audio(*request.audio_input, "emotion reference"); + out.emotion_audio = request.audio_input; + } + if (out.emotion_vector.has_value() || out.use_emotion_text) { + out.emotion_audio = std::nullopt; + } + + if (const auto value = runtime::find_option(request.options, {"do_sample"})) { + out.generation.do_sample = runtime::parse_bool_option(*value, "do_sample"); + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"top_p"})) { + out.generation.top_p = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"top_k"})) { + out.generation.top_k = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"temperature"})) { + out.generation.temperature = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"length_penalty"})) { + out.generation.length_penalty = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"num_beams"})) { + out.generation.num_beams = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"repetition_penalty"})) { + out.generation.repetition_penalty = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { + if (*value <= 0) { + throw std::runtime_error("IndexTTS2.5 max_tokens must be positive"); + } + out.generation.max_mel_tokens = *value; + } + if (const auto value = runtime::parse_u32_option(request.options, {"seed"})) { + out.generation.seed = *value; + } else { + out.generation.seed = runtime::random_u32_seed(); + } + + if (out.generation.top_k <= 0) { + throw std::runtime_error("IndexTTS2.5 top_k must be positive"); + } + if (!(out.generation.top_p > 0.0F && out.generation.top_p <= 1.0F)) { + throw std::runtime_error("IndexTTS2.5 top_p must be in (0, 1]"); + } + if (!(out.generation.temperature > 0.0F)) { + throw std::runtime_error("IndexTTS2.5 temperature must be positive"); + } + if (out.generation.num_beams <= 0) { + throw std::runtime_error("IndexTTS2.5 num_beams must be positive"); + } + return out; +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/s2mel.cpp b/src/models/index_tts2_5/s2mel.cpp new file mode 100644 index 00000000..d6890b00 --- /dev/null +++ b/src/models/index_tts2_5/s2mel.cpp @@ -0,0 +1,1271 @@ +#include "engine/models/index_tts2_5/s2mel.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/torch_random.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kMelChannels = 80; +constexpr int64_t kContentDim = 1024; +constexpr int64_t kGptDim = 1280; +constexpr int64_t kHidden = 512; +constexpr int64_t kStyleDim = 192; +constexpr int64_t kDitLayers = 13; +constexpr int64_t kWavenetLayers = 8; +constexpr int64_t kWavenetKernel = 5; +constexpr int64_t kTimeFreqDim = 128; +constexpr int64_t kTimeEmbeddingDim = 256; +constexpr int64_t kDitFfnDim = 1536; +constexpr int64_t kDitHeads = 8; +constexpr int64_t kDitHeadDim = kHidden / kDitHeads; +constexpr float kLayerNormEps = 1.0e-6F; +constexpr float kRmsNormEps = 1.0e-5F; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +core::TensorValue sub(core::ModuleBuildContext & ctx, const core::TensorValue & lhs, const core::TensorValue & rhs) { + core::validate_shape(rhs, lhs.shape, "Sub rhs"); + return core::wrap_tensor(ggml_sub(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32); +} + +core::TensorValue scale(core::ModuleBuildContext & ctx, const core::TensorValue & input, float value) { + return core::wrap_tensor(ggml_scale(ctx.ggml, input.tensor, value), input.shape, GGML_TYPE_F32); +} + +core::TensorValue add_one(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + return core::wrap_tensor(ggml_scale_bias(ctx.ggml, input.tensor, 1.0F, 1.0F), input.shape, GGML_TYPE_F32); +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t dim) { + return core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, input), + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); +} + +core::TensorValue slice_last( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t start, + int64_t length) { + return modules::SliceModule({static_cast(input.shape.rank - 1), start, length}).build(ctx, input); +} + +core::TensorValue apply_channel_affine( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & weight, + const core::TensorValue & bias, + int64_t channels) { + core::TensorShape broadcast_shape = {}; + broadcast_shape.rank = input.shape.rank; + for (size_t axis = 0; axis < broadcast_shape.rank; ++axis) { + broadcast_shape.dims[axis] = 1; + } + broadcast_shape.dims[1] = channels; + auto weight_view = core::reshape_tensor(ctx, weight, broadcast_shape); + auto bias_view = core::reshape_tensor(ctx, bias, broadcast_shape); + auto weight_rep = modules::RepeatModule({input.shape}).build(ctx, weight_view); + auto bias_rep = modules::RepeatModule({input.shape}).build(ctx, bias_view); + return modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, input, weight_rep), bias_rep); +} + +core::TensorValue broadcast_batch_time( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t batch, + int64_t frames, + int64_t dims) { + auto shaped = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, input), core::TensorShape::from_dims({batch, 1, dims})); + return modules::RepeatModule({core::TensorShape::from_dims({batch, frames, dims})}).build(ctx, shaped); +} + +core::TensorValue group_norm_1_group( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::NormWeights & weights, + int64_t channels) { + if (!weights.weight.has_value() || !weights.bias.has_value()) { + throw std::runtime_error("IndexTTS2.5 S2Mel length regulator group norm requires affine weights"); + } + const auto input4 = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, input), + core::TensorShape::from_dims({input.shape.dims[0], channels, 1, input.shape.dims[2]})); + auto normalized = core::wrap_tensor(ggml_group_norm(ctx.ggml, input4.tensor, 1, 1.0e-5F), input4.shape, GGML_TYPE_F32); + normalized = apply_channel_affine(ctx, normalized, *weights.weight, *weights.bias, channels); + return core::reshape_tensor(ctx, normalized, input.shape); +} + +core::TensorValue mish(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + const auto softplus = core::wrap_tensor(ggml_softplus(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); + const auto tanh = core::wrap_tensor(ggml_tanh(ctx.ggml, softplus.tensor), input.shape, GGML_TYPE_F32); + return core::wrap_tensor(ggml_mul(ctx.ggml, input.tensor, tanh.tensor), input.shape, GGML_TYPE_F32); +} + +core::TensorValue timestep_embedding( + core::ModuleBuildContext & ctx, + const core::TensorValue & timestep, + const core::TensorValue & freqs, + const modules::LinearWeights & linear0, + const modules::LinearWeights & linear2) { + const int64_t batch = timestep.shape.dims[0]; + auto t = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, timestep), core::TensorShape::from_dims({batch, 1})); + auto freqs_batched = modules::RepeatModule({core::TensorShape::from_dims({batch, kTimeFreqDim})}) + .build(ctx, core::reshape_tensor(ctx, freqs, core::TensorShape::from_dims({1, kTimeFreqDim}))); + auto args = modules::MulModule{}.build(ctx, modules::RepeatModule({freqs_batched.shape}).build(ctx, t), freqs_batched); + args = scale(ctx, args, 1000.0F); + auto cos_part = core::wrap_tensor(ggml_cos(ctx.ggml, core::ensure_backend_addressable_layout(ctx, args).tensor), args.shape, GGML_TYPE_F32); + auto sin_part = core::wrap_tensor(ggml_sin(ctx.ggml, args.tensor), args.shape, GGML_TYPE_F32); + auto emb = modules::ConcatModule({1}).build(ctx, cos_part, sin_part); + emb = modules::LinearModule({kTimeEmbeddingDim, kHidden, true, GGML_PREC_F32}).build(ctx, emb, linear0); + emb = modules::SiluModule{}.build(ctx, emb); + return modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, emb, linear2); +} + +core::TensorValue adaptive_rms_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & embedding, + const IndexTTS25AdaLayerNormWeights & weights) { + auto projected = modules::LinearModule({kHidden, 2 * kHidden, true, GGML_PREC_F32}).build(ctx, embedding, weights.project); + auto weight = broadcast_batch_time(ctx, slice_last(ctx, projected, 0, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); + auto bias = broadcast_batch_time(ctx, slice_last(ctx, projected, kHidden, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); + auto normed = modules::RMSNormModule({kHidden, kRmsNormEps, true, false}).build(ctx, input, {weights.norm_weight, std::nullopt}); + return modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, normed, weight), bias); +} + +core::TensorValue cfm_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const IndexTTS25DitLayerWeights & weights) { + auto qkv = modules::LinearModule({kHidden, 3 * kHidden, false, GGML_PREC_F32}).build(ctx, input, weights.qkv); + auto q = slice_last(ctx, qkv, 0, kHidden); + auto k = slice_last(ctx, qkv, kHidden, kHidden); + auto v = slice_last(ctx, qkv, 2 * kHidden, kHidden); + q = modules::RoPEModule({kDitHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, reshape_heads(ctx, q, kDitHeads, kDitHeadDim), positions); + k = modules::RoPEModule({kDitHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, reshape_heads(ctx, k, kDitHeads, kDitHeadDim), positions); + v = reshape_heads(ctx, v, kDitHeads, kDitHeadDim); + auto qh = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, q); + auto kh = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k); + auto vh = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v); + auto * flash = ggml_flash_attn_ext( + ctx.ggml, + core::ensure_backend_addressable_layout(ctx, qh).tensor, + core::ensure_backend_addressable_layout(ctx, kh).tensor, + core::ensure_backend_addressable_layout(ctx, vh).tensor, + nullptr, + 1.0F / std::sqrt(static_cast(kDitHeadDim)), + 0.0F, + 0.0F); + ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); + auto context = core::wrap_tensor( + flash, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], kDitHeads, kDitHeadDim}), + GGML_TYPE_F32); + context = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, context), input.shape); + return modules::LinearModule({kHidden, kHidden, false, GGML_PREC_F32}).build(ctx, context, weights.attention_out); +} + +core::TensorValue cfm_ffn( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const IndexTTS25DitLayerWeights & weights) { + auto gate = modules::LinearModule({kHidden, kDitFfnDim, false, GGML_PREC_F32}).build(ctx, input, weights.ffn_w1); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule({kHidden, kDitFfnDim, false, GGML_PREC_F32}).build(ctx, input, weights.ffn_w3); + auto hidden = modules::MulModule{}.build(ctx, gate, up); + return modules::LinearModule({kDitFfnDim, kHidden, false, GGML_PREC_F32}).build(ctx, hidden, weights.ffn_w2); +} + +core::TensorValue cfm_transformer_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & timestep, + const core::TensorValue & positions, + const IndexTTS25DitLayerWeights & weights, + const core::TensorValue * skip) { + auto x = input; + if (skip != nullptr) { + x = modules::LinearModule({2 * kHidden, kHidden, true, GGML_PREC_F32}) + .build(ctx, modules::ConcatModule({2}).build(ctx, x, *skip), weights.skip_in); + } + auto attn = cfm_attention(ctx, adaptive_rms_norm(ctx, x, timestep, weights.attention_norm), positions, weights); + auto h = modules::AddModule{}.build(ctx, x, attn); + auto ff = cfm_ffn(ctx, adaptive_rms_norm(ctx, h, timestep, weights.ffn_norm), weights); + return modules::AddModule{}.build(ctx, h, ff); +} + +core::TensorValue cfm_wavenet( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const core::TensorValue & timestep_b, + const IndexTTS25S2MelCfmWeights & weights) { + auto g = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, timestep_b), core::TensorShape::from_dims({timestep_b.shape.dims[0], kHidden, 1})); + g = modules::Conv1dModule({kHidden, 2 * kHidden * kWavenetLayers, 1, 1, 0, 1, true}).build(ctx, g, weights.wavenet_cond); + auto output = sub(ctx, input_bct, input_bct); + auto x = input_bct; + for (int64_t i = 0; i < kWavenetLayers; ++i) { + const int64_t dilation = 1; + const int64_t padding = (kWavenetKernel * dilation - dilation) / 2; + auto x_padded = modules::ReflectPad1dModule({padding, padding}).build(ctx, core::ensure_backend_addressable_layout(ctx, x)); + auto x_in = modules::Conv1dModule( + {kHidden, 2 * kHidden, kWavenetKernel, 1, 0, static_cast(dilation), true}) + .build(ctx, x_padded, weights.wavenet_layers[static_cast(i)].in_layer); + auto g_l = modules::SliceModule({1, i * 2 * kHidden, 2 * kHidden}).build(ctx, g); + g_l = modules::RepeatModule({x_in.shape}).build(ctx, g_l); + auto acts = modules::AddModule{}.build(ctx, x_in, g_l); + auto tanh_part = modules::SliceModule({1, 0, kHidden}).build(ctx, acts); + tanh_part = modules::TanhModule{}.build(ctx, tanh_part); + auto sigmoid_part = modules::SliceModule({1, kHidden, kHidden}).build(ctx, acts); + sigmoid_part = modules::SigmoidModule{}.build(ctx, sigmoid_part); + acts = modules::MulModule{}.build(ctx, tanh_part, sigmoid_part); + const int64_t res_skip_channels = i < kWavenetLayers - 1 ? 2 * kHidden : kHidden; + auto res_skip = modules::Conv1dModule({kHidden, res_skip_channels, 1, 1, 0, 1, true}) + .build(ctx, acts, weights.wavenet_layers[static_cast(i)].res_skip_layer); + if (i < kWavenetLayers - 1) { + auto res = modules::SliceModule({1, 0, kHidden}).build(ctx, res_skip); + auto skip = modules::SliceModule({1, kHidden, kHidden}).build(ctx, res_skip); + x = modules::AddModule{}.build(ctx, x, res); + output = modules::AddModule{}.build(ctx, output, skip); + } else { + output = modules::AddModule{}.build(ctx, output, res_skip); + } + } + return output; +} + +core::TensorValue cfm_final_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & timestep, + const IndexTTS25S2MelCfmWeights & weights) { + auto mod = modules::SiluModule{}.build(ctx, timestep); + mod = modules::LinearModule({kHidden, 2 * kHidden, true, GGML_PREC_F32}).build(ctx, mod, weights.final_modulation); + auto shift = broadcast_batch_time(ctx, slice_last(ctx, mod, 0, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); + auto scale_v = broadcast_batch_time(ctx, slice_last(ctx, mod, kHidden, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); + auto normed = modules::LayerNormModule({kHidden, kLayerNormEps, false, false}).build(ctx, input, {std::nullopt, std::nullopt}); + normed = modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, normed, add_one(ctx, scale_v)), shift); + return modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, normed, weights.final_linear); +} + +core::TensorValue build_cfm_estimator( + core::ModuleBuildContext & ctx, + const core::TensorValue & x_bct, + const core::TensorValue & prompt_bct, + const core::TensorValue & cond_btc, + const core::TensorValue & style_bc, + const core::TensorValue & timestep_b, + const core::TensorValue & positions, + const IndexTTS25S2MelCfmWeights & weights) { + const int64_t batch = x_bct.shape.dims[0]; + const int64_t frames = x_bct.shape.dims[2]; + auto t1 = timestep_embedding(ctx, timestep_b, weights.time_freqs, weights.time_mlp0, weights.time_mlp2); + auto cond = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, cond_btc, weights.cond_projection); + auto x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x_bct); + auto prompt = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, prompt_bct); + auto style = broadcast_batch_time(ctx, style_bc, batch, frames, kStyleDim); + auto hidden = modules::ConcatModule({2}).build(ctx, x, prompt); + hidden = modules::ConcatModule({2}).build(ctx, hidden, cond); + hidden = modules::ConcatModule({2}).build(ctx, hidden, style); + hidden = modules::LinearModule({kHidden + 2 * kMelChannels + kStyleDim, kHidden, true, GGML_PREC_F32}) + .build(ctx, hidden, weights.cond_x_merge); + + std::vector skips; + skips.reserve(static_cast(kDitLayers / 2)); + for (int64_t i = 0; i < kDitLayers; ++i) { + const core::TensorValue * skip = nullptr; + if (i > kDitLayers / 2) { + skip = &skips.back(); + } + hidden = cfm_transformer_layer(ctx, hidden, t1, positions, weights.dit_layers[static_cast(i)], skip); + if (i > kDitLayers / 2) { + skips.pop_back(); + } else if (i < kDitLayers / 2) { + skips.push_back(hidden); + } + } + hidden = adaptive_rms_norm(ctx, hidden, t1, weights.dit_norm); + hidden = modules::LinearModule({kHidden + kMelChannels, kHidden, true, GGML_PREC_F32}) + .build(ctx, modules::ConcatModule({2}).build(ctx, hidden, x), weights.skip_linear); + auto wavenet_x = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, weights.conv1); + wavenet_x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, wavenet_x); + auto t2 = timestep_embedding(ctx, timestep_b, weights.time2_freqs, weights.time2_mlp0, weights.time2_mlp2); + wavenet_x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, cfm_wavenet(ctx, wavenet_x, t2, weights)); + auto projected = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, weights.res_projection); + hidden = modules::AddModule{}.build(ctx, wavenet_x, projected); + hidden = cfm_final_layer(ctx, hidden, t1, weights); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, hidden); + return modules::Conv1dModule({kHidden, kMelChannels, 1, 1, 0, 1, true}).build(ctx, hidden, weights.conv2); +} + +std::vector fuse_weight_norm_linear( + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t out_features, + int64_t in_features) { + const auto g = source.require_f32(prefix + ".weight_g", {out_features, 1}); + const auto v = source.require_f32(prefix + ".weight_v", {out_features, in_features}); + std::vector weight(v.size(), 0.0F); + for (int64_t out = 0; out < out_features; ++out) { + double norm = 0.0; + for (int64_t in = 0; in < in_features; ++in) { + const float value = v[static_cast(out * in_features + in)]; + norm += static_cast(value) * static_cast(value); + } + const float scale = g[static_cast(out)] / static_cast(std::sqrt(norm)); + for (int64_t in = 0; in < in_features; ++in) { + const size_t index = static_cast(out * in_features + in); + weight[index] = v[index] * scale; + } + } + return weight; +} + +std::vector fuse_weight_norm_conv1d( + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size) { + const auto g = source.require_f32(prefix + ".weight_g", {out_channels, 1, 1}); + const auto v = source.require_f32(prefix + ".weight_v", {out_channels, in_channels, kernel_size}); + std::vector weight(v.size(), 0.0F); + for (int64_t out = 0; out < out_channels; ++out) { + double norm = 0.0; + for (int64_t in = 0; in < in_channels; ++in) { + for (int64_t k = 0; k < kernel_size; ++k) { + const float value = v[static_cast((out * in_channels + in) * kernel_size + k)]; + norm += static_cast(value) * static_cast(value); + } + } + const float scale = g[static_cast(out)] / static_cast(std::sqrt(norm)); + for (int64_t in = 0; in < in_channels; ++in) { + for (int64_t k = 0; k < kernel_size; ++k) { + const size_t index = static_cast((out * in_channels + in) * kernel_size + k); + weight[index] = v[index] * scale; + } + } + } + return weight; +} + +engine::modules::LinearWeights load_weight_norm_linear( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type, + int64_t out_features, + int64_t in_features) { + engine::modules::LinearWeights weights; + weights.weight = store.make_from_f32( + engine::core::TensorShape::from_dims({out_features, in_features}), + storage_type, + fuse_weight_norm_linear(source, prefix, out_features, in_features)); + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_features}); + return weights; +} + +engine::modules::Conv1dWeights load_weight_norm_conv1d( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size) { + engine::modules::Conv1dWeights weights; + weights.weight = store.make_from_f32( + engine::core::TensorShape::from_dims({out_channels, in_channels, kernel_size}), + storage_type, + fuse_weight_norm_conv1d(source, prefix, out_channels, in_channels, kernel_size)); + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); + return weights; +} + +IndexTTS25AdaLayerNormWeights load_ada_norm( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type) { + return { + store.load_f32_tensor(source, prefix + ".norm.weight", {kHidden}), + binding::linear_from_source(store, source, prefix + ".project_layer", storage_type, 2 * kHidden, kHidden, true), + }; +} + +IndexTTS25DitLayerWeights load_dit_layer( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + int64_t layer_index, + engine::assets::TensorStorageType storage_type) { + const std::string prefix = "cfm.estimator.transformer.layers." + std::to_string(layer_index); + IndexTTS25DitLayerWeights layer; + layer.attention_norm = load_ada_norm(store, source, prefix + ".attention_norm", storage_type); + layer.qkv = binding::linear_from_source(store, source, prefix + ".attention.wqkv", storage_type, 3 * kHidden, kHidden, false); + layer.attention_out = binding::linear_from_source(store, source, prefix + ".attention.wo", storage_type, kHidden, kHidden, false); + layer.ffn_norm = load_ada_norm(store, source, prefix + ".ffn_norm", storage_type); + layer.ffn_w1 = binding::linear_from_source(store, source, prefix + ".feed_forward.w1", storage_type, kDitFfnDim, kHidden, false); + layer.ffn_w2 = binding::linear_from_source(store, source, prefix + ".feed_forward.w2", storage_type, kHidden, kDitFfnDim, false); + layer.ffn_w3 = binding::linear_from_source(store, source, prefix + ".feed_forward.w3", storage_type, kDitFfnDim, kHidden, false); + layer.skip_in = binding::linear_from_source(store, source, prefix + ".skip_in_linear", storage_type, kHidden, 2 * kHidden, true); + return layer; +} + +IndexTTS25LengthRegulatorWeights load_length_regulator( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) { + IndexTTS25LengthRegulatorWeights weights; + weights.content_projection = binding::linear_from_source( + store, + source, + "length_regulator.content_in_proj", + matmul_storage_type, + kHidden, + kContentDim, + true); + for (int64_t i : {0, 3, 6, 9}) { + weights.convs.push_back(binding::conv1d_from_source( + store, + source, + "length_regulator.model." + std::to_string(i), + conv_storage_type, + kHidden, + kHidden, + 3, + true)); + } + for (int64_t i : {1, 4, 7, 10}) { + weights.norms.push_back({ + store.load_f32_tensor(source, "length_regulator.model." + std::to_string(i) + ".weight", {kHidden}), + store.load_f32_tensor(source, "length_regulator.model." + std::to_string(i) + ".bias", {kHidden}), + }); + } + weights.output = binding::conv1d_from_source( + store, + source, + "length_regulator.model.12", + conv_storage_type, + kHidden, + kHidden, + 1, + true); + return weights; +} + +IndexTTS25S2MelCfmWeights load_cfm( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) { + IndexTTS25S2MelCfmWeights weights; + weights.x_embedder = load_weight_norm_linear(store, source, "cfm.estimator.x_embedder", matmul_storage_type, kHidden, kMelChannels); + weights.cond_projection = binding::linear_from_source(store, source, "cfm.estimator.cond_projection", matmul_storage_type, kHidden, kHidden, true); + weights.cond_x_merge = binding::linear_from_source( + store, + source, + "cfm.estimator.cond_x_merge_linear", + matmul_storage_type, + kHidden, + kHidden + 2 * kMelChannels + kStyleDim, + true); + weights.skip_linear = binding::linear_from_source( + store, + source, + "cfm.estimator.skip_linear", + matmul_storage_type, + kHidden, + kHidden + kMelChannels, + true); + weights.time_freqs = store.load_f32_tensor(source, "cfm.estimator.t_embedder.freqs", {kTimeFreqDim}); + weights.time_mlp0 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder.mlp.0", matmul_storage_type, kHidden, kTimeEmbeddingDim, true); + weights.time_mlp2 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder.mlp.2", matmul_storage_type, kHidden, kHidden, true); + weights.time2_freqs = store.load_f32_tensor(source, "cfm.estimator.t_embedder2.freqs", {kTimeFreqDim}); + weights.time2_mlp0 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder2.mlp.0", matmul_storage_type, kHidden, kTimeEmbeddingDim, true); + weights.time2_mlp2 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder2.mlp.2", matmul_storage_type, kHidden, kHidden, true); + weights.dit_layers.reserve(static_cast(kDitLayers)); + for (int64_t i = 0; i < kDitLayers; ++i) { + weights.dit_layers.push_back(load_dit_layer(store, source, i, matmul_storage_type)); + } + weights.dit_norm = load_ada_norm(store, source, "cfm.estimator.transformer.norm", matmul_storage_type); + weights.conv1 = binding::linear_from_source(store, source, "cfm.estimator.conv1", matmul_storage_type, kHidden, kHidden, true); + weights.res_projection = binding::linear_from_source(store, source, "cfm.estimator.res_projection", matmul_storage_type, kHidden, kHidden, true); + weights.wavenet_cond = load_weight_norm_conv1d( + store, + source, + "cfm.estimator.wavenet.cond_layer.conv.conv", + conv_storage_type, + 2 * kHidden * kWavenetLayers, + kHidden, + 1); + weights.wavenet_layers.reserve(static_cast(kWavenetLayers)); + for (int64_t i = 0; i < kWavenetLayers; ++i) { + const int64_t res_skip_channels = i < kWavenetLayers - 1 ? 2 * kHidden : kHidden; + weights.wavenet_layers.push_back({ + load_weight_norm_conv1d( + store, + source, + "cfm.estimator.wavenet.in_layers." + std::to_string(i) + ".conv.conv", + conv_storage_type, + 2 * kHidden, + kHidden, + kWavenetKernel), + load_weight_norm_conv1d( + store, + source, + "cfm.estimator.wavenet.res_skip_layers." + std::to_string(i) + ".conv.conv", + conv_storage_type, + res_skip_channels, + kHidden, + 1), + }); + } + weights.final_modulation = binding::linear_from_source( + store, + source, + "cfm.estimator.final_layer.adaLN_modulation.1", + matmul_storage_type, + 2 * kHidden, + kHidden, + true); + weights.final_linear = load_weight_norm_linear( + store, + source, + "cfm.estimator.final_layer.linear", + matmul_storage_type, + kHidden, + kHidden); + weights.conv2 = binding::conv1d_from_source(store, source, "cfm.estimator.conv2", conv_storage_type, kMelChannels, kHidden, 1, true); + return weights; +} + +} // namespace + +std::shared_ptr load_index_tts2_5_s2mel_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes) { + if (assets.s2mel_weights == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel requires tensor source"); + } + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "index_tts2_5.s2mel.weights", + weight_context_bytes); + + const auto & source = *assets.s2mel_weights; + weights->gpt_layer.linear0 = binding::linear_from_source(*weights->store, source, "gpt_layer.0", matmul_storage_type, 256, 1280, true); + weights->gpt_layer.linear1 = binding::linear_from_source(*weights->store, source, "gpt_layer.1", matmul_storage_type, 128, 256, true); + weights->gpt_layer.linear2 = binding::linear_from_source(*weights->store, source, "gpt_layer.2", matmul_storage_type, kContentDim, 128, true); + weights->length_regulator = load_length_regulator(*weights->store, source, matmul_storage_type, conv_storage_type); + weights->cfm = load_cfm(*weights->store, source, matmul_storage_type, conv_storage_type); + + weights->store->upload(); + assets.s2mel_weights->release_storage(); + return weights; +} + +class IndexTTS25S2MelRuntime::GptLayerGraph { +public: + GptLayerGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t frames, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + frames_(frames) { + if (frames_ <= 0) { + throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph requires positive frame count"); + } + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph requires weights"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel GPT layer graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel GPT layer input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.s2mel.gpt_layer", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.s2mel.gpt_layer.inputs", + execution_.backend_type()}; + input_ = + core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kGptDim})).tensor; + ggml_set_input(input_); + auto x = core::wrap_tensor(input_, core::TensorShape::from_dims({1, frames_, kGptDim}), GGML_TYPE_F32); + x = modules::LinearModule({kGptDim, 256, true}).build(ctx, x, weights_->gpt_layer.linear0); + x = modules::LinearModule({256, 128, true}).build(ctx, x, weights_->gpt_layer.linear1); + x = modules::LinearModule({128, kContentDim, true}).build(ctx, x, weights_->gpt_layer.linear2); + output_ = core::ensure_backend_addressable_layout(ctx, x).tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(8192, frames_ * 128)), false); + ggml_build_forward_expand(graph_, output_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel GPT layer input buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel GPT layer graph"); + } + debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.s2mel.gpt_layer.frames", frames_); + } + + ~GptLayerGraph() { + clear_graph(); + } + + int64_t frames() const noexcept { + return frames_; + } + + IndexTTS25S2MelSequence run(const std::vector & latent) { + if (static_cast(latent.size()) != frames_ * kGptDim) { + throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer latent size mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(input_, latent.data(), 0, latent.size() * sizeof(float)); + debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph compute failed"); + } + IndexTTS25S2MelSequence output; + output.frames = frames_; + output.dims = kContentDim; + timing_start = Clock::now(); + output.values = core::read_tensor_f32(output_); + debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + return output; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t frames_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; +}; + +class IndexTTS25S2MelRuntime::LengthRegulatorGraph { +public: + LengthRegulatorGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t input_frames, + int64_t output_frames, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + input_frames_(input_frames), + output_frames_(output_frames) { + if (input_frames_ <= 0 || output_frames_ <= 0) { + throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph requires positive frame counts"); + } + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph requires weights"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel length regulator graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel length regulator input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.s2mel.length_regulator", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.s2mel.length_regulator.inputs", + execution_.backend_type()}; + input_ = core::make_tensor( + input_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, input_frames_, kContentDim})) + .tensor; + mask_ = + core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, output_frames_, kHidden})).tensor; + ggml_set_input(input_); + ggml_set_input(mask_); + auto x = core::wrap_tensor(input_, core::TensorShape::from_dims({1, input_frames_, kContentDim}), GGML_TYPE_F32); + x = modules::LinearModule({kContentDim, kHidden, true}).build(ctx, x, weights_->length_regulator.content_projection); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = modules::Interpolate1dModule({output_frames_, modules::Interpolate1dMode::Nearest}).build(ctx, x); + for (size_t layer = 0; layer < weights_->length_regulator.convs.size(); ++layer) { + x = modules::Conv1dModule({kHidden, kHidden, 3, 1, 1, 1, true}).build(ctx, x, weights_->length_regulator.convs[layer]); + x = group_norm_1_group(ctx, x, weights_->length_regulator.norms[layer], kHidden); + x = mish(ctx, x); + } + x = modules::Conv1dModule({kHidden, kHidden, 1, 1, 0, 1, true}).build(ctx, x, weights_->length_regulator.output); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = core::ensure_backend_addressable_layout(ctx, x); + auto mask = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, output_frames_, kHidden}), GGML_TYPE_F32); + auto out = modules::MulModule{}.build(ctx, x, mask); + output_ = core::ensure_backend_addressable_layout(ctx, out).tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(32768, output_frames_ * 512)), false); + ggml_build_forward_expand(graph_, output_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel length regulator input buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel length regulator graph"); + } + mask_values_.assign(static_cast(output_frames_ * kHidden), 1.0F); + core::write_tensor_f32( + core::wrap_tensor(mask_, core::TensorShape::from_dims({1, output_frames_, kHidden}), GGML_TYPE_F32), + mask_values_); + debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.s2mel.length_regulator.input_frames", input_frames_); + debug::trace_log_scalar("index_tts2_5.s2mel.length_regulator.output_frames", output_frames_); + } + + ~LengthRegulatorGraph() { + clear_graph(); + } + + bool matches(int64_t input_frames, int64_t output_frames) const noexcept { + return input_frames_ == input_frames && output_frames_ == output_frames; + } + + IndexTTS25S2MelSequence run(const std::vector & content) { + if (static_cast(content.size()) != input_frames_ * kContentDim) { + throw std::runtime_error("IndexTTS2.5 S2Mel length regulator content size mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(input_, content.data(), 0, content.size() * sizeof(float)); + debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph compute failed"); + } + IndexTTS25S2MelSequence output; + output.frames = output_frames_; + output.dims = kHidden; + timing_start = Clock::now(); + output.values = core::read_tensor_f32(output_); + debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + return output; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t input_frames_ = 0; + int64_t output_frames_ = 0; + std::vector mask_values_; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; +}; + +class IndexTTS25S2MelRuntime::CfmGraph { +public: + CfmGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t frames, + bool use_cfg, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + frames_(frames), + use_cfg_(use_cfg), + batch_(use_cfg ? 2 : 1) { + if (frames_ <= 0) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph requires positive frame count"); + } + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph requires weights"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel CFM graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel CFM input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.s2mel.cfm", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{input_ctx_.get(), "index_tts2_5.s2mel.cfm.inputs", execution_.backend_type()}; + x_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, kMelChannels, frames_})) + .tensor; + prompt_ = + core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, kMelChannels, frames_})).tensor; + cond_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, frames_, kHidden})).tensor; + style_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, kStyleDim})).tensor; + timestep_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_})).tensor; + positions_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({frames_})).tensor; + ggml_set_input(x_); + ggml_set_input(prompt_); + ggml_set_input(cond_); + ggml_set_input(style_); + ggml_set_input(timestep_); + ggml_set_input(positions_); + auto output = build_cfm_estimator( + ctx, + core::wrap_tensor(x_, core::TensorShape::from_dims({batch_, kMelChannels, frames_}), GGML_TYPE_F32), + core::wrap_tensor(prompt_, core::TensorShape::from_dims({batch_, kMelChannels, frames_}), GGML_TYPE_F32), + core::wrap_tensor(cond_, core::TensorShape::from_dims({batch_, frames_, kHidden}), GGML_TYPE_F32), + core::wrap_tensor(style_, core::TensorShape::from_dims({batch_, kStyleDim}), GGML_TYPE_F32), + core::wrap_tensor(timestep_, core::TensorShape::from_dims({batch_}), GGML_TYPE_F32), + core::wrap_tensor(positions_, core::TensorShape::from_dims({frames_}), GGML_TYPE_I32), + weights_->cfm); + output_ = core::ensure_backend_addressable_layout(ctx, output).tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(131072, frames_ * 4096)), false); + ggml_build_forward_expand(graph_, output_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel CFM input buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel CFM graph"); + } + positions_values_.assign(static_cast(frames_), 0); + for (int64_t i = 0; i < frames_; ++i) { + positions_values_[static_cast(i)] = static_cast(i); + } + ggml_backend_tensor_set(positions_, positions_values_.data(), 0, positions_values_.size() * sizeof(int32_t)); + debug::timing_log_scalar("index_tts2_5.s2mel.cfm.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.s2mel.cfm.frames", frames_); + debug::trace_log_scalar("index_tts2_5.s2mel.cfm.batch", batch_); + } + + ~CfmGraph() { + clear_graph(); + } + + bool matches(int64_t frames, bool use_cfg) const noexcept { + return frames_ == frames && use_cfg_ == use_cfg; + } + + std::vector run( + const std::vector & x, + const std::vector & prompt, + const std::vector & cond, + const std::vector & style, + const std::vector & timestep) { + const int64_t mel_values = batch_ * kMelChannels * frames_; + if (static_cast(x.size()) != mel_values || + static_cast(prompt.size()) != mel_values || + static_cast(cond.size()) != batch_ * frames_ * kHidden || + static_cast(style.size()) != batch_ * kStyleDim || + static_cast(timestep.size()) != batch_) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph input shape mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(x_, x.data(), 0, x.size() * sizeof(float)); + ggml_backend_tensor_set(prompt_, prompt.data(), 0, prompt.size() * sizeof(float)); + ggml_backend_tensor_set(cond_, cond.data(), 0, cond.size() * sizeof(float)); + ggml_backend_tensor_set(style_, style.data(), 0, style.size() * sizeof(float)); + ggml_backend_tensor_set(timestep_, timestep.data(), 0, timestep.size() * sizeof(float)); + debug::timing_log_scalar("index_tts2_5.s2mel.cfm.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar("index_tts2_5.s2mel.cfm.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph compute failed"); + } + std::vector out(static_cast(mel_values), 0.0F); + timing_start = Clock::now(); + ggml_backend_tensor_get(output_, out.data(), 0, out.size() * sizeof(float)); + debug::timing_log_scalar("index_tts2_5.s2mel.cfm.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + return out; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t frames_ = 0; + bool use_cfg_ = false; + int64_t batch_ = 1; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * x_ = nullptr; + ggml_tensor * prompt_ = nullptr; + ggml_tensor * cond_ = nullptr; + ggml_tensor * style_ = nullptr; + ggml_tensor * timestep_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; + std::vector positions_values_; +}; + +void copy_row( + const std::vector & src, + int64_t src_row, + std::vector & dst, + int64_t dst_row, + int64_t row_values) { + std::copy_n( + src.data() + static_cast(src_row * row_values), + static_cast(row_values), + dst.data() + static_cast(dst_row * row_values)); +} + +std::vector repeat_or_zero_rows(const std::vector & values, int64_t row_values, bool use_cfg, bool zero_second) { + if (!use_cfg) { + return values; + } + std::vector out(static_cast(2 * row_values), 0.0F); + copy_row(values, 0, out, 0, row_values); + if (!zero_second) { + copy_row(values, 0, out, 1, row_values); + } + return out; +} + +void zero_prompt_region(std::vector & values, int64_t channels, int64_t frames, int64_t prompt_frames) { + if (prompt_frames < 0 || prompt_frames > frames) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM prompt frame count is out of range"); + } + for (int64_t c = 0; c < channels; ++c) { + auto begin = values.begin() + static_cast(c * frames); + std::fill(begin, begin + static_cast(prompt_frames), 0.0F); + } +} + +std::vector make_prompt_x( + const std::vector & prompt, + int64_t channels, + int64_t frames, + int64_t prompt_frames) { + std::vector out(static_cast(channels * frames), 0.0F); + if (static_cast(prompt.size()) != channels * prompt_frames) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM reference mel shape mismatch"); + } + for (int64_t c = 0; c < channels; ++c) { + std::copy_n( + prompt.data() + static_cast(c * prompt_frames), + static_cast(prompt_frames), + out.data() + static_cast(c * frames)); + } + return out; +} + +std::vector make_condition_with_prompt( + const std::vector & condition, + int64_t frames) { + if (static_cast(condition.size()) != frames * kHidden) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM condition shape mismatch"); + } + return condition; +} + +IndexTTS25S2MelRuntime::IndexTTS25S2MelRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) + : assets_(std::move(assets)), + execution_(&execution), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel runtime requires assets"); + } + if (graph_arena_bytes_ == 0) { + throw std::runtime_error("IndexTTS2.5 S2Mel graph arena must be non-zero"); + } + weights_ = load_index_tts2_5_s2mel_weights( + *assets_, + execution.backend(), + execution.backend_type(), + matmul_storage_type, + conv_storage_type, + weight_context_bytes); +} + +IndexTTS25S2MelRuntime::~IndexTTS25S2MelRuntime() = default; + +void IndexTTS25S2MelRuntime::prepare_gpt_layer(int64_t frames) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel runtime execution context is missing"); + } + if (gpt_layer_graph_ != nullptr && gpt_layer_graph_->frames() == frames) { + return; + } + gpt_layer_graph_.reset(); + gpt_layer_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); +} + +void IndexTTS25S2MelRuntime::prepare_length_regulator(int64_t input_frames, int64_t output_frames) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel runtime execution context is missing"); + } + if (length_regulator_graph_ != nullptr && length_regulator_graph_->matches(input_frames, output_frames)) { + return; + } + length_regulator_graph_.reset(); + length_regulator_graph_ = std::make_unique( + *execution_, + weights_, + input_frames, + output_frames, + graph_arena_bytes_); +} + +void IndexTTS25S2MelRuntime::prepare_cfm(int64_t total_frames, bool use_cfg) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 S2Mel runtime execution context is missing"); + } + if (total_frames <= 0) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM prepare requires positive frame count"); + } + if (cfm_graph_ != nullptr && cfm_graph_->matches(total_frames, use_cfg)) { + return; + } + cfm_graph_.reset(); + cfm_graph_ = std::make_unique(*execution_, weights_, total_frames, use_cfg, graph_arena_bytes_); +} + +void IndexTTS25S2MelRuntime::release_pre_cfm_graphs() { + gpt_layer_graph_.reset(); + length_regulator_graph_.reset(); +} + +void IndexTTS25S2MelRuntime::release_cfm_graph() { + cfm_graph_.reset(); +} + +IndexTTS25S2MelSequence IndexTTS25S2MelRuntime::project_gpt_latent(const std::vector & latent, int64_t frames) { + if (gpt_layer_graph_ == nullptr || gpt_layer_graph_->frames() != frames) { + throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph was not prepared for this latent length"); + } + return gpt_layer_graph_->run(latent); +} + +IndexTTS25S2MelSequence IndexTTS25S2MelRuntime::regulate_length( + const std::vector & content, + int64_t input_frames, + int64_t output_frames) { + if (length_regulator_graph_ == nullptr || !length_regulator_graph_->matches(input_frames, output_frames)) { + throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph was not prepared for this shape"); + } + return length_regulator_graph_->run(content); +} + +IndexTTS25S2MelMel IndexTTS25S2MelRuntime::infer_mel( + const std::vector & condition, + int64_t total_frames, + const std::vector & reference_mel, + int64_t reference_frames, + const std::vector & style, + int64_t diffusion_steps, + float cfg_rate, + uint32_t seed, + uint64_t rng_offset_blocks) { + if (condition.empty() || reference_mel.empty() || style.empty()) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM requires non-empty condition, reference mel, and style"); + } + if (total_frames <= 0 || reference_frames <= 0 || reference_frames > total_frames) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM frame counts are invalid"); + } + if (diffusion_steps <= 0) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM diffusion steps must be positive"); + } + if (static_cast(style.size()) != kStyleDim) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM style shape mismatch"); + } + const bool use_cfg = cfg_rate > 0.0F; + if (cfm_graph_ == nullptr || !cfm_graph_->matches(total_frames, use_cfg)) { + throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph was not prepared for this shape"); + } + const auto rng_policy = engine::sampling::resolve_torch_cuda_sampling_policy( + execution_->backend_type(), + execution_->config().device, + "index_tts2_5.s2mel.cuda_sampling_policy", + "IndexTTS2.5", + engine::sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda); + std::vector x = engine::sampling::generate_torch_cuda_tensor_iterator_randn( + static_cast(kMelChannels * total_frames), + seed, + rng_offset_blocks, + rng_policy, + engine::sampling::TorchRandnPrecision::Float32); + auto prompt_x = make_prompt_x(reference_mel, kMelChannels, total_frames, reference_frames); + zero_prompt_region(x, kMelChannels, total_frames, reference_frames); + auto mu = make_condition_with_prompt(condition, total_frames); + + const auto cfm_start = Clock::now(); + double graph_ms = 0.0; + float t = 0.0F; + float dt = 1.0F / static_cast(diffusion_steps); + for (int64_t step = 1; step <= diffusion_steps; ++step) { + const auto graph_start = Clock::now(); + auto x_batched = repeat_or_zero_rows(x, kMelChannels * total_frames, use_cfg, false); + auto prompt_batched = repeat_or_zero_rows(prompt_x, kMelChannels * total_frames, use_cfg, true); + auto cond_batched = repeat_or_zero_rows(mu, total_frames * kHidden, use_cfg, true); + auto style_batched = repeat_or_zero_rows(style, kStyleDim, use_cfg, true); + std::vector timestep(static_cast(use_cfg ? 2 : 1), t); + const auto velocity = cfm_graph_->run(x_batched, prompt_batched, cond_batched, style_batched, timestep); + graph_ms += engine::debug::elapsed_ms(graph_start); + const int64_t row_values = kMelChannels * total_frames; + for (int64_t i = 0; i < row_values; ++i) { + float dphi = velocity[static_cast(i)]; + if (use_cfg) { + dphi = (1.0F + cfg_rate) * dphi - cfg_rate * velocity[static_cast(row_values + i)]; + } + x[static_cast(i)] += dt * dphi; + } + t += dt; + if (step < diffusion_steps) { + dt = (static_cast(step + 1) / static_cast(diffusion_steps)) - t; + } + zero_prompt_region(x, kMelChannels, total_frames, reference_frames); + } + + IndexTTS25S2MelMel out; + out.frames = total_frames - reference_frames; + out.channels = kMelChannels; + out.values.resize(static_cast(out.channels * out.frames)); + for (int64_t c = 0; c < kMelChannels; ++c) { + std::copy_n( + x.data() + static_cast(c * total_frames + reference_frames), + static_cast(out.frames), + out.values.data() + static_cast(c * out.frames)); + } + debug::timing_log_scalar("index_tts2_5.s2mel.cfm.euler_graph_ms", graph_ms); + debug::timing_log_scalar("index_tts2_5.s2mel.cfm.euler_total_ms", engine::debug::elapsed_ms(cfm_start)); + return out; +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/semantic_codec.cpp b/src/models/index_tts2_5/semantic_codec.cpp new file mode 100644 index 00000000..e1982a98 --- /dev/null +++ b/src/models/index_tts2_5/semantic_codec.cpp @@ -0,0 +1,723 @@ +#include "engine/models/index_tts2_5/semantic_codec.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +namespace binding = engine::modules::binding; +namespace core = engine::core; +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kHidden = 1024; +constexpr int64_t kVocosDim = 384; +constexpr int64_t kVocosIntermediate = 2048; +constexpr int64_t kVocosLayers = 12; +constexpr int64_t kConvNeXtKernel = 7; +constexpr int64_t kCodebookSize = 8192; +constexpr int64_t kCodebookDim = 8; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +std::vector normalized_codebook(const std::vector & codebook) { + std::vector out(codebook.size(), 0.0F); + for (int64_t row = 0; row < kCodebookSize; ++row) { + double norm = 0.0; + for (int64_t dim = 0; dim < kCodebookDim; ++dim) { + const float value = codebook[static_cast(row * kCodebookDim + dim)]; + norm += static_cast(value) * static_cast(value); + } + const float inv_norm = 1.0F / static_cast(std::sqrt(norm)); + for (int64_t dim = 0; dim < kCodebookDim; ++dim) { + const size_t index = static_cast(row * kCodebookDim + dim); + out[index] = codebook[index] * inv_norm; + } + } + return out; +} + +core::TensorValue div( + core::ModuleBuildContext & ctx, + const core::TensorValue & lhs, + const core::TensorValue & rhs) { + core::validate_shape(rhs, lhs.shape, "Div rhs"); + return core::wrap_tensor(ggml_div(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32); +} + +core::TensorValue vocos_backbone( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const IndexTTS25VocosBackboneWeights & weights) { + auto x = modules::Conv1dModule({input_bct.shape.dims[1], kVocosDim, kConvNeXtKernel, 1, 3, 1, true}) + .build(ctx, input_bct, weights.embed); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-6F, true, true}).build(ctx, x, weights.norm); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + for (const auto & block : weights.blocks) { + const auto residual = x; + x = modules::DepthwiseConv1dModule({kVocosDim, kConvNeXtKernel, 1, 3, 1, true}).build(ctx, x, block.depthwise); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-6F, true, true}).build(ctx, x, block.norm); + x = modules::LinearModule({kVocosDim, kVocosIntermediate, true}).build(ctx, x, block.pointwise_in); + x = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, x); + x = modules::LinearModule({kVocosIntermediate, kVocosDim, true}).build(ctx, x, block.pointwise_out); + const auto gamma = modules::RepeatModule({x.shape}).build( + ctx, + core::reshape_tensor(ctx, block.gamma, core::TensorShape::from_dims({1, 1, kVocosDim}))); + x = modules::MulModule{}.build(ctx, x, gamma); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = modules::AddModule{}.build(ctx, residual, x); + } + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + return modules::LayerNormModule({x.shape.last_dim(), 1.0e-6F, true, true}).build(ctx, x, weights.final_norm); +} + +core::TensorValue quantizer_in_project( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const IndexTTS25SemanticCodecWeights & weights) { + return modules::Conv1dModule({kHidden, kCodebookDim, 1, 1, 0, 1, true}).build(ctx, input_bct, weights.quantizer_in); +} + +core::TensorValue quantizer_out_project( + core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const IndexTTS25SemanticCodecWeights & weights) { + return modules::Conv1dModule({kCodebookDim, kHidden, 1, 1, 0, 1, true}).build(ctx, input_bct, weights.quantizer_out); +} + +core::TensorValue normalize_code_latents(core::ModuleBuildContext & ctx, const core::TensorValue & latents_bdt) { + const auto squared = core::wrap_tensor(ggml_sqr(ctx.ggml, latents_bdt.tensor), latents_bdt.shape, GGML_TYPE_F32); + auto sum = modules::ReduceSumModule({1}).build(ctx, squared); + sum = core::wrap_tensor(ggml_sqrt(ctx.ggml, sum.tensor), sum.shape, GGML_TYPE_F32); + return div(ctx, latents_bdt, modules::RepeatModule({latents_bdt.shape}).build(ctx, sum)); +} + +core::TensorValue argmax_last_dim(core::ModuleBuildContext & ctx, const core::TensorValue & logits) { + auto flat = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, logits), + core::TensorShape::from_dims({logits.shape.num_elements() / logits.shape.last_dim(), logits.shape.last_dim()})); + auto argmax = core::wrap_tensor( + ggml_argmax(ctx.ggml, flat.tensor), + core::TensorShape::from_dims({flat.shape.dims[0]}), + GGML_TYPE_I32); + return core::reshape_tensor(ctx, argmax, core::TensorShape::from_dims({logits.shape.dims[0], logits.shape.dims[1]})); +} + +core::TensorValue embed_codes_bct( + core::ModuleBuildContext & ctx, + const core::TensorValue & codes_bt, + const IndexTTS25SemanticCodecWeights & weights) { + auto emb_btd = modules::EmbeddingModule({kCodebookSize, kCodebookDim}).build(ctx, codes_bt, weights.codebook); + auto emb_bdt = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, emb_btd); + return quantizer_out_project(ctx, emb_bdt, weights); +} + +std::vector fuse_weight_norm_conv1d( + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size) { + const auto g = source.require_f32(prefix + ".weight_g", {out_channels, 1, 1}); + const auto v = source.require_f32(prefix + ".weight_v", {out_channels, in_channels, kernel_size}); + std::vector weight(v.size(), 0.0F); + for (int64_t out = 0; out < out_channels; ++out) { + double norm = 0.0; + for (int64_t in = 0; in < in_channels; ++in) { + for (int64_t k = 0; k < kernel_size; ++k) { + const float value = v[static_cast((out * in_channels + in) * kernel_size + k)]; + norm += static_cast(value) * static_cast(value); + } + } + const float scale = g[static_cast(out)] / static_cast(std::sqrt(norm)); + for (int64_t in = 0; in < in_channels; ++in) { + for (int64_t k = 0; k < kernel_size; ++k) { + const size_t index = static_cast((out * in_channels + in) * kernel_size + k); + weight[index] = v[index] * scale; + } + } + } + return weight; +} + +engine::modules::Conv1dWeights load_weight_norm_conv1d( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size) { + engine::modules::Conv1dWeights weights; + weights.weight = store.make_from_f32( + engine::core::TensorShape::from_dims({out_channels, in_channels, kernel_size}), + storage_type, + fuse_weight_norm_conv1d(source, prefix, out_channels, in_channels, kernel_size)); + weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); + return weights; +} + +IndexTTS25VocosConvNeXtBlockWeights load_convnext_block( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) { + IndexTTS25VocosConvNeXtBlockWeights block; + block.depthwise = binding::depthwise_conv1d_from_source( + store, + source, + prefix + ".dwconv", + conv_storage_type, + kVocosDim, + kConvNeXtKernel, + true); + block.norm = binding::norm_from_source(store, source, prefix + ".norm", kVocosDim); + block.pointwise_in = binding::linear_from_source( + store, + source, + prefix + ".pwconv1", + matmul_storage_type, + kVocosIntermediate, + kVocosDim, + true); + block.pointwise_out = binding::linear_from_source( + store, + source, + prefix + ".pwconv2", + matmul_storage_type, + kVocosDim, + kVocosIntermediate, + true); + block.gamma = store.load_f32_tensor(source, prefix + ".gamma", {kVocosDim}); + return block; +} + +IndexTTS25VocosBackboneWeights load_vocos_backbone( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + int64_t input_channels, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) { + IndexTTS25VocosBackboneWeights backbone; + backbone.embed = binding::conv1d_from_source( + store, + source, + prefix + ".embed", + conv_storage_type, + kVocosDim, + input_channels, + kConvNeXtKernel, + true); + backbone.norm = binding::norm_from_source(store, source, prefix + ".norm", kVocosDim); + backbone.blocks.reserve(static_cast(kVocosLayers)); + for (int64_t i = 0; i < kVocosLayers; ++i) { + backbone.blocks.push_back(load_convnext_block( + store, + source, + prefix + ".convnext." + std::to_string(i), + matmul_storage_type, + conv_storage_type)); + } + backbone.final_norm = binding::norm_from_source(store, source, prefix + ".final_layer_norm", kVocosDim); + return backbone; +} + +} // namespace + +class IndexTTS25SemanticCodecRuntime::QuantizeGraph { +public: + QuantizeGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t frames, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + frames_(frames) { + if (frames_ <= 0) { + throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph requires positive frame count"); + } + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph requires weights"); + } + + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec quantize graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec quantize input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.semantic_codec.quantize", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.semantic_codec.quantize.inputs", + execution_.backend_type()}; + + semantic_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kHidden})).tensor; + ggml_set_input(semantic_); + auto x = core::wrap_tensor(semantic_, core::TensorShape::from_dims({1, frames_, kHidden}), GGML_TYPE_F32); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = vocos_backbone(ctx, x, weights_->encoder_backbone); + x = modules::LinearModule({kVocosDim, kHidden, true}).build(ctx, x, weights_->encoder_projection); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + auto latents = normalize_code_latents(ctx, quantizer_in_project(ctx, x, *weights_)); + const auto latents_btd = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, latents); + auto logits = modules::LinearModule({kCodebookDim, kCodebookSize, false}) + .build(ctx, latents_btd, {weights_->normalized_codebook, std::nullopt}); + codes_ = argmax_last_dim(ctx, logits).tensor; + auto embedding = embed_codes_bct( + ctx, + core::wrap_tensor(codes_, core::TensorShape::from_dims({1, frames_}), GGML_TYPE_I32), + *weights_); + embedding_ = core::ensure_backend_addressable_layout(ctx, embedding).tensor; + ggml_set_output(codes_); + ggml_set_output(embedding_); + + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(65536, frames_ * 2048 + 4096)), false); + ggml_build_forward_expand(graph_, codes_); + ggml_build_forward_expand(graph_, embedding_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec quantize input buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec quantize graph"); + } + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.quantize.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.semantic_codec.quantize.frames", frames_); + } + + ~QuantizeGraph() { + clear_graph(); + } + + int64_t frames() const noexcept { + return frames_; + } + + IndexTTS25SemanticCodecOutput run(const IndexTTS25SemanticEmbedding & semantic) { + if (semantic.frames != frames_ || semantic.dims != kHidden) { + throw std::runtime_error("IndexTTS2.5 semantic codec semantic input shape does not match prepared graph"); + } + if (static_cast(semantic.values.size()) != frames_ * kHidden) { + throw std::runtime_error("IndexTTS2.5 semantic codec semantic input value count mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(semantic_, semantic.values.data(), 0, semantic.values.size() * sizeof(float)); + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.quantize.input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.quantize.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph compute failed"); + } + + IndexTTS25SemanticCodecOutput output; + output.frames = frames_; + output.dims = kHidden; + output.codes.resize(static_cast(frames_)); + output.embedding_channel_first.resize(static_cast(kHidden * frames_)); + timing_start = Clock::now(); + ggml_backend_tensor_get(codes_, output.codes.data(), 0, output.codes.size() * sizeof(int32_t)); + ggml_backend_tensor_get( + embedding_, + output.embedding_channel_first.data(), + 0, + output.embedding_channel_first.size() * sizeof(float)); + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.quantize.output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return output; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t frames_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * semantic_ = nullptr; + ggml_tensor * codes_ = nullptr; + ggml_tensor * embedding_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; +}; + +class IndexTTS25SemanticCodecRuntime::CodesGraph { +public: + CodesGraph( + core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t frames, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + frames_(frames) { + if (frames_ <= 0) { + throw std::runtime_error("IndexTTS2.5 semantic codec code graph requires positive frame count"); + } + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 semantic codec code graph requires weights"); + } + + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec code graph context"); + } + ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec code input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.semantic_codec.codes", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.semantic_codec.codes.inputs", + execution_.backend_type()}; + codes_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, frames_})).tensor; + upsample_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 2 * frames_); + ggml_set_input(codes_); + auto embedding = embed_codes_bct( + ctx, + core::wrap_tensor(codes_, core::TensorShape::from_dims({1, frames_}), GGML_TYPE_I32), + *weights_); + // EnhancedCodec.decode (codec/models.py): decoder backbone + projection, + // then 2x nearest upsample along time and the `up` conv. + auto x = vocos_backbone(ctx, embedding, weights_->decoder_backbone); + x = modules::LinearModule({kVocosDim, kHidden, true}).build(ctx, x, weights_->decoder_projection); + x = core::ensure_backend_addressable_layout(ctx, x); + x = core::wrap_tensor( + ggml_get_rows(ctx.ggml, x.tensor, upsample_ids_), + core::TensorShape::from_dims({1, 2 * frames_, kHidden}), + GGML_TYPE_F32); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = modules::Conv1dModule({kHidden, kHidden, 3, 1, 1, 1, true}).build(ctx, x, weights_->up); + embedding_ = core::ensure_backend_addressable_layout(ctx, x).tensor; + ggml_set_output(embedding_); + + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(4096, frames_ * 64 + 1024)), false); + ggml_build_forward_expand(graph_, embedding_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec code input buffer"); + } + std::vector upsample_ids(static_cast(2 * frames_)); + for (int64_t frame = 0; frame < frames_; ++frame) { + upsample_ids[static_cast(2 * frame)] = static_cast(frame); + upsample_ids[static_cast(2 * frame + 1)] = static_cast(frame); + } + ggml_backend_tensor_set(upsample_ids_, upsample_ids.data(), 0, upsample_ids.size() * sizeof(int32_t)); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec code graph"); + } + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.codes.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.semantic_codec.codes.frames", frames_); + } + + ~CodesGraph() { + clear_graph(); + } + + int64_t frames() const noexcept { + return frames_; + } + + IndexTTS25SemanticCodecOutput run(const std::vector & codes) { + if (static_cast(codes.size()) != frames_) { + throw std::runtime_error("IndexTTS2.5 semantic codec code count does not match prepared graph"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(codes_, codes.data(), 0, codes.size() * sizeof(int32_t)); + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.codes.input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.codes.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 semantic codec code graph compute failed"); + } + + IndexTTS25SemanticCodecOutput output; + output.frames = 2 * frames_; + output.dims = kHidden; + output.codes = codes; + output.embedding_channel_first.resize(static_cast(kHidden * output.frames)); + timing_start = Clock::now(); + ggml_backend_tensor_get( + embedding_, + output.embedding_channel_first.data(), + 0, + output.embedding_channel_first.size() * sizeof(float)); + debug::timing_log_scalar( + "index_tts2_5.semantic_codec.codes.output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return output; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t frames_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * codes_ = nullptr; + ggml_tensor * upsample_ids_ = nullptr; + ggml_tensor * embedding_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; +}; + +std::shared_ptr load_index_tts2_5_semantic_codec_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes) { + if (assets.semantic_codec_weights == nullptr) { + throw std::runtime_error("IndexTTS2.5 semantic codec requires tensor source"); + } + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "index_tts2_5.semantic_codec.weights", + weight_context_bytes); + + const auto & source = *assets.semantic_codec_weights; + weights->encoder_backbone = load_vocos_backbone( + *weights->store, + source, + "encoder.0", + kHidden, + matmul_storage_type, + conv_storage_type); + weights->encoder_projection = binding::linear_from_source( + *weights->store, + source, + "encoder.1", + matmul_storage_type, + kHidden, + kVocosDim, + true); + weights->quantizer_in = load_weight_norm_conv1d( + *weights->store, + source, + "quantizer.quantizers.0.in_project", + conv_storage_type, + kCodebookDim, + kHidden, + 1); + const auto codebook = source.require_f32("quantizer.quantizers.0.codebook.weight", {kCodebookSize, kCodebookDim}); + weights->codebook = weights->store->make_from_f32( + engine::core::TensorShape::from_dims({kCodebookSize, kCodebookDim}), + matmul_storage_type, + codebook); + weights->normalized_codebook = weights->store->make_from_f32( + engine::core::TensorShape::from_dims({kCodebookSize, kCodebookDim}), + matmul_storage_type, + normalized_codebook(codebook)); + weights->quantizer_out = load_weight_norm_conv1d( + *weights->store, + source, + "quantizer.quantizers.0.out_project", + conv_storage_type, + kHidden, + kCodebookDim, + 1); + weights->decoder_backbone = load_vocos_backbone( + *weights->store, + source, + "decoder.0", + kHidden, + matmul_storage_type, + conv_storage_type); + weights->decoder_projection = binding::linear_from_source( + *weights->store, + source, + "decoder.1", + matmul_storage_type, + kHidden, + kVocosDim, + true); + weights->up = binding::conv1d_from_source( + *weights->store, + source, + "up", + conv_storage_type, + kHidden, + kHidden, + 3, + true); + + weights->store->upload(); + assets.semantic_codec_weights->release_storage(); + return weights; +} + +IndexTTS25SemanticCodecRuntime::IndexTTS25SemanticCodecRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) + : assets_(std::move(assets)), + execution_(&execution), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 semantic codec runtime requires assets"); + } + if (graph_arena_bytes_ == 0) { + throw std::runtime_error("IndexTTS2.5 semantic codec graph arena must be non-zero"); + } + weights_ = load_index_tts2_5_semantic_codec_weights( + *assets_, + execution.backend(), + execution.backend_type(), + matmul_storage_type, + conv_storage_type, + weight_context_bytes); +} + +IndexTTS25SemanticCodecRuntime::~IndexTTS25SemanticCodecRuntime() = default; + +void IndexTTS25SemanticCodecRuntime::prepare_quantize(int64_t frames) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 semantic codec runtime execution context is missing"); + } + if (frames <= 0) { + throw std::runtime_error("IndexTTS2.5 semantic codec quantize prepare requires positive frames"); + } + if (quantize_graph_ != nullptr && quantize_graph_->frames() == frames) { + return; + } + quantize_graph_.reset(); + quantize_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); +} + +void IndexTTS25SemanticCodecRuntime::prepare_codes(int64_t frames) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 semantic codec runtime execution context is missing"); + } + if (frames <= 0) { + throw std::runtime_error("IndexTTS2.5 semantic codec code prepare requires positive frames"); + } + if (codes_graph_ != nullptr && codes_graph_->frames() == frames) { + return; + } + codes_graph_.reset(); + codes_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); +} + +IndexTTS25SemanticCodecOutput IndexTTS25SemanticCodecRuntime::quantize(const IndexTTS25SemanticEmbedding & semantic) { + if (quantize_graph_ == nullptr || quantize_graph_->frames() != semantic.frames) { + throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph was not prepared for this reference length"); + } + return quantize_graph_->run(semantic); +} + +IndexTTS25SemanticCodecOutput IndexTTS25SemanticCodecRuntime::codes_to_embedding( + const std::vector & codes, + int64_t frames) { + if (codes_graph_ == nullptr || codes_graph_->frames() != frames) { + throw std::runtime_error("IndexTTS2.5 semantic codec code graph was not prepared for this generation length"); + } + return codes_graph_->run(codes); +} + +void IndexTTS25SemanticCodecRuntime::release_graphs() { + quantize_graph_.reset(); + codes_graph_.reset(); +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/semantic_encoder.cpp b/src/models/index_tts2_5/semantic_encoder.cpp new file mode 100644 index 00000000..bc1b2e3f --- /dev/null +++ b/src/models/index_tts2_5/semantic_encoder.cpp @@ -0,0 +1,622 @@ +#include "engine/models/index_tts2_5/semantic_encoder.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +namespace binding = engine::modules::binding; +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kFeatureDim = 160; +constexpr int64_t kHidden = 1024; +constexpr int64_t kIntermediate = 4096; +constexpr int64_t kLayers = 24; +constexpr int64_t kSemanticOutputHiddenStateIndex = 17; +constexpr int64_t kHeads = 16; +constexpr int64_t kHeadDim = kHidden / kHeads; +constexpr int64_t kRelativePositions = 73; +constexpr int64_t kRelativeLeft = 64; +constexpr int64_t kRelativeRight = 8; +constexpr int64_t kConvKernel = 31; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +core::TensorValue scale_tensor(core::ModuleBuildContext & ctx, const core::TensorValue & input, float scale) { + return core::wrap_tensor(ggml_scale(ctx.ggml, input.tensor, scale), input.shape, GGML_TYPE_F32); +} + +core::TensorValue add_scaled_residual( + core::ModuleBuildContext & ctx, + const core::TensorValue & residual, + const core::TensorValue & update, + float update_scale) { + return modules::AddModule{}.build(ctx, residual, scale_tensor(ctx, update, update_scale)); +} + +core::TensorValue reshape_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + return core::reshape_tensor(ctx, contiguous, core::TensorShape::from_dims({1, input.shape.dims[1], kHeads, kHeadDim})); +} + +core::TensorValue repeat_last_dim_mask( + core::ModuleBuildContext & ctx, + const core::TensorValue & mask, + const core::TensorValue & like) { + auto mask_view = core::reshape_tensor(ctx, mask, core::TensorShape::from_dims({1, mask.shape.dims[0], 1})); + return modules::RepeatModule({like.shape}).build(ctx, mask_view); +} + +core::TensorValue apply_keep_mask( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & keep_mask) { + return modules::MulModule{}.build(ctx, input, repeat_last_dim_mask(ctx, keep_mask, input)); +} + +core::TensorValue broadcast_vector( + core::ModuleBuildContext & ctx, + const core::TensorValue & vector, + const core::TensorValue & like) { + auto view = core::reshape_tensor(ctx, vector, core::TensorShape::from_dims({1, 1, vector.shape.dims[0]})); + return modules::RepeatModule({like.shape}).build(ctx, view); +} + +core::TensorValue build_relative_key_bias( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & distance_ids, + const core::TensorValue & distance_embedding) { + const int64_t frames = q_heads.shape.dims[2]; + auto positions = modules::EmbeddingModule({kRelativePositions, kHeadDim}).build(ctx, distance_ids, distance_embedding); + auto q_by_row = modules::TransposeModule({{2, 1, 0, 3}, 4}).build(ctx, q_heads); + q_by_row = core::ensure_backend_addressable_layout(ctx, q_by_row); + q_by_row = core::reshape_tensor(ctx, q_by_row, core::TensorShape::from_dims({frames, kHeads, kHeadDim})); + auto positions_by_row = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, positions); + positions_by_row = core::ensure_backend_addressable_layout(ctx, positions_by_row); + auto by_row = modules::MatMulModule{}.build(ctx, q_by_row, positions_by_row); + by_row = core::ensure_backend_addressable_layout(ctx, by_row); + by_row = core::reshape_tensor(ctx, by_row, core::TensorShape::from_dims({frames, kHeads, 1, frames})); + return core::ensure_backend_addressable_layout(ctx, modules::TransposeModule({{2, 1, 0, 3}, 4}).build(ctx, by_row)); +} + +core::TensorValue wav2vec2bert_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & attention_mask, + const core::TensorValue & distance_ids, + const IndexTTS25Wav2Vec2BertAttentionWeights & weights) { + auto q = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, input, weights.q); + auto k = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, input, weights.k); + auto v = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, input, weights.v); + q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q)); + k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, k)); + v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, v)); + + auto scores = modules::MatMulModule{}.build( + ctx, + q, + modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k)); + scores = scale_tensor(ctx, scores, 1.0F / std::sqrt(static_cast(kHeadDim))); + auto relative = build_relative_key_bias(ctx, q, distance_ids, weights.distance_embedding); + relative = scale_tensor(ctx, relative, 1.0F / std::sqrt(static_cast(kHeadDim))); + scores = modules::AddModule{}.build(ctx, scores, relative); + scores = modules::AddModule{}.build(ctx, scores, attention_mask); + auto probs = core::wrap_tensor(ggml_soft_max(ctx.ggml, scores.tensor), scores.shape, GGML_TYPE_F32); + auto out = modules::MatMulModule{}.build(ctx, probs, v); + out = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, out); + out = core::ensure_backend_addressable_layout(ctx, out); + out = core::reshape_tensor(ctx, out, core::TensorShape::from_dims({1, input.shape.dims[1], kHidden})); + return modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, out, weights.out); +} + +core::TensorValue wav2vec2bert_feed_forward( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::LinearWeights & in, + const modules::LinearWeights & out) { + auto hidden = modules::LinearModule({kHidden, kIntermediate, true, GGML_PREC_F32}).build(ctx, input, in); + hidden = modules::SiluModule{}.build(ctx, hidden); + return modules::LinearModule({kIntermediate, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, out); +} + +core::TensorValue wav2vec2bert_conv( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & keep_mask, + const IndexTTS25Wav2Vec2BertConvWeights & weights) { + auto hidden = modules::LayerNormModule({input.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, input, weights.layer_norm); + hidden = apply_keep_mask(ctx, hidden, keep_mask); + hidden = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); + hidden = modules::Conv1dModule({kHidden, 2 * kHidden, 1, 1, 0, 1, false}).build(ctx, hidden, weights.pointwise_in); + auto gate = modules::SliceModule({1, 0, kHidden}).build(ctx, hidden); + auto value = modules::SliceModule({1, kHidden, kHidden}).build(ctx, hidden); + value = modules::SigmoidModule{}.build(ctx, value); + hidden = modules::MulModule{}.build(ctx, gate, value); + + auto first = modules::SliceModule({2, 0, 1}).build(ctx, hidden); + auto left_pad = modules::RepeatModule({core::TensorShape::from_dims({1, kHidden, kConvKernel - 1})}).build(ctx, first); + left_pad = scale_tensor(ctx, left_pad, 0.0F); + hidden = modules::ConcatModule({2}).build(ctx, left_pad, hidden); + hidden = modules::DepthwiseConv1dModule({kHidden, kConvKernel, 1, 0, 1, false}).build(ctx, hidden, weights.depthwise); + hidden = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); + hidden = modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.depthwise_layer_norm); + hidden = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); + hidden = modules::SiluModule{}.build(ctx, hidden); + hidden = modules::Conv1dModule({kHidden, kHidden, 1, 1, 0, 1, false}).build(ctx, hidden, weights.pointwise_out); + return modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); +} + +core::TensorValue wav2vec2bert_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & keep_mask, + const core::TensorValue & attention_mask, + const core::TensorValue & distance_ids, + const IndexTTS25Wav2Vec2BertLayerWeights & weights) { + auto hidden = modules::LayerNormModule({input.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, input, weights.ffn1_norm); + hidden = wav2vec2bert_feed_forward(ctx, hidden, weights.ffn1_in, weights.ffn1_out); + hidden = add_scaled_residual(ctx, input, hidden, 0.5F); + + auto attn = modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.self_attn_norm); + attn = wav2vec2bert_attention(ctx, attn, attention_mask, distance_ids, weights.self_attn); + hidden = modules::AddModule{}.build(ctx, hidden, attn); + + auto conv = wav2vec2bert_conv(ctx, hidden, keep_mask, weights.conv); + hidden = modules::AddModule{}.build(ctx, hidden, conv); + + auto ffn2 = modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.ffn2_norm); + ffn2 = wav2vec2bert_feed_forward(ctx, ffn2, weights.ffn2_in, weights.ffn2_out); + hidden = add_scaled_residual(ctx, hidden, ffn2, 0.5F); + return modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.final_norm); +} + +IndexTTS25Wav2Vec2BertAttentionWeights load_attention( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type) { + IndexTTS25Wav2Vec2BertAttentionWeights weights; + weights.q = binding::linear_from_source(store, source, prefix + ".linear_q", storage_type, kHidden, kHidden, true); + weights.k = binding::linear_from_source(store, source, prefix + ".linear_k", storage_type, kHidden, kHidden, true); + weights.v = binding::linear_from_source(store, source, prefix + ".linear_v", storage_type, kHidden, kHidden, true); + weights.out = binding::linear_from_source(store, source, prefix + ".linear_out", storage_type, kHidden, kHidden, true); + weights.distance_embedding = store.load_tensor( + source, + prefix + ".distance_embedding.weight", + storage_type, + {kRelativePositions, kHeadDim}); + return weights; +} + +IndexTTS25Wav2Vec2BertConvWeights load_conv_module( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + const std::string & prefix, + engine::assets::TensorStorageType storage_type) { + IndexTTS25Wav2Vec2BertConvWeights weights; + weights.layer_norm = binding::norm_from_source(store, source, prefix + ".layer_norm", kHidden); + weights.pointwise_in = binding::conv1d_from_source( + store, + source, + prefix + ".pointwise_conv1", + storage_type, + 2 * kHidden, + kHidden, + 1, + false); + weights.depthwise = binding::depthwise_conv1d_from_source( + store, + source, + prefix + ".depthwise_conv", + storage_type, + kHidden, + kConvKernel, + false); + weights.depthwise_layer_norm = binding::norm_from_source(store, source, prefix + ".depthwise_layer_norm", kHidden); + weights.pointwise_out = binding::conv1d_from_source( + store, + source, + prefix + ".pointwise_conv2", + storage_type, + kHidden, + kHidden, + 1, + false); + return weights; +} + +IndexTTS25Wav2Vec2BertLayerWeights load_layer( + engine::core::BackendWeightStore & store, + const engine::assets::TensorSource & source, + int64_t layer_index, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) { + const std::string prefix = "encoder.layers." + std::to_string(layer_index); + IndexTTS25Wav2Vec2BertLayerWeights layer; + layer.ffn1_norm = binding::norm_from_source(store, source, prefix + ".ffn1_layer_norm", kHidden); + layer.ffn1_in = binding::linear_from_source( + store, + source, + prefix + ".ffn1.intermediate_dense", + matmul_storage_type, + kIntermediate, + kHidden, + true); + layer.ffn1_out = binding::linear_from_source( + store, + source, + prefix + ".ffn1.output_dense", + matmul_storage_type, + kHidden, + kIntermediate, + true); + layer.self_attn_norm = binding::norm_from_source(store, source, prefix + ".self_attn_layer_norm", kHidden); + layer.self_attn = load_attention(store, source, prefix + ".self_attn", matmul_storage_type); + layer.conv = load_conv_module(store, source, prefix + ".conv_module", conv_storage_type); + layer.ffn2_norm = binding::norm_from_source(store, source, prefix + ".ffn2_layer_norm", kHidden); + layer.ffn2_in = binding::linear_from_source( + store, + source, + prefix + ".ffn2.intermediate_dense", + matmul_storage_type, + kIntermediate, + kHidden, + true); + layer.ffn2_out = binding::linear_from_source( + store, + source, + prefix + ".ffn2.output_dense", + matmul_storage_type, + kHidden, + kIntermediate, + true); + layer.final_norm = binding::norm_from_source(store, source, prefix + ".final_layer_norm", kHidden); + return layer; +} + +std::vector wav2vec2bert_std(const engine::assets::TensorSource & source) { + const auto var = source.require_f32("var", {kHidden}); + std::vector stddev(static_cast(kHidden), 0.0F); + for (int64_t i = 0; i < kHidden; ++i) { + stddev[static_cast(i)] = std::sqrt(var[static_cast(i)]); + } + return stddev; +} + +std::vector make_distance_ids(int64_t frames) { + std::vector ids(static_cast(frames * frames), 0); + for (int64_t row = 0; row < frames; ++row) { + for (int64_t col = 0; col < frames; ++col) { + const int64_t distance = std::max(-kRelativeLeft, std::min(kRelativeRight, col - row)); + ids[static_cast(row * frames + col)] = static_cast(distance + kRelativeLeft); + } + } + return ids; +} + +std::vector make_attention_mask(const std::vector & keep_mask, int64_t frames) { + if (static_cast(keep_mask.size()) != frames) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT attention mask length mismatch"); + } + std::vector mask(static_cast(kHeads * frames * frames), 0.0F); + const float hidden = std::numeric_limits::lowest(); + for (int64_t head = 0; head < kHeads; ++head) { + for (int64_t row = 0; row < frames; ++row) { + for (int64_t col = 0; col < frames; ++col) { + if (keep_mask[static_cast(col)] == 0) { + mask[static_cast((head * frames + row) * frames + col)] = hidden; + } + } + } + } + return mask; +} + +std::vector make_keep_mask_f32(const std::vector & keep_mask) { + std::vector out(keep_mask.size(), 0.0F); + for (size_t i = 0; i < keep_mask.size(); ++i) { + out[i] = keep_mask[i] == 0 ? 0.0F : 1.0F; + } + return out; +} + +} // namespace + +class IndexTTS25Wav2Vec2BertRuntime::Graph { +public: + Graph( + engine::core::ExecutionContext & execution, + std::shared_ptr weights, + int64_t frames, + size_t graph_arena_bytes) + : execution_(execution), + weights_(std::move(weights)), + frames_(frames) { + if (frames_ <= 0) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph requires positive frame count"); + } + if (weights_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph requires weights"); + } + + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Wav2Vec2-BERT graph context"); + } + ggml_init_params input_params{64ull * 1024ull * 1024ull, nullptr, true}; + input_ctx_.reset(ggml_init(input_params)); + if (input_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 Wav2Vec2-BERT input context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.wav2vec2bert", execution_.backend_type()}; + core::ModuleBuildContext input_ctx{ + input_ctx_.get(), + "index_tts2_5.wav2vec2bert.inputs", + execution_.backend_type()}; + + features_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kFeatureDim})).tensor; + keep_mask_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({frames_})).tensor; + attention_mask_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kHeads, frames_, frames_})).tensor; + distance_ids_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({frames_, frames_})).tensor; + ggml_set_input(features_); + ggml_set_input(keep_mask_); + ggml_set_input(attention_mask_); + ggml_set_input(distance_ids_); + + auto x = core::wrap_tensor(features_, core::TensorShape::from_dims({1, frames_, kFeatureDim}), GGML_TYPE_F32); + auto keep = core::wrap_tensor(keep_mask_, core::TensorShape::from_dims({frames_}), GGML_TYPE_F32); + auto attn_mask = core::wrap_tensor(attention_mask_, core::TensorShape::from_dims({1, kHeads, frames_, frames_}), GGML_TYPE_F32); + auto distances = core::wrap_tensor(distance_ids_, core::TensorShape::from_dims({frames_, frames_}), GGML_TYPE_I32); + + x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights_->feature_norm); + x = modules::LinearModule({kFeatureDim, kHidden, true, GGML_PREC_F32}).build(ctx, x, weights_->feature_projection); + x = apply_keep_mask(ctx, x, keep); + for (int64_t layer = 0; layer < kSemanticOutputHiddenStateIndex; ++layer) { + x = wav2vec2bert_layer(ctx, x, keep, attn_mask, distances, weights_->layers[static_cast(layer)]); + } + const auto mean = broadcast_vector(ctx, weights_->semantic_mean, x); + const auto std = broadcast_vector(ctx, weights_->semantic_std, x); + x = core::wrap_tensor(ggml_sub(ctx.ggml, x.tensor, mean.tensor), x.shape, GGML_TYPE_F32); + x = core::wrap_tensor(ggml_div(ctx.ggml, x.tensor, std.tensor), x.shape, GGML_TYPE_F32); + output_ = core::ensure_backend_addressable_layout(ctx, x).tensor; + ggml_set_output(output_); + + const int64_t graph_nodes = std::max( + 65536, + kSemanticOutputHiddenStateIndex * (frames_ * kHeads * 8 + frames_ * 16 + 1024) + 8192); + graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(graph_nodes), false); + ggml_build_forward_expand(graph_, output_); + input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); + if (input_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate IndexTTS2.5 Wav2Vec2-BERT input buffer"); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 Wav2Vec2-BERT graph"); + } + + const auto distances_host = make_distance_ids(frames_); + ggml_backend_tensor_set(distance_ids_, distances_host.data(), 0, distances_host.size() * sizeof(int32_t)); + debug::timing_log_scalar( + "index_tts2_5.wav2vec2bert.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("index_tts2_5.wav2vec2bert.frames", frames_); + } + + ~Graph() { + clear_graph(); + } + + int64_t frames() const noexcept { + return frames_; + } + + IndexTTS25SemanticEmbedding run(const IndexTTS25SemanticFeatureOutput & features) { + if (features.frames <= 0 || features.frames > frames_ || features.dims != kFeatureDim) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT feature shape does not match prepared graph"); + } + if (static_cast(features.values.size()) != features.frames * kFeatureDim) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT feature value count mismatch"); + } + if (static_cast(features.attention_mask.size()) != features.frames) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT feature attention mask count mismatch"); + } + + auto timing_start = Clock::now(); + std::vector padded_features(static_cast(frames_ * kFeatureDim), 0.0F); + std::copy(features.values.begin(), features.values.end(), padded_features.begin()); + std::vector padded_mask(static_cast(frames_), 0); + std::copy(features.attention_mask.begin(), features.attention_mask.end(), padded_mask.begin()); + const auto keep = make_keep_mask_f32(padded_mask); + const auto attention = make_attention_mask(padded_mask, frames_); + ggml_backend_tensor_set(features_, padded_features.data(), 0, padded_features.size() * sizeof(float)); + ggml_backend_tensor_set(keep_mask_, keep.data(), 0, keep.size() * sizeof(float)); + ggml_backend_tensor_set(attention_mask_, attention.data(), 0, attention.size() * sizeof(float)); + debug::timing_log_scalar( + "index_tts2_5.wav2vec2bert.input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + + core::set_backend_threads(execution_.backend(), execution_.config().threads); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(execution_.backend(), graph_); + ggml_backend_synchronize(execution_.backend()); + debug::timing_log_scalar( + "index_tts2_5.wav2vec2bert.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph compute failed"); + } + IndexTTS25SemanticEmbedding out; + out.frames = features.frames; + out.dims = kHidden; + out.values.resize(static_cast(features.frames * kHidden)); + timing_start = Clock::now(); + ggml_backend_tensor_get(output_, out.values.data(), 0, out.values.size() * sizeof(float)); + debug::timing_log_scalar( + "index_tts2_5.wav2vec2bert.output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return out; + } + +private: + void clear_graph() { + if (graph_ != nullptr) { + engine::core::release_backend_graph_resources(execution_.backend(), graph_); + graph_ = nullptr; + } + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + if (input_buffer_ != nullptr) { + ggml_backend_buffer_free(input_buffer_); + input_buffer_ = nullptr; + } + } + + engine::core::ExecutionContext & execution_; + std::shared_ptr weights_; + int64_t frames_ = 0; + std::unique_ptr input_ctx_; + std::unique_ptr ctx_; + ggml_tensor * features_ = nullptr; + ggml_tensor * keep_mask_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + ggml_tensor * distance_ids_ = nullptr; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t input_buffer_ = nullptr; +}; + +std::shared_ptr load_index_tts2_5_wav2vec2bert_weights( + const IndexTTS25Assets & assets, + ggml_backend_t backend, + engine::core::BackendType backend_type, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type, + size_t weight_context_bytes) { + if (assets.wav2vec2bert_weights == nullptr || assets.wav2vec2bert_stats == nullptr) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT requires model and stats tensor sources"); + } + auto weights = std::make_shared(); + weights->store = std::make_shared( + backend, + backend_type, + "index_tts2_5.wav2vec2bert.weights", + weight_context_bytes); + + const auto & source = *assets.wav2vec2bert_weights; + weights->feature_norm = binding::norm_from_source(*weights->store, source, "feature_projection.layer_norm", kFeatureDim); + weights->feature_projection = binding::linear_from_source( + *weights->store, + source, + "feature_projection.projection", + matmul_storage_type, + kHidden, + kFeatureDim, + true); + weights->layers.reserve(static_cast(kLayers)); + for (int64_t layer_index = 0; layer_index < kLayers; ++layer_index) { + weights->layers.push_back(load_layer( + *weights->store, + source, + layer_index, + matmul_storage_type, + conv_storage_type)); + } + weights->semantic_mean = weights->store->make_from_f32( + engine::core::TensorShape::from_dims({kHidden}), + matmul_storage_type, + assets.wav2vec2bert_stats->require_f32("mean", {kHidden})); + weights->semantic_std = weights->store->make_from_f32( + engine::core::TensorShape::from_dims({kHidden}), + matmul_storage_type, + wav2vec2bert_std(*assets.wav2vec2bert_stats)); + + weights->store->upload(); + assets.wav2vec2bert_weights->release_storage(); + assets.wav2vec2bert_stats->release_storage(); + return weights; +} + +IndexTTS25Wav2Vec2BertRuntime::IndexTTS25Wav2Vec2BertRuntime( + std::shared_ptr assets, + engine::core::ExecutionContext & execution, + size_t graph_arena_bytes, + size_t weight_context_bytes, + engine::assets::TensorStorageType matmul_storage_type, + engine::assets::TensorStorageType conv_storage_type) + : assets_(std::move(assets)), + execution_(&execution), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT runtime requires assets"); + } + if (graph_arena_bytes_ == 0) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph arena must be non-zero"); + } + weights_ = load_index_tts2_5_wav2vec2bert_weights( + *assets_, + execution.backend(), + execution.backend_type(), + matmul_storage_type, + conv_storage_type, + weight_context_bytes); +} + +IndexTTS25Wav2Vec2BertRuntime::~IndexTTS25Wav2Vec2BertRuntime() = default; + +void IndexTTS25Wav2Vec2BertRuntime::prepare(int64_t frames) { + if (execution_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT runtime execution context is missing"); + } + if (frames <= 0) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT prepare requires positive frames"); + } + if (graph_ != nullptr && graph_->frames() >= frames) { + return; + } + graph_.reset(); + graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); +} + +IndexTTS25SemanticEmbedding IndexTTS25Wav2Vec2BertRuntime::encode(const IndexTTS25SemanticFeatureOutput & features) { + if (graph_ == nullptr || graph_->frames() < features.frames) { + throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph was not prepared for this reference length"); + } + return graph_->run(features); +} + +void IndexTTS25Wav2Vec2BertRuntime::release_graph() { + graph_.reset(); +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/session.cpp b/src/models/index_tts2_5/session.cpp new file mode 100644 index 00000000..8fd34df7 --- /dev/null +++ b/src/models/index_tts2_5/session.cpp @@ -0,0 +1,838 @@ +#include "engine/models/index_tts2_5/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/io/text.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/text/chunking.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +using Clock = std::chrono::steady_clock; +constexpr int64_t kConditionDim = 512; +constexpr int64_t kGptDim = 1280; +constexpr int64_t kStyleDim = 192; +constexpr int64_t kEmotionCount = 8; +constexpr int64_t kDiffusionSteps = 25; +constexpr float kInferenceCfgRate = 0.7F; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("IndexTTS2.5 session requires assets"); + } + return assets; +} + +void validate_matmul_weight_storage(engine::assets::TensorStorageType storage_type, const char * option_name) { + if (storage_type == engine::assets::TensorStorageType::Native || + storage_type == engine::assets::TensorStorageType::F32 || + storage_type == engine::assets::TensorStorageType::F16 || + storage_type == engine::assets::TensorStorageType::BF16 || + storage_type == engine::assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " supports only native, f32, f16, bf16, and q8_0"); +} + +void validate_conv_weight_storage(engine::assets::TensorStorageType storage_type, const char * option_name) { + if (storage_type == engine::assets::TensorStorageType::Native || + storage_type == engine::assets::TensorStorageType::F32 || + storage_type == engine::assets::TensorStorageType::F16) { + return; + } + throw std::runtime_error(std::string(option_name) + " supports only native, f32, and f16"); +} + +uint64_t fnv1a_mix(uint64_t hash, const void * data, size_t size) { + const auto * bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + hash ^= bytes[i]; + hash *= 1099511628211ull; + } + return hash; +} + +uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { + uint64_t hash = 1469598103934665603ull; + for (const float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + hash = fnv1a_mix(hash, &bits, sizeof(bits)); + } + return hash; +} + +IndexTTS25AudioIdentity audio_identity(const runtime::AudioBuffer & audio) { + return { + audio.sample_rate, + audio.channels, + static_cast(audio.samples.size()), + hash_audio_samples(audio), + }; +} + +bool same_identity(const IndexTTS25AudioIdentity & lhs, const IndexTTS25AudioIdentity & rhs) { + return lhs.sample_rate == rhs.sample_rate && + lhs.channels == rhs.channels && + lhs.sample_count == rhs.sample_count && + lhs.sample_hash == rhs.sample_hash; +} + +std::size_t resolve_cache_slots( + const runtime::SessionOptions & options, + std::initializer_list keys, + const char * option_name) { + constexpr int64_t kDefaultCacheSlots = 1; + const int64_t slots = runtime::parse_i64_option(options.options, keys) + .value_or(kDefaultCacheSlots); + if (slots < 0) { + throw std::runtime_error(std::string(option_name) + " must be non-negative"); + } + if (static_cast(slots) > static_cast(std::numeric_limits::max())) { + throw std::runtime_error(std::string(option_name) + " is too large"); + } + return static_cast(slots); +} + +std::vector channel_first_to_time_major( + const std::vector & values, + int64_t channels, + int64_t frames) { + if (static_cast(values.size()) != channels * frames) { + throw std::runtime_error("IndexTTS2.5 channel-first tensor size mismatch"); + } + std::vector out(static_cast(frames * channels)); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t channel = 0; channel < channels; ++channel) { + out[static_cast(frame * channels + channel)] = + values[static_cast(channel * frames + frame)]; + } + } + return out; +} + +std::vector concat_conditions( + const IndexTTS25S2MelSequence & prompt, + const IndexTTS25S2MelSequence & generated) { + if (prompt.dims != kConditionDim || generated.dims != kConditionDim) { + throw std::runtime_error("IndexTTS2.5 condition dimension mismatch"); + } + std::vector out; + out.reserve(prompt.values.size() + generated.values.size()); + out.insert(out.end(), prompt.values.begin(), prompt.values.end()); + out.insert(out.end(), generated.values.begin(), generated.values.end()); + return out; +} + +void append_silence(runtime::AudioBuffer & audio, int ms) { + if (ms <= 0 || audio.sample_rate <= 0 || audio.channels <= 0) { + return; + } + const int64_t samples = static_cast(audio.sample_rate) * ms / 1000; + audio.samples.insert( + audio.samples.end(), + static_cast(samples * audio.channels), + 0.0F); +} + +std::vector scaled_emotion_weights(const std::vector & values, float alpha) { + if (static_cast(values.size()) != kEmotionCount) { + throw std::runtime_error("IndexTTS2.5 emotion vector must contain exactly 8 values"); + } + std::vector out = values; + const float scale = std::clamp(alpha, 0.0F, 1.0F); + if (scale != 1.0F) { + for (float & value : out) { + value = static_cast(static_cast(value * scale * 10000.0F)) / 10000.0F; + } + } + return out; +} + +bool mem_saver_from_options(const runtime::SessionOptions & options) { + if (const auto value = runtime::find_option(options.options, {"index_tts2_5.mem_saver", "mem_saver"})) { + return runtime::parse_bool_option(*value, "index_tts2_5.mem_saver"); + } + return false; +} + +// Debug intermediate dumps, enabled by setting INDEXTTS25_DUMP_DIR. Files are +// written as NPY v1.0 so they can be compared against the official Python +// golden tensors with numpy directly. +std::string dump_dir_from_env() { + const char * dir = std::getenv("INDEXTTS25_DUMP_DIR"); + if (dir == nullptr || *dir == '\0') { + return {}; + } + return std::string(dir); +} + +void write_npy( + const std::filesystem::path & path, + const char * descr, + const std::vector & shape, + const void * data, + size_t byte_count) { + std::string header = "{'descr': '"; + header += descr; + header += "', 'fortran_order': False, 'shape': ("; + for (const int64_t dim : shape) { + header += std::to_string(dim); + header += ", "; + } + header += "), }"; + const size_t prefix = 10; // magic(6) + version(2) + header_len(2) + const size_t total = prefix + header.size() + 1; + const size_t padded = (total + 63) / 64 * 64; + header.append(padded - prefix - header.size() - 1, ' '); + header.push_back('\n'); + std::ofstream out(path, std::ios::binary | std::ios::trunc); + if (!out) { + throw std::runtime_error("IndexTTS2.5 failed to open dump file: " + path.string()); + } + out.write("\x93NUMPY\x01\x00", 8); + const uint16_t header_len = static_cast(header.size()); + out.write(reinterpret_cast(&header_len), sizeof(header_len)); + out.write(header.data(), static_cast(header.size())); + out.write(static_cast(data), static_cast(byte_count)); +} + +void dump_f32( + const std::string & dir, + const std::string & name, + std::vector shape, + const std::vector & values) { + if (dir.empty()) { + return; + } + int64_t elements = 1; + for (const int64_t dim : shape) { + elements *= dim; + } + if (elements != static_cast(values.size())) { + throw std::runtime_error("IndexTTS2.5 dump shape does not match value count: " + name); + } + write_npy(std::filesystem::path(dir) / (name + ".npy"), " shape, + const std::vector & values) { + if (dir.empty()) { + return; + } + int64_t elements = 1; + for (const int64_t dim : shape) { + elements *= dim; + } + if (elements != static_cast(values.size())) { + throw std::runtime_error("IndexTTS2.5 dump shape does not match value count: " + name); + } + write_npy(std::filesystem::path(dir) / (name + ".npy"), " assets) + : RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + tokenizer_(assets_), + speaker_cache_(resolve_cache_slots(this->options(), {"index_tts2_5.speaker_cache_slots"}, "index_tts2_5.speaker_cache_slots")), + emotion_cache_(resolve_cache_slots(this->options(), {"index_tts2_5.emotion_cache_slots"}, "index_tts2_5.emotion_cache_slots")), + emotion_text_weights_cache_(resolve_cache_slots(this->options(), {"index_tts2_5.emotion_text_cache_slots"}, "index_tts2_5.emotion_text_cache_slots")) { + gpt_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"index_tts2_5.gpt_graph_arena_mb"}, gpt_graph_arena_bytes_); + s2mel_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"index_tts2_5.s2mel_graph_arena_mb"}, s2mel_graph_arena_bytes_); + reference_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"index_tts2_5.reference_graph_arena_mb"}, reference_graph_arena_bytes_); + emotion_text_prefill_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"index_tts2_5.emotion_text_prefill_graph_arena_mb"}, emotion_text_prefill_graph_arena_bytes_); + emotion_text_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"index_tts2_5.emotion_text_decode_graph_arena_mb"}, emotion_text_decode_graph_arena_bytes_); + weight_context_bytes_ = runtime::parse_size_mb_option( + options.options, {"index_tts2_5.weight_context_mb"}, weight_context_bytes_); + if (const auto value = runtime::parse_int_option(options.options, {"index_tts2_5.emotion_text_max_new_tokens"})) { + if (*value <= 0) { + throw std::runtime_error("index_tts2_5.emotion_text_max_new_tokens must be positive"); + } + emotion_text_max_new_tokens_ = *value; + } + if (const auto it = options.options.find("index_tts2_5.weight_type"); it != options.options.end()) { + matmul_weight_storage_type_ = engine::assets::parse_tensor_storage_type(it->second); + validate_matmul_weight_storage(matmul_weight_storage_type_, "index_tts2_5.weight_type"); + } + if (const auto it = options.options.find("index_tts2_5.conv_weight_type"); it != options.options.end()) { + conv_weight_storage_type_ = engine::assets::parse_tensor_storage_type(it->second); + validate_conv_weight_storage(conv_weight_storage_type_, "index_tts2_5.conv_weight_type"); + } + mem_saver_ = mem_saver_from_options(options); + for (const auto & [key, _] : options.options) { + if (key.rfind("index_tts2_5.", 0) == 0 && + key != "index_tts2_5.gpt_graph_arena_mb" && + key != "index_tts2_5.s2mel_graph_arena_mb" && + key != "index_tts2_5.reference_graph_arena_mb" && + key != "index_tts2_5.emotion_text_prefill_graph_arena_mb" && + key != "index_tts2_5.emotion_text_decode_graph_arena_mb" && + key != "index_tts2_5.emotion_text_max_new_tokens" && + key != "index_tts2_5.weight_context_mb" && + key != "index_tts2_5.weight_type" && + key != "index_tts2_5.conv_weight_type" && + key != "index_tts2_5.mem_saver" && + key != "index_tts2_5.speaker_cache_slots" && + key != "index_tts2_5.emotion_cache_slots" && + key != "index_tts2_5.emotion_text_cache_slots" && + key != "index_tts2_5.gpt.cuda_sampling_policy" && + key != "index_tts2_5.s2mel.cuda_sampling_policy") { + throw std::runtime_error("unknown IndexTTS2.5 session option: " + key); + } + } + if (task_.mode != runtime::RunMode::Offline || + (task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning)) { + throw std::runtime_error("IndexTTS2.5 currently supports offline TTS and voice-cloning sessions"); + } + + semantic_encoder_ = std::make_unique( + assets_, + execution_context(), + reference_graph_arena_bytes_, + weight_context_bytes_, + matmul_weight_storage_type_, + conv_weight_storage_type_); + semantic_codec_ = std::make_unique( + assets_, + execution_context(), + reference_graph_arena_bytes_, + weight_context_bytes_, + matmul_weight_storage_type_, + conv_weight_storage_type_); + style_encoder_ = std::make_unique( + assets_, + options.backend, + conv_weight_storage_type_); + gpt_ = std::make_unique( + assets_, + execution_context(), + gpt_graph_arena_bytes_, + weight_context_bytes_, + matmul_weight_storage_type_, + conv_weight_storage_type_); + s2mel_ = std::make_unique( + assets_, + execution_context(), + s2mel_graph_arena_bytes_, + weight_context_bytes_, + matmul_weight_storage_type_, + conv_weight_storage_type_); + vocoder_ = std::make_unique( + assets_, + options.backend, + conv_weight_storage_type_); + qwen_emotion_ = std::make_unique( + assets_, + execution_context(), + emotion_text_prefill_graph_arena_bytes_, + emotion_text_decode_graph_arena_bytes_, + weight_context_bytes_, + matmul_weight_storage_type_); + int64_t matrix_rows = 0; + for (const int64_t count : assets_->config.emo_num) { + matrix_rows += count; + } + speaker_matrix_ = assets_->speaker_matrix->require_f32("tensor", {matrix_rows, kStyleDim}); + emotion_matrix_ = assets_->emotion_matrix->require_f32("tensor", {matrix_rows, kGptDim}); + assets_->speaker_matrix->release_storage(); + assets_->emotion_matrix->release_storage(); +} + +bool IndexTTS25Session::AudioIdentityEqual::operator()( + const IndexTTS25AudioIdentity & lhs, + const IndexTTS25AudioIdentity & rhs) const { + return same_identity(lhs, rhs); +} + +std::string IndexTTS25Session::family() const { + return "index_tts2_5"; +} + +runtime::VoiceTaskKind IndexTTS25Session::task_kind() const { + return task_.task; +} + +runtime::RunMode IndexTTS25Session::run_mode() const { + return task_.mode; +} + +void IndexTTS25Session::prepare(const runtime::SessionPreparationRequest & request) { + if (request.text.has_value() || runtime::find_option(request.options, {"text", "prompt"}).has_value()) { + std::string text; + if (request.text.has_value()) { + text = request.text->text; + } else if (const auto value = runtime::find_option(request.options, {"text", "prompt"})) { + text = *value; + } + text = engine::io::trim_ascii_whitespace(text); + if (text.empty()) { + throw std::runtime_error("IndexTTS2.5 request requires text_input or text option"); + } + IndexTTS25GenerationOptions generation; + if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { + if (*value <= 0) { + throw std::runtime_error("IndexTTS2.5 max_tokens must be positive"); + } + generation.max_mel_tokens = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"num_beams"})) { + if (*value <= 0) { + throw std::runtime_error("IndexTTS2.5 num_beams must be positive"); + } + generation.num_beams = *value; + } + std::string lang; + if (const auto value = runtime::find_option(request.options, {"lang"})) { + lang = normalize_index_tts2_5_lang(*value); + } + const auto text_encoding = tokenizer_.encode_for_inference( + text, + IndexTTS25Request{}.max_text_tokens_per_segment, + lang); + for (const auto & segment : text_encoding.segment_token_ids) { + gpt_->prepare_generation( + static_cast(align_index_tts2_5_gpt_text_tokens(segment).size()), + generation.max_mel_tokens, + generation.num_beams); + } + } + mark_prepared(); +} + +const IndexTTS25Session::SpeakerState & IndexTTS25Session::resolve_speaker_state(const runtime::AudioBuffer & audio) { + const auto identity = audio_identity(audio); + if (const auto * cached = speaker_cache_.find(identity)) { + debug::trace_log_scalar("index_tts2_5.speaker_cache.hit", 1); + debug::trace_log_scalar("index_tts2_5.speaker_cache.slots", static_cast(speaker_cache_.capacity())); + debug::trace_log_scalar("index_tts2_5.speaker_cache.entries", static_cast(speaker_cache_.size())); + debug::trace_log_scalar("index_tts2_5.speaker_cache.evicted", 0); + return *cached; + } + const bool will_evict = speaker_cache_.capacity() > 0 && speaker_cache_.size() >= speaker_cache_.capacity(); + const auto start = Clock::now(); + const auto prepared = prepare_index_tts2_5_reference_audio( + audio.samples, + audio.sample_rate, + audio.channels, + assets_->config.s2mel, + static_cast(std::max(1, options().backend.threads)), + true); + semantic_encoder_->prepare(prepared.semantic_features.frames); + auto semantic = semantic_encoder_->encode(prepared.semantic_features); + // Official 2.5 regulates the raw (normalized) w2v-bert semantic directly; + // the semantic codec is only used to decode generated codes. + debug::trace_log_scalar("index_tts2_5.s2mel.reference_mel_frames", static_cast(prepared.mel.frames)); + s2mel_->prepare_length_regulator(semantic.frames, prepared.mel.frames); + auto prompt_condition = s2mel_->regulate_length( + semantic.values, + semantic.frames, + prepared.mel.frames); + + SpeakerState state; + state.identity = identity; + state.semantic = std::move(semantic); + state.reference_mel = prepared.mel; + state.style = style_encoder_->embed_fbank( + prepared.campplus_fbank.values, + prepared.campplus_fbank.frames, + prepared.campplus_fbank.dims); + state.prompt_condition = std::move(prompt_condition); + if (speaker_cache_.capacity() == 0) { + uncached_speaker_state_ = std::move(state); + } else { + speaker_cache_.put(identity, std::move(state)); + } + if (mem_saver_) { + semantic_encoder_->release_graph(); + semantic_codec_->release_graphs(); + s2mel_->release_pre_cfm_graphs(); + style_encoder_->release_graph(); + } + debug::trace_log_scalar("index_tts2_5.speaker_cache.hit", 0); + debug::trace_log_scalar("index_tts2_5.speaker_cache.slots", static_cast(speaker_cache_.capacity())); + debug::trace_log_scalar("index_tts2_5.speaker_cache.entries", static_cast(speaker_cache_.size())); + debug::trace_log_scalar("index_tts2_5.speaker_cache.evicted", will_evict ? 1 : 0); + debug::timing_log_scalar("index_tts2_5.speaker_state_ms", engine::debug::elapsed_ms(start)); + if (speaker_cache_.capacity() == 0) { + return *uncached_speaker_state_; + } + const auto * cached = speaker_cache_.find(identity); + if (cached == nullptr) { + throw std::runtime_error("IndexTTS2.5 speaker cache insert failed"); + } + return *cached; +} + +const IndexTTS25Session::EmotionState & IndexTTS25Session::resolve_emotion_state(const runtime::AudioBuffer & audio) { + const auto identity = audio_identity(audio); + if (const auto * cached = emotion_cache_.find(identity)) { + debug::trace_log_scalar("index_tts2_5.emotion_cache.hit", 1); + debug::trace_log_scalar("index_tts2_5.emotion_cache.slots", static_cast(emotion_cache_.capacity())); + debug::trace_log_scalar("index_tts2_5.emotion_cache.entries", static_cast(emotion_cache_.size())); + debug::trace_log_scalar("index_tts2_5.emotion_cache.evicted", 0); + return *cached; + } + const bool will_evict = emotion_cache_.capacity() > 0 && emotion_cache_.size() >= emotion_cache_.capacity(); + const auto start = Clock::now(); + const auto prepared = prepare_index_tts2_5_reference_audio( + audio.samples, + audio.sample_rate, + audio.channels, + assets_->config.s2mel, + static_cast(std::max(1, options().backend.threads)), + false); + semantic_encoder_->prepare(prepared.semantic_features.frames); + EmotionState state; + state.identity = identity; + state.semantic = semantic_encoder_->encode(prepared.semantic_features); + if (emotion_cache_.capacity() == 0) { + uncached_emotion_state_ = std::move(state); + } else { + emotion_cache_.put(identity, std::move(state)); + } + if (mem_saver_) { + semantic_encoder_->release_graph(); + } + debug::trace_log_scalar("index_tts2_5.emotion_cache.hit", 0); + debug::trace_log_scalar("index_tts2_5.emotion_cache.slots", static_cast(emotion_cache_.capacity())); + debug::trace_log_scalar("index_tts2_5.emotion_cache.entries", static_cast(emotion_cache_.size())); + debug::trace_log_scalar("index_tts2_5.emotion_cache.evicted", will_evict ? 1 : 0); + debug::timing_log_scalar("index_tts2_5.emotion_state_ms", engine::debug::elapsed_ms(start)); + if (emotion_cache_.capacity() == 0) { + return *uncached_emotion_state_; + } + const auto * cached = emotion_cache_.find(identity); + if (cached == nullptr) { + throw std::runtime_error("IndexTTS2.5 emotion cache insert failed"); + } + return *cached; +} + +std::vector IndexTTS25Session::explicit_emotion_matrix_vector( + const std::vector & emotion_weights, + const IndexTTS25StyleEmbedding & style, + bool use_random, + uint32_t seed) const { + if (static_cast(emotion_weights.size()) != kEmotionCount || + style.dims != kStyleDim || + static_cast(style.values.size()) != kStyleDim) { + throw std::runtime_error("IndexTTS2.5 explicit emotion matrix input shape mismatch"); + } + std::vector out(static_cast(kGptDim), 0.0F); + int64_t row_offset = 0; + std::mt19937 rng(seed); + for (int64_t emotion = 0; emotion < kEmotionCount; ++emotion) { + const int64_t rows = assets_->config.emo_num[static_cast(emotion)]; + int64_t best_row = 0; + if (rows <= 0) { + throw std::runtime_error("IndexTTS2.5 emo_num contains non-positive group size"); + } + if (use_random) { + std::uniform_int_distribution distribution(0, rows - 1); + best_row = distribution(rng); + } else { + float best_score = -std::numeric_limits::infinity(); + for (int64_t row = 0; row < rows; ++row) { + const float * matrix_row = speaker_matrix_.data() + static_cast((row_offset + row) * kStyleDim); + float dot = 0.0F; + float lhs_norm = 0.0F; + float rhs_norm = 0.0F; + for (int64_t dim = 0; dim < kStyleDim; ++dim) { + const float lhs = style.values[static_cast(dim)]; + const float rhs = matrix_row[dim]; + dot += lhs * rhs; + lhs_norm += lhs * lhs; + rhs_norm += rhs * rhs; + } + const float score = dot / (std::sqrt(lhs_norm) * std::sqrt(rhs_norm) + 1.0e-12F); + if (score > best_score) { + best_score = score; + best_row = row; + } + } + } + const float * emotion_row = emotion_matrix_.data() + static_cast((row_offset + best_row) * kGptDim); + for (int64_t dim = 0; dim < kGptDim; ++dim) { + out[static_cast(dim)] += emotion_weights[static_cast(emotion)] * emotion_row[dim]; + } + row_offset += rows; + } + return out; +} + +std::vector IndexTTS25Session::resolve_emotion_vector( + const IndexTTS25Request & request, + const SpeakerState & speaker, + const EmotionState & emotion) { + std::vector explicit_weights; + if (request.use_emotion_text) { + const auto emotion_text = request.emotion_text.value_or(request.text); + if (const auto * cached = emotion_text_weights_cache_.find(emotion_text)) { + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.hit", 1); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.slots", static_cast(emotion_text_weights_cache_.capacity())); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.entries", static_cast(emotion_text_weights_cache_.size())); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.evicted", 0); + explicit_weights = *cached; + } else { + const bool will_evict = + emotion_text_weights_cache_.capacity() > 0 && + emotion_text_weights_cache_.size() >= emotion_text_weights_cache_.capacity(); + explicit_weights = qwen_emotion_->infer(emotion_text, emotion_text_max_new_tokens_).values; + if (mem_saver_) { + qwen_emotion_->release_graphs(); + } + emotion_text_weights_cache_.put(emotion_text, explicit_weights); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.hit", 0); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.slots", static_cast(emotion_text_weights_cache_.capacity())); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.entries", static_cast(emotion_text_weights_cache_.size())); + debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.evicted", will_evict ? 1 : 0); + } + } else if (request.emotion_vector.has_value()) { + explicit_weights = *request.emotion_vector; + } + + if (explicit_weights.empty()) { + return gpt_->merge_emotion_vector( + speaker.semantic.values, + speaker.semantic.frames, + emotion.semantic.values, + emotion.semantic.frames, + request.emotion_alpha); + } + + auto scaled_weights = scaled_emotion_weights(explicit_weights, request.emotion_alpha); + auto matrix = explicit_emotion_matrix_vector( + scaled_weights, + speaker.style, + request.use_random_emotion, + request.generation.seed); + const float weight_sum = std::accumulate(scaled_weights.begin(), scaled_weights.end(), 0.0F); + if (weight_sum != 1.0F) { + auto base = gpt_->merge_emotion_vector( + speaker.semantic.values, + speaker.semantic.frames, + emotion.semantic.values, + emotion.semantic.frames, + 1.0F); + for (size_t i = 0; i < matrix.size(); ++i) { + matrix[i] += (1.0F - weight_sum) * base[i]; + } + } + return matrix; +} + +runtime::AudioBuffer IndexTTS25Session::synthesize_segment( + const std::vector & text_tokens, + int32_t lang_id, + size_t segment_index, + const std::string & dump_dir, + const SpeakerState & speaker, + const EmotionState & emotion, + const std::vector & emotion_vector, + const IndexTTS25GenerationOptions & options, + uint32_t segment_seed) { + const std::string seg = "_seg" + std::to_string(segment_index); + dump_i32(dump_dir, "text_tokens" + seg, {1, static_cast(text_tokens.size())}, text_tokens); + IndexTTS25GptGenerationRequest generation; + generation.text_tokens = text_tokens; + generation.speaker_style = speaker.style.values; + generation.lang_id = lang_id; + generation.emotion_semantic = emotion.semantic.values; + generation.emotion_frames = emotion.semantic.frames; + generation.emotion_vector = emotion_vector; + generation.top_p = options.top_p; + generation.top_k = options.top_k; + generation.temperature = options.temperature; + generation.repetition_penalty = options.repetition_penalty; + generation.do_sample = options.do_sample; + generation.length_penalty = options.length_penalty; + generation.num_beams = options.num_beams; + generation.max_mel_tokens = options.max_mel_tokens; + generation.seed = segment_seed; + + const auto gen_start = Clock::now(); + auto generated = gpt_->generate_speech(generation); + if (mem_saver_) { + gpt_->release_conditioning_graphs(); + } + debug::timing_log_scalar("index_tts2_5.gpt.generate_ms", engine::debug::elapsed_ms(gen_start)); + dump_i32(dump_dir, "gpt_codes" + seg, {1, static_cast(generated.codes.size())}, generated.codes); + if (generated.codes.empty()) { + throw std::runtime_error("IndexTTS2.5 GPT generated no acoustic codes"); + } + const int64_t code_frames = static_cast(generated.codes.size()); + + const auto s2mel_start = Clock::now(); + semantic_codec_->prepare_codes(code_frames); + auto semantic = semantic_codec_->codes_to_embedding(generated.codes, code_frames); + if (mem_saver_) { + semantic_codec_->release_graphs(); + } + auto content = channel_first_to_time_major( + semantic.embedding_channel_first, + semantic.dims, + semantic.frames); + dump_f32(dump_dir, "s_infer" + seg, {1, semantic.frames, semantic.dims}, content); + const int64_t target_frames = static_cast(static_cast(semantic.frames) * 1.72F); + const int64_t total_frames = speaker.prompt_condition.frames + target_frames; + s2mel_->prepare_length_regulator(semantic.frames, target_frames); + auto generated_condition = s2mel_->regulate_length(content, semantic.frames, target_frames); + dump_f32( + dump_dir, + "lr_gen" + seg, + {1, generated_condition.frames, generated_condition.dims}, + generated_condition.values); + auto condition = concat_conditions(speaker.prompt_condition, generated_condition); + dump_f32(dump_dir, "cat_condition" + seg, {1, total_frames, kConditionDim}, condition); + if (mem_saver_) { + gpt_->release_generation_graphs(); + s2mel_->release_pre_cfm_graphs(); + } + s2mel_->prepare_cfm(total_frames, kInferenceCfgRate > 0.0F); + auto mel = s2mel_->infer_mel( + condition, + total_frames, + speaker.reference_mel.values, + speaker.reference_mel.frames, + speaker.style.values, + kDiffusionSteps, + kInferenceCfgRate, + segment_seed, + generated.rng_offset_blocks); + dump_f32(dump_dir, "mel_out" + seg, {1, mel.channels, mel.frames}, mel.values); + debug::timing_log_scalar("index_tts2_5.s2mel.total_ms", engine::debug::elapsed_ms(s2mel_start)); + if (mem_saver_) { + s2mel_->release_cfm_graph(); + } + + const auto vocoder_start = Clock::now(); + auto audio = vocoder_->synthesize(mel.values, mel.frames); + debug::timing_log_scalar("index_tts2_5.vocoder_ms", engine::debug::elapsed_ms(vocoder_start)); + if (mem_saver_) { + vocoder_->release_runtime_graph(); + } + runtime::AudioBuffer out; + out.sample_rate = audio.sample_rate; + out.channels = 1; + out.samples = std::move(audio.waveform); + return out; +} + +runtime::TaskResult IndexTTS25Session::run(const runtime::TaskRequest & request) { + require_prepared("IndexTTS2.5 run"); + const auto wall_start = Clock::now(); + auto parsed = parse_index_tts2_5_request(request); + if (!parsed.speaker_audio.has_value()) { + throw std::runtime_error("IndexTTS2.5 request requires speaker audio"); + } + const runtime::AudioBuffer * emotion_audio = parsed.emotion_audio.has_value() + ? &*parsed.emotion_audio + : &*parsed.speaker_audio; + if (parsed.emotion_vector.has_value() || parsed.use_emotion_text) { + emotion_audio = &*parsed.speaker_audio; + } + + const auto & speaker = resolve_speaker_state(*parsed.speaker_audio); + const bool emotion_same_as_speaker = same_identity(audio_identity(*emotion_audio), speaker.identity); + debug::trace_log_scalar("index_tts2_5.emotion_audio.same_as_speaker", emotion_same_as_speaker); + const auto & emotion = resolve_emotion_state(*emotion_audio); + const auto emotion_vector = resolve_emotion_vector(parsed, speaker, emotion); + if (mem_saver_) { + gpt_->release_conditioning_graphs(); + } + const std::string dump_dir = dump_dir_from_env(); + dump_f32(dump_dir, "campplus_embedding", {1, speaker.style.dims}, speaker.style.values); + dump_f32(dump_dir, "speech_condition", {1, speaker.semantic.frames, speaker.semantic.dims}, speaker.semantic.values); + dump_f32( + dump_dir, + "prompt_condition", + {1, speaker.prompt_condition.frames, speaker.prompt_condition.dims}, + speaker.prompt_condition.values); + dump_f32(dump_dir, "ref_mel", {1, speaker.reference_mel.channels, speaker.reference_mel.frames}, speaker.reference_mel.values); + dump_f32(dump_dir, "emo_vec", {1, static_cast(emotion_vector.size())}, emotion_vector); + const auto text_chunk_size = engine::text::parse_text_chunk_size_override(request.options); + const auto text_chunk_mode = engine::text::parse_text_chunk_mode_override(request.options) + .value_or(engine::text::TextChunkMode::Default); + const std::vector text_chunks = text_chunk_size.has_value() + ? engine::text::split_text_chunks(parsed.text, *text_chunk_size, text_chunk_mode) + : std::vector{parsed.text}; + if (text_chunks.empty()) { + throw std::runtime_error("IndexTTS2.5 text chunking produced no chunks"); + } + + std::vector> segment_token_ids; + std::vector segment_lang_ids; + for (const auto & text_chunk : text_chunks) { + const auto text_encoding = tokenizer_.encode_for_inference( + text_chunk, + parsed.max_text_tokens_per_segment, + parsed.lang); + const int32_t lang_id = IndexTTS25TextTokenizer::lang_to_id(text_encoding.lang); + for (const auto & ids : text_encoding.segment_token_ids) { + segment_token_ids.push_back(ids); + segment_lang_ids.push_back(lang_id); + } + } + + runtime::AudioBuffer merged; + for (size_t i = 0; i < segment_token_ids.size(); ++i) { + if (i > 0) { + append_silence(merged, parsed.interval_silence_ms); + } + auto segment_audio = synthesize_segment( + segment_token_ids[i], + segment_lang_ids[i], + i, + dump_dir, + speaker, + emotion, + emotion_vector, + parsed.generation, + parsed.generation.seed + static_cast(i)); + runtime::append_audio_buffer(merged, segment_audio); + } + + runtime::TaskResult result; + result.audio_output = std::move(merged); + debug::trace_log_scalar("index_tts2_5.path.use_emotion_text", parsed.use_emotion_text); + debug::trace_log_scalar("index_tts2_5.path.has_emotion_vector", parsed.emotion_vector.has_value()); + debug::trace_log_scalar("index_tts2_5.path.has_emotion_audio", parsed.emotion_audio.has_value()); + if (text_chunk_size.has_value()) { + debug::trace_log_scalar("index_tts2_5.text_chunk_size", *text_chunk_size); + debug::trace_log_scalar("index_tts2_5.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + } + debug::trace_log_scalar("index_tts2_5.text_chunk_count", static_cast(text_chunks.size())); + debug::trace_log_scalar("index_tts2_5.text.segment_count", static_cast(segment_token_ids.size())); + debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + return result; +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/style_encoder.cpp b/src/models/index_tts2_5/style_encoder.cpp new file mode 100644 index 00000000..bb679438 --- /dev/null +++ b/src/models/index_tts2_5/style_encoder.cpp @@ -0,0 +1,38 @@ +#include "engine/models/index_tts2_5/style_encoder.h" + +#include +#include + +namespace engine::models::index_tts2_5 { + +IndexTTS25StyleEncoder::IndexTTS25StyleEncoder( + std::shared_ptr assets, + core::BackendConfig backend, + engine::assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 style encoder requires assets"); + } + engine::modules::CampplusEncoderConfig config; + config.feat_dim = assets_->config.s2mel.n_mels; + config.embedding_size = assets_->config.s2mel.style_dim; + config.weight_storage_type = weight_storage_type; + component_ = engine::modules::CampplusEncoderComponent::load_from_tensor_source( + assets_->campplus_weights, + std::move(backend), + config); +} + +IndexTTS25StyleEmbedding IndexTTS25StyleEncoder::embed_fbank( + const std::vector & features, + int64_t frames, + int64_t dims) const { + const auto out = component_.embed_from_features(features, frames, dims); + return {out.embedding, out.embedding_size}; +} + +void IndexTTS25StyleEncoder::release_graph() { + component_.release_runtime_graph(); +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/tokenizer_text.cpp b/src/models/index_tts2_5/tokenizer_text.cpp new file mode 100644 index 00000000..1c061d3e --- /dev/null +++ b/src/models/index_tts2_5/tokenizer_text.cpp @@ -0,0 +1,691 @@ +#include "engine/models/index_tts2_5/tokenizer_text.h" + +#include "engine/framework/text/chinese_normalization.h" +#include "engine/framework/text/text_normalization.h" + +#include "bpe-core.h" +#include "unicode.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +namespace vendor = llama_tokenizer_vendor; + +// IndexTTS-2.5 pads every text segment with a trailing token id 1. +constexpr int32_t kSegmentPadTokenId = 1; + +std::string decode_base64(const std::string & input) { + static const std::array table = [] { + std::array values{}; + values.fill(-1); + for (int i = 0; i < 26; ++i) { + values[static_cast('A' + i)] = static_cast(i); + values[static_cast('a' + i)] = static_cast(26 + i); + } + for (int i = 0; i < 10; ++i) { + values[static_cast('0' + i)] = static_cast(52 + i); + } + values[static_cast('+')] = 62; + values[static_cast('/')] = 63; + return values; + }(); + + std::string out; + int bits = 0; + int value = 0; + for (const unsigned char ch : input) { + if (ch == '=') { + break; + } + const int8_t digit = table[ch]; + if (digit < 0) { + throw std::runtime_error("IndexTTS2.5 tiktoken vocabulary contains invalid base64 token bytes"); + } + value = (value << 6) | digit; + bits += 6; + if (bits >= 8) { + bits -= 8; + out.push_back(static_cast((value >> bits) & 0xff)); + } + } + return out; +} + +// The vendored llama BPE runtime works in the GPT-2 byte-to-unicode domain +// (e.g. space becomes U+0120) so that byte-level merges, including tokens that +// split a UTF-8 codepoint, are reproduced exactly. tiktoken ranks are keyed by +// raw bytes, so every token is mapped once at load time. +std::string map_token_bytes(const std::string & bytes) { + std::string mapped; + for (const unsigned char byte : bytes) { + mapped += unicode_byte_to_utf8(byte); + } + return mapped; +} + +std::string pair_key(const std::string & left, const std::string & right) { + std::string key = left; + key.push_back('\0'); + key += right; + return key; +} + +// Language codes in the LANGUAGES order of indextts/utils/tokenizer.py. The +// first 99 entries double as the <|lang|> special tokens below; the remaining +// codes (plus the fallback "common") only index the GPT lang_embedding table. +const std::array kLanguages = { + "en", "zh", "de", "es", "ru", "ko", "fr", "ja", "pt", "tr", + "pl", "ca", "nl", "ar", "sv", "it", "id", "hi", "fi", "vi", + "he", "uk", "el", "ms", "cs", "ro", "da", "hu", "ta", "no", + "th", "ur", "hr", "bg", "lt", "la", "mi", "ml", "cy", "sk", + "te", "fa", "lv", "bn", "sr", "az", "sl", "kn", "et", "mk", + "br", "eu", "is", "hy", "ne", "mn", "bs", "kk", "sq", "sw", + "gl", "mr", "pa", "si", "km", "sn", "yo", "so", "af", "oc", + "ka", "be", "tg", "sd", "gu", "am", "yi", "lo", "uz", "fo", + "ht", "ps", "tk", "nn", "mt", "sa", "lb", "my", "bo", "tl", + "mg", "as", "tt", "haw", "ln", "ha", "ba", "jw", "su", +}; +const std::array kEmbeddingOnlyLanguages = { + "yue", "minnan", "wuyu", "dialect", "zh/en", "en/zh", "common", +}; +constexpr int32_t kCommonLangId = 105; + +void add_special_token(vendor::BpeVocabulary & vocab, const std::string & text, int32_t id) { + vocab.token_to_id.emplace(text, id); + vocab.id_to_token.emplace(id, vendor::TokenData{text, vendor::TOKEN_ATTR_CONTROL}); +} + +// Special token order must match indextts/utils/tokenizer.py exactly: +// ids are assigned sequentially starting right after the mergeable ranks. +void register_special_tokens(vendor::BpeVocabulary & vocab, int32_t base_id) { + static const std::array kAudioEvents = { + "ASR", "AED", "SER", "Speech", "/Speech", "BGM", "/BGM", + "Laughter", "/Laughter", "Applause", "/Applause", + }; + static const std::array kEmotions = { + "HAPPY", "SAD", "ANGRY", "NEUTRAL", + }; + static const std::array kTasks = { + "translate", "transcribe", "startoflm", "startofprev", "nospeech", "notimestamps", + }; + static const std::array kTtsVocal = { + "TTS/B", "TTS/O", "TTS/Q", "TTS/A", "TTS/CO", "TTS/CL", "TTS/H", + }; + + int32_t id = base_id; + add_special_token(vocab, "<|endoftext|>", id++); + add_special_token(vocab, "<|startoftranscript|>", id++); + for (const char * lang : kLanguages) { + add_special_token(vocab, "<|" + std::string(lang) + "|>", id++); + } + for (const char * event : kAudioEvents) { + add_special_token(vocab, "<|" + std::string(event) + "|>", id++); + } + for (const char * emotion : kEmotions) { + add_special_token(vocab, "<|" + std::string(emotion) + "|>", id++); + } + for (const char * task : kTasks) { + add_special_token(vocab, "<|" + std::string(task) + "|>", id++); + } + for (int i = 1; i <= 30; ++i) { + add_special_token(vocab, "<|SPECIAL_TOKEN_" + std::to_string(i) + "|>", id++); + } + for (const char * vocal : kTtsVocal) { + add_special_token(vocab, "<|" + std::string(vocal) + "|>", id++); + } + for (int i = 1; i <= 13; ++i) { + char name[32]; + std::snprintf(name, sizeof(name), "<|TTS/SP%02d|>", i); + add_special_token(vocab, name, id++); + } + // Timestamps <|0.00|> .. <|30.00|> in 0.02 steps; i * 0.02 == i / 50. + for (int i = 0; i <= 1500; ++i) { + char name[32]; + std::snprintf(name, sizeof(name), "<|%d.%02d|>", i / 50, (i * 2) % 100); + add_special_token(vocab, name, id++); + } +} + +std::shared_ptr load_tiktoken_vocabulary(const std::filesystem::path & vocab_path) { + std::ifstream input(vocab_path, std::ios::binary); + if (!input) { + throw std::runtime_error("IndexTTS2.5 failed to open tiktoken vocabulary: " + vocab_path.string()); + } + + auto vocab = std::make_shared(); + vocab->pre_type = vendor::PreTokenizerType::Gpt2; + + std::string line; + int64_t mergeable_count = 0; + while (std::getline(input, line)) { + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + if (line.empty()) { + continue; + } + std::istringstream parts(line); + std::string token_base64; + int64_t rank = -1; + if (!(parts >> token_base64 >> rank) || rank < 0 || rank > INT32_MAX) { + throw std::runtime_error("IndexTTS2.5 tiktoken vocabulary has an invalid line: " + line); + } + const std::string bytes = decode_base64(token_base64); + const auto token_id = static_cast(rank); + const std::string mapped = map_token_bytes(bytes); + vocab->token_to_id.emplace(mapped, token_id); + vocab->id_to_token.emplace(token_id, vendor::TokenData{mapped, 0}); + // tiktoken ranks double as merge priorities: an adjacent pair merges + // iff its concatenation is a token, with that token's rank. Register + // every split so find_bpe_rank(left, right) == rank(left + right). + for (size_t split = 1; split < bytes.size(); ++split) { + vocab->bpe_ranks.emplace( + pair_key(map_token_bytes(bytes.substr(0, split)), map_token_bytes(bytes.substr(split))), + token_id); + } + ++mergeable_count; + } + if (mergeable_count == 0) { + throw std::runtime_error("IndexTTS2.5 tiktoken vocabulary is empty: " + vocab_path.string()); + } + + register_special_tokens(*vocab, static_cast(mergeable_count)); + vendor::rebuild_special_tokens_cache(*vocab); + return vocab; +} + +size_t utf8_codepoint_size(unsigned char byte) { + if ((byte & 0x80U) == 0U) { + return 1; + } + if ((byte & 0xE0U) == 0xC0U) { + return 2; + } + if ((byte & 0xF0U) == 0xE0U) { + return 3; + } + if ((byte & 0xF8U) == 0xF0U) { + return 4; + } + return 1; +} + +uint32_t decode_utf8_codepoint(const std::string & text, size_t offset, size_t size) { + const auto byte = [&](size_t i) { return static_cast(text[offset + i]); }; + if (size == 1) { + return byte(0); + } + if (size == 2 && offset + 1 < text.size()) { + return ((byte(0) & 0x1FU) << 6U) | (byte(1) & 0x3FU); + } + if (size == 3 && offset + 2 < text.size()) { + return ((byte(0) & 0x0FU) << 12U) | ((byte(1) & 0x3FU) << 6U) | (byte(2) & 0x3FU); + } + if (size == 4 && offset + 3 < text.size()) { + return ((byte(0) & 0x07U) << 18U) | ((byte(1) & 0x3FU) << 12U) | ((byte(2) & 0x3FU) << 6U) | (byte(3) & 0x3FU); + } + return byte(0); +} + +uint32_t next_utf8_codepoint(const std::string & text, size_t & offset) { + const size_t size = std::min(utf8_codepoint_size(static_cast(text[offset])), text.size() - offset); + const uint32_t cp = decode_utf8_codepoint(text, offset, size); + offset += size; + return cp; +} + +bool is_han_codepoint(uint32_t cp) { + return cp >= 0x4E00U && cp <= 0x9FFFU; +} + +bool contains_han(const std::string & text) { + for (size_t i = 0; i < text.size();) { + if (is_han_codepoint(next_utf8_codepoint(text, i))) { + return true; + } + } + return false; +} + +std::string lowercase_ascii(std::string text) { + for (char & ch : text) { + ch = static_cast(std::tolower(static_cast(ch))); + } + return text; +} + +std::string uppercase_ascii(std::string text) { + for (char & ch : text) { + ch = static_cast(std::toupper(static_cast(ch))); + } + return text; +} + +bool is_kana(const std::string & text) { + if (text.empty()) { + return false; + } + bool all_hiragana = true; + bool all_katakana = true; + for (size_t i = 0; i < text.size();) { + const uint32_t cp = next_utf8_codepoint(text, i); + if (cp < 0x3040U || cp > 0x309FU) { + all_hiragana = false; + } + if (cp < 0x30A0U || cp > 0x30FFU) { + all_katakana = false; + } + } + return all_hiragana || all_katakana; +} + +struct AnnotationMatch { + size_t end = 0; // one past the match; 0 when there is no match at pos + size_t word_begin = 0; + size_t word_end = 0; + size_t pron_begin = 0; + size_t pron_end = 0; +}; + +// Matches <([^|>\n]+)\|([^>\n]+)> anchored at pos. +AnnotationMatch match_pronunciation_annotation(const std::string & text, size_t pos) { + AnnotationMatch match; + if (text[pos] != '<') { + return match; + } + size_t cursor = pos + 1; + const size_t word_begin = cursor; + while (cursor < text.size() && text[cursor] != '|' && text[cursor] != '>' && text[cursor] != '\n') { + ++cursor; + } + if (cursor == word_begin || cursor >= text.size() || text[cursor] != '|') { + return match; + } + match.word_begin = word_begin; + match.word_end = cursor; + const size_t pron_begin = ++cursor; + while (cursor < text.size() && text[cursor] != '>' && text[cursor] != '\n') { + ++cursor; + } + if (cursor == pron_begin || cursor >= text.size()) { + return AnnotationMatch{}; + } + match.pron_begin = pron_begin; + match.pron_end = cursor; + match.end = cursor + 1; + return match; +} + +// Base-26 spreadsheet-style index ("a".."z", "aa"..), mirroring the official +// TextNormalizer._protect_pronunciation_annotations placeholder naming. +std::string alpha_placeholder_index(size_t n) { + std::string s; + while (true) { + s.insert(s.begin(), static_cast('a' + (n % 26))); + const size_t q = n / 26; + if (q == 0) { + break; + } + n = q - 1; + } + return s; +} + +using PronunciationPlaceholders = std::vector>; + +// Replaces annotations with letter-only placeholders so +// text normalization cannot rewrite their digits/symbols (e.g. XING2). +std::pair protect_pronunciation_annotations(const std::string & text) { + std::string out; + out.reserve(text.size()); + PronunciationPlaceholders placeholders; + size_t pos = 0; + while (pos < text.size()) { + const auto match = match_pronunciation_annotation(text, pos); + if (match.end == 0) { + out.push_back(text[pos++]); + continue; + } + std::string key = "PRONPLACEHOLDER" + alpha_placeholder_index(placeholders.size()) + "PRONPLACEHOLDER"; + placeholders.emplace_back(key, text.substr(pos, match.end - pos)); + out += key; + pos = match.end; + } + return {out, placeholders}; +} + +std::string restore_pronunciation_annotations(std::string text, const PronunciationPlaceholders & placeholders) { + for (const auto & [key, original] : placeholders) { + size_t at = 0; + while ((at = text.find(key, at)) != std::string::npos) { + text.replace(at, key.size(), original); + at += original.size(); + } + } + return text; +} + +// Expands annotations (see infer_v2_5.py +// apply_pronunciation_annotations): +// Chinese word -> <|SPECIAL_TOKEN_2|>PRON<|SPECIAL_TOKEN_2|> +// other word -> <|SPECIAL_TOKEN_1|>PRON<|SPECIAL_TOKEN_1|> +// kana pron -> inlined as " PRON " +std::string apply_pronunciation_annotations(const std::string & text) { + std::string out; + out.reserve(text.size()); + size_t pos = 0; + while (pos < text.size()) { + const auto match = match_pronunciation_annotation(text, pos); + if (match.end == 0) { + out.push_back(text[pos++]); + continue; + } + const std::string word = text.substr(match.word_begin, match.word_end - match.word_begin); + const std::string pron = uppercase_ascii(text.substr(match.pron_begin, match.pron_end - match.pron_begin)); + if (is_kana(pron)) { + out.push_back(' '); + out += pron; + out.push_back(' '); + } else { + const char * wrapper = contains_han(word) ? "<|SPECIAL_TOKEN_2|>" : "<|SPECIAL_TOKEN_1|>"; + out += wrapper; + out += pron; + out += wrapper; + } + pos = match.end; + } + return out; +} + +// Uppercases the name inside <|...|> markers: re.sub(r'<\|([^|]+)\|>', upper). +std::string uppercase_special_token_names(const std::string & text) { + std::string out; + out.reserve(text.size()); + size_t pos = 0; + while (pos < text.size()) { + if (text[pos] != '<' || pos + 1 >= text.size() || text[pos + 1] != '|') { + out.push_back(text[pos++]); + continue; + } + size_t cursor = pos + 2; + while (cursor < text.size() && text[cursor] != '|') { + ++cursor; + } + if (cursor == pos + 2 || cursor + 1 >= text.size() || text[cursor + 1] != '>') { + out.push_back(text[pos++]); + continue; + } + out += "<|"; + out += uppercase_ascii(text.substr(pos + 2, cursor - (pos + 2))); + out += "|>"; + pos = cursor + 2; + } + return out; +} + +bool is_segment_delimiter(uint32_t cp) { + switch (cp) { + case U',': + case U'.': + case U'!': + case U'?': + case U';': + case U':': + case U'\n': + case 0xFF0CU: // ， + case 0x3002U: // 。 + case 0xFF01U: // ！ + case 0xFF1FU: // ？ + case 0x3001U: // 、 + case 0xFF1BU: // ； + case 0xFF1AU: // ： + return true; + default: + return false; + } +} + +// re.split(r'(?<=[，。！？、；：,\.!\?;:\n])', piece): split after each delimiter. +std::vector split_after_delimiters(const std::string & piece) { + std::vector parts; + std::string current; + for (size_t i = 0; i < piece.size();) { + const size_t begin = i; + const uint32_t cp = next_utf8_codepoint(piece, i); + current.append(piece, begin, i - begin); + if (is_segment_delimiter(cp)) { + parts.push_back(std::move(current)); + current.clear(); + } + } + if (!current.empty()) { + parts.push_back(std::move(current)); + } + return parts; +} + +// Matches "<|SPECIAL_TOKEN_|>" at pos; returns the match length or 0. +size_t match_special_token_marker(const std::string & text, size_t pos) { + static const std::string kPrefix = "<|SPECIAL_TOKEN_"; + if (text.compare(pos, kPrefix.size(), kPrefix) != 0) { + return 0; + } + size_t cursor = pos + kPrefix.size(); + const size_t digits_begin = cursor; + while (cursor < text.size() && std::isdigit(static_cast(text[cursor])) != 0) { + ++cursor; + } + if (cursor == digits_begin || cursor + 1 >= text.size() || text[cursor] != '|' || text[cursor + 1] != '>') { + return 0; + } + return cursor + 2 - pos; +} + +// SPLIT_PROTECTED_PATTERN spans (<|SPECIAL_TOKEN_n|>...<|SPECIAL_TOKEN_n|>) +// are kept atomic during segmentation. +std::vector> split_atomic_pieces(const std::string & text) { + std::vector> pieces; + size_t pos = 0; + while (pos < text.size()) { + size_t opener = std::string::npos; + size_t opener_len = 0; + for (size_t i = pos; i < text.size(); ++i) { + const size_t len = match_special_token_marker(text, i); + if (len > 0) { + opener = i; + opener_len = len; + break; + } + } + if (opener == std::string::npos) { + break; + } + size_t closer = std::string::npos; + size_t closer_len = 0; + for (size_t i = opener + opener_len; i < text.size(); ++i) { + const size_t len = match_special_token_marker(text, i); + if (len > 0) { + closer = i; + closer_len = len; + break; + } + } + if (closer == std::string::npos) { + break; + } + if (opener > pos) { + pieces.emplace_back(text.substr(pos, opener - pos), false); + } + pieces.emplace_back(text.substr(opener, closer + closer_len - opener), true); + pos = closer + closer_len; + } + if (pos < text.size()) { + pieces.emplace_back(text.substr(pos), false); + } + return pieces; +} + +} // namespace + +IndexTTS25TextTokenizer::IndexTTS25TextTokenizer(std::shared_ptr assets) + : assets_(std::move(assets)) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 text tokenizer requires assets"); + } + vocab_ = load_tiktoken_vocabulary(assets_->resources.require_file("tiktoken")); +} + +std::string IndexTTS25TextTokenizer::normalize_english(const std::string & text) const { + engine::text::EnglishTextNormalizationOptions options; + options.expand_common_contractions = true; + options.index_tts_punctuation = true; + options.uppercase_ascii = false; + return engine::text::normalize_english_text(text, options); +} + +std::string IndexTTS25TextTokenizer::normalize_chinese(const std::string & text) const { + return engine::text::normalize_chinese_text( + text, + engine::text::ChineseTextNormalizationTarget::IndexTTS); +} + +std::vector IndexTTS25TextTokenizer::encode(const std::string & text) const { + return vendor::tokenize_bpe(*vocab_, text, true); +} + +int32_t IndexTTS25TextTokenizer::special_token_id(const std::string & token_text) const { + const auto it = vocab_->token_to_id.find(token_text); + return it == vocab_->token_to_id.end() ? -1 : it->second; +} + +int32_t IndexTTS25TextTokenizer::lang_to_id(const std::string & lang) { + const std::string normalized = lowercase_ascii(lang); + for (size_t i = 0; i < kLanguages.size(); ++i) { + if (normalized == kLanguages[i]) { + return static_cast(i); + } + } + for (size_t i = 0; i < kEmbeddingOnlyLanguages.size(); ++i) { + if (normalized == kEmbeddingOnlyLanguages[i]) { + return static_cast(kLanguages.size() + i); + } + } + return kCommonLangId; +} + +IndexTTS25TextEncoding IndexTTS25TextTokenizer::encode_for_inference( + const std::string & text, + int max_text_tokens_per_segment, + const std::string & lang) const { + if (max_text_tokens_per_segment <= 0) { + throw std::runtime_error("IndexTTS2.5 max_text_tokens_per_segment must be positive"); + } + + std::string resolved_lang = lowercase_ascii(lang); + if (resolved_lang.empty()) { + resolved_lang = contains_han(text) ? "zh" : "en"; + } + + std::string processed = text; + if (resolved_lang == "zh" || resolved_lang == "en") { + // Protect annotations from the normalizer, as the + // official TextNormalizer does inside normalize(). + auto protected_text = protect_pronunciation_annotations(processed); + protected_text.first = resolved_lang == "zh" + ? normalize_chinese(protected_text.first) + : normalize_english(protected_text.first); + processed = restore_pronunciation_annotations(std::move(protected_text.first), protected_text.second); + } + // ja/es/ar and other languages currently pass through without TN. + if (resolved_lang == "zh" || resolved_lang == "ja" || resolved_lang == "en") { + processed = lowercase_ascii(std::move(processed)); + } else if (resolved_lang == "es") { + processed = uppercase_ascii(std::move(processed)); + } + processed = apply_pronunciation_annotations(processed); + processed = uppercase_special_token_names(processed); + + const std::string lang_prefix = "<|" + resolved_lang + "|> "; + const auto prefix_tokens = static_cast(encode(lang_prefix).size()); + const int64_t capacity = assets_->config.gpt.max_text_tokens; + int64_t budget = std::min(max_text_tokens_per_segment, capacity - 2) - prefix_tokens; + budget = std::max(budget, 1); + + std::vector segments; + const auto token_len = [this](const std::string & value) { + return static_cast(encode(value).size()); + }; + if (token_len(processed) <= budget) { + segments.push_back(processed); + } else { + std::vector chunks; + for (const auto & [piece, atomic] : split_atomic_pieces(processed)) { + if (atomic) { + chunks.push_back(piece); + continue; + } + for (const auto & part : split_after_delimiters(piece)) { + if (token_len(part) <= budget) { + chunks.push_back(part); + continue; + } + std::string current; + for (size_t i = 0; i < part.size();) { + const size_t begin = i; + next_utf8_codepoint(part, i); + const std::string ch = part.substr(begin, i - begin); + if (!current.empty() && token_len(current + ch) > budget) { + chunks.push_back(std::move(current)); + current = ch; + } else { + current += ch; + } + } + if (!current.empty()) { + chunks.push_back(std::move(current)); + } + } + } + std::string current; + for (const auto & chunk : chunks) { + if (!current.empty() && token_len(current + chunk) > budget) { + segments.push_back(std::move(current)); + current = chunk; + } else { + current += chunk; + } + } + if (!current.empty()) { + segments.push_back(std::move(current)); + } + if (segments.empty()) { + segments.push_back(processed); + } + } + + IndexTTS25TextEncoding encoding; + encoding.lang = resolved_lang; + encoding.normalized_text = processed; + encoding.segments = segments; + encoding.segment_token_ids.reserve(segments.size()); + for (const auto & segment : segments) { + std::vector ids = encode(lang_prefix + segment); + ids.push_back(kSegmentPadTokenId); + encoding.segment_token_ids.push_back(std::move(ids)); + } + return encoding; +} + +} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/vocoder.cpp b/src/models/index_tts2_5/vocoder.cpp new file mode 100644 index 00000000..52f9dbe0 --- /dev/null +++ b/src/models/index_tts2_5/vocoder.cpp @@ -0,0 +1,69 @@ +#include "engine/models/index_tts2_5/vocoder.h" + +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::index_tts2_5 { +namespace { + +constexpr int64_t kChunkedVocoderFrames = 768; +constexpr int64_t kChunkedVocoderOverlapFrames = 32; + +engine::modules::BigVganVocoderConfig parse_bigvgan_config( + const IndexTTS25Assets & assets, + engine::assets::TensorStorageType weight_storage_type) { + const auto root = assets.resources.parse_json("bigvgan_config"); + engine::modules::BigVganVocoderConfig config; + config.sampling_rate = engine::io::json::require_i64(root, "sampling_rate"); + config.num_mels = engine::io::json::require_i64(root, "num_mels"); + config.n_fft = engine::io::json::require_i64(root, "n_fft"); + config.hop_size = engine::io::json::require_i64(root, "hop_size"); + config.win_size = engine::io::json::require_i64(root, "win_size"); + config.upsample_initial_channel = engine::io::json::require_i64(root, "upsample_initial_channel"); + config.snake_logscale = engine::io::json::require_bool(root, "snake_logscale"); + config.upsample_rates = engine::io::json::require_i64_array(root, "upsample_rates"); + config.upsample_kernel_sizes = engine::io::json::require_i64_array(root, "upsample_kernel_sizes"); + config.resblock_kernel_sizes = engine::io::json::require_i64_array(root, "resblock_kernel_sizes"); + config.weight_storage_type = weight_storage_type; + if (config.sampling_rate != assets.config.s2mel.sample_rate || + config.num_mels != assets.config.s2mel.n_mels || + config.n_fft != assets.config.s2mel.n_fft || + config.hop_size != assets.config.s2mel.hop_length || + config.win_size != assets.config.s2mel.win_length) { + throw std::runtime_error("IndexTTS2.5 BigVGAN config does not match S2Mel mel config"); + } + return config; +} + +} // namespace + +IndexTTS25BigVganVocoder::IndexTTS25BigVganVocoder( + std::shared_ptr assets, + core::BackendConfig backend, + engine::assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)) { + if (assets_ == nullptr) { + throw std::runtime_error("IndexTTS2.5 BigVGAN vocoder requires assets"); + } + component_ = engine::modules::BigVganVocoderComponent::load_from_tensor_source( + assets_->bigvgan_weights, + std::move(backend), + parse_bigvgan_config(*assets_, weight_storage_type)); +} + +IndexTTS25VocoderOutput IndexTTS25BigVganVocoder::synthesize( + const std::vector & mel, + int64_t frames) const { + const auto out = frames > kChunkedVocoderFrames + ? component_.synthesize_chunked(mel, frames, kChunkedVocoderFrames, kChunkedVocoderOverlapFrames) + : component_.synthesize(mel, frames); + return {out.waveform, out.samples, static_cast(out.sample_rate)}; +} + +void IndexTTS25BigVganVocoder::release_runtime_graph() { + component_.release_runtime_graph(); +} + +} // namespace engine::models::index_tts2_5 diff --git a/tests/index_tts2_5/index_tts2_5_warm_bench.cpp b/tests/index_tts2_5/index_tts2_5_warm_bench.cpp new file mode 100644 index 00000000..f5868862 --- /dev/null +++ b/tests/index_tts2_5/index_tts2_5_warm_bench.cpp @@ -0,0 +1,319 @@ +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/audio/wav_writer.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/io/json.h" +#include "engine/framework/runtime/registry.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using Clock = std::chrono::steady_clock; + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +int int_arg(int argc, char ** argv, const std::string & name, int fallback) { + return std::stoi(arg_value(argc, argv, name, std::to_string(fallback))); +} + +engine::core::BackendType parse_backend(const std::string & value) { + if (value == "cuda") { + return engine::core::BackendType::Cuda; + } + if (value == "cpu") { + return engine::core::BackendType::Cpu; + } + throw std::runtime_error("IndexTTS2.5 warmbench backend must be cuda or cpu"); +} + +std::vector> parse_session_options(int argc, char ** argv) { + std::vector> out; + for (int i = 1; i + 1 < argc; ++i) { + if (std::string(argv[i]) != "--session-option") { + continue; + } + const std::string option = argv[i + 1]; + const size_t eq = option.find('='); + if (eq == std::string::npos || eq == 0) { + throw std::runtime_error("invalid IndexTTS2.5 --session-option: " + option); + } + out.emplace_back(option.substr(0, eq), option.substr(eq + 1)); + } + return out; +} + +std::string option_text(const engine::io::json::Value & value) { + if (value.is_bool()) { + return value.as_bool() ? "true" : "false"; + } + if (value.is_number()) { + return engine::io::json::stringify_number(value.as_number()); + } + if (value.is_array()) { + std::string out; + const auto & items = value.as_array(); + for (size_t i = 0; i < items.size(); ++i) { + if (i != 0) { + out += ","; + } + if (!items[i].is_number()) { + throw std::runtime_error("IndexTTS2.5 warmbench option arrays must contain only numbers"); + } + out += engine::io::json::stringify_number(items[i].as_number()); + } + return out; + } + return value.as_string(); +} + +std::string required_string(const engine::io::json::Value & object, const std::string & key) { + const auto * value = object.find(key); + if (value == nullptr || value->is_null()) { + throw std::runtime_error("IndexTTS2.5 warmbench request missing " + key); + } + return value->as_string(); +} + +std::string optional_string(const engine::io::json::Value & object, const std::string & key) { + const auto * value = object.find(key); + return value == nullptr || value->is_null() ? std::string{} : value->as_string(); +} + +void set_optional_option(engine::runtime::TaskRequest & request, const engine::io::json::Value & object, const std::string & key) { + const auto * value = object.find(key); + if (value != nullptr && !value->is_null()) { + request.options[key] = option_text(*value); + } +} + +engine::runtime::AudioBuffer read_audio_buffer(const std::filesystem::path & path) { + const auto wav = engine::audio::read_wav_f32(path); + return engine::runtime::AudioBuffer{wav.sample_rate, wav.channels, wav.samples}; +} + +engine::runtime::TaskRequest make_request(const engine::io::json::Value & object) { + engine::runtime::TaskRequest request; + const std::string language = optional_string(object, "language"); + request.text_input = engine::runtime::Transcript{required_string(object, "text"), language.empty() ? "en" : language}; + engine::runtime::VoiceCondition voice; + engine::runtime::VoiceReference speaker; + speaker.audio = read_audio_buffer(required_string(object, "voice_ref")); + voice.speaker = std::move(speaker); + request.voice = std::move(voice); + + const auto emotion_audio = optional_string(object, "audio"); + if (!emotion_audio.empty()) { + request.audio_input = read_audio_buffer(emotion_audio); + } + + set_optional_option(request, object, "lang"); + set_optional_option(request, object, "emotion_alpha"); + set_optional_option(request, object, "emotion_vector"); + set_optional_option(request, object, "use_emotion_text"); + set_optional_option(request, object, "emotion_text"); + set_optional_option(request, object, "use_random_emotion"); + set_optional_option(request, object, "interval_silence_ms"); + set_optional_option(request, object, "text_chunk_size"); + set_optional_option(request, object, "do_sample"); + set_optional_option(request, object, "top_p"); + set_optional_option(request, object, "top_k"); + set_optional_option(request, object, "temperature"); + set_optional_option(request, object, "length_penalty"); + set_optional_option(request, object, "num_beams"); + set_optional_option(request, object, "repetition_penalty"); + set_optional_option(request, object, "max_tokens"); + set_optional_option(request, object, "seed"); + return request; +} + +std::vector parse_requests(const std::string & request_sequence_json) { + if (request_sequence_json.empty()) { + throw std::runtime_error("IndexTTS2.5 warmbench requires --request-sequence-json"); + } + const auto root = engine::io::json::parse(request_sequence_json); + std::vector requests; + for (const auto & item : root.as_array()) { + requests.push_back(make_request(item)); + } + if (requests.empty()) { + throw std::runtime_error("IndexTTS2.5 warmbench request sequence is empty"); + } + return requests; +} + +engine::io::json::Value number(double value) { + return engine::io::json::Value::make_number(value); +} + +engine::io::json::Value string(std::string value) { + return engine::io::json::Value::make_string(std::move(value)); +} + +engine::io::json::Value audio_summary_json(const engine::runtime::AudioBuffer & audio) { + if (audio.samples.empty()) { + throw std::runtime_error("IndexTTS2.5 warmbench received empty audio output"); + } + double sum = 0.0; + double abs_sum = 0.0; + double sq_sum = 0.0; + float min_value = audio.samples.front(); + float max_value = audio.samples.front(); + for (const float sample : audio.samples) { + sum += static_cast(sample); + abs_sum += std::abs(static_cast(sample)); + sq_sum += static_cast(sample) * static_cast(sample); + min_value = std::min(min_value, sample); + max_value = std::max(max_value, sample); + } + const int channels = std::max(1, audio.channels); + const double frames = static_cast(audio.samples.size() / static_cast(channels)); + const double count = static_cast(audio.samples.size()); + return engine::io::json::Value::make_object({ + {"sample_rate", number(static_cast(audio.sample_rate))}, + {"channels", number(static_cast(audio.channels))}, + {"samples", number(count)}, + {"frames", number(frames)}, + {"duration_sec", number(audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0)}, + {"sum", number(sum)}, + {"mean_abs", number(abs_sum / count)}, + {"rms", number(std::sqrt(sq_sum / count))}, + {"min", number(min_value)}, + {"max", number(max_value)}, + }); +} + +} // namespace + +int main(int argc, char ** argv) { + try { + const std::filesystem::path model_path = arg_value(argc, argv, "--model", "models/IndexTTS-2.5"); + const std::string backend_name = arg_value(argc, argv, "--backend", "cuda"); + const int device = int_arg(argc, argv, "--device", 0); + const int threads = int_arg(argc, argv, "--threads", 8); + const int warmup = int_arg(argc, argv, "--warmup", 0); + const int iterations = int_arg(argc, argv, "--iterations", 1); + const std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", ""); + const std::filesystem::path output_dir = arg_value(argc, argv, "--output-dir", ""); + const std::filesystem::path timing_path = arg_value(argc, argv, "--timing-file", "/tmp/index_tts2_5_warm_bench_timing.log"); + engine::debug::configure_logging(engine::debug::LoggingConfig{true, timing_path.string()}); + + engine::runtime::ModelLoadRequest load_request; + load_request.model_path = model_path; + load_request.family_hint = "index_tts2_5"; + auto registry = engine::runtime::make_default_registry(); + auto model = registry.load(load_request); + + engine::runtime::TaskSpec task; + task.task = engine::runtime::VoiceTaskKind::Tts; + task.mode = engine::runtime::RunMode::Offline; + engine::runtime::SessionOptions session_options; + session_options.backend.type = parse_backend(backend_name); + session_options.backend.device = device; + session_options.backend.threads = threads; + for (const auto & [key, value] : parse_session_options(argc, argv)) { + session_options.options[key] = value; + } + auto requests = parse_requests(request_sequence_json); + auto session_base = model->create_task_session(task, session_options); + auto * session = dynamic_cast(session_base.get()); + if (session == nullptr) { + throw std::runtime_error("IndexTTS2.5 model did not create an offline voice task session"); + } + engine::runtime::SessionPreparationRequest preparation; + preparation.text = requests.front().text_input; + preparation.voice = requests.front().voice; + preparation.audio = requests.front().audio_input.has_value() + ? std::optional({ + requests.front().audio_input->sample_rate, + requests.front().audio_input->channels, + static_cast(requests.front().audio_input->samples.size()), + }) + : std::nullopt; + preparation.options = requests.front().options; + session->prepare(preparation); + + std::vector steps; + std::vector timing_lines; + timing_lines.push_back("index_tts2_5.backend " + backend_name); + timing_lines.push_back("index_tts2_5.model_root " + model_path.string()); + for (int i = 0; i < warmup; ++i) { + (void) session->run(requests.front()); + } + + for (size_t request_index = 0; request_index < requests.size(); ++request_index) { + double total_ms = 0.0; + engine::runtime::TaskResult last_result; + for (int iteration = 0; iteration < std::max(1, iterations); ++iteration) { + const auto start = Clock::now(); + last_result = session->run(requests[request_index]); + const auto end = Clock::now(); + const double wall_ms = std::chrono::duration(end - start).count(); + total_ms += wall_ms; + timing_lines.push_back("index_tts2_5.wall_ms " + engine::io::json::stringify_number(wall_ms)); + } + if (!last_result.audio_output.has_value()) { + throw std::runtime_error("IndexTTS2.5 warmbench expected audio output"); + } + const double avg_ms = total_ms / static_cast(std::max(1, iterations)); + std::filesystem::path audio_path; + if (!output_dir.empty()) { + std::filesystem::create_directories(output_dir); + audio_path = output_dir / ("request_" + std::to_string(request_index) + ".wav"); + engine::audio::write_pcm16_wav( + audio_path, + last_result.audio_output->sample_rate, + last_result.audio_output->channels, + last_result.audio_output->samples); + } + engine::io::json::Value::Object step{ + {"request_index", number(static_cast(request_index))}, + {"stems", engine::io::json::Value::make_array({ + engine::io::json::Value::make_object({ + {"name", string("audio")}, + {"summary", audio_summary_json(*last_result.audio_output)}, + {"audio", string(audio_path.string())}, + }), + })}, + {"metrics", engine::io::json::Value::make_object({{"wall_ms", number(avg_ms)}})}, + }; + steps.push_back(engine::io::json::Value::make_object(std::move(step))); + std::cout << "index_tts2_5.wall_ms=" << avg_ms << "\n"; + } + + if (!timing_path.empty()) { + std::filesystem::create_directories(timing_path.parent_path()); + std::ofstream timing(timing_path, std::ios::app); + for (const auto & line : timing_lines) { + timing << line << "\n"; + } + } + + const auto summary = engine::io::json::Value::make_object({ + {"family", string("index_tts2_5")}, + {"backend", string(backend_name)}, + {"sequence_steps", engine::io::json::Value::make_array(std::move(steps))}, + }); + std::cout << "summary_json=" << engine::io::json::stringify(summary) << "\n"; + return 0; + } catch (const std::exception & ex) { + std::cerr << "index_tts2_5_warm_bench failed: " << ex.what() << "\n"; + return 1; + } +} diff --git a/tests/index_tts2_5/index_tts2_5_warm_bench_cases.json b/tests/index_tts2_5/index_tts2_5_warm_bench_cases.json new file mode 100644 index 00000000..f2992421 --- /dev/null +++ b/tests/index_tts2_5/index_tts2_5_warm_bench_cases.json @@ -0,0 +1,131 @@ +{ + "voice_clone": { + "requests": [ + { + "text": "The palace is strict, no false rumors, Lady Qi!", + "voice_ref": "resources/index_tts2_5/official_examples/voice_02.wav", + "seed": 1234 + } + ] + }, + "chinese_voice_clone": { + "requests": [ + { + "text": "这个呀，就是我们精心制作准备的纪念品，大家可以看到这个色泽和这个材质啊，哎呀多么的光彩照人。", + "language": "zh", + "lang": "zh", + "voice_ref": "resources/index_tts2_5/official_examples/voice_03.wav", + "seed": 1245 + } + ] + }, + "chinese_emotion_text": { + "requests": [ + { + "text": "快躲起来！是他要来了！他要来抓我们了！", + "language": "zh", + "lang": "zh", + "voice_ref": "resources/index_tts2_5/official_examples/voice_12.wav", + "use_emotion_text": true, + "emotion_text": "你吓死我了！你是鬼吗？", + "emotion_alpha": 0.6, + "use_random_emotion": false, + "seed": 1246 + } + ] + }, + "emotion_reference": { + "requests": [ + { + "text": "The old theater was empty, but every chair still seemed to remember the audience.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_07.wav", + "audio": "resources/index_tts2_5/official_examples/emo_sad.wav", + "seed": 1235 + } + ] + }, + "emotion_reference_alpha": { + "requests": [ + { + "text": "I tried to sound calm, but the storm outside made every word feel heavier.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_07.wav", + "audio": "resources/index_tts2_5/official_examples/emo_sad.wav", + "emotion_alpha": 0.9, + "seed": 1236 + } + ] + }, + "emotion_vector": { + "requests": [ + { + "text": "I'm sorry, I really did forget, but I promise I will remember the important things.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_09.wav", + "emotion_vector": [0.0, 0.0, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0], + "use_random_emotion": false, + "seed": 1237 + } + ] + }, + "emotion_text": { + "requests": [ + { + "text": "Hide quickly. Someone is coming, and I do not think they are here to help us.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_12.wav", + "use_emotion_text": true, + "emotion_alpha": 0.6, + "use_random_emotion": false, + "seed": 1238 + } + ] + }, + "emotion_text_description": { + "requests": [ + { + "text": "Hide quickly. Someone is coming, and I do not think they are here to help us.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_12.wav", + "use_emotion_text": true, + "emotion_text": "You scared me. Are you a ghost?", + "emotion_alpha": 0.6, + "use_random_emotion": false, + "seed": 1239 + } + ] + }, + "long_session_changed_requests": { + "use_all_requests_by_default": true, + "requests": [ + { + "text": "The palace is strict, no false rumors, Lady Qi!", + "voice_ref": "resources/index_tts2_5/official_examples/voice_02.wav", + "seed": 1240 + }, + { + "text": "Please lower your voice. The guards are already listening outside the door.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_02.wav", + "seed": 1241 + }, + { + "text": "I am trying to stay brave, but every shadow in this room feels alive.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_02.wav", + "use_emotion_text": true, + "emotion_text": "You scared me. Are you a ghost?", + "emotion_alpha": 0.6, + "use_random_emotion": false, + "seed": 1242 + }, + { + "text": "The old theater was empty, but every chair still seemed to remember the audience.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_07.wav", + "audio": "resources/index_tts2_5/official_examples/emo_sad.wav", + "seed": 1243 + }, + { + "text": "I tried to sound calm, but the storm outside made every word feel heavier.", + "voice_ref": "resources/index_tts2_5/official_examples/voice_07.wav", + "audio": "resources/index_tts2_5/official_examples/emo_sad.wav", + "emotion_alpha": 0.9, + "seed": 1244 + } + ] + } +} diff --git a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json index 6f214309..be28add6 100644 --- a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json @@ -288,6 +288,26 @@ } ] }, + { + "id": "index_tts2_5_voice_clone_longform", + "coverage": "IndexTTS2.5 voice clone with shared long-form text for chunking and RTF measurement", + "family": "index_tts2_5", + "model": "models/IndexTTS2.5-GGUF", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "clone_longform", + "text": "At dawn the harbor station opens its tall windows and the first clerk begins a careful report for the day. She notes the weather above the river, the slow cargo boats beyond the bridge, and the market voices arriving from the eastern road. A brass clock marks each quarter hour while porters stack wooden crates, bakers carry warm bread across the square, and a violinist practices the same bright phrase under the stone archway. By midmorning the keeper of the lighthouse sends a message about shifting currents, the museum guide unlocks a cabinet of maps, and a teacher leads a quiet line of students toward the ferry. In the afternoon a painter describes the silver color of the water, a mechanic jokes with the tram driver, and the station master reads an announcement that asks every traveler to keep close watch over letters, tickets, and parcels. After sunset the same clerk continues the report because new visitors keep arriving from the inland road. She explains that a florist carries pale roses past the fountain, two carpenters compare measurements beside the warehouse door, and the watchman checks each lock before the tide reaches its highest mark. A child laughs when the tram bell rings, a cook lowers a basket of fruit to the cellar, and three sailors unfold a chart that shows old channels, sandbars, and safe turning points for the morning crossing. Near midnight the lamps still glow on wet stone, the last cart rattles toward the market gate, and the report ends by saying that the harbor remains orderly, the wind has softened, the ferries are secure, and the town can rest until the next sunrise returns over the water. On the following morning the clerk resumes the record with even greater care because a week of inspections is about to begin. She writes that a ferry captain checks the mooring ropes one by one, a bookseller arranges travel guides beside the station cafe, and a pair of gardeners lift wet soil into bright clay pots near the west entrance. The bakery sends out trays of seed bread, the telegraph operator copies three official notices, and a tailor unfolds navy cloth across a polished wooden counter while customers wait in a line that bends toward the fountain. Before noon a surveyor compares bridge numbers against an old ledger, two cousins argue cheerfully about the best route to the fish market, and a choir director rehearses a patient scale that echoes against the warehouse wall. The lighthouse keeper reports that the northern channel is calmer than expected, the harbor pilot recommends a slower turn near the sandbar, and the customs officer stamps a packet of forms before waving a cart through the side gate. Later the schoolteacher returns with another group of students, asking them to observe the colors of rope, paint, stone, and water so they can write more exact descriptions in the classroom. A photographer kneels beside a rain barrel to capture the reflection of the clock tower, a mechanic tightens a brass hinge on the tram door, and an elderly traveler asks the clerk whether the evening ferry still stops at the orchard village beyond the marsh. As dusk arrives, lamps are trimmed again, shutters are tested against the wind, and the station kitchen sends bowls of soup to workers who remain on the late shift. The report continues with notes about a carpenter measuring floorboards in the east hall, a florist tying silver ribbon around the last stems of the day, and a violin case resting open on a bench beside the ticket window while its owner copies melody marks into a notebook. Long after the market gate closes, the clerk still writes that the harbor road stays busy, the river glints beneath scattered lamps, and the town maintains its patient rhythm of signals, footsteps, voices, bells, and distant engines. On the third day the clerk decides the record should be more precise, so she marks each event by the quarter hour and notes which sounds carry farthest through the station concourse. At first light she hears broom bristles on the stone steps, kettle lids in the cafe kitchen, and the slow scrape of crates being nudged across a loading cart beside the river wall. A messenger in a green coat delivers two canvas pouches, the ticket agent counts rolled coins into a brass tray, and a mother reads directions aloud while her son traces the painted ferry schedule with one curious finger. Midmorning brings a burst of sunlight across the waiting hall, making every brass handle shine while the museum guide escorts visitors toward the gallery of maps and navigational instruments. A porter pauses to describe the oldest compass in the display, a student sketches the harbor outline in graphite, and an apprentice clockmaker compares the station bell to a pocket watch that once belonged to his grandfather. By noon the fish market sends salt and seaweed scents through the open doors, tram wheels hiss at the curb, and the baker from the square exchanges a laugh with the florist who is carrying fresh lilies to the hotel veranda. The clerk writes that a cooper rolls three narrow barrels toward the cellar ramp, a translator copies weather bulletins for inland travelers, and a painter in a blue scarf studies the changing color of the tide as if each small wave might explain a different part of the sky. In the late afternoon the station master reviews freight tags, the customs officer checks a parcel of glassware, and a choir of children crosses the square singing a phrase so soft that the watchman removes his cap to listen. Evening settles slowly; lamps brighten in sequence, a cook inventories apples and onions in the pantry, and two sailors spread a faded chart on a crate so they can debate whether the shoals have shifted since the previous autumn. Before sleep the clerk closes the day with a final note that every vessel is accounted for, every platform has been swept, every lock has been tested twice, and the harbor seems ready to welcome another tide, another market, and another patient stream of voices at sunrise.", + "language": "en", + "voice_ref": "resources/index_tts2/official_examples/voice_12.wav", + "seed": 1234 + } + ] + }, { "id": "supertonic_tts_longform", "coverage": "Supertonic preset voice TTS with shared long-form text for chunking and RTF measurement", @@ -382,6 +402,36 @@ "seed": 1234 } ] + }, + { + "id": "index_tts2_5_longform_voice_clone_6000_emotion_text", + "coverage": "IndexTTS2.5 longform voice clone with 6000-character text plus long emotion-text conditioning", + "family": "index_tts2_5", + "model": "models/IndexTTS2.5-GGUF", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "longform_voice_clone_emotion_text", + "text_repeat": { + "text": "The archivist kept her voice steady while reading the witness report aloud, pausing after each sentence so the council could follow the chain of events without losing the emotional thread. ", + "chars": 6000 + }, + "voice_ref": "resources/index_tts2/official_examples/voice_12.wav", + "emotion_repeat": { + "text": "Speak with restrained fear, controlled urgency, and a careful softness, as if protecting a secret while trying not to alarm the listener. ", + "chars": 2400 + }, + "options": { + "emotion_alpha": 0.6, + "use_random_emotion": false + }, + "seed": 1234 + } + ] } ] } diff --git a/webui/configs/model_params.json b/webui/configs/model_params.json index 407bc115..7335811b 100644 --- a/webui/configs/model_params.json +++ b/webui/configs/model_params.json @@ -171,6 +171,15 @@ {"name": "interval_silence_ms", "type": "number", "label": "interval_silence_ms（分段间静音）", "label_en": "interval_silence_ms", "default": 200, "minimum": 0, "step": 50, "precision": 0} ], + "index_tts2_5": [ + {"name": "lang", "type": "choice", "label": "lang（语种提示）", "label_en": "lang (language hint)", "default": "auto", "choices": ["auto", "zh", "en", "ja", "es", "ar"], "info": "auto：含汉字按中文，否则按英文；日/西/阿建议显式选择", "info_en": "auto: zh when the text contains Han characters, otherwise en; set ja/es/ar explicitly"}, + {"name": "emotion_text", "type": "text", "label": "emotion_text（情绪参考文本）", "label_en": "emotion_text (emotion reference text)", "default": "", "placeholder": "例：你吓死我了！你是鬼吗？", "placeholder_en": "e.g. You scared me to death!", "info": "填写后自动开启情感条件（use_emotion_text）", "info_en": "Setting this enables emotion conditioning."}, + {"name": "emotion_alpha", "type": "slider", "label": "emotion_alpha（情感强度）", "label_en": "emotion_alpha", "default": 1.0, "minimum": 0.0, "maximum": 1.0, "step": 0.05}, + {"name": "use_emotion_text", "type": "bool", "label": "use_emotion_text（从朗读文本推断情感）", "label_en": "use_emotion_text (infer from text)", "default": false}, + {"name": "use_random_emotion", "type": "bool", "label": "use_random_emotion（随机情感）", "label_en": "use_random_emotion", "default": false}, + {"name": "interval_silence_ms", "type": "number", "label": "interval_silence_ms（分段间静音）", "label_en": "interval_silence_ms", "default": 200, "minimum": 0, "step": 50, "precision": 0} + ], + "irodori_tts": [ {"name": "num_inference_steps", "type": "number", "label": "num_inference_steps（RF 扩散步数）", "label_en": "num_inference_steps", "default": 40, "minimum": 1, "step": 1, "precision": 0}, {"name": "duration_sec", "type": "number", "label": "duration_sec（0=模型自动预测时长）", "label_en": "duration_sec (0 = auto)", "default": 0, "minimum": 0, "step": 0.5}, diff --git a/webui/configs/models_catalog.json b/webui/configs/models_catalog.json index 671c8b7b..c347973f 100644 --- a/webui/configs/models_catalog.json +++ b/webui/configs/models_catalog.json @@ -15,6 +15,9 @@ { "id": "voxcpm2", "display_name": "VoxCPM2 (tts)", "family": "voxcpm2", "path": "models/VoxCPM2", "task": "tts", "mode": "offline", "download_id": "voxcpm2", "session_options": { "voxcpm2.weight_type": "q8_0" }, "min_vram_gb": 6 }, { "id": "vibevoice", "display_name": "VibeVoice 1.5B (tts, long-form/multi-speaker)", "family": "vibevoice", "path": "models/VibeVoice-1.5B", "task": "tts", "mode": "offline", "download_id": "vibevoice_1_5b", "min_vram_gb": 7 }, { "id": "index-tts2", "display_name": "IndexTTS2 (tts 中英克隆+情感)", "display_name_en": "IndexTTS2 (tts, zh/en clone + emotion)", "family": "index_tts2", "path": "models/IndexTTS-2", "task": "tts", "mode": "offline", "download_id": "index_tts2", "min_vram_gb": 8 }, + { "id": "index-tts2.5", "display_name": "IndexTTS2.5 (tts 多语种克隆+情感, GGUF Q8)", "display_name_en": "IndexTTS2.5 (tts, zh/en/ja/es/ar clone + emotion, GGUF Q8)", "family": "index_tts2_5", "path": "models/IndexTTS2.5-GGUF", "task": "tts", "mode": "offline", "download_id": "index_tts2_5_q8_0", "min_vram_gb": 8, + "input_hint": "**IndexTTS2.5**：中/英/日/西/阿零样本克隆；上传参考音色即克隆；可在『其它参数(JSON)』里传 `lang`（默认 auto：含汉字按中文，否则按英文）与情感选项。许可证为 bilibili Model Use License（非 OSI），商用前请确认条款。", + "input_hint_en": "**IndexTTS2.5**: zero-shot cloning in zh/en/ja/es/ar. Upload a reference voice to clone; pass `lang` (default auto: zh when the text contains Han characters, otherwise en) and emotion options through the JSON box. Weights are under the bilibili Model Use License (not OSI-approved) — check terms before commercial use." }, { "id": "irodori-tts", "display_name": "Irodori-TTS v4 Small (tts 日语, GGUF Q8)", "display_name_en": "Irodori-TTS v4 Small (ja tts, GGUF Q8)", "family": "irodori_tts", "path": "models/Irodori-TTS-v4-Small-GGUF", "task": "tts", "mode": "offline", "download_id": "irodori_tts_v4_small_q8_0", "min_vram_gb": 4, "input_hint": "**Irodori-TTS v4 Small**：日语 TTS；可不上传参考音色直接生成，也可上传参考音色进行克隆；可在声音设计页用日语 caption 描述音色。", "input_hint_en": "**Irodori-TTS v4 Small**: Japanese TTS. Generate without a reference voice, clone from an uploaded reference, or use the voice-design page with a Japanese voice caption." }, diff --git a/webui/configs/required_files.json b/webui/configs/required_files.json index dbc337e4..fc556abc 100644 --- a/webui/configs/required_files.json +++ b/webui/configs/required_files.json @@ -366,6 +366,15 @@ "qwen0.6bemo4-merge/merges.txt", "qwen0.6bemo4-merge/model.safetensors" ], + "index_tts2_5_q8_0": [ + "index-tts2_5-q8_0.gguf" + ], + "index_tts2_5_f16": [ + "index-tts2_5-f16.gguf" + ], + "index_tts2_5_orig": [ + "index-tts2_5-orig.gguf" + ], "mel_band_roformer": [ "config.json", "model.safetensors" From 77e53f5c562ed1db63fff622a4aab2959bdc1641 Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Tue, 11 Aug 2026 18:01:18 +0800 Subject: [PATCH 2/8] index_tts2_5: fix CPU-backend noise in s2mel CFM wavenet The zero skip-connections accumulator in cfm_wavenet was built with sub(ctx, input_bct, input_bct) where input_bct is a permuted (transposed) view. The ggml CPU binary-op kernels only require contiguous rows for src1 and silently misread a permuted src0 (the CUDA/HIP kernels handle it), so on CPU the accumulator started as garbage and the mel output collapsed into noise. Materialize the tensor before the sub. Also adds env-gated debug taps used for the investigation: INDEXTTS25_DUMP_DIR (NPY dumps of GPT/codec/LR/CFM stages) and INDEXTTS25_CFM_NOISE_NPY (diffusion noise injection); zero cost when unset. Verified: CPU/CUDA/HIP backends all pass ASR transcription checks against the official Python reference. --- src/models/index_tts2_5/s2mel.cpp | 222 +++++++++++++++++++++++++++++- 1 file changed, 217 insertions(+), 5 deletions(-) diff --git a/src/models/index_tts2_5/s2mel.cpp b/src/models/index_tts2_5/s2mel.cpp index d6890b00..4be40939 100644 --- a/src/models/index_tts2_5/s2mel.cpp +++ b/src/models/index_tts2_5/s2mel.cpp @@ -15,6 +15,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -52,6 +55,109 @@ struct GgmlContextDeleter { } }; +// Debug helpers for backend divergence investigation, enabled by environment +// variables: INDEXTTS25_DUMP_DIR dumps the initial CFM noise and the first +// diffusion velocity as NPY files; INDEXTTS25_CFM_NOISE_NPY replaces the RNG +// initial noise with the contents of an NPY float32 file. +struct CfmDebugTap { + std::string name; + core::TensorValue value; +}; + +std::string cfm_debug_dump_dir() { + const char * dir = std::getenv("INDEXTTS25_DUMP_DIR"); + if (dir == nullptr || *dir == '\0') { + return {}; + } + return std::string(dir); +} + +void cfm_write_npy_f32( + const std::string & dir, + const std::string & name, + const std::vector & shape, + const std::vector & values) { + if (dir.empty()) { + return; + } + std::string header = "{'descr': '(header.size()); + out.write(reinterpret_cast(&header_len), sizeof(header_len)); + out.write(header.data(), static_cast(header.size())); + out.write(reinterpret_cast(values.data()), static_cast(values.size() * sizeof(float))); +} + +std::vector cfm_read_npy_f32(const std::string & path, size_t expected_count) { + std::ifstream in(path, std::ios::binary); + if (!in) { + throw std::runtime_error("IndexTTS2.5 failed to open noise file: " + path); + } + char magic[8]; + in.read(magic, 8); + if (!in || std::memcmp(magic, "\x93NUMPY\x01\x00", 8) != 0) { + throw std::runtime_error("IndexTTS2.5 noise file is not an NPY v1 file: " + path); + } + uint16_t header_len = 0; + in.read(reinterpret_cast(&header_len), sizeof(header_len)); + std::string header(header_len, '\0'); + in.read(header.data(), header_len); + if (!in || header.find("'(std::stoll(token.substr(first))); + } + if (comma == std::string::npos) { + break; + } + begin = comma + 1; + } + if (count != expected_count) { + throw std::runtime_error("IndexTTS2.5 noise file element count mismatch: " + path); + } + std::vector out(count); + in.read(reinterpret_cast(out.data()), static_cast(count * sizeof(float))); + if (!in) { + throw std::runtime_error("IndexTTS2.5 noise file data is truncated: " + path); + } + return out; +} + +std::string cfm_noise_path_from_env() { + const char * path = std::getenv("INDEXTTS25_CFM_NOISE_NPY"); + if (path == nullptr || *path == '\0') { + return {}; + } + return std::string(path); +} + + core::TensorValue sub(core::ModuleBuildContext & ctx, const core::TensorValue & lhs, const core::TensorValue & rhs) { core::validate_shape(rhs, lhs.shape, "Sub rhs"); return core::wrap_tensor(ggml_sub(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32); @@ -234,10 +340,22 @@ core::TensorValue cfm_wavenet( core::ModuleBuildContext & ctx, const core::TensorValue & input_bct, const core::TensorValue & timestep_b, - const IndexTTS25S2MelCfmWeights & weights) { + const IndexTTS25S2MelCfmWeights & weights, + std::vector * debug_taps = nullptr) { + const auto tap = [&](const std::string & name, const core::TensorValue & value) { + if (debug_taps != nullptr) { + debug_taps->push_back({name, core::ensure_backend_addressable_layout(ctx, value)}); + } + }; auto g = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, timestep_b), core::TensorShape::from_dims({timestep_b.shape.dims[0], kHidden, 1})); g = modules::Conv1dModule({kHidden, 2 * kHidden * kWavenetLayers, 1, 1, 0, 1, true}).build(ctx, g, weights.wavenet_cond); - auto output = sub(ctx, input_bct, input_bct); + tap("wn_g", g); + // The zero accumulator must come from a contiguous tensor: input_bct is a + // permuted (transposed) view, and the ggml CPU binary-op kernels miscompute + // permuted src operands (the CUDA kernels handle them). + const auto zeros_base = core::ensure_backend_addressable_layout(ctx, input_bct); + auto output = sub(ctx, zeros_base, zeros_base); + tap("wn_zeros", output); auto x = input_bct; for (int64_t i = 0; i < kWavenetLayers; ++i) { const int64_t dilation = 1; @@ -246,6 +364,13 @@ core::TensorValue cfm_wavenet( auto x_in = modules::Conv1dModule( {kHidden, 2 * kHidden, kWavenetKernel, 1, 0, static_cast(dilation), true}) .build(ctx, x_padded, weights.wavenet_layers[static_cast(i)].in_layer); + if (i == 0) { + tap("wn_x_padded0", x_padded); + tap("wn_xin0", x_in); + } + if (i == 1) { + tap("wn_xin1", x_in); + } auto g_l = modules::SliceModule({1, i * 2 * kHidden, 2 * kHidden}).build(ctx, g); g_l = modules::RepeatModule({x_in.shape}).build(ctx, g_l); auto acts = modules::AddModule{}.build(ctx, x_in, g_l); @@ -254,9 +379,15 @@ core::TensorValue cfm_wavenet( auto sigmoid_part = modules::SliceModule({1, kHidden, kHidden}).build(ctx, acts); sigmoid_part = modules::SigmoidModule{}.build(ctx, sigmoid_part); acts = modules::MulModule{}.build(ctx, tanh_part, sigmoid_part); + if (i == 0) { + tap("wn_acts0", acts); + } const int64_t res_skip_channels = i < kWavenetLayers - 1 ? 2 * kHidden : kHidden; auto res_skip = modules::Conv1dModule({kHidden, res_skip_channels, 1, 1, 0, 1, true}) .build(ctx, acts, weights.wavenet_layers[static_cast(i)].res_skip_layer); + if (i == 0) { + tap("wn_res_skip0", res_skip); + } if (i < kWavenetLayers - 1) { auto res = modules::SliceModule({1, 0, kHidden}).build(ctx, res_skip); auto skip = modules::SliceModule({1, kHidden, kHidden}).build(ctx, res_skip); @@ -265,7 +396,18 @@ core::TensorValue cfm_wavenet( } else { output = modules::AddModule{}.build(ctx, output, res_skip); } + if (i == 0) { + tap("wn_x1", x); + tap("wn_output1", output); + } + if (i == 1) { + tap("wn_output2", output); + } + if (i == 2) { + tap("wn_output3", output); + } } + tap("wn_final", output); return output; } @@ -291,10 +433,17 @@ core::TensorValue build_cfm_estimator( const core::TensorValue & style_bc, const core::TensorValue & timestep_b, const core::TensorValue & positions, - const IndexTTS25S2MelCfmWeights & weights) { + const IndexTTS25S2MelCfmWeights & weights, + std::vector * debug_taps = nullptr) { + const auto tap = [&](const std::string & name, const core::TensorValue & value) { + if (debug_taps != nullptr) { + debug_taps->push_back({name, core::ensure_backend_addressable_layout(ctx, value)}); + } + }; const int64_t batch = x_bct.shape.dims[0]; const int64_t frames = x_bct.shape.dims[2]; auto t1 = timestep_embedding(ctx, timestep_b, weights.time_freqs, weights.time_mlp0, weights.time_mlp2); + tap("t1", t1); auto cond = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, cond_btc, weights.cond_projection); auto x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x_bct); auto prompt = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, prompt_bct); @@ -304,6 +453,7 @@ core::TensorValue build_cfm_estimator( hidden = modules::ConcatModule({2}).build(ctx, hidden, style); hidden = modules::LinearModule({kHidden + 2 * kMelChannels + kStyleDim, kHidden, true, GGML_PREC_F32}) .build(ctx, hidden, weights.cond_x_merge); + tap("merged", hidden); std::vector skips; skips.reserve(static_cast(kDitLayers / 2)); @@ -313,6 +463,9 @@ core::TensorValue build_cfm_estimator( skip = &skips.back(); } hidden = cfm_transformer_layer(ctx, hidden, t1, positions, weights.dit_layers[static_cast(i)], skip); + if (i == 0) { + tap("dit_layer0", hidden); + } if (i > kDitLayers / 2) { skips.pop_back(); } else if (i < kDitLayers / 2) { @@ -320,15 +473,18 @@ core::TensorValue build_cfm_estimator( } } hidden = adaptive_rms_norm(ctx, hidden, t1, weights.dit_norm); + tap("dit_out", hidden); hidden = modules::LinearModule({kHidden + kMelChannels, kHidden, true, GGML_PREC_F32}) .build(ctx, modules::ConcatModule({2}).build(ctx, hidden, x), weights.skip_linear); auto wavenet_x = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, weights.conv1); wavenet_x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, wavenet_x); auto t2 = timestep_embedding(ctx, timestep_b, weights.time2_freqs, weights.time2_mlp0, weights.time2_mlp2); - wavenet_x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, cfm_wavenet(ctx, wavenet_x, t2, weights)); + wavenet_x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, cfm_wavenet(ctx, wavenet_x, t2, weights, debug_taps)); auto projected = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, weights.res_projection); hidden = modules::AddModule{}.build(ctx, wavenet_x, projected); + tap("wavenet_out", hidden); hidden = cfm_final_layer(ctx, hidden, t1, weights); + tap("final_hidden", hidden); hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, hidden); return modules::Conv1dModule({kHidden, kMelChannels, 1, 1, 0, 1, true}).build(ctx, hidden, weights.conv2); } @@ -898,6 +1054,11 @@ class IndexTTS25S2MelRuntime::CfmGraph { if (input_ctx_ == nullptr) { throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel CFM input context"); } + ggml_init_params output_params{16ull * 1024ull * 1024ull, nullptr, true}; + output_ctx_.reset(ggml_init(output_params)); + if (output_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel CFM output context"); + } core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.s2mel.cfm", execution_.backend_type()}; core::ModuleBuildContext input_ctx{input_ctx_.get(), "index_tts2_5.s2mel.cfm.inputs", execution_.backend_type()}; x_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, kMelChannels, frames_})) @@ -914,6 +1075,7 @@ class IndexTTS25S2MelRuntime::CfmGraph { ggml_set_input(style_); ggml_set_input(timestep_); ggml_set_input(positions_); + std::vector taps; auto output = build_cfm_estimator( ctx, core::wrap_tensor(x_, core::TensorShape::from_dims({batch_, kMelChannels, frames_}), GGML_TYPE_F32), @@ -922,15 +1084,29 @@ class IndexTTS25S2MelRuntime::CfmGraph { core::wrap_tensor(style_, core::TensorShape::from_dims({batch_, kStyleDim}), GGML_TYPE_F32), core::wrap_tensor(timestep_, core::TensorShape::from_dims({batch_}), GGML_TYPE_F32), core::wrap_tensor(positions_, core::TensorShape::from_dims({frames_}), GGML_TYPE_I32), - weights_->cfm); + weights_->cfm, + cfm_debug_dump_dir().empty() ? nullptr : &taps); output_ = core::ensure_backend_addressable_layout(ctx, output).tensor; ggml_set_output(output_); graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(131072, frames_ * 4096)), false); ggml_build_forward_expand(graph_, output_); + core::ModuleBuildContext output_ctx{output_ctx_.get(), "index_tts2_5.s2mel.cfm.outputs", execution_.backend_type()}; + for (const auto & tap : taps) { + auto * tap_output = core::make_tensor(output_ctx, GGML_TYPE_F32, tap.value.shape).tensor; + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), tap.value.tensor, tap_output)); + debug_taps_.push_back({tap.name, tap_output, tap.value.shape}); + } input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); if (input_buffer_ == nullptr) { throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel CFM input buffer"); } + if (!debug_taps_.empty()) { + output_buffer_ = ggml_backend_alloc_ctx_tensors(output_ctx_.get(), execution_.backend()); + if (output_buffer_ == nullptr) { + clear_graph(); + throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel CFM output buffer"); + } + } gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) || @@ -989,6 +1165,16 @@ class IndexTTS25S2MelRuntime::CfmGraph { timing_start = Clock::now(); ggml_backend_tensor_get(output_, out.data(), 0, out.size() * sizeof(float)); debug::timing_log_scalar("index_tts2_5.s2mel.cfm.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); + const std::string dump_dir = cfm_debug_dump_dir(); + if (!dump_dir.empty() && !taps_dumped_) { + taps_dumped_ = true; + for (const auto & tap : debug_taps_) { + std::vector values(tap.shape.num_elements()); + ggml_backend_tensor_get(tap.tensor, values.data(), 0, values.size() * sizeof(float)); + const std::vector dims(tap.shape.dims.begin(), tap.shape.dims.begin() + tap.shape.rank); + cfm_write_npy_f32(dump_dir, "cfm_tap_" + tap.name, dims, values); + } + } return out; } @@ -1006,6 +1192,10 @@ class IndexTTS25S2MelRuntime::CfmGraph { ggml_backend_buffer_free(input_buffer_); input_buffer_ = nullptr; } + if (output_buffer_ != nullptr) { + ggml_backend_buffer_free(output_buffer_); + output_buffer_ = nullptr; + } } core::ExecutionContext & execution_; @@ -1014,6 +1204,7 @@ class IndexTTS25S2MelRuntime::CfmGraph { bool use_cfg_ = false; int64_t batch_ = 1; std::unique_ptr input_ctx_; + std::unique_ptr output_ctx_; std::unique_ptr ctx_; ggml_tensor * x_ = nullptr; ggml_tensor * prompt_ = nullptr; @@ -1022,9 +1213,17 @@ class IndexTTS25S2MelRuntime::CfmGraph { ggml_tensor * timestep_ = nullptr; ggml_tensor * positions_ = nullptr; ggml_tensor * output_ = nullptr; + struct DebugTapTensor { + std::string name; + ggml_tensor * tensor = nullptr; + core::TensorShape shape; + }; + std::vector debug_taps_; + bool taps_dumped_ = false; ggml_cgraph * graph_ = nullptr; ggml_gallocr_t gallocr_ = nullptr; ggml_backend_buffer_t input_buffer_ = nullptr; + ggml_backend_buffer_t output_buffer_ = nullptr; std::vector positions_values_; }; @@ -1221,9 +1420,15 @@ IndexTTS25S2MelMel IndexTTS25S2MelRuntime::infer_mel( rng_offset_blocks, rng_policy, engine::sampling::TorchRandnPrecision::Float32); + const std::string noise_path = cfm_noise_path_from_env(); + if (!noise_path.empty()) { + x = cfm_read_npy_f32(noise_path, static_cast(kMelChannels * total_frames)); + } auto prompt_x = make_prompt_x(reference_mel, kMelChannels, total_frames, reference_frames); zero_prompt_region(x, kMelChannels, total_frames, reference_frames); auto mu = make_condition_with_prompt(condition, total_frames); + const std::string cfm_dump_dir = cfm_debug_dump_dir(); + cfm_write_npy_f32(cfm_dump_dir, "cfm_noise", {kMelChannels, total_frames}, x); const auto cfm_start = Clock::now(); double graph_ms = 0.0; @@ -1237,6 +1442,13 @@ IndexTTS25S2MelMel IndexTTS25S2MelRuntime::infer_mel( auto style_batched = repeat_or_zero_rows(style, kStyleDim, use_cfg, true); std::vector timestep(static_cast(use_cfg ? 2 : 1), t); const auto velocity = cfm_graph_->run(x_batched, prompt_batched, cond_batched, style_batched, timestep); + if (step == 1) { + cfm_write_npy_f32( + cfm_dump_dir, + "cfm_velocity0", + {use_cfg ? 2 : 1, kMelChannels, total_frames}, + velocity); + } graph_ms += engine::debug::elapsed_ms(graph_start); const int64_t row_values = kMelChannels * total_frames; for (int64_t i = 0; i < row_values; ++i) { From 90a4df951259f6ce6d698f12574e855d64b9d17c Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Tue, 11 Aug 2026 20:32:50 +0800 Subject: [PATCH 3/8] index_tts2_5: add upstream-to-GGUF conversion tool and docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tools/convert_index_tts2_5.py stages an official IndexTeam/IndexTTS-2.5 snapshot (.pth checkpoints) into the Safetensors layout the engine expects — unwrapping the s2mel/codec container keys, prefixing CAMPPlus tensors with speaker_encoder., stripping BigVGAN's generator. prefix, wrapping the feat matrices — assembles the sidecar root (config, tiktoken vocabulary, auxiliary configs), and prints or runs the audiocpp_gguf command. Verified end-to-end on CUDA: fresh staging -> f16 GGUF -> inference -> ASR transcript matches the official Python reference; staged tensor name sets match the hand-verified staging for all 10 namespaces. --- docs/gguf.md | 1 + docs/tts.md | 13 ++ tools/convert_index_tts2_5.py | 263 ++++++++++++++++++++++++++++++++++ 3 files changed, 277 insertions(+) create mode 100644 tools/convert_index_tts2_5.py diff --git a/docs/gguf.md b/docs/gguf.md index 86873359..62d2e17e 100644 --- a/docs/gguf.md +++ b/docs/gguf.md @@ -70,6 +70,7 @@ Status labels: | `hviske_asr` | Done | Pass | --- | --- | Pass | | `inflect_v2` | Done | Pass | Pass | --- | --- | | `index_tts2` | Done | Pass | Pass | Pass (drift) | Pass (ASR match, drift) | +| `index_tts2_5` | Done | Pass | Pass | Pass | Pass (ASR match, drift) | | `irodori_tts` | Done | Pass | --- | Pass | Pass (ASR match, drift) | | `kroko_asr` | Done | Pass | --- | --- | Pass | | `marblenet_vad` | Bundled (tiny model) | Pass | --- | --- | --- | diff --git a/docs/tts.md b/docs/tts.md index 470d26c1..1a2ee594 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -613,6 +613,19 @@ License: IndexTTS-2.5 weights are distributed under the bilibili Model Use Licen | `--session-option index_tts2_5.emotion_text_max_new_tokens=` | tokens | `256` | Maximum generated tokens for emotion-text classification. | | `--session-option index_tts2_5.weight_context_mb=` | MB | `32` | Shared ggml weight metadata context size. | +### Converting From Upstream Weights + +`tools/convert_index_tts2_5.py` turns an official `IndexTeam/IndexTTS-2.5` snapshot (the `.pth` checkpoints) into the Safetensors staging layout the engine expects, and prints (or runs) the matching `audiocpp_gguf` command. The w2v-bert-2.0, CAMPPlus, and BigVGAN checkpoints are auto-detected under `/hf_cache/` (run the official inference once to populate it) and each has an explicit override flag: + +```bash +python tools/convert_index_tts2_5.py \ + --model-dir /path/to/IndexTTS-2.5 \ + --output-dir /path/to/staging \ + --run-converter /path/to/audiocpp_gguf --type q8_0 +``` + +The script repackages the checkpoints the loader needs (unwraps the `s2mel.pth`/`codec.pth` container keys, prefixes CAMPPlus tensors with `speaker_encoder.`, strips BigVGAN's `generator.` prefix, wraps the `feat1/feat2.pt` matrices as a single `tensor`) and assembles the sidecar `root/` (config, tiktoken vocabulary, auxiliary model configs) that gets embedded into the GGUF. + ## Irodori-TTS Irodori-TTS is Japanese TTS under `--family irodori_tts`. v4 Small is the preferred GGUF-first package and supports no-reference speech, reference-conditioned speech, and caption-based voice design in one checkpoint. The older 500M v3 and 600M v3 VoiceDesign packages remain supported for existing users. See [Irodori-TTS](models/irodori_tts.md) for v3/v4 differences, GGUF variants, options, and compatibility aliases. diff --git a/tools/convert_index_tts2_5.py b/tools/convert_index_tts2_5.py new file mode 100644 index 00000000..5a6c24b7 --- /dev/null +++ b/tools/convert_index_tts2_5.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python3 +"""Prepare an official IndexTTS-2.5 checkpoint for audio.cpp's GGUF converter. + +Point --model-dir at a complete IndexTeam/IndexTTS-2.5 snapshot. The script writes +a staging directory whose Safetensors layout matches the tensor namespaces the +index_tts2_5 engine expects, plus a root/ directory with the sidecar files +(config, tokenizer, auxiliary model configs) that audiocpp_gguf embeds into the +final GGUF. + +The w2v-bert-2.0, CAMPPlus and BigVGAN checkpoints are not part of the official +snapshot; they are auto-detected under /hf_cache/ (where the official +downloader places them) and each can be overridden explicitly. + +This tool does not download anything and never writes into --model-dir. + +Example: + python tools/convert_index_tts2_5.py \ + --model-dir /path/to/IndexTTS-2.5 \ + --output-dir /path/to/staging + +Then run the printed audiocpp_gguf command, or let the script run it: + + python tools/convert_index_tts2_5.py \ + --model-dir /path/to/IndexTTS-2.5 \ + --output-dir /path/to/staging \ + --run-converter /path/to/audiocpp_gguf --type f16 +""" + +from __future__ import annotations + +import argparse +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Dict + +import torch +from safetensors.torch import save_file + +# GGUF tensor namespaces (must match the index_tts2_5 model spec) and the +# staging file each one is produced from. +TENSOR_OUTPUTS = [ + ("gpt", "gpt.safetensors"), + ("s2mel", "s2mel.safetensors"), + ("speaker_matrix", "speaker_matrix.safetensors"), + ("emotion_matrix", "emotion_matrix.safetensors"), + ("wav2vec2bert_stats", "wav2vec2bert_stats.safetensors"), + ("wav2vec2bert", "wav2vec2bert.safetensors"), + ("semantic_codec", "semantic_codec.safetensors"), + ("campplus", "campplus.safetensors"), + ("bigvgan", "bigvgan.safetensors"), + ("qwen_emotion", "qwen_emotion.safetensors"), +] + +QWEN_SIDECARS = ( + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", +) + + +def _require_file(path: Path, label: str) -> Path: + if not path.is_file(): + raise FileNotFoundError(f"missing {label}: {path}") + return path + + +def _load_checkpoint(path: Path): + return torch.load(path, map_location="cpu", weights_only=False) + + +def _flatten(obj, prefix: str = "", out: Dict[str, torch.Tensor] | None = None) -> Dict[str, torch.Tensor]: + if out is None: + out = {} + if isinstance(obj, dict): + for key, value in obj.items(): + _flatten(value, f"{prefix}{key}.", out) + elif hasattr(obj, "shape"): + out[prefix.rstrip(".")] = obj.contiguous() + else: + raise TypeError(f"unexpected non-tensor leaf at {prefix!r}: {type(obj)}") + return out + + +def _save_safetensors(tensors: Dict[str, torch.Tensor], path: Path) -> None: + save_file(tensors, str(path)) + print(f"wrote {path} ({len(tensors)} tensors)") + + +def _convert_gpt(model_dir: Path, output_dir: Path) -> None: + # gpt.pth is already a flat tensor dict (includes spk_emb_proj and + # lang_embedding; the 2.5 checkpoint has no conditioning_encoder/speed_emb). + obj = _load_checkpoint(_require_file(model_dir / "gpt.pth", "gpt.pth")) + _save_safetensors(_flatten(obj), output_dir / "gpt.safetensors") + + +def _convert_s2mel(model_dir: Path, output_dir: Path) -> None: + # s2mel.pth wraps the state dict under "net" (cfm.*/length_regulator.*/gpt_layer.*). + obj = _load_checkpoint(_require_file(model_dir / "s2mel.pth", "s2mel.pth")) + if isinstance(obj, dict) and isinstance(obj.get("net"), dict): + obj = obj["net"] + _save_safetensors(_flatten(obj), output_dir / "s2mel.safetensors") + + +def _convert_semantic_codec(model_dir: Path, output_dir: Path) -> None: + # codec.pth wraps the state dict under "model" (encoder.*/decoder.*/quantizer.*/down/up). + obj = _load_checkpoint(_require_file(model_dir / "codec.pth", "codec.pth")) + if isinstance(obj, dict) and isinstance(obj.get("model"), dict): + obj = obj["model"] + _save_safetensors(_flatten(obj), output_dir / "semantic_codec.safetensors") + + +def _convert_emotion_matrices(model_dir: Path, output_dir: Path) -> None: + # feat1.pt/feat2.pt hold a single root-level tensor each: (73, 192) speaker + # matrix and (73, 1280) emotion matrix. + speaker = _load_checkpoint(_require_file(model_dir / "feat1.pt", "feat1.pt")) + emotion = _load_checkpoint(_require_file(model_dir / "feat2.pt", "feat2.pt")) + _save_safetensors({"tensor": speaker.float().contiguous()}, output_dir / "speaker_matrix.safetensors") + _save_safetensors({"tensor": emotion.float().contiguous()}, output_dir / "emotion_matrix.safetensors") + + +def _convert_wav2vec2bert_stats(model_dir: Path, output_dir: Path) -> None: + obj = _load_checkpoint(_require_file(model_dir / "wav2vec2bert_stats.pt", "wav2vec2bert_stats.pt")) + flat = {key: value.float().contiguous() for key, value in _flatten(obj).items()} + _save_safetensors(flat, output_dir / "wav2vec2bert_stats.safetensors") + + +def _convert_campplus(campplus_checkpoint: Path, output_dir: Path) -> None: + # The engine binds CAMPPlus weights under the "speaker_encoder." prefix. + obj = _load_checkpoint(_require_file(campplus_checkpoint, "campplus checkpoint")) + flat = _flatten(obj) + flat = {key if key.startswith("speaker_encoder.") else f"speaker_encoder.{key}": value for key, value in flat.items()} + _save_safetensors(flat, output_dir / "campplus.safetensors") + + +def _convert_bigvgan(bigvgan_dir: Path, output_dir: Path) -> None: + # bigvgan_generator.pt stores keys with a "generator." prefix; the engine + # expects bare names (conv_pre/ups.N/...). + obj = _load_checkpoint(_require_file(bigvgan_dir / "bigvgan_generator.pt", "bigvgan generator")) + flat = _flatten(obj) + flat = {key[len("generator."):] if key.startswith("generator.") else key: value for key, value in flat.items()} + _save_safetensors(flat, output_dir / "bigvgan.safetensors") + + +def _copy(src: Path, dst: Path, label: str) -> None: + _require_file(src, label) + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(src, dst) + print(f"copied {src} -> {dst}") + + +def build_converter_command(output_dir: Path, converter: str, quant_type: str) -> list[str]: + command = [converter] + for namespace, filename in TENSOR_OUTPUTS: + command += ["--input", f"{namespace}={output_dir / filename}"] + command += [ + "--root", str(output_dir / "root"), + "--family", "index_tts2_5", + "--type", quant_type, + "--output", str(output_dir / f"index-tts2_5-{quant_type}.gguf"), + ] + return command + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Stage an official IndexTTS-2.5 snapshot for audio.cpp's GGUF converter.", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--model-dir", type=Path, required=True, + help="Path to the official IndexTTS-2.5 snapshot (gpt.pth, s2mel.pth, ...).") + parser.add_argument("--output-dir", type=Path, required=True, + help="Staging directory to write; created if missing. Never inside --model-dir.") + parser.add_argument("--w2v-bert-dir", type=Path, default=None, + help="Directory with w2v-bert-2.0 model.safetensors/config.json/preprocessor_config.json. " + "Default: /hf_cache/w2v-bert-2.0") + parser.add_argument("--campplus-checkpoint", type=Path, default=None, + help="Path to campplus_cn_common.bin. Default: /hf_cache/campplus/campplus_cn_common.bin") + parser.add_argument("--bigvgan-dir", type=Path, default=None, + help="Directory with bigvgan_generator.pt and config.json. Default: /hf_cache/bigvgan") + parser.add_argument("--run-converter", type=str, default=None, + help="Path to the audiocpp_gguf executable; when given, run the GGUF conversion right away.") + parser.add_argument("--type", dest="quant_type", default="f16", + choices=("orig", "f16", "bf16", "q8_0", "q2_k", "q3_k", "q4_k", "q5_k", "q6_k"), + help="GGUF weight type used with --run-converter and in the printed command (default: f16).") + args = parser.parse_args() + + model_dir = args.model_dir.resolve() + output_dir = args.output_dir.resolve() + if not model_dir.is_dir(): + print(f"error: --model-dir does not exist: {model_dir}", file=sys.stderr) + return 1 + if output_dir == model_dir or model_dir in output_dir.parents: + print("error: --output-dir must not be inside --model-dir", file=sys.stderr) + return 1 + + hf_cache = model_dir / "hf_cache" + w2v_bert_dir = (args.w2v_bert_dir or hf_cache / "w2v-bert-2.0").resolve() + campplus_checkpoint = (args.campplus_checkpoint or hf_cache / "campplus" / "campplus_cn_common.bin").resolve() + bigvgan_dir = (args.bigvgan_dir or hf_cache / "bigvgan").resolve() + + output_dir.mkdir(parents=True, exist_ok=True) + root_dir = output_dir / "root" + root_dir.mkdir(parents=True, exist_ok=True) + + _convert_gpt(model_dir, output_dir) + _convert_s2mel(model_dir, output_dir) + _convert_semantic_codec(model_dir, output_dir) + _convert_emotion_matrices(model_dir, output_dir) + _convert_wav2vec2bert_stats(model_dir, output_dir) + _convert_campplus(campplus_checkpoint, output_dir) + _convert_bigvgan(bigvgan_dir, output_dir) + + # Auxiliary checkpoints ship as Safetensors already; copy them through. + _copy(w2v_bert_dir / "model.safetensors", output_dir / "wav2vec2bert.safetensors", "w2v-bert-2.0 weights") + _copy(model_dir / "qwen0.6bemo4-merge" / "model.safetensors", output_dir / "qwen_emotion.safetensors", + "qwen emotion weights") + + # Sidecar files embedded into the GGUF via --root. + _copy(model_dir / "config.yaml", root_dir / "config.yaml", "config.yaml") + _copy(model_dir / "multilingual_zh_ja_yue_char_del.tiktoken", + root_dir / "multilingual_zh_ja_yue_char_del.tiktoken", "tiktoken vocabulary") + _copy(w2v_bert_dir / "config.json", root_dir / "w2v-bert-2.0" / "config.json", "w2v-bert-2.0 config") + _copy(w2v_bert_dir / "preprocessor_config.json", root_dir / "w2v-bert-2.0" / "preprocessor_config.json", + "w2v-bert-2.0 preprocessor config") + _copy(bigvgan_dir / "config.json", root_dir / "bigvgan" / "config.json", "bigvgan config") + for name in QWEN_SIDECARS: + _copy(model_dir / "qwen0.6bemo4-merge" / name, root_dir / "qwen0.6bemo4-merge" / name, f"qwen sidecar {name}") + + converter = args.run_converter or "audiocpp_gguf" + command = build_converter_command(output_dir, converter, args.quant_type) + lines = [command[0]] + index = 1 + while index < len(command): + flag = command[index] + if flag.startswith("--") and index + 1 < len(command) and not command[index + 1].startswith("--"): + lines.append(f" {flag} {command[index + 1]} \\") + index += 2 + else: + lines.append(f" {flag} \\") + index += 1 + lines[-1] = lines[-1].rstrip(" \\") + print() + print("staging complete. Convert to GGUF with:") + print("\n".join(lines)) + + if args.run_converter is not None: + print() + print("running converter...") + result = subprocess.run(command) + if result.returncode != 0: + print(f"error: converter exited with {result.returncode}", file=sys.stderr) + return result.returncode + print(f"GGUF written to {output_dir / f'index-tts2_5-{args.quant_type}.gguf'}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 8c51a9e5f2e495efd827645bc852607244475fae Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Tue, 11 Aug 2026 20:38:37 +0800 Subject: [PATCH 4/8] index_tts2_5: emit native Safetensors layout from the conversion tool --native-dir assembles a directly loadable model directory (the spec's safetensors source layout: feat1/feat2.safetensors, semantic_codec_model.safetensors, bigvgan/ and w2v-bert-2.0/ subdirectories, sidecar configs) hardlinked from the staging output. Verified on CUDA: script-produced native directory loads and its ASR transcript matches the official reference. --- docs/tts.md | 2 ++ tools/convert_index_tts2_5.py | 59 +++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/docs/tts.md b/docs/tts.md index 1a2ee594..48506430 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -624,6 +624,8 @@ python tools/convert_index_tts2_5.py \ --run-converter /path/to/audiocpp_gguf --type q8_0 ``` +Pass `--native-dir /path/to/IndexTTS-2.5-native` to also emit a directly loadable native Safetensors model directory (hardlinked from the staging files), no GGUF conversion required. + The script repackages the checkpoints the loader needs (unwraps the `s2mel.pth`/`codec.pth` container keys, prefixes CAMPPlus tensors with `speaker_encoder.`, strips BigVGAN's `generator.` prefix, wraps the `feat1/feat2.pt` matrices as a single `tensor`) and assembles the sidecar `root/` (config, tiktoken vocabulary, auxiliary model configs) that gets embedded into the GGUF. ## Irodori-TTS diff --git a/tools/convert_index_tts2_5.py b/tools/convert_index_tts2_5.py index 5a6c24b7..10957c34 100644 --- a/tools/convert_index_tts2_5.py +++ b/tools/convert_index_tts2_5.py @@ -24,11 +24,15 @@ --model-dir /path/to/IndexTTS-2.5 \ --output-dir /path/to/staging \ --run-converter /path/to/audiocpp_gguf --type f16 + +Add --native-dir /path/to/native-model to also emit a directly loadable +native Safetensors model directory (no GGUF conversion needed). """ from __future__ import annotations import argparse +import os import shutil import subprocess import sys @@ -166,6 +170,48 @@ def build_converter_command(output_dir: Path, converter: str, quant_type: str) - return command +# Native Safetensors model directory layout: the spec's safetensors source maps +# logical tensor groups to these paths under the model root. +NATIVE_TENSOR_LAYOUT = [ + # (staging filename, relative path in the native model directory) + ("gpt.safetensors", "gpt.safetensors"), + ("s2mel.safetensors", "s2mel.safetensors"), + ("speaker_matrix.safetensors", "feat1.safetensors"), + ("emotion_matrix.safetensors", "feat2.safetensors"), + ("wav2vec2bert_stats.safetensors", "wav2vec2bert_stats.safetensors"), + ("wav2vec2bert.safetensors", "w2v-bert-2.0/model.safetensors"), + ("semantic_codec.safetensors", "semantic_codec_model.safetensors"), + ("campplus.safetensors", "campplus.safetensors"), + ("bigvgan.safetensors", "bigvgan/model.safetensors"), + ("qwen_emotion.safetensors", "qwen0.6bemo4-merge/model.safetensors"), +] + + +def _link_or_copy(src: Path, dst: Path) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + if dst.exists(): + dst.unlink() + try: + os.link(src, dst) + except OSError: + shutil.copyfile(src, dst) + + +def write_native_layout(output_dir: Path, native_dir: Path) -> None: + """Assemble the directly loadable native Safetensors model directory.""" + for staging_name, relative in NATIVE_TENSOR_LAYOUT: + _link_or_copy(_require_file(output_dir / staging_name, staging_name), native_dir / relative) + root_dir = output_dir / "root" + for name in ("config.yaml", "multilingual_zh_ja_yue_char_del.tiktoken", + "w2v-bert-2.0/config.json", "w2v-bert-2.0/preprocessor_config.json", + "bigvgan/config.json"): + _link_or_copy(_require_file(root_dir / name, name), native_dir / name) + for name in QWEN_SIDECARS: + _link_or_copy(_require_file(root_dir / "qwen0.6bemo4-merge" / name, name), + native_dir / "qwen0.6bemo4-merge" / name) + print(f"native model directory written to {native_dir}") + + def main() -> int: parser = argparse.ArgumentParser( description="Stage an official IndexTTS-2.5 snapshot for audio.cpp's GGUF converter.", @@ -184,6 +230,11 @@ def main() -> int: help="Directory with bigvgan_generator.pt and config.json. Default: /hf_cache/bigvgan") parser.add_argument("--run-converter", type=str, default=None, help="Path to the audiocpp_gguf executable; when given, run the GGUF conversion right away.") + parser.add_argument("--native-dir", type=Path, default=None, + help="Also write a directly loadable native Safetensors model directory (the layout of " + "the spec's safetensors source: feat1/feat2.safetensors, semantic_codec_model.safetensors, " + "bigvgan/model.safetensors, w2v-bert-2.0/, qwen0.6bemo4-merge/). Files are hardlinked " + "from the staging directory when possible, copied otherwise.") parser.add_argument("--type", dest="quant_type", default="f16", choices=("orig", "f16", "bf16", "q8_0", "q2_k", "q3_k", "q4_k", "q5_k", "q6_k"), help="GGUF weight type used with --run-converter and in the printed command (default: f16).") @@ -231,6 +282,14 @@ def main() -> int: for name in QWEN_SIDECARS: _copy(model_dir / "qwen0.6bemo4-merge" / name, root_dir / "qwen0.6bemo4-merge" / name, f"qwen sidecar {name}") + if args.native_dir is not None: + native_dir = args.native_dir.resolve() + if native_dir == model_dir or model_dir in native_dir.parents: + print("error: --native-dir must not be inside --model-dir", file=sys.stderr) + return 1 + native_dir.mkdir(parents=True, exist_ok=True) + write_native_layout(output_dir, native_dir) + converter = args.run_converter or "audiocpp_gguf" command = build_converter_command(output_dir, converter, args.quant_type) lines = [command[0]] From 95fba32ec6355740d88075e88a4531d9f09ee70b Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Wed, 12 Aug 2026 11:03:46 +0800 Subject: [PATCH 5/8] index_tts2: merge IndexTTS-2.5 as a variant of the index_tts2 family Replace the copied index_tts2_5 model family with an additive variant of the existing index_tts2 implementation, selected from the model config version field ("2.5") instead of a separate family or tensor-name probing: - one IndexTTS2TextTokenizer interface with per-variant backends: v2 keeps the SentencePiece behavior, v2.5 uses the multilingual tiktoken/BPE vocabulary with language special tokens and language-id handling - GPT speaker conditioning split into explicit variant modes: v2 keeps the conditioning_encoder + perceiver_encoder path with speed embeddings, v2.5 loads spk_emb_proj + lang_embedding and builds the campplus conditioning prefix inside the prefill graph; decode/cache logic stays shared - semantic codec gains a variant decode path: v2.5 runs the full EnhancedCodec decoder with the 2x nearest upsample and up convolution - v2.5 regulates the raw w2v-bert semantic directly (no codec quantize, no GPT latent projection) while v2 keeps its existing flow - 2.5 packages and the tiktoken tokenizer file join the index_tts2 model spec (tokenizer files are optional per variant); the conversion tool normalizes the staged config version to "2.5" and targets index_tts2 - the lang request option is parsed for the v2.5 tokenizer - drop the debug-only dump paths, noise injection and environment parity hooks; the s2mel CFM wavenet keeps the contiguous zero accumulator so the CPU backend computes permuted operands correctly Greedy (do_sample=false, seed-pinned) outputs are bit-identical to the previous split implementation for both variants on CUDA: v2.5 zh/ja GGUF and v2 zh native Safetensors. --- CMakeLists.txt | 22 - README.md | 3 +- docs/gguf.md | 3 +- docs/tts.md | 38 +- include/engine/models/index_tts2/gpt.h | 18 + include/engine/models/index_tts2/request.h | 4 + .../engine/models/index_tts2/semantic_codec.h | 2 + include/engine/models/index_tts2/session.h | 1 + .../engine/models/index_tts2/tokenizer_text.h | 57 +- include/engine/models/index_tts2/types.h | 21 + include/engine/models/index_tts2_5/assets.h | 29 - .../models/index_tts2_5/audio_features.h | 54 - include/engine/models/index_tts2_5/gpt.h | 186 -- include/engine/models/index_tts2_5/loader.h | 33 - .../engine/models/index_tts2_5/qwen_emotion.h | 82 - include/engine/models/index_tts2_5/request.h | 14 - include/engine/models/index_tts2_5/s2mel.h | 151 -- .../models/index_tts2_5/semantic_codec.h | 94 - .../models/index_tts2_5/semantic_encoder.h | 100 - include/engine/models/index_tts2_5/session.h | 122 - .../models/index_tts2_5/style_encoder.h | 34 - .../models/index_tts2_5/tokenizer_text.h | 58 - include/engine/models/index_tts2_5/types.h | 153 -- include/engine/models/index_tts2_5/vocoder.h | 34 - model_specs/index_tts2.json | 50 +- model_specs/index_tts2_5.json | 186 -- src/models/index_tts2/assets.cpp | 40 +- src/models/index_tts2/gpt.cpp | 263 +- src/models/index_tts2/loader.cpp | 18 +- src/models/index_tts2/request.cpp | 16 + src/models/index_tts2/s2mel.cpp | 6 +- src/models/index_tts2/semantic_codec.cpp | 53 +- src/models/index_tts2/session.cpp | 125 +- src/models/index_tts2/tokenizer_text.cpp | 697 +++++- src/models/index_tts2_5/assets.cpp | 259 -- src/models/index_tts2_5/audio_features.cpp | 547 ---- src/models/index_tts2_5/gpt.cpp | 2217 ----------------- src/models/index_tts2_5/loader.cpp | 157 -- src/models/index_tts2_5/qwen_emotion.cpp | 794 ------ src/models/index_tts2_5/request.cpp | 192 -- src/models/index_tts2_5/s2mel.cpp | 1483 ----------- src/models/index_tts2_5/semantic_codec.cpp | 723 ------ src/models/index_tts2_5/semantic_encoder.cpp | 622 ----- src/models/index_tts2_5/session.cpp | 838 ------- src/models/index_tts2_5/style_encoder.cpp | 38 - src/models/index_tts2_5/tokenizer_text.cpp | 691 ----- src/models/index_tts2_5/vocoder.cpp | 69 - .../index_tts2_5_warm_bench_cases.json | 0 tests/index_tts2/index_tts2_warm_bench.cpp | 11 + .../index_tts2_5/index_tts2_5_warm_bench.cpp | 319 --- ...audiocpp_cli_longform_tts_clone_cases.json | 4 +- tools/convert_index_tts2_5.py | 28 +- webui/configs/model_params.json | 10 +- webui/configs/models_catalog.json | 2 +- 54 files changed, 1283 insertions(+), 10488 deletions(-) delete mode 100644 include/engine/models/index_tts2_5/assets.h delete mode 100644 include/engine/models/index_tts2_5/audio_features.h delete mode 100644 include/engine/models/index_tts2_5/gpt.h delete mode 100644 include/engine/models/index_tts2_5/loader.h delete mode 100644 include/engine/models/index_tts2_5/qwen_emotion.h delete mode 100644 include/engine/models/index_tts2_5/request.h delete mode 100644 include/engine/models/index_tts2_5/s2mel.h delete mode 100644 include/engine/models/index_tts2_5/semantic_codec.h delete mode 100644 include/engine/models/index_tts2_5/semantic_encoder.h delete mode 100644 include/engine/models/index_tts2_5/session.h delete mode 100644 include/engine/models/index_tts2_5/style_encoder.h delete mode 100644 include/engine/models/index_tts2_5/tokenizer_text.h delete mode 100644 include/engine/models/index_tts2_5/types.h delete mode 100644 include/engine/models/index_tts2_5/vocoder.h delete mode 100644 model_specs/index_tts2_5.json delete mode 100644 src/models/index_tts2_5/assets.cpp delete mode 100644 src/models/index_tts2_5/audio_features.cpp delete mode 100644 src/models/index_tts2_5/gpt.cpp delete mode 100644 src/models/index_tts2_5/loader.cpp delete mode 100644 src/models/index_tts2_5/qwen_emotion.cpp delete mode 100644 src/models/index_tts2_5/request.cpp delete mode 100644 src/models/index_tts2_5/s2mel.cpp delete mode 100644 src/models/index_tts2_5/semantic_codec.cpp delete mode 100644 src/models/index_tts2_5/semantic_encoder.cpp delete mode 100644 src/models/index_tts2_5/session.cpp delete mode 100644 src/models/index_tts2_5/style_encoder.cpp delete mode 100644 src/models/index_tts2_5/tokenizer_text.cpp delete mode 100644 src/models/index_tts2_5/vocoder.cpp rename tests/{index_tts2_5 => index_tts2}/index_tts2_5_warm_bench_cases.json (100%) delete mode 100644 tests/index_tts2_5/index_tts2_5_warm_bench.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index aa6da691..664cde5c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -812,27 +812,6 @@ audiocpp_add_model(index_tts2 engine::models::index_tts2::make_index_tts2_loader ) -audiocpp_add_model(index_tts2_5 - SOURCES - src/models/index_tts2_5/assets.cpp - src/models/index_tts2_5/audio_features.cpp - src/models/index_tts2_5/gpt.cpp - src/models/index_tts2_5/loader.cpp - src/models/index_tts2_5/qwen_emotion.cpp - src/models/index_tts2_5/request.cpp - src/models/index_tts2_5/s2mel.cpp - src/models/index_tts2_5/semantic_codec.cpp - src/models/index_tts2_5/semantic_encoder.cpp - src/models/index_tts2_5/session.cpp - src/models/index_tts2_5/style_encoder.cpp - src/models/index_tts2_5/tokenizer_text.cpp - src/models/index_tts2_5/vocoder.cpp - INCLUDES - engine/models/index_tts2_5/loader.h - LOADERS - engine::models::index_tts2_5::make_index_tts2_5_loader -) - audiocpp_add_model(nemotron_asr SOURCES src/models/nemotron_asr/assets.cpp @@ -1378,7 +1357,6 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(higgs_audio_tts_warm_bench tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp) add_engine_warmbench(hviske_asr_warm_bench tests/hviske_asr/hviske_asr_warm_bench.cpp) add_engine_warmbench(index_tts2_warm_bench tests/index_tts2/index_tts2_warm_bench.cpp) - add_engine_warmbench(index_tts2_5_warm_bench tests/index_tts2_5/index_tts2_5_warm_bench.cpp) add_engine_warmbench(irodori_tts_warm_bench tests/irodori_tts/irodori_tts_warm_bench.cpp) add_engine_warmbench(marblenet_vad_warm_bench tests/marblenet_vad/marblenet_vad_warm_bench.cpp) add_engine_warmbench(miocodec_warm_bench tests/miocodec/miocodec_warm_bench.cpp) diff --git a/README.md b/README.md index 39a85c51..457e7a38 100644 --- a/README.md +++ b/README.md @@ -102,8 +102,7 @@ Runtime tags: safetensors is the default model loading path. `GGUF 16/Q8/Q4` mea | **vibevoice_asr** | ASR | auto | VibeVoice ASR | GGUF 16/Q8 | | **voxtral_realtime** | ASR | auto | Voxtral-Mini-4B-Realtime-2602 | GGUF 16/Q8/Q4, Stream | | **voxcpm2** | TTS, Clone, Design, Ctrl | ar, da, de, el, en, es, fi, fr, he, hi, id, it, ja, km, ko, lo, ms, my, nl, no, pl, pt, ru, sv, sw, th, tl, tr, vi, zh | VoxCPM2-2B, 48 kHz | GGUF 16/Q8, Stream | -| **index_tts2** | TTS, Clone, Ctrl | zh, en | IndexTTS-2 | GGUF 16/Q8 | -| **index_tts2_5** | TTS, Clone, Ctrl | zh, en, ja, es, ar | IndexTTS-2.5 | GGUF 16/Q8 | +| **index_tts2** | TTS, Clone, Ctrl | zh, en, ja, es, ar | IndexTTS-2, IndexTTS-2.5 (variant) | GGUF 16/Q8 | | **irodori_tts** | TTS, Clone, Design, Ctrl | ja | Irodori-TTS-v4-Small, Irodori-TTS-500M-v3, Irodori-TTS-600M-v3-VoiceDesign | GGUF 16/Q8 | | **moss_tts_nano** | TTS, Clone | auto | MOSS-TTS-Nano-100M | GGUF 16/Q8 | | **moss_tts_local** | TTS, Clone, Ctrl | auto, optional language hint | MOSS-TTS-Local-Transformer-v1.5 | GGUF 16/Q8 | diff --git a/docs/gguf.md b/docs/gguf.md index 62d2e17e..37f8a9ee 100644 --- a/docs/gguf.md +++ b/docs/gguf.md @@ -69,8 +69,7 @@ Status labels: | `htdemucs` | Done | Pass | --- | Pass | Pass (drift) | | `hviske_asr` | Done | Pass | --- | --- | Pass | | `inflect_v2` | Done | Pass | Pass | --- | --- | -| `index_tts2` | Done | Pass | Pass | Pass (drift) | Pass (ASR match, drift) | -| `index_tts2_5` | Done | Pass | Pass | Pass | Pass (ASR match, drift) | +| `index_tts2` | Done (v2 + v2.5 variant) | Pass | Pass | Pass (drift) | Pass (ASR match, drift) | | `irodori_tts` | Done | Pass | --- | Pass | Pass (ASR match, drift) | | `kroko_asr` | Done | Pass | --- | --- | Pass | | `marblenet_vad` | Bundled (tiny model) | Pass | --- | --- | --- | diff --git a/docs/tts.md b/docs/tts.md index 48506430..d0d68999 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -15,7 +15,7 @@ | Higgs Audio v3 TTS | `higgs_audio_tts` | `tts` | [Higgs Audio v3 TTS](#higgs-audio-v3-tts) | | Fish Audio S2 Pro | `fish_audio` | `tts` | [Fish Audio S2 Pro](#fish-audio-s2-pro) | | IndexTTS2 | `index_tts2` | `tts` | [IndexTTS2](#indextts2) | -| IndexTTS2.5 | `index_tts2_5` | `tts` | [IndexTTS2.5](#indextts25) | +| IndexTTS2.5 | `index_tts2` (variant `2.5`) | `tts` | [IndexTTS2.5](#indextts25) | | Irodori-TTS | `irodori_tts` | `tts`, `vdes` | [Irodori-TTS](#irodori-tts) | | GLM-TTS | `glm_tts` | `tts`, `clon` | [GLM-TTS](#glm-tts) | | Inflect Micro v2 | `inflect_v2` | `tts` | [Inflect v2](#inflect-v2) | @@ -551,9 +551,11 @@ audiocpp_cli --task tts --family index_tts2 --model /path/to/IndexTTS-2 --backen IndexTTS2.5 is IndexTeam/bilibili's multilingual zero-shot TTS model (released 2026-07): a 0.8B GPT (autoregressive) + DiT CFM + BigVGAN stack that keeps IndexTTS2's timbre-emotion decoupling and adds Japanese, Spanish, and Arabic on top of Chinese and English. It requires a speaker reference through the framework `--voice-ref` path. Inline `<文字|发音>` pronunciation overrides (pinyin, CMU phonemes, or kana) are supported. Upstream weights live at [IndexTeam/IndexTTS-2.5](https://huggingface.co/IndexTeam/IndexTTS-2.5); the reference implementation is [index-tts/index-tts](https://github.com/index-tts/index-tts) branch `indextts-2.5`. +IndexTTS2.5 is implemented as a variant of the `index_tts2` family rather than a separate family: both variants share the audio features, wav2vec2bert, Qwen emotion, style encoder, BigVGAN vocoder, S2Mel, and the GPT decode/cache code, while the tokenizer (SentencePiece vs multilingual tiktoken), GPT speaker conditioning (conditioning encoder + perceiver vs CAMPPlus `spk_emb_proj` + `lang_embedding`), and the semantic-codec decode path (v2.5 adds a 2x nearest upsample + `up` conv) are selected per variant from the model config `version` field (`"2.5"`). All IndexTTS2 session options (`index_tts2.*`) apply to both variants. + | Field | Value | |---|---| -| Family | `index_tts2_5` | +| Family | `index_tts2` (the `2.5` variant is selected from the model config `version` field; no separate family) | | Model directory | `models/IndexTTS2.5-GGUF` (default GGUF package `index_tts2_5_q8_0`; `index_tts2_5_f16` and `index_tts2_5_orig` also available) | | Task | `tts`, `clon` | | Modes | `offline` | @@ -564,13 +566,13 @@ IndexTTS2.5 is IndexTeam/bilibili's multilingual zero-shot TTS model (released 2 Voice clone: ```bash -audiocpp_cli --task clon --family index_tts2_5 --model /path/to/IndexTTS2.5-GGUF --backend cuda --text "Hello from IndexTTS2.5." --voice-ref /path/to/reference.wav --out out.wav +audiocpp_cli --task clon --family index_tts2 --model /path/to/IndexTTS2.5-GGUF --backend cuda --text "Hello from IndexTTS2.5." --voice-ref /path/to/reference.wav --out out.wav ``` Emotion text: ```bash -audiocpp_cli --task tts --family index_tts2_5 --model /path/to/IndexTTS2.5-GGUF --backend cuda --text "今天的演示会更有情绪。" --voice-ref /path/to/reference.wav --emotion "你吓死我了！你是鬼吗？" --request-option emotion_alpha=0.6 --out out.wav +audiocpp_cli --task tts --family index_tts2 --model /path/to/IndexTTS2.5-GGUF --backend cuda --text "今天的演示会更有情绪。" --voice-ref /path/to/reference.wav --emotion "你吓死我了！你是鬼吗？" --request-option emotion_alpha=0.6 --out out.wav ``` The `lang` request option selects the text language (`auto`, `zh`, `en`, `ja`, `es`, `ar`, or any tokenizer language code). The default `auto` picks `zh` when the text contains Han characters and `en` otherwise, so mixed Japanese/Spanish/Arabic text should set `--request-option lang=ja|es|ar` explicitly. @@ -599,19 +601,19 @@ License: IndexTTS-2.5 weights are distributed under the bilibili Model Use Licen | `--do-sample` | `true`, `false` | `true` | Enable stochastic GPT sampling. | | `--request-option length_penalty=` | float | `0.0` | GPT beam-search length penalty. | | `--request-option num_beams=` | integer | `3` | GPT beam count. | -| `--session-option index_tts2_5.mem_saver=true|false` | bool | `false` | Release staged reference and conditioning graphs after request phases. | -| `--session-option index_tts2_5.weight_type=native|f32|f16|bf16|q8_0` | enum | `native` | Matmul weight storage type. | -| `--session-option index_tts2_5.conv_weight_type=native|f32|f16` | enum | `native` | Convolution weight storage type. | -| `--session-option index_tts2_5.speaker_cache_slots=` | integer slots | `1` | Prepared speaker-reference cache slots; set `0` to disable reuse. | -| `--session-option index_tts2_5.emotion_cache_slots=` | integer slots | `1` | Prepared emotion-reference cache slots; set `0` to disable reuse. | -| `--session-option index_tts2_5.emotion_text_cache_slots=` | integer slots | `1` | Emotion-text weight cache slots; set `0` to disable reuse. | -| `--session-option index_tts2_5.gpt_graph_arena_mb=` | MB | model default | GPT graph arena size. | -| `--session-option index_tts2_5.s2mel_graph_arena_mb=` | MB | model default | S2Mel graph arena size. | -| `--session-option index_tts2_5.reference_graph_arena_mb=` | MB | model default | Reference encoder and codec graph arena size. | -| `--session-option index_tts2_5.emotion_text_prefill_graph_arena_mb=` | MB | model default | Emotion-text prefill graph arena size. | -| `--session-option index_tts2_5.emotion_text_decode_graph_arena_mb=` | MB | model default | Emotion-text cached-step graph arena size. | -| `--session-option index_tts2_5.emotion_text_max_new_tokens=` | tokens | `256` | Maximum generated tokens for emotion-text classification. | -| `--session-option index_tts2_5.weight_context_mb=` | MB | `32` | Shared ggml weight metadata context size. | +| `--session-option index_tts2.mem_saver=true|false` | bool | `false` | Release staged reference and conditioning graphs after request phases. | +| `--session-option index_tts2.weight_type=native|f32|f16|bf16|q8_0` | enum | `native` | Matmul weight storage type. | +| `--session-option index_tts2.conv_weight_type=native|f32|f16` | enum | `native` | Convolution weight storage type. | +| `--session-option index_tts2.speaker_cache_slots=` | integer slots | `1` | Prepared speaker-reference cache slots; set `0` to disable reuse. | +| `--session-option index_tts2.emotion_cache_slots=` | integer slots | `1` | Prepared emotion-reference cache slots; set `0` to disable reuse. | +| `--session-option index_tts2.emotion_text_cache_slots=` | integer slots | `1` | Emotion-text weight cache slots; set `0` to disable reuse. | +| `--session-option index_tts2.gpt_graph_arena_mb=` | MB | model default | GPT graph arena size. | +| `--session-option index_tts2.s2mel_graph_arena_mb=` | MB | model default | S2Mel graph arena size. | +| `--session-option index_tts2.reference_graph_arena_mb=` | MB | model default | Reference encoder and codec graph arena size. | +| `--session-option index_tts2.emotion_text_prefill_graph_arena_mb=` | MB | model default | Emotion-text prefill graph arena size. | +| `--session-option index_tts2.emotion_text_decode_graph_arena_mb=` | MB | model default | Emotion-text cached-step graph arena size. | +| `--session-option index_tts2.emotion_text_max_new_tokens=` | tokens | `256` | Maximum generated tokens for emotion-text classification. | +| `--session-option index_tts2.weight_context_mb=` | MB | `32` | Shared ggml weight metadata context size. | ### Converting From Upstream Weights @@ -626,7 +628,7 @@ python tools/convert_index_tts2_5.py \ Pass `--native-dir /path/to/IndexTTS-2.5-native` to also emit a directly loadable native Safetensors model directory (hardlinked from the staging files), no GGUF conversion required. -The script repackages the checkpoints the loader needs (unwraps the `s2mel.pth`/`codec.pth` container keys, prefixes CAMPPlus tensors with `speaker_encoder.`, strips BigVGAN's `generator.` prefix, wraps the `feat1/feat2.pt` matrices as a single `tensor`) and assembles the sidecar `root/` (config, tiktoken vocabulary, auxiliary model configs) that gets embedded into the GGUF. +The script repackages the checkpoints the loader needs (unwraps the `s2mel.pth`/`codec.pth` container keys, prefixes CAMPPlus tensors with `speaker_encoder.`, strips BigVGAN's `generator.` prefix, wraps the `feat1/feat2.pt` matrices as a single `tensor`) and assembles the sidecar `root/` (config, tiktoken vocabulary, auxiliary model configs) that gets embedded into the GGUF. The staged `config.yaml` has its `version` field normalized to `"2.5"` (the official snapshot ships `version: 2.0`); the engine uses that field to select the IndexTTS2 family variant. ## Irodori-TTS diff --git a/include/engine/models/index_tts2/gpt.h b/include/engine/models/index_tts2/gpt.h index 7df7b34c..93472fe0 100644 --- a/include/engine/models/index_tts2/gpt.h +++ b/include/engine/models/index_tts2/gpt.h @@ -88,6 +88,12 @@ struct IndexTTS2GptWeights { engine::modules::LinearWeights emotion_vec_projection; engine::modules::LinearWeights emotion_layer; std::vector speed_embedding_values; + // v2.5 campplus speaker conditioning (spk_cond_mode="campplus" in the + // official model_v2.py): projects the 192-dim CAMPPlus speaker/style + // embedding into a GPT speaker token. + engine::modules::LinearWeights spk_emb_proj; + // v2.5: row of this table is added to every text embedding during prefill. + engine::core::TensorValue lang_embedding; std::vector gpt_layers; engine::modules::NormWeights gpt_final_norm; engine::modules::NormWeights final_norm; @@ -109,8 +115,14 @@ struct IndexTTS2GptGeneration { struct IndexTTS2GptGenerationRequest { std::vector text_tokens; + // v2 speaker conditioning: wav2vec2bert semantic features of the reference. std::vector speaker_semantic; int64_t speaker_frames = 0; + // v2.5 speaker conditioning: 192-dim CAMPPlus speaker embedding, projected + // by spk_emb_proj inside the prefill graph. + std::vector speaker_style; + // v2.5: row of the GPT lang_embedding table added to every text embedding. + int32_t lang_id = 0; std::vector emotion_semantic; int64_t emotion_frames = 0; std::vector emotion_vector; @@ -125,6 +137,12 @@ struct IndexTTS2GptGenerationRequest { uint32_t seed = 0; }; +// Mirrors the valid_mask filtering in the official v2.5 prepare_gpt_inputs: +// any start/stop text tokens in the segment (including the trailing pad +// appended by the tokenizer) are dropped before the start/stop pair is +// re-added around it. +std::vector align_index_tts2_gpt_text_tokens(const std::vector & text_tokens); + std::shared_ptr load_index_tts2_gpt_weights( const IndexTTS2Assets & assets, ggml_backend_t backend, diff --git a/include/engine/models/index_tts2/request.h b/include/engine/models/index_tts2/request.h index a63d8ca7..8f28737c 100644 --- a/include/engine/models/index_tts2/request.h +++ b/include/engine/models/index_tts2/request.h @@ -5,6 +5,10 @@ namespace engine::models::index_tts2 { +// Normalizes the "lang" request option (v2.5): trims, lowercases, and maps +// "auto" to an empty string (tokenizer-side language inference). +std::string normalize_index_tts2_lang(const std::string & value); + IndexTTS2Request parse_index_tts2_request(const runtime::TaskRequest & request); } // namespace engine::models::index_tts2 diff --git a/include/engine/models/index_tts2/semantic_codec.h b/include/engine/models/index_tts2/semantic_codec.h index 0c3e4c84..dcece494 100644 --- a/include/engine/models/index_tts2/semantic_codec.h +++ b/include/engine/models/index_tts2/semantic_codec.h @@ -41,6 +41,8 @@ struct IndexTTS2SemanticCodecWeights { engine::modules::Conv1dWeights quantizer_out; IndexTTS2VocosBackboneWeights decoder_backbone; engine::modules::LinearWeights decoder_projection; + // v2.5 only: conv applied after the 2x nearest upsample in the decode path. + engine::modules::Conv1dWeights up; }; struct IndexTTS2SemanticCodecOutput { diff --git a/include/engine/models/index_tts2/session.h b/include/engine/models/index_tts2/session.h index 8a2a8db6..5f99a22b 100644 --- a/include/engine/models/index_tts2/session.h +++ b/include/engine/models/index_tts2/session.h @@ -73,6 +73,7 @@ class IndexTTS2Session final const EmotionState & emotion); runtime::AudioBuffer synthesize_segment( const std::vector & text_tokens, + int32_t lang_id, const SpeakerState & speaker, const EmotionState & emotion, const std::vector & emotion_vector, diff --git a/include/engine/models/index_tts2/tokenizer_text.h b/include/engine/models/index_tts2/tokenizer_text.h index a20d144e..4999fcfd 100644 --- a/include/engine/models/index_tts2/tokenizer_text.h +++ b/include/engine/models/index_tts2/tokenizer_text.h @@ -9,39 +9,92 @@ #include #include +namespace llama_tokenizer_vendor { +struct BpeVocabulary; +} // namespace llama_tokenizer_vendor + namespace engine::models::index_tts2 { struct IndexTTS2TextEncoding { + // v2.5: resolved language code; empty for v2. + std::string lang; std::string normalized_text; + // v2 SentencePiece pieces; empty for v2.5. std::vector pieces; std::vector token_ids; + // v2: piece strings per segment; v2.5: processed text per segment. std::vector> segments; std::vector> segment_token_ids; }; +// Variant-aware IndexTTS2 text tokenizer. v2 keeps the SentencePiece bpe.model +// behavior; v2.5 uses the multilingual tiktoken/BPE vocabulary with language +// special tokens and language-id handling. The variant is selected from the +// model config version, not from probing tokenizer files. class IndexTTS2TextTokenizer { public: explicit IndexTTS2TextTokenizer(std::shared_ptr assets); + IndexTTS2Variant variant() const noexcept { + return variant_; + } + std::string normalize_english(const std::string & text) const; std::string normalize_chinese(const std::string & text) const; - std::string normalize_text(const std::string & text) const; + + // v2: SentencePiece encode of the normalized text. + // v2.5: raw tiktoken encode with allowed_special="all"; does not apply any + // text normalization. Special tokens present in the text are recognized + // directly. std::vector encode(const std::string & text) const; + + // v2 only: normalize then tokenize helpers kept for parity/debug. + std::string normalize_text(const std::string & text) const; std::vector tokenize_to_pieces(const std::string & text) const; + + // v2.5 only: returns the id of an exact token text (e.g. "<|zh|>"), or -1 + // when unknown. + int32_t special_token_id(const std::string & token_text) const; + + // v2.5 only: maps a language code to the GPT lang_embedding row, following + // the LANGUAGES order of indextts/utils/tokenizer.py (en=0, zh=1, ...). + // Unknown codes map to "common". + static int32_t lang_to_id(const std::string & lang); + + // v2: normalize -> SentencePiece encode -> segment pieces by token budget. + // v2.5: normalize -> case rules -> pronunciation annotations -> + // special-token name uppercasing -> segment by token budget. Each segment + // is encoded as encode("<|{lang}|> " + segment) plus a trailing pad token + // id 1. When lang is empty, it is inferred (Han -> zh, else en). IndexTTS2TextEncoding encode_for_inference( const std::string & text, - int max_text_tokens_per_segment) const; + int max_text_tokens_per_segment, + const std::string & lang = "") const; private: + // v2 SentencePiece path. + IndexTTS2TextEncoding encode_for_inference_v2( + const std::string & text, + int max_text_tokens_per_segment) const; int32_t piece_to_id(const std::string & piece) const; std::string id_to_piece(int32_t id) const; std::vector> split_segments( const std::vector & pieces, int max_text_tokens_per_segment) const; + // v2.5 tiktoken path. + IndexTTS2TextEncoding encode_for_inference_v2_5( + const std::string & text, + int max_text_tokens_per_segment, + const std::string & lang) const; + std::shared_ptr assets_; + IndexTTS2Variant variant_ = IndexTTS2Variant::kV2; + // v2 SentencePiece model. std::vector pieces_; std::unordered_map piece_to_id_; + // v2.5 tiktoken vocabulary. + std::shared_ptr vocab_; }; } // namespace engine::models::index_tts2 diff --git a/include/engine/models/index_tts2/types.h b/include/engine/models/index_tts2/types.h index a89a1327..0171e776 100644 --- a/include/engine/models/index_tts2/types.h +++ b/include/engine/models/index_tts2/types.h @@ -5,10 +5,28 @@ #include #include #include +#include #include namespace engine::models::index_tts2 { +// Model variant, selected from the model config "version" field ("2.5" -> +// kV2_5, anything else -> kV2). Variant branches must be driven by this value, +// never by probing weight tensor names in hot paths. +enum class IndexTTS2Variant { + kV2, + kV2_5, +}; + +inline IndexTTS2Variant index_tts2_variant_from_version(std::string_view version) { + return version == "2.5" ? IndexTTS2Variant::kV2_5 : IndexTTS2Variant::kV2; +} + +// Rows in the v2.5 GPT lang_embedding table. indextts/utils/tokenizer.py +// defines 106 language codes (including "common"); the checkpoint table has +// one extra unused row. +constexpr int64_t kIndexTTS2LangEmbeddingRows = 107; + struct IndexTTS2GptConfig { int64_t model_dim = 1280; int64_t max_mel_tokens = 1815; @@ -132,6 +150,9 @@ struct IndexTTS2Request { std::string text; std::optional speaker_audio = std::nullopt; std::optional emotion_audio = std::nullopt; + // Text language hint (v2.5 only); empty means auto (zh when the text + // contains Han characters, otherwise en). + std::string lang; float emotion_alpha = 1.0F; std::optional> emotion_vector = std::nullopt; bool use_emotion_text = false; diff --git a/include/engine/models/index_tts2_5/assets.h b/include/engine/models/index_tts2_5/assets.h deleted file mode 100644 index 6b0fc014..00000000 --- a/include/engine/models/index_tts2_5/assets.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include "engine/framework/assets/resource_bundle.h" -#include "engine/framework/assets/tensor_source.h" -#include "engine/models/index_tts2_5/types.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25Assets { - assets::ResourceBundle resources; - IndexTTS25Config config; - std::shared_ptr gpt_weights; - std::shared_ptr s2mel_weights; - std::shared_ptr speaker_matrix; - std::shared_ptr emotion_matrix; - std::shared_ptr wav2vec2bert_stats; - std::shared_ptr wav2vec2bert_weights; - std::shared_ptr semantic_codec_weights; - std::shared_ptr campplus_weights; - std::shared_ptr bigvgan_weights; - std::shared_ptr qwen_emotion_weights; -}; - -std::shared_ptr load_index_tts2_5_assets(const std::filesystem::path & model_path); - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/audio_features.h b/include/engine/models/index_tts2_5/audio_features.h deleted file mode 100644 index 3022fdb3..00000000 --- a/include/engine/models/index_tts2_5/audio_features.h +++ /dev/null @@ -1,54 +0,0 @@ -#pragma once - -#include "engine/models/index_tts2_5/types.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25MelOutput { - std::vector values; - int64_t channels = 0; - int64_t frames = 0; -}; - -struct IndexTTS25FbankOutput { - std::vector values; - int64_t frames = 0; - int64_t dims = 0; -}; - -struct IndexTTS25SemanticFeatureOutput { - std::vector values; - std::vector attention_mask; - int64_t frames = 0; - int64_t dims = 0; -}; - -struct IndexTTS25PreparedReferenceAudio { - std::vector waveform_16k; - std::vector waveform_22k; - IndexTTS25MelOutput mel; - IndexTTS25FbankOutput campplus_fbank; - IndexTTS25SemanticFeatureOutput semantic_features; -}; - -IndexTTS25PreparedReferenceAudio prepare_index_tts2_5_reference_audio( - const std::vector & samples, - int sample_rate, - int channels, - const IndexTTS25S2MelConfig & mel_config, - size_t threads, - bool speaker_load_semantic = true); - -IndexTTS25MelOutput compute_index_tts2_5_mel_spectrogram( - const std::vector & waveform, - const IndexTTS25S2MelConfig & config, - size_t threads); - -IndexTTS25FbankOutput compute_index_tts2_5_campplus_fbank_16k(const std::vector & waveform_16k); - -IndexTTS25SemanticFeatureOutput compute_index_tts2_5_semantic_features_16k(const std::vector & waveform_16k); - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/gpt.h b/include/engine/models/index_tts2_5/gpt.h deleted file mode 100644 index a8f9b93e..00000000 --- a/include/engine/models/index_tts2_5/gpt.h +++ /dev/null @@ -1,186 +0,0 @@ -#pragma once - -#include "engine/framework/core/backend_weight_store.h" -#include "engine/framework/core/execution_context.h" -#include "engine/framework/modules/attention/types.h" -#include "engine/framework/modules/conv_modules.h" -#include "engine/framework/modules/linear_module.h" -#include "engine/framework/modules/norm_modules.h" -#include "engine/framework/runtime/kv_cache.h" -#include "engine/framework/modules/streaming_conv_modules.h" -#include "engine/models/index_tts2_5/assets.h" - -#include "ggml-backend.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25GptConditionSubsamplingWeights { - engine::modules::Conv2dWeights conv; - engine::modules::LinearWeights out; - engine::core::TensorValue pos_enc; -}; - -struct IndexTTS25GptConditionLayerWeights { - engine::modules::NormWeights norm_ff; - engine::modules::NormWeights norm_mha; - engine::modules::NormWeights norm_conv; - engine::modules::NormWeights norm_final; - engine::modules::LinearWeights feed_forward_in; - engine::modules::LinearWeights feed_forward_out; - engine::modules::RelativeAttentionWeights self_attn; - engine::modules::Conv1dWeights conv_pointwise_in; - engine::modules::DepthwiseConv1dWeights conv_depthwise; - engine::modules::NormWeights conv_norm; - engine::modules::Conv1dWeights conv_pointwise_out; -}; - -struct IndexTTS25GptConditionEncoderWeights { - IndexTTS25GptConditionSubsamplingWeights subsampling; - std::vector layers; - engine::modules::NormWeights after_norm; -}; - -struct IndexTTS25PerceiverAttentionWeights { - engine::modules::LinearWeights q; - engine::modules::LinearWeights kv; - engine::modules::LinearWeights out; -}; - -struct IndexTTS25PerceiverFeedForwardWeights { - engine::modules::LinearWeights in; - engine::modules::LinearWeights out; -}; - -struct IndexTTS25PerceiverLayerWeights { - IndexTTS25PerceiverAttentionWeights attention; - IndexTTS25PerceiverFeedForwardWeights feed_forward; -}; - -struct IndexTTS25PerceiverWeights { - engine::core::TensorValue latents; - engine::modules::LinearWeights project_context; - std::vector layers; - engine::core::TensorValue norm_gamma; -}; - -struct IndexTTS25Gpt2LayerWeights { - engine::modules::NormWeights attn_norm; - engine::modules::LinearWeights qkv; - engine::modules::LinearWeights attn_out; - engine::modules::NormWeights mlp_norm; - engine::modules::LinearWeights mlp_in; - engine::modules::LinearWeights mlp_out; -}; - -struct IndexTTS25GptWeights { - std::shared_ptr store; - IndexTTS25GptConditionEncoderWeights emotion_conditioner; - IndexTTS25PerceiverWeights emotion_perceiver; - engine::modules::LinearWeights spk_emb_proj; - engine::core::TensorValue lang_embedding; - engine::core::TensorValue text_embedding; - engine::core::TensorValue mel_embedding; - engine::core::TensorValue text_pos_embedding; - engine::core::TensorValue mel_pos_embedding; - engine::modules::LinearWeights emotion_vec_projection; - engine::modules::LinearWeights emotion_layer; - std::vector gpt_layers; - engine::modules::NormWeights gpt_final_norm; - engine::modules::NormWeights final_norm; - engine::modules::LinearWeights mel_head; - engine::modules::LinearWeights text_head; -}; - -struct IndexTTS25GptLatent { - std::vector values; - int64_t frames = 0; - int64_t dims = 0; -}; - -struct IndexTTS25GptGeneration { - std::vector codes; - uint64_t rng_offset_blocks = 0; -}; - -struct IndexTTS25GptGenerationRequest { - std::vector text_tokens; - // 192-dim CAMPPlus speaker embedding, projected by spk_emb_proj inside the - // prefill graph (spk_cond_mode="campplus" in the official model_v2.py). - std::vector speaker_style; - // Row of the GPT lang_embedding table added to every text embedding. - int32_t lang_id = 0; - std::vector emotion_semantic; - int64_t emotion_frames = 0; - std::vector emotion_vector; - float top_p = 0.8F; - int top_k = 30; - float temperature = 0.8F; - float repetition_penalty = 10.0F; - bool do_sample = true; - float length_penalty = 0.0F; - int num_beams = 3; - int max_mel_tokens = 1500; - uint32_t seed = 0; -}; - -// Mirrors the valid_mask filtering in the official prepare_gpt_inputs: any -// start/stop text tokens in the segment (including the trailing pad appended by -// the tokenizer) are dropped before the start/stop pair is re-added around it. -std::vector align_index_tts2_5_gpt_text_tokens(const std::vector & text_tokens); - -std::shared_ptr load_index_tts2_5_gpt_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes); - -class IndexTTS25GptRuntime { -public: - IndexTTS25GptRuntime( - std::shared_ptr assets, - engine::core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type); - ~IndexTTS25GptRuntime(); - - IndexTTS25GptRuntime(const IndexTTS25GptRuntime &) = delete; - IndexTTS25GptRuntime & operator=(const IndexTTS25GptRuntime &) = delete; - - void prepare_emotion_conditioning(int64_t frames); - void prepare_generation(int64_t text_tokens, int64_t max_mel_tokens, int64_t num_beams); - IndexTTS25GptLatent emotion_conditioning(const std::vector & semantic_btc, int64_t frames); - std::vector project_emotion_vector(const IndexTTS25GptLatent & emotion_conditioning); - std::vector merge_emotion_vector( - const std::vector & speaker_semantic, - int64_t speaker_frames, - const std::vector & emotion_semantic, - int64_t emotion_frames, - float alpha); - IndexTTS25GptGeneration generate_speech(const IndexTTS25GptGenerationRequest & request); - void release_conditioning_graphs(); - void release_generation_graphs(); - -private: - class ConditioningGraph; - class EmotionVectorGraph; - class PrefillGraph; - class DecodeGraph; - - std::shared_ptr assets_; - engine::core::ExecutionContext * execution_ = nullptr; - size_t graph_arena_bytes_ = 0; - std::shared_ptr weights_; - std::unique_ptr emotion_conditioning_graph_; - std::unique_ptr emotion_vector_graph_; - std::unique_ptr prefill_graph_; - std::unique_ptr decode_graph_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/loader.h b/include/engine/models/index_tts2_5/loader.h deleted file mode 100644 index e063f716..00000000 --- a/include/engine/models/index_tts2_5/loader.h +++ /dev/null @@ -1,33 +0,0 @@ -#pragma once - -#include "engine/framework/runtime/model.h" -#include "engine/models/index_tts2_5/assets.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -class IndexTTS25LoadedModel final : public runtime::ILoadedVoiceModel { -public: - IndexTTS25LoadedModel( - 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_index_tts2_5_model(const std::filesystem::path & model_path); -std::shared_ptr make_index_tts2_5_loader(); - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/qwen_emotion.h b/include/engine/models/index_tts2_5/qwen_emotion.h deleted file mode 100644 index 211a6cbb..00000000 --- a/include/engine/models/index_tts2_5/qwen_emotion.h +++ /dev/null @@ -1,82 +0,0 @@ -#pragma once - -#include "engine/framework/core/backend_weight_store.h" -#include "engine/framework/core/execution_context.h" -#include "engine/framework/modules/transformers/qwen_decoder.h" -#include "engine/framework/modules/norm_modules.h" -#include "engine/framework/tokenizers/llama_bpe.h" -#include "engine/models/index_tts2_5/assets.h" - -#include "ggml-backend.h" - -#include -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25QwenEmotionWeights { - std::shared_ptr store; - engine::core::TensorValue token_embedding; - engine::modules::QwenDecoderStackWeights decoder; - engine::modules::NormWeights final_norm; -}; - -struct IndexTTS25EmotionVector { - std::vector values; -}; - -std::shared_ptr load_index_tts2_5_qwen_emotion_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType storage_type, - size_t weight_context_bytes); - -class IndexTTS25QwenEmotionTokenizer { -public: - explicit IndexTTS25QwenEmotionTokenizer(std::shared_ptr assets); - - std::vector encode_chat_prompt(const std::string & text) const; - std::string decode(const std::vector & token_ids, bool skip_special_tokens) const; - int32_t eos_token_id() const noexcept; - int32_t think_end_token_id() const noexcept; - -private: - std::shared_ptr tokenizer_; - int32_t eos_token_id_ = 151643; - int32_t think_end_token_id_ = 151668; -}; - -class IndexTTS25QwenEmotionRuntime { -public: - IndexTTS25QwenEmotionRuntime( - std::shared_ptr assets, - engine::core::ExecutionContext & execution, - size_t prefill_graph_arena_bytes, - size_t decode_graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType storage_type); - ~IndexTTS25QwenEmotionRuntime(); - - IndexTTS25QwenEmotionRuntime(const IndexTTS25QwenEmotionRuntime &) = delete; - IndexTTS25QwenEmotionRuntime & operator=(const IndexTTS25QwenEmotionRuntime &) = delete; - - IndexTTS25EmotionVector infer(const std::string & text, int64_t max_new_tokens = 256); - void release_graphs(); - -private: - class PrefillGraph; - class DecodeGraph; - - std::shared_ptr assets_; - engine::core::ExecutionContext * execution_ = nullptr; - size_t prefill_graph_arena_bytes_ = 0; - size_t decode_graph_arena_bytes_ = 0; - std::shared_ptr weights_; - IndexTTS25QwenEmotionTokenizer tokenizer_; - std::unique_ptr prefill_graph_; - std::unique_ptr decode_graph_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/request.h b/include/engine/models/index_tts2_5/request.h deleted file mode 100644 index 474e30b0..00000000 --- a/include/engine/models/index_tts2_5/request.h +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "engine/framework/runtime/session.h" -#include "engine/models/index_tts2_5/types.h" - -namespace engine::models::index_tts2_5 { - -// Normalizes the "lang" request option: trims, lowercases, and maps "auto" to -// an empty string (tokenizer-side language inference). -std::string normalize_index_tts2_5_lang(const std::string & value); - -IndexTTS25Request parse_index_tts2_5_request(const runtime::TaskRequest & request); - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/s2mel.h b/include/engine/models/index_tts2_5/s2mel.h deleted file mode 100644 index 6f015502..00000000 --- a/include/engine/models/index_tts2_5/s2mel.h +++ /dev/null @@ -1,151 +0,0 @@ -#pragma once - -#include "engine/framework/core/backend_weight_store.h" -#include "engine/framework/core/execution_context.h" -#include "engine/framework/modules/conv_modules.h" -#include "engine/framework/modules/linear_module.h" -#include "engine/framework/modules/norm_modules.h" -#include "engine/framework/modules/streaming_conv_modules.h" -#include "engine/models/index_tts2_5/assets.h" - -#include "ggml-backend.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25LengthRegulatorWeights { - engine::modules::LinearWeights content_projection; - std::vector convs; - std::vector norms; - engine::modules::Conv1dWeights output; -}; - -struct IndexTTS25S2MelGptLayerWeights { - engine::modules::LinearWeights linear0; - engine::modules::LinearWeights linear1; - engine::modules::LinearWeights linear2; -}; - -struct IndexTTS25AdaLayerNormWeights { - engine::core::TensorValue norm_weight; - engine::modules::LinearWeights project; -}; - -struct IndexTTS25DitLayerWeights { - IndexTTS25AdaLayerNormWeights attention_norm; - engine::modules::LinearWeights qkv; - engine::modules::LinearWeights attention_out; - IndexTTS25AdaLayerNormWeights ffn_norm; - engine::modules::LinearWeights ffn_w1; - engine::modules::LinearWeights ffn_w2; - engine::modules::LinearWeights ffn_w3; - engine::modules::LinearWeights skip_in; -}; - -struct IndexTTS25WaveNetLayerWeights { - engine::modules::Conv1dWeights in_layer; - engine::modules::Conv1dWeights res_skip_layer; -}; - -struct IndexTTS25S2MelCfmWeights { - engine::modules::LinearWeights x_embedder; - engine::modules::LinearWeights cond_projection; - engine::modules::LinearWeights cond_x_merge; - engine::modules::LinearWeights skip_linear; - engine::modules::LinearWeights time_mlp0; - engine::modules::LinearWeights time_mlp2; - engine::core::TensorValue time_freqs; - engine::modules::LinearWeights time2_mlp0; - engine::modules::LinearWeights time2_mlp2; - engine::core::TensorValue time2_freqs; - std::vector dit_layers; - IndexTTS25AdaLayerNormWeights dit_norm; - engine::modules::LinearWeights conv1; - engine::modules::LinearWeights res_projection; - engine::modules::Conv1dWeights wavenet_cond; - std::vector wavenet_layers; - engine::modules::LinearWeights final_modulation; - engine::modules::LinearWeights final_linear; - engine::modules::Conv1dWeights conv2; -}; - -struct IndexTTS25S2MelWeights { - std::shared_ptr store; - IndexTTS25S2MelGptLayerWeights gpt_layer; - IndexTTS25LengthRegulatorWeights length_regulator; - IndexTTS25S2MelCfmWeights cfm; -}; - -struct IndexTTS25S2MelSequence { - std::vector values; - int64_t frames = 0; - int64_t dims = 0; -}; - -struct IndexTTS25S2MelMel { - std::vector values; - int64_t frames = 0; - int64_t channels = 80; -}; - -std::shared_ptr load_index_tts2_5_s2mel_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes); - -class IndexTTS25S2MelRuntime { -public: - IndexTTS25S2MelRuntime( - std::shared_ptr assets, - engine::core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type); - ~IndexTTS25S2MelRuntime(); - - IndexTTS25S2MelRuntime(const IndexTTS25S2MelRuntime &) = delete; - IndexTTS25S2MelRuntime & operator=(const IndexTTS25S2MelRuntime &) = delete; - - void prepare_gpt_layer(int64_t frames); - void prepare_length_regulator(int64_t input_frames, int64_t output_frames); - void prepare_cfm(int64_t total_frames, bool use_cfg); - void release_pre_cfm_graphs(); - void release_cfm_graph(); - - IndexTTS25S2MelSequence project_gpt_latent(const std::vector & latent, int64_t frames); - IndexTTS25S2MelSequence regulate_length( - const std::vector & content, - int64_t input_frames, - int64_t output_frames); - IndexTTS25S2MelMel infer_mel( - const std::vector & condition, - int64_t total_frames, - const std::vector & reference_mel, - int64_t reference_frames, - const std::vector & style, - int64_t diffusion_steps, - float cfg_rate, - uint32_t seed, - uint64_t rng_offset_blocks); - -private: - class GptLayerGraph; - class LengthRegulatorGraph; - class CfmGraph; - - std::shared_ptr assets_; - engine::core::ExecutionContext * execution_ = nullptr; - size_t graph_arena_bytes_ = 0; - std::shared_ptr weights_; - std::unique_ptr gpt_layer_graph_; - std::unique_ptr length_regulator_graph_; - std::unique_ptr cfm_graph_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/semantic_codec.h b/include/engine/models/index_tts2_5/semantic_codec.h deleted file mode 100644 index 2762d7f7..00000000 --- a/include/engine/models/index_tts2_5/semantic_codec.h +++ /dev/null @@ -1,94 +0,0 @@ -#pragma once - -#include "engine/framework/core/backend_weight_store.h" -#include "engine/framework/core/execution_context.h" -#include "engine/framework/modules/conv_modules.h" -#include "engine/framework/modules/linear_module.h" -#include "engine/framework/modules/norm_modules.h" -#include "engine/framework/modules/streaming_conv_modules.h" -#include "engine/models/index_tts2_5/assets.h" -#include "engine/models/index_tts2_5/semantic_encoder.h" - -#include "ggml-backend.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25VocosConvNeXtBlockWeights { - engine::modules::DepthwiseConv1dWeights depthwise; - engine::modules::NormWeights norm; - engine::modules::LinearWeights pointwise_in; - engine::modules::LinearWeights pointwise_out; - engine::core::TensorValue gamma; -}; - -struct IndexTTS25VocosBackboneWeights { - engine::modules::Conv1dWeights embed; - engine::modules::NormWeights norm; - std::vector blocks; - engine::modules::NormWeights final_norm; -}; - -struct IndexTTS25SemanticCodecWeights { - std::shared_ptr store; - IndexTTS25VocosBackboneWeights encoder_backbone; - engine::modules::LinearWeights encoder_projection; - engine::modules::Conv1dWeights quantizer_in; - engine::core::TensorValue codebook; - engine::core::TensorValue normalized_codebook; - engine::modules::Conv1dWeights quantizer_out; - IndexTTS25VocosBackboneWeights decoder_backbone; - engine::modules::LinearWeights decoder_projection; - engine::modules::Conv1dWeights up; -}; - -struct IndexTTS25SemanticCodecOutput { - std::vector codes; - std::vector embedding_channel_first; - int64_t frames = 0; - int64_t dims = 0; -}; - -std::shared_ptr load_index_tts2_5_semantic_codec_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes); - -class IndexTTS25SemanticCodecRuntime { -public: - IndexTTS25SemanticCodecRuntime( - std::shared_ptr assets, - engine::core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type); - ~IndexTTS25SemanticCodecRuntime(); - - IndexTTS25SemanticCodecRuntime(const IndexTTS25SemanticCodecRuntime &) = delete; - IndexTTS25SemanticCodecRuntime & operator=(const IndexTTS25SemanticCodecRuntime &) = delete; - - void prepare_quantize(int64_t frames); - void prepare_codes(int64_t frames); - IndexTTS25SemanticCodecOutput quantize(const IndexTTS25SemanticEmbedding & semantic); - IndexTTS25SemanticCodecOutput codes_to_embedding(const std::vector & codes, int64_t frames); - void release_graphs(); - -private: - class QuantizeGraph; - class CodesGraph; - - std::shared_ptr assets_; - engine::core::ExecutionContext * execution_ = nullptr; - size_t graph_arena_bytes_ = 0; - std::shared_ptr weights_; - std::unique_ptr quantize_graph_; - std::unique_ptr codes_graph_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/semantic_encoder.h b/include/engine/models/index_tts2_5/semantic_encoder.h deleted file mode 100644 index 5ffb21d9..00000000 --- a/include/engine/models/index_tts2_5/semantic_encoder.h +++ /dev/null @@ -1,100 +0,0 @@ -#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/modules/conv_modules.h" -#include "engine/framework/modules/linear_module.h" -#include "engine/framework/modules/norm_modules.h" -#include "engine/framework/modules/streaming_conv_modules.h" -#include "engine/models/index_tts2_5/assets.h" -#include "engine/models/index_tts2_5/audio_features.h" - -#include "ggml-backend.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25Wav2Vec2BertAttentionWeights { - engine::modules::LinearWeights q; - engine::modules::LinearWeights k; - engine::modules::LinearWeights v; - engine::modules::LinearWeights out; - engine::core::TensorValue distance_embedding; -}; - -struct IndexTTS25Wav2Vec2BertConvWeights { - engine::modules::NormWeights layer_norm; - engine::modules::Conv1dWeights pointwise_in; - engine::modules::DepthwiseConv1dWeights depthwise; - engine::modules::NormWeights depthwise_layer_norm; - engine::modules::Conv1dWeights pointwise_out; -}; - -struct IndexTTS25Wav2Vec2BertLayerWeights { - engine::modules::NormWeights ffn1_norm; - engine::modules::LinearWeights ffn1_in; - engine::modules::LinearWeights ffn1_out; - engine::modules::NormWeights self_attn_norm; - IndexTTS25Wav2Vec2BertAttentionWeights self_attn; - IndexTTS25Wav2Vec2BertConvWeights conv; - engine::modules::NormWeights ffn2_norm; - engine::modules::LinearWeights ffn2_in; - engine::modules::LinearWeights ffn2_out; - engine::modules::NormWeights final_norm; -}; - -struct IndexTTS25Wav2Vec2BertWeights { - std::shared_ptr store; - engine::modules::NormWeights feature_norm; - engine::modules::LinearWeights feature_projection; - std::vector layers; - engine::core::TensorValue semantic_mean; - engine::core::TensorValue semantic_std; -}; - -struct IndexTTS25SemanticEmbedding { - std::vector values; - int64_t frames = 0; - int64_t dims = 0; -}; - -std::shared_ptr load_index_tts2_5_wav2vec2bert_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes); - -class IndexTTS25Wav2Vec2BertRuntime { -public: - IndexTTS25Wav2Vec2BertRuntime( - std::shared_ptr assets, - engine::core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type); - ~IndexTTS25Wav2Vec2BertRuntime(); - - IndexTTS25Wav2Vec2BertRuntime(const IndexTTS25Wav2Vec2BertRuntime &) = delete; - IndexTTS25Wav2Vec2BertRuntime & operator=(const IndexTTS25Wav2Vec2BertRuntime &) = delete; - - void prepare(int64_t frames); - IndexTTS25SemanticEmbedding encode(const IndexTTS25SemanticFeatureOutput & features); - void release_graph(); - -private: - class Graph; - - std::shared_ptr assets_; - engine::core::ExecutionContext * execution_ = nullptr; - size_t graph_arena_bytes_ = 0; - std::shared_ptr weights_; - std::unique_ptr graph_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/session.h b/include/engine/models/index_tts2_5/session.h deleted file mode 100644 index 65ab86ca..00000000 --- a/include/engine/models/index_tts2_5/session.h +++ /dev/null @@ -1,122 +0,0 @@ -#pragma once - -#include "engine/framework/runtime/cache_slots.h" -#include "engine/framework/runtime/session_base.h" -#include "engine/models/index_tts2_5/assets.h" -#include "engine/models/index_tts2_5/audio_features.h" -#include "engine/models/index_tts2_5/gpt.h" -#include "engine/models/index_tts2_5/qwen_emotion.h" -#include "engine/models/index_tts2_5/request.h" -#include "engine/models/index_tts2_5/s2mel.h" -#include "engine/models/index_tts2_5/semantic_codec.h" -#include "engine/models/index_tts2_5/semantic_encoder.h" -#include "engine/models/index_tts2_5/style_encoder.h" -#include "engine/models/index_tts2_5/tokenizer_text.h" -#include "engine/models/index_tts2_5/vocoder.h" - -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25AudioIdentity { - int sample_rate = 0; - int channels = 0; - uint64_t sample_count = 0; - uint64_t sample_hash = 0; -}; - -class IndexTTS25Session final - : public runtime::RuntimeSessionBase - , public runtime::IOfflineVoiceTaskSession { -public: - IndexTTS25Session( - runtime::TaskSpec task, - runtime::SessionOptions options, - std::shared_ptr assets); - - 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: - struct SpeakerState { - IndexTTS25AudioIdentity identity; - IndexTTS25SemanticEmbedding semantic; - IndexTTS25MelOutput reference_mel; - IndexTTS25StyleEmbedding style; - IndexTTS25S2MelSequence prompt_condition; - }; - - struct EmotionState { - IndexTTS25AudioIdentity identity; - IndexTTS25SemanticEmbedding semantic; - }; - - struct AudioIdentityEqual { - bool operator()( - const IndexTTS25AudioIdentity & lhs, - const IndexTTS25AudioIdentity & rhs) const; - }; - - const SpeakerState & resolve_speaker_state(const runtime::AudioBuffer & audio); - const EmotionState & resolve_emotion_state(const runtime::AudioBuffer & audio); - std::vector resolve_emotion_vector( - const IndexTTS25Request & request, - const SpeakerState & speaker, - const EmotionState & emotion); - runtime::AudioBuffer synthesize_segment( - const std::vector & text_tokens, - int32_t lang_id, - size_t segment_index, - const std::string & dump_dir, - const SpeakerState & speaker, - const EmotionState & emotion, - const std::vector & emotion_vector, - const IndexTTS25GenerationOptions & options, - uint32_t segment_seed); - - std::vector explicit_emotion_matrix_vector( - const std::vector & emotion_weights, - const IndexTTS25StyleEmbedding & style, - bool use_random, - uint32_t seed) const; - - runtime::TaskSpec task_; - std::shared_ptr assets_; - size_t gpt_graph_arena_bytes_ = 2048ull * 1024ull * 1024ull; - size_t s2mel_graph_arena_bytes_ = 2048ull * 1024ull * 1024ull; - size_t reference_graph_arena_bytes_ = 512ull * 1024ull * 1024ull; - size_t emotion_text_prefill_graph_arena_bytes_ = 2048ull * 1024ull * 1024ull; - size_t emotion_text_decode_graph_arena_bytes_ = 512ull * 1024ull * 1024ull; - size_t weight_context_bytes_ = 32ull * 1024ull * 1024ull; - int64_t emotion_text_max_new_tokens_ = 256; - engine::assets::TensorStorageType matmul_weight_storage_type_ = engine::assets::TensorStorageType::Native; - engine::assets::TensorStorageType conv_weight_storage_type_ = engine::assets::TensorStorageType::Native; - bool mem_saver_ = false; - - IndexTTS25TextTokenizer tokenizer_; - std::unique_ptr semantic_encoder_; - std::unique_ptr semantic_codec_; - std::unique_ptr style_encoder_; - std::unique_ptr gpt_; - std::unique_ptr s2mel_; - std::unique_ptr vocoder_; - std::unique_ptr qwen_emotion_; - - std::vector speaker_matrix_; - std::vector emotion_matrix_; - runtime::CacheSlots speaker_cache_; - runtime::CacheSlots emotion_cache_; - runtime::CacheSlots> emotion_text_weights_cache_; - std::optional uncached_speaker_state_; - std::optional uncached_emotion_state_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/style_encoder.h b/include/engine/models/index_tts2_5/style_encoder.h deleted file mode 100644 index 9ed6bc0f..00000000 --- a/include/engine/models/index_tts2_5/style_encoder.h +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include "engine/framework/modules/speech_encoders/campplus_encoder.h" -#include "engine/models/index_tts2_5/assets.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25StyleEmbedding { - std::vector values; - int64_t dims = 0; -}; - -class IndexTTS25StyleEncoder { -public: - IndexTTS25StyleEncoder( - std::shared_ptr assets, - core::BackendConfig backend, - engine::assets::TensorStorageType weight_storage_type); - - IndexTTS25StyleEmbedding embed_fbank( - const std::vector & features, - int64_t frames, - int64_t dims) const; - void release_graph(); - -private: - std::shared_ptr assets_; - engine::modules::CampplusEncoderComponent component_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/tokenizer_text.h b/include/engine/models/index_tts2_5/tokenizer_text.h deleted file mode 100644 index 65258fd2..00000000 --- a/include/engine/models/index_tts2_5/tokenizer_text.h +++ /dev/null @@ -1,58 +0,0 @@ -#pragma once - -#include "engine/models/index_tts2_5/assets.h" - -#include -#include -#include -#include - -namespace llama_tokenizer_vendor { -struct BpeVocabulary; -} // namespace llama_tokenizer_vendor - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25TextEncoding { - std::string lang; - std::string normalized_text; - std::vector segments; - std::vector> segment_token_ids; -}; - -// Whisper-style tiktoken BPE text tokenizer for IndexTTS-2.5 (vocab size 60509). -// Replaces the SentencePiece tokenizer used by index_tts2. -class IndexTTS25TextTokenizer { -public: - explicit IndexTTS25TextTokenizer(std::shared_ptr assets); - - std::string normalize_english(const std::string & text) const; - std::string normalize_chinese(const std::string & text) const; - - // Raw tiktoken encode with allowed_special="all"; does not apply any text - // normalization. Special tokens present in the text are recognized directly. - std::vector encode(const std::string & text) const; - - // Returns the id of an exact token text (e.g. "<|zh|>"), or -1 when unknown. - int32_t special_token_id(const std::string & token_text) const; - - // Maps a language code to the GPT lang_embedding row, following the - // LANGUAGES order of indextts/utils/tokenizer.py (en=0, zh=1, ...). - // Unknown codes map to "common". - static int32_t lang_to_id(const std::string & lang); - - // Full inference pipeline: normalize -> case rules -> pronunciation - // annotations -> special-token name uppercasing -> segment by token budget. - // Each segment is encoded as encode("<|{lang}|> " + segment) plus a trailing - // pad token id 1. When lang is empty, it is inferred (Han -> zh, else en). - IndexTTS25TextEncoding encode_for_inference( - const std::string & text, - int max_text_tokens_per_segment, - const std::string & lang = "") const; - -private: - std::shared_ptr assets_; - std::shared_ptr vocab_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/types.h b/include/engine/models/index_tts2_5/types.h deleted file mode 100644 index 07f105b9..00000000 --- a/include/engine/models/index_tts2_5/types.h +++ /dev/null @@ -1,153 +0,0 @@ -#pragma once - -#include "engine/framework/runtime/session.h" - -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { - -// Rows in the GPT lang_embedding table. indextts/utils/tokenizer.py defines -// 106 language codes (including "common"); the checkpoint table has one extra -// unused row. -constexpr int64_t kIndexTTS25LangEmbeddingRows = 107; - -struct IndexTTS25GptConfig { - int64_t model_dim = 1280; - int64_t max_mel_tokens = 1815; - int64_t max_text_tokens = 600; - int64_t heads = 20; - bool use_mel_codes_as_input = true; - int64_t mel_length_compression = 1024; - int64_t layers = 24; - int64_t number_text_tokens = 12000; - int64_t number_mel_codes = 8194; - int64_t start_mel_token = 8192; - int64_t stop_mel_token = 8193; - int64_t start_text_token = 0; - int64_t stop_text_token = 1; - bool train_solo_embeddings = false; - std::string condition_type = "conformer_perceiver"; - int64_t condition_output_size = 512; - int64_t condition_linear_units = 2048; - int64_t condition_attention_heads = 8; - int64_t condition_num_blocks = 6; - std::string condition_input_layer = "conv2d2"; - int64_t condition_perceiver_mult = 2; - int64_t emo_condition_output_size = 512; - int64_t emo_condition_linear_units = 1024; - int64_t emo_condition_attention_heads = 4; - int64_t emo_condition_num_blocks = 4; - std::string emo_condition_input_layer = "conv2d2"; - int64_t emo_condition_perceiver_mult = 2; -}; - -struct IndexTTS25SemanticCodecConfig { - int64_t codebook_size = 8192; - int64_t hidden_size = 1024; - int64_t codebook_dim = 8; - int64_t vocos_dim = 384; - int64_t vocos_intermediate_dim = 2048; - int64_t vocos_num_layers = 12; -}; - -struct IndexTTS25S2MelConfig { - int sample_rate = 22050; - int64_t n_fft = 1024; - int64_t win_length = 1024; - int64_t hop_length = 256; - int64_t n_mels = 80; - float fmin = 0.0F; - std::optional fmax = std::nullopt; - std::string dit_type = "DiT"; - std::string reg_loss_type = "l1"; - int64_t style_dim = 192; - int64_t length_regulator_channels = 512; - bool length_regulator_is_discrete = false; - int64_t length_regulator_in_channels = 1024; - int64_t length_regulator_content_codebook_size = 2048; - std::vector length_regulator_sampling_ratios; - bool length_regulator_vector_quantize = false; - int64_t length_regulator_n_codebooks = 1; - float length_regulator_quantizer_dropout = 0.0F; - bool length_regulator_f0_condition = false; - int64_t length_regulator_n_f0_bins = 512; - int64_t dit_hidden_dim = 512; - int64_t dit_num_heads = 8; - int64_t dit_depth = 13; - float dit_class_dropout_prob = 0.1F; - int64_t dit_block_size = 8192; - int64_t dit_in_channels = 80; - bool dit_style_condition = true; - std::string dit_final_layer_type = "wavenet"; - std::string dit_target = "mel"; - int64_t dit_content_dim = 512; - int64_t dit_content_codebook_size = 1024; - std::string dit_content_type = "discrete"; - bool dit_f0_condition = false; - int64_t dit_n_f0_bins = 512; - int64_t dit_content_codebooks = 1; - bool dit_is_causal = false; - bool dit_long_skip_connection = true; - bool dit_zero_prompt_speech_token = false; - bool dit_time_as_token = false; - bool dit_style_as_token = false; - bool dit_uvit_skip_connection = true; - bool dit_add_resblock_in_transformer = false; - int64_t wavenet_hidden_dim = 512; - int64_t wavenet_num_layers = 8; - int64_t wavenet_kernel_size = 5; - int64_t wavenet_dilation_rate = 1; - float wavenet_dropout = 0.2F; - bool wavenet_style_condition = true; -}; - -struct IndexTTS25Config { - std::string version = "2.0"; - int dataset_sample_rate = 24000; - bool dataset_squeeze = false; - int dataset_mel_sample_rate = 24000; - int64_t dataset_mel_n_fft = 1024; - int64_t dataset_mel_hop_length = 256; - int64_t dataset_mel_win_length = 1024; - int64_t dataset_mel_n_mels = 100; - float dataset_mel_fmin = 0.0F; - bool dataset_mel_normalize = false; - IndexTTS25GptConfig gpt; - IndexTTS25SemanticCodecConfig semantic_codec; - IndexTTS25S2MelConfig s2mel; - std::vector emo_num; -}; - -struct IndexTTS25GenerationOptions { - bool do_sample = true; - float top_p = 0.8F; - int top_k = 30; - float temperature = 0.8F; - float length_penalty = 0.0F; - int num_beams = 3; - float repetition_penalty = 10.0F; - int max_mel_tokens = 1500; - uint32_t seed = 0; -}; - -struct IndexTTS25Request { - std::string text; - std::optional speaker_audio = std::nullopt; - std::optional emotion_audio = std::nullopt; - // Text language hint; empty means auto (zh when the text contains Han - // characters, otherwise en). - std::string lang; - float emotion_alpha = 1.0F; - std::optional> emotion_vector = std::nullopt; - bool use_emotion_text = false; - std::optional emotion_text = std::nullopt; - bool use_random_emotion = false; - int interval_silence_ms = 200; - int max_text_tokens_per_segment = 120; - IndexTTS25GenerationOptions generation; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/include/engine/models/index_tts2_5/vocoder.h b/include/engine/models/index_tts2_5/vocoder.h deleted file mode 100644 index f2ab1060..00000000 --- a/include/engine/models/index_tts2_5/vocoder.h +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include "engine/framework/modules/vocoders/bigvgan_vocoder.h" -#include "engine/models/index_tts2_5/assets.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -struct IndexTTS25VocoderOutput { - std::vector waveform; - int64_t samples = 0; - int sample_rate = 0; -}; - -class IndexTTS25BigVganVocoder { -public: - IndexTTS25BigVganVocoder( - std::shared_ptr assets, - core::BackendConfig backend, - engine::assets::TensorStorageType weight_storage_type); - - IndexTTS25VocoderOutput synthesize( - const std::vector & mel, - int64_t frames) const; - void release_runtime_graph(); - -private: - std::shared_ptr assets_; - engine::modules::BigVganVocoderComponent component_; -}; - -} // namespace engine::models::index_tts2_5 diff --git a/model_specs/index_tts2.json b/model_specs/index_tts2.json index 3f6c737c..72c680eb 100644 --- a/model_specs/index_tts2.json +++ b/model_specs/index_tts2.json @@ -1,7 +1,7 @@ { "family": "index_tts2", "display_name": "IndexTTS2", - "description": "Zero-shot TTS system for Chinese and English speech synthesis with voice cloning, emotion-speaker decoupling, text or audio emotion control, and explicit duration control.", + "description": "Zero-shot TTS system with voice cloning, emotion-speaker decoupling, text or audio emotion control, and explicit duration control. Variant v2 covers Chinese and English; variant v2.5 (config version 2.5) adds multilingual tiktoken tokenization with Japanese, Spanish and Arabic support.", "category": "tts", "status": "supported", "tasks": [ @@ -13,7 +13,10 @@ ], "languages": [ "zh", - "en" + "en", + "ja", + "es", + "ar" ], "capabilities": { "tts": [ @@ -99,6 +102,39 @@ "kind": "huggingface_snapshot", "repo": "mlx-community/index-tts2-mlx" } + }, + { + "id": "index_tts2_5_q8_0", + "display_name": "IndexTTS2.5 Q8_0 GGUF", + "format": "gguf", + "precision": "q8_0", + "target_directory": "IndexTTS2.5-GGUF", + "files": [ + "IndexTTS2.5-GGUF/index-tts2_5-q8_0.gguf" + ], + "strip_prefix": "IndexTTS2.5-GGUF" + }, + { + "id": "index_tts2_5_f16", + "display_name": "IndexTTS2.5 F16 GGUF", + "format": "gguf", + "precision": "f16", + "target_directory": "IndexTTS2.5-GGUF", + "files": [ + "IndexTTS2.5-GGUF/index-tts2_5-f16.gguf" + ], + "strip_prefix": "IndexTTS2.5-GGUF" + }, + { + "id": "index_tts2_5_orig", + "display_name": "IndexTTS2.5 Original-Dtype GGUF", + "format": "gguf", + "precision": "orig", + "target_directory": "IndexTTS2.5-GGUF", + "files": [ + "IndexTTS2.5-GGUF/index-tts2_5-orig.gguf" + ], + "strip_prefix": "IndexTTS2.5-GGUF" } ], "sources": [ @@ -110,7 +146,6 @@ }, "files": { "config": "model:config.yaml", - "bpe": "model:bpe.model", "wav2vec2bert_config": "model:w2v-bert-2.0/config.json", "wav2vec2bert_preprocessor_config": "model:w2v-bert-2.0/preprocessor_config.json", "bigvgan_config": "model:bigvgan/config.json", @@ -121,6 +156,10 @@ "qwen_emotion_vocab": "model:qwen0.6bemo4-merge/vocab.json", "qwen_emotion_merges": "model:qwen0.6bemo4-merge/merges.txt" }, + "optional_files": { + "bpe": "model:bpe.model", + "tiktoken": "model:multilingual_zh_ja_yue_char_del.tiktoken" + }, "tensors": { "gpt": { "source": "weights:", @@ -171,7 +210,6 @@ }, "files": { "config": "model:config.yaml", - "bpe": "model:bpe.model", "wav2vec2bert_config": "model:w2v-bert-2.0/config.json", "wav2vec2bert_preprocessor_config": "model:w2v-bert-2.0/preprocessor_config.json", "bigvgan_config": "model:bigvgan/config.json", @@ -182,6 +220,10 @@ "qwen_emotion_vocab": "model:qwen0.6bemo4-merge/vocab.json", "qwen_emotion_merges": "model:qwen0.6bemo4-merge/merges.txt" }, + "optional_files": { + "bpe": "model:bpe.model", + "tiktoken": "model:multilingual_zh_ja_yue_char_del.tiktoken" + }, "tensors": { "gpt": "model:gpt.safetensors", "s2mel": "model:s2mel.safetensors", diff --git a/model_specs/index_tts2_5.json b/model_specs/index_tts2_5.json deleted file mode 100644 index 5b4d8661..00000000 --- a/model_specs/index_tts2_5.json +++ /dev/null @@ -1,186 +0,0 @@ -{ - "family": "index_tts2_5", - "display_name": "IndexTTS2.5", - "description": "Multilingual zero-shot TTS system for Chinese, English, Japanese, Spanish and Arabic speech synthesis with voice cloning, timbre-emotion decoupling, text or audio emotion control, and explicit duration control.", - "category": "tts", - "status": "supported", - "tasks": [ - "tts", - "clone" - ], - "modes": [ - "offline" - ], - "languages": [ - "zh", - "en", - "ja", - "es", - "ar" - ], - "capabilities": { - "tts": [ - "emotion_control" - ], - "clone": [ - "speaker_reference", - "emotion_control" - ] - }, - "runtime": { - "tags": [ - "gguf" - ] - }, - "ui": { - "recommended_package": "index_tts2_5_q8_0", - "tags": [ - "TTS", - "Clone", - "GGUF" - ], - "docs": [ - "docs/tts.md", - "docs/gguf.md" - ] - }, - "package_defaults": { - "download": { - "kind": "huggingface_snapshot", - "repo": "audio-cpp/audio.cpp-gguf", - "revision": "main", - "gated": false - } - }, - "packages": [ - { - "id": "index_tts2_5_q8_0", - "display_name": "IndexTTS2.5 Q8_0 GGUF", - "default": true, - "format": "gguf", - "precision": "q8_0", - "target_directory": "IndexTTS2.5-GGUF", - "files": [ - "IndexTTS2.5-GGUF/index-tts2_5-q8_0.gguf" - ], - "strip_prefix": "IndexTTS2.5-GGUF" - }, - { - "id": "index_tts2_5_f16", - "display_name": "IndexTTS2.5 F16 GGUF", - "format": "gguf", - "precision": "f16", - "target_directory": "IndexTTS2.5-GGUF", - "files": [ - "IndexTTS2.5-GGUF/index-tts2_5-f16.gguf" - ], - "strip_prefix": "IndexTTS2.5-GGUF" - }, - { - "id": "index_tts2_5_orig", - "display_name": "IndexTTS2.5 Original-Dtype GGUF", - "format": "gguf", - "precision": "orig", - "target_directory": "IndexTTS2.5-GGUF", - "files": [ - "IndexTTS2.5-GGUF/index-tts2_5-orig.gguf" - ], - "strip_prefix": "IndexTTS2.5-GGUF" - } - ], - "sources": [ - { - "format": "gguf", - "roots": { - "model": ".", - "weights": "$gguf" - }, - "files": { - "config": "model:config.yaml", - "tiktoken": "model:multilingual_zh_ja_yue_char_del.tiktoken", - "wav2vec2bert_config": "model:w2v-bert-2.0/config.json", - "wav2vec2bert_preprocessor_config": "model:w2v-bert-2.0/preprocessor_config.json", - "bigvgan_config": "model:bigvgan/config.json", - "qwen_emotion_config": "model:qwen0.6bemo4-merge/config.json", - "qwen_emotion_generation_config": "model:qwen0.6bemo4-merge/generation_config.json", - "qwen_emotion_tokenizer": "model:qwen0.6bemo4-merge/tokenizer.json", - "qwen_emotion_tokenizer_config": "model:qwen0.6bemo4-merge/tokenizer_config.json", - "qwen_emotion_vocab": "model:qwen0.6bemo4-merge/vocab.json", - "qwen_emotion_merges": "model:qwen0.6bemo4-merge/merges.txt" - }, - "tensors": { - "gpt": { - "source": "weights:", - "prefix": "gpt" - }, - "s2mel": { - "source": "weights:", - "prefix": "s2mel" - }, - "speaker_matrix": { - "source": "weights:", - "prefix": "speaker_matrix" - }, - "emotion_matrix": { - "source": "weights:", - "prefix": "emotion_matrix" - }, - "wav2vec2bert_stats": { - "source": "weights:", - "prefix": "wav2vec2bert_stats" - }, - "wav2vec2bert": { - "source": "weights:", - "prefix": "wav2vec2bert" - }, - "semantic_codec": { - "source": "weights:", - "prefix": "semantic_codec" - }, - "campplus": { - "source": "weights:", - "prefix": "campplus" - }, - "bigvgan": { - "source": "weights:", - "prefix": "bigvgan" - }, - "qwen_emotion": { - "source": "weights:", - "prefix": "qwen_emotion" - } - } - }, - { - "format": "safetensors", - "roots": { - "model": "." - }, - "files": { - "config": "model:config.yaml", - "tiktoken": "model:multilingual_zh_ja_yue_char_del.tiktoken", - "wav2vec2bert_config": "model:w2v-bert-2.0/config.json", - "wav2vec2bert_preprocessor_config": "model:w2v-bert-2.0/preprocessor_config.json", - "bigvgan_config": "model:bigvgan/config.json", - "qwen_emotion_config": "model:qwen0.6bemo4-merge/config.json", - "qwen_emotion_generation_config": "model:qwen0.6bemo4-merge/generation_config.json", - "qwen_emotion_tokenizer": "model:qwen0.6bemo4-merge/tokenizer.json", - "qwen_emotion_tokenizer_config": "model:qwen0.6bemo4-merge/tokenizer_config.json", - "qwen_emotion_vocab": "model:qwen0.6bemo4-merge/vocab.json", - "qwen_emotion_merges": "model:qwen0.6bemo4-merge/merges.txt" - }, - "tensors": { - "gpt": "model:gpt.safetensors", - "s2mel": "model:s2mel.safetensors", - "speaker_matrix": "model:feat1.safetensors", - "emotion_matrix": "model:feat2.safetensors", - "wav2vec2bert_stats": "model:wav2vec2bert_stats.safetensors", - "wav2vec2bert": "model:w2v-bert-2.0/model.safetensors", - "semantic_codec": "model:semantic_codec_model.safetensors", - "campplus": "model:campplus.safetensors", - "bigvgan": "model:bigvgan/model.safetensors", - "qwen_emotion": "model:qwen0.6bemo4-merge/model.safetensors" - } - } - ] -} diff --git a/src/models/index_tts2/assets.cpp b/src/models/index_tts2/assets.cpp index 35ad21ea..9a5b4ac5 100644 --- a/src/models/index_tts2/assets.cpp +++ b/src/models/index_tts2/assets.cpp @@ -17,13 +17,27 @@ IndexTTS2Config parse_config(const assets::ResourceBundle & resources) { const auto document = resources.parse_flattened_yaml("config"); IndexTTS2Config config; config.version = yaml::optional_string(document, "version", config.version); - config.dataset_sample_rate = static_cast(yaml::require_i64(document, "dataset.sample_rate")); + // The official IndexTTS-2.5 config has no dataset section; these values are + // parsed for compatibility but not used at inference time. + if (const auto value = yaml::optional_int(document, "dataset.sample_rate")) { + config.dataset_sample_rate = *value; + } config.dataset_squeeze = yaml::optional_bool(document, "dataset.squeeze", config.dataset_squeeze); - config.dataset_mel_sample_rate = static_cast(yaml::require_i64(document, "dataset.mel.sample_rate")); - config.dataset_mel_n_fft = yaml::require_i64(document, "dataset.mel.n_fft"); - config.dataset_mel_hop_length = yaml::require_i64(document, "dataset.mel.hop_length"); - config.dataset_mel_win_length = yaml::require_i64(document, "dataset.mel.win_length"); - config.dataset_mel_n_mels = yaml::require_i64(document, "dataset.mel.n_mels"); + if (const auto value = yaml::optional_int(document, "dataset.mel.sample_rate")) { + config.dataset_mel_sample_rate = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.n_fft")) { + config.dataset_mel_n_fft = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.hop_length")) { + config.dataset_mel_hop_length = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.win_length")) { + config.dataset_mel_win_length = *value; + } + if (const auto value = yaml::optional_int(document, "dataset.mel.n_mels")) { + config.dataset_mel_n_mels = *value; + } config.dataset_mel_fmin = yaml::optional_f32(document, "dataset.mel.mel_fmin", config.dataset_mel_fmin); config.dataset_mel_normalize = yaml::optional_bool(document, "dataset.mel.normalize", config.dataset_mel_normalize); @@ -154,7 +168,14 @@ void validate_gpt_weights(const IndexTTS2Config & config, const assets::TensorSo assets::require_tensor_shape(source, "gpt.h.0.attn.c_proj.weight", {config.gpt.model_dim, config.gpt.model_dim}); assets::require_tensor_shape(source, "gpt.h.0.mlp.c_fc.weight", {config.gpt.model_dim, config.gpt.model_dim * 4}); assets::require_tensor_shape(source, "gpt.h.0.mlp.c_proj.weight", {config.gpt.model_dim * 4, config.gpt.model_dim}); - assets::require_tensor_shape(source, "conditioning_encoder.after_norm.weight", {config.gpt.condition_output_size}); + if (index_tts2_variant_from_version(config.version) == IndexTTS2Variant::kV2_5) { + // v2.5 campplus speaker conditioning: projected CAMPPlus embedding plus + // the language embedding table; no conditioning_encoder/speed_emb. + assets::require_tensor_shape(source, "spk_emb_proj.weight", {config.gpt.model_dim, config.s2mel.style_dim}); + assets::require_tensor_shape(source, "lang_embedding.weight", {kIndexTTS2LangEmbeddingRows, config.gpt.model_dim}); + } else { + assets::require_tensor_shape(source, "conditioning_encoder.after_norm.weight", {config.gpt.condition_output_size}); + } assets::require_tensor_shape(source, "emo_conditioning_encoder.after_norm.weight", {config.gpt.emo_condition_output_size}); } @@ -195,6 +216,11 @@ void validate_semantic_codec_weights(const IndexTTS2Config & config, const asset assets::require_tensor_shape(source, "quantizer.quantizers.0.codebook.weight", {config.semantic_codec.codebook_size, config.semantic_codec.codebook_dim}); assets::require_tensor_shape(source, "encoder.1.weight", {config.semantic_codec.hidden_size, config.semantic_codec.vocos_dim}); assets::require_tensor_shape(source, "decoder.1.weight", {config.semantic_codec.hidden_size, config.semantic_codec.vocos_dim}); + if (index_tts2_variant_from_version(config.version) == IndexTTS2Variant::kV2_5) { + // v2.5 decodes codes through the full EnhancedCodec path ending in the + // 2x nearest upsample and the `up` conv. + assets::require_tensor_shape(source, "up.weight", {config.semantic_codec.hidden_size, config.semantic_codec.hidden_size, 3}); + } } void validate_qwen_weights(const assets::TensorSource & source) { diff --git a/src/models/index_tts2/gpt.cpp b/src/models/index_tts2/gpt.cpp index 45e885d9..2647efac 100644 --- a/src/models/index_tts2/gpt.cpp +++ b/src/models/index_tts2/gpt.cpp @@ -39,7 +39,6 @@ constexpr int64_t kSpeakerConditionLayers = 6; constexpr int64_t kEmotionConditionLayers = 4; constexpr int64_t kGptLayers = 24; constexpr int64_t kGptMlpDim = 5120; -constexpr int64_t kTextTokens = 12001; constexpr int64_t kMelCodes = 8194; constexpr int64_t kMelPositions = 1818; constexpr int64_t kTextPositions = 602; @@ -47,12 +46,22 @@ constexpr int64_t kConditionPosFrames = 5000; constexpr int64_t kConditionConvKernel = 15; constexpr int64_t kGptHeads = 20; constexpr int64_t kGptHeadDim = kModelDim / kGptHeads; -constexpr int64_t kConditionTokens = 34; +// v2: 32 perceiver speaker tokens + two speed embedding tokens. +constexpr int64_t kV2ConditionTokens = 34; +constexpr int64_t kV2TextTokens = 12001; +// v2.5 spk_cond_mode="campplus": the projected 192-dim CAMPPlus embedding +// forms a single speaker token, followed by two all-zero tokens. +constexpr int64_t kCampplusStyleDim = 192; +constexpr int64_t kV2_5ConditionTokens = 3; constexpr int32_t kStartTextToken = 0; constexpr int32_t kStopTextToken = 1; constexpr int32_t kStartMelToken = 8192; constexpr int32_t kStopMelToken = 8193; +int64_t gpt_text_vocab_size(const IndexTTS2Config & config) { + return config.gpt.number_text_tokens + 1; +} + struct GgmlContextDeleter { void operator()(ggml_context * ctx) const noexcept { if (ctx != nullptr) { @@ -920,15 +929,20 @@ std::shared_ptr load_index_tts2_gpt_weights( weight_context_bytes); const auto & source = *assets.gpt_weights; - weights->speaker_conditioner = load_condition_encoder( - *weights->store, - source, - "conditioning_encoder", - kSpeakerConditionLayers, - 2048, - 8, - matmul_storage_type, - conv_storage_type); + const bool campplus_conditioning = + index_tts2_variant_from_version(assets.config.version) == IndexTTS2Variant::kV2_5; + const int64_t text_vocab = gpt_text_vocab_size(assets.config); + if (!campplus_conditioning) { + weights->speaker_conditioner = load_condition_encoder( + *weights->store, + source, + "conditioning_encoder", + kSpeakerConditionLayers, + 2048, + 8, + matmul_storage_type, + conv_storage_type); + } weights->emotion_conditioner = load_condition_encoder( *weights->store, source, @@ -938,16 +952,18 @@ std::shared_ptr load_index_tts2_gpt_weights( 4, matmul_storage_type, conv_storage_type); - weights->speaker_perceiver = load_perceiver( - *weights->store, - source, - "perceiver_encoder", - 32, - kModelDim, - kConditionDim, - 512, - 3412, - matmul_storage_type); + if (!campplus_conditioning) { + weights->speaker_perceiver = load_perceiver( + *weights->store, + source, + "perceiver_encoder", + 32, + kModelDim, + kConditionDim, + 512, + 3412, + matmul_storage_type); + } weights->emotion_perceiver = load_perceiver( *weights->store, source, @@ -962,7 +978,7 @@ std::shared_ptr load_index_tts2_gpt_weights( source, "text_embedding.weight", matmul_storage_type, - {kTextTokens, kModelDim}); + {text_vocab, kModelDim}); weights->mel_embedding = weights->store->load_tensor( source, "mel_embedding.weight", @@ -992,7 +1008,23 @@ std::shared_ptr load_index_tts2_gpt_weights( kModelDim, kModelDim, true); - weights->speed_embedding_values = source.require_f32("speed_emb.weight", {2, kModelDim}); + if (campplus_conditioning) { + weights->spk_emb_proj = binding::linear_from_source( + *weights->store, + source, + "spk_emb_proj", + matmul_storage_type, + kModelDim, + kCampplusStyleDim, + true); + weights->lang_embedding = weights->store->load_tensor( + source, + "lang_embedding.weight", + matmul_storage_type, + {kIndexTTS2LangEmbeddingRows, kModelDim}); + } else { + weights->speed_embedding_values = source.require_f32("speed_emb.weight", {2, kModelDim}); + } weights->gpt_layers.reserve(static_cast(kGptLayers)); for (int64_t i = 0; i < kGptLayers; ++i) { weights->gpt_layers.push_back(load_gpt2_layer(*weights->store, source, i, matmul_storage_type)); @@ -1012,7 +1044,7 @@ std::shared_ptr load_index_tts2_gpt_weights( source, "text_head", matmul_storage_type, - kTextTokens, + text_vocab, kModelDim, true); @@ -1298,13 +1330,18 @@ class IndexTTS2GptRuntime::PrefillGraph { core::ExecutionContext & execution, std::shared_ptr weights, int64_t text_tokens, + bool campplus_conditioning, + int64_t text_vocab, size_t graph_arena_bytes) : execution_(execution), weights_(std::move(weights)), + campplus_conditioning_(campplus_conditioning), + condition_tokens_(campplus_conditioning ? kV2_5ConditionTokens : kV2ConditionTokens), + text_vocab_(text_vocab), text_tokens_(text_tokens), text_steps_(text_tokens + 2), - prompt_steps_(kConditionTokens + text_tokens + 3) { - if (weights_ == nullptr || text_tokens_ <= 0) { + prompt_steps_(condition_tokens_ + text_tokens + 3) { + if (weights_ == nullptr || text_tokens_ < 0 || (!campplus_conditioning_ && text_tokens_ == 0)) { throw std::runtime_error("IndexTTS2 GPT prefill graph requires weights and text tokens"); } const auto build_start = Clock::now(); @@ -1332,17 +1369,43 @@ class IndexTTS2GptRuntime::PrefillGraph { output_ctx_.get(), "index_tts2.gpt.prefill.outputs", execution_.backend_type()}; - conds_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kConditionTokens, kModelDim})).tensor; + core::TensorValue conds; + if (campplus_conditioning_) { + style_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kCampplusStyleDim})).tensor; + emo_vec_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kModelDim})).tensor; + lang_id_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 1); + ggml_set_input(style_); + ggml_set_input(emo_vec_); + ggml_set_input(lang_id_); + } else { + conds_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, condition_tokens_, kModelDim})).tensor; + ggml_set_input(conds_); + conds = core::wrap_tensor(conds_, core::TensorShape::from_dims({1, condition_tokens_, kModelDim}), GGML_TYPE_F32); + } text_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, text_steps_); start_mel_id_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 1); - ggml_set_input(conds_); ggml_set_input(text_ids_); ggml_set_input(start_mel_id_); - auto conds = core::wrap_tensor(conds_, core::TensorShape::from_dims({1, kConditionTokens, kModelDim}), GGML_TYPE_F32); auto text_ids = core::wrap_tensor(text_ids_, core::TensorShape::from_dims({text_steps_}), GGML_TYPE_I32); - auto text = modules::EmbeddingModule({kTextTokens, kModelDim}).build(ctx, text_ids, weights_->text_embedding); + auto text = modules::EmbeddingModule({text_vocab_, kModelDim}).build(ctx, text_ids, weights_->text_embedding); auto text_pos = modules::SliceModule({0, 0, text_steps_}).build(ctx, weights_->text_pos_embedding); text = modules::AddModule{}.build(ctx, text, text_pos); + if (campplus_conditioning_) { + // v2.5 campplus conditioning prefix (model_v2.py inference_speech): + // conds = [spk_emb_proj(style) + emo_vec, zeros, zeros]. + auto style = core::wrap_tensor(style_, core::TensorShape::from_dims({1, kCampplusStyleDim}), GGML_TYPE_F32); + auto speaker_token = build_biased_gpt_projection(ctx, style, kCampplusStyleDim, kModelDim, weights_->spk_emb_proj); + speaker_token = core::reshape_tensor(ctx, speaker_token, core::TensorShape::from_dims({1, 1, kModelDim})); + auto emo_vec = core::wrap_tensor(emo_vec_, core::TensorShape::from_dims({1, kModelDim}), GGML_TYPE_F32); + emo_vec = core::reshape_tensor(ctx, emo_vec, core::TensorShape::from_dims({1, 1, kModelDim})); + conds = modules::AddModule{}.build(ctx, speaker_token, emo_vec); + auto zero_token = modules::RepeatModule({core::TensorShape::from_dims({1, condition_tokens_ - 1, kModelDim})}) + .build(ctx, scale(ctx, conds, 0.0F)); + conds = modules::ConcatModule({1}).build(ctx, conds, zero_token); + auto lang_id = core::wrap_tensor(lang_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto lang = modules::EmbeddingModule({kIndexTTS2LangEmbeddingRows, kModelDim}).build(ctx, lang_id, weights_->lang_embedding); + text = modules::AddModule{}.build(ctx, text, modules::RepeatModule({text.shape}).build(ctx, lang)); + } text = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, text), core::TensorShape::from_dims({1, text_steps_, kModelDim})); auto mel_id = core::wrap_tensor(start_mel_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); auto mel = modules::EmbeddingModule({kMelCodes, kModelDim}).build(ctx, mel_id, weights_->mel_embedding); @@ -1412,17 +1475,48 @@ class IndexTTS2GptRuntime::PrefillGraph { } GptPrefillOutput run(const std::vector & conds, const std::vector & text_tokens) { - if (static_cast(conds.size()) != kConditionTokens * kModelDim || + if (campplus_conditioning_) { + throw std::runtime_error("IndexTTS2 GPT prefill conds input requires the v2 speaker-conditioning mode"); + } + if (static_cast(conds.size()) != condition_tokens_ * kModelDim || + static_cast(text_tokens.size()) != text_tokens_) { + throw std::runtime_error("IndexTTS2 GPT prefill input shape mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(conds_, conds.data(), 0, conds.size() * sizeof(float)); + return run_with_text_tokens(text_tokens, timing_start); + } + + GptPrefillOutput run( + const std::vector & speaker_style, + const std::vector & emotion_vector, + int32_t lang_id, + const std::vector & text_tokens) { + if (!campplus_conditioning_) { + throw std::runtime_error("IndexTTS2 GPT prefill style/lang inputs require the v2.5 campplus speaker-conditioning mode"); + } + if (static_cast(speaker_style.size()) != kCampplusStyleDim || + static_cast(emotion_vector.size()) != kModelDim || static_cast(text_tokens.size()) != text_tokens_) { throw std::runtime_error("IndexTTS2 GPT prefill input shape mismatch"); } + if (lang_id < 0 || lang_id >= kIndexTTS2LangEmbeddingRows) { + throw std::runtime_error("IndexTTS2 GPT prefill lang id is out of range"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set(style_, speaker_style.data(), 0, speaker_style.size() * sizeof(float)); + ggml_backend_tensor_set(emo_vec_, emotion_vector.data(), 0, emotion_vector.size() * sizeof(float)); + ggml_backend_tensor_set(lang_id_, &lang_id, 0, sizeof(int32_t)); + return run_with_text_tokens(text_tokens, timing_start); + } + +private: + GptPrefillOutput run_with_text_tokens(const std::vector & text_tokens, Clock::time_point timing_start) { std::vector ids; ids.reserve(static_cast(text_steps_)); ids.push_back(kStartTextToken); ids.insert(ids.end(), text_tokens.begin(), text_tokens.end()); ids.push_back(kStopTextToken); - auto timing_start = Clock::now(); - ggml_backend_tensor_set(conds_, conds.data(), 0, conds.size() * sizeof(float)); ggml_backend_tensor_set(text_ids_, ids.data(), 0, ids.size() * sizeof(int32_t)); debug::timing_log_scalar("index_tts2.gpt.prefill.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); core::set_backend_threads(execution_.backend(), execution_.config().threads); @@ -1452,7 +1546,6 @@ class IndexTTS2GptRuntime::PrefillGraph { return out; } -private: void clear_graph() { if (graph_ != nullptr) { core::release_backend_graph_resources(execution_.backend(), graph_); @@ -1474,6 +1567,9 @@ class IndexTTS2GptRuntime::PrefillGraph { core::ExecutionContext & execution_; std::shared_ptr weights_; + bool campplus_conditioning_ = false; + int64_t condition_tokens_ = kV2ConditionTokens; + int64_t text_vocab_ = 0; int64_t text_tokens_ = 0; int64_t text_steps_ = 0; int64_t prompt_steps_ = 0; @@ -1481,6 +1577,9 @@ class IndexTTS2GptRuntime::PrefillGraph { std::unique_ptr output_ctx_; std::unique_ptr ctx_; ggml_tensor * conds_ = nullptr; + ggml_tensor * style_ = nullptr; + ggml_tensor * emo_vec_ = nullptr; + ggml_tensor * lang_id_ = nullptr; ggml_tensor * text_ids_ = nullptr; ggml_tensor * start_mel_id_ = nullptr; ggml_tensor * latent_ = nullptr; @@ -1526,16 +1625,16 @@ class IndexTTS2GptRuntime::ForwardGraph { input_ctx_.get(), "index_tts2.gpt.forward.inputs", execution_.backend_type()}; - conds_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kConditionTokens, kModelDim})).tensor; + conds_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kV2ConditionTokens, kModelDim})).tensor; text_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, text_steps_); mel_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, mel_steps_); ggml_set_input(conds_); ggml_set_input(text_ids_); ggml_set_input(mel_ids_); - auto conds = core::wrap_tensor(conds_, core::TensorShape::from_dims({1, kConditionTokens, kModelDim}), GGML_TYPE_F32); + auto conds = core::wrap_tensor(conds_, core::TensorShape::from_dims({1, kV2ConditionTokens, kModelDim}), GGML_TYPE_F32); auto text_ids = core::wrap_tensor(text_ids_, core::TensorShape::from_dims({text_steps_}), GGML_TYPE_I32); - auto text = modules::EmbeddingModule({kTextTokens, kModelDim}).build(ctx, text_ids, weights_->text_embedding); + auto text = modules::EmbeddingModule({kV2TextTokens, kModelDim}).build(ctx, text_ids, weights_->text_embedding); auto text_pos = modules::SliceModule({0, 0, text_steps_}).build(ctx, weights_->text_pos_embedding); text = modules::AddModule{}.build(ctx, text, text_pos); text = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, text), core::TensorShape::from_dims({1, text_steps_, kModelDim})); @@ -1553,14 +1652,14 @@ class IndexTTS2GptRuntime::ForwardGraph { x = out.output; } x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->gpt_final_norm); - x = modules::SliceModule({1, kConditionTokens + text_steps_, code_count_}).build(ctx, x); + x = modules::SliceModule({1, kV2ConditionTokens + text_steps_, code_count_}).build(ctx, x); x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->final_norm); output_ = core::ensure_backend_addressable_layout(ctx, x).tensor; ggml_set_output(output_); graph_ = ggml_new_graph_custom( ctx_.get(), - static_cast(std::max(65536, (kConditionTokens + text_steps_ + mel_steps_) * 8192)), + static_cast(std::max(65536, (kV2ConditionTokens + text_steps_ + mel_steps_) * 8192)), false); ggml_build_forward_expand(graph_, output_); input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); @@ -1591,7 +1690,7 @@ class IndexTTS2GptRuntime::ForwardGraph { const std::vector & conds, const std::vector & text_tokens, const std::vector & codes) { - if (static_cast(conds.size()) != kConditionTokens * kModelDim || + if (static_cast(conds.size()) != kV2ConditionTokens * kModelDim || static_cast(text_tokens.size()) != text_tokens_ || static_cast(codes.size()) != code_count_) { throw std::runtime_error("IndexTTS2 GPT forward graph input shape mismatch"); @@ -2022,6 +2121,9 @@ void IndexTTS2GptRuntime::prepare_speaker_conditioning(int64_t frames) { if (execution_ == nullptr) { throw std::runtime_error("IndexTTS2 GPT runtime execution context is missing"); } + if (index_tts2_variant_from_version(assets_->config.version) != IndexTTS2Variant::kV2) { + throw std::runtime_error("IndexTTS2 GPT speaker conditioning is only used by the v2 variant; v2.5 uses campplus conditioning inside the prefill graph"); + } if (frames <= 0) { throw std::runtime_error("IndexTTS2 GPT speaker conditioning prepare requires positive frames"); } @@ -2064,7 +2166,9 @@ void IndexTTS2GptRuntime::prepare_generation(int64_t text_tokens, int64_t max_me if (execution_ == nullptr) { throw std::runtime_error("IndexTTS2 GPT runtime execution context is missing"); } - if (text_tokens <= 0 || max_mel_tokens <= 0) { + const bool campplus_conditioning = + index_tts2_variant_from_version(assets_->config.version) == IndexTTS2Variant::kV2_5; + if (text_tokens < 0 || (!campplus_conditioning && text_tokens == 0) || max_mel_tokens <= 0) { throw std::runtime_error("IndexTTS2 GPT generation prepare requires positive lengths"); } if (num_beams != 1) { @@ -2072,7 +2176,13 @@ void IndexTTS2GptRuntime::prepare_generation(int64_t text_tokens, int64_t max_me } if (prefill_graph_ == nullptr || !prefill_graph_->matches(text_tokens)) { prefill_graph_.reset(); - prefill_graph_ = std::make_unique(*execution_, weights_, text_tokens, graph_arena_bytes_); + prefill_graph_ = std::make_unique( + *execution_, + weights_, + text_tokens, + campplus_conditioning, + gpt_text_vocab_size(assets_->config), + graph_arena_bytes_); } const int64_t required_cache_steps = prefill_graph_->prompt_steps() + max_mel_tokens + 1; const int64_t required_beam_slots = 2 * std::max(1, num_beams); @@ -2113,11 +2223,28 @@ std::vector IndexTTS2GptRuntime::merge_emotion_vector( } IndexTTS2GptGeneration IndexTTS2GptRuntime::generate_speech(const IndexTTS2GptGenerationRequest & request) { - if (request.text_tokens.empty()) { - throw std::runtime_error("IndexTTS2 GPT generation requires text tokens"); + const bool campplus_conditioning = + index_tts2_variant_from_version(assets_->config.version) == IndexTTS2Variant::kV2_5; + std::vector text_tokens = request.text_tokens; + IndexTTS2GptLatent speech_conditioning; + if (campplus_conditioning) { + // v2.5: drop embedded start/stop text tokens (mirrors the valid_mask + // filtering in the official prepare_gpt_inputs); the speaker token is + // built inside the prefill graph from the CAMPPlus style embedding. + text_tokens = align_index_tts2_gpt_text_tokens(request.text_tokens); + if (static_cast(request.speaker_style.size()) != kCampplusStyleDim) { + throw std::runtime_error("IndexTTS2 GPT generation speaker style shape mismatch"); + } + if (request.lang_id < 0 || request.lang_id >= kIndexTTS2LangEmbeddingRows) { + throw std::runtime_error("IndexTTS2 GPT generation lang id is out of range"); + } + } else { + if (request.text_tokens.empty()) { + throw std::runtime_error("IndexTTS2 GPT generation requires text tokens"); + } + prepare_speaker_conditioning(request.speaker_frames); + speech_conditioning = speaker_conditioning(request.speaker_semantic, request.speaker_frames); } - prepare_speaker_conditioning(request.speaker_frames); - const auto speech_conditioning = speaker_conditioning(request.speaker_semantic, request.speaker_frames); std::vector emotion_vector = request.emotion_vector; if (emotion_vector.empty()) { prepare_emotion_conditioning(request.emotion_frames); @@ -2127,21 +2254,26 @@ IndexTTS2GptGeneration IndexTTS2GptRuntime::generate_speech(const IndexTTS2GptGe throw std::runtime_error("IndexTTS2 GPT generation emotion vector shape mismatch"); } prepare_generation( - static_cast(request.text_tokens.size()), + static_cast(text_tokens.size()), request.max_mel_tokens, request.num_beams); - std::vector conds(static_cast(kConditionTokens * kModelDim), 0.0F); - for (int64_t token = 0; token < 32; ++token) { - for (int64_t dim = 0; dim < kModelDim; ++dim) { - conds[static_cast(token * kModelDim + dim)] = - speech_conditioning.values[static_cast(token * kModelDim + dim)] + - emotion_vector[static_cast(dim)]; + GptPrefillOutput prefill; + if (campplus_conditioning) { + prefill = prefill_graph_->run(request.speaker_style, emotion_vector, request.lang_id, text_tokens); + } else { + std::vector conds(static_cast(kV2ConditionTokens * kModelDim), 0.0F); + for (int64_t token = 0; token < 32; ++token) { + for (int64_t dim = 0; dim < kModelDim; ++dim) { + conds[static_cast(token * kModelDim + dim)] = + speech_conditioning.values[static_cast(token * kModelDim + dim)] + + emotion_vector[static_cast(dim)]; + } } + const auto & speed = weights_->speed_embedding_values; + std::copy_n(speed.data() + static_cast(kModelDim), static_cast(kModelDim), conds.data() + static_cast(32 * kModelDim)); + std::copy_n(speed.data(), static_cast(kModelDim), conds.data() + static_cast(33 * kModelDim)); + prefill = prefill_graph_->run(conds, request.text_tokens); } - const auto & speed = weights_->speed_embedding_values; - std::copy_n(speed.data() + static_cast(kModelDim), static_cast(kModelDim), conds.data() + static_cast(32 * kModelDim)); - std::copy_n(speed.data(), static_cast(kModelDim), conds.data() + static_cast(33 * kModelDim)); - auto prefill = prefill_graph_->run(conds, request.text_tokens); const auto sampling_policy = engine::sampling::resolve_torch_cuda_sampling_policy( execution_->backend_type(), execution_->config().device, @@ -2400,6 +2532,9 @@ IndexTTS2GptLatent IndexTTS2GptRuntime::forward_latent( const std::vector & emotion_semantic, int64_t emotion_frames, const std::vector & emotion_vector) { + if (index_tts2_variant_from_version(assets_->config.version) != IndexTTS2Variant::kV2) { + throw std::runtime_error("IndexTTS2 GPT latent forward is only used by the v2 variant; v2.5 decodes codes through the semantic codec"); + } if (speech_conditioning_latent.frames != 32 || speech_conditioning_latent.dims != kModelDim || static_cast(speech_conditioning_latent.values.size()) != 32 * kModelDim) { @@ -2416,7 +2551,7 @@ IndexTTS2GptLatent IndexTTS2GptRuntime::forward_latent( if (static_cast(emo.size()) != kModelDim) { throw std::runtime_error("IndexTTS2 GPT latent forward emotion vector shape mismatch"); } - std::vector conds(static_cast(kConditionTokens * kModelDim), 0.0F); + std::vector conds(static_cast(kV2ConditionTokens * kModelDim), 0.0F); for (int64_t token = 0; token < 32; ++token) { for (int64_t dim = 0; dim < kModelDim; ++dim) { conds[static_cast(token * kModelDim + dim)] = @@ -2449,4 +2584,16 @@ void IndexTTS2GptRuntime::release_generation_graphs() { forward_graph_.reset(); } +std::vector align_index_tts2_gpt_text_tokens(const std::vector & text_tokens) { + std::vector out; + out.reserve(text_tokens.size()); + for (const int32_t token : text_tokens) { + if (token == kStartTextToken || token == kStopTextToken) { + continue; + } + out.push_back(token); + } + return out; +} + } // namespace engine::models::index_tts2 diff --git a/src/models/index_tts2/loader.cpp b/src/models/index_tts2/loader.cpp index f99f3915..cf45822b 100644 --- a/src/models/index_tts2/loader.cpp +++ b/src/models/index_tts2/loader.cpp @@ -13,11 +13,13 @@ runtime::ModelMetadata metadata(const IndexTTS2Assets & assets) { runtime::ModelMetadata out; out.family = "index_tts2"; out.variant = assets.config.version; - out.description = "IndexTTS2 loaded from local extracted assets."; + out.description = index_tts2_variant_from_version(assets.config.version) == IndexTTS2Variant::kV2_5 + ? "IndexTTS2.5 (index_tts2 family variant) loaded from local extracted assets." + : "IndexTTS2 loaded from local extracted assets."; return out; } -runtime::CapabilitySet capabilities(const IndexTTS2Assets &) { +runtime::CapabilitySet capabilities(const IndexTTS2Assets & assets) { runtime::CapabilitySet out; out.supported_tasks = { {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, @@ -25,11 +27,15 @@ runtime::CapabilitySet capabilities(const IndexTTS2Assets &) { }; out.supports_speaker_reference = true; out.supports_style_condition = true; - out.languages = {"English", "Chinese"}; + if (index_tts2_variant_from_version(assets.config.version) == IndexTTS2Variant::kV2_5) { + out.languages = {"Chinese", "English", "Japanese", "Spanish", "Arabic"}; + } else { + out.languages = {"English", "Chinese"}; + } return out; } -runtime::ModelCliInterface cli(const IndexTTS2Assets &) { +runtime::ModelCliInterface cli(const IndexTTS2Assets & assets) { runtime::ModelCliInterface out; out.request_options = { {"emotion_alpha", "float", "Blend strength for explicit emotion conditioning."}, @@ -42,6 +48,10 @@ runtime::ModelCliInterface cli(const IndexTTS2Assets &) { {"length_penalty", "float", "GPT beam-search length penalty."}, {"num_beams", "n", "GPT beam count."}, }; + if (index_tts2_variant_from_version(assets.config.version) == IndexTTS2Variant::kV2_5) { + out.request_options.push_back( + {"lang", "auto|zh|en|ja|es|ar|...", "Text language hint; auto infers zh when the text contains Han characters, otherwise en."}); + } out.session_options = { {"index_tts2.weight_type", "native|f32|f16|bf16|q8_0", "Matmul weight storage type."}, {"index_tts2.conv_weight_type", "native|f32|f16", "Convolution weight storage type."}, diff --git a/src/models/index_tts2/request.cpp b/src/models/index_tts2/request.cpp index 25dfda23..abf9af22 100644 --- a/src/models/index_tts2/request.cpp +++ b/src/models/index_tts2/request.cpp @@ -3,6 +3,8 @@ #include "engine/framework/io/text.h" #include "engine/framework/runtime/options.h" +#include +#include #include #include #include @@ -57,6 +59,17 @@ void require_valid_audio(const runtime::AudioBuffer & audio, const char * label) } // namespace +std::string normalize_index_tts2_lang(const std::string & value) { + std::string lang = engine::io::trim_ascii_whitespace(value); + std::transform(lang.begin(), lang.end(), lang.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + if (lang == "auto") { + lang.clear(); + } + return lang; +} + IndexTTS2Request parse_index_tts2_request(const runtime::TaskRequest & request) { IndexTTS2Request out; if (request.text_input.has_value()) { @@ -74,6 +87,9 @@ IndexTTS2Request parse_index_tts2_request(const runtime::TaskRequest & request) } else { throw std::runtime_error("IndexTTS2 request requires --voice-ref or voice.speaker.audio"); } + if (const auto value = runtime::find_option(request.options, {"lang"})) { + out.lang = normalize_index_tts2_lang(*value); + } if (const auto value = runtime::parse_finite_float_option(request.options, {"emotion_alpha"})) { if (*value < 0.0F || *value > 1.0F) { diff --git a/src/models/index_tts2/s2mel.cpp b/src/models/index_tts2/s2mel.cpp index ab6a80e6..9209d503 100644 --- a/src/models/index_tts2/s2mel.cpp +++ b/src/models/index_tts2/s2mel.cpp @@ -237,7 +237,11 @@ core::TensorValue cfm_wavenet( const IndexTTS2S2MelCfmWeights & weights) { auto g = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, timestep_b), core::TensorShape::from_dims({timestep_b.shape.dims[0], kHidden, 1})); g = modules::Conv1dModule({kHidden, 2 * kHidden * kWavenetLayers, 1, 1, 0, 1, true}).build(ctx, g, weights.wavenet_cond); - auto output = sub(ctx, input_bct, input_bct); + // The zero accumulator must come from a contiguous tensor: input_bct is a + // permuted (transposed) view, and the ggml CPU binary-op kernels miscompute + // permuted src operands (the CUDA kernels handle them). + const auto zeros_base = core::ensure_backend_addressable_layout(ctx, input_bct); + auto output = sub(ctx, zeros_base, zeros_base); auto x = input_bct; for (int64_t i = 0; i < kWavenetLayers; ++i) { const int64_t dilation = 1; diff --git a/src/models/index_tts2/semantic_codec.cpp b/src/models/index_tts2/semantic_codec.cpp index c1aff416..f026e601 100644 --- a/src/models/index_tts2/semantic_codec.cpp +++ b/src/models/index_tts2/semantic_codec.cpp @@ -408,10 +408,12 @@ class IndexTTS2SemanticCodecRuntime::CodesGraph { core::ExecutionContext & execution, std::shared_ptr weights, int64_t frames, + bool upsample_decode, size_t graph_arena_bytes) : execution_(execution), weights_(std::move(weights)), - frames_(frames) { + frames_(frames), + upsample_decode_(upsample_decode) { if (frames_ <= 0) { throw std::runtime_error("IndexTTS2 semantic codec code graph requires positive frame count"); } @@ -436,11 +438,28 @@ class IndexTTS2SemanticCodecRuntime::CodesGraph { "index_tts2.semantic_codec.codes.inputs", execution_.backend_type()}; codes_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, frames_})).tensor; + if (upsample_decode_) { + upsample_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 2 * frames_); + } ggml_set_input(codes_); auto embedding = embed_codes_bct( ctx, core::wrap_tensor(codes_, core::TensorShape::from_dims({1, frames_}), GGML_TYPE_I32), *weights_); + if (upsample_decode_) { + // v2.5 EnhancedCodec.decode (codec/models.py): decoder backbone + + // projection, then 2x nearest upsample along time and the `up` conv. + auto x = vocos_backbone(ctx, embedding, weights_->decoder_backbone); + x = modules::LinearModule({kVocosDim, kHidden, true}).build(ctx, x, weights_->decoder_projection); + x = core::ensure_backend_addressable_layout(ctx, x); + x = core::wrap_tensor( + ggml_get_rows(ctx.ggml, x.tensor, upsample_ids_), + core::TensorShape::from_dims({1, 2 * frames_, kHidden}), + GGML_TYPE_F32); + x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); + x = modules::Conv1dModule({kHidden, kHidden, 3, 1, 1, 1, true}).build(ctx, x, weights_->up); + embedding = x; + } embedding_ = core::ensure_backend_addressable_layout(ctx, embedding).tensor; ggml_set_output(embedding_); @@ -450,6 +469,14 @@ class IndexTTS2SemanticCodecRuntime::CodesGraph { if (input_buffer_ == nullptr) { throw std::runtime_error("failed to allocate IndexTTS2 semantic codec code input buffer"); } + if (upsample_decode_) { + std::vector upsample_ids(static_cast(2 * frames_)); + for (int64_t frame = 0; frame < frames_; ++frame) { + upsample_ids[static_cast(2 * frame)] = static_cast(frame); + upsample_ids[static_cast(2 * frame + 1)] = static_cast(frame); + } + ggml_backend_tensor_set(upsample_ids_, upsample_ids.data(), 0, upsample_ids.size() * sizeof(int32_t)); + } gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); if (gallocr_ == nullptr || !ggml_gallocr_reserve(gallocr_, graph_) || @@ -493,10 +520,10 @@ class IndexTTS2SemanticCodecRuntime::CodesGraph { } IndexTTS2SemanticCodecOutput output; - output.frames = frames_; + output.frames = upsample_decode_ ? 2 * frames_ : frames_; output.dims = kHidden; output.codes = codes; - output.embedding_channel_first.resize(static_cast(kHidden * frames_)); + output.embedding_channel_first.resize(static_cast(kHidden * output.frames)); timing_start = Clock::now(); ggml_backend_tensor_get( embedding_, @@ -528,9 +555,11 @@ class IndexTTS2SemanticCodecRuntime::CodesGraph { core::ExecutionContext & execution_; std::shared_ptr weights_; int64_t frames_ = 0; + bool upsample_decode_ = false; std::unique_ptr input_ctx_; std::unique_ptr ctx_; ggml_tensor * codes_ = nullptr; + ggml_tensor * upsample_ids_ = nullptr; ggml_tensor * embedding_ = nullptr; ggml_cgraph * graph_ = nullptr; ggml_gallocr_t gallocr_ = nullptr; @@ -610,6 +639,17 @@ std::shared_ptr load_index_tts2_semantic_co kHidden, kVocosDim, true); + if (index_tts2_variant_from_version(assets.config.version) == IndexTTS2Variant::kV2_5) { + weights->up = binding::conv1d_from_source( + *weights->store, + source, + "up", + conv_storage_type, + kHidden, + kHidden, + 3, + true); + } weights->store->upload(); assets.semantic_codec_weights->release_storage(); @@ -668,7 +708,12 @@ void IndexTTS2SemanticCodecRuntime::prepare_codes(int64_t frames) { return; } codes_graph_.reset(); - codes_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); + codes_graph_ = std::make_unique( + *execution_, + weights_, + frames, + index_tts2_variant_from_version(assets_->config.version) == IndexTTS2Variant::kV2_5, + graph_arena_bytes_); } IndexTTS2SemanticCodecOutput IndexTTS2SemanticCodecRuntime::quantize(const IndexTTS2SemanticEmbedding & semantic) { diff --git a/src/models/index_tts2/session.cpp b/src/models/index_tts2/session.cpp index 98b025db..8e51de63 100644 --- a/src/models/index_tts2/session.cpp +++ b/src/models/index_tts2/session.cpp @@ -36,6 +36,10 @@ std::shared_ptr require_assets(std::shared_ptr(align_index_tts2_gpt_text_tokens(segment).size()) + : static_cast(segment.size()); gpt_->prepare_generation( - static_cast(segment.size()), + gpt_text_tokens, generation.max_mel_tokens, generation.num_beams); } @@ -384,18 +398,30 @@ const IndexTTS2Session::SpeakerState & IndexTTS2Session::resolve_speaker_state(c true); semantic_encoder_->prepare(prepared.semantic_features.frames); auto semantic = semantic_encoder_->encode(prepared.semantic_features); - semantic_codec_->prepare_quantize(semantic.frames); - auto reference_codes = semantic_codec_->quantize(semantic); - const auto reference_content = channel_first_to_time_major( - reference_codes.embedding_channel_first, - reference_codes.dims, - reference_codes.frames); - debug::trace_log_scalar("index_tts2.s2mel.reference_mel_frames", static_cast(prepared.mel.frames)); - s2mel_->prepare_length_regulator(reference_codes.frames, prepared.mel.frames); - auto prompt_condition = s2mel_->regulate_length( - reference_content, - reference_codes.frames, - prepared.mel.frames); + IndexTTS2S2MelSequence prompt_condition; + if (is_v2_5_variant(*assets_)) { + // Official v2.5 regulates the raw (normalized) w2v-bert semantic + // directly; the semantic codec is only used to decode generated codes. + debug::trace_log_scalar("index_tts2.s2mel.reference_mel_frames", static_cast(prepared.mel.frames)); + s2mel_->prepare_length_regulator(semantic.frames, prepared.mel.frames); + prompt_condition = s2mel_->regulate_length( + semantic.values, + semantic.frames, + prepared.mel.frames); + } else { + semantic_codec_->prepare_quantize(semantic.frames); + auto reference_codes = semantic_codec_->quantize(semantic); + const auto reference_content = channel_first_to_time_major( + reference_codes.embedding_channel_first, + reference_codes.dims, + reference_codes.frames); + debug::trace_log_scalar("index_tts2.s2mel.reference_mel_frames", static_cast(prepared.mel.frames)); + s2mel_->prepare_length_regulator(reference_codes.frames, prepared.mel.frames); + prompt_condition = s2mel_->regulate_length( + reference_content, + reference_codes.frames, + prepared.mel.frames); + } SpeakerState state; state.identity = identity; @@ -592,15 +618,22 @@ std::vector IndexTTS2Session::resolve_emotion_vector( runtime::AudioBuffer IndexTTS2Session::synthesize_segment( const std::vector & text_tokens, + int32_t lang_id, const SpeakerState & speaker, const EmotionState & emotion, const std::vector & emotion_vector, const IndexTTS2GenerationOptions & options, uint32_t segment_seed) { + const bool v2_5 = is_v2_5_variant(*assets_); IndexTTS2GptGenerationRequest generation; generation.text_tokens = text_tokens; - generation.speaker_semantic = speaker.semantic.values; - generation.speaker_frames = speaker.semantic.frames; + if (v2_5) { + generation.speaker_style = speaker.style.values; + generation.lang_id = lang_id; + } else { + generation.speaker_semantic = speaker.semantic.values; + generation.speaker_frames = speaker.semantic.frames; + } generation.emotion_semantic = emotion.semantic.values; generation.emotion_frames = emotion.semantic.frames; generation.emotion_vector = emotion_vector; @@ -624,29 +657,42 @@ runtime::AudioBuffer IndexTTS2Session::synthesize_segment( throw std::runtime_error("IndexTTS2 GPT generated no acoustic codes"); } const int64_t code_frames = static_cast(generated.codes.size()); - const int64_t target_frames = static_cast(static_cast(code_frames) * 1.72F); - const int64_t total_frames = speaker.prompt_condition.frames + target_frames; - const auto forward_start = Clock::now(); - auto latent = gpt_->forward_latent( - generated.speech_conditioning_latent, - text_tokens, - generated.codes, - emotion.semantic.values, - emotion.semantic.frames, - emotion_vector); - debug::timing_log_scalar("index_tts2.gpt.forward_ms", engine::debug::elapsed_ms(forward_start)); const auto s2mel_start = Clock::now(); - s2mel_->prepare_gpt_layer(latent.frames); - auto projected = s2mel_->project_gpt_latent(latent.values, latent.frames); semantic_codec_->prepare_codes(code_frames); auto semantic = semantic_codec_->codes_to_embedding(generated.codes, code_frames); if (mem_saver_) { semantic_codec_->release_graphs(); } - auto content = add_latent_to_semantic(semantic, projected); - s2mel_->prepare_length_regulator(code_frames, target_frames); - auto generated_condition = s2mel_->regulate_length(content, code_frames, target_frames); + std::vector content; + int64_t content_frames = 0; + if (v2_5) { + // v2.5: the codec decode (with its 2x upsample) yields the S2Mel + // content directly; there is no GPT latent projection in this variant. + content = channel_first_to_time_major( + semantic.embedding_channel_first, + semantic.dims, + semantic.frames); + content_frames = semantic.frames; + } else { + const auto forward_start = Clock::now(); + auto latent = gpt_->forward_latent( + generated.speech_conditioning_latent, + text_tokens, + generated.codes, + emotion.semantic.values, + emotion.semantic.frames, + emotion_vector); + debug::timing_log_scalar("index_tts2.gpt.forward_ms", engine::debug::elapsed_ms(forward_start)); + s2mel_->prepare_gpt_layer(latent.frames); + auto projected = s2mel_->project_gpt_latent(latent.values, latent.frames); + content = add_latent_to_semantic(semantic, projected); + content_frames = code_frames; + } + const int64_t target_frames = static_cast(static_cast(content_frames) * 1.72F); + const int64_t total_frames = speaker.prompt_condition.frames + target_frames; + s2mel_->prepare_length_regulator(content_frames, target_frames); + auto generated_condition = s2mel_->regulate_length(content, content_frames, target_frames); auto condition = concat_conditions(speaker.prompt_condition, generated_condition); if (mem_saver_) { gpt_->release_generation_graphs(); @@ -713,15 +759,19 @@ runtime::TaskResult IndexTTS2Session::run(const runtime::TaskRequest & request) throw std::runtime_error("IndexTTS2 text chunking produced no chunks"); } + const bool v2_5 = is_v2_5_variant(*assets_); std::vector> segment_token_ids; + std::vector segment_lang_ids; for (const auto & text_chunk : text_chunks) { const auto text_encoding = tokenizer_.encode_for_inference( text_chunk, - parsed.max_text_tokens_per_segment); - segment_token_ids.insert( - segment_token_ids.end(), - text_encoding.segment_token_ids.begin(), - text_encoding.segment_token_ids.end()); + parsed.max_text_tokens_per_segment, + parsed.lang); + const int32_t lang_id = v2_5 ? IndexTTS2TextTokenizer::lang_to_id(text_encoding.lang) : 0; + for (const auto & ids : text_encoding.segment_token_ids) { + segment_token_ids.push_back(ids); + segment_lang_ids.push_back(lang_id); + } } runtime::AudioBuffer merged; @@ -731,6 +781,7 @@ runtime::TaskResult IndexTTS2Session::run(const runtime::TaskRequest & request) } auto segment_audio = synthesize_segment( segment_token_ids[i], + segment_lang_ids[i], speaker, emotion, emotion_vector, diff --git a/src/models/index_tts2/tokenizer_text.cpp b/src/models/index_tts2/tokenizer_text.cpp index b7b18d80..fd2ec69b 100644 --- a/src/models/index_tts2/tokenizer_text.cpp +++ b/src/models/index_tts2/tokenizer_text.cpp @@ -3,26 +3,30 @@ #include "engine/framework/text/chinese_normalization.h" #include "engine/framework/text/text_normalization.h" +#include "bpe-core.h" +#include "unicode.h" + #include +#include #include +#include +#include +#include +#include #include #include #include #include +#include namespace engine::models::index_tts2 { namespace { -bool contains_token( - const std::vector & values, - const std::vector & needles) { - for (const auto & value : values) { - if (std::find(needles.begin(), needles.end(), value) != needles.end()) { - return true; - } - } - return false; -} +namespace vendor = llama_tokenizer_vendor; + +// --------------------------------------------------------------------------- +// Shared UTF-8 / text helpers +// --------------------------------------------------------------------------- size_t utf8_codepoint_size(unsigned char byte) { if ((byte & 0x80U) == 0U) { @@ -57,27 +61,22 @@ uint32_t decode_utf8_codepoint(const std::string & text, size_t offset, size_t s return byte(0); } -bool is_han_codepoint(uint32_t cp) { - return (cp >= 0x4E00U && cp <= 0x9FFFU); +uint32_t next_utf8_codepoint(const std::string & text, size_t & offset) { + const size_t size = std::min(utf8_codepoint_size(static_cast(text[offset])), text.size() - offset); + const uint32_t cp = decode_utf8_codepoint(text, offset, size); + offset += size; + return cp; } -bool is_cjk_codepoint(uint32_t cp) { - return (cp >= 0x1100U && cp <= 0x11FFU) - || (cp >= 0x2E80U && cp <= 0xA4CFU) - || (cp >= 0xA840U && cp <= 0xD7AFU) - || (cp >= 0xF900U && cp <= 0xFAFFU) - || (cp >= 0xFE30U && cp <= 0xFE4FU) - || (cp >= 0xFF65U && cp <= 0xFFDCU) - || (cp >= 0x20000U && cp <= 0x2FFFFU); +bool is_han_codepoint(uint32_t cp) { + return cp >= 0x4E00U && cp <= 0x9FFFU; } bool contains_han(const std::string & text) { for (size_t i = 0; i < text.size();) { - const size_t size = std::min(utf8_codepoint_size(static_cast(text[i])), text.size() - i); - if (is_han_codepoint(decode_utf8_codepoint(text, i, size))) { + if (is_han_codepoint(next_utf8_codepoint(text, i))) { return true; } - i += size; } return false; } @@ -101,6 +100,38 @@ std::string uppercase_ascii(std::string text) { return text; } +std::string lowercase_ascii(std::string text) { + for (char & ch : text) { + ch = static_cast(std::tolower(static_cast(ch))); + } + return text; +} + +// --------------------------------------------------------------------------- +// v2 SentencePiece helpers +// --------------------------------------------------------------------------- + +bool contains_token( + const std::vector & values, + const std::vector & needles) { + for (const auto & value : values) { + if (std::find(needles.begin(), needles.end(), value) != needles.end()) { + return true; + } + } + return false; +} + +bool is_cjk_codepoint(uint32_t cp) { + return (cp >= 0x1100U && cp <= 0x11FFU) + || (cp >= 0x2E80U && cp <= 0xA4CFU) + || (cp >= 0xA840U && cp <= 0xD7AFU) + || (cp >= 0xF900U && cp <= 0xFAFFU) + || (cp >= 0xFE30U && cp <= 0xFE4FU) + || (cp >= 0xFF65U && cp <= 0xFFDCU) + || (cp >= 0x20000U && cp <= 0x2FFFFU); +} + std::string tokenize_by_cjk_char(const std::string & text) { std::vector tokens; std::string pending; @@ -207,6 +238,457 @@ std::vector> split_segments_by_token( return merged_segments; } +// --------------------------------------------------------------------------- +// v2.5 tiktoken helpers +// --------------------------------------------------------------------------- + +// IndexTTS-2.5 pads every text segment with a trailing token id 1. +constexpr int32_t kSegmentPadTokenId = 1; + +std::string decode_base64(const std::string & input) { + static const std::array table = [] { + std::array values{}; + values.fill(-1); + for (int i = 0; i < 26; ++i) { + values[static_cast('A' + i)] = static_cast(i); + values[static_cast('a' + i)] = static_cast(26 + i); + } + for (int i = 0; i < 10; ++i) { + values[static_cast('0' + i)] = static_cast(52 + i); + } + values[static_cast('+')] = 62; + values[static_cast('/')] = 63; + return values; + }(); + + std::string out; + int bits = 0; + int value = 0; + for (const unsigned char ch : input) { + if (ch == '=') { + break; + } + const int8_t digit = table[ch]; + if (digit < 0) { + throw std::runtime_error("IndexTTS2 tiktoken vocabulary contains invalid base64 token bytes"); + } + value = (value << 6) | digit; + bits += 6; + if (bits >= 8) { + bits -= 8; + out.push_back(static_cast((value >> bits) & 0xff)); + } + } + return out; +} + +// The vendored llama BPE runtime works in the GPT-2 byte-to-unicode domain +// (e.g. space becomes U+0120) so that byte-level merges, including tokens that +// split a UTF-8 codepoint, are reproduced exactly. tiktoken ranks are keyed by +// raw bytes, so every token is mapped once at load time. +std::string map_token_bytes(const std::string & bytes) { + std::string mapped; + for (const unsigned char byte : bytes) { + mapped += unicode_byte_to_utf8(byte); + } + return mapped; +} + +std::string pair_key(const std::string & left, const std::string & right) { + std::string key = left; + key.push_back('\0'); + key += right; + return key; +} + +// Language codes in the LANGUAGES order of indextts/utils/tokenizer.py. The +// first 99 entries double as the <|lang|> special tokens below; the remaining +// codes (plus the fallback "common") only index the GPT lang_embedding table. +const std::array kLanguages = { + "en", "zh", "de", "es", "ru", "ko", "fr", "ja", "pt", "tr", + "pl", "ca", "nl", "ar", "sv", "it", "id", "hi", "fi", "vi", + "he", "uk", "el", "ms", "cs", "ro", "da", "hu", "ta", "no", + "th", "ur", "hr", "bg", "lt", "la", "mi", "ml", "cy", "sk", + "te", "fa", "lv", "bn", "sr", "az", "sl", "kn", "et", "mk", + "br", "eu", "is", "hy", "ne", "mn", "bs", "kk", "sq", "sw", + "gl", "mr", "pa", "si", "km", "sn", "yo", "so", "af", "oc", + "ka", "be", "tg", "sd", "gu", "am", "yi", "lo", "uz", "fo", + "ht", "ps", "tk", "nn", "mt", "sa", "lb", "my", "bo", "tl", + "mg", "as", "tt", "haw", "ln", "ha", "ba", "jw", "su", +}; +const std::array kEmbeddingOnlyLanguages = { + "yue", "minnan", "wuyu", "dialect", "zh/en", "en/zh", "common", +}; +constexpr int32_t kCommonLangId = 105; + +void add_special_token(vendor::BpeVocabulary & vocab, const std::string & text, int32_t id) { + vocab.token_to_id.emplace(text, id); + vocab.id_to_token.emplace(id, vendor::TokenData{text, vendor::TOKEN_ATTR_CONTROL}); +} + +// Special token order must match indextts/utils/tokenizer.py exactly: +// ids are assigned sequentially starting right after the mergeable ranks. +void register_special_tokens(vendor::BpeVocabulary & vocab, int32_t base_id) { + static const std::array kAudioEvents = { + "ASR", "AED", "SER", "Speech", "/Speech", "BGM", "/BGM", + "Laughter", "/Laughter", "Applause", "/Applause", + }; + static const std::array kEmotions = { + "HAPPY", "SAD", "ANGRY", "NEUTRAL", + }; + static const std::array kTasks = { + "translate", "transcribe", "startoflm", "startofprev", "nospeech", "notimestamps", + }; + static const std::array kTtsVocal = { + "TTS/B", "TTS/O", "TTS/Q", "TTS/A", "TTS/CO", "TTS/CL", "TTS/H", + }; + + int32_t id = base_id; + add_special_token(vocab, "<|endoftext|>", id++); + add_special_token(vocab, "<|startoftranscript|>", id++); + for (const char * lang : kLanguages) { + add_special_token(vocab, "<|" + std::string(lang) + "|>", id++); + } + for (const char * event : kAudioEvents) { + add_special_token(vocab, "<|" + std::string(event) + "|>", id++); + } + for (const char * emotion : kEmotions) { + add_special_token(vocab, "<|" + std::string(emotion) + "|>", id++); + } + for (const char * task : kTasks) { + add_special_token(vocab, "<|" + std::string(task) + "|>", id++); + } + for (int i = 1; i <= 30; ++i) { + add_special_token(vocab, "<|SPECIAL_TOKEN_" + std::to_string(i) + "|>", id++); + } + for (const char * vocal : kTtsVocal) { + add_special_token(vocab, "<|" + std::string(vocal) + "|>", id++); + } + for (int i = 1; i <= 13; ++i) { + char name[32]; + std::snprintf(name, sizeof(name), "<|TTS/SP%02d|>", i); + add_special_token(vocab, name, id++); + } + // Timestamps <|0.00|> .. <|30.00|> in 0.02 steps; i * 0.02 == i / 50. + for (int i = 0; i <= 1500; ++i) { + char name[32]; + std::snprintf(name, sizeof(name), "<|%d.%02d|>", i / 50, (i * 2) % 100); + add_special_token(vocab, name, id++); + } +} + +std::shared_ptr load_tiktoken_vocabulary(const std::filesystem::path & vocab_path) { + std::ifstream input(vocab_path, std::ios::binary); + if (!input) { + throw std::runtime_error("IndexTTS2 failed to open tiktoken vocabulary: " + vocab_path.string()); + } + + auto vocab = std::make_shared(); + vocab->pre_type = vendor::PreTokenizerType::Gpt2; + + std::string line; + int64_t mergeable_count = 0; + while (std::getline(input, line)) { + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + if (line.empty()) { + continue; + } + std::istringstream parts(line); + std::string token_base64; + int64_t rank = -1; + if (!(parts >> token_base64 >> rank) || rank < 0 || rank > INT32_MAX) { + throw std::runtime_error("IndexTTS2 tiktoken vocabulary has an invalid line: " + line); + } + const std::string bytes = decode_base64(token_base64); + const auto token_id = static_cast(rank); + const std::string mapped = map_token_bytes(bytes); + vocab->token_to_id.emplace(mapped, token_id); + vocab->id_to_token.emplace(token_id, vendor::TokenData{mapped, 0}); + // tiktoken ranks double as merge priorities: an adjacent pair merges + // iff its concatenation is a token, with that token's rank. Register + // every split so find_bpe_rank(left, right) == rank(left + right). + for (size_t split = 1; split < bytes.size(); ++split) { + vocab->bpe_ranks.emplace( + pair_key(map_token_bytes(bytes.substr(0, split)), map_token_bytes(bytes.substr(split))), + token_id); + } + ++mergeable_count; + } + if (mergeable_count == 0) { + throw std::runtime_error("IndexTTS2 tiktoken vocabulary is empty: " + vocab_path.string()); + } + + register_special_tokens(*vocab, static_cast(mergeable_count)); + vendor::rebuild_special_tokens_cache(*vocab); + return vocab; +} + +bool is_kana(const std::string & text) { + if (text.empty()) { + return false; + } + bool all_hiragana = true; + bool all_katakana = true; + for (size_t i = 0; i < text.size();) { + const uint32_t cp = next_utf8_codepoint(text, i); + if (cp < 0x3040U || cp > 0x309FU) { + all_hiragana = false; + } + if (cp < 0x30A0U || cp > 0x30FFU) { + all_katakana = false; + } + } + return all_hiragana || all_katakana; +} + +struct AnnotationMatch { + size_t end = 0; // one past the match; 0 when there is no match at pos + size_t word_begin = 0; + size_t word_end = 0; + size_t pron_begin = 0; + size_t pron_end = 0; +}; + +// Matches <([^|>\n]+)\|([^>\n]+)> anchored at pos. +AnnotationMatch match_pronunciation_annotation(const std::string & text, size_t pos) { + AnnotationMatch match; + if (text[pos] != '<') { + return match; + } + size_t cursor = pos + 1; + const size_t word_begin = cursor; + while (cursor < text.size() && text[cursor] != '|' && text[cursor] != '>' && text[cursor] != '\n') { + ++cursor; + } + if (cursor == word_begin || cursor >= text.size() || text[cursor] != '|') { + return match; + } + match.word_begin = word_begin; + match.word_end = cursor; + const size_t pron_begin = ++cursor; + while (cursor < text.size() && text[cursor] != '>' && text[cursor] != '\n') { + ++cursor; + } + if (cursor == pron_begin || cursor >= text.size()) { + return AnnotationMatch{}; + } + match.pron_begin = pron_begin; + match.pron_end = cursor; + match.end = cursor + 1; + return match; +} + +// Base-26 spreadsheet-style index ("a".."z", "aa"..), mirroring the official +// TextNormalizer._protect_pronunciation_annotations placeholder naming. +std::string alpha_placeholder_index(size_t n) { + std::string s; + while (true) { + s.insert(s.begin(), static_cast('a' + (n % 26))); + const size_t q = n / 26; + if (q == 0) { + break; + } + n = q - 1; + } + return s; +} + +using PronunciationPlaceholders = std::vector>; + +// Replaces annotations with letter-only placeholders so +// text normalization cannot rewrite their digits/symbols (e.g. XING2). +std::pair protect_pronunciation_annotations(const std::string & text) { + std::string out; + out.reserve(text.size()); + PronunciationPlaceholders placeholders; + size_t pos = 0; + while (pos < text.size()) { + const auto match = match_pronunciation_annotation(text, pos); + if (match.end == 0) { + out.push_back(text[pos++]); + continue; + } + std::string key = "PRONPLACEHOLDER" + alpha_placeholder_index(placeholders.size()) + "PRONPLACEHOLDER"; + placeholders.emplace_back(key, text.substr(pos, match.end - pos)); + out += key; + pos = match.end; + } + return {out, placeholders}; +} + +std::string restore_pronunciation_annotations(std::string text, const PronunciationPlaceholders & placeholders) { + for (const auto & [key, original] : placeholders) { + size_t at = 0; + while ((at = text.find(key, at)) != std::string::npos) { + text.replace(at, key.size(), original); + at += original.size(); + } + } + return text; +} + +// Expands annotations (see infer_v2_5.py +// apply_pronunciation_annotations): +// Chinese word -> <|SPECIAL_TOKEN_2|>PRON<|SPECIAL_TOKEN_2|> +// other word -> <|SPECIAL_TOKEN_1|>PRON<|SPECIAL_TOKEN_1|> +// kana pron -> inlined as " PRON " +std::string apply_pronunciation_annotations(const std::string & text) { + std::string out; + out.reserve(text.size()); + size_t pos = 0; + while (pos < text.size()) { + const auto match = match_pronunciation_annotation(text, pos); + if (match.end == 0) { + out.push_back(text[pos++]); + continue; + } + const std::string word = text.substr(match.word_begin, match.word_end - match.word_begin); + const std::string pron = uppercase_ascii(text.substr(match.pron_begin, match.pron_end - match.pron_begin)); + if (is_kana(pron)) { + out.push_back(' '); + out += pron; + out.push_back(' '); + } else { + const char * wrapper = contains_han(word) ? "<|SPECIAL_TOKEN_2|>" : "<|SPECIAL_TOKEN_1|>"; + out += wrapper; + out += pron; + out += wrapper; + } + pos = match.end; + } + return out; +} + +// Uppercases the name inside <|...|> markers: re.sub(r'<\|([^|]+)\|>', upper). +std::string uppercase_special_token_names(const std::string & text) { + std::string out; + out.reserve(text.size()); + size_t pos = 0; + while (pos < text.size()) { + if (text[pos] != '<' || pos + 1 >= text.size() || text[pos + 1] != '|') { + out.push_back(text[pos++]); + continue; + } + size_t cursor = pos + 2; + while (cursor < text.size() && text[cursor] != '|') { + ++cursor; + } + if (cursor == pos + 2 || cursor + 1 >= text.size() || text[cursor + 1] != '>') { + out.push_back(text[pos++]); + continue; + } + out += "<|"; + out += uppercase_ascii(text.substr(pos + 2, cursor - (pos + 2))); + out += "|>"; + pos = cursor + 2; + } + return out; +} + +bool is_segment_delimiter(uint32_t cp) { + switch (cp) { + case U',': + case U'.': + case U'!': + case U'?': + case U';': + case U':': + case U'\n': + case 0xFF0CU: // ， + case 0x3002U: // 。 + case 0xFF01U: // ！ + case 0xFF1FU: // ？ + case 0x3001U: // 、 + case 0xFF1BU: // ； + case 0xFF1AU: // ： + return true; + default: + return false; + } +} + +// re.split(r'(?<=[，。！？、；：,\.!\?;:\n])', piece): split after each delimiter. +std::vector split_after_delimiters(const std::string & piece) { + std::vector parts; + std::string current; + for (size_t i = 0; i < piece.size();) { + const size_t begin = i; + const uint32_t cp = next_utf8_codepoint(piece, i); + current.append(piece, begin, i - begin); + if (is_segment_delimiter(cp)) { + parts.push_back(std::move(current)); + current.clear(); + } + } + if (!current.empty()) { + parts.push_back(std::move(current)); + } + return parts; +} + +// Matches "<|SPECIAL_TOKEN_|>" at pos; returns the match length or 0. +size_t match_special_token_marker(const std::string & text, size_t pos) { + static const std::string kPrefix = "<|SPECIAL_TOKEN_"; + if (text.compare(pos, kPrefix.size(), kPrefix) != 0) { + return 0; + } + size_t cursor = pos + kPrefix.size(); + const size_t digits_begin = cursor; + while (cursor < text.size() && std::isdigit(static_cast(text[cursor])) != 0) { + ++cursor; + } + if (cursor == digits_begin || cursor + 1 >= text.size() || text[cursor] != '|' || text[cursor + 1] != '>') { + return 0; + } + return cursor + 2 - pos; +} + +// SPLIT_PROTECTED_PATTERN spans (<|SPECIAL_TOKEN_n|>...<|SPECIAL_TOKEN_n|>) +// are kept atomic during segmentation. +std::vector> split_atomic_pieces(const std::string & text) { + std::vector> pieces; + size_t pos = 0; + while (pos < text.size()) { + size_t opener = std::string::npos; + size_t opener_len = 0; + for (size_t i = pos; i < text.size(); ++i) { + const size_t len = match_special_token_marker(text, i); + if (len > 0) { + opener = i; + opener_len = len; + break; + } + } + if (opener == std::string::npos) { + break; + } + size_t closer = std::string::npos; + size_t closer_len = 0; + for (size_t i = opener + opener_len; i < text.size(); ++i) { + const size_t len = match_special_token_marker(text, i); + if (len > 0) { + closer = i; + closer_len = len; + break; + } + } + if (closer == std::string::npos) { + break; + } + if (opener > pos) { + pieces.emplace_back(text.substr(pos, opener - pos), false); + } + pieces.emplace_back(text.substr(opener, closer + closer_len - opener), true); + pos = closer + closer_len; + } + if (pos < text.size()) { + pieces.emplace_back(text.substr(pos), false); + } + return pieces; +} + } // namespace IndexTTS2TextTokenizer::IndexTTS2TextTokenizer(std::shared_ptr assets) @@ -214,10 +696,15 @@ IndexTTS2TextTokenizer::IndexTTS2TextTokenizer(std::shared_ptrresources.require_file("bpe")); - piece_to_id_.reserve(pieces_.size()); - for (const auto & piece : pieces_) { - piece_to_id_.emplace(piece.text, static_cast(piece.id)); + variant_ = index_tts2_variant_from_version(assets_->config.version); + if (variant_ == IndexTTS2Variant::kV2_5) { + vocab_ = load_tiktoken_vocabulary(assets_->resources.require_file("tiktoken")); + } else { + pieces_ = engine::tokenizers::load_sentencepiece_model(assets_->resources.require_file("bpe")); + piece_to_id_.reserve(pieces_.size()); + for (const auto & piece : pieces_) { + piece_to_id_.emplace(piece.text, static_cast(piece.id)); + } } } @@ -225,7 +712,7 @@ std::string IndexTTS2TextTokenizer::normalize_english(const std::string & text) engine::text::EnglishTextNormalizationOptions options; options.expand_common_contractions = true; options.index_tts_punctuation = true; - options.uppercase_ascii = true; + options.uppercase_ascii = variant_ == IndexTTS2Variant::kV2; return engine::text::normalize_english_text(text, options); } @@ -235,15 +722,24 @@ std::string IndexTTS2TextTokenizer::normalize_chinese(const std::string & text) engine::text::ChineseTextNormalizationTarget::IndexTTS); } -std::string IndexTTS2TextTokenizer::normalize_text(const std::string & text) const { - return contains_han(text) ? tokenize_by_cjk_char(normalize_chinese(text)) : normalize_english(text); -} - std::vector IndexTTS2TextTokenizer::encode(const std::string & text) const { + if (variant_ == IndexTTS2Variant::kV2_5) { + return vendor::tokenize_bpe(*vocab_, text, true); + } return engine::tokenizers::tokenize_sentencepiece(pieces_, normalize_text(text)); } +std::string IndexTTS2TextTokenizer::normalize_text(const std::string & text) const { + if (variant_ != IndexTTS2Variant::kV2) { + throw std::runtime_error("IndexTTS2 normalize_text is only available for the v2 SentencePiece tokenizer"); + } + return contains_han(text) ? tokenize_by_cjk_char(normalize_chinese(text)) : normalize_english(text); +} + std::vector IndexTTS2TextTokenizer::tokenize_to_pieces(const std::string & text) const { + if (variant_ != IndexTTS2Variant::kV2) { + throw std::runtime_error("IndexTTS2 tokenize_to_pieces is only available for the v2 SentencePiece tokenizer"); + } const auto ids = encode(text); std::vector out; out.reserve(ids.size()); @@ -253,12 +749,45 @@ std::vector IndexTTS2TextTokenizer::tokenize_to_pieces(const std::s return out; } +int32_t IndexTTS2TextTokenizer::special_token_id(const std::string & token_text) const { + if (variant_ != IndexTTS2Variant::kV2_5) { + throw std::runtime_error("IndexTTS2 special_token_id is only available for the v2.5 tiktoken tokenizer"); + } + const auto it = vocab_->token_to_id.find(token_text); + return it == vocab_->token_to_id.end() ? -1 : it->second; +} + +int32_t IndexTTS2TextTokenizer::lang_to_id(const std::string & lang) { + const std::string normalized = lowercase_ascii(lang); + for (size_t i = 0; i < kLanguages.size(); ++i) { + if (normalized == kLanguages[i]) { + return static_cast(i); + } + } + for (size_t i = 0; i < kEmbeddingOnlyLanguages.size(); ++i) { + if (normalized == kEmbeddingOnlyLanguages[i]) { + return static_cast(kLanguages.size() + i); + } + } + return kCommonLangId; +} + IndexTTS2TextEncoding IndexTTS2TextTokenizer::encode_for_inference( const std::string & text, - int max_text_tokens_per_segment) const { + int max_text_tokens_per_segment, + const std::string & lang) const { if (max_text_tokens_per_segment <= 0) { throw std::runtime_error("IndexTTS2 max_text_tokens_per_segment must be positive"); } + if (variant_ == IndexTTS2Variant::kV2_5) { + return encode_for_inference_v2_5(text, max_text_tokens_per_segment, lang); + } + return encode_for_inference_v2(text, max_text_tokens_per_segment); +} + +IndexTTS2TextEncoding IndexTTS2TextTokenizer::encode_for_inference_v2( + const std::string & text, + int max_text_tokens_per_segment) const { IndexTTS2TextEncoding encoding; encoding.normalized_text = normalize_text(text); encoding.token_ids = engine::tokenizers::tokenize_sentencepiece(pieces_, encoding.normalized_text); @@ -279,6 +808,108 @@ IndexTTS2TextEncoding IndexTTS2TextTokenizer::encode_for_inference( return encoding; } +IndexTTS2TextEncoding IndexTTS2TextTokenizer::encode_for_inference_v2_5( + const std::string & text, + int max_text_tokens_per_segment, + const std::string & lang) const { + std::string resolved_lang = lowercase_ascii(lang); + if (resolved_lang.empty()) { + resolved_lang = contains_han(text) ? "zh" : "en"; + } + + std::string processed = text; + if (resolved_lang == "zh" || resolved_lang == "en") { + // Protect annotations from the normalizer, as the + // official TextNormalizer does inside normalize(). + auto protected_text = protect_pronunciation_annotations(processed); + protected_text.first = resolved_lang == "zh" + ? normalize_chinese(protected_text.first) + : normalize_english(protected_text.first); + processed = restore_pronunciation_annotations(std::move(protected_text.first), protected_text.second); + } + // ja/es/ar and other languages currently pass through without TN. + if (resolved_lang == "zh" || resolved_lang == "ja" || resolved_lang == "en") { + processed = lowercase_ascii(std::move(processed)); + } else if (resolved_lang == "es") { + processed = uppercase_ascii(std::move(processed)); + } + processed = apply_pronunciation_annotations(processed); + processed = uppercase_special_token_names(processed); + + const std::string lang_prefix = "<|" + resolved_lang + "|> "; + const auto prefix_tokens = static_cast(encode(lang_prefix).size()); + const int64_t capacity = assets_->config.gpt.max_text_tokens; + int64_t budget = std::min(max_text_tokens_per_segment, capacity - 2) - prefix_tokens; + budget = std::max(budget, 1); + + std::vector segments; + const auto token_len = [this](const std::string & value) { + return static_cast(encode(value).size()); + }; + if (token_len(processed) <= budget) { + segments.push_back(processed); + } else { + std::vector chunks; + for (const auto & [piece, atomic] : split_atomic_pieces(processed)) { + if (atomic) { + chunks.push_back(piece); + continue; + } + for (const auto & part : split_after_delimiters(piece)) { + if (token_len(part) <= budget) { + chunks.push_back(part); + continue; + } + std::string current; + for (size_t i = 0; i < part.size();) { + const size_t begin = i; + next_utf8_codepoint(part, i); + const std::string ch = part.substr(begin, i - begin); + if (!current.empty() && token_len(current + ch) > budget) { + chunks.push_back(std::move(current)); + current = ch; + } else { + current += ch; + } + } + if (!current.empty()) { + chunks.push_back(std::move(current)); + } + } + } + std::string current; + for (const auto & chunk : chunks) { + if (!current.empty() && token_len(current + chunk) > budget) { + segments.push_back(std::move(current)); + current = chunk; + } else { + current += chunk; + } + } + if (!current.empty()) { + segments.push_back(std::move(current)); + } + if (segments.empty()) { + segments.push_back(processed); + } + } + + IndexTTS2TextEncoding encoding; + encoding.lang = resolved_lang; + encoding.normalized_text = processed; + encoding.segments.reserve(segments.size()); + for (const auto & segment : segments) { + encoding.segments.push_back({segment}); + } + encoding.segment_token_ids.reserve(segments.size()); + for (const auto & segment : segments) { + std::vector ids = encode(lang_prefix + segment); + ids.push_back(kSegmentPadTokenId); + encoding.segment_token_ids.push_back(std::move(ids)); + } + return encoding; +} + int32_t IndexTTS2TextTokenizer::piece_to_id(const std::string & piece) const { const auto it = piece_to_id_.find(piece); if (it != piece_to_id_.end()) { diff --git a/src/models/index_tts2_5/assets.cpp b/src/models/index_tts2_5/assets.cpp deleted file mode 100644 index cb3d0b11..00000000 --- a/src/models/index_tts2_5/assets.cpp +++ /dev/null @@ -1,259 +0,0 @@ -#include "engine/models/index_tts2_5/assets.h" - -#include "engine/framework/model_spec/package.h" -#include "engine/framework/io/config.h" -#include "engine/framework/io/json.h" -#include "engine/framework/io/yaml.h" - -#include - -namespace engine::models::index_tts2_5 { -namespace { - -namespace json = engine::io::json; -namespace yaml = engine::io::yaml; - -IndexTTS25Config parse_config(const assets::ResourceBundle & resources) { - const auto document = resources.parse_flattened_yaml("config"); - IndexTTS25Config config; - config.version = yaml::optional_string(document, "version", config.version); - // The official IndexTTS-2.5 config_v2_5.yaml has no dataset section; these values - // are parsed for compatibility but not used at inference time. - if (const auto value = yaml::optional_int(document, "dataset.sample_rate")) { - config.dataset_sample_rate = *value; - } - config.dataset_squeeze = yaml::optional_bool(document, "dataset.squeeze", config.dataset_squeeze); - if (const auto value = yaml::optional_int(document, "dataset.mel.sample_rate")) { - config.dataset_mel_sample_rate = *value; - } - if (const auto value = yaml::optional_int(document, "dataset.mel.n_fft")) { - config.dataset_mel_n_fft = *value; - } - if (const auto value = yaml::optional_int(document, "dataset.mel.hop_length")) { - config.dataset_mel_hop_length = *value; - } - if (const auto value = yaml::optional_int(document, "dataset.mel.win_length")) { - config.dataset_mel_win_length = *value; - } - if (const auto value = yaml::optional_int(document, "dataset.mel.n_mels")) { - config.dataset_mel_n_mels = *value; - } - config.dataset_mel_fmin = yaml::optional_f32(document, "dataset.mel.mel_fmin", config.dataset_mel_fmin); - config.dataset_mel_normalize = yaml::optional_bool(document, "dataset.mel.normalize", config.dataset_mel_normalize); - - config.gpt.model_dim = yaml::require_i64(document, "gpt.model_dim"); - config.gpt.max_mel_tokens = yaml::require_i64(document, "gpt.max_mel_tokens"); - config.gpt.max_text_tokens = yaml::require_i64(document, "gpt.max_text_tokens"); - config.gpt.heads = yaml::require_i64(document, "gpt.heads"); - config.gpt.use_mel_codes_as_input = yaml::optional_bool(document, "gpt.use_mel_codes_as_input", config.gpt.use_mel_codes_as_input); - config.gpt.mel_length_compression = yaml::require_i64(document, "gpt.mel_length_compression"); - config.gpt.layers = yaml::require_i64(document, "gpt.layers"); - config.gpt.number_text_tokens = yaml::require_i64(document, "gpt.number_text_tokens"); - config.gpt.number_mel_codes = yaml::require_i64(document, "gpt.number_mel_codes"); - config.gpt.start_mel_token = yaml::require_i64(document, "gpt.start_mel_token"); - config.gpt.stop_mel_token = yaml::require_i64(document, "gpt.stop_mel_token"); - config.gpt.start_text_token = yaml::require_i64(document, "gpt.start_text_token"); - config.gpt.stop_text_token = yaml::require_i64(document, "gpt.stop_text_token"); - config.gpt.train_solo_embeddings = yaml::optional_bool(document, "gpt.train_solo_embeddings", config.gpt.train_solo_embeddings); - config.gpt.condition_type = yaml::require_string(document, "gpt.condition_type"); - config.gpt.condition_output_size = yaml::require_i64(document, "gpt.condition_module.output_size"); - config.gpt.condition_linear_units = yaml::require_i64(document, "gpt.condition_module.linear_units"); - config.gpt.condition_attention_heads = yaml::require_i64(document, "gpt.condition_module.attention_heads"); - config.gpt.condition_num_blocks = yaml::require_i64(document, "gpt.condition_module.num_blocks"); - config.gpt.condition_input_layer = yaml::require_string(document, "gpt.condition_module.input_layer"); - config.gpt.condition_perceiver_mult = yaml::require_i64(document, "gpt.condition_module.perceiver_mult"); - config.gpt.emo_condition_output_size = yaml::require_i64(document, "gpt.emo_condition_module.output_size"); - config.gpt.emo_condition_linear_units = yaml::require_i64(document, "gpt.emo_condition_module.linear_units"); - config.gpt.emo_condition_attention_heads = yaml::require_i64(document, "gpt.emo_condition_module.attention_heads"); - config.gpt.emo_condition_num_blocks = yaml::require_i64(document, "gpt.emo_condition_module.num_blocks"); - config.gpt.emo_condition_input_layer = yaml::require_string(document, "gpt.emo_condition_module.input_layer"); - config.gpt.emo_condition_perceiver_mult = yaml::require_i64(document, "gpt.emo_condition_module.perceiver_mult"); - - config.semantic_codec.codebook_size = yaml::require_i64(document, "semantic_codec.codebook_size"); - config.semantic_codec.hidden_size = yaml::require_i64(document, "semantic_codec.hidden_size"); - config.semantic_codec.codebook_dim = yaml::require_i64(document, "semantic_codec.codebook_dim"); - config.semantic_codec.vocos_dim = yaml::require_i64(document, "semantic_codec.vocos_dim"); - config.semantic_codec.vocos_intermediate_dim = yaml::require_i64(document, "semantic_codec.vocos_intermediate_dim"); - config.semantic_codec.vocos_num_layers = yaml::require_i64(document, "semantic_codec.vocos_num_layers"); - - config.s2mel.sample_rate = static_cast(yaml::require_i64(document, "s2mel.preprocess_params.sr")); - config.s2mel.n_fft = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.n_fft"); - config.s2mel.win_length = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.win_length"); - config.s2mel.hop_length = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.hop_length"); - config.s2mel.n_mels = yaml::require_i64(document, "s2mel.preprocess_params.spect_params.n_mels"); - config.s2mel.fmin = yaml::optional_f32(document, "s2mel.preprocess_params.spect_params.fmin", config.s2mel.fmin); - config.s2mel.fmax = yaml::optional_nullable_f32(document, "s2mel.preprocess_params.spect_params.fmax"); - config.s2mel.dit_type = yaml::require_string(document, "s2mel.dit_type"); - config.s2mel.reg_loss_type = yaml::require_string(document, "s2mel.reg_loss_type"); - config.s2mel.style_dim = yaml::require_i64(document, "s2mel.style_encoder.dim"); - config.s2mel.length_regulator_channels = yaml::require_i64(document, "s2mel.length_regulator.channels"); - config.s2mel.length_regulator_is_discrete = yaml::optional_bool(document, "s2mel.length_regulator.is_discrete", config.s2mel.length_regulator_is_discrete); - config.s2mel.length_regulator_in_channels = yaml::require_i64(document, "s2mel.length_regulator.in_channels"); - config.s2mel.length_regulator_content_codebook_size = yaml::require_i64(document, "s2mel.length_regulator.content_codebook_size"); - config.s2mel.length_regulator_sampling_ratios = yaml::require_list_i64(document, "s2mel.length_regulator.sampling_ratios"); - config.s2mel.length_regulator_vector_quantize = yaml::optional_bool(document, "s2mel.length_regulator.vector_quantize", config.s2mel.length_regulator_vector_quantize); - config.s2mel.length_regulator_n_codebooks = yaml::require_i64(document, "s2mel.length_regulator.n_codebooks"); - config.s2mel.length_regulator_quantizer_dropout = yaml::optional_f32(document, "s2mel.length_regulator.quantizer_dropout", config.s2mel.length_regulator_quantizer_dropout); - config.s2mel.length_regulator_f0_condition = yaml::optional_bool(document, "s2mel.length_regulator.f0_condition", config.s2mel.length_regulator_f0_condition); - config.s2mel.length_regulator_n_f0_bins = yaml::require_i64(document, "s2mel.length_regulator.n_f0_bins"); - config.s2mel.dit_hidden_dim = yaml::require_i64(document, "s2mel.DiT.hidden_dim"); - config.s2mel.dit_num_heads = yaml::require_i64(document, "s2mel.DiT.num_heads"); - config.s2mel.dit_depth = yaml::require_i64(document, "s2mel.DiT.depth"); - config.s2mel.dit_class_dropout_prob = yaml::optional_f32(document, "s2mel.DiT.class_dropout_prob", config.s2mel.dit_class_dropout_prob); - config.s2mel.dit_block_size = yaml::require_i64(document, "s2mel.DiT.block_size"); - config.s2mel.dit_in_channels = yaml::require_i64(document, "s2mel.DiT.in_channels"); - config.s2mel.dit_style_condition = yaml::optional_bool(document, "s2mel.DiT.style_condition", config.s2mel.dit_style_condition); - config.s2mel.dit_final_layer_type = yaml::require_string(document, "s2mel.DiT.final_layer_type"); - config.s2mel.dit_target = yaml::require_string(document, "s2mel.DiT.target"); - config.s2mel.dit_content_dim = yaml::require_i64(document, "s2mel.DiT.content_dim"); - config.s2mel.dit_content_codebook_size = yaml::require_i64(document, "s2mel.DiT.content_codebook_size"); - config.s2mel.dit_content_type = yaml::require_string(document, "s2mel.DiT.content_type"); - config.s2mel.dit_f0_condition = yaml::optional_bool(document, "s2mel.DiT.f0_condition", config.s2mel.dit_f0_condition); - config.s2mel.dit_n_f0_bins = yaml::require_i64(document, "s2mel.DiT.n_f0_bins"); - config.s2mel.dit_content_codebooks = yaml::require_i64(document, "s2mel.DiT.content_codebooks"); - config.s2mel.dit_is_causal = yaml::optional_bool(document, "s2mel.DiT.is_causal", config.s2mel.dit_is_causal); - config.s2mel.dit_long_skip_connection = yaml::optional_bool(document, "s2mel.DiT.long_skip_connection", config.s2mel.dit_long_skip_connection); - config.s2mel.dit_zero_prompt_speech_token = yaml::optional_bool(document, "s2mel.DiT.zero_prompt_speech_token", config.s2mel.dit_zero_prompt_speech_token); - config.s2mel.dit_time_as_token = yaml::optional_bool(document, "s2mel.DiT.time_as_token", config.s2mel.dit_time_as_token); - config.s2mel.dit_style_as_token = yaml::optional_bool(document, "s2mel.DiT.style_as_token", config.s2mel.dit_style_as_token); - config.s2mel.dit_uvit_skip_connection = yaml::optional_bool(document, "s2mel.DiT.uvit_skip_connection", config.s2mel.dit_uvit_skip_connection); - config.s2mel.dit_add_resblock_in_transformer = yaml::optional_bool(document, "s2mel.DiT.add_resblock_in_transformer", config.s2mel.dit_add_resblock_in_transformer); - config.s2mel.wavenet_hidden_dim = yaml::require_i64(document, "s2mel.wavenet.hidden_dim"); - config.s2mel.wavenet_num_layers = yaml::require_i64(document, "s2mel.wavenet.num_layers"); - config.s2mel.wavenet_kernel_size = yaml::require_i64(document, "s2mel.wavenet.kernel_size"); - config.s2mel.wavenet_dilation_rate = yaml::require_i64(document, "s2mel.wavenet.dilation_rate"); - config.s2mel.wavenet_dropout = yaml::optional_f32(document, "s2mel.wavenet.p_dropout", config.s2mel.wavenet_dropout); - config.s2mel.wavenet_style_condition = yaml::optional_bool(document, "s2mel.wavenet.style_condition", config.s2mel.wavenet_style_condition); - - config.emo_num = yaml::require_list_i64(document, "emo_num"); - return config; -} - -void validate_qwen_emotion_config(const assets::ResourceBundle & resources) { - const auto root = resources.parse_json("qwen_emotion_config"); - if (json::optional_string(root, "model_type", "") != "qwen3") { - throw std::runtime_error("IndexTTS2.5 Qwen emotion model must have model_type=qwen3"); - } - if (json::optional_i64(root, "hidden_size", 0) != 1024 || - json::optional_i64(root, "num_hidden_layers", 0) != 28 || - json::optional_i64(root, "num_attention_heads", 0) != 16) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion model config does not match expected 0.6B architecture"); - } -} - -void validate_config(const IndexTTS25Config & config, const assets::ResourceBundle & resources) { - engine::io::require_positive(config.dataset_sample_rate, "dataset.sample_rate"); - engine::io::require_positive(config.dataset_mel_sample_rate, "dataset.mel.sample_rate"); - engine::io::require_positive(config.dataset_mel_n_fft, "dataset.mel.n_fft"); - engine::io::require_positive(config.gpt.model_dim, "gpt.model_dim"); - engine::io::require_positive(config.gpt.layers, "gpt.layers"); - engine::io::require_divisible(config.gpt.model_dim, config.gpt.heads, "gpt.model_dim / gpt.heads"); - engine::io::require_positive(config.semantic_codec.codebook_size, "semantic_codec.codebook_size"); - engine::io::require_positive(config.s2mel.sample_rate, "s2mel.sample_rate"); - engine::io::require_positive(config.s2mel.n_mels, "s2mel.n_mels"); - engine::io::require_positive(config.s2mel.dit_hidden_dim, "s2mel.DiT.hidden_dim"); - engine::io::require_divisible(config.s2mel.dit_hidden_dim, config.s2mel.dit_num_heads, "s2mel.DiT.hidden_dim / num_heads"); - engine::io::require_nonnegative(config.s2mel.length_regulator_quantizer_dropout, "length_regulator.quantizer_dropout"); - engine::io::require_positive(config.s2mel.wavenet_dropout + 1.0F, "wavenet.p_dropout"); - if (config.emo_num.empty()) { - throw std::runtime_error("IndexTTS2.5 config emo_num must not be empty"); - } - validate_qwen_emotion_config(resources); -} - -void validate_gpt_weights(const IndexTTS25Config & config, const assets::TensorSource & source) { - assets::require_tensor_shape(source, "text_embedding.weight", {config.gpt.number_text_tokens + 1, config.gpt.model_dim}); - assets::require_tensor_shape(source, "mel_embedding.weight", {config.gpt.number_mel_codes, config.gpt.model_dim}); - assets::require_tensor_shape(source, "gpt.h.0.attn.c_attn.weight", {config.gpt.model_dim, config.gpt.model_dim * 3}); - assets::require_tensor_shape(source, "gpt.h.0.attn.c_proj.weight", {config.gpt.model_dim, config.gpt.model_dim}); - assets::require_tensor_shape(source, "gpt.h.0.mlp.c_fc.weight", {config.gpt.model_dim, config.gpt.model_dim * 4}); - assets::require_tensor_shape(source, "gpt.h.0.mlp.c_proj.weight", {config.gpt.model_dim * 4, config.gpt.model_dim}); - assets::require_tensor_shape(source, "spk_emb_proj.weight", {config.gpt.model_dim, config.s2mel.style_dim}); - assets::require_tensor_shape(source, "lang_embedding.weight", {kIndexTTS25LangEmbeddingRows, config.gpt.model_dim}); - assets::require_tensor_shape(source, "emo_conditioning_encoder.after_norm.weight", {config.gpt.emo_condition_output_size}); -} - -void validate_s2mel_weights(const IndexTTS25Config & config, const assets::TensorSource & source) { - assets::require_tensor_shape(source, "gpt_layer.0.weight", {256, config.gpt.model_dim}); - assets::require_tensor_shape(source, "gpt_layer.2.weight", {config.s2mel.length_regulator_in_channels, 128}); - assets::require_tensor_shape(source, "length_regulator.model.0.weight", {config.s2mel.length_regulator_channels, config.s2mel.length_regulator_channels, 3}); - assets::require_tensor_shape(source, "cfm.estimator.x_embedder.weight_v", {config.s2mel.dit_hidden_dim, config.s2mel.dit_in_channels}); - assets::require_tensor_shape(source, "cfm.estimator.transformer.layers.0.attention.wqkv.weight", {config.s2mel.dit_hidden_dim * 3, config.s2mel.dit_hidden_dim}); - assets::require_tensor_shape(source, "cfm.estimator.final_layer.adaLN_modulation.1.weight", {config.s2mel.dit_hidden_dim * 2, config.s2mel.dit_hidden_dim}); -} - -void validate_matrix_weights( - const IndexTTS25Config & config, - const assets::TensorSource & speaker_matrix, - const assets::TensorSource & emotion_matrix) { - int64_t total = 0; - for (const int64_t count : config.emo_num) { - engine::io::require_positive(count, "emo_num item"); - total += count; - } - assets::require_tensor_shape(speaker_matrix, "tensor", {total, config.s2mel.style_dim}); - assets::require_tensor_shape(emotion_matrix, "tensor", {total, config.gpt.model_dim}); -} - -void validate_w2v_stats(const IndexTTS25Config & config, const assets::TensorSource & source) { - assets::require_tensor_shape(source, "mean", {config.semantic_codec.hidden_size}); - assets::require_tensor_shape(source, "var", {config.semantic_codec.hidden_size}); -} - -void validate_w2v_weights(const assets::TensorSource & source) { - assets::require_tensor_shape(source, "feature_projection.projection.weight", {1024, 160}); - assets::require_tensor_shape(source, "encoder.layers.0.self_attn.linear_k.weight", {1024, 1024}); - assets::require_tensor_shape(source, "encoder.layers.0.conv_module.depthwise_conv.weight", {1024, 1, 31}); -} - -void validate_semantic_codec_weights(const IndexTTS25Config & config, const assets::TensorSource & source) { - assets::require_tensor_shape(source, "quantizer.quantizers.0.codebook.weight", {config.semantic_codec.codebook_size, config.semantic_codec.codebook_dim}); - assets::require_tensor_shape(source, "encoder.1.weight", {config.semantic_codec.hidden_size, config.semantic_codec.vocos_dim}); - assets::require_tensor_shape(source, "decoder.1.weight", {config.semantic_codec.hidden_size, config.semantic_codec.vocos_dim}); - assets::require_tensor_shape(source, "up.weight", {config.semantic_codec.hidden_size, config.semantic_codec.hidden_size, 3}); -} - -void validate_qwen_weights(const assets::TensorSource & source) { - assets::require_tensor_shape(source, "model.embed_tokens.weight", {151936, 1024}); - assets::require_tensor_shape(source, "model.layers.0.self_attn.q_proj.weight", {2048, 1024}); - assets::require_tensor_shape(source, "model.layers.0.self_attn.k_proj.weight", {1024, 1024}); - assets::require_tensor_shape(source, "model.layers.0.mlp.gate_proj.weight", {3072, 1024}); - assets::require_tensor_shape(source, "model.norm.weight", {1024}); -} - -void validate_weight_anchors(const IndexTTS25Assets & assets) { - validate_gpt_weights(assets.config, *assets.gpt_weights); - validate_s2mel_weights(assets.config, *assets.s2mel_weights); - validate_matrix_weights(assets.config, *assets.speaker_matrix, *assets.emotion_matrix); - validate_w2v_stats(assets.config, *assets.wav2vec2bert_stats); - validate_w2v_weights(*assets.wav2vec2bert_weights); - validate_semantic_codec_weights(assets.config, *assets.semantic_codec_weights); - validate_qwen_weights(*assets.qwen_emotion_weights); -} - -} // namespace - -std::shared_ptr load_index_tts2_5_assets(const std::filesystem::path & model_path) { - auto assets = std::make_shared(); - assets->resources = engine::model_spec::load_resource_bundle( - model_path, - engine::model_spec::default_spec_path("index_tts2_5")); - assets->config = parse_config(assets->resources); - validate_config(assets->config, assets->resources); - - assets->gpt_weights = assets->resources.open_tensor_source("gpt"); - assets->s2mel_weights = assets->resources.open_tensor_source("s2mel"); - assets->speaker_matrix = assets->resources.open_tensor_source("speaker_matrix"); - assets->emotion_matrix = assets->resources.open_tensor_source("emotion_matrix"); - assets->wav2vec2bert_stats = assets->resources.open_tensor_source("wav2vec2bert_stats"); - assets->wav2vec2bert_weights = assets->resources.open_tensor_source("wav2vec2bert"); - assets->semantic_codec_weights = assets->resources.open_tensor_source("semantic_codec"); - assets->campplus_weights = assets->resources.open_tensor_source("campplus"); - assets->bigvgan_weights = assets->resources.open_tensor_source("bigvgan"); - assets->qwen_emotion_weights = assets->resources.open_tensor_source("qwen_emotion"); - - validate_weight_anchors(*assets); - return assets; -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/audio_features.cpp b/src/models/index_tts2_5/audio_features.cpp deleted file mode 100644 index 2e585500..00000000 --- a/src/models/index_tts2_5/audio_features.cpp +++ /dev/null @@ -1,547 +0,0 @@ -#include "engine/models/index_tts2_5/audio_features.h" - -#include "engine/framework/audio/conversion.h" -#include "engine/framework/audio/dsp.h" -#include "engine/framework/audio/resampling.h" -#include "engine/framework/audio/waveform_ops.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -struct MelFilterbankKey { - int64_t sample_rate = 0; - int64_t n_fft = 0; - int64_t num_mels = 0; - float fmin = 0.0F; - float fmax = 0.0F; - - bool operator==(const MelFilterbankKey & other) const noexcept { - return sample_rate == other.sample_rate && n_fft == other.n_fft && num_mels == other.num_mels && - fmin == other.fmin && fmax == other.fmax; - } -}; - -struct MelFilterbankKeyHash { - size_t operator()(const MelFilterbankKey & key) const noexcept { - size_t seed = std::hash{}(key.sample_rate); - seed ^= std::hash{}(key.n_fft) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - seed ^= std::hash{}(key.num_mels) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - seed ^= std::hash{}(key.fmin) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - seed ^= std::hash{}(key.fmax) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - return seed; - } -}; - -struct KaldiFilterbankKey { - int64_t sample_rate = 0; - int64_t padded_window_size = 0; - int64_t num_mels = 0; - float low_freq = 0.0F; - float high_freq = 0.0F; - - bool operator==(const KaldiFilterbankKey & other) const noexcept { - return sample_rate == other.sample_rate && - padded_window_size == other.padded_window_size && - num_mels == other.num_mels && - low_freq == other.low_freq && - high_freq == other.high_freq; - } -}; - -struct KaldiFilterbankKeyHash { - size_t operator()(const KaldiFilterbankKey & key) const noexcept { - size_t seed = 0; - seed ^= std::hash{}(key.sample_rate) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - seed ^= std::hash{}(key.padded_window_size) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - seed ^= std::hash{}(key.num_mels) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - seed ^= std::hash{}(key.low_freq) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - seed ^= std::hash{}(key.high_freq) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - return seed; - } -}; - -std::vector make_povey_window(int64_t window_size) { - std::vector window(static_cast(window_size), 0.0F); - constexpr float kPi = 3.14159265358979323846F; - for (int64_t i = 0; i < window_size; ++i) { - const float hann = - 0.5F - 0.5F * std::cos(2.0F * kPi * static_cast(i) / static_cast(window_size - 1)); - window[static_cast(i)] = std::pow(hann, 0.85F); - } - return window; -} - -std::vector make_kaldi_mel_filterbank( - int64_t sample_rate, - int64_t n_fft, - int64_t n_mels, - float low_freq, - float high_freq) { - const int64_t num_fft_bins = n_fft / 2 + 1; - const float nyquist = 0.5F * static_cast(sample_rate); - if (high_freq <= 0.0F) { - high_freq += nyquist; - } - const float fft_bin_width = static_cast(sample_rate) / static_cast(n_fft); - const float mel_low = 1127.0F * std::log(1.0F + low_freq / 700.0F); - const float mel_high = 1127.0F * std::log(1.0F + high_freq / 700.0F); - const float mel_delta = (mel_high - mel_low) / static_cast(n_mels + 1); - - std::vector filterbank(static_cast(n_mels * num_fft_bins), 0.0F); - for (int64_t mel_bin = 0; mel_bin < n_mels; ++mel_bin) { - const float left_mel = mel_low + static_cast(mel_bin) * mel_delta; - const float center_mel = mel_low + static_cast(mel_bin + 1) * mel_delta; - const float right_mel = mel_low + static_cast(mel_bin + 2) * mel_delta; - for (int64_t fft_bin = 0; fft_bin < num_fft_bins; ++fft_bin) { - const float freq = fft_bin_width * static_cast(fft_bin); - const float mel = 1127.0F * std::log(1.0F + freq / 700.0F); - const float up_slope = (mel - left_mel) / std::max(center_mel - left_mel, 1.0e-12F); - const float down_slope = (right_mel - mel) / std::max(right_mel - center_mel, 1.0e-12F); - filterbank[static_cast(mel_bin * num_fft_bins + fft_bin)] = - std::max(0.0F, std::min(up_slope, down_slope)); - } - } - return filterbank; -} - -const std::vector & cached_mel_filterbank(const IndexTTS25S2MelConfig & config) { - static std::mutex mutex; - static std::unordered_map, MelFilterbankKeyHash> cache; - const float fmax = config.fmax.value_or(static_cast(config.sample_rate) / 2.0F); - const MelFilterbankKey key{config.sample_rate, config.n_fft, config.n_mels, config.fmin, fmax}; - std::lock_guard lock(mutex); - const auto it = cache.find(key); - if (it != cache.end()) { - return it->second; - } - const auto filterbank = engine::audio::MelFilterbank().build({ - config.sample_rate, - config.n_fft, - config.n_mels, - config.fmin, - fmax, - true, - }); - return cache.emplace( - key, - filterbank.values).first->second; -} - -const std::vector & cached_povey_window(int64_t window_size) { - static std::mutex mutex; - static std::unordered_map> cache; - std::lock_guard lock(mutex); - const auto it = cache.find(window_size); - if (it != cache.end()) { - return it->second; - } - return cache.emplace(window_size, make_povey_window(window_size)).first->second; -} - -const std::vector & cached_kaldi_mel_filterbank( - int64_t sample_rate, - int64_t padded_window_size, - int64_t num_mels, - float low_freq, - float high_freq) { - static std::mutex mutex; - static std::unordered_map, KaldiFilterbankKeyHash> cache; - const KaldiFilterbankKey key{sample_rate, padded_window_size, num_mels, low_freq, high_freq}; - std::lock_guard lock(mutex); - const auto it = cache.find(key); - if (it != cache.end()) { - return it->second; - } - return cache.emplace( - key, - make_kaldi_mel_filterbank(sample_rate, padded_window_size, num_mels, low_freq, high_freq)).first->second; -} - -struct RealDftTables { - std::vector cos; - std::vector sin; -}; - -const RealDftTables & cached_real_dft_tables_512() { - static const RealDftTables tables = [] { - constexpr int64_t kFft = 512; - constexpr int64_t kFreqBins = kFft / 2 + 1; - constexpr double kPi = 3.14159265358979323846264338327950288; - RealDftTables out; - out.cos.resize(static_cast(kFreqBins * kFft)); - out.sin.resize(static_cast(kFreqBins * kFft)); - for (int64_t freq = 0; freq < kFreqBins; ++freq) { - for (int64_t n = 0; n < kFft; ++n) { - const double angle = 2.0 * kPi * static_cast(freq * n) / static_cast(kFft); - out.cos[static_cast(freq * kFft + n)] = std::cos(angle); - out.sin[static_cast(freq * kFft + n)] = std::sin(angle); - } - } - return out; - }(); - return tables; -} - -std::vector require_mono_samples(const std::vector & samples, int channels) { - if (channels <= 0) { - throw std::runtime_error("IndexTTS2.5 audio channel count must be positive"); - } - if (samples.empty()) { - throw std::runtime_error("IndexTTS2.5 audio must not be empty"); - } - if (samples.size() % static_cast(channels) != 0) { - throw std::runtime_error("IndexTTS2.5 audio sample count must be divisible by channel count"); - } - if (channels == 1) { - return samples; - } - return engine::audio::mixdown_interleaved_to_mono_average(samples, channels); -} - -std::vector resample_mono(const std::vector & input, int input_sample_rate, int output_sample_rate) { - if (input_sample_rate <= 0 || output_sample_rate <= 0) { - throw std::runtime_error("IndexTTS2.5 resampling requires positive sample rates"); - } - if (input_sample_rate == output_sample_rate || input.empty()) { - return input; - } - engine::audio::TorchaudioSincHannResampleOptions options; - options.kernel_mode = engine::audio::TorchaudioSincHannKernelMode::Float32ComputationStoredAsFloat32; - options.accumulation = engine::audio::TorchaudioSincHannAccumulation::Float32; - return engine::audio::resample_mono_torchaudio_sinc_hann(input, input_sample_rate, output_sample_rate, options); -} - -std::vector resample_mono_librosa(const std::vector & input, int input_sample_rate, int output_sample_rate) { - if (input_sample_rate <= 0 || output_sample_rate <= 0) { - throw std::runtime_error("IndexTTS2.5 librosa-style resampling requires positive sample rates"); - } - if (input_sample_rate == output_sample_rate || input.empty()) { - return input; - } - engine::audio::SoxrResampleOptions options; - options.profile = engine::audio::SoxrResampleProfile::ExplicitFloat32Runtime; - options.output_length_policy = engine::audio::SoxrOutputLengthPolicy::ExactExpected; - options.require_full_input = true; - if (auto output = engine::audio::try_resample_mono_soxr(input, input_sample_rate, output_sample_rate, options)) { - return *output; - } - return engine::audio::resample_mono_torchaudio_sinc_hann(input, input_sample_rate, output_sample_rate); -} - -} // namespace - -IndexTTS25MelOutput compute_index_tts2_5_mel_spectrogram( - const std::vector & waveform, - const IndexTTS25S2MelConfig & config, - size_t threads) { - if (config.sample_rate <= 0 || config.n_fft <= 0 || config.win_length <= 0 || - config.hop_length <= 0 || config.n_mels <= 0) { - throw std::runtime_error("IndexTTS2.5 mel spectrogram config is invalid"); - } - if (waveform.empty()) { - throw std::runtime_error("IndexTTS2.5 mel spectrogram requires non-empty waveform"); - } - - const int64_t pad = (config.n_fft - config.hop_length) / 2; - const auto padded = engine::audio::reflect_pad_samples(waveform, pad, pad); - const engine::audio::STFTConfig stft_config{ - config.n_fft, - config.hop_length, - config.win_length, - false, - engine::audio::STFTPadMode::Reflect, - engine::audio::STFTFamily::Kokoro, - }; - const auto & window = engine::audio::get_cached_stft_window(stft_config); - auto magnitude = engine::audio::STFT().compute_magnitude( - padded, - window, - 1, - static_cast(padded.size()), - stft_config, - threads); - for (float & value : magnitude.values) { - value = std::sqrt(value * value + 1.0e-9F); - } - - const auto & filterbank = cached_mel_filterbank(config); - const int64_t freq_bins = magnitude.shape[1]; - const int64_t frames = magnitude.shape[2]; - if (static_cast(filterbank.size()) != config.n_mels * freq_bins) { - throw std::runtime_error("IndexTTS2.5 mel filterbank shape mismatch"); - } - - IndexTTS25MelOutput output; - output.channels = config.n_mels; - output.frames = frames; - output.values.assign(static_cast(config.n_mels * frames), 0.0F); -#ifdef _OPENMP -#pragma omp parallel for collapse(2) if(config.n_mels * frames >= 4096) -#endif - for (int64_t mel = 0; mel < config.n_mels; ++mel) { - for (int64_t frame = 0; frame < frames; ++frame) { - float sum = 0.0F; - for (int64_t freq = 0; freq < freq_bins; ++freq) { - sum += filterbank[static_cast(mel * freq_bins + freq)] * - magnitude.values[static_cast(freq * frames + frame)]; - } - output.values[static_cast(mel * frames + frame)] = std::log(std::max(sum, 1.0e-5F)); - } - } - return output; -} - -IndexTTS25FbankOutput compute_index_tts2_5_campplus_fbank_16k(const std::vector & waveform_16k) { - constexpr int64_t kSampleRate = 16000; - constexpr int64_t kWindowSize = 400; - constexpr int64_t kWindowShift = 160; - constexpr int64_t kPaddedWindowSize = 512; - constexpr int64_t kNumMels = 80; - constexpr float kLowFreq = 20.0F; - constexpr float kHighFreq = 0.0F; - constexpr float kPreemphasis = 0.97F; - constexpr float kEpsilon = std::numeric_limits::epsilon(); - - if (static_cast(waveform_16k.size()) < kWindowSize) { - throw std::runtime_error("IndexTTS2.5 CAMPPlus fbank requires at least one 25 ms frame"); - } - - const int64_t frames = 1 + (static_cast(waveform_16k.size()) - kWindowSize) / kWindowShift; - const auto & window = cached_povey_window(kWindowSize); - const auto & mel_filterbank = cached_kaldi_mel_filterbank( - kSampleRate, - kPaddedWindowSize, - kNumMels, - kLowFreq, - kHighFreq); - - std::vector frame(static_cast(kWindowSize), 0.0F); - std::vector stft_batch(static_cast(frames * kPaddedWindowSize), 0.0F); - for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { - const int64_t start = frame_index * kWindowShift; - float mean = 0.0F; - for (int64_t i = 0; i < kWindowSize; ++i) { - const float sample = waveform_16k[static_cast(start + i)]; - frame[static_cast(i)] = sample; - mean += sample; - } - mean /= static_cast(kWindowSize); - for (int64_t i = 0; i < kWindowSize; ++i) { - frame[static_cast(i)] -= mean; - } - for (int64_t i = kWindowSize - 1; i > 0; --i) { - frame[static_cast(i)] -= kPreemphasis * frame[static_cast(i - 1)]; - } - frame[0] -= kPreemphasis * frame[0]; - for (int64_t i = 0; i < kWindowSize; ++i) { - stft_batch[static_cast(frame_index * kPaddedWindowSize + i)] = - frame[static_cast(i)] * window[static_cast(i)]; - } - } - - std::vector stft_window(static_cast(kPaddedWindowSize), 1.0F); - const engine::audio::STFTConfig stft_config{ - kPaddedWindowSize, - kPaddedWindowSize, - kPaddedWindowSize, - false, - engine::audio::STFTPadMode::Constant, - engine::audio::STFTFamily::Default, - }; - const auto magnitude = engine::audio::STFT().compute_magnitude( - stft_batch, - stft_window, - frames, - kPaddedWindowSize, - stft_config); - - const int64_t freq_bins = (kPaddedWindowSize / 2) + 1; - IndexTTS25FbankOutput output; - output.frames = frames; - output.dims = kNumMels; - output.values.assign(static_cast(frames * kNumMels), 0.0F); - for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { - for (int64_t mel_bin = 0; mel_bin < kNumMels; ++mel_bin) { - float energy = 0.0F; - for (int64_t freq = 0; freq < freq_bins; ++freq) { - const float mag = magnitude.values[static_cast(frame_index * freq_bins + freq)]; - energy += (mag * mag) * mel_filterbank[static_cast(mel_bin * freq_bins + freq)]; - } - output.values[static_cast(frame_index * kNumMels + mel_bin)] = - std::log(std::max(energy, kEpsilon)); - } - } - - for (int64_t mel_bin = 0; mel_bin < kNumMels; ++mel_bin) { - float mean = 0.0F; - for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { - mean += output.values[static_cast(frame_index * kNumMels + mel_bin)]; - } - mean /= static_cast(frames); - for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { - output.values[static_cast(frame_index * kNumMels + mel_bin)] -= mean; - } - } - return output; -} - -IndexTTS25SemanticFeatureOutput compute_index_tts2_5_semantic_features_16k(const std::vector & waveform_16k) { - constexpr int64_t kSampleRate = 16000; - constexpr int64_t kWindowSize = 400; - constexpr int64_t kWindowShift = 160; - constexpr int64_t kFftSize = 512; - constexpr int64_t kFreqBins = kFftSize / 2 + 1; - constexpr int64_t kNumMels = 80; - constexpr float kLowFreq = 20.0F; - constexpr float kHighFreq = 8000.0F; - constexpr float kPreemphasis = 0.97F; - constexpr double kInputScale = 32768.0; - constexpr double kMelFloor = 1.192092955078125e-07; - - if (static_cast(waveform_16k.size()) < kWindowSize) { - throw std::runtime_error("IndexTTS2.5 semantic fbank requires at least one 25 ms frame"); - } - - const int64_t frames = 1 + (static_cast(waveform_16k.size()) - kWindowSize) / kWindowShift; - const auto & window = cached_povey_window(kWindowSize); - const auto & mel_filterbank = cached_kaldi_mel_filterbank( - kSampleRate, - kFftSize, - kNumMels, - kLowFreq, - kHighFreq); - const auto & dft = cached_real_dft_tables_512(); - - IndexTTS25FbankOutput fbank; - fbank.frames = frames; - fbank.dims = kNumMels; - fbank.values.assign(static_cast(frames * kNumMels), 0.0F); - -#ifdef _OPENMP -#pragma omp parallel for if(frames >= 8) -#endif - for (int64_t frame_index = 0; frame_index < frames; ++frame_index) { - const int64_t start = frame_index * kWindowShift; - double buffer[kFftSize] = {}; - double mean = 0.0; - for (int64_t i = 0; i < kWindowSize; ++i) { - const double sample = static_cast(waveform_16k[static_cast(start + i)]) * kInputScale; - buffer[i] = sample; - mean += sample; - } - mean /= static_cast(kWindowSize); - for (int64_t i = 0; i < kWindowSize; ++i) { - buffer[i] -= mean; - } - for (int64_t i = kWindowSize - 1; i > 0; --i) { - buffer[i] -= kPreemphasis * buffer[i - 1]; - } - buffer[0] *= (1.0 - kPreemphasis); - for (int64_t i = 0; i < kWindowSize; ++i) { - buffer[i] *= static_cast(window[static_cast(i)]); - } - - double power[kFreqBins] = {}; - for (int64_t freq = 0; freq < kFreqBins; ++freq) { - double re = 0.0; - double im = 0.0; - const size_t table_offset = static_cast(freq * kFftSize); - for (int64_t n = 0; n < kFftSize; ++n) { - const double sample = buffer[n]; - re += sample * dft.cos[table_offset + static_cast(n)]; - im -= sample * dft.sin[table_offset + static_cast(n)]; - } - power[freq] = re * re + im * im; - } - - for (int64_t mel_bin = 0; mel_bin < kNumMels; ++mel_bin) { - double energy = 0.0; - for (int64_t freq = 0; freq < kFreqBins; ++freq) { - energy += static_cast(mel_filterbank[static_cast(mel_bin * kFreqBins + freq)]) * - power[freq]; - } - fbank.values[static_cast(frame_index * kNumMels + mel_bin)] = - static_cast(std::log(std::max(energy, kMelFloor))); - } - } - - for (int64_t mel_bin = 0; mel_bin < fbank.dims; ++mel_bin) { - double mean = 0.0; - for (int64_t frame = 0; frame < fbank.frames; ++frame) { - mean += static_cast(fbank.values[static_cast(frame * fbank.dims + mel_bin)]); - } - mean /= static_cast(fbank.frames); - - double variance = 0.0; - for (int64_t frame = 0; frame < fbank.frames; ++frame) { - const double diff = - static_cast(fbank.values[static_cast(frame * fbank.dims + mel_bin)]) - mean; - variance += diff * diff; - } - variance = fbank.frames > 1 ? variance / static_cast(fbank.frames - 1) : 0.0; - const float scale = static_cast(1.0 / std::sqrt(variance + 1.0e-7)); - for (int64_t frame = 0; frame < fbank.frames; ++frame) { - float & value = fbank.values[static_cast(frame * fbank.dims + mel_bin)]; - value = (value - static_cast(mean)) * scale; - } - } - - const int64_t padded_frames = fbank.frames + (fbank.frames % 2); - std::vector padded(static_cast(padded_frames * fbank.dims), 1.0F); - std::copy(fbank.values.begin(), fbank.values.end(), padded.begin()); - - IndexTTS25SemanticFeatureOutput output; - output.frames = padded_frames / 2; - output.dims = fbank.dims * 2; - output.values.assign(static_cast(output.frames * output.dims), 0.0F); - output.attention_mask.assign(static_cast(output.frames), 0); - for (int64_t pair = 0; pair < output.frames; ++pair) { - const int64_t first_frame = pair * 2; - const int64_t second_frame = first_frame + 1; - std::copy_n( - padded.data() + static_cast(first_frame * fbank.dims), - static_cast(fbank.dims), - output.values.data() + static_cast(pair * output.dims)); - std::copy_n( - padded.data() + static_cast(second_frame * fbank.dims), - static_cast(fbank.dims), - output.values.data() + static_cast(pair * output.dims + fbank.dims)); - output.attention_mask[static_cast(pair)] = second_frame < fbank.frames ? 1 : 0; - } - return output; -} - -IndexTTS25PreparedReferenceAudio prepare_index_tts2_5_reference_audio( - const std::vector & samples, - int sample_rate, - int channels, - const IndexTTS25S2MelConfig & mel_config, - size_t threads, - bool speaker_load_semantic) { - constexpr int64_t kMaxReferenceSeconds = 15; - auto mono = require_mono_samples(samples, channels); - const int64_t max_input_samples = static_cast(sample_rate) * kMaxReferenceSeconds; - if (max_input_samples > 0 && static_cast(mono.size()) > max_input_samples) { - engine::audio::truncate_samples_to_count(mono, static_cast(max_input_samples)); - } - - IndexTTS25PreparedReferenceAudio output; - output.waveform_22k = resample_mono_librosa(mono, sample_rate, mel_config.sample_rate); - output.waveform_16k = speaker_load_semantic - ? resample_mono(output.waveform_22k, mel_config.sample_rate, 16000) - : resample_mono_librosa(mono, sample_rate, 16000); - output.mel = compute_index_tts2_5_mel_spectrogram(output.waveform_22k, mel_config, threads); - output.campplus_fbank = compute_index_tts2_5_campplus_fbank_16k(output.waveform_16k); - output.semantic_features = compute_index_tts2_5_semantic_features_16k(output.waveform_16k); - return output; -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/gpt.cpp b/src/models/index_tts2_5/gpt.cpp deleted file mode 100644 index 1637a8ec..00000000 --- a/src/models/index_tts2_5/gpt.cpp +++ /dev/null @@ -1,2217 +0,0 @@ -#include "engine/models/index_tts2_5/gpt.h" - -#include "engine/framework/core/backend.h" -#include "engine/framework/debug/profiler.h" -#include "engine/framework/modules/activation_modules.h" -#include "engine/framework/modules/attention/relative_attention.h" -#include "engine/framework/modules/lookup_modules.h" -#include "engine/framework/modules/optimizations/fast_kv_modules.h" -#include "engine/framework/modules/optimizations/fast_projection_modules.h" -#include "engine/framework/modules/primitive_modules.h" -#include "engine/framework/modules/structural_modules.h" -#include "engine/framework/modules/weight_binding.h" -#include "engine/framework/sampling/torch_random.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -namespace binding = engine::modules::binding; -namespace core = engine::core; -namespace modules = engine::modules; -using Clock = std::chrono::steady_clock; - -constexpr int64_t kSemanticHidden = 1024; -constexpr int64_t kModelDim = 1280; -constexpr int64_t kConditionDim = 512; -constexpr int64_t kEmotionConditionLayers = 4; -constexpr int64_t kGptLayers = 24; -constexpr int64_t kGptMlpDim = 5120; -constexpr int64_t kTextTokens = 60510; -constexpr int64_t kMelCodes = 8194; -constexpr int64_t kMelPositions = 1818; -constexpr int64_t kTextPositions = 602; -constexpr int64_t kConditionPosFrames = 5000; -constexpr int64_t kConditionConvKernel = 15; -constexpr int64_t kGptHeads = 20; -constexpr int64_t kGptHeadDim = kModelDim / kGptHeads; -// spk_cond_mode="campplus": the projected 192-dim CAMPPlus embedding forms a -// single speaker token, followed by two all-zero tokens. -constexpr int64_t kCampplusStyleDim = 192; -constexpr int64_t kConditionTokens = 3; -constexpr int32_t kStartTextToken = 0; -constexpr int32_t kStopTextToken = 1; -constexpr int32_t kStartMelToken = 8192; -constexpr int32_t kStopMelToken = 8193; - -struct GgmlContextDeleter { - void operator()(ggml_context * ctx) const noexcept { - if (ctx != nullptr) { - ggml_free(ctx); - } - } -}; - -core::TensorValue div(core::ModuleBuildContext & ctx, const core::TensorValue & lhs, const core::TensorValue & rhs) { - core::validate_shape(rhs, lhs.shape, "Div rhs"); - return core::wrap_tensor(ggml_div(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32); -} - -core::TensorValue scale(core::ModuleBuildContext & ctx, const core::TensorValue & input, float value) { - return core::wrap_tensor(ggml_scale(ctx.ggml, input.tensor, value), input.shape, GGML_TYPE_F32); -} - -core::TensorValue transpose_btc_bct(core::ModuleBuildContext & ctx, const core::TensorValue & input) { - return modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, input); -} - -core::TensorValue build_biased_gpt_projection( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - int64_t in_features, - int64_t out_features, - const modules::LinearWeights & weights) { - if (ctx.backend_type != core::BackendType::Cuda) { - return modules::LinearModule({in_features, out_features, true, GGML_PREC_F32}).build(ctx, input, weights); - } - if (out_features % 4 != 0) { - throw std::runtime_error("IndexTTS2.5 GPT fast projection requires output features divisible by 4"); - } - auto projected = modules::FastPackedProjection4Module({in_features, out_features, GGML_PREC_F32}) - .build(ctx, input, {weights.weight, std::nullopt}); - if (!weights.bias.has_value()) { - throw std::runtime_error("IndexTTS2.5 GPT linear bias is missing"); - } - const auto matrix_shape = core::TensorShape::from_dims({projected.shape.prefix_elements(), out_features}); - auto matrix = core::reshape_tensor(ctx, projected, matrix_shape); - matrix = core::wrap_tensor(ggml_add(ctx.ggml, matrix.tensor, weights.bias->tensor), matrix_shape, GGML_TYPE_F32); - return core::reshape_tensor(ctx, matrix, input.shape.with_last_dim(out_features)); -} - -core::TensorValue repeat_bias( - core::ModuleBuildContext & ctx, - const core::TensorValue & bias, - const core::TensorValue & like, - int64_t heads, - int64_t dim) { - auto view = core::reshape_tensor(ctx, bias, core::TensorShape::from_dims({1, heads, 1, dim})); - return modules::RepeatModule({like.shape}).build(ctx, view); -} - -core::TensorValue reshape_heads( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - int64_t heads, - int64_t dim) { - const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); - return core::reshape_tensor(ctx, contiguous, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); -} - -core::TensorValue gelu_geglu( - core::ModuleBuildContext & ctx, - const core::TensorValue & input) { - const int64_t half = input.shape.last_dim() / 2; - auto x = modules::SliceModule({static_cast(input.shape.rank - 1), 0, half}).build(ctx, input); - auto gate = modules::SliceModule({static_cast(input.shape.rank - 1), half, half}).build(ctx, input); - gate = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, gate); - return modules::MulModule{}.build(ctx, x, gate); -} - -core::TensorValue glu_axis( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - int axis) { - if (axis < 0 || axis >= static_cast(input.shape.rank) || input.shape.dims[static_cast(axis)] % 2 != 0) { - throw std::runtime_error("IndexTTS2.5 GLU axis shape mismatch"); - } - const int64_t half = input.shape.dims[static_cast(axis)] / 2; - auto value = modules::SliceModule({axis, 0, half}).build(ctx, input); - auto gate = modules::SliceModule({axis, half, half}).build(ctx, input); - gate = modules::SigmoidModule{}.build(ctx, gate); - return modules::MulModule{}.build(ctx, value, gate); -} - -core::TensorValue rms_norm( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & gamma) { - const auto squared = core::wrap_tensor(ggml_sqr(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); - auto sum = modules::ReduceSumModule({static_cast(input.shape.rank - 1)}).build(ctx, squared); - sum = core::wrap_tensor(ggml_sqrt(ctx.ggml, sum.tensor), sum.shape, GGML_TYPE_F32); - auto normed = div(ctx, input, modules::RepeatModule({input.shape}).build(ctx, sum)); - normed = scale(ctx, normed, std::sqrt(static_cast(input.shape.last_dim()))); - auto gamma_view = core::reshape_tensor(ctx, gamma, core::TensorShape::from_dims({1, 1, gamma.shape.dims[0]})); - return modules::MulModule{}.build(ctx, normed, modules::RepeatModule({input.shape}).build(ctx, gamma_view)); -} - -core::TensorValue condition_subsample( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25GptConditionEncoderWeights & weights) { - const int64_t frames = input.shape.dims[1]; - const int64_t frames_after = (frames - 3) / 2 + 1; - auto x = core::reshape_tensor(ctx, input, core::TensorShape::from_dims({1, 1, frames, kSemanticHidden})); - x = modules::Conv2dModule({1, kConditionDim, 3, 3, 2, 2, 0, 0, 1, 1, true}).build(ctx, x, weights.subsampling.conv); - x = modules::ReluModule{}.build(ctx, x); - x = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, x); - x = core::wrap_tensor(ggml_cont(ctx.ggml, x.tensor), x.shape, x.type); - x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, frames_after, kConditionDim * ((kSemanticHidden - 1) / 2)})); - x = modules::LinearModule({kConditionDim * ((kSemanticHidden - 1) / 2), kConditionDim, true}) - .build(ctx, x, weights.subsampling.out); - return scale(ctx, x, std::sqrt(static_cast(kConditionDim))); -} - -core::TensorValue condition_rel_attention( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & pos_emb, - const IndexTTS25GptConditionLayerWeights & weights, - int64_t heads) { - const int64_t dim = kConditionDim / heads; - auto q = modules::LinearModule({kConditionDim, kConditionDim, true}) - .build(ctx, input, {weights.self_attn.attention.q_weight, weights.self_attn.attention.q_bias}); - auto k = modules::LinearModule({kConditionDim, kConditionDim, true}) - .build(ctx, input, {weights.self_attn.attention.k_weight, weights.self_attn.attention.k_bias}); - auto v = modules::LinearModule({kConditionDim, kConditionDim, true}) - .build(ctx, input, {weights.self_attn.attention.v_weight, weights.self_attn.attention.v_bias}); - auto p = modules::LinearModule({kConditionDim, kConditionDim, false}) - .build(ctx, pos_emb, {weights.self_attn.pos_weight, std::nullopt}); - - q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, heads, dim)); - k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, k, heads, dim)); - v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, v, heads, dim)); - p = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, p, heads, dim)); - - const auto q_u = modules::AddModule{}.build(ctx, q, repeat_bias(ctx, weights.self_attn.pos_bias_u, q, heads, dim)); - const auto q_v = modules::AddModule{}.build(ctx, q, repeat_bias(ctx, weights.self_attn.pos_bias_v, q, heads, dim)); - const auto k_t = modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k); - const auto p_t = modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, p); - auto scores = modules::AddModule{}.build(ctx, modules::MatMulModule{}.build(ctx, q_u, k_t), modules::MatMulModule{}.build(ctx, q_v, p_t)); - scores = scale(ctx, scores, 1.0F / std::sqrt(static_cast(dim))); - auto attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, core::ensure_backend_addressable_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); - auto context = modules::MatMulModule{}.build(ctx, attn, v); - context = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context); - context = core::ensure_backend_addressable_layout(ctx, context); - context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], kConditionDim})); - return modules::LinearModule({kConditionDim, kConditionDim, true}) - .build(ctx, context, {weights.self_attn.attention.out_weight, weights.self_attn.attention.out_bias}); -} - -core::TensorValue condition_conv_module( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25GptConditionLayerWeights & weights) { - auto x = transpose_btc_bct(ctx, input); - x = modules::Conv1dModule({kConditionDim, 2 * kConditionDim, 1, 1, 0, 1, true}).build(ctx, x, weights.conv_pointwise_in); - x = glu_axis(ctx, x, 1); - x = modules::DepthwiseConv1dModule({kConditionDim, kConditionConvKernel, 1, 7, 1, true}).build(ctx, x, weights.conv_depthwise); - x = transpose_btc_bct(ctx, x); - x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.conv_norm); - x = modules::SiluModule{}.build(ctx, x); - x = transpose_btc_bct(ctx, x); - x = modules::Conv1dModule({kConditionDim, kConditionDim, 1, 1, 0, 1, true}).build(ctx, x, weights.conv_pointwise_out); - return transpose_btc_bct(ctx, x); -} - -core::TensorValue condition_layer( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & pos_emb, - const IndexTTS25GptConditionLayerWeights & weights, - int64_t heads) { - auto x = input; - auto y = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_mha); - y = condition_rel_attention(ctx, y, pos_emb, weights, heads); - x = modules::AddModule{}.build(ctx, x, y); - - y = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_conv); - y = condition_conv_module(ctx, y, weights); - x = modules::AddModule{}.build(ctx, x, y); - - y = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_ff); - y = modules::LinearModule({kConditionDim, weights.feed_forward_in.weight.shape.dims[0], true}).build(ctx, y, weights.feed_forward_in); - y = modules::SiluModule{}.build(ctx, y); - y = modules::LinearModule({weights.feed_forward_in.weight.shape.dims[0], kConditionDim, true}).build(ctx, y, weights.feed_forward_out); - x = modules::AddModule{}.build(ctx, x, y); - return modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights.norm_final); -} - -core::TensorValue perceiver_attention( - core::ModuleBuildContext & ctx, - const core::TensorValue & latents, - const core::TensorValue & context, - const IndexTTS25PerceiverAttentionWeights & weights, - int64_t dim, - int64_t heads, - int64_t inner) { - const int64_t head_dim = inner / heads; - const auto full_context = modules::ConcatModule({1}).build(ctx, latents, context); - auto q = modules::LinearModule({dim, inner, false}).build(ctx, latents, weights.q); - auto kv = modules::LinearModule({dim, 2 * inner, false}).build(ctx, full_context, weights.kv); - auto k = modules::SliceModule({2, 0, inner}).build(ctx, kv); - auto v = modules::SliceModule({2, inner, inner}).build(ctx, kv); - q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, heads, head_dim)); - k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, k, heads, head_dim)); - v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, v, heads, head_dim)); - auto scores = modules::MatMulModule{}.build(ctx, q, modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k)); - scores = scale(ctx, scores, 1.0F / std::sqrt(static_cast(head_dim))); - auto attn = core::wrap_tensor(ggml_soft_max(ctx.ggml, core::ensure_backend_addressable_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); - auto output = modules::MatMulModule{}.build(ctx, attn, v); - output = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, output); - output = core::ensure_backend_addressable_layout(ctx, output); - output = core::reshape_tensor(ctx, output, core::TensorShape::from_dims({latents.shape.dims[0], latents.shape.dims[1], inner})); - return modules::LinearModule({inner, dim, false}).build(ctx, output, weights.out); -} - -core::TensorValue perceiver_ff( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25PerceiverFeedForwardWeights & weights, - int64_t dim, - int64_t ff_in) { - auto hidden = modules::LinearModule({dim, ff_in, true}).build(ctx, input, weights.in); - hidden = gelu_geglu(ctx, hidden); - return modules::LinearModule({ff_in / 2, dim, true}).build(ctx, hidden, weights.out); -} - -core::TensorValue perceiver( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25PerceiverWeights & weights, - int64_t latents_count, - int64_t dim, - int64_t heads, - int64_t inner, - int64_t ff_in) { - auto context = modules::LinearModule({kConditionDim, dim, true}).build(ctx, input, weights.project_context); - auto latents = modules::RepeatModule({core::TensorShape::from_dims({1, latents_count, dim})}) - .build(ctx, core::reshape_tensor(ctx, weights.latents, core::TensorShape::from_dims({1, latents_count, dim}))); - for (const auto & layer : weights.layers) { - latents = modules::AddModule{}.build(ctx, latents, perceiver_attention(ctx, latents, context, layer.attention, dim, heads, inner)); - latents = modules::AddModule{}.build(ctx, latents, perceiver_ff(ctx, latents, layer.feed_forward, dim, ff_in)); - } - return rms_norm(ctx, latents, weights.norm_gamma); -} - -core::TensorValue condition_encoder( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25GptConditionEncoderWeights & encoder, - const IndexTTS25PerceiverWeights & perceiver_weights, - int64_t encoder_heads, - int64_t perceiver_latents, - int64_t perceiver_dim, - int64_t perceiver_heads, - int64_t perceiver_inner, - int64_t perceiver_ff_in) { - auto x = condition_subsample(ctx, input, encoder); - const int64_t frames = x.shape.dims[1]; - auto pos_emb = modules::SliceModule({1, 0, frames}).build(ctx, encoder.subsampling.pos_enc); - for (const auto & layer : encoder.layers) { - x = condition_layer(ctx, x, pos_emb, layer, encoder_heads); - } - x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, encoder.after_norm); - return perceiver(ctx, x, perceiver_weights, perceiver_latents, perceiver_dim, perceiver_heads, perceiver_inner, perceiver_ff_in); -} - -engine::modules::LinearWeights load_hf_conv1d_linear( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type, - int64_t in_features, - int64_t out_features, - bool use_bias) { - engine::modules::LinearWeights weights; - const auto source_weight = source.require_f32(prefix + ".weight", {in_features, out_features}); - std::vector transposed(static_cast(out_features * in_features)); - for (int64_t in = 0; in < in_features; ++in) { - for (int64_t out = 0; out < out_features; ++out) { - transposed[static_cast(out * in_features + in)] = - source_weight[static_cast(in * out_features + out)]; - } - } - weights.weight = store.make_from_f32( - engine::core::TensorShape::from_dims({out_features, in_features}), - storage_type, - std::move(transposed)); - if (use_bias) { - weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_features}); - } - return weights; -} - -engine::modules::LinearWeights load_biasless_linear( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type, - int64_t out_features, - int64_t in_features) { - return binding::linear_from_source(store, source, prefix, storage_type, out_features, in_features, false); -} - -engine::modules::RelativeAttentionWeights load_condition_relative_attention( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type, - int64_t heads) { - engine::modules::RelativeAttentionWeights weights; - weights.attention.q_weight = store.load_tensor(source, prefix + ".linear_q.weight", storage_type, {kConditionDim, kConditionDim}); - weights.attention.q_bias = store.load_f32_tensor(source, prefix + ".linear_q.bias", {kConditionDim}); - weights.attention.k_weight = store.load_tensor(source, prefix + ".linear_k.weight", storage_type, {kConditionDim, kConditionDim}); - weights.attention.k_bias = store.load_f32_tensor(source, prefix + ".linear_k.bias", {kConditionDim}); - weights.attention.v_weight = store.load_tensor(source, prefix + ".linear_v.weight", storage_type, {kConditionDim, kConditionDim}); - weights.attention.v_bias = store.load_f32_tensor(source, prefix + ".linear_v.bias", {kConditionDim}); - weights.attention.out_weight = store.load_tensor(source, prefix + ".linear_out.weight", storage_type, {kConditionDim, kConditionDim}); - weights.attention.out_bias = store.load_f32_tensor(source, prefix + ".linear_out.bias", {kConditionDim}); - weights.pos_weight = store.load_tensor(source, prefix + ".linear_pos.weight", storage_type, {kConditionDim, kConditionDim}); - weights.pos_bias_u = store.load_f32_tensor(source, prefix + ".pos_bias_u", {heads, kConditionDim / heads}); - weights.pos_bias_v = store.load_f32_tensor(source, prefix + ".pos_bias_v", {heads, kConditionDim / heads}); - return weights; -} - -IndexTTS25GptConditionLayerWeights load_condition_layer( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t linear_units, - int64_t heads, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) { - IndexTTS25GptConditionLayerWeights layer; - layer.norm_ff = binding::norm_from_source(store, source, prefix + ".norm_ff", kConditionDim); - layer.norm_mha = binding::norm_from_source(store, source, prefix + ".norm_mha", kConditionDim); - layer.norm_conv = binding::norm_from_source(store, source, prefix + ".norm_conv", kConditionDim); - layer.norm_final = binding::norm_from_source(store, source, prefix + ".norm_final", kConditionDim); - layer.feed_forward_in = binding::linear_from_source( - store, - source, - prefix + ".feed_forward.w_1", - matmul_storage_type, - linear_units, - kConditionDim, - true); - layer.feed_forward_out = binding::linear_from_source( - store, - source, - prefix + ".feed_forward.w_2", - matmul_storage_type, - kConditionDim, - linear_units, - true); - layer.self_attn = load_condition_relative_attention(store, source, prefix + ".self_attn", matmul_storage_type, heads); - layer.conv_pointwise_in = binding::conv1d_from_source( - store, - source, - prefix + ".conv_module.pointwise_conv1", - conv_storage_type, - 2 * kConditionDim, - kConditionDim, - 1, - true); - layer.conv_depthwise = binding::depthwise_conv1d_from_source( - store, - source, - prefix + ".conv_module.depthwise_conv", - conv_storage_type, - kConditionDim, - kConditionConvKernel, - true); - layer.conv_norm = binding::norm_from_source(store, source, prefix + ".conv_module.norm", kConditionDim); - layer.conv_pointwise_out = binding::conv1d_from_source( - store, - source, - prefix + ".conv_module.pointwise_conv2", - conv_storage_type, - kConditionDim, - kConditionDim, - 1, - true); - return layer; -} - -IndexTTS25GptConditionEncoderWeights load_condition_encoder( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t layers, - int64_t linear_units, - int64_t heads, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) { - IndexTTS25GptConditionEncoderWeights encoder; - encoder.subsampling.conv = binding::conv2d_from_source( - store, - source, - prefix + ".embed.conv.0", - conv_storage_type, - kConditionDim, - 1, - 3, - 3, - true); - encoder.subsampling.out = binding::linear_from_source( - store, - source, - prefix + ".embed.out.0", - matmul_storage_type, - kConditionDim, - kConditionDim * ((kSemanticHidden - 1) / 2), - true); - encoder.subsampling.pos_enc = store.load_f32_tensor( - source, - prefix + ".embed.pos_enc.pe", - {1, kConditionPosFrames, kConditionDim}); - encoder.layers.reserve(static_cast(layers)); - for (int64_t i = 0; i < layers; ++i) { - encoder.layers.push_back(load_condition_layer( - store, - source, - prefix + ".encoders." + std::to_string(i), - linear_units, - heads, - matmul_storage_type, - conv_storage_type)); - } - encoder.after_norm = binding::norm_from_source(store, source, prefix + ".after_norm", kConditionDim); - return encoder; -} - -IndexTTS25PerceiverLayerWeights load_perceiver_layer( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t dim, - int64_t inner, - int64_t ff_in, - engine::assets::TensorStorageType storage_type) { - IndexTTS25PerceiverLayerWeights layer; - layer.attention.q = load_biasless_linear(store, source, prefix + ".0.to_q", storage_type, inner, dim); - layer.attention.kv = load_biasless_linear(store, source, prefix + ".0.to_kv", storage_type, inner * 2, dim); - layer.attention.out = load_biasless_linear(store, source, prefix + ".0.to_out", storage_type, dim, inner); - layer.feed_forward.in = binding::linear_from_source(store, source, prefix + ".1.0", storage_type, ff_in, dim, true); - layer.feed_forward.out = binding::linear_from_source(store, source, prefix + ".1.2", storage_type, dim, ff_in / 2, true); - return layer; -} - -IndexTTS25PerceiverWeights load_perceiver( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t latents, - int64_t dim, - int64_t context_dim, - int64_t inner, - int64_t ff_in, - engine::assets::TensorStorageType storage_type) { - IndexTTS25PerceiverWeights weights; - weights.latents = store.load_f32_tensor(source, prefix + ".latents", {latents, dim}); - weights.project_context = binding::linear_from_source( - store, - source, - prefix + ".proj_context", - storage_type, - dim, - context_dim, - true); - weights.layers.reserve(2); - for (int64_t i = 0; i < 2; ++i) { - weights.layers.push_back(load_perceiver_layer( - store, - source, - prefix + ".layers." + std::to_string(i), - dim, - inner, - ff_in, - storage_type)); - } - weights.norm_gamma = store.load_f32_tensor(source, prefix + ".norm.gamma", {dim}); - return weights; -} - -IndexTTS25Gpt2LayerWeights load_gpt2_layer( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - int64_t layer_index, - engine::assets::TensorStorageType storage_type) { - const std::string prefix = "gpt.h." + std::to_string(layer_index); - IndexTTS25Gpt2LayerWeights layer; - layer.attn_norm = binding::norm_from_source(store, source, prefix + ".ln_1", kModelDim); - layer.qkv = load_hf_conv1d_linear(store, source, prefix + ".attn.c_attn", storage_type, kModelDim, 3 * kModelDim, true); - layer.attn_out = load_hf_conv1d_linear(store, source, prefix + ".attn.c_proj", storage_type, kModelDim, kModelDim, true); - layer.mlp_norm = binding::norm_from_source(store, source, prefix + ".ln_2", kModelDim); - layer.mlp_in = load_hf_conv1d_linear(store, source, prefix + ".mlp.c_fc", storage_type, kModelDim, kGptMlpDim, true); - layer.mlp_out = load_hf_conv1d_linear(store, source, prefix + ".mlp.c_proj", storage_type, kGptMlpDim, kModelDim, true); - return layer; -} - -struct Gpt2LayerOutput { - core::TensorValue output; - core::TensorValue key; - core::TensorValue value; -}; - -core::TensorValue gpt_attention_from_heads( - core::ModuleBuildContext & ctx, - const core::TensorValue & q_heads, - const core::TensorValue & k_heads, - const core::TensorValue & v_heads, - const std::optional & attention_mask) { - if (attention_mask.has_value()) { - auto q_contiguous = core::ensure_backend_addressable_layout(ctx, q_heads); - auto * flash = ggml_flash_attn_ext( - ctx.ggml, - q_contiguous.tensor, - k_heads.tensor, - v_heads.tensor, - attention_mask->tensor, - 1.0F / std::sqrt(static_cast(kGptHeadDim)), - 0.0F, - 0.0F); - ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); - return core::wrap_tensor( - flash, - core::TensorShape::from_dims({q_contiguous.shape.dims[0], q_contiguous.shape.dims[2], q_contiguous.shape.dims[1], kGptHeadDim}), - GGML_TYPE_F32); - } - 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(kGptHeadDim))), - scores.shape, - GGML_TYPE_F32); - scores = core::wrap_tensor(ggml_diag_mask_inf(ctx.ggml, scores.tensor, 0), scores.shape, GGML_TYPE_F32); - scores = core::wrap_tensor(ggml_soft_max(ctx.ggml, core::ensure_backend_addressable_layout(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); - return modules::MatMulModule{}.build(ctx, scores, v_heads); -} - -core::TensorValue gpt_mlp( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25Gpt2LayerWeights & weights) { - auto hidden = build_biased_gpt_projection(ctx, input, kModelDim, kGptMlpDim, weights.mlp_in); - hidden = modules::GeluModule({modules::GeluApproximation::Tanh}).build(ctx, hidden); - return build_biased_gpt_projection(ctx, hidden, kGptMlpDim, kModelDim, weights.mlp_out); -} - -Gpt2LayerOutput gpt2_layer_full( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25Gpt2LayerWeights & weights, - const std::optional & attention_mask = std::nullopt) { - auto normed = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, input, weights.attn_norm); - auto qkv = build_biased_gpt_projection(ctx, normed, kModelDim, 3 * kModelDim, weights.qkv); - auto q = modules::SliceModule({2, 0, kModelDim}).build(ctx, qkv); - auto k = modules::SliceModule({2, kModelDim, kModelDim}).build(ctx, qkv); - auto v = modules::SliceModule({2, 2 * kModelDim, kModelDim}).build(ctx, qkv); - auto k_cache = reshape_heads(ctx, k, kGptHeads, kGptHeadDim); - auto v_cache = reshape_heads(ctx, v, kGptHeads, kGptHeadDim); - q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, kGptHeads, kGptHeadDim)); - auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k_cache); - auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v_cache); - auto context = gpt_attention_from_heads(ctx, q, k_heads, v_heads, attention_mask); - context = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context); - context = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, context), input.shape); - auto x = modules::AddModule{}.build(ctx, input, build_biased_gpt_projection(ctx, context, kModelDim, kModelDim, weights.attn_out)); - auto mlp_in = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights.mlp_norm); - return {modules::AddModule{}.build(ctx, x, gpt_mlp(ctx, mlp_in, weights)), k_cache, v_cache}; -} - -Gpt2LayerOutput gpt2_layer_cached_tail( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25Gpt2LayerWeights & weights, - const core::TensorValue & cache_key, - const core::TensorValue & cache_value, - const core::TensorValue & cache_slots, - const core::TensorValue & attention_mask) { - if (cache_key.shape.dims[0] != input.shape.dims[0] || - cache_value.shape.dims[0] != input.shape.dims[0] || - cache_key.shape.dims[1] != cache_value.shape.dims[1] || - cache_key.shape.dims[2] != cache_value.shape.dims[2] || - cache_key.shape.dims[3] != cache_value.shape.dims[3] || - cache_slots.shape.dims[0] != input.shape.dims[0]) { - throw std::runtime_error("IndexTTS2.5 GPT cached layer batch cache shape mismatch"); - } - auto normed = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, input, weights.attn_norm); - auto qkv = build_biased_gpt_projection(ctx, normed, kModelDim, 3 * kModelDim, weights.qkv); - auto q = modules::SliceModule({2, 0, kModelDim}).build(ctx, qkv); - auto k = modules::SliceModule({2, kModelDim, kModelDim}).build(ctx, qkv); - auto v = modules::SliceModule({2, 2 * kModelDim, kModelDim}).build(ctx, qkv); - q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q, kGptHeads, kGptHeadDim)); - k = reshape_heads(ctx, k, kGptHeads, kGptHeadDim); - v = reshape_heads(ctx, v, kGptHeads, kGptHeadDim); - - const modules::FastKVSetRowsModule set_rows; - auto updated_key = set_rows.build(ctx, cache_key, k, cache_slots); - auto updated_value = set_rows.build(ctx, cache_value, v, cache_slots); - - auto context = gpt_attention_from_heads( - ctx, - q, - modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, updated_key), - modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, updated_value), - attention_mask); - context = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, context), input.shape); - auto x = modules::AddModule{}.build(ctx, input, build_biased_gpt_projection(ctx, context, kModelDim, kModelDim, weights.attn_out)); - auto mlp_in = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights.mlp_norm); - return {modules::AddModule{}.build(ctx, x, gpt_mlp(ctx, mlp_in, weights)), k, v}; -} - -struct TopPItem { - size_t index = 0; - float score = 0.0F; - float weight = 0.0F; -}; - -struct SampleScore { - size_t flat_index = 0; - float score = 0.0F; -}; - -struct RankedSample { - size_t score_index = 0; - size_t flat_index = 0; - double rank = 0.0; -}; - -struct IndexTTS25SamplerWorkspace { - std::vector scores; - std::vector top_k_heap; - std::vector top_p_items; - std::vector seen_tokens; - uint32_t seen_generation = 1; - std::vector finite_score_indices; - std::vector sample_scores; - std::vector ranked_samples; - std::vector selected_scores; -}; - -void apply_repetition_penalty( - std::vector & logits, - const std::vector & codes, - float penalty, - IndexTTS25SamplerWorkspace & workspace) { - if (penalty == 1.0F) { - return; - } - if (!(penalty > 0.0F)) { - throw std::runtime_error("IndexTTS2.5 GPT repetition_penalty must be positive"); - } - if (workspace.seen_tokens.size() != logits.size()) { - workspace.seen_tokens.assign(logits.size(), 0); - workspace.seen_generation = 1; - } else if (workspace.seen_generation == 0) { - std::fill(workspace.seen_tokens.begin(), workspace.seen_tokens.end(), 0); - workspace.seen_generation = 1; - } - const uint32_t generation = workspace.seen_generation++; - const auto apply_token = [&](int32_t token) { - if (token < 0 || static_cast(token) >= logits.size() || workspace.seen_tokens[static_cast(token)] == generation) { - return; - } - workspace.seen_tokens[static_cast(token)] = generation; - float & value = logits[static_cast(token)]; - value = value < 0.0F ? value * penalty : value / penalty; - }; - apply_token(kStopTextToken); - apply_token(kStartMelToken); - for (const int32_t token : codes) { - apply_token(token); - } -} - -float kth_largest_threshold(const std::vector & scores, size_t keep_count, std::vector & heap) { - heap.clear(); - heap.reserve(keep_count); - const auto greater = std::greater{}; - for (const float score : scores) { - if (heap.size() < keep_count) { - heap.push_back(score); - std::push_heap(heap.begin(), heap.end(), greater); - } else if (score > heap.front()) { - std::pop_heap(heap.begin(), heap.end(), greater); - heap.back() = score; - std::push_heap(heap.begin(), heap.end(), greater); - } - } - return heap.front(); -} - -void index_tts2_5_log_probs( - const std::vector & logits, - const std::vector & codes, - float repetition_penalty, - int top_k, - float top_p, - float temperature, - IndexTTS25SamplerWorkspace & workspace) { - if (!(temperature > 0.0F)) { - throw std::runtime_error("IndexTTS2.5 GPT temperature must be positive"); - } - float max_logit = -std::numeric_limits::infinity(); - for (float logit : logits) { - max_logit = std::max(max_logit, logit); - } - float total = 0.0F; - for (float logit : logits) { - total += std::exp(logit - max_logit); - } - if (!(total > 0.0F)) { - throw std::runtime_error("IndexTTS2.5 GPT sampler invalid logit mass"); - } - const float log_total = std::log(total); - auto & scores = workspace.scores; - scores.resize(logits.size()); - for (size_t i = 0; i < logits.size(); ++i) { - scores[i] = logits[i] - max_logit - log_total; - } - apply_repetition_penalty(scores, codes, repetition_penalty, workspace); - for (float & score : scores) { - score /= temperature; - } - const size_t min_tokens_to_keep = 2; - auto & finite_indices = workspace.finite_score_indices; - finite_indices.clear(); - if (top_k > 0 && static_cast(top_k) < scores.size()) { - const size_t keep_count = std::max(static_cast(top_k), min_tokens_to_keep); - const float threshold = kth_largest_threshold(scores, keep_count, workspace.top_k_heap); - float max_score = -std::numeric_limits::infinity(); - for (size_t i = 0; i < scores.size(); ++i) { - if (scores[i] < threshold) { - scores[i] = -std::numeric_limits::infinity(); - } else { - finite_indices.push_back(i); - max_score = std::max(max_score, scores[i]); - } - } - if (!std::isfinite(max_score)) { - throw std::runtime_error("IndexTTS2.5 GPT sampler has no finite score"); - } - } else { - finite_indices.reserve(scores.size()); - float max_score = -std::numeric_limits::infinity(); - for (size_t i = 0; i < scores.size(); ++i) { - if (std::isfinite(scores[i])) { - finite_indices.push_back(i); - max_score = std::max(max_score, scores[i]); - } - } - if (!std::isfinite(max_score)) { - throw std::runtime_error("IndexTTS2.5 GPT sampler has no finite score"); - } - } - if (top_p > 0.0F && top_p < 1.0F) { - auto & sorted = workspace.top_p_items; - sorted.clear(); - sorted.reserve(finite_indices.size()); - float total = 0.0F; - float max_score = -std::numeric_limits::infinity(); - for (const size_t i : finite_indices) { - max_score = std::max(max_score, scores[i]); - } - for (const size_t i : finite_indices) { - const float weight = std::exp(scores[i] - max_score); - sorted.push_back({i, scores[i], weight}); - total += weight; - } - if (!(total > 0.0F)) { - throw std::runtime_error("IndexTTS2.5 GPT sampler invalid top-p mass"); - } - std::sort(sorted.begin(), sorted.end(), [](const TopPItem & lhs, const TopPItem & rhs) { - if (lhs.score == rhs.score) { - return lhs.index < rhs.index; - } - return lhs.score < rhs.score; - }); - float cumulative = 0.0F; - const float remove_mass = 1.0F - top_p; - const size_t keep_from = sorted.size() > min_tokens_to_keep ? sorted.size() - min_tokens_to_keep : 0; - finite_indices.clear(); - for (size_t i = 0; i < sorted.size(); ++i) { - cumulative += sorted[i].weight / total; - if (i < keep_from && cumulative <= remove_mass) { - scores[sorted[i].index] = -std::numeric_limits::infinity(); - } else { - finite_indices.push_back(sorted[i].index); - } - } - } -} - -void sample_index_tts2_5_indices( - const std::vector & scores, - size_t total_score_count, - size_t count, - uint64_t seed, - uint64_t step, - const engine::sampling::TorchCudaSamplingPolicy & policy, - std::vector & ranked, - std::vector & selected) { - float max_score = -std::numeric_limits::infinity(); - for (const auto & score : scores) { - max_score = std::max(max_score, score.score); - } - if (!std::isfinite(max_score)) { - throw std::runtime_error("IndexTTS2.5 GPT sampler has no finite beam score"); - } - ranked.clear(); - ranked.reserve(scores.size()); - for (size_t i = 0; i < scores.size(); ++i) { - const auto & score = scores[i]; - const float probability = std::exp(score.score - max_score); - const float exponential = engine::sampling::torch_cuda_tensor_iterator_exponential_element( - seed, - static_cast(total_score_count), - static_cast(score.flat_index), - step, - policy.multiprocessor_count, - policy.max_threads_per_multiprocessor); - ranked.push_back({i, score.flat_index, static_cast(probability) / static_cast(exponential)}); - } - if (ranked.empty()) { - throw std::runtime_error("IndexTTS2.5 GPT sampler failed to select beam candidates"); - } - const size_t keep = std::min(count, ranked.size()); - std::partial_sort( - ranked.begin(), - ranked.begin() + static_cast(keep), - ranked.end(), - [](const RankedSample & lhs, const RankedSample & rhs) { - if (lhs.rank == rhs.rank) { - return lhs.flat_index < rhs.flat_index; - } - return lhs.rank > rhs.rank; - }); - selected.clear(); - selected.reserve(keep); - for (size_t i = 0; i < keep; ++i) { - selected.push_back(ranked[i].score_index); - } -} - -} // namespace - -std::shared_ptr load_index_tts2_5_gpt_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes) { - if (assets.gpt_weights == nullptr) { - throw std::runtime_error("IndexTTS2.5 GPT requires tensor source"); - } - auto weights = std::make_shared(); - weights->store = std::make_shared( - backend, - backend_type, - "index_tts2_5.gpt.weights", - weight_context_bytes); - - const auto & source = *assets.gpt_weights; - weights->emotion_conditioner = load_condition_encoder( - *weights->store, - source, - "emo_conditioning_encoder", - kEmotionConditionLayers, - 1024, - 4, - matmul_storage_type, - conv_storage_type); - weights->emotion_perceiver = load_perceiver( - *weights->store, - source, - "emo_perceiver_encoder", - 1, - kSemanticHidden, - kConditionDim, - 256, - 2730, - matmul_storage_type); - weights->spk_emb_proj = binding::linear_from_source( - *weights->store, - source, - "spk_emb_proj", - matmul_storage_type, - kModelDim, - kCampplusStyleDim, - true); - weights->lang_embedding = weights->store->load_tensor( - source, - "lang_embedding.weight", - matmul_storage_type, - {kIndexTTS25LangEmbeddingRows, kModelDim}); - weights->text_embedding = weights->store->load_tensor( - source, - "text_embedding.weight", - matmul_storage_type, - {kTextTokens, kModelDim}); - weights->mel_embedding = weights->store->load_tensor( - source, - "mel_embedding.weight", - matmul_storage_type, - {kMelCodes, kModelDim}); - weights->text_pos_embedding = weights->store->load_f32_tensor( - source, - "text_pos_embedding.emb.weight", - {kTextPositions, kModelDim}); - weights->mel_pos_embedding = weights->store->load_f32_tensor( - source, - "mel_pos_embedding.emb.weight", - {kMelPositions, kModelDim}); - weights->emotion_vec_projection = binding::linear_from_source( - *weights->store, - source, - "emovec_layer", - matmul_storage_type, - kModelDim, - kSemanticHidden, - true); - weights->emotion_layer = binding::linear_from_source( - *weights->store, - source, - "emo_layer", - matmul_storage_type, - kModelDim, - kModelDim, - true); - weights->gpt_layers.reserve(static_cast(kGptLayers)); - for (int64_t i = 0; i < kGptLayers; ++i) { - weights->gpt_layers.push_back(load_gpt2_layer(*weights->store, source, i, matmul_storage_type)); - } - weights->gpt_final_norm = binding::norm_from_source(*weights->store, source, "gpt.ln_f", kModelDim); - weights->final_norm = binding::norm_from_source(*weights->store, source, "final_norm", kModelDim); - weights->mel_head = binding::linear_from_source( - *weights->store, - source, - "mel_head", - matmul_storage_type, - kMelCodes, - kModelDim, - true); - weights->text_head = binding::linear_from_source( - *weights->store, - source, - "text_head", - matmul_storage_type, - kTextTokens, - kModelDim, - true); - - weights->store->upload(); - assets.gpt_weights->release_storage(); - return weights; -} - -class IndexTTS25GptRuntime::ConditioningGraph { -public: - ConditioningGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t frames, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - frames_(frames) { - if (frames_ <= 0) { - throw std::runtime_error("IndexTTS2.5 GPT conditioning graph requires positive frame count"); - } - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 GPT conditioning graph requires weights"); - } - - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT conditioning graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT conditioning input context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.emo_condition", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.gpt.emo_condition.inputs", - execution_.backend_type()}; - semantic_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kSemanticHidden})).tensor; - ggml_set_input(semantic_); - auto input = core::wrap_tensor(semantic_, core::TensorShape::from_dims({1, frames_, kSemanticHidden}), GGML_TYPE_F32); - auto out = condition_encoder( - ctx, - input, - weights_->emotion_conditioner, - weights_->emotion_perceiver, - 4, - 1, - kSemanticHidden, - 4, - 256, - 2730); - output_ = core::ensure_backend_addressable_layout(ctx, out).tensor; - output_frames_ = out.shape.dims[1]; - output_dims_ = out.shape.dims[2]; - ggml_set_output(output_); - - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(65536, frames_ * 4096 + 8192)), false); - ggml_build_forward_expand(graph_, output_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT conditioning input buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT conditioning graph"); - } - debug::timing_log_scalar( - "index_tts2_5.gpt.emo_condition.graph.build_ms", - engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.gpt.emo_condition.frames", frames_); - } - - ~ConditioningGraph() { - clear_graph(); - } - - int64_t frames() const noexcept { - return frames_; - } - - IndexTTS25GptLatent run(const std::vector & semantic_btc) { - if (static_cast(semantic_btc.size()) != frames_ * kSemanticHidden) { - throw std::runtime_error("IndexTTS2.5 GPT conditioning input value count mismatch"); - } - auto timing_start = Clock::now(); - ggml_backend_tensor_set(semantic_, semantic_btc.data(), 0, semantic_btc.size() * sizeof(float)); - debug::timing_log_scalar( - "index_tts2_5.gpt.emo_condition.input_upload_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar( - "index_tts2_5.gpt.emo_condition.graph.compute_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 GPT conditioning graph compute failed"); - } - - IndexTTS25GptLatent output; - output.frames = output_frames_; - output.dims = output_dims_; - output.values.resize(static_cast(output.frames * output.dims)); - timing_start = Clock::now(); - ggml_backend_tensor_get(output_, output.values.data(), 0, output.values.size() * sizeof(float)); - debug::timing_log_scalar( - "index_tts2_5.gpt.emo_condition.output_read_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - return output; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t frames_ = 0; - int64_t output_frames_ = 0; - int64_t output_dims_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr ctx_; - ggml_tensor * semantic_ = nullptr; - ggml_tensor * output_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; - }; - -class IndexTTS25GptRuntime::EmotionVectorGraph { -public: - EmotionVectorGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)) { - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 GPT emotion vector graph requires weights"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT emotion vector graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT emotion vector input context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.emotion_vector", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.gpt.emotion_vector.inputs", - execution_.backend_type()}; - input_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kSemanticHidden})).tensor; - ggml_set_input(input_); - auto x = core::wrap_tensor(input_, core::TensorShape::from_dims({1, kSemanticHidden}), GGML_TYPE_F32); - x = modules::LinearModule({kSemanticHidden, kModelDim, true, GGML_PREC_F32}).build(ctx, x, weights_->emotion_vec_projection); - x = modules::LinearModule({kModelDim, kModelDim, true, GGML_PREC_F32}).build(ctx, x, weights_->emotion_layer); - output_ = core::ensure_backend_addressable_layout(ctx, x).tensor; - ggml_set_output(output_); - graph_ = ggml_new_graph_custom(ctx_.get(), 8192, false); - ggml_build_forward_expand(graph_, output_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT emotion vector input buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT emotion vector graph"); - } - debug::timing_log_scalar("index_tts2_5.gpt.emotion_vector.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - } - - ~EmotionVectorGraph() { - clear_graph(); - } - - std::vector run(const IndexTTS25GptLatent & emotion_conditioning) { - if (emotion_conditioning.frames != 1 || - emotion_conditioning.dims != kSemanticHidden || - static_cast(emotion_conditioning.values.size()) != kSemanticHidden) { - throw std::runtime_error("IndexTTS2.5 GPT emotion vector input shape mismatch"); - } - ggml_backend_tensor_set(input_, emotion_conditioning.values.data(), 0, emotion_conditioning.values.size() * sizeof(float)); - core::set_backend_threads(execution_.backend(), execution_.config().threads); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 GPT emotion vector graph compute failed"); - } - std::vector out(static_cast(kModelDim)); - ggml_backend_tensor_get(output_, out.data(), 0, out.size() * sizeof(float)); - return out; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - std::unique_ptr input_ctx_; - std::unique_ptr ctx_; - ggml_tensor * input_ = nullptr; - ggml_tensor * output_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; -}; - -struct GptPrefillOutput { - std::vector logits; - std::vector latent; - runtime::TransformerKVState kv_state; -}; - -class IndexTTS25GptRuntime::PrefillGraph { -public: - PrefillGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t text_tokens, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - text_tokens_(text_tokens), - text_steps_(text_tokens + 2), - prompt_steps_(kConditionTokens + text_tokens + 3) { - if (weights_ == nullptr || text_tokens_ < 0) { - throw std::runtime_error("IndexTTS2.5 GPT prefill graph requires weights and non-negative text tokens"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT prefill graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT prefill input context"); - } - ggml_init_params output_params{16ull * 1024ull * 1024ull, nullptr, true}; - output_ctx_.reset(ggml_init(output_params)); - if (output_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT prefill output context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.prefill", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.gpt.prefill.inputs", - execution_.backend_type()}; - core::ModuleBuildContext output_ctx{ - output_ctx_.get(), - "index_tts2_5.gpt.prefill.outputs", - execution_.backend_type()}; - style_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kCampplusStyleDim})).tensor; - emo_vec_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kModelDim})).tensor; - lang_id_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 1); - text_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, text_steps_); - start_mel_id_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 1); - ggml_set_input(style_); - ggml_set_input(emo_vec_); - ggml_set_input(lang_id_); - ggml_set_input(text_ids_); - ggml_set_input(start_mel_id_); - // campplus conditioning prefix (model_v2.py inference_speech): - // conds = [spk_emb_proj(style) + emo_vec, zeros, zeros]. - auto style = core::wrap_tensor(style_, core::TensorShape::from_dims({1, kCampplusStyleDim}), GGML_TYPE_F32); - auto speaker_token = build_biased_gpt_projection(ctx, style, kCampplusStyleDim, kModelDim, weights_->spk_emb_proj); - speaker_token = core::reshape_tensor(ctx, speaker_token, core::TensorShape::from_dims({1, 1, kModelDim})); - auto emo_vec = core::wrap_tensor(emo_vec_, core::TensorShape::from_dims({1, kModelDim}), GGML_TYPE_F32); - emo_vec = core::reshape_tensor(ctx, emo_vec, core::TensorShape::from_dims({1, 1, kModelDim})); - auto conds = modules::AddModule{}.build(ctx, speaker_token, emo_vec); - auto zero_token = modules::RepeatModule({core::TensorShape::from_dims({1, kConditionTokens - 1, kModelDim})}) - .build(ctx, scale(ctx, conds, 0.0F)); - conds = modules::ConcatModule({1}).build(ctx, conds, zero_token); - auto text_ids = core::wrap_tensor(text_ids_, core::TensorShape::from_dims({text_steps_}), GGML_TYPE_I32); - auto text = modules::EmbeddingModule({kTextTokens, kModelDim}).build(ctx, text_ids, weights_->text_embedding); - auto text_pos = modules::SliceModule({0, 0, text_steps_}).build(ctx, weights_->text_pos_embedding); - text = modules::AddModule{}.build(ctx, text, text_pos); - auto lang_id = core::wrap_tensor(lang_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); - auto lang = modules::EmbeddingModule({kIndexTTS25LangEmbeddingRows, kModelDim}).build(ctx, lang_id, weights_->lang_embedding); - text = modules::AddModule{}.build(ctx, text, modules::RepeatModule({text.shape}).build(ctx, lang)); - text = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, text), core::TensorShape::from_dims({1, text_steps_, kModelDim})); - auto mel_id = core::wrap_tensor(start_mel_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); - auto mel = modules::EmbeddingModule({kMelCodes, kModelDim}).build(ctx, mel_id, weights_->mel_embedding); - auto mel_pos = modules::SliceModule({0, 0, 1}).build(ctx, weights_->mel_pos_embedding); - mel = modules::AddModule{}.build(ctx, mel, mel_pos); - mel = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, mel), core::TensorShape::from_dims({1, 1, kModelDim})); - auto x = modules::ConcatModule({1}).build(ctx, conds, text); - x = modules::ConcatModule({1}).build(ctx, x, mel); - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(65536, prompt_steps_ * 8192)), false); - for (const auto & layer : weights_->gpt_layers) { - auto out = gpt2_layer_full(ctx, x, layer); - x = out.output; - auto key = core::ensure_backend_addressable_layout(ctx, out.key); - auto value = core::ensure_backend_addressable_layout(ctx, out.value); - auto * key_output = core::make_tensor(output_ctx, GGML_TYPE_F32, key.shape).tensor; - auto * value_output = core::make_tensor(output_ctx, GGML_TYPE_F32, value.shape).tensor; - keys_.push_back(key_output); - values_.push_back(value_output); - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), key.tensor, key_output)); - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), value.tensor, value_output)); - } - x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->gpt_final_norm); - x = modules::SliceModule({1, prompt_steps_ - 1, 1}).build(ctx, x); - x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->final_norm); - auto latent = core::ensure_backend_addressable_layout(ctx, x); - auto logits = core::ensure_backend_addressable_layout( - ctx, - modules::LinearModule({kModelDim, kMelCodes, true, GGML_PREC_F32}).build(ctx, x, weights_->mel_head)); - latent_ = core::make_tensor(output_ctx, GGML_TYPE_F32, latent.shape).tensor; - logits_ = core::make_tensor(output_ctx, GGML_TYPE_F32, logits.shape).tensor; - ggml_set_output(latent_); - ggml_set_output(logits_); - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), latent.tensor, latent_)); - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), logits.tensor, logits_)); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT prefill input buffer"); - } - output_buffer_ = ggml_backend_alloc_ctx_tensors(output_ctx_.get(), execution_.backend()); - if (output_buffer_ == nullptr) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT prefill output buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT prefill graph"); - } - const int32_t start_mel = kStartMelToken; - ggml_backend_tensor_set(start_mel_id_, &start_mel, 0, sizeof(int32_t)); - debug::timing_log_scalar("index_tts2_5.gpt.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.gpt.prefill.prompt_steps", prompt_steps_); - } - - ~PrefillGraph() { - clear_graph(); - } - - bool matches(int64_t text_tokens) const noexcept { - return text_tokens_ == text_tokens; - } - - int64_t prompt_steps() const noexcept { - return prompt_steps_; - } - - GptPrefillOutput run( - const std::vector & speaker_style, - const std::vector & emotion_vector, - int32_t lang_id, - const std::vector & text_tokens) { - if (static_cast(speaker_style.size()) != kCampplusStyleDim || - static_cast(emotion_vector.size()) != kModelDim || - static_cast(text_tokens.size()) != text_tokens_) { - throw std::runtime_error("IndexTTS2.5 GPT prefill input shape mismatch"); - } - if (lang_id < 0 || lang_id >= kIndexTTS25LangEmbeddingRows) { - throw std::runtime_error("IndexTTS2.5 GPT prefill lang id is out of range"); - } - std::vector ids; - ids.reserve(static_cast(text_steps_)); - ids.push_back(kStartTextToken); - ids.insert(ids.end(), text_tokens.begin(), text_tokens.end()); - ids.push_back(kStopTextToken); - auto timing_start = Clock::now(); - ggml_backend_tensor_set(style_, speaker_style.data(), 0, speaker_style.size() * sizeof(float)); - ggml_backend_tensor_set(emo_vec_, emotion_vector.data(), 0, emotion_vector.size() * sizeof(float)); - ggml_backend_tensor_set(lang_id_, &lang_id, 0, sizeof(int32_t)); - ggml_backend_tensor_set(text_ids_, ids.data(), 0, ids.size() * sizeof(int32_t)); - debug::timing_log_scalar("index_tts2_5.gpt.prefill.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar("index_tts2_5.gpt.prefill.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 GPT prefill graph compute failed"); - } - GptPrefillOutput out; - out.logits.resize(static_cast(kMelCodes)); - out.latent.resize(static_cast(kModelDim)); - ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); - ggml_backend_tensor_get(latent_, out.latent.data(), 0, out.latent.size() * sizeof(float)); - out.kv_state.current_end = prompt_steps_; - out.kv_state.layers.resize(keys_.size()); - const size_t layer_values = static_cast(prompt_steps_ * kGptHeads * kGptHeadDim); - for (size_t layer = 0; layer < keys_.size(); ++layer) { - auto & state = out.kv_state.layers[layer]; - state.valid_steps = prompt_steps_; - state.key.resize(layer_values); - state.value.resize(layer_values); - ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); - ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); - } - return out; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - if (output_buffer_ != nullptr) { - ggml_backend_buffer_free(output_buffer_); - output_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t text_tokens_ = 0; - int64_t text_steps_ = 0; - int64_t prompt_steps_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr output_ctx_; - std::unique_ptr ctx_; - ggml_tensor * style_ = nullptr; - ggml_tensor * emo_vec_ = nullptr; - ggml_tensor * lang_id_ = nullptr; - ggml_tensor * text_ids_ = nullptr; - ggml_tensor * start_mel_id_ = nullptr; - ggml_tensor * latent_ = nullptr; - ggml_tensor * logits_ = nullptr; - std::vector keys_; - std::vector values_; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; - ggml_backend_buffer_t output_buffer_ = nullptr; -}; - -class IndexTTS25GptRuntime::DecodeGraph { -public: - struct StepOutput { - std::vector logits; - }; - - struct BatchOutput { - std::vector steps; - }; - - DecodeGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t cache_steps, - int64_t beam_count, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - cache_steps_(cache_steps), - beam_count_(beam_count), - beam_slots_(2 * beam_count) { - if (weights_ == nullptr || cache_steps_ <= 0 || beam_count_ <= 0) { - throw std::runtime_error("IndexTTS2.5 GPT decode graph requires weights and cache steps"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT decode graph context"); - } - ggml_init_params state_params{256ull * 1024ull * 1024ull, nullptr, true}; - state_ctx_.reset(ggml_init(state_params)); - if (state_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 GPT decode state context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.gpt.decode", execution_.backend_type()}; - core::ModuleBuildContext state_ctx{ - state_ctx_.get(), - "index_tts2_5.gpt.decode.state", - execution_.backend_type()}; - token_ids_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, beam_count_); - mel_positions_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, beam_count_); - cache_slots_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, beam_count_); - attention_mask_ = ggml_new_tensor_4d(state_ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, beam_count_); - ggml_set_input(token_ids_); - ggml_set_input(mel_positions_); - ggml_set_input(cache_slots_); - ggml_set_input(attention_mask_); - for (int64_t bank = 0; bank < 2; ++bank) { - auto & keys = bank_keys_[static_cast(bank)]; - auto & values = bank_values_[static_cast(bank)]; - keys.reserve(weights_->gpt_layers.size()); - values.reserve(weights_->gpt_layers.size()); - for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { - keys.push_back(core::make_tensor( - state_ctx, - GGML_TYPE_F32, - core::TensorShape::from_dims({beam_count_, cache_steps_, kGptHeads, kGptHeadDim}))); - values.push_back(core::make_tensor( - state_ctx, - GGML_TYPE_F32, - core::TensorShape::from_dims({beam_count_, cache_steps_, kGptHeads, kGptHeadDim}))); - } - } - build_prefix_views(); - build_bank_graph(ctx, 0); - build_bank_graph(ctx, 1); - state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), execution_.backend()); - if (state_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT decode state buffer"); - } - debug::trace_log_scalar( - "index_tts2_5.gpt.decode.state_buffer_mib", - static_cast(ggml_backend_buffer_get_size(state_buffer_)) / (1024.0 * 1024.0)); - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, bank_graphs_[0].graph) || - !ggml_gallocr_reserve(gallocr_, bank_graphs_[1].graph) || - !ggml_gallocr_alloc_graph(gallocr_, bank_graphs_[0].graph)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT decode graph"); - } - attention_mask_values_.assign(static_cast(beam_count_ * cache_steps_), ggml_fp32_to_fp16(-INFINITY)); - token_values_.assign(static_cast(beam_count_), 0); - position_values_.assign(static_cast(beam_count_), 0); - cache_slot_values_.assign(static_cast(beam_count_), 0); - debug::timing_log_scalar("index_tts2_5.gpt.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.gpt.decode.cache_steps", cache_steps_); - debug::trace_log_scalar("index_tts2_5.gpt.decode.beam_batch", beam_count_); - } - - ~DecodeGraph() { - clear_graph(); - } - - void clear_graph() { - for (auto & graph : bank_graphs_) { - if (graph.graph != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph.graph); - graph.graph = nullptr; - } - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (state_buffer_ != nullptr) { - ggml_backend_buffer_free(state_buffer_); - state_buffer_ = nullptr; - } - } - - bool can_run(int64_t required_steps, int64_t required_beam_slots) const noexcept { - return cache_steps_ >= required_steps && beam_slots_ >= required_beam_slots && beam_count_ * 2 == required_beam_slots; - } - - void initialize_beam_slot(int64_t slot, const runtime::TransformerKVState & state) { - if (slot < 0 || slot >= beam_slots_) { - throw std::runtime_error("IndexTTS2.5 GPT beam slot is out of range"); - } - if (state.layers.size() != weights_->gpt_layers.size()) { - throw std::runtime_error("IndexTTS2.5 GPT beam state layer count mismatch"); - } - for (size_t layer = 0; layer < state.layers.size(); ++layer) { - const auto & layer_state = state.layers[layer]; - if (!state.layers.empty() && layer_state.valid_steps != state.layers.front().valid_steps) { - throw std::runtime_error("IndexTTS2.5 GPT beam state valid step mismatch"); - } - ggml_backend_tensor_set( - beam_key_prefix_views_[static_cast(slot)][static_cast(layer_state.valid_steps)][layer], - layer_state.key.data(), - 0, - layer_state.key.size() * sizeof(float)); - ggml_backend_tensor_set( - beam_value_prefix_views_[static_cast(slot)][static_cast(layer_state.valid_steps)][layer], - layer_state.value.data(), - 0, - layer_state.value.size() * sizeof(float)); - } - } - - BatchOutput run_batch_from_beams( - const std::vector & parent_slots, - const std::vector & child_slots, - int64_t valid_steps, - const std::vector & tokens, - int32_t mel_position) { - const size_t active = parent_slots.size(); - if (active == 0 || child_slots.size() != active || tokens.size() != active) { - throw std::runtime_error("IndexTTS2.5 GPT batched decode input shape mismatch"); - } - if (active > static_cast(beam_count_)) { - throw std::runtime_error("IndexTTS2.5 GPT batched decode exceeds beam batch"); - } - if (valid_steps >= cache_steps_) { - throw std::runtime_error("IndexTTS2.5 GPT decode cache exhausted"); - } - const int64_t child_bank = child_slots.front() / beam_count_; - if (child_bank < 0 || child_bank > 1) { - throw std::runtime_error("IndexTTS2.5 GPT child beam bank is out of range"); - } - for (size_t row = 0; row < active; ++row) { - if (parent_slots[row] < 0 || parent_slots[row] >= beam_slots_ || - child_slots[row] < 0 || child_slots[row] >= beam_slots_ || - child_slots[row] / beam_count_ != child_bank || - child_slots[row] % beam_count_ != static_cast(row)) { - throw std::runtime_error("IndexTTS2.5 GPT beam slot layout mismatch"); - } - for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { - if (parent_slots[row] != child_slots[row]) { - copy_beam_prefix(parent_slots[row], child_slots[row], valid_steps, layer); - } - } - token_values_[row] = tokens[row]; - position_values_[row] = mel_position; - } - for (size_t row = active; row < static_cast(beam_count_); ++row) { - const int64_t child_slot = child_bank * beam_count_ + static_cast(row); - for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { - copy_beam_prefix(child_slots.front(), child_slot, valid_steps, layer); - } - token_values_[row] = tokens.front(); - position_values_[row] = mel_position; - } - - const auto masked = ggml_fp32_to_fp16(-INFINITY); - const auto visible = ggml_fp32_to_fp16(0.0F); - std::fill(attention_mask_values_.begin(), attention_mask_values_.end(), masked); - for (int64_t row = 0; row < beam_count_; ++row) { - auto * row_values = attention_mask_values_.data() + static_cast(row * cache_steps_); - for (int64_t step = 0; step <= valid_steps; ++step) { - row_values[static_cast(step)] = visible; - } - cache_slot_values_[static_cast(row)] = static_cast(row * cache_steps_ + valid_steps); - } - ggml_backend_tensor_set(token_ids_, token_values_.data(), 0, token_values_.size() * sizeof(int32_t)); - ggml_backend_tensor_set(mel_positions_, position_values_.data(), 0, position_values_.size() * sizeof(int32_t)); - ggml_backend_tensor_set(cache_slots_, cache_slot_values_.data(), 0, cache_slot_values_.size() * sizeof(int32_t)); - ggml_backend_tensor_set(attention_mask_, attention_mask_values_.data(), 0, attention_mask_values_.size() * sizeof(ggml_fp16_t)); - - auto & graph = bank_graphs_[static_cast(child_bank)]; - if (!ggml_gallocr_alloc_graph(gallocr_, graph.graph)) { - throw std::runtime_error("failed to allocate IndexTTS2.5 GPT decode bank graph"); - } - core::set_backend_threads(execution_.backend(), execution_.config().threads); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph.graph); - ggml_backend_synchronize(execution_.backend()); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 GPT decode graph compute failed"); - } - - std::vector logits(static_cast(beam_count_ * kMelCodes)); - ggml_backend_tensor_get(graph.logits, logits.data(), 0, logits.size() * sizeof(float)); - BatchOutput out; - out.steps.reserve(active); - for (size_t row = 0; row < active; ++row) { - StepOutput step; - step.logits.assign( - logits.begin() + static_cast(row * static_cast(kMelCodes)), - logits.begin() + static_cast((row + 1) * static_cast(kMelCodes))); - out.steps.push_back(std::move(step)); - } - return out; - } - -private: - struct BankGraph { - ggml_cgraph * graph = nullptr; - ggml_tensor * logits = nullptr; - }; - - void build_prefix_views() { - beam_key_prefix_views_.assign(static_cast(beam_slots_), {}); - beam_value_prefix_views_.assign(static_cast(beam_slots_), {}); - for (int64_t slot = 0; slot < beam_slots_; ++slot) { - const int64_t bank = slot / beam_count_; - const int64_t row = slot % beam_count_; - auto & key_steps = beam_key_prefix_views_[static_cast(slot)]; - auto & value_steps = beam_value_prefix_views_[static_cast(slot)]; - key_steps.assign(static_cast(cache_steps_ + 1), {}); - value_steps.assign(static_cast(cache_steps_ + 1), {}); - for (int64_t steps = 1; steps <= cache_steps_; ++steps) { - auto & key_layers = key_steps[static_cast(steps)]; - auto & value_layers = value_steps[static_cast(steps)]; - key_layers.reserve(weights_->gpt_layers.size()); - value_layers.reserve(weights_->gpt_layers.size()); - const int64_t elems = steps * kGptHeads * kGptHeadDim; - const size_t byte_offset = static_cast(row * cache_steps_ * kGptHeads * kGptHeadDim) * sizeof(float); - for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { - key_layers.push_back(ggml_view_1d(state_ctx_.get(), bank_keys_[static_cast(bank)][layer].tensor, elems, byte_offset)); - value_layers.push_back(ggml_view_1d(state_ctx_.get(), bank_values_[static_cast(bank)][layer].tensor, elems, byte_offset)); - } - } - } - } - - void copy_beam_prefix(int64_t parent_slot, int64_t child_slot, int64_t valid_steps, size_t layer) { - if (valid_steps <= 0 || parent_slot == child_slot) { - return; - } - if (valid_steps > cache_steps_) { - throw std::runtime_error("IndexTTS2.5 GPT beam prefix copy exceeds cache capacity"); - } - ggml_backend_tensor_copy( - beam_key_prefix_views_[static_cast(parent_slot)][static_cast(valid_steps)][layer], - beam_key_prefix_views_[static_cast(child_slot)][static_cast(valid_steps)][layer]); - ggml_backend_tensor_copy( - beam_value_prefix_views_[static_cast(parent_slot)][static_cast(valid_steps)][layer], - beam_value_prefix_views_[static_cast(child_slot)][static_cast(valid_steps)][layer]); - } - - void build_bank_graph(core::ModuleBuildContext & ctx, int64_t bank) { - BankGraph graph; - graph.graph = ggml_new_graph_custom(ctx_.get(), 65536, false); - auto token = core::wrap_tensor(token_ids_, core::TensorShape::from_dims({beam_count_}), GGML_TYPE_I32); - auto x = modules::EmbeddingModule({kMelCodes, kModelDim}).build(ctx, token, weights_->mel_embedding); - auto pos = core::wrap_tensor(mel_positions_, core::TensorShape::from_dims({beam_count_}), GGML_TYPE_I32); - auto pos_emb = modules::EmbeddingModule({kMelPositions, kModelDim}).build(ctx, pos, weights_->mel_pos_embedding); - x = modules::AddModule{}.build(ctx, x, pos_emb); - x = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, x), core::TensorShape::from_dims({beam_count_, 1, kModelDim})); - auto mask = core::wrap_tensor(attention_mask_, core::TensorShape::from_dims({beam_count_, 1, 1, cache_steps_}), GGML_TYPE_F16); - auto cache_slots = core::wrap_tensor(cache_slots_, core::TensorShape::from_dims({beam_count_}), GGML_TYPE_I32); - for (size_t layer = 0; layer < weights_->gpt_layers.size(); ++layer) { - auto out = gpt2_layer_cached_tail( - ctx, - x, - weights_->gpt_layers[layer], - bank_keys_[static_cast(bank)][layer], - bank_values_[static_cast(bank)][layer], - cache_slots, - mask); - x = out.output; - } - x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->gpt_final_norm); - x = modules::LayerNormModule({kModelDim, 1.0e-5F, true, true}).build(ctx, x, weights_->final_norm); - graph.logits = modules::LinearModule({kModelDim, kMelCodes, true, GGML_PREC_F32}).build(ctx, x, weights_->mel_head).tensor; - ggml_set_output(graph.logits); - ggml_build_forward_expand(graph.graph, graph.logits); - bank_graphs_[static_cast(bank)] = graph; - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t cache_steps_ = 0; - int64_t beam_count_ = 0; - int64_t beam_slots_ = 0; - std::unique_ptr state_ctx_; - std::unique_ptr ctx_; - ggml_tensor * token_ids_ = nullptr; - ggml_tensor * mel_positions_ = nullptr; - ggml_tensor * cache_slots_ = nullptr; - ggml_tensor * attention_mask_ = nullptr; - std::array bank_graphs_; - std::array, 2> bank_keys_; - std::array, 2> bank_values_; - std::vector>> beam_key_prefix_views_; - std::vector>> beam_value_prefix_views_; - std::vector attention_mask_values_; - std::vector token_values_; - std::vector position_values_; - std::vector cache_slot_values_; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t state_buffer_ = nullptr; -}; - -IndexTTS25GptRuntime::IndexTTS25GptRuntime( - std::shared_ptr assets, - core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) - : assets_(std::move(assets)), - execution_(&execution), - graph_arena_bytes_(graph_arena_bytes) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 GPT runtime requires assets"); - } - if (graph_arena_bytes_ == 0) { - throw std::runtime_error("IndexTTS2.5 GPT graph arena must be non-zero"); - } - weights_ = load_index_tts2_5_gpt_weights( - *assets_, - execution.backend(), - execution.backend_type(), - matmul_storage_type, - conv_storage_type, - weight_context_bytes); -} - -IndexTTS25GptRuntime::~IndexTTS25GptRuntime() = default; - -void IndexTTS25GptRuntime::prepare_emotion_conditioning(int64_t frames) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 GPT runtime execution context is missing"); - } - if (frames <= 0) { - throw std::runtime_error("IndexTTS2.5 GPT emotion conditioning prepare requires positive frames"); - } - if (emotion_conditioning_graph_ != nullptr && emotion_conditioning_graph_->frames() == frames) { - return; - } - emotion_conditioning_graph_.reset(); - emotion_conditioning_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); -} - -IndexTTS25GptLatent IndexTTS25GptRuntime::emotion_conditioning(const std::vector & semantic_btc, int64_t frames) { - if (emotion_conditioning_graph_ == nullptr || emotion_conditioning_graph_->frames() != frames) { - throw std::runtime_error("IndexTTS2.5 GPT emotion conditioning graph was not prepared for this reference length"); - } - return emotion_conditioning_graph_->run(semantic_btc); -} - -void IndexTTS25GptRuntime::prepare_generation(int64_t text_tokens, int64_t max_mel_tokens, int64_t num_beams) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 GPT runtime execution context is missing"); - } - if (text_tokens < 0 || max_mel_tokens <= 0) { - throw std::runtime_error("IndexTTS2.5 GPT generation prepare requires non-negative text tokens and positive mel tokens"); - } - if (num_beams != 1) { - debug::trace_log_scalar("index_tts2_5.gpt.generation.num_beams", num_beams); - } - if (prefill_graph_ == nullptr || !prefill_graph_->matches(text_tokens)) { - prefill_graph_.reset(); - prefill_graph_ = std::make_unique(*execution_, weights_, text_tokens, graph_arena_bytes_); - } - const int64_t required_cache_steps = prefill_graph_->prompt_steps() + max_mel_tokens + 1; - const int64_t required_beam_slots = 2 * std::max(1, num_beams); - if (decode_graph_ == nullptr || !decode_graph_->can_run(required_cache_steps, required_beam_slots)) { - decode_graph_.reset(); - decode_graph_ = std::make_unique( - *execution_, - weights_, - required_cache_steps, - std::max(1, num_beams), - graph_arena_bytes_); - } -} - -std::vector IndexTTS25GptRuntime::project_emotion_vector(const IndexTTS25GptLatent & emotion_conditioning) { - if (emotion_vector_graph_ == nullptr) { - emotion_vector_graph_ = std::make_unique(*execution_, weights_, graph_arena_bytes_); - } - return emotion_vector_graph_->run(emotion_conditioning); -} - -std::vector IndexTTS25GptRuntime::merge_emotion_vector( - const std::vector & speaker_semantic, - int64_t speaker_frames, - const std::vector & emotion_semantic, - int64_t emotion_frames, - float alpha) { - prepare_emotion_conditioning(speaker_frames); - const auto base_condition = emotion_conditioning(speaker_semantic, speaker_frames); - prepare_emotion_conditioning(emotion_frames); - const auto emotion_condition = emotion_conditioning(emotion_semantic, emotion_frames); - auto base = project_emotion_vector(base_condition); - const auto emotion = project_emotion_vector(emotion_condition); - for (size_t i = 0; i < base.size(); ++i) { - base[i] = base[i] + alpha * (emotion[i] - base[i]); - } - return base; -} - -IndexTTS25GptGeneration IndexTTS25GptRuntime::generate_speech(const IndexTTS25GptGenerationRequest & request) { - if (request.text_tokens.empty()) { - throw std::runtime_error("IndexTTS2.5 GPT generation requires text tokens"); - } - const auto text_tokens = align_index_tts2_5_gpt_text_tokens(request.text_tokens); - if (static_cast(request.speaker_style.size()) != kCampplusStyleDim) { - throw std::runtime_error("IndexTTS2.5 GPT generation speaker style shape mismatch"); - } - if (request.lang_id < 0 || request.lang_id >= kIndexTTS25LangEmbeddingRows) { - throw std::runtime_error("IndexTTS2.5 GPT generation lang id is out of range"); - } - std::vector emotion_vector = request.emotion_vector; - if (emotion_vector.empty()) { - prepare_emotion_conditioning(request.emotion_frames); - emotion_vector = project_emotion_vector(emotion_conditioning(request.emotion_semantic, request.emotion_frames)); - } - if (static_cast(emotion_vector.size()) != kModelDim) { - throw std::runtime_error("IndexTTS2.5 GPT generation emotion vector shape mismatch"); - } - prepare_generation( - static_cast(text_tokens.size()), - request.max_mel_tokens, - request.num_beams); - auto prefill = prefill_graph_->run(request.speaker_style, emotion_vector, request.lang_id, text_tokens); - const auto sampling_policy = engine::sampling::resolve_torch_cuda_sampling_policy( - execution_->backend_type(), - execution_->config().device, - "index_tts2_5.gpt.cuda_sampling_policy", - "IndexTTS2.5", - engine::sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault); - - struct Beam { - std::vector codes; - std::vector logits; - int64_t slot = 0; - int64_t valid_steps = 0; - int64_t current_end = 0; - float score = 0.0F; - bool finished = false; - }; - auto normalized_score = [&](const Beam & beam) { - if (request.length_penalty == 0.0F) { - return beam.score; - } - const float length = static_cast(std::max(1, beam.codes.size())); - return beam.score / std::pow(length, request.length_penalty); - }; - struct BeamCandidate { - size_t parent = 0; - int32_t token = 0; - float score = 0.0F; - bool finished = false; - }; - const int beam_count = std::max(1, request.num_beams); - const int64_t prefill_valid_steps = prefill.kv_state.layers.empty() ? 0 : prefill.kv_state.layers.front().valid_steps; - std::vector beams; - beams.reserve(static_cast(beam_count)); - for (int beam = 0; beam < beam_count; ++beam) { - decode_graph_->initialize_beam_slot(beam, prefill.kv_state); - Beam initial; - initial.logits = prefill.logits; - initial.slot = beam; - initial.valid_steps = prefill_valid_steps; - initial.current_end = prefill.kv_state.current_end; - initial.score = beam == 0 ? 0.0F : -1.0e9F; - beams.push_back(std::move(initial)); - } - - std::vector completed; - auto add_completed = [&](Beam beam) { - completed.push_back(std::move(beam)); - if (static_cast(completed.size()) > beam_count) { - const auto worst = std::min_element(completed.begin(), completed.end(), [&](const Beam & lhs, const Beam & rhs) { - return normalized_score(lhs) < normalized_score(rhs); - }); - completed.erase(worst); - } - }; - auto completed_worst_score = [&]() { - if (completed.empty()) { - return -std::numeric_limits::infinity(); - } - const auto worst = std::min_element(completed.begin(), completed.end(), [&](const Beam & lhs, const Beam & rhs) { - return normalized_score(lhs) < normalized_score(rhs); - }); - return normalized_score(*worst); - }; - auto should_stop = [&](float best_sum_logprobs, int64_t cur_generated_len) { - if (static_cast(completed.size()) < beam_count) { - return false; - } - const float length = static_cast(std::max(1, cur_generated_len)); - const float highest_attainable = request.length_penalty == 0.0F - ? best_sum_logprobs - : best_sum_logprobs / std::pow(length, request.length_penalty); - return completed_worst_score() >= highest_attainable; - }; - bool first_decode_timing_logged = false; - int active_bank = 0; - bool beam_search_done = false; - uint64_t sample_call_index = 0; - uint64_t rng_offset_blocks = 0; - double sampling_ms = 0.0; - double decode_run_ms = 0.0; - IndexTTS25SamplerWorkspace sampler_workspace; - std::vector candidates; - std::vector next_beams; - std::vector parent_slots; - std::vector child_slots; - std::vector next_tokens; - candidates.reserve(static_cast(2 * beam_count)); - next_beams.reserve(static_cast(beam_count)); - parent_slots.reserve(static_cast(beam_count)); - child_slots.reserve(static_cast(beam_count)); - next_tokens.reserve(static_cast(beam_count)); - for (int step = 0; step < request.max_mel_tokens && !beams.empty(); ++step) { - const auto sampling_start = Clock::now(); - candidates.clear(); - sampler_workspace.sample_scores.clear(); - sampler_workspace.sample_scores.reserve(beams.size() * static_cast(std::max(request.top_k, 1))); - for (size_t beam_index = 0; beam_index < beams.size(); ++beam_index) { - index_tts2_5_log_probs( - beams[beam_index].logits, - beams[beam_index].codes, - request.repetition_penalty, - request.top_k, - request.top_p, - request.temperature, - sampler_workspace); - const size_t beam_offset = beam_index * static_cast(kMelCodes); - for (const size_t token : sampler_workspace.finite_score_indices) { - const float log_prob = sampler_workspace.scores[token]; - sampler_workspace.sample_scores.push_back({beam_offset + token, beams[beam_index].score + log_prob}); - } - } - const size_t flat_score_count = beams.size() * static_cast(kMelCodes); - const size_t keep = std::min(static_cast(2 * beam_count), flat_score_count); - if (request.do_sample) { - rng_offset_blocks += engine::sampling::torch_cuda_tensor_iterator_offset_blocks( - static_cast(flat_score_count), - sampling_policy); - sample_index_tts2_5_indices( - sampler_workspace.sample_scores, - flat_score_count, - keep, - request.seed, - sample_call_index++, - sampling_policy, - sampler_workspace.ranked_samples, - sampler_workspace.selected_scores); - std::sort(sampler_workspace.selected_scores.begin(), sampler_workspace.selected_scores.end(), [&](size_t lhs, size_t rhs) { - const auto & lhs_score = sampler_workspace.sample_scores[lhs]; - const auto & rhs_score = sampler_workspace.sample_scores[rhs]; - if (lhs_score.score == rhs_score.score) { - return lhs_score.flat_index < rhs_score.flat_index; - } - return lhs_score.score > rhs_score.score; - }); - } else { - sampler_workspace.selected_scores.resize(sampler_workspace.sample_scores.size()); - std::iota(sampler_workspace.selected_scores.begin(), sampler_workspace.selected_scores.end(), 0); - const size_t finite_keep = std::min(keep, sampler_workspace.selected_scores.size()); - std::partial_sort( - sampler_workspace.selected_scores.begin(), - sampler_workspace.selected_scores.begin() + static_cast(finite_keep), - sampler_workspace.selected_scores.end(), - [&](size_t lhs, size_t rhs) { - const auto & lhs_score = sampler_workspace.sample_scores[lhs]; - const auto & rhs_score = sampler_workspace.sample_scores[rhs]; - if (lhs_score.score == rhs_score.score) { - return lhs_score.flat_index < rhs_score.flat_index; - } - return lhs_score.score > rhs_score.score; - }); - sampler_workspace.selected_scores.resize(finite_keep); - } - for (size_t rank = 0; rank < sampler_workspace.selected_scores.size(); ++rank) { - const auto & selected = sampler_workspace.sample_scores[sampler_workspace.selected_scores[rank]]; - const size_t flat_index = selected.flat_index; - const size_t parent = flat_index / static_cast(kMelCodes); - const auto token = static_cast(flat_index % static_cast(kMelCodes)); - candidates.push_back({ - parent, - token, - selected.score, - token == kStopMelToken}); - } - sampling_ms += engine::debug::elapsed_ms(sampling_start, Clock::now()); - next_beams.clear(); - const int next_bank = 1 - active_bank; - parent_slots.clear(); - child_slots.clear(); - next_tokens.clear(); - for (size_t rank = 0; rank < candidates.size(); ++rank) { - const auto & candidate = candidates[rank]; - const Beam & parent = beams[candidate.parent]; - Beam next; - next.codes = parent.codes; - next.score = candidate.score; - if (candidate.finished) { - if (static_cast(rank) < beam_count) { - next.finished = true; - add_completed(std::move(next)); - } - continue; - } - next.codes.push_back(candidate.token); - next.slot = static_cast(next_bank * beam_count + static_cast(next_beams.size())); - next.valid_steps = parent.valid_steps + 1; - next.current_end = parent.current_end + 1; - parent_slots.push_back(parent.slot); - child_slots.push_back(next.slot); - next_tokens.push_back(candidate.token); - next_beams.push_back(std::move(next)); - if (static_cast(next_beams.size()) == beam_count) { - break; - } - } - if (!next_beams.empty()) { - const auto run_start = Clock::now(); - const int64_t parent_valid_steps = next_beams.front().valid_steps - 1; - const auto batch_out = decode_graph_->run_batch_from_beams( - parent_slots, - child_slots, - parent_valid_steps, - next_tokens, - static_cast(next_beams.front().codes.size() + 1)); - const auto run_ms = engine::debug::elapsed_ms(run_start, Clock::now()); - decode_run_ms += run_ms; - if (batch_out.steps.size() != next_beams.size()) { - throw std::runtime_error("IndexTTS2.5 GPT batched decode output size mismatch"); - } - for (size_t beam = 0; beam < next_beams.size(); ++beam) { - next_beams[beam].logits = batch_out.steps[beam].logits; - } - if (!first_decode_timing_logged) { - debug::timing_log_scalar("index_tts2_5.gpt.decode.first_run_ms", run_ms); - first_decode_timing_logged = true; - } - } - if (!candidates.empty() && should_stop(candidates.front().score, static_cast(step + 1))) { - beam_search_done = true; - break; - } - beams.swap(next_beams); - active_bank = next_bank; - } - debug::timing_log_scalar("index_tts2_5.gpt.sampling_ms", sampling_ms); - debug::timing_log_scalar("index_tts2_5.gpt.decode.run_ms", decode_run_ms); - if (!beam_search_done) { - for (auto & beam : beams) { - add_completed(std::move(beam)); - } - } - if (completed.empty()) { - throw std::runtime_error("IndexTTS2.5 GPT generation produced no beam candidates"); - } - const auto best = std::max_element(completed.begin(), completed.end(), [&](const Beam & lhs, const Beam & rhs) { - return normalized_score(lhs) < normalized_score(rhs); - }); - IndexTTS25GptGeneration out; - out.codes = best->codes; - out.rng_offset_blocks = rng_offset_blocks; - bool stop_seen = false; - for (size_t i = 0; i < out.codes.size(); ++i) { - if (out.codes[i] == kStopMelToken) { - stop_seen = true; - } - } - debug::trace_log_scalar("index_tts2_5.gpt.generated_code_count", static_cast(out.codes.size())); - debug::trace_log_scalar("index_tts2_5.gpt.generated_stop_seen", stop_seen); - return out; -} - -void IndexTTS25GptRuntime::release_conditioning_graphs() { - emotion_conditioning_graph_.reset(); - emotion_vector_graph_.reset(); -} - -void IndexTTS25GptRuntime::release_generation_graphs() { - prefill_graph_.reset(); - decode_graph_.reset(); -} - -std::vector align_index_tts2_5_gpt_text_tokens(const std::vector & text_tokens) { - std::vector out; - out.reserve(text_tokens.size()); - for (const int32_t token : text_tokens) { - if (token == kStartTextToken || token == kStopTextToken) { - continue; - } - out.push_back(token); - } - return out; -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/loader.cpp b/src/models/index_tts2_5/loader.cpp deleted file mode 100644 index dd2ab65c..00000000 --- a/src/models/index_tts2_5/loader.cpp +++ /dev/null @@ -1,157 +0,0 @@ -#include "engine/models/index_tts2_5/loader.h" - -#include "engine/framework/model_spec/package.h" -#include "engine/models/index_tts2_5/session.h" - -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -runtime::ModelMetadata metadata(const IndexTTS25Assets & assets) { - runtime::ModelMetadata out; - out.family = "index_tts2_5"; - out.variant = assets.config.version; - out.description = "IndexTTS2.5 loaded from local extracted assets."; - return out; -} - -runtime::CapabilitySet capabilities(const IndexTTS25Assets &) { - runtime::CapabilitySet out; - out.supported_tasks = { - {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, - {runtime::VoiceTaskKind::VoiceCloning, {runtime::RunMode::Offline}}, - }; - out.supports_speaker_reference = true; - out.supports_style_condition = true; - out.languages = {"Chinese", "English", "Japanese", "Spanish", "Arabic"}; - return out; -} - -runtime::ModelCliInterface cli(const IndexTTS25Assets &) { - runtime::ModelCliInterface out; - out.request_options = { - {"lang", "auto|zh|en|ja|es|ar|...", "Text language hint; auto infers zh when the text contains Han characters, otherwise en."}, - {"emotion_alpha", "float", "Blend strength for explicit emotion conditioning."}, - {"emotion_vector", "float[,float...]", "Eight-value explicit emotion vector."}, - {"use_emotion_text", "bool", "Infer emotion from text instead of reference audio."}, - {"emotion_text", "text", "Text used when emotion-text conditioning is enabled."}, - {"use_random_emotion", "bool", "Use random emotion weights in the emotion mixer."}, - {"interval_silence_ms", "n", "Silence inserted between generated text chunks."}, - {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode used when text_chunk_size is set."}, - {"length_penalty", "float", "GPT beam-search length penalty."}, - {"num_beams", "n", "GPT beam count."}, - }; - out.session_options = { - {"index_tts2_5.weight_type", "native|f32|f16|bf16|q8_0", "Matmul weight storage type."}, - {"index_tts2_5.conv_weight_type", "native|f32|f16", "Convolution weight storage type."}, - {"index_tts2_5.gpt_graph_arena_mb", "n", "GPT graph arena size."}, - {"index_tts2_5.s2mel_graph_arena_mb", "n", "S2Mel graph arena size."}, - {"index_tts2_5.reference_graph_arena_mb", "n", "Reference encoder and codec graph arena size."}, - {"index_tts2_5.emotion_text_prefill_graph_arena_mb", "n", "Emotion-text prefill graph arena size."}, - {"index_tts2_5.emotion_text_decode_graph_arena_mb", "n", "Emotion-text cached-step graph arena size."}, - {"index_tts2_5.emotion_text_max_new_tokens", "n", "Maximum generated tokens for emotion-text classification; default 256."}, - {"index_tts2_5.weight_context_mb", "n", "Shared weight context size."}, - {"index_tts2_5.mem_saver", "true|false", "Release staged reference and conditioning graphs after request phases; default false."}, - {"index_tts2_5.speaker_cache_slots", "n", "Prepared speaker-reference cache slots; default 1."}, - {"index_tts2_5.emotion_cache_slots", "n", "Prepared emotion-reference cache slots; default 1."}, - {"index_tts2_5.emotion_text_cache_slots", "n", "Emotion-text weight cache slots; default 1."}, - }; - return out; -} - -class IndexTTS25Loader final : public runtime::IVoiceModelLoader { -public: - std::string family() const override { - return "index_tts2_5"; - } - - runtime::CapabilitySet advertised_capabilities() const override { - runtime::CapabilitySet out; - out.supported_tasks = { - {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, - {runtime::VoiceTaskKind::VoiceCloning, {runtime::RunMode::Offline}}, - }; - out.supports_speaker_reference = true; - out.supports_style_condition = true; - return out; - } - - bool can_load(const runtime::ModelLoadRequest & request) const override { - try { - const auto package_spec = engine::model_spec::default_spec_path(family()); - (void) engine::model_spec::load_resource_bundle( - request.model_path, - package_spec); - 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_index_tts2_5_assets(request.model_path); - runtime::ModelInspection inspection; - inspection.model_root = assets->resources.model_root(); - inspection.metadata = metadata(*assets); - inspection.capabilities = capabilities(*assets); - inspection.cli = cli(*assets); - const auto package_spec = engine::model_spec::default_spec_path(family()); - inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( - request.model_path, - package_spec, - engine::model_spec::ResourceKind::Files); - inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( - request.model_path, - package_spec, - engine::model_spec::ResourceKind::Tensors); - return inspection; - } - - std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { - return load_index_tts2_5_model(request.model_path); - } -}; - -} // namespace - -IndexTTS25LoadedModel::IndexTTS25LoadedModel( - runtime::ModelMetadata metadata, - runtime::CapabilitySet capabilities, - std::shared_ptr assets) - : metadata_(std::move(metadata)), - capabilities_(std::move(capabilities)), - assets_(std::move(assets)) {} - -const runtime::ModelMetadata & IndexTTS25LoadedModel::metadata() const noexcept { - return metadata_; -} - -const runtime::CapabilitySet & IndexTTS25LoadedModel::capabilities() const noexcept { - return capabilities_; -} - -std::unique_ptr IndexTTS25LoadedModel::create_task_session( - const runtime::TaskSpec & task, - const runtime::SessionOptions & options) const { - if (task.mode != runtime::RunMode::Offline || - (task.task != runtime::VoiceTaskKind::Tts && task.task != runtime::VoiceTaskKind::VoiceCloning)) { - throw std::runtime_error("IndexTTS2.5 only supports offline TTS and voice-cloning sessions"); - } - return std::make_unique(task, options, assets_); -} - -std::unique_ptr load_index_tts2_5_model(const std::filesystem::path & model_path) { - auto assets = load_index_tts2_5_assets(model_path); - return std::make_unique( - metadata(*assets), - capabilities(*assets), - std::move(assets)); -} - -std::shared_ptr make_index_tts2_5_loader() { - return std::make_shared(); -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/qwen_emotion.cpp b/src/models/index_tts2_5/qwen_emotion.cpp deleted file mode 100644 index 77955d47..00000000 --- a/src/models/index_tts2_5/qwen_emotion.cpp +++ /dev/null @@ -1,794 +0,0 @@ -#include "engine/models/index_tts2_5/qwen_emotion.h" - -#include "engine/framework/core/backend.h" -#include "engine/framework/debug/profiler.h" -#include "engine/framework/modules/linear_module.h" -#include "engine/framework/modules/lookup_modules.h" -#include "engine/framework/modules/structural_modules.h" -#include "engine/framework/modules/weight_binding.h" -#include "engine/framework/runtime/kv_cache.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -namespace binding = engine::modules::binding; -namespace core = engine::core; -namespace modules = engine::modules; -using Clock = std::chrono::steady_clock; - -constexpr int64_t kHidden = 1024; -constexpr int64_t kIntermediate = 3072; -constexpr int64_t kLayers = 28; -constexpr int64_t kAttentionHeads = 16; -constexpr int64_t kKvHeads = 8; -constexpr int64_t kHeadDim = 128; -constexpr int64_t kVocab = 151936; -constexpr float kRmsEps = 1.0e-6F; -constexpr float kRopeTheta = 1000000.0F; - -struct GgmlContextDeleter { - void operator()(ggml_context * ctx) const noexcept { - if (ctx != nullptr) { - ggml_free(ctx); - } - } -}; - -modules::QwenDecoderStackConfig qwen_config() { - modules::QwenDecoderStackConfig config; - config.hidden_size = kHidden; - config.num_attention_heads = kAttentionHeads; - config.num_key_value_heads = kKvHeads; - config.head_dim = kHeadDim; - config.intermediate_size = kIntermediate; - config.layers = kLayers; - config.rms_norm_eps = kRmsEps; - config.rope_theta = kRopeTheta; - config.attention_precision = GGML_PREC_F32; - config.projection_precision = GGML_PREC_DEFAULT; - return config; -} - -std::vector causal_mask(int64_t steps) { - std::vector mask(static_cast(steps * steps), 0.0F); - for (int64_t q = 0; q < steps; ++q) { - for (int64_t k = q + 1; k < steps; ++k) { - mask[static_cast(q * steps + k)] = -std::numeric_limits::infinity(); - } - } - return mask; -} - -float clamp_emotion(float value) { - return std::clamp(value, 0.0F, 1.2F); -} - -float parse_named_score(const std::string & content, const std::string & key) { - const std::regex pattern("\"?" + key + "\"?\\s*:\\s*([0-9]+(?:\\.[0-9]+)?)"); - std::smatch match; - if (!std::regex_search(content, match, pattern)) { - return 0.0F; - } - return clamp_emotion(std::stof(match[1].str())); -} - -IndexTTS25EmotionVector convert_emotion_json(const std::string & content, const std::string & source_text) { - IndexTTS25EmotionVector out; - out.values = { - parse_named_score(content, "高兴"), - parse_named_score(content, "愤怒"), - parse_named_score(content, "悲伤"), - parse_named_score(content, "恐惧"), - parse_named_score(content, "反感"), - parse_named_score(content, "低落"), - parse_named_score(content, "惊讶"), - parse_named_score(content, "自然"), - }; - std::string lower = source_text; - std::transform(lower.begin(), lower.end(), lower.begin(), [](unsigned char ch) { return static_cast(std::tolower(ch)); }); - if (lower.find("低落") != std::string::npos || - lower.find("melancholy") != std::string::npos || - lower.find("melancholic") != std::string::npos || - lower.find("depression") != std::string::npos || - lower.find("depressed") != std::string::npos || - lower.find("gloomy") != std::string::npos) { - std::swap(out.values[2], out.values[5]); - } - const bool all_zero = std::all_of(out.values.begin(), out.values.end(), [](float value) { return value <= 0.0F; }); - if (all_zero) { - out.values[7] = 1.0F; - } - return out; -} - -engine::modules::QwenDecoderLayerWeights load_qwen_layer( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - int64_t layer_index, - engine::assets::TensorStorageType storage_type) { - const std::string prefix = "model.layers." + std::to_string(layer_index); - engine::modules::QwenDecoderLayerWeights layer; - layer.input_norm = binding::norm_weight_from_source(store, source, prefix + ".input_layernorm", kHidden); - layer.self_attention.q_weight = store.load_tensor( - source, - prefix + ".self_attn.q_proj.weight", - storage_type, - {kAttentionHeads * kHeadDim, kHidden}); - layer.self_attention.k_weight = store.load_tensor( - source, - prefix + ".self_attn.k_proj.weight", - storage_type, - {kKvHeads * kHeadDim, kHidden}); - layer.self_attention.v_weight = store.load_tensor( - source, - prefix + ".self_attn.v_proj.weight", - storage_type, - {kKvHeads * kHeadDim, kHidden}); - layer.self_attention.out_weight = store.load_tensor( - source, - prefix + ".self_attn.o_proj.weight", - storage_type, - {kHidden, kAttentionHeads * kHeadDim}); - layer.q_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.q_norm", kHeadDim); - layer.k_norm = binding::norm_weight_from_source(store, source, prefix + ".self_attn.k_norm", kHeadDim); - layer.post_norm = binding::norm_weight_from_source(store, source, prefix + ".post_attention_layernorm", kHidden); - layer.mlp.gate_proj = binding::linear_from_source( - store, - source, - prefix + ".mlp.gate_proj", - storage_type, - kIntermediate, - kHidden, - false); - layer.mlp.up_proj = binding::linear_from_source( - store, - source, - prefix + ".mlp.up_proj", - storage_type, - kIntermediate, - kHidden, - false); - layer.mlp.down_proj = binding::linear_from_source( - store, - source, - prefix + ".mlp.down_proj", - storage_type, - kHidden, - kIntermediate, - false); - return layer; -} - -} // namespace - -std::shared_ptr load_index_tts2_5_qwen_emotion_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType storage_type, - size_t weight_context_bytes) { - if (assets.qwen_emotion_weights == nullptr) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion requires tensor source"); - } - auto weights = std::make_shared(); - weights->store = std::make_shared( - backend, - backend_type, - "index_tts2_5.qwen_emotion.weights", - weight_context_bytes); - - const auto & source = *assets.qwen_emotion_weights; - weights->token_embedding = weights->store->load_tensor( - source, - "model.embed_tokens.weight", - storage_type, - {kVocab, kHidden}); - weights->decoder.layers.reserve(static_cast(kLayers)); - for (int64_t layer = 0; layer < kLayers; ++layer) { - weights->decoder.layers.push_back(load_qwen_layer(*weights->store, source, layer, storage_type)); - } - weights->final_norm = binding::norm_weight_from_source(*weights->store, source, "model.norm", kHidden); - weights->store->upload(); - assets.qwen_emotion_weights->release_storage(); - return weights; -} - -IndexTTS25QwenEmotionTokenizer::IndexTTS25QwenEmotionTokenizer(std::shared_ptr assets) { - if (assets == nullptr) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion tokenizer requires assets"); - } - engine::tokenizers::LlamaBpeTokenizerSpec spec; - spec.vocab_path = assets->resources.require_file("qwen_emotion_vocab"); - spec.merges_path = assets->resources.require_file("qwen_emotion_merges"); - spec.tokenizer_config_path = assets->resources.require_file("qwen_emotion_tokenizer_config"); - spec.tokenizer_json_path = assets->resources.require_file("qwen_emotion_tokenizer"); - spec.pre_type = engine::tokenizers::LlamaBpePreTokenizer::Qwen2; - tokenizer_ = engine::tokenizers::load_llama_bpe_tokenizer(spec); - if (const auto id = tokenizer_->find_token_id("<|endoftext|>"); id.has_value()) { - eos_token_id_ = *id; - } - if (const auto id = tokenizer_->find_token_id(""); id.has_value()) { - think_end_token_id_ = *id; - } -} - -std::vector IndexTTS25QwenEmotionTokenizer::encode_chat_prompt(const std::string & text) const { - std::vector ids = tokenizer_->encode("System: 文本情感分类", true); - ids.push_back(eos_token_id_); - std::string user_text = "\nHuman: " + text; - size_t trailing_spaces = 0; - while (trailing_spaces < user_text.size() && user_text[user_text.size() - trailing_spaces - 1] == ' ') { - ++trailing_spaces; - } - if (trailing_spaces > 0) { - user_text.resize(user_text.size() - trailing_spaces); - } - auto user = tokenizer_->encode(user_text, true); - ids.insert(ids.end(), user.begin(), user.end()); - if (trailing_spaces > 0) { - const auto space_token = tokenizer_->find_token_id("Ġ"); - if (!space_token.has_value()) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion tokenizer missing space token"); - } - ids.insert(ids.end(), trailing_spaces, *space_token); - } - ids.push_back(eos_token_id_); - auto assistant = tokenizer_->encode("\nAssistant:", true); - ids.insert(ids.end(), assistant.begin(), assistant.end()); - return ids; -} - -std::string IndexTTS25QwenEmotionTokenizer::decode(const std::vector & token_ids, bool skip_special_tokens) const { - return tokenizer_->decode(token_ids, skip_special_tokens); -} - -int32_t IndexTTS25QwenEmotionTokenizer::eos_token_id() const noexcept { - return eos_token_id_; -} - -int32_t IndexTTS25QwenEmotionTokenizer::think_end_token_id() const noexcept { - return think_end_token_id_; -} - -class IndexTTS25QwenEmotionRuntime::PrefillGraph { -public: - PrefillGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t prompt_steps, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - prompt_steps_(prompt_steps) { - if (prompt_steps_ <= 0) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion prefill graph requires prompt tokens"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion prefill graph context"); - } - ggml_init_params input_params{64ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion prefill input context"); - } - ggml_init_params output_params{16ull * 1024ull * 1024ull, nullptr, true}; - output_ctx_.reset(ggml_init(output_params)); - if (output_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion prefill output context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.qwen_emotion.prefill", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.qwen_emotion.prefill.inputs", - execution_.backend_type()}; - core::ModuleBuildContext output_ctx{ - output_ctx_.get(), - "index_tts2_5.qwen_emotion.prefill.outputs", - execution_.backend_type()}; - token_ids_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, prompt_steps_})).tensor; - positions_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({prompt_steps_})).tensor; - mask_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_steps_})).tensor; - ggml_set_input(token_ids_); - ggml_set_input(positions_); - ggml_set_input(mask_); - auto x = modules::EmbeddingModule({kVocab, kHidden}).build( - ctx, - core::wrap_tensor(token_ids_, core::TensorShape::from_dims({1, prompt_steps_}), GGML_TYPE_I32), - weights_->token_embedding); - auto outputs = modules::QwenDecoderStackModule(qwen_config()).build( - ctx, - x, - core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32), - weights_->decoder, - std::nullopt, - core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_steps_}), GGML_TYPE_F32)); - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(131072, prompt_steps_ * 4096)), false); - for (const auto & layer : outputs.state.layers) { - auto key = core::ensure_backend_addressable_layout(ctx, *layer.key); - auto value = core::ensure_backend_addressable_layout(ctx, *layer.value); - auto * key_output = core::make_tensor(output_ctx, GGML_TYPE_F32, key.shape).tensor; - auto * value_output = core::make_tensor(output_ctx, GGML_TYPE_F32, value.shape).tensor; - keys_.push_back(key_output); - values_.push_back(value_output); - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), key.tensor, key_output)); - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), value.tensor, value_output)); - } - auto last = modules::SliceModule({1, prompt_steps_ - 1, 1}).build(ctx, outputs.output); - last = modules::RMSNormModule({kHidden, kRmsEps, true, false}).build(ctx, last, weights_->final_norm); - auto logits = modules::LinearModule({kHidden, kVocab, false}).build(ctx, last, {weights_->token_embedding, std::nullopt}); - auto flat_logits = core::reshape_tensor( - ctx, - core::ensure_backend_addressable_layout(ctx, logits), - core::TensorShape::from_dims({1, kVocab})); - auto * next_token_source = ggml_argmax(ctx.ggml, flat_logits.tensor); - next_token_ = core::make_tensor(output_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1})).tensor; - ggml_set_output(next_token_); - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), next_token_source, next_token_)); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion prefill input buffer"); - } - output_buffer_ = ggml_backend_alloc_ctx_tensors(output_ctx_.get(), execution_.backend()); - if (output_buffer_ == nullptr) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion prefill output buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion prefill graph"); - } - std::vector positions(static_cast(prompt_steps_)); - for (int64_t i = 0; i < prompt_steps_; ++i) { - positions[static_cast(i)] = static_cast(i); - } - core::write_tensor_i32(core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32), positions); - const auto mask = causal_mask(prompt_steps_); - core::write_tensor_f32(core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_steps_}), GGML_TYPE_F32), mask); - debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.prefill.prompt_tokens", prompt_steps_); - } - - ~PrefillGraph() { - clear_graph(); - } - - bool matches(const IndexTTS25QwenEmotionWeights & weights, ggml_backend_t backend, int64_t prompt_steps) const noexcept { - return weights_.get() == &weights && execution_.backend() == backend && prompt_steps_ == prompt_steps; - } - - int32_t run(const std::vector & prompt_ids) { - if (static_cast(prompt_ids.size()) != prompt_steps_) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion prompt length mismatch"); - } - auto timing_start = Clock::now(); - ggml_backend_tensor_set(token_ids_, prompt_ids.data(), 0, prompt_ids.size() * sizeof(int32_t)); - debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_, nullptr, "IndexTTS2.5 Qwen emotion prefill"); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion prefill graph compute failed"); - } - timing_start = Clock::now(); - const auto next_token = core::read_tensor_i32(next_token_); - debug::timing_log_scalar("index_tts2_5.qwen_emotion.prefill.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - if (next_token.size() != 1) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion prefill argmax output shape mismatch"); - } - return next_token.front(); - } - - int64_t prompt_steps() const noexcept { - return prompt_steps_; - } - - size_t layer_count() const noexcept { - return keys_.size(); - } - - void copy_layer_state_to(size_t layer, ggml_tensor * key_destination, ggml_tensor * value_destination) const { - const int64_t values = prompt_steps_ * kKvHeads * kHeadDim; - std::vector key(static_cast(values)); - std::vector value(static_cast(values)); - ggml_backend_tensor_get(keys_.at(layer), key.data(), 0, key.size() * sizeof(float)); - ggml_backend_tensor_get(values_.at(layer), value.data(), 0, value.size() * sizeof(float)); - ggml_backend_tensor_set(key_destination, key.data(), 0, key.size() * sizeof(float)); - ggml_backend_tensor_set(value_destination, value.data(), 0, value.size() * sizeof(float)); - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - if (output_buffer_ != nullptr) { - ggml_backend_buffer_free(output_buffer_); - output_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t prompt_steps_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr output_ctx_; - std::unique_ptr ctx_; - ggml_tensor * token_ids_ = nullptr; - ggml_tensor * positions_ = nullptr; - ggml_tensor * mask_ = nullptr; - ggml_tensor * next_token_ = nullptr; - std::vector keys_; - std::vector values_; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; - ggml_backend_buffer_t output_buffer_ = nullptr; -}; - -class IndexTTS25QwenEmotionRuntime::DecodeGraph { -public: - DecodeGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t cache_steps, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - cache_steps_(cache_steps) { - if (cache_steps_ <= 0) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion decode graph requires cache capacity"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion decode graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion decode input context"); - } - ggml_init_params state_params{128ull * 1024ull * 1024ull, nullptr, true}; - state_ctx_.reset(ggml_init(state_params)); - if (state_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Qwen emotion decode state context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.qwen_emotion.decode", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.qwen_emotion.decode.inputs", - execution_.backend_type()}; - core::ModuleBuildContext state_ctx{ - state_ctx_.get(), - "index_tts2_5.qwen_emotion.decode.state", - execution_.backend_type()}; - token_id_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, 1})).tensor; - position_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1})).tensor; - mask_ = core::make_tensor(input_ctx, GGML_TYPE_F16, core::TensorShape::from_dims({1, 1, 1, cache_steps_ + 1})).tensor; - ggml_set_input(token_id_); - ggml_set_input(position_); - ggml_set_input(mask_); - graph_ = ggml_new_graph_custom(ctx_.get(), 131072, false); - std::vector cache_keys; - std::vector cache_values; - auto x = modules::EmbeddingModule({kVocab, kHidden}).build( - ctx, - core::wrap_tensor(token_id_, core::TensorShape::from_dims({1, 1}), GGML_TYPE_I32), - weights_->token_embedding); - const auto cfg = qwen_config(); - const modules::QwenDecoderLayerModule layer_module(modules::qwen_decoder_layer_config_from_stack(cfg)); - const auto mask = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_ + 1}), GGML_TYPE_F16); - for (const auto & layer : weights_->decoder.layers) { - cache_keys.push_back(core::make_tensor( - state_ctx, - GGML_TYPE_F32, - core::TensorShape::from_dims({1, cache_steps_ + 1, kKvHeads, kHeadDim}))); - cache_values.push_back(core::make_tensor( - state_ctx, - GGML_TYPE_F32, - core::TensorShape::from_dims({1, cache_steps_ + 1, kKvHeads, kHeadDim}))); - auto out = layer_module.build_with_static_cache_tail( - ctx, - graph_, - x, - core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32), - layer, - cache_keys.back(), - cache_values.back(), - std::nullopt, - mask); - x = out.output; - } - kv_cache_ = engine::runtime::TransformerKVCache(cache_steps_ + 1, kKvHeads * kHeadDim, std::move(cache_keys), std::move(cache_values)); - build_transfer_views(); - x = modules::RMSNormModule({kHidden, kRmsEps, true, false}).build(ctx, x, weights_->final_norm); - auto logits = modules::LinearModule({kHidden, kVocab, false}).build(ctx, x, {weights_->token_embedding, std::nullopt}); - auto flat_logits = core::reshape_tensor( - ctx, - core::ensure_backend_addressable_layout(ctx, logits), - core::TensorShape::from_dims({1, kVocab})); - next_token_ = ggml_argmax(ctx.ggml, flat_logits.tensor); - ggml_set_output(next_token_); - ggml_build_forward_expand(graph_, next_token_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion decode input buffer"); - } - state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), execution_.backend()); - if (state_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion decode state buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 Qwen emotion decode graph"); - } - mask_values_.assign(static_cast(cache_steps_ + 1), ggml_fp32_to_fp16(-std::numeric_limits::infinity())); - debug::timing_log_scalar("index_tts2_5.qwen_emotion.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.decode.cache_steps", cache_steps_); - } - - ~DecodeGraph() { - clear_graph(); - } - - bool can_run(const IndexTTS25QwenEmotionWeights & weights, ggml_backend_t backend, int64_t required_steps) const noexcept { - return weights_.get() == &weights && execution_.backend() == backend && cache_steps_ >= required_steps; - } - - void import_state(const PrefillGraph & prefill) { - if (prefill.layer_count() != key_sources_.size() || prefill.prompt_steps() > cache_steps_) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion decode prefill state shape mismatch"); - } - kv_cache_.retain_prefix(0); - const size_t prefix = static_cast(prefill.prompt_steps()); - for (size_t layer = 0; layer < key_sources_.size(); ++layer) { - prefill.copy_layer_state_to(layer, key_prefix_destinations_[prefix][layer], value_prefix_destinations_[prefix][layer]); - } - kv_cache_.advance_after_direct_append(prefill.prompt_steps()); - } - - int32_t run_step(int32_t token) { - if (kv_cache_.valid_steps() >= cache_steps_) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion decode cache exhausted"); - } - ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t)); - const int32_t position = static_cast(kv_cache_.current_end()); - ggml_backend_tensor_set(position_, &position, 0, sizeof(int32_t)); - std::fill(mask_values_.begin(), mask_values_.end(), ggml_fp32_to_fp16(-std::numeric_limits::infinity())); - for (int64_t i = 0; i < kv_cache_.valid_steps(); ++i) { - mask_values_[static_cast(i)] = ggml_fp32_to_fp16(0.0F); - } - mask_values_[static_cast(cache_steps_)] = ggml_fp32_to_fp16(0.0F); - ggml_backend_tensor_set(mask_, mask_values_.data(), 0, mask_values_.size() * sizeof(ggml_fp16_t)); - core::set_backend_threads(execution_.backend(), execution_.config().threads); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_, nullptr, "IndexTTS2.5 Qwen emotion decode"); - ggml_backend_synchronize(execution_.backend()); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion decode graph compute failed"); - } - const auto next_token = core::read_tensor_i32(next_token_); - if (next_token.size() != 1) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion decode argmax output shape mismatch"); - } - const size_t dst_slot = static_cast(kv_cache_.valid_steps()); - for (size_t layer = 0; layer < key_sources_.size(); ++layer) { - ggml_backend_tensor_copy(key_sources_[layer], key_destinations_[dst_slot][layer]); - ggml_backend_tensor_copy(value_sources_[layer], value_destinations_[dst_slot][layer]); - } - kv_cache_.advance_after_direct_append(1); - return next_token.front(); - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (state_buffer_ != nullptr) { - ggml_backend_buffer_free(state_buffer_); - state_buffer_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - void build_transfer_views() { - const int64_t step_elems = kKvHeads * kHeadDim; - const size_t scratch_offset = static_cast(cache_steps_ * step_elems) * sizeof(float); - key_sources_.clear(); - value_sources_.clear(); - key_sources_.reserve(weights_->decoder.layers.size()); - value_sources_.reserve(weights_->decoder.layers.size()); - for (size_t layer = 0; layer < weights_->decoder.layers.size(); ++layer) { - key_sources_.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.key_tensor(layer).tensor, step_elems, scratch_offset)); - value_sources_.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.value_tensor(layer).tensor, step_elems, scratch_offset)); - } - key_destinations_.assign(static_cast(cache_steps_), {}); - value_destinations_.assign(static_cast(cache_steps_), {}); - key_prefix_destinations_.assign(static_cast(cache_steps_ + 1), {}); - value_prefix_destinations_.assign(static_cast(cache_steps_ + 1), {}); - for (int64_t slot = 0; slot < cache_steps_; ++slot) { - const size_t byte_offset = static_cast(slot * step_elems) * sizeof(float); - auto & key_slot = key_destinations_[static_cast(slot)]; - auto & value_slot = value_destinations_[static_cast(slot)]; - key_slot.reserve(key_sources_.size()); - value_slot.reserve(value_sources_.size()); - for (size_t layer = 0; layer < key_sources_.size(); ++layer) { - key_slot.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.key_tensor(layer).tensor, step_elems, byte_offset)); - value_slot.push_back(ggml_view_1d(state_ctx_.get(), kv_cache_.value_tensor(layer).tensor, step_elems, byte_offset)); - } - } - for (int64_t prefix = 1; prefix <= cache_steps_; ++prefix) { - auto & key_prefix = key_prefix_destinations_[static_cast(prefix)]; - auto & value_prefix = value_prefix_destinations_[static_cast(prefix)]; - key_prefix.reserve(key_sources_.size()); - value_prefix.reserve(value_sources_.size()); - for (size_t layer = 0; layer < key_sources_.size(); ++layer) { - key_prefix.push_back(ggml_view_1d( - state_ctx_.get(), - kv_cache_.key_tensor(layer).tensor, - prefix * step_elems, - 0)); - value_prefix.push_back(ggml_view_1d( - state_ctx_.get(), - kv_cache_.value_tensor(layer).tensor, - prefix * step_elems, - 0)); - } - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t cache_steps_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr state_ctx_; - std::unique_ptr ctx_; - ggml_tensor * token_id_ = nullptr; - ggml_tensor * position_ = nullptr; - ggml_tensor * mask_ = nullptr; - ggml_tensor * next_token_ = nullptr; - std::vector key_sources_; - std::vector value_sources_; - std::vector> key_destinations_; - std::vector> value_destinations_; - std::vector> key_prefix_destinations_; - std::vector> value_prefix_destinations_; - std::vector mask_values_; - engine::runtime::TransformerKVCache kv_cache_; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; - ggml_backend_buffer_t state_buffer_ = nullptr; -}; - -IndexTTS25QwenEmotionRuntime::IndexTTS25QwenEmotionRuntime( - std::shared_ptr assets, - core::ExecutionContext & execution, - size_t prefill_graph_arena_bytes, - size_t decode_graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType storage_type) - : assets_(std::move(assets)), - execution_(&execution), - prefill_graph_arena_bytes_(prefill_graph_arena_bytes), - decode_graph_arena_bytes_(decode_graph_arena_bytes), - tokenizer_(assets_) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion runtime requires assets"); - } - if (prefill_graph_arena_bytes_ == 0 || decode_graph_arena_bytes_ == 0) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion graph arenas must be non-zero"); - } - weights_ = load_index_tts2_5_qwen_emotion_weights( - *assets_, - execution.backend(), - execution.backend_type(), - storage_type, - weight_context_bytes); -} - -IndexTTS25QwenEmotionRuntime::~IndexTTS25QwenEmotionRuntime() = default; - -IndexTTS25EmotionVector IndexTTS25QwenEmotionRuntime::infer(const std::string & text, int64_t max_new_tokens) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion runtime execution context is missing"); - } - if (max_new_tokens <= 0) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion max_new_tokens must be positive"); - } - const auto prompt_ids = tokenizer_.encode_chat_prompt(text); - const int64_t prompt_steps = static_cast(prompt_ids.size()); - if (prefill_graph_ == nullptr || !prefill_graph_->matches(*weights_, execution_->backend(), prompt_steps)) { - prefill_graph_.reset(); - prefill_graph_ = std::make_unique(*execution_, weights_, prompt_steps, prefill_graph_arena_bytes_); - } - const int64_t required_cache_steps = prompt_steps + max_new_tokens; - if (decode_graph_ == nullptr || !decode_graph_->can_run(*weights_, execution_->backend(), required_cache_steps)) { - decode_graph_.reset(); - decode_graph_ = std::make_unique(*execution_, weights_, required_cache_steps, decode_graph_arena_bytes_); - } - const int32_t prefill_token = prefill_graph_->run(prompt_ids); - decode_graph_->import_state(*prefill_graph_); - - std::vector generated; - generated.reserve(static_cast(max_new_tokens)); - int32_t token = prefill_token; - double decode_run_ms = 0.0; - bool saw_eos = false; - for (int64_t step = 0; step < max_new_tokens; ++step) { - if (token == tokenizer_.eos_token_id()) { - saw_eos = true; - break; - } - generated.push_back(token); - if (step + 1 >= max_new_tokens) { - break; - } - const auto decode_start = Clock::now(); - token = decode_graph_->run_step(token); - decode_run_ms += engine::debug::elapsed_ms(decode_start, Clock::now()); - } - if (!saw_eos && static_cast(generated.size()) >= max_new_tokens) { - throw std::runtime_error("IndexTTS2.5 Qwen emotion decode reached max_new_tokens before EOS"); - } - - size_t start = 0; - for (size_t i = generated.size(); i > 0; --i) { - if (generated[i - 1] == tokenizer_.think_end_token_id()) { - start = i; - break; - } - } - const std::vector answer(generated.begin() + static_cast(start), generated.end()); - const std::string content = tokenizer_.decode(answer, true); - debug::timing_log_scalar("index_tts2_5.qwen_emotion.decode.run_ms", decode_run_ms); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.generated_tokens", generated.size()); - return convert_emotion_json(content, text); -} - -void IndexTTS25QwenEmotionRuntime::release_graphs() { - prefill_graph_.reset(); - decode_graph_.reset(); -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/request.cpp b/src/models/index_tts2_5/request.cpp deleted file mode 100644 index 9c792973..00000000 --- a/src/models/index_tts2_5/request.cpp +++ /dev/null @@ -1,192 +0,0 @@ -#include "engine/models/index_tts2_5/request.h" - -#include "engine/framework/io/text.h" -#include "engine/framework/runtime/options.h" - -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -std::vector parse_emotion_vector(const std::string & value) { - std::string normalized = engine::io::trim_ascii_whitespace(value); - if (!normalized.empty() && normalized.front() == '[') { - normalized.erase(normalized.begin()); - } - if (!normalized.empty() && normalized.back() == ']') { - normalized.pop_back(); - } - std::vector values; - std::stringstream stream(normalized); - std::string item; - while (std::getline(stream, item, ',')) { - item = engine::io::trim_ascii_whitespace(item); - if (item.empty()) { - throw std::runtime_error("IndexTTS2.5 emotion_vector contains an empty item"); - } - size_t parsed = 0; - const float parsed_value = std::stof(item, &parsed); - if (parsed != item.size() || !std::isfinite(parsed_value)) { - throw std::runtime_error("IndexTTS2.5 emotion_vector must contain finite floats"); - } - values.push_back(parsed_value); - } - if (values.size() != 8) { - throw std::runtime_error("IndexTTS2.5 emotion_vector must contain exactly 8 values"); - } - return values; -} - -const runtime::AudioBuffer * speaker_audio_from_request(const runtime::TaskRequest & request) { - if (request.voice.has_value() && - request.voice->speaker.has_value() && - request.voice->speaker->audio.has_value()) { - return &*request.voice->speaker->audio; - } - return nullptr; -} - -void require_valid_audio(const runtime::AudioBuffer & audio, const char * label) { - if (audio.sample_rate <= 0 || audio.channels <= 0 || audio.samples.empty()) { - throw std::runtime_error(std::string("IndexTTS2.5 ") + label + " audio is empty or invalid"); - } -} - -} // namespace - -std::string normalize_index_tts2_5_lang(const std::string & value) { - std::string lang = engine::io::trim_ascii_whitespace(value); - std::transform(lang.begin(), lang.end(), lang.begin(), [](unsigned char ch) { - return static_cast(std::tolower(ch)); - }); - if (lang == "auto") { - lang.clear(); - } - return lang; -} - -IndexTTS25Request parse_index_tts2_5_request(const runtime::TaskRequest & request) { - IndexTTS25Request out; - if (request.text_input.has_value()) { - out.text = engine::io::trim_ascii_whitespace(request.text_input->text); - } else if (const auto value = runtime::find_option(request.options, {"text", "prompt"})) { - out.text = engine::io::trim_ascii_whitespace(*value); - } - if (out.text.empty()) { - throw std::runtime_error("IndexTTS2.5 request requires text_input or text option"); - } - - if (const auto * speaker = speaker_audio_from_request(request)) { - require_valid_audio(*speaker, "speaker reference"); - out.speaker_audio = *speaker; - } else { - throw std::runtime_error("IndexTTS2.5 request requires --voice-ref or voice.speaker.audio"); - } - - if (const auto value = runtime::find_option(request.options, {"lang"})) { - out.lang = normalize_index_tts2_5_lang(*value); - } - if (const auto value = runtime::parse_finite_float_option(request.options, {"emotion_alpha"})) { - if (*value < 0.0F || *value > 1.0F) { - throw std::runtime_error("IndexTTS2.5 emotion_alpha must be in [0, 1]"); - } - out.emotion_alpha = *value; - } - if (const auto value = runtime::find_option(request.options, {"emotion_vector"})) { - out.emotion_vector = parse_emotion_vector(*value); - } - const auto use_emotion_text_value = runtime::find_option(request.options, {"use_emotion_text"}); - if (use_emotion_text_value.has_value()) { - out.use_emotion_text = runtime::parse_bool_option(*use_emotion_text_value, "use_emotion_text"); - } - if (const auto value = runtime::find_option(request.options, {"emotion_text"})) { - if (!engine::io::trim_ascii_whitespace(*value).empty()) { - out.emotion_text = *value; - } - } - if (request.voice.has_value() && - request.voice->style.has_value() && - request.voice->style->emotion.has_value()) { - const auto & text = *request.voice->style->emotion; - if (!engine::io::trim_ascii_whitespace(text).empty()) { - if (use_emotion_text_value.has_value() && !out.use_emotion_text) { - throw std::runtime_error("IndexTTS2.5 --emotion conflicts with use_emotion_text=false"); - } - if (out.emotion_text.has_value() && *out.emotion_text != text) { - throw std::runtime_error("IndexTTS2.5 --emotion conflicts with emotion_text"); - } - out.use_emotion_text = true; - out.emotion_text = text; - } - } - if (const auto value = runtime::find_option(request.options, {"use_random_emotion"})) { - out.use_random_emotion = runtime::parse_bool_option(*value, "use_random_emotion"); - } - if (const auto value = runtime::parse_int_option(request.options, {"interval_silence_ms"})) { - if (*value < 0) { - throw std::runtime_error("IndexTTS2.5 interval_silence_ms must be non-negative"); - } - out.interval_silence_ms = *value; - } - if (request.audio_input.has_value()) { - require_valid_audio(*request.audio_input, "emotion reference"); - out.emotion_audio = request.audio_input; - } - if (out.emotion_vector.has_value() || out.use_emotion_text) { - out.emotion_audio = std::nullopt; - } - - if (const auto value = runtime::find_option(request.options, {"do_sample"})) { - out.generation.do_sample = runtime::parse_bool_option(*value, "do_sample"); - } - if (const auto value = runtime::parse_finite_float_option(request.options, {"top_p"})) { - out.generation.top_p = *value; - } - if (const auto value = runtime::parse_int_option(request.options, {"top_k"})) { - out.generation.top_k = *value; - } - if (const auto value = runtime::parse_finite_float_option(request.options, {"temperature"})) { - out.generation.temperature = *value; - } - if (const auto value = runtime::parse_finite_float_option(request.options, {"length_penalty"})) { - out.generation.length_penalty = *value; - } - if (const auto value = runtime::parse_int_option(request.options, {"num_beams"})) { - out.generation.num_beams = *value; - } - if (const auto value = runtime::parse_finite_float_option(request.options, {"repetition_penalty"})) { - out.generation.repetition_penalty = *value; - } - if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { - if (*value <= 0) { - throw std::runtime_error("IndexTTS2.5 max_tokens must be positive"); - } - out.generation.max_mel_tokens = *value; - } - if (const auto value = runtime::parse_u32_option(request.options, {"seed"})) { - out.generation.seed = *value; - } else { - out.generation.seed = runtime::random_u32_seed(); - } - - if (out.generation.top_k <= 0) { - throw std::runtime_error("IndexTTS2.5 top_k must be positive"); - } - if (!(out.generation.top_p > 0.0F && out.generation.top_p <= 1.0F)) { - throw std::runtime_error("IndexTTS2.5 top_p must be in (0, 1]"); - } - if (!(out.generation.temperature > 0.0F)) { - throw std::runtime_error("IndexTTS2.5 temperature must be positive"); - } - if (out.generation.num_beams <= 0) { - throw std::runtime_error("IndexTTS2.5 num_beams must be positive"); - } - return out; -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/s2mel.cpp b/src/models/index_tts2_5/s2mel.cpp deleted file mode 100644 index 4be40939..00000000 --- a/src/models/index_tts2_5/s2mel.cpp +++ /dev/null @@ -1,1483 +0,0 @@ -#include "engine/models/index_tts2_5/s2mel.h" - -#include "engine/framework/core/backend.h" -#include "engine/framework/debug/profiler.h" -#include "engine/framework/modules/activation_modules.h" -#include "engine/framework/modules/positional_modules.h" -#include "engine/framework/modules/primitive_modules.h" -#include "engine/framework/modules/structural_modules.h" -#include "engine/framework/modules/weight_binding.h" -#include "engine/framework/sampling/torch_random.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -namespace binding = engine::modules::binding; -namespace core = engine::core; -namespace modules = engine::modules; -using Clock = std::chrono::steady_clock; - -constexpr int64_t kMelChannels = 80; -constexpr int64_t kContentDim = 1024; -constexpr int64_t kGptDim = 1280; -constexpr int64_t kHidden = 512; -constexpr int64_t kStyleDim = 192; -constexpr int64_t kDitLayers = 13; -constexpr int64_t kWavenetLayers = 8; -constexpr int64_t kWavenetKernel = 5; -constexpr int64_t kTimeFreqDim = 128; -constexpr int64_t kTimeEmbeddingDim = 256; -constexpr int64_t kDitFfnDim = 1536; -constexpr int64_t kDitHeads = 8; -constexpr int64_t kDitHeadDim = kHidden / kDitHeads; -constexpr float kLayerNormEps = 1.0e-6F; -constexpr float kRmsNormEps = 1.0e-5F; - -struct GgmlContextDeleter { - void operator()(ggml_context * ctx) const noexcept { - if (ctx != nullptr) { - ggml_free(ctx); - } - } -}; - -// Debug helpers for backend divergence investigation, enabled by environment -// variables: INDEXTTS25_DUMP_DIR dumps the initial CFM noise and the first -// diffusion velocity as NPY files; INDEXTTS25_CFM_NOISE_NPY replaces the RNG -// initial noise with the contents of an NPY float32 file. -struct CfmDebugTap { - std::string name; - core::TensorValue value; -}; - -std::string cfm_debug_dump_dir() { - const char * dir = std::getenv("INDEXTTS25_DUMP_DIR"); - if (dir == nullptr || *dir == '\0') { - return {}; - } - return std::string(dir); -} - -void cfm_write_npy_f32( - const std::string & dir, - const std::string & name, - const std::vector & shape, - const std::vector & values) { - if (dir.empty()) { - return; - } - std::string header = "{'descr': '(header.size()); - out.write(reinterpret_cast(&header_len), sizeof(header_len)); - out.write(header.data(), static_cast(header.size())); - out.write(reinterpret_cast(values.data()), static_cast(values.size() * sizeof(float))); -} - -std::vector cfm_read_npy_f32(const std::string & path, size_t expected_count) { - std::ifstream in(path, std::ios::binary); - if (!in) { - throw std::runtime_error("IndexTTS2.5 failed to open noise file: " + path); - } - char magic[8]; - in.read(magic, 8); - if (!in || std::memcmp(magic, "\x93NUMPY\x01\x00", 8) != 0) { - throw std::runtime_error("IndexTTS2.5 noise file is not an NPY v1 file: " + path); - } - uint16_t header_len = 0; - in.read(reinterpret_cast(&header_len), sizeof(header_len)); - std::string header(header_len, '\0'); - in.read(header.data(), header_len); - if (!in || header.find("'(std::stoll(token.substr(first))); - } - if (comma == std::string::npos) { - break; - } - begin = comma + 1; - } - if (count != expected_count) { - throw std::runtime_error("IndexTTS2.5 noise file element count mismatch: " + path); - } - std::vector out(count); - in.read(reinterpret_cast(out.data()), static_cast(count * sizeof(float))); - if (!in) { - throw std::runtime_error("IndexTTS2.5 noise file data is truncated: " + path); - } - return out; -} - -std::string cfm_noise_path_from_env() { - const char * path = std::getenv("INDEXTTS25_CFM_NOISE_NPY"); - if (path == nullptr || *path == '\0') { - return {}; - } - return std::string(path); -} - - -core::TensorValue sub(core::ModuleBuildContext & ctx, const core::TensorValue & lhs, const core::TensorValue & rhs) { - core::validate_shape(rhs, lhs.shape, "Sub rhs"); - return core::wrap_tensor(ggml_sub(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32); -} - -core::TensorValue scale(core::ModuleBuildContext & ctx, const core::TensorValue & input, float value) { - return core::wrap_tensor(ggml_scale(ctx.ggml, input.tensor, value), input.shape, GGML_TYPE_F32); -} - -core::TensorValue add_one(core::ModuleBuildContext & ctx, const core::TensorValue & input) { - return core::wrap_tensor(ggml_scale_bias(ctx.ggml, input.tensor, 1.0F, 1.0F), input.shape, GGML_TYPE_F32); -} - -core::TensorValue reshape_heads( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - int64_t heads, - int64_t dim) { - return core::reshape_tensor( - ctx, - core::ensure_backend_addressable_layout(ctx, input), - core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim})); -} - -core::TensorValue slice_last( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - int64_t start, - int64_t length) { - return modules::SliceModule({static_cast(input.shape.rank - 1), start, length}).build(ctx, input); -} - -core::TensorValue apply_channel_affine( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & weight, - const core::TensorValue & bias, - int64_t channels) { - core::TensorShape broadcast_shape = {}; - broadcast_shape.rank = input.shape.rank; - for (size_t axis = 0; axis < broadcast_shape.rank; ++axis) { - broadcast_shape.dims[axis] = 1; - } - broadcast_shape.dims[1] = channels; - auto weight_view = core::reshape_tensor(ctx, weight, broadcast_shape); - auto bias_view = core::reshape_tensor(ctx, bias, broadcast_shape); - auto weight_rep = modules::RepeatModule({input.shape}).build(ctx, weight_view); - auto bias_rep = modules::RepeatModule({input.shape}).build(ctx, bias_view); - return modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, input, weight_rep), bias_rep); -} - -core::TensorValue broadcast_batch_time( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - int64_t batch, - int64_t frames, - int64_t dims) { - auto shaped = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, input), core::TensorShape::from_dims({batch, 1, dims})); - return modules::RepeatModule({core::TensorShape::from_dims({batch, frames, dims})}).build(ctx, shaped); -} - -core::TensorValue group_norm_1_group( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const modules::NormWeights & weights, - int64_t channels) { - if (!weights.weight.has_value() || !weights.bias.has_value()) { - throw std::runtime_error("IndexTTS2.5 S2Mel length regulator group norm requires affine weights"); - } - const auto input4 = core::reshape_tensor( - ctx, - core::ensure_backend_addressable_layout(ctx, input), - core::TensorShape::from_dims({input.shape.dims[0], channels, 1, input.shape.dims[2]})); - auto normalized = core::wrap_tensor(ggml_group_norm(ctx.ggml, input4.tensor, 1, 1.0e-5F), input4.shape, GGML_TYPE_F32); - normalized = apply_channel_affine(ctx, normalized, *weights.weight, *weights.bias, channels); - return core::reshape_tensor(ctx, normalized, input.shape); -} - -core::TensorValue mish(core::ModuleBuildContext & ctx, const core::TensorValue & input) { - const auto softplus = core::wrap_tensor(ggml_softplus(ctx.ggml, input.tensor), input.shape, GGML_TYPE_F32); - const auto tanh = core::wrap_tensor(ggml_tanh(ctx.ggml, softplus.tensor), input.shape, GGML_TYPE_F32); - return core::wrap_tensor(ggml_mul(ctx.ggml, input.tensor, tanh.tensor), input.shape, GGML_TYPE_F32); -} - -core::TensorValue timestep_embedding( - core::ModuleBuildContext & ctx, - const core::TensorValue & timestep, - const core::TensorValue & freqs, - const modules::LinearWeights & linear0, - const modules::LinearWeights & linear2) { - const int64_t batch = timestep.shape.dims[0]; - auto t = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, timestep), core::TensorShape::from_dims({batch, 1})); - auto freqs_batched = modules::RepeatModule({core::TensorShape::from_dims({batch, kTimeFreqDim})}) - .build(ctx, core::reshape_tensor(ctx, freqs, core::TensorShape::from_dims({1, kTimeFreqDim}))); - auto args = modules::MulModule{}.build(ctx, modules::RepeatModule({freqs_batched.shape}).build(ctx, t), freqs_batched); - args = scale(ctx, args, 1000.0F); - auto cos_part = core::wrap_tensor(ggml_cos(ctx.ggml, core::ensure_backend_addressable_layout(ctx, args).tensor), args.shape, GGML_TYPE_F32); - auto sin_part = core::wrap_tensor(ggml_sin(ctx.ggml, args.tensor), args.shape, GGML_TYPE_F32); - auto emb = modules::ConcatModule({1}).build(ctx, cos_part, sin_part); - emb = modules::LinearModule({kTimeEmbeddingDim, kHidden, true, GGML_PREC_F32}).build(ctx, emb, linear0); - emb = modules::SiluModule{}.build(ctx, emb); - return modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, emb, linear2); -} - -core::TensorValue adaptive_rms_norm( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & embedding, - const IndexTTS25AdaLayerNormWeights & weights) { - auto projected = modules::LinearModule({kHidden, 2 * kHidden, true, GGML_PREC_F32}).build(ctx, embedding, weights.project); - auto weight = broadcast_batch_time(ctx, slice_last(ctx, projected, 0, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); - auto bias = broadcast_batch_time(ctx, slice_last(ctx, projected, kHidden, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); - auto normed = modules::RMSNormModule({kHidden, kRmsNormEps, true, false}).build(ctx, input, {weights.norm_weight, std::nullopt}); - return modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, normed, weight), bias); -} - -core::TensorValue cfm_attention( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & positions, - const IndexTTS25DitLayerWeights & weights) { - auto qkv = modules::LinearModule({kHidden, 3 * kHidden, false, GGML_PREC_F32}).build(ctx, input, weights.qkv); - auto q = slice_last(ctx, qkv, 0, kHidden); - auto k = slice_last(ctx, qkv, kHidden, kHidden); - auto v = slice_last(ctx, qkv, 2 * kHidden, kHidden); - q = modules::RoPEModule({kDitHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, reshape_heads(ctx, q, kDitHeads, kDitHeadDim), positions); - k = modules::RoPEModule({kDitHeadDim, GGML_ROPE_TYPE_NORMAL, 10000.0F}).build(ctx, reshape_heads(ctx, k, kDitHeads, kDitHeadDim), positions); - v = reshape_heads(ctx, v, kDitHeads, kDitHeadDim); - auto qh = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, q); - auto kh = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k); - auto vh = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v); - auto * flash = ggml_flash_attn_ext( - ctx.ggml, - core::ensure_backend_addressable_layout(ctx, qh).tensor, - core::ensure_backend_addressable_layout(ctx, kh).tensor, - core::ensure_backend_addressable_layout(ctx, vh).tensor, - nullptr, - 1.0F / std::sqrt(static_cast(kDitHeadDim)), - 0.0F, - 0.0F); - ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); - auto context = core::wrap_tensor( - flash, - core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], kDitHeads, kDitHeadDim}), - GGML_TYPE_F32); - context = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, context), input.shape); - return modules::LinearModule({kHidden, kHidden, false, GGML_PREC_F32}).build(ctx, context, weights.attention_out); -} - -core::TensorValue cfm_ffn( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const IndexTTS25DitLayerWeights & weights) { - auto gate = modules::LinearModule({kHidden, kDitFfnDim, false, GGML_PREC_F32}).build(ctx, input, weights.ffn_w1); - gate = modules::SiluModule{}.build(ctx, gate); - auto up = modules::LinearModule({kHidden, kDitFfnDim, false, GGML_PREC_F32}).build(ctx, input, weights.ffn_w3); - auto hidden = modules::MulModule{}.build(ctx, gate, up); - return modules::LinearModule({kDitFfnDim, kHidden, false, GGML_PREC_F32}).build(ctx, hidden, weights.ffn_w2); -} - -core::TensorValue cfm_transformer_layer( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & timestep, - const core::TensorValue & positions, - const IndexTTS25DitLayerWeights & weights, - const core::TensorValue * skip) { - auto x = input; - if (skip != nullptr) { - x = modules::LinearModule({2 * kHidden, kHidden, true, GGML_PREC_F32}) - .build(ctx, modules::ConcatModule({2}).build(ctx, x, *skip), weights.skip_in); - } - auto attn = cfm_attention(ctx, adaptive_rms_norm(ctx, x, timestep, weights.attention_norm), positions, weights); - auto h = modules::AddModule{}.build(ctx, x, attn); - auto ff = cfm_ffn(ctx, adaptive_rms_norm(ctx, h, timestep, weights.ffn_norm), weights); - return modules::AddModule{}.build(ctx, h, ff); -} - -core::TensorValue cfm_wavenet( - core::ModuleBuildContext & ctx, - const core::TensorValue & input_bct, - const core::TensorValue & timestep_b, - const IndexTTS25S2MelCfmWeights & weights, - std::vector * debug_taps = nullptr) { - const auto tap = [&](const std::string & name, const core::TensorValue & value) { - if (debug_taps != nullptr) { - debug_taps->push_back({name, core::ensure_backend_addressable_layout(ctx, value)}); - } - }; - auto g = core::reshape_tensor(ctx, core::ensure_backend_addressable_layout(ctx, timestep_b), core::TensorShape::from_dims({timestep_b.shape.dims[0], kHidden, 1})); - g = modules::Conv1dModule({kHidden, 2 * kHidden * kWavenetLayers, 1, 1, 0, 1, true}).build(ctx, g, weights.wavenet_cond); - tap("wn_g", g); - // The zero accumulator must come from a contiguous tensor: input_bct is a - // permuted (transposed) view, and the ggml CPU binary-op kernels miscompute - // permuted src operands (the CUDA kernels handle them). - const auto zeros_base = core::ensure_backend_addressable_layout(ctx, input_bct); - auto output = sub(ctx, zeros_base, zeros_base); - tap("wn_zeros", output); - auto x = input_bct; - for (int64_t i = 0; i < kWavenetLayers; ++i) { - const int64_t dilation = 1; - const int64_t padding = (kWavenetKernel * dilation - dilation) / 2; - auto x_padded = modules::ReflectPad1dModule({padding, padding}).build(ctx, core::ensure_backend_addressable_layout(ctx, x)); - auto x_in = modules::Conv1dModule( - {kHidden, 2 * kHidden, kWavenetKernel, 1, 0, static_cast(dilation), true}) - .build(ctx, x_padded, weights.wavenet_layers[static_cast(i)].in_layer); - if (i == 0) { - tap("wn_x_padded0", x_padded); - tap("wn_xin0", x_in); - } - if (i == 1) { - tap("wn_xin1", x_in); - } - auto g_l = modules::SliceModule({1, i * 2 * kHidden, 2 * kHidden}).build(ctx, g); - g_l = modules::RepeatModule({x_in.shape}).build(ctx, g_l); - auto acts = modules::AddModule{}.build(ctx, x_in, g_l); - auto tanh_part = modules::SliceModule({1, 0, kHidden}).build(ctx, acts); - tanh_part = modules::TanhModule{}.build(ctx, tanh_part); - auto sigmoid_part = modules::SliceModule({1, kHidden, kHidden}).build(ctx, acts); - sigmoid_part = modules::SigmoidModule{}.build(ctx, sigmoid_part); - acts = modules::MulModule{}.build(ctx, tanh_part, sigmoid_part); - if (i == 0) { - tap("wn_acts0", acts); - } - const int64_t res_skip_channels = i < kWavenetLayers - 1 ? 2 * kHidden : kHidden; - auto res_skip = modules::Conv1dModule({kHidden, res_skip_channels, 1, 1, 0, 1, true}) - .build(ctx, acts, weights.wavenet_layers[static_cast(i)].res_skip_layer); - if (i == 0) { - tap("wn_res_skip0", res_skip); - } - if (i < kWavenetLayers - 1) { - auto res = modules::SliceModule({1, 0, kHidden}).build(ctx, res_skip); - auto skip = modules::SliceModule({1, kHidden, kHidden}).build(ctx, res_skip); - x = modules::AddModule{}.build(ctx, x, res); - output = modules::AddModule{}.build(ctx, output, skip); - } else { - output = modules::AddModule{}.build(ctx, output, res_skip); - } - if (i == 0) { - tap("wn_x1", x); - tap("wn_output1", output); - } - if (i == 1) { - tap("wn_output2", output); - } - if (i == 2) { - tap("wn_output3", output); - } - } - tap("wn_final", output); - return output; -} - -core::TensorValue cfm_final_layer( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & timestep, - const IndexTTS25S2MelCfmWeights & weights) { - auto mod = modules::SiluModule{}.build(ctx, timestep); - mod = modules::LinearModule({kHidden, 2 * kHidden, true, GGML_PREC_F32}).build(ctx, mod, weights.final_modulation); - auto shift = broadcast_batch_time(ctx, slice_last(ctx, mod, 0, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); - auto scale_v = broadcast_batch_time(ctx, slice_last(ctx, mod, kHidden, kHidden), input.shape.dims[0], input.shape.dims[1], kHidden); - auto normed = modules::LayerNormModule({kHidden, kLayerNormEps, false, false}).build(ctx, input, {std::nullopt, std::nullopt}); - normed = modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, normed, add_one(ctx, scale_v)), shift); - return modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, normed, weights.final_linear); -} - -core::TensorValue build_cfm_estimator( - core::ModuleBuildContext & ctx, - const core::TensorValue & x_bct, - const core::TensorValue & prompt_bct, - const core::TensorValue & cond_btc, - const core::TensorValue & style_bc, - const core::TensorValue & timestep_b, - const core::TensorValue & positions, - const IndexTTS25S2MelCfmWeights & weights, - std::vector * debug_taps = nullptr) { - const auto tap = [&](const std::string & name, const core::TensorValue & value) { - if (debug_taps != nullptr) { - debug_taps->push_back({name, core::ensure_backend_addressable_layout(ctx, value)}); - } - }; - const int64_t batch = x_bct.shape.dims[0]; - const int64_t frames = x_bct.shape.dims[2]; - auto t1 = timestep_embedding(ctx, timestep_b, weights.time_freqs, weights.time_mlp0, weights.time_mlp2); - tap("t1", t1); - auto cond = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, cond_btc, weights.cond_projection); - auto x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x_bct); - auto prompt = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, prompt_bct); - auto style = broadcast_batch_time(ctx, style_bc, batch, frames, kStyleDim); - auto hidden = modules::ConcatModule({2}).build(ctx, x, prompt); - hidden = modules::ConcatModule({2}).build(ctx, hidden, cond); - hidden = modules::ConcatModule({2}).build(ctx, hidden, style); - hidden = modules::LinearModule({kHidden + 2 * kMelChannels + kStyleDim, kHidden, true, GGML_PREC_F32}) - .build(ctx, hidden, weights.cond_x_merge); - tap("merged", hidden); - - std::vector skips; - skips.reserve(static_cast(kDitLayers / 2)); - for (int64_t i = 0; i < kDitLayers; ++i) { - const core::TensorValue * skip = nullptr; - if (i > kDitLayers / 2) { - skip = &skips.back(); - } - hidden = cfm_transformer_layer(ctx, hidden, t1, positions, weights.dit_layers[static_cast(i)], skip); - if (i == 0) { - tap("dit_layer0", hidden); - } - if (i > kDitLayers / 2) { - skips.pop_back(); - } else if (i < kDitLayers / 2) { - skips.push_back(hidden); - } - } - hidden = adaptive_rms_norm(ctx, hidden, t1, weights.dit_norm); - tap("dit_out", hidden); - hidden = modules::LinearModule({kHidden + kMelChannels, kHidden, true, GGML_PREC_F32}) - .build(ctx, modules::ConcatModule({2}).build(ctx, hidden, x), weights.skip_linear); - auto wavenet_x = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, weights.conv1); - wavenet_x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, wavenet_x); - auto t2 = timestep_embedding(ctx, timestep_b, weights.time2_freqs, weights.time2_mlp0, weights.time2_mlp2); - wavenet_x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, cfm_wavenet(ctx, wavenet_x, t2, weights, debug_taps)); - auto projected = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, weights.res_projection); - hidden = modules::AddModule{}.build(ctx, wavenet_x, projected); - tap("wavenet_out", hidden); - hidden = cfm_final_layer(ctx, hidden, t1, weights); - tap("final_hidden", hidden); - hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, hidden); - return modules::Conv1dModule({kHidden, kMelChannels, 1, 1, 0, 1, true}).build(ctx, hidden, weights.conv2); -} - -std::vector fuse_weight_norm_linear( - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t out_features, - int64_t in_features) { - const auto g = source.require_f32(prefix + ".weight_g", {out_features, 1}); - const auto v = source.require_f32(prefix + ".weight_v", {out_features, in_features}); - std::vector weight(v.size(), 0.0F); - for (int64_t out = 0; out < out_features; ++out) { - double norm = 0.0; - for (int64_t in = 0; in < in_features; ++in) { - const float value = v[static_cast(out * in_features + in)]; - norm += static_cast(value) * static_cast(value); - } - const float scale = g[static_cast(out)] / static_cast(std::sqrt(norm)); - for (int64_t in = 0; in < in_features; ++in) { - const size_t index = static_cast(out * in_features + in); - weight[index] = v[index] * scale; - } - } - return weight; -} - -std::vector fuse_weight_norm_conv1d( - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t out_channels, - int64_t in_channels, - int64_t kernel_size) { - const auto g = source.require_f32(prefix + ".weight_g", {out_channels, 1, 1}); - const auto v = source.require_f32(prefix + ".weight_v", {out_channels, in_channels, kernel_size}); - std::vector weight(v.size(), 0.0F); - for (int64_t out = 0; out < out_channels; ++out) { - double norm = 0.0; - for (int64_t in = 0; in < in_channels; ++in) { - for (int64_t k = 0; k < kernel_size; ++k) { - const float value = v[static_cast((out * in_channels + in) * kernel_size + k)]; - norm += static_cast(value) * static_cast(value); - } - } - const float scale = g[static_cast(out)] / static_cast(std::sqrt(norm)); - for (int64_t in = 0; in < in_channels; ++in) { - for (int64_t k = 0; k < kernel_size; ++k) { - const size_t index = static_cast((out * in_channels + in) * kernel_size + k); - weight[index] = v[index] * scale; - } - } - } - return weight; -} - -engine::modules::LinearWeights load_weight_norm_linear( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type, - int64_t out_features, - int64_t in_features) { - engine::modules::LinearWeights weights; - weights.weight = store.make_from_f32( - engine::core::TensorShape::from_dims({out_features, in_features}), - storage_type, - fuse_weight_norm_linear(source, prefix, out_features, in_features)); - weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_features}); - return weights; -} - -engine::modules::Conv1dWeights load_weight_norm_conv1d( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type, - int64_t out_channels, - int64_t in_channels, - int64_t kernel_size) { - engine::modules::Conv1dWeights weights; - weights.weight = store.make_from_f32( - engine::core::TensorShape::from_dims({out_channels, in_channels, kernel_size}), - storage_type, - fuse_weight_norm_conv1d(source, prefix, out_channels, in_channels, kernel_size)); - weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); - return weights; -} - -IndexTTS25AdaLayerNormWeights load_ada_norm( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type) { - return { - store.load_f32_tensor(source, prefix + ".norm.weight", {kHidden}), - binding::linear_from_source(store, source, prefix + ".project_layer", storage_type, 2 * kHidden, kHidden, true), - }; -} - -IndexTTS25DitLayerWeights load_dit_layer( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - int64_t layer_index, - engine::assets::TensorStorageType storage_type) { - const std::string prefix = "cfm.estimator.transformer.layers." + std::to_string(layer_index); - IndexTTS25DitLayerWeights layer; - layer.attention_norm = load_ada_norm(store, source, prefix + ".attention_norm", storage_type); - layer.qkv = binding::linear_from_source(store, source, prefix + ".attention.wqkv", storage_type, 3 * kHidden, kHidden, false); - layer.attention_out = binding::linear_from_source(store, source, prefix + ".attention.wo", storage_type, kHidden, kHidden, false); - layer.ffn_norm = load_ada_norm(store, source, prefix + ".ffn_norm", storage_type); - layer.ffn_w1 = binding::linear_from_source(store, source, prefix + ".feed_forward.w1", storage_type, kDitFfnDim, kHidden, false); - layer.ffn_w2 = binding::linear_from_source(store, source, prefix + ".feed_forward.w2", storage_type, kHidden, kDitFfnDim, false); - layer.ffn_w3 = binding::linear_from_source(store, source, prefix + ".feed_forward.w3", storage_type, kDitFfnDim, kHidden, false); - layer.skip_in = binding::linear_from_source(store, source, prefix + ".skip_in_linear", storage_type, kHidden, 2 * kHidden, true); - return layer; -} - -IndexTTS25LengthRegulatorWeights load_length_regulator( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) { - IndexTTS25LengthRegulatorWeights weights; - weights.content_projection = binding::linear_from_source( - store, - source, - "length_regulator.content_in_proj", - matmul_storage_type, - kHidden, - kContentDim, - true); - for (int64_t i : {0, 3, 6, 9}) { - weights.convs.push_back(binding::conv1d_from_source( - store, - source, - "length_regulator.model." + std::to_string(i), - conv_storage_type, - kHidden, - kHidden, - 3, - true)); - } - for (int64_t i : {1, 4, 7, 10}) { - weights.norms.push_back({ - store.load_f32_tensor(source, "length_regulator.model." + std::to_string(i) + ".weight", {kHidden}), - store.load_f32_tensor(source, "length_regulator.model." + std::to_string(i) + ".bias", {kHidden}), - }); - } - weights.output = binding::conv1d_from_source( - store, - source, - "length_regulator.model.12", - conv_storage_type, - kHidden, - kHidden, - 1, - true); - return weights; -} - -IndexTTS25S2MelCfmWeights load_cfm( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) { - IndexTTS25S2MelCfmWeights weights; - weights.x_embedder = load_weight_norm_linear(store, source, "cfm.estimator.x_embedder", matmul_storage_type, kHidden, kMelChannels); - weights.cond_projection = binding::linear_from_source(store, source, "cfm.estimator.cond_projection", matmul_storage_type, kHidden, kHidden, true); - weights.cond_x_merge = binding::linear_from_source( - store, - source, - "cfm.estimator.cond_x_merge_linear", - matmul_storage_type, - kHidden, - kHidden + 2 * kMelChannels + kStyleDim, - true); - weights.skip_linear = binding::linear_from_source( - store, - source, - "cfm.estimator.skip_linear", - matmul_storage_type, - kHidden, - kHidden + kMelChannels, - true); - weights.time_freqs = store.load_f32_tensor(source, "cfm.estimator.t_embedder.freqs", {kTimeFreqDim}); - weights.time_mlp0 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder.mlp.0", matmul_storage_type, kHidden, kTimeEmbeddingDim, true); - weights.time_mlp2 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder.mlp.2", matmul_storage_type, kHidden, kHidden, true); - weights.time2_freqs = store.load_f32_tensor(source, "cfm.estimator.t_embedder2.freqs", {kTimeFreqDim}); - weights.time2_mlp0 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder2.mlp.0", matmul_storage_type, kHidden, kTimeEmbeddingDim, true); - weights.time2_mlp2 = binding::linear_from_source(store, source, "cfm.estimator.t_embedder2.mlp.2", matmul_storage_type, kHidden, kHidden, true); - weights.dit_layers.reserve(static_cast(kDitLayers)); - for (int64_t i = 0; i < kDitLayers; ++i) { - weights.dit_layers.push_back(load_dit_layer(store, source, i, matmul_storage_type)); - } - weights.dit_norm = load_ada_norm(store, source, "cfm.estimator.transformer.norm", matmul_storage_type); - weights.conv1 = binding::linear_from_source(store, source, "cfm.estimator.conv1", matmul_storage_type, kHidden, kHidden, true); - weights.res_projection = binding::linear_from_source(store, source, "cfm.estimator.res_projection", matmul_storage_type, kHidden, kHidden, true); - weights.wavenet_cond = load_weight_norm_conv1d( - store, - source, - "cfm.estimator.wavenet.cond_layer.conv.conv", - conv_storage_type, - 2 * kHidden * kWavenetLayers, - kHidden, - 1); - weights.wavenet_layers.reserve(static_cast(kWavenetLayers)); - for (int64_t i = 0; i < kWavenetLayers; ++i) { - const int64_t res_skip_channels = i < kWavenetLayers - 1 ? 2 * kHidden : kHidden; - weights.wavenet_layers.push_back({ - load_weight_norm_conv1d( - store, - source, - "cfm.estimator.wavenet.in_layers." + std::to_string(i) + ".conv.conv", - conv_storage_type, - 2 * kHidden, - kHidden, - kWavenetKernel), - load_weight_norm_conv1d( - store, - source, - "cfm.estimator.wavenet.res_skip_layers." + std::to_string(i) + ".conv.conv", - conv_storage_type, - res_skip_channels, - kHidden, - 1), - }); - } - weights.final_modulation = binding::linear_from_source( - store, - source, - "cfm.estimator.final_layer.adaLN_modulation.1", - matmul_storage_type, - 2 * kHidden, - kHidden, - true); - weights.final_linear = load_weight_norm_linear( - store, - source, - "cfm.estimator.final_layer.linear", - matmul_storage_type, - kHidden, - kHidden); - weights.conv2 = binding::conv1d_from_source(store, source, "cfm.estimator.conv2", conv_storage_type, kMelChannels, kHidden, 1, true); - return weights; -} - -} // namespace - -std::shared_ptr load_index_tts2_5_s2mel_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes) { - if (assets.s2mel_weights == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel requires tensor source"); - } - auto weights = std::make_shared(); - weights->store = std::make_shared( - backend, - backend_type, - "index_tts2_5.s2mel.weights", - weight_context_bytes); - - const auto & source = *assets.s2mel_weights; - weights->gpt_layer.linear0 = binding::linear_from_source(*weights->store, source, "gpt_layer.0", matmul_storage_type, 256, 1280, true); - weights->gpt_layer.linear1 = binding::linear_from_source(*weights->store, source, "gpt_layer.1", matmul_storage_type, 128, 256, true); - weights->gpt_layer.linear2 = binding::linear_from_source(*weights->store, source, "gpt_layer.2", matmul_storage_type, kContentDim, 128, true); - weights->length_regulator = load_length_regulator(*weights->store, source, matmul_storage_type, conv_storage_type); - weights->cfm = load_cfm(*weights->store, source, matmul_storage_type, conv_storage_type); - - weights->store->upload(); - assets.s2mel_weights->release_storage(); - return weights; -} - -class IndexTTS25S2MelRuntime::GptLayerGraph { -public: - GptLayerGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t frames, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - frames_(frames) { - if (frames_ <= 0) { - throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph requires positive frame count"); - } - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph requires weights"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel GPT layer graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel GPT layer input context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.s2mel.gpt_layer", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.s2mel.gpt_layer.inputs", - execution_.backend_type()}; - input_ = - core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kGptDim})).tensor; - ggml_set_input(input_); - auto x = core::wrap_tensor(input_, core::TensorShape::from_dims({1, frames_, kGptDim}), GGML_TYPE_F32); - x = modules::LinearModule({kGptDim, 256, true}).build(ctx, x, weights_->gpt_layer.linear0); - x = modules::LinearModule({256, 128, true}).build(ctx, x, weights_->gpt_layer.linear1); - x = modules::LinearModule({128, kContentDim, true}).build(ctx, x, weights_->gpt_layer.linear2); - output_ = core::ensure_backend_addressable_layout(ctx, x).tensor; - ggml_set_output(output_); - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(8192, frames_ * 128)), false); - ggml_build_forward_expand(graph_, output_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel GPT layer input buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel GPT layer graph"); - } - debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.s2mel.gpt_layer.frames", frames_); - } - - ~GptLayerGraph() { - clear_graph(); - } - - int64_t frames() const noexcept { - return frames_; - } - - IndexTTS25S2MelSequence run(const std::vector & latent) { - if (static_cast(latent.size()) != frames_ * kGptDim) { - throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer latent size mismatch"); - } - auto timing_start = Clock::now(); - ggml_backend_tensor_set(input_, latent.data(), 0, latent.size() * sizeof(float)); - debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph compute failed"); - } - IndexTTS25S2MelSequence output; - output.frames = frames_; - output.dims = kContentDim; - timing_start = Clock::now(); - output.values = core::read_tensor_f32(output_); - debug::timing_log_scalar("index_tts2_5.s2mel.gpt_layer.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - return output; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t frames_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr ctx_; - ggml_tensor * input_ = nullptr; - ggml_tensor * output_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; -}; - -class IndexTTS25S2MelRuntime::LengthRegulatorGraph { -public: - LengthRegulatorGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t input_frames, - int64_t output_frames, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - input_frames_(input_frames), - output_frames_(output_frames) { - if (input_frames_ <= 0 || output_frames_ <= 0) { - throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph requires positive frame counts"); - } - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph requires weights"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel length regulator graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel length regulator input context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.s2mel.length_regulator", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.s2mel.length_regulator.inputs", - execution_.backend_type()}; - input_ = core::make_tensor( - input_ctx, - GGML_TYPE_F32, - core::TensorShape::from_dims({1, input_frames_, kContentDim})) - .tensor; - mask_ = - core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, output_frames_, kHidden})).tensor; - ggml_set_input(input_); - ggml_set_input(mask_); - auto x = core::wrap_tensor(input_, core::TensorShape::from_dims({1, input_frames_, kContentDim}), GGML_TYPE_F32); - x = modules::LinearModule({kContentDim, kHidden, true}).build(ctx, x, weights_->length_regulator.content_projection); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - x = modules::Interpolate1dModule({output_frames_, modules::Interpolate1dMode::Nearest}).build(ctx, x); - for (size_t layer = 0; layer < weights_->length_regulator.convs.size(); ++layer) { - x = modules::Conv1dModule({kHidden, kHidden, 3, 1, 1, 1, true}).build(ctx, x, weights_->length_regulator.convs[layer]); - x = group_norm_1_group(ctx, x, weights_->length_regulator.norms[layer], kHidden); - x = mish(ctx, x); - } - x = modules::Conv1dModule({kHidden, kHidden, 1, 1, 0, 1, true}).build(ctx, x, weights_->length_regulator.output); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - x = core::ensure_backend_addressable_layout(ctx, x); - auto mask = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, output_frames_, kHidden}), GGML_TYPE_F32); - auto out = modules::MulModule{}.build(ctx, x, mask); - output_ = core::ensure_backend_addressable_layout(ctx, out).tensor; - ggml_set_output(output_); - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(32768, output_frames_ * 512)), false); - ggml_build_forward_expand(graph_, output_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel length regulator input buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel length regulator graph"); - } - mask_values_.assign(static_cast(output_frames_ * kHidden), 1.0F); - core::write_tensor_f32( - core::wrap_tensor(mask_, core::TensorShape::from_dims({1, output_frames_, kHidden}), GGML_TYPE_F32), - mask_values_); - debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.s2mel.length_regulator.input_frames", input_frames_); - debug::trace_log_scalar("index_tts2_5.s2mel.length_regulator.output_frames", output_frames_); - } - - ~LengthRegulatorGraph() { - clear_graph(); - } - - bool matches(int64_t input_frames, int64_t output_frames) const noexcept { - return input_frames_ == input_frames && output_frames_ == output_frames; - } - - IndexTTS25S2MelSequence run(const std::vector & content) { - if (static_cast(content.size()) != input_frames_ * kContentDim) { - throw std::runtime_error("IndexTTS2.5 S2Mel length regulator content size mismatch"); - } - auto timing_start = Clock::now(); - ggml_backend_tensor_set(input_, content.data(), 0, content.size() * sizeof(float)); - debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph compute failed"); - } - IndexTTS25S2MelSequence output; - output.frames = output_frames_; - output.dims = kHidden; - timing_start = Clock::now(); - output.values = core::read_tensor_f32(output_); - debug::timing_log_scalar("index_tts2_5.s2mel.length_regulator.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - return output; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t input_frames_ = 0; - int64_t output_frames_ = 0; - std::vector mask_values_; - std::unique_ptr input_ctx_; - std::unique_ptr ctx_; - ggml_tensor * input_ = nullptr; - ggml_tensor * mask_ = nullptr; - ggml_tensor * output_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; -}; - -class IndexTTS25S2MelRuntime::CfmGraph { -public: - CfmGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t frames, - bool use_cfg, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - frames_(frames), - use_cfg_(use_cfg), - batch_(use_cfg ? 2 : 1) { - if (frames_ <= 0) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph requires positive frame count"); - } - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph requires weights"); - } - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel CFM graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel CFM input context"); - } - ggml_init_params output_params{16ull * 1024ull * 1024ull, nullptr, true}; - output_ctx_.reset(ggml_init(output_params)); - if (output_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 S2Mel CFM output context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.s2mel.cfm", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{input_ctx_.get(), "index_tts2_5.s2mel.cfm.inputs", execution_.backend_type()}; - x_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, kMelChannels, frames_})) - .tensor; - prompt_ = - core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, kMelChannels, frames_})).tensor; - cond_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, frames_, kHidden})).tensor; - style_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_, kStyleDim})).tensor; - timestep_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_})).tensor; - positions_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({frames_})).tensor; - ggml_set_input(x_); - ggml_set_input(prompt_); - ggml_set_input(cond_); - ggml_set_input(style_); - ggml_set_input(timestep_); - ggml_set_input(positions_); - std::vector taps; - auto output = build_cfm_estimator( - ctx, - core::wrap_tensor(x_, core::TensorShape::from_dims({batch_, kMelChannels, frames_}), GGML_TYPE_F32), - core::wrap_tensor(prompt_, core::TensorShape::from_dims({batch_, kMelChannels, frames_}), GGML_TYPE_F32), - core::wrap_tensor(cond_, core::TensorShape::from_dims({batch_, frames_, kHidden}), GGML_TYPE_F32), - core::wrap_tensor(style_, core::TensorShape::from_dims({batch_, kStyleDim}), GGML_TYPE_F32), - core::wrap_tensor(timestep_, core::TensorShape::from_dims({batch_}), GGML_TYPE_F32), - core::wrap_tensor(positions_, core::TensorShape::from_dims({frames_}), GGML_TYPE_I32), - weights_->cfm, - cfm_debug_dump_dir().empty() ? nullptr : &taps); - output_ = core::ensure_backend_addressable_layout(ctx, output).tensor; - ggml_set_output(output_); - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(131072, frames_ * 4096)), false); - ggml_build_forward_expand(graph_, output_); - core::ModuleBuildContext output_ctx{output_ctx_.get(), "index_tts2_5.s2mel.cfm.outputs", execution_.backend_type()}; - for (const auto & tap : taps) { - auto * tap_output = core::make_tensor(output_ctx, GGML_TYPE_F32, tap.value.shape).tensor; - ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), tap.value.tensor, tap_output)); - debug_taps_.push_back({tap.name, tap_output, tap.value.shape}); - } - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel CFM input buffer"); - } - if (!debug_taps_.empty()) { - output_buffer_ = ggml_backend_alloc_ctx_tensors(output_ctx_.get(), execution_.backend()); - if (output_buffer_ == nullptr) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel CFM output buffer"); - } - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 S2Mel CFM graph"); - } - positions_values_.assign(static_cast(frames_), 0); - for (int64_t i = 0; i < frames_; ++i) { - positions_values_[static_cast(i)] = static_cast(i); - } - ggml_backend_tensor_set(positions_, positions_values_.data(), 0, positions_values_.size() * sizeof(int32_t)); - debug::timing_log_scalar("index_tts2_5.s2mel.cfm.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.s2mel.cfm.frames", frames_); - debug::trace_log_scalar("index_tts2_5.s2mel.cfm.batch", batch_); - } - - ~CfmGraph() { - clear_graph(); - } - - bool matches(int64_t frames, bool use_cfg) const noexcept { - return frames_ == frames && use_cfg_ == use_cfg; - } - - std::vector run( - const std::vector & x, - const std::vector & prompt, - const std::vector & cond, - const std::vector & style, - const std::vector & timestep) { - const int64_t mel_values = batch_ * kMelChannels * frames_; - if (static_cast(x.size()) != mel_values || - static_cast(prompt.size()) != mel_values || - static_cast(cond.size()) != batch_ * frames_ * kHidden || - static_cast(style.size()) != batch_ * kStyleDim || - static_cast(timestep.size()) != batch_) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph input shape mismatch"); - } - auto timing_start = Clock::now(); - ggml_backend_tensor_set(x_, x.data(), 0, x.size() * sizeof(float)); - ggml_backend_tensor_set(prompt_, prompt.data(), 0, prompt.size() * sizeof(float)); - ggml_backend_tensor_set(cond_, cond.data(), 0, cond.size() * sizeof(float)); - ggml_backend_tensor_set(style_, style.data(), 0, style.size() * sizeof(float)); - ggml_backend_tensor_set(timestep_, timestep.data(), 0, timestep.size() * sizeof(float)); - debug::timing_log_scalar("index_tts2_5.s2mel.cfm.input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar("index_tts2_5.s2mel.cfm.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph compute failed"); - } - std::vector out(static_cast(mel_values), 0.0F); - timing_start = Clock::now(); - ggml_backend_tensor_get(output_, out.data(), 0, out.size() * sizeof(float)); - debug::timing_log_scalar("index_tts2_5.s2mel.cfm.output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); - const std::string dump_dir = cfm_debug_dump_dir(); - if (!dump_dir.empty() && !taps_dumped_) { - taps_dumped_ = true; - for (const auto & tap : debug_taps_) { - std::vector values(tap.shape.num_elements()); - ggml_backend_tensor_get(tap.tensor, values.data(), 0, values.size() * sizeof(float)); - const std::vector dims(tap.shape.dims.begin(), tap.shape.dims.begin() + tap.shape.rank); - cfm_write_npy_f32(dump_dir, "cfm_tap_" + tap.name, dims, values); - } - } - return out; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - if (output_buffer_ != nullptr) { - ggml_backend_buffer_free(output_buffer_); - output_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t frames_ = 0; - bool use_cfg_ = false; - int64_t batch_ = 1; - std::unique_ptr input_ctx_; - std::unique_ptr output_ctx_; - std::unique_ptr ctx_; - ggml_tensor * x_ = nullptr; - ggml_tensor * prompt_ = nullptr; - ggml_tensor * cond_ = nullptr; - ggml_tensor * style_ = nullptr; - ggml_tensor * timestep_ = nullptr; - ggml_tensor * positions_ = nullptr; - ggml_tensor * output_ = nullptr; - struct DebugTapTensor { - std::string name; - ggml_tensor * tensor = nullptr; - core::TensorShape shape; - }; - std::vector debug_taps_; - bool taps_dumped_ = false; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; - ggml_backend_buffer_t output_buffer_ = nullptr; - std::vector positions_values_; -}; - -void copy_row( - const std::vector & src, - int64_t src_row, - std::vector & dst, - int64_t dst_row, - int64_t row_values) { - std::copy_n( - src.data() + static_cast(src_row * row_values), - static_cast(row_values), - dst.data() + static_cast(dst_row * row_values)); -} - -std::vector repeat_or_zero_rows(const std::vector & values, int64_t row_values, bool use_cfg, bool zero_second) { - if (!use_cfg) { - return values; - } - std::vector out(static_cast(2 * row_values), 0.0F); - copy_row(values, 0, out, 0, row_values); - if (!zero_second) { - copy_row(values, 0, out, 1, row_values); - } - return out; -} - -void zero_prompt_region(std::vector & values, int64_t channels, int64_t frames, int64_t prompt_frames) { - if (prompt_frames < 0 || prompt_frames > frames) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM prompt frame count is out of range"); - } - for (int64_t c = 0; c < channels; ++c) { - auto begin = values.begin() + static_cast(c * frames); - std::fill(begin, begin + static_cast(prompt_frames), 0.0F); - } -} - -std::vector make_prompt_x( - const std::vector & prompt, - int64_t channels, - int64_t frames, - int64_t prompt_frames) { - std::vector out(static_cast(channels * frames), 0.0F); - if (static_cast(prompt.size()) != channels * prompt_frames) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM reference mel shape mismatch"); - } - for (int64_t c = 0; c < channels; ++c) { - std::copy_n( - prompt.data() + static_cast(c * prompt_frames), - static_cast(prompt_frames), - out.data() + static_cast(c * frames)); - } - return out; -} - -std::vector make_condition_with_prompt( - const std::vector & condition, - int64_t frames) { - if (static_cast(condition.size()) != frames * kHidden) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM condition shape mismatch"); - } - return condition; -} - -IndexTTS25S2MelRuntime::IndexTTS25S2MelRuntime( - std::shared_ptr assets, - core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) - : assets_(std::move(assets)), - execution_(&execution), - graph_arena_bytes_(graph_arena_bytes) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel runtime requires assets"); - } - if (graph_arena_bytes_ == 0) { - throw std::runtime_error("IndexTTS2.5 S2Mel graph arena must be non-zero"); - } - weights_ = load_index_tts2_5_s2mel_weights( - *assets_, - execution.backend(), - execution.backend_type(), - matmul_storage_type, - conv_storage_type, - weight_context_bytes); -} - -IndexTTS25S2MelRuntime::~IndexTTS25S2MelRuntime() = default; - -void IndexTTS25S2MelRuntime::prepare_gpt_layer(int64_t frames) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel runtime execution context is missing"); - } - if (gpt_layer_graph_ != nullptr && gpt_layer_graph_->frames() == frames) { - return; - } - gpt_layer_graph_.reset(); - gpt_layer_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); -} - -void IndexTTS25S2MelRuntime::prepare_length_regulator(int64_t input_frames, int64_t output_frames) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel runtime execution context is missing"); - } - if (length_regulator_graph_ != nullptr && length_regulator_graph_->matches(input_frames, output_frames)) { - return; - } - length_regulator_graph_.reset(); - length_regulator_graph_ = std::make_unique( - *execution_, - weights_, - input_frames, - output_frames, - graph_arena_bytes_); -} - -void IndexTTS25S2MelRuntime::prepare_cfm(int64_t total_frames, bool use_cfg) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 S2Mel runtime execution context is missing"); - } - if (total_frames <= 0) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM prepare requires positive frame count"); - } - if (cfm_graph_ != nullptr && cfm_graph_->matches(total_frames, use_cfg)) { - return; - } - cfm_graph_.reset(); - cfm_graph_ = std::make_unique(*execution_, weights_, total_frames, use_cfg, graph_arena_bytes_); -} - -void IndexTTS25S2MelRuntime::release_pre_cfm_graphs() { - gpt_layer_graph_.reset(); - length_regulator_graph_.reset(); -} - -void IndexTTS25S2MelRuntime::release_cfm_graph() { - cfm_graph_.reset(); -} - -IndexTTS25S2MelSequence IndexTTS25S2MelRuntime::project_gpt_latent(const std::vector & latent, int64_t frames) { - if (gpt_layer_graph_ == nullptr || gpt_layer_graph_->frames() != frames) { - throw std::runtime_error("IndexTTS2.5 S2Mel GPT layer graph was not prepared for this latent length"); - } - return gpt_layer_graph_->run(latent); -} - -IndexTTS25S2MelSequence IndexTTS25S2MelRuntime::regulate_length( - const std::vector & content, - int64_t input_frames, - int64_t output_frames) { - if (length_regulator_graph_ == nullptr || !length_regulator_graph_->matches(input_frames, output_frames)) { - throw std::runtime_error("IndexTTS2.5 S2Mel length regulator graph was not prepared for this shape"); - } - return length_regulator_graph_->run(content); -} - -IndexTTS25S2MelMel IndexTTS25S2MelRuntime::infer_mel( - const std::vector & condition, - int64_t total_frames, - const std::vector & reference_mel, - int64_t reference_frames, - const std::vector & style, - int64_t diffusion_steps, - float cfg_rate, - uint32_t seed, - uint64_t rng_offset_blocks) { - if (condition.empty() || reference_mel.empty() || style.empty()) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM requires non-empty condition, reference mel, and style"); - } - if (total_frames <= 0 || reference_frames <= 0 || reference_frames > total_frames) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM frame counts are invalid"); - } - if (diffusion_steps <= 0) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM diffusion steps must be positive"); - } - if (static_cast(style.size()) != kStyleDim) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM style shape mismatch"); - } - const bool use_cfg = cfg_rate > 0.0F; - if (cfm_graph_ == nullptr || !cfm_graph_->matches(total_frames, use_cfg)) { - throw std::runtime_error("IndexTTS2.5 S2Mel CFM graph was not prepared for this shape"); - } - const auto rng_policy = engine::sampling::resolve_torch_cuda_sampling_policy( - execution_->backend_type(), - execution_->config().device, - "index_tts2_5.s2mel.cuda_sampling_policy", - "IndexTTS2.5", - engine::sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda); - std::vector x = engine::sampling::generate_torch_cuda_tensor_iterator_randn( - static_cast(kMelChannels * total_frames), - seed, - rng_offset_blocks, - rng_policy, - engine::sampling::TorchRandnPrecision::Float32); - const std::string noise_path = cfm_noise_path_from_env(); - if (!noise_path.empty()) { - x = cfm_read_npy_f32(noise_path, static_cast(kMelChannels * total_frames)); - } - auto prompt_x = make_prompt_x(reference_mel, kMelChannels, total_frames, reference_frames); - zero_prompt_region(x, kMelChannels, total_frames, reference_frames); - auto mu = make_condition_with_prompt(condition, total_frames); - const std::string cfm_dump_dir = cfm_debug_dump_dir(); - cfm_write_npy_f32(cfm_dump_dir, "cfm_noise", {kMelChannels, total_frames}, x); - - const auto cfm_start = Clock::now(); - double graph_ms = 0.0; - float t = 0.0F; - float dt = 1.0F / static_cast(diffusion_steps); - for (int64_t step = 1; step <= diffusion_steps; ++step) { - const auto graph_start = Clock::now(); - auto x_batched = repeat_or_zero_rows(x, kMelChannels * total_frames, use_cfg, false); - auto prompt_batched = repeat_or_zero_rows(prompt_x, kMelChannels * total_frames, use_cfg, true); - auto cond_batched = repeat_or_zero_rows(mu, total_frames * kHidden, use_cfg, true); - auto style_batched = repeat_or_zero_rows(style, kStyleDim, use_cfg, true); - std::vector timestep(static_cast(use_cfg ? 2 : 1), t); - const auto velocity = cfm_graph_->run(x_batched, prompt_batched, cond_batched, style_batched, timestep); - if (step == 1) { - cfm_write_npy_f32( - cfm_dump_dir, - "cfm_velocity0", - {use_cfg ? 2 : 1, kMelChannels, total_frames}, - velocity); - } - graph_ms += engine::debug::elapsed_ms(graph_start); - const int64_t row_values = kMelChannels * total_frames; - for (int64_t i = 0; i < row_values; ++i) { - float dphi = velocity[static_cast(i)]; - if (use_cfg) { - dphi = (1.0F + cfg_rate) * dphi - cfg_rate * velocity[static_cast(row_values + i)]; - } - x[static_cast(i)] += dt * dphi; - } - t += dt; - if (step < diffusion_steps) { - dt = (static_cast(step + 1) / static_cast(diffusion_steps)) - t; - } - zero_prompt_region(x, kMelChannels, total_frames, reference_frames); - } - - IndexTTS25S2MelMel out; - out.frames = total_frames - reference_frames; - out.channels = kMelChannels; - out.values.resize(static_cast(out.channels * out.frames)); - for (int64_t c = 0; c < kMelChannels; ++c) { - std::copy_n( - x.data() + static_cast(c * total_frames + reference_frames), - static_cast(out.frames), - out.values.data() + static_cast(c * out.frames)); - } - debug::timing_log_scalar("index_tts2_5.s2mel.cfm.euler_graph_ms", graph_ms); - debug::timing_log_scalar("index_tts2_5.s2mel.cfm.euler_total_ms", engine::debug::elapsed_ms(cfm_start)); - return out; -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/semantic_codec.cpp b/src/models/index_tts2_5/semantic_codec.cpp deleted file mode 100644 index e1982a98..00000000 --- a/src/models/index_tts2_5/semantic_codec.cpp +++ /dev/null @@ -1,723 +0,0 @@ -#include "engine/models/index_tts2_5/semantic_codec.h" - -#include "engine/framework/core/backend.h" -#include "engine/framework/debug/profiler.h" -#include "engine/framework/modules/activation_modules.h" -#include "engine/framework/modules/lookup_modules.h" -#include "engine/framework/modules/primitive_modules.h" -#include "engine/framework/modules/structural_modules.h" -#include "engine/framework/modules/weight_binding.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -namespace binding = engine::modules::binding; -namespace core = engine::core; -namespace modules = engine::modules; -using Clock = std::chrono::steady_clock; - -constexpr int64_t kHidden = 1024; -constexpr int64_t kVocosDim = 384; -constexpr int64_t kVocosIntermediate = 2048; -constexpr int64_t kVocosLayers = 12; -constexpr int64_t kConvNeXtKernel = 7; -constexpr int64_t kCodebookSize = 8192; -constexpr int64_t kCodebookDim = 8; - -struct GgmlContextDeleter { - void operator()(ggml_context * ctx) const noexcept { - if (ctx != nullptr) { - ggml_free(ctx); - } - } -}; - -std::vector normalized_codebook(const std::vector & codebook) { - std::vector out(codebook.size(), 0.0F); - for (int64_t row = 0; row < kCodebookSize; ++row) { - double norm = 0.0; - for (int64_t dim = 0; dim < kCodebookDim; ++dim) { - const float value = codebook[static_cast(row * kCodebookDim + dim)]; - norm += static_cast(value) * static_cast(value); - } - const float inv_norm = 1.0F / static_cast(std::sqrt(norm)); - for (int64_t dim = 0; dim < kCodebookDim; ++dim) { - const size_t index = static_cast(row * kCodebookDim + dim); - out[index] = codebook[index] * inv_norm; - } - } - return out; -} - -core::TensorValue div( - core::ModuleBuildContext & ctx, - const core::TensorValue & lhs, - const core::TensorValue & rhs) { - core::validate_shape(rhs, lhs.shape, "Div rhs"); - return core::wrap_tensor(ggml_div(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32); -} - -core::TensorValue vocos_backbone( - core::ModuleBuildContext & ctx, - const core::TensorValue & input_bct, - const IndexTTS25VocosBackboneWeights & weights) { - auto x = modules::Conv1dModule({input_bct.shape.dims[1], kVocosDim, kConvNeXtKernel, 1, 3, 1, true}) - .build(ctx, input_bct, weights.embed); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-6F, true, true}).build(ctx, x, weights.norm); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - for (const auto & block : weights.blocks) { - const auto residual = x; - x = modules::DepthwiseConv1dModule({kVocosDim, kConvNeXtKernel, 1, 3, 1, true}).build(ctx, x, block.depthwise); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-6F, true, true}).build(ctx, x, block.norm); - x = modules::LinearModule({kVocosDim, kVocosIntermediate, true}).build(ctx, x, block.pointwise_in); - x = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, x); - x = modules::LinearModule({kVocosIntermediate, kVocosDim, true}).build(ctx, x, block.pointwise_out); - const auto gamma = modules::RepeatModule({x.shape}).build( - ctx, - core::reshape_tensor(ctx, block.gamma, core::TensorShape::from_dims({1, 1, kVocosDim}))); - x = modules::MulModule{}.build(ctx, x, gamma); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - x = modules::AddModule{}.build(ctx, residual, x); - } - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - return modules::LayerNormModule({x.shape.last_dim(), 1.0e-6F, true, true}).build(ctx, x, weights.final_norm); -} - -core::TensorValue quantizer_in_project( - core::ModuleBuildContext & ctx, - const core::TensorValue & input_bct, - const IndexTTS25SemanticCodecWeights & weights) { - return modules::Conv1dModule({kHidden, kCodebookDim, 1, 1, 0, 1, true}).build(ctx, input_bct, weights.quantizer_in); -} - -core::TensorValue quantizer_out_project( - core::ModuleBuildContext & ctx, - const core::TensorValue & input_bct, - const IndexTTS25SemanticCodecWeights & weights) { - return modules::Conv1dModule({kCodebookDim, kHidden, 1, 1, 0, 1, true}).build(ctx, input_bct, weights.quantizer_out); -} - -core::TensorValue normalize_code_latents(core::ModuleBuildContext & ctx, const core::TensorValue & latents_bdt) { - const auto squared = core::wrap_tensor(ggml_sqr(ctx.ggml, latents_bdt.tensor), latents_bdt.shape, GGML_TYPE_F32); - auto sum = modules::ReduceSumModule({1}).build(ctx, squared); - sum = core::wrap_tensor(ggml_sqrt(ctx.ggml, sum.tensor), sum.shape, GGML_TYPE_F32); - return div(ctx, latents_bdt, modules::RepeatModule({latents_bdt.shape}).build(ctx, sum)); -} - -core::TensorValue argmax_last_dim(core::ModuleBuildContext & ctx, const core::TensorValue & logits) { - auto flat = core::reshape_tensor( - ctx, - core::ensure_backend_addressable_layout(ctx, logits), - core::TensorShape::from_dims({logits.shape.num_elements() / logits.shape.last_dim(), logits.shape.last_dim()})); - auto argmax = core::wrap_tensor( - ggml_argmax(ctx.ggml, flat.tensor), - core::TensorShape::from_dims({flat.shape.dims[0]}), - GGML_TYPE_I32); - return core::reshape_tensor(ctx, argmax, core::TensorShape::from_dims({logits.shape.dims[0], logits.shape.dims[1]})); -} - -core::TensorValue embed_codes_bct( - core::ModuleBuildContext & ctx, - const core::TensorValue & codes_bt, - const IndexTTS25SemanticCodecWeights & weights) { - auto emb_btd = modules::EmbeddingModule({kCodebookSize, kCodebookDim}).build(ctx, codes_bt, weights.codebook); - auto emb_bdt = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, emb_btd); - return quantizer_out_project(ctx, emb_bdt, weights); -} - -std::vector fuse_weight_norm_conv1d( - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t out_channels, - int64_t in_channels, - int64_t kernel_size) { - const auto g = source.require_f32(prefix + ".weight_g", {out_channels, 1, 1}); - const auto v = source.require_f32(prefix + ".weight_v", {out_channels, in_channels, kernel_size}); - std::vector weight(v.size(), 0.0F); - for (int64_t out = 0; out < out_channels; ++out) { - double norm = 0.0; - for (int64_t in = 0; in < in_channels; ++in) { - for (int64_t k = 0; k < kernel_size; ++k) { - const float value = v[static_cast((out * in_channels + in) * kernel_size + k)]; - norm += static_cast(value) * static_cast(value); - } - } - const float scale = g[static_cast(out)] / static_cast(std::sqrt(norm)); - for (int64_t in = 0; in < in_channels; ++in) { - for (int64_t k = 0; k < kernel_size; ++k) { - const size_t index = static_cast((out * in_channels + in) * kernel_size + k); - weight[index] = v[index] * scale; - } - } - } - return weight; -} - -engine::modules::Conv1dWeights load_weight_norm_conv1d( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type, - int64_t out_channels, - int64_t in_channels, - int64_t kernel_size) { - engine::modules::Conv1dWeights weights; - weights.weight = store.make_from_f32( - engine::core::TensorShape::from_dims({out_channels, in_channels, kernel_size}), - storage_type, - fuse_weight_norm_conv1d(source, prefix, out_channels, in_channels, kernel_size)); - weights.bias = store.load_f32_tensor(source, prefix + ".bias", {out_channels}); - return weights; -} - -IndexTTS25VocosConvNeXtBlockWeights load_convnext_block( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) { - IndexTTS25VocosConvNeXtBlockWeights block; - block.depthwise = binding::depthwise_conv1d_from_source( - store, - source, - prefix + ".dwconv", - conv_storage_type, - kVocosDim, - kConvNeXtKernel, - true); - block.norm = binding::norm_from_source(store, source, prefix + ".norm", kVocosDim); - block.pointwise_in = binding::linear_from_source( - store, - source, - prefix + ".pwconv1", - matmul_storage_type, - kVocosIntermediate, - kVocosDim, - true); - block.pointwise_out = binding::linear_from_source( - store, - source, - prefix + ".pwconv2", - matmul_storage_type, - kVocosDim, - kVocosIntermediate, - true); - block.gamma = store.load_f32_tensor(source, prefix + ".gamma", {kVocosDim}); - return block; -} - -IndexTTS25VocosBackboneWeights load_vocos_backbone( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - int64_t input_channels, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) { - IndexTTS25VocosBackboneWeights backbone; - backbone.embed = binding::conv1d_from_source( - store, - source, - prefix + ".embed", - conv_storage_type, - kVocosDim, - input_channels, - kConvNeXtKernel, - true); - backbone.norm = binding::norm_from_source(store, source, prefix + ".norm", kVocosDim); - backbone.blocks.reserve(static_cast(kVocosLayers)); - for (int64_t i = 0; i < kVocosLayers; ++i) { - backbone.blocks.push_back(load_convnext_block( - store, - source, - prefix + ".convnext." + std::to_string(i), - matmul_storage_type, - conv_storage_type)); - } - backbone.final_norm = binding::norm_from_source(store, source, prefix + ".final_layer_norm", kVocosDim); - return backbone; -} - -} // namespace - -class IndexTTS25SemanticCodecRuntime::QuantizeGraph { -public: - QuantizeGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t frames, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - frames_(frames) { - if (frames_ <= 0) { - throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph requires positive frame count"); - } - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph requires weights"); - } - - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec quantize graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec quantize input context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.semantic_codec.quantize", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.semantic_codec.quantize.inputs", - execution_.backend_type()}; - - semantic_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kHidden})).tensor; - ggml_set_input(semantic_); - auto x = core::wrap_tensor(semantic_, core::TensorShape::from_dims({1, frames_, kHidden}), GGML_TYPE_F32); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - x = vocos_backbone(ctx, x, weights_->encoder_backbone); - x = modules::LinearModule({kVocosDim, kHidden, true}).build(ctx, x, weights_->encoder_projection); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - auto latents = normalize_code_latents(ctx, quantizer_in_project(ctx, x, *weights_)); - const auto latents_btd = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, latents); - auto logits = modules::LinearModule({kCodebookDim, kCodebookSize, false}) - .build(ctx, latents_btd, {weights_->normalized_codebook, std::nullopt}); - codes_ = argmax_last_dim(ctx, logits).tensor; - auto embedding = embed_codes_bct( - ctx, - core::wrap_tensor(codes_, core::TensorShape::from_dims({1, frames_}), GGML_TYPE_I32), - *weights_); - embedding_ = core::ensure_backend_addressable_layout(ctx, embedding).tensor; - ggml_set_output(codes_); - ggml_set_output(embedding_); - - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(65536, frames_ * 2048 + 4096)), false); - ggml_build_forward_expand(graph_, codes_); - ggml_build_forward_expand(graph_, embedding_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec quantize input buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec quantize graph"); - } - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.quantize.graph.build_ms", - engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.semantic_codec.quantize.frames", frames_); - } - - ~QuantizeGraph() { - clear_graph(); - } - - int64_t frames() const noexcept { - return frames_; - } - - IndexTTS25SemanticCodecOutput run(const IndexTTS25SemanticEmbedding & semantic) { - if (semantic.frames != frames_ || semantic.dims != kHidden) { - throw std::runtime_error("IndexTTS2.5 semantic codec semantic input shape does not match prepared graph"); - } - if (static_cast(semantic.values.size()) != frames_ * kHidden) { - throw std::runtime_error("IndexTTS2.5 semantic codec semantic input value count mismatch"); - } - auto timing_start = Clock::now(); - ggml_backend_tensor_set(semantic_, semantic.values.data(), 0, semantic.values.size() * sizeof(float)); - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.quantize.input_upload_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.quantize.graph.compute_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph compute failed"); - } - - IndexTTS25SemanticCodecOutput output; - output.frames = frames_; - output.dims = kHidden; - output.codes.resize(static_cast(frames_)); - output.embedding_channel_first.resize(static_cast(kHidden * frames_)); - timing_start = Clock::now(); - ggml_backend_tensor_get(codes_, output.codes.data(), 0, output.codes.size() * sizeof(int32_t)); - ggml_backend_tensor_get( - embedding_, - output.embedding_channel_first.data(), - 0, - output.embedding_channel_first.size() * sizeof(float)); - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.quantize.output_read_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - return output; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t frames_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr ctx_; - ggml_tensor * semantic_ = nullptr; - ggml_tensor * codes_ = nullptr; - ggml_tensor * embedding_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; -}; - -class IndexTTS25SemanticCodecRuntime::CodesGraph { -public: - CodesGraph( - core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t frames, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - frames_(frames) { - if (frames_ <= 0) { - throw std::runtime_error("IndexTTS2.5 semantic codec code graph requires positive frame count"); - } - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 semantic codec code graph requires weights"); - } - - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec code graph context"); - } - ggml_init_params input_params{16ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 semantic codec code input context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.semantic_codec.codes", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.semantic_codec.codes.inputs", - execution_.backend_type()}; - codes_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, frames_})).tensor; - upsample_ids_ = ggml_new_tensor_1d(input_ctx_.get(), GGML_TYPE_I32, 2 * frames_); - ggml_set_input(codes_); - auto embedding = embed_codes_bct( - ctx, - core::wrap_tensor(codes_, core::TensorShape::from_dims({1, frames_}), GGML_TYPE_I32), - *weights_); - // EnhancedCodec.decode (codec/models.py): decoder backbone + projection, - // then 2x nearest upsample along time and the `up` conv. - auto x = vocos_backbone(ctx, embedding, weights_->decoder_backbone); - x = modules::LinearModule({kVocosDim, kHidden, true}).build(ctx, x, weights_->decoder_projection); - x = core::ensure_backend_addressable_layout(ctx, x); - x = core::wrap_tensor( - ggml_get_rows(ctx.ggml, x.tensor, upsample_ids_), - core::TensorShape::from_dims({1, 2 * frames_, kHidden}), - GGML_TYPE_F32); - x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); - x = modules::Conv1dModule({kHidden, kHidden, 3, 1, 1, 1, true}).build(ctx, x, weights_->up); - embedding_ = core::ensure_backend_addressable_layout(ctx, x).tensor; - ggml_set_output(embedding_); - - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(std::max(4096, frames_ * 64 + 1024)), false); - ggml_build_forward_expand(graph_, embedding_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec code input buffer"); - } - std::vector upsample_ids(static_cast(2 * frames_)); - for (int64_t frame = 0; frame < frames_; ++frame) { - upsample_ids[static_cast(2 * frame)] = static_cast(frame); - upsample_ids[static_cast(2 * frame + 1)] = static_cast(frame); - } - ggml_backend_tensor_set(upsample_ids_, upsample_ids.data(), 0, upsample_ids.size() * sizeof(int32_t)); - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 semantic codec code graph"); - } - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.codes.graph.build_ms", - engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.semantic_codec.codes.frames", frames_); - } - - ~CodesGraph() { - clear_graph(); - } - - int64_t frames() const noexcept { - return frames_; - } - - IndexTTS25SemanticCodecOutput run(const std::vector & codes) { - if (static_cast(codes.size()) != frames_) { - throw std::runtime_error("IndexTTS2.5 semantic codec code count does not match prepared graph"); - } - auto timing_start = Clock::now(); - ggml_backend_tensor_set(codes_, codes.data(), 0, codes.size() * sizeof(int32_t)); - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.codes.input_upload_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.codes.graph.compute_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 semantic codec code graph compute failed"); - } - - IndexTTS25SemanticCodecOutput output; - output.frames = 2 * frames_; - output.dims = kHidden; - output.codes = codes; - output.embedding_channel_first.resize(static_cast(kHidden * output.frames)); - timing_start = Clock::now(); - ggml_backend_tensor_get( - embedding_, - output.embedding_channel_first.data(), - 0, - output.embedding_channel_first.size() * sizeof(float)); - debug::timing_log_scalar( - "index_tts2_5.semantic_codec.codes.output_read_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - return output; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t frames_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr ctx_; - ggml_tensor * codes_ = nullptr; - ggml_tensor * upsample_ids_ = nullptr; - ggml_tensor * embedding_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; -}; - -std::shared_ptr load_index_tts2_5_semantic_codec_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes) { - if (assets.semantic_codec_weights == nullptr) { - throw std::runtime_error("IndexTTS2.5 semantic codec requires tensor source"); - } - auto weights = std::make_shared(); - weights->store = std::make_shared( - backend, - backend_type, - "index_tts2_5.semantic_codec.weights", - weight_context_bytes); - - const auto & source = *assets.semantic_codec_weights; - weights->encoder_backbone = load_vocos_backbone( - *weights->store, - source, - "encoder.0", - kHidden, - matmul_storage_type, - conv_storage_type); - weights->encoder_projection = binding::linear_from_source( - *weights->store, - source, - "encoder.1", - matmul_storage_type, - kHidden, - kVocosDim, - true); - weights->quantizer_in = load_weight_norm_conv1d( - *weights->store, - source, - "quantizer.quantizers.0.in_project", - conv_storage_type, - kCodebookDim, - kHidden, - 1); - const auto codebook = source.require_f32("quantizer.quantizers.0.codebook.weight", {kCodebookSize, kCodebookDim}); - weights->codebook = weights->store->make_from_f32( - engine::core::TensorShape::from_dims({kCodebookSize, kCodebookDim}), - matmul_storage_type, - codebook); - weights->normalized_codebook = weights->store->make_from_f32( - engine::core::TensorShape::from_dims({kCodebookSize, kCodebookDim}), - matmul_storage_type, - normalized_codebook(codebook)); - weights->quantizer_out = load_weight_norm_conv1d( - *weights->store, - source, - "quantizer.quantizers.0.out_project", - conv_storage_type, - kHidden, - kCodebookDim, - 1); - weights->decoder_backbone = load_vocos_backbone( - *weights->store, - source, - "decoder.0", - kHidden, - matmul_storage_type, - conv_storage_type); - weights->decoder_projection = binding::linear_from_source( - *weights->store, - source, - "decoder.1", - matmul_storage_type, - kHidden, - kVocosDim, - true); - weights->up = binding::conv1d_from_source( - *weights->store, - source, - "up", - conv_storage_type, - kHidden, - kHidden, - 3, - true); - - weights->store->upload(); - assets.semantic_codec_weights->release_storage(); - return weights; -} - -IndexTTS25SemanticCodecRuntime::IndexTTS25SemanticCodecRuntime( - std::shared_ptr assets, - core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) - : assets_(std::move(assets)), - execution_(&execution), - graph_arena_bytes_(graph_arena_bytes) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 semantic codec runtime requires assets"); - } - if (graph_arena_bytes_ == 0) { - throw std::runtime_error("IndexTTS2.5 semantic codec graph arena must be non-zero"); - } - weights_ = load_index_tts2_5_semantic_codec_weights( - *assets_, - execution.backend(), - execution.backend_type(), - matmul_storage_type, - conv_storage_type, - weight_context_bytes); -} - -IndexTTS25SemanticCodecRuntime::~IndexTTS25SemanticCodecRuntime() = default; - -void IndexTTS25SemanticCodecRuntime::prepare_quantize(int64_t frames) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 semantic codec runtime execution context is missing"); - } - if (frames <= 0) { - throw std::runtime_error("IndexTTS2.5 semantic codec quantize prepare requires positive frames"); - } - if (quantize_graph_ != nullptr && quantize_graph_->frames() == frames) { - return; - } - quantize_graph_.reset(); - quantize_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); -} - -void IndexTTS25SemanticCodecRuntime::prepare_codes(int64_t frames) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 semantic codec runtime execution context is missing"); - } - if (frames <= 0) { - throw std::runtime_error("IndexTTS2.5 semantic codec code prepare requires positive frames"); - } - if (codes_graph_ != nullptr && codes_graph_->frames() == frames) { - return; - } - codes_graph_.reset(); - codes_graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); -} - -IndexTTS25SemanticCodecOutput IndexTTS25SemanticCodecRuntime::quantize(const IndexTTS25SemanticEmbedding & semantic) { - if (quantize_graph_ == nullptr || quantize_graph_->frames() != semantic.frames) { - throw std::runtime_error("IndexTTS2.5 semantic codec quantize graph was not prepared for this reference length"); - } - return quantize_graph_->run(semantic); -} - -IndexTTS25SemanticCodecOutput IndexTTS25SemanticCodecRuntime::codes_to_embedding( - const std::vector & codes, - int64_t frames) { - if (codes_graph_ == nullptr || codes_graph_->frames() != frames) { - throw std::runtime_error("IndexTTS2.5 semantic codec code graph was not prepared for this generation length"); - } - return codes_graph_->run(codes); -} - -void IndexTTS25SemanticCodecRuntime::release_graphs() { - quantize_graph_.reset(); - codes_graph_.reset(); -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/semantic_encoder.cpp b/src/models/index_tts2_5/semantic_encoder.cpp deleted file mode 100644 index bc1b2e3f..00000000 --- a/src/models/index_tts2_5/semantic_encoder.cpp +++ /dev/null @@ -1,622 +0,0 @@ -#include "engine/models/index_tts2_5/semantic_encoder.h" - -#include "engine/framework/core/backend.h" -#include "engine/framework/debug/profiler.h" -#include "engine/framework/modules/activation_modules.h" -#include "engine/framework/modules/weight_binding.h" -#include "engine/framework/modules/lookup_modules.h" -#include "engine/framework/modules/primitive_modules.h" -#include "engine/framework/modules/structural_modules.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -namespace binding = engine::modules::binding; -namespace modules = engine::modules; -using Clock = std::chrono::steady_clock; - -constexpr int64_t kFeatureDim = 160; -constexpr int64_t kHidden = 1024; -constexpr int64_t kIntermediate = 4096; -constexpr int64_t kLayers = 24; -constexpr int64_t kSemanticOutputHiddenStateIndex = 17; -constexpr int64_t kHeads = 16; -constexpr int64_t kHeadDim = kHidden / kHeads; -constexpr int64_t kRelativePositions = 73; -constexpr int64_t kRelativeLeft = 64; -constexpr int64_t kRelativeRight = 8; -constexpr int64_t kConvKernel = 31; - -struct GgmlContextDeleter { - void operator()(ggml_context * ctx) const noexcept { - if (ctx != nullptr) { - ggml_free(ctx); - } - } -}; - -core::TensorValue scale_tensor(core::ModuleBuildContext & ctx, const core::TensorValue & input, float scale) { - return core::wrap_tensor(ggml_scale(ctx.ggml, input.tensor, scale), input.shape, GGML_TYPE_F32); -} - -core::TensorValue add_scaled_residual( - core::ModuleBuildContext & ctx, - const core::TensorValue & residual, - const core::TensorValue & update, - float update_scale) { - return modules::AddModule{}.build(ctx, residual, scale_tensor(ctx, update, update_scale)); -} - -core::TensorValue reshape_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input) { - const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); - return core::reshape_tensor(ctx, contiguous, core::TensorShape::from_dims({1, input.shape.dims[1], kHeads, kHeadDim})); -} - -core::TensorValue repeat_last_dim_mask( - core::ModuleBuildContext & ctx, - const core::TensorValue & mask, - const core::TensorValue & like) { - auto mask_view = core::reshape_tensor(ctx, mask, core::TensorShape::from_dims({1, mask.shape.dims[0], 1})); - return modules::RepeatModule({like.shape}).build(ctx, mask_view); -} - -core::TensorValue apply_keep_mask( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & keep_mask) { - return modules::MulModule{}.build(ctx, input, repeat_last_dim_mask(ctx, keep_mask, input)); -} - -core::TensorValue broadcast_vector( - core::ModuleBuildContext & ctx, - const core::TensorValue & vector, - const core::TensorValue & like) { - auto view = core::reshape_tensor(ctx, vector, core::TensorShape::from_dims({1, 1, vector.shape.dims[0]})); - return modules::RepeatModule({like.shape}).build(ctx, view); -} - -core::TensorValue build_relative_key_bias( - core::ModuleBuildContext & ctx, - const core::TensorValue & q_heads, - const core::TensorValue & distance_ids, - const core::TensorValue & distance_embedding) { - const int64_t frames = q_heads.shape.dims[2]; - auto positions = modules::EmbeddingModule({kRelativePositions, kHeadDim}).build(ctx, distance_ids, distance_embedding); - auto q_by_row = modules::TransposeModule({{2, 1, 0, 3}, 4}).build(ctx, q_heads); - q_by_row = core::ensure_backend_addressable_layout(ctx, q_by_row); - q_by_row = core::reshape_tensor(ctx, q_by_row, core::TensorShape::from_dims({frames, kHeads, kHeadDim})); - auto positions_by_row = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, positions); - positions_by_row = core::ensure_backend_addressable_layout(ctx, positions_by_row); - auto by_row = modules::MatMulModule{}.build(ctx, q_by_row, positions_by_row); - by_row = core::ensure_backend_addressable_layout(ctx, by_row); - by_row = core::reshape_tensor(ctx, by_row, core::TensorShape::from_dims({frames, kHeads, 1, frames})); - return core::ensure_backend_addressable_layout(ctx, modules::TransposeModule({{2, 1, 0, 3}, 4}).build(ctx, by_row)); -} - -core::TensorValue wav2vec2bert_attention( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & attention_mask, - const core::TensorValue & distance_ids, - const IndexTTS25Wav2Vec2BertAttentionWeights & weights) { - auto q = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, input, weights.q); - auto k = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, input, weights.k); - auto v = modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, input, weights.v); - q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, q)); - k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, k)); - v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, reshape_heads(ctx, v)); - - auto scores = modules::MatMulModule{}.build( - ctx, - q, - modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k)); - scores = scale_tensor(ctx, scores, 1.0F / std::sqrt(static_cast(kHeadDim))); - auto relative = build_relative_key_bias(ctx, q, distance_ids, weights.distance_embedding); - relative = scale_tensor(ctx, relative, 1.0F / std::sqrt(static_cast(kHeadDim))); - scores = modules::AddModule{}.build(ctx, scores, relative); - scores = modules::AddModule{}.build(ctx, scores, attention_mask); - auto probs = core::wrap_tensor(ggml_soft_max(ctx.ggml, scores.tensor), scores.shape, GGML_TYPE_F32); - auto out = modules::MatMulModule{}.build(ctx, probs, v); - out = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, out); - out = core::ensure_backend_addressable_layout(ctx, out); - out = core::reshape_tensor(ctx, out, core::TensorShape::from_dims({1, input.shape.dims[1], kHidden})); - return modules::LinearModule({kHidden, kHidden, true, GGML_PREC_F32}).build(ctx, out, weights.out); -} - -core::TensorValue wav2vec2bert_feed_forward( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const modules::LinearWeights & in, - const modules::LinearWeights & out) { - auto hidden = modules::LinearModule({kHidden, kIntermediate, true, GGML_PREC_F32}).build(ctx, input, in); - hidden = modules::SiluModule{}.build(ctx, hidden); - return modules::LinearModule({kIntermediate, kHidden, true, GGML_PREC_F32}).build(ctx, hidden, out); -} - -core::TensorValue wav2vec2bert_conv( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & keep_mask, - const IndexTTS25Wav2Vec2BertConvWeights & weights) { - auto hidden = modules::LayerNormModule({input.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, input, weights.layer_norm); - hidden = apply_keep_mask(ctx, hidden, keep_mask); - hidden = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); - hidden = modules::Conv1dModule({kHidden, 2 * kHidden, 1, 1, 0, 1, false}).build(ctx, hidden, weights.pointwise_in); - auto gate = modules::SliceModule({1, 0, kHidden}).build(ctx, hidden); - auto value = modules::SliceModule({1, kHidden, kHidden}).build(ctx, hidden); - value = modules::SigmoidModule{}.build(ctx, value); - hidden = modules::MulModule{}.build(ctx, gate, value); - - auto first = modules::SliceModule({2, 0, 1}).build(ctx, hidden); - auto left_pad = modules::RepeatModule({core::TensorShape::from_dims({1, kHidden, kConvKernel - 1})}).build(ctx, first); - left_pad = scale_tensor(ctx, left_pad, 0.0F); - hidden = modules::ConcatModule({2}).build(ctx, left_pad, hidden); - hidden = modules::DepthwiseConv1dModule({kHidden, kConvKernel, 1, 0, 1, false}).build(ctx, hidden, weights.depthwise); - hidden = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); - hidden = modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.depthwise_layer_norm); - hidden = modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); - hidden = modules::SiluModule{}.build(ctx, hidden); - hidden = modules::Conv1dModule({kHidden, kHidden, 1, 1, 0, 1, false}).build(ctx, hidden, weights.pointwise_out); - return modules::TransposeModule({{0, 2, 1}, 3}).build(ctx, hidden); -} - -core::TensorValue wav2vec2bert_layer( - core::ModuleBuildContext & ctx, - const core::TensorValue & input, - const core::TensorValue & keep_mask, - const core::TensorValue & attention_mask, - const core::TensorValue & distance_ids, - const IndexTTS25Wav2Vec2BertLayerWeights & weights) { - auto hidden = modules::LayerNormModule({input.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, input, weights.ffn1_norm); - hidden = wav2vec2bert_feed_forward(ctx, hidden, weights.ffn1_in, weights.ffn1_out); - hidden = add_scaled_residual(ctx, input, hidden, 0.5F); - - auto attn = modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.self_attn_norm); - attn = wav2vec2bert_attention(ctx, attn, attention_mask, distance_ids, weights.self_attn); - hidden = modules::AddModule{}.build(ctx, hidden, attn); - - auto conv = wav2vec2bert_conv(ctx, hidden, keep_mask, weights.conv); - hidden = modules::AddModule{}.build(ctx, hidden, conv); - - auto ffn2 = modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.ffn2_norm); - ffn2 = wav2vec2bert_feed_forward(ctx, ffn2, weights.ffn2_in, weights.ffn2_out); - hidden = add_scaled_residual(ctx, hidden, ffn2, 0.5F); - return modules::LayerNormModule({hidden.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, hidden, weights.final_norm); -} - -IndexTTS25Wav2Vec2BertAttentionWeights load_attention( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type) { - IndexTTS25Wav2Vec2BertAttentionWeights weights; - weights.q = binding::linear_from_source(store, source, prefix + ".linear_q", storage_type, kHidden, kHidden, true); - weights.k = binding::linear_from_source(store, source, prefix + ".linear_k", storage_type, kHidden, kHidden, true); - weights.v = binding::linear_from_source(store, source, prefix + ".linear_v", storage_type, kHidden, kHidden, true); - weights.out = binding::linear_from_source(store, source, prefix + ".linear_out", storage_type, kHidden, kHidden, true); - weights.distance_embedding = store.load_tensor( - source, - prefix + ".distance_embedding.weight", - storage_type, - {kRelativePositions, kHeadDim}); - return weights; -} - -IndexTTS25Wav2Vec2BertConvWeights load_conv_module( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - const std::string & prefix, - engine::assets::TensorStorageType storage_type) { - IndexTTS25Wav2Vec2BertConvWeights weights; - weights.layer_norm = binding::norm_from_source(store, source, prefix + ".layer_norm", kHidden); - weights.pointwise_in = binding::conv1d_from_source( - store, - source, - prefix + ".pointwise_conv1", - storage_type, - 2 * kHidden, - kHidden, - 1, - false); - weights.depthwise = binding::depthwise_conv1d_from_source( - store, - source, - prefix + ".depthwise_conv", - storage_type, - kHidden, - kConvKernel, - false); - weights.depthwise_layer_norm = binding::norm_from_source(store, source, prefix + ".depthwise_layer_norm", kHidden); - weights.pointwise_out = binding::conv1d_from_source( - store, - source, - prefix + ".pointwise_conv2", - storage_type, - kHidden, - kHidden, - 1, - false); - return weights; -} - -IndexTTS25Wav2Vec2BertLayerWeights load_layer( - engine::core::BackendWeightStore & store, - const engine::assets::TensorSource & source, - int64_t layer_index, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) { - const std::string prefix = "encoder.layers." + std::to_string(layer_index); - IndexTTS25Wav2Vec2BertLayerWeights layer; - layer.ffn1_norm = binding::norm_from_source(store, source, prefix + ".ffn1_layer_norm", kHidden); - layer.ffn1_in = binding::linear_from_source( - store, - source, - prefix + ".ffn1.intermediate_dense", - matmul_storage_type, - kIntermediate, - kHidden, - true); - layer.ffn1_out = binding::linear_from_source( - store, - source, - prefix + ".ffn1.output_dense", - matmul_storage_type, - kHidden, - kIntermediate, - true); - layer.self_attn_norm = binding::norm_from_source(store, source, prefix + ".self_attn_layer_norm", kHidden); - layer.self_attn = load_attention(store, source, prefix + ".self_attn", matmul_storage_type); - layer.conv = load_conv_module(store, source, prefix + ".conv_module", conv_storage_type); - layer.ffn2_norm = binding::norm_from_source(store, source, prefix + ".ffn2_layer_norm", kHidden); - layer.ffn2_in = binding::linear_from_source( - store, - source, - prefix + ".ffn2.intermediate_dense", - matmul_storage_type, - kIntermediate, - kHidden, - true); - layer.ffn2_out = binding::linear_from_source( - store, - source, - prefix + ".ffn2.output_dense", - matmul_storage_type, - kHidden, - kIntermediate, - true); - layer.final_norm = binding::norm_from_source(store, source, prefix + ".final_layer_norm", kHidden); - return layer; -} - -std::vector wav2vec2bert_std(const engine::assets::TensorSource & source) { - const auto var = source.require_f32("var", {kHidden}); - std::vector stddev(static_cast(kHidden), 0.0F); - for (int64_t i = 0; i < kHidden; ++i) { - stddev[static_cast(i)] = std::sqrt(var[static_cast(i)]); - } - return stddev; -} - -std::vector make_distance_ids(int64_t frames) { - std::vector ids(static_cast(frames * frames), 0); - for (int64_t row = 0; row < frames; ++row) { - for (int64_t col = 0; col < frames; ++col) { - const int64_t distance = std::max(-kRelativeLeft, std::min(kRelativeRight, col - row)); - ids[static_cast(row * frames + col)] = static_cast(distance + kRelativeLeft); - } - } - return ids; -} - -std::vector make_attention_mask(const std::vector & keep_mask, int64_t frames) { - if (static_cast(keep_mask.size()) != frames) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT attention mask length mismatch"); - } - std::vector mask(static_cast(kHeads * frames * frames), 0.0F); - const float hidden = std::numeric_limits::lowest(); - for (int64_t head = 0; head < kHeads; ++head) { - for (int64_t row = 0; row < frames; ++row) { - for (int64_t col = 0; col < frames; ++col) { - if (keep_mask[static_cast(col)] == 0) { - mask[static_cast((head * frames + row) * frames + col)] = hidden; - } - } - } - } - return mask; -} - -std::vector make_keep_mask_f32(const std::vector & keep_mask) { - std::vector out(keep_mask.size(), 0.0F); - for (size_t i = 0; i < keep_mask.size(); ++i) { - out[i] = keep_mask[i] == 0 ? 0.0F : 1.0F; - } - return out; -} - -} // namespace - -class IndexTTS25Wav2Vec2BertRuntime::Graph { -public: - Graph( - engine::core::ExecutionContext & execution, - std::shared_ptr weights, - int64_t frames, - size_t graph_arena_bytes) - : execution_(execution), - weights_(std::move(weights)), - frames_(frames) { - if (frames_ <= 0) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph requires positive frame count"); - } - if (weights_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph requires weights"); - } - - const auto build_start = Clock::now(); - ggml_init_params params{graph_arena_bytes, nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Wav2Vec2-BERT graph context"); - } - ggml_init_params input_params{64ull * 1024ull * 1024ull, nullptr, true}; - input_ctx_.reset(ggml_init(input_params)); - if (input_ctx_ == nullptr) { - throw std::runtime_error("failed to initialize IndexTTS2.5 Wav2Vec2-BERT input context"); - } - core::ModuleBuildContext ctx{ctx_.get(), "index_tts2_5.wav2vec2bert", execution_.backend_type()}; - core::ModuleBuildContext input_ctx{ - input_ctx_.get(), - "index_tts2_5.wav2vec2bert.inputs", - execution_.backend_type()}; - - features_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, frames_, kFeatureDim})).tensor; - keep_mask_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({frames_})).tensor; - attention_mask_ = core::make_tensor(input_ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kHeads, frames_, frames_})).tensor; - distance_ids_ = core::make_tensor(input_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({frames_, frames_})).tensor; - ggml_set_input(features_); - ggml_set_input(keep_mask_); - ggml_set_input(attention_mask_); - ggml_set_input(distance_ids_); - - auto x = core::wrap_tensor(features_, core::TensorShape::from_dims({1, frames_, kFeatureDim}), GGML_TYPE_F32); - auto keep = core::wrap_tensor(keep_mask_, core::TensorShape::from_dims({frames_}), GGML_TYPE_F32); - auto attn_mask = core::wrap_tensor(attention_mask_, core::TensorShape::from_dims({1, kHeads, frames_, frames_}), GGML_TYPE_F32); - auto distances = core::wrap_tensor(distance_ids_, core::TensorShape::from_dims({frames_, frames_}), GGML_TYPE_I32); - - x = modules::LayerNormModule({x.shape.last_dim(), 1.0e-5F, true, true}).build(ctx, x, weights_->feature_norm); - x = modules::LinearModule({kFeatureDim, kHidden, true, GGML_PREC_F32}).build(ctx, x, weights_->feature_projection); - x = apply_keep_mask(ctx, x, keep); - for (int64_t layer = 0; layer < kSemanticOutputHiddenStateIndex; ++layer) { - x = wav2vec2bert_layer(ctx, x, keep, attn_mask, distances, weights_->layers[static_cast(layer)]); - } - const auto mean = broadcast_vector(ctx, weights_->semantic_mean, x); - const auto std = broadcast_vector(ctx, weights_->semantic_std, x); - x = core::wrap_tensor(ggml_sub(ctx.ggml, x.tensor, mean.tensor), x.shape, GGML_TYPE_F32); - x = core::wrap_tensor(ggml_div(ctx.ggml, x.tensor, std.tensor), x.shape, GGML_TYPE_F32); - output_ = core::ensure_backend_addressable_layout(ctx, x).tensor; - ggml_set_output(output_); - - const int64_t graph_nodes = std::max( - 65536, - kSemanticOutputHiddenStateIndex * (frames_ * kHeads * 8 + frames_ * 16 + 1024) + 8192); - graph_ = ggml_new_graph_custom(ctx_.get(), static_cast(graph_nodes), false); - ggml_build_forward_expand(graph_, output_); - input_buffer_ = ggml_backend_alloc_ctx_tensors(input_ctx_.get(), execution_.backend()); - if (input_buffer_ == nullptr) { - throw std::runtime_error("failed to allocate IndexTTS2.5 Wav2Vec2-BERT input buffer"); - } - gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(execution_.backend())); - if (gallocr_ == nullptr || - !ggml_gallocr_reserve(gallocr_, graph_) || - !ggml_gallocr_alloc_graph(gallocr_, graph_)) { - clear_graph(); - throw std::runtime_error("failed to allocate IndexTTS2.5 Wav2Vec2-BERT graph"); - } - - const auto distances_host = make_distance_ids(frames_); - ggml_backend_tensor_set(distance_ids_, distances_host.data(), 0, distances_host.size() * sizeof(int32_t)); - debug::timing_log_scalar( - "index_tts2_5.wav2vec2bert.graph.build_ms", - engine::debug::elapsed_ms(build_start, Clock::now())); - debug::trace_log_scalar("index_tts2_5.wav2vec2bert.frames", frames_); - } - - ~Graph() { - clear_graph(); - } - - int64_t frames() const noexcept { - return frames_; - } - - IndexTTS25SemanticEmbedding run(const IndexTTS25SemanticFeatureOutput & features) { - if (features.frames <= 0 || features.frames > frames_ || features.dims != kFeatureDim) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT feature shape does not match prepared graph"); - } - if (static_cast(features.values.size()) != features.frames * kFeatureDim) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT feature value count mismatch"); - } - if (static_cast(features.attention_mask.size()) != features.frames) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT feature attention mask count mismatch"); - } - - auto timing_start = Clock::now(); - std::vector padded_features(static_cast(frames_ * kFeatureDim), 0.0F); - std::copy(features.values.begin(), features.values.end(), padded_features.begin()); - std::vector padded_mask(static_cast(frames_), 0); - std::copy(features.attention_mask.begin(), features.attention_mask.end(), padded_mask.begin()); - const auto keep = make_keep_mask_f32(padded_mask); - const auto attention = make_attention_mask(padded_mask, frames_); - ggml_backend_tensor_set(features_, padded_features.data(), 0, padded_features.size() * sizeof(float)); - ggml_backend_tensor_set(keep_mask_, keep.data(), 0, keep.size() * sizeof(float)); - ggml_backend_tensor_set(attention_mask_, attention.data(), 0, attention.size() * sizeof(float)); - debug::timing_log_scalar( - "index_tts2_5.wav2vec2bert.input_upload_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - - core::set_backend_threads(execution_.backend(), execution_.config().threads); - timing_start = Clock::now(); - const ggml_status status = engine::core::compute_backend_graph(execution_.backend(), graph_); - ggml_backend_synchronize(execution_.backend()); - debug::timing_log_scalar( - "index_tts2_5.wav2vec2bert.graph.compute_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - if (status != GGML_STATUS_SUCCESS) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph compute failed"); - } - IndexTTS25SemanticEmbedding out; - out.frames = features.frames; - out.dims = kHidden; - out.values.resize(static_cast(features.frames * kHidden)); - timing_start = Clock::now(); - ggml_backend_tensor_get(output_, out.values.data(), 0, out.values.size() * sizeof(float)); - debug::timing_log_scalar( - "index_tts2_5.wav2vec2bert.output_read_ms", - engine::debug::elapsed_ms(timing_start, Clock::now())); - return out; - } - -private: - void clear_graph() { - if (graph_ != nullptr) { - engine::core::release_backend_graph_resources(execution_.backend(), graph_); - graph_ = nullptr; - } - if (gallocr_ != nullptr) { - ggml_gallocr_free(gallocr_); - gallocr_ = nullptr; - } - if (input_buffer_ != nullptr) { - ggml_backend_buffer_free(input_buffer_); - input_buffer_ = nullptr; - } - } - - engine::core::ExecutionContext & execution_; - std::shared_ptr weights_; - int64_t frames_ = 0; - std::unique_ptr input_ctx_; - std::unique_ptr ctx_; - ggml_tensor * features_ = nullptr; - ggml_tensor * keep_mask_ = nullptr; - ggml_tensor * attention_mask_ = nullptr; - ggml_tensor * distance_ids_ = nullptr; - ggml_tensor * output_ = nullptr; - ggml_cgraph * graph_ = nullptr; - ggml_gallocr_t gallocr_ = nullptr; - ggml_backend_buffer_t input_buffer_ = nullptr; -}; - -std::shared_ptr load_index_tts2_5_wav2vec2bert_weights( - const IndexTTS25Assets & assets, - ggml_backend_t backend, - engine::core::BackendType backend_type, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type, - size_t weight_context_bytes) { - if (assets.wav2vec2bert_weights == nullptr || assets.wav2vec2bert_stats == nullptr) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT requires model and stats tensor sources"); - } - auto weights = std::make_shared(); - weights->store = std::make_shared( - backend, - backend_type, - "index_tts2_5.wav2vec2bert.weights", - weight_context_bytes); - - const auto & source = *assets.wav2vec2bert_weights; - weights->feature_norm = binding::norm_from_source(*weights->store, source, "feature_projection.layer_norm", kFeatureDim); - weights->feature_projection = binding::linear_from_source( - *weights->store, - source, - "feature_projection.projection", - matmul_storage_type, - kHidden, - kFeatureDim, - true); - weights->layers.reserve(static_cast(kLayers)); - for (int64_t layer_index = 0; layer_index < kLayers; ++layer_index) { - weights->layers.push_back(load_layer( - *weights->store, - source, - layer_index, - matmul_storage_type, - conv_storage_type)); - } - weights->semantic_mean = weights->store->make_from_f32( - engine::core::TensorShape::from_dims({kHidden}), - matmul_storage_type, - assets.wav2vec2bert_stats->require_f32("mean", {kHidden})); - weights->semantic_std = weights->store->make_from_f32( - engine::core::TensorShape::from_dims({kHidden}), - matmul_storage_type, - wav2vec2bert_std(*assets.wav2vec2bert_stats)); - - weights->store->upload(); - assets.wav2vec2bert_weights->release_storage(); - assets.wav2vec2bert_stats->release_storage(); - return weights; -} - -IndexTTS25Wav2Vec2BertRuntime::IndexTTS25Wav2Vec2BertRuntime( - std::shared_ptr assets, - engine::core::ExecutionContext & execution, - size_t graph_arena_bytes, - size_t weight_context_bytes, - engine::assets::TensorStorageType matmul_storage_type, - engine::assets::TensorStorageType conv_storage_type) - : assets_(std::move(assets)), - execution_(&execution), - graph_arena_bytes_(graph_arena_bytes) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT runtime requires assets"); - } - if (graph_arena_bytes_ == 0) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph arena must be non-zero"); - } - weights_ = load_index_tts2_5_wav2vec2bert_weights( - *assets_, - execution.backend(), - execution.backend_type(), - matmul_storage_type, - conv_storage_type, - weight_context_bytes); -} - -IndexTTS25Wav2Vec2BertRuntime::~IndexTTS25Wav2Vec2BertRuntime() = default; - -void IndexTTS25Wav2Vec2BertRuntime::prepare(int64_t frames) { - if (execution_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT runtime execution context is missing"); - } - if (frames <= 0) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT prepare requires positive frames"); - } - if (graph_ != nullptr && graph_->frames() >= frames) { - return; - } - graph_.reset(); - graph_ = std::make_unique(*execution_, weights_, frames, graph_arena_bytes_); -} - -IndexTTS25SemanticEmbedding IndexTTS25Wav2Vec2BertRuntime::encode(const IndexTTS25SemanticFeatureOutput & features) { - if (graph_ == nullptr || graph_->frames() < features.frames) { - throw std::runtime_error("IndexTTS2.5 Wav2Vec2-BERT graph was not prepared for this reference length"); - } - return graph_->run(features); -} - -void IndexTTS25Wav2Vec2BertRuntime::release_graph() { - graph_.reset(); -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/session.cpp b/src/models/index_tts2_5/session.cpp deleted file mode 100644 index 8fd34df7..00000000 --- a/src/models/index_tts2_5/session.cpp +++ /dev/null @@ -1,838 +0,0 @@ -#include "engine/models/index_tts2_5/session.h" - -#include "engine/framework/debug/profiler.h" -#include "engine/framework/io/text.h" -#include "engine/framework/runtime/options.h" -#include "engine/framework/text/chunking.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -using Clock = std::chrono::steady_clock; -constexpr int64_t kConditionDim = 512; -constexpr int64_t kGptDim = 1280; -constexpr int64_t kStyleDim = 192; -constexpr int64_t kEmotionCount = 8; -constexpr int64_t kDiffusionSteps = 25; -constexpr float kInferenceCfgRate = 0.7F; - -std::shared_ptr require_assets(std::shared_ptr assets) { - if (assets == nullptr) { - throw std::runtime_error("IndexTTS2.5 session requires assets"); - } - return assets; -} - -void validate_matmul_weight_storage(engine::assets::TensorStorageType storage_type, const char * option_name) { - if (storage_type == engine::assets::TensorStorageType::Native || - storage_type == engine::assets::TensorStorageType::F32 || - storage_type == engine::assets::TensorStorageType::F16 || - storage_type == engine::assets::TensorStorageType::BF16 || - storage_type == engine::assets::TensorStorageType::Q8_0) { - return; - } - throw std::runtime_error(std::string(option_name) + " supports only native, f32, f16, bf16, and q8_0"); -} - -void validate_conv_weight_storage(engine::assets::TensorStorageType storage_type, const char * option_name) { - if (storage_type == engine::assets::TensorStorageType::Native || - storage_type == engine::assets::TensorStorageType::F32 || - storage_type == engine::assets::TensorStorageType::F16) { - return; - } - throw std::runtime_error(std::string(option_name) + " supports only native, f32, and f16"); -} - -uint64_t fnv1a_mix(uint64_t hash, const void * data, size_t size) { - const auto * bytes = static_cast(data); - for (size_t i = 0; i < size; ++i) { - hash ^= bytes[i]; - hash *= 1099511628211ull; - } - return hash; -} - -uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { - uint64_t hash = 1469598103934665603ull; - for (const float sample : audio.samples) { - uint32_t bits = 0; - std::memcpy(&bits, &sample, sizeof(bits)); - hash = fnv1a_mix(hash, &bits, sizeof(bits)); - } - return hash; -} - -IndexTTS25AudioIdentity audio_identity(const runtime::AudioBuffer & audio) { - return { - audio.sample_rate, - audio.channels, - static_cast(audio.samples.size()), - hash_audio_samples(audio), - }; -} - -bool same_identity(const IndexTTS25AudioIdentity & lhs, const IndexTTS25AudioIdentity & rhs) { - return lhs.sample_rate == rhs.sample_rate && - lhs.channels == rhs.channels && - lhs.sample_count == rhs.sample_count && - lhs.sample_hash == rhs.sample_hash; -} - -std::size_t resolve_cache_slots( - const runtime::SessionOptions & options, - std::initializer_list keys, - const char * option_name) { - constexpr int64_t kDefaultCacheSlots = 1; - const int64_t slots = runtime::parse_i64_option(options.options, keys) - .value_or(kDefaultCacheSlots); - if (slots < 0) { - throw std::runtime_error(std::string(option_name) + " must be non-negative"); - } - if (static_cast(slots) > static_cast(std::numeric_limits::max())) { - throw std::runtime_error(std::string(option_name) + " is too large"); - } - return static_cast(slots); -} - -std::vector channel_first_to_time_major( - const std::vector & values, - int64_t channels, - int64_t frames) { - if (static_cast(values.size()) != channels * frames) { - throw std::runtime_error("IndexTTS2.5 channel-first tensor size mismatch"); - } - std::vector out(static_cast(frames * channels)); - for (int64_t frame = 0; frame < frames; ++frame) { - for (int64_t channel = 0; channel < channels; ++channel) { - out[static_cast(frame * channels + channel)] = - values[static_cast(channel * frames + frame)]; - } - } - return out; -} - -std::vector concat_conditions( - const IndexTTS25S2MelSequence & prompt, - const IndexTTS25S2MelSequence & generated) { - if (prompt.dims != kConditionDim || generated.dims != kConditionDim) { - throw std::runtime_error("IndexTTS2.5 condition dimension mismatch"); - } - std::vector out; - out.reserve(prompt.values.size() + generated.values.size()); - out.insert(out.end(), prompt.values.begin(), prompt.values.end()); - out.insert(out.end(), generated.values.begin(), generated.values.end()); - return out; -} - -void append_silence(runtime::AudioBuffer & audio, int ms) { - if (ms <= 0 || audio.sample_rate <= 0 || audio.channels <= 0) { - return; - } - const int64_t samples = static_cast(audio.sample_rate) * ms / 1000; - audio.samples.insert( - audio.samples.end(), - static_cast(samples * audio.channels), - 0.0F); -} - -std::vector scaled_emotion_weights(const std::vector & values, float alpha) { - if (static_cast(values.size()) != kEmotionCount) { - throw std::runtime_error("IndexTTS2.5 emotion vector must contain exactly 8 values"); - } - std::vector out = values; - const float scale = std::clamp(alpha, 0.0F, 1.0F); - if (scale != 1.0F) { - for (float & value : out) { - value = static_cast(static_cast(value * scale * 10000.0F)) / 10000.0F; - } - } - return out; -} - -bool mem_saver_from_options(const runtime::SessionOptions & options) { - if (const auto value = runtime::find_option(options.options, {"index_tts2_5.mem_saver", "mem_saver"})) { - return runtime::parse_bool_option(*value, "index_tts2_5.mem_saver"); - } - return false; -} - -// Debug intermediate dumps, enabled by setting INDEXTTS25_DUMP_DIR. Files are -// written as NPY v1.0 so they can be compared against the official Python -// golden tensors with numpy directly. -std::string dump_dir_from_env() { - const char * dir = std::getenv("INDEXTTS25_DUMP_DIR"); - if (dir == nullptr || *dir == '\0') { - return {}; - } - return std::string(dir); -} - -void write_npy( - const std::filesystem::path & path, - const char * descr, - const std::vector & shape, - const void * data, - size_t byte_count) { - std::string header = "{'descr': '"; - header += descr; - header += "', 'fortran_order': False, 'shape': ("; - for (const int64_t dim : shape) { - header += std::to_string(dim); - header += ", "; - } - header += "), }"; - const size_t prefix = 10; // magic(6) + version(2) + header_len(2) - const size_t total = prefix + header.size() + 1; - const size_t padded = (total + 63) / 64 * 64; - header.append(padded - prefix - header.size() - 1, ' '); - header.push_back('\n'); - std::ofstream out(path, std::ios::binary | std::ios::trunc); - if (!out) { - throw std::runtime_error("IndexTTS2.5 failed to open dump file: " + path.string()); - } - out.write("\x93NUMPY\x01\x00", 8); - const uint16_t header_len = static_cast(header.size()); - out.write(reinterpret_cast(&header_len), sizeof(header_len)); - out.write(header.data(), static_cast(header.size())); - out.write(static_cast(data), static_cast(byte_count)); -} - -void dump_f32( - const std::string & dir, - const std::string & name, - std::vector shape, - const std::vector & values) { - if (dir.empty()) { - return; - } - int64_t elements = 1; - for (const int64_t dim : shape) { - elements *= dim; - } - if (elements != static_cast(values.size())) { - throw std::runtime_error("IndexTTS2.5 dump shape does not match value count: " + name); - } - write_npy(std::filesystem::path(dir) / (name + ".npy"), " shape, - const std::vector & values) { - if (dir.empty()) { - return; - } - int64_t elements = 1; - for (const int64_t dim : shape) { - elements *= dim; - } - if (elements != static_cast(values.size())) { - throw std::runtime_error("IndexTTS2.5 dump shape does not match value count: " + name); - } - write_npy(std::filesystem::path(dir) / (name + ".npy"), " assets) - : RuntimeSessionBase(options), - task_(task), - assets_(require_assets(std::move(assets))), - tokenizer_(assets_), - speaker_cache_(resolve_cache_slots(this->options(), {"index_tts2_5.speaker_cache_slots"}, "index_tts2_5.speaker_cache_slots")), - emotion_cache_(resolve_cache_slots(this->options(), {"index_tts2_5.emotion_cache_slots"}, "index_tts2_5.emotion_cache_slots")), - emotion_text_weights_cache_(resolve_cache_slots(this->options(), {"index_tts2_5.emotion_text_cache_slots"}, "index_tts2_5.emotion_text_cache_slots")) { - gpt_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"index_tts2_5.gpt_graph_arena_mb"}, gpt_graph_arena_bytes_); - s2mel_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"index_tts2_5.s2mel_graph_arena_mb"}, s2mel_graph_arena_bytes_); - reference_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"index_tts2_5.reference_graph_arena_mb"}, reference_graph_arena_bytes_); - emotion_text_prefill_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"index_tts2_5.emotion_text_prefill_graph_arena_mb"}, emotion_text_prefill_graph_arena_bytes_); - emotion_text_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"index_tts2_5.emotion_text_decode_graph_arena_mb"}, emotion_text_decode_graph_arena_bytes_); - weight_context_bytes_ = runtime::parse_size_mb_option( - options.options, {"index_tts2_5.weight_context_mb"}, weight_context_bytes_); - if (const auto value = runtime::parse_int_option(options.options, {"index_tts2_5.emotion_text_max_new_tokens"})) { - if (*value <= 0) { - throw std::runtime_error("index_tts2_5.emotion_text_max_new_tokens must be positive"); - } - emotion_text_max_new_tokens_ = *value; - } - if (const auto it = options.options.find("index_tts2_5.weight_type"); it != options.options.end()) { - matmul_weight_storage_type_ = engine::assets::parse_tensor_storage_type(it->second); - validate_matmul_weight_storage(matmul_weight_storage_type_, "index_tts2_5.weight_type"); - } - if (const auto it = options.options.find("index_tts2_5.conv_weight_type"); it != options.options.end()) { - conv_weight_storage_type_ = engine::assets::parse_tensor_storage_type(it->second); - validate_conv_weight_storage(conv_weight_storage_type_, "index_tts2_5.conv_weight_type"); - } - mem_saver_ = mem_saver_from_options(options); - for (const auto & [key, _] : options.options) { - if (key.rfind("index_tts2_5.", 0) == 0 && - key != "index_tts2_5.gpt_graph_arena_mb" && - key != "index_tts2_5.s2mel_graph_arena_mb" && - key != "index_tts2_5.reference_graph_arena_mb" && - key != "index_tts2_5.emotion_text_prefill_graph_arena_mb" && - key != "index_tts2_5.emotion_text_decode_graph_arena_mb" && - key != "index_tts2_5.emotion_text_max_new_tokens" && - key != "index_tts2_5.weight_context_mb" && - key != "index_tts2_5.weight_type" && - key != "index_tts2_5.conv_weight_type" && - key != "index_tts2_5.mem_saver" && - key != "index_tts2_5.speaker_cache_slots" && - key != "index_tts2_5.emotion_cache_slots" && - key != "index_tts2_5.emotion_text_cache_slots" && - key != "index_tts2_5.gpt.cuda_sampling_policy" && - key != "index_tts2_5.s2mel.cuda_sampling_policy") { - throw std::runtime_error("unknown IndexTTS2.5 session option: " + key); - } - } - if (task_.mode != runtime::RunMode::Offline || - (task_.task != runtime::VoiceTaskKind::Tts && task_.task != runtime::VoiceTaskKind::VoiceCloning)) { - throw std::runtime_error("IndexTTS2.5 currently supports offline TTS and voice-cloning sessions"); - } - - semantic_encoder_ = std::make_unique( - assets_, - execution_context(), - reference_graph_arena_bytes_, - weight_context_bytes_, - matmul_weight_storage_type_, - conv_weight_storage_type_); - semantic_codec_ = std::make_unique( - assets_, - execution_context(), - reference_graph_arena_bytes_, - weight_context_bytes_, - matmul_weight_storage_type_, - conv_weight_storage_type_); - style_encoder_ = std::make_unique( - assets_, - options.backend, - conv_weight_storage_type_); - gpt_ = std::make_unique( - assets_, - execution_context(), - gpt_graph_arena_bytes_, - weight_context_bytes_, - matmul_weight_storage_type_, - conv_weight_storage_type_); - s2mel_ = std::make_unique( - assets_, - execution_context(), - s2mel_graph_arena_bytes_, - weight_context_bytes_, - matmul_weight_storage_type_, - conv_weight_storage_type_); - vocoder_ = std::make_unique( - assets_, - options.backend, - conv_weight_storage_type_); - qwen_emotion_ = std::make_unique( - assets_, - execution_context(), - emotion_text_prefill_graph_arena_bytes_, - emotion_text_decode_graph_arena_bytes_, - weight_context_bytes_, - matmul_weight_storage_type_); - int64_t matrix_rows = 0; - for (const int64_t count : assets_->config.emo_num) { - matrix_rows += count; - } - speaker_matrix_ = assets_->speaker_matrix->require_f32("tensor", {matrix_rows, kStyleDim}); - emotion_matrix_ = assets_->emotion_matrix->require_f32("tensor", {matrix_rows, kGptDim}); - assets_->speaker_matrix->release_storage(); - assets_->emotion_matrix->release_storage(); -} - -bool IndexTTS25Session::AudioIdentityEqual::operator()( - const IndexTTS25AudioIdentity & lhs, - const IndexTTS25AudioIdentity & rhs) const { - return same_identity(lhs, rhs); -} - -std::string IndexTTS25Session::family() const { - return "index_tts2_5"; -} - -runtime::VoiceTaskKind IndexTTS25Session::task_kind() const { - return task_.task; -} - -runtime::RunMode IndexTTS25Session::run_mode() const { - return task_.mode; -} - -void IndexTTS25Session::prepare(const runtime::SessionPreparationRequest & request) { - if (request.text.has_value() || runtime::find_option(request.options, {"text", "prompt"}).has_value()) { - std::string text; - if (request.text.has_value()) { - text = request.text->text; - } else if (const auto value = runtime::find_option(request.options, {"text", "prompt"})) { - text = *value; - } - text = engine::io::trim_ascii_whitespace(text); - if (text.empty()) { - throw std::runtime_error("IndexTTS2.5 request requires text_input or text option"); - } - IndexTTS25GenerationOptions generation; - if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { - if (*value <= 0) { - throw std::runtime_error("IndexTTS2.5 max_tokens must be positive"); - } - generation.max_mel_tokens = *value; - } - if (const auto value = runtime::parse_int_option(request.options, {"num_beams"})) { - if (*value <= 0) { - throw std::runtime_error("IndexTTS2.5 num_beams must be positive"); - } - generation.num_beams = *value; - } - std::string lang; - if (const auto value = runtime::find_option(request.options, {"lang"})) { - lang = normalize_index_tts2_5_lang(*value); - } - const auto text_encoding = tokenizer_.encode_for_inference( - text, - IndexTTS25Request{}.max_text_tokens_per_segment, - lang); - for (const auto & segment : text_encoding.segment_token_ids) { - gpt_->prepare_generation( - static_cast(align_index_tts2_5_gpt_text_tokens(segment).size()), - generation.max_mel_tokens, - generation.num_beams); - } - } - mark_prepared(); -} - -const IndexTTS25Session::SpeakerState & IndexTTS25Session::resolve_speaker_state(const runtime::AudioBuffer & audio) { - const auto identity = audio_identity(audio); - if (const auto * cached = speaker_cache_.find(identity)) { - debug::trace_log_scalar("index_tts2_5.speaker_cache.hit", 1); - debug::trace_log_scalar("index_tts2_5.speaker_cache.slots", static_cast(speaker_cache_.capacity())); - debug::trace_log_scalar("index_tts2_5.speaker_cache.entries", static_cast(speaker_cache_.size())); - debug::trace_log_scalar("index_tts2_5.speaker_cache.evicted", 0); - return *cached; - } - const bool will_evict = speaker_cache_.capacity() > 0 && speaker_cache_.size() >= speaker_cache_.capacity(); - const auto start = Clock::now(); - const auto prepared = prepare_index_tts2_5_reference_audio( - audio.samples, - audio.sample_rate, - audio.channels, - assets_->config.s2mel, - static_cast(std::max(1, options().backend.threads)), - true); - semantic_encoder_->prepare(prepared.semantic_features.frames); - auto semantic = semantic_encoder_->encode(prepared.semantic_features); - // Official 2.5 regulates the raw (normalized) w2v-bert semantic directly; - // the semantic codec is only used to decode generated codes. - debug::trace_log_scalar("index_tts2_5.s2mel.reference_mel_frames", static_cast(prepared.mel.frames)); - s2mel_->prepare_length_regulator(semantic.frames, prepared.mel.frames); - auto prompt_condition = s2mel_->regulate_length( - semantic.values, - semantic.frames, - prepared.mel.frames); - - SpeakerState state; - state.identity = identity; - state.semantic = std::move(semantic); - state.reference_mel = prepared.mel; - state.style = style_encoder_->embed_fbank( - prepared.campplus_fbank.values, - prepared.campplus_fbank.frames, - prepared.campplus_fbank.dims); - state.prompt_condition = std::move(prompt_condition); - if (speaker_cache_.capacity() == 0) { - uncached_speaker_state_ = std::move(state); - } else { - speaker_cache_.put(identity, std::move(state)); - } - if (mem_saver_) { - semantic_encoder_->release_graph(); - semantic_codec_->release_graphs(); - s2mel_->release_pre_cfm_graphs(); - style_encoder_->release_graph(); - } - debug::trace_log_scalar("index_tts2_5.speaker_cache.hit", 0); - debug::trace_log_scalar("index_tts2_5.speaker_cache.slots", static_cast(speaker_cache_.capacity())); - debug::trace_log_scalar("index_tts2_5.speaker_cache.entries", static_cast(speaker_cache_.size())); - debug::trace_log_scalar("index_tts2_5.speaker_cache.evicted", will_evict ? 1 : 0); - debug::timing_log_scalar("index_tts2_5.speaker_state_ms", engine::debug::elapsed_ms(start)); - if (speaker_cache_.capacity() == 0) { - return *uncached_speaker_state_; - } - const auto * cached = speaker_cache_.find(identity); - if (cached == nullptr) { - throw std::runtime_error("IndexTTS2.5 speaker cache insert failed"); - } - return *cached; -} - -const IndexTTS25Session::EmotionState & IndexTTS25Session::resolve_emotion_state(const runtime::AudioBuffer & audio) { - const auto identity = audio_identity(audio); - if (const auto * cached = emotion_cache_.find(identity)) { - debug::trace_log_scalar("index_tts2_5.emotion_cache.hit", 1); - debug::trace_log_scalar("index_tts2_5.emotion_cache.slots", static_cast(emotion_cache_.capacity())); - debug::trace_log_scalar("index_tts2_5.emotion_cache.entries", static_cast(emotion_cache_.size())); - debug::trace_log_scalar("index_tts2_5.emotion_cache.evicted", 0); - return *cached; - } - const bool will_evict = emotion_cache_.capacity() > 0 && emotion_cache_.size() >= emotion_cache_.capacity(); - const auto start = Clock::now(); - const auto prepared = prepare_index_tts2_5_reference_audio( - audio.samples, - audio.sample_rate, - audio.channels, - assets_->config.s2mel, - static_cast(std::max(1, options().backend.threads)), - false); - semantic_encoder_->prepare(prepared.semantic_features.frames); - EmotionState state; - state.identity = identity; - state.semantic = semantic_encoder_->encode(prepared.semantic_features); - if (emotion_cache_.capacity() == 0) { - uncached_emotion_state_ = std::move(state); - } else { - emotion_cache_.put(identity, std::move(state)); - } - if (mem_saver_) { - semantic_encoder_->release_graph(); - } - debug::trace_log_scalar("index_tts2_5.emotion_cache.hit", 0); - debug::trace_log_scalar("index_tts2_5.emotion_cache.slots", static_cast(emotion_cache_.capacity())); - debug::trace_log_scalar("index_tts2_5.emotion_cache.entries", static_cast(emotion_cache_.size())); - debug::trace_log_scalar("index_tts2_5.emotion_cache.evicted", will_evict ? 1 : 0); - debug::timing_log_scalar("index_tts2_5.emotion_state_ms", engine::debug::elapsed_ms(start)); - if (emotion_cache_.capacity() == 0) { - return *uncached_emotion_state_; - } - const auto * cached = emotion_cache_.find(identity); - if (cached == nullptr) { - throw std::runtime_error("IndexTTS2.5 emotion cache insert failed"); - } - return *cached; -} - -std::vector IndexTTS25Session::explicit_emotion_matrix_vector( - const std::vector & emotion_weights, - const IndexTTS25StyleEmbedding & style, - bool use_random, - uint32_t seed) const { - if (static_cast(emotion_weights.size()) != kEmotionCount || - style.dims != kStyleDim || - static_cast(style.values.size()) != kStyleDim) { - throw std::runtime_error("IndexTTS2.5 explicit emotion matrix input shape mismatch"); - } - std::vector out(static_cast(kGptDim), 0.0F); - int64_t row_offset = 0; - std::mt19937 rng(seed); - for (int64_t emotion = 0; emotion < kEmotionCount; ++emotion) { - const int64_t rows = assets_->config.emo_num[static_cast(emotion)]; - int64_t best_row = 0; - if (rows <= 0) { - throw std::runtime_error("IndexTTS2.5 emo_num contains non-positive group size"); - } - if (use_random) { - std::uniform_int_distribution distribution(0, rows - 1); - best_row = distribution(rng); - } else { - float best_score = -std::numeric_limits::infinity(); - for (int64_t row = 0; row < rows; ++row) { - const float * matrix_row = speaker_matrix_.data() + static_cast((row_offset + row) * kStyleDim); - float dot = 0.0F; - float lhs_norm = 0.0F; - float rhs_norm = 0.0F; - for (int64_t dim = 0; dim < kStyleDim; ++dim) { - const float lhs = style.values[static_cast(dim)]; - const float rhs = matrix_row[dim]; - dot += lhs * rhs; - lhs_norm += lhs * lhs; - rhs_norm += rhs * rhs; - } - const float score = dot / (std::sqrt(lhs_norm) * std::sqrt(rhs_norm) + 1.0e-12F); - if (score > best_score) { - best_score = score; - best_row = row; - } - } - } - const float * emotion_row = emotion_matrix_.data() + static_cast((row_offset + best_row) * kGptDim); - for (int64_t dim = 0; dim < kGptDim; ++dim) { - out[static_cast(dim)] += emotion_weights[static_cast(emotion)] * emotion_row[dim]; - } - row_offset += rows; - } - return out; -} - -std::vector IndexTTS25Session::resolve_emotion_vector( - const IndexTTS25Request & request, - const SpeakerState & speaker, - const EmotionState & emotion) { - std::vector explicit_weights; - if (request.use_emotion_text) { - const auto emotion_text = request.emotion_text.value_or(request.text); - if (const auto * cached = emotion_text_weights_cache_.find(emotion_text)) { - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.hit", 1); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.slots", static_cast(emotion_text_weights_cache_.capacity())); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.entries", static_cast(emotion_text_weights_cache_.size())); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.evicted", 0); - explicit_weights = *cached; - } else { - const bool will_evict = - emotion_text_weights_cache_.capacity() > 0 && - emotion_text_weights_cache_.size() >= emotion_text_weights_cache_.capacity(); - explicit_weights = qwen_emotion_->infer(emotion_text, emotion_text_max_new_tokens_).values; - if (mem_saver_) { - qwen_emotion_->release_graphs(); - } - emotion_text_weights_cache_.put(emotion_text, explicit_weights); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.hit", 0); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.slots", static_cast(emotion_text_weights_cache_.capacity())); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.entries", static_cast(emotion_text_weights_cache_.size())); - debug::trace_log_scalar("index_tts2_5.qwen_emotion.text_cache.evicted", will_evict ? 1 : 0); - } - } else if (request.emotion_vector.has_value()) { - explicit_weights = *request.emotion_vector; - } - - if (explicit_weights.empty()) { - return gpt_->merge_emotion_vector( - speaker.semantic.values, - speaker.semantic.frames, - emotion.semantic.values, - emotion.semantic.frames, - request.emotion_alpha); - } - - auto scaled_weights = scaled_emotion_weights(explicit_weights, request.emotion_alpha); - auto matrix = explicit_emotion_matrix_vector( - scaled_weights, - speaker.style, - request.use_random_emotion, - request.generation.seed); - const float weight_sum = std::accumulate(scaled_weights.begin(), scaled_weights.end(), 0.0F); - if (weight_sum != 1.0F) { - auto base = gpt_->merge_emotion_vector( - speaker.semantic.values, - speaker.semantic.frames, - emotion.semantic.values, - emotion.semantic.frames, - 1.0F); - for (size_t i = 0; i < matrix.size(); ++i) { - matrix[i] += (1.0F - weight_sum) * base[i]; - } - } - return matrix; -} - -runtime::AudioBuffer IndexTTS25Session::synthesize_segment( - const std::vector & text_tokens, - int32_t lang_id, - size_t segment_index, - const std::string & dump_dir, - const SpeakerState & speaker, - const EmotionState & emotion, - const std::vector & emotion_vector, - const IndexTTS25GenerationOptions & options, - uint32_t segment_seed) { - const std::string seg = "_seg" + std::to_string(segment_index); - dump_i32(dump_dir, "text_tokens" + seg, {1, static_cast(text_tokens.size())}, text_tokens); - IndexTTS25GptGenerationRequest generation; - generation.text_tokens = text_tokens; - generation.speaker_style = speaker.style.values; - generation.lang_id = lang_id; - generation.emotion_semantic = emotion.semantic.values; - generation.emotion_frames = emotion.semantic.frames; - generation.emotion_vector = emotion_vector; - generation.top_p = options.top_p; - generation.top_k = options.top_k; - generation.temperature = options.temperature; - generation.repetition_penalty = options.repetition_penalty; - generation.do_sample = options.do_sample; - generation.length_penalty = options.length_penalty; - generation.num_beams = options.num_beams; - generation.max_mel_tokens = options.max_mel_tokens; - generation.seed = segment_seed; - - const auto gen_start = Clock::now(); - auto generated = gpt_->generate_speech(generation); - if (mem_saver_) { - gpt_->release_conditioning_graphs(); - } - debug::timing_log_scalar("index_tts2_5.gpt.generate_ms", engine::debug::elapsed_ms(gen_start)); - dump_i32(dump_dir, "gpt_codes" + seg, {1, static_cast(generated.codes.size())}, generated.codes); - if (generated.codes.empty()) { - throw std::runtime_error("IndexTTS2.5 GPT generated no acoustic codes"); - } - const int64_t code_frames = static_cast(generated.codes.size()); - - const auto s2mel_start = Clock::now(); - semantic_codec_->prepare_codes(code_frames); - auto semantic = semantic_codec_->codes_to_embedding(generated.codes, code_frames); - if (mem_saver_) { - semantic_codec_->release_graphs(); - } - auto content = channel_first_to_time_major( - semantic.embedding_channel_first, - semantic.dims, - semantic.frames); - dump_f32(dump_dir, "s_infer" + seg, {1, semantic.frames, semantic.dims}, content); - const int64_t target_frames = static_cast(static_cast(semantic.frames) * 1.72F); - const int64_t total_frames = speaker.prompt_condition.frames + target_frames; - s2mel_->prepare_length_regulator(semantic.frames, target_frames); - auto generated_condition = s2mel_->regulate_length(content, semantic.frames, target_frames); - dump_f32( - dump_dir, - "lr_gen" + seg, - {1, generated_condition.frames, generated_condition.dims}, - generated_condition.values); - auto condition = concat_conditions(speaker.prompt_condition, generated_condition); - dump_f32(dump_dir, "cat_condition" + seg, {1, total_frames, kConditionDim}, condition); - if (mem_saver_) { - gpt_->release_generation_graphs(); - s2mel_->release_pre_cfm_graphs(); - } - s2mel_->prepare_cfm(total_frames, kInferenceCfgRate > 0.0F); - auto mel = s2mel_->infer_mel( - condition, - total_frames, - speaker.reference_mel.values, - speaker.reference_mel.frames, - speaker.style.values, - kDiffusionSteps, - kInferenceCfgRate, - segment_seed, - generated.rng_offset_blocks); - dump_f32(dump_dir, "mel_out" + seg, {1, mel.channels, mel.frames}, mel.values); - debug::timing_log_scalar("index_tts2_5.s2mel.total_ms", engine::debug::elapsed_ms(s2mel_start)); - if (mem_saver_) { - s2mel_->release_cfm_graph(); - } - - const auto vocoder_start = Clock::now(); - auto audio = vocoder_->synthesize(mel.values, mel.frames); - debug::timing_log_scalar("index_tts2_5.vocoder_ms", engine::debug::elapsed_ms(vocoder_start)); - if (mem_saver_) { - vocoder_->release_runtime_graph(); - } - runtime::AudioBuffer out; - out.sample_rate = audio.sample_rate; - out.channels = 1; - out.samples = std::move(audio.waveform); - return out; -} - -runtime::TaskResult IndexTTS25Session::run(const runtime::TaskRequest & request) { - require_prepared("IndexTTS2.5 run"); - const auto wall_start = Clock::now(); - auto parsed = parse_index_tts2_5_request(request); - if (!parsed.speaker_audio.has_value()) { - throw std::runtime_error("IndexTTS2.5 request requires speaker audio"); - } - const runtime::AudioBuffer * emotion_audio = parsed.emotion_audio.has_value() - ? &*parsed.emotion_audio - : &*parsed.speaker_audio; - if (parsed.emotion_vector.has_value() || parsed.use_emotion_text) { - emotion_audio = &*parsed.speaker_audio; - } - - const auto & speaker = resolve_speaker_state(*parsed.speaker_audio); - const bool emotion_same_as_speaker = same_identity(audio_identity(*emotion_audio), speaker.identity); - debug::trace_log_scalar("index_tts2_5.emotion_audio.same_as_speaker", emotion_same_as_speaker); - const auto & emotion = resolve_emotion_state(*emotion_audio); - const auto emotion_vector = resolve_emotion_vector(parsed, speaker, emotion); - if (mem_saver_) { - gpt_->release_conditioning_graphs(); - } - const std::string dump_dir = dump_dir_from_env(); - dump_f32(dump_dir, "campplus_embedding", {1, speaker.style.dims}, speaker.style.values); - dump_f32(dump_dir, "speech_condition", {1, speaker.semantic.frames, speaker.semantic.dims}, speaker.semantic.values); - dump_f32( - dump_dir, - "prompt_condition", - {1, speaker.prompt_condition.frames, speaker.prompt_condition.dims}, - speaker.prompt_condition.values); - dump_f32(dump_dir, "ref_mel", {1, speaker.reference_mel.channels, speaker.reference_mel.frames}, speaker.reference_mel.values); - dump_f32(dump_dir, "emo_vec", {1, static_cast(emotion_vector.size())}, emotion_vector); - const auto text_chunk_size = engine::text::parse_text_chunk_size_override(request.options); - const auto text_chunk_mode = engine::text::parse_text_chunk_mode_override(request.options) - .value_or(engine::text::TextChunkMode::Default); - const std::vector text_chunks = text_chunk_size.has_value() - ? engine::text::split_text_chunks(parsed.text, *text_chunk_size, text_chunk_mode) - : std::vector{parsed.text}; - if (text_chunks.empty()) { - throw std::runtime_error("IndexTTS2.5 text chunking produced no chunks"); - } - - std::vector> segment_token_ids; - std::vector segment_lang_ids; - for (const auto & text_chunk : text_chunks) { - const auto text_encoding = tokenizer_.encode_for_inference( - text_chunk, - parsed.max_text_tokens_per_segment, - parsed.lang); - const int32_t lang_id = IndexTTS25TextTokenizer::lang_to_id(text_encoding.lang); - for (const auto & ids : text_encoding.segment_token_ids) { - segment_token_ids.push_back(ids); - segment_lang_ids.push_back(lang_id); - } - } - - runtime::AudioBuffer merged; - for (size_t i = 0; i < segment_token_ids.size(); ++i) { - if (i > 0) { - append_silence(merged, parsed.interval_silence_ms); - } - auto segment_audio = synthesize_segment( - segment_token_ids[i], - segment_lang_ids[i], - i, - dump_dir, - speaker, - emotion, - emotion_vector, - parsed.generation, - parsed.generation.seed + static_cast(i)); - runtime::append_audio_buffer(merged, segment_audio); - } - - runtime::TaskResult result; - result.audio_output = std::move(merged); - debug::trace_log_scalar("index_tts2_5.path.use_emotion_text", parsed.use_emotion_text); - debug::trace_log_scalar("index_tts2_5.path.has_emotion_vector", parsed.emotion_vector.has_value()); - debug::trace_log_scalar("index_tts2_5.path.has_emotion_audio", parsed.emotion_audio.has_value()); - if (text_chunk_size.has_value()) { - debug::trace_log_scalar("index_tts2_5.text_chunk_size", *text_chunk_size); - debug::trace_log_scalar("index_tts2_5.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); - } - debug::trace_log_scalar("index_tts2_5.text_chunk_count", static_cast(text_chunks.size())); - debug::trace_log_scalar("index_tts2_5.text.segment_count", static_cast(segment_token_ids.size())); - debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); - return result; -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/style_encoder.cpp b/src/models/index_tts2_5/style_encoder.cpp deleted file mode 100644 index bb679438..00000000 --- a/src/models/index_tts2_5/style_encoder.cpp +++ /dev/null @@ -1,38 +0,0 @@ -#include "engine/models/index_tts2_5/style_encoder.h" - -#include -#include - -namespace engine::models::index_tts2_5 { - -IndexTTS25StyleEncoder::IndexTTS25StyleEncoder( - std::shared_ptr assets, - core::BackendConfig backend, - engine::assets::TensorStorageType weight_storage_type) - : assets_(std::move(assets)) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 style encoder requires assets"); - } - engine::modules::CampplusEncoderConfig config; - config.feat_dim = assets_->config.s2mel.n_mels; - config.embedding_size = assets_->config.s2mel.style_dim; - config.weight_storage_type = weight_storage_type; - component_ = engine::modules::CampplusEncoderComponent::load_from_tensor_source( - assets_->campplus_weights, - std::move(backend), - config); -} - -IndexTTS25StyleEmbedding IndexTTS25StyleEncoder::embed_fbank( - const std::vector & features, - int64_t frames, - int64_t dims) const { - const auto out = component_.embed_from_features(features, frames, dims); - return {out.embedding, out.embedding_size}; -} - -void IndexTTS25StyleEncoder::release_graph() { - component_.release_runtime_graph(); -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/tokenizer_text.cpp b/src/models/index_tts2_5/tokenizer_text.cpp deleted file mode 100644 index 1c061d3e..00000000 --- a/src/models/index_tts2_5/tokenizer_text.cpp +++ /dev/null @@ -1,691 +0,0 @@ -#include "engine/models/index_tts2_5/tokenizer_text.h" - -#include "engine/framework/text/chinese_normalization.h" -#include "engine/framework/text/text_normalization.h" - -#include "bpe-core.h" -#include "unicode.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -namespace vendor = llama_tokenizer_vendor; - -// IndexTTS-2.5 pads every text segment with a trailing token id 1. -constexpr int32_t kSegmentPadTokenId = 1; - -std::string decode_base64(const std::string & input) { - static const std::array table = [] { - std::array values{}; - values.fill(-1); - for (int i = 0; i < 26; ++i) { - values[static_cast('A' + i)] = static_cast(i); - values[static_cast('a' + i)] = static_cast(26 + i); - } - for (int i = 0; i < 10; ++i) { - values[static_cast('0' + i)] = static_cast(52 + i); - } - values[static_cast('+')] = 62; - values[static_cast('/')] = 63; - return values; - }(); - - std::string out; - int bits = 0; - int value = 0; - for (const unsigned char ch : input) { - if (ch == '=') { - break; - } - const int8_t digit = table[ch]; - if (digit < 0) { - throw std::runtime_error("IndexTTS2.5 tiktoken vocabulary contains invalid base64 token bytes"); - } - value = (value << 6) | digit; - bits += 6; - if (bits >= 8) { - bits -= 8; - out.push_back(static_cast((value >> bits) & 0xff)); - } - } - return out; -} - -// The vendored llama BPE runtime works in the GPT-2 byte-to-unicode domain -// (e.g. space becomes U+0120) so that byte-level merges, including tokens that -// split a UTF-8 codepoint, are reproduced exactly. tiktoken ranks are keyed by -// raw bytes, so every token is mapped once at load time. -std::string map_token_bytes(const std::string & bytes) { - std::string mapped; - for (const unsigned char byte : bytes) { - mapped += unicode_byte_to_utf8(byte); - } - return mapped; -} - -std::string pair_key(const std::string & left, const std::string & right) { - std::string key = left; - key.push_back('\0'); - key += right; - return key; -} - -// Language codes in the LANGUAGES order of indextts/utils/tokenizer.py. The -// first 99 entries double as the <|lang|> special tokens below; the remaining -// codes (plus the fallback "common") only index the GPT lang_embedding table. -const std::array kLanguages = { - "en", "zh", "de", "es", "ru", "ko", "fr", "ja", "pt", "tr", - "pl", "ca", "nl", "ar", "sv", "it", "id", "hi", "fi", "vi", - "he", "uk", "el", "ms", "cs", "ro", "da", "hu", "ta", "no", - "th", "ur", "hr", "bg", "lt", "la", "mi", "ml", "cy", "sk", - "te", "fa", "lv", "bn", "sr", "az", "sl", "kn", "et", "mk", - "br", "eu", "is", "hy", "ne", "mn", "bs", "kk", "sq", "sw", - "gl", "mr", "pa", "si", "km", "sn", "yo", "so", "af", "oc", - "ka", "be", "tg", "sd", "gu", "am", "yi", "lo", "uz", "fo", - "ht", "ps", "tk", "nn", "mt", "sa", "lb", "my", "bo", "tl", - "mg", "as", "tt", "haw", "ln", "ha", "ba", "jw", "su", -}; -const std::array kEmbeddingOnlyLanguages = { - "yue", "minnan", "wuyu", "dialect", "zh/en", "en/zh", "common", -}; -constexpr int32_t kCommonLangId = 105; - -void add_special_token(vendor::BpeVocabulary & vocab, const std::string & text, int32_t id) { - vocab.token_to_id.emplace(text, id); - vocab.id_to_token.emplace(id, vendor::TokenData{text, vendor::TOKEN_ATTR_CONTROL}); -} - -// Special token order must match indextts/utils/tokenizer.py exactly: -// ids are assigned sequentially starting right after the mergeable ranks. -void register_special_tokens(vendor::BpeVocabulary & vocab, int32_t base_id) { - static const std::array kAudioEvents = { - "ASR", "AED", "SER", "Speech", "/Speech", "BGM", "/BGM", - "Laughter", "/Laughter", "Applause", "/Applause", - }; - static const std::array kEmotions = { - "HAPPY", "SAD", "ANGRY", "NEUTRAL", - }; - static const std::array kTasks = { - "translate", "transcribe", "startoflm", "startofprev", "nospeech", "notimestamps", - }; - static const std::array kTtsVocal = { - "TTS/B", "TTS/O", "TTS/Q", "TTS/A", "TTS/CO", "TTS/CL", "TTS/H", - }; - - int32_t id = base_id; - add_special_token(vocab, "<|endoftext|>", id++); - add_special_token(vocab, "<|startoftranscript|>", id++); - for (const char * lang : kLanguages) { - add_special_token(vocab, "<|" + std::string(lang) + "|>", id++); - } - for (const char * event : kAudioEvents) { - add_special_token(vocab, "<|" + std::string(event) + "|>", id++); - } - for (const char * emotion : kEmotions) { - add_special_token(vocab, "<|" + std::string(emotion) + "|>", id++); - } - for (const char * task : kTasks) { - add_special_token(vocab, "<|" + std::string(task) + "|>", id++); - } - for (int i = 1; i <= 30; ++i) { - add_special_token(vocab, "<|SPECIAL_TOKEN_" + std::to_string(i) + "|>", id++); - } - for (const char * vocal : kTtsVocal) { - add_special_token(vocab, "<|" + std::string(vocal) + "|>", id++); - } - for (int i = 1; i <= 13; ++i) { - char name[32]; - std::snprintf(name, sizeof(name), "<|TTS/SP%02d|>", i); - add_special_token(vocab, name, id++); - } - // Timestamps <|0.00|> .. <|30.00|> in 0.02 steps; i * 0.02 == i / 50. - for (int i = 0; i <= 1500; ++i) { - char name[32]; - std::snprintf(name, sizeof(name), "<|%d.%02d|>", i / 50, (i * 2) % 100); - add_special_token(vocab, name, id++); - } -} - -std::shared_ptr load_tiktoken_vocabulary(const std::filesystem::path & vocab_path) { - std::ifstream input(vocab_path, std::ios::binary); - if (!input) { - throw std::runtime_error("IndexTTS2.5 failed to open tiktoken vocabulary: " + vocab_path.string()); - } - - auto vocab = std::make_shared(); - vocab->pre_type = vendor::PreTokenizerType::Gpt2; - - std::string line; - int64_t mergeable_count = 0; - while (std::getline(input, line)) { - if (!line.empty() && line.back() == '\r') { - line.pop_back(); - } - if (line.empty()) { - continue; - } - std::istringstream parts(line); - std::string token_base64; - int64_t rank = -1; - if (!(parts >> token_base64 >> rank) || rank < 0 || rank > INT32_MAX) { - throw std::runtime_error("IndexTTS2.5 tiktoken vocabulary has an invalid line: " + line); - } - const std::string bytes = decode_base64(token_base64); - const auto token_id = static_cast(rank); - const std::string mapped = map_token_bytes(bytes); - vocab->token_to_id.emplace(mapped, token_id); - vocab->id_to_token.emplace(token_id, vendor::TokenData{mapped, 0}); - // tiktoken ranks double as merge priorities: an adjacent pair merges - // iff its concatenation is a token, with that token's rank. Register - // every split so find_bpe_rank(left, right) == rank(left + right). - for (size_t split = 1; split < bytes.size(); ++split) { - vocab->bpe_ranks.emplace( - pair_key(map_token_bytes(bytes.substr(0, split)), map_token_bytes(bytes.substr(split))), - token_id); - } - ++mergeable_count; - } - if (mergeable_count == 0) { - throw std::runtime_error("IndexTTS2.5 tiktoken vocabulary is empty: " + vocab_path.string()); - } - - register_special_tokens(*vocab, static_cast(mergeable_count)); - vendor::rebuild_special_tokens_cache(*vocab); - return vocab; -} - -size_t utf8_codepoint_size(unsigned char byte) { - if ((byte & 0x80U) == 0U) { - return 1; - } - if ((byte & 0xE0U) == 0xC0U) { - return 2; - } - if ((byte & 0xF0U) == 0xE0U) { - return 3; - } - if ((byte & 0xF8U) == 0xF0U) { - return 4; - } - return 1; -} - -uint32_t decode_utf8_codepoint(const std::string & text, size_t offset, size_t size) { - const auto byte = [&](size_t i) { return static_cast(text[offset + i]); }; - if (size == 1) { - return byte(0); - } - if (size == 2 && offset + 1 < text.size()) { - return ((byte(0) & 0x1FU) << 6U) | (byte(1) & 0x3FU); - } - if (size == 3 && offset + 2 < text.size()) { - return ((byte(0) & 0x0FU) << 12U) | ((byte(1) & 0x3FU) << 6U) | (byte(2) & 0x3FU); - } - if (size == 4 && offset + 3 < text.size()) { - return ((byte(0) & 0x07U) << 18U) | ((byte(1) & 0x3FU) << 12U) | ((byte(2) & 0x3FU) << 6U) | (byte(3) & 0x3FU); - } - return byte(0); -} - -uint32_t next_utf8_codepoint(const std::string & text, size_t & offset) { - const size_t size = std::min(utf8_codepoint_size(static_cast(text[offset])), text.size() - offset); - const uint32_t cp = decode_utf8_codepoint(text, offset, size); - offset += size; - return cp; -} - -bool is_han_codepoint(uint32_t cp) { - return cp >= 0x4E00U && cp <= 0x9FFFU; -} - -bool contains_han(const std::string & text) { - for (size_t i = 0; i < text.size();) { - if (is_han_codepoint(next_utf8_codepoint(text, i))) { - return true; - } - } - return false; -} - -std::string lowercase_ascii(std::string text) { - for (char & ch : text) { - ch = static_cast(std::tolower(static_cast(ch))); - } - return text; -} - -std::string uppercase_ascii(std::string text) { - for (char & ch : text) { - ch = static_cast(std::toupper(static_cast(ch))); - } - return text; -} - -bool is_kana(const std::string & text) { - if (text.empty()) { - return false; - } - bool all_hiragana = true; - bool all_katakana = true; - for (size_t i = 0; i < text.size();) { - const uint32_t cp = next_utf8_codepoint(text, i); - if (cp < 0x3040U || cp > 0x309FU) { - all_hiragana = false; - } - if (cp < 0x30A0U || cp > 0x30FFU) { - all_katakana = false; - } - } - return all_hiragana || all_katakana; -} - -struct AnnotationMatch { - size_t end = 0; // one past the match; 0 when there is no match at pos - size_t word_begin = 0; - size_t word_end = 0; - size_t pron_begin = 0; - size_t pron_end = 0; -}; - -// Matches <([^|>\n]+)\|([^>\n]+)> anchored at pos. -AnnotationMatch match_pronunciation_annotation(const std::string & text, size_t pos) { - AnnotationMatch match; - if (text[pos] != '<') { - return match; - } - size_t cursor = pos + 1; - const size_t word_begin = cursor; - while (cursor < text.size() && text[cursor] != '|' && text[cursor] != '>' && text[cursor] != '\n') { - ++cursor; - } - if (cursor == word_begin || cursor >= text.size() || text[cursor] != '|') { - return match; - } - match.word_begin = word_begin; - match.word_end = cursor; - const size_t pron_begin = ++cursor; - while (cursor < text.size() && text[cursor] != '>' && text[cursor] != '\n') { - ++cursor; - } - if (cursor == pron_begin || cursor >= text.size()) { - return AnnotationMatch{}; - } - match.pron_begin = pron_begin; - match.pron_end = cursor; - match.end = cursor + 1; - return match; -} - -// Base-26 spreadsheet-style index ("a".."z", "aa"..), mirroring the official -// TextNormalizer._protect_pronunciation_annotations placeholder naming. -std::string alpha_placeholder_index(size_t n) { - std::string s; - while (true) { - s.insert(s.begin(), static_cast('a' + (n % 26))); - const size_t q = n / 26; - if (q == 0) { - break; - } - n = q - 1; - } - return s; -} - -using PronunciationPlaceholders = std::vector>; - -// Replaces annotations with letter-only placeholders so -// text normalization cannot rewrite their digits/symbols (e.g. XING2). -std::pair protect_pronunciation_annotations(const std::string & text) { - std::string out; - out.reserve(text.size()); - PronunciationPlaceholders placeholders; - size_t pos = 0; - while (pos < text.size()) { - const auto match = match_pronunciation_annotation(text, pos); - if (match.end == 0) { - out.push_back(text[pos++]); - continue; - } - std::string key = "PRONPLACEHOLDER" + alpha_placeholder_index(placeholders.size()) + "PRONPLACEHOLDER"; - placeholders.emplace_back(key, text.substr(pos, match.end - pos)); - out += key; - pos = match.end; - } - return {out, placeholders}; -} - -std::string restore_pronunciation_annotations(std::string text, const PronunciationPlaceholders & placeholders) { - for (const auto & [key, original] : placeholders) { - size_t at = 0; - while ((at = text.find(key, at)) != std::string::npos) { - text.replace(at, key.size(), original); - at += original.size(); - } - } - return text; -} - -// Expands annotations (see infer_v2_5.py -// apply_pronunciation_annotations): -// Chinese word -> <|SPECIAL_TOKEN_2|>PRON<|SPECIAL_TOKEN_2|> -// other word -> <|SPECIAL_TOKEN_1|>PRON<|SPECIAL_TOKEN_1|> -// kana pron -> inlined as " PRON " -std::string apply_pronunciation_annotations(const std::string & text) { - std::string out; - out.reserve(text.size()); - size_t pos = 0; - while (pos < text.size()) { - const auto match = match_pronunciation_annotation(text, pos); - if (match.end == 0) { - out.push_back(text[pos++]); - continue; - } - const std::string word = text.substr(match.word_begin, match.word_end - match.word_begin); - const std::string pron = uppercase_ascii(text.substr(match.pron_begin, match.pron_end - match.pron_begin)); - if (is_kana(pron)) { - out.push_back(' '); - out += pron; - out.push_back(' '); - } else { - const char * wrapper = contains_han(word) ? "<|SPECIAL_TOKEN_2|>" : "<|SPECIAL_TOKEN_1|>"; - out += wrapper; - out += pron; - out += wrapper; - } - pos = match.end; - } - return out; -} - -// Uppercases the name inside <|...|> markers: re.sub(r'<\|([^|]+)\|>', upper). -std::string uppercase_special_token_names(const std::string & text) { - std::string out; - out.reserve(text.size()); - size_t pos = 0; - while (pos < text.size()) { - if (text[pos] != '<' || pos + 1 >= text.size() || text[pos + 1] != '|') { - out.push_back(text[pos++]); - continue; - } - size_t cursor = pos + 2; - while (cursor < text.size() && text[cursor] != '|') { - ++cursor; - } - if (cursor == pos + 2 || cursor + 1 >= text.size() || text[cursor + 1] != '>') { - out.push_back(text[pos++]); - continue; - } - out += "<|"; - out += uppercase_ascii(text.substr(pos + 2, cursor - (pos + 2))); - out += "|>"; - pos = cursor + 2; - } - return out; -} - -bool is_segment_delimiter(uint32_t cp) { - switch (cp) { - case U',': - case U'.': - case U'!': - case U'?': - case U';': - case U':': - case U'\n': - case 0xFF0CU: // ， - case 0x3002U: // 。 - case 0xFF01U: // ！ - case 0xFF1FU: // ？ - case 0x3001U: // 、 - case 0xFF1BU: // ； - case 0xFF1AU: // ： - return true; - default: - return false; - } -} - -// re.split(r'(?<=[，。！？、；：,\.!\?;:\n])', piece): split after each delimiter. -std::vector split_after_delimiters(const std::string & piece) { - std::vector parts; - std::string current; - for (size_t i = 0; i < piece.size();) { - const size_t begin = i; - const uint32_t cp = next_utf8_codepoint(piece, i); - current.append(piece, begin, i - begin); - if (is_segment_delimiter(cp)) { - parts.push_back(std::move(current)); - current.clear(); - } - } - if (!current.empty()) { - parts.push_back(std::move(current)); - } - return parts; -} - -// Matches "<|SPECIAL_TOKEN_|>" at pos; returns the match length or 0. -size_t match_special_token_marker(const std::string & text, size_t pos) { - static const std::string kPrefix = "<|SPECIAL_TOKEN_"; - if (text.compare(pos, kPrefix.size(), kPrefix) != 0) { - return 0; - } - size_t cursor = pos + kPrefix.size(); - const size_t digits_begin = cursor; - while (cursor < text.size() && std::isdigit(static_cast(text[cursor])) != 0) { - ++cursor; - } - if (cursor == digits_begin || cursor + 1 >= text.size() || text[cursor] != '|' || text[cursor + 1] != '>') { - return 0; - } - return cursor + 2 - pos; -} - -// SPLIT_PROTECTED_PATTERN spans (<|SPECIAL_TOKEN_n|>...<|SPECIAL_TOKEN_n|>) -// are kept atomic during segmentation. -std::vector> split_atomic_pieces(const std::string & text) { - std::vector> pieces; - size_t pos = 0; - while (pos < text.size()) { - size_t opener = std::string::npos; - size_t opener_len = 0; - for (size_t i = pos; i < text.size(); ++i) { - const size_t len = match_special_token_marker(text, i); - if (len > 0) { - opener = i; - opener_len = len; - break; - } - } - if (opener == std::string::npos) { - break; - } - size_t closer = std::string::npos; - size_t closer_len = 0; - for (size_t i = opener + opener_len; i < text.size(); ++i) { - const size_t len = match_special_token_marker(text, i); - if (len > 0) { - closer = i; - closer_len = len; - break; - } - } - if (closer == std::string::npos) { - break; - } - if (opener > pos) { - pieces.emplace_back(text.substr(pos, opener - pos), false); - } - pieces.emplace_back(text.substr(opener, closer + closer_len - opener), true); - pos = closer + closer_len; - } - if (pos < text.size()) { - pieces.emplace_back(text.substr(pos), false); - } - return pieces; -} - -} // namespace - -IndexTTS25TextTokenizer::IndexTTS25TextTokenizer(std::shared_ptr assets) - : assets_(std::move(assets)) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 text tokenizer requires assets"); - } - vocab_ = load_tiktoken_vocabulary(assets_->resources.require_file("tiktoken")); -} - -std::string IndexTTS25TextTokenizer::normalize_english(const std::string & text) const { - engine::text::EnglishTextNormalizationOptions options; - options.expand_common_contractions = true; - options.index_tts_punctuation = true; - options.uppercase_ascii = false; - return engine::text::normalize_english_text(text, options); -} - -std::string IndexTTS25TextTokenizer::normalize_chinese(const std::string & text) const { - return engine::text::normalize_chinese_text( - text, - engine::text::ChineseTextNormalizationTarget::IndexTTS); -} - -std::vector IndexTTS25TextTokenizer::encode(const std::string & text) const { - return vendor::tokenize_bpe(*vocab_, text, true); -} - -int32_t IndexTTS25TextTokenizer::special_token_id(const std::string & token_text) const { - const auto it = vocab_->token_to_id.find(token_text); - return it == vocab_->token_to_id.end() ? -1 : it->second; -} - -int32_t IndexTTS25TextTokenizer::lang_to_id(const std::string & lang) { - const std::string normalized = lowercase_ascii(lang); - for (size_t i = 0; i < kLanguages.size(); ++i) { - if (normalized == kLanguages[i]) { - return static_cast(i); - } - } - for (size_t i = 0; i < kEmbeddingOnlyLanguages.size(); ++i) { - if (normalized == kEmbeddingOnlyLanguages[i]) { - return static_cast(kLanguages.size() + i); - } - } - return kCommonLangId; -} - -IndexTTS25TextEncoding IndexTTS25TextTokenizer::encode_for_inference( - const std::string & text, - int max_text_tokens_per_segment, - const std::string & lang) const { - if (max_text_tokens_per_segment <= 0) { - throw std::runtime_error("IndexTTS2.5 max_text_tokens_per_segment must be positive"); - } - - std::string resolved_lang = lowercase_ascii(lang); - if (resolved_lang.empty()) { - resolved_lang = contains_han(text) ? "zh" : "en"; - } - - std::string processed = text; - if (resolved_lang == "zh" || resolved_lang == "en") { - // Protect annotations from the normalizer, as the - // official TextNormalizer does inside normalize(). - auto protected_text = protect_pronunciation_annotations(processed); - protected_text.first = resolved_lang == "zh" - ? normalize_chinese(protected_text.first) - : normalize_english(protected_text.first); - processed = restore_pronunciation_annotations(std::move(protected_text.first), protected_text.second); - } - // ja/es/ar and other languages currently pass through without TN. - if (resolved_lang == "zh" || resolved_lang == "ja" || resolved_lang == "en") { - processed = lowercase_ascii(std::move(processed)); - } else if (resolved_lang == "es") { - processed = uppercase_ascii(std::move(processed)); - } - processed = apply_pronunciation_annotations(processed); - processed = uppercase_special_token_names(processed); - - const std::string lang_prefix = "<|" + resolved_lang + "|> "; - const auto prefix_tokens = static_cast(encode(lang_prefix).size()); - const int64_t capacity = assets_->config.gpt.max_text_tokens; - int64_t budget = std::min(max_text_tokens_per_segment, capacity - 2) - prefix_tokens; - budget = std::max(budget, 1); - - std::vector segments; - const auto token_len = [this](const std::string & value) { - return static_cast(encode(value).size()); - }; - if (token_len(processed) <= budget) { - segments.push_back(processed); - } else { - std::vector chunks; - for (const auto & [piece, atomic] : split_atomic_pieces(processed)) { - if (atomic) { - chunks.push_back(piece); - continue; - } - for (const auto & part : split_after_delimiters(piece)) { - if (token_len(part) <= budget) { - chunks.push_back(part); - continue; - } - std::string current; - for (size_t i = 0; i < part.size();) { - const size_t begin = i; - next_utf8_codepoint(part, i); - const std::string ch = part.substr(begin, i - begin); - if (!current.empty() && token_len(current + ch) > budget) { - chunks.push_back(std::move(current)); - current = ch; - } else { - current += ch; - } - } - if (!current.empty()) { - chunks.push_back(std::move(current)); - } - } - } - std::string current; - for (const auto & chunk : chunks) { - if (!current.empty() && token_len(current + chunk) > budget) { - segments.push_back(std::move(current)); - current = chunk; - } else { - current += chunk; - } - } - if (!current.empty()) { - segments.push_back(std::move(current)); - } - if (segments.empty()) { - segments.push_back(processed); - } - } - - IndexTTS25TextEncoding encoding; - encoding.lang = resolved_lang; - encoding.normalized_text = processed; - encoding.segments = segments; - encoding.segment_token_ids.reserve(segments.size()); - for (const auto & segment : segments) { - std::vector ids = encode(lang_prefix + segment); - ids.push_back(kSegmentPadTokenId); - encoding.segment_token_ids.push_back(std::move(ids)); - } - return encoding; -} - -} // namespace engine::models::index_tts2_5 diff --git a/src/models/index_tts2_5/vocoder.cpp b/src/models/index_tts2_5/vocoder.cpp deleted file mode 100644 index 52f9dbe0..00000000 --- a/src/models/index_tts2_5/vocoder.cpp +++ /dev/null @@ -1,69 +0,0 @@ -#include "engine/models/index_tts2_5/vocoder.h" - -#include "engine/framework/io/json.h" - -#include -#include - -namespace engine::models::index_tts2_5 { -namespace { - -constexpr int64_t kChunkedVocoderFrames = 768; -constexpr int64_t kChunkedVocoderOverlapFrames = 32; - -engine::modules::BigVganVocoderConfig parse_bigvgan_config( - const IndexTTS25Assets & assets, - engine::assets::TensorStorageType weight_storage_type) { - const auto root = assets.resources.parse_json("bigvgan_config"); - engine::modules::BigVganVocoderConfig config; - config.sampling_rate = engine::io::json::require_i64(root, "sampling_rate"); - config.num_mels = engine::io::json::require_i64(root, "num_mels"); - config.n_fft = engine::io::json::require_i64(root, "n_fft"); - config.hop_size = engine::io::json::require_i64(root, "hop_size"); - config.win_size = engine::io::json::require_i64(root, "win_size"); - config.upsample_initial_channel = engine::io::json::require_i64(root, "upsample_initial_channel"); - config.snake_logscale = engine::io::json::require_bool(root, "snake_logscale"); - config.upsample_rates = engine::io::json::require_i64_array(root, "upsample_rates"); - config.upsample_kernel_sizes = engine::io::json::require_i64_array(root, "upsample_kernel_sizes"); - config.resblock_kernel_sizes = engine::io::json::require_i64_array(root, "resblock_kernel_sizes"); - config.weight_storage_type = weight_storage_type; - if (config.sampling_rate != assets.config.s2mel.sample_rate || - config.num_mels != assets.config.s2mel.n_mels || - config.n_fft != assets.config.s2mel.n_fft || - config.hop_size != assets.config.s2mel.hop_length || - config.win_size != assets.config.s2mel.win_length) { - throw std::runtime_error("IndexTTS2.5 BigVGAN config does not match S2Mel mel config"); - } - return config; -} - -} // namespace - -IndexTTS25BigVganVocoder::IndexTTS25BigVganVocoder( - std::shared_ptr assets, - core::BackendConfig backend, - engine::assets::TensorStorageType weight_storage_type) - : assets_(std::move(assets)) { - if (assets_ == nullptr) { - throw std::runtime_error("IndexTTS2.5 BigVGAN vocoder requires assets"); - } - component_ = engine::modules::BigVganVocoderComponent::load_from_tensor_source( - assets_->bigvgan_weights, - std::move(backend), - parse_bigvgan_config(*assets_, weight_storage_type)); -} - -IndexTTS25VocoderOutput IndexTTS25BigVganVocoder::synthesize( - const std::vector & mel, - int64_t frames) const { - const auto out = frames > kChunkedVocoderFrames - ? component_.synthesize_chunked(mel, frames, kChunkedVocoderFrames, kChunkedVocoderOverlapFrames) - : component_.synthesize(mel, frames); - return {out.waveform, out.samples, static_cast(out.sample_rate)}; -} - -void IndexTTS25BigVganVocoder::release_runtime_graph() { - component_.release_runtime_graph(); -} - -} // namespace engine::models::index_tts2_5 diff --git a/tests/index_tts2_5/index_tts2_5_warm_bench_cases.json b/tests/index_tts2/index_tts2_5_warm_bench_cases.json similarity index 100% rename from tests/index_tts2_5/index_tts2_5_warm_bench_cases.json rename to tests/index_tts2/index_tts2_5_warm_bench_cases.json diff --git a/tests/index_tts2/index_tts2_warm_bench.cpp b/tests/index_tts2/index_tts2_warm_bench.cpp index 8c32172b..01c49e53 100644 --- a/tests/index_tts2/index_tts2_warm_bench.cpp +++ b/tests/index_tts2/index_tts2_warm_bench.cpp @@ -127,6 +127,7 @@ engine::runtime::TaskRequest make_request(const engine::io::json::Value & object set_optional_option(request, object, "emotion_vector"); set_optional_option(request, object, "use_emotion_text"); set_optional_option(request, object, "emotion_text"); + set_optional_option(request, object, "lang"); set_optional_option(request, object, "use_random_emotion"); set_optional_option(request, object, "interval_silence_ms"); set_optional_option(request, object, "text_chunk_size"); @@ -217,7 +218,10 @@ int main(int argc, char ** argv) { load_request.model_path = model_path; load_request.family_hint = "index_tts2"; auto registry = engine::runtime::make_default_registry(); + const auto load_start = Clock::now(); auto model = registry.load(load_request); + const auto load_end = Clock::now(); + const double load_ms = std::chrono::duration(load_end - load_start).count(); engine::runtime::TaskSpec task; task.task = engine::runtime::VoiceTaskKind::Tts; @@ -230,6 +234,7 @@ int main(int argc, char ** argv) { session_options.options[key] = value; } auto requests = parse_requests(request_sequence_json); + const auto session_start = Clock::now(); auto session_base = model->create_task_session(task, session_options); auto * session = dynamic_cast(session_base.get()); if (session == nullptr) { @@ -247,11 +252,17 @@ int main(int argc, char ** argv) { : std::nullopt; preparation.options = requests.front().options; session->prepare(preparation); + const auto session_end = Clock::now(); + const double session_ms = std::chrono::duration(session_end - session_start).count(); std::vector steps; std::vector timing_lines; timing_lines.push_back("index_tts2.backend " + backend_name); timing_lines.push_back("index_tts2.model_root " + model_path.string()); + timing_lines.push_back("index_tts2.load_ms " + engine::io::json::stringify_number(load_ms)); + timing_lines.push_back("index_tts2.session_prepare_ms " + engine::io::json::stringify_number(session_ms)); + std::cout << "index_tts2.load_ms=" << load_ms << "\n"; + std::cout << "index_tts2.session_prepare_ms=" << session_ms << "\n"; for (int i = 0; i < warmup; ++i) { (void) session->run(requests.front()); } diff --git a/tests/index_tts2_5/index_tts2_5_warm_bench.cpp b/tests/index_tts2_5/index_tts2_5_warm_bench.cpp deleted file mode 100644 index f5868862..00000000 --- a/tests/index_tts2_5/index_tts2_5_warm_bench.cpp +++ /dev/null @@ -1,319 +0,0 @@ -#include "engine/framework/audio/wav_reader.h" -#include "engine/framework/audio/wav_writer.h" -#include "engine/framework/debug/profiler.h" -#include "engine/framework/io/json.h" -#include "engine/framework/runtime/registry.h" -#include "engine/framework/runtime/session.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -using Clock = std::chrono::steady_clock; - -std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { - for (int i = 1; i + 1 < argc; ++i) { - if (argv[i] == name) { - return argv[i + 1]; - } - } - return fallback; -} - -int int_arg(int argc, char ** argv, const std::string & name, int fallback) { - return std::stoi(arg_value(argc, argv, name, std::to_string(fallback))); -} - -engine::core::BackendType parse_backend(const std::string & value) { - if (value == "cuda") { - return engine::core::BackendType::Cuda; - } - if (value == "cpu") { - return engine::core::BackendType::Cpu; - } - throw std::runtime_error("IndexTTS2.5 warmbench backend must be cuda or cpu"); -} - -std::vector> parse_session_options(int argc, char ** argv) { - std::vector> out; - for (int i = 1; i + 1 < argc; ++i) { - if (std::string(argv[i]) != "--session-option") { - continue; - } - const std::string option = argv[i + 1]; - const size_t eq = option.find('='); - if (eq == std::string::npos || eq == 0) { - throw std::runtime_error("invalid IndexTTS2.5 --session-option: " + option); - } - out.emplace_back(option.substr(0, eq), option.substr(eq + 1)); - } - return out; -} - -std::string option_text(const engine::io::json::Value & value) { - if (value.is_bool()) { - return value.as_bool() ? "true" : "false"; - } - if (value.is_number()) { - return engine::io::json::stringify_number(value.as_number()); - } - if (value.is_array()) { - std::string out; - const auto & items = value.as_array(); - for (size_t i = 0; i < items.size(); ++i) { - if (i != 0) { - out += ","; - } - if (!items[i].is_number()) { - throw std::runtime_error("IndexTTS2.5 warmbench option arrays must contain only numbers"); - } - out += engine::io::json::stringify_number(items[i].as_number()); - } - return out; - } - return value.as_string(); -} - -std::string required_string(const engine::io::json::Value & object, const std::string & key) { - const auto * value = object.find(key); - if (value == nullptr || value->is_null()) { - throw std::runtime_error("IndexTTS2.5 warmbench request missing " + key); - } - return value->as_string(); -} - -std::string optional_string(const engine::io::json::Value & object, const std::string & key) { - const auto * value = object.find(key); - return value == nullptr || value->is_null() ? std::string{} : value->as_string(); -} - -void set_optional_option(engine::runtime::TaskRequest & request, const engine::io::json::Value & object, const std::string & key) { - const auto * value = object.find(key); - if (value != nullptr && !value->is_null()) { - request.options[key] = option_text(*value); - } -} - -engine::runtime::AudioBuffer read_audio_buffer(const std::filesystem::path & path) { - const auto wav = engine::audio::read_wav_f32(path); - return engine::runtime::AudioBuffer{wav.sample_rate, wav.channels, wav.samples}; -} - -engine::runtime::TaskRequest make_request(const engine::io::json::Value & object) { - engine::runtime::TaskRequest request; - const std::string language = optional_string(object, "language"); - request.text_input = engine::runtime::Transcript{required_string(object, "text"), language.empty() ? "en" : language}; - engine::runtime::VoiceCondition voice; - engine::runtime::VoiceReference speaker; - speaker.audio = read_audio_buffer(required_string(object, "voice_ref")); - voice.speaker = std::move(speaker); - request.voice = std::move(voice); - - const auto emotion_audio = optional_string(object, "audio"); - if (!emotion_audio.empty()) { - request.audio_input = read_audio_buffer(emotion_audio); - } - - set_optional_option(request, object, "lang"); - set_optional_option(request, object, "emotion_alpha"); - set_optional_option(request, object, "emotion_vector"); - set_optional_option(request, object, "use_emotion_text"); - set_optional_option(request, object, "emotion_text"); - set_optional_option(request, object, "use_random_emotion"); - set_optional_option(request, object, "interval_silence_ms"); - set_optional_option(request, object, "text_chunk_size"); - set_optional_option(request, object, "do_sample"); - set_optional_option(request, object, "top_p"); - set_optional_option(request, object, "top_k"); - set_optional_option(request, object, "temperature"); - set_optional_option(request, object, "length_penalty"); - set_optional_option(request, object, "num_beams"); - set_optional_option(request, object, "repetition_penalty"); - set_optional_option(request, object, "max_tokens"); - set_optional_option(request, object, "seed"); - return request; -} - -std::vector parse_requests(const std::string & request_sequence_json) { - if (request_sequence_json.empty()) { - throw std::runtime_error("IndexTTS2.5 warmbench requires --request-sequence-json"); - } - const auto root = engine::io::json::parse(request_sequence_json); - std::vector requests; - for (const auto & item : root.as_array()) { - requests.push_back(make_request(item)); - } - if (requests.empty()) { - throw std::runtime_error("IndexTTS2.5 warmbench request sequence is empty"); - } - return requests; -} - -engine::io::json::Value number(double value) { - return engine::io::json::Value::make_number(value); -} - -engine::io::json::Value string(std::string value) { - return engine::io::json::Value::make_string(std::move(value)); -} - -engine::io::json::Value audio_summary_json(const engine::runtime::AudioBuffer & audio) { - if (audio.samples.empty()) { - throw std::runtime_error("IndexTTS2.5 warmbench received empty audio output"); - } - double sum = 0.0; - double abs_sum = 0.0; - double sq_sum = 0.0; - float min_value = audio.samples.front(); - float max_value = audio.samples.front(); - for (const float sample : audio.samples) { - sum += static_cast(sample); - abs_sum += std::abs(static_cast(sample)); - sq_sum += static_cast(sample) * static_cast(sample); - min_value = std::min(min_value, sample); - max_value = std::max(max_value, sample); - } - const int channels = std::max(1, audio.channels); - const double frames = static_cast(audio.samples.size() / static_cast(channels)); - const double count = static_cast(audio.samples.size()); - return engine::io::json::Value::make_object({ - {"sample_rate", number(static_cast(audio.sample_rate))}, - {"channels", number(static_cast(audio.channels))}, - {"samples", number(count)}, - {"frames", number(frames)}, - {"duration_sec", number(audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0)}, - {"sum", number(sum)}, - {"mean_abs", number(abs_sum / count)}, - {"rms", number(std::sqrt(sq_sum / count))}, - {"min", number(min_value)}, - {"max", number(max_value)}, - }); -} - -} // namespace - -int main(int argc, char ** argv) { - try { - const std::filesystem::path model_path = arg_value(argc, argv, "--model", "models/IndexTTS-2.5"); - const std::string backend_name = arg_value(argc, argv, "--backend", "cuda"); - const int device = int_arg(argc, argv, "--device", 0); - const int threads = int_arg(argc, argv, "--threads", 8); - const int warmup = int_arg(argc, argv, "--warmup", 0); - const int iterations = int_arg(argc, argv, "--iterations", 1); - const std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", ""); - const std::filesystem::path output_dir = arg_value(argc, argv, "--output-dir", ""); - const std::filesystem::path timing_path = arg_value(argc, argv, "--timing-file", "/tmp/index_tts2_5_warm_bench_timing.log"); - engine::debug::configure_logging(engine::debug::LoggingConfig{true, timing_path.string()}); - - engine::runtime::ModelLoadRequest load_request; - load_request.model_path = model_path; - load_request.family_hint = "index_tts2_5"; - auto registry = engine::runtime::make_default_registry(); - auto model = registry.load(load_request); - - engine::runtime::TaskSpec task; - task.task = engine::runtime::VoiceTaskKind::Tts; - task.mode = engine::runtime::RunMode::Offline; - engine::runtime::SessionOptions session_options; - session_options.backend.type = parse_backend(backend_name); - session_options.backend.device = device; - session_options.backend.threads = threads; - for (const auto & [key, value] : parse_session_options(argc, argv)) { - session_options.options[key] = value; - } - auto requests = parse_requests(request_sequence_json); - auto session_base = model->create_task_session(task, session_options); - auto * session = dynamic_cast(session_base.get()); - if (session == nullptr) { - throw std::runtime_error("IndexTTS2.5 model did not create an offline voice task session"); - } - engine::runtime::SessionPreparationRequest preparation; - preparation.text = requests.front().text_input; - preparation.voice = requests.front().voice; - preparation.audio = requests.front().audio_input.has_value() - ? std::optional({ - requests.front().audio_input->sample_rate, - requests.front().audio_input->channels, - static_cast(requests.front().audio_input->samples.size()), - }) - : std::nullopt; - preparation.options = requests.front().options; - session->prepare(preparation); - - std::vector steps; - std::vector timing_lines; - timing_lines.push_back("index_tts2_5.backend " + backend_name); - timing_lines.push_back("index_tts2_5.model_root " + model_path.string()); - for (int i = 0; i < warmup; ++i) { - (void) session->run(requests.front()); - } - - for (size_t request_index = 0; request_index < requests.size(); ++request_index) { - double total_ms = 0.0; - engine::runtime::TaskResult last_result; - for (int iteration = 0; iteration < std::max(1, iterations); ++iteration) { - const auto start = Clock::now(); - last_result = session->run(requests[request_index]); - const auto end = Clock::now(); - const double wall_ms = std::chrono::duration(end - start).count(); - total_ms += wall_ms; - timing_lines.push_back("index_tts2_5.wall_ms " + engine::io::json::stringify_number(wall_ms)); - } - if (!last_result.audio_output.has_value()) { - throw std::runtime_error("IndexTTS2.5 warmbench expected audio output"); - } - const double avg_ms = total_ms / static_cast(std::max(1, iterations)); - std::filesystem::path audio_path; - if (!output_dir.empty()) { - std::filesystem::create_directories(output_dir); - audio_path = output_dir / ("request_" + std::to_string(request_index) + ".wav"); - engine::audio::write_pcm16_wav( - audio_path, - last_result.audio_output->sample_rate, - last_result.audio_output->channels, - last_result.audio_output->samples); - } - engine::io::json::Value::Object step{ - {"request_index", number(static_cast(request_index))}, - {"stems", engine::io::json::Value::make_array({ - engine::io::json::Value::make_object({ - {"name", string("audio")}, - {"summary", audio_summary_json(*last_result.audio_output)}, - {"audio", string(audio_path.string())}, - }), - })}, - {"metrics", engine::io::json::Value::make_object({{"wall_ms", number(avg_ms)}})}, - }; - steps.push_back(engine::io::json::Value::make_object(std::move(step))); - std::cout << "index_tts2_5.wall_ms=" << avg_ms << "\n"; - } - - if (!timing_path.empty()) { - std::filesystem::create_directories(timing_path.parent_path()); - std::ofstream timing(timing_path, std::ios::app); - for (const auto & line : timing_lines) { - timing << line << "\n"; - } - } - - const auto summary = engine::io::json::Value::make_object({ - {"family", string("index_tts2_5")}, - {"backend", string(backend_name)}, - {"sequence_steps", engine::io::json::Value::make_array(std::move(steps))}, - }); - std::cout << "summary_json=" << engine::io::json::stringify(summary) << "\n"; - return 0; - } catch (const std::exception & ex) { - std::cerr << "index_tts2_5_warm_bench failed: " << ex.what() << "\n"; - return 1; - } -} diff --git a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json index be28add6..4ad10c62 100644 --- a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json @@ -291,7 +291,7 @@ { "id": "index_tts2_5_voice_clone_longform", "coverage": "IndexTTS2.5 voice clone with shared long-form text for chunking and RTF measurement", - "family": "index_tts2_5", + "family": "index_tts2", "model": "models/IndexTTS2.5-GGUF", "task": "tts", "mode": "offline", @@ -406,7 +406,7 @@ { "id": "index_tts2_5_longform_voice_clone_6000_emotion_text", "coverage": "IndexTTS2.5 longform voice clone with 6000-character text plus long emotion-text conditioning", - "family": "index_tts2_5", + "family": "index_tts2", "model": "models/IndexTTS2.5-GGUF", "task": "tts", "mode": "offline", diff --git a/tools/convert_index_tts2_5.py b/tools/convert_index_tts2_5.py index 10957c34..59cbe6e9 100644 --- a/tools/convert_index_tts2_5.py +++ b/tools/convert_index_tts2_5.py @@ -42,7 +42,7 @@ import torch from safetensors.torch import save_file -# GGUF tensor namespaces (must match the index_tts2_5 model spec) and the +# GGUF tensor namespaces (must match the index_tts2 model spec) and the # staging file each one is produced from. TENSOR_OUTPUTS = [ ("gpt", "gpt.safetensors"), @@ -157,13 +157,35 @@ def _copy(src: Path, dst: Path, label: str) -> None: print(f"copied {src} -> {dst}") +def _stage_config_v2_5(src: Path, dst: Path) -> None: + """Stage config.yaml with the version field normalized to "2.5". + + The official IndexTTS-2.5 snapshot ships config.yaml with `version: 2.0` + (inherited from IndexTTS-2). audio.cpp selects the IndexTTS2 family variant + from this field, so the staged copy must declare 2.5 explicitly. + """ + _require_file(src, "config.yaml") + dst.parent.mkdir(parents=True, exist_ok=True) + lines = src.read_text(encoding="utf-8").splitlines() + replaced = False + for i, line in enumerate(lines): + if line.strip().startswith("version:"): + lines[i] = 'version: "2.5"' + replaced = True + break + if not replaced: + lines.append('version: "2.5"') + dst.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"staged {src} -> {dst} (version normalized to \"2.5\")") + + def build_converter_command(output_dir: Path, converter: str, quant_type: str) -> list[str]: command = [converter] for namespace, filename in TENSOR_OUTPUTS: command += ["--input", f"{namespace}={output_dir / filename}"] command += [ "--root", str(output_dir / "root"), - "--family", "index_tts2_5", + "--family", "index_tts2", "--type", quant_type, "--output", str(output_dir / f"index-tts2_5-{quant_type}.gguf"), ] @@ -272,7 +294,7 @@ def main() -> int: "qwen emotion weights") # Sidecar files embedded into the GGUF via --root. - _copy(model_dir / "config.yaml", root_dir / "config.yaml", "config.yaml") + _stage_config_v2_5(model_dir / "config.yaml", root_dir / "config.yaml") _copy(model_dir / "multilingual_zh_ja_yue_char_del.tiktoken", root_dir / "multilingual_zh_ja_yue_char_del.tiktoken", "tiktoken vocabulary") _copy(w2v_bert_dir / "config.json", root_dir / "w2v-bert-2.0" / "config.json", "w2v-bert-2.0 config") diff --git a/webui/configs/model_params.json b/webui/configs/model_params.json index 7335811b..e414f7e7 100644 --- a/webui/configs/model_params.json +++ b/webui/configs/model_params.json @@ -164,15 +164,7 @@ ], "index_tts2": [ - {"name": "emotion_text", "type": "text", "label": "emotion_text（情绪参考文本）", "label_en": "emotion_text (emotion reference text)", "default": "", "placeholder": "例：你吓死我了！你是鬼吗？", "placeholder_en": "e.g. You scared me to death!", "info": "填写后自动开启情感条件（use_emotion_text）", "info_en": "Setting this enables emotion conditioning."}, - {"name": "emotion_alpha", "type": "slider", "label": "emotion_alpha（情感强度）", "label_en": "emotion_alpha", "default": 1.0, "minimum": 0.0, "maximum": 1.0, "step": 0.05}, - {"name": "use_emotion_text", "type": "bool", "label": "use_emotion_text（从朗读文本推断情感）", "label_en": "use_emotion_text (infer from text)", "default": false}, - {"name": "use_random_emotion", "type": "bool", "label": "use_random_emotion（随机情感）", "label_en": "use_random_emotion", "default": false}, - {"name": "interval_silence_ms", "type": "number", "label": "interval_silence_ms（分段间静音）", "label_en": "interval_silence_ms", "default": 200, "minimum": 0, "step": 50, "precision": 0} - ], - - "index_tts2_5": [ - {"name": "lang", "type": "choice", "label": "lang（语种提示）", "label_en": "lang (language hint)", "default": "auto", "choices": ["auto", "zh", "en", "ja", "es", "ar"], "info": "auto：含汉字按中文，否则按英文；日/西/阿建议显式选择", "info_en": "auto: zh when the text contains Han characters, otherwise en; set ja/es/ar explicitly"}, + {"name": "lang", "type": "choice", "label": "lang（语种提示, 仅 IndexTTS2.5 模型）", "label_en": "lang (language hint, IndexTTS2.5 models only)", "default": "auto", "choices": ["auto", "zh", "en", "ja", "es", "ar"], "info": "仅对 IndexTTS2.5（多语种）模型生效：auto 含汉字按中文，否则按英文；日/西/阿建议显式选择", "info_en": "Only applies to IndexTTS2.5 (multilingual) models: auto picks zh when the text contains Han characters, otherwise en; set ja/es/ar explicitly"}, {"name": "emotion_text", "type": "text", "label": "emotion_text（情绪参考文本）", "label_en": "emotion_text (emotion reference text)", "default": "", "placeholder": "例：你吓死我了！你是鬼吗？", "placeholder_en": "e.g. You scared me to death!", "info": "填写后自动开启情感条件（use_emotion_text）", "info_en": "Setting this enables emotion conditioning."}, {"name": "emotion_alpha", "type": "slider", "label": "emotion_alpha（情感强度）", "label_en": "emotion_alpha", "default": 1.0, "minimum": 0.0, "maximum": 1.0, "step": 0.05}, {"name": "use_emotion_text", "type": "bool", "label": "use_emotion_text（从朗读文本推断情感）", "label_en": "use_emotion_text (infer from text)", "default": false}, diff --git a/webui/configs/models_catalog.json b/webui/configs/models_catalog.json index c347973f..6e19b9bd 100644 --- a/webui/configs/models_catalog.json +++ b/webui/configs/models_catalog.json @@ -15,7 +15,7 @@ { "id": "voxcpm2", "display_name": "VoxCPM2 (tts)", "family": "voxcpm2", "path": "models/VoxCPM2", "task": "tts", "mode": "offline", "download_id": "voxcpm2", "session_options": { "voxcpm2.weight_type": "q8_0" }, "min_vram_gb": 6 }, { "id": "vibevoice", "display_name": "VibeVoice 1.5B (tts, long-form/multi-speaker)", "family": "vibevoice", "path": "models/VibeVoice-1.5B", "task": "tts", "mode": "offline", "download_id": "vibevoice_1_5b", "min_vram_gb": 7 }, { "id": "index-tts2", "display_name": "IndexTTS2 (tts 中英克隆+情感)", "display_name_en": "IndexTTS2 (tts, zh/en clone + emotion)", "family": "index_tts2", "path": "models/IndexTTS-2", "task": "tts", "mode": "offline", "download_id": "index_tts2", "min_vram_gb": 8 }, - { "id": "index-tts2.5", "display_name": "IndexTTS2.5 (tts 多语种克隆+情感, GGUF Q8)", "display_name_en": "IndexTTS2.5 (tts, zh/en/ja/es/ar clone + emotion, GGUF Q8)", "family": "index_tts2_5", "path": "models/IndexTTS2.5-GGUF", "task": "tts", "mode": "offline", "download_id": "index_tts2_5_q8_0", "min_vram_gb": 8, + { "id": "index-tts2.5", "display_name": "IndexTTS2.5 (tts 多语种克隆+情感, GGUF Q8)", "display_name_en": "IndexTTS2.5 (tts, zh/en/ja/es/ar clone + emotion, GGUF Q8)", "family": "index_tts2", "path": "models/IndexTTS2.5-GGUF", "task": "tts", "mode": "offline", "download_id": "index_tts2_5_q8_0", "min_vram_gb": 8, "input_hint": "**IndexTTS2.5**：中/英/日/西/阿零样本克隆；上传参考音色即克隆；可在『其它参数(JSON)』里传 `lang`（默认 auto：含汉字按中文，否则按英文）与情感选项。许可证为 bilibili Model Use License（非 OSI），商用前请确认条款。", "input_hint_en": "**IndexTTS2.5**: zero-shot cloning in zh/en/ja/es/ar. Upload a reference voice to clone; pass `lang` (default auto: zh when the text contains Han characters, otherwise en) and emotion options through the JSON box. Weights are under the bilibili Model Use License (not OSI-approved) — check terms before commercial use." }, { "id": "irodori-tts", "display_name": "Irodori-TTS v4 Small (tts 日语, GGUF Q8)", "display_name_en": "Irodori-TTS v4 Small (ja tts, GGUF Q8)", "family": "irodori_tts", "path": "models/Irodori-TTS-v4-Small-GGUF", "task": "tts", "mode": "offline", "download_id": "irodori_tts_v4_small_q8_0", "min_vram_gb": 4, From c23d1c32efd5a8d2ac7250fb716dd9841b16fa4b Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Wed, 12 Aug 2026 14:15:16 +0800 Subject: [PATCH 6/8] sampling: add ENGINE_TORCH_SAMPLING_POLICY to pin the TensorIterator RNG layout The TensorIterator-style Philox layout is normally probed from the CUDA device (multiprocessor count / max threads per SM), so the same seed yields different noise realizations on different GPUs, and HIP/CPU always fall back to the default 1x256 layout. The new environment variable pins the policy explicitly ("default" or "x") with the CUDA fast path disabled, making seeded generation bit-comparable across backends and machines (verified: index_tts2 v2.5 fp32, greedy and sampled, CUDA vs HIP waveform corr > 0.999). --- src/framework/sampling/torch_random.cpp | 47 +++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/src/framework/sampling/torch_random.cpp b/src/framework/sampling/torch_random.cpp index 70b1bb7d..82367d9a 100644 --- a/src/framework/sampling/torch_random.cpp +++ b/src/framework/sampling/torch_random.cpp @@ -8,7 +8,9 @@ #include #include +#include #include +#include #include #include @@ -299,6 +301,48 @@ void log_default_policy(std::string_view category, std::string_view reason) { + "(multiprocessor_count=1, max_threads_per_multiprocessor=256): " + std::string(reason)); } +// ENGINE_TORCH_SAMPLING_POLICY pins the TensorIterator RNG layout instead of +// probing the CUDA device, making the noise realization identical across +// backends (CUDA/HIP/CPU) and machines. Accepted values: "default" (1x256) +// or "x" (e.g. +// "68x1024"). Unset keeps the legacy behavior (device probe on CUDA, default +// layout elsewhere). The pinned layout never uses the CUDA fast path so every +// backend computes the same Philox element mapping on the host. +std::optional pinned_policy_from_env(std::string_view log_category) { + const char * value = std::getenv("ENGINE_TORCH_SAMPLING_POLICY"); + if (value == nullptr || *value == '\0') { + return std::nullopt; + } + TorchCudaSamplingPolicy policy; + std::string text(value); + if (text != "default") { + const auto cross = text.find('x'); + if (cross == std::string::npos) { + throw std::runtime_error( + "ENGINE_TORCH_SAMPLING_POLICY must be \"default\" or \"x\", got: " + text); + } + try { + policy.multiprocessor_count = std::stoll(text.substr(0, cross)); + policy.max_threads_per_multiprocessor = std::stoll(text.substr(cross + 1)); + } catch (const std::exception &) { + throw std::runtime_error( + "ENGINE_TORCH_SAMPLING_POLICY must be \"default\" or \"x\", got: " + text); + } + if (policy.multiprocessor_count <= 0 || policy.max_threads_per_multiprocessor <= 0) { + throw std::runtime_error("ENGINE_TORCH_SAMPLING_POLICY values must be positive: " + text); + } + } + policy.cuda_fast_path = false; + engine::debug::log_message( + engine::debug::LogLevel::Warning, + log_category, + "using pinned Torch RNG layout policy from ENGINE_TORCH_SAMPLING_POLICY " + "(multiprocessor_count=" + std::to_string(policy.multiprocessor_count) + + ", max_threads_per_multiprocessor=" + std::to_string(policy.max_threads_per_multiprocessor) + + ")"); + return policy; +} + } // namespace TorchCudaSamplingPolicy resolve_torch_cuda_sampling_policy( @@ -307,6 +351,9 @@ TorchCudaSamplingPolicy resolve_torch_cuda_sampling_policy( std::string_view log_category, std::string_view model_name, TorchCudaSamplingPolicyFailureMode failure_mode) { + if (const auto pinned = pinned_policy_from_env(log_category)) { + return *pinned; + } TorchCudaSamplingPolicy policy; if (backend_type != engine::core::BackendType::Cuda) { log_default_policy(log_category, "backend is not CUDA"); From 4f7db3df37cedba9af830eaf5956f94d89fd24e9 Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Wed, 12 Aug 2026 14:42:33 +0800 Subject: [PATCH 7/8] text: verbalize digits attached to letters in English normalization The official wetext English normalizer expands letter-digit compounds ("DS4" -> "DS four", "R2D2" -> "R two D two", "4K" -> "four K", "H264" -> "H two hundred sixty four"), but the standalone cardinal rule \b(\d+)\b skipped digits glued to letters, so the model received raw "ds4" tokens and read them as "D-S-A". Split letter<->digit boundaries after the date/decimal/ordinal passes and before the cardinal pass so each digit group is spelled out the same way as the official pipeline. --- src/framework/text/text_normalization.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/framework/text/text_normalization.cpp b/src/framework/text/text_normalization.cpp index 8c5c494e..9f7c3cfa 100644 --- a/src/framework/text/text_normalization.cpp +++ b/src/framework/text/text_normalization.cpp @@ -462,6 +462,13 @@ std::string normalize_english_numbers(std::string text) { [](const std::smatch & match) { return english_ordinal_from_digits(match[1].str()); }); + // Split letter<->digit boundaries so digits attached to letters verbalize + // like the official wetext English normalizer ("DS4" -> "DS four", + // "R2D2" -> "R two D two", "4K" -> "four K"). Runs after dates/decimals/ + // ordinals so "1st", "2.5" and friends have already been expanded; the + // standalone cardinal rule below then spells out each digit group. + text = std::regex_replace(std::move(text), std::regex(R"(([A-Za-z])(\d))"), "$1 $2"); + text = std::regex_replace(std::move(text), std::regex(R"((\d)([A-Za-z]))"), "$1 $2"); text = normalize_english_regex( std::move(text), std::regex(R"(\b(\d+)\b)"), From 17222073b40d786f697aaf62cd83aa667839617e Mon Sep 17 00:00:00 2001 From: IIIIIllllIIIIIlllll <2432896620@qq.com> Date: Wed, 12 Aug 2026 18:04:18 +0800 Subject: [PATCH 8/8] text: verbalize ASCII symbols in English normalization (index_tts2) The official wetext English normalizer verbalizes standalone ASCII symbols ("invalid_request_error" -> "invalid underscore request underscore error", "C++" -> "C plus plus", "a=b" -> "a equal sign b", "~" -> "tilde"). Keeping them raw fed untrained token sequences to the model and mangled the following word. Add a verbalize_symbols option (off by default to leave other models untouched) and enable it for index_tts2. --- .../framework/text/text_normalization.h | 3 +++ src/framework/text/text_normalization.cpp | 21 +++++++++++++++++++ src/models/index_tts2/tokenizer_text.cpp | 1 + 3 files changed, 25 insertions(+) diff --git a/include/engine/framework/text/text_normalization.h b/include/engine/framework/text/text_normalization.h index ec026cde..2b10c051 100644 --- a/include/engine/framework/text/text_normalization.h +++ b/include/engine/framework/text/text_normalization.h @@ -10,6 +10,9 @@ struct EnglishTextNormalizationOptions { bool spell_numbers = true; bool index_tts_punctuation = false; bool uppercase_ascii = false; + // Verbalize standalone ASCII symbols like the official wetext English + // normalizer ("a_b" -> "a underscore b", "C++" -> "C plus plus"). + bool verbalize_symbols = false; }; std::string replace_all(std::string text, std::string_view from, std::string_view to); diff --git a/src/framework/text/text_normalization.cpp b/src/framework/text/text_normalization.cpp index 9f7c3cfa..be6b7512 100644 --- a/src/framework/text/text_normalization.cpp +++ b/src/framework/text/text_normalization.cpp @@ -489,6 +489,27 @@ std::string normalize_english_text(std::string_view text, const EnglishTextNorma if (options.spell_numbers) { out = normalize_english_numbers(std::move(out)); } + if (options.verbalize_symbols) { + // Match the official wetext English normalizer: standalone ASCII + // symbols are verbalized ("a_b" -> "a underscore b", + // "C++" -> "C plus plus", "a=b" -> "a equal sign b"). Runs after + // number spelling so "50%" has already become "fifty percent". + const std::pair symbol_words[] = { + {"_", " underscore "}, + {"+", " plus "}, + {"=", " equal sign "}, + {"*", " asterisk "}, + {"&", " and "}, + {"#", " hash "}, + {"%", " percent "}, + {"|", " vertical bar "}, + {"~", " tilde "}, + }; + for (const auto & [from, to] : symbol_words) { + out = replace_all(std::move(out), from, to); + } + out = collapse_ascii_whitespace(out); + } if (options.index_tts_punctuation) { out = apply_index_tts_punctuation_map(std::move(out)); } diff --git a/src/models/index_tts2/tokenizer_text.cpp b/src/models/index_tts2/tokenizer_text.cpp index fd2ec69b..8b87547c 100644 --- a/src/models/index_tts2/tokenizer_text.cpp +++ b/src/models/index_tts2/tokenizer_text.cpp @@ -713,6 +713,7 @@ std::string IndexTTS2TextTokenizer::normalize_english(const std::string & text) options.expand_common_contractions = true; options.index_tts_punctuation = true; options.uppercase_ascii = variant_ == IndexTTS2Variant::kV2; + options.verbalize_symbols = true; return engine::text::normalize_english_text(text, options); }