diff --git a/CMakeLists.txt b/CMakeLists.txt index 39d1dfcb..881f75a3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -348,6 +348,14 @@ add_library(engine_runtime STATIC src/models/voxtral_realtime/text_decoder.cpp src/models/voxtral_realtime/session.cpp src/models/voxtral_realtime/loader.cpp + src/models/fish_audio/ar.cpp + src/models/fish_audio/assets.cpp + src/models/fish_audio/codec.cpp + src/models/fish_audio/generator.cpp + src/models/fish_audio/loader.cpp + src/models/fish_audio/prompt_builder.cpp + src/models/fish_audio/session.cpp + src/models/fish_audio/tokenizer_text.cpp src/models/heartmula/assets.cpp src/models/heartmula/codec.cpp src/models/heartmula/generator.cpp @@ -371,6 +379,15 @@ add_library(engine_runtime STATIC src/models/higgs_audio_stt/postprocess.cpp src/models/higgs_audio_stt/session.cpp src/models/higgs_audio_stt/loader.cpp + src/models/higgs_audio_tts/ar.cpp + src/models/higgs_audio_tts/assets.cpp + src/models/higgs_audio_tts/codec.cpp + src/models/higgs_audio_tts/codebooks.cpp + src/models/higgs_audio_tts/generator.cpp + src/models/higgs_audio_tts/loader.cpp + src/models/higgs_audio_tts/sampler.cpp + src/models/higgs_audio_tts/session.cpp + src/models/higgs_audio_tts/tokenizer_text.cpp src/models/irodori_tts/assets.cpp src/models/irodori_tts/codec.cpp src/models/irodori_tts/condition_encoder.cpp @@ -723,6 +740,7 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(chatterbox_warm_bench tests/chatterbox/chatterbox_warm_bench.cpp) add_engine_warmbench(citrinet_asr_warm_bench tests/citrinet_asr/citrinet_asr_warm_bench.cpp) add_engine_warmbench(higgs_audio_stt_warm_bench tests/higgs_audio_stt/higgs_audio_stt_warm_bench.cpp) + 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(irodori_tts_warm_bench tests/irodori_tts/irodori_tts_warm_bench.cpp) @@ -881,6 +899,16 @@ if (ENGINE_BUILD_TESTS) COMMAND encoder_module_test ) + add_engine_unittest( + qwen_decoder_packed_projection_test + tests/unittests/test_qwen_decoder_packed_projections.cpp + ) + + add_test( + NAME qwen_decoder_packed_projection_test + COMMAND qwen_decoder_packed_projection_test + ) + add_engine_unittest(conv_transpose_fast_path_test tests/unittests/test_conv_transpose_fast_path.cpp) add_test( diff --git a/README.md b/README.md index 3d2c6b72..b980d19d 100644 --- a/README.md +++ b/README.md @@ -51,8 +51,10 @@ audio.cpp would not be moving this quickly without generous contributors bringin | **ace_step** | music generation, music editing | 50+ langs | ACE-Step 1.5 Turbo and Base with acestep-5Hz-lm-1.7B | | **chatterbox** | TTS, voice cloning, voice conversion | ar, da, de, el, en, es, fi, fr, hi, it, ko, ms, nl, no, pl, pt, sv, sw, tr | Chatterbox with 0.5B backbone | | **citrinet_asr** | ASR | en | Citrinet-256 | +| **fish_audio** | TTS, voice cloning | auto, en, zh | Fish Audio S2 Pro | | **heartmula** | music generation | zh, en, ja, ko, es | HeartMuLa-oss-3B with HeartCodec-oss | | **higgs_audio_stt** | ASR | en | Higgs Audio v3 STT | +| **higgs_audio_tts** | TTS, voice cloning | auto | Higgs Audio v3 TTS 4B | | **htdemucs** | source separation | lang agnostic | HTDemucs, HTDemucs_ft | | **hviske_asr** | ASR | da | Hviske v5.3 | | **marblenet_vad** | VAD | lang agnostic | MarbleNet VAD | @@ -90,8 +92,6 @@ Community model ports live under `community_models` to make the ownership bounda | **outetts** | TTS, voice cloning | en, ar, zh, nl, fr, de, it, ja, ko, lt, ru, es, pt, be, bn, ka, hu, lv, fa, pl, sw, ta, uk | Mirek [@mirek190](https://github.com/mirek190) | Llama-OuteTTS-1.0-1B TTS and voice cloning support | | **vietneu_tts** | TTS, voice cloning | vi, en | Phuoc [@phuocnguyen90](https://github.com/phuocnguyen90) | [VieNeu-TTS-v3-Turbo](vietneu_tts.md) TTS and voice cloning support | -WIP (loaders not registered in this release tree — catalog entries are `UnsupportedSource`): Kokoro 82M bf16, Higgs Audio v3 TTS 4B, Parakeet TDT 0.6B v3, Fish Audio S2 Pro. See [docs/maintainers/loader_and_catalog.md](docs/maintainers/loader_and_catalog.md). - PocketTTS language selection is a model-load option. When the model path points at the PocketTTS root, the loader uses `english` unless you pass `--load-option language=`. Kyutai's normal non-English PocketTTS releases are smaller distilled language models intended for the fast PocketTTS path. The `_24l` variants are larger 24-layer, undistilled preview models that can sound better but are slower. Kyutai currently publishes French only as `french_24l`, not as a normal distilled `french` language directory, so French is not listed as a normal PocketTTS language here. ## Docker @@ -373,20 +373,22 @@ Recommended top-level install packages: `Yes` means Hugging Face has a ready-to-use repo that the framework can download as-is. `No` means the tool must assemble, convert, or post-process files before the framework can use them. Packages whose loaders are not registered in this release tree are listed as **Unavailable** (see [docs/maintainers/loader_and_catalog.md](docs/maintainers/loader_and_catalog.md)). +For shared audio.cpp GGUF packages, the model manager installs the default Q8_0 GGUF. Other precision variants can be downloaded directly from [audio-cpp/audio.cpp-gguf](https://huggingface.co/audio-cpp/audio.cpp-gguf); see [docs/gguf.md](docs/gguf.md) for GGUF support status. + | Package id | Model | HF ready-to-use repo | |---|---|---| | `ace_step` | ACE-Step 1.5 Turbo/Base | No | | `chatterbox` | Chatterbox | **Yes** | | `citrinet_asr` | Citrinet ASR converted layout | No | +| `fish_audio_s2_pro` | Fish Audio S2 Pro GGUF Q8_0 | **Yes** | | `heartmula` | HeartMuLa | No | | `higgs_audio_stt` | Higgs Audio STT | No | -| `higgs_audio_v3_tts_4b` | Higgs Audio v3 TTS 4B | Unavailable (loader not in this tree) | +| `higgs_audio_v3_tts_4b` | Higgs Audio v3 TTS 4B GGUF Q8_0 | **Yes** | | `htdemucs` | HTDemucs | No | | `hviske_asr` | Hviske ASR | **Yes** | | `irodori_tts_500m_v3` | Irodori-TTS 500M v3 | No | | `irodori_tts_600m_v3_voice_design` | Irodori-TTS 600M v3 VoiceDesign | No | | `index_tts2` | IndexTTS-2 | **Yes** | -| `kokoro_82m_bf16` | Kokoro 82M bf16 | Unavailable (loader not in this tree) | | `marblenet_vad` | MarbleNet VAD converted layout | No | | `mel_band_roformer` | Mel-Band RoFormer MLX | **Yes** | | `miocodec_25hz_44k_v2` | MioCodec 25Hz 44.1kHz v2 | No | @@ -399,7 +401,6 @@ Recommended top-level install packages: | `nemotron_asr` | Nemotron ASR | **Yes** | | `omnivoice` | OmniVoice | **Yes** | | `outetts_1_0_1b` | OuteTTS 1.0 1B with IBM DAC codec and Qwen3-aligned voice cloning | No | -| `parakeet_tdt_0_6b_v3` | Parakeet TDT 0.6B v3 | Unavailable (loader not in this tree) | | `pocket_tts` | PocketTTS | **Yes** | | `qwen3_asr_0_6b` | Qwen3 ASR 0.6B | **Yes** | | `qwen3_asr_1_7b_hf` | Qwen3 ASR 1.7B HF | **Yes** | @@ -419,7 +420,7 @@ Recommended top-level install packages: | `vibevoice_1_5b` | VibeVoice 1.5B | No | | `vibevoice_7b` | VibeVoice 7B | No | | `vibevoice_asr` | VibeVoice ASR | No | -| `voxtral_realtime` | Voxtral Mini 4B Realtime | **Yes** | +| `voxtral_realtime` | Voxtral Mini 4B Realtime GGUF Q8_0 | **Yes** | | `voxcpm2` | VoxCPM2 | No | > [!WARNING] @@ -612,7 +613,6 @@ For TTS-family models, the measured one-shot RTF is: | model | audio len (s) | wall time (s) | RTF | x faster than real time | |---|---:|---:|---:|---:| | chatterbox | 9.72 | 2.45 | 0.252 | 3.97x | -| kokoro tts | 10.15 | 0.64 | 0.063 | 15.90x | | miotts | 20.40 | 3.30 | 0.162 | 6.18x | | moss_tts_local | 9.60 | 0.97 | 0.101 | 9.91x | | omnivoice | 9.00 | 1.32 | 0.146 | 6.84x | @@ -627,7 +627,6 @@ For long-form TTS tests, each run uses the same 6,026-character, 1,028-word inpu | model | audio len (s) | wall time (s) | RTF | x faster than real time | |---|---:|---:|---:|---:| | chatterbox | 391.24 | 58.57 | 0.150 | 6.68x | -| kokoro tts | 371.17 | 7.19 | 0.019 | 51.60x | | index tts2 | 422.12 | 139.95 | 0.332 | 3.02x | | miotts | 399.16 | 66.59 | 0.167 | 5.99x | | moss_tts_nano | 391.20 | 43.16 | 0.110 | 9.06x | diff --git a/docs/asr.md b/docs/asr.md index e08bcb52..65da80fd 100644 --- a/docs/asr.md +++ b/docs/asr.md @@ -255,12 +255,12 @@ audiocpp_cli --task asr --family vibevoice_asr --model models/VibeVoice-ASR --ba ## Voxtral Realtime -Voxtral Realtime is a Mistral realtime ASR model with offline and streaming sessions. It accepts the native Hugging Face model directory and standalone audio.cpp GGUF packages. +Voxtral Realtime is a Mistral realtime ASR model with offline and streaming sessions. The model manager installs the Q8_0 standalone GGUF package by default; native Hugging Face directories and other standalone GGUF variants can also be used when provided directly. | Field | Value | |---|---| | Family | `voxtral_realtime` | -| Model directory | `models/Voxtral-Mini-4B-Realtime-2602` or a standalone Voxtral GGUF | +| Model path | `models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf` when installed through the model manager | | Task | `asr` | | Modes | `offline`, `streaming` | | Output | Transcription text | @@ -270,19 +270,19 @@ Voxtral Realtime is a Mistral realtime ASR model with offline and streaming sess Offline CLI: ```bash -audiocpp_cli --task asr --family voxtral_realtime --model --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt +audiocpp_cli --task asr --family voxtral_realtime --model models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt ``` Sampling and token-cap options can be passed through request options: ```bash -audiocpp_cli --task asr --family voxtral_realtime --model --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt --request-option max_new_tokens=256 --do-sample false --temperature 1.0 --top-p 1.0 --top-k 50 --seed 1234 +audiocpp_cli --task asr --family voxtral_realtime --model models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt --request-option max_new_tokens=256 --do-sample false --temperature 1.0 --top-p 1.0 --top-k 50 --seed 1234 ``` Streaming CLI: ```bash -audiocpp_cli --task asr --family voxtral_realtime --model --backend cuda --threads 8 --mode streaming --audio assets/resources/sample.wav --text-out transcript.txt +audiocpp_cli --task asr --family voxtral_realtime --model models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf --backend cuda --threads 8 --mode streaming --audio assets/resources/sample.wav --text-out transcript.txt ``` Streaming server config: @@ -299,7 +299,7 @@ Streaming server config: { "id": "voxtral-stream", "family": "voxtral_realtime", - "path": "/path/to/Voxtral-Mini-4B-Realtime-2602", + "path": "/path/to/voxtral-mini-4b-realtime-2602-q8_0.gguf", "task": "asr", "mode": "streaming" } diff --git a/docs/gguf.md b/docs/gguf.md index 68911b34..3b9474c8 100644 --- a/docs/gguf.md +++ b/docs/gguf.md @@ -259,8 +259,10 @@ Status labels: | `ace_step` | No | --- | --- | --- | --- | | `chatterbox` | No | --- | --- | --- | --- | | `citrinet_asr` | Done | Pass | --- | --- | Pass | +| `fish_audio` | Done | Pass | --- | Pass | Pass | | `heartmula` | No | --- | --- | --- | --- | | `higgs_audio_stt` | Done | Pass | --- | Pass | Pass | +| `higgs_audio_tts` | Done | Pass | --- | Pass | Pass | | `htdemucs` | Done | Pass | --- | Pass | Pass (drift) | | `hviske_asr` | Done | Pass | --- | --- | Pass | | `index_tts2` | Done | Pass | Pass | Pass (drift) | No (similarity drift, frame drift, text minor drift) | diff --git a/docs/tts.md b/docs/tts.md index 7262ed3e..8f350f45 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -10,6 +10,8 @@ | OmniVoice | `omnivoice` | `tts` | [OmniVoice](#omnivoice) | | PocketTTS | `pocket_tts` | `tts` | [PocketTTS](#pockettts) | | VoxCPM2 | `voxcpm2` | `tts`, `vdes` | [VoxCPM2](#voxcpm2) | +| 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) | | Irodori-TTS | `irodori_tts` | `tts`, `vdes` | [Irodori-TTS](#irodori-tts) | | OuteTTS | `outetts` | `tts`, `clon` | [OuteTTS](#outetts) | @@ -297,6 +299,89 @@ audiocpp_cli --task tts --family voxcpm2 --model models/VoxCPM2 --backend cuda - | `--num-inference-steps` | integer | `10` | Flow matching steps. | | `--guidance-scale` | float | `2.0` | CFG strength. | +## Higgs Audio v3 TTS + +Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the framework chunker for long text and keeps the reference prompt state in the model session. + +| Field | Value | +|---|---| +| Family | `higgs_audio_tts` | +| Model path | `models/Higgs-Audio-v3-TTS-4B-GGUF/higgs-audio-v3-tts-4b-q8_0.gguf` when installed through the model manager | +| Task | `tts` | +| Modes | `offline` | +| Languages | Model auto-handles supported languages | +| Voice input | Reference WAV through `--voice-ref`; transcript through `--reference-text` when known | +| Built-in voices | Not exposed | + +```bash +audiocpp_cli --task tts --family higgs_audio_tts --model models/Higgs-Audio-v3-TTS-4B-GGUF/higgs-audio-v3-tts-4b-q8_0.gguf --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +``` + +The model manager installs the Q8_0 standalone GGUF package by default: + +```bash +python3 tools/model_manager.py install --models-root models higgs_audio_v3_tts_4b +``` + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--voice-ref` | WAV path | required | Reference speaker audio. | +| `--reference-text` | text | empty string | Transcript for reference audio. | +| `--text-chunk-size` | integer chars | `1024` | Long-form chunk size. | +| `--max-tokens` | integer | `2048` | Maximum generated AR tokens per chunk. | +| `--temperature` | float | `0.8` | AR sampling temperature. | +| `--top-k` | integer | `30` | AR top-k sampling limit. The narrower default is less prone to premature EOC than the Python client's `50`. | +| `--top-p` | float | `0.8` | AR nucleus sampling limit. The Python client's unfiltered equivalent is `1.0`. | +| `--repetition-penalty` | float | `1.1` | Accepted for Python API compatibility; Higgs audio-code sampling does not consume it. | + +## Fish Audio S2 Pro + +Fish Audio S2 Pro is a TTS and reference voice-clone model. The integration uses the framework text chunker for long-form input, caches prepared reference audio in the session, and supports GGUF loading through the package spec path. + +| Field | Value | +|---|---| +| Family | `fish_audio` | +| Model path | `models/Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf` when installed through the model manager | +| Task | `tts` | +| Modes | `offline` | +| Languages | Model auto-handles language; tested paths cover English and Chinese-style prompts | +| Voice input | Optional reference WAV through `--voice-ref`; transcript through `--reference-text` when known | +| Built-in voices | Not exposed | + +Text-to-speech: + +```bash +audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf --backend cuda --text "Hello from Fish Audio." --out out.wav +``` + +Reference voice clone: + +```bash +audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf --backend cuda --text "The final render is ready for review." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +``` + +The model manager installs the Q8_0 standalone GGUF package by default: + +```bash +python3 tools/model_manager.py install --models-root models fish_audio_s2_pro +``` + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--voice-ref` | WAV path | not set | Reference speaker audio for voice cloning. | +| `--reference-text` | text | empty string | Transcript for reference audio. | +| `--max-new-tokens` | integer | `1024` | Maximum generated semantic tokens per chunk. `0` uses the default. | +| `--text-chunk-size` | integer chars | `200` | Long-form chunk size. | +| `--text-chunk-mode` | `default`, `tag_aware`, `japanese`, `endline` | `default` | Framework text chunking mode. | +| `--temperature` | float | `0.8` | Sampling temperature. | +| `--top-k` | integer | `30` | Top-k sampling limit. | +| `--top-p` | float | `0.8` | Nucleus sampling limit. | +| `--seed` | integer | random when omitted | Sampling seed for reproducible output. | +| `--session-option fish_audio.mem_saver=true|false` | bool | `false` | Release cached AR runtime graphs after each request. | +| `--session-option fish_audio.reference_cache_slots=` | integer | `1` | Prepared reference-audio cache slots. | +| `--session-option fish_audio.weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | AR matmul weight storage type. | +| `--session-option fish_audio.codec_weight_type=` | `native`, `f32`, `f16`, `q8_0` | `native` | Codec conv/matmul weight storage type. | + ## IndexTTS2 IndexTTS2 is a Chinese and English TTS model with voice cloning and expressive emotion controls. It requires a speaker reference through the framework `--voice-ref` path. diff --git a/include/engine/framework/core/backend.h b/include/engine/framework/core/backend.h index ac8c0277..8c186826 100644 --- a/include/engine/framework/core/backend.h +++ b/include/engine/framework/core/backend.h @@ -70,10 +70,16 @@ void write_tensor_f32_slice(const TensorValue & tensor, size_t element_offset, c void write_tensor_f32(const TensorValue & tensor, const std::vector & values); void write_tensor_f16(const TensorValue & tensor, const float * values, size_t count); void write_tensor_f16(const TensorValue & tensor, const std::vector & values); +void write_tensor_bf16(const TensorValue & tensor, const float * values, size_t count); +void write_tensor_bf16(const TensorValue & tensor, const std::vector & values); void write_tensor_i32(const TensorValue & tensor, const int32_t * values, size_t count); void write_tensor_i32(const TensorValue & tensor, const std::vector & values); void read_tensor_f32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_f32(const ggml_tensor * tensor); +void read_tensor_f16_into(const ggml_tensor * tensor, std::vector & values); +std::vector read_tensor_f16(const ggml_tensor * tensor); +void read_tensor_bf16_into(const ggml_tensor * tensor, std::vector & values); +std::vector read_tensor_bf16(const ggml_tensor * tensor); void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_i32(const ggml_tensor * tensor); diff --git a/include/engine/framework/modules/attention/qwen_causal_decoder.h b/include/engine/framework/modules/attention/qwen_causal_decoder.h index ffc51353..0ee5ebe7 100644 --- a/include/engine/framework/modules/attention/qwen_causal_decoder.h +++ b/include/engine/framework/modules/attention/qwen_causal_decoder.h @@ -79,6 +79,11 @@ std::vector qwen_position_ids(int64_t steps, int64_t offset = 0); std::vector qwen_causal_prefill_mask_values(int64_t batch_size, int64_t steps); +std::vector qwen_causal_suffix_mask_values( + int64_t batch_size, + int64_t query_steps, + int64_t prefix_steps); + void write_qwen_causal_prefill_mask( ggml_tensor * tensor, int64_t batch_size, diff --git a/include/engine/framework/modules/attention/qwen_decoder.h b/include/engine/framework/modules/attention/qwen_decoder.h index 6107ae74..56b85603 100644 --- a/include/engine/framework/modules/attention/qwen_decoder.h +++ b/include/engine/framework/modules/attention/qwen_decoder.h @@ -24,11 +24,27 @@ enum class QwenDecoderStaticCacheUpdateMode { DirectSetRows, }; +enum class QwenDecoderStaticCacheSetRowsMode { + Exact, + BackendViewOptimized, +}; + enum class QwenDecoderQKVLayout { Separate, PackedQKV, }; +enum class QwenDecoderMLPMode { + Exact, + FusedSwiGLU, + PackedGateUp, +}; + +enum class QwenDecoderPrefixAttentionMode { + Exact, + FlashWithPrefix, +}; + enum class QwenDecoderPositionEncoding { Rotary, None, @@ -55,16 +71,23 @@ struct QwenDecoderActivationCastPolicy { struct QwenDecoderAttentionPolicy { QwenDecoderAttentionMode prefill_mode = QwenDecoderAttentionMode::ManualRepeat; QwenDecoderAttentionMode static_mode = QwenDecoderAttentionMode::FlashGrouped; + QwenDecoderPrefixAttentionMode prefix_mode = QwenDecoderPrefixAttentionMode::Exact; int64_t grouped_query_min_steps = 0; }; struct QwenDecoderStaticCachePolicy { QwenDecoderStaticCacheUpdateMode update_mode = QwenDecoderStaticCacheUpdateMode::ScratchTail; + QwenDecoderStaticCacheSetRowsMode set_rows_mode = QwenDecoderStaticCacheSetRowsMode::Exact; +}; + +struct QwenDecoderMLPPolicy { + QwenDecoderMLPMode mode = QwenDecoderMLPMode::Exact; }; struct QwenDecoderRuntimePolicy { QwenDecoderAttentionPolicy attention; QwenDecoderStaticCachePolicy static_cache; + QwenDecoderMLPPolicy mlp; }; struct QwenDecoderLayerConfig { @@ -88,6 +111,7 @@ struct QwenDecoderLayerConfig { struct QwenMLPWeights { LinearWeights gate_proj; LinearWeights up_proj; + std::optional gate_up_proj; LinearWeights down_proj; }; diff --git a/include/engine/framework/modules/optimizations/fast_kv_modules.h b/include/engine/framework/modules/optimizations/fast_kv_modules.h index 75175b2a..247a8cc2 100644 --- a/include/engine/framework/modules/optimizations/fast_kv_modules.h +++ b/include/engine/framework/modules/optimizations/fast_kv_modules.h @@ -4,8 +4,19 @@ namespace engine::modules { +enum class FastKVSetRowsMode { + Exact, + BackendViewOptimized, +}; + +struct FastKVSetRowsConfig { + FastKVSetRowsMode mode = FastKVSetRowsMode::Exact; +}; + class FastKVSetRowsModule { public: + explicit FastKVSetRowsModule(FastKVSetRowsConfig config = {}); + const core::ModuleSchema & schema() const noexcept; core::TensorValue build( @@ -15,6 +26,9 @@ class FastKVSetRowsModule { const core::TensorValue & row_index) const; static const core::ModuleSchema & static_schema() noexcept; + +private: + FastKVSetRowsConfig config_; }; } // namespace engine::modules diff --git a/include/engine/framework/runtime/kv_cache.h b/include/engine/framework/runtime/kv_cache.h index e38a00e7..81156423 100644 --- a/include/engine/framework/runtime/kv_cache.h +++ b/include/engine/framework/runtime/kv_cache.h @@ -19,6 +19,11 @@ struct TransformerKVState { std::vector layers; }; +struct TransformerKVCacheOptions { + bool allow_f16_storage = false; + bool allow_bf16_storage = false; +}; + class TransformerKVCache { public: TransformerKVCache() = default; @@ -27,6 +32,12 @@ class TransformerKVCache { int64_t step_elems, std::vector keys, std::vector values); + TransformerKVCache( + int64_t cache_steps, + int64_t step_elems, + std::vector keys, + std::vector values, + TransformerKVCacheOptions options); void import_state(const TransformerKVState & state); TransformerKVState export_state() const; @@ -55,6 +66,7 @@ class TransformerKVCache { int64_t step_elems_ = 0; int64_t valid_steps_ = 0; int64_t current_end_ = 0; + TransformerKVCacheOptions options_; std::vector layers_; }; @@ -65,6 +77,7 @@ core::TensorValue view_transformer_kv_cache_steps( int64_t steps, int64_t heads, int64_t head_dim, - const char * label); + const char * label, + ggml_type view_type = GGML_TYPE_F32); } // namespace engine::runtime diff --git a/include/engine/models/fish_audio/ar.h b/include/engine/models/fish_audio/ar.h new file mode 100644 index 00000000..db0e0a0c --- /dev/null +++ b/include/engine/models/fish_audio/ar.h @@ -0,0 +1,31 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/models/fish_audio/assets.h" +#include "engine/models/fish_audio/types.h" + +#include + +namespace engine::models::fish_audio { + +class FishAudioARRuntime { +public: + FishAudioARRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~FishAudioARRuntime(); + + FishAudioCodes generate(const FishAudioPrompt & prompt, const FishAudioGenerationOptions & options); + void release_runtime_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/assets.h b/include/engine/models/fish_audio/assets.h new file mode 100644 index 00000000..338ce866 --- /dev/null +++ b/include/engine/models/fish_audio/assets.h @@ -0,0 +1,24 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/models/fish_audio/types.h" + +#include +#include + +namespace engine::assets { +class TensorSource; +} + +namespace engine::models::fish_audio { + +struct FishAudioAssets { + assets::ResourceBundle resources; + FishAudioConfig config; + std::shared_ptr model_weights; + std::shared_ptr codec_weights; +}; + +std::shared_ptr load_fish_audio_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/codec.h b/include/engine/models/fish_audio/codec.h new file mode 100644 index 00000000..2da7556e --- /dev/null +++ b/include/engine/models/fish_audio/codec.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/models/fish_audio/assets.h" +#include "engine/models/fish_audio/types.h" + +#include + +namespace engine::models::fish_audio { + +class FishAudioCodecRuntime { +public: + FishAudioCodecRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType matmul_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type); + ~FishAudioCodecRuntime(); + + FishAudioCodes encode_reference(const runtime::AudioBuffer & audio); + runtime::AudioBuffer decode(const FishAudioCodes & codes); + void release_encode_graph(); + void release_runtime_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/generator.h b/include/engine/models/fish_audio/generator.h new file mode 100644 index 00000000..4f9c7800 --- /dev/null +++ b/include/engine/models/fish_audio/generator.h @@ -0,0 +1,41 @@ +#pragma once + +#include "engine/models/fish_audio/ar.h" +#include "engine/models/fish_audio/codec.h" +#include "engine/models/fish_audio/prompt_builder.h" +#include "engine/models/fish_audio/tokenizer_text.h" + +#include +#include + +namespace engine::models::fish_audio { + +struct FishAudioGenerationResult { + runtime::AudioBuffer audio; + FishAudioCodes codes; +}; + +class FishAudioGenerator { +public: + FishAudioGenerator( + std::shared_ptr assets, + std::unique_ptr ar, + std::unique_ptr codec); + ~FishAudioGenerator(); + + FishAudioCodes encode_reference(const runtime::AudioBuffer & audio); + FishAudioGenerationResult generate( + const FishAudioRequest & request, + const std::optional & reference_codes, + const std::optional & previous_turn, + bool mem_saver); + +private: + std::shared_ptr assets_; + FishAudioTextTokenizer tokenizer_; + FishAudioPromptBuilder prompt_builder_; + std::unique_ptr ar_; + std::unique_ptr codec_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/loader.h b/include/engine/models/fish_audio/loader.h new file mode 100644 index 00000000..65e49f4e --- /dev/null +++ b/include/engine/models/fish_audio/loader.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/models/fish_audio/assets.h" + +#include +#include + +namespace engine::models::fish_audio { + +class FishAudioLoadedModel final : public runtime::ILoadedVoiceModel { +public: + FishAudioLoadedModel( + 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_fish_audio_model(const std::filesystem::path & model_path); +std::shared_ptr make_fish_audio_loader(); + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/prompt_builder.h b/include/engine/models/fish_audio/prompt_builder.h new file mode 100644 index 00000000..6ed69e49 --- /dev/null +++ b/include/engine/models/fish_audio/prompt_builder.h @@ -0,0 +1,22 @@ +#pragma once + +#include "engine/models/fish_audio/tokenizer_text.h" +#include "engine/models/fish_audio/types.h" + +namespace engine::models::fish_audio { + +class FishAudioPromptBuilder { +public: + FishAudioPromptBuilder(std::shared_ptr assets, FishAudioTextTokenizer tokenizer); + + FishAudioPrompt build( + const FishAudioRequest & request, + const std::optional & reference_codes, + const std::optional & previous_turn) const; + +private: + std::shared_ptr assets_; + FishAudioTextTokenizer tokenizer_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/session.h b/include/engine/models/fish_audio/session.h new file mode 100644 index 00000000..40e6bb5f --- /dev/null +++ b/include/engine/models/fish_audio/session.h @@ -0,0 +1,58 @@ +#pragma once + +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/fish_audio/assets.h" +#include "engine/models/fish_audio/generator.h" + +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { + +class FishAudioSession final : public runtime::RuntimeSessionBase, public runtime::IOfflineVoiceTaskSession { +public: + FishAudioSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + ~FishAudioSession() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + struct ReferenceCacheKey { + std::string source_id; + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; + }; + + struct ReferenceCacheKeyEqual { + bool operator()(const ReferenceCacheKey & lhs, const ReferenceCacheKey & rhs) const; + }; + + struct ReferenceCacheEntry { + FishAudioCodes codes; + }; + + FishAudioRequest make_request(const runtime::TaskRequest & request) const; + const FishAudioCodes & resolve_reference_codes(const FishAudioReference & reference); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::unique_ptr generator_; + std::optional defaults_; + runtime::CacheSlots reference_cache_; + std::optional uncached_reference_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/tokenizer_text.h b/include/engine/models/fish_audio/tokenizer_text.h new file mode 100644 index 00000000..b457fb2f --- /dev/null +++ b/include/engine/models/fish_audio/tokenizer_text.h @@ -0,0 +1,28 @@ +#pragma once + +#include "engine/models/fish_audio/assets.h" + +#include +#include +#include +#include + +namespace engine::models::fish_audio { + +class FishAudioTextTokenizer { +public: + struct Impl; + + explicit FishAudioTextTokenizer(std::shared_ptr assets); + + std::vector encode(const std::string & text) const; + int32_t token_id(const std::string & token) const; + int32_t im_end_id() const noexcept; + int32_t semantic_begin_id() const noexcept; + int32_t semantic_end_id() const noexcept; + +private: + std::shared_ptr impl_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/types.h b/include/engine/models/fish_audio/types.h new file mode 100644 index 00000000..d76523ea --- /dev/null +++ b/include/engine/models/fish_audio/types.h @@ -0,0 +1,105 @@ +#pragma once + +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include + +namespace engine::models::fish_audio { + +struct FishAudioGenerationOptions { + int64_t max_new_tokens = 1024; + int64_t text_chunk_size = 200; + float top_p = 0.8F; + int top_k = 30; + float temperature = 0.8F; + uint32_t seed = 1234; +}; + +struct FishAudioReference { + std::optional audio = std::nullopt; + std::string text; + std::string cache_id; +}; + +struct FishAudioRequest { + std::string text; + std::optional reference = std::nullopt; + FishAudioGenerationOptions generation; +}; + +struct FishAudioCodes { + std::vector codes; + int64_t codebooks = 0; + int64_t frames = 0; +}; + +struct FishAudioConversationTurn { + std::string text; + FishAudioCodes codes; +}; + +struct FishAudioPrompt { + std::vector matrix; + int64_t codebook_rows = 0; + int64_t steps = 0; + std::string text; +}; + +struct FishAudioTextConfig { + int64_t vocab_size = 0; + int64_t n_layer = 0; + int64_t dim = 0; + int64_t intermediate_size = 0; + int64_t n_head = 0; + int64_t n_local_heads = 0; + int64_t head_dim = 0; + int64_t max_seq_len = 0; + float rope_base = 1000000.0F; + float norm_eps = 1.0e-6F; + bool tie_word_embeddings = true; + bool attention_qk_norm = true; +}; + +struct FishAudioFastConfig { + int64_t vocab_size = 0; + int64_t num_codebooks = 0; + int64_t n_layer = 0; + int64_t dim = 0; + int64_t intermediate_size = 0; + int64_t n_head = 0; + int64_t n_local_heads = 0; + int64_t head_dim = 0; + int64_t max_seq_len = 0; + float rope_base = 1000000.0F; + float norm_eps = 1.0e-6F; + bool tie_word_embeddings = false; + bool attention_qk_norm = false; +}; + +struct FishAudioCodecConfig { + int sample_rate = 44100; + int64_t semantic_codebook_size = 4096; + int64_t residual_codebook_size = 1024; + int64_t quantizer_codebooks = 9; + int64_t total_codebooks = 10; + int64_t codebook_dim = 8; + int64_t latent_dim = 1024; + int64_t frame_length = 2048; +}; + +struct FishAudioConfig { + std::string model_type; + std::string torch_dtype; + int64_t semantic_start_token_id = 0; + int64_t semantic_end_token_id = 0; + int64_t im_end_token_id = 0; + bool norm_fastlayer_input = false; + FishAudioTextConfig text; + FishAudioFastConfig fast; + FishAudioCodecConfig codec; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/higgs_audio_tts/ar.h b/include/engine/models/higgs_audio_tts/ar.h new file mode 100644 index 00000000..195bfdf3 --- /dev/null +++ b/include/engine/models/higgs_audio_tts/ar.h @@ -0,0 +1,168 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/attention/qwen_decoder.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/models/higgs_audio_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::core { +class BackendWeightStore; +} + +namespace engine::models::higgs_audio_tts { + +struct HiggsQwenDecoderStackWeights { + std::vector layers; +}; + +struct HiggsARWeights { + std::shared_ptr store; + core::TensorValue text_embedding; + core::TensorValue modality_embedding; + HiggsQwenDecoderStackWeights decoder; + core::TensorValue norm; + bool packed_qkv = false; +}; + +HiggsARWeights load_higgs_ar_weights( + const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + +class HiggsARRuntime { +public: + HiggsARRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + + const HiggsAssets & assets() const noexcept; + const HiggsARWeights & weights() const noexcept; + ggml_backend_t backend() const noexcept; + core::BackendType backend_type() const noexcept; + int device() const noexcept; + int threads() const noexcept; + +private: + std::shared_ptr assets_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int device_ = 0; + int threads_ = 1; + std::shared_ptr weights_; +}; + +struct HiggsARDecodeInput { + std::vector last_codes; + bool use_last_codes = false; +}; + +struct HiggsARDecodeOutput { + std::vector codebook_logits; +}; + +struct HiggsARDecodeTiming { + double input_upload_ms = 0.0; + double mask_upload_ms = 0.0; + double graph_compute_ms = 0.0; + double output_read_ms = 0.0; + int64_t steps = 0; + + void add(const HiggsARDecodeTiming & other) noexcept; +}; + +struct HiggsARPrefillInput { + std::vector text_tokens; + std::vector fused_code_ids; + std::vector text_gate; + std::vector code_gate; + int64_t steps = 0; +}; + +struct HiggsARPrefillOutput { + HiggsARDecodeOutput output; + runtime::TransformerKVState kv_state; + bool wrote_cache = false; +}; + +class HiggsARKVCache { +public: + HiggsARKVCache(std::shared_ptr runtime, int64_t cache_steps); + ~HiggsARKVCache(); + + HiggsARKVCache(const HiggsARKVCache &) = delete; + HiggsARKVCache & operator=(const HiggsARKVCache &) = delete; + + bool can_run(const HiggsARRuntime & runtime, int64_t required_steps) const; + int64_t cache_steps() const; + int64_t valid_steps() const; + int64_t current_end() const; + void reset(); + void retain_prefix(int64_t prefix_steps); + void import_state(const runtime::TransformerKVState & state); + runtime::TransformerKVState export_state() const; + void advance_after_direct_append(int64_t steps); + const core::TensorValue & key_tensor(size_t layer) const; + const core::TensorValue & value_tensor(size_t layer) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +class HiggsARPrefillGraph { +public: + HiggsARPrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t start_step, + HiggsARKVCache * cache, + size_t graph_arena_bytes); + ~HiggsARPrefillGraph(); + + HiggsARPrefillGraph(const HiggsARPrefillGraph &) = delete; + HiggsARPrefillGraph & operator=(const HiggsARPrefillGraph &) = delete; + + bool matches(const HiggsARRuntime & runtime, int64_t prompt_steps, int64_t start_step) const; + HiggsARPrefillOutput run(const HiggsARPrefillInput & input, int64_t start_step = 0); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +class HiggsARDecodeGraph { +public: + HiggsARDecodeGraph( + std::shared_ptr runtime, + int64_t cache_steps, + HiggsARKVCache & cache, + size_t graph_arena_bytes); + ~HiggsARDecodeGraph(); + + HiggsARDecodeGraph(const HiggsARDecodeGraph &) = delete; + HiggsARDecodeGraph & operator=(const HiggsARDecodeGraph &) = delete; + + bool can_run(const HiggsARRuntime & runtime, int64_t required_steps) const; + int64_t cache_steps() const; + void import_prefill_state(const runtime::TransformerKVState & state); + void begin_decode_run(); + HiggsARDecodeTiming timing() const; + void run_step_into(const HiggsARDecodeInput & input, HiggsARDecodeOutput & output, bool log_timing = false); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/assets.h b/include/engine/models/higgs_audio_tts/assets.h new file mode 100644 index 00000000..aadc195e --- /dev/null +++ b/include/engine/models/higgs_audio_tts/assets.h @@ -0,0 +1,65 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" + +#include +#include +#include +#include + +namespace engine::assets { +class TensorSource; +} + +namespace engine::models::higgs_audio_tts { + +struct HiggsTextConfig { + std::string model_type; + int64_t vocab_size = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; + int64_t head_dim = 0; + int64_t max_position_embeddings = 0; + int64_t bos_token_id = 0; + int64_t eos_token_id = 0; + int64_t pad_token_id = -1; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 1000000.0F; + bool tie_word_embeddings = true; +}; + +struct HiggsAudioEncoderConfig { + std::string model_type; + std::string encoder_type; + int64_t num_codebooks = 0; + int64_t vocab_size = 0; + int64_t out_dim = 0; + int64_t mel_per_sample = 0; + int64_t max_chunk_size = 0; + bool tie_word_embeddings = true; + bool use_delay_pattern = true; +}; + +struct HiggsConfig { + std::string model_type; + std::string architecture; + int64_t hidden_size = 0; + int64_t vocab_size = 0; + int64_t audio_token_id = -100; + int64_t ignore_index = -100; + HiggsTextConfig text; + HiggsAudioEncoderConfig audio; +}; + +struct HiggsAssets { + assets::ResourceBundle resources; + HiggsConfig config; + std::shared_ptr weights; +}; + +std::shared_ptr load_higgs_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/codebooks.h b/include/engine/models/higgs_audio_tts/codebooks.h new file mode 100644 index 00000000..06024ab9 --- /dev/null +++ b/include/engine/models/higgs_audio_tts/codebooks.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include + +namespace engine::models::higgs_audio_tts { + +constexpr int32_t kHiggsBocId = 1024; +constexpr int32_t kHiggsEocId = 1025; +constexpr int32_t kHiggsStopCode = -1; + +int64_t higgs_delayed_frame_count(int64_t raw_frames, int64_t codebooks); + +std::vector apply_higgs_delay_pattern( + const std::vector & raw_codes, + int64_t raw_frames, + int64_t codebooks); + +std::vector reverse_higgs_delay_pattern( + const std::vector & delayed_codes, + int64_t delayed_frames, + int64_t codebooks); + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/codec.h b/include/engine/models/higgs_audio_tts/codec.h new file mode 100644 index 00000000..28516eea --- /dev/null +++ b/include/engine/models/higgs_audio_tts/codec.h @@ -0,0 +1,138 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/higgs_audio_tts/assets.h" + +#include +#include +#include +#include +#include + +namespace engine::core { +class BackendWeightStore; +} + +namespace engine::models::higgs_audio_tts { + +class HiggsCodecDecodeGraph; +class HiggsCodecEncodeGraph; + +struct HiggsCodecVectorQuantizerWeights { + core::TensorValue codebook; + modules::LinearWeights project_in; + modules::LinearWeights project_out; +}; + +struct HiggsCodecResidualUnitWeights { + modules::Snake1dWeights snake1; + modules::Conv1dWeights conv1; + modules::Snake1dWeights snake2; + modules::Conv1dWeights conv2; +}; + +struct HiggsCodecDecoderBlockWeights { + modules::Snake1dWeights snake; + modules::ConvTranspose1dWeights conv_transpose; + std::vector residual_units; +}; + +struct HiggsCodecEncoderBlockWeights { + modules::Snake1dWeights snake; + modules::Conv1dWeights conv; + std::vector residual_units; +}; + +struct HiggsCodecSemanticResidualUnitWeights { + modules::Conv1dWeights conv1; + modules::Conv1dWeights conv2; +}; + +struct HiggsCodecSemanticEncoderBlockWeights { + std::vector residual_units; + modules::Conv1dWeights conv; +}; + +struct HiggsCodecWeights { + std::shared_ptr store; + std::unordered_map semantic_model; + std::vector quantizers; + modules::Conv1dWeights acoustic_encoder_input; + std::vector acoustic_encoder_blocks; + modules::Snake1dWeights acoustic_encoder_output_snake; + modules::Conv1dWeights acoustic_encoder_output; + modules::Conv1dWeights semantic_encoder_input; + std::vector semantic_encoder_blocks; + modules::LinearWeights codec_project; + modules::LinearWeights acoustic_project; + modules::Conv1dWeights acoustic_decoder_input; + std::vector acoustic_decoder_blocks; + modules::Snake1dWeights acoustic_decoder_output_snake; + modules::Conv1dWeights acoustic_decoder_output; +}; + +struct HiggsCodecDecodeOutput { + int sample_rate = 24000; + int channels = 1; + int64_t samples = 0; + std::vector values; +}; + +struct HiggsCodecEncodeOutput { + std::vector codes; + int64_t frames = 0; + int64_t codebooks = 0; +}; + +class HiggsCodecRuntime { +public: + HiggsCodecRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t decode_graph_arena_bytes, + size_t encode_graph_arena_bytes, + assets::TensorStorageType weight_storage_type); + ~HiggsCodecRuntime(); + + const HiggsCodecWeights & weights() const noexcept; + ggml_backend_t backend() const noexcept; + core::BackendType backend_type() const noexcept; + int threads() const noexcept; + size_t decode_graph_arena_bytes() const noexcept; + size_t encode_graph_arena_bytes() const noexcept; + + HiggsCodecEncodeOutput encode_reference(const runtime::AudioBuffer & audio) const; + HiggsCodecDecodeOutput decode_codes( + const std::vector & codes, + int64_t frames, + int64_t codebooks) const; + void release_encode_graph(); + void release_runtime_graphs(); + +private: + std::shared_ptr assets_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + size_t decode_graph_arena_bytes_ = 0; + size_t encode_graph_arena_bytes_ = 0; + std::shared_ptr weights_; + mutable std::unique_ptr encode_graph_; + mutable std::unique_ptr decode_graph_; +}; + +HiggsCodecWeights load_higgs_codec_decode_weights( + const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/generator.h b/include/engine/models/higgs_audio_tts/generator.h new file mode 100644 index 00000000..aeb15f9b --- /dev/null +++ b/include/engine/models/higgs_audio_tts/generator.h @@ -0,0 +1,77 @@ +#pragma once + +#include "engine/models/higgs_audio_tts/ar.h" +#include "engine/models/higgs_audio_tts/codec.h" +#include "engine/models/higgs_audio_tts/sampler.h" +#include "engine/models/higgs_audio_tts/tokenizer_text.h" + +#include +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { + +struct HiggsGenerationOptions { + int64_t max_tokens = 2048; + float temperature = 1.0F; + std::optional top_p; + std::optional top_k; + float repetition_penalty = 1.0F; + std::optional seed; +}; + +struct HiggsGenerationRequest { + std::string text; + std::string reference_text; + std::vector reference_codes; + int64_t reference_frames = 0; + int64_t reference_codebooks = 0; + HiggsGenerationOptions options; +}; + +struct HiggsGenerationResult { + HiggsCodecDecodeOutput audio; + std::vector delayed_codes; + int64_t delayed_frames = 0; + std::vector raw_codes; + int64_t raw_frames = 0; +}; + +class HiggsGenerator { +public: + HiggsGenerator(std::shared_ptr assets, + std::shared_ptr ar, + std::shared_ptr codec, + size_t ar_decode_graph_arena_bytes); + + void prepare(const HiggsGenerationRequest & request); + HiggsGenerationResult generate(const HiggsGenerationRequest & request); + +private: + struct ReferencePrefixCache { + std::string reference_text; + std::vector reference_codes; + int64_t reference_frames = 0; + int64_t reference_codebooks = 0; + std::vector delayed_reference_codes; + int64_t delayed_reference_frames = 0; + std::vector prefix_tokens; + int64_t prefix_steps = 0; + }; + + std::shared_ptr assets_; + std::shared_ptr ar_; + std::shared_ptr codec_; + HiggsTextTokenizer tokenizer_; + size_t ar_decode_graph_arena_bytes_ = 0; + std::optional reference_prefix_cache_; + bool reference_kv_ready_ = false; + std::optional cuda_sampling_policy_; + std::unique_ptr ar_kv_cache_; + std::unique_ptr prefill_graph_; + std::unique_ptr decode_graph_; +}; + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/loader.h b/include/engine/models/higgs_audio_tts/loader.h new file mode 100644 index 00000000..b5fb9d58 --- /dev/null +++ b/include/engine/models/higgs_audio_tts/loader.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/models/higgs_audio_tts/assets.h" + +#include +#include + +namespace engine::models::higgs_audio_tts { + +class HiggsTTSLoadedModel final : public runtime::ILoadedVoiceModel { +public: + HiggsTTSLoadedModel( + 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_higgs_audio_tts_model(const std::filesystem::path & model_path); +std::shared_ptr make_higgs_audio_tts_loader(); + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/sampler.h b/include/engine/models/higgs_audio_tts/sampler.h new file mode 100644 index 00000000..76943d75 --- /dev/null +++ b/include/engine/models/higgs_audio_tts/sampler.h @@ -0,0 +1,56 @@ +#pragma once + +#include "engine/framework/sampling/torch_random.h" +#include "engine/models/higgs_audio_tts/codebooks.h" + +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { + +constexpr int64_t kHiggsMaxTopK = 1026; + +using HiggsCudaSamplingPolicy = engine::sampling::TorchCudaSamplingPolicy; + +struct HiggsSamplingOptions { + float temperature = 1.0F; + std::optional top_p; + std::optional top_k; + bool has_seed = false; + uint64_t seed = 0; + HiggsCudaSamplingPolicy cuda_policy; + std::mt19937 * fallback_rng = nullptr; +}; + +struct HiggsSamplerState { + int64_t num_codebooks = 0; + int64_t delay_count = 0; + std::optional eoc_countdown; + bool generation_done = false; + int64_t step_count = 0; + std::vector last_codes; +}; + +class HiggsCodebookSampler { +public: + explicit HiggsCodebookSampler(int64_t num_codebooks, int64_t codebook_vocab_size); + + HiggsSamplerState make_state() const; + const std::vector & step(const float * logits, + int64_t logits_count, + HiggsSamplerState & state, + HiggsSamplingOptions & options); + +private: + int64_t num_codebooks_ = 0; + int64_t codebook_vocab_size_ = 0; + std::vector scratch_scores_; + std::vector scratch_probs_; + std::vector scratch_order_; + std::vector scratch_kept_; + std::vector scratch_codes_; +}; + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/session.h b/include/engine/models/higgs_audio_tts/session.h new file mode 100644 index 00000000..6b19ff2a --- /dev/null +++ b/include/engine/models/higgs_audio_tts/session.h @@ -0,0 +1,79 @@ +#pragma once + +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/higgs_audio_tts/assets.h" +#include "engine/models/higgs_audio_tts/ar.h" +#include "engine/models/higgs_audio_tts/codec.h" +#include "engine/models/higgs_audio_tts/generator.h" + +#include +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { + +class HiggsTTSSession final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + HiggsTTSSession( + 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 ReferenceCacheEntry { + HiggsCodecEncodeOutput codes; + }; + + struct ReferenceCacheKey { + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; + std::string reference_text; + }; + + struct ReferenceCacheKeyEqual { + bool operator()(const ReferenceCacheKey & lhs, const ReferenceCacheKey & rhs) const noexcept { + return lhs.sample_rate == rhs.sample_rate && + lhs.channels == rhs.channels && + lhs.sample_count == rhs.sample_count && + lhs.sample_hash == rhs.sample_hash && + lhs.reference_text == rhs.reference_text; + } + }; + + HiggsGenerationRequest make_generation_request( + const runtime::TaskRequest & request, + const HiggsCodecEncodeOutput * resolved_reference_codes = nullptr); + const HiggsCodecEncodeOutput & resolve_reference_codes( + const runtime::AudioBuffer & audio, + const std::string & reference_text); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + size_t ar_weight_context_bytes_ = 4096ull * 1024ull * 1024ull; + size_t codec_weight_context_bytes_ = 1536ull * 1024ull * 1024ull; + size_t ar_decode_graph_arena_bytes_ = 512ull * 1024ull * 1024ull; + size_t codec_decode_graph_arena_bytes_ = 128ull * 1024ull * 1024ull; + size_t codec_encode_graph_arena_bytes_ = 256ull * 1024ull * 1024ull; + assets::TensorStorageType ar_weight_storage_type_ = assets::TensorStorageType::Native; + assets::TensorStorageType codec_weight_storage_type_ = assets::TensorStorageType::Native; + std::shared_ptr ar_; + std::shared_ptr codec_; + std::unique_ptr generator_; + runtime::CacheSlots reference_cache_; + std::optional uncached_reference_; +}; + +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_audio_tts/tokenizer_text.h b/include/engine/models/higgs_audio_tts/tokenizer_text.h new file mode 100644 index 00000000..33cb41a1 --- /dev/null +++ b/include/engine/models/higgs_audio_tts/tokenizer_text.h @@ -0,0 +1,37 @@ +#pragma once + +#include "engine/models/higgs_audio_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { + +struct HiggsPromptRequest { + std::string text; + std::string reference_text; + int64_t delayed_reference_tokens = 0; +}; + +struct HiggsPromptEncoding { + std::vector token_ids; + std::vector text_ids; + std::vector reference_text_ids; +}; + +class HiggsTextTokenizer { +public: + struct Impl; + + explicit HiggsTextTokenizer(std::shared_ptr assets); + + std::vector encode(const std::string & text) const; + HiggsPromptEncoding encode_prompt(const HiggsPromptRequest & request) const; + +private: + std::shared_ptr impl_; +}; + +} // namespace engine::models::higgs_audio_tts diff --git a/model_specs/fish_audio.json b/model_specs/fish_audio.json new file mode 100644 index 00000000..73e7a869 --- /dev/null +++ b/model_specs/fish_audio.json @@ -0,0 +1,42 @@ +{ + "family": "fish_audio", + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json", + "tokenizer_config": "model:tokenizer_config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "model_weights": { + "source": "weights:", + "prefix": "model_weights" + }, + "codec_weights": { + "source": "weights:", + "prefix": "codec_weights" + } + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.json", + "tokenizer_config": "model:tokenizer_config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "model_weights": "model:model_audio_cpp.safetensors.index.json", + "codec_weights": "model:codec.safetensors" + } + } + ] +} diff --git a/model_specs/higgs_audio_tts.json b/model_specs/higgs_audio_tts.json new file mode 100644 index 00000000..798d4d6d --- /dev/null +++ b/model_specs/higgs_audio_tts.json @@ -0,0 +1,36 @@ +{ + "family": "higgs_audio_tts", + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json", + "tokenizer_json": "model:tokenizer.json", + "tokenizer_config": "model:tokenizer_config.json", + "chat_template": "model:chat_template.jinja" + }, + "tensors": { + "weights": "weights:" + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.json", + "tokenizer_json": "model:tokenizer.json", + "tokenizer_config": "model:tokenizer_config.json", + "chat_template": "model:chat_template.jinja" + }, + "tensors": { + "weights": "model:model.safetensors.index.json" + } + } + ] +} diff --git a/src/framework/assets/model_package.cpp b/src/framework/assets/model_package.cpp index d93b5819..f025ff60 100644 --- a/src/framework/assets/model_package.cpp +++ b/src/framework/assets/model_package.cpp @@ -56,6 +56,41 @@ bool is_gguf_file(const std::filesystem::path & path) { return extension == ".gguf"; } +std::vector directory_gguf_files(const std::filesystem::path & path) { + std::vector files; + if (!engine::io::is_existing_directory(path)) { + return files; + } + for (const auto & entry : std::filesystem::directory_iterator(path)) { + const auto candidate = entry.path(); + if (is_gguf_file(candidate)) { + files.push_back(candidate.filename().string()); + } + } + std::sort(files.begin(), files.end()); + return files; +} + +std::string directory_gguf_hint(std::string_view family) { + if (!active_model_path.has_value()) { + return {}; + } + const auto files = directory_gguf_files(*active_model_path); + if (files.empty()) { + return {}; + } + std::string message = "model directory has no default GGUF for family '" + std::string(family) + + "': " + active_model_path->string() + "; found: "; + for (size_t i = 0; i < files.size(); ++i) { + if (i != 0) { + message += ", "; + } + message += files[i]; + } + message += "; pass the GGUF file directly with --model, or rename it to model.gguf"; + return message; +} + std::optional active_gguf_path() { if (!active_model_path.has_value()) return std::nullopt; @@ -378,6 +413,9 @@ std::filesystem::path default_model_package_spec_path(std::string_view family) { if (const auto external = discover_external_model_spec(family)) { return *external; } + if (const auto hint = directory_gguf_hint(family); !hint.empty()) { + throw std::runtime_error(hint); + } throw std::runtime_error("model package spec not found for family '" + std::string(family) + "' (provide --model-spec-override, embed it in the GGUF, enable " "AUDIOCPP_DEPLOYMENT_BUILD, or install model_specs/" + diff --git a/src/framework/core/backend.cpp b/src/framework/core/backend.cpp index b34cdea3..62cf854f 100644 --- a/src/framework/core/backend.cpp +++ b/src/framework/core/backend.cpp @@ -469,6 +469,26 @@ void write_tensor_f16(const TensorValue & tensor, const std::vector & val write_tensor_f16(tensor, values.data(), values.size()); } +void write_tensor_bf16(const TensorValue & tensor, const float * values, size_t count) { + if (tensor.type != GGML_TYPE_BF16) { + throw std::runtime_error("write_tensor_bf16 requires GGML_TYPE_BF16 tensor"); + } + if (tensor.shape.num_elements() != static_cast(count)) { + throw std::runtime_error( + "write_tensor_bf16 value count does not match tensor shape for tensor '" + + std::string(tensor.tensor != nullptr ? tensor.tensor->name : "") + + "': expected " + std::to_string(tensor.shape.num_elements()) + + ", got " + std::to_string(count)); + } + std::vector bf16_values(count); + ggml_fp32_to_bf16_row(values, bf16_values.data(), static_cast(count)); + ggml_backend_tensor_set(tensor.tensor, bf16_values.data(), 0, count * sizeof(ggml_bf16_t)); +} + +void write_tensor_bf16(const TensorValue & tensor, const std::vector & values) { + write_tensor_bf16(tensor, values.data(), values.size()); +} + void write_tensor_i32(const TensorValue & tensor, const int32_t * values, size_t count) { if (tensor.type != GGML_TYPE_I32) { throw std::runtime_error("write_tensor_i32 requires GGML_TYPE_I32 tensor"); @@ -527,13 +547,37 @@ void read_tensor_f32_into(const ggml_tensor * tensor, std::vector & value read_tensor_typed_into(tensor, GGML_TYPE_F32, values); } -std::vector read_tensor_f32(const ggml_tensor * tensor) { - return read_tensor_typed(tensor, GGML_TYPE_F32); -} - -void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values) { - read_tensor_typed_into(tensor, GGML_TYPE_I32, values); -} +std::vector read_tensor_f32(const ggml_tensor * tensor) { + return read_tensor_typed(tensor, GGML_TYPE_F32); +} + +void read_tensor_f16_into(const ggml_tensor * tensor, std::vector & values) { + const auto fp16_values = read_tensor_typed(tensor, GGML_TYPE_F16); + values.resize(fp16_values.size()); + ggml_fp16_to_fp32_row(fp16_values.data(), values.data(), static_cast(values.size())); +} + +std::vector read_tensor_f16(const ggml_tensor * tensor) { + std::vector values; + read_tensor_f16_into(tensor, values); + return values; +} + +void read_tensor_bf16_into(const ggml_tensor * tensor, std::vector & values) { + const auto bf16_values = read_tensor_typed(tensor, GGML_TYPE_BF16); + values.resize(bf16_values.size()); + ggml_bf16_to_fp32_row(bf16_values.data(), values.data(), static_cast(values.size())); +} + +std::vector read_tensor_bf16(const ggml_tensor * tensor) { + std::vector values; + read_tensor_bf16_into(tensor, values); + return values; +} + +void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values) { + read_tensor_typed_into(tensor, GGML_TYPE_I32, values); +} std::vector read_tensor_i32(const ggml_tensor * tensor) { return read_tensor_typed(tensor, GGML_TYPE_I32); diff --git a/src/framework/modules/attention/qwen_causal_decoder.cpp b/src/framework/modules/attention/qwen_causal_decoder.cpp index 229243c1..63c5fcd8 100644 --- a/src/framework/modules/attention/qwen_causal_decoder.cpp +++ b/src/framework/modules/attention/qwen_causal_decoder.cpp @@ -175,6 +175,36 @@ std::vector qwen_causal_prefill_mask_values(int64_t batch_size, int return out; } +std::vector qwen_causal_suffix_mask_values( + int64_t batch_size, + int64_t query_steps, + int64_t prefix_steps) { + if (batch_size <= 0) { + throw std::runtime_error("qwen_causal_suffix_mask_values requires positive batch size"); + } + validate_steps(query_steps, "qwen_causal_suffix_mask_values"); + if (prefix_steps < 0) { + throw std::runtime_error("qwen_causal_suffix_mask_values requires non-negative prefix steps"); + } + const int64_t key_steps = prefix_steps + query_steps; + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + std::vector one(static_cast(query_steps * key_steps), masked); + for (int64_t row = 0; row < query_steps; ++row) { + const size_t row_offset = static_cast(row * key_steps); + std::fill_n( + one.begin() + static_cast(row_offset), + prefix_steps + row + 1, + visible); + } + std::vector out; + out.reserve(static_cast(batch_size) * one.size()); + for (int64_t batch = 0; batch < batch_size; ++batch) { + out.insert(out.end(), one.begin(), one.end()); + } + return out; +} + void write_qwen_causal_prefill_mask( ggml_tensor * tensor, int64_t batch_size, diff --git a/src/framework/modules/attention/qwen_decoder.cpp b/src/framework/modules/attention/qwen_decoder.cpp index d85f7da4..188fbc11 100644 --- a/src/framework/modules/attention/qwen_decoder.cpp +++ b/src/framework/modules/attention/qwen_decoder.cpp @@ -334,35 +334,124 @@ core::TensorValue build_mlp( const core::TensorValue & input, const QwenDecoderLayerConfig & config, const QwenMLPWeights & weights) { - auto gate = LinearModule( - { - config.hidden_size, - config.intermediate_size, - weights.gate_proj.bias.has_value(), - config.projection_precision, - }) - .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { - gate = activation_cast(ctx, gate, config.activation_cast); - } - gate = SiluModule{}.build(ctx, gate); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { - gate = activation_cast(ctx, gate, config.activation_cast); - } - auto up = LinearModule( - { - config.hidden_size, - config.intermediate_size, - weights.up_proj.bias.has_value(), - config.projection_precision, - }) - .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { - up = activation_cast(ctx, up, config.activation_cast); + if (config.runtime.mlp.mode == QwenDecoderMLPMode::Exact) { + auto gate = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.gate_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + gate = activation_cast(ctx, gate, config.activation_cast); + } + gate = SiluModule{}.build(ctx, gate); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { + gate = activation_cast(ctx, gate, config.activation_cast); + } + auto up = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.up_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + up = activation_cast(ctx, up, config.activation_cast); + } + auto gated = MulModule{}.build(ctx, gate, up); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { + gated = activation_cast(ctx, gated, config.activation_cast); + } + auto down = LinearModule( + { + config.intermediate_size, + config.hidden_size, + weights.down_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, gated, require_linear(weights.down_proj, false, "QwenMLPWeights.down_proj")); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + down = activation_cast(ctx, down, config.activation_cast); + } + return down; } - auto gated = MulModule{}.build(ctx, gate, up); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { - gated = activation_cast(ctx, gated, config.activation_cast); + + core::TensorValue gate; + core::TensorValue up; + std::optional packed_gate_up; + const auto mlp_mode = config.runtime.mlp.mode; + if (mlp_mode == QwenDecoderMLPMode::PackedGateUp) { + if (!weights.gate_up_proj.has_value()) { + throw std::runtime_error("QwenMLPWeights.gate_up_proj is required for packed gate/up mode"); + } + auto gate_up = LinearModule( + { + config.hidden_size, + config.intermediate_size * 2, + weights.gate_up_proj->bias.has_value(), + config.projection_precision, + }) + .build( + ctx, + input, + require_linear(*weights.gate_up_proj, false, "QwenMLPWeights.gate_up_proj")); + packed_gate_up = gate_up; + gate = SliceModule({2, 0, config.intermediate_size}).build(ctx, gate_up); + up = SliceModule({2, config.intermediate_size, config.intermediate_size}).build(ctx, gate_up); + } else { + gate = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.gate_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); + up = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.up_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); + } + const bool can_use_fused_swiglu = + !config.activation_cast.enabled || + (!config.activation_cast.after_mlp_projection && + !config.activation_cast.after_mlp_silu && + !config.activation_cast.after_mlp_mul); + core::TensorValue gated; + if (can_use_fused_swiglu && mlp_mode == QwenDecoderMLPMode::PackedGateUp) { + gated = core::wrap_tensor( + ggml_swiglu(ctx.ggml, packed_gate_up->tensor), + core::TensorShape::from_dims({ + input.shape.dims[0], + input.shape.dims[1], + config.intermediate_size, + }), + packed_gate_up->type); + } else if (can_use_fused_swiglu && mlp_mode == QwenDecoderMLPMode::FusedSwiGLU) { + gated = core::wrap_tensor( + ggml_swiglu_split(ctx.ggml, gate.tensor, up.tensor), + gate.shape, + gate.type); + } else { + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + gate = activation_cast(ctx, gate, config.activation_cast); + up = activation_cast(ctx, up, config.activation_cast); + } + gate = SiluModule{}.build(ctx, gate); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { + gate = activation_cast(ctx, gate, config.activation_cast); + } + gated = MulModule{}.build(ctx, gate, up); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { + gated = activation_cast(ctx, gated, config.activation_cast); + } } auto down = LinearModule( { @@ -445,8 +534,33 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( v = core::ensure_backend_addressable_layout(ctx, v); auto q_heads = TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); - auto all_k = prefix_key.has_value() ? ConcatModule({1}).build(ctx, *prefix_key, k) : k; - auto all_v = prefix_value.has_value() ? ConcatModule({1}).build(ctx, *prefix_value, v) : v; + const bool use_prefix_flash = + prefix_key.has_value() && + config_.runtime.attention.prefix_mode == QwenDecoderPrefixAttentionMode::FlashWithPrefix && + config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV; + core::TensorValue all_k = k; + core::TensorValue all_v = v; + if (use_prefix_flash) { + auto attention_prefix_key = prefix_key; + auto attention_prefix_value = prefix_value; + if (attention_prefix_key->type != k.type) { + attention_prefix_key = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_key->tensor, k.type), + attention_prefix_key->shape, + k.type); + } + if (attention_prefix_value->type != v.type) { + attention_prefix_value = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_value->tensor, v.type), + attention_prefix_value->shape, + v.type); + } + all_k = ConcatModule({1}).build(ctx, *attention_prefix_key, k); + all_v = ConcatModule({1}).build(ctx, *attention_prefix_value, v); + } else if (prefix_key.has_value()) { + all_k = ConcatModule({1}).build(ctx, *prefix_key, k); + all_v = ConcatModule({1}).build(ctx, *prefix_value, v); + } core::TensorValue context; if (!prefix_key.has_value() && attention_mask.has_value() && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { @@ -461,8 +575,10 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( dim, *attention_mask, config_.attention_precision); - } else if (!prefix_key.has_value() && attention_mask.has_value() && - config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) { + } else if (attention_mask.has_value() && + ((!prefix_key.has_value() && + config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) || + use_prefix_flash)) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); auto v_heads = TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v); @@ -588,7 +704,11 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( if (!cache_slot.has_value()) { throw std::runtime_error("Qwen decoder direct static-cache update requires cache_slot"); } - const FastKVSetRowsModule set_rows; + const FastKVSetRowsModule set_rows({ + config_.runtime.static_cache.set_rows_mode == QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized + ? FastKVSetRowsMode::BackendViewOptimized + : FastKVSetRowsMode::Exact, + }); attention_key_cache = set_rows.build(ctx, cache_key, k, *cache_slot); attention_value_cache = set_rows.build(ctx, cache_value, v, *cache_slot); if (config_.activation_cast.enabled && config_.activation_cast.after_static_cache_update) { @@ -629,8 +749,6 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery) { k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); - k_heads = core::wrap_tensor(ggml_cont(ctx.ggml, k_heads.tensor), k_heads.shape, k_heads.type); - v_heads = core::wrap_tensor(ggml_cont(ctx.ggml, v_heads.tensor), v_heads.shape, v_heads.type); context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { context = flash_attention_from_grouped_heads_view_kv( diff --git a/src/framework/modules/optimizations/fast_kv_modules.cpp b/src/framework/modules/optimizations/fast_kv_modules.cpp index ce9e08eb..aa879b0e 100644 --- a/src/framework/modules/optimizations/fast_kv_modules.cpp +++ b/src/framework/modules/optimizations/fast_kv_modules.cpp @@ -30,6 +30,8 @@ const core::ModuleSchema kFastKVSetRowsSchema = { } // namespace +FastKVSetRowsModule::FastKVSetRowsModule(FastKVSetRowsConfig config) : config_(config) {} + const core::ModuleSchema & FastKVSetRowsModule::schema() const noexcept { return static_schema(); } @@ -51,8 +53,14 @@ core::TensorValue FastKVSetRowsModule::build( if (row_index.shape.rank != 1 || (row_index.shape.dims[0] != 1 && row_index.shape.dims[0] != batch)) { throw std::runtime_error("FastKVSetRowsModule row_index must have shape {1} or {batch}"); } - if (cache.type != GGML_TYPE_F32 || row.type != GGML_TYPE_F32) { - throw std::runtime_error("FastKVSetRowsModule requires f32 cache and row tensors"); + const bool optimized = config_.mode == FastKVSetRowsMode::BackendViewOptimized; + if (((!optimized && cache.type != GGML_TYPE_F32) || + (optimized && cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16 && cache.type != GGML_TYPE_BF16)) || + row.type != GGML_TYPE_F32) { + throw std::runtime_error( + optimized + ? "FastKVSetRowsModule requires an f32/f16/bf16 cache and an f32 row tensor" + : "FastKVSetRowsModule requires f32 cache and row tensors"); } if (row_index.type != GGML_TYPE_I32 && row_index.type != GGML_TYPE_I64) { throw std::runtime_error("FastKVSetRowsModule requires i32 or i64 row_index tensor"); @@ -69,16 +77,44 @@ core::TensorValue FastKVSetRowsModule::build( } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({1, row_elems})); + auto flat_row = optimized + ? core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + 1, + contiguous_row.tensor->nb[2], + 0), + core::TensorShape::from_dims({1, row_elems}), + row.type) + : core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({1, row_elems})); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); + if (optimized) { + updated->src[2] = cache.tensor; + } auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({batch * steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({batch, row_elems})); + auto flat_row = optimized + ? core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + batch, + contiguous_row.tensor->nb[3], + 0), + core::TensorShape::from_dims({batch, row_elems}), + row.type) + : core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({batch, row_elems})); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); + if (optimized) { + updated->src[2] = cache.tensor; + } auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 24f9925c..555cfd44 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -9,13 +9,76 @@ namespace engine::runtime { +namespace { + +void validate_cache_tensor(const core::TensorValue & tensor, const TransformerKVCacheOptions & options) { + if (tensor.type == GGML_TYPE_F32) { + return; + } + if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { + return; + } + if (options.allow_bf16_storage && tensor.type == GGML_TYPE_BF16) { + return; + } + throw std::runtime_error( + options.allow_f16_storage || options.allow_bf16_storage + ? "TransformerKVCache supports only f32/f16/bf16 cache tensors when enabled" + : "TransformerKVCache requires f32 cache tensors"); +} + +void write_cache_tensor( + const core::TensorValue & tensor, + const std::vector & values, + const TransformerKVCacheOptions & options) { + validate_cache_tensor(tensor, options); + if (tensor.type == GGML_TYPE_F32) { + core::write_tensor_f32(tensor, values); + return; + } + if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { + core::write_tensor_f16(tensor, values); + return; + } + if (options.allow_bf16_storage && tensor.type == GGML_TYPE_BF16) { + core::write_tensor_bf16(tensor, values); + return; + } + throw std::runtime_error("TransformerKVCache requires f32 cache tensors"); +} + +std::vector read_cache_tensor(const core::TensorValue & tensor, const TransformerKVCacheOptions & options) { + validate_cache_tensor(tensor, options); + if (tensor.type == GGML_TYPE_F32) { + return core::read_tensor_f32(tensor.tensor); + } + if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { + return core::read_tensor_f16(tensor.tensor); + } + if (options.allow_bf16_storage && tensor.type == GGML_TYPE_BF16) { + return core::read_tensor_bf16(tensor.tensor); + } + throw std::runtime_error("TransformerKVCache requires f32 cache tensors"); +} + +} // namespace + TransformerKVCache::TransformerKVCache( int64_t cache_steps, int64_t step_elems, std::vector keys, std::vector values) + : TransformerKVCache(cache_steps, step_elems, std::move(keys), std::move(values), {}) {} + +TransformerKVCache::TransformerKVCache( + int64_t cache_steps, + int64_t step_elems, + std::vector keys, + std::vector values, + TransformerKVCacheOptions options) : cache_steps_(std::max(0, cache_steps)), - step_elems_(std::max(0, step_elems)) { + step_elems_(std::max(0, step_elems)), + options_(options) { if (step_elems_ <= 0) { throw std::runtime_error("TransformerKVCache requires positive step_elems"); } @@ -25,6 +88,8 @@ TransformerKVCache::TransformerKVCache( const size_t cache_elems = static_cast(cache_steps_ * step_elems_); layers_.reserve(keys.size()); for (size_t layer = 0; layer < keys.size(); ++layer) { + validate_cache_tensor(keys[layer], options_); + validate_cache_tensor(values[layer], options_); layers_.push_back(LayerCache{ std::move(keys[layer]), std::move(values[layer]), @@ -68,8 +133,8 @@ void TransformerKVCache::import_state(const TransformerKVState & state) { std::copy(source.key.begin(), source.key.end(), cache.import_key_scratch.begin()); std::copy(source.value.begin(), source.value.end(), cache.import_value_scratch.begin()); } - core::write_tensor_f32(cache.key_tensor, cache.import_key_scratch); - core::write_tensor_f32(cache.value_tensor, cache.import_value_scratch); + write_cache_tensor(cache.key_tensor, cache.import_key_scratch, options_); + write_cache_tensor(cache.value_tensor, cache.import_value_scratch, options_); } } } @@ -85,8 +150,8 @@ TransformerKVState TransformerKVCache::export_state() const { if (keep_elems == 0) { continue; } - const auto key_values = core::read_tensor_f32(layers_[layer].key_tensor.tensor); - const auto value_values = core::read_tensor_f32(layers_[layer].value_tensor.tensor); + const auto key_values = read_cache_tensor(layers_[layer].key_tensor, options_); + const auto value_values = read_cache_tensor(layers_[layer].value_tensor, options_); out.key.assign(key_values.begin(), key_values.begin() + static_cast(keep_elems)); out.value.assign(value_values.begin(), value_values.begin() + static_cast(keep_elems)); } @@ -141,11 +206,11 @@ void TransformerKVCache::trace_log_state(const std::string & name, int64_t num_h return; } const size_t keep_elems = static_cast(valid_steps_ * step_elems_); - const auto first_key = core::read_tensor_f32(layers_.front().key_tensor.tensor); + const auto first_key = read_cache_tensor(layers_.front().key_tensor, options_); std::vector first_key_keep(first_key.begin(), first_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer0.key", {1, valid_steps_, num_heads, head_dim}, first_key_keep); if (layers_.size() > 1) { - const auto last_key = core::read_tensor_f32(layers_.back().key_tensor.tensor); + const auto last_key = read_cache_tensor(layers_.back().key_tensor, options_); std::vector last_key_keep(last_key.begin(), last_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer_last.key", {1, valid_steps_, num_heads, head_dim}, last_key_keep); } @@ -158,7 +223,8 @@ core::TensorValue view_transformer_kv_cache_steps( int64_t steps, int64_t heads, int64_t head_dim, - const char * label) { + const char * label, + ggml_type view_type) { if (start < 0 || steps <= 0 || start + steps > cache.shape.dims[1]) { throw std::runtime_error(std::string(label) + " cache view range is invalid"); } @@ -175,7 +241,7 @@ core::TensorValue view_transformer_kv_cache_steps( cache.tensor->nb[3], static_cast(start) * cache.tensor->nb[2]), core::TensorShape::from_dims({1, steps, heads, head_dim}), - GGML_TYPE_F32); + view_type); } } // namespace engine::runtime diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index ba6f560a..571f3de5 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -4,18 +4,14 @@ #include "engine/framework/assets/model_package.h" #include "engine/framework/io/config.h" #include "engine/framework/io/filesystem.h" -// Parked loaders (sources not in this release tree). When commenting these out, -// also mark matching ModelPackage entries UnsupportedSource — see -// docs/maintainers/loader_and_catalog.md and tools/check_loader_catalog_sync.py. -// #include "engine/models/higgs_tts/loader.h" -// #include "engine/models/kokoro_tts/loader.h" -// #include "engine/models/parakeet_tdt/loader.h" #include "engine/models/ace_step/loader.h" #include "engine/models/chatterbox/loader.h" #include "engine/models/citrinet_asr/session.h" #include "engine/models/demucs/loader.h" +#include "engine/models/fish_audio/loader.h" #include "engine/models/heartmula/loader.h" #include "engine/models/higgs_audio_stt/loader.h" +#include "engine/models/higgs_audio_tts/loader.h" #include "engine/models/hviske_asr/loader.h" #include "engine/models/index_tts2/loader.h" #include "engine/models/irodori_tts/loader.h" @@ -244,11 +240,6 @@ ModelRegistry make_registry_from_config( ModelRegistry make_default_registry(const std::optional & config_path) { const std::vector> available_loaders = { - // Parked loaders — keep catalog packages UnsupportedSource while these stay commented. - // See docs/maintainers/loader_and_catalog.md. - // engine::models::kokoro_tts::make_kokoro_tts_loader(), - // engine::models::higgs_tts::make_higgs_tts_loader(), - // engine::models::parakeet_tdt::make_parakeet_tdt_loader(), engine::models::ace_step::make_ace_step_loader(), engine::models::demucs::make_htdemucs_loader(), engine::models::roformer::make_mel_band_roformer_loader(), @@ -262,8 +253,10 @@ ModelRegistry make_default_registry(const std::optional & engine::models::vibevoice::make_vibevoice_loader(), engine::models::vibevoice_asr::make_vibevoice_asr_loader(), engine::models::voxtral_realtime::make_voxtral_realtime_loader(), + engine::models::fish_audio::make_fish_audio_loader(), engine::models::heartmula::make_heartmula_loader(), engine::models::higgs_audio_stt::make_higgs_audio_stt_loader(), + engine::models::higgs_audio_tts::make_higgs_audio_tts_loader(), engine::models::hviske_asr::make_hviske_asr_loader(), engine::models::irodori_tts::make_irodori_tts_loader(), engine::models::nemotron_asr::make_nemotron_asr_loader(), diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp new file mode 100644 index 00000000..945495fc --- /dev/null +++ b/src/models/fish_audio/ar.cpp @@ -0,0 +1,1560 @@ +#include "engine/models/fish_audio/ar.h" + +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/attention/qwen_causal_decoder.h" +#include "engine/framework/modules/attention/qwen_decoder.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/torch_random.h" + +#include "../common/constant_tensor_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +namespace binding = engine::modules::binding; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kRasWindow = 10; +constexpr float kRasHighTemperature = 1.0F; +constexpr float kRasHighTopP = 0.9F; + +struct FishARProfile { + double graph_build_prefill_ms = 0.0; + double graph_build_step_ms = 0.0; + double graph_build_fast_ms = 0.0; + double slow_embedding_ms = 0.0; + double fast_embedding_ms = 0.0; + double prefill_input_upload_ms = 0.0; + double prefill_graph_ms = 0.0; + double prefill_output_read_ms = 0.0; + double step_input_upload_ms = 0.0; + double step_mask_upload_ms = 0.0; + double step_graph_ms = 0.0; + double step_output_read_ms = 0.0; + double fast_input_upload_ms = 0.0; + double fast_mask_upload_ms = 0.0; + double fast_graph_ms = 0.0; + double fast_output_read_ms = 0.0; + double sample_bias_ms = 0.0; + double sample_main_ms = 0.0; + double sample_high_ms = 0.0; + double sample_fast_ms = 0.0; + int64_t prefill_runs = 0; + int64_t step_runs = 0; + int64_t fast_runs = 0; + int64_t generated_frames = 0; +}; + +struct SampleCandidate { + int32_t index = 0; + float probability = 0.0F; +}; + +struct SampleDistribution { + size_t source_size = 0; + std::vector candidates; +}; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct FishLayerWeights { + assets::TensorDataF32 input_norm; + core::TensorValue qkv_proj; + core::TensorValue o_proj; + std::optional q_norm; + std::optional k_norm; + assets::TensorDataF32 post_norm; + core::TensorValue gate_up_proj; + core::TensorValue down_proj; +}; + +struct FishARWeights { + std::shared_ptr store; + assets::TensorData text_embedding_host; + assets::TensorData codebook_embedding_host; + assets::TensorData fast_embedding_host; + core::TensorValue text_embedding; + std::vector slow_layers; + assets::TensorDataF32 slow_norm; + std::vector fast_layers; + assets::TensorDataF32 fast_norm; + core::TensorValue fast_output; +}; + +struct SlowForwardOutput { + std::vector logits; + std::vector hidden; +}; + +struct SlowPrefillOutput { + SlowForwardOutput forward; +}; + +struct FishPrefillCacheTarget { + std::vector keys; + std::vector values; +}; + +modules::QwenDecoderActivationCastPolicy fish_activation_cast_policy(core::BackendType backend_type) { + modules::QwenDecoderActivationCastPolicy policy; + if (backend_type == core::BackendType::Vulkan) { + return policy; + } + policy.enabled = true; + policy.type = GGML_TYPE_BF16; + policy.after_input_norm = true; + policy.after_qkv_projection = true; + policy.after_qk_norm = true; + policy.after_rope = true; + policy.after_static_cache_update = true; + policy.after_attention = true; + policy.after_attention_output = true; + policy.after_residual = true; + policy.after_ffn_norm = true; + policy.after_mlp_projection = true; + policy.after_mlp_silu = true; + policy.after_mlp_mul = true; + policy.after_output = true; + return policy; +} + +modules::QwenCausalDecoderConfig make_slow_decoder_config( + const FishAudioTextConfig & config, + core::BackendType backend_type) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.dim; + out.stack.num_attention_heads = config.n_head; + out.stack.num_key_value_heads = config.n_local_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.n_layer; + out.stack.rms_norm_eps = config.norm_eps; + out.stack.rope_theta = config.rope_base; + out.stack.rope_type = GGML_ROPE_TYPE_NORMAL; + out.stack.attention_precision = GGML_PREC_F32; + out.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; + out.stack.use_qk_norm = config.attention_qk_norm; + out.stack.activation_cast = fish_activation_cast_policy(backend_type); + out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + out.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.lm_head_precision = GGML_PREC_F32; + return out; +} + +modules::QwenCausalDecoderConfig make_fast_decoder_config( + const FishAudioFastConfig & config, + core::BackendType backend_type) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.dim; + out.stack.num_attention_heads = config.n_head; + out.stack.num_key_value_heads = config.n_local_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.n_layer; + out.stack.rms_norm_eps = config.norm_eps; + out.stack.rope_theta = config.rope_base; + out.stack.rope_type = GGML_ROPE_TYPE_NORMAL; + out.stack.attention_precision = GGML_PREC_F32; + out.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; + out.stack.use_qk_norm = config.attention_qk_norm; + out.stack.activation_cast = fish_activation_cast_policy(backend_type); + out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + out.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.lm_head_precision = GGML_PREC_F32; + return out; +} + +modules::QwenDecoderLayerWeights bind_layer( + common::ConstantTensorCache & constants, + const FishLayerWeights & weights, + bool use_qk_norm) { + modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_data(constants, weights.input_norm); + out.self_attention.qkv_weight = weights.qkv_proj; + out.self_attention.out_weight = weights.o_proj; + if (use_qk_norm) { + if (!weights.q_norm.has_value() || !weights.k_norm.has_value()) { + throw std::runtime_error("Fish Audio q/k norm weights are missing"); + } + out.q_norm = binding::norm_data(constants, *weights.q_norm); + out.k_norm = binding::norm_data(constants, *weights.k_norm); + } + out.post_norm = binding::norm_data(constants, weights.post_norm); + out.mlp.gate_up_proj = binding::linear_data(constants, weights.gate_up_proj); + out.mlp.down_proj = binding::linear_data(constants, weights.down_proj); + return out; +} + +modules::QwenCausalDecoderWeights bind_slow_weights( + common::ConstantTensorCache & constants, + const FishARWeights & weights, + const FishAudioTextConfig & config) { + modules::QwenCausalDecoderWeights out; + out.stack.layers.reserve(weights.slow_layers.size()); + for (const auto & layer : weights.slow_layers) { + out.stack.layers.push_back(bind_layer(constants, layer, config.attention_qk_norm)); + } + out.final_norm = binding::norm_data(constants, weights.slow_norm); + out.lm_head = binding::linear_data(constants, weights.text_embedding); + return out; +} + +modules::QwenDecoderLayerWeights bind_fast_layer( + common::ConstantTensorCache & constants, + const FishLayerWeights & weights, + const FishAudioFastConfig & config) { + return bind_layer(constants, weights, config.attention_qk_norm); +} + +void copy_tensor_row_to_f32(const assets::TensorData & table, int64_t row, int64_t width, float * out) { + if (row < 0 || width <= 0 || table.shape.rank != 2 || table.shape.dims[1] != width || + row >= table.shape.dims[0]) { + throw std::runtime_error("Fish Audio embedding row lookup shape mismatch"); + } + const size_t row_bytes = ggml_row_size(table.type, width); + const size_t offset = static_cast(row) * row_bytes; + if (offset + row_bytes > table.bytes.size()) { + throw std::runtime_error("Fish Audio embedding row lookup exceeded tensor storage"); + } + const auto * bytes = reinterpret_cast(table.bytes.data()) + offset; + if (table.type == GGML_TYPE_F32) { + std::memcpy(out, bytes, static_cast(width) * sizeof(float)); + } else if (table.type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row(reinterpret_cast(bytes), out, width); + } else if (table.type == GGML_TYPE_BF16) { + ggml_bf16_to_fp32_row(reinterpret_cast(bytes), out, width); + } else { + throw std::runtime_error("Fish Audio host embedding lookup requires f32/f16/bf16 native embeddings"); + } +} + +std::vector lookup_row(const assets::TensorData & table, int64_t row, int64_t width) { + std::vector out(static_cast(width), 0.0F); + copy_tensor_row_to_f32(table, row, width, out.data()); + return out; +} + +void add_row(const assets::TensorData & table, int64_t row, int64_t width, std::vector & out) { + std::vector tmp(static_cast(width), 0.0F); + copy_tensor_row_to_f32(table, row, width, tmp.data()); + for (int64_t i = 0; i < width; ++i) { + out[static_cast(i)] += tmp[static_cast(i)]; + } +} + +bool is_semantic_token(const FishAudioConfig & config, int32_t token) { + return token >= config.semantic_start_token_id && token <= config.semantic_end_token_id; +} + +std::vector build_slow_embeddings( + const FishAudioConfig & config, + const FishARWeights & weights, + const int32_t * matrix, + int64_t steps) { + const int64_t rows = config.fast.num_codebooks + 1; + const int64_t hidden = config.text.dim; + std::vector out(static_cast(steps * hidden), 0.0F); + const float semantic_scale = 1.0F / std::sqrt(static_cast(rows)); + for (int64_t step = 0; step < steps; ++step) { + const int32_t token = matrix[step]; + auto row = lookup_row(weights.text_embedding_host, token, hidden); + if (is_semantic_token(config, token)) { + for (int64_t codebook = 0; codebook < config.fast.num_codebooks; ++codebook) { + const int32_t code = matrix[(codebook + 1) * steps + step]; + add_row( + weights.codebook_embedding_host, + codebook * config.fast.vocab_size + code, + hidden, + row); + } + for (float & value : row) { + value *= semantic_scale; + } + } + std::copy(row.begin(), row.end(), out.begin() + static_cast(step * hidden)); + } + return out; +} + +std::vector build_slow_embedding_for_frame( + const FishAudioConfig & config, + const FishARWeights & weights, + const std::vector & frame) { + if (static_cast(frame.size()) != config.fast.num_codebooks + 1) { + throw std::runtime_error("Fish Audio frame size mismatch"); + } + return build_slow_embeddings(config, weights, frame.data(), 1); +} + +std::vector build_fast_embedding( + const FishAudioConfig & config, + const FishARWeights & weights, + int32_t code) { + return lookup_row(weights.fast_embedding_host, code, config.fast.dim); +} + +FishLayerWeights load_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden, + int64_t heads, + int64_t kv_heads, + int64_t head_dim, + int64_t intermediate, + bool qk_norm, + assets::TensorStorageType storage_type) { + FishLayerWeights w; + w.input_norm = source.require_f32_tensor(prefix + ".attention_norm.weight", {hidden}); + { + const auto q = source.require_tensor(prefix + ".attention.q_proj.weight", storage_type, {heads * head_dim, hidden}); + const auto k = source.require_tensor(prefix + ".attention.k_proj.weight", storage_type, {kv_heads * head_dim, hidden}); + const auto v = source.require_tensor(prefix + ".attention.v_proj.weight", storage_type, {kv_heads * head_dim, hidden}); + if (q.type != k.type || q.type != v.type) { + throw std::runtime_error("Fish Audio packed QKV weights require matching storage types"); + } + std::vector packed; + packed.reserve(q.bytes.size() + k.bytes.size() + v.bytes.size()); + packed.insert(packed.end(), q.bytes.begin(), q.bytes.end()); + packed.insert(packed.end(), k.bytes.begin(), k.bytes.end()); + packed.insert(packed.end(), v.bytes.begin(), v.bytes.end()); + w.qkv_proj = store.make_tensor( + core::TensorShape::from_dims({(heads + 2 * kv_heads) * head_dim, hidden}), + q.type, + packed.data(), + packed.size()); + } + w.o_proj = store.load_tensor(source, prefix + ".attention.wo.weight", storage_type, {hidden, heads * head_dim}); + if (qk_norm) { + w.q_norm = source.require_f32_tensor(prefix + ".attention.q_norm.weight", {head_dim}); + w.k_norm = source.require_f32_tensor(prefix + ".attention.k_norm.weight", {head_dim}); + } + w.post_norm = source.require_f32_tensor(prefix + ".ffn_norm.weight", {hidden}); + w.down_proj = store.load_tensor(source, prefix + ".feed_forward.w2.weight", storage_type, {hidden, intermediate}); + { + const auto gate = source.require_tensor(prefix + ".feed_forward.w1.weight", storage_type, {intermediate, hidden}); + const auto up = source.require_tensor(prefix + ".feed_forward.w3.weight", storage_type, {intermediate, hidden}); + if (gate.type != up.type) { + throw std::runtime_error("Fish Audio packed gate/up weights require matching storage types"); + } + std::vector packed; + packed.reserve(gate.bytes.size() + up.bytes.size()); + packed.insert(packed.end(), gate.bytes.begin(), gate.bytes.end()); + packed.insert(packed.end(), up.bytes.begin(), up.bytes.end()); + w.gate_up_proj = store.make_tensor( + core::TensorShape::from_dims({intermediate * 2, hidden}), + gate.type, + packed.data(), + packed.size()); + } + return w; +} + +FishARWeights load_ar_weights( + const FishAudioAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + const auto & source = *assets.model_weights; + const auto & config = assets.config; + FishARWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "fish_audio.ar.weights", + weight_context_bytes); + weights.text_embedding_host = source.require_tensor( + "embeddings.weight", + assets::TensorStorageType::Native, + {config.text.vocab_size, config.text.dim}); + weights.codebook_embedding_host = source.require_tensor( + "codebook_embeddings.weight", + assets::TensorStorageType::Native, + {config.fast.vocab_size * config.fast.num_codebooks, config.text.dim}); + weights.fast_embedding_host = source.require_tensor( + "fast_embeddings.weight", + assets::TensorStorageType::Native, + {config.fast.vocab_size, config.fast.dim}); + weights.text_embedding = weights.store->load_tensor( + source, + "embeddings.weight", + storage_type, + {config.text.vocab_size, config.text.dim}); + weights.slow_layers.reserve(static_cast(config.text.n_layer)); + for (int64_t i = 0; i < config.text.n_layer; ++i) { + weights.slow_layers.push_back(load_layer( + *weights.store, + source, + "layers." + std::to_string(i), + config.text.dim, + config.text.n_head, + config.text.n_local_heads, + config.text.head_dim, + config.text.intermediate_size, + config.text.attention_qk_norm, + storage_type)); + } + weights.slow_norm = source.require_f32_tensor("norm.weight", {config.text.dim}); + weights.fast_layers.reserve(static_cast(config.fast.n_layer)); + for (int64_t i = 0; i < config.fast.n_layer; ++i) { + weights.fast_layers.push_back(load_layer( + *weights.store, + source, + "fast_layers." + std::to_string(i), + config.fast.dim, + config.fast.n_head, + config.fast.n_local_heads, + config.fast.head_dim, + config.fast.intermediate_size, + config.fast.attention_qk_norm, + storage_type)); + } + weights.fast_norm = source.require_f32_tensor("fast_norm.weight", {config.fast.dim}); + weights.fast_output = weights.store->load_tensor( + source, + "fast_output.weight", + storage_type, + {config.fast.vocab_size, config.fast.dim}); + weights.store->upload(); + return weights; +} + +struct SampleState { + uint64_t seed = 0; + uint64_t call_index = 0; + std::mt19937 rng; + std::vector previous_main; +}; + +SampleDistribution logits_to_distribution( + const std::vector & logits, + float temperature, + float top_p, + int top_k) { + if (logits.empty()) { + throw std::runtime_error("Fish Audio sampling requires non-empty logits"); + } + std::vector order; + order.reserve(logits.size()); + float max_logit = -std::numeric_limits::infinity(); + for (size_t i = 0; i < logits.size(); ++i) { + const float logit = logits[i]; + if (!std::isfinite(logit)) { + continue; + } + order.push_back(static_cast(i)); + max_logit = std::max(max_logit, logit); + } + double denom = 0.0; + for (const int32_t index : order) { + denom += std::exp(logits[static_cast(index)] - max_logit); + } + if (denom <= 0.0) { + throw std::runtime_error("Fish Audio sampling logits produced zero probability mass"); + } + const size_t candidate_count = std::min(order.size(), static_cast(std::max(top_k, 1))); + const auto by_logit_desc = [&](int32_t lhs, int32_t rhs) { + return logits[static_cast(lhs)] > logits[static_cast(rhs)]; + }; + if (candidate_count < order.size()) { + std::partial_sort(order.begin(), order.begin() + static_cast(candidate_count), order.end(), by_logit_desc); + order.resize(candidate_count); + } else { + std::sort(order.begin(), order.end(), by_logit_desc); + } + double cumulative = 0.0; + std::vector kept; + kept.reserve(candidate_count); + for (size_t i = 0; i < order.size(); ++i) { + const int32_t index = order[i]; + const float logit = logits[static_cast(index)]; + const float prob = static_cast(std::exp(logit - max_logit) / denom); + cumulative += prob; + const bool remove = cumulative > static_cast(top_p) && i != 0; + if (!remove) { + kept.push_back({index, 0.0F}); + } + } + float filtered_max = -std::numeric_limits::infinity(); + const float temperature_scale = std::max(temperature, 1.0e-5F); + for (const auto & candidate : kept) { + filtered_max = std::max(filtered_max, logits[static_cast(candidate.index)] / temperature_scale); + } + double filtered_denom = 0.0; + for (auto & candidate : kept) { + candidate.probability = + std::exp(logits[static_cast(candidate.index)] / temperature_scale - filtered_max); + filtered_denom += candidate.probability; + } + if (filtered_denom <= 0.0) { + throw std::runtime_error("Fish Audio sampling filter produced zero probability mass"); + } + for (auto & candidate : kept) { + candidate.probability = static_cast(static_cast(candidate.probability) / filtered_denom); + } + return {logits.size(), std::move(kept)}; +} + +int32_t sample_from_logits( + const std::vector & logits, + float temperature, + float top_p, + int top_k, + SampleState & state, + const sampling::TorchCudaSamplingPolicy & policy) { + const auto distribution = logits_to_distribution(logits, temperature, top_p, top_k); + const uint64_t call_index = state.call_index++; + if (!policy.cuda_fast_path) { + std::vector weights; + weights.reserve(distribution.candidates.size()); + for (const auto & candidate : distribution.candidates) { + weights.push_back(static_cast(std::max(candidate.probability, 0.0F))); + } + std::discrete_distribution sampler(weights.begin(), weights.end()); + return distribution.candidates[sampler(state.rng)].index; + } + int32_t best = 0; + double best_score = -std::numeric_limits::infinity(); + for (const auto & candidate : distribution.candidates) { + if (!(candidate.probability > 0.0F)) { + continue; + } + const float exponential = sampling::torch_cuda_tensor_iterator_exponential_element( + state.seed, + static_cast(distribution.source_size), + static_cast(candidate.index), + call_index, + policy.multiprocessor_count, + policy.max_threads_per_multiprocessor); + const float uniform = std::exp(-exponential); + const float uniform_bf16 = ggml_bf16_to_fp32(ggml_fp32_to_bf16(uniform)); + const float exponential_bf16 = ggml_bf16_to_fp32(ggml_fp32_to_bf16(-std::log(uniform_bf16))); + const double score = static_cast(candidate.probability) / static_cast(exponential_bf16); + if (score > best_score) { + best_score = score; + best = candidate.index; + } + } + return best; +} + +std::vector apply_semantic_bias( + const FishAudioConfig & config, + int32_t im_end_id, + const std::vector & logits) { + std::vector out(logits.size(), -std::numeric_limits::infinity()); + const int64_t begin = std::max(0, config.semantic_start_token_id); + const int64_t end = std::min(static_cast(logits.size()) - 1, config.semantic_end_token_id); + for (int64_t i = begin; i <= end; ++i) { + out[static_cast(i)] = logits[static_cast(i)]; + } + if (im_end_id >= 0 && static_cast(im_end_id) < logits.size()) { + out[static_cast(im_end_id)] = logits[static_cast(im_end_id)]; + } + return out; +} + +core::TensorValue make_fish_causal_mask( + core::ModuleBuildContext &, + common::ConstantTensorCache & constants, + int64_t steps) { + auto values = modules::qwen_causal_prefill_mask_values(1, steps); + return constants.make_tensor( + core::TensorShape::from_dims({1, 1, steps, steps}), + GGML_TYPE_F16, + values.data(), + values.size() * sizeof(ggml_fp16_t)); +} + +struct FishCausalDecoderOutputs { + core::TensorValue hidden; + core::TensorValue logits; + modules::QwenDecoderStackState state; +}; + +FishCausalDecoderOutputs build_fish_causal_decoder( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenCausalDecoderWeights & weights, + const modules::QwenCausalDecoderConfig & config, + bool norm_fastlayer_input) { + auto mask = make_fish_causal_mask(ctx, constants, input.shape.dims[1]); + auto x = input; + modules::QwenDecoderStackState state; + state.layers.reserve(weights.stack.layers.size()); + const auto layer_config = modules::qwen_decoder_layer_config_from_stack(config.stack); + const modules::QwenDecoderLayerModule layer_module(layer_config); + for (const auto & layer : weights.stack.layers) { + auto out = layer_module.build(ctx, x, positions, layer, std::nullopt, std::nullopt, mask); + x = out.output; + auto state_key = core::wrap_tensor(ggml_dup(ctx.ggml, out.key.tensor), out.key.shape, out.key.type); + auto state_value = core::wrap_tensor(ggml_dup(ctx.ggml, out.value.tensor), out.value.shape, out.value.type); + state.layers.push_back({state_key, state_value}); + } + auto hidden_sequence = modules::RMSNormModule({config.stack.hidden_size, config.stack.rms_norm_eps, true, false}) + .build(ctx, x, weights.final_norm); + const int64_t steps = hidden_sequence.shape.dims[1]; + auto fast_hidden_source = norm_fastlayer_input ? hidden_sequence : x; + auto hidden = modules::SliceModule({1, steps - 1, 1}).build(ctx, fast_hidden_source); + auto logits = modules::LinearModule({config.stack.hidden_size, config.logits_size, false, config.lm_head_precision}) + .build(ctx, modules::SliceModule({1, steps - 1, 1}).build(ctx, hidden_sequence), weights.lm_head); + auto hidden_out = core::wrap_tensor(ggml_dup(ctx.ggml, hidden.tensor), hidden.shape, hidden.type); + auto logits_out = core::wrap_tensor(ggml_dup(ctx.ggml, logits.tensor), logits.shape, logits.type); + return {hidden_out, logits_out, std::move(state)}; +} + +struct FishStaticDecoderOutputs { + core::TensorValue hidden; + core::TensorValue logits; + runtime::TransformerKVCache cache; +}; + +FishStaticDecoderOutputs build_fish_static_decoder( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenCausalDecoderWeights & weights, + const modules::QwenCausalDecoderConfig & config, + int64_t cache_steps, + const core::TensorValue & attention_mask, + const core::TensorValue & cache_slot, + std::vector cache_keys, + std::vector cache_values, + bool norm_fastlayer_input) { + if (cache_keys.size() != weights.stack.layers.size() || cache_values.size() != weights.stack.layers.size()) { + throw std::runtime_error("Fish Audio static decoder cache layer count mismatch"); + } + const int64_t step_elems = config.stack.num_key_value_heads * config.stack.head_dim; + auto x = input; + const auto layer_config = modules::qwen_decoder_layer_config_from_stack(config.stack); + const modules::QwenDecoderLayerModule layer_module(layer_config); + for (size_t layer_index = 0; layer_index < weights.stack.layers.size(); ++layer_index) { + auto out = layer_module.build_with_static_cache_tail( + ctx, + graph, + x, + positions, + weights.stack.layers[layer_index], + cache_keys[layer_index], + cache_values[layer_index], + cache_slot, + attention_mask); + x = out.output; + } + auto hidden = modules::RMSNormModule({config.stack.hidden_size, config.stack.rms_norm_eps, true, false}) + .build(ctx, x, weights.final_norm); + const auto logits = modules::LinearModule({ + config.stack.hidden_size, + config.logits_size, + config.use_lm_head_bias, + config.lm_head_precision, + }) + .build(ctx, hidden, weights.lm_head); + auto fast_hidden = norm_fastlayer_input ? hidden : x; + runtime::TransformerKVCacheOptions cache_options; + cache_options.allow_bf16_storage = !cache_keys.empty() && cache_keys.front().type == GGML_TYPE_BF16; + return { + fast_hidden, + logits, + runtime::TransformerKVCache( + cache_steps, + step_elems, + std::move(cache_keys), + std::move(cache_values), + cache_options), + }; +} + +} // namespace + +class FishARWeightsRuntime { +public: + FishARWeightsRuntime( + std::shared_ptr assets, + core::BackendConfig backend_config, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)), + threads_(threads), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("Fish Audio AR weights runtime requires assets"); + } + backend_config.threads = threads_; + backend_ = core::init_backend(backend_config); + backend_type_ = core::backend_type(backend_); + weights_ = std::make_shared( + load_ar_weights(*assets_, backend_, backend_type_, weight_context_bytes, weight_storage_type)); + slow_step_constants_ = std::make_unique( + backend_, + threads_, + "fish_audio.ar.step.constants", + 256ull * 1024ull * 1024ull); + fast_constants_ = std::make_unique( + backend_, + threads_, + "fish_audio.ar.fast.constants", + 256ull * 1024ull * 1024ull); + } + + ~FishARWeightsRuntime() { + fast_constants_.reset(); + slow_step_constants_.reset(); + weights_.reset(); + if (backend_ != nullptr) { + ggml_backend_free(backend_); + } + } + + FishARWeightsRuntime(const FishARWeightsRuntime &) = delete; + FishARWeightsRuntime & operator=(const FishARWeightsRuntime &) = delete; + + const FishAudioAssets & assets() const noexcept { + return *assets_; + } + + const FishARWeights & weights() const noexcept { + return *weights_; + } + + int threads() const noexcept { + return threads_; + } + + size_t graph_arena_bytes() const noexcept { + return graph_arena_bytes_; + } + + ggml_backend_t backend() const noexcept { + return backend_; + } + + core::BackendType backend_type() const noexcept { + return backend_type_; + } + + common::ConstantTensorCache & slow_step_constants() const noexcept { + return *slow_step_constants_; + } + + common::ConstantTensorCache & fast_constants() const noexcept { + return *fast_constants_; + } + +private: + std::shared_ptr assets_; + std::shared_ptr weights_; + int threads_ = 1; + size_t graph_arena_bytes_ = 0; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + std::unique_ptr slow_step_constants_; + std::unique_ptr fast_constants_; +}; + +class FishAudioARRuntime::Impl { +public: + Impl( + std::shared_ptr assets, + core::BackendConfig backend_config, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : runtime_(std::make_shared( + std::move(assets), + backend_config, + threads, + graph_arena_bytes, + weight_context_bytes, + weight_storage_type)), + sampling_policy_(sampling::resolve_torch_cuda_sampling_policy( + runtime_->backend_type(), + backend_config.device, + "fish_audio.ar.cuda_sampling_policy", + "Fish Audio", + sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault)) {} + + ~Impl() { + step_graph_.reset(); + prefill_graph_.reset(); + fast_graph_.reset(); + runtime_.reset(); + } + + FishAudioCodes generate(const FishAudioPrompt & prompt, const FishAudioGenerationOptions & options) { + FishARProfile profile; + const auto & assets = runtime_->assets(); + const auto & weights = runtime_->weights(); + if (prompt.codebook_rows != assets.config.fast.num_codebooks + 1 || + static_cast(prompt.matrix.size()) != prompt.codebook_rows * prompt.steps) { + throw std::runtime_error("Fish Audio AR prompt shape mismatch"); + } + const int64_t max_new_tokens = std::min(options.max_new_tokens, assets.config.text.max_seq_len - prompt.steps); + if (max_new_tokens <= 0) { + throw std::runtime_error("Fish Audio prompt leaves no room for generated tokens"); + } + ensure_step_graph(prompt.steps + max_new_tokens, profile); + ensure_prefill_graph(prompt.steps, profile); + ensure_fast_graph(profile); + SampleState sample; + sample.seed = options.seed; + sample.rng.seed(options.seed); + sample.previous_main.assign(static_cast(kRasWindow), 0); + auto timing_start = Clock::now(); + auto embeddings = build_slow_embeddings(assets.config, weights, prompt.matrix.data(), prompt.steps); + profile.slow_embedding_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + auto prefill = prefill_graph_->run(embeddings, profile); + std::vector generated_frame_major; + generated_frame_major.reserve(static_cast(max_new_tokens * assets.config.fast.num_codebooks)); + auto frame = sample_frame(prefill.forward.logits, prefill.forward.hidden, options, sample, false, profile); + if (frame.front() == im_end_id()) { + log_profile(profile); + return FishAudioCodes{{}, assets.config.fast.num_codebooks, 0}; + } + append_frame(generated_frame_major, frame); + ++profile.generated_frames; + step_graph_->finish_prefill(prompt.steps); + bool ended_by_im_end = false; + for (int64_t step = 1; step < max_new_tokens; ++step) { + timing_start = Clock::now(); + const auto input = build_slow_embedding_for_frame(assets.config, weights, frame); + profile.slow_embedding_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + auto step_out = step_graph_->run(input, profile); + frame = sample_frame(step_out.logits, step_out.hidden, options, sample, true, profile); + if (frame.front() == im_end_id()) { + ended_by_im_end = true; + break; + } + append_frame(generated_frame_major, frame); + ++profile.generated_frames; + } + if (!ended_by_im_end && !generated_frame_major.empty()) { + generated_frame_major.resize(generated_frame_major.size() - static_cast(assets.config.fast.num_codebooks)); + --profile.generated_frames; + } + FishAudioCodes out; + out.codebooks = assets.config.fast.num_codebooks; + out.frames = static_cast(generated_frame_major.size()) / out.codebooks; + out.codes.assign(static_cast(out.codebooks * out.frames), 0); + for (int64_t frame_index = 0; frame_index < out.frames; ++frame_index) { + for (int64_t codebook = 0; codebook < out.codebooks; ++codebook) { + out.codes[static_cast(codebook * out.frames + frame_index)] = + generated_frame_major[static_cast(frame_index * out.codebooks + codebook)]; + } + } + log_profile(profile); + return out; + } + + void release_runtime_graphs() { + step_graph_.reset(); + prefill_graph_.reset(); + fast_graph_.reset(); + } + +private: + class PrefillGraph { + public: + PrefillGraph( + std::shared_ptr runtime, + int64_t steps, + FishPrefillCacheTarget target_cache) + : runtime_(std::move(runtime)), + steps_(steps), + target_cache_(std::move(target_cache)) { + const auto & assets = runtime_->assets(); + const auto & config = assets.config.text; + if (target_cache_.keys.size() != runtime_->weights().slow_layers.size() || + target_cache_.values.size() != runtime_->weights().slow_layers.size()) { + throw std::runtime_error("Fish Audio prefill target cache layer count mismatch"); + } + ggml_init_params params{runtime_->graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio AR prefill context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.ar.prefill", runtime_->backend_type()}; + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps_, config.dim})); + input_ = input.tensor; + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, steps_); + auto positions_value = core::wrap_tensor(positions_, core::TensorShape::from_dims({steps_}), GGML_TYPE_I32); + constants_ = std::make_unique( + runtime_->backend(), + runtime_->threads(), + "fish_audio.ar.prefill.constants", + 256ull * 1024ull * 1024ull); + constants_->begin_graph(); + auto decoder = build_fish_causal_decoder( + ctx, + *constants_, + input, + positions_value, + bind_slow_weights(*constants_, runtime_->weights(), config), + make_slow_decoder_config(config, runtime_->backend_type()), + assets.config.norm_fastlayer_input); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + for (size_t layer_index = 0; layer_index < decoder.state.layers.size(); ++layer_index) { + const auto & layer = decoder.state.layers[layer_index]; + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("Fish Audio prefill decoder did not produce K/V state"); + } + auto key_dest = runtime::view_transformer_kv_cache_steps( + ctx, + target_cache_.keys[layer_index], + 0, + steps_, + config.n_local_heads, + config.head_dim, + "Fish Audio prefill key cache", + target_cache_.keys[layer_index].type); + auto value_dest = runtime::view_transformer_kv_cache_steps( + ctx, + target_cache_.values[layer_index], + 0, + steps_, + config.n_local_heads, + config.head_dim, + "Fish Audio prefill value cache", + target_cache_.values[layer_index].type); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), layer.key->tensor, key_dest.tensor)); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), layer.value->tensor, value_dest.tensor)); + } + hidden_ = decoder.hidden.tensor; + logits_ = decoder.logits.tensor; + ggml_set_output(hidden_); + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + ggml_build_forward_expand(graph_, hidden_); + constants_->finish_graph(); + constants_->ensure_uploaded(); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("failed to allocate Fish Audio AR prefill graph"); + } + auto positions = modules::qwen_position_ids(steps_); + ggml_backend_tensor_set(positions_, positions.data(), 0, positions.size() * sizeof(int32_t)); + } + + ~PrefillGraph() { + core::release_backend_graph_resources(runtime_->backend(), graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + SlowPrefillOutput run(const std::vector & embeddings, FishARProfile & profile) { + const auto & config = runtime_->assets().config.text; + if (static_cast(embeddings.size()) != steps_ * config.dim) { + throw std::runtime_error("Fish Audio prefill embedding size mismatch"); + } + ++profile.prefill_runs; + auto timing_start = Clock::now(); + ggml_backend_tensor_set(input_, embeddings.data(), 0, embeddings.size() * sizeof(float)); + profile.prefill_input_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(runtime_->backend(), graph_, nullptr, "fish_audio.ar.prefill"); + ggml_backend_synchronize(runtime_->backend()); + profile.prefill_graph_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio AR prefill graph compute failed"); + } + SlowPrefillOutput out; + out.forward.logits.resize(static_cast(config.vocab_size)); + out.forward.hidden.resize(static_cast(config.dim)); + timing_start = Clock::now(); + ggml_backend_tensor_get(logits_, out.forward.logits.data(), 0, out.forward.logits.size() * sizeof(float)); + ggml_backend_tensor_get(hidden_, out.forward.hidden.data(), 0, out.forward.hidden.size() * sizeof(float)); + profile.prefill_output_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + return out; + } + + int64_t steps() const noexcept { return steps_; } + + private: + std::shared_ptr runtime_; + int64_t steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * hidden_ = nullptr; + ggml_tensor * logits_ = nullptr; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + std::unique_ptr constants_; + FishPrefillCacheTarget target_cache_; + }; + + class StepGraph { + public: + StepGraph(std::shared_ptr runtime, int64_t cache_steps) + : runtime_(std::move(runtime)), + cache_steps_(cache_steps) { + ggml_init_params state_params{8ull * 1024ull * 1024ull, nullptr, true}; + state_ctx_.reset(ggml_init(state_params)); + if (state_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio AR step state context"); + } + ggml_init_params graph_params{runtime_->graph_arena_bytes(), nullptr, true}; + graph_ctx_.reset(ggml_init(graph_params)); + if (graph_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio AR step context"); + } + const auto & assets = runtime_->assets(); + const auto & config = assets.config.text; + input_ = ggml_new_tensor_3d(state_ctx_.get(), GGML_TYPE_F32, config.dim, 1, 1); + position_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, 1); + cache_slot_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, 1); + mask_ = ggml_new_tensor_4d(state_ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + std::vector cache_keys; + std::vector cache_values; + cache_keys.reserve(runtime_->weights().slow_layers.size()); + cache_values.reserve(runtime_->weights().slow_layers.size()); + const ggml_type cache_type = + runtime_->backend_type() == core::BackendType::Vulkan ? GGML_TYPE_F32 : GGML_TYPE_BF16; + for (size_t layer = 0; layer < runtime_->weights().slow_layers.size(); ++layer) { + cache_keys.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + cache_type, + config.head_dim, + config.n_local_heads, + cache_steps_, + 1), + core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), + cache_type)); + cache_values.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + cache_type, + config.head_dim, + config.n_local_heads, + cache_steps_, + 1), + core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), + cache_type)); + } + state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); + if (state_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Fish Audio AR step state tensors"); + } + + core::ModuleBuildContext ctx{graph_ctx_.get(), "fish_audio.ar.step", runtime_->backend_type()}; + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.dim})); + input = core::wrap_tensor(ggml_cpy(ctx.ggml, input_, input.tensor), input.shape, input.type); + auto position_value = core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto cache_slot_value = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto mask_value = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_}), GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(graph_ctx_.get(), 65536, false); + auto & constants = runtime_->slow_step_constants(); + constants.begin_graph(); + auto decoder = build_fish_static_decoder( + ctx, + graph_, + input, + position_value, + bind_slow_weights(constants, runtime_->weights(), config), + make_slow_decoder_config(config, runtime_->backend_type()), + cache_steps_, + mask_value, + cache_slot_value, + std::move(cache_keys), + std::move(cache_values), + assets.config.norm_fastlayer_input); + cache_ = std::move(decoder.cache); + hidden_ = decoder.hidden.tensor; + logits_ = decoder.logits.tensor; + ggml_set_output(hidden_); + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + ggml_build_forward_expand(graph_, hidden_); + constants.finish_graph(); + constants.ensure_uploaded(); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("failed to allocate Fish Audio AR step tensors"); + } + mask_scratch_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + } + + ~StepGraph() { + core::release_backend_graph_resources(runtime_->backend(), graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + if (state_buffer_ != nullptr) { + ggml_backend_buffer_free(state_buffer_); + } + } + + int64_t cache_steps() const noexcept { return cache_steps_; } + + FishPrefillCacheTarget prefill_target_cache() const { + FishPrefillCacheTarget out; + out.keys.reserve(runtime_->weights().slow_layers.size()); + out.values.reserve(runtime_->weights().slow_layers.size()); + for (size_t layer = 0; layer < runtime_->weights().slow_layers.size(); ++layer) { + out.keys.push_back(cache_.key_tensor(layer)); + out.values.push_back(cache_.value_tensor(layer)); + } + return out; + } + + void finish_prefill(int64_t steps) { + cache_.retain_prefix(0); + cache_.advance_after_direct_append(steps); + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + std::fill(mask_scratch_.begin(), mask_scratch_.end(), masked); + for (int64_t i = 0; i < cache_.valid_steps(); ++i) { + mask_scratch_[static_cast(i)] = visible; + } + ggml_backend_tensor_set(mask_, mask_scratch_.data(), 0, mask_scratch_.size() * sizeof(ggml_fp16_t)); + } + + SlowForwardOutput run(const std::vector & embedding, FishARProfile & profile) { + const auto & config = runtime_->assets().config.text; + if (static_cast(embedding.size()) != config.dim) { + throw std::runtime_error("Fish Audio step embedding size mismatch"); + } + if (cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("Fish Audio step cache exceeds capacity"); + } + ++profile.step_runs; + auto timing_start = Clock::now(); + const int32_t pos = static_cast(cache_.current_end()); + ggml_backend_tensor_set(position_, &pos, 0, sizeof(pos)); + const int32_t cache_slot = static_cast(cache_.valid_steps()); + ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(cache_slot)); + const auto visible = ggml_fp32_to_fp16(0.0F); + mask_scratch_[static_cast(cache_.valid_steps())] = visible; + ggml_backend_tensor_set( + mask_, + &visible, + static_cast(cache_.valid_steps()) * sizeof(ggml_fp16_t), + sizeof(ggml_fp16_t)); + profile.step_mask_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + ggml_backend_tensor_set(input_, embedding.data(), 0, embedding.size() * sizeof(float)); + profile.step_input_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(runtime_->backend(), graph_, nullptr, "fish_audio.ar.step"); + ggml_backend_synchronize(runtime_->backend()); + profile.step_graph_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio AR step graph compute failed"); + } + cache_.advance_after_direct_append(1); + SlowForwardOutput out; + out.logits.resize(static_cast(config.vocab_size)); + out.hidden.resize(static_cast(config.dim)); + timing_start = Clock::now(); + ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + ggml_backend_tensor_get(hidden_, out.hidden.data(), 0, out.hidden.size() * sizeof(float)); + profile.step_output_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + return out; + } + + private: + std::shared_ptr runtime_; + int64_t cache_steps_ = 0; + std::unique_ptr state_ctx_; + std::unique_ptr graph_ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * position_ = nullptr; + ggml_tensor * cache_slot_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * hidden_ = nullptr; + ggml_tensor * logits_ = nullptr; + runtime::TransformerKVCache cache_; + std::vector mask_scratch_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t state_buffer_ = nullptr; + }; + + class FastGraph { + public: + explicit FastGraph(std::shared_ptr runtime) + : runtime_(std::move(runtime)) { + ggml_init_params state_params{8ull * 1024ull * 1024ull, nullptr, true}; + state_ctx_.reset(ggml_init(state_params)); + if (state_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio fast AR state context"); + } + ggml_init_params graph_params{runtime_->graph_arena_bytes(), nullptr, true}; + graph_ctx_.reset(ggml_init(graph_params)); + if (graph_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio fast AR context"); + } + const auto & config = runtime_->assets().config.fast; + const auto & weights = runtime_->weights(); + input_ = ggml_new_tensor_3d(state_ctx_.get(), GGML_TYPE_F32, config.dim, 1, 1); + position_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, 1); + mask_ = ggml_new_tensor_4d(state_ctx_.get(), GGML_TYPE_F16, config.num_codebooks, 1, 1, 1); + std::vector cache_keys; + std::vector cache_values; + cache_keys.reserve(weights.fast_layers.size()); + cache_values.reserve(weights.fast_layers.size()); + const ggml_type cache_type = + runtime_->backend_type() == core::BackendType::Vulkan ? GGML_TYPE_F32 : GGML_TYPE_BF16; + for (size_t layer = 0; layer < weights.fast_layers.size(); ++layer) { + cache_keys.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + cache_type, + config.head_dim, + config.n_local_heads, + config.num_codebooks, + 1), + core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), + cache_type)); + cache_values.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + cache_type, + config.head_dim, + config.n_local_heads, + config.num_codebooks, + 1), + core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), + cache_type)); + } + state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); + if (state_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Fish Audio fast AR state tensors"); + } + for (const auto & cache : cache_keys) { + std::vector zeros(static_cast(ggml_nbytes(cache.tensor)), 0); + ggml_backend_tensor_set(cache.tensor, zeros.data(), 0, zeros.size()); + } + for (const auto & cache : cache_values) { + std::vector zeros(static_cast(ggml_nbytes(cache.tensor)), 0); + ggml_backend_tensor_set(cache.tensor, zeros.data(), 0, zeros.size()); + } + + core::ModuleBuildContext ctx{graph_ctx_.get(), "fish_audio.ar.fast", runtime_->backend_type()}; + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.dim})); + input = core::wrap_tensor(ggml_cpy(ctx.ggml, input_, input.tensor), input.shape, input.type); + auto position_value = core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto mask_value = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, config.num_codebooks}), GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(graph_ctx_.get(), 32768, false); + auto & constants = runtime_->fast_constants(); + constants.begin_graph(); + modules::QwenCausalDecoderWeights decoder_weights; + decoder_weights.stack.layers.reserve(weights.fast_layers.size()); + for (const auto & layer : weights.fast_layers) { + decoder_weights.stack.layers.push_back(bind_fast_layer(constants, layer, config)); + } + decoder_weights.final_norm = binding::norm_data(constants, weights.fast_norm); + decoder_weights.lm_head = binding::linear_data(constants, weights.fast_output); + auto decoder = build_fish_static_decoder( + ctx, + graph_, + input, + position_value, + decoder_weights, + make_fast_decoder_config(config, runtime_->backend_type()), + config.num_codebooks, + mask_value, + position_value, + std::move(cache_keys), + std::move(cache_values), + true); + logits_ = decoder.logits.tensor; + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + constants.finish_graph(); + constants.ensure_uploaded(); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("failed to allocate Fish Audio fast AR graph"); + } + mask_scratch_.assign(static_cast(config.num_codebooks), ggml_fp32_to_fp16(-INFINITY)); + } + + ~FastGraph() { + core::release_backend_graph_resources(runtime_->backend(), graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + if (state_buffer_ != nullptr) { + ggml_backend_buffer_free(state_buffer_); + } + } + + std::vector run(const std::vector & input, int64_t position, FishARProfile & profile) { + const auto & config = runtime_->assets().config.fast; + if (static_cast(input.size()) != config.dim) { + throw std::runtime_error("Fish Audio fast AR input size mismatch"); + } + ++profile.fast_runs; + auto timing_start = Clock::now(); + const int32_t pos = static_cast(position); + ggml_backend_tensor_set(position_, &pos, 0, sizeof(pos)); + const auto visible = ggml_fp32_to_fp16(0.0F); + if (position == 0) { + std::fill(mask_scratch_.begin(), mask_scratch_.end(), ggml_fp32_to_fp16(-INFINITY)); + mask_scratch_[0] = visible; + ggml_backend_tensor_set(mask_, mask_scratch_.data(), 0, mask_scratch_.size() * sizeof(ggml_fp16_t)); + } else { + mask_scratch_[static_cast(position)] = visible; + ggml_backend_tensor_set( + mask_, + &visible, + static_cast(position) * sizeof(ggml_fp16_t), + sizeof(ggml_fp16_t)); + } + profile.fast_mask_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + ggml_backend_tensor_set(input_, input.data(), 0, input.size() * sizeof(float)); + profile.fast_input_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(runtime_->backend(), graph_, nullptr, "fish_audio.ar.fast"); + ggml_backend_synchronize(runtime_->backend()); + profile.fast_graph_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio fast AR graph compute failed"); + } + std::vector logits(static_cast(config.vocab_size), 0.0F); + timing_start = Clock::now(); + ggml_backend_tensor_get(logits_, logits.data(), 0, logits.size() * sizeof(float)); + profile.fast_output_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + return logits; + } + + private: + std::shared_ptr runtime_; + std::unique_ptr state_ctx_; + std::unique_ptr graph_ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * position_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector mask_scratch_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t state_buffer_ = nullptr; + }; + + void ensure_prefill_graph(int64_t steps, FishARProfile & profile) { + if (!prefill_graph_ || prefill_graph_->steps() != steps) { + const auto build_start = Clock::now(); + if (!step_graph_) { + throw std::runtime_error("Fish Audio AR prefill requires a step graph"); + } + prefill_graph_ = std::make_unique(runtime_, steps, step_graph_->prefill_target_cache()); + profile.graph_build_prefill_ms += engine::debug::elapsed_ms(build_start, Clock::now()); + } + } + + void ensure_step_graph(int64_t cache_steps, FishARProfile & profile) { + if (!step_graph_ || step_graph_->cache_steps() < cache_steps) { + const auto build_start = Clock::now(); + step_graph_ = std::make_unique(runtime_, cache_steps); + prefill_graph_.reset(); + profile.graph_build_step_ms += engine::debug::elapsed_ms(build_start, Clock::now()); + } + } + + void ensure_fast_graph(FishARProfile & profile) { + if (!fast_graph_) { + const auto build_start = Clock::now(); + fast_graph_ = std::make_unique(runtime_); + profile.graph_build_fast_ms += engine::debug::elapsed_ms(build_start, Clock::now()); + } + } + + int32_t im_end_id() const { + return static_cast(runtime_->assets().config.im_end_token_id); + } + + void append_frame(std::vector & out, const std::vector & frame) const { + if (static_cast(frame.size()) != runtime_->assets().config.fast.num_codebooks + 1) { + throw std::runtime_error("Fish Audio generated frame shape mismatch"); + } + out.insert(out.end(), frame.begin() + 1, frame.end()); + } + + std::vector sample_frame( + const std::vector & slow_logits, + const std::vector & slow_hidden, + const FishAudioGenerationOptions & options, + SampleState & sample, + bool apply_ras, + FishARProfile & profile) { + const auto & config = runtime_->assets().config; + const auto & weights = runtime_->weights(); + auto timing_start = Clock::now(); + const auto biased = apply_semantic_bias(config, im_end_id(), slow_logits); + profile.sample_bias_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + int32_t main_token = sample_from_logits( + biased, + options.temperature, + options.top_p, + options.top_k, + sample, + sampling_policy_); + profile.sample_main_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + const int32_t high_token = sample_from_logits( + biased, + kRasHighTemperature, + kRasHighTopP, + options.top_k, + sample, + sampling_policy_); + profile.sample_high_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (apply_ras && is_semantic_token(config, main_token) && + std::find(sample.previous_main.begin(), sample.previous_main.end(), main_token) != sample.previous_main.end()) { + main_token = high_token; + } + std::rotate(sample.previous_main.begin(), sample.previous_main.begin() + 1, sample.previous_main.end()); + sample.previous_main.back() = main_token; + + std::vector frame(static_cast(config.fast.num_codebooks + 1), 0); + frame[0] = main_token; + if (!is_semantic_token(config, main_token)) { + return frame; + } + const auto fast0_logits = fast_graph_->run(slow_hidden, 0, profile); + int32_t code = std::clamp( + main_token - static_cast(config.semantic_start_token_id), + 0, + static_cast(config.fast.vocab_size - 1)); + frame[1] = code; + for (int64_t codebook = 1; codebook < config.fast.num_codebooks; ++codebook) { + timing_start = Clock::now(); + const auto embedding = build_fast_embedding(config, weights, code); + profile.fast_embedding_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + const auto logits = fast_graph_->run(embedding, codebook, profile); + timing_start = Clock::now(); + code = sample_from_logits( + logits, + options.temperature, + options.top_p, + options.top_k, + sample, + sampling_policy_); + profile.sample_fast_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + frame[static_cast(codebook + 1)] = code; + } + return frame; + } + + void log_profile(const FishARProfile & profile) const { + engine::debug::timing_log_scalar("fish_audio.ar.profile.graph_build_prefill_ms", profile.graph_build_prefill_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.graph_build_step_ms", profile.graph_build_step_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.graph_build_fast_ms", profile.graph_build_fast_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.slow_embedding_ms", profile.slow_embedding_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_embedding_ms", profile.fast_embedding_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_input_upload_ms", profile.prefill_input_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_graph_ms", profile.prefill_graph_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_output_read_ms", profile.prefill_output_read_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_input_upload_ms", profile.step_input_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_mask_upload_ms", profile.step_mask_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_graph_ms", profile.step_graph_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_output_read_ms", profile.step_output_read_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_input_upload_ms", profile.fast_input_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_mask_upload_ms", profile.fast_mask_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_graph_ms", profile.fast_graph_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_output_read_ms", profile.fast_output_read_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_bias_ms", profile.sample_bias_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_main_ms", profile.sample_main_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_high_ms", profile.sample_high_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_fast_ms", profile.sample_fast_ms); + engine::debug::trace_log_scalar("fish_audio.ar.profile.prefill_runs", profile.prefill_runs); + engine::debug::trace_log_scalar("fish_audio.ar.profile.step_runs", profile.step_runs); + engine::debug::trace_log_scalar("fish_audio.ar.profile.fast_runs", profile.fast_runs); + engine::debug::trace_log_scalar("fish_audio.ar.profile.generated_frames", profile.generated_frames); + } + + std::shared_ptr runtime_; + sampling::TorchCudaSamplingPolicy sampling_policy_; + std::unique_ptr prefill_graph_; + std::unique_ptr step_graph_; + std::unique_ptr fast_graph_; +}; + +FishAudioARRuntime::FishAudioARRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + std::move(assets), + backend, + threads, + graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +FishAudioARRuntime::~FishAudioARRuntime() = default; + +FishAudioCodes FishAudioARRuntime::generate( + const FishAudioPrompt & prompt, + const FishAudioGenerationOptions & options) { + return impl_->generate(prompt, options); +} + +void FishAudioARRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/assets.cpp b/src/models/fish_audio/assets.cpp new file mode 100644 index 00000000..7870a300 --- /dev/null +++ b/src/models/fish_audio/assets.cpp @@ -0,0 +1,127 @@ +#include "engine/models/fish_audio/assets.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/framework/io/config.h" +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::fish_audio { +namespace json = engine::io::json; +namespace { + +FishAudioTextConfig parse_text_config(const json::Value & value) { + if (json::optional_string(value, "model_type", "") != "fish_qwen3") { + throw std::runtime_error("Fish Audio text_config.model_type mismatch"); + } + FishAudioTextConfig config; + config.vocab_size = json::require_i64(value, "vocab_size"); + config.n_layer = json::require_i64(value, "n_layer"); + config.dim = json::require_i64(value, "dim"); + config.intermediate_size = json::require_i64(value, "intermediate_size"); + config.n_head = json::require_i64(value, "n_head"); + config.n_local_heads = json::optional_i64(value, "n_local_heads", config.n_head); + config.head_dim = json::require_i64(value, "head_dim"); + config.max_seq_len = json::require_i64(value, "max_seq_len"); + config.rope_base = json::optional_f32(value, "rope_base", config.rope_base); + config.norm_eps = json::optional_f32(value, "norm_eps", config.norm_eps); + config.tie_word_embeddings = json::optional_bool(value, "tie_word_embeddings", config.tie_word_embeddings); + config.attention_qk_norm = json::optional_bool(value, "attention_qk_norm", config.attention_qk_norm); + engine::io::require_positive(config.vocab_size, "text vocab_size"); + engine::io::require_positive(config.n_layer, "text n_layer"); + engine::io::require_positive(config.dim, "text dim"); + engine::io::require_positive(config.intermediate_size, "text intermediate_size"); + engine::io::require_positive(config.n_head, "text n_head"); + engine::io::require_positive(config.n_local_heads, "text n_local_heads"); + engine::io::require_positive(config.head_dim, "text head_dim"); + engine::io::require_positive(config.max_seq_len, "text max_seq_len"); + engine::io::require_divisible(config.n_head, config.n_local_heads, "text n_head / n_local_heads"); + return config; +} + +FishAudioFastConfig parse_fast_config(const json::Value & value) { + if (json::optional_string(value, "model_type", "") != "fish_qwen3_audio_decoder") { + throw std::runtime_error("Fish Audio audio_decoder_config.model_type mismatch"); + } + FishAudioFastConfig config; + config.vocab_size = json::require_i64(value, "vocab_size"); + config.num_codebooks = json::require_i64(value, "num_codebooks"); + config.n_layer = json::require_i64(value, "n_layer"); + config.dim = json::require_i64(value, "dim"); + config.intermediate_size = json::require_i64(value, "intermediate_size"); + config.n_head = json::require_i64(value, "n_head"); + config.n_local_heads = json::optional_i64(value, "n_local_heads", config.n_head); + config.head_dim = json::require_i64(value, "head_dim"); + config.max_seq_len = json::optional_i64(value, "max_seq_len", config.num_codebooks + 1); + config.rope_base = json::optional_f32(value, "rope_base", config.rope_base); + config.norm_eps = json::optional_f32(value, "norm_eps", config.norm_eps); + config.tie_word_embeddings = json::optional_bool(value, "tie_word_embeddings", config.tie_word_embeddings); + config.attention_qk_norm = json::optional_bool(value, "attention_qk_norm", config.attention_qk_norm); + engine::io::require_positive(config.vocab_size, "fast vocab_size"); + engine::io::require_positive(config.num_codebooks, "fast num_codebooks"); + engine::io::require_positive(config.n_layer, "fast n_layer"); + engine::io::require_positive(config.dim, "fast dim"); + engine::io::require_positive(config.intermediate_size, "fast intermediate_size"); + engine::io::require_positive(config.n_head, "fast n_head"); + engine::io::require_positive(config.n_local_heads, "fast n_local_heads"); + engine::io::require_positive(config.head_dim, "fast head_dim"); + engine::io::require_divisible(config.n_head, config.n_local_heads, "fast n_head / n_local_heads"); + return config; +} + +FishAudioConfig parse_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + FishAudioConfig config; + config.model_type = json::optional_string(root, "model_type", ""); + if (config.model_type != "fish_qwen3_omni") { + throw std::runtime_error("Fish Audio model_type mismatch"); + } + config.torch_dtype = json::optional_string(root, "torch_dtype", config.torch_dtype); + config.semantic_start_token_id = json::require_i64(root, "semantic_start_token_id"); + config.semantic_end_token_id = json::require_i64(root, "semantic_end_token_id"); + config.im_end_token_id = json::require_i64(root, "eos_token_id"); + config.norm_fastlayer_input = json::optional_bool(root, "norm_fastlayer_input", config.model_type == "fish_qwen3_omni"); + config.text = parse_text_config(root.require("text_config")); + config.fast = parse_fast_config(root.require("audio_decoder_config")); + config.codec.total_codebooks = config.fast.num_codebooks; + if (config.fast.dim != config.text.dim) { + throw std::runtime_error("Fish Audio fast dim must match text dim for S2-Pro"); + } + if (!config.text.tie_word_embeddings) { + throw std::runtime_error("Fish Audio S2-Pro expects tied text embeddings"); + } + if (config.semantic_start_token_id <= 0 || config.semantic_end_token_id < config.semantic_start_token_id) { + throw std::runtime_error("Fish Audio semantic token range is invalid"); + } + return config; +} + +void validate_weight_anchors(const FishAudioAssets & assets) { + assets.model_weights->require_metadata("embeddings.weight"); + assets.model_weights->require_metadata("codebook_embeddings.weight"); + assets.model_weights->require_metadata("layers.0.attention.q_proj.weight"); + assets.model_weights->require_metadata("layers.0.attention.k_proj.weight"); + assets.model_weights->require_metadata("layers.0.attention.v_proj.weight"); + assets.model_weights->require_metadata("fast_layers.0.attention.q_proj.weight"); + assets.model_weights->require_metadata("fast_embeddings.weight"); + assets.model_weights->require_metadata("fast_output.weight"); + assets.codec_weights->require_metadata("quantizer.semantic_quantizer.quantizers.0.codebook.weight"); + assets.codec_weights->require_metadata("decoder.model.0.conv.weight"); +} + +} // namespace + +std::shared_ptr load_fish_audio_assets(const std::filesystem::path & model_path) { + FishAudioAssets assets; + assets.resources = assets::load_resource_bundle_from_package_spec( + model_path, + assets::default_model_package_spec_path("fish_audio")); + assets.config = parse_config(assets.resources); + assets.model_weights = assets.resources.open_tensor_source("model_weights"); + assets.codec_weights = assets.resources.open_tensor_source("codec_weights"); + validate_weight_anchors(assets); + return std::make_shared(std::move(assets)); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/codec.cpp b/src/models/fish_audio/codec.cpp new file mode 100644 index 00000000..f63c6978 --- /dev/null +++ b/src/models/fish_audio/codec.cpp @@ -0,0 +1,1148 @@ +#include "engine/models/fish_audio/codec.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention_modules.h" +#include "engine/framework/modules/conditioning_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include "../common/constant_tensor_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +namespace binding = engine::modules::binding; + +constexpr int64_t kCodecDim = 1024; +constexpr int64_t kCodecTransformerHeads = 16; +constexpr int64_t kCodecHeadDim = 64; +constexpr int64_t kCodecIntermediate = 3072; +constexpr int64_t kCodecTransformerLayers = 8; +constexpr float kCodecNormEps = 1.0e-5F; +constexpr float kConvNextNormEps = 1.0e-6F; +constexpr float kCodecRopeTheta = 10000.0F; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +std::vector dims_vector(const core::TensorShape & shape) { + std::vector out; + out.reserve(shape.rank); + for (size_t i = 0; i < shape.rank; ++i) { + out.push_back(shape.dims[i]); + } + return out; +} + +std::vector prepare_codec_mono( + const runtime::AudioBuffer & audio, + int target_sample_rate_hz) { + auto mono = engine::audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); + if (audio.sample_rate != target_sample_rate_hz) { + mono = engine::audio::resample_mono_torchaudio_sinc_hann( + mono, + audio.sample_rate, + target_sample_rate_hz); + } + return mono; +} + +struct CodecTransformerLayerWeights { + modules::NormWeights attention_norm; + modules::AttentionWeights attention; + modules::LayerScaleWeights attention_scale; + modules::NormWeights ffn_norm; + modules::QwenMLPWeights feed_forward; + modules::LayerScaleWeights ffn_scale; +}; + +struct CodecTransformerWeights { + std::vector layers; + modules::NormWeights norm; +}; + +struct ResidualUnitWeights { + modules::Snake1dWeights snake1; + modules::Conv1dWeights conv1; + modules::Snake1dWeights snake2; + modules::Conv1dWeights conv2; +}; + +struct EncoderBlockWeights { + ResidualUnitWeights residual1; + ResidualUnitWeights residual3; + ResidualUnitWeights residual9; + modules::Snake1dWeights snake; + modules::Conv1dWeights conv; + std::optional transformer; +}; + +struct DecoderBlockWeights { + modules::Snake1dWeights snake; + modules::ConvTranspose1dWeights conv; + ResidualUnitWeights residual1; + ResidualUnitWeights residual3; + ResidualUnitWeights residual9; +}; + +struct ConvNeXtBlockWeights { + modules::DepthwiseConv1dWeights dwconv; + modules::NormWeights norm; + modules::LinearWeights pwconv1; + modules::LinearWeights pwconv2; + modules::LayerScaleWeights gamma; +}; + +struct QuantizerUnitWeights { + modules::Conv1dWeights in_proj; + modules::Conv1dWeights out_proj; + core::TensorValue codebook; + core::TensorValue normalized_codebook; +}; + +struct FishCodecWeights { + std::shared_ptr store; + modules::Conv1dWeights encoder_first; + std::vector encoder_blocks; + modules::Snake1dWeights encoder_final_snake; + modules::Conv1dWeights encoder_final; + + std::vector> downsample; + CodecTransformerWeights pre_module; + QuantizerUnitWeights semantic_quantizer; + std::vector residual_quantizers; + CodecTransformerWeights post_module; + std::vector> upsample; + + modules::Conv1dWeights decoder_first; + std::vector decoder_blocks; + modules::Snake1dWeights decoder_final_snake; + modules::Conv1dWeights decoder_final; +}; + +int64_t ceil_div(int64_t a, int64_t b) { + return (a + b - 1) / b; +} + +std::vector normalized_rows(const std::vector & values, int64_t rows, int64_t cols) { + if (static_cast(values.size()) != rows * cols) { + throw std::runtime_error("Fish Audio normalized_rows shape mismatch"); + } + std::vector out(values.size(), 0.0F); + for (int64_t row = 0; row < rows; ++row) { + double sum = 0.0; + for (int64_t col = 0; col < cols; ++col) { + const float value = values[static_cast(row * cols + col)]; + sum += static_cast(value) * static_cast(value); + } + const float inv = sum > 0.0 ? static_cast(1.0 / std::sqrt(sum)) : 0.0F; + for (int64_t col = 0; col < cols; ++col) { + const size_t index = static_cast(row * cols + col); + out[index] = values[index] * inv; + } + } + return out; +} + +core::TensorValue slice_frames(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t start, int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("Fish Audio codec slice_frames requires positive frames"); + } + return modules::SliceModule({2, start, frames}).build(ctx, input); +} + +core::TensorValue zero_prefix_like(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("Fish Audio zero_prefix_like requires positive frames"); + } + auto first = modules::SliceModule({2, 0, 1}).build(ctx, input); + if (ctx.backend_type == core::BackendType::Cpu) { + first = core::ensure_backend_addressable_layout(ctx, first); + } + first = core::wrap_tensor(ggml_scale(ctx.ggml, first.tensor, 0.0F), first.shape, GGML_TYPE_F32); + return modules::RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], frames})}) + .build(ctx, first); +} + +core::TensorValue zero_suffix_like(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("Fish Audio zero_suffix_like requires positive frames"); + } + auto last = modules::SliceModule({2, input.shape.dims[2] - 1, 1}).build(ctx, input); + if (ctx.backend_type == core::BackendType::Cpu) { + last = core::ensure_backend_addressable_layout(ctx, last); + } + last = core::wrap_tensor(ggml_scale(ctx.ggml, last.tensor, 0.0F), last.shape, GGML_TYPE_F32); + return modules::RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], frames})}) + .build(ctx, last); +} + +int64_t extra_padding_for_conv1d(int64_t frames, int64_t effective_kernel, int64_t stride, int64_t left_pad) { + const double n_frames = (static_cast(frames - effective_kernel + left_pad) / static_cast(stride)) + 1.0; + const int64_t ideal_length = + (static_cast(std::ceil(n_frames)) - 1) * stride + (effective_kernel - left_pad); + return ideal_length - frames; +} + +core::TensorValue causal_pad(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t left_pad, int64_t right_pad) { + if (left_pad < 0) { + throw std::runtime_error("Fish Audio causal conv requires non-negative left padding"); + } + if (right_pad < 0) { + throw std::runtime_error("Fish Audio causal conv requires non-negative right padding"); + } + if (left_pad == 0 && right_pad == 0) { + return input; + } + auto out = input; + if (left_pad > 0) { + out = modules::ConcatModule({2}).build(ctx, zero_prefix_like(ctx, input, left_pad), out); + } + if (right_pad > 0) { + out = modules::ConcatModule({2}).build(ctx, out, zero_suffix_like(ctx, input, right_pad)); + } + return out; +} + +core::TensorValue causal_conv1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::Conv1dWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t kernel, + int stride, + int dilation, + bool use_bias) { + const int64_t effective_kernel = (kernel - 1) * dilation + 1; + const int64_t left_pad = effective_kernel - stride; + const int64_t right_pad = extra_padding_for_conv1d(input.shape.dims[2], effective_kernel, stride, left_pad); + auto padded = causal_pad(ctx, input, left_pad, right_pad); + return modules::Conv1dModule({ + in_channels, + out_channels, + kernel, + stride, + 0, + dilation, + use_bias, + }).build(ctx, padded, weights); +} + +core::TensorValue causal_depthwise_conv1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::DepthwiseConv1dWeights & weights, + int64_t channels, + int64_t kernel, + int stride, + int dilation, + bool use_bias) { + const int64_t effective_kernel = (kernel - 1) * dilation + 1; + const int64_t left_pad = effective_kernel - stride; + const int64_t right_pad = extra_padding_for_conv1d(input.shape.dims[2], effective_kernel, stride, left_pad); + auto padded = causal_pad(ctx, input, left_pad, right_pad); + return modules::DepthwiseConv1dModule({ + channels, + kernel, + stride, + 0, + dilation, + use_bias, + }).build(ctx, padded, weights); +} + +core::TensorValue causal_conv_transpose1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::ConvTranspose1dWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t kernel, + int stride, + bool use_bias) { + auto out = modules::ConvTranspose1dModule({ + in_channels, + out_channels, + kernel, + stride, + 0, + 1, + use_bias, + }).build(ctx, input, weights); + const int64_t pad = kernel - stride; + const int64_t padding_right = static_cast(std::ceil(static_cast(pad))); + const int64_t padding_left = pad - padding_right; + return slice_frames(ctx, out, padding_left, out.shape.dims[2] - padding_left - padding_right); +} + +core::TensorValue l2_normalize_last(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + const bool materialize_input = ctx.backend_type == core::BackendType::Metal; + const auto normalized_input = materialize_input + ? core::ensure_backend_addressable_layout(ctx, input) + : input; + auto squared = modules::MulModule{}.build(ctx, normalized_input, normalized_input); + auto sum = modules::ReduceSumModule({static_cast(input.shape.rank - 1)}).build(ctx, squared); + auto shifted = core::wrap_tensor(ggml_scale_bias(ctx.ggml, sum.tensor, 1.0F, 1.0e-12F), sum.shape, GGML_TYPE_F32); + auto denom = modules::SqrtModule{}.build(ctx, shifted); + auto repeated = modules::RepeatModule({normalized_input.shape}).build(ctx, denom); + return core::wrap_tensor(ggml_div(ctx.ggml, normalized_input.tensor, repeated.tensor), normalized_input.shape, GGML_TYPE_F32); +} + +core::TensorValue build_mlp( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::QwenMLPWeights & weights) { + auto gate = modules::LinearModule({kCodecDim, kCodecIntermediate, false, GGML_PREC_F32}) + .build(ctx, input, weights.gate_proj); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule({kCodecDim, kCodecIntermediate, false, GGML_PREC_F32}) + .build(ctx, input, weights.up_proj); + auto hidden = modules::MulModule{}.build(ctx, gate, up); + return modules::LinearModule({kCodecIntermediate, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, hidden, weights.down_proj); +} + +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({input.shape.dims[0], input.shape.dims[1], kCodecTransformerHeads, kCodecHeadDim})); +} + +core::TensorValue attention_from_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + const core::TensorValue & attention_mask) { + auto q = modules::TransposeModule({{0, 2, 1, 3}, q_heads.shape.rank}).build(ctx, q_heads); + auto k = modules::TransposeModule({{0, 2, 1, 3}, k_heads.shape.rank}).build(ctx, k_heads); + auto v = modules::TransposeModule({{0, 2, 1, 3}, v_heads.shape.rank}).build(ctx, v_heads); + q = core::wrap_tensor(ggml_cont(ctx.ggml, q.tensor), q.shape, q.type); + k = core::wrap_tensor(ggml_cont(ctx.ggml, k.tensor), k.shape, k.type); + v = core::wrap_tensor(ggml_cont(ctx.ggml, v.tensor), v.shape, v.type); + auto * flash = ggml_flash_attn_ext( + ctx.ggml, + q.tensor, + k.tensor, + v.tensor, + attention_mask.tensor, + 1.0F / std::sqrt(static_cast(kCodecHeadDim)), + 0.0F, + 0.0F); + ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); + return core::wrap_tensor( + flash, + core::TensorShape::from_dims({q.shape.dims[0], q.shape.dims[2], q.shape.dims[1], kCodecHeadDim}), + GGML_TYPE_F32); +} + +core::TensorValue build_transformer_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const core::TensorValue & attention_mask, + const CodecTransformerLayerWeights & weights) { + auto normed = modules::RMSNormModule({kCodecDim, kCodecNormEps, true, false}).build(ctx, input, weights.attention_norm); + auto q = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.q_weight, std::nullopt}); + auto k = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.k_weight, std::nullopt}); + auto v = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.v_weight, std::nullopt}); + q = modules::RoPEModule({kCodecHeadDim, GGML_ROPE_TYPE_NORMAL, kCodecRopeTheta}).build(ctx, reshape_heads(ctx, q), positions); + k = modules::RoPEModule({kCodecHeadDim, GGML_ROPE_TYPE_NORMAL, kCodecRopeTheta}).build(ctx, reshape_heads(ctx, k), positions); + v = reshape_heads(ctx, v); + auto context = attention_from_heads(ctx, q, k, v, attention_mask); + 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], kCodecDim})); + auto attn = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, context, {weights.attention.out_weight, std::nullopt}); + attn = modules::LayerScaleModule{}.build(ctx, attn, weights.attention_scale); + auto hidden = modules::AddModule{}.build(ctx, input, attn); + auto ffn_in = modules::RMSNormModule({kCodecDim, kCodecNormEps, true, false}).build(ctx, hidden, weights.ffn_norm); + auto ff = build_mlp(ctx, ffn_in, weights.feed_forward); + ff = modules::LayerScaleModule{}.build(ctx, ff, weights.ffn_scale); + return modules::AddModule{}.build(ctx, hidden, ff); +} + +core::TensorValue make_positions( + core::ModuleBuildContext &, + common::ConstantTensorCache & constants, + int64_t frames) { + std::vector values(static_cast(frames)); + for (int64_t i = 0; i < frames; ++i) { + values[static_cast(i)] = static_cast(i); + } + return constants.make_tensor(core::TensorShape::from_dims({frames}), GGML_TYPE_I32, values.data(), values.size() * sizeof(int32_t)); +} + +core::TensorValue make_causal_mask( + core::ModuleBuildContext &, + common::ConstantTensorCache & constants, + int64_t frames, + int64_t window_size) { + std::vector values(static_cast(frames * frames), ggml_fp32_to_fp16(-std::numeric_limits::infinity())); + for (int64_t row = 0; row < frames; ++row) { + const int64_t begin = window_size > 0 ? std::max(0, row - window_size + 1) : 0; + for (int64_t col = begin; col <= row; ++col) { + values[static_cast(row * frames + col)] = ggml_fp32_to_fp16(0.0F); + } + } + return constants.make_tensor(core::TensorShape::from_dims({frames, frames}), GGML_TYPE_F16, values.data(), values.size() * sizeof(ggml_fp16_t)); +} + +core::TensorValue build_window_transformer( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & input_bct, + const CodecTransformerWeights & weights, + int64_t window_size) { + auto x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, input_bct); + auto positions = make_positions(ctx, constants, x.shape.dims[1]); + auto mask = make_causal_mask(ctx, constants, x.shape.dims[1], window_size); + for (const auto & layer : weights.layers) { + x = build_transformer_layer(ctx, x, positions, mask, layer); + } + x = modules::RMSNormModule({kCodecDim, kCodecNormEps, true, false}).build(ctx, x, weights.norm); + return modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); +} + +core::TensorValue build_residual_unit( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ResidualUnitWeights & weights, + int64_t channels, + int dilation) { + auto y = modules::Snake1dModule({channels}).build(ctx, input, weights.snake1); + y = causal_conv1d(ctx, y, weights.conv1, channels, channels, 7, 1, dilation, true); + y = modules::Snake1dModule({channels}).build(ctx, y, weights.snake2); + y = causal_conv1d(ctx, y, weights.conv2, channels, channels, 1, 1, 1, true); + core::TensorValue x = input; + if (x.shape.dims[2] != y.shape.dims[2]) { + x = slice_frames(ctx, x, 0, y.shape.dims[2]); + } + return modules::AddModule{}.build(ctx, x, y); +} + +core::TensorValue build_convnext( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ConvNeXtBlockWeights & weights, + int64_t channels) { + auto y = causal_depthwise_conv1d(ctx, input, weights.dwconv, channels, 7, 1, 1, true); + y = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, y); + y = modules::LayerNormModule({channels, kConvNextNormEps, true, true}).build(ctx, y, weights.norm); + y = modules::LinearModule({channels, channels * 4, true, GGML_PREC_F32}).build(ctx, y, weights.pwconv1); + y = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, y); + y = modules::LinearModule({channels * 4, channels, true, GGML_PREC_F32}).build(ctx, y, weights.pwconv2); + y = modules::LayerScaleModule{}.build(ctx, y, weights.gamma); + y = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, y); + core::TensorValue x = input; + if (x.shape.dims[2] != y.shape.dims[2]) { + x = slice_frames(ctx, x, 0, y.shape.dims[2]); + } + return modules::AddModule{}.build(ctx, x, y); +} + +core::TensorValue build_encoder( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & input, + const FishCodecWeights & weights) { + auto x = causal_conv1d(ctx, input, weights.encoder_first, 1, 64, 7, 1, 1, true); + int64_t channels = 64; + const int strides[] = {2, 4, 8, 8}; + for (size_t index = 0; index < weights.encoder_blocks.size(); ++index) { + const auto & block = weights.encoder_blocks[index]; + x = build_residual_unit(ctx, x, block.residual1, channels, 1); + x = build_residual_unit(ctx, x, block.residual3, channels, 3); + x = build_residual_unit(ctx, x, block.residual9, channels, 9); + x = modules::Snake1dModule({channels}).build(ctx, x, block.snake); + x = causal_conv1d(ctx, x, block.conv, channels, channels * 2, 2 * strides[index], strides[index], 1, true); + channels *= 2; + if (block.transformer.has_value()) { + x = build_window_transformer(ctx, constants, x, *block.transformer, 512); + } + } + x = modules::Snake1dModule({channels}).build(ctx, x, weights.encoder_final_snake); + return causal_conv1d(ctx, x, weights.encoder_final, channels, kCodecDim, 3, 1, 1, true); +} + +core::TensorValue build_decoder( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const FishCodecWeights & weights) { + auto x = causal_conv1d(ctx, input, weights.decoder_first, kCodecDim, 1536, 7, 1, 1, true); + int64_t channels = 1536; + const int strides[] = {8, 8, 4, 2}; + for (size_t index = 0; index < weights.decoder_blocks.size(); ++index) { + const auto & block = weights.decoder_blocks[index]; + x = modules::Snake1dModule({channels}).build(ctx, x, block.snake); + x = causal_conv_transpose1d(ctx, x, block.conv, channels, channels / 2, 2 * strides[index], strides[index], true); + channels /= 2; + x = build_residual_unit(ctx, x, block.residual1, channels, 1); + x = build_residual_unit(ctx, x, block.residual3, channels, 3); + x = build_residual_unit(ctx, x, block.residual9, channels, 9); + } + x = modules::Snake1dModule({channels}).build(ctx, x, weights.decoder_final_snake); + x = causal_conv1d(ctx, x, weights.decoder_final, channels, 1, 7, 1, 1, true); + return modules::TanhModule{}.build(ctx, x); +} + +core::TensorValue build_quantizer_out( + core::ModuleBuildContext & ctx, + const core::TensorValue & ids_bt, + const QuantizerUnitWeights & weights, + int64_t codebook_size) { + auto emb_btd = modules::CodebookLookupModule({codebook_size, 8}).build(ctx, ids_bt, weights.codebook); + auto emb_bdt = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, emb_btd); + return modules::Conv1dModule({8, kCodecDim, 1, 1, 0, 1, true}).build(ctx, emb_bdt, weights.out_proj); +} + +core::TensorValue build_decode_quantizer( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const std::vector & code_inputs, + const FishCodecWeights & weights) { + auto latent = build_quantizer_out(ctx, code_inputs[0], weights.semantic_quantizer, 4096); + for (size_t index = 0; index < weights.residual_quantizers.size(); ++index) { + auto residual = build_quantizer_out(ctx, code_inputs[index + 1], weights.residual_quantizers[index], 1024); + latent = modules::AddModule{}.build(ctx, latent, residual); + } + latent = build_window_transformer(ctx, constants, latent, weights.post_module, 128); + for (const auto & stage : weights.upsample) { + latent = causal_conv_transpose1d(ctx, latent, stage.first, kCodecDim, kCodecDim, 2, 2, true); + latent = build_convnext(ctx, latent, stage.second, kCodecDim); + } + return latent; +} + +core::TensorValue build_encode_quantizer( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & encoder_latent, + const FishCodecWeights & weights, + std::vector & code_outputs, + std::vector> & trace_outputs) { + auto x = encoder_latent; + for (const auto & stage : weights.downsample) { + x = causal_conv1d(ctx, x, stage.first, kCodecDim, kCodecDim, 2, 2, 1, true); + x = build_convnext(ctx, x, stage.second, kCodecDim); + } + trace_outputs.push_back({"fish_audio.codec.after_downsample", x}); + x = build_window_transformer(ctx, constants, x, weights.pre_module, 128); + trace_outputs.push_back({"fish_audio.codec.after_pre_module", x}); + + auto residual = x; + auto quantize_one = [&](const QuantizerUnitWeights & quantizer, int64_t codebook_size) { + auto projected = modules::Conv1dModule({kCodecDim, 8, 1, 1, 0, 1, true}).build(ctx, residual, quantizer.in_proj); + auto projected_btd = l2_normalize_last(ctx, modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, projected)); + auto logits = modules::LinearModule({8, codebook_size, false, GGML_PREC_F32}) + .build(ctx, projected_btd, {quantizer.normalized_codebook, std::nullopt}); + auto flat_logits = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, logits), + core::TensorShape::from_dims({logits.shape.dims[1], codebook_size})); + auto * ids_raw = ggml_argmax(ctx.ggml, flat_logits.tensor); + ggml_set_output(ids_raw); + code_outputs.push_back(ids_raw); + auto ids = core::reshape_tensor( + ctx, + core::wrap_tensor(ids_raw, core::TensorShape::from_dims({logits.shape.dims[1]}), GGML_TYPE_I32), + core::TensorShape::from_dims({1, logits.shape.dims[1]})); + auto quantized = build_quantizer_out(ctx, ids, quantizer, codebook_size); + residual = core::wrap_tensor(ggml_sub(ctx.ggml, residual.tensor, quantized.tensor), residual.shape, GGML_TYPE_F32); + }; + quantize_one(weights.semantic_quantizer, 4096); + for (const auto & quantizer : weights.residual_quantizers) { + quantize_one(quantizer, 1024); + } + return x; +} + +modules::Snake1dWeights load_snake(core::BackendWeightStore & store, const assets::TensorSource & source, const std::string & name, int64_t channels) { + return {store.make_f32( + core::TensorShape::from_dims({channels}), + source.require_f32(name + ".alpha", {1, channels, 1}))}; +} + +CodecTransformerWeights load_transformer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t layers) { + CodecTransformerWeights out; + out.layers.reserve(static_cast(layers)); + for (int64_t layer = 0; layer < layers; ++layer) { + const std::string layer_prefix = prefix + ".layers." + std::to_string(layer); + CodecTransformerLayerWeights weights; + weights.attention_norm = binding::norm_weight_from_source(store, source, layer_prefix + ".attention_norm", kCodecDim); + weights.attention.q_weight = store.load_tensor(source, layer_prefix + ".attention.q_proj.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention.k_weight = store.load_tensor(source, layer_prefix + ".attention.k_proj.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention.v_weight = store.load_tensor(source, layer_prefix + ".attention.v_proj.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention.out_weight = store.load_tensor(source, layer_prefix + ".attention.wo.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention_scale = binding::layer_scale_from_named_source(store, source, layer_prefix + ".attention_layer_scale.gamma"); + weights.ffn_norm = binding::norm_weight_from_source(store, source, layer_prefix + ".ffn_norm", kCodecDim); + weights.feed_forward.gate_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w1.weight", storage_type, {kCodecIntermediate, kCodecDim}); + weights.feed_forward.down_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w2.weight", storage_type, {kCodecDim, kCodecIntermediate}); + weights.feed_forward.up_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w3.weight", storage_type, {kCodecIntermediate, kCodecDim}); + weights.ffn_scale = binding::layer_scale_from_named_source(store, source, layer_prefix + ".ffn_layer_scale.gamma"); + out.layers.push_back(std::move(weights)); + } + out.norm = binding::norm_weight_from_source(store, source, prefix + ".norm", kCodecDim); + return out; +} + +ResidualUnitWeights load_residual_unit( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels) { + ResidualUnitWeights out; + out.snake1 = load_snake(store, source, prefix + ".block.0", channels); + out.conv1 = binding::conv1d_from_named_source( + store, + source, + prefix + ".block.1.conv.weight", + prefix + ".block.1.conv.bias", + storage_type); + out.snake2 = load_snake(store, source, prefix + ".block.2", channels); + out.conv2 = binding::conv1d_from_named_source( + store, + source, + prefix + ".block.3.conv.weight", + prefix + ".block.3.conv.bias", + storage_type); + return out; +} + +ConvNeXtBlockWeights load_convnext( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels) { + ConvNeXtBlockWeights out; + out.dwconv = binding::depthwise_conv1d_from_source( + store, + source, + prefix + ".dwconv.conv", + storage_type, + channels, + 7, + true); + out.norm = binding::norm_from_source(store, source, prefix + ".norm", channels); + out.pwconv1 = binding::linear_from_source(store, source, prefix + ".pwconv1", storage_type, channels * 4, channels, true); + out.pwconv2 = binding::linear_from_source(store, source, prefix + ".pwconv2", storage_type, channels, channels * 4, true); + out.gamma = binding::layer_scale_from_named_source(store, source, prefix + ".gamma"); + return out; +} + +QuantizerUnitWeights load_quantizer_unit( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t codebook_size) { + QuantizerUnitWeights out; + out.in_proj = binding::conv1d_from_named_source( + store, + source, + prefix + ".in_proj.weight", + prefix + ".in_proj.bias", + storage_type); + out.out_proj = binding::conv1d_from_named_source( + store, + source, + prefix + ".out_proj.weight", + prefix + ".out_proj.bias", + storage_type); + const auto codebook = source.require_f32(prefix + ".codebook.weight", {codebook_size, 8}); + out.codebook = store.make_from_f32(core::TensorShape::from_dims({codebook_size, 8}), storage_type, codebook); + out.normalized_codebook = store.make_from_f32( + core::TensorShape::from_dims({codebook_size, 8}), + storage_type, + normalized_rows(codebook, codebook_size, 8)); + return out; +} + +std::shared_ptr load_weights( + const FishAudioAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType matmul_storage_type, + assets::TensorStorageType conv_storage_type) { + auto weights = std::make_shared(); + weights->store = std::make_shared(backend, backend_type, "Fish Audio codec", weight_context_bytes); + auto & store = *weights->store; + const auto & source = *assets.codec_weights; + + weights->encoder_first = binding::conv1d_from_named_source( + store, + source, + "encoder.block.0.conv.weight", + "encoder.block.0.conv.bias", + conv_storage_type); + int64_t encoder_channels = 64; + for (int64_t block_index = 0; block_index < 4; ++block_index) { + const std::string prefix = "encoder.block." + std::to_string(block_index + 1) + ".block"; + EncoderBlockWeights block; + block.residual1 = load_residual_unit(store, source, prefix + ".0", conv_storage_type, encoder_channels); + block.residual3 = load_residual_unit(store, source, prefix + ".1", conv_storage_type, encoder_channels); + block.residual9 = load_residual_unit(store, source, prefix + ".2", conv_storage_type, encoder_channels); + block.snake = load_snake(store, source, prefix + ".3", encoder_channels); + block.conv = binding::conv1d_from_named_source( + store, + source, + prefix + ".4.conv.weight", + prefix + ".4.conv.bias", + conv_storage_type); + encoder_channels *= 2; + if (block_index == 3) { + block.transformer = load_transformer(store, source, prefix + ".5", matmul_storage_type, 4); + } + weights->encoder_blocks.push_back(std::move(block)); + } + weights->encoder_final_snake = load_snake(store, source, "encoder.block.5", kCodecDim); + weights->encoder_final = binding::conv1d_from_named_source( + store, + source, + "encoder.block.6.conv.weight", + "encoder.block.6.conv.bias", + conv_storage_type); + + for (int64_t i = 0; i < 2; ++i) { + const std::string prefix = "quantizer.downsample." + std::to_string(i); + weights->downsample.push_back({ + binding::conv1d_from_named_source( + store, + source, + prefix + ".0.conv.weight", + prefix + ".0.conv.bias", + conv_storage_type), + load_convnext(store, source, prefix + ".1", matmul_storage_type, kCodecDim), + }); + } + weights->pre_module = load_transformer(store, source, "quantizer.pre_module", matmul_storage_type, kCodecTransformerLayers); + weights->semantic_quantizer = load_quantizer_unit(store, source, "quantizer.semantic_quantizer.quantizers.0", matmul_storage_type, 4096); + for (int64_t i = 0; i < assets.config.codec.quantizer_codebooks; ++i) { + weights->residual_quantizers.push_back( + load_quantizer_unit(store, source, "quantizer.quantizer.quantizers." + std::to_string(i), matmul_storage_type, 1024)); + } + weights->post_module = load_transformer(store, source, "quantizer.post_module", matmul_storage_type, kCodecTransformerLayers); + for (int64_t i = 0; i < 2; ++i) { + const std::string prefix = "quantizer.upsample." + std::to_string(i); + weights->upsample.push_back({ + binding::conv_transpose1d_from_named_source( + store, + source, + prefix + ".0.conv.weight", + prefix + ".0.conv.bias", + conv_storage_type), + load_convnext(store, source, prefix + ".1", matmul_storage_type, kCodecDim), + }); + } + + weights->decoder_first = binding::conv1d_from_named_source( + store, + source, + "decoder.model.0.conv.weight", + "decoder.model.0.conv.bias", + conv_storage_type); + int64_t decoder_channels = 1536; + for (int64_t block_index = 0; block_index < 4; ++block_index) { + const std::string prefix = "decoder.model." + std::to_string(block_index + 1) + ".block"; + DecoderBlockWeights block; + block.snake = load_snake(store, source, prefix + ".0", decoder_channels); + block.conv = binding::conv_transpose1d_from_named_source( + store, + source, + prefix + ".1.conv.weight", + prefix + ".1.conv.bias", + conv_storage_type); + decoder_channels /= 2; + block.residual1 = load_residual_unit(store, source, prefix + ".2", conv_storage_type, decoder_channels); + block.residual3 = load_residual_unit(store, source, prefix + ".3", conv_storage_type, decoder_channels); + block.residual9 = load_residual_unit(store, source, prefix + ".4", conv_storage_type, decoder_channels); + weights->decoder_blocks.push_back(std::move(block)); + } + weights->decoder_final_snake = load_snake(store, source, "decoder.model.5", 96); + weights->decoder_final = binding::conv1d_from_named_source( + store, + source, + "decoder.model.6.conv.weight", + "decoder.model.6.conv.bias", + conv_storage_type); + + store.upload(); + return weights; +} + +struct DecodeGraph { + DecodeGraph( + std::shared_ptr assets, + std::shared_ptr weights, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + int64_t frames) + : assets_(std::move(assets)), + weights_(std::move(weights)), + backend_(execution_context.backend()), + backend_type_(execution_context.backend_type()), + threads_(std::max(1, execution_context.config().threads)), + frame_capacity_(frames), + constants_(backend_, threads_, "Fish Audio codec decode constants") { + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio codec decode graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.codec.decode", backend_type_}; + constants_.begin_graph(); + for (int64_t codebook = 0; codebook < assets_->config.codec.total_codebooks; ++codebook) { + auto ids = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, frame_capacity_})); + ggml_set_input(ids.tensor); + code_inputs_.push_back(ids); + } + auto latent = build_decode_quantizer(ctx, constants_, code_inputs_, *weights_); + auto waveform = build_decoder(ctx, latent, *weights_); + output_ = waveform.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 1048576, false); + ggml_build_forward_expand(graph_, output_); + constants_.finish_graph(); + constants_.ensure_uploaded(); + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_))); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_.get(), graph_)) { + throw std::runtime_error("failed to allocate Fish Audio codec decode graph"); + } + } + + ~DecodeGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + } + + bool matches(int64_t frames, ggml_backend_t backend, int threads) const { + return frame_capacity_ >= frames && backend_ == backend && threads_ == std::max(1, threads); + } + + runtime::AudioBuffer run(const FishAudioCodes & codes) { + const int64_t codebooks = assets_->config.codec.total_codebooks; + if (codes.codebooks != codebooks || codes.frames <= 0 || + static_cast(codes.codes.size()) != codebooks * codes.frames) { + std::ostringstream oss; + oss << "Fish Audio codec decode code shape mismatch: expected_codebooks=" << codebooks + << " actual_codebooks=" << codes.codebooks + << " frames=" << codes.frames + << " values=" << codes.codes.size() + << " expected_values=" << (codebooks * codes.frames); + throw std::runtime_error(oss.str()); + } + if (codes.frames > frame_capacity_) { + throw std::runtime_error("Fish Audio codec decode request exceeds graph capacity"); + } + for (int64_t codebook = 0; codebook < codebooks; ++codebook) { + std::vector padded(static_cast(frame_capacity_), 0); + for (int64_t frame = 0; frame < codes.frames; ++frame) { + int32_t value = codes.codes[static_cast(codebook * codes.frames + frame)]; + if (codebook == 0) { + value = std::clamp(value, 0, 4095); + } else { + value = std::clamp(value, 0, 1023); + } + padded[static_cast(frame)] = value; + } + core::write_tensor_i32(code_inputs_[static_cast(codebook)], padded); + } + core::set_backend_threads(backend_, threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio codec decode graph compute failed"); + } + auto values = core::read_tensor_f32(output_); + const int64_t expected_samples = codes.frames * assets_->config.codec.frame_length; + if (static_cast(values.size()) > expected_samples) { + values.resize(static_cast(expected_samples)); + } + return runtime::AudioBuffer{assets_->config.codec.sample_rate, 1, std::move(values)}; + } + +private: + std::shared_ptr assets_; + std::shared_ptr weights_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + int64_t frame_capacity_ = 0; + std::unique_ptr ctx_; + std::vector code_inputs_; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; + common::ConstantTensorCache constants_; +}; + +struct EncodeGraph { + EncodeGraph( + std::shared_ptr assets, + std::shared_ptr weights, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + int64_t samples, + int64_t frames) + : assets_(std::move(assets)), + weights_(std::move(weights)), + backend_(execution_context.backend()), + backend_type_(execution_context.backend_type()), + threads_(std::max(1, execution_context.config().threads)), + sample_capacity_(samples), + frame_capacity_(frames), + constants_(backend_, threads_, "Fish Audio codec encode constants") { + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio codec encode graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.codec.encode", backend_type_}; + constants_.begin_graph(); + input_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, sample_capacity_})); + ggml_set_input(input_.tensor); + auto encoded = build_encoder(ctx, constants_, input_, *weights_); + trace_outputs_.push_back({"fish_audio.codec.encoder_latent", encoded}); + build_encode_quantizer(ctx, constants_, encoded, *weights_, code_outputs_, trace_outputs_); + graph_ = ggml_new_graph_custom(ctx_.get(), 1048576, false); + for (const auto & trace_output : trace_outputs_) { + ggml_set_output(trace_output.second.tensor); + ggml_build_forward_expand(graph_, trace_output.second.tensor); + } + for (ggml_tensor * code_output : code_outputs_) { + ggml_build_forward_expand(graph_, code_output); + } + constants_.finish_graph(); + constants_.ensure_uploaded(); + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_))); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_.get(), graph_)) { + throw std::runtime_error("failed to allocate Fish Audio codec encode graph"); + } + } + + ~EncodeGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + } + + bool matches(int64_t samples, int64_t frames, ggml_backend_t backend, int threads) const { + return sample_capacity_ >= samples && + frame_capacity_ >= frames && + backend_ == backend && + threads_ == std::max(1, threads); + } + + FishAudioCodes run(const runtime::AudioBuffer & audio) { + auto mono = prepare_codec_mono(audio, assets_->config.codec.sample_rate); + const int64_t original_samples = static_cast(mono.size()); + const int64_t padded_samples = ceil_div(original_samples, assets_->config.codec.frame_length) * assets_->config.codec.frame_length; + const int64_t frames = ceil_div(original_samples, assets_->config.codec.frame_length); + if (padded_samples > sample_capacity_ || frames > frame_capacity_) { + throw std::runtime_error("Fish Audio codec encode request exceeds graph capacity"); + } + mono.resize(static_cast(sample_capacity_), 0.0F); + core::write_tensor_f32(input_, mono); + core::set_backend_threads(backend_, threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio codec encode graph compute failed"); + } + if (engine::debug::trace_log_enabled()) { + for (const auto & trace_output : trace_outputs_) { + engine::debug::trace_log_f32( + trace_output.first, + dims_vector(trace_output.second.shape), + core::read_tensor_f32(trace_output.second.tensor)); + } + } + FishAudioCodes out; + out.codebooks = static_cast(code_outputs_.size()); + out.frames = frames; + out.codes.resize(static_cast(out.codebooks * out.frames)); + for (int64_t codebook = 0; codebook < out.codebooks; ++codebook) { + auto values = core::read_tensor_i32(code_outputs_[static_cast(codebook)]); + for (int64_t frame = 0; frame < out.frames; ++frame) { + out.codes[static_cast(codebook * out.frames + frame)] = values[static_cast(frame)]; + } + } + engine::debug::trace_log_i32( + "fish_audio.codec.reference_codes", + {out.codebooks, out.frames}, + out.codes); + return out; + } + +private: + std::shared_ptr assets_; + std::shared_ptr weights_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + int64_t sample_capacity_ = 0; + int64_t frame_capacity_ = 0; + std::unique_ptr ctx_; + core::TensorValue input_; + std::vector code_outputs_; + std::vector> trace_outputs_; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; + common::ConstantTensorCache constants_; +}; + +} // namespace + +class FishAudioCodecRuntime::Impl { +public: + Impl( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType matmul_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) + : assets_(std::move(assets)), + execution_(std::move(backend)), + threads_(std::max(1, threads)), + graph_arena_bytes_(graph_arena_bytes) { + weights_ = load_weights( + *assets_, + execution_.backend(), + execution_.backend_type(), + weight_context_bytes, + matmul_weight_storage_type, + conv_weight_storage_type); + } + + FishAudioCodes encode_reference(const runtime::AudioBuffer & audio) { + auto mono = prepare_codec_mono(audio, assets_->config.codec.sample_rate); + const int64_t samples = ceil_div(static_cast(mono.size()), assets_->config.codec.frame_length) * + assets_->config.codec.frame_length; + const int64_t frames = ceil_div(static_cast(mono.size()), assets_->config.codec.frame_length); + if (encode_graph_ == nullptr || !encode_graph_->matches(samples, frames, execution_.backend(), threads_)) { + encode_graph_ = std::make_unique(assets_, weights_, execution_, graph_arena_bytes_, samples, frames); + } + return encode_graph_->run(audio); + } + + runtime::AudioBuffer decode(const FishAudioCodes & codes) { + if (decode_graph_ == nullptr || !decode_graph_->matches(codes.frames, execution_.backend(), threads_)) { + decode_graph_ = std::make_unique(assets_, weights_, execution_, graph_arena_bytes_, codes.frames); + } + return decode_graph_->run(codes); + } + + void release_encode_graph() { + encode_graph_.reset(); + } + + void release_runtime_graphs() { + encode_graph_.reset(); + decode_graph_.reset(); + } + +private: + std::shared_ptr assets_; + core::ExecutionContext execution_; + int threads_ = 1; + size_t graph_arena_bytes_ = 0; + std::shared_ptr weights_; + std::unique_ptr encode_graph_; + std::unique_ptr decode_graph_; +}; + +FishAudioCodecRuntime::FishAudioCodecRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType matmul_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) + : impl_(std::make_unique( + std::move(assets), + std::move(backend), + threads, + graph_arena_bytes, + weight_context_bytes, + matmul_weight_storage_type, + conv_weight_storage_type)) {} + +FishAudioCodecRuntime::~FishAudioCodecRuntime() = default; + +FishAudioCodes FishAudioCodecRuntime::encode_reference(const runtime::AudioBuffer & audio) { + return impl_->encode_reference(audio); +} + +runtime::AudioBuffer FishAudioCodecRuntime::decode(const FishAudioCodes & codes) { + return impl_->decode(codes); +} + +void FishAudioCodecRuntime::release_encode_graph() { + impl_->release_encode_graph(); +} + +void FishAudioCodecRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/generator.cpp b/src/models/fish_audio/generator.cpp new file mode 100644 index 00000000..a5aa4586 --- /dev/null +++ b/src/models/fish_audio/generator.cpp @@ -0,0 +1,74 @@ +#include "engine/models/fish_audio/generator.h" + +#include "engine/framework/debug/profiler.h" + +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +using Clock = std::chrono::steady_clock; + +} // namespace + +FishAudioGenerator::FishAudioGenerator( + std::shared_ptr assets, + std::unique_ptr ar, + std::unique_ptr codec) + : assets_(std::move(assets)), + tokenizer_(assets_), + prompt_builder_(assets_, tokenizer_), + ar_(std::move(ar)), + codec_(std::move(codec)) { + if (assets_ == nullptr || ar_ == nullptr || codec_ == nullptr) { + throw std::runtime_error("Fish Audio generator requires assets, AR runtime, and codec runtime"); + } +} + +FishAudioGenerator::~FishAudioGenerator() = default; + +FishAudioCodes FishAudioGenerator::encode_reference(const runtime::AudioBuffer & audio) { + auto codes = codec_->encode_reference(audio); + codec_->release_encode_graph(); + return codes; +} + +FishAudioGenerationResult FishAudioGenerator::generate( + const FishAudioRequest & request, + const std::optional & reference_codes, + const std::optional & previous_turn, + bool mem_saver) { + engine::debug::trace_log_scalar("fish_audio.request.has_reference", request.reference.has_value()); + engine::debug::trace_log_scalar("fish_audio.request.text_chars", static_cast(request.text.size())); + engine::debug::trace_log_scalar("fish_audio.request.has_previous_turn", previous_turn.has_value()); + engine::debug::trace_log_scalar("fish_audio.sampler.seed", request.generation.seed); + const auto prompt_start = Clock::now(); + const auto prompt = prompt_builder_.build(request, reference_codes, previous_turn); + engine::debug::timing_log_scalar( + "fish_audio.prompt_build_ms", + engine::debug::elapsed_ms(prompt_start, Clock::now())); + + const auto ar_start = Clock::now(); + FishAudioGenerationResult result; + result.codes = ar_->generate(prompt, request.generation); + engine::debug::trace_log_scalar("fish_audio.generated.frames", result.codes.frames); + engine::debug::trace_log_scalar("fish_audio.generated.codebooks", result.codes.codebooks); + engine::debug::timing_log_scalar( + "fish_audio.ar_generate_ms", + engine::debug::elapsed_ms(ar_start, Clock::now())); + + const auto decode_start = Clock::now(); + result.audio = codec_->decode(result.codes); + engine::debug::timing_log_scalar( + "fish_audio.codec_decode_ms", + engine::debug::elapsed_ms(decode_start, Clock::now())); + codec_->release_runtime_graphs(); + if (mem_saver) { + ar_->release_runtime_graphs(); + } + return result; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp new file mode 100644 index 00000000..f8f6a3d7 --- /dev/null +++ b/src/models/fish_audio/loader.cpp @@ -0,0 +1,146 @@ +#include "engine/models/fish_audio/loader.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/models/fish_audio/session.h" + +#include +#include + +namespace engine::models::fish_audio { +namespace { + +runtime::ModelMetadata metadata(const FishAudioAssets &) { + runtime::ModelMetadata out; + out.family = "fish_audio"; + out.variant = "s2-pro"; + out.description = "Fish Audio S2-Pro loaded from prepared local assets."; + out.config_candidates = {"config.json", "tokenizer_config.json", "tokenizer.json"}; + out.weight_candidates = {"model_audio_cpp.safetensors.index.json", "codec.safetensors", "model.gguf"}; + return out; +} + +runtime::CapabilitySet capabilities(const FishAudioAssets &) { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.languages = {"en", "zh", "auto"}; + out.supports_speaker_reference = true; + out.supports_style_condition = true; + return out; +} + +runtime::ModelCliInterface cli(const FishAudioAssets &) { + runtime::ModelCliInterface out; + out.request_options = { + {"reference_text", "TEXT", "Reference transcript used with speaker reference audio."}, + {"max_new_tokens", "N", "Maximum Fish Audio semantic tokens to generate; default 1024, 0 uses the default."}, + {"text_chunk_size", "N", "Long-form text chunk size; default 200."}, + {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode."}, + {"top_p", "FLOAT", "Top-p sampling value."}, + {"top_k", "N", "Top-k sampling value."}, + {"temperature", "FLOAT", "Sampling temperature."}, + {"seed", "N", "Sampling seed for reproducible output; omitted uses a random seed."}, + }; + out.session_options = { + {"fish_audio.mem_saver", "true|false", "Release cached AR runtime graphs after each request; default false."}, + {"fish_audio.reference_cache_slots", "n", "Prepared reference-audio cache slots; default 1."}, + {"fish_audio.weight_type", "native|f32|f16|bf16|q8_0", "AR matmul weight storage type; default native."}, + {"fish_audio.codec_weight_type", "native|f32|f16|q8_0", "Codec conv/matmul weight storage type; default native."}, + }; + return out; +} + +class FishAudioLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "fish_audio"; + } + + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + out.supports_style_condition = true; + return out; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + if (request.family_hint.has_value() && *request.family_hint != family()) { + return false; + } + try { + const auto package_spec = engine::assets::default_model_package_spec_path(family()); + (void) engine::assets::load_resource_bundle_from_package_spec(request.model_path, package_spec); + return true; + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto assets = load_fish_audio_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::assets::default_model_package_spec_path(family()); + inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Files); + inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Tensors); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_fish_audio_model(request.model_path); + } +}; + +} // namespace + +FishAudioLoadedModel::FishAudioLoadedModel( + 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 & FishAudioLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & FishAudioLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr FishAudioLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.task != runtime::VoiceTaskKind::Tts || task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Fish Audio S2-Pro supports offline TTS sessions"); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_fish_audio_model(const std::filesystem::path & model_path) { + auto assets = load_fish_audio_assets(model_path); + return std::make_unique( + metadata(*assets), + capabilities(*assets), + std::move(assets)); +} + +std::shared_ptr make_fish_audio_loader() { + return std::make_shared(); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/prompt_builder.cpp b/src/models/fish_audio/prompt_builder.cpp new file mode 100644 index 00000000..c4da2492 --- /dev/null +++ b/src/models/fish_audio/prompt_builder.cpp @@ -0,0 +1,124 @@ +#include "engine/models/fish_audio/prompt_builder.h" + +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +void append_tokens(std::vector & out, const std::vector & tokens) { + out.insert(out.end(), tokens.begin(), tokens.end()); +} + +struct CodeSpan { + int64_t start = 0; + const FishAudioCodes * codes = nullptr; +}; + +std::string reference_text_with_speakers(const std::string & text) { + static const std::regex speaker_re(R"(<\|speaker:\d+\|>)"); + if (std::regex_search(text, speaker_re)) { + return text; + } + return "<|speaker:0|>" + text; +} + +void append_code_span( + std::vector & row0, + std::vector & spans, + const FishAudioTextTokenizer & tokenizer, + const FishAudioCodes & codes, + int64_t expected_codebooks) { + if (codes.codebooks != expected_codebooks) { + throw std::runtime_error("Fish Audio prompt codebook count mismatch"); + } + const int64_t start = static_cast(row0.size()); + const int32_t semantic_begin = tokenizer.semantic_begin_id(); + for (int64_t frame = 0; frame < codes.frames; ++frame) { + row0.push_back(semantic_begin + codes.codes[static_cast(frame)]); + } + spans.push_back({start, &codes}); +} + +} // namespace + +FishAudioPromptBuilder::FishAudioPromptBuilder( + std::shared_ptr assets, + FishAudioTextTokenizer tokenizer) + : assets_(std::move(assets)), + tokenizer_(std::move(tokenizer)) { + if (assets_ == nullptr) { + throw std::runtime_error("Fish Audio prompt builder requires assets"); + } +} + +FishAudioPrompt FishAudioPromptBuilder::build( + const FishAudioRequest & request, + const std::optional & reference_codes, + const std::optional & previous_turn) const { + if (request.text.empty()) { + throw std::runtime_error("Fish Audio request text must not be empty"); + } + const int64_t rows = assets_->config.fast.num_codebooks + 1; + if (rows <= 1) { + throw std::runtime_error("Fish Audio prompt rows are invalid"); + } + + std::vector row0; + std::vector code_spans; + if (request.reference.has_value()) { + if (!reference_codes.has_value()) { + throw std::runtime_error("Fish Audio reference request requires encoded reference codes"); + } + append_tokens(row0, tokenizer_.encode("<|im_start|>system\n")); + append_tokens(row0, tokenizer_.encode("convert the provided text to speech reference to the following:\n\nText:\n")); + append_tokens(row0, tokenizer_.encode(reference_text_with_speakers(request.reference->text))); + append_tokens(row0, tokenizer_.encode("\n\nSpeech:\n")); + append_code_span(row0, code_spans, tokenizer_, *reference_codes, assets_->config.fast.num_codebooks); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + } else { + append_tokens(row0, tokenizer_.encode("<|im_start|>system\n")); + append_tokens(row0, tokenizer_.encode("convert the provided text to speech")); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + } + if (previous_turn.has_value()) { + append_tokens(row0, tokenizer_.encode("<|im_start|>user\n")); + append_tokens(row0, tokenizer_.encode(previous_turn->text)); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + append_tokens(row0, tokenizer_.encode("<|im_start|>assistant\n<|voice|>")); + append_code_span(row0, code_spans, tokenizer_, previous_turn->codes, assets_->config.fast.num_codebooks); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + } + append_tokens(row0, tokenizer_.encode("<|im_start|>user\n")); + append_tokens(row0, tokenizer_.encode(request.text)); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + append_tokens(row0, tokenizer_.encode("<|im_start|>assistant\n<|voice|>")); + + FishAudioPrompt prompt; + prompt.codebook_rows = rows; + prompt.steps = static_cast(row0.size()); + prompt.text = request.text; + prompt.matrix.assign(static_cast(rows * prompt.steps), 0); + for (int64_t step = 0; step < prompt.steps; ++step) { + prompt.matrix[static_cast(step)] = row0[static_cast(step)]; + } + for (const auto & span : code_spans) { + if (span.codes == nullptr) { + throw std::runtime_error("Fish Audio prompt code span is missing codes"); + } + for (int64_t frame = 0; frame < span.codes->frames; ++frame) { + const int64_t step = span.start + frame; + if (step < 0 || step >= prompt.steps) { + throw std::runtime_error("Fish Audio prompt code span exceeds prompt length"); + } + for (int64_t codebook = 0; codebook < span.codes->codebooks; ++codebook) { + prompt.matrix[static_cast((codebook + 1) * prompt.steps + step)] = + span.codes->codes[static_cast(codebook * span.codes->frames + frame)]; + } + } + } + return prompt; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/session.cpp b/src/models/fish_audio/session.cpp new file mode 100644 index 00000000..f23c6d28 --- /dev/null +++ b/src/models/fish_audio/session.cpp @@ -0,0 +1,459 @@ +#include "engine/models/fish_audio/session.h" + +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/session.h" +#include "engine/framework/text/chunking.h" +#include "engine/models/fish_audio/ar.h" +#include "engine/models/fish_audio/codec.h" +#include "engine/models/fish_audio/generator.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +using Clock = std::chrono::steady_clock; +namespace fs = std::filesystem; + +constexpr size_t kDefaultArGraphArenaBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecGraphArenaBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kDefaultArWeightContextBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecWeightContextBytes = 512ull * 1024ull * 1024ull; +constexpr int64_t kDefaultReferenceCacheSlots = 1; +constexpr const char * kReferenceTextOption = "reference_text"; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Fish Audio session requires assets"); + } + return assets; +} + +assets::TensorStorageType option_weight_type( + const runtime::SessionOptions & options, + const char * key, + assets::TensorStorageType fallback) { + const auto it = options.options.find(key); + if (it == options.options.end()) { + return fallback; + } + return assets::parse_tensor_storage_type(it->second); +} + +void validate_ar_weight_storage(assets::TensorStorageType type, const char * option_name) { + if (type == assets::TensorStorageType::Native || + type == assets::TensorStorageType::F32 || + type == assets::TensorStorageType::F16 || + type == assets::TensorStorageType::BF16 || + type == assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " supports native/f32/f16/bf16/q8_0"); +} + +void validate_codec_weight_storage(assets::TensorStorageType type, const char * option_name) { + if (type == assets::TensorStorageType::Native || + type == assets::TensorStorageType::F32 || + type == assets::TensorStorageType::F16 || + type == assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " supports native/f32/f16/q8_0"); +} + +bool mem_saver_from_options(const runtime::SessionOptions & options) { + if (const auto value = runtime::find_option(options.options, {"fish_audio.mem_saver", "mem_saver"})) { + return runtime::parse_bool_option(*value, "fish_audio.mem_saver"); + } + return false; +} + +std::size_t resolve_reference_cache_slots(const runtime::SessionOptions & options) { + const int64_t slots = runtime::parse_i64_option( + options.options, + {"fish_audio.reference_cache_slots", "reference_cache_slots"}) + .value_or(kDefaultReferenceCacheSlots); + if (slots < 0) { + throw std::runtime_error("fish_audio.reference_cache_slots must be non-negative"); + } + if (static_cast(slots) > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("fish_audio.reference_cache_slots is too large"); + } + return static_cast(slots); +} + +uint64_t mix_reference_key(uint64_t key, uint64_t value) { + key ^= value; + key *= 1099511628211ull; + return key; +} + +uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { + uint64_t key = 1469598103934665603ull; + for (const float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + key = mix_reference_key(key, static_cast(bits)); + } + return key; +} + +FishAudioGenerationOptions generation_options_from_request(const runtime::TaskRequest & request) { + FishAudioGenerationOptions options; + if (const auto value = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"})) { + if (*value < 0) { + throw std::runtime_error("Fish Audio max_new_tokens must be non-negative"); + } + if (*value > 0) { + options.max_new_tokens = *value; + } + } + options.text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(options.text_chunk_size); + options.top_p = runtime::parse_float_option(request.options, {"top_p"}).value_or(options.top_p); + options.top_k = runtime::parse_int_option(request.options, {"top_k"}).value_or(options.top_k); + options.temperature = runtime::parse_float_option(request.options, {"temperature"}).value_or(options.temperature); + options.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(runtime::random_u32_seed()); + if (options.max_new_tokens <= 0) { + throw std::runtime_error("Fish Audio max_new_tokens must be positive after default resolution"); + } + if (options.text_chunk_size <= 0) { + throw std::runtime_error("Fish Audio text_chunk_size must be positive"); + } + if (!(options.top_p > 0.0F && options.top_p <= 1.0F)) { + throw std::runtime_error("Fish Audio top_p must be in (0, 1]"); + } + if (options.top_k <= 0) { + throw std::runtime_error("Fish Audio top_k must be positive"); + } + if (!(options.temperature > 0.0F && options.temperature < 2.0F)) { + throw std::runtime_error("Fish Audio temperature must be in (0, 2)"); + } + return options; +} + +std::string lower_ascii(std::string value) { + std::transform( + value.begin(), + value.end(), + value.begin(), + [](unsigned char ch) { return static_cast(std::tolower(ch)); }); + return value; +} + +bool valid_reference_id_char(unsigned char ch) { + return std::isalnum(ch) != 0 || ch == '-' || ch == '_' || ch == ' '; +} + +void validate_reference_id(const std::string & id) { + if (id.empty() || id.size() > 255) { + throw std::runtime_error( + "Fish Audio cached_voice_id must be 1-255 characters"); + } + for (const unsigned char ch : id) { + if (!valid_reference_id_char(ch)) { + throw std::runtime_error( + "Fish Audio cached_voice_id may only contain alphanumeric characters, hyphens, underscores, and spaces"); + } + } +} + +bool is_supported_saved_reference_audio(const fs::path & path) { + return lower_ascii(path.extension().string()) == ".wav"; +} + +std::vector collect_saved_reference_audio_files(const fs::path & directory) { + std::vector files; + for (const auto & entry : fs::recursive_directory_iterator(directory)) { + if (!entry.is_regular_file() || !is_supported_saved_reference_audio(entry.path())) { + continue; + } + auto lab_path = entry.path(); + lab_path.replace_extension(".lab"); + if (engine::io::is_existing_file(lab_path)) { + files.push_back(entry.path()); + } + } + std::sort(files.begin(), files.end()); + return files; +} + +runtime::AudioBuffer read_saved_reference_audio(const fs::path & path) { + auto wav = engine::audio::read_wav_f32(path); + return runtime::AudioBuffer{wav.sample_rate, wav.channels, std::move(wav.samples)}; +} + +FishAudioReference load_saved_reference( + const FishAudioAssets & assets, + const std::string & reference_id) { + validate_reference_id(reference_id); + const auto reference_dir = engine::io::require_directory( + assets.resources.model_root() / "references" / reference_id, + "Fish Audio cached voice reference"); + const auto audio_files = collect_saved_reference_audio_files(reference_dir); + if (audio_files.empty()) { + throw std::runtime_error( + "Fish Audio cached_voice_id '" + reference_id + + "' requires one WAV reference with a matching .lab file under " + + reference_dir.string()); + } + if (audio_files.size() > 1) { + throw std::runtime_error( + "Fish Audio cached_voice_id '" + reference_id + + "' has multiple WAV references with .lab files; the C++ session expects exactly one reference pair"); + } + auto lab_path = audio_files.front(); + lab_path.replace_extension(".lab"); + return FishAudioReference{ + read_saved_reference_audio(audio_files.front()), + engine::io::read_text_file(lab_path), + reference_id}; +} + +std::string reference_cache_id_from_voice(const std::optional & voice) { + if (voice.has_value() && + voice->speaker.has_value() && + voice->speaker->cached_voice_id.has_value() && + !voice->speaker->cached_voice_id->empty()) { + return *voice->speaker->cached_voice_id; + } + return {}; +} + +bool has_reference_selector(const std::optional & voice) { + if (!voice.has_value() || !voice->speaker.has_value()) { + return false; + } + const auto & speaker = *voice->speaker; + return speaker.audio.has_value() || + (speaker.cached_voice_id.has_value() && !speaker.cached_voice_id->empty()); +} + +std::optional reference_from_voice( + const FishAudioAssets & assets, + const std::optional & voice, + const std::unordered_map & options, + const char * role) { + if (!has_reference_selector(voice)) { + return std::nullopt; + } + const auto & speaker = *voice->speaker; + if (speaker.audio.has_value()) { + auto reference_text = runtime::find_option(options, {kReferenceTextOption}); + if (!reference_text.has_value()) { + throw std::runtime_error( + std::string(role) + " with inline reference audio requires reference_text option"); + } + return FishAudioReference{ + speaker.audio, + *reference_text, + reference_cache_id_from_voice(voice)}; + } + return load_saved_reference(assets, *speaker.cached_voice_id); +} + +} // namespace + +FishAudioSession::FishAudioSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + reference_cache_(resolve_reference_cache_slots(this->options())) { + if (task_.task != runtime::VoiceTaskKind::Tts || task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Fish Audio only supports offline TTS sessions"); + } + const auto ar_weight_type = + option_weight_type(options, "fish_audio.weight_type", assets::TensorStorageType::Native); + const auto codec_weight_type = + option_weight_type(options, "fish_audio.codec_weight_type", assets::TensorStorageType::Native); + validate_ar_weight_storage(ar_weight_type, "fish_audio.weight_type"); + validate_codec_weight_storage(codec_weight_type, "fish_audio.codec_weight_type"); + const int threads = options.backend.threads > 0 ? options.backend.threads : 1; + auto ar = std::make_unique( + assets_, + options.backend, + threads, + runtime::parse_size_mb_option(options.options, {"fish_audio.ar_graph_arena_mb"}, kDefaultArGraphArenaBytes), + runtime::parse_size_mb_option(options.options, {"fish_audio.ar_weight_context_mb"}, kDefaultArWeightContextBytes), + ar_weight_type); + auto codec = std::make_unique( + assets_, + options.backend, + threads, + runtime::parse_size_mb_option(options.options, {"fish_audio.codec_graph_arena_mb"}, kDefaultCodecGraphArenaBytes), + runtime::parse_size_mb_option(options.options, {"fish_audio.codec_weight_context_mb"}, kDefaultCodecWeightContextBytes), + codec_weight_type, + codec_weight_type); + generator_ = std::make_unique( + assets_, + std::move(ar), + std::move(codec)); + assets_->model_weights->release_storage(); + assets_->codec_weights->release_storage(); +} + +FishAudioSession::~FishAudioSession() = default; + +std::string FishAudioSession::family() const { + return "fish_audio"; +} + +runtime::VoiceTaskKind FishAudioSession::task_kind() const { + return task_.task; +} + +runtime::RunMode FishAudioSession::run_mode() const { + return task_.mode; +} + +bool FishAudioSession::ReferenceCacheKeyEqual::operator()( + const ReferenceCacheKey & lhs, + const ReferenceCacheKey & rhs) const { + return lhs.source_id == rhs.source_id && + lhs.sample_rate == rhs.sample_rate && + lhs.channels == rhs.channels && + lhs.sample_count == rhs.sample_count && + lhs.sample_hash == rhs.sample_hash; +} + +void FishAudioSession::prepare(const runtime::SessionPreparationRequest & request) { + defaults_.reset(); + FishAudioRequest defaults; + bool has_defaults = false; + if (request.text.has_value()) { + defaults.text = request.text->text; + has_defaults = true; + } + if (const auto value = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"})) { + if (*value < 0) { + throw std::runtime_error("Fish Audio max_new_tokens must be non-negative"); + } + if (*value > 0) { + defaults.generation.max_new_tokens = *value; + } + } + if (auto reference = reference_from_voice(*assets_, request.voice, request.options, "Fish Audio prepare"); + reference.has_value()) { + defaults.reference = std::move(*reference); + has_defaults = true; + } + if (has_defaults) { + defaults_ = std::move(defaults); + } + mark_prepared(); +} + +FishAudioRequest FishAudioSession::make_request(const runtime::TaskRequest & request) const { + FishAudioRequest out = defaults_.value_or(FishAudioRequest{}); + if (request.text_input.has_value()) { + out.text = request.text_input->text; + } + out.generation = generation_options_from_request(request); + if (auto reference = reference_from_voice(*assets_, request.voice, request.options, "Fish Audio request"); + reference.has_value()) { + out.reference = std::move(*reference); + } else if (request.text_input.has_value()) { + out.reference = std::nullopt; + } + if (out.text.empty()) { + throw std::runtime_error("Fish Audio request text must not be empty"); + } + return out; +} + +const FishAudioCodes & FishAudioSession::resolve_reference_codes(const FishAudioReference & reference) { + ReferenceCacheKey key; + key.source_id = reference.cache_id; + if (reference.cache_id.empty() && !reference.audio.has_value()) { + throw std::runtime_error("Fish Audio cached reference requires reference audio or a reference id"); + } + if (reference.audio.has_value() && reference.cache_id.empty()) { + key.sample_rate = reference.audio->sample_rate; + key.channels = reference.audio->channels; + key.sample_count = static_cast(reference.audio->samples.size()); + key.sample_hash = hash_audio_samples(*reference.audio); + } + if (const auto * cached = reference_cache_.find(key)) { + engine::debug::trace_log_scalar("fish_audio.reference_cache.hit", 1); + engine::debug::trace_log_scalar("fish_audio.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.entries", static_cast(reference_cache_.size())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.evicted", 0); + return cached->codes; + } + if (!reference.audio.has_value()) { + throw std::runtime_error("Fish Audio reference id is not cached and no reference audio was provided"); + } + const bool will_evict = reference_cache_.capacity() > 0 && reference_cache_.size() >= reference_cache_.capacity(); + const auto start = Clock::now(); + ReferenceCacheEntry entry; + entry.codes = generator_->encode_reference(*reference.audio); + engine::debug::trace_log_scalar("fish_audio.reference.frames", entry.codes.frames); + engine::debug::trace_log_scalar("fish_audio.reference.codebooks", entry.codes.codebooks); + if (reference_cache_.capacity() == 0) { + uncached_reference_ = std::move(entry); + } else { + reference_cache_.put(key, std::move(entry)); + } + engine::debug::trace_log_scalar("fish_audio.reference_cache.hit", 0); + engine::debug::trace_log_scalar("fish_audio.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.entries", static_cast(reference_cache_.size())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.evicted", will_evict ? 1 : 0); + engine::debug::timing_log_scalar("fish_audio.reference_encode_ms", engine::debug::elapsed_ms(start, Clock::now())); + if (reference_cache_.capacity() == 0) { + return uncached_reference_->codes; + } + const auto * cached = reference_cache_.find(key); + if (cached == nullptr) { + throw std::runtime_error("Fish Audio reference cache insert failed"); + } + return cached->codes; +} + +runtime::TaskResult FishAudioSession::run(const runtime::TaskRequest & request) { + require_prepared("Fish Audio run()"); + const auto wall_start = Clock::now(); + const bool mem_saver = mem_saver_from_options(options()); + const auto request_options = generation_options_from_request(request); + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + const auto chunk_requests = runtime::chunk_text_request(request, request_options.text_chunk_size, text_chunk_mode); + engine::debug::trace_log_scalar("fish_audio.text_chunk_size", request_options.text_chunk_size); + engine::debug::trace_log_scalar("fish_audio.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + engine::debug::trace_log_scalar("fish_audio.text_chunk_count", static_cast(chunk_requests.size())); + + runtime::AudioBuffer merged_audio; + std::optional reference_codes = std::nullopt; + std::optional previous_turn = std::nullopt; + for (size_t chunk_index = 0; chunk_index < chunk_requests.size(); ++chunk_index) { + const auto & chunk_request = chunk_requests[chunk_index]; + auto fish_request = make_request(chunk_request); + if (fish_request.reference.has_value() && !reference_codes.has_value()) { + reference_codes = resolve_reference_codes(*fish_request.reference); + } + auto generated = generator_->generate(fish_request, reference_codes, previous_turn, mem_saver); + runtime::append_audio_buffer(merged_audio, generated.audio); + if (chunk_requests.size() > 1) { + previous_turn = FishAudioConversationTurn{fish_request.text, std::move(generated.codes)}; + } + } + runtime::TaskResult result; + result.audio_output = std::move(merged_audio); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); + return result; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/tokenizer_text.cpp b/src/models/fish_audio/tokenizer_text.cpp new file mode 100644 index 00000000..43839f3f --- /dev/null +++ b/src/models/fish_audio/tokenizer_text.cpp @@ -0,0 +1,69 @@ +#include "engine/models/fish_audio/tokenizer_text.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::models::fish_audio { +namespace { + +int32_t require_token_id(const engine::tokenizers::LlamaBpeTokenizer & tokenizer, const std::string & token) { + const auto id = tokenizer.find_token_id(token); + if (!id.has_value()) { + throw std::runtime_error("Fish Audio tokenizer missing token: " + token); + } + return *id; +} + +} // namespace + +struct FishAudioTextTokenizer::Impl { + explicit Impl(std::shared_ptr input_assets) + : assets(std::move(input_assets)), + tokenizer(engine::tokenizers::LlamaBpeTokenizerSpec{ + {}, + {}, + assets->resources.require_file("tokenizer_config"), + assets->resources.require_file("tokenizer_json"), + engine::tokenizers::LlamaBpePreTokenizer::Qwen2, + }), + im_end(require_token_id(tokenizer, "<|im_end|>")), + semantic_begin(static_cast(assets->config.semantic_start_token_id)), + semantic_end(static_cast(assets->config.semantic_end_token_id)) {} + + std::shared_ptr assets; + engine::tokenizers::LlamaBpeTokenizer tokenizer; + int32_t im_end = 0; + int32_t semantic_begin = 0; + int32_t semantic_end = 0; +}; + +FishAudioTextTokenizer::FishAudioTextTokenizer(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Fish Audio text tokenizer requires assets"); + } + impl_ = std::make_shared(std::move(assets)); +} + +std::vector FishAudioTextTokenizer::encode(const std::string & text) const { + return impl_->tokenizer.encode(text, true); +} + +int32_t FishAudioTextTokenizer::token_id(const std::string & token) const { + return require_token_id(impl_->tokenizer, token); +} + +int32_t FishAudioTextTokenizer::im_end_id() const noexcept { + return impl_->im_end; +} + +int32_t FishAudioTextTokenizer::semantic_begin_id() const noexcept { + return impl_->semantic_begin; +} + +int32_t FishAudioTextTokenizer::semantic_end_id() const noexcept { + return impl_->semantic_end; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/higgs_audio_tts/ar.cpp b/src/models/higgs_audio_tts/ar.cpp new file mode 100644 index 00000000..c113cb04 --- /dev/null +++ b/src/models/higgs_audio_tts/ar.cpp @@ -0,0 +1,1367 @@ +#include "engine/models/higgs_audio_tts/ar.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/attention/qwen_causal_decoder.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; +constexpr int64_t kLayerwisePrefillMinSteps = 2048; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConfig & config) { + modules::QwenDecoderStackConfig out; + out.hidden_size = config.hidden_size; + out.num_attention_heads = config.num_attention_heads; + out.num_key_value_heads = config.num_key_value_heads; + out.head_dim = config.head_dim; + out.intermediate_size = config.intermediate_size; + out.layers = config.num_hidden_layers; + out.rms_norm_eps = config.rms_norm_eps; + out.rope_theta = config.rope_theta; + out.attention_precision = GGML_PREC_F32; + out.projection_precision = GGML_PREC_DEFAULT; + out.qkv_layout = modules::QwenDecoderQKVLayout::Separate; + out.use_qk_norm = true; + out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::FlashWithPrefix; + out.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + out.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; + return out; +} + +class HiggsQwenDecoderComponent { +public: + HiggsQwenDecoderComponent(const HiggsTextConfig & config, bool packed_qkv) + : stack_config_(make_higgs_qwen_stack_config(config)), + layer_config_(modules::qwen_decoder_layer_config_from_stack(stack_config_)), + layer_module_([&] { + layer_config_.qkv_layout = packed_qkv + ? modules::QwenDecoderQKVLayout::PackedQKV + : modules::QwenDecoderQKVLayout::Separate; + return layer_config_; + }()) {} + + modules::QwenDecoderLayerOutputs build_prefill_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenDecoderLayerWeights & weights, + const core::TensorValue & attention_mask, + const std::optional & prefix_key = std::nullopt, + const std::optional & prefix_value = std::nullopt) const { + return layer_module_.build(ctx, input, positions, weights, prefix_key, prefix_value, attention_mask); + } + + modules::QwenDecoderLayerOutputs build_decode_layer( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask) const { + return layer_module_.build_with_static_cache_tail( + ctx, + graph, + input, + positions, + weights, + cache_key, + cache_value, + cache_slot, + attention_mask); + } + +private: + modules::QwenDecoderStackConfig stack_config_; + modules::QwenDecoderLayerConfig layer_config_; + modules::QwenDecoderLayerModule layer_module_; +}; + +core::TensorValue higgs_cache_view( + core::ModuleBuildContext & ctx, + const core::TensorValue & cache, + int64_t start, + int64_t steps, + int64_t heads, + int64_t head_dim) { + if (start < 0 || steps <= 0 || start + steps > cache.shape.dims[1]) { + throw std::runtime_error("Higgs TTS AR cache view range is invalid"); + } + return core::wrap_tensor( + ggml_view_4d( + ctx.ggml, + cache.tensor, + head_dim, + heads, + steps, + 1, + cache.tensor->nb[1], + cache.tensor->nb[2], + cache.tensor->nb[3], + static_cast(start) * cache.tensor->nb[2]), + core::TensorShape::from_dims({1, steps, heads, head_dim}), + cache.type); +} + +modules::QwenDecoderLayerWeights load_layer_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const HiggsTextConfig & config, + int64_t layer_index, + assets::TensorStorageType storage_type) { + const std::string prefix = "body.layers." + std::to_string(layer_index); + const int64_t q_out = config.num_attention_heads * config.head_dim; + const int64_t kv_out = config.num_key_value_heads * config.head_dim; + modules::QwenDecoderLayerWeights weights; + weights.input_norm = { + store.load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size}), + std::nullopt, + }; + { + const auto q = source.require_tensor( + prefix + ".self_attn.q_proj.weight", + storage_type, + {q_out, config.hidden_size}); + const auto k = source.require_tensor( + prefix + ".self_attn.k_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + const auto v = source.require_tensor( + prefix + ".self_attn.v_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + if (q.type != k.type || q.type != v.type) { + throw std::runtime_error("Higgs TTS packed QKV weights require matching storage types"); + } + std::vector packed; + packed.reserve(q.bytes.size() + k.bytes.size() + v.bytes.size()); + packed.insert(packed.end(), q.bytes.begin(), q.bytes.end()); + packed.insert(packed.end(), k.bytes.begin(), k.bytes.end()); + packed.insert(packed.end(), v.bytes.begin(), v.bytes.end()); + weights.self_attention.qkv_weight = store.make_tensor( + core::TensorShape::from_dims({q_out + 2 * kv_out, config.hidden_size}), + q.type, + packed.data(), + packed.size()); + } + weights.self_attention.out_weight = store.load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + storage_type, + {config.hidden_size, q_out}); + weights.q_norm = { + store.load_f32_tensor(source, prefix + ".self_attn.q_norm.weight", {config.head_dim}), + std::nullopt, + }; + weights.k_norm = { + store.load_f32_tensor(source, prefix + ".self_attn.k_norm.weight", {config.head_dim}), + std::nullopt, + }; + weights.post_norm = { + store.load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.hidden_size}), + std::nullopt, + }; + { + const auto gate = source.require_tensor( + prefix + ".mlp.gate_proj.weight", + storage_type, + {config.intermediate_size, config.hidden_size}); + const auto up = source.require_tensor( + prefix + ".mlp.up_proj.weight", + storage_type, + {config.intermediate_size, config.hidden_size}); + if (gate.type != up.type) { + throw std::runtime_error("Higgs TTS packed gate/up weights require matching storage types"); + } + std::vector packed; + packed.reserve(gate.bytes.size() + up.bytes.size()); + packed.insert(packed.end(), gate.bytes.begin(), gate.bytes.end()); + packed.insert(packed.end(), up.bytes.begin(), up.bytes.end()); + weights.mlp.gate_up_proj = modules::LinearWeights{ + store.make_tensor( + core::TensorShape::from_dims({config.intermediate_size * 2, config.hidden_size}), + gate.type, + packed.data(), + packed.size()), + std::nullopt, + }; + } + weights.mlp.down_proj = { + store.load_tensor( + source, + prefix + ".mlp.down_proj.weight", + storage_type, + {config.hidden_size, config.intermediate_size}), + std::nullopt, + }; + return weights; +} + +HiggsQwenDecoderStackWeights load_decoder_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const HiggsTextConfig & config, + assets::TensorStorageType storage_type) { + HiggsQwenDecoderStackWeights weights; + weights.layers.reserve(static_cast(config.num_hidden_layers)); + for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { + weights.layers.push_back(load_layer_weights(store, source, config, layer, storage_type)); + } + return weights; +} + +core::TensorValue build_higgs_decode_code_embedding( + core::ModuleBuildContext & ctx, + const HiggsARWeights & weights, + const HiggsConfig & config, + ggml_tensor * fused_code_ids) { + auto code_ids = core::wrap_tensor( + fused_code_ids, + core::TensorShape::from_dims({config.audio.num_codebooks}), + GGML_TYPE_I32); + auto code = modules::EmbeddingModule( + {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}) + .build(ctx, code_ids, weights.modality_embedding); + code = modules::ReduceSumModule({0}).build(ctx, code); + return core::reshape_tensor(ctx, code, core::TensorShape::from_dims({1, 1, config.text.hidden_size})); +} + +core::TensorValue build_higgs_prefill_input_embedding( + core::ModuleBuildContext & ctx, + const HiggsARWeights & weights, + const HiggsConfig & config, + ggml_tensor * text_tokens, + ggml_tensor * fused_code_ids, + ggml_tensor * text_gate, + ggml_tensor * code_gate, + int64_t steps) { + auto text_ids = core::wrap_tensor(text_tokens, core::TensorShape::from_dims({steps}), GGML_TYPE_I32); + auto text = modules::EmbeddingModule({config.text.vocab_size, config.text.hidden_size}) + .build(ctx, text_ids, weights.text_embedding); + text = core::reshape_tensor(ctx, text, core::TensorShape::from_dims({1, steps, config.text.hidden_size})); + + auto code_ids = core::wrap_tensor( + fused_code_ids, + core::TensorShape::from_dims({steps, config.audio.num_codebooks}), + GGML_TYPE_I32); + auto code = modules::EmbeddingModule( + {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}) + .build(ctx, code_ids, weights.modality_embedding); + code = modules::ReduceSumModule({1}).build(ctx, code); + code = core::reshape_tensor(ctx, code, core::TensorShape::from_dims({1, steps, config.text.hidden_size})); + + auto text_gate_value = core::wrap_tensor( + text_gate, + core::TensorShape::from_dims({1, steps, 1}), + GGML_TYPE_F32); + auto code_gate_value = core::wrap_tensor( + code_gate, + core::TensorShape::from_dims({1, steps, 1}), + GGML_TYPE_F32); + text_gate_value = core::wrap_tensor( + ggml_repeat(ctx.ggml, text_gate_value.tensor, text.tensor), text.shape, GGML_TYPE_F32); + code_gate_value = core::wrap_tensor( + ggml_repeat(ctx.ggml, code_gate_value.tensor, code.tensor), code.shape, GGML_TYPE_F32); + return modules::AddModule{}.build( + ctx, + modules::MulModule{}.build(ctx, text, text_gate_value), + modules::MulModule{}.build(ctx, code, code_gate_value)); +} + +core::TensorValue build_modality_logits( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const HiggsARWeights & weights, + const HiggsConfig & config) { + const int64_t out_features = config.audio.num_codebooks * config.audio.vocab_size; + const bool use_fast_projection = + ctx.backend_type == core::BackendType::Cuda && hidden.shape.rank == 3 && + hidden.shape.dims[1] == 1 && out_features % 4 == 0; + auto logits = + use_fast_projection + ? modules::FastPackedProjection4Module({config.text.hidden_size, out_features, GGML_PREC_DEFAULT}) + .build(ctx, hidden, {weights.modality_embedding, std::nullopt}) + : modules::LinearModule({config.text.hidden_size, out_features, false}) + .build(ctx, hidden, {weights.modality_embedding, std::nullopt}); + return core::reshape_tensor( + ctx, + logits, + core::TensorShape::from_dims({config.audio.num_codebooks, config.audio.vocab_size})); +} + +} // namespace + +HiggsARWeights load_higgs_ar_weights( + const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) { + if (assets.weights == nullptr) { + throw std::runtime_error("Higgs TTS AR weights require tensor source"); + } + if (backend == nullptr) { + throw std::runtime_error("Higgs TTS AR backend is not initialized"); + } + const auto & config = assets.config; + const auto & source = *assets.weights; + HiggsARWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "higgs_audio_tts.ar.weights", + weight_context_bytes); + weights.text_embedding = weights.store->load_tensor( + source, + "tied.embedding.text_embedding.weight", + weight_storage_type, + {config.text.vocab_size, config.text.hidden_size}); + weights.modality_embedding = weights.store->load_tensor( + source, + "tied.embedding.modality_embeddings.0.embedding.weight", + weight_storage_type, + {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}); + weights.decoder = load_decoder_weights(*weights.store, source, config.text, weight_storage_type); + weights.norm = weights.store->load_f32_tensor(source, "body.norm.weight", {config.text.hidden_size}); + weights.packed_qkv = true; + weights.store->upload(); + return weights; +} + +void HiggsARDecodeTiming::add(const HiggsARDecodeTiming & other) noexcept { + input_upload_ms += other.input_upload_ms; + mask_upload_ms += other.mask_upload_ms; + graph_compute_ms += other.graph_compute_ms; + output_read_ms += other.output_read_ms; + steps += other.steps; +} + +HiggsARRuntime::HiggsARRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + device_(execution.config().device), + threads_(std::max(1, execution.config().threads)) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS AR runtime requires assets"); + } + if (assets_->weights == nullptr) { + throw std::runtime_error("Higgs TTS AR runtime requires tensor source"); + } + weights_ = std::make_shared( + load_higgs_ar_weights(*assets_, backend_, backend_type_, weight_context_bytes, weight_storage_type)); +} + +const HiggsAssets & HiggsARRuntime::assets() const noexcept { + return *assets_; +} + +const HiggsARWeights & HiggsARRuntime::weights() const noexcept { + return *weights_; +} + +ggml_backend_t HiggsARRuntime::backend() const noexcept { + return backend_; +} + +core::BackendType HiggsARRuntime::backend_type() const noexcept { + return backend_type_; +} + +int HiggsARRuntime::device() const noexcept { + return device_; +} + +int HiggsARRuntime::threads() const noexcept { + return threads_; +} + +struct HiggsARKVCache::Impl { + Impl(std::shared_ptr input_runtime, int64_t input_cache_steps) + : runtime(std::move(input_runtime)), + cache_steps(input_cache_steps) { + if (runtime == nullptr) { + throw std::runtime_error("Higgs TTS AR KV cache requires runtime"); + } + if (cache_steps <= 0) { + throw std::runtime_error("Higgs TTS AR KV cache requires positive capacity"); + } + const auto & config = runtime->assets().config; + const auto & tensor_weights = runtime->weights(); + const int64_t dim = config.text.head_dim; + cache_layer_count = tensor_weights.decoder.layers.size(); + ggml_init_params params{4 * 1024 * 1024, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR KV cache context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.kv_cache", runtime->backend_type()}; + std::vector key_tensors; + std::vector value_tensors; + key_tensors.reserve(tensor_weights.decoder.layers.size()); + value_tensors.reserve(tensor_weights.decoder.layers.size()); + for (size_t layer = 0; layer < tensor_weights.decoder.layers.size(); ++layer) { + key_tensors.push_back(core::make_tensor( + build_ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); + value_tensors.push_back(core::make_tensor( + build_ctx, + GGML_TYPE_F16, + core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); + } + runtime::TransformerKVCacheOptions cache_options; + cache_options.allow_f16_storage = true; + cache = runtime::TransformerKVCache( + cache_steps, + config.text.num_key_value_heads * dim, + std::move(key_tensors), + std::move(value_tensors), + cache_options); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime->backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR KV cache"); + } + } + + ~Impl() { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + bool can_run(const HiggsARRuntime & candidate_runtime, int64_t required_steps) const { + return runtime.get() == &candidate_runtime && cache_steps >= required_steps; + } + + void reset() { + runtime::TransformerKVState state; + state.current_end = 0; + state.layers.resize(cache_layer_count); + cache.import_state(state); + } + + void retain_prefix(int64_t prefix_steps) { + cache.retain_prefix(prefix_steps); + } + + void import_state(const runtime::TransformerKVState & state) { + cache.import_state(state); + } + + runtime::TransformerKVState export_state() const { + return cache.export_state(); + } + + void advance_after_direct_append(int64_t steps) { + cache.advance_after_direct_append(steps); + } + + std::shared_ptr runtime; + int64_t cache_steps = 0; + size_t cache_layer_count = 0; + std::unique_ptr ctx; + runtime::TransformerKVCache cache; + ggml_backend_buffer_t buffer = nullptr; +}; + +HiggsARKVCache::HiggsARKVCache(std::shared_ptr runtime, int64_t cache_steps) + : impl_(std::make_unique(std::move(runtime), cache_steps)) {} + +HiggsARKVCache::~HiggsARKVCache() = default; + +bool HiggsARKVCache::can_run(const HiggsARRuntime & runtime, int64_t required_steps) const { + return impl_->can_run(runtime, required_steps); +} + +int64_t HiggsARKVCache::cache_steps() const { + return impl_->cache.cache_steps(); +} + +int64_t HiggsARKVCache::valid_steps() const { + return impl_->cache.valid_steps(); +} + +int64_t HiggsARKVCache::current_end() const { + return impl_->cache.current_end(); +} + +void HiggsARKVCache::reset() { + impl_->reset(); +} + +void HiggsARKVCache::retain_prefix(int64_t prefix_steps) { + impl_->retain_prefix(prefix_steps); +} + +void HiggsARKVCache::import_state(const runtime::TransformerKVState & state) { + impl_->import_state(state); +} + +runtime::TransformerKVState HiggsARKVCache::export_state() const { + return impl_->export_state(); +} + +void HiggsARKVCache::advance_after_direct_append(int64_t steps) { + impl_->advance_after_direct_append(steps); +} + +const core::TensorValue & HiggsARKVCache::key_tensor(size_t layer) const { + return impl_->cache.key_tensor(layer); +} + +const core::TensorValue & HiggsARKVCache::value_tensor(size_t layer) const { + return impl_->cache.value_tensor(layer); +} + +struct HiggsARDecodeGraph::Impl { + Impl( + std::shared_ptr input_runtime, + int64_t input_cache_steps, + HiggsARKVCache & input_cache, + size_t graph_arena_bytes) + : runtime(std::move(input_runtime)), + cache(&input_cache), + cache_steps(input_cache_steps) { + if (runtime == nullptr) { + throw std::runtime_error("Higgs TTS AR decode graph requires runtime"); + } + if (cache == nullptr || !cache->can_run(*runtime, cache_steps)) { + throw std::runtime_error("Higgs TTS AR decode graph requires matching KV cache"); + } + if (cache_steps <= 0) { + throw std::runtime_error("Higgs TTS AR decode graph requires positive 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 Higgs TTS AR decode graph context"); + } + const auto & config = runtime->assets().config; + const auto & tensor_weights = runtime->weights(); + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.decode", runtime->backend_type()}; + + fused_code_ids = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks); + auto x = build_higgs_decode_code_embedding( + build_ctx, + tensor_weights, + config, + fused_code_ids); + + positions = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, 1); + auto positions_value = core::wrap_tensor(positions, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + cache_slot = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I64, 1); + auto cache_slot_value = core::wrap_tensor(cache_slot, core::TensorShape::from_dims({1}), GGML_TYPE_I64); + attention_mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, cache_steps, 1, 1, 1); + auto attention_mask_value = core::wrap_tensor( + attention_mask, + core::TensorShape::from_dims({1, 1, 1, cache_steps}), + GGML_TYPE_F16); + + graph = ggml_new_graph_custom(ctx.get(), 65536, false); + const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); + for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { + auto out = decoder.build_decode_layer( + build_ctx, + graph, + x, + positions_value, + tensor_weights.decoder.layers[layer_index], + cache->key_tensor(layer_index), + cache->value_tensor(layer_index), + cache_slot_value, + attention_mask_value); + x = out.output; + } + + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(build_ctx, x, {tensor_weights.norm, std::nullopt}); + auto logits = build_modality_logits(build_ctx, x, tensor_weights, config); + logits_output = logits.tensor; + ggml_set_output(logits_output); + ggml_build_forward_expand(graph, logits_output); + + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime->backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR decode graph"); + } + core::set_backend_threads(runtime->backend(), runtime->threads()); + fused_code_ids_values.assign(static_cast(config.audio.num_codebooks), 0); + attention_mask_values.assign(static_cast(cache_steps), ggml_fp32_to_fp16(-INFINITY)); + engine::debug::timing_log_scalar( + "higgs_audio_tts.ar.decode.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~Impl() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + bool can_run(const HiggsARRuntime & candidate_runtime, int64_t required_steps) const { + return runtime.get() == &candidate_runtime && cache != nullptr && cache->can_run(candidate_runtime, required_steps); + } + + int64_t cache_steps_value() const { + return cache_steps; + } + + void import_prefill_state(const runtime::TransformerKVState & state) { + cache->import_state(state); + } + + void reset_timing() noexcept { + input_upload_ms = 0.0; + mask_upload_ms = 0.0; + graph_compute_ms = 0.0; + output_read_ms = 0.0; + steps = 0; + } + + void begin_decode_run() { + reset_timing(); + std::fill(attention_mask_values.begin(), attention_mask_values.end(), ggml_fp32_to_fp16(-INFINITY)); + const int64_t visible_steps = std::min(cache->valid_steps(), cache_steps); + std::fill( + attention_mask_values.begin(), + attention_mask_values.begin() + static_cast(visible_steps), + ggml_fp32_to_fp16(0.0F)); + ggml_backend_tensor_set( + attention_mask, + attention_mask_values.data(), + 0, + attention_mask_values.size() * sizeof(ggml_fp16_t)); + } + + HiggsARDecodeTiming timing() const { + return {input_upload_ms, mask_upload_ms, graph_compute_ms, output_read_ms, steps}; + } + + void run_step_into(const HiggsARDecodeInput & input, HiggsARDecodeOutput & output, bool log_timing) { + const auto & config = runtime->assets().config; + if (cache->valid_steps() >= cache_steps) { + throw std::runtime_error("Higgs TTS AR decode cache exhausted"); + } + if (!input.use_last_codes) { + throw std::runtime_error("Higgs TTS AR decode graph expects code-token steps after prefill"); + } + if (static_cast(input.last_codes.size()) != config.audio.num_codebooks) { + throw std::runtime_error("Higgs TTS AR decode last codebook row shape mismatch"); + } + + auto timing_start = Clock::now(); + for (int64_t codebook = 0; codebook < config.audio.num_codebooks; ++codebook) { + const int32_t code = input.last_codes[static_cast(codebook)]; + if (code < 0 || code >= config.audio.vocab_size) { + throw std::runtime_error("Higgs TTS AR decode codebook token is outside vocabulary"); + } + fused_code_ids_values[static_cast(codebook)] = + static_cast(code + codebook * config.audio.vocab_size); + } + ggml_backend_tensor_set( + fused_code_ids, + fused_code_ids_values.data(), + 0, + fused_code_ids_values.size() * sizeof(int32_t)); + + const int32_t position = static_cast(cache->current_end()); + ggml_backend_tensor_set(positions, &position, 0, sizeof(int32_t)); + const int64_t cache_slot_value = cache->valid_steps(); + ggml_backend_tensor_set(cache_slot, &cache_slot_value, 0, sizeof(int64_t)); + const double input_upload_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + input_upload_ms += input_upload_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.input_upload_ms", input_upload_delta_ms); + } + timing_start = Clock::now(); + attention_mask_values[static_cast(cache_slot_value)] = ggml_fp32_to_fp16(0.0F); + ggml_backend_tensor_set( + attention_mask, + attention_mask_values.data() + cache_slot_value, + static_cast(cache_slot_value) * sizeof(ggml_fp16_t), + sizeof(ggml_fp16_t)); + const double mask_upload_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + mask_upload_ms += mask_upload_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.mask_upload_ms", mask_upload_delta_ms); + } + + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (log_timing || engine::debug::timing_log_enabled()) { + ggml_backend_synchronize(runtime->backend()); + } + const double graph_compute_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + graph_compute_ms += graph_compute_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.graph.compute_ms", graph_compute_delta_ms); + } + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR decode graph compute failed"); + } + + output.codebook_logits.resize(static_cast(config.audio.num_codebooks * config.audio.vocab_size)); + timing_start = Clock::now(); + ggml_backend_tensor_get( + logits_output, + output.codebook_logits.data(), + 0, + output.codebook_logits.size() * sizeof(float)); + const double output_read_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + output_read_ms += output_read_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.output_read_ms", output_read_delta_ms); + } + + cache->advance_after_direct_append(1); + ++steps; + } + + std::shared_ptr runtime; + HiggsARKVCache * cache = nullptr; + int64_t cache_steps = 0; + std::unique_ptr ctx; + ggml_tensor * fused_code_ids = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * cache_slot = nullptr; + ggml_tensor * attention_mask = nullptr; + ggml_tensor * logits_output = nullptr; + std::vector fused_code_ids_values; + std::vector attention_mask_values; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + double input_upload_ms = 0.0; + double mask_upload_ms = 0.0; + double graph_compute_ms = 0.0; + double output_read_ms = 0.0; + int64_t steps = 0; +}; + +struct HiggsARPrefillGraph::Impl { + Impl( + std::shared_ptr input_runtime, + int64_t input_prompt_steps, + int64_t input_start_step, + HiggsARKVCache * input_cache, + size_t graph_arena_bytes) + : runtime(std::move(input_runtime)), + target_cache(input_cache), + prompt_steps(input_prompt_steps), + start_step(input_start_step), + run_steps(input_prompt_steps - input_start_step), + prefill_cache_steps(input_prompt_steps), + layerwise(input_prompt_steps >= kLayerwisePrefillMinSteps && input_start_step == 0), + graph_arena_bytes(graph_arena_bytes) { + if (runtime == nullptr) { + throw std::runtime_error("Higgs TTS AR prefill graph requires runtime"); + } + if (prompt_steps <= 0) { + throw std::runtime_error("Higgs TTS AR prefill graph requires positive prompt length"); + } + if (start_step < 0 || start_step >= prompt_steps) { + throw std::runtime_error("Higgs TTS AR prefill graph start step is outside the prompt"); + } + if (start_step > 0 && target_cache == nullptr) { + throw std::runtime_error("Higgs TTS AR suffix prefill requires a target KV cache"); + } + if (layerwise) { + engine::debug::timing_log_scalar("higgs_audio_tts.ar.prefill.graph.build_ms", 0.0); + return; + } + 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 Higgs TTS AR prefill graph context"); + } + const auto & config = runtime->assets().config; + const auto & tensor_weights = runtime->weights(); + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill", runtime->backend_type()}; + + text_tokens = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, run_steps); + fused_code_ids = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks, run_steps); + text_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, run_steps, 1); + code_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, run_steps, 1); + auto x = build_higgs_prefill_input_embedding( + build_ctx, + tensor_weights, + config, + text_tokens, + fused_code_ids, + text_gate, + code_gate, + run_steps); + positions = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, run_steps); + auto positions_value = + core::wrap_tensor(positions, core::TensorShape::from_dims({run_steps}), GGML_TYPE_I32); + attention_mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, prefill_cache_steps, run_steps, 1, 1); + auto attention_mask_value = core::wrap_tensor( + attention_mask, + core::TensorShape::from_dims({1, 1, run_steps, prefill_cache_steps}), + GGML_TYPE_F16); + graph = ggml_new_graph_custom(ctx.get(), 262144, false); + keys.reserve(tensor_weights.decoder.layers.size()); + values.reserve(tensor_weights.decoder.layers.size()); + const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); + for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { + std::optional prefix_key; + std::optional prefix_value; + if (start_step > 0) { + prefix_key = higgs_cache_view( + build_ctx, + target_cache->key_tensor(layer_index), + 0, + start_step, + config.text.num_key_value_heads, + config.text.head_dim); + prefix_value = higgs_cache_view( + build_ctx, + target_cache->value_tensor(layer_index), + 0, + start_step, + config.text.num_key_value_heads, + config.text.head_dim); + } + auto out = decoder.build_prefill_layer( + build_ctx, + x, + positions_value, + tensor_weights.decoder.layers[layer_index], + attention_mask_value, + prefix_key, + prefix_value); + x = out.output; + if (target_cache != nullptr) { + auto key_dest = higgs_cache_view( + build_ctx, + target_cache->key_tensor(layer_index), + start_step, + run_steps, + config.text.num_key_value_heads, + config.text.head_dim); + auto value_dest = higgs_cache_view( + build_ctx, + target_cache->value_tensor(layer_index), + start_step, + run_steps, + config.text.num_key_value_heads, + config.text.head_dim); + ggml_build_forward_expand(graph, ggml_cpy(ctx.get(), out.key.tensor, key_dest.tensor)); + ggml_build_forward_expand(graph, ggml_cpy(ctx.get(), out.value.tensor, value_dest.tensor)); + } else { + keys.push_back(out.key.tensor); + values.push_back(out.value.tensor); + } + } + x = modules::SliceModule({1, run_steps - 1, 1}).build(build_ctx, x); + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(build_ctx, x, {tensor_weights.norm, std::nullopt}); + auto logits = build_modality_logits(build_ctx, x, tensor_weights, config); + logits_output = logits.tensor; + ggml_set_output(logits_output); + ggml_build_forward_expand(graph, logits_output); + + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime->backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR prefill graph"); + } + text_token_values.assign(static_cast(run_steps), 0); + fused_code_id_values.assign(static_cast(run_steps * config.audio.num_codebooks), 0); + text_gate_values.assign(static_cast(run_steps), 0.0F); + code_gate_values.assign(static_cast(run_steps), 0.0F); + positions_values = modules::qwen_position_ids(run_steps, start_step); + attention_mask_values = modules::qwen_causal_suffix_mask_values(1, run_steps, start_step); + engine::debug::timing_log_scalar( + "higgs_audio_tts.ar.prefill.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~Impl() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + bool matches( + const HiggsARRuntime & candidate_runtime, + int64_t candidate_prompt_steps, + int64_t candidate_start_step) const { + return runtime.get() == &candidate_runtime && + prompt_steps == candidate_prompt_steps && + start_step == candidate_start_step; + } + + struct EmbeddingGraph { + EmbeddingGraph(const HiggsARRuntime & runtime, int64_t steps, size_t arena_bytes) + : runtime(&runtime), steps(steps) { + const auto & config = runtime.assets().config; + ggml_init_params params{arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR embedding graph context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill.embedding", runtime.backend_type()}; + text_tokens = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, steps); + fused_code_ids = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks, steps); + text_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, steps, 1); + code_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, steps, 1); + output = build_higgs_prefill_input_embedding( + build_ctx, + runtime.weights(), + config, + text_tokens, + fused_code_ids, + text_gate, + code_gate, + steps) + .tensor; + graph = ggml_new_graph_custom(ctx.get(), 32768, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime.backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR embedding graph"); + } + } + + ~EmbeddingGraph() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + std::vector run(const HiggsARPrefillInput & input) { + const auto & config = runtime->assets().config; + ggml_backend_tensor_set(text_tokens, input.text_tokens.data(), 0, input.text_tokens.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + fused_code_ids, + input.fused_code_ids.data(), + 0, + input.fused_code_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(text_gate, input.text_gate.data(), 0, input.text_gate.size() * sizeof(float)); + ggml_backend_tensor_set(code_gate, input.code_gate.data(), 0, input.code_gate.size() * sizeof(float)); + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR embedding graph compute failed"); + } + std::vector hidden(static_cast(steps * config.text.hidden_size)); + ggml_backend_tensor_get(output, hidden.data(), 0, hidden.size() * sizeof(float)); + return hidden; + } + + const HiggsARRuntime * runtime = nullptr; + int64_t steps = 0; + std::unique_ptr ctx; + ggml_tensor * text_tokens = nullptr; + ggml_tensor * fused_code_ids = nullptr; + ggml_tensor * text_gate = nullptr; + ggml_tensor * code_gate = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + }; + + struct LayerGraph { + LayerGraph( + const HiggsARRuntime & runtime, + const modules::QwenDecoderLayerWeights & layer, + int64_t steps, + size_t arena_bytes) + : runtime(&runtime), steps(steps) { + const auto & config = runtime.assets().config; + ggml_init_params params{arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR layer prefill graph context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill.layer", runtime.backend_type()}; + auto x = core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, steps, config.text.hidden_size})); + input = x.tensor; + positions = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, steps); + auto positions_value = + core::wrap_tensor(positions, core::TensorShape::from_dims({steps}), GGML_TYPE_I32); + attention_mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, steps, steps, 1, 1); + auto attention_mask_value = core::wrap_tensor( + attention_mask, + core::TensorShape::from_dims({1, 1, steps, steps}), + GGML_TYPE_F16); + const HiggsQwenDecoderComponent decoder(config.text, runtime.weights().packed_qkv); + auto out = decoder.build_prefill_layer( + build_ctx, + x, + positions_value, + layer, + attention_mask_value); + output = out.output.tensor; + key = out.key.tensor; + value = out.value.tensor; + graph = ggml_new_graph_custom(ctx.get(), 65536, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime.backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR layer prefill graph"); + } + + const auto position_values = modules::qwen_position_ids(steps); + ggml_backend_tensor_set(positions, position_values.data(), 0, position_values.size() * sizeof(int32_t)); + auto mask = modules::qwen_causal_prefill_mask_values(1, steps); + ggml_backend_tensor_set(attention_mask, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); + } + + ~LayerGraph() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + struct Output { + std::vector hidden; + std::vector key; + std::vector value; + }; + + Output run(const std::vector & hidden) { + const auto & config = runtime->assets().config; + const size_t hidden_values = static_cast(steps * config.text.hidden_size); + if (hidden.size() != hidden_values) { + throw std::runtime_error("Higgs TTS AR layer prefill input size mismatch"); + } + ggml_backend_tensor_set(input, hidden.data(), 0, hidden.size() * sizeof(float)); + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR layer prefill graph compute failed"); + } + Output out; + out.hidden.resize(hidden_values); + ggml_backend_tensor_get(output, out.hidden.data(), 0, out.hidden.size() * sizeof(float)); + const size_t layer_values = static_cast( + steps * config.text.num_key_value_heads * config.text.head_dim); + out.key.resize(layer_values); + out.value.resize(layer_values); + ggml_backend_tensor_get(key, out.key.data(), 0, out.key.size() * sizeof(float)); + ggml_backend_tensor_get(value, out.value.data(), 0, out.value.size() * sizeof(float)); + return out; + } + + const HiggsARRuntime * runtime = nullptr; + int64_t steps = 0; + std::unique_ptr ctx; + ggml_tensor * input = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * attention_mask = nullptr; + ggml_tensor * output = nullptr; + ggml_tensor * key = nullptr; + ggml_tensor * value = nullptr; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + }; + + struct FinalGraph { + FinalGraph(const HiggsARRuntime & runtime, size_t arena_bytes) : runtime(&runtime) { + const auto & config = runtime.assets().config; + ggml_init_params params{arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR final prefill graph context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill.final", runtime.backend_type()}; + auto x = core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, 1, config.text.hidden_size})); + input = x.tensor; + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(build_ctx, x, {runtime.weights().norm, std::nullopt}); + auto logits = build_modality_logits(build_ctx, x, runtime.weights(), config); + output = logits.tensor; + graph = ggml_new_graph_custom(ctx.get(), 8192, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime.backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR final prefill graph"); + } + } + + ~FinalGraph() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + std::vector run(const std::vector & hidden) { + const auto & config = runtime->assets().config; + if (static_cast(hidden.size()) != config.text.hidden_size) { + throw std::runtime_error("Higgs TTS AR final prefill input size mismatch"); + } + ggml_backend_tensor_set(input, hidden.data(), 0, hidden.size() * sizeof(float)); + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR final prefill graph compute failed"); + } + std::vector logits(static_cast( + config.audio.num_codebooks * config.audio.vocab_size)); + ggml_backend_tensor_get(output, logits.data(), 0, logits.size() * sizeof(float)); + return logits; + } + + const HiggsARRuntime * runtime = nullptr; + std::unique_ptr ctx; + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + }; + + HiggsARPrefillOutput run_layerwise(const HiggsARPrefillInput & input) { + const auto & config = runtime->assets().config; + auto hidden = EmbeddingGraph(*runtime, prompt_steps, graph_arena_bytes).run(input); + HiggsARPrefillOutput out; + out.kv_state.current_end = prompt_steps; + out.kv_state.layers.resize(runtime->weights().decoder.layers.size()); + for (size_t layer = 0; layer < runtime->weights().decoder.layers.size(); ++layer) { + LayerGraph graph( + *runtime, + runtime->weights().decoder.layers[layer], + prompt_steps, + graph_arena_bytes); + auto layer_out = graph.run(hidden); + hidden = std::move(layer_out.hidden); + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps; + state.key = std::move(layer_out.key); + state.value = std::move(layer_out.value); + } + std::vector last_hidden(static_cast(config.text.hidden_size)); + const auto last_begin = hidden.begin() + + static_cast((prompt_steps - 1) * config.text.hidden_size); + std::copy( + last_begin, + last_begin + static_cast(config.text.hidden_size), + last_hidden.begin()); + out.output.codebook_logits = FinalGraph(*runtime, graph_arena_bytes).run(last_hidden); + return out; + } + + HiggsARPrefillOutput run(const HiggsARPrefillInput & input, int64_t candidate_start_step) { + const auto & config = runtime->assets().config; + if (input.steps != prompt_steps || + static_cast(input.text_tokens.size()) != prompt_steps || + static_cast(input.fused_code_ids.size()) != prompt_steps * config.audio.num_codebooks || + static_cast(input.text_gate.size()) != prompt_steps || + static_cast(input.code_gate.size()) != prompt_steps) { + throw std::runtime_error("Higgs TTS AR prefill graph input shape mismatch"); + } + if (candidate_start_step != start_step) { + throw std::runtime_error("Higgs TTS AR prefill graph start step mismatch"); + } + if (start_step > 0 && + (target_cache == nullptr || target_cache->valid_steps() < start_step || + target_cache->current_end() != start_step)) { + throw std::runtime_error("Higgs TTS AR suffix prefill requires the retained prefix in KV cache"); + } + if (layerwise) { + return run_layerwise(input); + } + for (int64_t step = 0; step < run_steps; ++step) { + const int64_t source_step = start_step + step; + text_token_values[static_cast(step)] = + input.text_tokens[static_cast(source_step)]; + text_gate_values[static_cast(step)] = + input.text_gate[static_cast(source_step)]; + code_gate_values[static_cast(step)] = + input.code_gate[static_cast(source_step)]; + for (int64_t codebook = 0; codebook < config.audio.num_codebooks; ++codebook) { + fused_code_id_values[static_cast(step * config.audio.num_codebooks + codebook)] = + input.fused_code_ids[static_cast(source_step * config.audio.num_codebooks + codebook)]; + } + } + ggml_backend_tensor_set(text_tokens, text_token_values.data(), 0, text_token_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + fused_code_ids, + fused_code_id_values.data(), + 0, + fused_code_id_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set(text_gate, text_gate_values.data(), 0, text_gate_values.size() * sizeof(float)); + ggml_backend_tensor_set(code_gate, code_gate_values.data(), 0, code_gate_values.size() * sizeof(float)); + ggml_backend_tensor_set(positions, positions_values.data(), 0, positions_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + attention_mask, + attention_mask_values.data(), + 0, + attention_mask_values.size() * sizeof(ggml_fp16_t)); + + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR prefill graph compute failed"); + } + HiggsARPrefillOutput out; + out.output.codebook_logits.resize(static_cast(config.audio.num_codebooks * config.audio.vocab_size)); + ggml_backend_tensor_get( + logits_output, + out.output.codebook_logits.data(), + 0, + out.output.codebook_logits.size() * sizeof(float)); + if (target_cache != nullptr) { + target_cache->advance_after_direct_append(run_steps); + out.wrote_cache = true; + out.kv_state.current_end = prompt_steps; + return out; + } + out.kv_state.current_end = prompt_steps; + out.kv_state.layers.resize(keys.size()); + const size_t layer_values = static_cast( + prompt_steps * config.text.num_key_value_heads * config.text.head_dim); + 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; + } + + std::shared_ptr runtime; + HiggsARKVCache * target_cache = nullptr; + int64_t prompt_steps = 0; + int64_t start_step = 0; + int64_t run_steps = 0; + int64_t prefill_cache_steps = 0; + bool layerwise = false; + size_t graph_arena_bytes = 0; + std::unique_ptr ctx; + ggml_tensor * text_tokens = nullptr; + ggml_tensor * fused_code_ids = nullptr; + ggml_tensor * text_gate = nullptr; + ggml_tensor * code_gate = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * attention_mask = nullptr; + ggml_tensor * logits_output = nullptr; + std::vector keys; + std::vector values; + std::vector text_token_values; + std::vector fused_code_id_values; + std::vector text_gate_values; + std::vector code_gate_values; + std::vector positions_values; + std::vector attention_mask_values; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; +}; + +HiggsARPrefillGraph::HiggsARPrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t start_step, + HiggsARKVCache * cache, + size_t graph_arena_bytes) + : impl_(std::make_unique(std::move(runtime), prompt_steps, start_step, cache, graph_arena_bytes)) {} + +HiggsARPrefillGraph::~HiggsARPrefillGraph() = default; + +bool HiggsARPrefillGraph::matches( + const HiggsARRuntime & runtime, + int64_t prompt_steps, + int64_t start_step) const { + return impl_->matches(runtime, prompt_steps, start_step); +} + +HiggsARPrefillOutput HiggsARPrefillGraph::run(const HiggsARPrefillInput & input, int64_t start_step) { + return impl_->run(input, start_step); +} + + +HiggsARDecodeGraph::HiggsARDecodeGraph( + std::shared_ptr runtime, + int64_t cache_steps, + HiggsARKVCache & cache, + size_t graph_arena_bytes) + : impl_(std::make_unique(std::move(runtime), cache_steps, cache, graph_arena_bytes)) {} + +HiggsARDecodeGraph::~HiggsARDecodeGraph() = default; + +bool HiggsARDecodeGraph::can_run(const HiggsARRuntime & runtime, int64_t required_steps) const { + return impl_->can_run(runtime, required_steps); +} + +int64_t HiggsARDecodeGraph::cache_steps() const { + return impl_->cache_steps_value(); +} + +void HiggsARDecodeGraph::import_prefill_state(const runtime::TransformerKVState & state) { + impl_->import_prefill_state(state); +} + +void HiggsARDecodeGraph::begin_decode_run() { + impl_->begin_decode_run(); +} + +HiggsARDecodeTiming HiggsARDecodeGraph::timing() const { + return impl_->timing(); +} + +void HiggsARDecodeGraph::run_step_into( + const HiggsARDecodeInput & input, + HiggsARDecodeOutput & output, + bool log_timing) { + impl_->run_step_into(input, output, log_timing); +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/assets.cpp b/src/models/higgs_audio_tts/assets.cpp new file mode 100644 index 00000000..7a163d45 --- /dev/null +++ b/src/models/higgs_audio_tts/assets.cpp @@ -0,0 +1,171 @@ +#include "engine/models/higgs_audio_tts/assets.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/framework/io/config.h" +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace json = engine::io::json; +namespace { + +constexpr const char * kExpectedArchitecture = "HiggsMultimodalQwen3ForConditionalGeneration"; + +float parse_rope_theta(const engine::io::json::Value & text_config) { + const auto * rope_parameters = text_config.find("rope_parameters"); + if (rope_parameters != nullptr && rope_parameters->is_object()) { + return json::optional_f32(*rope_parameters, "rope_theta", 1000000.0F); + } + const auto * rope_theta = text_config.find("rope_theta"); + if (rope_theta != nullptr && rope_theta->is_number()) { + return rope_theta->as_f32(); + } + return 1000000.0F; +} + +std::string parse_architecture(const engine::io::json::Value & root) { + const auto * architectures = root.find("architectures"); + if (architectures == nullptr || !architectures->is_array() || architectures->as_array().empty()) { + throw std::runtime_error("Higgs TTS config must provide architectures[0]"); + } + return architectures->as_array().front().as_string(); +} + +HiggsTextConfig parse_text_config(const engine::io::json::Value & root) { + HiggsTextConfig config; + config.model_type = json::require_string(root, "model_type"); + if (config.model_type != "qwen3") { + throw std::runtime_error("Higgs TTS text_config.model_type mismatch"); + } + config.vocab_size = json::require_i64(root, "vocab_size"); + config.hidden_size = json::require_i64(root, "hidden_size"); + config.intermediate_size = json::require_i64(root, "intermediate_size"); + config.num_hidden_layers = json::require_i64(root, "num_hidden_layers"); + config.num_attention_heads = json::require_i64(root, "num_attention_heads"); + config.num_key_value_heads = json::require_i64(root, "num_key_value_heads"); + config.head_dim = json::optional_i64(root, "head_dim", config.hidden_size / config.num_attention_heads); + config.max_position_embeddings = json::require_i64(root, "max_position_embeddings"); + config.bos_token_id = json::require_i64(root, "bos_token_id"); + config.eos_token_id = json::require_i64(root, "eos_token_id"); + config.pad_token_id = json::optional_nullable_i64(root, "pad_token_id", -1); + config.rms_norm_eps = json::optional_f32(root, "rms_norm_eps", config.rms_norm_eps); + config.rope_theta = parse_rope_theta(root); + config.tie_word_embeddings = json::optional_bool(root, "tie_word_embeddings", config.tie_word_embeddings); + + engine::io::require_positive(config.vocab_size, "text vocab_size"); + engine::io::require_positive(config.hidden_size, "text hidden_size"); + engine::io::require_positive(config.intermediate_size, "text intermediate_size"); + engine::io::require_positive(config.num_hidden_layers, "text num_hidden_layers"); + engine::io::require_positive(config.num_attention_heads, "text num_attention_heads"); + engine::io::require_positive(config.num_key_value_heads, "text num_key_value_heads"); + engine::io::require_positive(config.head_dim, "text head_dim"); + engine::io::require_positive(config.max_position_embeddings, "text max_position_embeddings"); + engine::io::require_divisible(config.num_attention_heads, config.num_key_value_heads, "text grouped-query attention"); + return config; +} + +HiggsAudioEncoderConfig parse_audio_config(const engine::io::json::Value & root) { + HiggsAudioEncoderConfig config; + config.model_type = json::require_string(root, "model_type"); + if (config.model_type != "higgs_audio_encoder") { + throw std::runtime_error("Higgs TTS audio_encoder_config.model_type mismatch"); + } + config.encoder_type = json::require_string(root, "encoder_type"); + if (config.encoder_type != "discrete") { + throw std::runtime_error("Higgs TTS audio_encoder_config.encoder_type mismatch"); + } + config.num_codebooks = json::require_i64(root, "num_codebooks"); + config.vocab_size = json::require_i64(root, "vocab_size"); + config.out_dim = json::require_i64(root, "out_dim"); + config.mel_per_sample = json::require_i64(root, "mel_per_sample"); + config.max_chunk_size = json::require_i64(root, "max_chunk_size"); + config.tie_word_embeddings = json::optional_bool(root, "tie_word_embeddings", config.tie_word_embeddings); + config.use_delay_pattern = json::optional_bool(root, "use_delay_pattern", config.use_delay_pattern); + + engine::io::require_positive(config.num_codebooks, "audio num_codebooks"); + engine::io::require_positive(config.vocab_size, "audio vocab_size"); + engine::io::require_positive(config.out_dim, "audio out_dim"); + engine::io::require_positive(config.mel_per_sample, "audio mel_per_sample"); + engine::io::require_positive(config.max_chunk_size, "audio max_chunk_size"); + if (!config.tie_word_embeddings) { + throw std::runtime_error("Higgs TTS currently expects tied modality embedding/head weights"); + } + if (!config.use_delay_pattern) { + throw std::runtime_error("Higgs TTS v3 requires delay-pattern audio codebooks"); + } + return config; +} + +HiggsConfig parse_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + HiggsConfig config; + config.model_type = json::require_string(root, "model_type"); + if (config.model_type != "higgs_multimodal_qwen3") { + throw std::runtime_error("Higgs TTS model_type mismatch"); + } + config.architecture = parse_architecture(root); + if (config.architecture != kExpectedArchitecture) { + throw std::runtime_error("Higgs TTS architecture mismatch"); + } + config.hidden_size = json::require_i64(root, "_hidden_size"); + config.vocab_size = json::require_i64(root, "_vocab_size"); + config.audio_token_id = json::optional_i64(root, "audio_token_id", config.audio_token_id); + config.ignore_index = json::optional_i64(root, "ignore_index", config.ignore_index); + config.text = parse_text_config(root.require("text_config")); + config.audio = parse_audio_config(root.require("audio_encoder_config")); + + if (config.hidden_size != config.text.hidden_size) { + throw std::runtime_error("Higgs TTS _hidden_size must match text hidden_size"); + } + if (config.vocab_size != config.text.vocab_size) { + throw std::runtime_error("Higgs TTS _vocab_size must match text vocab_size"); + } + if (config.audio.out_dim != config.text.hidden_size) { + throw std::runtime_error("Higgs TTS audio out_dim must match text hidden_size"); + } + if (config.audio_token_id != -100) { + throw std::runtime_error("Higgs TTS audio_token_id mismatch"); + } + return config; +} + +void validate_weight_anchors(const HiggsAssets & assets) { + const auto & config = assets.config; + const auto & weights = *assets.weights; + const int64_t hidden = config.text.hidden_size; + const int64_t audio_fused_vocab = config.audio.num_codebooks * config.audio.vocab_size; + assets::require_tensor_shape(weights, "tied.embedding.text_embedding.weight", {config.text.vocab_size, hidden}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.embedding.weight", {audio_fused_vocab, hidden}); + assets::require_tensor_shape(weights, "body.norm.weight", {hidden}); + assets::require_tensor_shape(weights, "body.layers.0.input_layernorm.weight", {hidden}); + assets::require_tensor_shape(weights, "body.layers.0.post_attention_layernorm.weight", {hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.q_proj.weight", {config.text.num_attention_heads * config.text.head_dim, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.k_proj.weight", {config.text.num_key_value_heads * config.text.head_dim, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.v_proj.weight", {config.text.num_key_value_heads * config.text.head_dim, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.o_proj.weight", {hidden, config.text.num_attention_heads * config.text.head_dim}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.q_norm.weight", {config.text.head_dim}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.k_norm.weight", {config.text.head_dim}); + assets::require_tensor_shape(weights, "body.layers.0.mlp.gate_proj.weight", {config.text.intermediate_size, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.mlp.up_proj.weight", {config.text.intermediate_size, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.mlp.down_proj.weight", {hidden, config.text.intermediate_size}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.model.acoustic_encoder.conv1.weight", {64, 1, 7}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.model.acoustic_decoder.conv2.weight", {1, 32, 7}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.model.quantizer.quantizers.0.codebook.embed", {1024, 64}); +} + +} // namespace + +std::shared_ptr load_higgs_assets(const std::filesystem::path & model_path) { + HiggsAssets assets; + assets.resources = assets::load_resource_bundle_from_package_spec( + model_path, + assets::default_model_package_spec_path("higgs_audio_tts")); + assets.config = parse_config(assets.resources); + assets.weights = assets.resources.open_tensor_source("weights"); + validate_weight_anchors(assets); + return std::make_shared(std::move(assets)); +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/codebooks.cpp b/src/models/higgs_audio_tts/codebooks.cpp new file mode 100644 index 00000000..54c15d38 --- /dev/null +++ b/src/models/higgs_audio_tts/codebooks.cpp @@ -0,0 +1,79 @@ +#include "engine/models/higgs_audio_tts/codebooks.h" + +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +void require_codebook_matrix( + const std::vector & codes, + int64_t frames, + int64_t codebooks, + const char * label) { + if (frames <= 0 || codebooks <= 0) { + throw std::runtime_error(std::string("Higgs TTS ") + label + " requires positive frames and codebooks"); + } + if (static_cast(codes.size()) != frames * codebooks) { + throw std::runtime_error(std::string("Higgs TTS ") + label + " code matrix shape mismatch"); + } +} + +size_t flat_index(int64_t frame, int64_t codebook, int64_t codebooks) { + return static_cast(frame * codebooks + codebook); +} + +} // namespace + +int64_t higgs_delayed_frame_count(int64_t raw_frames, int64_t codebooks) { + if (raw_frames <= 0 || codebooks <= 0) { + throw std::runtime_error("Higgs TTS delayed frame count requires positive dimensions"); + } + return raw_frames + codebooks - 1; +} + +std::vector apply_higgs_delay_pattern( + const std::vector & raw_codes, + int64_t raw_frames, + int64_t codebooks) { + require_codebook_matrix(raw_codes, raw_frames, codebooks, "delay pattern input"); + const int64_t delayed_frames = higgs_delayed_frame_count(raw_frames, codebooks); + std::vector delayed(static_cast(delayed_frames * codebooks), kHiggsEocId); +#ifdef _OPENMP +#pragma omp parallel for if (raw_frames * codebooks > 1024) +#endif + for (int64_t codebook = 0; codebook < codebooks; ++codebook) { + for (int64_t frame = 0; frame < codebook; ++frame) { + delayed[flat_index(frame, codebook, codebooks)] = kHiggsBocId; + } + for (int64_t frame = 0; frame < raw_frames; ++frame) { + delayed[flat_index(codebook + frame, codebook, codebooks)] = + raw_codes[flat_index(frame, codebook, codebooks)]; + } + } + return delayed; +} + +std::vector reverse_higgs_delay_pattern( + const std::vector & delayed_codes, + int64_t delayed_frames, + int64_t codebooks) { + require_codebook_matrix(delayed_codes, delayed_frames, codebooks, "reverse delay pattern input"); + const int64_t raw_frames = delayed_frames - (codebooks - 1); + if (raw_frames <= 0) { + throw std::runtime_error("Higgs TTS delayed codes must include at least one recoverable raw frame"); + } + std::vector raw(static_cast(raw_frames * codebooks), 0); +#ifdef _OPENMP +#pragma omp parallel for if (raw_frames * codebooks > 1024) +#endif + for (int64_t codebook = 0; codebook < codebooks; ++codebook) { + for (int64_t frame = 0; frame < raw_frames; ++frame) { + raw[flat_index(frame, codebook, codebooks)] = + delayed_codes[flat_index(codebook + frame, codebook, codebooks)]; + } + } + return raw; +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/codec.cpp b/src/models/higgs_audio_tts/codec.cpp new file mode 100644 index 00000000..658919d3 --- /dev/null +++ b/src/models/higgs_audio_tts/codec.cpp @@ -0,0 +1,1635 @@ +#include "engine/models/higgs_audio_tts/codec.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr const char * kCodecPrefix = "tied.embedding.modality_embeddings.0.model."; +constexpr int64_t kCodecCodebooks = 8; +constexpr int64_t kCodecCodebookSize = 1024; +constexpr int64_t kCodecCodebookDim = 64; +constexpr int64_t kCodecHiddenSize = 1024; +constexpr int64_t kAcousticHiddenSize = 256; +constexpr int64_t kSemanticHiddenSize = 768; +constexpr int64_t kCodecProjectInputSize = kAcousticHiddenSize + kSemanticHiddenSize; +constexpr int64_t kResidualUnitsPerBlock = 3; +constexpr int64_t kSemanticResidualUnitsPerBlock = 2; +constexpr int64_t kSemanticIntermediateSize = 3072; +constexpr int64_t kSemanticAttentionHeads = 12; +constexpr int64_t kSemanticLayers = 12; +constexpr int64_t kSemanticConvLayers = 7; +constexpr int64_t kSemanticSampleRate = 16000; +constexpr int64_t kSemanticPadSamples = 160; +constexpr float kSemanticLayerNormEps = 1.0e-5F; + +const int64_t kUpsampleRatios[] = {8, 5, 4, 2, 3}; +const int64_t kDecoderChannels[] = {1024, 512, 256, 128, 64, 32}; +const int64_t kEncoderChannels[] = {64, 128, 256, 512, 1024, 2048}; +const int64_t kSemanticConvDim[] = {512, 512, 512, 512, 512, 512, 512}; +const int64_t kSemanticConvKernel[] = {10, 3, 3, 3, 3, 2, 2}; +const int64_t kSemanticConvStride[] = {5, 2, 2, 2, 2, 2, 2}; +constexpr int64_t kDecoderBlockCount = 5; +constexpr int64_t kEncoderBlockCount = 5; +constexpr int64_t kSemanticBlockCount = 2; +constexpr int kCodecSampleRate = 24000; +constexpr int64_t kCodecDecodeCapacityBucketFrames = 32; +constexpr int64_t kCodecHopLength = 960; +constexpr int64_t kCodecPadSamples = kCodecHopLength / 2; +constexpr int64_t kCodecDecodeWindowFrames = 128; +constexpr int64_t kCodecDecodeOverlapFrames = 8; +constexpr int64_t kCodecFullDecodeMaxFrames = 512; +constexpr int64_t kResidualDilations[] = {1, 3, 9}; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +std::string codec_name(const std::string & name) { return std::string(kCodecPrefix) + name; } + +modules::LinearWeights load_linear(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + assets::TensorStorageType storage_type, + int64_t out_features, + int64_t in_features, + bool use_bias) { + modules::LinearWeights weights; + weights.weight = store.load_tensor( + source, codec_name(name + ".weight"), storage_type, {out_features, in_features}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, codec_name(name + ".bias"), {out_features}); + } + return weights; +} + +modules::Conv1dWeights load_conv1d(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size, + bool use_bias) { + modules::Conv1dWeights weights; + weights.weight = store.load_tensor(source, + codec_name(name + ".weight"), + storage_type, + {out_channels, in_channels, kernel_size}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, codec_name(name + ".bias"), {out_channels}); + } + return weights; +} + +modules::ConvTranspose1dWeights load_conv_transpose1d(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + assets::TensorStorageType storage_type, + int64_t in_channels, + int64_t out_channels, + int64_t kernel_size, + bool use_bias) { + modules::ConvTranspose1dWeights weights; + weights.weight = store.load_tensor(source, + codec_name(name + ".weight"), + storage_type, + {in_channels, out_channels, kernel_size}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, codec_name(name + ".bias"), {out_channels}); + } + return weights; +} + +modules::Snake1dWeights load_snake(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + int64_t channels) { + const auto raw = source.require_f32(codec_name(name), std::vector{1, channels, 1}); + return {store.make_f32(core::TensorShape::from_dims({channels}), raw)}; +} + +core::TensorValue require_semantic_tensor(const HiggsCodecWeights & weights, + const std::string & name) { + const auto it = weights.semantic_model.find(name); + if (it == weights.semantic_model.end()) { + throw std::runtime_error("Higgs TTS codec missing semantic tensor: " + name); + } + return it->second; +} + +modules::NormWeights semantic_norm_weights(const HiggsCodecWeights & weights, + const std::string & prefix) { + return modules::NormWeights{require_semantic_tensor(weights, prefix + ".weight"), + require_semantic_tensor(weights, prefix + ".bias")}; +} + +modules::LinearWeights semantic_linear_weights(const HiggsCodecWeights & weights, + const std::string & prefix) { + return modules::LinearWeights{require_semantic_tensor(weights, prefix + ".weight"), + require_semantic_tensor(weights, prefix + ".bias")}; +} + +modules::Conv1dWeights semantic_conv_weights(const HiggsCodecWeights & weights, + const std::string & prefix, + bool use_bias) { + modules::Conv1dWeights out; + out.weight = require_semantic_tensor(weights, prefix + ".weight"); + if (use_bias) { + out.bias = require_semantic_tensor(weights, prefix + ".bias"); + } + return out; +} + +int64_t +conv1d_output_frames(int64_t input_frames, int64_t kernel, int64_t stride, int64_t padding) { + return (input_frames + 2 * padding - kernel) / stride + 1; +} + +int64_t ceil_div(int64_t value, int64_t divisor) { return (value + divisor - 1) / divisor; } + +int64_t semantic_feature_frames(int64_t input_samples) { + int64_t frames = input_samples; + for (int64_t index = 0; index < kSemanticConvLayers; ++index) { + frames = conv1d_output_frames(frames, + kSemanticConvKernel[static_cast(index)], + kSemanticConvStride[static_cast(index)], + 0); + if (frames <= 0) { + throw std::runtime_error("Higgs TTS semantic encoder input is too short"); + } + } + return ceil_div(frames, 2); +} + +int64_t acoustic_encoder_frames(int64_t input_samples) { + int64_t frames = conv1d_output_frames(input_samples, 7, 1, 3); + for (int64_t block = 0; block < kEncoderBlockCount; ++block) { + const int64_t ratio = kUpsampleRatios[static_cast(block)]; + frames = conv1d_output_frames(frames, 2 * ratio, ratio, (ratio + 1) / 2); + if (frames <= 0) { + throw std::runtime_error("Higgs TTS acoustic encoder input is too short"); + } + } + frames = conv1d_output_frames(frames, 3, 1, 1); + if (frames <= 0) { + throw std::runtime_error("Higgs TTS acoustic encoder input is too short"); + } + return frames; +} + +std::vector pad_zeros(const std::vector & input, int64_t left, int64_t right) { + std::vector out(static_cast(left + static_cast(input.size()) + right), + 0.0F); + std::copy(input.begin(), input.end(), out.begin() + left); + return out; +} + +std::vector resample_mono_if_needed(const std::vector & mono, + int source_sample_rate, + int target_sample_rate) { + if (source_sample_rate == target_sample_rate) { + return mono; + } + auto out = + audio::resample_mono_torchaudio_sinc_hann(mono, source_sample_rate, target_sample_rate); + if (out.empty()) { + throw std::runtime_error("Higgs TTS codec resampling produced no samples"); + } + return out; +} + +std::vector prepare_mono_audio(const runtime::AudioBuffer & audio) { + if (audio.sample_rate <= 0) { + throw std::runtime_error("Higgs TTS codec reference audio sample rate must be positive"); + } + if (audio.channels <= 0) { + throw std::runtime_error("Higgs TTS codec reference audio channel count must be positive"); + } + if (audio.samples.empty()) { + throw std::runtime_error("Higgs TTS codec reference audio is empty"); + } + if ((audio.samples.size() % static_cast(audio.channels)) != 0) { + throw std::runtime_error("Higgs TTS codec reference audio sample count is " + "not divisible by channels"); + } + return audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); +} + +std::vector prepare_codec_audio_24k(const std::vector & mono, + int source_sample_rate) { + auto audio_24k = resample_mono_if_needed(mono, source_sample_rate, kCodecSampleRate); + if (static_cast(audio_24k.size()) < kCodecSampleRate) { + audio_24k.resize(static_cast(kCodecSampleRate), 0.0F); + } + return audio_24k; +} + +std::vector prepare_semantic_audio_16k(const std::vector & mono, + int source_sample_rate) { + auto semantic_16k = resample_mono_if_needed(mono, source_sample_rate, kSemanticSampleRate); + return pad_zeros(semantic_16k, kSemanticPadSamples, kSemanticPadSamples); +} + +std::vector effective_semantic_pos_conv_weight(const assets::TensorSource & source, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size) { + const auto g = source.require_f32(codec_name("semantic_model.encoder.pos_conv_embed." + "conv.parametrizations.weight.original0"), + {1, 1, kernel_size}); + const auto v = source.require_f32(codec_name("semantic_model.encoder.pos_conv_embed." + "conv.parametrizations.weight.original1"), + {out_channels, in_channels, kernel_size}); + std::vector weight(v.size()); + for (int64_t k = 0; k < kernel_size; ++k) { + double sum = 0.0; + for (int64_t out = 0; out < out_channels; ++out) { + for (int64_t in = 0; in < in_channels; ++in) { + const size_t index = + static_cast((out * in_channels + in) * kernel_size + k); + sum += static_cast(v[index]) * static_cast(v[index]); + } + } + const double norm = std::sqrt(sum); + if (norm == 0.0) { + throw std::runtime_error("Higgs TTS semantic positional-conv weight norm is zero"); + } + const float scale_value = + static_cast(static_cast(g[static_cast(k)]) / norm); + for (int64_t out = 0; out < out_channels; ++out) { + for (int64_t in = 0; in < in_channels; ++in) { + const size_t index = + static_cast((out * in_channels + in) * kernel_size + k); + weight[index] = v[index] * scale_value; + } + } + } + return weight; +} + +void load_semantic_tensor(HiggsCodecWeights & weights, + const assets::TensorSource & source, + const std::string & name, + const std::vector & shape, + assets::TensorStorageType storage_type) { + weights.semantic_model.emplace( + name, + weights.store->load_tensor( + source, codec_name("semantic_model." + name), storage_type, shape)); +} + +void load_semantic_f32_tensor(HiggsCodecWeights & weights, + const assets::TensorSource & source, + const std::string & name, + const std::vector & shape) { + weights.semantic_model.emplace( + name, weights.store->load_f32_tensor(source, codec_name("semantic_model." + name), shape)); +} + +void load_hubert_semantic_model_weights(HiggsCodecWeights & weights, + const assets::TensorSource & source, + assets::TensorStorageType storage_type) { + for (int64_t layer = 0; layer < kSemanticConvLayers; ++layer) { + const std::string prefix = "feature_extractor.conv_layers." + std::to_string(layer); + load_semantic_tensor(weights, + source, + prefix + ".conv.weight", + {kSemanticConvDim[static_cast(layer)], + layer == 0 ? 1 : kSemanticConvDim[static_cast(layer - 1)], + kSemanticConvKernel[static_cast(layer)]}, + storage_type); + if (layer == 0) { + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.weight", {kSemanticConvDim[0]}); + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.bias", {kSemanticConvDim[0]}); + } + } + load_semantic_f32_tensor( + weights, source, "feature_projection.layer_norm.weight", {kSemanticConvDim[6]}); + load_semantic_f32_tensor( + weights, source, "feature_projection.layer_norm.bias", {kSemanticConvDim[6]}); + load_semantic_tensor(weights, + source, + "feature_projection.projection.weight", + {kSemanticHiddenSize, kSemanticConvDim[6]}, + storage_type); + load_semantic_f32_tensor( + weights, source, "feature_projection.projection.bias", {kSemanticHiddenSize}); + load_semantic_f32_tensor(weights, source, "encoder.layer_norm.weight", {kSemanticHiddenSize}); + load_semantic_f32_tensor(weights, source, "encoder.layer_norm.bias", {kSemanticHiddenSize}); + weights.semantic_model.emplace( + "encoder.pos_conv_embed.conv.weight", + weights.store->make_f32( + core::TensorShape::from_dims({kSemanticHiddenSize, kSemanticHiddenSize / 16, 128}), + effective_semantic_pos_conv_weight( + source, kSemanticHiddenSize, kSemanticHiddenSize / 16, 128))); + load_semantic_f32_tensor( + weights, source, "encoder.pos_conv_embed.conv.bias", {kSemanticHiddenSize}); + for (int64_t layer = 0; layer < kSemanticLayers; ++layer) { + const std::string prefix = "encoder.layers." + std::to_string(layer); + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.weight", {kSemanticHiddenSize}); + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.bias", {kSemanticHiddenSize}); + load_semantic_f32_tensor( + weights, source, prefix + ".final_layer_norm.weight", {kSemanticHiddenSize}); + load_semantic_f32_tensor( + weights, source, prefix + ".final_layer_norm.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.q_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.q_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.k_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.k_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.v_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.v_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.out_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.out_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".feed_forward.intermediate_dense.weight", + {kSemanticIntermediateSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor(weights, + source, + prefix + ".feed_forward.intermediate_dense.bias", + {kSemanticIntermediateSize}); + load_semantic_tensor(weights, + source, + prefix + ".feed_forward.output_dense.weight", + {kSemanticHiddenSize, kSemanticIntermediateSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".feed_forward.output_dense.bias", {kSemanticHiddenSize}); + } +} + +HiggsCodecResidualUnitWeights load_residual_unit(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels) { + HiggsCodecResidualUnitWeights weights; + weights.snake1 = load_snake(store, source, prefix + ".snake1.alpha", channels); + weights.conv1 = + load_conv1d(store, source, prefix + ".conv1", storage_type, channels, channels, 7, true); + weights.snake2 = load_snake(store, source, prefix + ".snake2.alpha", channels); + weights.conv2 = + load_conv1d(store, source, prefix + ".conv2", storage_type, channels, channels, 1, true); + return weights; +} + +HiggsCodecDecoderBlockWeights load_decoder_block(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t block_index, + assets::TensorStorageType storage_type) { + if (block_index < 0 || block_index >= kDecoderBlockCount) { + throw std::runtime_error("Higgs TTS codec decoder block index is out of range"); + } + const int64_t in_channels = kDecoderChannels[static_cast(block_index)]; + const int64_t out_channels = kDecoderChannels[static_cast(block_index + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block_index)]; + const std::string prefix = "acoustic_decoder.block." + std::to_string(block_index); + + HiggsCodecDecoderBlockWeights weights; + weights.snake = load_snake(store, source, prefix + ".snake1.alpha", in_channels); + weights.conv_transpose = load_conv_transpose1d(store, + source, + prefix + ".conv_t1", + storage_type, + in_channels, + out_channels, + 2 * ratio, + true); + weights.residual_units.reserve(kResidualUnitsPerBlock); + for (int64_t unit = 0; unit < kResidualUnitsPerBlock; ++unit) { + weights.residual_units.push_back( + load_residual_unit(store, + source, + prefix + ".res_unit" + std::to_string(unit + 1), + storage_type, + out_channels)); + } + return weights; +} + +HiggsCodecEncoderBlockWeights load_encoder_block(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t block_index, + assets::TensorStorageType storage_type) { + if (block_index < 0 || block_index >= kEncoderBlockCount) { + throw std::runtime_error("Higgs TTS codec encoder block index is out of range"); + } + const int64_t in_channels = kEncoderChannels[static_cast(block_index)]; + const int64_t out_channels = kEncoderChannels[static_cast(block_index + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block_index)]; + const std::string prefix = "acoustic_encoder.block." + std::to_string(block_index); + + HiggsCodecEncoderBlockWeights weights; + weights.snake = load_snake(store, source, prefix + ".snake1.alpha", in_channels); + weights.conv = load_conv1d( + store, source, prefix + ".conv1", storage_type, out_channels, in_channels, 2 * ratio, true); + weights.residual_units.reserve(kResidualUnitsPerBlock); + for (int64_t unit = 0; unit < kResidualUnitsPerBlock; ++unit) { + weights.residual_units.push_back( + load_residual_unit(store, + source, + prefix + ".res_unit" + std::to_string(unit + 1), + storage_type, + in_channels)); + } + return weights; +} + +HiggsCodecSemanticResidualUnitWeights +load_semantic_residual_unit(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type) { + HiggsCodecSemanticResidualUnitWeights weights; + weights.conv1 = load_conv1d(store, + source, + prefix + ".conv1", + storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 3, + false); + weights.conv2 = load_conv1d(store, + source, + prefix + ".conv2", + storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 1, + false); + return weights; +} + +HiggsCodecSemanticEncoderBlockWeights +load_semantic_encoder_block(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t block_index, + assets::TensorStorageType storage_type) { + if (block_index < 0 || block_index >= kSemanticBlockCount) { + throw std::runtime_error("Higgs TTS codec semantic encoder block index is out of range"); + } + const std::string prefix = "encoder_semantic.conv_blocks." + std::to_string(block_index); + HiggsCodecSemanticEncoderBlockWeights weights; + weights.residual_units.reserve(kSemanticResidualUnitsPerBlock); + for (int64_t unit = 0; unit < kSemanticResidualUnitsPerBlock; ++unit) { + weights.residual_units.push_back(load_semantic_residual_unit( + store, source, prefix + ".res_units." + std::to_string(unit), storage_type)); + } + weights.conv = load_conv1d(store, + source, + prefix + ".conv", + storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 3, + true); + return weights; +} + +HiggsCodecVectorQuantizerWeights load_quantizer(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t index, + assets::TensorStorageType storage_type) { + const std::string prefix = "quantizer.quantizers." + std::to_string(index); + HiggsCodecVectorQuantizerWeights weights; + weights.codebook = store.load_tensor(source, + codec_name(prefix + ".codebook.embed"), + storage_type, + {kCodecCodebookSize, kCodecCodebookDim}); + weights.project_in = load_linear(store, + source, + prefix + ".project_in", + storage_type, + kCodecCodebookDim, + kCodecHiddenSize, + true); + weights.project_out = load_linear(store, + source, + prefix + ".project_out", + storage_type, + kCodecHiddenSize, + kCodecCodebookDim, + true); + return weights; +} + +core::TensorValue +conv_transpose_with_adjusted_output_padding(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::ConvTranspose1dWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t ratio) { + const int64_t kernel = 2 * ratio; + const int64_t padding = (ratio + 1) / 2; + const int64_t output_padding = ratio % 2; + auto full = modules::ConvTranspose1dModule({ + in_channels, + out_channels, + kernel, + static_cast(ratio), + 0, + 1, + weights.bias.has_value(), + }) + .build(ctx, input, weights); + const int64_t cropped_frames = + (input.shape.dims[2] - 1) * ratio - 2 * padding + kernel + output_padding; + if (cropped_frames <= 0 || cropped_frames > full.shape.dims[2]) { + throw std::runtime_error( + "Higgs TTS codec adjusted ConvTranspose1d output length is invalid"); + } + return modules::SliceModule({2, padding, cropped_frames}).build(ctx, full); +} + +core::TensorValue dac_snake(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::Snake1dWeights & weights, + int64_t channels) { + const auto input_ready = core::ensure_backend_addressable_layout(ctx, input); + const auto input_f32 = + input_ready.type == GGML_TYPE_F32 + ? input_ready + : core::wrap_tensor(ggml_cast(ctx.ggml, input_ready.tensor, GGML_TYPE_F32), + input_ready.shape, + GGML_TYPE_F32); + auto alpha = + core::reshape_tensor(ctx, weights.alpha, core::TensorShape::from_dims({1, channels, 1})); + auto ax = core::wrap_tensor( + ggml_mul(ctx.ggml, input_f32.tensor, alpha.tensor), input_f32.shape, GGML_TYPE_F32); + auto s = core::wrap_tensor(ggml_sin(ctx.ggml, ax.tensor), input_f32.shape, GGML_TYPE_F32); + auto s2 = + core::wrap_tensor(ggml_mul(ctx.ggml, s.tensor, s.tensor), input_f32.shape, GGML_TYPE_F32); + auto denom = core::wrap_tensor( + ggml_scale_bias(ctx.ggml, alpha.tensor, 1.0F, 1.0e-9F), alpha.shape, GGML_TYPE_F32); + auto periodic = core::wrap_tensor( + ggml_div(ctx.ggml, s2.tensor, denom.tensor), input_f32.shape, GGML_TYPE_F32); + return core::wrap_tensor( + ggml_add(ctx.ggml, input_f32.tensor, periodic.tensor), input_f32.shape, GGML_TYPE_F32); +} + +core::TensorValue residual_unit(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const HiggsCodecResidualUnitWeights & weights, + int64_t channels, + int64_t dilation) { + auto hidden = dac_snake(ctx, input, weights.snake1, channels); + hidden = modules::Conv1dModule({channels, + channels, + 7, + 1, + static_cast(3 * dilation), + static_cast(dilation), + true}) + .build(ctx, hidden, weights.conv1); + hidden = dac_snake(ctx, hidden, weights.snake2, channels); + hidden = modules::Conv1dModule({channels, channels, 1, 1, 0, 1, true}) + .build(ctx, hidden, weights.conv2); + return modules::AddModule{}.build(ctx, input, hidden); +} + +core::TensorValue contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return core::ensure_backend_addressable_layout(ctx, value); +} + +core::TensorValue transpose_bct_btc(core::ModuleBuildContext & ctx, + const core::TensorValue & value) { + return modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, value); +} + +core::TensorValue add_same(core::ModuleBuildContext & ctx, + const core::TensorValue & lhs, + const core::TensorValue & rhs) { + return modules::AddModule{}.build(ctx, lhs, rhs); +} + +core::TensorValue +scale(core::ModuleBuildContext & ctx, const core::TensorValue & value, float factor) { + return core::wrap_tensor( + ggml_scale(ctx.ggml, contiguous(ctx, value).tensor, factor), value.shape, GGML_TYPE_F32); +} + +core::TensorValue group_norm_affine(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t groups, + float eps, + const modules::NormWeights & weights) { + core::TensorValue output; + if (input.shape.rank == 3 && groups == input.shape.dims[1]) { + auto input_f32 = input.type == GGML_TYPE_F32 + ? input + : core::wrap_tensor(ggml_cast(ctx.ggml, input.tensor, GGML_TYPE_F32), + input.shape, + GGML_TYPE_F32); + auto mean = modules::ReduceMeanModule({2}).build(ctx, input_f32); + auto mean_rep = core::wrap_tensor( + ggml_repeat(ctx.ggml, mean.tensor, input_f32.tensor), input_f32.shape, GGML_TYPE_F32); + auto centered = core::wrap_tensor( + ggml_sub(ctx.ggml, input_f32.tensor, mean_rep.tensor), input_f32.shape, GGML_TYPE_F32); + auto variance = modules::ReduceMeanModule({2}).build( + ctx, modules::MulModule().build(ctx, centered, centered)); + auto stddev = core::wrap_tensor( + ggml_sqrt(ctx.ggml, ggml_scale_bias(ctx.ggml, variance.tensor, 1.0F, eps)), + variance.shape, + GGML_TYPE_F32); + auto stddev_rep = core::wrap_tensor( + ggml_repeat(ctx.ggml, stddev.tensor, input_f32.tensor), input_f32.shape, GGML_TYPE_F32); + output = core::wrap_tensor( + ggml_div(ctx.ggml, centered.tensor, stddev_rep.tensor), input_f32.shape, GGML_TYPE_F32); + } else { + output = core::wrap_tensor( + ggml_group_norm(ctx.ggml, input.tensor, groups, eps), input.shape, GGML_TYPE_F32); + } + if (weights.weight.has_value()) { + auto weight = core::reshape_tensor( + ctx, *weights.weight, core::TensorShape::from_dims({1, input.shape.dims[1], 1})); + auto repeated = core::wrap_tensor( + ggml_repeat(ctx.ggml, weight.tensor, output.tensor), output.shape, GGML_TYPE_F32); + output = core::wrap_tensor( + ggml_mul(ctx.ggml, output.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); + } + if (weights.bias.has_value()) { + auto bias = core::reshape_tensor( + ctx, *weights.bias, core::TensorShape::from_dims({1, input.shape.dims[1], 1})); + auto repeated = core::wrap_tensor( + ggml_repeat(ctx.ggml, bias.tensor, output.tensor), output.shape, GGML_TYPE_F32); + output = core::wrap_tensor( + ggml_add(ctx.ggml, output.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); + } + return output; +} + +core::TensorValue grouped_pos_conv(core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const modules::Conv1dWeights & weights) { + constexpr int64_t groups = 16; + constexpr int64_t channels_per_group = kSemanticHiddenSize / groups; + const auto input_contiguous = contiguous(ctx, input_bct); + core::TensorValue out; + for (int64_t group = 0; group < groups; ++group) { + auto input_group = modules::SliceModule({1, group * channels_per_group, channels_per_group}) + .build(ctx, input_contiguous); + auto weight_group = + modules::SliceModule({0, group * channels_per_group, channels_per_group}) + .build(ctx, weights.weight); + modules::Conv1dWeights group_weights{weight_group, std::nullopt}; + if (weights.bias.has_value()) { + group_weights.bias = + modules::SliceModule({0, group * channels_per_group, channels_per_group}) + .build(ctx, *weights.bias); + } + auto group_out = + modules::Conv1dModule( + {channels_per_group, channels_per_group, 128, 1, 64, 1, weights.bias.has_value()}) + .build(ctx, input_group, group_weights); + out = out.valid() ? modules::ConcatModule({1}).build(ctx, out, group_out) : group_out; + } + out = modules::SliceModule({2, 0, input_bct.shape.dims[2]}).build(ctx, out); + return modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, out); +} + +core::TensorValue semantic_self_attention(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + const HiggsCodecWeights & weights, + int64_t layer_index) { + constexpr int64_t head_dim = kSemanticHiddenSize / kSemanticAttentionHeads; + const std::string prefix = "encoder.layers." + std::to_string(layer_index) + ".attention"; + auto q = modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".q_proj")); + auto k = modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".k_proj")); + auto v = modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".v_proj")); + q = core::reshape_tensor(ctx, + contiguous(ctx, q), + core::TensorShape::from_dims({hidden_btc.shape.dims[0], + hidden_btc.shape.dims[1], + kSemanticAttentionHeads, + head_dim})); + k = core::reshape_tensor(ctx, + contiguous(ctx, k), + core::TensorShape::from_dims({hidden_btc.shape.dims[0], + hidden_btc.shape.dims[1], + kSemanticAttentionHeads, + head_dim})); + v = core::reshape_tensor(ctx, + contiguous(ctx, v), + core::TensorShape::from_dims({hidden_btc.shape.dims[0], + hidden_btc.shape.dims[1], + kSemanticAttentionHeads, + head_dim})); + q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, q); + k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k); + v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v); + const auto k_t = modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k); + auto scores = modules::MatMulModule().build(ctx, q, k_t); + scores = scale(ctx, scores, static_cast(1.0 / std::sqrt(static_cast(head_dim)))); + auto attn = core::wrap_tensor( + ggml_soft_max(ctx.ggml, contiguous(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); + auto context = modules::MatMulModule().build(ctx, attn, v); + context = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context); + context = core::reshape_tensor( + ctx, + contiguous(ctx, context), + core::TensorShape::from_dims( + {hidden_btc.shape.dims[0], hidden_btc.shape.dims[1], kSemanticHiddenSize})); + return modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, context, semantic_linear_weights(weights, prefix + ".out_proj")); +} + +core::TensorValue semantic_feed_forward(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + const HiggsCodecWeights & weights, + int64_t layer_index) { + const std::string prefix = "encoder.layers." + std::to_string(layer_index) + ".feed_forward"; + auto x = + modules::LinearModule({kSemanticHiddenSize, kSemanticIntermediateSize, true, GGML_PREC_F32}) + .build( + ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".intermediate_dense")); + x = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, x); + return modules::LinearModule( + {kSemanticIntermediateSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, x, semantic_linear_weights(weights, prefix + ".output_dense")); +} + +core::TensorValue downsample_time_by_2(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + int64_t target_frames) { + core::TensorValue out; + for (int64_t frame = 0; frame < target_frames; ++frame) { + auto slice = modules::SliceModule({1, frame * 2, 1}).build(ctx, hidden_btc); + out = out.valid() ? modules::ConcatModule({1}).build(ctx, out, slice) : slice; + } + return out; +} + +struct HiggsCodecEncodeGraphValues { + std::array(kCodecCodebooks)> codes = {}; +}; + +core::TensorValue hubert_hidden_state_mean(core::ModuleBuildContext & ctx, + const core::TensorValue & input_values, + const HiggsCodecWeights & weights, + int64_t target_frames) { + auto hidden = core::reshape_tensor( + ctx, + input_values, + core::TensorShape::from_dims({input_values.shape.dims[0], 1, input_values.shape.dims[1]})); + int64_t in_channels = 1; + for (int64_t index = 0; index < kSemanticConvLayers; ++index) { + const std::string prefix = "feature_extractor.conv_layers." + std::to_string(index); + hidden = modules::Conv1dModule( + {in_channels, + kSemanticConvDim[static_cast(index)], + kSemanticConvKernel[static_cast(index)], + static_cast(kSemanticConvStride[static_cast(index)]), + 0, + 1, + false}) + .build(ctx, hidden, semantic_conv_weights(weights, prefix + ".conv", false)); + if (index == 0) { + hidden = group_norm_affine(ctx, + hidden, + kSemanticConvDim[0], + kSemanticLayerNormEps, + semantic_norm_weights(weights, prefix + ".layer_norm")); + } + hidden = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, hidden); + in_channels = kSemanticConvDim[static_cast(index)]; + } + hidden = transpose_bct_btc(ctx, hidden); + hidden = + modules::LayerNormModule({kSemanticConvDim[6], kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, "feature_projection.layer_norm")); + hidden = + modules::LinearModule({kSemanticConvDim[6], kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden, semantic_linear_weights(weights, "feature_projection.projection")); + + auto pos = + grouped_pos_conv(ctx, + transpose_bct_btc(ctx, hidden), + semantic_conv_weights(weights, "encoder.pos_conv_embed.conv", true)); + hidden = add_same(ctx, hidden, transpose_bct_btc(ctx, pos)); + hidden = modules::LayerNormModule({kSemanticHiddenSize, kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, "encoder.layer_norm")); + + auto sum = hidden; + for (int64_t layer = 0; layer < kSemanticLayers; ++layer) { + const std::string prefix = "encoder.layers." + std::to_string(layer); + const auto attn_residual = hidden; + hidden = semantic_self_attention(ctx, hidden, weights, layer); + hidden = add_same(ctx, attn_residual, hidden); + hidden = modules::LayerNormModule({kSemanticHiddenSize, kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, prefix + ".layer_norm")); + hidden = add_same(ctx, hidden, semantic_feed_forward(ctx, hidden, weights, layer)); + hidden = + modules::LayerNormModule({kSemanticHiddenSize, kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, prefix + ".final_layer_norm")); + sum = add_same(ctx, sum, hidden); + } + auto hidden_mean = scale(ctx, sum, 1.0F / static_cast(kSemanticLayers + 1)); + auto features = downsample_time_by_2(ctx, hidden_mean, target_frames); + return features; +} + +core::TensorValue semantic_residual_unit(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const HiggsCodecSemanticResidualUnitWeights & weights) { + auto hidden = modules::EluModule().build(ctx, input); + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 3, 1, 1, 1, false}) + .build(ctx, hidden, weights.conv1); + hidden = modules::EluModule().build(ctx, hidden); + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 1, 1, 0, 1, false}) + .build(ctx, hidden, weights.conv2); + return modules::AddModule{}.build(ctx, input, hidden); +} + +core::TensorValue semantic_encoder(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + const HiggsCodecWeights & weights) { + auto hidden = transpose_bct_btc(ctx, hidden_btc); + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 3, 1, 1, 1, false}) + .build(ctx, hidden, weights.semantic_encoder_input); + for (const auto & block : weights.semantic_encoder_blocks) { + for (const auto & unit : block.residual_units) { + hidden = semantic_residual_unit(ctx, hidden, unit); + } + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 3, 1, 1, 1, true}) + .build(ctx, hidden, block.conv); + } + return hidden; +} + +core::TensorValue acoustic_encoder(core::ModuleBuildContext & ctx, + const core::TensorValue & waveform, + const HiggsCodecWeights & weights, + int64_t target_frames) { + auto hidden = core::reshape_tensor( + ctx, waveform, core::TensorShape::from_dims({1, 1, waveform.shape.dims[0]})); + hidden = modules::Conv1dModule({1, kEncoderChannels[0], 7, 1, 3, 1, true}) + .build(ctx, hidden, weights.acoustic_encoder_input); + for (int64_t block = 0; block < kEncoderBlockCount; ++block) { + const int64_t in_channels = kEncoderChannels[static_cast(block)]; + const int64_t out_channels = kEncoderChannels[static_cast(block + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block)]; + const auto & block_weights = weights.acoustic_encoder_blocks[static_cast(block)]; + for (size_t unit_index = 0; unit_index < block_weights.residual_units.size(); + ++unit_index) { + hidden = residual_unit(ctx, + hidden, + block_weights.residual_units[unit_index], + in_channels, + kResidualDilations[unit_index]); + } + hidden = dac_snake(ctx, hidden, block_weights.snake, in_channels); + hidden = modules::Conv1dModule({in_channels, + out_channels, + 2 * ratio, + static_cast(ratio), + static_cast((ratio + 1) / 2), + 1, + true}) + .build(ctx, hidden, block_weights.conv); + } + hidden = dac_snake(ctx, hidden, weights.acoustic_encoder_output_snake, kEncoderChannels[5]); + hidden = modules::Conv1dModule({kEncoderChannels[5], kAcousticHiddenSize, 3, 1, 1, 1, true}) + .build(ctx, hidden, weights.acoustic_encoder_output); + if (hidden.shape.dims[2] != target_frames) { + hidden = modules::SliceModule({2, 0, target_frames}).build(ctx, hidden); + } + return hidden; +} + +std::array(kCodecCodebooks)> +quantizer_encode(core::ModuleBuildContext & ctx, + const core::TensorValue & embeddings_bct, + const HiggsCodecWeights & weights) { + auto residual = embeddings_bct; + std::array(kCodecCodebooks)> all_codes{}; + for (int64_t codebook = 0; codebook < kCodecCodebooks; ++codebook) { + auto hidden = transpose_bct_btc(ctx, residual); + hidden = + modules::LinearModule({kCodecHiddenSize, kCodecCodebookDim, true, GGML_PREC_F32}) + .build(ctx, hidden, weights.quantizers[static_cast(codebook)].project_in); + auto flat = core::reshape_tensor( + ctx, + contiguous(ctx, hidden), + core::TensorShape::from_dims({hidden.shape.dims[1], kCodecCodebookDim})); + auto codebook_weight = weights.quantizers[static_cast(codebook)].codebook; + auto codebook_t = modules::TransposeModule({{1, 0, 2, 3}, 2}).build(ctx, codebook_weight); + auto dot = modules::MatMulModule().build(ctx, flat, codebook_t); + auto x2 = + modules::ReduceSumModule({1}).build(ctx, modules::MulModule().build(ctx, flat, flat)); + x2 = core::wrap_tensor( + ggml_repeat(ctx.ggml, x2.tensor, dot.tensor), dot.shape, GGML_TYPE_F32); + const auto codebook_weight_f32 = + codebook_weight.type == GGML_TYPE_F32 + ? codebook_weight + : core::wrap_tensor(ggml_cast(ctx.ggml, codebook_weight.tensor, GGML_TYPE_F32), + codebook_weight.shape, + GGML_TYPE_F32); + auto e2 = modules::ReduceSumModule({1}).build( + ctx, modules::MulModule().build(ctx, codebook_weight_f32, codebook_weight_f32)); + e2 = core::reshape_tensor(ctx, e2, core::TensorShape::from_dims({1, kCodecCodebookSize})); + e2 = core::wrap_tensor( + ggml_repeat(ctx.ggml, e2.tensor, dot.tensor), dot.shape, GGML_TYPE_F32); + auto logits = + core::wrap_tensor(ggml_scale(ctx.ggml, dot.tensor, 2.0F), dot.shape, GGML_TYPE_F32); + logits = core::wrap_tensor( + ggml_sub(ctx.ggml, logits.tensor, x2.tensor), logits.shape, GGML_TYPE_F32); + logits = core::wrap_tensor( + ggml_sub(ctx.ggml, logits.tensor, e2.tensor), logits.shape, GGML_TYPE_F32); + auto ids = core::wrap_tensor(ggml_argmax(ctx.ggml, contiguous(ctx, logits).tensor), + core::TensorShape::from_dims({embeddings_bct.shape.dims[2]}), + GGML_TYPE_I32); + all_codes[static_cast(codebook)] = ids; + auto quantized = modules::EmbeddingModule({kCodecCodebookSize, kCodecCodebookDim}) + .build(ctx, ids, codebook_weight); + quantized = + modules::LinearModule({kCodecCodebookDim, kCodecHiddenSize, true, GGML_PREC_F32}) + .build( + ctx, quantized, weights.quantizers[static_cast(codebook)].project_out); + quantized = core::reshape_tensor( + ctx, + contiguous(ctx, quantized), + core::TensorShape::from_dims({1, embeddings_bct.shape.dims[2], kCodecHiddenSize})); + quantized = transpose_bct_btc(ctx, quantized); + residual = core::wrap_tensor( + ggml_sub(ctx.ggml, residual.tensor, quantized.tensor), residual.shape, GGML_TYPE_F32); + } + return all_codes; +} + +HiggsCodecEncodeGraphValues codec_encode(core::ModuleBuildContext & ctx, + const core::TensorValue & waveform_24k, + const core::TensorValue & semantic_waveform_16k, + const HiggsCodecWeights & weights, + int64_t target_frames) { + HiggsCodecEncodeGraphValues out; + auto semantic = + hubert_hidden_state_mean(ctx, semantic_waveform_16k, weights, target_frames); + semantic = semantic_encoder(ctx, semantic, weights); + if (semantic.shape.dims[2] != target_frames) { + semantic = modules::SliceModule({2, 0, target_frames}).build(ctx, semantic); + } + auto acoustic = acoustic_encoder(ctx, waveform_24k, weights, target_frames); + auto concat = modules::ConcatModule({1}).build(ctx, acoustic, semantic); + auto hidden = transpose_bct_btc(ctx, concat); + hidden = modules::LinearModule({kCodecProjectInputSize, kCodecHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden, weights.codec_project); + hidden = transpose_bct_btc(ctx, hidden); + out.codes = quantizer_encode(ctx, hidden, weights); + return out; +} + +core::TensorValue quantizer_decode(core::ModuleBuildContext & ctx, + ggml_tensor * codes, + const HiggsCodecWeights & weights, + int64_t frames) { + auto codes_value = core::wrap_tensor( + codes, core::TensorShape::from_dims({frames, kCodecCodebooks}), GGML_TYPE_I32); + std::vector projected; + projected.reserve(weights.quantizers.size()); + for (int64_t codebook = 0; codebook < kCodecCodebooks; ++codebook) { + auto ids = modules::SliceModule({1, codebook, 1}).build(ctx, codes_value); + ids = core::reshape_tensor(ctx, + core::ensure_backend_addressable_layout(ctx, ids), + core::TensorShape::from_dims({frames})); + auto hidden = + modules::EmbeddingModule({kCodecCodebookSize, kCodecCodebookDim}) + .build(ctx, ids, weights.quantizers[static_cast(codebook)].codebook); + hidden = + modules::LinearModule({kCodecCodebookDim, kCodecHiddenSize, true}) + .build(ctx, hidden, weights.quantizers[static_cast(codebook)].project_out); + projected.push_back(hidden); + } + auto sum = projected.front(); + for (size_t index = 1; index < projected.size(); ++index) { + sum = modules::AddModule{}.build(ctx, sum, projected[index]); + } + return sum; +} + +core::TensorValue acoustic_decoder(core::ModuleBuildContext & ctx, + const core::TensorValue & quantized, + const HiggsCodecWeights & weights) { + auto hidden = modules::LinearModule({kCodecHiddenSize, kAcousticHiddenSize, true}) + .build(ctx, quantized, weights.acoustic_project); + hidden = core::reshape_tensor( + ctx, hidden, core::TensorShape::from_dims({1, hidden.shape.dims[0], kAcousticHiddenSize})); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, hidden); + hidden = modules::Conv1dModule({kAcousticHiddenSize, kDecoderChannels[0], 7, 1, 3, 1, true}) + .build(ctx, hidden, weights.acoustic_decoder_input); + for (int64_t block = 0; block < kDecoderBlockCount; ++block) { + const int64_t in_channels = kDecoderChannels[static_cast(block)]; + const int64_t out_channels = kDecoderChannels[static_cast(block + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block)]; + const auto & block_weights = weights.acoustic_decoder_blocks[static_cast(block)]; + hidden = dac_snake(ctx, hidden, block_weights.snake, in_channels); + hidden = conv_transpose_with_adjusted_output_padding( + ctx, hidden, block_weights.conv_transpose, in_channels, out_channels, ratio); + for (size_t unit_index = 0; unit_index < block_weights.residual_units.size(); + ++unit_index) { + hidden = residual_unit(ctx, + hidden, + block_weights.residual_units[unit_index], + out_channels, + kResidualDilations[unit_index]); + } + } + hidden = dac_snake(ctx, hidden, weights.acoustic_decoder_output_snake, kDecoderChannels[5]); + return modules::Conv1dModule({kDecoderChannels[5], 1, 7, 1, 3, 1, true}) + .build(ctx, hidden, weights.acoustic_decoder_output); +} + +} // namespace + +class HiggsCodecEncodeGraph { +public: + HiggsCodecEncodeGraph(const HiggsCodecRuntime * runtime, + int64_t acoustic_samples, + int64_t semantic_samples, + int64_t frames) + : runtime_(runtime), acoustic_samples_(acoustic_samples), + semantic_samples_(semantic_samples), frames_(frames) { + if (runtime_ == nullptr) { + throw std::runtime_error("Higgs TTS codec encode graph requires runtime"); + } + if (acoustic_samples_ <= 0 || semantic_samples_ <= 0 || frames_ <= 0) { + throw std::runtime_error("Higgs TTS codec encode graph requires positive dimensions"); + } + const auto build_start = Clock::now(); + ggml_init_params params{runtime_->encode_graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS codec encode graph context"); + } + core::ModuleBuildContext build_ctx{ + ctx_.get(), "higgs_audio_tts.codec.encode", runtime_->backend_type()}; + acoustic_input_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_F32, acoustic_samples_); + semantic_input_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_F32, semantic_samples_); + ggml_set_input(acoustic_input_); + ggml_set_input(semantic_input_); + auto acoustic = core::wrap_tensor( + acoustic_input_, core::TensorShape::from_dims({acoustic_samples_}), GGML_TYPE_F32); + auto semantic = core::wrap_tensor( + semantic_input_, core::TensorShape::from_dims({1, semantic_samples_}), GGML_TYPE_F32); + auto encoded = codec_encode( + build_ctx, acoustic, semantic, runtime_->weights(), frames_); + graph_ = ggml_new_graph_custom(ctx_.get(), 262144, false); + for (size_t codebook = 0; codebook < outputs_.size(); ++codebook) { + outputs_[codebook] = encoded.codes[codebook].tensor; + ggml_set_output(outputs_[codebook]); + ggml_build_forward_expand(graph_, outputs_[codebook]); + } + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + throw std::runtime_error("failed to allocate Higgs TTS codec encode graph"); + } + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~HiggsCodecEncodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(const HiggsCodecRuntime & runtime, + int64_t acoustic_samples, + int64_t semantic_samples, + int64_t frames) const { + return runtime_ == &runtime && acoustic_samples_ == acoustic_samples && + semantic_samples_ == semantic_samples && frames_ == frames; + } + + HiggsCodecEncodeOutput + run(const std::vector & acoustic, const std::vector & semantic, int64_t frames) { + if (static_cast(acoustic.size()) != acoustic_samples_ || + static_cast(semantic.size()) != semantic_samples_ || frames != frames_) { + throw std::runtime_error("Higgs TTS codec encode graph shape mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set( + acoustic_input_, acoustic.data(), 0, acoustic.size() * sizeof(float)); + ggml_backend_tensor_set( + semantic_input_, semantic.data(), 0, semantic.size() * sizeof(float)); + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode_input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS codec encode graph compute failed"); + } + timing_start = Clock::now(); + HiggsCodecEncodeOutput out; + out.frames = frames; + out.codebooks = kCodecCodebooks; + out.codes.resize(static_cast(frames * kCodecCodebooks)); + std::vector codebook_codes(static_cast(frames_)); + for (int64_t codebook = 0; codebook < kCodecCodebooks; ++codebook) { + ggml_backend_tensor_get(outputs_[static_cast(codebook)], + codebook_codes.data(), + 0, + codebook_codes.size() * sizeof(int32_t)); + for (int64_t frame = 0; frame < frames_; ++frame) { + out.codes[static_cast(frame * kCodecCodebooks + codebook)] = + codebook_codes[static_cast(frame)]; + } + } + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode_output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return out; + } + +private: + const HiggsCodecRuntime * runtime_ = nullptr; + int64_t acoustic_samples_ = 0; + int64_t semantic_samples_ = 0; + int64_t frames_ = 0; + std::unique_ptr ctx_; + ggml_tensor * acoustic_input_ = nullptr; + ggml_tensor * semantic_input_ = nullptr; + std::array(kCodecCodebooks)> outputs_ = {}; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +class HiggsCodecDecodeGraph { +public: + HiggsCodecDecodeGraph(const HiggsCodecRuntime * runtime, int64_t frames) + : runtime_(runtime), capacity_frames_(frames) { + if (runtime_ == nullptr) { + throw std::runtime_error("Higgs TTS codec decode graph requires runtime"); + } + if (capacity_frames_ <= 0) { + throw std::runtime_error("Higgs TTS codec decode graph requires positive frame count"); + } + const auto build_start = Clock::now(); + ggml_init_params params{runtime_->decode_graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS codec decode graph context"); + } + core::ModuleBuildContext build_ctx{ + ctx_.get(), "higgs_audio_tts.codec.decode", runtime_->backend_type()}; + codes_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_I32, kCodecCodebooks, capacity_frames_); + frame_mask_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, 1, capacity_frames_); + ggml_set_input(codes_); + ggml_set_input(frame_mask_); + auto hidden = quantizer_decode(build_ctx, codes_, runtime_->weights(), capacity_frames_); + const auto mask = core::wrap_tensor( + frame_mask_, core::TensorShape::from_dims({capacity_frames_, 1}), GGML_TYPE_F32); + hidden = modules::MulModule{}.build( + build_ctx, + hidden, + core::wrap_tensor(ggml_repeat(build_ctx.ggml, mask.tensor, hidden.tensor), + hidden.shape, + GGML_TYPE_F32)); + auto audio = acoustic_decoder(build_ctx, hidden, runtime_->weights()); + output_ = audio.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, output_); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } + throw std::runtime_error("failed to allocate Higgs TTS codec decode graph"); + } + code_scratch_.assign(static_cast(capacity_frames_ * kCodecCodebooks), 0); + frame_mask_values_.assign(static_cast(capacity_frames_), 0.0F); + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~HiggsCodecDecodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + bool matches(const HiggsCodecRuntime & runtime, int64_t frames) const { + return runtime_ == &runtime && frames <= capacity_frames_; + } + + int64_t capacity_frames() const { return capacity_frames_; } + + HiggsCodecDecodeOutput run(const std::vector & codes, int64_t frames) { + if (frames <= 0 || frames > capacity_frames_) { + throw std::runtime_error("Higgs TTS codec decode frame count exceeds graph capacity"); + } + if (static_cast(codes.size()) != frames * kCodecCodebooks) { + throw std::runtime_error("Higgs TTS codec decode code matrix shape mismatch"); + } + for (const int32_t code : codes) { + if (code < 0 || code >= kCodecCodebookSize) { + throw std::runtime_error("Higgs TTS codec decode code is outside codebook range"); + } + } + std::fill(code_scratch_.begin(), code_scratch_.end(), 0); + for (int64_t frame = 0; frame < frames; ++frame) { + const auto src = codes.begin() + static_cast(frame * kCodecCodebooks); + const auto dst = + code_scratch_.begin() + static_cast(frame * kCodecCodebooks); + std::copy_n(src, static_cast(kCodecCodebooks), dst); + } + std::fill(frame_mask_values_.begin(), frame_mask_values_.end(), 0.0F); + std::fill(frame_mask_values_.begin(), + frame_mask_values_.begin() + static_cast(frames), + 1.0F); + auto timing_start = Clock::now(); + ggml_backend_tensor_set( + codes_, code_scratch_.data(), 0, code_scratch_.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + frame_mask_, frame_mask_values_.data(), 0, frame_mask_values_.size() * sizeof(float)); + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode_input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS codec decode graph compute failed"); + } + HiggsCodecDecodeOutput out; + out.sample_rate = kCodecSampleRate; + out.channels = 1; + out.samples = frames; + for (int64_t block = 0; block < kDecoderBlockCount; ++block) { + out.samples *= kUpsampleRatios[static_cast(block)]; + } + out.values.resize(static_cast(out.samples)); + timing_start = Clock::now(); + ggml_backend_tensor_get(output_, out.values.data(), 0, out.values.size() * sizeof(float)); + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode_output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return out; + } + +private: + const HiggsCodecRuntime * runtime_ = nullptr; + int64_t capacity_frames_ = 0; + std::unique_ptr ctx_; + ggml_tensor * codes_ = nullptr; + ggml_tensor * frame_mask_ = nullptr; + ggml_tensor * output_ = nullptr; + std::vector code_scratch_; + std::vector frame_mask_values_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; +}; + +HiggsCodecWeights load_higgs_codec_decode_weights(const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) { + if (assets.weights == nullptr) { + throw std::runtime_error("Higgs TTS codec weights require tensor source"); + } + if (backend == nullptr) { + throw std::runtime_error("Higgs TTS codec backend is not initialized"); + } + HiggsCodecWeights weights; + weights.store = std::make_shared( + backend, backend_type, "higgs_audio_tts.codec.weights", weight_context_bytes); + const auto & source = *assets.weights; + load_hubert_semantic_model_weights(weights, source, weight_storage_type); + weights.quantizers.reserve(kCodecCodebooks); + for (int64_t index = 0; index < kCodecCodebooks; ++index) { + weights.quantizers.push_back( + load_quantizer(*weights.store, source, index, weight_storage_type)); + } + weights.acoustic_encoder_input = load_conv1d(*weights.store, + source, + "acoustic_encoder.conv1", + weight_storage_type, + kEncoderChannels[0], + 1, + 7, + true); + weights.acoustic_encoder_blocks.reserve(kEncoderBlockCount); + for (int64_t block = 0; block < kEncoderBlockCount; ++block) { + weights.acoustic_encoder_blocks.push_back( + load_encoder_block(*weights.store, source, block, weight_storage_type)); + } + weights.acoustic_encoder_output_snake = + load_snake(*weights.store, source, "acoustic_encoder.snake1.alpha", kEncoderChannels[5]); + weights.acoustic_encoder_output = load_conv1d(*weights.store, + source, + "acoustic_encoder.conv2", + weight_storage_type, + kAcousticHiddenSize, + kEncoderChannels[5], + 3, + true); + weights.semantic_encoder_input = load_conv1d(*weights.store, + source, + "encoder_semantic.conv", + weight_storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 3, + false); + weights.semantic_encoder_blocks.reserve(kSemanticBlockCount); + for (int64_t block = 0; block < kSemanticBlockCount; ++block) { + weights.semantic_encoder_blocks.push_back( + load_semantic_encoder_block(*weights.store, source, block, weight_storage_type)); + } + weights.codec_project = load_linear(*weights.store, + source, + "fc", + weight_storage_type, + kCodecHiddenSize, + kCodecProjectInputSize, + true); + weights.acoustic_project = load_linear(*weights.store, + source, + "fc2", + weight_storage_type, + kAcousticHiddenSize, + kCodecHiddenSize, + true); + weights.acoustic_decoder_input = load_conv1d(*weights.store, + source, + "acoustic_decoder.conv1", + weight_storage_type, + kDecoderChannels[0], + kAcousticHiddenSize, + 7, + true); + weights.acoustic_decoder_blocks.reserve(kDecoderBlockCount); + for (int64_t block = 0; block < kDecoderBlockCount; ++block) { + weights.acoustic_decoder_blocks.push_back( + load_decoder_block(*weights.store, source, block, weight_storage_type)); + } + weights.acoustic_decoder_output_snake = + load_snake(*weights.store, source, "acoustic_decoder.snake1.alpha", kDecoderChannels[5]); + weights.acoustic_decoder_output = load_conv1d(*weights.store, + source, + "acoustic_decoder.conv2", + weight_storage_type, + 1, + kDecoderChannels[5], + 7, + true); + weights.store->upload(); + return weights; +} + +HiggsCodecRuntime::HiggsCodecRuntime(std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t decode_graph_arena_bytes, + size_t encode_graph_arena_bytes, + assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)), backend_(execution.backend()), + backend_type_(execution.backend_type()), threads_(std::max(1, execution.config().threads)), + decode_graph_arena_bytes_(decode_graph_arena_bytes), + encode_graph_arena_bytes_(encode_graph_arena_bytes), + weights_(nullptr) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS codec runtime requires assets"); + } + if (assets_->weights == nullptr) { + throw std::runtime_error("Higgs TTS codec runtime requires tensor source"); + } + weights_ = std::make_shared(load_higgs_codec_decode_weights( + *assets_, backend_, backend_type_, weight_context_bytes, weight_storage_type)); + if (decode_graph_arena_bytes_ == 0) { + throw std::runtime_error("Higgs TTS codec decode graph arena bytes must be non-zero"); + } + if (encode_graph_arena_bytes_ == 0) { + throw std::runtime_error("Higgs TTS codec encode graph arena bytes must be non-zero"); + } +} + +HiggsCodecRuntime::~HiggsCodecRuntime() = default; + +const HiggsCodecWeights & HiggsCodecRuntime::weights() const noexcept { return *weights_; } + +ggml_backend_t HiggsCodecRuntime::backend() const noexcept { return backend_; } + +core::BackendType HiggsCodecRuntime::backend_type() const noexcept { return backend_type_; } + +int HiggsCodecRuntime::threads() const noexcept { return threads_; } + +size_t HiggsCodecRuntime::decode_graph_arena_bytes() const noexcept { + return decode_graph_arena_bytes_; +} + +size_t HiggsCodecRuntime::encode_graph_arena_bytes() const noexcept { + return encode_graph_arena_bytes_; +} + +HiggsCodecEncodeOutput +HiggsCodecRuntime::encode_reference(const runtime::AudioBuffer & audio) const { + const auto mono = prepare_mono_audio(audio); + const auto acoustic_24k_base = prepare_codec_audio_24k(mono, audio.sample_rate); + auto semantic_16k = prepare_semantic_audio_16k(mono, audio.sample_rate); + const int64_t frames = semantic_feature_frames(static_cast(semantic_16k.size())); + + std::vector acoustic_24k = acoustic_24k_base; + const int64_t acoustic_frames = + acoustic_encoder_frames(static_cast(acoustic_24k.size())); + if (acoustic_frames != frames) { + acoustic_24k = pad_zeros(acoustic_24k_base, kCodecPadSamples, kCodecPadSamples); + const int64_t padded_frames = + acoustic_encoder_frames(static_cast(acoustic_24k.size())); + if (padded_frames != frames) { + throw std::runtime_error("Higgs TTS codec acoustic and semantic encoder " + "frame counts do not match"); + } + } + + if (encode_graph_ == nullptr || + !encode_graph_->matches(*this, + static_cast(acoustic_24k.size()), + static_cast(semantic_16k.size()), + frames)) { + encode_graph_.reset(); + encode_graph_ = + std::make_unique(this, + static_cast(acoustic_24k.size()), + static_cast(semantic_16k.size()), + frames); + } + engine::debug::trace_log_scalar("higgs_audio_tts.codec.encode.input_frames", frames); + engine::debug::trace_log_f32("higgs_audio_tts.codec.encode.input_acoustic_24k", + {static_cast(acoustic_24k.size())}, + acoustic_24k); + engine::debug::trace_log_f32("higgs_audio_tts.codec.encode.input_semantic_16k", + {static_cast(semantic_16k.size())}, + semantic_16k); + return encode_graph_->run(acoustic_24k, semantic_16k, frames); +} + +HiggsCodecDecodeOutput HiggsCodecRuntime::decode_codes(const std::vector & codes, + int64_t frames, + int64_t codebooks) const { + if (frames <= 0) { + throw std::runtime_error("Higgs TTS codec decode requires positive frame count"); + } + if (codebooks != kCodecCodebooks) { + throw std::runtime_error("Higgs TTS codec decode requires exactly 8 codebooks"); + } + if (static_cast(codes.size()) != frames * codebooks) { + throw std::runtime_error("Higgs TTS codec decode code count mismatch"); + } + engine::debug::trace_log_scalar("higgs_audio_tts.codec.decode.input_frames", frames); + engine::debug::trace_log_scalar("higgs_audio_tts.codec.decode.input_codebooks", codebooks); + engine::debug::trace_log_i32("higgs_audio_tts.codec.decode.input_codes", + {frames, codebooks}, + codes); + + auto run_window = [&](const std::vector & window_codes, + int64_t window_frames, + int64_t min_capacity_frames) -> HiggsCodecDecodeOutput { + const int64_t bucketed_frames = + ((std::max(window_frames, min_capacity_frames) + + kCodecDecodeCapacityBucketFrames - 1) / + kCodecDecodeCapacityBucketFrames) * + kCodecDecodeCapacityBucketFrames; + encode_graph_.reset(); + if (decode_graph_ == nullptr || !decode_graph_->matches(*this, window_frames)) { + decode_graph_.reset(); + decode_graph_ = std::make_unique(this, bucketed_frames); + } + return decode_graph_->run(window_codes, window_frames); + }; + + if (frames <= kCodecFullDecodeMaxFrames) { + return run_window(codes, frames, frames); + } + + HiggsCodecDecodeOutput out; + out.sample_rate = kCodecSampleRate; + out.channels = 1; + out.samples = frames * kCodecHopLength; + out.values.reserve(static_cast(out.samples)); + + std::vector window_codes; + int64_t emitted_frames = 0; + while (emitted_frames < frames) { + const int64_t window_begin = + std::max(0, emitted_frames - kCodecDecodeOverlapFrames); + const int64_t emit_end = + std::min(frames, emitted_frames + kCodecDecodeWindowFrames); + const int64_t window_frames = emit_end - window_begin; + window_codes.resize(static_cast(window_frames * kCodecCodebooks)); + for (int64_t frame = 0; frame < window_frames; ++frame) { + const auto src = + codes.begin() + + static_cast((window_begin + frame) * kCodecCodebooks); + auto dst = + window_codes.begin() + static_cast(frame * kCodecCodebooks); + std::copy_n(src, static_cast(kCodecCodebooks), dst); + } + + const auto window = run_window( + window_codes, + window_frames, + kCodecDecodeWindowFrames + kCodecDecodeOverlapFrames); + const int64_t trim_frames = emitted_frames - window_begin; + const int64_t emit_frames = emit_end - emitted_frames; + const int64_t sample_begin = trim_frames * kCodecHopLength; + const int64_t sample_count = emit_frames * kCodecHopLength; + if (sample_begin < 0 || sample_count <= 0 || + sample_begin + sample_count > static_cast(window.values.size())) { + throw std::runtime_error("Higgs TTS codec decode window produced invalid length"); + } + out.values.insert(out.values.end(), + window.values.begin() + static_cast(sample_begin), + window.values.begin() + + static_cast(sample_begin + sample_count)); + emitted_frames = emit_end; + } + if (static_cast(out.values.size()) != out.samples) { + throw std::runtime_error("Higgs TTS codec chunked decode output length mismatch"); + } + return out; +} + +void HiggsCodecRuntime::release_encode_graph() { + encode_graph_.reset(); +} + +void HiggsCodecRuntime::release_runtime_graphs() { + release_encode_graph(); + decode_graph_.reset(); +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/generator.cpp b/src/models/higgs_audio_tts/generator.cpp new file mode 100644 index 00000000..2820067e --- /dev/null +++ b/src/models/higgs_audio_tts/generator.cpp @@ -0,0 +1,553 @@ +#include "engine/models/higgs_audio_tts/generator.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/sampling/torch_random.h" +#include "engine/models/higgs_audio_tts/codebooks.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr int64_t kInitialGeneratedCacheSteps = 128; +constexpr int64_t kMinimumCacheBucketSteps = 128; + +int64_t bucketed_initial_cache_steps(int64_t prompt_steps, int64_t max_tokens) { + const int64_t maximum = prompt_steps + max_tokens; + const int64_t required = prompt_steps + std::min(max_tokens, kInitialGeneratedCacheSteps); + int64_t bucket = kMinimumCacheBucketSteps; + while (bucket < required && bucket <= maximum / 2) { + bucket *= 2; + } + if (bucket < required) { + bucket = required; + } + return std::min(bucket, maximum); +} + +void validate_generation_options(const HiggsGenerationOptions & options) { + if (options.max_tokens <= 0) { + throw std::runtime_error("Higgs TTS max_tokens must be positive after default resolution"); + } + if (!(options.temperature > 0.0F)) { + throw std::runtime_error("Higgs TTS temperature must be positive"); + } + if (options.top_p.has_value() && !(*options.top_p > 0.0F)) { + throw std::runtime_error("Higgs TTS top_p must be positive"); + } + if (options.top_k.has_value() && *options.top_k < 0) { + throw std::runtime_error("Higgs TTS top_k must be non-negative"); + } + if (!(options.repetition_penalty > 0.0F) || !std::isfinite(options.repetition_penalty)) { + throw std::runtime_error("Higgs TTS repetition_penalty must be finite and positive"); + } +} + +size_t flat_index(int64_t frame, int64_t codebook, int64_t codebooks) { + return static_cast(frame * codebooks + codebook); +} + +struct HiggsPromptInput { + std::vector token_ids; + std::vector reference_positions; +}; + +struct HiggsPreparedPrompt { + HiggsPromptInput prompt; + HiggsARPrefillInput ar_input; + int64_t prefix_steps = 0; +}; + +HiggsPromptInput make_prompt_input(const HiggsPromptEncoding & prompt, + int64_t delayed_reference_frames, + const HiggsConfig & config) { + HiggsPromptInput input; + input.token_ids = prompt.token_ids; + input.reference_positions.reserve(static_cast(delayed_reference_frames)); + for (int64_t position = 0; position < static_cast(input.token_ids.size()); + ++position) { + if (input.token_ids[static_cast(position)] == config.audio_token_id) { + input.reference_positions.push_back(static_cast(position)); + input.token_ids[static_cast(position)] = 0; + } + } + if (static_cast(input.reference_positions.size()) != delayed_reference_frames) { + throw std::runtime_error("Higgs TTS prompt audio placeholder count does " + "not match delayed reference codes"); + } + return input; +} + +HiggsPreparedPrompt make_prepared_prompt(const HiggsPromptEncoding & prompt, + const std::vector & delayed_reference_codes, + int64_t delayed_reference_frames, + const HiggsConfig & config) { + HiggsPreparedPrompt prepared; + prepared.prompt = make_prompt_input(prompt, delayed_reference_frames, config); + const int64_t prompt_steps = static_cast(prepared.prompt.token_ids.size()); + prepared.ar_input.steps = prompt_steps; + prepared.ar_input.text_tokens = prepared.prompt.token_ids; + prepared.ar_input.fused_code_ids.assign( + static_cast(prompt_steps * config.audio.num_codebooks), 0); + prepared.ar_input.text_gate.assign(static_cast(prompt_steps), 0.0F); + prepared.ar_input.code_gate.assign(static_cast(prompt_steps), 0.0F); + size_t reference_row = 0; + for (int64_t position = 0; position < prompt_steps; ++position) { + if (reference_row < prepared.prompt.reference_positions.size() && + prepared.prompt.reference_positions[reference_row] == position) { + for (int64_t codebook = 0; codebook < config.audio.num_codebooks; ++codebook) { + const int32_t code = delayed_reference_codes[flat_index( + static_cast(reference_row), codebook, config.audio.num_codebooks)]; + if (code < 0 || code >= config.audio.vocab_size) { + throw std::runtime_error( + "Higgs TTS AR prefill codebook token is outside vocabulary"); + } + prepared.ar_input + .fused_code_ids[flat_index(position, codebook, config.audio.num_codebooks)] = + static_cast(code + codebook * config.audio.vocab_size); + } + prepared.ar_input.code_gate[static_cast(position)] = 1.0F; + ++reference_row; + } else { + const int32_t text_token = prepared.prompt.token_ids[static_cast(position)]; + if (text_token < 0 || text_token >= config.text.vocab_size) { + throw std::runtime_error("Higgs TTS AR prefill text token is outside vocabulary"); + } + prepared.ar_input.text_gate[static_cast(position)] = 1.0F; + } + } + if (reference_row != prepared.prompt.reference_positions.size()) { + throw std::runtime_error( + "Higgs TTS prompt prefill did not consume all reference code rows"); + } + if (!prepared.prompt.reference_positions.empty()) { + prepared.prefix_steps = + static_cast(prepared.prompt.reference_positions.back()) + 1; + } + return prepared; +} + +} // namespace + +HiggsGenerator::HiggsGenerator(std::shared_ptr assets, + std::shared_ptr ar, + std::shared_ptr codec, + size_t ar_decode_graph_arena_bytes) + : assets_([&]() { + if (assets == nullptr) { + throw std::runtime_error("Higgs TTS generator requires assets"); + } + return std::move(assets); + }()), + ar_([&]() { + if (ar == nullptr) { + throw std::runtime_error("Higgs TTS generator requires AR runtime"); + } + return std::move(ar); + }()), + codec_([&]() { + if (codec == nullptr) { + throw std::runtime_error("Higgs TTS generator requires codec runtime"); + } + return std::move(codec); + }()), + tokenizer_(assets_), + ar_decode_graph_arena_bytes_(ar_decode_graph_arena_bytes) { + if (ar_decode_graph_arena_bytes_ == 0) { + throw std::runtime_error("Higgs TTS generator graph arena bytes must be non-zero"); + } +} + +void HiggsGenerator::prepare(const HiggsGenerationRequest & request) { + const auto & config = assets_->config; + const bool has_reference = request.reference_frames > 0 || !request.reference_codes.empty(); + if (has_reference) { + if (request.reference_frames <= 0 || + request.reference_codebooks != config.audio.num_codebooks) { + throw std::runtime_error("Higgs TTS generation requires reference codes " + "shaped [frames, num_codebooks]"); + } + if (static_cast(request.reference_codes.size()) != + request.reference_frames * request.reference_codebooks) { + throw std::runtime_error("Higgs TTS generation reference code count mismatch"); + } + } else if (request.reference_codebooks != 0) { + throw std::runtime_error( + "Higgs TTS generation got reference codebooks without reference codes"); + } + validate_generation_options(request.options); + const int64_t delayed_reference_frames = + has_reference + ? higgs_delayed_frame_count(request.reference_frames, request.reference_codebooks) + : 0; + std::vector delayed_reference_codes; + if (has_reference) { + delayed_reference_codes = apply_higgs_delay_pattern( + request.reference_codes, request.reference_frames, request.reference_codebooks); + } + const HiggsPromptEncoding prompt = tokenizer_.encode_prompt({ + request.text, + request.reference_text, + delayed_reference_frames, + }); + const auto prepared = + make_prepared_prompt(prompt, delayed_reference_codes, delayed_reference_frames, config); + const int64_t prompt_steps = prepared.ar_input.steps; + if (prompt_steps + request.options.max_tokens > config.text.max_position_embeddings) { + throw std::runtime_error("Higgs TTS generation exceeds text model max_position_embeddings"); + } + if (has_reference && prepared.prefix_steps > 0) { + ReferencePrefixCache cache; + cache.reference_text = request.reference_text; + cache.reference_codes = request.reference_codes; + cache.reference_frames = request.reference_frames; + cache.reference_codebooks = request.reference_codebooks; + cache.delayed_reference_codes = std::move(delayed_reference_codes); + cache.delayed_reference_frames = delayed_reference_frames; + cache.prefix_steps = prepared.prefix_steps; + cache.prefix_tokens.assign(prepared.prompt.token_ids.begin(), + prepared.prompt.token_ids.begin() + + static_cast(prepared.prefix_steps)); + const bool same_reference = + reference_prefix_cache_.has_value() && + reference_prefix_cache_->reference_text == cache.reference_text && + reference_prefix_cache_->reference_codes == cache.reference_codes && + reference_prefix_cache_->reference_frames == cache.reference_frames && + reference_prefix_cache_->reference_codebooks == cache.reference_codebooks && + reference_prefix_cache_->prefix_steps == cache.prefix_steps && + reference_prefix_cache_->prefix_tokens == cache.prefix_tokens; + reference_prefix_cache_ = std::move(cache); + if (!same_reference) { + reference_kv_ready_ = false; + } + } else { + reference_prefix_cache_.reset(); + reference_kv_ready_ = false; + } +} + +HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & request) { + const auto & config = assets_->config; + const bool has_reference = request.reference_frames > 0 || !request.reference_codes.empty(); + if (has_reference) { + if (request.reference_frames <= 0 || + request.reference_codebooks != config.audio.num_codebooks) { + throw std::runtime_error("Higgs TTS generation requires reference codes " + "shaped [frames, num_codebooks]"); + } + if (static_cast(request.reference_codes.size()) != + request.reference_frames * request.reference_codebooks) { + throw std::runtime_error("Higgs TTS generation reference code count mismatch"); + } + } else if (request.reference_codebooks != 0) { + throw std::runtime_error( + "Higgs TTS generation got reference codebooks without reference codes"); + } + validate_generation_options(request.options); + engine::debug::trace_log_scalar("higgs_audio_tts.request.text", request.text); + engine::debug::trace_log_scalar("higgs_audio_tts.request.reference_text", request.reference_text); + engine::debug::trace_log_scalar("higgs_audio_tts.request.text_chars", request.text.size()); + engine::debug::trace_log_scalar("higgs_audio_tts.request.reference_text_chars", + request.reference_text.size()); + engine::debug::trace_log_scalar("higgs_audio_tts.request.max_tokens", + request.options.max_tokens); + engine::debug::trace_log_scalar("higgs_audio_tts.request.temperature", request.options.temperature); + engine::debug::trace_log_scalar("higgs_audio_tts.request.top_p", + request.options.top_p.has_value() ? + std::to_string(*request.options.top_p) : "none"); + engine::debug::trace_log_scalar("higgs_audio_tts.request.top_k", + request.options.top_k.has_value() ? + std::to_string(*request.options.top_k) : "none"); + engine::debug::trace_log_scalar("higgs_audio_tts.request.repetition_penalty", + request.options.repetition_penalty); + engine::debug::trace_log_scalar("higgs_audio_tts.request.has_seed", request.options.seed.has_value()); + engine::debug::trace_log_scalar("higgs_audio_tts.request.seed", + request.options.seed.has_value() ? + std::to_string(*request.options.seed) : "none"); + if (has_reference) { + engine::debug::trace_log_i32("higgs_audio_tts.request.reference_codes", + {request.reference_frames, request.reference_codebooks}, + request.reference_codes); + } + + std::vector delayed_reference_codes_storage; + const std::vector * delayed_reference_codes = &delayed_reference_codes_storage; + int64_t delayed_reference_frames = 0; + const ReferencePrefixCache * matching_reference_cache = nullptr; + if (has_reference) { + if (reference_prefix_cache_.has_value() && + reference_prefix_cache_->reference_text == request.reference_text && + reference_prefix_cache_->reference_frames == request.reference_frames && + reference_prefix_cache_->reference_codebooks == request.reference_codebooks && + reference_prefix_cache_->reference_codes == request.reference_codes) { + matching_reference_cache = &*reference_prefix_cache_; + delayed_reference_codes = &matching_reference_cache->delayed_reference_codes; + delayed_reference_frames = matching_reference_cache->delayed_reference_frames; + } else { + delayed_reference_codes_storage = apply_higgs_delay_pattern( + request.reference_codes, request.reference_frames, request.reference_codebooks); + delayed_reference_frames = + higgs_delayed_frame_count(request.reference_frames, request.reference_codebooks); + } + } + const HiggsPromptEncoding prompt = tokenizer_.encode_prompt({ + request.text, + request.reference_text, + delayed_reference_frames, + }); + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.text_tokens", prompt.text_ids.size()); + engine::debug::trace_log_i32("higgs_audio_tts.prompt.text_ids", + {static_cast(prompt.text_ids.size())}, + prompt.text_ids); + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.reference_text_tokens", + prompt.reference_text_ids.size()); + engine::debug::trace_log_i32("higgs_audio_tts.prompt.reference_text_ids", + {static_cast(prompt.reference_text_ids.size())}, + prompt.reference_text_ids); + const auto prepared = + make_prepared_prompt(prompt, *delayed_reference_codes, delayed_reference_frames, config); + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.tokens", prepared.prompt.token_ids.size()); + engine::debug::trace_log_i32("higgs_audio_tts.prompt.token_ids", + {static_cast(prepared.prompt.token_ids.size())}, + prepared.prompt.token_ids); + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.delayed_reference_rows", + delayed_reference_frames); + engine::debug::trace_log_i32("higgs_audio_tts.prompt.delayed_reference_codes", + {delayed_reference_frames, config.audio.num_codebooks}, + *delayed_reference_codes); + engine::debug::trace_log_i32("higgs_audio_tts.ar.prefill.text_tokens", + {prepared.ar_input.steps}, + prepared.ar_input.text_tokens); + engine::debug::trace_log_i32("higgs_audio_tts.ar.prefill.fused_code_ids", + {prepared.ar_input.steps, config.audio.num_codebooks}, + prepared.ar_input.fused_code_ids); + engine::debug::trace_log_f32("higgs_audio_tts.ar.prefill.text_gate", + {prepared.ar_input.steps}, + prepared.ar_input.text_gate); + engine::debug::trace_log_f32("higgs_audio_tts.ar.prefill.code_gate", + {prepared.ar_input.steps}, + prepared.ar_input.code_gate); + const int64_t prompt_steps = prepared.ar_input.steps; + if (prompt_steps + request.options.max_tokens > config.text.max_position_embeddings) { + throw std::runtime_error("Higgs TTS generation exceeds text model max_position_embeddings"); + } + + const auto prefill_start = Clock::now(); + const bool reference_cache_hit = + has_reference && prepared.prefix_steps > 0 && matching_reference_cache != nullptr && + matching_reference_cache->prefix_steps == prepared.prefix_steps && + static_cast(matching_reference_cache->prefix_tokens.size()) == prepared.prefix_steps && + std::equal(matching_reference_cache->prefix_tokens.begin(), + matching_reference_cache->prefix_tokens.end(), + prepared.prompt.token_ids.begin()); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_prefix_cache_hit", reference_cache_hit); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_prefix_steps", prepared.prefix_steps); + const int64_t max_cache_steps = prompt_steps + request.options.max_tokens; + const int64_t initial_cache_steps = bucketed_initial_cache_steps(prompt_steps, request.options.max_tokens); + const bool cache_rebuild = + ar_kv_cache_ == nullptr || !ar_kv_cache_->can_run(*ar_, initial_cache_steps) || + ar_kv_cache_->cache_steps() != initial_cache_steps; + if (cache_rebuild) { + decode_graph_.reset(); + ar_kv_cache_ = std::make_unique(ar_, initial_cache_steps); + reference_kv_ready_ = false; + } + const bool reference_kv_cache_hit = + reference_cache_hit && reference_kv_ready_ && + ar_kv_cache_->valid_steps() >= prepared.prefix_steps; + const int64_t prefill_start_step = reference_kv_cache_hit ? prepared.prefix_steps : 0; + if (reference_kv_cache_hit) { + ar_kv_cache_->retain_prefix(prefill_start_step); + } else { + ar_kv_cache_->reset(); + } + engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_kv_cache_hit", reference_kv_cache_hit); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.prefill_start_step", prefill_start_step); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.prefill_run_steps", prompt_steps - prefill_start_step); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.kv_cache_steps", initial_cache_steps); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.kv_cache_rebuild", cache_rebuild); + if (prefill_graph_ == nullptr || + !prefill_graph_->matches(*ar_, prompt_steps, prefill_start_step)) { + prefill_graph_.reset(); + prefill_graph_ = std::make_unique( + ar_, prompt_steps, prefill_start_step, ar_kv_cache_.get(), ar_decode_graph_arena_bytes_); + } + auto prefill_output = prefill_graph_->run(prepared.ar_input, prefill_start_step); + prefill_graph_.reset(); + reference_kv_ready_ = reference_cache_hit; + + if (decode_graph_ == nullptr || !decode_graph_->can_run(*ar_, ar_kv_cache_->cache_steps())) { + decode_graph_ = std::make_unique( + ar_, ar_kv_cache_->cache_steps(), *ar_kv_cache_, ar_decode_graph_arena_bytes_); + } + if (!prefill_output.wrote_cache) { + decode_graph_->import_prefill_state(prefill_output.kv_state); + } + HiggsARDecodeOutput prefill = std::move(prefill_output.output); + engine::debug::timing_log_scalar("higgs_audio_tts.generator.prefill_ms", + engine::debug::elapsed_ms(prefill_start, Clock::now())); + engine::debug::trace_log_f32("higgs_audio_tts.sampler.prefill_logits", + {config.audio.num_codebooks, config.audio.vocab_size}, + prefill.codebook_logits); + + HiggsCodebookSampler sampler(config.audio.num_codebooks, config.audio.vocab_size); + HiggsSamplerState state = sampler.make_state(); + HiggsSamplingOptions sampling; + sampling.temperature = request.options.temperature; + sampling.top_p = request.options.top_p; + sampling.top_k = request.options.top_k; + // Python accepts repetition_penalty on the public speech request, but the + // Higgs audio-codebook sampler does not consume it. + sampling.has_seed = request.options.seed.has_value(); + sampling.seed = request.options.seed.value_or(runtime::random_u64_seed()); + std::mt19937 fallback_rng(static_cast(sampling.seed)); + sampling.fallback_rng = &fallback_rng; + if (!sampling.has_seed && !cuda_sampling_policy_.has_value()) { + cuda_sampling_policy_ = engine::sampling::resolve_torch_cuda_sampling_policy( + ar_->backend_type(), + ar_->device(), + "higgs_audio_tts.cuda_sampling_policy", + "Higgs TTS", + engine::sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault); + } + if (cuda_sampling_policy_.has_value()) { + sampling.cuda_policy = *cuda_sampling_policy_; + } + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.temperature", sampling.temperature); + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.has_seed", sampling.has_seed); + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.seed", sampling.seed); + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.top_p", + sampling.top_p.has_value() ? std::to_string(*sampling.top_p) + : "none"); + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.top_k", + sampling.top_k.has_value() ? std::to_string(*sampling.top_k) + : "none"); + HiggsGenerationResult result; + result.delayed_codes.reserve( + static_cast(request.options.max_tokens * config.audio.num_codebooks)); + const auto & first_sampled = sampler.step(prefill.codebook_logits.data(), + static_cast(prefill.codebook_logits.size()), + state, + sampling); + engine::debug::trace_log_i32("higgs_audio_tts.sampler.first_output_codes", + {static_cast(first_sampled.size())}, + first_sampled); + result.delayed_codes.insert( + result.delayed_codes.end(), first_sampled.begin(), first_sampled.end()); + result.delayed_frames += 1; + + const auto decode_start = Clock::now(); + HiggsARDecodeOutput decoded; + decoded.codebook_logits.reserve( + static_cast(config.audio.num_codebooks * config.audio.vocab_size)); + bool logged_decode_step_timing = false; + double sampler_total_ms = 0.0; + HiggsARDecodeTiming decode_timing_total; + decode_graph_->begin_decode_run(); + while (!state.generation_done && result.delayed_frames < request.options.max_tokens) { + if (ar_kv_cache_->valid_steps() >= ar_kv_cache_->cache_steps()) { + decode_timing_total.add(decode_graph_->timing()); + const auto kv_state = ar_kv_cache_->export_state(); + const int64_t grown_cache_steps = + std::min(max_cache_steps, + std::max(ar_kv_cache_->cache_steps() * 2, ar_kv_cache_->valid_steps() + 1)); + if (grown_cache_steps <= ar_kv_cache_->cache_steps()) { + throw std::runtime_error("Higgs TTS AR cache cannot grow"); + } + decode_graph_.reset(); + ar_kv_cache_ = std::make_unique(ar_, grown_cache_steps); + ar_kv_cache_->import_state(kv_state); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.kv_cache_grown_steps", grown_cache_steps); + decode_graph_ = std::make_unique( + ar_, ar_kv_cache_->cache_steps(), *ar_kv_cache_, ar_decode_graph_arena_bytes_); + decode_graph_->begin_decode_run(); + } + HiggsARDecodeInput input; + input.use_last_codes = state.delay_count > 0; + input.last_codes = state.last_codes; + const auto step_start = Clock::now(); + decode_graph_->run_step_into(input, decoded, !logged_decode_step_timing); + if (!logged_decode_step_timing) { + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode.step0.ar_ms", + engine::debug::elapsed_ms(step_start, Clock::now())); + } + const auto sample_start = Clock::now(); + const auto & sampled = sampler.step(decoded.codebook_logits.data(), + static_cast(decoded.codebook_logits.size()), + state, + sampling); + sampler_total_ms += engine::debug::elapsed_ms(sample_start, Clock::now()); + if (!logged_decode_step_timing) { + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode.step0.sampler_ms", + sampler_total_ms); + logged_decode_step_timing = true; + } + if (!sampled.empty() && sampled.front() != kHiggsStopCode) { + result.delayed_codes.insert(result.delayed_codes.end(), sampled.begin(), sampled.end()); + result.delayed_frames += 1; + } + } + decode_timing_total.add(decode_graph_->timing()); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.steps", decode_timing_total.steps); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.input_upload_ms", decode_timing_total.input_upload_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.mask_upload_ms", decode_timing_total.mask_upload_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.graph.compute_ms", decode_timing_total.graph_compute_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.output_read_ms", decode_timing_total.output_read_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode.sampler_ms", sampler_total_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode_ms", + engine::debug::elapsed_ms(decode_start, Clock::now())); + if (!state.generation_done) { + throw std::runtime_error("Higgs TTS generation reached max_tokens before EOC"); + } + + result.raw_codes = reverse_higgs_delay_pattern( + result.delayed_codes, result.delayed_frames, config.audio.num_codebooks); + result.raw_frames = result.delayed_frames - (config.audio.num_codebooks - 1); + engine::debug::trace_log_i32("higgs_audio_tts.generator.delayed_codes", + {result.delayed_frames, config.audio.num_codebooks}, + result.delayed_codes); + const int64_t delayed_head_rows = std::min(result.delayed_frames, 8); + engine::debug::trace_log_i32("higgs_audio_tts.generator.delayed_codes_head8", + {delayed_head_rows, config.audio.num_codebooks}, + std::vector( + result.delayed_codes.begin(), + result.delayed_codes.begin() + + static_cast( + delayed_head_rows * config.audio.num_codebooks))); + const int32_t codec_vocab = static_cast(config.audio.vocab_size - 2); +#ifdef _OPENMP +#pragma omp parallel for if (static_cast(result.raw_codes.size()) > 1024) +#endif + for (int64_t index = 0; index < static_cast(result.raw_codes.size()); ++index) { + int32_t & code = result.raw_codes[static_cast(index)]; + if (code >= codec_vocab) { + code = 0; + } + } + engine::debug::trace_log_i32("higgs_audio_tts.generator.raw_codes_for_codec", + {result.raw_frames, config.audio.num_codebooks}, + result.raw_codes); + const auto codec_start = Clock::now(); + result.audio = + codec_->decode_codes(result.raw_codes, result.raw_frames, config.audio.num_codebooks); + engine::debug::trace_log_f32("higgs_audio_tts.codec.decode.output_audio", + {result.audio.samples}, + result.audio.values); + engine::debug::timing_log_scalar("higgs_audio_tts.generator.codec_decode_ms", + engine::debug::elapsed_ms(codec_start, Clock::now())); + codec_->release_runtime_graphs(); + return result; +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/loader.cpp b/src/models/higgs_audio_tts/loader.cpp new file mode 100644 index 00000000..d14b1db9 --- /dev/null +++ b/src/models/higgs_audio_tts/loader.cpp @@ -0,0 +1,157 @@ +#include "engine/models/higgs_audio_tts/loader.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/models/higgs_audio_tts/session.h" + +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +runtime::ModelMetadata metadata(const HiggsAssets &) { + runtime::ModelMetadata out; + out.family = "higgs_audio_tts"; + out.variant = "v3-4b"; + out.description = "Higgs Audio v3 TTS loaded from local SGLang-Omni compatible assets."; + out.config_candidates = { + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + }; + out.weight_candidates = {"model.safetensors.index.json", "model.gguf"}; + return out; +} + +runtime::CapabilitySet capabilities(const HiggsAssets &) { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + out.languages = {"Auto"}; + return out; +} + +runtime::ModelCliInterface cli(const HiggsAssets &) { + runtime::ModelCliInterface out; + out.request_options = { + {"max_tokens", "n", "Maximum generated AR tokens; default 2048, 0 uses the default."}, + {"temperature", "float", "AR sampling temperature."}, + {"top_k", "n", "AR top-k sampling limit."}, + {"top_p", "float", "AR nucleus sampling probability."}, + {"repetition_penalty", "float", "Accepted for Python API compatibility; Higgs audio sampling does not consume it."}, + {"seed", "n", "Torch RNG seed."}, + {"text_chunk_size", "n", "Long-form text chunk size; default 1024."}, + {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode."}, + }; + out.session_options = { + {"higgs_audio_tts.weight_type", "native|f32|f16|bf16|q8_0", "AR and codec weight storage type."}, + {"higgs_audio_tts.ar_weight_type", "native|f32|f16|bf16|q8_0", "Autoregressive decoder weight storage type."}, + {"higgs_audio_tts.codec_weight_type", "native|f32|f16|bf16|q8_0", "Codec weight storage type."}, + {"higgs_audio_tts.ar_weight_context_mb", "n", "AR weight context size."}, + {"higgs_audio_tts.codec_weight_context_mb", "n", "Codec weight context size."}, + {"higgs_audio_tts.ar_decode_graph_arena_mb", "n", "AR decode graph arena size."}, + {"higgs_audio_tts.codec_decode_graph_arena_mb", "n", "Codec decode graph arena size."}, + {"higgs_audio_tts.codec_encode_graph_arena_mb", "n", "Codec encode graph arena size."}, + {"higgs_audio_tts.reference_cache_slots", "n", "Encoded reference-audio cache slots; default 1."}, + }; + return out; +} + +class HiggsTTSLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "higgs_audio_tts"; + } + + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + return out; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + if (request.family_hint.has_value() && *request.family_hint != family()) { + return false; + } + try { + const auto package_spec = engine::assets::default_model_package_spec_path(family()); + (void) engine::assets::load_resource_bundle_from_package_spec(request.model_path, package_spec); + return true; + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto assets = load_higgs_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::assets::default_model_package_spec_path(family()); + inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Files); + inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Tensors); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_higgs_audio_tts_model(request.model_path); + } +}; + +} // namespace + +HiggsTTSLoadedModel::HiggsTTSLoadedModel( + 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 & HiggsTTSLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & HiggsTTSLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr HiggsTTSLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Higgs TTS only supports offline sessions"); + } + if (task.task != runtime::VoiceTaskKind::Tts) { + throw std::runtime_error("Higgs TTS only supports the Tts task"); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_higgs_audio_tts_model(const std::filesystem::path & model_path) { + auto assets = load_higgs_assets(model_path); + return std::make_unique( + metadata(*assets), + capabilities(*assets), + std::move(assets)); +} + +std::shared_ptr make_higgs_audio_tts_loader() { + return std::make_shared(); +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/sampler.cpp b/src/models/higgs_audio_tts/sampler.cpp new file mode 100644 index 00000000..9238fb7c --- /dev/null +++ b/src/models/higgs_audio_tts/sampler.cpp @@ -0,0 +1,480 @@ +#include "engine/models/higgs_audio_tts/sampler.h" + +#include "engine/framework/sampling/torch_random.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +constexpr float kGreedyTemperatureThreshold = 1.0e-5F; + +struct SamplerScratch { + std::vector & scores; + std::vector & probs; + std::vector & order; + std::vector & kept; +}; + +void require_cuda_sampling_policy(const HiggsCudaSamplingPolicy & policy) { + if (policy.multiprocessor_count <= 0 || policy.max_threads_per_multiprocessor <= 0) { + throw std::runtime_error("Higgs TTS stochastic sampler requires CUDA " + "sampling device properties"); + } +} + +uint32_t rotl32(uint32_t value, int shift) { + return static_cast((value << shift) | (value >> (32 - shift))); +} + +uint32_t fmix32(uint32_t value) { + value ^= value >> 16; + value *= 0x85EBCA6Bu; + value ^= value >> 13; + value *= 0xC2B2AE35u; + value ^= value >> 16; + return value; +} + +uint32_t murmur3_mix(uint32_t hash, uint32_t key) { + key *= 0xCC9E2D51u; + key = rotl32(key, 15); + key *= 0x1B873593u; + hash ^= key; + hash = rotl32(hash, 13); + hash = hash * 5u + 0xE6546B64u; + return hash; +} + +uint32_t sglang_murmur_hash32(uint64_t seed, uint32_t position, uint32_t column) { + uint32_t hash = 0; + hash = murmur3_mix(hash, static_cast(seed & 0xFFFFFFFFull)); + hash = murmur3_mix(hash, static_cast((seed >> 32) & 0xFFFFFFFFull)); + hash = murmur3_mix(hash, position); + hash = murmur3_mix(hash, column); + hash ^= 16u; + return fmix32(hash); +} + +int32_t argmax_row(const float * logits, int64_t vocab_size) { + int32_t best = 0; + float best_value = logits[0]; + for (int64_t i = 1; i < vocab_size; ++i) { + const float value = logits[i]; + if (value > best_value) { + best_value = value; + best = static_cast(i); + } + } + return best; +} + +float finite_max(const std::vector & scores, const std::vector & candidates) { + float max_score = -std::numeric_limits::infinity(); + if (candidates.empty()) { + for (const float score : scores) { + if (std::isfinite(score)) { + max_score = std::max(max_score, score); + } + } + } else { + for (const int64_t index : candidates) { + const float score = scores[static_cast(index)]; + if (std::isfinite(score)) { + max_score = std::max(max_score, score); + } + } + } + if (!std::isfinite(max_score)) { + throw std::runtime_error("Higgs TTS sampler kept no finite logits"); + } + return max_score; +} + +void scores_to_probs(SamplerScratch & scratch, int64_t vocab_size) { + const float max_score = finite_max(scratch.scores, scratch.kept); + double total = 0.0; + scratch.probs.assign(static_cast(vocab_size), 0.0F); + if (scratch.kept.empty()) { + for (int64_t i = 0; i < vocab_size; ++i) { + const float score = scratch.scores[static_cast(i)]; + if (std::isfinite(score)) { + const float value = + static_cast(std::exp(static_cast(score - max_score))); + scratch.probs[static_cast(i)] = value; + total += static_cast(value); + } + } + } else { + for (const int64_t index : scratch.kept) { + const float value = static_cast(std::exp( + static_cast(scratch.scores[static_cast(index)] - max_score))); + scratch.probs[static_cast(index)] = value; + total += static_cast(value); + } + } + if (!(total > 0.0) || !std::isfinite(total)) { + throw std::runtime_error("Higgs TTS sampler probability mass is invalid"); + } + const float inv_total = static_cast(1.0 / total); + for (float & prob : scratch.probs) { + prob *= inv_total; + } +} + +void renormalize_probs(SamplerScratch & scratch) { + double total = 0.0; + if (scratch.kept.empty()) { + for (const float prob : scratch.probs) { + total += static_cast(prob); + } + } else { + for (const int64_t index : scratch.kept) { + total += static_cast(scratch.probs[static_cast(index)]); + } + } + if (!(total > 0.0) || !std::isfinite(total)) { + throw std::runtime_error("Higgs TTS sampler probability mass is invalid"); + } + const float inv_total = static_cast(1.0 / total); + if (scratch.kept.empty()) { + for (float & prob : scratch.probs) { + prob *= inv_total; + } + } else { + for (const int64_t index : scratch.kept) { + scratch.probs[static_cast(index)] *= inv_total; + } + } +} + +void apply_top_k_to_probs(SamplerScratch & scratch, int64_t vocab_size, int64_t top_k) { + scratch.kept.clear(); + if (top_k <= 0 || top_k >= vocab_size) { + return; + } + scratch.order.resize(static_cast(vocab_size)); + std::iota(scratch.order.begin(), scratch.order.end(), 0); + auto kth = scratch.order.begin() + static_cast(top_k - 1); + std::nth_element( + scratch.order.begin(), kth, scratch.order.end(), [&](int64_t lhs, int64_t rhs) { + return scratch.probs[static_cast(lhs)] > + scratch.probs[static_cast(rhs)]; + }); + const float threshold = scratch.probs[static_cast(*kth)]; + scratch.kept.reserve(static_cast(top_k)); + for (int64_t i = 0; i < vocab_size; ++i) { + auto & prob = scratch.probs[static_cast(i)]; + if (prob < threshold) { + prob = 0.0F; + } else { + scratch.kept.push_back(i); + } + } + renormalize_probs(scratch); +} + +void apply_top_k_to_scores(SamplerScratch & scratch, int64_t vocab_size, int64_t top_k) { + scratch.kept.clear(); + if (top_k <= 0 || top_k >= vocab_size) { + return; + } + scratch.order.resize(static_cast(vocab_size)); + std::iota(scratch.order.begin(), scratch.order.end(), 0); + auto kth = scratch.order.begin() + static_cast(top_k - 1); + std::nth_element( + scratch.order.begin(), kth, scratch.order.end(), [&](int64_t lhs, int64_t rhs) { + return scratch.scores[static_cast(lhs)] > + scratch.scores[static_cast(rhs)]; + }); + const float threshold = scratch.scores[static_cast(*kth)]; + scratch.kept.reserve(static_cast(top_k)); + for (int64_t i = 0; i < vocab_size; ++i) { + if (scratch.scores[static_cast(i)] >= threshold) { + scratch.kept.push_back(i); + } + } +} + +void apply_top_p_to_probs(SamplerScratch & scratch, int64_t vocab_size, float top_p) { + if (!(top_p < 1.0F)) { + return; + } + if (!(top_p > 0.0F)) { + throw std::runtime_error("Higgs TTS sampler top_p must be positive"); + } + if (scratch.kept.empty()) { + scratch.kept.reserve(static_cast(vocab_size)); + for (int64_t i = 0; i < vocab_size; ++i) { + if (scratch.probs[static_cast(i)] > 0.0F) { + scratch.kept.push_back(i); + } + } + } + if (scratch.kept.empty()) { + throw std::runtime_error("Higgs TTS sampler top-p kept no probabilities"); + } + std::stable_sort(scratch.kept.begin(), scratch.kept.end(), [&](int64_t lhs, int64_t rhs) { + return scratch.probs[static_cast(lhs)] > scratch.probs[static_cast(rhs)]; + }); + + float cumulative = 0.0F; + float threshold = scratch.probs[static_cast(scratch.kept.front())]; + for (const int64_t index : scratch.kept) { + threshold = scratch.probs[static_cast(index)]; + cumulative += threshold; + if (cumulative >= top_p) { + break; + } + } + + size_t kept_count = 0; + for (const int64_t index : scratch.kept) { + auto & prob = scratch.probs[static_cast(index)]; + if (prob < threshold) { + prob = 0.0F; + } else { + scratch.kept[kept_count++] = index; + } + } + scratch.kept.resize(kept_count); + renormalize_probs(scratch); +} + +double sglang_gumbel_from_hash(uint32_t hashed) { + constexpr double kUint32Max = static_cast(std::numeric_limits::max()); + const double x = static_cast(hashed) / kUint32Max; + const double log_x = std::max(std::log(x), std::numeric_limits::lowest()); + return -std::log(-log_x); +} + +int32_t sample_seeded_sglang_gumbel(const std::vector & probs, + const std::vector & candidates, + uint64_t seed, + uint64_t position) { + double best_rank = -std::numeric_limits::infinity(); + int32_t best = -1; + const auto sample_one = [&](int64_t i) { + const float prob = probs[static_cast(i)]; + if (!(prob > 0.0F)) { + return; + } + const double logprob = std::log(static_cast(prob)); + const uint32_t hashed = sglang_murmur_hash32( + seed, static_cast(position & 0xFFFFFFFFull), static_cast(i)); + const double rank = logprob + sglang_gumbel_from_hash(hashed); + if (rank > best_rank) { + best_rank = rank; + best = static_cast(i); + } + }; + if (candidates.empty()) { + for (int64_t i = 0; i < static_cast(probs.size()); ++i) { + sample_one(i); + } + } else { + for (const int64_t index : candidates) { + sample_one(index); + } + } + if (best < 0) { + throw std::runtime_error("Higgs TTS sampler failed to select a codebook token"); + } + return best; +} + +int32_t sample_unseeded_torch_multinomial(const std::vector & probs, + const std::vector & candidates, + uint64_t seed, + uint64_t call_index, + const HiggsCudaSamplingPolicy & policy, + std::mt19937 & fallback_rng) { + const int64_t vocab_size = static_cast(probs.size()); + if (!policy.cuda_fast_path) { + std::vector weights; + if (candidates.empty()) { + weights.reserve(probs.size()); + for (float prob : probs) { + weights.push_back(static_cast(std::max(prob, 0.0F))); + } + std::discrete_distribution distribution(weights.begin(), weights.end()); + return static_cast(distribution(fallback_rng)); + } + weights.reserve(candidates.size()); + for (const int64_t index : candidates) { + weights.push_back(static_cast(std::max(probs[static_cast(index)], 0.0F))); + } + std::discrete_distribution distribution(weights.begin(), weights.end()); + return static_cast(candidates[distribution(fallback_rng)]); + } + require_cuda_sampling_policy(policy); + + double best_rank = -std::numeric_limits::infinity(); + int32_t best = -1; + const auto sample_one = [&](int64_t i) { + const float prob = probs[static_cast(i)]; + if (!(prob > 0.0F)) { + return; + } + const float exponential = engine::sampling::torch_cuda_tensor_iterator_exponential_element( + seed, + static_cast(vocab_size), + static_cast(i), + call_index, + policy.multiprocessor_count, + policy.max_threads_per_multiprocessor); + const double rank = static_cast(prob) / static_cast(exponential); + if (rank > best_rank) { + best_rank = rank; + best = static_cast(i); + } + }; + if (candidates.empty()) { + for (int64_t i = 0; i < vocab_size; ++i) { + sample_one(i); + } + } else { + for (const int64_t index : candidates) { + sample_one(index); + } + } + if (best < 0) { + throw std::runtime_error("Higgs TTS sampler failed to select a codebook token"); + } + return best; +} + +int32_t sample_codebook_row(const float * logits, + int64_t vocab_size, + const HiggsSamplingOptions & options, + uint64_t call_index, + SamplerScratch & scratch) { + if (logits == nullptr || vocab_size <= 0) { + throw std::runtime_error("Higgs TTS sampler requires logits"); + } + if (options.temperature <= kGreedyTemperatureThreshold || + (options.top_k.has_value() && *options.top_k == 1)) { + return argmax_row(logits, vocab_size); + } + if (!(options.temperature > 0.0F) || !std::isfinite(options.temperature)) { + throw std::runtime_error("Higgs TTS sampler temperature must be finite and positive"); + } + + scratch.kept.clear(); + scratch.scores.resize(static_cast(vocab_size)); + for (int64_t i = 0; i < vocab_size; ++i) { + scratch.scores[static_cast(i)] = logits[i] / options.temperature; + } + bool top_k_applied_to_scores = false; + if (options.top_k.has_value()) { + if (*options.top_k < 0) { + throw std::runtime_error("Higgs TTS sampler top_k must be non-negative"); + } + const int64_t top_k = std::min(*options.top_k, vocab_size); + if (top_k > 0 && top_k < vocab_size) { + apply_top_k_to_scores(scratch, vocab_size, top_k); + top_k_applied_to_scores = true; + } + } + scores_to_probs(scratch, vocab_size); + if (options.top_k.has_value() && !top_k_applied_to_scores) { + apply_top_k_to_probs(scratch, vocab_size, std::min(*options.top_k, vocab_size)); + } + if (options.top_p.has_value()) { + apply_top_p_to_probs(scratch, vocab_size, *options.top_p); + } + if (options.has_seed) { + return sample_seeded_sglang_gumbel( + scratch.probs, scratch.kept, options.seed & 0x7FFFFFFFull, call_index); + } + if (options.fallback_rng == nullptr) { + throw std::runtime_error("Higgs TTS sampler fallback RNG is missing"); + } + return sample_unseeded_torch_multinomial( + scratch.probs, scratch.kept, options.seed, call_index, options.cuda_policy, *options.fallback_rng); +} + +} // namespace + +HiggsCodebookSampler::HiggsCodebookSampler(int64_t num_codebooks, int64_t codebook_vocab_size) + : num_codebooks_(num_codebooks), codebook_vocab_size_(codebook_vocab_size) { + if (num_codebooks_ <= 0 || codebook_vocab_size_ <= 0) { + throw std::runtime_error("Higgs TTS sampler requires positive codebook dimensions"); + } + scratch_scores_.reserve(static_cast(codebook_vocab_size_)); + scratch_probs_.reserve(static_cast(codebook_vocab_size_)); + scratch_order_.reserve(static_cast(codebook_vocab_size_)); + scratch_kept_.reserve(static_cast(codebook_vocab_size_)); + scratch_codes_.reserve(static_cast(num_codebooks_)); +} + +HiggsSamplerState HiggsCodebookSampler::make_state() const { + HiggsSamplerState state; + state.num_codebooks = num_codebooks_; + state.last_codes.assign(static_cast(num_codebooks_), 0); + return state; +} + +const std::vector & HiggsCodebookSampler::step(const float * logits, + int64_t logits_count, + HiggsSamplerState & state, + HiggsSamplingOptions & options) { + if (state.num_codebooks != num_codebooks_) { + throw std::runtime_error("Higgs TTS sampler state codebook count mismatch"); + } + if (logits_count != num_codebooks_ * codebook_vocab_size_) { + throw std::runtime_error("Higgs TTS sampler logits shape mismatch"); + } + + if (state.generation_done) { + scratch_codes_.assign(static_cast(num_codebooks_), kHiggsStopCode); + return scratch_codes_; + } + + scratch_codes_.assign(static_cast(num_codebooks_), 0); + SamplerScratch scratch{scratch_scores_, scratch_probs_, scratch_order_, scratch_kept_}; + for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { + const float * row = logits + static_cast(codebook * codebook_vocab_size_); + scratch_codes_[static_cast(codebook)] = + sample_codebook_row(row, + codebook_vocab_size_, + options, + static_cast(state.step_count * num_codebooks_ + codebook), + scratch); + } + + if (state.delay_count < num_codebooks_) { + const int64_t next_codebook = state.delay_count + 1; + if (next_codebook < num_codebooks_) { + for (int64_t codebook = next_codebook; codebook < num_codebooks_; ++codebook) { + scratch_codes_[static_cast(codebook)] = kHiggsBocId; + } + } + state.delay_count += 1; + } else if (state.eoc_countdown.has_value()) { + *state.eoc_countdown -= 1; + if (*state.eoc_countdown <= 0) { + state.generation_done = true; + } + } else if (scratch_codes_.front() == kHiggsEocId) { + if (num_codebooks_ <= 2) { + state.generation_done = true; + } else { + state.eoc_countdown = num_codebooks_ - 2; + } + } + + state.step_count += 1; + if (!state.generation_done) { + state.last_codes = scratch_codes_; + } + return scratch_codes_; +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/session.cpp b/src/models/higgs_audio_tts/session.cpp new file mode 100644 index 00000000..0a5d4779 --- /dev/null +++ b/src/models/higgs_audio_tts/session.cpp @@ -0,0 +1,323 @@ +#include "engine/models/higgs_audio_tts/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/text/chunking.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr int64_t kDefaultTextChunkSize = 1024; +constexpr int64_t kDefaultReferenceCacheSlots = 1; + +void validate_matmul_weight_storage(assets::TensorStorageType storage_type, const char * option_name) { + if (storage_type == assets::TensorStorageType::Native || + storage_type == assets::TensorStorageType::F32 || + storage_type == assets::TensorStorageType::F16 || + storage_type == assets::TensorStorageType::BF16 || + storage_type == assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " currently supports only native, f32, f16, bf16, and q8_0"); +} + +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; +} + +std::size_t resolve_reference_cache_slots(const runtime::SessionOptions & options) { + const int64_t slots = runtime::parse_i64_option( + options.options, + {"higgs_audio_tts.reference_cache_slots", "reference_cache_slots"}) + .value_or(kDefaultReferenceCacheSlots); + if (slots < 0) { + throw std::runtime_error("higgs_audio_tts.reference_cache_slots must be non-negative"); + } + if (static_cast(slots) > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("higgs_audio_tts.reference_cache_slots is too large"); + } + return static_cast(slots); +} + +const runtime::AudioBuffer * find_reference_audio(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; + } + if (request.audio_input.has_value()) { + return &*request.audio_input; + } + return nullptr; +} + +HiggsGenerationOptions generation_options_from_request( + const runtime::TaskRequest & request, + const HiggsConfig & config) { + HiggsGenerationOptions options; + options.max_tokens = 2048; + options.temperature = 0.8F; + options.top_p = 0.8F; + options.top_k = 30; + options.repetition_penalty = 1.1F; + if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { + if (*value < 0) { + throw std::runtime_error("Higgs TTS max_tokens must be non-negative"); + } + if (*value > 0) { + options.max_tokens = *value; + } + } + if (const auto value = runtime::parse_float_option(request.options, {"temperature"})) { + options.temperature = *value; + } + if (const auto value = runtime::parse_float_option(request.options, {"top_p"})) { + options.top_p = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"top_k"})) { + options.top_k = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"repetition_penalty"})) { + options.repetition_penalty = *value; + } + if (const auto value = runtime::parse_u64_option(request.options, {"seed"})) { + options.seed = *value; + } + if (options.max_tokens > config.text.max_position_embeddings) { + throw std::runtime_error("Higgs TTS max_tokens exceeds model max_position_embeddings"); + } + if (!(options.repetition_penalty > 0.0F) || !std::isfinite(options.repetition_penalty)) { + throw std::runtime_error("Higgs TTS repetition_penalty must be finite and positive"); + } + return options; +} + +} // namespace + +HiggsTTSSession::HiggsTTSSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : RuntimeSessionBase(options), + task_(task), + assets_(std::move(assets)), + reference_cache_(resolve_reference_cache_slots(this->options())) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS session requires assets"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Higgs TTS currently supports offline sessions"); + } + if (task_.task != runtime::VoiceTaskKind::Tts) { + throw std::runtime_error("Higgs TTS only supports the Tts task"); + } + ar_weight_context_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_audio_tts.ar_weight_context_mb"}, ar_weight_context_bytes_); + codec_weight_context_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_audio_tts.codec_weight_context_mb"}, codec_weight_context_bytes_); + ar_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_audio_tts.ar_decode_graph_arena_mb"}, ar_decode_graph_arena_bytes_); + codec_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_audio_tts.codec_decode_graph_arena_mb"}, codec_decode_graph_arena_bytes_); + codec_encode_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_audio_tts.codec_encode_graph_arena_mb"}, codec_encode_graph_arena_bytes_); + + if (const auto it = options.options.find("higgs_audio_tts.weight_type"); it != options.options.end()) { + const auto storage_type = assets::parse_tensor_storage_type(it->second); + validate_matmul_weight_storage(storage_type, "higgs_audio_tts.weight_type"); + ar_weight_storage_type_ = storage_type; + codec_weight_storage_type_ = storage_type; + } + if (const auto it = options.options.find("higgs_audio_tts.ar_weight_type"); it != options.options.end()) { + ar_weight_storage_type_ = assets::parse_tensor_storage_type(it->second); + validate_matmul_weight_storage(ar_weight_storage_type_, "higgs_audio_tts.ar_weight_type"); + } + if (const auto it = options.options.find("higgs_audio_tts.codec_weight_type"); it != options.options.end()) { + codec_weight_storage_type_ = assets::parse_tensor_storage_type(it->second); + validate_matmul_weight_storage(codec_weight_storage_type_, "higgs_audio_tts.codec_weight_type"); + } + for (const auto & [key, _] : options.options) { + if (key.rfind("higgs_audio_tts.", 0) == 0 && + key != "higgs_audio_tts.ar_weight_context_mb" && + key != "higgs_audio_tts.codec_weight_context_mb" && + key != "higgs_audio_tts.ar_decode_graph_arena_mb" && + key != "higgs_audio_tts.codec_decode_graph_arena_mb" && + key != "higgs_audio_tts.codec_encode_graph_arena_mb" && + key != "higgs_audio_tts.reference_cache_slots" && + key != "higgs_audio_tts.weight_type" && + key != "higgs_audio_tts.ar_weight_type" && + key != "higgs_audio_tts.codec_weight_type") { + throw std::runtime_error("unknown Higgs TTS session option: " + key); + } + } + + ar_ = std::make_shared( + assets_, + execution_context(), + ar_weight_context_bytes_, + ar_weight_storage_type_); + codec_ = std::make_shared( + assets_, + execution_context(), + codec_weight_context_bytes_, + codec_decode_graph_arena_bytes_, + codec_encode_graph_arena_bytes_, + codec_weight_storage_type_); + generator_ = std::make_unique( + assets_, + ar_, + codec_, + ar_decode_graph_arena_bytes_); +} + +std::string HiggsTTSSession::family() const { + return "higgs_audio_tts"; +} + +runtime::VoiceTaskKind HiggsTTSSession::task_kind() const { + return task_.task; +} + +runtime::RunMode HiggsTTSSession::run_mode() const { + return task_.mode; +} + +void HiggsTTSSession::prepare(const runtime::SessionPreparationRequest & request) { + if (request.text.has_value()) { + runtime::TaskRequest task_request; + task_request.text_input = request.text; + task_request.voice = request.voice; + task_request.options = request.options; + const auto generation_request = make_generation_request(task_request); + generator_->prepare(generation_request); + } + mark_prepared(); +} + +runtime::TaskResult HiggsTTSSession::run(const runtime::TaskRequest & request) { + require_prepared("Higgs TTS run"); + const auto wall_start = Clock::now(); + const int64_t text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + const auto chunk_requests = runtime::chunk_text_request(request, text_chunk_size, text_chunk_mode); + const std::string reference_text = runtime::find_option(request.options, {"reference_text"}).value_or(""); + const auto * reference_audio = find_reference_audio(request); + const HiggsCodecEncodeOutput * reference_codes = + reference_audio != nullptr ? &resolve_reference_codes(*reference_audio, reference_text) : nullptr; + debug::trace_log_scalar("higgs_audio_tts.text_chunk_size", text_chunk_size); + debug::trace_log_scalar("higgs_audio_tts.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + debug::trace_log_scalar("higgs_audio_tts.text_chunk_count", static_cast(chunk_requests.size())); + + runtime::AudioBuffer merged_audio; + for (const auto & chunk_request : chunk_requests) { + const auto generation_request = make_generation_request(chunk_request, reference_codes); + auto result = generator_->generate(generation_request); + runtime::append_audio_buffer(merged_audio, runtime::AudioBuffer{ + result.audio.sample_rate, + result.audio.channels, + std::move(result.audio.values), + }); + } + + runtime::TaskResult out; + out.audio_output = std::move(merged_audio); + debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); + return out; +} + +const HiggsCodecEncodeOutput & HiggsTTSSession::resolve_reference_codes( + const runtime::AudioBuffer & audio, + const std::string & reference_text) { + const uint64_t sample_count = static_cast(audio.samples.size()); + const uint64_t sample_hash = hash_audio_samples(audio); + ReferenceCacheKey key; + key.reference_text = reference_text; + key.sample_rate = audio.sample_rate; + key.channels = audio.channels; + key.sample_count = sample_count; + key.sample_hash = sample_hash; + debug::trace_log_scalar("higgs_audio_tts.reference_audio.sample_rate", audio.sample_rate); + debug::trace_log_scalar("higgs_audio_tts.reference_audio.channels", audio.channels); + debug::trace_log_f32("higgs_audio_tts.reference_audio.samples", + {static_cast(audio.samples.size())}, + audio.samples); + debug::trace_log_scalar("higgs_audio_tts.reference_cache.capacity", static_cast(reference_cache_.capacity())); + debug::trace_log_scalar("higgs_audio_tts.reference_cache.size", static_cast(reference_cache_.size())); + if (const auto * cached = reference_cache_.find(key)) { + debug::trace_log_scalar("higgs_audio_tts.reference_cache.hit", 1); + return cached->codes; + } + debug::trace_log_scalar("higgs_audio_tts.reference_cache.hit", 0); + + const auto encode_start = Clock::now(); + ReferenceCacheEntry entry; + entry.codes = codec_->encode_reference(audio); + codec_->release_encode_graph(); + debug::trace_log_scalar("higgs_audio_tts.reference_codes.frames", entry.codes.frames); + debug::trace_log_scalar("higgs_audio_tts.reference_codes.codebooks", entry.codes.codebooks); + debug::trace_log_i32("higgs_audio_tts.reference_codes.values", + {entry.codes.frames, entry.codes.codebooks}, + entry.codes.codes); + if (reference_cache_.capacity() == 0) { + uncached_reference_ = std::move(entry); + debug::timing_log_scalar("higgs_audio_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); + return uncached_reference_->codes; + } + reference_cache_.put(key, std::move(entry)); + debug::timing_log_scalar("higgs_audio_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); + return reference_cache_.find(key)->codes; +} + +HiggsGenerationRequest HiggsTTSSession::make_generation_request( + const runtime::TaskRequest & request, + const HiggsCodecEncodeOutput * resolved_reference_codes) { + if (!request.text_input.has_value()) { + throw std::runtime_error("Higgs TTS requires text input"); + } + const std::string reference_text = runtime::find_option(request.options, {"reference_text"}).value_or(""); + + HiggsGenerationRequest out; + out.text = request.text_input->text; + out.reference_text = reference_text; + out.options = generation_options_from_request(request, assets_->config); + if (resolved_reference_codes != nullptr) { + out.reference_codes = resolved_reference_codes->codes; + out.reference_frames = resolved_reference_codes->frames; + out.reference_codebooks = resolved_reference_codes->codebooks; + } else if (const auto * reference_audio = find_reference_audio(request)) { + const auto & reference_codes = resolve_reference_codes(*reference_audio, reference_text); + out.reference_codes = reference_codes.codes; + out.reference_frames = reference_codes.frames; + out.reference_codebooks = reference_codes.codebooks; + } + return out; +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_audio_tts/tokenizer_text.cpp b/src/models/higgs_audio_tts/tokenizer_text.cpp new file mode 100644 index 00000000..9cfd9739 --- /dev/null +++ b/src/models/higgs_audio_tts/tokenizer_text.cpp @@ -0,0 +1,94 @@ +#include "engine/models/higgs_audio_tts/tokenizer_text.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include +#include +#include + +namespace engine::models::higgs_audio_tts { +namespace { + +int32_t require_token_id(const engine::tokenizers::LlamaBpeTokenizer & tokenizer, const std::string & token) { + const auto id = tokenizer.find_token_id(token); + if (!id.has_value()) { + throw std::runtime_error("Higgs TTS tokenizer missing required token: " + token); + } + return *id; +} + +} // namespace + +struct HiggsTextTokenizer::Impl { + explicit Impl(std::shared_ptr input_assets) + : assets(std::move(input_assets)), + tokenizer(engine::tokenizers::LlamaBpeTokenizerSpec{ + {}, + {}, + assets->resources.require_file("tokenizer_config"), + assets->resources.require_file("tokenizer_json"), + engine::tokenizers::LlamaBpePreTokenizer::Qwen2, + }), + tts_id(require_token_id(tokenizer, "<|tts|>")), + ref_audio_id(require_token_id(tokenizer, "<|ref_audio|>")), + ref_text_id(require_token_id(tokenizer, "<|ref_text|>")), + text_id(require_token_id(tokenizer, "<|text|>")), + audio_id(require_token_id(tokenizer, "<|audio|>")), + audio_placeholder_id(static_cast(assets->config.audio_token_id)) {} + + std::shared_ptr assets; + engine::tokenizers::LlamaBpeTokenizer tokenizer; + int32_t tts_id = 0; + int32_t ref_audio_id = 0; + int32_t ref_text_id = 0; + int32_t text_id = 0; + int32_t audio_id = 0; + int32_t audio_placeholder_id = -100; +}; + +HiggsTextTokenizer::HiggsTextTokenizer(std::shared_ptr assets) + : impl_([&]() { + if (assets == nullptr) { + throw std::runtime_error("Higgs TTS text tokenizer requires assets"); + } + return std::make_shared(std::move(assets)); + }()) {} + +std::vector HiggsTextTokenizer::encode(const std::string & text) const { + return impl_->tokenizer.encode(text, true); +} + +HiggsPromptEncoding HiggsTextTokenizer::encode_prompt(const HiggsPromptRequest & request) const { + if (request.delayed_reference_tokens < 0) { + throw std::runtime_error("Higgs TTS delayed_reference_tokens must be non-negative"); + } + + HiggsPromptEncoding encoding; + encoding.text_ids = encode(request.text); + if (!request.reference_text.empty() && request.delayed_reference_tokens > 0) { + encoding.reference_text_ids = encode(request.reference_text); + } + + encoding.token_ids.push_back(impl_->tts_id); + if (!encoding.reference_text_ids.empty()) { + encoding.token_ids.push_back(impl_->ref_text_id); + encoding.token_ids.insert( + encoding.token_ids.end(), + encoding.reference_text_ids.begin(), + encoding.reference_text_ids.end()); + } + if (request.delayed_reference_tokens > 0) { + encoding.token_ids.push_back(impl_->ref_audio_id); + encoding.token_ids.insert( + encoding.token_ids.end(), + static_cast(request.delayed_reference_tokens), + impl_->audio_placeholder_id); + } + encoding.token_ids.push_back(impl_->text_id); + encoding.token_ids.insert(encoding.token_ids.end(), encoding.text_ids.begin(), encoding.text_ids.end()); + encoding.token_ids.push_back(impl_->audio_id); + return encoding; +} + +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/pocket_tts/flow_lm.cpp b/src/models/pocket_tts/flow_lm.cpp index e5b5e5f1..21398831 100644 --- a/src/models/pocket_tts/flow_lm.cpp +++ b/src/models/pocket_tts/flow_lm.cpp @@ -14,6 +14,8 @@ namespace engine::models::pocket_tts { namespace { +constexpr size_t kPromptGraphNodeCapacity = 262144; + modules::TransformerEncoderBlockWeights make_transformer_layer_weights( core::ModuleBuildContext & ctx, const models::pocket_tts::PocketTTSBackendWeights & weights, @@ -383,7 +385,8 @@ class FlowLMStepRuntime { core::write_tensor_f32(attention_mask_, attention_mask_buffer_); core::set_backend_threads(backend_, threads_); if (prompt_steps_ > 0) { - prompt_graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false); + // Long or dense prompts add per-step KV transfer nodes beyond the default graph capacity. + prompt_graph_ = ggml_new_graph_custom(ggml_ctx_, kPromptGraphNodeCapacity, false); ggml_build_forward_expand(prompt_graph_, prompt_output_.tensor); for (size_t step = 0; step < prompt_step_key_sources_.size(); ++step) { for (size_t layer = 0; layer < prompt_step_key_sources_[step].size(); ++layer) { diff --git a/src/models/qwen3_tts/talker.cpp b/src/models/qwen3_tts/talker.cpp index fb352f08..00ed0779 100644 --- a/src/models/qwen3_tts/talker.cpp +++ b/src/models/qwen3_tts/talker.cpp @@ -1627,6 +1627,7 @@ class Qwen3TalkerStepRuntime::Impl { !talker_prefill_equal(*cached_prompt_prefill_, request)) { cached_prompt_state_ = build_prompt_state(request, weights_->assets().config, weights_->weights()); cached_prompt_prefill_ = request; + cached_prefill_output_.reset(); } const auto & state = *cached_prompt_state_; const auto prompt_state_end = Clock::now(); @@ -1636,9 +1637,13 @@ class Qwen3TalkerStepRuntime::Impl { throw std::runtime_error("Qwen3 talker prompt exceeds step runtime capacity"); } const auto prefill_start = Clock::now(); - auto prefill_output = run_prefill_embeddings_with_state(state.prompt, prompt_steps); + const bool prefill_cache_hit = cached_prefill_output_.has_value(); + if (!prefill_cache_hit) { + cached_prefill_output_ = run_prefill_embeddings_with_state(state.prompt, prompt_steps); + } + const auto & prefill_output = *cached_prefill_output_; const auto prefill_end = Clock::now(); - auto current = std::move(prefill_output.result); + auto current = prefill_output.result; double code_predictor_build_ms = 0.0; if (code_predictor_graph_ == nullptr) { const auto build_start = Clock::now(); @@ -1648,12 +1653,14 @@ class Qwen3TalkerStepRuntime::Impl { double cached_step_build_ms = 0.0; double import_prefill_state_ms = 0.0; int64_t cached_step_capacity = 0; - auto cached_state = std::move(prefill_output.state); + runtime::TransformerKVState exported_cached_state; + const runtime::TransformerKVState * cached_state = &prefill_output.state; bool cached_graph_has_state = false; auto ensure_cached_step_capacity = [&](int64_t required_capacity) { if (cached_step_graph_ != nullptr && cached_graph_has_state && !cached_step_graph_->can_run(*weights_, required_capacity)) { - cached_state = cached_step_graph_->export_state(); + exported_cached_state = cached_step_graph_->export_state(); + cached_state = &exported_cached_state; cached_graph_has_state = false; } if (cached_step_graph_ == nullptr || !cached_step_graph_->can_run(*weights_, required_capacity)) { @@ -1675,7 +1682,7 @@ class Qwen3TalkerStepRuntime::Impl { } if (!cached_graph_has_state) { const auto import_start = Clock::now(); - cached_step_graph_->import_prefill_state(cached_state); + cached_step_graph_->import_prefill_state(*cached_state); import_prefill_state_ms += engine::debug::elapsed_ms(import_start, Clock::now()); cached_graph_has_state = true; } @@ -1760,6 +1767,7 @@ class Qwen3TalkerStepRuntime::Impl { out.decoder_input_codes.frames += out.generated_codes.frames; debug::timing_log_scalar("qwen3_tts.talker.prompt_state_ms", engine::debug::elapsed_ms(prompt_state_start, prompt_state_end)); debug::timing_log_scalar("qwen3_tts.talker.prefill_ms", engine::debug::elapsed_ms(prefill_start, prefill_end)); + debug::timing_log_scalar("qwen3_tts.talker.prefill_cache.hit", prefill_cache_hit); debug::timing_log_scalar("qwen3_tts.talker.code_predictor_build_ms", code_predictor_build_ms); debug::timing_log_scalar("qwen3_tts.talker.cached_step_build_ms", cached_step_build_ms); debug::timing_log_scalar("qwen3_tts.talker.import_prefill_state_ms", import_prefill_state_ms); @@ -1812,6 +1820,7 @@ class Qwen3TalkerStepRuntime::Impl { std::unique_ptr code_predictor_graph_; std::optional cached_prompt_prefill_; std::optional cached_prompt_state_; + std::optional cached_prefill_output_; }; Qwen3TalkerStepRuntime::Qwen3TalkerStepRuntime(std::unique_ptr impl) : impl_(std::move(impl)) { diff --git a/tests/higgs_audio_tts/.gitignore b/tests/higgs_audio_tts/.gitignore new file mode 100644 index 00000000..b45f5596 --- /dev/null +++ b/tests/higgs_audio_tts/.gitignore @@ -0,0 +1,2 @@ +results/ +__pycache__/ diff --git a/tests/higgs_audio_tts/README.md b/tests/higgs_audio_tts/README.md new file mode 100644 index 00000000..dbf8e476 --- /dev/null +++ b/tests/higgs_audio_tts/README.md @@ -0,0 +1,38 @@ +# Higgs Audio v3 TTS tests + +The focused framework unit test validates that packed QKV/gate-up projections +match the separate projections, suffix causal masks are correct, F16 KV writes +preserve their values, and the decode graph exposes the intended CUDA paths: +grouped FlashAttention, packed SwiGLU, direct KV updates, and the +`ROPE -> VIEW -> SET_ROWS` fusion pattern. + +```powershell +cmake --build build/windows-cuda-release --config Release --target qwen_decoder_packed_projection_test higgs_audio_tts_warm_bench -j 8 +ctest --test-dir build/windows-cuda-release -C Release -R qwen_decoder_packed_projection_test --output-on-failure +``` + +Run the fixed-seed, five-request CUDA benchmark and save every generated WAV: + +```powershell +tests/higgs_audio_tts/run_cuda_performance.ps1 ` + -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` + -Label candidate +``` + +Compare a candidate run with a prior result directory request by request: + +```powershell +tests/higgs_audio_tts/run_cuda_performance.ps1 ` + -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` + -Label candidate ` + -Baseline tests/higgs_audio_tts/results/baseline +``` + +The comparison reports frame counts, wall time, RTF, speedup, waveform cosine, +and 80-band log-mel cosine. Result WAVs, logs, and JSON reports are written below +`tests/higgs_audio_tts/results/`, which is intentionally ignored by Git. +The comparison helper requires Python 3 with NumPy. + +Add `-RequireSameFrames` when comparing paths that are expected to be +deterministic and numerically identical. Sampled or mixed-precision paths still +report their frame drift and similarity metrics without hiding the results. diff --git a/tests/higgs_audio_tts/compare_warmbench_results.py b/tests/higgs_audio_tts/compare_warmbench_results.py new file mode 100644 index 00000000..1ebc8a43 --- /dev/null +++ b/tests/higgs_audio_tts/compare_warmbench_results.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +"""Compare Higgs warmbench runs request by request. + +Each result directory is expected to contain timing.log and audio/audio_N.wav, +as emitted by higgs_audio_tts_warm_bench. The report includes exact frame counts, +wall time, RTF, speedup, waveform cosine, and log-mel cosine per request. +""" + +from __future__ import annotations + +import argparse +import json +import math +import wave +from pathlib import Path + +import numpy as np + + +def read_wav(path: Path) -> tuple[int, np.ndarray]: + with wave.open(str(path), "rb") as wav: + if wav.getsampwidth() != 2: + raise ValueError(f"{path}: expected PCM16 WAV") + channels = wav.getnchannels() + sample_rate = wav.getframerate() + samples = np.frombuffer(wav.readframes(wav.getnframes()), dtype=" 1: + samples = samples.reshape(-1, channels).mean(axis=1) + return sample_rate, samples / 32768.0 + + +def cosine(a: np.ndarray, b: np.ndarray) -> float: + count = min(a.size, b.size) + if count == 0: + return math.nan + a = a[:count].astype(np.float64, copy=False) + b = b[:count].astype(np.float64, copy=False) + denom = np.linalg.norm(a) * np.linalg.norm(b) + return float(np.dot(a, b) / denom) if denom > 0.0 else math.nan + + +def hz_to_mel(hz: np.ndarray | float) -> np.ndarray | float: + return 2595.0 * np.log10(1.0 + np.asarray(hz) / 700.0) + + +def mel_to_hz(mel: np.ndarray | float) -> np.ndarray | float: + return 700.0 * (np.power(10.0, np.asarray(mel) / 2595.0) - 1.0) + + +def log_mel(samples: np.ndarray, sample_rate: int, n_fft: int = 1024, hop: int = 256, bands: int = 80) -> np.ndarray: + if samples.size < n_fft: + samples = np.pad(samples, (0, n_fft - samples.size)) + frame_count = 1 + (samples.size - n_fft) // hop + shape = (frame_count, n_fft) + strides = (samples.strides[0] * hop, samples.strides[0]) + frames = np.lib.stride_tricks.as_strided(samples, shape=shape, strides=strides) + spectrum = np.abs(np.fft.rfft(frames * np.hanning(n_fft), axis=1)) ** 2 + + mel_points = np.linspace(hz_to_mel(0.0), hz_to_mel(sample_rate / 2.0), bands + 2) + bins = np.floor((n_fft + 1) * mel_to_hz(mel_points) / sample_rate).astype(np.int64) + bins = np.clip(bins, 0, spectrum.shape[1] - 1) + filters = np.zeros((bands, spectrum.shape[1]), dtype=np.float64) + for band in range(bands): + left, center, right = bins[band : band + 3] + if center > left: + filters[band, left:center] = np.arange(center - left) / (center - left) + if right > center: + filters[band, center:right] = np.arange(right - center, 0, -1) / (right - center) + return np.log(np.maximum(spectrum @ filters.T, 1.0e-10)).astype(np.float32) + + +def timings(path: Path) -> dict[int, float]: + result: dict[int, float] = {} + for line in (path / "timing.log").read_text(encoding="utf-8").splitlines(): + prefix = "higgs_audio_tts.cpp.request_" + if not line.startswith(prefix) or ".wall_ms=" not in line: + continue + index_text, value = line[len(prefix) :].split(".wall_ms=", 1) + result[int(index_text)] = float(value) + return result + + +def audio_files(path: Path) -> dict[int, Path]: + result: dict[int, Path] = {} + for wav_path in (path / "audio").glob("audio_*.wav"): + result[int(wav_path.stem.removeprefix("audio_"))] = wav_path + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--baseline", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--output", type=Path) + parser.add_argument("--require-same-frames", action="store_true") + parser.add_argument("--min-wav-cosine", type=float) + parser.add_argument("--min-logmel-cosine", type=float) + args = parser.parse_args() + + baseline_times = timings(args.baseline) + candidate_times = timings(args.candidate) + baseline_audio = audio_files(args.baseline) + candidate_audio = audio_files(args.candidate) + indices = sorted(set(baseline_times) & set(candidate_times) & set(baseline_audio) & set(candidate_audio)) + if not indices: + raise ValueError("no matching warmbench requests were found") + + failed = False + requests = [] + for index in indices: + baseline_rate, baseline_samples = read_wav(baseline_audio[index]) + candidate_rate, candidate_samples = read_wav(candidate_audio[index]) + if baseline_rate != candidate_rate: + raise ValueError(f"request {index}: sample rates differ") + same_frames = baseline_samples.size == candidate_samples.size + wav_cosine = cosine(baseline_samples, candidate_samples) + baseline_mel = log_mel(baseline_samples, baseline_rate) + candidate_mel = log_mel(candidate_samples, candidate_rate) + mel_frames = min(baseline_mel.shape[0], candidate_mel.shape[0]) + logmel_cosine = cosine(baseline_mel[:mel_frames].reshape(-1), candidate_mel[:mel_frames].reshape(-1)) + baseline_duration_sec = baseline_samples.size / baseline_rate + candidate_duration_sec = candidate_samples.size / candidate_rate + baseline_ms = baseline_times[index] + candidate_ms = candidate_times[index] + baseline_rtf = baseline_ms / 1000.0 / baseline_duration_sec + candidate_rtf = candidate_ms / 1000.0 / candidate_duration_sec + request = { + "request_index": index, + "baseline_frames": int(baseline_samples.size), + "candidate_frames": int(candidate_samples.size), + "same_frames": same_frames, + "baseline_wall_ms": baseline_ms, + "candidate_wall_ms": candidate_ms, + "wall_speedup": baseline_ms / candidate_ms, + "baseline_rtf": baseline_rtf, + "candidate_rtf": candidate_rtf, + "rtf_speedup": baseline_rtf / candidate_rtf, + "wav_cosine": wav_cosine, + "logmel_cosine": logmel_cosine, + } + requests.append(request) + failed = failed or (args.require_same_frames and not same_frames) + failed = failed or (args.min_wav_cosine is not None and wav_cosine < args.min_wav_cosine) + failed = failed or (args.min_logmel_cosine is not None and logmel_cosine < args.min_logmel_cosine) + print( + f"request={index} frames={baseline_samples.size}/{candidate_samples.size} " + f"wall_ms={baseline_ms:.3f}/{candidate_ms:.3f} " + f"rtf={baseline_rtf:.4f}/{candidate_rtf:.4f} speedup={baseline_rtf / candidate_rtf:.3f}x " + f"wav_cos={wav_cosine:.8f} mel_cos={logmel_cosine:.8f}" + ) + + payload = {"baseline": str(args.baseline), "candidate": str(args.candidate), "requests": requests} + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/higgs_audio_tts/higgs_audio_tts_cuda_bench_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_cuda_bench_cases.json new file mode 100644 index 00000000..cc4ff314 --- /dev/null +++ b/tests/higgs_audio_tts/higgs_audio_tts_cuda_bench_cases.json @@ -0,0 +1,35 @@ +[ + { + "id": "clone_prefix_first", + "text": "Hello. This is the first retained prefix test.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 1234 + }, + { + "id": "clone_prefix_second", + "text": "The second request should reuse the cloned voice prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 2234 + }, + { + "id": "clone_prefix_third", + "text": "A third short sentence checks stable repeated generation.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 3234 + } +] diff --git a/tests/higgs_audio_tts/higgs_audio_tts_cuda_mixed_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_cuda_mixed_cases.json new file mode 100644 index 00000000..345b8f3c --- /dev/null +++ b/tests/higgs_audio_tts/higgs_audio_tts_cuda_mixed_cases.json @@ -0,0 +1,44 @@ +[ + { + "id": "clone_first", + "text": "This cloned request prepares the reusable reference prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6134 + }, + { + "id": "clone_reuse", + "text": "This cloned request reuses the same reference prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6234 + }, + { + "id": "unconditioned", + "text": "This request intentionally generates an unconditioned random voice.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6334 + }, + { + "id": "clone_after_unconditioned", + "text": "The cloned voice is rebuilt safely after the unconditioned request.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6434 + } +] diff --git a/tests/higgs_audio_tts/higgs_audio_tts_cuda_perf_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_cuda_perf_cases.json new file mode 100644 index 00000000..83403b80 --- /dev/null +++ b/tests/higgs_audio_tts/higgs_audio_tts_cuda_perf_cases.json @@ -0,0 +1,57 @@ +[ + { + "id": "control_room_update", + "text": "The control room confirmed the overnight checks and the field team can restart the survey.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 1234 + }, + { + "id": "lab_briefing", + "text": "The lab briefing is ready. Please confirm the calibration notes before the afternoon check-in.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 2234 + }, + { + "id": "dispatch_note", + "text": "Dispatch logged the revised route. The north access road is clear and the receiver test can begin after lunch.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 3234 + }, + { + "id": "short_status", + "text": "All systems are ready for the next test.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 4234 + }, + { + "id": "weather_report", + "text": "Light rain is expected this evening, but tomorrow morning should remain calm and clear.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 5234 + } +] diff --git a/tests/higgs_tts/higgs_tts_python_warm_bench.py b/tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py similarity index 96% rename from tests/higgs_tts/higgs_tts_python_warm_bench.py rename to tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py index cd4c3b80..7eeba322 100644 --- a/tests/higgs_tts/higgs_tts_python_warm_bench.py +++ b/tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py @@ -35,7 +35,7 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Python reference Higgs Audio v3 TTS warmbench.") - parser.add_argument("--family", default="higgs_tts") + parser.add_argument("--family", default="higgs_audio_tts") parser.add_argument("--model", type=Path, default=DEFAULT_MODEL) parser.add_argument("--reference-root", type=Path, default=REFERENCE_ROOT) parser.add_argument("--backend", choices=("cuda",), default="cuda") @@ -57,8 +57,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--port", type=int, default=18180) parser.add_argument("--server-timeout-sec", type=float, default=300.0) parser.add_argument("--output-dir", type=Path, default=None) - parser.add_argument("--audio-out", type=Path, default=Path("higgs_tts_python_audio.wav")) - parser.add_argument("--timing-file", type=Path, default=Path("higgs_tts_python_timing.log")) + parser.add_argument("--audio-out", type=Path, default=Path("higgs_audio_tts_python_audio.wav")) + parser.add_argument("--timing-file", type=Path, default=Path("higgs_audio_tts_python_timing.log")) parser.add_argument("--summary-file", type=Path, default=None) return parser.parse_args() @@ -322,9 +322,9 @@ def main() -> int: raise RuntimeError("Higgs TTS warmbench request sequence is empty") output_dir = args.output_dir or args.audio_out.parent output_dir.mkdir(parents=True, exist_ok=True) - timing_lines = ["higgs_tts.python.model_load_excluded=1"] + timing_lines = ["higgs_audio_tts.python.model_load_excluded=1"] server = start_server(args) - timing_lines.append(f"higgs_tts.python.server_pid={server.pid}") + timing_lines.append(f"higgs_audio_tts.python.server_pid={server.pid}") try: wait_for_server(args, server) for _ in range(max(0, args.warmup)): @@ -338,8 +338,8 @@ def main() -> int: summary, wall_ms = run_request(args, request, audio_path) total_ms += wall_ms average_ms = total_ms / float(max(1, args.iterations)) - timing_lines.append(f"higgs_tts.python.request_{request_index}.wall_ms={average_ms:.6f}") - print(f"higgs_tts.python.wall_ms={average_ms}") + timing_lines.append(f"higgs_audio_tts.python.request_{request_index}.wall_ms={average_ms:.6f}") + print(f"higgs_audio_tts.python.wall_ms={average_ms}") steps.append( { "request_index": request_index, @@ -353,7 +353,7 @@ def main() -> int: "metrics": {"wall_ms": average_ms}, } ) - summary_payload = {"family": "higgs_tts", "backend": args.backend, "sequence_steps": steps} + summary_payload = {"family": "higgs_audio_tts", "backend": args.backend, "sequence_steps": steps} if args.summary_file: args.summary_file.parent.mkdir(parents=True, exist_ok=True) args.summary_file.write_text(json.dumps(summary_payload, ensure_ascii=False) + "\n", encoding="utf-8") diff --git a/tests/higgs_tts/higgs_tts_sampler_logits.bin b/tests/higgs_audio_tts/higgs_audio_tts_sampler_logits.bin similarity index 100% rename from tests/higgs_tts/higgs_tts_sampler_logits.bin rename to tests/higgs_audio_tts/higgs_audio_tts_sampler_logits.bin diff --git a/tests/higgs_tts/higgs_tts_warm_bench.cpp b/tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp similarity index 80% rename from tests/higgs_tts/higgs_tts_warm_bench.cpp rename to tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp index 665e3881..2ddb68f6 100644 --- a/tests/higgs_tts/higgs_tts_warm_bench.cpp +++ b/tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -106,13 +107,14 @@ void set_optional_option( } } -engine::runtime::AudioBuffer read_reference_audio(const engine::io::json::Value & object) { +std::optional read_reference_audio( + const engine::io::json::Value & object) { auto reference_path = optional_string(object, "reference_audio"); if (reference_path.empty()) { reference_path = optional_string(object, "voice_ref"); } if (reference_path.empty()) { - throw std::runtime_error("Higgs TTS warmbench request missing field: reference_audio"); + return std::nullopt; } const auto wav = engine::audio::read_wav_f32(resolve_path(reference_path)); return engine::runtime::AudioBuffer{wav.sample_rate, wav.channels, wav.samples}; @@ -121,10 +123,12 @@ engine::runtime::AudioBuffer read_reference_audio(const engine::io::json::Value engine::runtime::TaskRequest make_request(const engine::io::json::Value & object) { engine::runtime::TaskRequest request; request.text_input = engine::runtime::Transcript{required_string(object, "text"), ""}; - request.voice = engine::runtime::VoiceCondition{}; - request.voice->speaker = engine::runtime::VoiceReference{}; - request.voice->speaker->audio = read_reference_audio(object); - request.options["reference_text"] = required_string(object, "reference_text"); + if (auto reference_audio = read_reference_audio(object); reference_audio.has_value()) { + request.voice = engine::runtime::VoiceCondition{}; + request.voice->speaker = engine::runtime::VoiceReference{}; + request.voice->speaker->audio = std::move(*reference_audio); + request.options["reference_text"] = required_string(object, "reference_text"); + } set_optional_option(request, object, "max_tokens", "max_tokens"); set_optional_option(request, object, "temperature", "temperature"); set_optional_option(request, object, "top_p", "top_p"); @@ -205,10 +209,18 @@ engine::io::json::Value step_json( if (!audio_path.empty()) { stem.emplace("audio", string(audio_path.string())); } + const auto & audio = *result.audio_output; + const double frames = static_cast( + audio.samples.size() / static_cast(std::max(1, audio.channels))); + const double duration_sec = audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0; + const double rtf = duration_sec > 0.0 ? wall_ms / 1000.0 / duration_sec : 0.0; return engine::io::json::Value::make_object({ {"request_index", number(static_cast(request_index))}, {"stems", engine::io::json::Value::make_array({engine::io::json::Value::make_object(std::move(stem))})}, - {"metrics", engine::io::json::Value::make_object({{"wall_ms", number(wall_ms)}})}, + {"metrics", engine::io::json::Value::make_object({ + {"wall_ms", number(wall_ms)}, + {"rtf", number(rtf)}, + })}, }); } @@ -238,10 +250,22 @@ int main(int argc, char ** argv) { 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", ""); + std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", ""); + const std::filesystem::path request_sequence_file = + arg_value(argc, argv, "--request-sequence-file", ""); + if (request_sequence_json.empty() && !request_sequence_file.empty()) { + std::ifstream input(request_sequence_file, std::ios::binary); + if (!input) { + throw std::runtime_error( + "failed to open Higgs TTS request sequence: " + request_sequence_file.string()); + } + request_sequence_json.assign( + std::istreambuf_iterator(input), + std::istreambuf_iterator()); + } 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/higgs_tts_warm_bench_timing.log"); + arg_value(argc, argv, "--timing-file", "/tmp/higgs_audio_tts_warm_bench_timing.log"); const std::filesystem::path log_path = arg_value(argc, argv, "--log-file", ""); if (has_flag(argc, argv, "--enable-trace")) { if (log_path.empty()) { @@ -252,7 +276,7 @@ int main(int argc, char ** argv) { engine::runtime::ModelLoadRequest load_request; load_request.model_path = model_path; - load_request.family_hint = "higgs_tts"; + load_request.family_hint = "higgs_audio_tts"; auto registry = engine::runtime::make_default_registry(); auto model = registry.load(load_request); @@ -281,7 +305,7 @@ int main(int argc, char ** argv) { std::filesystem::create_directories(output_dir); } - std::vector timing_lines{"higgs_tts.cpp.model_load_excluded=1"}; + std::vector timing_lines{"higgs_audio_tts.cpp.model_load_excluded=1"}; engine::io::json::Value::Array steps; steps.reserve(requests.size()); for (size_t request_index = 0; request_index < requests.size(); ++request_index) { @@ -307,21 +331,30 @@ int main(int argc, char ** argv) { last_result.audio_output->samples); } timing_lines.push_back( - "higgs_tts.cpp.request_" + std::to_string(request_index) + ".wall_ms=" + std::to_string(wall_ms)); - std::cout << "higgs_tts.cpp.wall_ms=" << wall_ms << "\n"; + "higgs_audio_tts.cpp.request_" + std::to_string(request_index) + ".wall_ms=" + std::to_string(wall_ms)); + const auto & audio = *last_result.audio_output; + const double frames = static_cast( + audio.samples.size() / static_cast(std::max(1, audio.channels))); + const double duration_sec = audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0; + const double rtf = duration_sec > 0.0 ? wall_ms / 1000.0 / duration_sec : 0.0; + timing_lines.push_back( + "higgs_audio_tts.cpp.request_" + std::to_string(request_index) + ".rtf=" + std::to_string(rtf)); + std::cout << "higgs_audio_tts.cpp.request=" << request_index + << " wall_ms=" << wall_ms + << " rtf=" << rtf << "\n"; steps.push_back(step_json(last_result, static_cast(request_index), wall_ms, audio_path)); } write_timing(timing_path, timing_lines); const auto summary = engine::io::json::Value::make_object({ - {"family", string("higgs_tts")}, + {"family", string("higgs_audio_tts")}, {"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 << "higgs_tts_warm_bench failed: " << ex.what() << "\n"; + std::cerr << "higgs_audio_tts_warm_bench failed: " << ex.what() << "\n"; return 1; } } diff --git a/tests/higgs_tts/higgs_tts_warm_bench_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_warm_bench_cases.json similarity index 100% rename from tests/higgs_tts/higgs_tts_warm_bench_cases.json rename to tests/higgs_audio_tts/higgs_audio_tts_warm_bench_cases.json diff --git a/tests/higgs_audio_tts/run_cuda_performance.ps1 b/tests/higgs_audio_tts/run_cuda_performance.ps1 new file mode 100644 index 00000000..cf15c446 --- /dev/null +++ b/tests/higgs_audio_tts/run_cuda_performance.ps1 @@ -0,0 +1,69 @@ +param( + [Parameter(Mandatory = $true)] + [string]$Model, + [string]$BuildDir = "build/windows-cuda-release", + [string]$Label = (Get-Date -Format "yyyyMMdd-HHmmss"), + [string]$Baseline = "", + [switch]$RequireSameFrames, + [int]$Device = 0, + [int]$Threads = 8, + [int]$Warmup = 1, + [int]$Iterations = 1 +) + +$ErrorActionPreference = "Stop" +$RepoRoot = (Resolve-Path (Join-Path $PSScriptRoot "../..")).Path +$ModelPath = (Resolve-Path $Model).Path +$Bench = Join-Path $RepoRoot "$BuildDir/bin/higgs_audio_tts_warm_bench.exe" +$Cases = Join-Path $PSScriptRoot "higgs_audio_tts_cuda_perf_cases.json" +$ResultDir = Join-Path $PSScriptRoot "results/$Label" + +if (-not (Test-Path -LiteralPath $Bench)) { + throw "Warmbench binary does not exist: $Bench" +} + +New-Item -ItemType Directory -Force (Join-Path $ResultDir "audio") | Out-Null +$PreviousErrorActionPreference = $ErrorActionPreference +$ErrorActionPreference = "Continue" +& $Bench ` + --model $ModelPath ` + --backend cuda ` + --device $Device ` + --threads $Threads ` + --warmup $Warmup ` + --iterations $Iterations ` + --request-sequence-file $Cases ` + --output-dir (Join-Path $ResultDir "audio") ` + --timing-file (Join-Path $ResultDir "timing.log") 2>&1 | + ForEach-Object { + if ($_ -is [System.Management.Automation.ErrorRecord]) { + $_.Exception.Message + } else { + $_ + } + } | + Tee-Object -FilePath (Join-Path $ResultDir "console.log") +$BenchExitCode = $LASTEXITCODE +$ErrorActionPreference = $PreviousErrorActionPreference +if ($BenchExitCode -ne 0) { + exit $BenchExitCode +} + +if ($Baseline) { + $BaselinePath = (Resolve-Path $Baseline).Path + $CompareArgs = @( + (Join-Path $PSScriptRoot "compare_warmbench_results.py"), + "--baseline", $BaselinePath, + "--candidate", $ResultDir, + "--output", (Join-Path $ResultDir "comparison.json") + ) + if ($RequireSameFrames) { + $CompareArgs += "--require-same-frames" + } + python @CompareArgs + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } +} + +Write-Host "Higgs CUDA performance artifacts: $ResultDir" diff --git a/tests/unittests/test_qwen_decoder_packed_projections.cpp b/tests/unittests/test_qwen_decoder_packed_projections.cpp new file mode 100644 index 00000000..f66353ff --- /dev/null +++ b/tests/unittests/test_qwen_decoder_packed_projections.cpp @@ -0,0 +1,514 @@ +#include "engine/framework/core/backend.h" +#include "engine/framework/modules/attention/qwen_causal_decoder.h" +#include "engine/framework/modules/attention/qwen_decoder.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr size_t kGraphBytes = 16 * 1024 * 1024; +constexpr size_t kGraphNodes = 4096; + +std::vector patterned(size_t count, float phase, float scale) { + std::vector values(count); + for (size_t i = 0; i < count; ++i) { + const float x = static_cast(i); + values[i] = scale * (std::sin(phase + 0.19f * x) + 0.35f * std::cos(phase + 0.07f * x)); + } + return values; +} + +void require_allclose( + const std::vector & actual, + const std::vector & expected, + float tolerance, + const std::string & label) { + if (actual.size() != expected.size()) { + throw std::runtime_error(label + " size mismatch"); + } + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + if (diff > tolerance) { + std::ostringstream message; + message << label << " mismatch at " << i << ": expected " << expected[i] + << ", got " << actual[i] << ", diff=" << diff; + throw std::runtime_error(message.str()); + } + } +} + +struct LayerResult { + std::vector output; + std::vector key; + std::vector value; +}; + +LayerResult run_layer(bool packed) { + constexpr int64_t batch = 1; + constexpr int64_t steps = 3; + constexpr int64_t hidden = 8; + constexpr int64_t heads = 2; + constexpr int64_t kv_heads = 1; + constexpr int64_t head_dim = 4; + constexpr int64_t intermediate = 12; + constexpr int64_t q_out = heads * head_dim; + constexpr int64_t kv_out = kv_heads * head_dim; + + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "qwen_packed_projection_test", engine::core::BackendType::Cpu}; + auto make_f32 = [&](std::initializer_list dims) { + return engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims(dims)); + }; + + auto input = make_f32({batch, steps, hidden}); + auto positions = engine::core::make_tensor( + ctx, + GGML_TYPE_I32, + engine::core::TensorShape::from_dims({steps})); + + engine::modules::QwenDecoderLayerWeights weights; + weights.input_norm = {make_f32({hidden}), std::nullopt}; + weights.post_norm = {make_f32({hidden}), std::nullopt}; + weights.self_attention.out_weight = make_f32({hidden, hidden}); + weights.mlp.down_proj = {make_f32({hidden, intermediate}), std::nullopt}; + + const auto q_values = patterned(static_cast(q_out * hidden), 0.1f, 0.12f); + const auto k_values = patterned(static_cast(kv_out * hidden), 0.5f, 0.10f); + const auto v_values = patterned(static_cast(kv_out * hidden), 0.9f, 0.08f); + const auto gate_values = patterned(static_cast(intermediate * hidden), 1.3f, 0.11f); + const auto up_values = patterned(static_cast(intermediate * hidden), 1.7f, 0.09f); + + if (packed) { + weights.self_attention.qkv_weight = make_f32({q_out + 2 * kv_out, hidden}); + weights.mlp.gate_up_proj = engine::modules::LinearWeights{ + make_f32({intermediate * 2, hidden}), + std::nullopt, + }; + } else { + weights.self_attention.q_weight = make_f32({q_out, hidden}); + weights.self_attention.k_weight = make_f32({kv_out, hidden}); + weights.self_attention.v_weight = make_f32({kv_out, hidden}); + weights.mlp.gate_proj = {make_f32({intermediate, hidden}), std::nullopt}; + weights.mlp.up_proj = {make_f32({intermediate, hidden}), std::nullopt}; + } + + engine::modules::QwenDecoderLayerConfig config; + config.hidden_size = hidden; + config.num_attention_heads = heads; + config.num_key_value_heads = kv_heads; + config.head_dim = head_dim; + config.intermediate_size = intermediate; + config.rms_norm_eps = 1e-5f; + config.qkv_layout = packed + ? engine::modules::QwenDecoderQKVLayout::PackedQKV + : engine::modules::QwenDecoderQKVLayout::Separate; + config.runtime.mlp.mode = packed + ? engine::modules::QwenDecoderMLPMode::PackedGateUp + : engine::modules::QwenDecoderMLPMode::Exact; + config.use_qk_norm = false; + config.runtime.attention.prefill_mode = engine::modules::QwenDecoderAttentionMode::ManualRepeat; + + const auto outputs = engine::modules::QwenDecoderLayerModule(config).build( + ctx, + input, + positions, + weights); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, outputs.output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate test tensors"); + } + + engine::core::write_tensor_f32(input, patterned(static_cast(batch * steps * hidden), 2.1f, 0.20f)); + engine::core::write_tensor_i32(positions, {0, 1, 2}); + engine::core::write_tensor_f32(*weights.input_norm.weight, patterned(hidden, 0.3f, 0.7f)); + engine::core::write_tensor_f32(*weights.post_norm.weight, patterned(hidden, 0.7f, 0.8f)); + engine::core::write_tensor_f32( + weights.self_attention.out_weight, + patterned(static_cast(hidden * hidden), 1.1f, 0.10f)); + engine::core::write_tensor_f32( + weights.mlp.down_proj.weight, + patterned(static_cast(hidden * intermediate), 1.9f, 0.10f)); + + if (packed) { + std::vector qkv_values; + qkv_values.reserve(q_values.size() + k_values.size() + v_values.size()); + qkv_values.insert(qkv_values.end(), q_values.begin(), q_values.end()); + qkv_values.insert(qkv_values.end(), k_values.begin(), k_values.end()); + qkv_values.insert(qkv_values.end(), v_values.begin(), v_values.end()); + engine::core::write_tensor_f32(*weights.self_attention.qkv_weight, qkv_values); + + std::vector gate_up_values; + gate_up_values.reserve(gate_values.size() + up_values.size()); + gate_up_values.insert(gate_up_values.end(), gate_values.begin(), gate_values.end()); + gate_up_values.insert(gate_up_values.end(), up_values.begin(), up_values.end()); + engine::core::write_tensor_f32(weights.mlp.gate_up_proj->weight, gate_up_values); + } else { + engine::core::write_tensor_f32(weights.self_attention.q_weight, q_values); + engine::core::write_tensor_f32(weights.self_attention.k_weight, k_values); + engine::core::write_tensor_f32(weights.self_attention.v_weight, v_values); + engine::core::write_tensor_f32(weights.mlp.gate_proj.weight, gate_values); + engine::core::write_tensor_f32(weights.mlp.up_proj.weight, up_values); + } + + ggml_backend_graph_compute(backend, graph); + LayerResult result; + engine::core::read_tensor_f32_into(outputs.output.tensor, result.output); + engine::core::read_tensor_f32_into(outputs.key.tensor, result.key); + engine::core::read_tensor_f32_into(outputs.value.tensor, result.value); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + return result; + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +void test_packed_qkv_and_gate_up_match_separate_projections() { + const auto separate = run_layer(false); + const auto packed = run_layer(true); + require_allclose(packed.output, separate.output, 2.0e-5f, "decoder output"); + require_allclose(packed.key, separate.key, 2.0e-5f, "decoder key"); + require_allclose(packed.value, separate.value, 2.0e-5f, "decoder value"); +} + +void test_suffix_causal_mask() { + const auto values = engine::modules::qwen_causal_suffix_mask_values(2, 3, 2); + if (values.size() != 30) { + throw std::runtime_error("suffix causal mask size mismatch"); + } + const std::vector expected{ + true, true, true, false, false, + true, true, true, true, false, + true, true, true, true, true, + }; + for (int batch = 0; batch < 2; ++batch) { + for (size_t i = 0; i < expected.size(); ++i) { + const float actual = ggml_fp16_to_fp32(values[static_cast(batch) * expected.size() + i]); + if ((expected[i] && actual != 0.0F) || (!expected[i] && !std::isinf(actual))) { + throw std::runtime_error("suffix causal mask visibility mismatch"); + } + } + } +} + +void test_f16_kv_set_rows() { + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_test", engine::core::BackendType::Cpu}; + const auto cache = engine::core::make_tensor( + ctx, + GGML_TYPE_F16, + engine::core::TensorShape::from_dims({1, 3, 1, 2})); + const auto row = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({1, 1, 1, 2})); + const auto row_index = engine::core::make_tensor( + ctx, + GGML_TYPE_I64, + engine::core::TensorShape::from_dims({1})); + const auto output = engine::modules::FastKVSetRowsModule({ + engine::modules::FastKVSetRowsMode::BackendViewOptimized, + }).build(ctx, cache, row, row_index); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate f16 KV test tensors"); + } + + engine::core::write_tensor_f16(cache, std::vector(6, 0.0F)); + engine::core::write_tensor_f32(row, {1.25F, -2.5F}); + const int64_t index = 1; + ggml_backend_tensor_set(row_index.tensor, &index, 0, sizeof(index)); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("f16 KV set-rows graph compute failed"); + } + + const auto values = engine::core::read_tensor_f16(output.tensor); + require_allclose(values, {0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F}, 1.0e-3F, "f16 KV cache"); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +void test_f16_kv_set_rows_batched() { + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_batched_test", engine::core::BackendType::Cpu}; + const auto cache = engine::core::make_tensor( + ctx, + GGML_TYPE_F16, + engine::core::TensorShape::from_dims({2, 3, 1, 2})); + const auto rows = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({2, 1, 1, 2})); + const auto row_indices = engine::core::make_tensor( + ctx, + GGML_TYPE_I64, + engine::core::TensorShape::from_dims({2})); + const auto output = engine::modules::FastKVSetRowsModule({ + engine::modules::FastKVSetRowsMode::BackendViewOptimized, + }).build(ctx, cache, rows, row_indices); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate batched F16 KV test tensors"); + } + + engine::core::write_tensor_f16(cache, std::vector(12, 0.0F)); + engine::core::write_tensor_f32(rows, {1.25F, -2.5F, 3.5F, -4.5F}); + const std::vector indices{1, 4}; + ggml_backend_tensor_set(row_indices.tensor, indices.data(), 0, indices.size() * sizeof(int64_t)); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("batched F16 KV set-rows graph compute failed"); + } + + const auto values = engine::core::read_tensor_f16(output.tensor); + require_allclose( + values, + {0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F, + 0.0F, 0.0F, 3.5F, -4.5F, 0.0F, 0.0F}, + 1.0e-3F, + "batched f16 KV cache"); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +int count_graph_op(ggml_cgraph * graph, ggml_op op) { + int count = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + count += node != nullptr && node->op == op ? 1 : 0; + } + return count; +} + +bool graph_contains_sequence(ggml_cgraph * graph, std::initializer_list ops) { + if (ops.size() == 0 || static_cast(ops.size()) > ggml_graph_n_nodes(graph)) { + return false; + } + for (int start = 0; start + static_cast(ops.size()) <= ggml_graph_n_nodes(graph); ++start) { + bool matches = true; + int offset = 0; + for (const ggml_op op : ops) { + const ggml_tensor * node = ggml_graph_node(graph, start + offset++); + matches = matches && node != nullptr && node->op == op; + } + if (matches) { + return true; + } + } + return false; +} + +std::string graph_ops(ggml_cgraph * graph) { + std::ostringstream out; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + if (i != 0) { + out << ','; + } + out << (node != nullptr ? ggml_op_name(node->op) : "null"); + } + return out.str(); +} + +void test_higgs_decode_graph_exposes_cuda_fast_paths() { + constexpr int64_t hidden = 8; + constexpr int64_t heads = 2; + constexpr int64_t kv_heads = 1; + constexpr int64_t head_dim = 4; + constexpr int64_t intermediate = 12; + constexpr int64_t cache_steps = 8; + constexpr int64_t qkv_out = heads * head_dim + 2 * kv_heads * head_dim; + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + throw std::runtime_error("failed to initialize Higgs decode graph test context"); + } + + try { + engine::core::ModuleBuildContext ctx{ggml, "higgs_decode_fast_path_test", engine::core::BackendType::Cuda}; + auto make_tensor = [&](ggml_type type, std::initializer_list dims) { + return engine::core::make_tensor(ctx, type, engine::core::TensorShape::from_dims(dims)); + }; + + const auto input = make_tensor(GGML_TYPE_F32, {1, 1, hidden}); + const auto positions = make_tensor(GGML_TYPE_I32, {1}); + const auto cache_key = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim}); + const auto cache_value = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim}); + const auto cache_slot = make_tensor(GGML_TYPE_I64, {1}); + const auto attention_mask = make_tensor(GGML_TYPE_F16, {1, 1, 1, cache_steps}); + + engine::modules::QwenDecoderLayerWeights weights; + weights.input_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt}; + weights.q_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt}; + weights.k_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt}; + weights.post_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt}; + weights.self_attention.qkv_weight = make_tensor(GGML_TYPE_F32, {qkv_out, hidden}); + weights.self_attention.out_weight = make_tensor(GGML_TYPE_F32, {hidden, hidden}); + weights.mlp.gate_up_proj = engine::modules::LinearWeights{ + make_tensor(GGML_TYPE_F32, {intermediate * 2, hidden}), + std::nullopt, + }; + weights.mlp.down_proj = { + make_tensor(GGML_TYPE_F32, {hidden, intermediate}), + std::nullopt, + }; + + engine::modules::QwenDecoderLayerConfig config; + config.hidden_size = hidden; + config.num_attention_heads = heads; + config.num_key_value_heads = kv_heads; + config.head_dim = head_dim; + config.intermediate_size = intermediate; + config.qkv_layout = engine::modules::QwenDecoderQKVLayout::PackedQKV; + config.use_qk_norm = true; + config.runtime.attention.static_mode = + engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + config.runtime.static_cache.update_mode = + engine::modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + config.runtime.static_cache.set_rows_mode = + engine::modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + config.runtime.mlp.mode = engine::modules::QwenDecoderMLPMode::PackedGateUp; + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + const auto outputs = engine::modules::QwenDecoderLayerModule(config).build_with_static_cache_tail( + ctx, + graph, + input, + positions, + weights, + cache_key, + cache_value, + cache_slot, + attention_mask); + ggml_build_forward_expand(graph, outputs.output.tensor); + + if (count_graph_op(graph, GGML_OP_FLASH_ATTN_EXT) != 1) { + throw std::runtime_error("Higgs decode graph must contain one grouped FlashAttention op"); + } + if (count_graph_op(graph, GGML_OP_SET_ROWS) != 2) { + throw std::runtime_error("Higgs decode graph must update both F16 KV caches with set-rows"); + } + if (count_graph_op(graph, GGML_OP_GLU) != 1) { + throw std::runtime_error("Higgs decode graph must contain one packed SwiGLU op"); + } + if (count_graph_op(graph, GGML_OP_REPEAT) != 0) { + throw std::runtime_error("grouped FlashAttention must not materialize repeated KV heads"); + } + if (!graph_contains_sequence(graph, {GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS})) { + throw std::runtime_error( + "Higgs key-cache update must expose CUDA RoPE/view/set-rows fusion; graph=" + + graph_ops(graph)); + } + + ggml_free(ggml); + } catch (...) { + ggml_free(ggml); + throw; + } +} + +} // namespace + +int main() { + try { + test_packed_qkv_and_gate_up_match_separate_projections(); + test_suffix_causal_mask(); + test_f16_kv_set_rows(); + test_f16_kv_set_rows_batched(); + test_higgs_decode_graph_exposes_cuda_fast_paths(); + std::cout << "qwen_decoder_packed_projection_test: ok\n"; + return 0; + } catch (const std::exception & ex) { + std::cerr << "qwen_decoder_packed_projection_test: failed: " << ex.what() << "\n"; + return 1; + } +} diff --git a/tests/warmbench.py b/tests/warmbench.py index f1180f36..5067e124 100644 --- a/tests/warmbench.py +++ b/tests/warmbench.py @@ -460,21 +460,21 @@ "log_mel_cosine_min": 0.90, "cpp_session_options": ["heartmula.weight_type=f32"], }, - "higgs_tts": { - "kind": "higgs_tts", + "higgs_audio_tts": { + "kind": "higgs_audio_tts", "modes": ["offline"], - "cpp_bin": "build/debug/bin/higgs_tts_warm_bench", - "python_script": "tests/higgs_tts/higgs_tts_python_warm_bench.py", + "cpp_bin": "build/debug/bin/higgs_audio_tts_warm_bench", + "python_script": "tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py", "python_conda_env": "qwen3-tts", "model": "models/higgs-audio-v3-tts-4b", - "case_catalog": "tests/higgs_tts/higgs_tts_warm_bench_cases.json", + "case_catalog": "tests/higgs_audio_tts/higgs_audio_tts_warm_bench_cases.json", "default_case_name": "default", "default_requests_per_session": 1, "default_warmup": 0, "wav_cosine_min": 0.90, "log_mel_cosine_min": 0.90, "length_ratio_min": 0.98, - "cpp_session_options": ["higgs_tts.codec_weight_type=f32"], + "cpp_session_options": ["higgs_audio_tts.codec_weight_type=f32"], }, "index_tts2": { "kind": "index_tts2", @@ -4215,7 +4215,7 @@ def build_heartmula_commands( return python_command, cpp_command -def build_higgs_tts_commands( +def build_higgs_audio_tts_commands( config: dict[str, Any], backend: str, args: argparse.Namespace, @@ -4370,7 +4370,7 @@ def validate_sequence_result(summary: dict[str, Any], request_count: int, kind: and len(step.get("stems", [])) > 0 and isinstance(step.get("metrics", {}), dict) for step in steps) - elif kind in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_tts", "index_tts2"}: + elif kind in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_audio_tts", "index_tts2"}: payload_valid = all( isinstance(step.get("stems", []), list) and len(step.get("stems", [])) > 0 @@ -4514,10 +4514,10 @@ def run_scenario( irodori_requests, request_manifest = resolve_vevo2_case(config, args) args.requests_per_session = len(irodori_requests) python_command, cpp_command = build_irodori_tts_commands(scenario_config, backend, args, scenario_dir, irodori_requests) - elif scenario_config["kind"] == "higgs_tts": + elif scenario_config["kind"] == "higgs_audio_tts": higgs_requests, request_manifest = resolve_vevo2_case(config, args) args.requests_per_session = len(higgs_requests) - python_command, cpp_command = build_higgs_tts_commands(scenario_config, backend, args, scenario_dir, higgs_requests) + python_command, cpp_command = build_higgs_audio_tts_commands(scenario_config, backend, args, scenario_dir, higgs_requests) elif scenario_config["kind"] == "index_tts2": index_tts2_requests, request_manifest = resolve_vevo2_case(config, args) args.requests_per_session = len(index_tts2_requests) @@ -4865,7 +4865,7 @@ def run_scenario( cpp_step_path = cpp_step_paths[request_index] if request_index < len(cpp_step_paths) else "" append_log(master_log, f"PYTHON OUTPUT family={family} mode={mode} backend={backend} request={request_index} path={python_step_path} valid={int(file_is_nonempty(python_step_path))}") append_log(master_log, f"CPP OUTPUT family={family} mode={mode} backend={backend} request={request_index} path={cpp_step_path} valid={int(file_is_nonempty(cpp_step_path))}") - elif scenario_config["kind"] in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_tts", "index_tts2"}: + elif scenario_config["kind"] in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_audio_tts", "index_tts2"}: python_valid = validate_sequence_result(python_summary, args.requests_per_session, scenario_config["kind"]) cpp_valid = validate_sequence_result(cpp_summary, args.requests_per_session, scenario_config["kind"]) python_step_paths = write_sequence_step_artifacts(python_summary.get("sequence_steps", []), scenario_dir / "python_json", "python") 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 f1b9230a..3314e72f 100644 --- a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json @@ -225,9 +225,9 @@ ] }, { - "id": "higgs_tts_voice_clone_longform", + "id": "higgs_audio_tts_voice_clone_longform", "coverage": "Higgs Audio v3 voice clone with framework long-form text chunking, bounded AR generation, and codec decode", - "family": "higgs_tts", + "family": "higgs_audio_tts", "model": "models/higgs-audio-v3-tts-4b", "task": "tts", "mode": "offline", @@ -239,7 +239,32 @@ "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.", "voice_ref": "resources/a.wav", - "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication." + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "text_chunk_size": 512 + } + ] + }, + { + "id": "fish_audio_voice_clone_longform", + "coverage": "Fish Audio S2-Pro voice clone with framework long-form text chunking, bounded AR generation, and codec decode", + "family": "fish_audio", + "model": "models/s2-pro", + "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.", + "voice_ref": "resources/a.wav", + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "text_chunk_size": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 1234 } ] }, diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index 70429432..eddc048a 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -548,9 +548,9 @@ ] }, { - "id": "higgs_tts_voice_clone_chunked", + "id": "higgs_audio_tts_voice_clone_chunked", "coverage": "Higgs Audio v3 voice clone path with framework text chunking, AR generation, and codec decode", - "family": "higgs_tts", + "family": "higgs_audio_tts", "model": "models/higgs-audio-v3-tts-4b", "task": "tts", "mode": "offline", @@ -571,9 +571,9 @@ ] }, { - "id": "higgs_tts_voice_clone_cache_pollution", + "id": "higgs_audio_tts_voice_clone_cache_pollution", "coverage": "Higgs Audio v3 long-lived session cache pollution check with short, long, medium, and long requests", - "family": "higgs_tts", + "family": "higgs_audio_tts", "model": "models/higgs-audio-v3-tts-4b", "task": "tts", "mode": "offline", @@ -611,6 +611,53 @@ } ] }, + { + "id": "fish_audio_s2_pro_paths", + "coverage": "Fish Audio S2-Pro path coverage for default voice, reference voice clone, inline control tag, AR generation, and codec decode in one offline session", + "family": "fish_audio", + "model": "models/s2-pro", + "task": "tts", + "mode": "offline", + "session_options": { + "fish_audio.weight_type": "native", + "fish_audio.codec_weight_type": "native", + "fish_audio.reference_cache_slots": "1" + }, + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "official_auto_voice_english", + "text": "The field recorder captured a clean reference take, and the operator confirmed that every timestamp matched the written production notes.", + "text_chunk_size": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 1234 + }, + { + "id": "official_reference_voice_clone", + "text": "The studio engineer checked the short voice prompt, confirmed the take was clear, and started the final render.", + "voice_ref": "resources/sample.wav", + "reference_text": "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you.", + "text_chunk_size": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 2234 + }, + { + "id": "official_inline_control_tag", + "text": "[whisper in small voice] The prototype actually worked after the last reset, and the control room stayed quiet until every green light appeared.", + "text_chunk_size": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 3234 + } + ] + }, { "id": "heartmula_music_generation", "coverage": "HeartMuLa text-to-music generation using native model weight types and the warmbench reference request", diff --git a/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py b/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py index 19864d8d..9f3f90b0 100644 --- a/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py +++ b/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py @@ -11,6 +11,7 @@ import wave from pathlib import Path +import librosa import numpy as np @@ -122,29 +123,6 @@ def stft_magnitude(samples: np.ndarray, n_fft: int = 1024, hop: int = 256) -> np return spec -def hz_to_mel(freq: np.ndarray | float) -> np.ndarray | float: - return 2595.0 * np.log10(1.0 + np.asarray(freq) / 700.0) - - -def mel_to_hz(mel: np.ndarray) -> np.ndarray: - return 700.0 * (np.power(10.0, mel / 2595.0) - 1.0) - - -def mel_filterbank(sample_rate: int, n_fft: int, n_mels: int = 80) -> np.ndarray: - freq_bins = n_fft // 2 + 1 - mel_points = np.linspace(float(hz_to_mel(0.0)), float(hz_to_mel(sample_rate / 2.0)), n_mels + 2) - hz_points = mel_to_hz(mel_points) - bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(np.int64) - filters = np.zeros((n_mels, freq_bins), dtype=np.float32) - for mel_index in range(n_mels): - left, center, right = bins[mel_index : mel_index + 3] - if center > left: - filters[mel_index, left:center] = (np.arange(left, center) - left) / (center - left) - if right > center: - filters[mel_index, center:right] = (right - np.arange(center, right)) / (right - center) - return filters - - def wav_similarity_detail(src_path: Path, baseline_path: Path) -> tuple[float, str]: src_rate, src_audio = read_wav_f32(src_path) baseline_rate, baseline_audio = read_wav_f32(baseline_path) @@ -163,12 +141,13 @@ def wav_similarity_detail(src_path: Path, baseline_path: Path) -> tuple[float, s stft_cos = cosine_similarity(src_mag, baseline_mag) log_stft_cos = cosine_similarity(np.log1p(src_mag), np.log1p(baseline_mag)) - if src_rate == baseline_rate and freq_bins > 0 and frames > 0: - filters = mel_filterbank(src_rate, 1024) - filters = filters[:, :freq_bins] - src_mel = np.log1p(filters @ src_mag) - baseline_mel = np.log1p(filters @ baseline_mag) - log_mel_cos = cosine_similarity(src_mel, baseline_mel) + if src_rate == baseline_rate and src_mono.size > 0 and baseline_mono.size > 0: + src_mel = librosa.feature.melspectrogram(y=mono(src_audio), sr=src_rate) + baseline_mel = librosa.feature.melspectrogram(y=mono(baseline_audio), sr=baseline_rate) + mel_frames = min(src_mel.shape[1], baseline_mel.shape[1]) + src_log_mel = librosa.power_to_db(src_mel[:, :mel_frames], ref=1.0) + baseline_log_mel = librosa.power_to_db(baseline_mel[:, :mel_frames], ref=1.0) + log_mel_cos = cosine_similarity(src_log_mel, baseline_log_mel) log_mel_text = f"{log_mel_cos:.9f}" else: log_mel_text = "n/a" diff --git a/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py b/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py index ba66b1c5..8c45d9e0 100644 --- a/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py +++ b/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py @@ -349,7 +349,16 @@ def resolve_model_path(models_root: Path, value: str) -> Path: return models_root / path +def resolve_model_override(value: Path | None) -> Path | None: + if value is None: + return None + if value.is_absolute(): + return value + return REPO_ROOT / value + + def build_command(args: argparse.Namespace, case: dict[str, Any], case_dir: Path) -> list[str]: + model_path = resolve_model_override(args.model_path) or resolve_model_path(args.models_root, case["model"]) command = [ str(args.audiocpp_cli_bin), "--task", @@ -357,7 +366,7 @@ def build_command(args: argparse.Namespace, case: dict[str, Any], case_dir: Path "--family", case["family"], "--model", - str(resolve_model_path(args.models_root, case["model"])), + str(model_path), "--backend", case.get("backend", args.backend), "--mode", @@ -498,6 +507,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--cases", type=Path, default=DEFAULT_CASES) parser.add_argument("--audiocpp-cli-bin", type=Path, default=DEFAULT_AUDIOCPP_CLI_BIN) parser.add_argument("--models-root", type=Path, default=DEFAULT_MODELS_ROOT) + parser.add_argument("--model-path", type=Path, help="Override the model path for every selected case") parser.add_argument("--backend", default="cuda", choices=["cpu", "cuda", "vulkan", "metal", "best"]) parser.add_argument("--device", type=int, default=0) parser.add_argument("--threads", type=int, default=DEFAULT_THREADS) diff --git a/tools/model_manager.py b/tools/model_manager.py index 9af40949..8afd7514 100644 --- a/tools/model_manager.py +++ b/tools/model_manager.py @@ -107,6 +107,7 @@ class SnapshotSource: include_prefixes: tuple[str, ...] = () include_suffixes: tuple[str, ...] = () exclude_prefixes: tuple[str, ...] = () + strip_prefix: str = "" @dataclasses.dataclass(frozen=True) @@ -149,6 +150,7 @@ class ModelPackage: standalone: bool | None = None parent_package_id: str | None = None tasks: tuple[str, ...] = () + modes: tuple[str, ...] = () gated: bool | None = None @@ -216,22 +218,6 @@ def package_usage_examples(package: ModelPackage) -> list[str]: "vae/diffusion_pytorch_model.safetensors", ), ), - ModelPackage( - id="kokoro_82m_bf16", - display_name="Kokoro 82M bf16", - target_directory="Kokoro-82M-bf16", - source=UnsupportedSource( - reason=( - "kokoro_tts loader is not registered in this release tree yet " - "(commented out in src/framework/runtime/registry.cpp). " - "Re-enable the loader, add model_specs/kokoro_tts.json, then " - "restore a SnapshotSource here." - ), - ), - required_files=("config.json", "kokoro-v1_0.safetensors", "voices/af_heart.safetensors"), - family="kokoro_tts", - tasks=("tts",), - ), ModelPackage( id="moss_tts_nano_100m", display_name="MOSS-TTS-Nano 100M", @@ -396,30 +382,32 @@ def package_usage_examples(package: ModelPackage) -> list[str]: ), ModelPackage( id="voxtral_realtime", - display_name="Voxtral Mini 4B Realtime", - target_directory="Voxtral-Mini-4B-Realtime-2602", + display_name="Voxtral Mini 4B Realtime GGUF", + target_directory="Voxtral-Mini-4B-Realtime-2602-GGUF", source=SnapshotSource( - repo_id="mistralai/Voxtral-Mini-4B-Realtime-2602", - include_prefixes=( - "config.json", - "generation_config.json", - "model.safetensors", - "params.json", - "processor_config.json", - "tekken.json", - ), - ), - required_files=( - "config.json", - "generation_config.json", - "model.safetensors", - "params.json", - "processor_config.json", - "tekken.json", + repo_id="audio-cpp/audio.cpp-gguf", + include_prefixes=("Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf",), + strip_prefix="Voxtral-Mini-4B-Realtime-2602-GGUF/", ), + required_files=("voxtral-mini-4b-realtime-2602-q8_0.gguf",), family="voxtral_realtime", tasks=("asr",), - description="Native Hugging Face checkpoint for Voxtral realtime ASR; no conversion is required.", + modes=("offline", "streaming"), + description="Standalone audio.cpp Q8_0 GGUF package for Voxtral realtime ASR.", + ), + ModelPackage( + id="fish_audio_s2_pro", + display_name="Fish Audio S2 Pro GGUF", + target_directory="Fish-Audio-S2-Pro-GGUF", + source=SnapshotSource( + repo_id="audio-cpp/audio.cpp-gguf", + include_prefixes=("Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf",), + strip_prefix="Fish-Audio-S2-Pro-GGUF/", + ), + required_files=("fish-audio-s2-pro-q8_0.gguf",), + family="fish_audio", + tasks=("tts",), + description="Standalone audio.cpp Q8_0 GGUF package for Fish Audio S2 Pro.", ), ModelPackage( id="higgs_audio_stt", @@ -605,22 +593,6 @@ def package_usage_examples(package: ModelPackage) -> list[str]: source=SnapshotSource(repo_id="nvidia/diar_sortformer_4spk-v1"), required_files=("config.json", "model.safetensors", "processor_config.json"), ), - ModelPackage( - id="parakeet_tdt_0_6b_v3", - display_name="Parakeet TDT 0.6B v3", - target_directory="parakeet-tdt-0.6b-v3", - source=UnsupportedSource( - reason=( - "parakeet_tdt loader is not registered in this release tree yet " - "(commented out in src/framework/runtime/registry.cpp). " - "Re-enable the loader, add model_specs/parakeet_tdt.json, then " - "restore a SnapshotSource here." - ), - ), - required_files=("config.json", "model.safetensors", "processor_config.json", "tokenizer.json"), - family="parakeet_tdt", - tasks=("asr",), - ), ModelPackage( id="pocket_tts", display_name="PocketTTS", @@ -844,27 +816,17 @@ def package_usage_examples(package: ModelPackage) -> list[str]: ), ModelPackage( id="higgs_audio_v3_tts_4b", - display_name="Higgs Audio v3 TTS 4B", - target_directory="higgs-audio-v3-tts-4b", - source=UnsupportedSource( - reason=( - "higgs_tts / higgs_audio_tts loader is not registered in this " - "release tree yet (commented out in " - "src/framework/runtime/registry.cpp). Re-enable the loader with " - "one consistent family id, add the matching model_specs entry, " - "then restore a SnapshotSource here." - ), - ), - required_files=( - "chat_template.jinja", - "config.json", - "model.safetensors.index.json", - "model.safetensors", - "tokenizer.json", - "tokenizer_config.json", + display_name="Higgs Audio v3 TTS 4B GGUF", + target_directory="Higgs-Audio-v3-TTS-4B-GGUF", + source=SnapshotSource( + repo_id="audio-cpp/audio.cpp-gguf", + include_prefixes=("Higgs-Audio-v3-TTS-4B-GGUF/higgs-audio-v3-tts-4b-q8_0.gguf",), + strip_prefix="Higgs-Audio-v3-TTS-4B-GGUF/", ), + required_files=("higgs-audio-v3-tts-4b-q8_0.gguf",), family="higgs_audio_tts", tasks=("tts",), + description="Standalone audio.cpp Q8_0 GGUF package for Higgs Audio v3 TTS 4B.", ), ModelPackage( id="heartmula", @@ -1403,7 +1365,7 @@ def _default_tasks_from_family(family: str) -> list[str]: return [] if "forced_aligner" in key or key.endswith("_aligner") or key.endswith("_align"): return ["align"] - if key.endswith("_asr") or key.endswith("_stt") or key in {"parakeet_tdt", "whisper", "voxtral_realtime"}: + if key.endswith("_asr") or key.endswith("_stt") or key in {"whisper", "voxtral_realtime"}: return ["asr"] if "vad" in key: return ["vad"] @@ -1420,8 +1382,6 @@ def _default_tasks_from_family(family: str) -> list[str]: if key.endswith("_asr") or key.endswith("_stt"): return ["asr"] if "tts" in key or key in { - "kokoro", - "kokoro_tts", "chatterbox", "voxcpm2", "omnivoice", @@ -1484,6 +1444,7 @@ def package_payload(package: ModelPackage) -> dict[str, object]: "include_prefixes": list(source.include_prefixes), "include_suffixes": list(source.include_suffixes), "exclude_prefixes": list(source.exclude_prefixes), + "strip_prefix": source.strip_prefix, } installable = True elif isinstance(source, CompositeSnapshotSource): @@ -1498,6 +1459,7 @@ def package_payload(package: ModelPackage) -> dict[str, object]: "include_prefixes": list(placement.source.include_prefixes), "include_suffixes": list(placement.source.include_suffixes), "exclude_prefixes": list(placement.source.exclude_prefixes), + "strip_prefix": placement.source.strip_prefix, } for placement in source.placements ], @@ -1539,7 +1501,7 @@ def package_payload(package: ModelPackage) -> dict[str, object]: "source": source_payload, "family": family, "tasks": tasks, - "modes": ["offline"] if tasks else [], + "modes": list(package.modes) if package.modes else (["offline"] if tasks else []), "standalone": standalone, "parent_package_id": parent_package_id, "gated": _package_is_gated(package), @@ -1577,11 +1539,22 @@ def http_json(url: str) -> object: return json.load(response) -def list_hf_files(source: SnapshotSource) -> list[tuple[str, int | None]]: +def local_snapshot_path(source: SnapshotSource, remote_path: str) -> str: + if not source.strip_prefix: + return remote_path + if not remote_path.startswith(source.strip_prefix): + raise RuntimeError(f"snapshot path does not start with strip_prefix: {remote_path}") + local_path = remote_path[len(source.strip_prefix):] + if not local_path: + raise RuntimeError(f"snapshot strip_prefix removed full path: {remote_path}") + return local_path + + +def list_hf_files(source: SnapshotSource) -> list[tuple[str, str, int | None]]: payload = http_json(hf_tree_url(source)) if not isinstance(payload, list): raise RuntimeError(f"unexpected HuggingFace tree payload for {source.repo_id}") - files: list[tuple[str, int | None]] = [] + files: list[tuple[str, str, int | None]] = [] for entry in payload: if not isinstance(entry, dict): continue @@ -1596,7 +1569,7 @@ def list_hf_files(source: SnapshotSource) -> list[tuple[str, int | None]]: if any(path.startswith(prefix) for prefix in source.exclude_prefixes): continue size = entry.get("size") - files.append((path, size if isinstance(size, int) else None)) + files.append((path, local_snapshot_path(source, path), size if isinstance(size, int) else None)) if not files: raise RuntimeError(f"no installable files found for {source.repo_id}") return files @@ -1647,11 +1620,11 @@ def install_snapshot_into_dir( validate: bool = True, ) -> None: files = list_hf_files(source) - for relative, expected_size in files: + for remote, relative, expected_size in files: destination = destination_root / relative destination.parent.mkdir(parents=True, exist_ok=True) - print(f"download {relative}") - download_file(hf_resolve_url(source, relative), destination, expected_size) + print(f"download {remote}") + download_file(hf_resolve_url(source, remote), destination, expected_size) if validate: validate_required_files_list(required_files, destination_root, source.repo_id)