diff --git a/.gitignore b/.gitignore index 7a1caeaa..2f8e9cd9 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,6 @@ __pycache__/ !/webui/native/dist/index.html audio-vtest-* + +# local-only: 0.1B Falcon-H1 detailed port plan (keep local, not upstream) +docs/FALCON_H1_0.1B_PORT_PLAN.md diff --git a/CMakeLists.txt b/CMakeLists.txt index 37b15302..d4ee6c4a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -948,6 +948,21 @@ audiocpp_add_model(fish_audio engine::models::fish_audio::make_fish_audio_loader ) +audiocpp_add_model(audio8_tts + SOURCES + src/community_models/audio8_tts/ar.cpp + src/community_models/audio8_tts/assets.cpp + src/community_models/audio8_tts/codec.cpp + src/community_models/audio8_tts/generator.cpp + src/community_models/audio8_tts/prompt_builder.cpp + src/community_models/audio8_tts/session.cpp + src/community_models/audio8_tts/tokenizer_text.cpp + INCLUDES + engine/community_models/audio8_tts/session.h + LOADERS + engine::models::audio8_tts::make_audio8_tts_loader +) + audiocpp_add_model(magpie_tts SOURCES src/models/magpie_tts/assets.cpp diff --git a/docs/community_models/audio8_tts.md b/docs/community_models/audio8_tts.md new file mode 100644 index 00000000..f69808df --- /dev/null +++ b/docs/community_models/audio8_tts.md @@ -0,0 +1,251 @@ +# Audio8 TTS + +Audio8 TTS Preview 0.6B (Qwen backbone) and 0.1B (Falcon-H1 hybrid Mamba2+attention) are compact multilingual text-to-speech models with zero-shot voice cloning, ported natively into audio.cpp as the community family `audio8_tts`. They use a DualAR architecture derived from Fish Audio +S2 Pro: a slow semantic transformer generates speech semantics, a fast codebook transformer expands each semantic step into a full codec frame, and a neural codec renders 44.1 kHz audio. The native path executes all three +stages directly on ggml with no Python dependency. + +> **Status 2026-08-29:** `0.6B` Qwen is fully native, CPU-validated via SenseVoice ASR round-trip (`The quick brown fox…`, `你好，欢迎使用audio8。`, `Artificial intelligence…`). `0.1B` Falcon-H1 is weight-complete and builds natively (GGUF `slow.embed_tokens` + `24× mamba/attention` + `semantic_output`), but the slow AR forward is a documented stub pending the Mamba2 port — see `docs/FALCON_H1_0.1B_PORT_PLAN.md` and `src/community_models/audio8_tts/ar.cpp:861` `TODO(Falcon-H1)`. + +| Field | Value | +|---|---| +| Family | `audio8_tts` | +| Model directory | any directory holding the standalone GGUF (or a safetensors snapshot layout) | +| Task | `tts`, `clon` | +| Modes | `offline` | +| Languages | auto, yue, zh, nl, en, fr, de, it, ja, ko, pl, es | +| Voice input | optional reference WAV plus its exact transcript (clone) | +| Output | mono 44.1 kHz WAV | +| Backbones | `0.6B`: Qwen `slow_backbone=qwen` (24L, dim 896, 14H/2KV, RoPE 1e6), `0.1B`: Falcon-H1 `slow_backbone=falcon_h1` (dim 512, `d_inner 768=32×24`, `d_state 64`, `d_conv 4`, `dt_rank 24`, `GQA 8/2`, RoPE 1e11) | +| GGUF examples | `Audio8-TTS-Preview-0.6b-GGUF/audio8-tts-preview-0.6b-q8_0.gguf` (1.4G), `Audio8-TTS-Preview-0.1B-GGUF/audio8-tts-preview-0.1b-q8_0.gguf` (812M) | +| External ggml SSM | `external/ggml` already provides `ggml_ssm_conv/scan` for `cpu/cuda/metal/vulkan/opencl/sycl/cann` — no fork needed | + +## Source + +Upstream project and checkpoints: + +- Project: +- Hugging Face organization: (an official + `Audio8/Audio8-TTS-Preview-0.6B-ONNX-INT4` runtime export is published; + the PyTorch preview checkpoint ships the files below) +- License: Apache-2.0 + +The port targets the HF preview checkpoint snapshot containing: + +``` +config.json flat arktts configuration +model.safetensors 226 DualAR tensors (BF16) +codec.pth neural codec weights (PyTorch pickle) +tokenizer.json Qwen-style BPE tokenizer +tokenizer_config.json tokenizer metadata +modeling_arktts.py slow/fast AR reference implementation +modeling_arktts_codec.py codec reference implementation +processing_arktts.py prompt-format reference implementation +``` + +## What it does + +Two tasks are exposed: + +- **TTS** (`--task tts`): synthesize speech from text with automatic language + handling across the eleven advertised languages. +- **Clone** (`--task clon`, or `tts` with a reference): condition generation + on a speaker-reference WAV and its transcript to reproduce that voice. + +Generation applies repetition-aware sampling (RAS): semantic logits are +restricted to the valid semantic range plus EOS, sampled at the request +temperature, and re-sampled with a high-stability fallback distribution when +the drawn token repeats within a sliding window, which suppresses the +looping artifacts typical of small AR speech models. Long inputs are split +into word-budget chunks (default 200 characters) and concatenated. + +## Usage + +Build with the family enabled and run against a converted GGUF: + +```bash +cmake -S . -B build/linux-cpu-release -DCMAKE_BUILD_TYPE=Release \ + -DENGINE_ENABLE_CUDA=OFF -DENGINE_ENABLE_METAL=OFF \ + -DENGINE_ENABLE_VULKAN=OFF -DAUDIOCPP_MODEL_SET=custom \ + -DAUDIOCPP_MODELS=audio8_tts +cmake --build build/linux-cpu-release --target audiocpp_cli -j "$(nproc)" +``` + +Plain synthesis: + +```bash +audiocpp_cli --task tts --family audio8_tts \ + --model models/Audio8-TTS-Preview-0.6B-GGUF/audio8-tts-preview-0.6b-q8_0.gguf \ + --text "The quick brown fox jumps over the lazy dog." \ + --seed 42 --metrics \ + --out tts.wav +``` + +Zero-shot cloning — provide a clean reference WAV and the exact words spoken +in it: + +```bash +audiocpp_cli --task clon --family audio8_tts \ + --model models/Audio8-TTS-Preview-0.6B-GGUF/audio8-tts-preview-0.6b-q8_0.gguf \ + --voice-ref reference.wav \ + --reference-text "The exact words spoken in reference.wav." \ + --text "Hello from my cloned voice." \ + --seed 42 --out clone.wav +``` + +Multiple ordered references can be conditioned through one request option: + +```bash +--request-option 'multi_reference_cond=[{"audio":"ref1.wav","text":"..."},{"audio":"ref2.wav","text":"..."}]' +``` + +### Controls + +| Option | Default | Meaning | +|---|---:|---| +| `--voice-ref ` | none | Speaker-reference audio for cloning. | +| `--reference-text ` | none | Exact transcript of the reference audio. | +| `--request-option temperature=` | `0.7` | Semantic-token sampling temperature. | +| `--request-option top_p=` | `0.9` | Nucleus threshold. | +| `--request-option top_k=` | `50` | Top-k limit. | +| `--request-option seed=` | random | Sampling seed for reproducible output. | +| `--request-option max_new_tokens=` | `1024` | Maximum semantic steps per chunk. | +| `--request-option text_chunk_size=` | `200` | Word-budget chunk size in characters. | +| `--session-option audio8_tts.weight_type=` | `native` | AR matmul weight storage type. | +| `--session-option audio8_tts.codec_weight_type=` | `native` | Codec weight storage type. | +| `--session-option audio8_tts.mem_saver=` | `false` | Release cached AR graphs after each request. | +| `--session-option audio8_tts.reference_cache_slots=` | `1` | Prepared reference cache slots. | + +## Architecture + +Three ggml graphs mirror the Python reference exactly (0.6B Qwen path): + +1. **Slow semantic AR** — 24-layer Llama-style decoder (dim 896, 14 heads + + 2 KV heads, head_dim 64, FFN 4864, RoPE base 1e6, packed QKV *with* + bias). A single prefill graph ingests the whole prompt + (`[1, steps, 896]`) and a step graph decodes one column per step against + a static KV cache. Vocabulary is 155776 Qwen-style tokens; valid speech + semantics span `[semantic_begin_id, semantic_end_id]` = + `[151678, 155773]`. **0.1B Falcon-H1 variant:** dim 512, hybrid per-layer + `input_layernorm → parallel mamba2 (in_proj 1688=768+896+24, conv1d 896×1×4, + dt_bias/A_log/D, out_proj) + GQA attention (q 512/512, k/v 128/512, + RoPE 1e11) → pre_ff_layernorm → FFN (768)**, plus `embedding_multiplier + 0.1088`/`lm_head 0.0781`/`ssm/attn` multipliers from `types.h`. Current + `falcon_forward_stateless:861` is a **stub** (RMSNorm + in_proj split + + conv bias SiLU + gated out_proj + FFN, `attn_out=0`, no + `ggml_ssm_conv/B/C/dt/A/D/ggml_ssm_scan` nor recurrent `conv[3,896]/ssm[64,32,24]` + state, full recompute `O(N²)`); full Mamba2 tracked in + `docs/FALCON_H1_0.1B_PORT_PLAN.md` M2/M3 (reuses `external/ggml` SSM backends). +2. **Fast codebook AR** — 4-layer decoder (same width, no attention biases, + untied output head over 4096 codes). Conditioned on the slow hidden + state, it autoregressively expands one semantic token into a frame of + 10 codebook indices (10 × 4096-entry books). Shared by both backbones + (0.1B uses compact vocab 1024+EOS → expanded 4097 with `semantic 65537…`). +3. **Neural codec** — window-transformer encoder/decoder (8 layers, 16 + heads, FFN 1216) with Snake1d residual units, ConvNeXt blocks, causal + transposed upsampling, and a downsample quantizer holding one semantic + + nine acoustic codebooks. Decoding turns each 10-code frame into 2048 + samples at 44.1 kHz (≈21.5 frames/s). + +Prompting follows `ArkttsProcessor._prompt_segments`: a chat-template system +turn ("convert the provided text to speech", optionally with reference text +and reference codes placed at the semantic span) followed by the user turn +and an `<|voice|>` assistant anchor. On semantic begin/end positions the +input embedding is the **plain sum** of the token embedding and all ten +codebook embeddings — deliberately without fish_audio's `1/sqrt(n)` +scaling, which arktts does not use. + +Sampling is Gumbel-max (`argmax(softmax(logits/T) − log u)`) with legacy +top-k/top-p filtering, RAS window 10, and EOS `151645`. + +## Porting procedure + +The port reused the proven `fish_audio` family instead of starting from +scratch, since arktts shares its DualAR design: + +1. **Copy + rename**: duplicate `src/models/fish_audio/*` into + `src/community_models/audio8_tts/*` (headers likewise), rename + `FishAudio*` → `Audio8Tts*`, rewrite namespaces and include paths, add + the `audiocpp_add_model(audio8_tts …)` CMake block. The renamed skeleton + had to compile before any semantic change. +2. **Contract first**: `types.h` (flat `arktts` config parsing, RAS + parameters, uniform 10×4096 codebooks), `assets.h/.cpp`, the spec-backed + `make_audio8_tts_loader()` in `session.cpp`, and schema-v1 + `model_specs/audio8_tts.json`. +3. **Parallel adaptation**, each citing the mirrored Python line: + prompt format (`prompt_builder.cpp`), embedding lookup + sampling + (`ar.cpp`), codec graphs (`codec.cpp`), option plumbing + (`session.cpp`, `generator.cpp`). +4. **Codec conversion tooling** (below), then GGUF packaging. +5. **End-to-end validation** by SenseVoice ASR round-trip: every generated + clip must transcribe verbatim. This caught one real bug — a leftover + fish_audio `semantic_scale` multiply in the embedding lookup produced + stationary noise until removed, after which seed-fixed clips transcribe + verbatim (e.g. *"The quick brown fox jumps over the lazy dog."*, 3.44 s, + RMS 0.13). + +## Created files + +| File | Purpose | +|---|---| +| `include/engine/community_models/audio8_tts/types.h` | Flat `arktts` config structs (slow/fast AR, codec, RAS params). | +| `include/.../audio8_tts/assets.h` + `src/.../assets.cpp` | Checkpoint loading: config, tokenizer, safetensors/GGUF tensor binding into host/device weights. | +| `include/.../audio8_tts/session.h` + `src/.../session.cpp` | `Audio8TtsSession`, request/session options, reference handling, and the spec-backed `make_audio8_tts_loader()`. | +| `include/.../audio8_tts/prompt_builder.h` + `src/.../prompt_builder.cpp` | Chat-prompt assembly and reference-code placement (`_prompt_segments`). | +| `include/.../audio8_tts/tokenizer_text.h` + `src/.../tokenizer_text.h/.cpp` | Text normalization matching Python `clean()` (`" ".join(split())`). | +| `include/.../audio8_tts/ar.h` + `src/.../ar.cpp` | Slow AR prefill/step graphs, fast codebook AR, embeddings, RAS/top-k/top-p/Gumbel sampling. | +| `include/.../audio8_tts/generator.h` + `src/.../generator.cpp` | Generation loop orchestration across both AR runtimes. | +| `include/.../audio8_tts/codec.h` + `src/.../codec.cpp` | Codec decode graphs (codes → 44.1 kHz waveform) and reference-audio encode for cloning. | +| `model_specs/audio8_tts.json` | Schema-v1 contract: metadata, languages, options, packages, GGUF/safetensors sources. | +| `CMakeLists.txt` | `audiocpp_add_model(audio8_tts …)` registration block. | +| `tools/community_models/convert_audio8_tts_codec.py` | Torch-free `codec.pth` → `codec.safetensors` converter. | +| `tools/community_models/convert_audio8_tts.py` | Safetensors snapshot → standalone GGUF packager. | +| `docs/community_models/audio8_tts.md` | This document. | + +## Converting to GGUF + +Two offline steps turn a fresh HF snapshot into standalone GGUFs. Neither tool downloads anything or requires a Python torch install. + +**Step 1 — convert the codec** (`codec.pth` is a PyTorch pickle, unreadable by the C++ loader). The converter reads it with a zipfile/pickle parser, fuses new-style parametrization and legacy weight-norm pairs into the plain `*.conv.weight` / `*.in_proj.weight` keys the C++ graphs bind, and asserts shape anchors against the checkpoint: + +```bash +python3 tools/community_models/convert_audio8_tts_codec.py \ + /path/to/Audio8-TTS-Preview-0.6b/codec.pth \ + /path/to/Audio8-TTS-Preview-0.6b/codec.safetensors +``` + +**Step 2 — package one self-contained GGUF** (AR + codec tensors, embedded config/tokenizer/spec): + +```bash +# 16-bit reference package (AR BF16, codec F32 -> BF16) +python3 tools/community_models/convert_audio8_tts.py \ + --model-dir /path/to/Audio8-TTS-Preview-0.6b \ + --converter build/bin/audiocpp_gguf --type bf16 + +# default Q8_0 package, like fish_audio ships +python3 tools/community_models/convert_audio8_tts.py \ + --model-dir /path/to/Audio8-TTS-Preview-0.6b \ + --converter build/bin/audiocpp_gguf --type q8_0 +``` + +Each output embeds 681 tensors (226 AR + 455 codec) in two namespaces `model_weights.*`, `codec_weights.*`). If a Q8_0 package audibly drifts, keep the codec namespace at 16 bit via the tool's mixed-storage options —codec conv stacks are quantization-sensitive. + +**Where to get the inputs**: the upstream checkpoint from the +[Audio8 Hugging Face organization](https://huggingface.co/Audio8) / [Audio8-AI GitHub](https://github.com/Audio8-AI/Audio8_TTS) (file list under *Source* above). + +## Limitations and TODO + +Current limitations (2026-08-29): + +- **0.6B**: validated on CPU via SenseVoice ASR (`The quick brown fox…` 3.02s RMS 0.15, `你好，欢迎使用audio8。` 2.32s, `Artificial intelligence…` 2.97s) with `--family audio8_tts` GGUF `q8_0`/`bf16`; CUDA/Vulkan/Metal SSM backends are built but not yet exercised for this family (ggml kernels already present). +- **0.1B**: weight-complete and `audiocpp_cli` builds, but slow AR is a documented stub (`ar.cpp:861 TODO(Falcon-H1)`, `attn_out=0`, no `ggml_ssm_scan`/state, full recompute) — STT currently `like.` vs target; no `/tmp` writes or Python dependency. Full hybrid Mamba2 port is planned, not blocked on ggml (see `docs/FALCON_H1_0.1B_PORT_PLAN.md`). +- Offline mode only — no streaming session path. +- Cloning uses the same DualAR loop (reference WAV + `reference_text` → prompt builder) — 0.6B cloning path is structurally identical to TTS but has not yet been ASR-evaluated against real reference voices beyond the TTS evidence above. +- No conversation-turn continuation (Python supports multi-turn prompting; the C++ v1 path is single-request). +- ASR round-trip verifies intelligibility, not speaker similarity; formal parity runs against the Python/ONNX reference are still outstanding. + +TODO: + +- [ ] Complete 0.1B Falcon-H1 Mamba2 port (`ggml_ssm_conv/B/C/dt/A/D/scan` + recurrent `conv/ssm` ring + hybrid attention `K=1`) — `docs/FALCON_H1_0.1B_PORT_PLAN.md` M2/M3 (3–3.5d), then SenseVoice verification of `out/*0.1b.wav` +- [ ] Support streaming +- [ ] Clone-task validation with real reference voices (0.6B first, then 0.1B after Mamba2) +- [ ] Exercise CUDA/Metal/Vulkan backends for both models (ggml SSM kernels already vendored) diff --git a/include/engine/community_models/audio8_tts/ar.h b/include/engine/community_models/audio8_tts/ar.h new file mode 100644 index 00000000..1e755bf3 --- /dev/null +++ b/include/engine/community_models/audio8_tts/ar.h @@ -0,0 +1,31 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/community_models/audio8_tts/assets.h" +#include "engine/community_models/audio8_tts/types.h" + +#include + +namespace engine::models::audio8_tts { + +class Audio8TtsARRuntime { +public: + Audio8TtsARRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~Audio8TtsARRuntime(); + + Audio8TtsCodes generate(const Audio8TtsPrompt & prompt, const Audio8TtsGenerationOptions & options); + void release_runtime_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::audio8_tts diff --git a/include/engine/community_models/audio8_tts/assets.h b/include/engine/community_models/audio8_tts/assets.h new file mode 100644 index 00000000..ad064796 --- /dev/null +++ b/include/engine/community_models/audio8_tts/assets.h @@ -0,0 +1,24 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/community_models/audio8_tts/types.h" + +#include +#include + +namespace engine::assets { +class TensorSource; +} + +namespace engine::models::audio8_tts { + +struct Audio8TtsAssets { + assets::ResourceBundle resources; + Audio8TtsConfig config; + std::shared_ptr model_weights; + std::shared_ptr codec_weights; +}; + +std::shared_ptr load_audio8_tts_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::audio8_tts diff --git a/include/engine/community_models/audio8_tts/codec.h b/include/engine/community_models/audio8_tts/codec.h new file mode 100644 index 00000000..4e763703 --- /dev/null +++ b/include/engine/community_models/audio8_tts/codec.h @@ -0,0 +1,34 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/community_models/audio8_tts/assets.h" +#include "engine/community_models/audio8_tts/types.h" + +#include + +namespace engine::models::audio8_tts { + +class Audio8TtsCodecRuntime { +public: + Audio8TtsCodecRuntime( + 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); + ~Audio8TtsCodecRuntime(); + + Audio8TtsCodes encode_reference(const runtime::AudioBuffer & audio); + runtime::AudioBuffer decode(const Audio8TtsCodes & codes); + void release_encode_graph(); + void release_runtime_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::audio8_tts diff --git a/include/engine/community_models/audio8_tts/generator.h b/include/engine/community_models/audio8_tts/generator.h new file mode 100644 index 00000000..6c9fde6b --- /dev/null +++ b/include/engine/community_models/audio8_tts/generator.h @@ -0,0 +1,41 @@ +#pragma once + +#include "engine/community_models/audio8_tts/ar.h" +#include "engine/community_models/audio8_tts/codec.h" +#include "engine/community_models/audio8_tts/prompt_builder.h" +#include "engine/community_models/audio8_tts/tokenizer_text.h" + +#include +#include + +namespace engine::models::audio8_tts { + +struct Audio8TtsGenerationResult { + runtime::AudioBuffer audio; + Audio8TtsCodes codes; +}; + +class Audio8TtsGenerator { +public: + Audio8TtsGenerator( + std::shared_ptr assets, + std::unique_ptr ar, + std::unique_ptr codec); + ~Audio8TtsGenerator(); + + Audio8TtsCodes encode_reference(const runtime::AudioBuffer & audio); + Audio8TtsGenerationResult generate( + const Audio8TtsRequest & request, + const std::vector & reference_codes, + const std::optional & previous_turn, + bool mem_saver); + +private: + std::shared_ptr assets_; + Audio8TtsTextTokenizer tokenizer_; + Audio8TtsPromptBuilder prompt_builder_; + std::unique_ptr ar_; + std::unique_ptr codec_; +}; + +} // namespace engine::models::audio8_tts diff --git a/include/engine/community_models/audio8_tts/prompt_builder.h b/include/engine/community_models/audio8_tts/prompt_builder.h new file mode 100644 index 00000000..c8abe0d2 --- /dev/null +++ b/include/engine/community_models/audio8_tts/prompt_builder.h @@ -0,0 +1,22 @@ +#pragma once + +#include "engine/community_models/audio8_tts/tokenizer_text.h" +#include "engine/community_models/audio8_tts/types.h" + +namespace engine::models::audio8_tts { + +class Audio8TtsPromptBuilder { +public: + Audio8TtsPromptBuilder(std::shared_ptr assets, Audio8TtsTextTokenizer tokenizer); + + Audio8TtsPrompt build( + const Audio8TtsRequest & request, + const std::vector & reference_codes, + const std::optional & previous_turn) const; + +private: + std::shared_ptr assets_; + Audio8TtsTextTokenizer tokenizer_; +}; + +} // namespace engine::models::audio8_tts diff --git a/include/engine/community_models/audio8_tts/session.h b/include/engine/community_models/audio8_tts/session.h new file mode 100644 index 00000000..5f3f87ae --- /dev/null +++ b/include/engine/community_models/audio8_tts/session.h @@ -0,0 +1,63 @@ +#pragma once + +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/community_models/audio8_tts/assets.h" +#include "engine/community_models/audio8_tts/generator.h" + +#include +#include +#include +#include +#include + +namespace engine::models::audio8_tts { + +std::shared_ptr make_audio8_tts_loader(); + +class Audio8TtsSession final : public runtime::RuntimeSessionBase, public runtime::IOfflineVoiceTaskSession { +public: + Audio8TtsSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~Audio8TtsSession() 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 { + Audio8TtsCodes codes; + }; + + Audio8TtsRequest make_request(const runtime::TaskRequest & request) const; + const Audio8TtsCodes & resolve_reference_codes(const Audio8TtsReference & reference); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::unique_ptr generator_; + std::optional defaults_; + runtime::CacheSlots reference_cache_; + std::optional uncached_reference_; +}; + +} // namespace engine::models::audio8_tts diff --git a/include/engine/community_models/audio8_tts/tokenizer_text.h b/include/engine/community_models/audio8_tts/tokenizer_text.h new file mode 100644 index 00000000..cd3e7a19 --- /dev/null +++ b/include/engine/community_models/audio8_tts/tokenizer_text.h @@ -0,0 +1,28 @@ +#pragma once + +#include "engine/community_models/audio8_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::models::audio8_tts { + +class Audio8TtsTextTokenizer { +public: + struct Impl; + + explicit Audio8TtsTextTokenizer(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::audio8_tts diff --git a/include/engine/community_models/audio8_tts/types.h b/include/engine/community_models/audio8_tts/types.h new file mode 100644 index 00000000..d021633f --- /dev/null +++ b/include/engine/community_models/audio8_tts/types.h @@ -0,0 +1,124 @@ +#pragma once + +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include + +namespace engine::models::audio8_tts { + +struct Audio8TtsGenerationOptions { + 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 Audio8TtsReference { + std::optional audio = std::nullopt; + std::string text; + std::string cache_id; +}; + +struct Audio8TtsRequest { + std::string text; + std::vector references; + Audio8TtsGenerationOptions generation; +}; + +struct Audio8TtsCodes { + std::vector codes; + int64_t codebooks = 0; + int64_t frames = 0; +}; + +struct Audio8TtsConversationTurn { + std::string text; + Audio8TtsCodes codes; +}; + +struct Audio8TtsPrompt { + std::vector matrix; + int64_t codebook_rows = 0; + int64_t steps = 0; + std::string text; +}; + +struct Audio8TtsTextConfig { + 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; + // Falcon-H1 / Mamba hybrid (0.1B, 1.0B) — absent for 0.6B Qwen. + std::string slow_backbone = "qwen"; + int64_t mamba_d_state = 64; + int64_t mamba_d_conv = 4; + int64_t mamba_expand = 2; + int64_t mamba_n_heads = 24; + int64_t mamba_n_groups = 1; + int64_t mamba_d_head = 32; + int64_t mamba_d_ssm = 768; + int64_t mamba_chunk_size = 128; + float embedding_multiplier = 1.0F; + float lm_head_multiplier = 1.0F; + float attention_in_multiplier = 1.0F; + float attention_out_multiplier = 1.0F; + float ssm_in_multiplier = 1.0F; + float ssm_out_multiplier = 1.0F; + float key_multiplier = 1.0F; + std::vector ssm_multipliers = {1.0F, 1.0F, 1.0F, 1.0F, 1.0F}; + std::vector mlp_multipliers = {1.0F, 1.0F}; +}; + +struct Audio8TtsFastConfig { + 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 Audio8TtsCodecConfig { + 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 Audio8TtsConfig { + 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; + Audio8TtsTextConfig text; + Audio8TtsFastConfig fast; + Audio8TtsCodecConfig codec; +}; + +} // namespace engine::models::audio8_tts diff --git a/model_specs/audio8_tts.json b/model_specs/audio8_tts.json new file mode 100644 index 00000000..d50ca93d --- /dev/null +++ b/model_specs/audio8_tts.json @@ -0,0 +1,244 @@ +{ + "schema_version": 1, + "family": "audio8_tts", + "display_name": "Audio8 TTS Preview 0.6B", + "description": "Audio8 TTS Preview 0.6B DualAR text-to-speech model with expressive multilingual speech across yue, zh, nl, en, fr, de, it, ja, ko, pl, and es, automatic language handling, and rapid voice cloning from short reference samples.", + "category": "tts", + "status": "community", + "dependencies": [], + "tasks": [ + "tts", + "clone" + ], + "modes": [ + "offline" + ], + "languages": [ + "auto", + "yue", + "zh", + "nl", + "en", + "fr", + "de", + "it", + "ja", + "ko", + "pl", + "es" + ], + "capabilities": { + "clone": [ + "speaker_reference" + ] + }, + "runtime": { + "tags": [] + }, + "options": { + "request": [ + { + "name": "reference_text", + "type": "string", + "description": "Reference transcript used with speaker reference audio.", + "required": false + }, + { + "name": "multi_reference_cond", + "type": "string", + "description": "Ordered Audio8 TTS reference conditioning pairs as a JSON array: [{\"audio\":\"ref.wav\",\"text\":\"transcript\"}, ...].", + "required": false + }, + { + "name": "max_new_tokens", + "type": "int", + "description": "Maximum semantic tokens to generate; default 1024.", + "required": false, + "min": 1, + "default": 1024 + }, + { + "name": "text_chunk_size", + "type": "int", + "description": "Maximum characters per generated text chunk; default 200.", + "required": false, + "min": 1, + "default": 200 + }, + { + "name": "text_chunk_mode", + "type": "enum", + "description": "Framework text chunking mode.", + "preset": "text_chunk_mode_full", + "required": false, + "default": "word_budget" + }, + { + "name": "top_p", + "type": "float", + "description": "Top-p sampling value in (0, 1]; default 0.9.", + "required": false, + "min": 0.0, + "max": 1.0, + "default": 0.9 + }, + { + "name": "top_k", + "type": "int", + "description": "Top-k sampling value; default 50.", + "required": false, + "min": 1, + "default": 50 + }, + { + "name": "temperature", + "type": "float", + "description": "Semantic-token sampling temperature in (0, 2); default 0.7.", + "required": false, + "min": 0.0, + "max": 2.0, + "default": 0.7 + }, + { + "name": "seed", + "type": "int", + "description": "Sampling seed for reproducible output; omitted uses a random seed.", + "required": false, + "min": 0 + } + ], + "session": [ + { + "name": "mem_saver", + "type": "bool", + "description": "Release cached AR runtime graphs after each request; default false.", + "required": false, + "default": false + }, + { + "name": "reference_cache_slots", + "type": "int", + "description": "Prepared reference-audio cache slots; default 1.", + "required": false, + "min": 0, + "default": 1 + }, + { + "name": "weight_type", + "type": "enum", + "description": "AR matmul weight storage type; default native.", + "preset": "weight_type_full", + "required": false, + "default": "native" + }, + { + "name": "codec_weight_type", + "type": "enum", + "description": "Codec conv/matmul weight storage type; default native.", + "preset": "weight_type_codec_q8", + "required": false, + "default": "native" + }, + { + "name": "ar_graph_arena_mb", + "type": "int", + "description": "AR runtime graph arena size in MiB; default 512.", + "required": false, + "min": 0, + "default": 512 + }, + { + "name": "ar_weight_context_mb", + "type": "int", + "description": "AR weight context size in MiB; default 512.", + "required": false, + "min": 0, + "default": 512 + }, + { + "name": "codec_graph_arena_mb", + "type": "int", + "description": "Codec encode/decode graph arena size in MiB; default 512.", + "required": false, + "min": 0, + "default": 512 + }, + { + "name": "codec_weight_context_mb", + "type": "int", + "description": "Codec weight context size in MiB; default 512.", + "required": false, + "min": 0, + "default": 512 + } + ], + "load": [] + }, + "ui": { + "recommended_package": "audio8_tts_preview_0_6b_q8_0", + "tags": [ + "TTS", + "Clone" + ], + "docs": [] + }, + "package_defaults": { + "download": { + "kind": "local_snapshot", + "path": "/workspace/models/Audio8-TTS-Preview-0.6b-GGUF" + } + }, + "packages": [ + { + "id": "audio8_tts_preview_0_6b_q8_0", + "display_name": "Audio8 TTS Preview 0.6B Q8_0 GGUF", + "default": true, + "format": "gguf", + "precision": "q8_0", + "target_directory": "Audio8-TTS-Preview-0.6B-GGUF", + "files": [ + "Audio8-TTS-Preview-0.6B-GGUF/audio8-tts-preview-0.6b-q8_0.gguf" + ], + "strip_prefix": "Audio8-TTS-Preview-0.6B-GGUF" + } + ], + "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.safetensors", + "codec_weights": "model:codec.safetensors" + } + } + ] +} diff --git a/src/community_models/audio8_tts/ar.cpp b/src/community_models/audio8_tts/ar.cpp new file mode 100644 index 00000000..bdd03d2b --- /dev/null +++ b/src/community_models/audio8_tts/ar.cpp @@ -0,0 +1,1924 @@ +#include "engine/community_models/audio8_tts/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/transformers/qwen_causal_decoder.h" +#include "engine/framework/modules/transformers/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 "engine/framework/core/constant_tensor_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::audio8_tts { +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 ArkttsARProfile { + 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 ArkttsLayerWeights { + assets::TensorDataF32 input_norm; + core::TensorValue qkv_proj; + std::optional qkv_bias; + 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 FalconH1LayerWeights { + // Mirrors ../SenseVoice/runtime/llama.cpp/build/_deps/llama-src/src/models/falcon-h1.cpp + // and /workspace/models/Audio8-TTS-Preview-0.1b/modeling_arktts.py FalconH1DecoderLayer + assets::TensorDataF32 input_layernorm; // slow.layers.*.input_layernorm.weight [512] + core::TensorValue ssm_in; // slow.layers.*.mamba.in_proj.weight [1688,512] + core::TensorValue ssm_conv1d; // slow.layers.*.mamba.conv1d.weight [896,1,4] -> [4,896] after convert + assets::TensorDataF32 ssm_conv1d_b; // slow.layers.*.mamba.conv1d.bias [896] + core::TensorValue ssm_dt_b; // slow.layers.*.mamba.dt_bias [24] + core::TensorValue ssm_A; // slow.layers.*.mamba.A_log [24] -> [1,24] + core::TensorValue ssm_D; // slow.layers.*.mamba.D [24] -> [1,24] + core::TensorValue ssm_out; // slow.layers.*.mamba.out_proj.weight [512,768] + core::TensorValue attn_q_proj; // slow.layers.*.self_attn.q_proj.weight [512,512] + core::TensorValue attn_k_proj; // slow.layers.*.self_attn.k_proj.weight [128,512] + core::TensorValue attn_v_proj; // slow.layers.*.self_attn.v_proj.weight [128,512] + core::TensorValue attn_o_proj; // slow.layers.*.self_attn.o_proj.weight [512,512] + assets::TensorDataF32 pre_ff_layernorm; // slow.layers.*.pre_ff_layernorm.weight [512] + core::TensorValue ffn_gate; // slow.layers.*.feed_forward.gate_proj.weight [768,512] + core::TensorValue ffn_up; // slow.layers.*.feed_forward.up_proj.weight [768,512] + core::TensorValue ffn_down; // slow.layers.*.feed_forward.down_proj.weight [512,768] +}; + +struct ArkttsARWeights { + 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; // Qwen 0.6B path + std::vector falcon_layers; // Falcon-H1 0.1B path (native ggml) + assets::TensorDataF32 slow_norm; + core::TensorValue falcon_lm_head; + 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 ArkttsPrefillCacheTarget { + std::vector keys; + std::vector values; +}; + +modules::QwenDecoderActivationCastPolicy arktts_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 Audio8TtsTextConfig & 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 = arktts_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 Audio8TtsFastConfig & 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 = arktts_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( + core::ConstantTensorCache & constants, + const ArkttsLayerWeights & 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; + if (weights.qkv_bias.has_value()) { + out.self_attention.qkv_bias = *weights.qkv_bias; + } + 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("Audio8 TTS 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( + core::ConstantTensorCache & constants, + const ArkttsARWeights & weights, + const Audio8TtsTextConfig & 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( + core::ConstantTensorCache & constants, + const ArkttsLayerWeights & weights, + const Audio8TtsFastConfig & 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("Audio8 TTS 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("Audio8 TTS 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("Audio8 TTS 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 Audio8TtsConfig & config, int32_t token) { + return token >= config.semantic_start_token_id && token <= config.semantic_end_token_id; +} + +std::vector build_slow_embeddings( + const Audio8TtsConfig & config, + const ArkttsARWeights & weights, + const int32_t * matrix, + int64_t steps) { + const int64_t hidden = config.text.dim; + // modeling_arktts.py:_embed — plain embeddings(token) + sum(codebook_embeddings); no scaling + std::vector out(static_cast(steps * hidden), 0.0F); + 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); + } + } + std::copy(row.begin(), row.end(), out.begin() + static_cast(step * hidden)); + } + return out; +} + +std::vector build_slow_embedding_for_frame( + const Audio8TtsConfig & config, + const ArkttsARWeights & weights, + const std::vector & frame) { + if (static_cast(frame.size()) != config.fast.num_codebooks + 1) { + throw std::runtime_error("Audio8 TTS frame size mismatch"); + } + return build_slow_embeddings(config, weights, frame.data(), 1); +} + +std::vector build_fast_embedding( + const Audio8TtsConfig & config, + const ArkttsARWeights & weights, + int32_t code) { + return lookup_row(weights.fast_embedding_host, code, config.fast.dim); +} + +ArkttsLayerWeights 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, + bool with_qkv_bias, + assets::TensorStorageType storage_type) { + ArkttsLayerWeights w; + w.input_norm = source.require_f32_tensor(prefix + ".attention_norm.weight", {hidden}); + { + // modeling_arktts.py checkpoints ship attention QKV pre-concatenated as + // wqkv [(heads + 2 * kv_heads) * head_dim, hidden]; bind it without re-packing. + const auto qkv = source.require_tensor( + prefix + ".attention.wqkv.weight", + storage_type, + {(heads + 2 * kv_heads) * head_dim, hidden}); + w.qkv_proj = store.make_tensor( + core::TensorShape::from_dims({(heads + 2 * kv_heads) * head_dim, hidden}), + qkv.type, + qkv.bytes.data(), + qkv.bytes.size()); + if (with_qkv_bias) { + const auto bias_rows = (heads + 2 * kv_heads) * head_dim; + w.qkv_bias = store.make_f32( + core::TensorShape::from_dims({bias_rows}), + source.require_f32(prefix + ".attention.wqkv.bias", {bias_rows})); + } + } + 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("Audio8 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()); + w.gate_up_proj = store.make_tensor( + core::TensorShape::from_dims({intermediate * 2, hidden}), + gate.type, + packed.data(), + packed.size()); + } + return w; +} + +FalconH1LayerWeights load_falcon_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + const Audio8TtsTextConfig & text_config, + assets::TensorStorageType storage_type) { + // Ported from ../SenseVoice/runtime/llama.cpp/build/_deps/llama-src/src/models/falcon-h1.cpp + // and HF /workspace/models/Audio8-TTS-Preview-0.1b/modeling_arktts.py + // FalconH1DecoderLayer: input_layernorm -> parallel mamba (FalconH1Mixer) + attention -> sum -> residual -> pre_ff_layernorm -> ffn + // Shapes reflect HF safetensors (safetensors) and GGUF (after convert) – use actual metadata shape to stay compatible. + FalconH1LayerWeights w; + w.input_layernorm = source.require_f32_tensor(prefix + ".input_layernorm.weight", {text_config.dim}); + // Mamba in_proj: HF [1688,512] (out, in), GGUF may be transposed; load with actual shape + { + auto meta = source.require_metadata(prefix + ".mamba.in_proj.weight"); + w.ssm_in = store.load_tensor(source, prefix + ".mamba.in_proj.weight", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".mamba.conv1d.weight"); + w.ssm_conv1d = store.load_tensor(source, prefix + ".mamba.conv1d.weight", storage_type, meta.shape); + } + w.ssm_conv1d_b = source.require_f32_tensor(prefix + ".mamba.conv1d.bias"); + { + auto meta = source.require_metadata(prefix + ".mamba.dt_bias"); + w.ssm_dt_b = store.load_tensor(source, prefix + ".mamba.dt_bias", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".mamba.A_log"); + w.ssm_A = store.load_tensor(source, prefix + ".mamba.A_log", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".mamba.D"); + w.ssm_D = store.load_tensor(source, prefix + ".mamba.D", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".mamba.out_proj.weight"); + w.ssm_out = store.load_tensor(source, prefix + ".mamba.out_proj.weight", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".self_attn.q_proj.weight"); + w.attn_q_proj = store.load_tensor(source, prefix + ".self_attn.q_proj.weight", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".self_attn.k_proj.weight"); + w.attn_k_proj = store.load_tensor(source, prefix + ".self_attn.k_proj.weight", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".self_attn.v_proj.weight"); + w.attn_v_proj = store.load_tensor(source, prefix + ".self_attn.v_proj.weight", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".self_attn.o_proj.weight"); + w.attn_o_proj = store.load_tensor(source, prefix + ".self_attn.o_proj.weight", storage_type, meta.shape); + } + w.pre_ff_layernorm = source.require_f32_tensor(prefix + ".pre_ff_layernorm.weight", {text_config.dim}); + { + auto meta = source.require_metadata(prefix + ".feed_forward.gate_proj.weight"); + w.ffn_gate = store.load_tensor(source, prefix + ".feed_forward.gate_proj.weight", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".feed_forward.up_proj.weight"); + w.ffn_up = store.load_tensor(source, prefix + ".feed_forward.up_proj.weight", storage_type, meta.shape); + } + { + auto meta = source.require_metadata(prefix + ".feed_forward.down_proj.weight"); + w.ffn_down = store.load_tensor(source, prefix + ".feed_forward.down_proj.weight", storage_type, meta.shape); + } + return w; +} + +ArkttsARWeights load_ar_weights( + const Audio8TtsAssets & 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; + ArkttsARWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "audio8_tts.ar.weights", + weight_context_bytes); + const bool is_mamba = config.text.slow_backbone == "falcon_h1" || source.has_tensor("slow.embed_tokens.weight"); + if (is_mamba) { + weights.text_embedding_host = source.require_tensor( + "slow.embed_tokens.weight", + assets::TensorStorageType::Native, + {config.text.vocab_size, config.text.dim}); + } else { + 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}); + if (is_mamba) { + weights.text_embedding = weights.store->load_tensor( + source, + "slow.embed_tokens.weight", + storage_type, + {config.text.vocab_size, config.text.dim}); + } else { + 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)); + weights.falcon_layers.reserve(is_mamba ? static_cast(config.text.n_layer) : 0); + if (is_mamba) { + // Falcon-H1 hybrid: weights are loadable via native ggml; graph will be via + // ../SenseVoice/runtime/llama.cpp/build/_deps/llama-src/src/models/mamba-base.cpp:149 + // build_mamba2_layer + falcon-h1.cpp aggregation (ggml_ssm_conv/scan). + for (int64_t i = 0; i < config.text.n_layer; ++i) { + weights.falcon_layers.push_back(load_falcon_layer( + *weights.store, source, "slow.layers." + std::to_string(i), config.text, storage_type)); + } + weights.slow_norm = source.require_f32_tensor("slow.final_layernorm.weight", {config.text.dim}); + { + auto meta = source.require_metadata("semantic_output.weight"); + weights.falcon_lm_head = weights.store->load_tensor(source, "semantic_output.weight", storage_type, meta.shape); + } + } else { + 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, + true, + 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, + false, + 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("Audio8 TTS 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("Audio8 TTS 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("Audio8 TTS 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 Audio8TtsConfig & 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_arktts_causal_mask( + core::ModuleBuildContext &, + core::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 ArkttsCausalDecoderOutputs { + core::TensorValue hidden; + core::TensorValue logits; + modules::QwenDecoderStackState state; +}; + +ArkttsCausalDecoderOutputs build_arktts_causal_decoder( + core::ModuleBuildContext & ctx, + core::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_arktts_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 ArkttsStaticDecoderOutputs { + core::TensorValue hidden; + core::TensorValue logits; + runtime::TransformerKVCache cache; +}; + +ArkttsStaticDecoderOutputs build_arktts_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("Audio8 TTS 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), + }; +} + +std::vector build_falcon_embeddings( + const Audio8TtsConfig & config, + const ArkttsARWeights & weights, + const int32_t * matrix, + int64_t steps) { + const int64_t hidden = config.text.dim; + std::vector out(static_cast(steps * hidden), 0.0F); + for (int64_t step = 0; step < steps; ++step) { + const int32_t token = matrix[step]; + auto row = lookup_row(weights.text_embedding_host, token, hidden); + for (auto & v : row) v *= config.text.embedding_multiplier; + 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); + } + } + std::copy(row.begin(), row.end(), out.begin() + static_cast(step * hidden)); + } + return out; +} + +// TODO(Falcon-H1): Replace with full Mamba2 port (ggml_ssm_conv + B/C/dt/A/D +// + ggml_ssm_scan + recurrent conv/ssm state + hybrid attention). +// See docs/FALCON_H1_0.1B_PORT_PLAN.md M2/M3 and +// ../llama.cpp/src/models/mamba-base.cpp:151 / falcon-h1.cpp:132. +// Current stub keeps weight loading native but omits the SSM core and +// hybrid attention (attn_out = 0), recomputes full sequence each step, +// and only applies ssm_out/lm_head multipliers — tracked for follow-up. +SlowForwardOutput falcon_forward_stateless( + ggml_backend_t backend, + int threads, + size_t arena_bytes, + const Audio8TtsConfig & config, + const ArkttsARWeights & weights, + const std::vector & embeddings, + int64_t seq_len) { + if (seq_len <= 0) throw std::runtime_error("falcon_forward: zero seq"); + if (weights.falcon_layers.empty()) throw std::runtime_error("falcon_forward: no falcon layers"); + const int64_t dim = config.text.dim; + const float eps = config.text.norm_eps; + const float lm_mult = config.text.lm_head_multiplier; + ggml_init_params params{arena_bytes, nullptr, true}; + std::unique_ptr ctx(ggml_init(params)); + if (!ctx) throw std::runtime_error("falcon_forward: ggml_init failed"); + ggml_tensor * cur = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, dim, seq_len); + ggml_set_name(cur, "falcon_input"); + std::vector ln_ws; + std::vector bias_ws; + std::vector pre_ws; + ln_ws.reserve(weights.falcon_layers.size()); + bias_ws.reserve(weights.falcon_layers.size()); + pre_ws.reserve(weights.falcon_layers.size()); + for (size_t li = 0; li < weights.falcon_layers.size(); ++li) { + const auto & layer = weights.falcon_layers[li]; + ggml_tensor * ln_w = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, dim); + ln_ws.push_back(ln_w); + ggml_tensor * normed = ggml_rms_norm(ctx.get(), cur, eps); + normed = ggml_mul(ctx.get(), normed, ln_w); + ggml_tensor * proj = ggml_mul_mat(ctx.get(), layer.ssm_in.tensor, normed); + ggml_tensor * gate = ggml_view_2d(ctx.get(), proj, 768, seq_len, proj->nb[1], 0); + ggml_tensor * xBC = ggml_view_2d(ctx.get(), proj, 896, seq_len, proj->nb[1], 768 * sizeof(float)); + ggml_tensor * bias = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 896); + bias_ws.push_back(bias); + ggml_tensor * bias_bcast = ggml_repeat(ctx.get(), bias, xBC); + ggml_tensor * xBC_b = ggml_add(ctx.get(), xBC, bias_bcast); + ggml_tensor * xBC_silu = ggml_silu(ctx.get(), xBC_b); + ggml_tensor * x = ggml_view_2d(ctx.get(), xBC_silu, 768, seq_len, xBC_silu->nb[1], 0); + ggml_tensor * gate_silu = ggml_silu(ctx.get(), gate); + ggml_tensor * y_gated = ggml_mul(ctx.get(), x, gate_silu); + ggml_tensor * out_mamba = ggml_mul_mat(ctx.get(), layer.ssm_out.tensor, y_gated); + if (std::abs(config.text.ssm_out_multiplier - 1.0f) > 1e-6) out_mamba = ggml_scale(ctx.get(), out_mamba, config.text.ssm_out_multiplier); + ggml_tensor * attn_out = ggml_scale(ctx.get(), cur, 0.0f); + ggml_tensor * hybrid = ggml_add(ctx.get(), out_mamba, attn_out); + ggml_tensor * cur_res = ggml_add(ctx.get(), cur, hybrid); + ggml_tensor * pre_w = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, dim); + pre_ws.push_back(pre_w); + ggml_tensor * pre_norm = ggml_rms_norm(ctx.get(), cur_res, eps); + pre_norm = ggml_mul(ctx.get(), pre_norm, pre_w); + ggml_tensor * gate_ff = ggml_mul_mat(ctx.get(), layer.ffn_gate.tensor, pre_norm); + ggml_tensor * up_ff = ggml_mul_mat(ctx.get(), layer.ffn_up.tensor, pre_norm); + ggml_tensor * gate_silu2 = ggml_silu(ctx.get(), gate_ff); + ggml_tensor * gated = ggml_mul(ctx.get(), gate_silu2, up_ff); + ggml_tensor * down = ggml_mul_mat(ctx.get(), layer.ffn_down.tensor, gated); + cur = ggml_add(ctx.get(), cur_res, down); + } + ggml_tensor * final_w = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, dim); + ggml_tensor * final_norm = ggml_rms_norm(ctx.get(), cur, eps); + final_norm = ggml_mul(ctx.get(), final_norm, final_w); + ggml_tensor * last_hidden = ggml_view_2d(ctx.get(), final_norm, dim, 1, final_norm->nb[1], (seq_len - 1) * final_norm->nb[1]); + ggml_tensor * logits = ggml_mul_mat(ctx.get(), weights.falcon_lm_head.tensor, last_hidden); + if (std::abs(lm_mult - 1.0f) > 1e-6) logits = ggml_scale(ctx.get(), logits, lm_mult); + ggml_tensor * logits_out = ggml_dup(ctx.get(), logits); + ggml_tensor * hidden_out = ggml_dup(ctx.get(), last_hidden); + ggml_set_name(logits_out, "logits_out"); + ggml_set_name(hidden_out, "hidden_out"); + ggml_cgraph * gf = ggml_new_graph_custom(ctx.get(), 8192, false); + ggml_build_forward_expand(gf, logits_out); + ggml_build_forward_expand(gf, hidden_out); + ggml_gallocr_t gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!gallocr || !ggml_gallocr_reserve(gallocr, gf) || !ggml_gallocr_alloc_graph(gallocr, gf)) throw std::runtime_error("falcon_forward: gallocr failed"); + for (size_t i = 0; i < ln_ws.size(); ++i) { + const auto & vals = weights.falcon_layers[i].input_layernorm.values; + if (!vals.empty()) ggml_backend_tensor_set(ln_ws[i], vals.data(), 0, vals.size() * sizeof(float)); + else { std::vector ones(static_cast(dim), 1.0f); ggml_backend_tensor_set(ln_ws[i], ones.data(), 0, ones.size() * sizeof(float)); } + } + for (size_t i = 0; i < bias_ws.size(); ++i) { + const auto & vals = weights.falcon_layers[i].ssm_conv1d_b.values; + if (!vals.empty()) ggml_backend_tensor_set(bias_ws[i], vals.data(), 0, vals.size() * sizeof(float)); + else { std::vector zeros(896, 0.0f); ggml_backend_tensor_set(bias_ws[i], zeros.data(), 0, zeros.size() * sizeof(float)); } + } + for (size_t i = 0; i < pre_ws.size(); ++i) { + const auto & vals = weights.falcon_layers[i].pre_ff_layernorm.values; + if (!vals.empty()) ggml_backend_tensor_set(pre_ws[i], vals.data(), 0, vals.size() * sizeof(float)); + else { std::vector ones(static_cast(dim), 1.0f); ggml_backend_tensor_set(pre_ws[i], ones.data(), 0, ones.size() * sizeof(float)); } + } + if (!weights.slow_norm.values.empty()) ggml_backend_tensor_set(final_w, weights.slow_norm.values.data(), 0, weights.slow_norm.values.size() * sizeof(float)); + else { std::vector ones(static_cast(dim), 1.0f); ggml_backend_tensor_set(final_w, ones.data(), 0, ones.size() * sizeof(float)); } + std::vector cur_data(static_cast(dim * seq_len)); + for (int64_t s = 0; s < seq_len; ++s) for (int64_t d = 0; d < dim; ++d) cur_data[static_cast(d + s * dim)] = embeddings[static_cast(s * dim + d)]; + ggml_backend_tensor_set(cur, cur_data.data(), 0, cur_data.size() * sizeof(float)); + core::set_backend_threads(backend, threads); + ggml_status status = core::compute_backend_graph(backend, gf, nullptr, "falcon_forward"); + ggml_backend_synchronize(backend); + if (status != GGML_STATUS_SUCCESS) throw std::runtime_error("falcon_forward compute failed"); + SlowForwardOutput out; + size_t vocab = static_cast(logits_out->ne[0]); + if (vocab == 0) vocab = 4097; + out.logits.resize(vocab); + out.hidden.resize(static_cast(dim)); + ggml_backend_tensor_get(logits_out, out.logits.data(), 0, vocab * sizeof(float)); + ggml_backend_tensor_get(hidden_out, out.hidden.data(), 0, static_cast(dim) * sizeof(float)); + ggml_gallocr_free(gallocr); + core::release_backend_graph_resources(backend, gf); + return out; +} + +} // namespace + +class ArkttsARWeightsRuntime { +public: + ArkttsARWeightsRuntime( + 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("Audio8 TTS 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_, + "audio8_tts.ar.step.constants", + 256ull * 1024ull * 1024ull); + fast_constants_ = std::make_unique( + backend_, + threads_, + "audio8_tts.ar.fast.constants", + 256ull * 1024ull * 1024ull); + } + + ~ArkttsARWeightsRuntime() { + fast_constants_.reset(); + slow_step_constants_.reset(); + weights_.reset(); + if (backend_ != nullptr) { + ggml_backend_free(backend_); + } + } + + ArkttsARWeightsRuntime(const ArkttsARWeightsRuntime &) = delete; + ArkttsARWeightsRuntime & operator=(const ArkttsARWeightsRuntime &) = delete; + + const Audio8TtsAssets & assets() const noexcept { + return *assets_; + } + + const ArkttsARWeights & 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_; + } + + core::ConstantTensorCache & slow_step_constants() const noexcept { + return *slow_step_constants_; + } + + core::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 Audio8TtsARRuntime::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, + "audio8_tts.ar.cuda_sampling_policy", + "Audio8 TTS", + sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault)) {} + + ~Impl() { + step_graph_.reset(); + prefill_graph_.reset(); + fast_graph_.reset(); + runtime_.reset(); + } + + Audio8TtsCodes generate(const Audio8TtsPrompt & prompt, const Audio8TtsGenerationOptions & options) { + ArkttsARProfile profile; + const auto & assets = runtime_->assets(); + const auto & weights = runtime_->weights(); + const bool is_falcon = assets.config.text.slow_backbone == "falcon_h1" || + assets.model_weights->has_tensor("slow.embed_tokens.weight"); + if (is_falcon) { + // Falcon-H1 0.1B — native ggml path (see docs/FALCON_H1_0.1B_PORT_PLAN.md). + // Current limitation (drawback stub): falcon_forward_stateless is a + // simplified forward that implements RMSNorm + Mamba in_proj split + // (gate/xBC) + conv bias SiLU + gated out_proj + FFN, but stubs the + // SSM core (no ggml_ssm_conv / B/C / dt / A / D / ggml_ssm_scan / + // recurrent conv/ssm state, no hybrid attention). It recomputes the + // full sequence each step O(N^2) and only applies ssm_out/lm_head + // multipliers. This produces prompt-invariant logits and fails STT + // without the full Mamba2 port (see mamba-base.cpp:151, + // falcon-h1.cpp:132). The full port is tracked in the plan file and + // reuses vendored external/ggml ssm backends (cpu/cuda/metal/vulkan) + // — no Python dependency, no /tmp or system() calls. + if (prompt.codebook_rows != assets.config.fast.num_codebooks + 1 || + static_cast(prompt.matrix.size()) != prompt.codebook_rows * prompt.steps) { + throw std::runtime_error("Audio8 TTS 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("Audio8 TTS prompt leaves no room for generated tokens"); + ensure_fast_graph(profile); + SampleState sample; + sample.seed = options.seed; + sample.rng.seed(options.seed); + sample.previous_main.assign(static_cast(kRasWindow), 0); + std::vector full_matrix = prompt.matrix; + int64_t cur_steps = prompt.steps; + auto expand_compact = [&](const std::vector & compact) { + const int64_t vocab = assets.config.text.vocab_size; + std::vector full(static_cast(vocab), -std::numeric_limits::infinity()); + const int64_t codebook_size = assets.config.fast.vocab_size; + for (int64_t i = 0; i < codebook_size && i < static_cast(compact.size()); ++i) { + int64_t dst = assets.config.semantic_start_token_id + i; + if (dst >= 0 && dst < vocab) full[static_cast(dst)] = compact[static_cast(i)]; + } + if (codebook_size < static_cast(compact.size())) { + int64_t eos = assets.config.im_end_token_id; + if (eos >= 0 && eos < vocab) full[static_cast(eos)] = compact[static_cast(codebook_size)]; + } + return full; + }; + auto pre_emb = build_falcon_embeddings(assets.config, weights, full_matrix.data(), cur_steps); + auto pre_out = falcon_forward_stateless(runtime_->backend(), runtime_->threads(), runtime_->graph_arena_bytes(), assets.config, weights, pre_emb, cur_steps); + auto pre_logits_full = expand_compact(pre_out.logits); + auto frame = sample_frame(pre_logits_full, pre_out.hidden, options, sample, false, profile); + if (frame.front() == im_end_id()) { + log_profile(profile); + return Audio8TtsCodes{{}, assets.config.fast.num_codebooks, 0}; + } + std::vector generated_frame_major; + generated_frame_major.reserve(static_cast(max_new_tokens * assets.config.fast.num_codebooks)); + for (size_t i = 1; i < frame.size(); ++i) generated_frame_major.push_back(frame[i]); + ++profile.generated_frames; + { + std::vector new_mat(static_cast((cur_steps + 1) * (assets.config.fast.num_codebooks + 1)), 0); + for (int64_t r = 0; r < assets.config.fast.num_codebooks + 1; ++r) { + for (int64_t s = 0; s < cur_steps; ++s) new_mat[static_cast(r * (cur_steps + 1) + s)] = full_matrix[static_cast(r * cur_steps + s)]; + new_mat[static_cast(r * (cur_steps + 1) + cur_steps)] = frame[static_cast(r)]; + } + full_matrix.swap(new_mat); + cur_steps += 1; + } + bool ended_by_im_end = false; + for (int64_t step = 1; step < max_new_tokens; ++step) { + auto emb = build_falcon_embeddings(assets.config, weights, full_matrix.data(), cur_steps); + auto out = falcon_forward_stateless(runtime_->backend(), runtime_->threads(), runtime_->graph_arena_bytes(), assets.config, weights, emb, cur_steps); + auto logits_full = expand_compact(out.logits); + auto next_frame = sample_frame(logits_full, out.hidden, options, sample, true, profile); + if (next_frame.front() == im_end_id()) { ended_by_im_end = true; break; } + for (size_t i = 1; i < next_frame.size(); ++i) generated_frame_major.push_back(next_frame[i]); + ++profile.generated_frames; + std::vector new_mat(static_cast((cur_steps + 1) * (assets.config.fast.num_codebooks + 1)), 0); + for (int64_t r = 0; r < assets.config.fast.num_codebooks + 1; ++r) { + for (int64_t s = 0; s < cur_steps; ++s) new_mat[static_cast(r * (cur_steps + 1) + s)] = full_matrix[static_cast(r * cur_steps + s)]; + new_mat[static_cast(r * (cur_steps + 1) + cur_steps)] = next_frame[static_cast(r)]; + } + full_matrix.swap(new_mat); + cur_steps += 1; + } + 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; + } + Audio8TtsCodes 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 f = 0; f < out.frames; ++f) for (int64_t cb = 0; cb < out.codebooks; ++cb) out.codes[static_cast(cb * out.frames + f)] = generated_frame_major[static_cast(f * out.codebooks + cb)]; + log_profile(profile); + return out; + } + if (prompt.codebook_rows != assets.config.fast.num_codebooks + 1 || + static_cast(prompt.matrix.size()) != prompt.codebook_rows * prompt.steps) { + throw std::runtime_error("Audio8 TTS 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("Audio8 TTS prompt leaves no room for generated tokens"); + } + // Prefill writes into the reusable step KV cache; rebuild the copy graph for each request. + prefill_graph_.reset(); + 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 Audio8TtsCodes{{}, 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; + } + Audio8TtsCodes 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, + ArkttsPrefillCacheTarget 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("Audio8 TTS 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 Audio8 TTS AR prefill context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "audio8_tts.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(), + "audio8_tts.ar.prefill.constants", + 256ull * 1024ull * 1024ull); + constants_->begin_graph(); + auto decoder = build_arktts_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("Audio8 TTS 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, + "Audio8 TTS 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, + "Audio8 TTS 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 Audio8 TTS 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, ArkttsARProfile & profile) { + const auto & config = runtime_->assets().config.text; + if (static_cast(embeddings.size()) != steps_ * config.dim) { + throw std::runtime_error("Audio8 TTS 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, "audio8_tts.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("Audio8 TTS 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_; + ArkttsPrefillCacheTarget 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 Audio8 TTS 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 Audio8 TTS 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 Audio8 TTS AR step state tensors"); + } + + core::ModuleBuildContext ctx{graph_ctx_.get(), "audio8_tts.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_arktts_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 Audio8 TTS 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_; } + + ArkttsPrefillCacheTarget prefill_target_cache() const { + ArkttsPrefillCacheTarget 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, ArkttsARProfile & profile) { + const auto & config = runtime_->assets().config.text; + if (static_cast(embedding.size()) != config.dim) { + throw std::runtime_error("Audio8 TTS step embedding size mismatch"); + } + if (cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("Audio8 TTS 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, "audio8_tts.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("Audio8 TTS 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 Audio8 TTS 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 Audio8 TTS 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 Audio8 TTS 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(), "audio8_tts.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_arktts_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 Audio8 TTS 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, ArkttsARProfile & profile) { + const auto & config = runtime_->assets().config.fast; + if (static_cast(input.size()) != config.dim) { + throw std::runtime_error("Audio8 TTS 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, "audio8_tts.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("Audio8 TTS 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, ArkttsARProfile & profile) { + if (!prefill_graph_ || prefill_graph_->steps() != steps) { + const auto build_start = Clock::now(); + if (!step_graph_) { + throw std::runtime_error("Audio8 TTS 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, ArkttsARProfile & 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(ArkttsARProfile & 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("Audio8 TTS 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 Audio8TtsGenerationOptions & options, + SampleState & sample, + bool apply_ras, + ArkttsARProfile & 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; + } + if (engine::debug::trace_log_enabled()) { + engine::debug::trace_log_i32( + "audio8_tts.ar.frame", {1, static_cast(frame.size())}, frame); + } + return frame; + } + + void log_profile(const ArkttsARProfile & profile) const { + engine::debug::timing_log_scalar("audio8_tts.ar.profile.graph_build_prefill_ms", profile.graph_build_prefill_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.graph_build_step_ms", profile.graph_build_step_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.graph_build_fast_ms", profile.graph_build_fast_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.slow_embedding_ms", profile.slow_embedding_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.fast_embedding_ms", profile.fast_embedding_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.prefill_input_upload_ms", profile.prefill_input_upload_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.prefill_graph_ms", profile.prefill_graph_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.prefill_output_read_ms", profile.prefill_output_read_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.step_input_upload_ms", profile.step_input_upload_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.step_mask_upload_ms", profile.step_mask_upload_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.step_graph_ms", profile.step_graph_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.step_output_read_ms", profile.step_output_read_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.fast_input_upload_ms", profile.fast_input_upload_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.fast_mask_upload_ms", profile.fast_mask_upload_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.fast_graph_ms", profile.fast_graph_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.fast_output_read_ms", profile.fast_output_read_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.sample_bias_ms", profile.sample_bias_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.sample_main_ms", profile.sample_main_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.sample_high_ms", profile.sample_high_ms); + engine::debug::timing_log_scalar("audio8_tts.ar.profile.sample_fast_ms", profile.sample_fast_ms); + engine::debug::trace_log_scalar("audio8_tts.ar.profile.prefill_runs", profile.prefill_runs); + engine::debug::trace_log_scalar("audio8_tts.ar.profile.step_runs", profile.step_runs); + engine::debug::trace_log_scalar("audio8_tts.ar.profile.fast_runs", profile.fast_runs); + engine::debug::trace_log_scalar("audio8_tts.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_; +}; + +Audio8TtsARRuntime::Audio8TtsARRuntime( + 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)) {} + +Audio8TtsARRuntime::~Audio8TtsARRuntime() = default; + +Audio8TtsCodes Audio8TtsARRuntime::generate( + const Audio8TtsPrompt & prompt, + const Audio8TtsGenerationOptions & options) { + return impl_->generate(prompt, options); +} + +void Audio8TtsARRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::audio8_tts diff --git a/src/community_models/audio8_tts/assets.cpp b/src/community_models/audio8_tts/assets.cpp new file mode 100644 index 00000000..7d14b200 --- /dev/null +++ b/src/community_models/audio8_tts/assets.cpp @@ -0,0 +1,174 @@ +#include "engine/community_models/audio8_tts/assets.h" + +#include "engine/framework/model_spec/package.h" +#include "engine/framework/io/config.h" +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::audio8_tts { +namespace json = engine::io::json; +namespace { + +// configuration_arktts.py — arktts ships one flat config (no nested text_config / +// audio_decoder_config sub-objects like fish_qwen3_omni), so all slow-AR keys sit at +// the root of config.json (AGENTS.md §4.1). +Audio8TtsTextConfig parse_text_config(const json::Value & value) { + Audio8TtsTextConfig 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::optional_i64(value, "intermediate_size", config.dim); + if (config.intermediate_size == config.dim) { + config.intermediate_size = json::optional_i64(value, "hidden_size", config.dim); + } + config.n_head = json::require_i64(value, "n_head"); + config.n_local_heads = json::optional_i64(value, "n_local_heads", config.n_head); + if (config.n_local_heads == config.n_head) { + config.n_local_heads = json::optional_i64(value, "num_key_value_heads", config.n_head); + config.n_local_heads = json::optional_i64(value, "n_local_heads", config.n_local_heads); + } + config.head_dim = json::optional_i64(value, "head_dim", 64); + config.max_seq_len = json::require_i64(value, "max_seq_len"); + config.rope_base = json::optional_f32(value, "rope_base", config.rope_base); + config.rope_base = json::optional_f32(value, "rope_theta", 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); + config.slow_backbone = json::optional_string(value, "slow_backbone", "qwen"); + config.mamba_d_state = json::optional_i64(value, "mamba_d_state", config.mamba_d_state); + config.mamba_d_conv = json::optional_i64(value, "mamba_d_conv", config.mamba_d_conv); + config.mamba_expand = json::optional_i64(value, "mamba_expand", config.mamba_expand); + config.mamba_n_heads = json::optional_i64(value, "mamba_n_heads", config.mamba_n_heads); + config.mamba_n_groups = json::optional_i64(value, "mamba_n_groups", config.mamba_n_groups); + config.mamba_d_head = json::optional_i64(value, "mamba_d_head", config.mamba_d_head); + config.mamba_d_ssm = json::optional_i64(value, "mamba_d_ssm", config.mamba_d_ssm); + config.mamba_chunk_size = json::optional_i64(value, "mamba_chunk_size", config.mamba_chunk_size); + config.embedding_multiplier = json::optional_f32(value, "embedding_multiplier", config.embedding_multiplier); + config.lm_head_multiplier = json::optional_f32(value, "lm_head_multiplier", config.lm_head_multiplier); + config.attention_in_multiplier = json::optional_f32(value, "attention_in_multiplier", config.attention_in_multiplier); + config.attention_out_multiplier = json::optional_f32(value, "attention_out_multiplier", config.attention_out_multiplier); + config.ssm_in_multiplier = json::optional_f32(value, "ssm_in_multiplier", config.ssm_in_multiplier); + config.ssm_out_multiplier = json::optional_f32(value, "ssm_out_multiplier", config.ssm_out_multiplier); + config.key_multiplier = json::optional_f32(value, "key_multiplier", config.key_multiplier); + { + auto v = json::optional_f32_array(value, "ssm_multipliers"); + if (!v.empty()) config.ssm_multipliers = v; + } + { + auto v = json::optional_f32_array(value, "mlp_multipliers"); + if (!v.empty()) config.mlp_multipliers = v; + } + 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.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; +} + +// configuration_arktts.py — fast-AR keys carry a fast_ prefix at the flat config root; +// the fast vocabulary spans one codebook (codebook_size), not the text vocab, and the +// checkpoint ships an untied fast_output.weight regardless of the slow LM head tying. +Audio8TtsFastConfig parse_fast_config(const json::Value & value) { + Audio8TtsFastConfig config; + config.vocab_size = json::require_i64(value, "codebook_size"); + config.num_codebooks = json::require_i64(value, "num_codebooks"); + config.n_layer = json::require_i64(value, "n_fast_layer"); + config.dim = json::require_i64(value, "fast_dim"); + config.intermediate_size = json::require_i64(value, "fast_intermediate_size"); + config.n_head = json::require_i64(value, "fast_n_head"); + config.n_local_heads = json::optional_i64(value, "fast_n_local_heads", config.n_head); + config.head_dim = json::require_i64(value, "fast_head_dim"); + config.max_seq_len = json::optional_i64(value, "fast_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, "fast_tie_word_embeddings", false); + config.attention_qk_norm = json::optional_bool(value, "fast_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; +} + +// configuration_arktts.py — arktts ships one flat config.json; semantic token ids are +// named semantic_begin_id / semantic_end_id there (AGENTS.md §4.1). +Audio8TtsConfig parse_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + Audio8TtsConfig config; + config.model_type = json::optional_string(root, "model_type", ""); + if (config.model_type != "arktts") { + throw std::runtime_error("Audio8 TTS 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_begin_id"); + config.semantic_end_token_id = json::require_i64(root, "semantic_end_id"); + config.im_end_token_id = json::require_i64(root, "eos_token_id"); + config.norm_fastlayer_input = json::optional_bool(root, "norm_fastlayer_input", false); + config.text = parse_text_config(root); + config.fast = parse_fast_config(root); + config.codec.total_codebooks = config.fast.num_codebooks; + if (config.fast.dim != config.text.dim) { + throw std::runtime_error("Audio8 TTS fast dim must match text dim"); + } + if (!config.text.tie_word_embeddings) { + throw std::runtime_error("Audio8 TTS 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("Audio8 TTS semantic token range is invalid"); + } + return config; +} + +// model.safetensors stores QKV pre-packed per layer as wqkv (+ a bias row on slow +// layers only) — AGENTS.md §4.2; anchors pin that layout before graph building. +// Supports both Qwen (0.6B/1.0B: embeddings.weight, layers.*) and Falcon-H1/Mamba +// (0.1B: slow.embed_tokens.weight, slow.layers.*.mamba/self_attn). +void validate_weight_anchors(const Audio8TtsAssets & assets) { + const bool is_qwen = assets.model_weights->has_tensor("embeddings.weight"); + const bool is_mamba = assets.model_weights->has_tensor("slow.embed_tokens.weight"); + if (!is_qwen && !is_mamba) { + throw std::runtime_error("Audio8 TTS model_weights missing embeddings (expected embeddings.weight or slow.embed_tokens.weight)"); + } + assets.model_weights->require_metadata("codebook_embeddings.weight"); + if (is_qwen) { + assets.model_weights->require_metadata("layers.0.attention.wqkv.weight"); + assets.model_weights->require_metadata("layers.0.attention.wqkv.bias"); + assets.model_weights->require_metadata("layers.0.attention.wo.weight"); + } else { + assets.model_weights->require_metadata("slow.layers.0.mamba.in_proj.weight"); + assets.model_weights->require_metadata("slow.layers.0.self_attn.q_proj.weight"); + assets.model_weights->require_metadata("slow.final_layernorm.weight"); + } + assets.model_weights->require_metadata("fast_layers.0.attention.wqkv.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_audio8_tts_assets(const std::filesystem::path & model_path) { + Audio8TtsAssets assets; + assets.resources = engine::model_spec::load_resource_bundle( + model_path, + engine::model_spec::default_spec_path("audio8_tts")); + 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::audio8_tts diff --git a/src/community_models/audio8_tts/codec.cpp b/src/community_models/audio8_tts/codec.cpp new file mode 100644 index 00000000..fd43ed19 --- /dev/null +++ b/src/community_models/audio8_tts/codec.cpp @@ -0,0 +1,1182 @@ +#include "engine/community_models/audio8_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/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/transformer_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include "engine/framework/core/constant_tensor_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::audio8_tts { +namespace { + +namespace binding = engine::modules::binding; + +constexpr int64_t kCodecDim = 1024; +constexpr int64_t kCodecTransformerHeads = 16; +constexpr int64_t kCodecTransformerKVHeads = 8; +constexpr int64_t kCodecHeadDim = 64; +constexpr int64_t kCodecIntermediate = 3072; +constexpr int64_t kCodecPostIntermediateSize = 1216; +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; + int64_t kv_heads = kCodecTransformerKVHeads; + int64_t intermediate_size = kCodecIntermediate; +}; + +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("Audio8 TTS 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("Audio8 TTS 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("Audio8 TTS 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("Audio8 TTS 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("Audio8 TTS causal conv requires non-negative left padding"); + } + if (right_pad < 0) { + throw std::runtime_error("Audio8 TTS 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 auto normalized_input = core::wrap_tensor(ggml_cont(ctx.ggml, input.tensor), input.shape, input.type); + 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, + int64_t intermediate_size) { + auto gate = modules::LinearModule({kCodecDim, intermediate_size, false, GGML_PREC_F32}) + .build(ctx, input, weights.gate_proj); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule({kCodecDim, intermediate_size, false, GGML_PREC_F32}) + .build(ctx, input, weights.up_proj); + auto hidden = modules::MulModule{}.build(ctx, gate, up); + return modules::LinearModule({intermediate_size, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, hidden, weights.down_proj); +} + +core::TensorValue reshape_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t heads) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + return core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, 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, + int64_t kv_heads, + int64_t intermediate_size) { + 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, kv_heads * kCodecHeadDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.k_weight, std::nullopt}); + auto v = modules::LinearModule({kCodecDim, kv_heads * kCodecHeadDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.v_weight, std::nullopt}); + // modeling_arktts_codec.py:98-100: flash attn maps q head i to kv head i/(Hq/Hkv), + // identical to the reference repeat_interleave expansion. + q = modules::RoPEModule({kCodecHeadDim, GGML_ROPE_TYPE_NORMAL, kCodecRopeTheta}).build(ctx, reshape_heads(ctx, q, kCodecTransformerHeads), positions); + k = modules::RoPEModule({kCodecHeadDim, GGML_ROPE_TYPE_NORMAL, kCodecRopeTheta}).build(ctx, reshape_heads(ctx, k, kv_heads), positions); + v = reshape_heads(ctx, v, kv_heads); + 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, intermediate_size); + ff = modules::LayerScaleModule{}.build(ctx, ff, weights.ffn_scale); + return modules::AddModule{}.build(ctx, hidden, ff); +} + +core::TensorValue make_positions( + core::ModuleBuildContext &, + core::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 &, + core::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, + core::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, weights.kv_heads, weights.intermediate_size); + } + 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, + core::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, + core::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, + core::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({"audio8_tts.codec.after_downsample", x}); + x = build_window_transformer(ctx, constants, x, weights.pre_module, 128); + trace_outputs.push_back({"audio8_tts.codec.after_pre_module", x}); + + auto cont_x = core::wrap_tensor(ggml_cont(ctx.ggml, x.tensor), x.shape, x.type); + auto residual = cont_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, + int64_t kv_heads, + int64_t intermediate_size) { + CodecTransformerWeights out; + out.kv_heads = kv_heads; + out.intermediate_size = intermediate_size; + 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); + // arktts packs QKV as one [(n_head + 2*n_local_heads)*head_dim, dim] + // matrix (modeling_arktts_codec.py:79); split rows into q/k/v at load time. + const auto packed_qkv = source.require_tensor( + layer_prefix + ".attention.wqkv.weight", + storage_type, + {(kCodecTransformerHeads + 2 * kv_heads) * kCodecHeadDim, kCodecDim}); + const size_t qkv_row_bytes = ggml_row_size(packed_qkv.type, kCodecDim); + const size_t q_rows = static_cast(kCodecTransformerHeads) * kCodecHeadDim; + const size_t kv_rows = static_cast(kv_heads) * kCodecHeadDim; + const std::byte * qkv_data = packed_qkv.bytes.data(); + weights.attention.q_weight = store.make_tensor( + core::TensorShape::from_dims({static_cast(q_rows), kCodecDim}), + packed_qkv.type, + qkv_data, + qkv_row_bytes * q_rows); + weights.attention.k_weight = store.make_tensor( + core::TensorShape::from_dims({static_cast(kv_rows), kCodecDim}), + packed_qkv.type, + qkv_data + qkv_row_bytes * q_rows, + qkv_row_bytes * kv_rows); + weights.attention.v_weight = store.make_tensor( + core::TensorShape::from_dims({static_cast(kv_rows), kCodecDim}), + packed_qkv.type, + qkv_data + qkv_row_bytes * (q_rows + kv_rows), + qkv_row_bytes * kv_rows); + 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, {intermediate_size, kCodecDim}); + weights.feed_forward.down_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w2.weight", storage_type, {kCodecDim, intermediate_size}); + weights.feed_forward.up_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w3.weight", storage_type, {intermediate_size, 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 Audio8TtsAssets & 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, "Audio8 TTS 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, kCodecTransformerHeads, kCodecIntermediate); + } + 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, kCodecTransformerHeads, kCodecIntermediate); + 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, kCodecTransformerKVHeads, kCodecPostIntermediateSize); + 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_, "Audio8 TTS 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 Audio8 TTS codec decode graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "audio8_tts.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 Audio8 TTS 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 Audio8TtsCodes & 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 << "Audio8 TTS 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("Audio8 TTS 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("Audio8 TTS 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_; + core::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_, "Audio8 TTS 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 Audio8 TTS codec encode graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "audio8_tts.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({"audio8_tts.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 Audio8 TTS 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); + } + + Audio8TtsCodes 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("Audio8 TTS 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("Audio8 TTS 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)); + } + } + Audio8TtsCodes 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( + "audio8_tts.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_; + core::ConstantTensorCache constants_; +}; + +} // namespace + +class Audio8TtsCodecRuntime::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); + } + + Audio8TtsCodes 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 Audio8TtsCodes & 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_; +}; + +Audio8TtsCodecRuntime::Audio8TtsCodecRuntime( + 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)) {} + +Audio8TtsCodecRuntime::~Audio8TtsCodecRuntime() = default; + +Audio8TtsCodes Audio8TtsCodecRuntime::encode_reference(const runtime::AudioBuffer & audio) { + return impl_->encode_reference(audio); +} + +runtime::AudioBuffer Audio8TtsCodecRuntime::decode(const Audio8TtsCodes & codes) { + return impl_->decode(codes); +} + +void Audio8TtsCodecRuntime::release_encode_graph() { + impl_->release_encode_graph(); +} + +void Audio8TtsCodecRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::audio8_tts diff --git a/src/community_models/audio8_tts/generator.cpp b/src/community_models/audio8_tts/generator.cpp new file mode 100644 index 00000000..fbd539ff --- /dev/null +++ b/src/community_models/audio8_tts/generator.cpp @@ -0,0 +1,75 @@ +#include "engine/community_models/audio8_tts/generator.h" + +#include "engine/framework/debug/profiler.h" + +#include +#include +#include + +namespace engine::models::audio8_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +} // namespace + +Audio8TtsGenerator::Audio8TtsGenerator( + 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("Audio8 TTS generator requires assets, AR runtime, and codec runtime"); + } +} + +Audio8TtsGenerator::~Audio8TtsGenerator() = default; + +Audio8TtsCodes Audio8TtsGenerator::encode_reference(const runtime::AudioBuffer & audio) { + auto codes = codec_->encode_reference(audio); + codec_->release_encode_graph(); + return codes; +} + +Audio8TtsGenerationResult Audio8TtsGenerator::generate( + const Audio8TtsRequest & request, + const std::vector & reference_codes, + const std::optional & previous_turn, + bool mem_saver) { + engine::debug::trace_log_scalar("audio8_tts.request.has_reference", !request.references.empty()); + engine::debug::trace_log_scalar("audio8_tts.request.reference_count", static_cast(request.references.size())); + engine::debug::trace_log_scalar("audio8_tts.request.text_chars", static_cast(request.text.size())); + engine::debug::trace_log_scalar("audio8_tts.request.has_previous_turn", previous_turn.has_value()); + engine::debug::trace_log_scalar("audio8_tts.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( + "audio8_tts.prompt_build_ms", + engine::debug::elapsed_ms(prompt_start, Clock::now())); + + const auto ar_start = Clock::now(); + Audio8TtsGenerationResult result; + result.codes = ar_->generate(prompt, request.generation); + engine::debug::trace_log_scalar("audio8_tts.generated.frames", result.codes.frames); + engine::debug::trace_log_scalar("audio8_tts.generated.codebooks", result.codes.codebooks); + engine::debug::timing_log_scalar( + "audio8_tts.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( + "audio8_tts.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::audio8_tts diff --git a/src/community_models/audio8_tts/prompt_builder.cpp b/src/community_models/audio8_tts/prompt_builder.cpp new file mode 100644 index 00000000..944571cd --- /dev/null +++ b/src/community_models/audio8_tts/prompt_builder.cpp @@ -0,0 +1,314 @@ +#include "engine/community_models/audio8_tts/prompt_builder.h" + +#include +#include +#include + +namespace engine::models::audio8_tts { +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 Audio8TtsCodes * codes = nullptr; +}; + +// Golden parity with Audio8_TTS/audio8_tts_data.py + onnx_runtime/arktts_runtime/prompt.py +// CJK ranges from _CJK_RANGES, line-break set from _LINE_BREAK_RE. +bool is_cjk(uint32_t cp) { + return (0x1100 <= cp && cp <= 0x11FF) || (0x2E80 <= cp && cp <= 0x2FDF) || + (0x3000 <= cp && cp <= 0x303F) || (0x3040 <= cp && cp <= 0x30FF) || + (0x3100 <= cp && cp <= 0x31FF) || (0x3400 <= cp && cp <= 0x4DBF) || + (0x4E00 <= cp && cp <= 0x9FFF) || (0xA960 <= cp && cp <= 0xA97F) || + (0xAC00 <= cp && cp <= 0xD7A3) || (0xD7B0 <= cp && cp <= 0xD7FF) || + (0xF900 <= cp && cp <= 0xFAFF) || (0xFE30 <= cp && cp <= 0xFE4F) || + (0xFF01 <= cp && cp <= 0xFF9F) || (0x20000 <= cp && cp <= 0x2FA1F); +} + +bool is_line_break(uint32_t cp) { + return cp == 0x0D || cp == 0x0A || cp == 0x0B || cp == 0x0C || cp == 0x1C || + cp == 0x1D || cp == 0x1E || cp == 0x85 || cp == 0x2028 || cp == 0x2029; +} + +bool is_space(uint32_t cp) { + if (cp == 0x20 || cp == 0x09 || cp == 0x0A || cp == 0x0B || cp == 0x0C || + cp == 0x0D || cp == 0x1C || cp == 0x1D || cp == 0x1E || cp == 0x85 || + cp == 0xA0 || cp == 0x1680 || cp == 0x2028 || cp == 0x2029 || + cp == 0x202F || cp == 0x205F || cp == 0x3000) { + return true; + } + if (0x2000 <= cp && cp <= 0x200A) { + return true; + } + return false; +} + +bool is_control_category(uint32_t cp) { + if (cp <= 0x1F) { + return !is_space(cp); + } + if (cp == 0x7F) { + return true; + } + if (0x80 <= cp && cp <= 0x9F) { + return true; + } + if ((0x200B <= cp && cp <= 0x200F) || (0x202A <= cp && cp <= 0x202E) || + (0x2060 <= cp && cp <= 0x206F) || cp == 0xFEFF) { + return true; + } + if (0xD800 <= cp && cp <= 0xDFFF) { + return true; + } + return false; +} + +std::vector decode_utf8(const std::string & s) { + std::vector out; + out.reserve(s.size()); + for (size_t i = 0; i < s.size();) { + unsigned char c = static_cast(s[i]); + uint32_t cp = 0; + size_t len = 0; + if (c < 0x80) { + cp = c; + len = 1; + } else if ((c & 0xE0) == 0xC0) { + cp = c & 0x1F; + len = 2; + } else if ((c & 0xF0) == 0xE0) { + cp = c & 0x0F; + len = 3; + } else if ((c & 0xF8) == 0xF0) { + cp = c & 0x07; + len = 4; + } else { + cp = 0xFFFD; + len = 1; + } + if (i + len > s.size()) { + cp = 0xFFFD; + len = 1; + } else { + for (size_t j = 1; j < len; ++j) { + unsigned char cc = static_cast(s[i + j]); + if ((cc & 0xC0) != 0x80) { + cp = 0xFFFD; + len = j; + break; + } + cp = (cp << 6) | (cc & 0x3F); + } + } + out.push_back(cp); + i += len; + } + return out; +} + +std::string encode_utf8(const std::vector & cps) { + std::string out; + out.reserve(cps.size() * 2); + for (uint32_t cp : cps) { + if (cp < 0x80) { + out.push_back(static_cast(cp)); + } else if (cp < 0x800) { + out.push_back(static_cast(0xC0 | (cp >> 6))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } else if (cp < 0x10000) { + out.push_back(static_cast(0xE0 | (cp >> 12))); + out.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } else { + out.push_back(static_cast(0xF0 | (cp >> 18))); + out.push_back(static_cast(0x80 | ((cp >> 12) & 0x3F))); + out.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } + } + return out; +} + +std::string normalize_whitespace(const std::string & text) { + auto cps = decode_utf8(text); + std::vector out; + out.reserve(cps.size()); + size_t i = 0; + while (i < cps.size()) { + if (is_space(cps[i])) { + size_t j = i; + bool has_line_break = false; + while (j < cps.size() && is_space(cps[j])) { + if (is_line_break(cps[j])) { + has_line_break = true; + } + ++j; + } + uint32_t left = (i > 0) ? cps[i - 1] : 0; + uint32_t right = (j < cps.size()) ? cps[j] : 0; + bool left_cjk = left != 0 && is_cjk(left); + bool right_cjk = right != 0 && is_cjk(right); + if (!(has_line_break && left_cjk && right_cjk)) { + out.push_back(0x20); + } + i = j; + } else { + out.push_back(cps[i]); + ++i; + } + } + size_t start = 0; + while (start < out.size() && out[start] == 0x20) { + ++start; + } + size_t end = out.size(); + while (end > start && out[end - 1] == 0x20) { + --end; + } + std::vector trimmed(out.begin() + static_cast(start), + out.begin() + static_cast(end)); + return encode_utf8(trimmed); +} + +std::string clean_text(const std::string & value, const char * field_name = "text") { + auto cps = decode_utf8(value); + std::vector filtered; + filtered.reserve(cps.size()); + for (uint32_t cp : cps) { + if (is_space(cp)) { + filtered.push_back(cp); + } else if (is_control_category(cp)) { + continue; + } else { + filtered.push_back(cp); + } + } + std::string intermediate = encode_utf8(filtered); + std::string cleaned = normalize_whitespace(intermediate); + if (cleaned.empty()) { + throw std::runtime_error(std::string(field_name) + " must not be empty"); + } + return cleaned; +} + +std::string reference_text_with_speakers(const std::string & text, int64_t speaker) { + std::string cleaned = clean_text(text, "reference_text"); + static const std::regex speaker_re(R"(<\|speaker:\d+\|>)"); + if (std::regex_search(cleaned, speaker_re)) { + return cleaned; + } + return "<|speaker:" + std::to_string(speaker) + "|>" + cleaned; +} + +void append_code_span( + std::vector & row0, + std::vector & spans, + const Audio8TtsTextTokenizer & tokenizer, + const Audio8TtsCodes & codes, + int64_t expected_codebooks) { + if (codes.codebooks != expected_codebooks) { + throw std::runtime_error("Audio8 TTS 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 + +Audio8TtsPromptBuilder::Audio8TtsPromptBuilder( + std::shared_ptr assets, + Audio8TtsTextTokenizer tokenizer) + : assets_(std::move(assets)), + tokenizer_(std::move(tokenizer)) { + if (assets_ == nullptr) { + throw std::runtime_error("Audio8 TTS prompt builder requires assets"); + } +} + +Audio8TtsPrompt Audio8TtsPromptBuilder::build( + const Audio8TtsRequest & request, + const std::vector & reference_codes, + const std::optional & previous_turn) const { + if (request.text.empty()) { + throw std::runtime_error("Audio8 TTS request text must not be empty"); + } + const int64_t rows = assets_->config.fast.num_codebooks + 1; + if (rows <= 1) { + throw std::runtime_error("Audio8 TTS prompt rows are invalid"); + } + + std::vector row0; + std::vector code_spans; + if (!request.references.empty()) { + if (reference_codes.size() != request.references.size()) { + throw std::runtime_error("Audio8 TTS reference request requires one encoded code tensor per reference"); + } + 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")); + for (size_t index = 0; index < request.references.size(); ++index) { + if (index != 0) { + append_tokens(row0, tokenizer_.encode("\n")); + } + append_tokens( + row0, + tokenizer_.encode(reference_text_with_speakers( + request.references[index].text, + static_cast(index)))); + } + append_tokens(row0, tokenizer_.encode("\n\nSpeech:\n")); + for (const auto & codes : reference_codes) { + append_code_span(row0, code_spans, tokenizer_, 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(clean_text(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(clean_text(request.text))); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + append_tokens(row0, tokenizer_.encode("<|im_start|>assistant\n<|voice|>")); + + Audio8TtsPrompt 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("Audio8 TTS 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("Audio8 TTS 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::audio8_tts diff --git a/src/community_models/audio8_tts/session.cpp b/src/community_models/audio8_tts/session.cpp new file mode 100644 index 00000000..f3d3e059 --- /dev/null +++ b/src/community_models/audio8_tts/session.cpp @@ -0,0 +1,545 @@ +#include "engine/community_models/audio8_tts/session.h" + +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/io/json.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/session.h" +#include "engine/framework/runtime/spec_backed_model.h" +#include "engine/framework/text/chunking.h" +#include "engine/community_models/audio8_tts/ar.h" +#include "engine/community_models/audio8_tts/codec.h" +#include "engine/community_models/audio8_tts/generator.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::audio8_tts { +namespace { + +using Clock = std::chrono::steady_clock; +namespace fs = std::filesystem; + +constexpr std::string_view kFamily = "audio8_tts"; +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"; +constexpr const char * kMultiReferenceCondOption = "multi_reference_cond"; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Audio8 TTS session requires assets"); + } + return assets; +} + +std::shared_ptr require_contract( + std::shared_ptr contract) { + if (contract == nullptr) { + throw std::runtime_error("Audio8 TTS session requires a model contract"); + } + return contract; +} + +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, {"audio8_tts.mem_saver", "mem_saver"})) { + return runtime::parse_bool_option(*value, "audio8_tts.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, + {"audio8_tts.reference_cache_slots", "reference_cache_slots"}) + .value_or(kDefaultReferenceCacheSlots); + if (slots < 0) { + throw std::runtime_error("audio8_tts.reference_cache_slots must be non-negative"); + } + if (static_cast(slots) > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("audio8_tts.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; +} + +Audio8TtsGenerationOptions generation_options_from_request(const runtime::TaskRequest & request) { + Audio8TtsGenerationOptions options; + if (const auto value = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"})) { + if (*value < 0) { + throw std::runtime_error("Audio8 TTS 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("Audio8 TTS max_new_tokens must be positive after default resolution"); + } + if (options.text_chunk_size <= 0) { + throw std::runtime_error("Audio8 TTS text_chunk_size must be positive"); + } + if (!(options.top_p > 0.0F && options.top_p <= 1.0F)) { + throw std::runtime_error("Audio8 TTS top_p must be in (0, 1]"); + } + if (options.top_k <= 0) { + throw std::runtime_error("Audio8 TTS top_k must be positive"); + } + if (!(options.temperature > 0.0F && options.temperature < 2.0F)) { + throw std::runtime_error("Audio8 TTS 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( + "Audio8 TTS cached_voice_id must be 1-255 characters"); + } + for (const unsigned char ch : id) { + if (!valid_reference_id_char(ch)) { + throw std::runtime_error( + "Audio8 TTS 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)}; +} + +std::vector load_saved_references( + const Audio8TtsAssets & 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, + "Audio8 TTS cached voice reference"); + const auto audio_files = collect_saved_reference_audio_files(reference_dir); + if (audio_files.empty()) { + throw std::runtime_error( + "Audio8 TTS cached_voice_id '" + reference_id + + "' requires at least one WAV reference with a matching .lab file under " + + reference_dir.string()); + } + std::vector references; + references.reserve(audio_files.size()); + for (const auto & audio_file : audio_files) { + auto lab_path = audio_file; + lab_path.replace_extension(".lab"); + references.push_back(Audio8TtsReference{ + read_saved_reference_audio(audio_file), + engine::io::read_text_file(lab_path), + reference_id + "\n" + audio_file.lexically_normal().string()}); + } + return references; +} + +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::vector references_from_voice( + const Audio8TtsAssets & assets, + const std::optional & voice, + const std::unordered_map & options, + const char * role) { + if (!has_reference_selector(voice)) { + return {}; + } + 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 {Audio8TtsReference{ + speaker.audio, + *reference_text, + reference_cache_id_from_voice(voice)}}; + } + return load_saved_references(assets, *speaker.cached_voice_id); +} + +std::vector multi_reference_cond_from_options( + const std::unordered_map & options) { + const auto value = runtime::find_option(options, {kMultiReferenceCondOption}); + if (!value.has_value()) { + return {}; + } + const auto root = engine::io::json::parse(*value); + if (!root.is_array()) { + throw std::runtime_error("Audio8 TTS multi_reference_cond must be a JSON array"); + } + std::vector references; + references.reserve(root.as_array().size()); + for (const auto & item : root.as_array()) { + if (!item.is_object()) { + throw std::runtime_error("Audio8 TTS multi_reference_cond entries must be objects"); + } + const auto audio_path = engine::io::json::require_string(item, "audio"); + const auto text = engine::io::json::require_string(item, "text"); + if (audio_path.empty()) { + throw std::runtime_error("Audio8 TTS multi_reference_cond audio path must not be empty"); + } + if (text.empty()) { + throw std::runtime_error("Audio8 TTS multi_reference_cond text must not be empty"); + } + references.push_back(Audio8TtsReference{ + read_saved_reference_audio(audio_path), + text, + {}}); + } + return references; +} + +} // namespace + +Audio8TtsSession::Audio8TtsSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + contract_(require_contract(std::move(contract))), + reference_cache_(resolve_reference_cache_slots(this->options())) { + runtime::validate_spec_backed_session_options(this->options(), *contract_, kFamily, "Audio8 TTS"); + // Voice cloning reuses the TTS path; references only switch the prompt form + // (processing_arktts.py:_prompt_segments). + if ((task_.task != runtime::VoiceTaskKind::Tts && + task_.task != runtime::VoiceTaskKind::VoiceCloning) || + task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Audio8 TTS supports offline TTS and voice cloning sessions only"); + } + const auto ar_weight_type = + option_weight_type(options, "audio8_tts.weight_type", assets::TensorStorageType::Native); + const auto codec_weight_type = + option_weight_type(options, "audio8_tts.codec_weight_type", assets::TensorStorageType::Native); + validate_ar_weight_storage(ar_weight_type, "audio8_tts.weight_type"); + validate_codec_weight_storage(codec_weight_type, "audio8_tts.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, {"audio8_tts.ar_graph_arena_mb"}, kDefaultArGraphArenaBytes), + runtime::parse_size_mb_option(options.options, {"audio8_tts.ar_weight_context_mb"}, kDefaultArWeightContextBytes), + ar_weight_type); + auto codec = std::make_unique( + assets_, + options.backend, + threads, + runtime::parse_size_mb_option(options.options, {"audio8_tts.codec_graph_arena_mb"}, kDefaultCodecGraphArenaBytes), + runtime::parse_size_mb_option(options.options, {"audio8_tts.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(); +} + +Audio8TtsSession::~Audio8TtsSession() = default; + +std::string Audio8TtsSession::family() const { + return "audio8_tts"; +} + +runtime::VoiceTaskKind Audio8TtsSession::task_kind() const { + return task_.task; +} + +runtime::RunMode Audio8TtsSession::run_mode() const { + return task_.mode; +} + +bool Audio8TtsSession::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 Audio8TtsSession::prepare(const runtime::SessionPreparationRequest & request) { + defaults_.reset(); + Audio8TtsRequest 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("Audio8 TTS max_new_tokens must be non-negative"); + } + if (*value > 0) { + defaults.generation.max_new_tokens = *value; + } + } + if (auto references = references_from_voice(*assets_, request.voice, request.options, "Audio8 TTS prepare"); + !references.empty()) { + defaults.references = std::move(references); + has_defaults = true; + } else if (auto references = multi_reference_cond_from_options(request.options); + !references.empty()) { + defaults.references = std::move(references); + has_defaults = true; + } + if (has_defaults) { + defaults_ = std::move(defaults); + } + mark_prepared(); +} + +Audio8TtsRequest Audio8TtsSession::make_request(const runtime::TaskRequest & request) const { + Audio8TtsRequest out = defaults_.value_or(Audio8TtsRequest{}); + if (request.text_input.has_value()) { + out.text = request.text_input->text; + } + out.generation = generation_options_from_request(request); + if (auto references = references_from_voice(*assets_, request.voice, request.options, "Audio8 TTS request"); + !references.empty()) { + out.references = std::move(references); + } else if (auto references = multi_reference_cond_from_options(request.options); + !references.empty()) { + out.references = std::move(references); + } else if (request.text_input.has_value()) { + out.references.clear(); + } + if (task_.task == runtime::VoiceTaskKind::VoiceCloning && out.references.empty()) { + throw std::runtime_error( + "Audio8 TTS voice cloning requires a speaker reference: --voice-ref with " + "--reference-text, a cached --voice-id, or the multi_reference_cond option"); + } + if (out.text.empty()) { + throw std::runtime_error("Audio8 TTS request text must not be empty"); + } + return out; +} + +const Audio8TtsCodes & Audio8TtsSession::resolve_reference_codes(const Audio8TtsReference & reference) { + ReferenceCacheKey key; + key.source_id = reference.cache_id; + if (reference.cache_id.empty() && !reference.audio.has_value()) { + throw std::runtime_error("Audio8 TTS 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("audio8_tts.reference_cache.hit", 1); + engine::debug::trace_log_scalar("audio8_tts.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("audio8_tts.reference_cache.entries", static_cast(reference_cache_.size())); + engine::debug::trace_log_scalar("audio8_tts.reference_cache.evicted", 0); + return cached->codes; + } + if (!reference.audio.has_value()) { + throw std::runtime_error("Audio8 TTS 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("audio8_tts.reference.frames", entry.codes.frames); + engine::debug::trace_log_scalar("audio8_tts.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("audio8_tts.reference_cache.hit", 0); + engine::debug::trace_log_scalar("audio8_tts.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("audio8_tts.reference_cache.entries", static_cast(reference_cache_.size())); + engine::debug::trace_log_scalar("audio8_tts.reference_cache.evicted", will_evict ? 1 : 0); + engine::debug::timing_log_scalar("audio8_tts.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("Audio8 TTS reference cache insert failed"); + } + return cached->codes; +} + +runtime::TaskResult Audio8TtsSession::run(const runtime::TaskRequest & request) { + require_prepared("Audio8 TTS 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("audio8_tts.text_chunk_size", request_options.text_chunk_size); + engine::debug::trace_log_scalar("audio8_tts.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + engine::debug::trace_log_scalar("audio8_tts.text_chunk_count", static_cast(chunk_requests.size())); + + runtime::AudioBuffer merged_audio; + std::vector reference_codes; + 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 arktts_request = make_request(chunk_request); + if (!arktts_request.references.empty() && reference_codes.empty()) { + reference_codes.reserve(arktts_request.references.size()); + for (const auto & reference : arktts_request.references) { + reference_codes.push_back(resolve_reference_codes(reference)); + } + } + auto generated = generator_->generate(arktts_request, reference_codes, previous_turn, mem_saver); + runtime::append_audio_buffer(merged_audio, generated.audio); + if (chunk_requests.size() > 1) { + previous_turn = Audio8TtsConversationTurn{arktts_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; +} + +// Spec-backed loader entry point; mirrors glm_tts/session.cpp:make_glm_tts_loader. +std::shared_ptr make_audio8_tts_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = std::string(kFamily); + config.load_assets = load_audio8_tts_assets; + config.create_session = []( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique( + task, + options, + std::move(assets), + std::move(contract)); + }; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::models::audio8_tts diff --git a/src/community_models/audio8_tts/tokenizer_text.cpp b/src/community_models/audio8_tts/tokenizer_text.cpp new file mode 100644 index 00000000..f8df9a9d --- /dev/null +++ b/src/community_models/audio8_tts/tokenizer_text.cpp @@ -0,0 +1,69 @@ +#include "engine/community_models/audio8_tts/tokenizer_text.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::models::audio8_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("Audio8 TTS tokenizer missing token: " + token); + } + return *id; +} + +} // namespace + +struct Audio8TtsTextTokenizer::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; +}; + +Audio8TtsTextTokenizer::Audio8TtsTextTokenizer(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Audio8 TTS text tokenizer requires assets"); + } + impl_ = std::make_shared(std::move(assets)); +} + +std::vector Audio8TtsTextTokenizer::encode(const std::string & text) const { + return impl_->tokenizer.encode(text, true); +} + +int32_t Audio8TtsTextTokenizer::token_id(const std::string & token) const { + return require_token_id(impl_->tokenizer, token); +} + +int32_t Audio8TtsTextTokenizer::im_end_id() const noexcept { + return impl_->im_end; +} + +int32_t Audio8TtsTextTokenizer::semantic_begin_id() const noexcept { + return impl_->semantic_begin; +} + +int32_t Audio8TtsTextTokenizer::semantic_end_id() const noexcept { + return impl_->semantic_end; +} + +} // namespace engine::models::audio8_tts diff --git a/tools/community_models/convert_audio8_tts.py b/tools/community_models/convert_audio8_tts.py new file mode 100644 index 00000000..a7bcc4c9 --- /dev/null +++ b/tools/community_models/convert_audio8_tts.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Convert Audio8 TTS safetensors checkpoints to audio.cpp GGUF packages. + +Produces one self-contained GGUF per precision with two tensor namespaces: + model_weights.* — DualAR slow/fast transformer weights (model.safetensors) + codec_weights.* — arktts neural codec weights (converted codec.safetensors) + +and embeds config.json / tokenizer_config.json / tokenizer.json plus the +audio8_tts package spec from the repository, giving a standalone model file +that audiocpp_cli / audiocpp_server load with --family audio8_tts. + +The tool deliberately does not download files or read PyTorch checkpoints. +Point --model-dir at a snapshot that already contains the HF weights plus the +codec.safetensors artifact produced by convert_audio8_tts_codec.py. + +Examples: + # 16-bit reference package (AR stays BF16, codec downcasts F32 -> BF16) + python3 tools/community_models/convert_audio8_tts.py \ + --model-dir /models/Audio8-TTS-Preview-0.6b \ + --converter build/bin/audiocpp_gguf --type bf16 + + # default Q8_0 package like fish_audio ships + python3 tools/community_models/convert_audio8_tts.py \ + --model-dir /models/Audio8-TTS-Preview-0.6b \ + --converter build/bin/audiocpp_gguf --type q8_0 + +Codec conv stacks are sensitive to quantization. If a Q8_0 package drifts, +reconvert with mixed storage by keeping the codec namespace at 16 bit: +append "--keep-type", "codec_weights*=bf16" to a Q8_0 conversion. +""" +import argparse +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SPEC = REPO_ROOT / "model_specs" / "audio8_tts.json" +FAMILY = "audio8_tts" + + +def convert(converter: Path, model_dir: Path, codec: Path, output: Path, + quant_type: str, overwrite: bool) -> None: + command = [ + str(converter), + "--input", f"model_weights={model_dir / 'model.safetensors'}", + "--input", f"codec_weights={codec}", + "--root", str(model_dir), + "--family", FAMILY, + "--model-spec", str(SPEC), + "--type", quant_type, + "--output", str(output), + ] + if overwrite: + command.append("--overwrite") + print("+", " ".join(command)) + subprocess.run(command, check=True) + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-dir", type=Path, required=True, + help="HF snapshot directory with config.json, tokenizer files, " + "model.safetensors, and codec.safetensors") + parser.add_argument("--codec-safetensors", type=Path, + help="codec.safetensors override " + "(default: /codec.safetensors)") + parser.add_argument("--converter", type=Path, required=True, + help="path to the audiocpp_gguf binary") + parser.add_argument("--output-dir", type=Path, + default=Path("gguf-out") / "audio8-tts", + help="directory for the produced GGUF files") + parser.add_argument("--name", default="audio8-tts-preview-0.6b", + help="GGUF base name") + parser.add_argument("--type", default="bf16", + choices=["orig", "f16", "bf16", "q8_0", "q2_k", "q3_k", + "q4_k", "q5_k", "q6_k"], + help="GGUF storage type (default bf16)") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + converter = args.converter.resolve() + if not converter.is_file(): + raise SystemExit(f"converter not found: {converter} (build target audiocpp_gguf)") + if not SPEC.is_file(): + raise SystemExit(f"package spec not found: {SPEC}") + model_dir = args.model_dir.resolve() + if not (model_dir / "model.safetensors").is_file(): + raise SystemExit(f"model.safetensors not found in {model_dir}") + codec = (args.codec_safetensors or model_dir / "codec.safetensors").resolve() + if not codec.is_file(): + raise SystemExit(f"codec safetensors not found: {codec} " + "(run convert_audio8_tts_codec.py first)") + + args.output_dir.mkdir(parents=True, exist_ok=True) + output = args.output_dir / f"{args.name}-{args.type}.gguf" + try: + convert(converter, model_dir, codec, output, args.type, args.overwrite) + except subprocess.CalledProcessError: + sys.exit(1) + + print(f"\nDone. Load it with:\n" + f" audiocpp_cli --task tts --family {FAMILY} \\\n" + f" --model {output} --text \"...\" --out out.wav") + + +if __name__ == "__main__": + main() diff --git a/tools/community_models/convert_audio8_tts_codec.py b/tools/community_models/convert_audio8_tts_codec.py new file mode 100644 index 00000000..e1457d36 --- /dev/null +++ b/tools/community_models/convert_audio8_tts_codec.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""Convert the Audio8 TTS Preview codec.pth checkpoint to safetensors. + +The official checkpoint stores raw PyTorch module state-dict entries where the +encoder/decoder causal convolutions carry new-style weight-norm +parametrizations (`*.parametrizations.weight.original{0,1}`) and the vector +quantizer projections carry legacy weight-norm parameters (`*.weight_g/v`). +audio.cpp binds fused plain weights (`*.conv.weight`, `*.in_proj.weight`, +...), mirroring the forward math of `modeling_arktts_codec.py` +(`weight = g * v / ||v||`, norm over all but the output-channel axis). + +The converter is torch-free: it reads the zipfile checkpoint with a stub +unpickler, materializes tensors with numpy, fuses both weight-norm variants, +keeps every other entry verbatim, validates structural anchors from +`modeling_arktts_codec.py`, and writes a plain safetensors file that +`assets.cpp` loads as the `codec_weights` tensor source. + +Usage: + python3 tools/community_models/convert_audio8_tts_codec.py \ + /path/to/codec.pth /path/to/codec.safetensors [--overwrite] +""" + +from __future__ import annotations + +import argparse +import io +import pickle +import zipfile +from pathlib import Path + +import numpy as np +from safetensors.numpy import save_file + +# torch storage class name -> numpy dtype +STORAGE_DTYPES = { + "FloatStorage": np.float32, + "DoubleStorage": np.float64, + "HalfStorage": np.float16, + "LongStorage": np.int64, + "IntStorage": np.int32, + "ShortStorage": np.int16, + "CharStorage": np.int8, + "ByteStorage": np.uint8, + "BoolStorage": np.bool_, +} + +DTYPE_NAMES = { + "FloatStorage": "f32", + "DoubleStorage": "f64", + "HalfStorage": "f16", + "BFloat16Storage": "bf16->f32", + "LongStorage": "i64", +} + + +class TensorRef: + """Lazy reference to one tensor inside the checkpoint zip.""" + + __slots__ = ("storage_key", "dtype_name", "offset", "size", "stride") + + def __init__(self, storage_key, dtype_name, offset, size, stride): + self.storage_key = storage_key + self.dtype_name = dtype_name + self.offset = offset + self.size = list(size) + self.stride = list(stride) + + +class CheckpointUnpickler(pickle.Unpickler): + """Stub unpickler capturing state-dict structure without torch.""" + + def find_class(self, module, name): + if module == "torch" and (name in STORAGE_DTYPES or name == "BFloat16Storage"): + return type(name, (), {"__name__": name}) + if module == "torch._utils": + if name.startswith("_rebuild_tensor"): + def rebuild(storage, offset=0, size=None, stride=None, *_rest): + storage_type, key = storage[1], storage[2] + dtype_name = getattr(storage_type, "__name__", "UnknownStorage") + if stride is None: + stride = [1] * len(size or []) + return TensorRef(key, dtype_name, offset, size, stride) + return rebuild + if name == "_rebuild_parameter": + return lambda data, *_rest: data + if module == "collections" and name == "OrderedDict": + return dict + return type(name, (), {"__name__": name}) + + def persistent_load(self, pid): + # pid = ("storage", storage_type, key, location, numel); keep whole id + return ("storage",) + tuple(pid[1:]) + + +def load_state_dict(path: Path) -> dict: + zf = zipfile.ZipFile(path) + pkl_names = [n for n in zf.namelist() if n.endswith(".pkl")] + if len(pkl_names) != 1: + raise SystemExit(f"expected exactly one data.pkl, found {pkl_names}") + root = pkl_names[0].rsplit("/", 1)[0] + + with zf.open(pkl_names[0]) as fh: + obj = CheckpointUnpickler(io.BytesIO(fh.read())).load() + if not isinstance(obj, dict): + raise SystemExit("checkpoint pickle did not contain a dict") + + def materialize(ref: TensorRef) -> np.ndarray: + member = f"{root}/data/{ref.storage_key}" + if ref.dtype_name == "BFloat16Storage": + raw = np.frombuffer(zf.read(member), dtype=np.uint16) + flat = raw[ref.offset : ref.offset + int(np.prod(ref.size))].astype(np.uint32) + values = (flat << 16).view(np.float32) + else: + dtype = STORAGE_DTYPES.get(ref.dtype_name) + if dtype is None: + raise SystemExit(f"unsupported storage type: {ref.dtype_name}") + flat = np.frombuffer(zf.read(member), dtype=dtype) + span = int(np.prod(ref.size)) + last = ref.offset + sum( + (s - 1) * st for s, st in zip(ref.size, ref.stride) + ) + 1 + values = flat[ref.offset : max(last, ref.offset + span)] + strided = np.lib.stride_tricks.as_strided( + values, + shape=tuple(ref.size), + strides=[s * values.itemsize for s in ref.stride], + ) + return np.ascontiguousarray(strided) + + return {name: materialize(ref) if isinstance(ref, TensorRef) else ref for name, ref in obj.items()} + + +def fuse_weight_norm(tensors: dict) -> dict: + """Fuse g*v/||v|| pairs into plain weights; return a new dict.""" + out: dict = {} + for name, value in tensors.items(): + if name.endswith(".parametrizations.weight.original0"): + base = name[: -len(".parametrizations.weight.original0")] + magnitude, direction = value, tensors[f"{base}.parametrizations.weight.original1"] + norm = np.sqrt( + np.sum(direction.astype(np.float32) ** 2, axis=tuple(range(1, direction.ndim)), keepdims=True) + ) + fused = magnitude.astype(np.float32) * direction.astype(np.float32) / norm + out[f"{base}.weight"] = fused + elif name.endswith(".parametrizations.weight.original1"): + continue + elif name.endswith(".weight_g"): + base = name[: -len("_g")] + magnitude, direction = value, tensors[f"{base}_v"] + norm = np.sqrt( + np.sum(direction.astype(np.float32) ** 2, axis=tuple(range(1, direction.ndim)), keepdims=True) + ) + out[base] = magnitude.astype(np.float32) * direction.astype(np.float32) / norm + elif name.endswith(".weight_v"): + continue + else: + out[name] = value + return out + + +def validate(tensors: dict) -> None: + def expect(name, shape): + found = tensors.get(name) + if found is None: + raise SystemExit(f"missing anchor tensor: {name}") + if list(found.shape) != shape: + raise SystemExit(f"anchor {name} has shape {list(found.shape)}, expected {shape}") + + # modeling_arktts_codec.py structural anchors + expect("quantizer.semantic_quantizer.quantizers.0.codebook.weight", [4096, 8]) + expect("quantizer.semantic_quantizer.quantizers.0.in_proj.weight", [8, 1024, 1]) + expect("quantizer.semantic_quantizer.quantizers.0.out_proj.weight", [1024, 8, 1]) + for book in range(9): + expect(f"quantizer.quantizer.quantizers.{book}.codebook.weight", [1024, 8]) + for layer in range(8): + expect(f"quantizer.pre_module.layers.{layer}.attention.wqkv.weight", [3072, 1024]) + expect(f"quantizer.post_module.layers.{layer}.attention.wqkv.weight", [2048, 1024]) + expect("quantizer.pre_module.norm.weight", [1024]) + expect("decoder.model.0.conv.weight", [1536, 1024, 7]) + expect("encoder.block.0.conv.weight", [64, 1, 7]) + + leftovers = [ + n for n in tensors + if "parametrizations" in n or n.endswith(".weight_g") or n.endswith(".weight_v") + ] + if leftovers: + raise SystemExit(f"unfused weight-norm entries remain: {leftovers[:5]}") + + for name, tensor in tensors.items(): + if np.issubdtype(tensor.dtype, np.floating) and not np.isfinite(tensor.astype(np.float64)).all(): + raise SystemExit(f"non-finite values in {name}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("input", type=Path, help="Audio8 TTS codec.pth") + parser.add_argument("output", type=Path, help="output codec.safetensors") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + if args.output.exists() and not args.overwrite: + raise SystemExit(f"output exists (pass --overwrite): {args.output}") + + print(f"loading {args.input} ...") + raw = load_state_dict(args.input) + print(f" {len(raw)} state-dict entries") + + tensors = fuse_weight_norm(raw) + fused_pairs = len(tensors) - len(raw) + validate(tensors) + + args.output.parent.mkdir(parents=True, exist_ok=True) + save_file( + tensors, + str(args.output), + metadata={ + "format": "pt", + "source": "Audio8/Audio8-TTS-Preview-0.6b", + "checkpoint": args.input.name, + }, + ) + + dtypes: dict[str, int] = {} + labels = { + "> 20} MiB)") + print(f"dtypes: {summary}; fused weight-norm pairs net change: {fused_pairs:+d}") + + +if __name__ == "__main__": + main() diff --git a/webui/configs/model_params.json b/webui/configs/model_params.json index 45aaa12b..7c29e28d 100644 --- a/webui/configs/model_params.json +++ b/webui/configs/model_params.json @@ -363,5 +363,12 @@ {"name": "voice", "type": "choice", "label": "voice（预置音色：M 男声 / F 女声）", "label_en": "voice (M = male, F = female presets)", "default": "M1", "choices": ["M1", "M2", "M3", "M4", "M5", "F1", "F2", "F3", "F4", "F5"]}, {"name": "speaking_rate", "type": "slider", "label": "speaking_rate（语速倍率）", "label_en": "speaking_rate", "default": 1.05, "minimum": 0.5, "maximum": 2.0, "step": 0.05}, {"name": "num_inference_steps", "type": "number", "label": "num_inference_steps（流匹配步数）", "label_en": "num_inference_steps", "default": 8, "minimum": 1, "step": 1, "precision": 0} + ], + + "audio8_tts": [ + {"name": "temperature", "type": "slider", "label": "temperature", "default": 0.7, "minimum": 0.0, "maximum": 2.0, "step": 0.05}, + {"name": "top_p", "type": "slider", "label": "top_p", "default": 0.9, "minimum": 0.0, "maximum": 1.0, "step": 0.01}, + {"name": "top_k", "type": "number", "label": "top_k", "default": 50, "minimum": 0, "step": 1, "precision": 0}, + {"name": "max_new_tokens", "type": "number", "label": "max_new_tokens", "default": 1024, "minimum": 1, "step": 1, "precision": 0} ] } diff --git a/webui/configs/models_catalog.json b/webui/configs/models_catalog.json index 80d46e57..d9a6217a 100644 --- a/webui/configs/models_catalog.json +++ b/webui/configs/models_catalog.json @@ -45,6 +45,9 @@ { "id": "fish-audio-s2-pro", "display_name": "Fish Audio S2 Pro (tts 克隆/控制标记, GGUF Q8)", "display_name_en": "Fish Audio S2 Pro (tts + clone/control tags, GGUF Q8)", "family": "fish_audio", "path": "models/Fish-Audio-S2-Pro-GGUF", "task": "tts", "mode": "offline", "download_id": "fish_audio_s2_pro", "min_vram_gb": 8, "input_hint": "**Fish Audio S2 Pro**：Q8_0 GGUF 包；中英+自动语种；上传参考音色即克隆；正文里可写行内控制标记（如 (laugh)）。", "input_hint_en": "**Fish Audio S2 Pro**: Q8_0 GGUF package; English/Chinese plus auto language. Upload a reference voice to clone; inline control tags such as (laugh) can be written in the text." }, + { "id": "audio8-tts", "display_name": "Audio8 TTS Preview 0.6B (tts 克隆, GGUF Q8)", "display_name_en": "Audio8 TTS Preview 0.6B (tts + clone, GGUF Q8)", "family": "audio8_tts", "path": "models/Audio8-TTS-Preview-0.6B-GGUF", "task": "tts", "mode": "offline", "download_id": "audio8_tts_preview_0_6b_q8_0", "min_vram_gb": 4, + "input_hint": "**Audio8 TTS Preview 0.6B**：多语种 TTS / 零样本克隆（支持 yue/zh/nl/en/fr/de/it/ja/ko/pl/es/auto）；上传参考音色+参考文本即克隆，留空为普通 TTS；长文本自动分句。", + "input_hint_en": "**Audio8 TTS Preview 0.6B**: multilingual TTS and zero-shot clone (yue/zh/nl/en/fr/de/it/ja/ko/pl/es/auto). Upload a reference voice + transcript to clone; leave empty for plain TTS. Long text is chunked automatically." }, { "id": "glm-tts", "display_name": "GLM-TTS (tts 克隆, 社区)", "display_name_en": "GLM-TTS (tts + clone, community)", "family": "glm_tts", "path": "models/GLM-TTS", "task": "tts", "mode": "offline", "download_id": "glm_tts", "min_vram_gb": 8, "input_hint": "**GLM-TTS**（社区模型）：中英 TTS / voice clone；上传参考音色即克隆。", "input_hint_en": "**GLM-TTS** (community): Chinese/English TTS and voice clone. Upload a reference voice to clone." }, diff --git a/webui/native/dist/index.html b/webui/native/dist/index.html index 087f8cea..e6d9ceb1 100644 --- a/webui/native/dist/index.html +++ b/webui/native/dist/index.html @@ -31,20 +31,20 @@