diff --git a/CMakeLists.txt b/CMakeLists.txt index 6c788a2c..8e9f5897 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -625,6 +625,24 @@ audiocpp_add_model(neutts engine::models::neutts::make_neutts_loader ) +audiocpp_add_model(echo_tts + SOURCES + src/community_models/echo_tts/session.cpp + src/community_models/echo_tts/tokenizer.cpp + src/community_models/echo_tts/latent_post.cpp + src/community_models/echo_tts/sampler.cpp + src/community_models/echo_tts/dit.cpp + INCLUDES + engine/community_models/echo_tts/session.h + engine/community_models/echo_tts/config.h + engine/community_models/echo_tts/tokenizer.h + engine/community_models/echo_tts/latent_post.h + engine/community_models/echo_tts/sampler.h + engine/community_models/echo_tts/dit.h + LOADERS + engine::models::echo_tts::make_echo_tts_loader +) + if (MSVC) set_source_files_properties( src/community_models/inflect_v2/frontend.cpp @@ -1753,6 +1771,17 @@ if (ENGINE_BUILD_TESTS) add_engine_unittest(dots_tts_vocoder_parity tests/dots_tts/dots_tts_vocoder_parity.cpp) + # Needs the GGUF and a PyTorch reference dump, so it is driven by hand + # rather than registered with add_test -- same as dots_tts_vocoder_parity. + add_engine_unittest(echo_tts_dit_parity tests/echo_tts/echo_tts_dit_parity.cpp) + + add_engine_unittest(echo_tts_host_units tests/echo_tts/echo_tts_host_units.cpp) + + add_test( + NAME echo_tts_host_units + COMMAND echo_tts_host_units + ) + add_engine_unittest(midi_file_test tests/unittests/test_midi_file.cpp) add_test( diff --git a/docs/community_models/echo_tts.md b/docs/community_models/echo_tts.md new file mode 100644 index 00000000..5f976515 --- /dev/null +++ b/docs/community_models/echo_tts.md @@ -0,0 +1,383 @@ +# Echo-TTS + +Echo-TTS is an English zero-shot voice-cloning TTS model. A 2.8B diffusion transformer (EchoDiT) +generates 80-dimensional latents in PCA space, which the Fish S1-DAC autoencoder decodes to 44.1 kHz +audio. Cloning takes a reference wav with **no transcript required**. + +Upstream: [jordand/echo-tts-base](https://huggingface.co/jordand/echo-tts-base) · +autoencoder: [jordand/fish-s1-dac-min](https://huggingface.co/jordand/fish-s1-dac-min) + +| Family | `echo_tts` | +|---|---| +| Tasks | `clon` | +| Modes | offline | +| Languages | en | +| Sample rate | 44 100 Hz | +| Model directory | `models/echo-tts` | + +## Status + +**Work in progress.** Every stage is *implemented*; the open question is how much of it is +*verified*. Those are tracked separately on purpose, because a clean build and plausible audio +prove neither. + +| Milestone | Scope | Implemented | Numerically verified | +|---|---|---|---| +| M0 | Family registration, model spec v1 | yes | n/a | +| M1 | GGUF conversion, DiT, PCA inverse, Fish decode | yes | **denoiser yes, trajectory no** | +| M2 | Native speaker encoding (Fish encoder + RVQ) | yes | folded into the denoiser probe below | +| M3 | Long-form via the framework text chunker | yes | **no** | +| M4 | Q8_0 conversion, RTF and memory evidence | partial | **no** | + +Cloning is self-contained — `session.cpp` calls `codec_->encode_zq` directly, so no pre-computed +speaker latent is required. + +### Numerical parity against PyTorch + +Measured on an RTX 3090 (sm_86), CUDA, F16 GGUF, by `tests/echo_tts/echo_tts_dit_parity.cpp` +against a dump from `tools/community_models/echo_tts_reference.py`. Gates are cosine over the +flattened tensors **and** max-absolute-error, never equality: cosine alone cannot see a uniform +scale error, and the host Philox stream matches CUDA to ~2 ULP rather than bit-exactly. + +The reference defaults to **bfloat16** while the GGUF here is **F16**. Both dumps are shown because +the difference between them is the single largest term in the table: + +| Check | vs bfloat16 reference | vs float16 reference | Gate | Verdict | +|---|---|---|---|---| +| Denoiser, one conditional forward at t = 0.7 | cosine 0.999977, max-abs 0.086 | **cosine 0.999999188, max-abs 0.010** | ≥ 0.999 | **PASS** | +| 40-step sampler, reference initial noise injected | cosine 0.913082 | cosine 0.988459, max-abs 1.07 | ≥ 0.999 | **below gate** | +| 40-step sampler, our own seeded draw | cosine 0.905481 | cosine 0.976972, max-abs 1.87 | ≥ 0.999 | **below gate** | + +**The denoiser probe passes and is the number that carries the port.** It is what settles the four +details that fail *silently* rather than loudly — half-head RoPE, the interleaved (not NEOX) rotary +pairing, the speaker patchify reshape, and the adaLN `shift/scale/gate` order. A 0 % WER cannot do +that job: it scores the words, not the speaker identity, so a wrong patchify reshape in particular +could yield fluent, correctly-worded speech in the wrong voice and still read as 0 %. + +Note what the probe does **not** isolate. The reference text ids and speaker latents are injected, +but `prepare_conditioning()` then runs this port's own text encoder, speaker encoder and KV +projections, so the number covers the combined conditioning-plus-denoiser path. That is enough to +catch a wrong block; it is not enough to localise one. Per-block activation dumps +(`echo_tts_pack_reference.py --blocks`) exist for that and have not been run. + +**The 40-step trajectory is below the gate, and that is reported as a failure of the check as +written.** What is established about it: + +- **It is not the RNG.** Injecting the reference's own initial noise scores no better than our + seeded draw (0.988 vs 0.977), so the Philox difference is not the mechanism. +- **It is dominated by dtype.** Re-dumping the reference at float16 to match the GGUF moves the + seeded trajectory from 0.905 to 0.977 and the denoiser from 0.999977 to 0.999999. +- **It compounds with step count.** At 4 steps the same comparison scores 0.9965/0.9966; at 40 it + scores 0.9131/0.9055 (bfloat16 reference). Monotonic degradation with step count is the signature + of accumulating per-step rounding, amplified every step by dual CFG at 3.0 and 8.0, rather than of + a structural defect. +- **An independent line-by-line review of the sampler found no defect** — schedule, inclusive CFG + bounds, single application of `truncation_factor`, three-lane CFG combination, the Euler update + and the speaker-KV boundary all agree with `inference.py`. + +That is an explanation, not a proof. Until the residual is closed or the gate is deliberately +restated, treat the trajectory as **unverified**. + +### Host-side checks + +`tests/echo_tts/echo_tts_host_units.cpp` is registered with `add_test` and needs neither a GPU nor +the checkpoint. It covers WhisperD normalisation (including the asymmetric quote rewrite and the +bare-`S1` tag suppression), full byte-token id vectors, truncation and padding, PCA projection and +inversion pinned independently against hand-computed values on a rectangular non-symmetric basis, +and the flattening-point crop including its standard-deviation *and* mean thresholds. + +Separately, and **by hand rather than in CI**: PCA project/unproject at 5.7e-06 against numpy on +the real (80, 1024) basis, the Euler dual-CFG update at 6.6e-07 against a numpy transcription of +`inference.py`, the timestep embedding at 0.0 diff, and the tokenizer at 140/140 ids on the parity +prompt. `combine_cfg_lanes` and `euler_timestep_schedule` have no registered coverage. + +### End-to-end run, RTX 3090 (sm_86), CUDA, F16 GGUF + +| Check | Result | +|---|---| +| Conversion | `manifest OK`; 1117 DiT tensors written, 219 blockwise tensors dropped, 495 codec tensors | +| GGUF verifier | pass — 1614 tensors (`dit_weights/` 1117, `ae/` 495, `pca/` 2), F16 1043 / F32 571 | +| `latent_scale` | 0.0555555559694767 (= 1/18), matching the reference | +| Generation | exit 0, 44 100 Hz mono, no NaNs, peak 0.80 (below the normalisation threshold) | +| ASR round-trip, 15 words | WER 0 % — the only diffs are Whisper writing spoken "dot" as punctuation | +| ASR round-trip, 32 words | WER 0.0 %, 0 edits | +| Throughput | 9.195 s of audio in 7.89 s wall — **RTF 0.86 cold**, including the 5.5 GB model load | + +Transcription used `faster-whisper-large-v3-turbo`. Generation cost is essentially constant across +those two runs (7.75 s vs 7.89 s) because the window is fixed at 640 frames, so longer text inside +one chunk is close to free. + +WER on 32 words is a small sample and scores intelligibility only. It shows the pipeline runs end to +end and produces the right words; the denoiser cosine above is what shows the graph is right. + +### What is still missing + +- **No regression test for `fish_audio` itself.** `build_decode_quantizer` was **restructured**, not + merely extended, so a supported core family's decode path changed with no coverage of its own. + This is the largest gap in the list. +- No per-block DiT activation dump, so the passing denoiser cosine proves correctness without + localising where any future regression lives. +- The 40-step trajectory residual above is explained but not closed. +- No A/B of the flash-attention path against `AUDIOCPP_ECHO_TTS_NO_FLASH=1` on a fixed seed. +- No listening comparison of F16 against Q8_0. +- Warm RTF and VRAM-stability-across-requests numbers. + +This PR stays in draft until that evidence exists. + +## Known limitations + +### Fixed 29.72-second generation window + +Echo is trained to generate at most **640 latents**, and 640 × 2048 ÷ 44100 = **29.7215 s**. This is +a property of the model, not of this port. + +Behaviour outside that window: + +- Text corresponding to more than ~30 s is **spoken faster** to fit, rather than truncated. This is + learned behaviour arising from global attention over the text, not an explicit compression step. +- The upstream tokenizer hard-truncates text past **768 UTF-8 bytes**. +- Requesting a shorter window does *not* compress the whole utterance into it — upstream documents + that the model generates a **prefix** of the utterance instead. + +`long_form` is therefore **not** claimed in `capabilities` at this stage. + +### Blockwise generation does not extend the window + +Upstream ships a blockwise sampler that generates in connected blocks and supports continuing from +existing audio. It **subdivides** the ≤30 s window rather than extending it: upstream requires +`sum(block_sizes) + continuation_length < 640` "to be in-distribution with training data", and +documents prefix plus continuation as "up to 30 seconds combined". Upstream also notes blockwise +"hasn't been thoroughly tested". + +## Licence — read before using output commercially + +Echo-TTS is **CC-BY-NC-SA-4.0**, and the restriction covers **generated audio, not only the +weights**. The output constraint is inherited from the Fish S1-DAC autoencoder — the same mechanism +that makes Fish Speech's own outputs non-commercial. + +Practically: **audio produced by this model may not be used commercially**, regardless of how the +rest of your stack is licensed. audio.cpp itself is Apache 2.0 and is unaffected; model weights are +a separate download. + +There is existing precedent in-tree — `fish_audio` (Fish Audio S2 Pro) carries the identical +output restriction from the identical dependency. + +## Why this model + +Selected by comparing every model tracked in [tts-bench](https://github.com/5uck1ess/tts-bench) — a +public benchmark covering **62 local TTS models** across speed, objective scores, and blind human +preference — against audio.cpp's existing support table. + +| Measure | Echo-TTS | Field | +|---|---|---| +| Blind cloning Elo | **1162** | #3 of 40 (35 games; 738 cloning votes total) | +| Speaker similarity (SIM) | **0.836** | 2nd of 41 scored models | +| UTMOS (naturalness) | 4.21 | — | +| WER (intelligibility) | 7.45 % | — | +| Frozen pairwise study | **21-1-6** | near-tied 1st of 28 | + +Two honest caveats: the cloning arena averages ~30 games per model, so gaps under ~100 Elo are +noise, and the ranking uses a single reference clip. Echo's standing is robust to both — it is +top-3 on human votes *and* 2nd on objective SIM, which are independent measurements. + +Compute profile suits a GGUF port: ~2.8 B parameters at 1.35× RTFx and 9.4 GB VRAM in PyTorch on an +RTX 3090, so there is real work to amortise. + +## Architecture + +| Component | Params | Role | +|---|---:|---| +| EchoDiT trunk, 24 blocks | 1.75 B | Joint attention + SwiGLU MLP, adaLN timestep modulation | +| Text encoder | 294 M | UTF-8 **byte** tokens (256 vocab) — no phonemizer or G2P | +| Speaker encoder | 294 M | Reference PCA latents → speaker states | +| Latent encoder | 294 M | Blockwise only; omitted in M1 | +| PCA state | 83 K | Fish 1024-D ↔ DiT 80-D, `latent_scale` = 1/18 | +| Fish S1-DAC | 391 M weights | Reference encoding and waveform decoding | + +Sampling is 40 Euler steps with **two independent CFG scales** — text (default 3.0) and speaker +(default 8.0) — gated to `t ∈ [0.5, 1.0]`. + +Note the Fish checkpoint stores an additional 303.6 M elements of `freqs_cis` and `causal_mask` +buffers. These are regenerated at runtime rather than shipped in the GGUF. + +## Options + +| Option | Type | Default | Description | +|---|---|---|---| +| `target_voice` | string | — | Reference wav for cloning. No transcript needed. | +| `cfg_scale_text` | float | 3.0 | Guidance scale on the text condition. | +| `cfg_scale_speaker` | float | 8.0 | Guidance scale on the speaker condition. | +| `num_steps` | int | 40 | Euler sampler steps. | +| `truncation_factor` | float | 0.8 | Initial-noise truncation. | +| `speaker_kv_scale` | float | 1.0 | Force-speaker KV scaling; 1.5 is upstream's default when enabled. Raise only if the model drifts to a different speaker on out-of-distribution text. | +| `seed` | int | 0 | RNG seed for the initial latent. | + +## Text format + +Prompts follow the [WhisperD](https://huggingface.co/jordand/whisper-d-v1a) transcription style: + +- `[S1] ` is prepended automatically when neither `[S1]` nor `[S2]` is present. +- Colons, semicolons, and em dashes are normalised to commas. +- Commas generally function as pauses. +- Exclamation points and other emphatic punctuation increase expressiveness but can reduce quality. + +Multi-speaker dialogue is expressed with `[S1]` / `[S2]` tags. + +## Reference audio + +Up to 5 minutes is accepted; 10 seconds or less works well. Audio is mixed to mono, resampled to +44.1 kHz, and peak-limited before encoding. + +## Running it + +``` +audiocpp_cli \ + --family echo_tts \ + --model /path/to/Echo-TTS-GGUF \ + --task clon \ + --voice-ref reference.wav \ + --text "[S1] Alright, I'm going to demo this new model." \ + --out out.wav +``` + +The speaker reference is `--voice-ref`, not `--target-voice`; the latter is for +path-based voice conversion. No transcript of the reference is needed. Useful +request options: `num_steps` (default 40), `cfg_scale_text` (3.0), +`cfg_scale_speaker` (8.0), `truncation_factor` (0.8), and `seed`. + +## Quantisation + +`--precision q8_0` produces a roughly half-size GGUF: + +| | F16 | Q8_0 | +| --- | ---: | ---: | +| DiT | 4.76 GB | 2.53 GB | +| codec | 0.78 GB | 0.50 GB | + +Q8_0 packs 32 weights per block behind one shared scale, so a tensor qualifies +only when its last logical dimension is a multiple of 32. The converter routes +each tensor accordingly rather than quantising blindly: + +* **Q8_0** -- every 2-D matmul weight with a conforming row length. That is all + but one DiT tensor, and 78% of codec weights. +* **F16** -- convolution kernels (`ggml_conv_1d` has no quantised path, which is + why `codec.cpp` takes matmul and conv storage types separately) and the one + non-conforming matmul, `in_proj.weight` at (2048, 80). +* **F32** -- norm weights, biases, snake alphas, LayerScale/ConvNeXt gammas and + the codebooks, exactly as at other precisions. + +Round-trip error is around 6e-05 RMSE with cosine similarity above 0.9999 on +weight-like distributions. Because the scale is per 32-weight block, the large +outliers this model carries in its late blocks and in `k_norm` degrade only +their own block rather than a whole row -- and `k_norm` is F32 regardless. + +Quality has not been compared against F16 on real audio. Start with `orig` and +treat Q8_0 as an experiment until someone listens to both. + +## Limiting the reference length + +`reference_max_seconds` trims the speaker reference before encoding. Shorter +references cost less and often clone better -- upstream's guidance favours +around 10 s, and a long clip averages timbre over more prosodic variation. + +Per request (bare name): + +``` +--request-option reference_max_seconds=30 +``` + +As a default for a CLI run or a server, in the session scope (family-prefixed, +which is how the framework namespaces session and load options): + +``` +--session-option echo_tts.reference_max_seconds=30 +``` + +In a server config file the same key goes under `session_options`, with a string +value. A request value overrides the session default. Values above the trained +maximum of 297.1 s are clamped rather than rejected. Trimming happens before +chunked encoding, so a cap also bounds encode time and VRAM. + +## Reference encoding cost + +Encoding the speaker reference is linear in its length: one Fish encode pass per +~29.7 s chunk, so a 4m29s clip is ten passes against one for a 28 s clip. At the +trained maximum that is roughly 22% of a request's arithmetic, before per-graph +launch overhead. + +The result depends only on the audio and the trim length, so it is cached across +requests. A server rotating a few voices pays the cost once per voice instead of +once per request: + +``` +--session-option echo_tts.reference_cache_slots=8 +``` + +Default 4; `0` disables it. Each slot holds only the projected latent, at most +2 MB. The cache lives with the session, so it helps a running server and not a +one-shot CLI invocation. Beyond caching, the levers are `reference_max_seconds` +and shorter references generally -- around 10 s is one chunk, the floor. + +## The Fish S1-DAC autoencoder + +Echo decodes its 80-D PCA latents through the Fish S1 DAC and encodes speaker +references with the same model. audio.cpp already implements that codec for the +`fish_audio` family, so Echo reuses the implementation -- but **not** the +weights. `fish_audio` ships Fish Audio S2 Pro; Echo is trained against the S1 +DAC (`jordand/fish-s1-dac-min`), and `pca_state.safetensors` is fitted to that +codec's latent space specifically. Pointing Echo at S2 Pro would produce +plausible-looking latents and wrong audio, with no error anywhere. + +The S1 weights are therefore packaged inside Echo's own GGUF, in the `ae` +namespace: + +``` +python3 tools/community_models/convert_echo_tts.py \ + --model-dir /path/to/echo-tts-base \ + --fish-dir /path/to/fish-s1-dac-min \ + --outfile Echo-TTS-GGUF/model.gguf +``` + +No companion model and no extra options are needed at run time. Two details the +converter handles: + +* **Weight normalisation is folded.** The checkpoint stores it in two forms -- + `conv.parametrizations.weight.original0/original1` on the convolutions and + legacy `weight_g`/`weight_v` on the quantiser projections -- and `codec.cpp` + expects plain `conv.weight`. Both reduce to `w = g * v / ||v||` with the norm + taken over every axis but the first. Note that for `ConvTranspose1d` axis 0 is + the *input* channel count, so `g` is sized by input channels there; the + decoder's four transposed convolutions are the only place this bites. +* **Fused qkv projections are split.** `autoencoder.py` keeps one `wqkv` + linear and splits its output into three equal blocks; `codec.cpp` loads + `attention.q_proj` / `k_proj` / `v_proj` separately, so the converter + partitions the weight rows in the same order. +* **Exact tensor shapes are carried in metadata.** `ggml_n_dims()` ignores + trailing dimensions of size 1, so a `(1, C, 1)` snake alpha would read back as + `(C, 1)` and fail `codec.cpp`'s `{1, C, 1}` shape check. The converter emits + `audiocpp.tensor_ranks` (INT32) and `audiocpp.tensor_shapes` (INT64) in tensor + order, which audio.cpp uses in preference to the lossy inference. +* **The GGUF embeds its own model spec.** `package.cpp` refuses to load a + published GGUF that does not, so the converter copies + `model_specs/echo_tts.json` into the `audiocpp.model_spec.*` metadata keys. + A distributed file is therefore self-describing and does not depend on the + reader having a matching `model_specs/` checkout. Use `--model-spec` to embed + a spec from elsewhere. +* **Namespaces are separated by `/`, not `.`** -- `dit_weights/...`, `pca/...`, + `ae/...`. `PrefixedTensorSourceView` matches on `prefix + "/"`, so a + dot-separated name is never routed and the loader reports the namespace as + non-existent rather than the tensor as missing. +* **The codec namespace is `ae`, not `codec_weights`.** ggml caps tensor + names at 64 characters (`GGML_MAX_NAME`) and rejects the whole file at load + time if any name reaches it. The longest name `codec.cpp` loads is already 60 + characters, so only a three-character prefix fits; `codec_weights.` would push + 157 of the 455 codec tensors over. The converter refuses to write a GGUF that + would trip this, and `verify_echo_gguf.py` re-checks it. +* **Registered buffers are dropped.** Two causal masks and three RoPE tables + account for 305 MB of the 1.87 GB checkpoint and are rebuilt at graph + construction, so they are not stored. + +That leaves roughly 1.57 GB of codec weights on top of the 4.76 GB DiT. +`docs/community_models/echo_tts_autoencoder_reuse.md` covers how the two +families share the codec implementation and where the seam sits in +`src/models/fish_audio/codec.cpp`. diff --git a/docs/community_models/echo_tts_autoencoder_reuse.md b/docs/community_models/echo_tts_autoencoder_reuse.md new file mode 100644 index 00000000..a18163e7 --- /dev/null +++ b/docs/community_models/echo_tts_autoencoder_reuse.md @@ -0,0 +1,129 @@ +# Echo-TTS: autoencoder reuse + +Status: verified against checkpoint sizes and upstream source. No weights were +downloaded to reach these conclusions; every number below is reproducible from +`autoencoder.py` plus the file sizes Hugging Face reports. + +## Summary + +Echo-TTS depends on the Fish S1-DAC autoencoder, and **audio.cpp already +implements that exact autoencoder** for the `fish_audio` family in +`src/models/fish_audio/codec.cpp`. The Echo port does not need a new decoder, +encoder, quantiser, or window-limited transformer. It needs a `z_q` seam on the +existing one. + +This changes the cost of milestones M1 and M2 substantially relative to the +original PR plan, which scoped "Fish decode" and "native speaker encoding +(Fish encoder + RVQ)" as separate pieces of work. + +## Evidence + +### Configuration + +`jordand/fish-s1-dac-min/config.json` reports: + + sample_rate 44100, encoder_dim 64, encoder_rates [2,4,8,8], latent_dim 1024, + decoder_dim 1536, decoder_rates [8,8,4,2], n_codebooks 9, codebook_size 1024, + codebook_dim 8, semantic_codebook_size 4096, causal true + +Every one of these matches the constants already compiled into +`fish_audio/codec.cpp`: `kCodecDim` 1024, semantic codebook 4096, nine residual +quantisers of 1024, codebook dim 8, a final decoder snake at 96 channels +(= 1536 / 2^4), and causal convolutions throughout. + +### Parameter budget + +Deriving the parameter count from `autoencoder.py` and comparing against the +1.87 GB Hugging Face reports for `pytorch_model.safetensors`: + +| Component | Parameters | +| --- | ---: | +| Encoder | 76,851,328 | +| Decoder | 54,102,722 | +| Quantiser (incl. pre/post transformers) | 260,475,040 | +| **Total** | **391,429,090** | + +At F32 that is 1.566 GB. The `Transformer` base class registers two buffers per +instance — a `freqs_cis` table and a `block_size^2` boolean `causal_mask` — which +for the three surviving transformer instances (encoder block 3 at block_size +16384, quantiser pre/post at 4096) comes to 305 MB. Together: **1.871 GB**, +against the 1.87 GB reported. This also reproduces the 303.6 MB +"regenerable buffers" figure noted on the PR. + +The match only holds once the decoder is counted **without** a transformer, which +leads to the next point. + +### The decoder has no transformer + +`build_ae` passes `decoder_transformer_layers=[4, 0, 0, 0]`, which reads as though +decoder block 0 carries a 4-layer transformer. It does not. `DecoderBlock.__init__` +constructs `transformer_module` into a local variable and then builds +`self.block = nn.Sequential(Snake1d, conv_trans, ResidualUnit x3)` without it. +The module is never assigned to `self`, so it is not a submodule, has no +parameters, and is absent from the checkpoint. `EncoderBlock`, by contrast, does +include `transformer_module` in its `Sequential`. + +Two independent checks agree: + +1. The 1.87 GB file size only reconciles when the decoder transformer is excluded + (including it predicts 2.05 GB, and adds a second 16384x16384 mask buffer that + would break the 303.6 MB figure). +2. `fish_audio/codec.cpp` already loads the encoder transformer conditionally at + `block_index == 3` and loads no transformer anywhere in the decoder path. + +The C++ was evidently written against the real checkpoint, and it agrees with +the source reading. Worth knowing before anyone "fixes" the apparent omission. + +## Integration seam + +Echo needs continuous `z_q` where `fish_audio` uses discrete codes. Both seams +sit at existing boundaries in `codec.cpp`: + +**Decode.** `DAC.decode_zq` is `post_module -> upsample -> decoder`. +`build_decode_quantizer` already performs exactly that chain; it just derives its +input by looking up codebook entries first: + + latent = build_quantizer_out(semantic) + sum(build_quantizer_out(residual_i)) + latent = build_window_transformer(..., post_module, 128) <- Echo enters here + for stage in upsample: ... + +Echo supplies `latent` directly from the PCA inverse and runs from the +`post_module` line onward. The refactor is to split the code-lookup prefix from +the `post_module`-onward suffix so both families can call the suffix. + +**Encode.** `DAC.encode_zq` quantises and then sums the dequantised results: +`z_q = z_q_semantic + z_q_residual`. `build_encode_quantizer` already computes +each `quantized` term internally on the way to emitting code indices; the sum is +available at that point and is currently discarded. Exposing it gives native +speaker encoding without new model code, which is most of milestone M2. + +## Consequences for packaging + +The *implementation* is shared; the *weights* are not. + +`fish_audio` ships Fish Audio S2 Pro. Echo is trained against the Fish S1 DAC +(`jordand/fish-s1-dac-min`, a mirror of `fishaudio/openaudio-s1-mini`), and its +PCA basis is fitted to that codec's latent space. The S2 technical report says S2 +retains S1's RVQ codec, and the shapes line up (10 codebooks, ~21 Hz), but +"retains the codec" in a report can mean the architecture rather than identical +weights -- and a retrained-but-isomorphic codec would yield wrong audio with no +error raised anywhere. That is not a risk worth taking to save a download. + +`convert_echo_tts.py` therefore packages the S1 codec into Echo's GGUF under the +`codec_weights` prefix, folding weight normalisation and dropping the 305 MB of +regenerable buffers. Echo constructs a minimal `FishAudioAssets` around that +tensor source: only four config fields (`sample_rate`, `frame_length`, +`total_codebooks`, `quantizer_codebooks`) ever reach the codec graphs, and their +defaults already describe S1-DAC. + +Verified against the real checkpoint manifest: the folded output supplies all 220 +tensor names `codec.cpp` loads, and the 541 stored tensors resolve to 455 after +folding and buffer removal. + +## Caveat + +Everything above is derived from source reading plus file-size arithmetic. The +parameter total agreeing with the reported size to three significant figures is +strong evidence, but it is not the same as having loaded the tensors. The +tensor-name check in `convert_echo_tts.py` and a parity run against +`echo_tts_reference.py` remain the gates before any of this is claimed as done. diff --git a/docs/community_models/echo_tts_dit_status.md b/docs/community_models/echo_tts_dit_status.md new file mode 100644 index 00000000..bbd6a9a0 --- /dev/null +++ b/docs/community_models/echo_tts_dit_status.md @@ -0,0 +1,133 @@ +# Echo-TTS DiT: implementation status + +## What exists + +| Component | State | Verification | +| --- | --- | --- | +| Byte tokenizer + WhisperD normalisation | complete | executed, output checked by hand | +| PCA forward / inverse | complete | executed, cross-checked against numpy | +| Flattening-point crop | complete | executed | +| Euler dual-CFG sampler | complete | executed, matches a numpy transcription of `inference.py` to 6.6e-07 | +| Timestep embedding | complete | matches `model.py` exactly (0.0 diff) | +| Attention mask construction | complete | layout checked against upstream `cat()` semantics and ggml constraints | +| DiT graph (encoders, joint attention, adaLN, blocks) | written | compiles against real headers; **never executed** | +| Weight loading (1,117 tensors) | written | compiles; **tensor names unconfirmed against a real checkpoint** | +| Conditioning / denoiser graph execution | written | compiles; **never executed** | +| Fish codec `z_q` seam | not started | — | +| Session integration | not started | — | + +The distinction in that last column is the important one. Everything above the +line was run and compared against a reference. Everything below it has only been +type-checked. A clean compile here means the framework APIs are used correctly; +it says nothing about whether the numbers are right. + +## Design decisions worth reviewing + +### Flash attention in the DiT joint attention + +`joint_attention` uses `ggml_flash_attn_ext`, which never materialises the +`(lanes, heads, seq, keys)` scores tensor. That tensor was the largest +per-request allocation in the model: + +| Case | Keys | Scores tensor removed | +| --- | ---: | ---: | +| Typical (64 text bytes, 10 s reference) | 793 | 97 MB per attention | +| Long text, 30 s reference | 1569 | 193 MB per attention | +| Worst case (768 text, 5 min reference) | 3008 | 370 MB per attention | + +Live across 24 blocks with `ggml_gallocr` reuse, the practical saving is a few +hundred MB to over a gigabyte, and flash attention is also faster. + +This was initially written with the explicit lowering on the belief that the +speaker-unconditional CFG lane produces fully masked rows, which would make +`-inf` softmax to NaN. That was wrong: `make_denoiser_mask` leaves the self block +of every row unmasked, so a query always attends to at least its own 640 +positions and no row can be fully masked. + +Two details the flash path requires. The mask must be F16, so the masked value +is `-65000` rather than `-1e9`; the latter converts to `-inf` in F16, which would +reintroduce exactly the NaN hazard the explicit path was chosen to avoid. And +`q->ne[2] % mask->ne[2]` and `q->ne[3] % mask->ne[3]` must both be zero, which +holds because the mask carries a singleton head axis and matches the lane count. + +Set `AUDIOCPP_ECHO_TTS_NO_FLASH=1` to fall back to the explicit lowering and F32 +mask, for A/B comparison without a rebuild. + +The two encoders still use the explicit lowering. Their sequences are short (a +few hundred tokens at most) so the scores tensors are small, and the speaker +encoder is causal with no explicit mask, which the flash path rejects. + +### Speaker references are encoded in chunks + +`encode_speaker` splits the reference into ~29.7 s chunks (640 latents x 2048 +samples), zero-pads the last one, and concatenates the per-chunk latents, +following `inference.py::get_speaker_latent_and_mask`. Upstream's comment calls +that the longest chunk seen in training, so this is a fidelity matter as much as +a memory one -- encoding several minutes in a single pass is a different +computation from what the model saw. + +The memory difference is large, because the Fish encoder's first stages run at +the full 44.1 kHz rate. A single 64-channel activation is 0.34 GB for one chunk +against 3.04 GB for a 4m29s reference encoded in one pass, and several such +tensors are live at once. Fixed-size chunks also mean one encode graph is built +and reused across all chunks. + +After chunking, the dominant per-request allocation at long reference lengths is +the persistent KV cache: 0.59 GB at 4m29s and 0.65 GB at the 297 s maximum, +stored F32. Halving it to F16 is the obvious next step if that ever matters. + +### KV cache as a separate backend buffer + +The conditioning encoders and the denoiser are separate graphs so the encoders +run once per request rather than once per sampler step. They share the cached +projections through tensors allocated in their own `ggml_context` and backend +buffer, referenced as leaves by both graphs. `ggml_gallocr` leaves +already-allocated tensors alone, so the conditioning graph writes into them with +`ggml_cpy` and the denoiser graph reads them directly. + +Consequence: changing text length or speaker length invalidates the cache and +every graph built against it. `prepare_conditioning` tears all of it down and +rebuilds, which is correct but means a request with new conditioning pays full +graph construction. Acceptable given that a 40-step sample dominates. + +### Speaker KV scaling round-trips through the host + +`scale_speaker_kv` reads the cached tensors back, scales, and re-uploads, because +the cache has no graph attached. This runs at most twice per request (once to +apply, once to undo at the threshold) and touches at most 24 layers x 2 tensors. +It is not on the per-step path. If it ever shows up in a profile, the fix is a +tiny scaling graph rather than a host round trip. + +## Things most likely to be wrong + +Listed in rough order of how much damage they would do and how hard they would +be to spot without a parity run: + +1. **Tensor names.** Derived from `model.py`'s module structure, corroborated by + a parameter count matching the published file size to ten digits, but never + resolved against an actual checkpoint. `convert_echo_tts.py --model-dir ...` + settles this in seconds and prints exactly what is wrong if anything is. +2. **Half-head RoPE.** Heads 0-7 rotate, 8-15 do not. Implemented as + slice/rope/concat on the head axis. Wrong here means plausible-sounding but + incorrect audio, with no shape error. +3. **Rotary pairing convention.** `GGML_ROPE_TYPE_NORMAL` (interleaved), matching + upstream's complex view of adjacent pairs. The in-tree `rf_dit.cpp` uses NEOX, + so copying from it would be wrong. +4. **Speaker patchify reshape.** Folding `patch_size` frames into the feature + axis assumes row-major frame-then-channel ordering. A transposed reading would + still produce correct shapes. +5. **adaLN chunk order.** `shift, scale, gate` from `cond_embed.chunk(3, -1)`. + A permutation here is silent. + +Items 2-5 are all caught by the per-block parity dumps from +`tools/community_models/echo_tts_reference.py`, which is why that script dumps +per-block activations at a fixed timestep rather than only the final output. + +## Next steps + +1. Run `convert_echo_tts.py` against the real checkpoint to confirm item 1. +2. Build on a machine with a GPU and run the parity comparison per block. +3. Split `fish_audio/codec.cpp`'s `build_decode_quantizer` at the `post_module` + boundary and expose the summed `quantized` term from + `build_encode_quantizer`, giving Echo decode and native speaker encoding. +4. Wire the session: tokenize, encode speaker, sample, PCA inverse, decode, crop. diff --git a/docs/community_models/echo_tts_parity_run1.md b/docs/community_models/echo_tts_parity_run1.md new file mode 100644 index 00000000..14aeebbc --- /dev/null +++ b/docs/community_models/echo_tts_parity_run1.md @@ -0,0 +1,91 @@ +# Echo-TTS parity run 1: findings + +Source: `echo_ref.npz`, 146 arrays, generated from `audio_prompts/musk1.wav`, +seed 0, 40 steps, sequence_length 640, model dtype bfloat16. + +## Components now verified against real data + +| Component | Result | +| --- | --- | +| Byte tokenizer + normalisation | **exact** — all 140 ids and the normalised string match byte-for-byte | +| PCA orientation | **confirmed** — `pca.components` is `(80, 1024)`, as assumed | +| `pca_unproject` (C++) | max abs 5.7e-06 against numpy on the real basis (z_q range 11.97) | +| PCA round trip (C++) | max abs 5.3e-06 against the real speaker latent (range 2.51) | +| Reconstructed z_q vs `ae.encode_zq` | min -10.2110 / max +11.9744 vs reference -10.2115 / +11.9740 | +| `find_flattening_point` (C++) | **exact** — 140 of 640 frames, matching the reference heuristic | +| `latent_scale` | float32(1/18) exactly | + +The flattening-point match is worth calling out: it ran on the real generated +latent, not a synthetic one, and 140 frames is 6.502 s of audio from a 29.72 s +window. Getting this wrong changes the output duration silently, and it is +sensitive to the variance convention — I checked that `ddof=0` also lands on 140 +here, so this particular case would not have caught a wrong choice. The +implementation uses the unbiased estimator to match `torch.std`, which is right +for the general case regardless. + +## RNG: same stream, not bit-exact + +`generate_torch_cuda_randn(51200, 0)` was compared against +`sampler.initial_noise`: + + cosine 1.000000000000 (1 - cos = 3.1e-14) + median error 2 ULP + p99 error 149 ULP + correlation 1.000000000000 to 12 digits + rank agreement 99.64% + +Same Philox stream and same normal transform; the residual is CUDA-vs-host libm +precision in the transcendental calls. **Seeded parity will be near-identical but +not bit-exact**, so a 40-step trajectory will diverge slightly from the +reference. Against the PR's cosine >= 0.999 gate this is irrelevant (the noise +alone passes with ~3e10 margin), but any test written to expect bit-equality +would fail for reasons that are not bugs. Write the gates as cosine plus +max-abs-error, not equality. + +Also confirmed: `dit.x_input` is bitwise identical to `sampler.initial_noise`, +so the fixed-timestep probe and the sampler share a starting draw. + +## Massive activations in the late blocks + +Activation magnitude grows monotonically through the stack: + +| Block | std | max | +| ---: | ---: | ---: | +| 0 | 0.266 | 5.69 | +| 12 | 0.393 | 14.88 | +| 20 | 1.147 | 38.75 | +| 22 | 2.619 | 97.50 | +| 23 | 5.893 | 187.00 | + +std grows 22x and max 33x from first block to last, with most of it in the final +four blocks. Separately, the layer-23 key caches for **both** text and speaker +peak at exactly 510.0 while layers 0 and 12 peak near 8-10. The two paths share +one `k_norm` weight per layer and carry unrelated inputs, so an identical maximum +points at a large element in that weight rather than at the data — the standard +massive-activation / attention-sink pattern. + +Consequences: + +1. **F16 activations are safe.** 510 and 187 are far below the 65504 F16 ceiling. + No overflow risk in the planned conversion. +2. **The converter's decision to keep norm weights in F32 was right for a reason + that was not known when it was made.** `KEEP_F32_SUBSTRINGS` already covers + `q_norm` and `k_norm`. Quantising a weight with a ~510 outlier to a + block-scaled int8 would destroy the small elements sharing its block. +3. **Q8_0 (milestone M3) needs care in the last four blocks.** A per-block int8 + scale resolves roughly 1/127 of the block maximum, which at block 23 is ~1.5 + absolute against a std of 5.89. Mixed precision — leaving blocks 20-23 at F16 + — is the obvious first thing to try if Q8_0 degrades quality. + +## Still unverified + +The DiT graph itself. Every check above exercises host-side code; nothing has run +the ggml graph, because that needs a build. The per-block dumps in this file are +exactly what the block-by-block comparison will consume once it can run, and the +growth table above doubles as a smoke test: a port that gets the residual stream +right should reproduce that monotone 22x growth, and one that gets adaLN gating +or the half-head RoPE wrong will not. + +Useful next request, if another run is cheap: `--full-blocks`, which dumps every +block activation in full rather than stats plus a 64-value prefix. Not needed +until there is a build to compare against. diff --git a/docs/community_models/echo_tts_performance.md b/docs/community_models/echo_tts_performance.md new file mode 100644 index 00000000..55975b46 --- /dev/null +++ b/docs/community_models/echo_tts_performance.md @@ -0,0 +1,137 @@ +### 1. Adaptive generation window + +`parse_sampler_options` hardcoded `sequence_length = max_sequence_length`, so +every chunk paid the full 640-latent / 29.72 s window and `find_flattening_point` +discarded the silent tail. Now estimated per chunk from the tokenized byte count, +with one automatic retry at full length if no flattening point is found. + +The rate is derived, not guessed: 640 frames span 29.7215 s → 21.53 frames/s; +`kDefaultTextChunkSize` is documented in-file as ~20 s at 300 codepoints → +~15 bytes/s; 21.53 / 15 = **1.435 frames per UTF-8 byte**, with a 1.30 margin. + +| chunk bytes | ≈ speech | window | denoiser cost | +|---|---|---|---| +| 60 | 4.0 s | 128 | 5.00x cheaper | +| 120 | 8.0 s | 256 | 2.50x | +| 200 | 13.3 s | 384 | 1.67x | +| 300 | 20.0 s | 576 | 1.11x | +| 340+ | 22.7 s+ | 640 | 1.00x (unchanged) | + +**Off by default.** The original claim here was that an under-estimate "costs +time, never fidelity", because the 640-frame retry draws bit-identical noise from +the sequential Philox stream. The noise claim is true and the fidelity conclusion +does not follow from it. Echo's generated self-attention is fully non-causal +(`self_mask = torch.ones((batch_size, seq_len))`, `model.py:249`), so every +latent position attends across the entire window. Shrinking 640 to 128 changes +the computation at every retained position, not just how many positions survive +-- and the retry fires only when no flattening point is found, so a short window +that yields a plausible flat tail is never corrected. + +Enable with `AUDIOCPP_ECHO_TTS_ADAPTIVE_WINDOW=1` once it has been A/B'd against +the full window on a fixed seed. The cost saving below is real; it is the +default that was wrong. + +Estimates snap to a 64-frame grid (`kWindowQuantum`) because denoiser graphs are +keyed on `sequence_length` and rebuilt when it changes. + +- Pin explicitly: `sequence_length` request option (skips the estimate). +- Enable: `AUDIOCPP_ECHO_TTS_ADAPTIVE_WINDOW=1` (off by default). + +Confirmed in your logs: 23 bytes → 128-frame window, `keys` 824 → 312. + +### 2. KV cache pre-expanded across CFG lanes + +`joint_attention` called `expand()` — a `RepeatModule` broadcasting the cached +text and speaker K/V across the 3 CFG lanes — every layer, every step. The cache +is now allocated at `kMaxCfgLanes` and broadcast once in the conditioning graph, +so `expand()` short-circuits. The single-lane graph reads lane 0 through +`kv_for_lanes()`, a zero-copy view (the lane axis is outermost). + +`kv_for_lanes()` reconciles both directions — it narrows when the cache is wider +than the graph, and returns the cache untouched when it is narrower, leaving +`expand()` to broadcast as before. That second case is what +`AUDIOCPP_ECHO_TTS_NO_KV_EXPAND=1` produces: + +| `NO_KV_EXPAND` | cache | graph | `kv_for_lanes` | `expand` | +|---|---|---|---|---| +| off | 3 | 1 | slice to 1 | no-op | +| off | 3 | 3 | passthrough | no-op | +| on | 1 | 1 | passthrough | no-op | +| on | 1 | 3 | passthrough | repeat to 3 | + +**Tradeoff:** the cache is 3x larger. Confirmed active in your logs by +`kv_text.0.k n=141312` = 23 × 2048 × 3. + +### 3. `reference_max_seconds` defaults to 15 s + +Previously unbounded to the trained maximum, so an untrimmed clip charged up to +1600 speaker tokens to `keys` in all 24 blocks at every step, plus a linear +encode pass per chunk of reference. Confirmed in your logs: 30 s → 161 tokens +became 15 s → 80 tokens, taking `keys` from 824 to 744 at the same window. + +### 4. `cfg_interval` (opt-in, default 1 = off) + +Echo guides every step in the t >= 0.5 window with three forward passes: + +``` +v_pred = v_cond + 3.0*(v_cond - v_text_uncond) + 8.0*(v_cond - v_speaker_uncond) +``` + +`v_cond` moves quickly in t; the *correction* does not. `cfg_interval` measures +the correction every Nth guided step and reuses it in between, so skipped steps +cost one forward pass instead of three. + +| num_steps | interval | guided | refreshes | lane-evals | vs 40/1 | +|---|---|---|---|---|---| +| 40 | 1 | 20 | 20 | 80 | 1.00x | +| 40 | 2 | 20 | 10 | 60 | 1.33x | +| 40 | 3 | 20 | 7 | 54 | 1.48x | +| 30 | 1 | 15 | 15 | 60 | 1.33x | +| 30 | 2 | 15 | 8 | 46 | 1.74x | +| 20 | 2 | 10 | 5 | 30 | 2.67x | +| 14 | 1 | 7 | 7 | 28 | 2.86x | +| 14 | 2 | 7 | 4 | 22 | 3.64x | + +**The number that matters is `refreshes`, not the speedup.** Ten refreshes +across the guided phase means each reused correction is one small t-step stale. +Four means the correction was already coarsely sampled before you subsampled it, +and the speaker term's weight of 8.0 multiplies any staleness straight into +timbre and pronunciation -- a failure mode you hear rather than see on a +waveform. + +So: **raise it to 2 at 30+ steps; leave it at 1 below ~20.** `interval=3` buys +1.48x against 1.33x at 40 steps -- most of the fidelity risk for a fraction of +the extra speed. The curve flattens because the unguided half of the schedule is +a floor: at 40 steps, 20 of the 80 lane-evals were never guided, so no interval +gets past 2.0x. + +Note that **40 steps at interval 2 and 30 steps at interval 1 +both cost 60 lane-evals.** Same compute, spent differently: fewer steps coarsens +the whole ODE trajectory, a longer interval leaves the trajectory intact and only +lets the guidance go stale. Neither dominates on paper. Compare them by +listening. + +The implementation holds the correction in absolute units rather than as a +ratio, so a stale value cannot amplify a small `v_cond`, and the first guided +step always refreshes so a stale delta is never applied before one exists. +Verified against a mock denoiser at 0.011% max deviation for interval 2 and +0.020% for interval 3 -- but that mock had deliberately smooth uncond offsets, +so treat those as a lower bound, not a measurement on real audio. + +--- + +## Suggested config + +```json +"session_options": { + "echo_tts.reference_cache_slots": "8", + "echo_tts.reference_max_seconds": "15" +}, +"default_request_options": { + "num_steps": 14, + "cfg_interval": 1 +} +``` + +Do **not** put `sequence_length` here — it pins the window and gives back the +whole of change 1. diff --git a/docs/community_models/models.md b/docs/community_models/models.md index 69060be2..5daf5f4a 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -16,6 +16,7 @@ Practical expectations: | Family | Task | Supported language(s) | Contributor | What They Added | |---|---|---|---|---| +| **echo_tts** | TTS, voice cloning | en | Tym [@5uck1ess](https://github.com/5uck1ess), [@dignome](https://github.com/dignome) | [Echo-TTS](echo_tts.md) 44.1 kHz zero-shot voice cloning: 2.8B diffusion transformer in 80-D PCA space, decoded by the Fish S1-DAC autoencoder. Byte-level text, no phonemiser, no reference transcript | | **glm_tts** | TTS, voice cloning | zh, en | Mirek [@mirek190](https://github.com/mirek190) | [GLM-TTS](glm_tts.md) zero-shot synthesis and voice cloning support | | **inflect_v2** | TTS | en | Community | [Inflect Micro v2 and Nano v2](inflect_v2.md) native FP32 offline synthesis | | **kroko_asr** | ASR | de, en, es, fr, it, he, nl, pt, sv, tr | Mirek [@mirek190](https://github.com/mirek190) | [Kroko Community ASR](kroko_asr.md) native offline/streaming Zipformer2/RNN-T transcription with word timestamps | diff --git a/include/engine/community_models/echo_tts/config.h b/include/engine/community_models/echo_tts/config.h new file mode 100644 index 00000000..aa46e1d9 --- /dev/null +++ b/include/engine/community_models/echo_tts/config.h @@ -0,0 +1,207 @@ +#pragma once + +#include "engine/framework/core/module.h" +#include "engine/framework/modules/linear_module.h" + +#include +#include +#include + +namespace engine::models::echo_tts { + +// Architecture constants. These mirror the EchoDiT constructor arguments in +// upstream inference.py::load_model_from_hf, which are not stored in the +// checkpoint. The converter re-emits them as GGUF metadata and the loader +// cross-checks the values it reads back against these defaults, so a future +// upstream config change surfaces as a load error rather than silent garbage. +struct EchoTtsConfig { + // Denoiser (EchoDiT). + int64_t latent_size = 80; + int64_t model_size = 2048; + int64_t num_layers = 24; + int64_t num_heads = 16; + int64_t intermediate_size = 5888; + float norm_eps = 1.0e-5F; + + // Text encoder. + int64_t text_vocab_size = 256; + int64_t text_model_size = 1280; + int64_t text_num_layers = 14; + int64_t text_num_heads = 10; + int64_t text_intermediate_size = 3328; + + // Speaker encoder. + int64_t speaker_patch_size = 4; + int64_t speaker_model_size = 1280; + int64_t speaker_num_layers = 14; + int64_t speaker_num_heads = 10; + int64_t speaker_intermediate_size = 3328; + + // Conditioning. + int64_t timestep_embed_size = 512; + int64_t adaln_rank = 256; + + // Sampling / windowing limits, all fixed by training. + int64_t max_sequence_length = 640; // 640 * 2048 / 44100 = 29.7215 s + int64_t max_text_length = 768; // hard truncation, UTF-8 bytes + int64_t max_speaker_latent_length = 6400; + int64_t speaker_chunk_latents = 640; // 640 * 2048 samples per encode chunk + + // Autoencoder. + int64_t ae_downsample_factor = 2048; + int64_t ae_latent_dim = 1024; // Fish S1-DAC z_q channel count + int64_t sample_rate = 44100; + + int64_t head_dim() const { return model_size / num_heads; } + int64_t text_head_dim() const { return text_model_size / text_num_heads; } + int64_t speaker_head_dim() const { return speaker_model_size / speaker_num_heads; } + + // The DiT applies RoPE to the first half of the heads only + // (model.py::JointAttention::_apply_rotary_half chunks along the head axis). + int64_t rope_heads() const { return num_heads / 2; } + + void validate() const; +}; + +// RMSNorm in this model always uses a weight and never a bias. Head-wise norms +// (q_norm / k_norm) carry a (num_heads, head_dim) weight applied after the +// reduction over head_dim. +struct EchoRmsNormWeights { + core::TensorValue weight; +}; + +struct EchoMlpWeights { + modules::LinearWeights w1; + modules::LinearWeights w2; + modules::LinearWeights w3; +}; + +// SelfAttention, used by both encoders. `gate` is applied as +// output * sigmoid(gate) before the output projection. +struct EchoSelfAttentionWeights { + modules::LinearWeights wq; + modules::LinearWeights wk; + modules::LinearWeights wv; + modules::LinearWeights wo; + modules::LinearWeights gate; + EchoRmsNormWeights q_norm; + EchoRmsNormWeights k_norm; +}; + +struct EchoEncoderBlockWeights { + EchoSelfAttentionWeights attention; + EchoMlpWeights mlp; + EchoRmsNormWeights attention_norm; + EchoRmsNormWeights mlp_norm; +}; + +struct EchoTextEncoderWeights { + core::TensorValue text_embedding; + std::vector blocks; +}; + +struct EchoSpeakerEncoderWeights { + modules::LinearWeights in_proj; // (latent_size * patch_size) -> speaker_model_size + std::vector blocks; +}; + +// LowRankAdaLN: each of shift/scale/gate is refined by a rank-256 residual +// MLP, `up(down(silu(v))) + v`. down has no bias, up does. +struct EchoAdaLnWeights { + modules::LinearWeights shift_down; + modules::LinearWeights scale_down; + modules::LinearWeights gate_down; + modules::LinearWeights shift_up; + modules::LinearWeights scale_up; + modules::LinearWeights gate_up; +}; + +struct EchoJointAttentionWeights { + modules::LinearWeights wq; + modules::LinearWeights wk; + modules::LinearWeights wv; + modules::LinearWeights wk_text; + modules::LinearWeights wv_text; + modules::LinearWeights wk_speaker; + modules::LinearWeights wv_speaker; + modules::LinearWeights gate; + modules::LinearWeights wo; + EchoRmsNormWeights q_norm; + EchoRmsNormWeights k_norm; +}; + +struct EchoDitBlockWeights { + EchoJointAttentionWeights attention; + EchoMlpWeights mlp; + EchoAdaLnWeights attention_adaln; + EchoAdaLnWeights mlp_adaln; +}; + +struct EchoDitWeights { + EchoTextEncoderWeights text_encoder; + EchoSpeakerEncoderWeights speaker_encoder; + EchoRmsNormWeights text_norm; + EchoRmsNormWeights speaker_norm; + + modules::LinearWeights cond_0; // timestep_embed_size -> model_size + modules::LinearWeights cond_2; // model_size -> model_size + modules::LinearWeights cond_4; // model_size -> model_size * 3 + + modules::LinearWeights in_proj; // latent_size -> model_size, bias + std::vector blocks; + EchoRmsNormWeights out_norm; + modules::LinearWeights out_proj; // model_size -> latent_size, bias +}; + +// PCA basis mapping the DiT's 80-D working space to the 1024-D Fish z_q space. +// Stored row-major as (latent_size, ae_latent_dim). +struct EchoPcaState { + std::vector components; // latent_size * ae_latent_dim + std::vector mean; // ae_latent_dim + float latent_scale = 1.0F; +}; + +// Request-level sampler configuration, parsed from spec options. +struct EchoSamplerOptions { + int num_steps = 40; + float cfg_scale_text = 3.0F; + float cfg_scale_speaker = 8.0F; + float cfg_min_t = 0.5F; + float cfg_max_t = 1.0F; + // Evaluate the two unconditional lanes only every Nth step inside the CFG + // window, reusing the previous guidance correction in between. 1 reproduces + // upstream exactly. + // + // The correction -- w_text*(v_cond - v_text) + w_speaker*(v_cond - v_speaker) + // -- varies slowly in t even though v_cond does not, so it tolerates being + // resampled. What matters is how many times it is actually measured across + // the guided phase, which is num_steps/2 rounded up, divided by this value: + // + // num_steps=40, interval=2 -> 10 refreshes, 80 -> 60 lane-evals (1.33x) + // num_steps=30, interval=2 -> 8 refreshes, 60 -> 46 lane-evals (1.30x) + // num_steps=14, interval=2 -> 4 refreshes, 28 -> 22 lane-evals (1.27x) + // + // Ten refreshes is dense enough that each reused correction is one small + // t-step stale; four is not. So this is worth raising at 30+ steps and is + // not worth it at 14, where the correction is already coarsely sampled and + // the speaker term's weight of 8.0 multiplies any staleness straight into + // timbre and pronunciation -- a failure mode you hear rather than see. + // + // Note that num_steps=40 with interval=2 and num_steps=30 with interval=1 + // both cost 60 lane-evals. They spend the same compute differently: fewer + // steps coarsens the whole ODE trajectory, while a longer interval leaves + // the trajectory intact and only lets the guidance go stale. Neither + // dominates on paper; compare them by listening before committing. + int cfg_interval = 1; + std::optional truncation_factor = 0.8F; + std::optional speaker_kv_scale; + std::optional speaker_kv_max_layers; + std::optional speaker_kv_min_t; + int64_t sequence_length = 640; + // True when the caller pinned sequence_length explicitly, which suppresses + // the per-chunk window estimate. + bool window_pinned = false; + uint64_t seed = 0; +}; + +} // namespace engine::models::echo_tts diff --git a/include/engine/community_models/echo_tts/dit.h b/include/engine/community_models/echo_tts/dit.h new file mode 100644 index 00000000..7392e0ec --- /dev/null +++ b/include/engine/community_models/echo_tts/dit.h @@ -0,0 +1,73 @@ +#pragma once + +#include "engine/community_models/echo_tts/config.h" +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include + +namespace engine::models::echo_tts { + +// Conditioning for one generation: the tokenized text and the speaker latent, +// both already on the host. Masks are 1.0 for real positions and 0.0 for +// padding; the runtime converts them to additive attention masks. +struct EchoConditioning { + std::vector text_input_ids; + std::vector text_mask; + int64_t text_length = 0; + + std::vector speaker_latent; // (speaker_frames, latent_size), row-major + std::vector speaker_mask; // (speaker_frames) + int64_t speaker_frames = 0; +}; + +// Owns the DiT weights and the two graphs that use them. +// +// The conditioning encoders run once per request and their per-block key/value +// projections are held in a persistent device buffer. The denoiser graph then +// reads those buffers as leaves, so the text and speaker stacks are not +// re-executed on every sampler step. Upstream gets the same effect by passing +// Python lists of cached tensors into the forward call. +class EchoDitRuntime { +public: + EchoDitRuntime( + const EchoTtsConfig & config, + const assets::TensorSource & source, + const std::string & tensor_prefix, + core::ExecutionContext & execution, + assets::TensorStorageType matmul_storage_type); + ~EchoDitRuntime(); + + EchoDitRuntime(const EchoDitRuntime &) = delete; + EchoDitRuntime & operator=(const EchoDitRuntime &) = delete; + + const EchoTtsConfig & config() const noexcept; + + // Runs the text and speaker encoders and populates the cached key/value + // projections. Must be called before sample(). + void prepare_conditioning(const EchoConditioning & conditioning); + + // Runs the dual-CFG Euler sampler and returns the final latent, shaped + // (sequence_length, latent_size) row-major. + std::vector sample(const EchoSamplerOptions & options); + + // One conditional denoiser forward at a fixed timestep, bypassing the + // sampler. Exists so a parity harness can isolate a wrong DiT block from a + // wrong integration step: feeding the reference's own x and t makes any + // difference in the result attributable to the graph alone. + // + // `x` is (sequence_length, latent_size) row-major -- always a SINGLE lane -- + // and sets the sequence length for this call. `lanes` selects how many + // velocity fields come back: 1 (conditional only) or 3 (cond, text-uncond, + // speaker-uncond, concatenated). Requires prepare_conditioning(). + std::vector denoise_once(const std::vector & x, float t, int lanes = 1); + + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::echo_tts diff --git a/include/engine/community_models/echo_tts/latent_post.h b/include/engine/community_models/echo_tts/latent_post.h new file mode 100644 index 00000000..e8c4cfc9 --- /dev/null +++ b/include/engine/community_models/echo_tts/latent_post.h @@ -0,0 +1,41 @@ +#pragma once + +#include "engine/community_models/echo_tts/config.h" + +#include +#include + +namespace engine::models::echo_tts { + +// Forward PCA: Fish z_q (frames, ae_latent_dim) -> DiT latents (frames, latent_size). +// Mirrors inference.py::ae_encode, minus the transposes, which the caller owns. +std::vector pca_project( + const EchoPcaState & pca, + const EchoTtsConfig & config, + const std::vector & z_q, + int64_t frames); + +// Inverse PCA: DiT latents (frames, latent_size) -> Fish z_q (frames, ae_latent_dim). +// Mirrors inference.py::ae_decode. +std::vector pca_unproject( + const EchoPcaState & pca, + const EchoTtsConfig & config, + const std::vector & latents, + int64_t frames); + +// Port of inference.py::find_flattening_point. `latents` is (frames, latent_size) +// row-major. Returns the number of leading frames to keep. +// +// The generated latent tail goes flat once the model has finished speaking, and +// this heuristic is what upstream uses to find that point. It is deliberately +// bit-for-bit faithful, including the unbiased (N-1) variance, because the crop +// index directly sets the output duration. +int64_t find_flattening_point( + const std::vector & latents, + int64_t frames, + int64_t latent_size, + int64_t window_size = 20, + float std_threshold = 0.05F, + float target_value = 0.0F); + +} // namespace engine::models::echo_tts diff --git a/include/engine/community_models/echo_tts/sampler.h b/include/engine/community_models/echo_tts/sampler.h new file mode 100644 index 00000000..1cbd7552 --- /dev/null +++ b/include/engine/community_models/echo_tts/sampler.h @@ -0,0 +1,52 @@ +#pragma once + +#include "engine/community_models/echo_tts/config.h" + +#include +#include +#include + +namespace engine::models::echo_tts { + +// One denoiser evaluation. `x` is (sequence_length * latent_size) and `lanes` is +// 1 (conditional only) or 3 (cond, text-uncond, speaker-uncond, concatenated +// along the batch axis). The result holds `lanes` velocity fields of the same +// per-lane size. +using EchoDenoiseFn = std::function( + const std::vector & x, float t, int lanes)>; + +// Returns the timestep schedule used by +// inference.py::sample_euler_cfg_independent_guidances: +// linspace(1, 0, num_steps + 1) * 0.999 +// The 0.999 scale exists so that temporal rescaling can be applied on the first +// step; it is not a rounding artifact and changes the trajectory if dropped. +std::vector euler_timestep_schedule(int num_steps); + +// True when classifier-free guidance is active at timestep t. Mirrors the +// upstream inclusive comparison on both ends. +bool cfg_active(float t, float cfg_min_t, float cfg_max_t); + +// Combines the three CFG lanes into a single velocity, following upstream's +// independent-guidance form: +// v = v_cond +// + w_text * (v_cond - v_uncond_text) +// + w_speaker * (v_cond - v_uncond_speaker) +std::vector combine_cfg_lanes( + const std::vector & lanes, + int64_t lane_elements, + float cfg_scale_text, + float cfg_scale_speaker); + +// Runs the sampler loop. `denoise` supplies the model evaluation and +// `initial_noise` the starting latent, both injected so this can be tested +// without a backend. `on_kv_rescale`, when set, is invoked at the timestep where +// upstream undoes speaker KV scaling. +std::vector run_euler_sampler( + const EchoSamplerOptions & options, + int64_t sequence_length, + int64_t latent_size, + std::vector initial_noise, + const EchoDenoiseFn & denoise, + const std::function & on_kv_rescale = {}); + +} // namespace engine::models::echo_tts diff --git a/include/engine/community_models/echo_tts/session.h b/include/engine/community_models/echo_tts/session.h new file mode 100644 index 00000000..70305888 --- /dev/null +++ b/include/engine/community_models/echo_tts/session.h @@ -0,0 +1,98 @@ +#pragma once + +#include "engine/community_models/echo_tts/config.h" +#include "engine/framework/assets/resource_bundle.h" +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/fish_audio/assets.h" +#include "engine/models/fish_audio/codec.h" + +#include +#include +#include +#include + +namespace engine::models::echo_tts { + +class EchoDitRuntime; + +struct EchoTtsAssets { + assets::ResourceBundle resources; + EchoTtsConfig config; + EchoPcaState pca; + std::shared_ptr dit_weights; + // The Fish S1-DAC autoencoder, packaged inside Echo's own GGUF. audio.cpp + // implements this codec for the fish_audio family; Echo reuses the + // implementation but supplies the S1 weights its PCA basis was fitted to. + // See docs/community_models/echo_tts_autoencoder_reuse.md. + std::shared_ptr codec_assets; +}; + +// Encoding a speaker reference is linear in its length -- a 4.5-minute clip is +// ten Fish encode passes -- and the result depends only on the audio and the +// trim length, so it is cached across requests. Servers reusing a handful of +// voices then pay it once per voice rather than once per request. +struct EchoReferenceIdentity { + std::string id; + int64_t max_samples = 0; +}; + +struct EchoReferenceIdentityEqual { + bool operator()(const EchoReferenceIdentity & a, const EchoReferenceIdentity & b) const { + return a.max_samples == b.max_samples && a.id == b.id; + } +}; + +struct EchoPreparedSpeaker { + std::vector latent; + int64_t frames = 0; +}; + +std::shared_ptr make_echo_tts_loader(); + +class EchoTtsSession final + : public runtime::RuntimeSessionBase, + public runtime::IOfflineVoiceTaskSession { +public: + EchoTtsSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~EchoTtsSession() 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; + void reset(); + +private: + // Reference trim limit: request option, else session default, else the + // trained maximum. Returned in samples at the codec rate. + int64_t resolve_reference_max_samples( + const std::unordered_map & request_options) const; + EchoSamplerOptions parse_sampler_options( + const std::unordered_map & options) const; + // Encodes reference audio to 80-D PCA latents, mirroring + // inference.py::get_speaker_latent_and_mask. + void encode_speaker(const runtime::AudioBuffer & audio); + runtime::AudioBuffer synthesize_chunk( + const std::string & text, + const EchoSamplerOptions & sampler); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::unique_ptr dit_; + std::unique_ptr codec_; + int64_t reference_max_samples_ = 0; + std::vector speaker_latent_; + int64_t speaker_frames_ = 0; + runtime::CacheSlots + reference_cache_; +}; + +} // namespace engine::models::echo_tts diff --git a/include/engine/community_models/echo_tts/tokenizer.h b/include/engine/community_models/echo_tts/tokenizer.h new file mode 100644 index 00000000..78538770 --- /dev/null +++ b/include/engine/community_models/echo_tts/tokenizer.h @@ -0,0 +1,31 @@ +#pragma once + +#include +#include +#include + +namespace engine::models::echo_tts { + +struct EchoTokenizedText { + std::vector input_ids; + std::vector mask; // 1.0 for real tokens, 0.0 for padding + std::string normalized_text; + bool truncated = false; +}; + +// Applies the WhisperD-style normalisation from inference.py::tokenizer_encode +// and returns the normalised string. Exposed separately because the session +// reports the normalised text back to the caller. +std::string normalize_echo_text(const std::string & text); + +// Byte-level tokenizer: a BOS 0 followed by the raw UTF-8 bytes of the +// normalised text. `max_length` is the hard cap (768 upstream) and counts the +// BOS. When `pad_to_max` is false the returned vectors are exactly as long as +// the encoded text. +EchoTokenizedText tokenize_echo_text( + const std::string & text, + int64_t max_length, + bool normalize = true, + bool pad_to_max = false); + +} // namespace engine::models::echo_tts diff --git a/include/engine/models/fish_audio/codec.h b/include/engine/models/fish_audio/codec.h index 2da7556e..e5547fbb 100644 --- a/include/engine/models/fish_audio/codec.h +++ b/include/engine/models/fish_audio/codec.h @@ -5,7 +5,9 @@ #include "engine/models/fish_audio/assets.h" #include "engine/models/fish_audio/types.h" +#include #include +#include namespace engine::models::fish_audio { @@ -23,6 +25,12 @@ class FishAudioCodecRuntime { FishAudioCodes encode_reference(const runtime::AudioBuffer & audio); runtime::AudioBuffer decode(const FishAudioCodes & codes); + + // Continuous-latent access to the same autoencoder. Echo-TTS conditions on + // and generates z_q directly and never materialises codebook indices. + // `values` is (frames, channels) row-major. + FishAudioLatents encode_zq(const runtime::AudioBuffer & audio); + runtime::AudioBuffer decode_zq(const std::vector & latents, int64_t frames); void release_encode_graph(); void release_runtime_graphs(); diff --git a/include/engine/models/fish_audio/types.h b/include/engine/models/fish_audio/types.h index 8062b3ea..12929059 100644 --- a/include/engine/models/fish_audio/types.h +++ b/include/engine/models/fish_audio/types.h @@ -9,6 +9,14 @@ namespace engine::models::fish_audio { +// Continuous quantiser latents (z_q), the boundary Echo-TTS shares with this +// autoencoder. Stored (frames, channels) row-major. +struct FishAudioLatents { + int64_t frames = 0; + int64_t channels = 0; + std::vector values; +}; + struct FishAudioGenerationOptions { int64_t max_new_tokens = 1024; int64_t text_chunk_size = 200; diff --git a/model_specs/echo_tts.json b/model_specs/echo_tts.json new file mode 100644 index 00000000..cfe19b31 --- /dev/null +++ b/model_specs/echo_tts.json @@ -0,0 +1,175 @@ +{ + "schema_version": 1, + "family": "echo_tts", + "display_name": "Echo-TTS", + "description": "Echo-TTS is an English zero-shot voice-cloning TTS model packaged for audio.cpp. A 2.8B diffusion transformer generates 80-D latents in PCA space which the Fish S1-DAC decodes to 44.1 kHz audio. Generation is a fixed 29.72 s window (640 latents).", + "category": "tts", + "status": "experimental", + "tasks": [ + "clone" + ], + "modes": [ + "offline" + ], + "languages": [ + "en" + ], + "runtime": { + "tags": [ + "gguf" + ] + }, + "capabilities": { + "clone": [ + "speaker_reference" + ] + }, + "options": { + "request": [ + { + "name": "target_voice", + "type": "audio_path", + "description": "Reference audio path for zero-shot cloning. Wav only; no transcript required.", + "required": false + }, + { + "name": "cfg_scale_text", + "type": "float", + "description": "Classifier-free guidance scale on the text condition.", + "required": false, + "min": 0.0, + "default": 3.0 + }, + { + "name": "cfg_scale_speaker", + "type": "float", + "description": "Classifier-free guidance scale on the speaker condition.", + "required": false, + "min": 0.0, + "default": 8.0 + }, + { + "name": "num_steps", + "type": "int", + "description": "Euler sampler steps.", + "required": false, + "min": 1, + "default": 40 + }, + { + "name": "truncation_factor", + "type": "float", + "description": "Initial-noise truncation factor.", + "required": false, + "min": 0.0, + "max": 1.0, + "default": 0.8 + }, + { + "name": "speaker_kv_scale", + "type": "float", + "description": "Force-speaker KV scaling. 1.0 disables; 1.5 is the upstream default when enabled.", + "required": false, + "min": 1.0, + "default": 1.0 + }, + { + "name": "seed", + "type": "int", + "description": "RNG seed for the initial latent.", + "required": false, + "default": 0 + }, + { + "name": "reference_max_seconds", + "type": "float", + "required": false, + "description": "Trim the speaker reference to at most this many seconds before encoding. Shorter references are cheaper and often clone better; upstream's guidance favours around 10 s. Defaults to 15 s. Values above the trained maximum of 297.1 s are clamped. Set per request, or as a session default from CLI or server config.", + "default": 15.0 + }, + { + "name": "sequence_length", + "type": "int", + "description": "Pin the generation window in latents (1..640; 640 = 29.72 s). Left unset, the window is estimated per chunk from text length and widened automatically if the utterance does not finish inside it, which is substantially cheaper for short text.", + "required": false, + "min": 1, + "max": 640 + }, + { + "name": "cfg_interval", + "type": "int", + "description": "Refresh the two unconditional CFG lanes only every Nth step inside the guidance window, reusing the previous guidance correction in between. 1 reproduces upstream exactly. Worth raising to 2 at 30+ steps, where the correction is still measured 8-10 times across the guided phase (roughly 1.3x fewer denoiser evaluations); not recommended below ~20 steps, where it is already coarsely sampled and staleness shows up as degraded timbre and text adherence.", + "required": false, + "min": 1, + "default": 1 + } + ], + "session": [ + { + "name": "reference_max_seconds", + "type": "float", + "required": false, + "description": "Trim the speaker reference to at most this many seconds before encoding. Shorter references are cheaper and often clone better; upstream's guidance favours around 10 s. Defaults to 15 s. Values above the trained maximum of 297.1 s are clamped. Set per request, or as a session default from CLI or server config.", + "default": 15.0 + }, + { + "name": "reference_cache_slots", + "type": "int", + "required": false, + "description": "How many encoded speaker references to keep. Encoding is linear in reference length, so reusing a voice across requests avoids repeating it. 0 disables caching; default 4." + } + ], + "load": [] + }, + "packages": [ + { + "id": "echo_tts_orig", + "display_name": "Echo-TTS Original-Dtype GGUF", + "default": true, + "format": "gguf", + "precision": "orig", + "target_directory": "Echo-TTS-GGUF", + "files": [ + "Echo-TTS-GGUF/model.gguf" + ], + "download": { + "kind": "unsupported", + "reason": "Echo-TTS model packaging is not implemented yet." + }, + "strip_prefix": "Echo-TTS-GGUF" + } + ], + "dependencies": [], + "ui": { + "recommended_package": "echo_tts_orig", + "tags": [ + "TTS", + "Clone", + "GGUF" + ], + "docs": [] + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "tensors": { + "dit_weights": { + "source": "weights:", + "prefix": "dit_weights" + }, + "pca": { + "source": "weights:", + "prefix": "pca" + }, + "codec_weights": { + "source": "weights:", + "prefix": "ae" + } + } + } + ] +} diff --git a/src/community_models/echo_tts/dit.cpp b/src/community_models/echo_tts/dit.cpp new file mode 100644 index 00000000..ca9cbc20 --- /dev/null +++ b/src/community_models/echo_tts/dit.cpp @@ -0,0 +1,919 @@ +#include "dit_blocks.inc" + +namespace engine::models::echo_tts { +namespace { + +// Joint attention over [self | text | speaker]. Two details differ from an +// ordinary cross-attention block and neither is caught by a shape check: +// +// 1. RoPE covers only the first half of the heads. Upstream's +// _apply_rotary_half chunks along dim=-2, which is the head axis, so heads +// 0..7 rotate and 8..15 do not. +// 2. The text and speaker keys arrive already k_norm'd from the cache. They +// must not be normalised again here, and they never receive RoPE. +core::TensorValue joint_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const core::TensorValue & k_text, + const core::TensorValue & v_text, + const core::TensorValue & k_speaker, + const core::TensorValue & v_speaker, + const core::TensorValue & mask, + const EchoJointAttentionWeights & weights, + const EchoTtsConfig & config) { + const int64_t D = config.model_size; + const int64_t heads = config.num_heads; + const int64_t head_dim = config.head_dim(); + const int64_t rope_heads = config.rope_heads(); + const int64_t batch = input.shape.dims[0]; + const int64_t seq = input.shape.dims[1]; + + auto q = modules::LinearModule({D, D, false, GGML_PREC_F32}).build(ctx, input, weights.wq); + auto k = modules::LinearModule({D, D, false, GGML_PREC_F32}).build(ctx, input, weights.wk); + auto v = modules::LinearModule({D, D, false, GGML_PREC_F32}).build(ctx, input, weights.wv); + + q = head_rms_norm(ctx, reshape_heads(ctx, q, heads, head_dim), weights.q_norm.weight, config.norm_eps); + k = head_rms_norm(ctx, reshape_heads(ctx, k, heads, head_dim), weights.k_norm.weight, config.norm_eps); + v = reshape_heads(ctx, v, heads, head_dim); + + const modules::RoPEModule rope({head_dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}); + auto rotate_half = [&](const core::TensorValue & value) { + auto front = modules::SliceModule({2, 0, rope_heads}).build(ctx, value); + auto back = modules::SliceModule({2, rope_heads, heads - rope_heads}).build(ctx, value); + front = rope.build(ctx, contiguous(ctx, front), positions); + return modules::ConcatModule({2}).build(ctx, contiguous(ctx, front), contiguous(ctx, back)); + }; + q = rotate_half(q); + k = rotate_half(k); + + // Broadcast the batch-1 cache across the CFG lanes, matching upstream's + // _concat_kv_caches(cond, cond, cond). + auto expand = [&](const core::TensorValue & value) { + if (value.shape.dims[0] == batch) { + return contiguous(ctx, value); + } + return contiguous( + ctx, + modules::RepeatModule({core::TensorShape::from_dims( + {batch, value.shape.dims[1], heads, head_dim})}) + .build(ctx, contiguous(ctx, value))); + }; + + // Sequence-axis order is self, text, speaker; the mask uses the same order. + auto k_all = modules::ConcatModule({1}).build(ctx, contiguous(ctx, k), expand(k_text)); + k_all = modules::ConcatModule({1}).build(ctx, contiguous(ctx, k_all), expand(k_speaker)); + auto v_all = modules::ConcatModule({1}).build(ctx, contiguous(ctx, v), expand(v_text)); + v_all = modules::ConcatModule({1}).build(ctx, contiguous(ctx, v_all), expand(v_speaker)); + + // Flash attention avoids materialising the (lanes, heads, seq, keys) scores + // tensor, which at 640 queries is 90-370 MB per attention and is the largest + // per-request allocation in the model. It is safe here because the self + // block of the mask is never masked, so no query row can be fully masked and + // the softmax is always well defined. Set AUDIOCPP_ECHO_TTS_NO_FLASH=1 to + // fall back to the explicit lowering for comparison. + modules::ScaledDotProductAttentionConfig attn_config; + attn_config.head_dim = head_dim; + attn_config.lowering = echo_flash_disabled() + ? modules::ScaledDotProductAttentionLowering::Explicit + : modules::ScaledDotProductAttentionLowering::Flash; + attn_config.precision = GGML_PREC_F32; + attn_config.causality = modules::AttentionCausality::NonCausal; + auto context = modules::ScaledDotProductAttentionModule(attn_config) + .build(ctx, to_bhsd(ctx, q), to_bhsd(ctx, k_all), to_bhsd(ctx, v_all), mask); + + context = core::reshape_tensor( + ctx, contiguous(ctx, context), core::TensorShape::from_dims({batch, seq, D})); + context = apply_attention_gate(ctx, context, input, weights.gate, D); + return modules::LinearModule({D, D, false, GGML_PREC_F32}).build(ctx, context, weights.wo); +} + +core::TensorValue dit_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & x, + const core::TensorValue & cond_embed, + const core::TensorValue & positions, + const core::TensorValue & k_text, + const core::TensorValue & v_text, + const core::TensorValue & k_speaker, + const core::TensorValue & v_speaker, + const core::TensorValue & mask, + const EchoDitBlockWeights & weights, + const EchoTtsConfig & config) { + auto attn_ada = adaln( + ctx, x, cond_embed, weights.attention_adaln, + config.model_size, config.adaln_rank, config.norm_eps); + auto attn = joint_attention( + ctx, attn_ada.normed, positions, k_text, v_text, k_speaker, v_speaker, + mask, weights.attention, config); + // gate is (lanes, 1, dim); it scales every sequence position. + attn = broadcast_mul(ctx, attn, attn_ada.gate); + auto hidden = modules::AddModule{}.build(ctx, x, attn); + + auto mlp_ada = adaln( + ctx, hidden, cond_embed, weights.mlp_adaln, + config.model_size, config.adaln_rank, config.norm_eps); + auto mlp_out = mlp(ctx, mlp_ada.normed, weights.mlp, config.model_size, config.intermediate_size); + mlp_out = broadcast_mul(ctx, mlp_out, mlp_ada.gate); + return modules::AddModule{}.build(ctx, hidden, mlp_out); +} + +// model.py::get_timestep_embedding, evaluated on the host because it depends +// only on t, which changes once per sampler step. +std::vector timestep_embedding(float t, int64_t embed_size, int64_t lanes) { + const int64_t half = embed_size / 2; + std::vector out(static_cast(lanes * embed_size)); + for (int64_t i = 0; i < half; ++i) { + const double freq = + 1000.0 * std::exp(-std::log(10000.0) * static_cast(i) / + static_cast(half)); + const double arg = static_cast(t) * freq; + const auto cos_v = static_cast(std::cos(arg)); + const auto sin_v = static_cast(std::sin(arg)); + for (int64_t lane = 0; lane < lanes; ++lane) { + float * row = out.data() + lane * embed_size; + row[i] = cos_v; + row[half + i] = sin_v; + } + } + return out; +} + +// Parity debugging. Set AUDIOCPP_ECHO_TTS_DEBUG=1 to tap every DiT block output +// and the two encoder outputs, printing the same mean/std/min/max summary that +// tools/community_models/echo_tts_reference.py emits, so a C++ run can be +// compared block by block against the reference dump. +bool echo_debug_enabled() { + static const bool enabled = [] { + const char * value = std::getenv("AUDIOCPP_ECHO_TTS_DEBUG"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return enabled; +} + +void print_tensor_stats(const std::string & label, const std::vector & values) { + if (values.empty()) { + std::fprintf(stderr, " %-28s \n", label.c_str()); + return; + } + double sum = 0.0; + double sum_sq = 0.0; + float low = values[0]; + float high = values[0]; + for (const float value : values) { + sum += value; + sum_sq += static_cast(value) * value; + low = std::min(low, value); + high = std::max(high, value); + } + const double mean = sum / static_cast(values.size()); + const double variance = sum_sq / static_cast(values.size()) - mean * mean; + std::fprintf( + stderr, + " %-28s mean=%+.6f std=%.6f min=%+.4f max=%+.4f n=%zu\n", + label.c_str(), mean, variance > 0.0 ? std::sqrt(variance) : 0.0, + static_cast(low), static_cast(high), values.size()); +} + +std::vector iota_positions(int64_t count) { + std::vector positions(static_cast(count)); + for (int64_t i = 0; i < count; ++i) { + positions[static_cast(i)] = static_cast(i); + } + return positions; +} + +} // namespace + +// --- runtime ------------------------------------------------------------ + +class EchoDitRuntime::Impl { +public: + Impl( + const EchoTtsConfig & config, + const assets::TensorSource & source, + const std::string & tensor_prefix, + core::ExecutionContext & execution, + assets::TensorStorageType matmul_storage_type) + : config_(config), + execution_(execution), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + store_(backend_, backend_type_, "Echo-TTS DiT weights", kWeightContextBytes) { + config_.validate(); + if (backend_ == nullptr) { + throw std::runtime_error("Echo-TTS DiT backend initialization failed"); + } + weights_ = load_dit_weights(config_, store_, source, tensor_prefix, matmul_storage_type); + store_.upload(); + } + + ~Impl() { release_all(); } + + const EchoTtsConfig & config() const noexcept { return config_; } + bool conditioning_ready() const noexcept { return conditioning_ready_; } + + void prepare_conditioning(const EchoConditioning & conditioning) { + validate_conditioning(conditioning); + + // Any change in conditioning length invalidates every cached graph, + // because all of them are built for fixed key counts. + release_denoiser_graphs(); + release_conditioning_graph(); + release_kv_cache(); + + text_length_ = conditioning.text_length; + speaker_frames_ = conditioning.speaker_frames; + speaker_tokens_ = speaker_frames_ / config_.speaker_patch_size; + + text_mask_ = conditioning.text_mask; + speaker_mask_.assign(static_cast(speaker_tokens_), 0.0F); + for (int64_t i = 0; i < speaker_tokens_; ++i) { + // model.py subsamples the speaker mask by the patch size before use. + speaker_mask_[static_cast(i)] = + conditioning.speaker_mask[static_cast(i * config_.speaker_patch_size)]; + } + + allocate_kv_cache(); + build_conditioning_graph(); + + core::write_tensor_i32(text_ids_, conditioning.text_input_ids); + core::write_tensor_f32(text_attn_mask_, make_text_self_mask()); + core::write_tensor_i32(text_positions_, iota_positions(text_length_)); + core::write_tensor_f32(speaker_latent_, conditioning.speaker_latent); + core::write_tensor_i32(speaker_positions_, iota_positions(speaker_tokens_)); + + core::set_backend_threads(backend_, threads_); + const auto status = core::compute_graph( + execution_, conditioning_graph_, conditioning_plan_, "echo_tts.conditioning"); + // compute_graph does not synchronise. On CUDA the copies into the + // persistent KV cache are still in flight when it returns, so the + // denoiser would read whatever the buffer happened to hold. + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Echo-TTS conditioning graph execution failed"); + } + + if (echo_debug_enabled()) { + std::fprintf(stderr, "\n[echo_tts] conditioning: text_len=%lld speaker_frames=%lld " + "speaker_tokens=%lld\n", + static_cast(text_length_), + static_cast(speaker_frames_), + static_cast(speaker_tokens_)); + print_tensor_stats("kv_text.0.k", core::read_tensor_f32(kv_.k_text[0].tensor)); + print_tensor_stats("kv_text.23.k", + core::read_tensor_f32(kv_.k_text[kv_.k_text.size() - 1].tensor)); + print_tensor_stats("kv_speaker.0.k", core::read_tensor_f32(kv_.k_speaker[0].tensor)); + print_tensor_stats("kv_speaker.23.k", + core::read_tensor_f32(kv_.k_speaker[kv_.k_speaker.size() - 1].tensor)); + std::fflush(stderr); + } + + // The encoders are not needed again for this request; only the cached + // projections they wrote into the persistent buffer are. + release_conditioning_graph(); + conditioning_ready_ = true; + debug_printed_ = false; + } + + std::vector denoise(const std::vector & x, float t, int lanes) { + if (!conditioning_ready_) { + throw std::runtime_error("Echo-TTS denoise() called before prepare_conditioning()"); + } + if (lanes != 1 && lanes != 3) { + throw std::runtime_error("Echo-TTS denoiser supports 1 or 3 CFG lanes"); + } + const int64_t elements = sequence_length_ * config_.latent_size; + if (static_cast(x.size()) != elements) { + throw std::runtime_error("Echo-TTS denoiser received a mis-shaped latent"); + } + + auto & graph = lanes == 1 ? single_ : triple_; + if (graph.graph == nullptr) { + build_denoiser_graph(graph, lanes); + } + + // x and the timestep embedding change every step and live in gallocr + // space, so they are rewritten here. Positions and the mask are constant + // and live in the graph's own persistent buffer; see DenoiserGraph. + for (int lane = 0; lane < lanes; ++lane) { + core::write_tensor_f32_slice( + graph.x, static_cast(lane * elements), x.data(), x.size()); + } + core::write_tensor_f32( + graph.timestep, timestep_embedding(t, config_.timestep_embed_size, lanes)); + + core::set_backend_threads(backend_, threads_); + const auto status = + core::compute_graph(execution_, graph.graph, graph.plan, "echo_tts.denoise"); + // Must complete before the output is read back to the host; without this + // the sampler integrates stale device memory. + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Echo-TTS denoiser graph execution failed"); + } + if (echo_debug_enabled() && !debug_printed_) { + debug_printed_ = true; + std::fprintf(stderr, "[echo_tts] denoiser t=%.4f lanes=%d keys=%lld\n", + t, lanes, + static_cast(sequence_length_ + text_length_ + speaker_tokens_)); + print_tensor_stats("dit.v_pred", core::read_tensor_f32(graph.output)); + std::fflush(stderr); + } + return core::read_tensor_f32(graph.output); + } + + void set_sequence_length(int64_t sequence_length) { + if (sequence_length <= 0 || sequence_length > config_.max_sequence_length) { + throw std::runtime_error("Echo-TTS sequence_length is out of range"); + } + if (sequence_length != sequence_length_) { + release_denoiser_graphs(); + sequence_length_ = sequence_length; + } + } + + std::vector initial_noise(const EchoSamplerOptions & options) const { + const size_t count = + static_cast(options.sequence_length * config_.latent_size); + // Reproduces torch's CUDA Philox stream so a fixed seed is comparable + // against the reference implementation. + auto noise = sampling::generate_torch_cuda_randn(count, options.seed); + if (echo_debug_enabled() && noise.size() >= 3) { + std::fprintf( + stderr, + "[echo_tts] seed=%llu steps=%d cfg_text=%.2f cfg_speaker=%.2f " + "truncation=%.2f\n initial_noise first3=[% .6f % .6f % .6f]\n", + static_cast(options.seed), + options.num_steps, + static_cast(options.cfg_scale_text), + static_cast(options.cfg_scale_speaker), + options.truncation_factor.has_value() + ? static_cast(*options.truncation_factor) : 1.0, + static_cast(noise[0]), static_cast(noise[1]), + static_cast(noise[2])); + std::fflush(stderr); + } + return noise; + } + + // Scales the cached speaker keys and values in place, matching + // inference.py::_multiply_kv_cache. Done on a host round trip because the + // cache is a plain backend buffer with no graph attached. + void scale_speaker_kv(float scale, std::optional max_layers) { + const int64_t limit = max_layers.has_value() + ? std::min(*max_layers, config_.num_layers) + : config_.num_layers; + for (int64_t layer = 0; layer < limit; ++layer) { + scale_tensor_in_place(kv_.k_speaker[static_cast(layer)], scale); + scale_tensor_in_place(kv_.v_speaker[static_cast(layer)], scale); + } + } + + void release_conditioning_graph() { + conditioning_plan_.reset(); + if (conditioning_graph_ != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, conditioning_graph_); + conditioning_graph_ = nullptr; + } + if (conditioning_alloc_ != nullptr) { + ggml_gallocr_free(conditioning_alloc_); + conditioning_alloc_ = nullptr; + } + conditioning_ctx_.reset(); + } + +private: + struct DenoiserGraph { + // Positions and the attention mask are constant once the graph is built, + // but ggml_gallocr reclaims an input's memory after its last consumer and + // reuses it for intermediates, so a value written into gallocr space does + // not survive the next compute. They live in their own context and + // backend buffer: pre-allocated tensors are skipped by gallocr, so they + // are written once and read by every sampler step. This also avoids + // rebuilding and re-uploading a multi-megabyte mask 40 times per chunk. + GgmlContextPtr const_ctx; + ggml_backend_buffer_t const_buffer = nullptr; + GgmlContextPtr ctx; + ggml_cgraph * graph = nullptr; + ggml_gallocr_t alloc = nullptr; + core::HostGraphPlan plan; + core::TensorValue x; + core::TensorValue timestep; + core::TensorValue positions; + core::TensorValue mask; + ggml_tensor * output = nullptr; + int lanes = 0; + }; + + struct KvCache { + std::vector k_text; + std::vector v_text; + std::vector k_speaker; + std::vector v_speaker; + }; + + void validate_conditioning(const EchoConditioning & c) const { + if (c.text_length <= 0 || c.text_length > config_.max_text_length) { + throw std::runtime_error("Echo-TTS text length out of range"); + } + if (static_cast(c.text_input_ids.size()) != c.text_length || + static_cast(c.text_mask.size()) != c.text_length) { + throw std::runtime_error("Echo-TTS text buffers disagree with text_length"); + } + if (c.speaker_frames < config_.speaker_patch_size || + c.speaker_frames % config_.speaker_patch_size != 0) { + throw std::runtime_error( + "Echo-TTS speaker latent length must be a positive multiple of the patch size"); + } + if (c.speaker_frames > config_.max_speaker_latent_length) { + throw std::runtime_error("Echo-TTS speaker latent exceeds the trained maximum"); + } + if (static_cast(c.speaker_latent.size()) != + c.speaker_frames * config_.latent_size) { + throw std::runtime_error("Echo-TTS speaker latent has an unexpected element count"); + } + if (static_cast(c.speaker_mask.size()) != c.speaker_frames) { + throw std::runtime_error("Echo-TTS speaker mask disagrees with speaker_frames"); + } + } + + // Bidirectional text encoder mask: padded key positions are suppressed for + // every query row. + std::vector make_text_self_mask() const { + std::vector mask(static_cast(text_length_ * text_length_), 0.0F); + for (int64_t q = 0; q < text_length_; ++q) { + float * row = mask.data() + q * text_length_; + for (int64_t k = 0; k < text_length_; ++k) { + row[k] = text_mask_[static_cast(k)] > 0.5F ? 0.0F : kMaskedBias; + } + } + return mask; + } + + // Denoiser mask, laid out per lane as [self | text | speaker]. Lane 0 is + // fully conditional, lane 1 drops text, lane 2 drops speaker, reproducing + // upstream's concatenated cond/uncond masks. + std::vector make_denoiser_mask(int lanes) const { + const int64_t keys = sequence_length_ + text_length_ + speaker_tokens_; + std::vector mask( + static_cast(static_cast(lanes) * sequence_length_ * keys), 0.0F); + for (int lane = 0; lane < lanes; ++lane) { + const bool text_on = lane != 1; + const bool speaker_on = lane != 2; + for (int64_t q = 0; q < sequence_length_; ++q) { + float * row = mask.data() + + (static_cast(lane) * sequence_length_ + q) * keys; + const float masked = echo_flash_disabled() ? kMaskedBias : kMaskedBiasF16; + for (int64_t i = 0; i < text_length_; ++i) { + const bool keep = text_on && text_mask_[static_cast(i)] > 0.5F; + row[sequence_length_ + i] = keep ? 0.0F : masked; + } + for (int64_t i = 0; i < speaker_tokens_; ++i) { + const bool keep = speaker_on && speaker_mask_[static_cast(i)] > 0.5F; + row[sequence_length_ + text_length_ + i] = keep ? 0.0F : masked; + } + } + } + return mask; + } + + void scale_tensor_in_place(const core::TensorValue & tensor, float scale) { + auto values = core::read_tensor_f32(tensor.tensor); + for (auto & value : values) { + value *= scale; + } + core::write_tensor_f32(tensor, values); + } + + // The cache lives in its own context and backend buffer so that both the + // conditioning graph (which writes it) and the denoiser graphs (which read + // it) can reference the same tensors as leaves. + void allocate_kv_cache() { + const int64_t heads = config_.num_heads; + const int64_t head_dim = config_.head_dim(); + const size_t tensor_count = static_cast(config_.num_layers) * 4; + ggml_init_params params{ + ggml_tensor_overhead() * (tensor_count + 16), nullptr, true}; + kv_ctx_.reset(ggml_init(params)); + if (kv_ctx_ == nullptr) { + throw std::runtime_error("Echo-TTS KV cache context initialization failed"); + } + core::ModuleBuildContext ctx{kv_ctx_.get(), "echo_tts.kv_cache", backend_type_}; + // The cache is allocated at the widest lane count any graph will ask + // for, so joint_attention's expand() finds dims[0] already equal to the + // batch and short-circuits to a passthrough. The single-lane graph + // reads lane 0, which is a contiguous prefix because the lane axis is + // outermost in ggml's layout. + kv_lanes_ = echo_kv_expand_disabled() ? 1 : kMaxCfgLanes; + auto make = [&](int64_t tokens) { + return core::make_tensor( + ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({kv_lanes_, tokens, heads, head_dim})); + }; + for (int64_t layer = 0; layer < config_.num_layers; ++layer) { + kv_.k_text.push_back(make(text_length_)); + kv_.v_text.push_back(make(text_length_)); + kv_.k_speaker.push_back(make(speaker_tokens_)); + kv_.v_speaker.push_back(make(speaker_tokens_)); + } + kv_buffer_ = ggml_backend_alloc_ctx_tensors(kv_ctx_.get(), backend_); + if (kv_buffer_ == nullptr) { + throw std::runtime_error("Echo-TTS KV cache buffer allocation failed"); + } + } + + void build_conditioning_graph() { + ggml_init_params params{kGraphArenaBytes, nullptr, true}; + conditioning_ctx_.reset(ggml_init(params)); + if (conditioning_ctx_ == nullptr) { + throw std::runtime_error("Echo-TTS conditioning graph context initialization failed"); + } + core::ModuleBuildContext ctx{conditioning_ctx_.get(), "echo_tts.conditioning", backend_type_}; + + const int64_t TD = config_.text_model_size; + const int64_t SD = config_.speaker_model_size; + const int64_t D = config_.model_size; + const int64_t heads = config_.num_heads; + const int64_t head_dim = config_.head_dim(); + + text_ids_ = core::make_tensor( + ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, text_length_})); + text_attn_mask_ = core::make_tensor( + ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, text_length_, text_length_})); + text_positions_ = core::make_tensor( + ctx, GGML_TYPE_I32, core::TensorShape::from_dims({text_length_})); + speaker_latent_ = core::make_tensor( + ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({1, speaker_frames_, config_.latent_size})); + speaker_positions_ = core::make_tensor( + ctx, GGML_TYPE_I32, core::TensorShape::from_dims({speaker_tokens_})); + for (auto * input : {text_ids_.tensor, text_attn_mask_.tensor, text_positions_.tensor, + speaker_latent_.tensor, speaker_positions_.tensor}) { + ggml_set_input(input); + } + + // Text encoder: byte embedding, then bidirectional blocks. + auto text_state = modules::EmbeddingModule({config_.text_vocab_size, TD}) + .build(ctx, text_ids_, weights_.text_encoder.text_embedding); + const std::optional text_mask_opt{text_attn_mask_}; + for (const auto & block : weights_.text_encoder.blocks) { + text_state = encoder_block( + ctx, text_state, text_positions_, text_mask_opt, block, + TD, config_.text_intermediate_size, config_.text_num_heads, + config_.norm_eps, false); + } + text_state = rms_norm(ctx, text_state, weights_.text_norm.weight, config_.norm_eps); + + // Speaker encoder: patchify by folding groups of `patch_size` frames into + // the feature axis, project, then causal blocks. The /6 scale is + // upstream's activation-dynamics fix, not a normalisation. + auto speaker_state = core::reshape_tensor( + ctx, + contiguous(ctx, speaker_latent_), + core::TensorShape::from_dims( + {1, speaker_tokens_, config_.latent_size * config_.speaker_patch_size})); + speaker_state = modules::LinearModule( + {config_.latent_size * config_.speaker_patch_size, SD, true, GGML_PREC_F32}) + .build(ctx, speaker_state, weights_.speaker_encoder.in_proj); + speaker_state = core::wrap_tensor( + ggml_scale(ctx.ggml, contiguous(ctx, speaker_state).tensor, 1.0F / 6.0F), + speaker_state.shape, + GGML_TYPE_F32); + const std::optional no_mask; + for (const auto & block : weights_.speaker_encoder.blocks) { + speaker_state = encoder_block( + ctx, speaker_state, speaker_positions_, no_mask, block, + SD, config_.speaker_intermediate_size, config_.speaker_num_heads, + config_.norm_eps, true); + } + speaker_state = rms_norm(ctx, speaker_state, weights_.speaker_norm.weight, config_.norm_eps); + + conditioning_graph_ = ggml_new_graph_custom(conditioning_ctx_.get(), 1048576, false); + + // Project the encoder outputs into each block's key/value space and copy + // the result into the persistent cache. Keys are k_norm'd here, exactly + // once, so the denoiser must not normalise them again. + for (int64_t layer = 0; layer < config_.num_layers; ++layer) { + const size_t index = static_cast(layer); + const auto & attn = weights_.blocks[index].attention; + auto project = [&](const core::TensorValue & state, + const modules::LinearWeights & weight, + int64_t in_dim, + int64_t tokens, + bool normalise) { + auto value = modules::LinearModule({in_dim, D, false, GGML_PREC_F32}) + .build(ctx, state, weight); + value = core::reshape_tensor( + ctx, contiguous(ctx, value), + core::TensorShape::from_dims({1, tokens, heads, head_dim})); + if (normalise) { + value = head_rms_norm(ctx, value, attn.k_norm.weight, config_.norm_eps); + } + if (kv_lanes_ > 1) { + // Broadcast to the cache's lane count once, here, instead of + // once per layer per sampler step inside joint_attention. + // Upstream's _concat_kv_caches(cond, cond, cond) is the same + // operation; only its position in the schedule changes. + value = modules::RepeatModule( + {core::TensorShape::from_dims( + {kv_lanes_, tokens, heads, head_dim})}) + .build(ctx, contiguous(ctx, value)); + } + return value; + }; + struct Slot { + core::TensorValue source; + core::TensorValue destination; + }; + const Slot slots[] = { + {project(text_state, attn.wk_text, TD, text_length_, true), kv_.k_text[index]}, + {project(text_state, attn.wv_text, TD, text_length_, false), kv_.v_text[index]}, + {project(speaker_state, attn.wk_speaker, SD, speaker_tokens_, true), kv_.k_speaker[index]}, + {project(speaker_state, attn.wv_speaker, SD, speaker_tokens_, false), kv_.v_speaker[index]}, + }; + for (const auto & slot : slots) { + auto * copy = ggml_cpy( + ctx.ggml, contiguous(ctx, slot.source).tensor, slot.destination.tensor); + ggml_build_forward_expand(conditioning_graph_, copy); + } + } + + conditioning_alloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (conditioning_alloc_ == nullptr || + !ggml_gallocr_reserve(conditioning_alloc_, conditioning_graph_) || + !ggml_gallocr_alloc_graph(conditioning_alloc_, conditioning_graph_)) { + throw std::runtime_error("Echo-TTS conditioning graph allocation failed"); + } + core::prepare_host_graph_plan(execution_, conditioning_graph_, conditioning_plan_); + } + + // Reconciles the cache's lane count with the graph's. + // + // With pre-expansion on, the cache is allocated at kMaxCfgLanes and the + // single-lane graph reads a prefix of it; the lane axis is outermost, so + // that is a view with no copy. + // + // With pre-expansion off (AUDIOCPP_ECHO_TTS_NO_KV_EXPAND=1) the cache is + // batch-1 and the three-lane graph is *wider* than it. Narrowing is not + // possible and not wanted: returning the cache unchanged leaves + // joint_attention's expand() to broadcast it per step, which is exactly the + // pre-patch behaviour the flag exists to restore. + core::TensorValue kv_for_lanes( + core::ModuleBuildContext & ctx, const core::TensorValue & cached, int lanes) const { + if (cached.shape.dims[0] <= static_cast(lanes)) { + return cached; + } + return modules::SliceModule({0, 0, static_cast(lanes)}).build(ctx, cached); + } + + void build_denoiser_graph(DenoiserGraph & target, int lanes) { + ggml_init_params params{kGraphArenaBytes, nullptr, true}; + target.ctx.reset(ggml_init(params)); + if (target.ctx == nullptr) { + throw std::runtime_error("Echo-TTS denoiser graph context initialization failed"); + } + core::ModuleBuildContext ctx{target.ctx.get(), "echo_tts.denoise", backend_type_}; + + const int64_t D = config_.model_size; + const int64_t lane_count = lanes; + const int64_t keys = sequence_length_ + text_length_ + speaker_tokens_; + + target.lanes = lanes; + target.x = core::make_tensor( + ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({lane_count, sequence_length_, config_.latent_size})); + target.timestep = core::make_tensor( + ctx, GGML_TYPE_F32, + core::TensorShape::from_dims({lane_count, config_.timestep_embed_size})); + ggml_init_params const_params{ggml_tensor_overhead() * 8, nullptr, true}; + target.const_ctx.reset(ggml_init(const_params)); + if (target.const_ctx == nullptr) { + throw std::runtime_error("Echo-TTS denoiser constant context initialization failed"); + } + core::ModuleBuildContext const_ctx{ + target.const_ctx.get(), "echo_tts.denoise.const", backend_type_}; + target.positions = core::make_tensor( + const_ctx, GGML_TYPE_I32, core::TensorShape::from_dims({sequence_length_})); + // ggml_flash_attn_ext requires an F16 mask; the explicit path takes F32. + target.mask = core::make_tensor( + const_ctx, echo_flash_disabled() ? GGML_TYPE_F32 : GGML_TYPE_F16, + core::TensorShape::from_dims({lane_count, 1, sequence_length_, keys})); + target.const_buffer = + ggml_backend_alloc_ctx_tensors(target.const_ctx.get(), backend_); + if (target.const_buffer == nullptr) { + throw std::runtime_error("Echo-TTS denoiser constant buffer allocation failed"); + } + core::write_tensor_i32(target.positions, iota_positions(sequence_length_)); + if (echo_flash_disabled()) { + core::write_tensor_f32(target.mask, make_denoiser_mask(lanes)); + } else { + core::write_tensor_f16(target.mask, make_denoiser_mask(lanes)); + } + + for (auto * input : {target.x.tensor, target.timestep.tensor}) { + ggml_set_input(input); + } + + // cond_module: Linear, SiLU, Linear, SiLU, Linear -> 3 * model_size. + auto cond = modules::LinearModule({config_.timestep_embed_size, D, false, GGML_PREC_F32}) + .build(ctx, target.timestep, weights_.cond_0); + cond = modules::SiluModule{}.build(ctx, cond); + cond = modules::LinearModule({D, D, false, GGML_PREC_F32}).build(ctx, cond, weights_.cond_2); + cond = modules::SiluModule{}.build(ctx, cond); + cond = modules::LinearModule({D, D * 3, false, GGML_PREC_F32}).build(ctx, cond, weights_.cond_4); + // Insert the sequence axis so the conditioning broadcasts over steps. + cond = core::reshape_tensor( + ctx, contiguous(ctx, cond), core::TensorShape::from_dims({lane_count, 1, D * 3})); + + auto hidden = modules::LinearModule({config_.latent_size, D, true, GGML_PREC_F32}) + .build(ctx, target.x, weights_.in_proj); + for (int64_t layer = 0; layer < config_.num_layers; ++layer) { + const size_t index = static_cast(layer); + hidden = dit_block( + ctx, hidden, cond, target.positions, + kv_for_lanes(ctx, kv_.k_text[index], lanes), + kv_for_lanes(ctx, kv_.v_text[index], lanes), + kv_for_lanes(ctx, kv_.k_speaker[index], lanes), + kv_for_lanes(ctx, kv_.v_speaker[index], lanes), + target.mask, weights_.blocks[index], config_); + } + hidden = rms_norm(ctx, hidden, weights_.out_norm.weight, config_.norm_eps); + hidden = modules::LinearModule({D, config_.latent_size, true, GGML_PREC_F32}) + .build(ctx, hidden, weights_.out_proj); + + target.output = contiguous(ctx, hidden).tensor; + ggml_set_output(target.output); + target.graph = ggml_new_graph_custom(target.ctx.get(), 1048576, false); + ggml_build_forward_expand(target.graph, target.output); + + target.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (target.alloc == nullptr || + !ggml_gallocr_reserve(target.alloc, target.graph) || + !ggml_gallocr_alloc_graph(target.alloc, target.graph)) { + throw std::runtime_error("Echo-TTS denoiser graph allocation failed"); + } + core::prepare_host_graph_plan(execution_, target.graph, target.plan); + } + + void release_denoiser_graph(DenoiserGraph & target) { + target.plan.reset(); + if (target.graph != nullptr) { + core::release_backend_graph_resources(backend_type_, backend_, target.graph); + target.graph = nullptr; + } + if (target.alloc != nullptr) { + ggml_gallocr_free(target.alloc); + target.alloc = nullptr; + } + target.ctx.reset(); + if (target.const_buffer != nullptr) { + ggml_backend_buffer_free(target.const_buffer); + target.const_buffer = nullptr; + } + target.const_ctx.reset(); + target.output = nullptr; + target.lanes = 0; + } + + void release_denoiser_graphs() { + release_denoiser_graph(single_); + release_denoiser_graph(triple_); + } + + void release_kv_cache() { + kv_ = KvCache{}; + if (kv_buffer_ != nullptr) { + ggml_backend_buffer_free(kv_buffer_); + kv_buffer_ = nullptr; + } + kv_ctx_.reset(); + conditioning_ready_ = false; + } + + void release_all() { + release_denoiser_graphs(); + release_conditioning_graph(); + release_kv_cache(); + } + + EchoTtsConfig config_; + core::ExecutionContext & execution_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + core::BackendWeightStore store_; + EchoDitWeights weights_; + + bool conditioning_ready_ = false; + bool debug_printed_ = false; + int64_t text_length_ = 0; + int64_t speaker_frames_ = 0; + int64_t speaker_tokens_ = 0; + int64_t sequence_length_ = 640; + int64_t kv_lanes_ = 1; + std::vector text_mask_; + std::vector speaker_mask_; + + GgmlContextPtr kv_ctx_; + ggml_backend_buffer_t kv_buffer_ = nullptr; + KvCache kv_; + + GgmlContextPtr conditioning_ctx_; + ggml_cgraph * conditioning_graph_ = nullptr; + ggml_gallocr_t conditioning_alloc_ = nullptr; + core::HostGraphPlan conditioning_plan_; + core::TensorValue text_ids_; + core::TensorValue text_attn_mask_; + core::TensorValue text_positions_; + core::TensorValue speaker_latent_; + core::TensorValue speaker_positions_; + + DenoiserGraph single_; + DenoiserGraph triple_; +}; + +EchoDitRuntime::EchoDitRuntime( + const EchoTtsConfig & config, + const assets::TensorSource & source, + const std::string & tensor_prefix, + core::ExecutionContext & execution, + assets::TensorStorageType matmul_storage_type) + : impl_(std::make_unique(config, source, tensor_prefix, execution, matmul_storage_type)) {} + +EchoDitRuntime::~EchoDitRuntime() = default; + +const EchoTtsConfig & EchoDitRuntime::config() const noexcept { return impl_->config(); } + +void EchoDitRuntime::prepare_conditioning(const EchoConditioning & conditioning) { + impl_->prepare_conditioning(conditioning); +} + +std::vector EchoDitRuntime::denoise_once(const std::vector & x, float t, int lanes) { + if (!impl_->conditioning_ready()) { + throw std::runtime_error("Echo-TTS denoise_once() called before prepare_conditioning()"); + } + if (lanes < 1) { + throw std::runtime_error("Echo-TTS denoise_once() requires at least one lane"); + } + // `x` is always ONE lane. `lanes` selects how many velocity fields the + // denoiser returns -- see sampler.cpp, which calls denoise(x_t, t, 3) with + // an x_t of exactly `elements`, then expects elements * 3 back. Dividing + // the input by `lanes` here would set a sequence length 3x too small. + const int64_t latent_size = impl_->config().latent_size; + if (latent_size <= 0 || static_cast(x.size()) % latent_size != 0) { + throw std::runtime_error("Echo-TTS denoise_once() received a mis-shaped latent buffer"); + } + impl_->set_sequence_length(static_cast(x.size()) / latent_size); + return impl_->denoise(x, t, lanes); +} + +std::vector EchoDitRuntime::sample(const EchoSamplerOptions & options) { + if (!impl_->conditioning_ready()) { + throw std::runtime_error("Echo-TTS sample() called before prepare_conditioning()"); + } + impl_->set_sequence_length(options.sequence_length); + + Impl * impl = impl_.get(); + if (options.speaker_kv_scale.has_value()) { + impl->scale_speaker_kv(*options.speaker_kv_scale, options.speaker_kv_max_layers); + } + std::function on_kv_rescale; + if (options.speaker_kv_scale.has_value()) { + const float inverse = 1.0F / *options.speaker_kv_scale; + const auto max_layers = options.speaker_kv_max_layers; + on_kv_rescale = [impl, inverse, max_layers]() { + impl->scale_speaker_kv(inverse, max_layers); + }; + } + + auto denoise = [impl](const std::vector & x, float t, int lanes) { + return impl->denoise(x, t, lanes); + }; + return run_euler_sampler( + options, + options.sequence_length, + impl->config().latent_size, + impl->initial_noise(options), + denoise, + on_kv_rescale); +} + +void EchoTtsConfig::validate() const { + if (model_size % num_heads != 0 || text_model_size % text_num_heads != 0 || + speaker_model_size % speaker_num_heads != 0) { + throw std::runtime_error("Echo-TTS head counts must divide their model sizes"); + } + if (num_heads % 2 != 0) { + throw std::runtime_error("Echo-TTS requires an even head count for half-rotary attention"); + } + if (timestep_embed_size % 2 != 0) { + throw std::runtime_error("Echo-TTS timestep embedding size must be even"); + } + if (latent_size <= 0 || speaker_patch_size <= 0) { + throw std::runtime_error("Echo-TTS latent and patch sizes must be positive"); + } +} + +} // namespace engine::models::echo_tts diff --git a/src/community_models/echo_tts/dit_blocks.inc b/src/community_models/echo_tts/dit_blocks.inc new file mode 100644 index 00000000..9dae6b90 --- /dev/null +++ b/src/community_models/echo_tts/dit_blocks.inc @@ -0,0 +1,450 @@ +#include "engine/community_models/echo_tts/dit.h" + +#include "engine/community_models/echo_tts/sampler.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention/scaled_dot_product_attention.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/structural_modules.h" +#include "engine/framework/sampling/torch_random.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::echo_tts { +namespace { + +constexpr size_t kWeightContextBytes = 6144ull * 1024ull * 1024ull; +constexpr size_t kGraphArenaBytes = 512ull * 1024ull * 1024ull; +constexpr float kRopeTheta = 10000.0F; + +// cond, text-uncond, speaker-uncond. The KV cache is sized to this so the +// three-lane graph never has to broadcast it at sampling time. +constexpr int64_t kMaxCfgLanes = 3; + +// Additive mask value for disallowed keys. -INFINITY would be exact but +// produces NaN when an entire row is masked, which happens for the +// speaker-uncond CFG lane; a large finite penalty is the standard ggml +// workaround and leaves the softmax well defined. +constexpr float kMaskedBias = -1.0e9F; + +// The denoiser mask is F16 for the flash-attention path, whose maximum +// magnitude is 65504, so -1e9 would become -inf on conversion. A large finite +// penalty keeps the softmax well defined even if a row were ever fully masked. +constexpr float kMaskedBiasF16 = -65000.0F; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; +using GgmlContextPtr = std::unique_ptr; + +// Opt-out for the flash-attention lowering in the DiT's joint attention, so the +// explicit path can be compared against it without a rebuild. +bool echo_flash_disabled() { + static const bool disabled = [] { + const char * value = std::getenv("AUDIOCPP_ECHO_TTS_NO_FLASH"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return disabled; +} + +// Pre-broadcasting the conditioning KV cache across the CFG lanes moves a +// per-layer, per-step RepeatModule out of the denoiser graph and into the +// once-per-request conditioning graph. At 24 layers x 40 steps the repeat it +// removes is tens of gigabytes of pure copy traffic on a long reference. +// Set AUDIOCPP_ECHO_TTS_NO_KV_EXPAND=1 to fall back to the batch-1 cache. +bool echo_kv_expand_disabled() { + static const bool disabled = [] { + const char * value = std::getenv("AUDIOCPP_ECHO_TTS_NO_KV_EXPAND"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return disabled; +} + +core::TensorValue contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return core::ensure_backend_addressable_layout(ctx, value); +} + +modules::LinearWeights load_linear( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + assets::TensorStorageType storage, + int64_t out_features, + int64_t in_features, + bool use_bias) { + modules::LinearWeights weights; + weights.weight = store.load_tensor(source, name + ".weight", storage, {out_features, in_features}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, name + ".bias", {out_features}); + } + return weights; +} + +EchoRmsNormWeights load_norm( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + std::initializer_list shape) { + return EchoRmsNormWeights{store.load_f32_tensor(source, name + ".weight", shape)}; +} + +EchoMlpWeights load_mlp( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage, + int64_t dim, + int64_t inter) { + EchoMlpWeights weights; + weights.w1 = load_linear(store, source, prefix + ".w1", storage, inter, dim, false); + weights.w3 = load_linear(store, source, prefix + ".w3", storage, inter, dim, false); + weights.w2 = load_linear(store, source, prefix + ".w2", storage, dim, inter, false); + return weights; +} + +EchoEncoderBlockWeights load_encoder_block( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage, + int64_t dim, + int64_t inter, + int64_t heads) { + const int64_t head_dim = dim / heads; + EchoEncoderBlockWeights block; + auto & attn = block.attention; + attn.wq = load_linear(store, source, prefix + ".attention.wq", storage, dim, dim, false); + attn.wk = load_linear(store, source, prefix + ".attention.wk", storage, dim, dim, false); + attn.wv = load_linear(store, source, prefix + ".attention.wv", storage, dim, dim, false); + attn.wo = load_linear(store, source, prefix + ".attention.wo", storage, dim, dim, false); + attn.gate = load_linear(store, source, prefix + ".attention.gate", storage, dim, dim, false); + attn.q_norm = load_norm(store, source, prefix + ".attention.q_norm", {heads, head_dim}); + attn.k_norm = load_norm(store, source, prefix + ".attention.k_norm", {heads, head_dim}); + block.mlp = load_mlp(store, source, prefix + ".mlp", storage, dim, inter); + block.attention_norm = load_norm(store, source, prefix + ".attention_norm", {dim}); + block.mlp_norm = load_norm(store, source, prefix + ".mlp_norm", {dim}); + return block; +} + +EchoAdaLnWeights load_adaln( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage, + int64_t dim, + int64_t rank) { + EchoAdaLnWeights weights; + weights.shift_down = load_linear(store, source, prefix + ".shift_down", storage, rank, dim, false); + weights.scale_down = load_linear(store, source, prefix + ".scale_down", storage, rank, dim, false); + weights.gate_down = load_linear(store, source, prefix + ".gate_down", storage, rank, dim, false); + weights.shift_up = load_linear(store, source, prefix + ".shift_up", storage, dim, rank, true); + weights.scale_up = load_linear(store, source, prefix + ".scale_up", storage, dim, rank, true); + weights.gate_up = load_linear(store, source, prefix + ".gate_up", storage, dim, rank, true); + return weights; +} + +EchoDitWeights load_dit_weights( + const EchoTtsConfig & config, + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & p, + assets::TensorStorageType storage) { + EchoDitWeights w; + const int64_t D = config.model_size; + const int64_t TD = config.text_model_size; + const int64_t SD = config.speaker_model_size; + + w.text_encoder.text_embedding = store.load_tensor( + source, p + "text_encoder.text_embedding.weight", storage, + {config.text_vocab_size, TD}); + for (int64_t i = 0; i < config.text_num_layers; ++i) { + w.text_encoder.blocks.push_back(load_encoder_block( + store, source, p + "text_encoder.blocks." + std::to_string(i), storage, + TD, config.text_intermediate_size, config.text_num_heads)); + } + + w.speaker_encoder.in_proj = load_linear( + store, source, p + "speaker_encoder.in_proj", storage, + SD, config.latent_size * config.speaker_patch_size, true); + for (int64_t i = 0; i < config.speaker_num_layers; ++i) { + w.speaker_encoder.blocks.push_back(load_encoder_block( + store, source, p + "speaker_encoder.blocks." + std::to_string(i), storage, + SD, config.speaker_intermediate_size, config.speaker_num_heads)); + } + + w.text_norm = load_norm(store, source, p + "text_norm", {TD}); + w.speaker_norm = load_norm(store, source, p + "speaker_norm", {SD}); + + w.cond_0 = load_linear(store, source, p + "cond_module.0", storage, D, config.timestep_embed_size, false); + w.cond_2 = load_linear(store, source, p + "cond_module.2", storage, D, D, false); + w.cond_4 = load_linear(store, source, p + "cond_module.4", storage, D * 3, D, false); + + w.in_proj = load_linear(store, source, p + "in_proj", storage, D, config.latent_size, true); + + const int64_t head_dim = config.head_dim(); + for (int64_t i = 0; i < config.num_layers; ++i) { + const std::string prefix = p + "blocks." + std::to_string(i); + EchoDitBlockWeights block; + auto & a = block.attention; + a.wq = load_linear(store, source, prefix + ".attention.wq", storage, D, D, false); + a.wk = load_linear(store, source, prefix + ".attention.wk", storage, D, D, false); + a.wv = load_linear(store, source, prefix + ".attention.wv", storage, D, D, false); + a.wk_text = load_linear(store, source, prefix + ".attention.wk_text", storage, D, TD, false); + a.wv_text = load_linear(store, source, prefix + ".attention.wv_text", storage, D, TD, false); + a.wk_speaker = load_linear(store, source, prefix + ".attention.wk_speaker", storage, D, SD, false); + a.wv_speaker = load_linear(store, source, prefix + ".attention.wv_speaker", storage, D, SD, false); + a.gate = load_linear(store, source, prefix + ".attention.gate", storage, D, D, false); + a.wo = load_linear(store, source, prefix + ".attention.wo", storage, D, D, false); + a.q_norm = load_norm(store, source, prefix + ".attention.q_norm", {config.num_heads, head_dim}); + a.k_norm = load_norm(store, source, prefix + ".attention.k_norm", {config.num_heads, head_dim}); + block.mlp = load_mlp(store, source, prefix + ".mlp", storage, D, config.intermediate_size); + block.attention_adaln = load_adaln(store, source, prefix + ".attention_adaln", storage, D, config.adaln_rank); + block.mlp_adaln = load_adaln(store, source, prefix + ".mlp_adaln", storage, D, config.adaln_rank); + w.blocks.push_back(std::move(block)); + } + + w.out_norm = load_norm(store, source, p + "out_norm", {D}); + w.out_proj = load_linear(store, source, p + "out_proj", storage, config.latent_size, D, true); + return w; +} + +// --- graph building blocks --------------------------------------------- + +core::TensorValue rms_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & weight, + float eps) { + return modules::RMSNormModule({input.shape.last_dim(), eps, true, false}) + .build(ctx, input, {weight, std::nullopt}); +} + +// Elementwise multiply where the right operand is broadcast over one or more +// leading axes. MulModule requires identical shapes, but adaLN produces a +// per-sequence-position-invariant gate of shape (lanes, 1, dim) that has to +// scale a (lanes, seq, dim) activation. ggml_mul itself broadcasts whenever +// every rhs dimension divides the lhs, which is exactly this case. +core::TensorValue broadcast_mul( + core::ModuleBuildContext & ctx, + const core::TensorValue & lhs, + const core::TensorValue & rhs) { + return core::wrap_tensor( + ggml_mul(ctx.ggml, contiguous(ctx, lhs).tensor, contiguous(ctx, rhs).tensor), + lhs.shape, + GGML_TYPE_F32); +} + +// Per-head RMS normalisation for q_norm / k_norm. +// +// Echo's q_norm and k_norm carry a (num_heads, head_dim) weight: each head has +// its own learned scale. RMSNormModule takes a 1-D weight sized to the last +// dimension, so it cannot express this -- the in-tree rf_dit.cpp uses that +// module because its q_norm really is 1-D {head_dim}, which is a different +// architecture, not a different spelling of the same one. +// +// The reduction is over head_dim only. In ggml layout the activation is +// (head_dim, heads, seq, batch) and the weight is (head_dim, heads, 1, 1), so a +// plain multiply broadcasts the per-head scales across sequence and batch. +core::TensorValue head_rms_norm( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & weight, + float eps) { + auto * normed = ggml_rms_norm(ctx.ggml, contiguous(ctx, input).tensor, eps); + return core::wrap_tensor( + ggml_mul(ctx.ggml, normed, contiguous(ctx, weight).tensor), + input.shape, + GGML_TYPE_F32); +} + +// RMS normalisation with no learned scale, used inside LowRankAdaLN where the +// scale arrives from the conditioning path instead. +core::TensorValue rms_norm_bare( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + float eps) { + return core::wrap_tensor( + ggml_rms_norm(ctx.ggml, contiguous(ctx, input).tensor, eps), + input.shape, + GGML_TYPE_F32); +} + +core::TensorValue reshape_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t heads, + int64_t head_dim) { + return core::reshape_tensor( + ctx, + contiguous(ctx, input), + core::TensorShape::from_dims( + {input.shape.dims[0], input.shape.dims[1], heads, head_dim})); +} + +core::TensorValue to_bhsd(core::ModuleBuildContext & ctx, const core::TensorValue & bshd) { + return modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, bshd); +} + +core::TensorValue mlp( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const EchoMlpWeights & weights, + int64_t dim, + int64_t inter) { + auto gate = modules::LinearModule({dim, inter, false, GGML_PREC_F32}) + .build(ctx, input, weights.w1); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule({dim, inter, false, GGML_PREC_F32}) + .build(ctx, input, weights.w3); + auto hidden = modules::MulModule{}.build(ctx, gate, up); + return modules::LinearModule({inter, dim, false, GGML_PREC_F32}) + .build(ctx, hidden, weights.w2); +} + +// Applies output * sigmoid(gate) before the output projection, as every +// attention block in this model does. +core::TensorValue apply_attention_gate( + core::ModuleBuildContext & ctx, + const core::TensorValue & context, + const core::TensorValue & input, + const modules::LinearWeights & gate_weights, + int64_t dim) { + auto gate = modules::LinearModule({dim, dim, false, GGML_PREC_F32}) + .build(ctx, input, gate_weights); + gate = modules::SigmoidModule{}.build(ctx, gate); + return modules::MulModule{}.build(ctx, context, gate); +} + +// SelfAttention shared by the text and speaker encoders. RoPE covers all heads +// here, unlike the DiT's joint attention. +core::TensorValue encoder_self_attention( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const std::optional & mask, + const EchoSelfAttentionWeights & weights, + int64_t dim, + int64_t heads, + float eps, + bool causal) { + const int64_t head_dim = dim / heads; + auto q = modules::LinearModule({dim, dim, false, GGML_PREC_F32}).build(ctx, input, weights.wq); + auto k = modules::LinearModule({dim, dim, false, GGML_PREC_F32}).build(ctx, input, weights.wk); + auto v = modules::LinearModule({dim, dim, false, GGML_PREC_F32}).build(ctx, input, weights.wv); + + q = head_rms_norm(ctx, reshape_heads(ctx, q, heads, head_dim), weights.q_norm.weight, eps); + k = head_rms_norm(ctx, reshape_heads(ctx, k, heads, head_dim), weights.k_norm.weight, eps); + v = reshape_heads(ctx, v, heads, head_dim); + + // Upstream builds freqs_cis by viewing adjacent pairs as complex, which is + // the interleaved convention -- GGML_ROPE_TYPE_NORMAL, not NEOX. + const modules::RoPEModule rope({head_dim, GGML_ROPE_TYPE_NORMAL, kRopeTheta}); + q = rope.build(ctx, q, positions); + k = rope.build(ctx, k, positions); + + modules::ScaledDotProductAttentionConfig attn_config; + attn_config.head_dim = head_dim; + attn_config.lowering = modules::ScaledDotProductAttentionLowering::Explicit; + attn_config.precision = GGML_PREC_F32; + attn_config.causality = causal ? modules::AttentionCausality::Causal + : modules::AttentionCausality::NonCausal; + auto context = modules::ScaledDotProductAttentionModule(attn_config) + .build(ctx, to_bhsd(ctx, q), to_bhsd(ctx, k), to_bhsd(ctx, v), mask); + + context = core::reshape_tensor( + ctx, + contiguous(ctx, context), + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], dim})); + context = apply_attention_gate(ctx, context, input, weights.gate, dim); + return modules::LinearModule({dim, dim, false, GGML_PREC_F32}).build(ctx, context, weights.wo); +} + +core::TensorValue encoder_block( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const std::optional & mask, + const EchoEncoderBlockWeights & weights, + int64_t dim, + int64_t inter, + int64_t heads, + float eps, + bool causal) { + auto normed = rms_norm(ctx, input, weights.attention_norm.weight, eps); + auto attn = encoder_self_attention(ctx, normed, positions, mask, weights.attention, dim, heads, eps, causal); + auto hidden = modules::AddModule{}.build(ctx, input, attn); + + auto mlp_normed = rms_norm(ctx, hidden, weights.mlp_norm.weight, eps); + auto mlp_out = mlp(ctx, mlp_normed, weights.mlp, dim, inter); + return modules::AddModule{}.build(ctx, hidden, mlp_out); +} + +// LowRankAdaLN. Returns the normalised activation and the tanh-bounded gate. +struct AdaLnResult { + core::TensorValue normed; + core::TensorValue gate; +}; + +AdaLnResult adaln( + core::ModuleBuildContext & ctx, + const core::TensorValue & x, + const core::TensorValue & cond_embed, + const EchoAdaLnWeights & weights, + int64_t dim, + int64_t rank, + float eps) { + // cond_embed is (batch, 1, dim * 3); chunk into shift, scale, gate. + auto shift = modules::SliceModule({2, 0, dim}).build(ctx, cond_embed); + auto scale = modules::SliceModule({2, dim, dim}).build(ctx, cond_embed); + auto gate = modules::SliceModule({2, 2 * dim, dim}).build(ctx, cond_embed); + + auto refine = [&](const core::TensorValue & value, + const modules::LinearWeights & down, + const modules::LinearWeights & up) { + auto hidden = modules::SiluModule{}.build(ctx, value); + hidden = modules::LinearModule({dim, rank, false, GGML_PREC_F32}).build(ctx, hidden, down); + hidden = modules::LinearModule({rank, dim, true, GGML_PREC_F32}).build(ctx, hidden, up); + return modules::AddModule{}.build(ctx, hidden, value); + }; + + shift = refine(shift, weights.shift_down, weights.shift_up); + scale = refine(scale, weights.scale_down, weights.scale_up); + gate = refine(gate, weights.gate_down, weights.gate_up); + + auto normed = rms_norm_bare(ctx, x, eps); + // x * (scale + 1) + shift, with the conditioning broadcast over sequence. + auto scale_plus_one = core::wrap_tensor( + ggml_scale_bias(ctx.ggml, contiguous(ctx, scale).tensor, 1.0F, 1.0F), + scale.shape, + GGML_TYPE_F32); + normed = core::wrap_tensor( + ggml_mul(ctx.ggml, contiguous(ctx, normed).tensor, contiguous(ctx, scale_plus_one).tensor), + normed.shape, + GGML_TYPE_F32); + normed = core::wrap_tensor( + ggml_add(ctx.ggml, normed.tensor, contiguous(ctx, shift).tensor), + normed.shape, + GGML_TYPE_F32); + + gate = modules::TanhModule{}.build(ctx, gate); + return AdaLnResult{normed, gate}; +} + +} // namespace +} // namespace engine::models::echo_tts diff --git a/src/community_models/echo_tts/latent_post.cpp b/src/community_models/echo_tts/latent_post.cpp new file mode 100644 index 00000000..90d9df89 --- /dev/null +++ b/src/community_models/echo_tts/latent_post.cpp @@ -0,0 +1,132 @@ +#include "engine/community_models/echo_tts/latent_post.h" + +#include +#include + +namespace engine::models::echo_tts { + +std::vector pca_project( + const EchoPcaState & pca, + const EchoTtsConfig & config, + const std::vector & z_q, + int64_t frames) { + const int64_t features = config.ae_latent_dim; + const int64_t components = config.latent_size; + if (static_cast(z_q.size()) != frames * features) { + throw std::runtime_error("Echo-TTS PCA projection received a mis-shaped z_q buffer"); + } + if (static_cast(pca.components.size()) != components * features || + static_cast(pca.mean.size()) != features) { + throw std::runtime_error("Echo-TTS PCA state has unexpected dimensions"); + } + + std::vector out(static_cast(frames * components), 0.0F); + for (int64_t f = 0; f < frames; ++f) { + const float * row = z_q.data() + f * features; + float * dst = out.data() + f * components; + for (int64_t c = 0; c < components; ++c) { + const float * basis = pca.components.data() + c * features; + double acc = 0.0; + for (int64_t k = 0; k < features; ++k) { + acc += static_cast(row[k] - pca.mean[static_cast(k)]) * + static_cast(basis[k]); + } + dst[c] = static_cast(acc) * pca.latent_scale; + } + } + return out; +} + +std::vector pca_unproject( + const EchoPcaState & pca, + const EchoTtsConfig & config, + const std::vector & latents, + int64_t frames) { + const int64_t features = config.ae_latent_dim; + const int64_t components = config.latent_size; + if (static_cast(latents.size()) != frames * components) { + throw std::runtime_error("Echo-TTS PCA inverse received a mis-shaped latent buffer"); + } + if (static_cast(pca.components.size()) != components * features || + static_cast(pca.mean.size()) != features) { + throw std::runtime_error("Echo-TTS PCA state has unexpected dimensions"); + } + if (pca.latent_scale == 0.0F) { + throw std::runtime_error("Echo-TTS PCA latent_scale must be non-zero"); + } + + std::vector out(static_cast(frames * features), 0.0F); + const float inv_scale = 1.0F / pca.latent_scale; + for (int64_t f = 0; f < frames; ++f) { + const float * row = latents.data() + f * components; + float * dst = out.data() + f * features; + for (int64_t k = 0; k < features; ++k) { + dst[k] = pca.mean[static_cast(k)]; + } + for (int64_t c = 0; c < components; ++c) { + const float coeff = row[c] * inv_scale; + if (coeff == 0.0F) { + continue; + } + const float * basis = pca.components.data() + c * features; + for (int64_t k = 0; k < features; ++k) { + dst[k] += coeff * basis[k]; + } + } + } + return out; +} + +int64_t find_flattening_point( + const std::vector & latents, + int64_t frames, + int64_t latent_size, + int64_t window_size, + float std_threshold, + float target_value) { + if (frames <= 0 || latent_size <= 0 || window_size <= 0) { + return frames; + } + if (static_cast(latents.size()) != frames * latent_size) { + throw std::runtime_error("Echo-TTS flattening search received a mis-shaped latent buffer"); + } + + // Upstream pads the sequence with `window_size` zero frames before scanning, + // so a generation that runs to the end of the window still terminates. + const int64_t padded_frames = frames + window_size; + const int64_t count = window_size * latent_size; + if (count < 2) { + return frames; + } + + auto value_at = [&](int64_t frame, int64_t channel) -> double { + if (frame >= frames) { + return 0.0; + } + return static_cast(latents[static_cast(frame * latent_size + channel)]); + }; + + for (int64_t start = 0; start < padded_frames - window_size; ++start) { + double sum = 0.0; + double sum_sq = 0.0; + for (int64_t f = start; f < start + window_size; ++f) { + for (int64_t c = 0; c < latent_size; ++c) { + const double v = value_at(f, c); + sum += v; + sum_sq += v * v; + } + } + const double mean = sum / static_cast(count); + // torch.std defaults to the unbiased estimator (correction = 1). + const double variance = + (sum_sq - sum * mean) / static_cast(count - 1); + const double stddev = variance > 0.0 ? std::sqrt(variance) : 0.0; + if (stddev < static_cast(std_threshold) && + std::abs(mean - static_cast(target_value)) < 0.1) { + return start; + } + } + return frames; +} + +} // namespace engine::models::echo_tts diff --git a/src/community_models/echo_tts/sampler.cpp b/src/community_models/echo_tts/sampler.cpp new file mode 100644 index 00000000..61e5749e --- /dev/null +++ b/src/community_models/echo_tts/sampler.cpp @@ -0,0 +1,152 @@ +#include "engine/community_models/echo_tts/sampler.h" + +#include +#include +#include + +namespace engine::models::echo_tts { +namespace { + +// inference.py::sample_euler_cfg_independent_guidances, INIT_SCALE. +constexpr float kInitScale = 0.999F; + +} // namespace + +std::vector euler_timestep_schedule(int num_steps) { + if (num_steps <= 0) { + throw std::runtime_error("Echo-TTS sampler requires at least one step"); + } + std::vector schedule(static_cast(num_steps) + 1); + for (int i = 0; i <= num_steps; ++i) { + // torch.linspace(1, 0, n + 1) puts exact endpoints at both ends. + const float ramp = + 1.0F - static_cast(i) / static_cast(num_steps); + schedule[static_cast(i)] = ramp * kInitScale; + } + return schedule; +} + +bool cfg_active(float t, float cfg_min_t, float cfg_max_t) { + return t >= cfg_min_t && t <= cfg_max_t; +} + +std::vector combine_cfg_lanes( + const std::vector & lanes, + int64_t lane_elements, + float cfg_scale_text, + float cfg_scale_speaker) { + if (lane_elements <= 0 || + static_cast(lanes.size()) != lane_elements * 3) { + throw std::runtime_error("Echo-TTS CFG combine expects exactly three lanes"); + } + const float * cond = lanes.data(); + const float * uncond_text = lanes.data() + lane_elements; + const float * uncond_speaker = lanes.data() + 2 * lane_elements; + + std::vector out(static_cast(lane_elements)); + for (int64_t i = 0; i < lane_elements; ++i) { + const float c = cond[i]; + out[static_cast(i)] = + c + cfg_scale_text * (c - uncond_text[i]) + + cfg_scale_speaker * (c - uncond_speaker[i]); + } + return out; +} + +std::vector run_euler_sampler( + const EchoSamplerOptions & options, + int64_t sequence_length, + int64_t latent_size, + std::vector initial_noise, + const EchoDenoiseFn & denoise, + const std::function & on_kv_rescale) { + const int64_t elements = sequence_length * latent_size; + if (static_cast(initial_noise.size()) != elements) { + throw std::runtime_error("Echo-TTS sampler received a mis-shaped noise buffer"); + } + if (!denoise) { + throw std::runtime_error("Echo-TTS sampler requires a denoise callback"); + } + + std::vector x_t = std::move(initial_noise); + if (options.truncation_factor.has_value()) { + const float factor = *options.truncation_factor; + for (auto & value : x_t) { + value *= factor; + } + } + + const auto schedule = euler_timestep_schedule(options.num_steps); + bool kv_scaled = options.speaker_kv_scale.has_value(); + + const int cfg_interval = std::max(1, options.cfg_interval); + // Guidance correction carried between refreshes when cfg_interval > 1. This + // is the whole additive term, w_text * (v_cond - v_text) + w_speaker * + // (v_cond - v_speaker), held in absolute units rather than as a ratio so a + // reused correction cannot amplify a small v_cond. + std::vector cfg_delta; + int steps_since_refresh = 0; + + for (int step = 0; step < options.num_steps; ++step) { + const float t = schedule[static_cast(step)]; + const float t_next = schedule[static_cast(step) + 1]; + const bool use_cfg = cfg_active(t, options.cfg_min_t, options.cfg_max_t); + + std::vector v_pred; + if (use_cfg) { + // The first CFG step always refreshes, so a stale delta is never + // applied before one has been measured. + const bool refresh = cfg_delta.empty() || steps_since_refresh >= cfg_interval - 1; + if (refresh) { + auto lanes = denoise(x_t, t, 3); + if (static_cast(lanes.size()) != elements * 3) { + throw std::runtime_error("Echo-TTS denoiser returned mis-shaped CFG lanes"); + } + v_pred = combine_cfg_lanes( + lanes, elements, options.cfg_scale_text, options.cfg_scale_speaker); + if (cfg_interval > 1) { + cfg_delta.resize(static_cast(elements)); + for (int64_t i = 0; i < elements; ++i) { + cfg_delta[static_cast(i)] = + v_pred[static_cast(i)] - lanes[static_cast(i)]; + } + } + steps_since_refresh = 0; + } else { + v_pred = denoise(x_t, t, 1); + if (static_cast(v_pred.size()) != elements) { + throw std::runtime_error("Echo-TTS denoiser returned a mis-shaped velocity"); + } + for (int64_t i = 0; i < elements; ++i) { + v_pred[static_cast(i)] += cfg_delta[static_cast(i)]; + } + ++steps_since_refresh; + } + } else { + v_pred = denoise(x_t, t, 1); + if (static_cast(v_pred.size()) != elements) { + throw std::runtime_error("Echo-TTS denoiser returned a mis-shaped velocity"); + } + } + + // Speaker KV scaling is undone once the schedule crosses below + // speaker_kv_min_t, matching upstream's boundary test on (t, t_next). + if (kv_scaled && options.speaker_kv_min_t.has_value()) { + const float threshold = *options.speaker_kv_min_t; + if (t_next < threshold && t >= threshold) { + if (on_kv_rescale) { + on_kv_rescale(); + } + kv_scaled = false; + } + } + + const float dt = t_next - t; + for (int64_t i = 0; i < elements; ++i) { + x_t[static_cast(i)] += v_pred[static_cast(i)] * dt; + } + } + return x_t; +} + +} // namespace engine::models::echo_tts diff --git a/src/community_models/echo_tts/session.cpp b/src/community_models/echo_tts/session.cpp new file mode 100644 index 00000000..b483d003 --- /dev/null +++ b/src/community_models/echo_tts/session.cpp @@ -0,0 +1,611 @@ +#include "engine/community_models/echo_tts/session.h" + +#include "engine/community_models/echo_tts/dit.h" +#include "engine/community_models/echo_tts/latent_post.h" +#include "engine/community_models/echo_tts/tokenizer.h" +#include "engine/framework/model_spec/package.h" +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/audio/waveform_ops.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" +#include "engine/framework/text/chunking.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::echo_tts { +namespace { + +// Mirrors the tap in dit.cpp so the whole pipeline can be traced with one flag. +bool echo_session_debug_enabled() { + static const bool enabled = [] { + const char * value = std::getenv("AUDIOCPP_ECHO_TTS_DEBUG"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return enabled; +} + +void report_stats(const char * label, const std::vector & values) { + if (!echo_session_debug_enabled()) { + return; + } + if (values.empty()) { + std::fprintf(stderr, " %-26s \n", label); + return; + } + double sum = 0.0; + double sum_sq = 0.0; + float low = values[0]; + float high = values[0]; + for (const float value : values) { + sum += value; + sum_sq += static_cast(value) * value; + low = std::min(low, value); + high = std::max(high, value); + } + const double mean = sum / static_cast(values.size()); + const double variance = sum_sq / static_cast(values.size()) - mean * mean; + std::fprintf(stderr, " %-26s mean=%+.6f std=%.6f min=%+.4f max=%+.4f n=%zu\n", + label, mean, variance > 0.0 ? std::sqrt(variance) : 0.0, + static_cast(low), static_cast(high), values.size()); +} + +constexpr const char * kFamily = "echo_tts"; +constexpr int kSampleRate = 44100; +// ~20 s of English, leaving headroom before the model starts compressing speech +// to fit the fixed 29.72 s window. Overridable per request. +// ~20 s of typical English, against a fixed 29.72 s generation window. Dense or +// fast-reading text can still overrun it, which is the main prompt-dependent +// failure mode; text_chunk_size overrides this per request. +constexpr int64_t kDefaultTextChunkSize = 300; +// Enough for a server rotating a few voices; each slot holds only the projected +// latent, at most 6400 frames x 80 floats = 2 MB. +constexpr std::size_t kDefaultReferenceCacheSlots = 4; + +// Default reference trim. Every speaker token stays resident in `keys` for +// every attention in every block at every sampler step, and the speaker encoder +// itself is linear in reference length, so an untrimmed 4.5-minute clip charges +// 1600 tokens to all 24 blocks x 40 steps. 15 s is ~81 tokens, and the model +// card's own guidance is that ~10 s clones at least as well. Override with +// reference_max_seconds per request or echo_tts.reference_max_seconds per +// session; the trained maximum is still reachable that way. +constexpr int64_t kDefaultReferenceMaxSamples = 15 * kSampleRate; + +// Adaptive generation window. +// +// Denoiser cost is linear in sequence_length for the projections and the MLP, +// and worse than linear for the self block of the attention, so generating the +// full 640-latent window for a six-second sentence pays roughly five times over +// for latents that find_flattening_point then discards. +// +// The byte-to-frame rate is fixed by the model: 640 frames span 29.7215 s, so +// one second is 21.53 frames. kDefaultTextChunkSize is documented in this file +// as ~20 s of typical English at 300 codepoints, i.e. ~15 bytes/s, which puts +// the ratio at 21.53 / 15 = 1.435 frames per UTF-8 byte. The margin covers +// slower delivery, and a short utterance still needs room for the leading +// silence and the flat tail the crop looks for. +constexpr float kFramesPerTextByte = 1.435F; +constexpr float kWindowMargin = 1.30F; +constexpr int64_t kMinWindowFrames = 128; +// Denoiser graphs are keyed on sequence_length and rebuilt whenever it changes, +// so estimates are snapped to a coarse grid: consecutive chunks of similar +// length then reuse the same graph and the same gallocr reservation. +constexpr int64_t kWindowQuantum = 64; + +// Opt-in, not opt-out. See the comment at the adaptive-window call site for why +// the default is the full trained window. +bool echo_adaptive_window_enabled() { + static const bool enabled = [] { + const char * value = std::getenv("AUDIOCPP_ECHO_TTS_ADAPTIVE_WINDOW"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return enabled; +} + +// Rounds `frames` up to the graph-reuse grid and clamps into range. +int64_t quantize_window(int64_t frames, int64_t max_frames) { + frames = std::max(frames, kMinWindowFrames); + frames = ((frames + kWindowQuantum - 1) / kWindowQuantum) * kWindowQuantum; + return std::min(frames, max_frames); +} + +// Predicts how many latents this chunk needs. Deliberately generous: a window +// that is too short costs a full-length retry, while one that is slightly too +// long only wastes the difference. +int64_t estimate_window_frames(int64_t text_bytes, int64_t max_frames) { + const auto predicted = static_cast( + std::ceil(static_cast(text_bytes) * kFramesPerTextByte * kWindowMargin)); + return quantize_window(predicted, max_frames); +} + +bool echo_debug_enabled() { + static const bool enabled = [] { + const char * value = std::getenv("AUDIOCPP_ECHO_TTS_DEBUG"); + return value != nullptr && value[0] != '\0' && value[0] != '0'; + }(); + return enabled; +} +constexpr size_t kDefaultDitWeightContextBytes = 6144ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecGraphArenaBytes = 1024ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecWeightContextBytes = 2048ull * 1024ull * 1024ull; + +EchoPcaState load_pca_state( + const assets::TensorSource & source, + const EchoTtsConfig & config) { + EchoPcaState pca; + // `source` is already a view scoped to the "pca" namespace, so lookups here + // are bare names; prefixing again would ask for "pca/pca.components". + pca.components = source.require_f32( + "components", {config.latent_size, config.ae_latent_dim}); + pca.mean = source.require_f32("mean", {config.ae_latent_dim}); + return pca; +} + +std::shared_ptr load_echo_tts_assets( + const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = engine::model_spec::load_resource_bundle_for_family(model_path, kFamily); + assets->dit_weights = assets->resources.open_tensor_source("dit_weights"); + auto pca_source = assets->resources.open_tensor_source("pca"); + assets->pca = load_pca_state(*pca_source, assets->config); + // Published as float32(1/18); reconstructed exactly rather than stored. + assets->pca.latent_scale = 1.0F / 18.0F; + assets->config.validate(); + + // The Fish S1-DAC travels inside Echo's own GGUF. audio.cpp implements this + // codec for the fish_audio family and Echo reuses that implementation, but + // not its weights: fish_audio ships S2 Pro, while Echo's PCA basis is fitted + // to the S1 DAC's latent space. Only four config fields reach the codec + // graphs, and their defaults already describe S1-DAC. + auto codec_assets = std::make_shared(); + codec_assets->codec_weights = assets->resources.open_tensor_source("codec_weights"); + codec_assets->config.codec.sample_rate = kSampleRate; + codec_assets->config.codec.frame_length = assets->config.ae_downsample_factor; + codec_assets->config.codec.total_codebooks = 10; + codec_assets->config.codec.quantizer_codebooks = 9; + assets->codec_assets = std::move(codec_assets); + return assets; +} + +} // namespace + +EchoTtsSession::EchoTtsSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : RuntimeSessionBase(std::move(options)), + task_(task), + assets_(std::move(assets)), + contract_(std::move(contract)) { + if (contract_ == nullptr) { + throw std::runtime_error("Echo-TTS session requires a model contract"); + } + // Without this a typo in server config is silently ignored. + runtime::validate_spec_backed_session_options( + RuntimeSessionBase::options(), *contract_, kFamily, "Echo-TTS"); + const auto slots = runtime::parse_int_option( + RuntimeSessionBase::options().options, {"echo_tts.reference_cache_slots"}); + if (slots.has_value()) { + if (*slots < 0) { + throw std::runtime_error("echo_tts.reference_cache_slots must be non-negative"); + } + reference_cache_.set_capacity(static_cast(*slots)); + } else { + reference_cache_.set_capacity(kDefaultReferenceCacheSlots); + } + if (assets_ == nullptr) { + throw std::runtime_error("Echo-TTS session requires loaded assets"); + } + if (task_.task != runtime::VoiceTaskKind::VoiceCloning || + task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Echo-TTS only supports offline voice cloning"); + } +} + +EchoTtsSession::~EchoTtsSession() = default; + +std::string EchoTtsSession::family() const { return kFamily; } +runtime::VoiceTaskKind EchoTtsSession::task_kind() const { return task_.task; } +runtime::RunMode EchoTtsSession::run_mode() const { return task_.mode; } + +EchoSamplerOptions EchoTtsSession::parse_sampler_options( + const std::unordered_map & options) const { + EchoSamplerOptions sampler; + sampler.num_steps = + runtime::parse_int_option(options, {"num_steps"}).value_or(sampler.num_steps); + sampler.cfg_scale_text = + runtime::parse_float_option(options, {"cfg_scale_text"}).value_or(sampler.cfg_scale_text); + sampler.cfg_scale_speaker = + runtime::parse_float_option(options, {"cfg_scale_speaker"}).value_or(sampler.cfg_scale_speaker); + if (const auto truncation = runtime::parse_float_option(options, {"truncation_factor"})) { + sampler.truncation_factor = *truncation; + } + if (const auto kv_scale = runtime::parse_float_option(options, {"speaker_kv_scale"})) { + // 1.0 is the documented "disabled" value, not a scale to apply. + if (*kv_scale != 1.0F) { + sampler.speaker_kv_scale = *kv_scale; + sampler.speaker_kv_min_t = 0.5F; + } + } + if (const auto seed = runtime::parse_int_option(options, {"seed"})) { + sampler.seed = static_cast(std::max(0, *seed)); + } + // The window defaults to the trained maximum. synthesize_chunk narrows it + // per chunk unless the caller pins it here, in which case the estimate is + // skipped entirely and the requested value is used verbatim. + sampler.sequence_length = assets_->config.max_sequence_length; + if (const auto window = runtime::parse_int_option(options, {"sequence_length"})) { + if (*window <= 0 || *window > assets_->config.max_sequence_length) { + throw std::runtime_error( + "Echo-TTS sequence_length must be in 1..max_sequence_length"); + } + sampler.sequence_length = *window; + sampler.window_pinned = true; + } + if (const auto interval = runtime::parse_int_option(options, {"cfg_interval"})) { + if (*interval < 1) { + throw std::runtime_error("Echo-TTS cfg_interval must be at least 1"); + } + sampler.cfg_interval = *interval; + } + if (sampler.num_steps <= 0) { + throw std::runtime_error("Echo-TTS num_steps must be positive"); + } + return sampler; +} + +void EchoTtsSession::prepare(const runtime::SessionPreparationRequest & request) { + (void)request; + if (dit_ == nullptr) { + dit_ = std::make_unique( + assets_->config, + *assets_->dit_weights, + // Namespace-scoped source: tensor names are already stripped of the + // "dit_weights/" prefix by the resource bundle. + "", + execution_context(), + assets::TensorStorageType::Native); + } + if (codec_ == nullptr) { + if (assets_->codec_assets == nullptr || + assets_->codec_assets->codec_weights == nullptr) { + throw std::runtime_error( + "Echo-TTS GGUF has no codec_weights; re-run convert_echo_tts.py with " + "--fish-dir pointing at the Fish S1-DAC checkpoint"); + } + const int threads = options().backend.threads > 0 ? options().backend.threads : 1; + codec_ = std::make_unique( + assets_->codec_assets, + options().backend, + threads, + kDefaultCodecGraphArenaBytes, + kDefaultCodecWeightContextBytes, + assets::TensorStorageType::Native, + assets::TensorStorageType::Native); + } + mark_prepared(); +} + +int64_t EchoTtsSession::resolve_reference_max_samples( + const std::unordered_map & request_options) const { + const auto & config = assets_->config; + const int64_t trained_max = config.max_speaker_latent_length * config.ae_downsample_factor; + + // A per-request value wins; otherwise the session default from CLI or server + // config; otherwise kDefaultReferenceMaxSamples (15 s), NOT the trained + // maximum -- see that constant for why, and model_specs/echo_tts.json which + // publishes 15.0 as the default in both scopes. Request options are bare names, + // session options carry the family prefix -- parse_cli_options adds it for + // the session and load scopes only. + auto seconds = runtime::parse_float_option(request_options, {"reference_max_seconds"}); + if (!seconds.has_value()) { + seconds = runtime::parse_float_option( + options().options, {"echo_tts.reference_max_seconds"}); + } + if (!seconds.has_value()) { + return kDefaultReferenceMaxSamples; + } + if (!(*seconds > 0.0F)) { + throw std::runtime_error("Echo-TTS reference_max_seconds must be positive"); + } + const auto requested = static_cast( + static_cast(*seconds) * static_cast(kSampleRate)); + // Clamped rather than rejected: asking for more than the model was trained + // on is a reasonable thing to type, and silently exceeding it is not. + return std::min(requested, trained_max); +} + +namespace { + +// Cheap content hash over the reference samples. A collision would swap one +// speaker for another, so it mixes length, rate, channels and every sample +// rather than sampling, and the cache key adds the trim length. +std::string reference_identity(const runtime::AudioBuffer & audio) { + std::uint64_t hash = 1469598103934665603ULL; + auto mix = [&hash](std::uint64_t value) { + hash ^= value; + hash *= 1099511628211ULL; + }; + mix(static_cast(audio.samples.size())); + mix(static_cast(audio.sample_rate)); + mix(static_cast(audio.channels)); + for (const float sample : audio.samples) { + std::uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + mix(bits); + } + return std::to_string(hash); +} + +} // namespace + +void EchoTtsSession::encode_speaker(const runtime::AudioBuffer & audio) { + const auto & config = assets_->config; + + const EchoReferenceIdentity identity{reference_identity(audio), reference_max_samples_}; + if (const auto * cached = reference_cache_.find(identity)) { + speaker_latent_ = cached->latent; + speaker_frames_ = cached->frames; + if (echo_debug_enabled()) { + std::fprintf(stderr, "[echo_tts] speaker reference cache hit (%lld frames)\n", + static_cast(speaker_frames_)); + } + return; + } + + // Mixed down and resampled once, so chunk boundaries land on exact codec + // frames rather than on pre-resample sample indices. + auto mono = engine::audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); + if (audio.sample_rate != kSampleRate) { + mono = engine::audio::resample_mono_torchaudio_sinc_hann( + mono, audio.sample_rate, kSampleRate); + } + + const int64_t chunk_samples = config.speaker_chunk_latents * config.ae_downsample_factor; + if (static_cast(mono.size()) > reference_max_samples_) { + if (echo_debug_enabled()) { + std::fprintf(stderr, "[echo_tts] reference trimmed %.2f s -> %.2f s\n", + static_cast(mono.size()) / kSampleRate, + static_cast(reference_max_samples_) / kSampleRate); + } + mono.resize(static_cast(reference_max_samples_)); + } + const int64_t actual_frames = static_cast(mono.size()) / config.ae_downsample_factor; + + // Encoded in ~30 s chunks, zero-padded to a fixed length, exactly as + // inference.py::get_speaker_latent_and_mask does. That is not only a memory + // measure: the chunk size is the longest span seen in training, and a single + // pass over several minutes of audio is a different computation. It also + // keeps every encode graph the same shape, so one graph is reused. + std::vector latents; + latents.reserve(static_cast(actual_frames + config.speaker_chunk_latents) * + static_cast(config.latent_size)); + for (int64_t offset = 0; offset < static_cast(mono.size()); offset += chunk_samples) { + const int64_t available = + std::min(chunk_samples, static_cast(mono.size()) - offset); + runtime::AudioBuffer chunk{kSampleRate, 1, std::vector( + static_cast(chunk_samples), 0.0F)}; + std::copy_n(mono.begin() + offset, available, chunk.samples.begin()); + + auto chunk_latents = codec_->encode_zq(chunk); + auto projected = pca_project( + assets_->pca, config, chunk_latents.values, chunk_latents.frames); + latents.insert(latents.end(), projected.begin(), projected.end()); + } + + // Trim the padding introduced by the final chunk, then crop to a multiple of + // the patch size the speaker encoder folds over. + int64_t frames = std::min( + actual_frames, static_cast(latents.size()) / config.latent_size); + frames = frames / config.speaker_patch_size * config.speaker_patch_size; + if (frames <= 0) { + throw std::runtime_error( + "Echo-TTS speaker reference is too short; at least " + "4 latent frames (~0.19 s) are required"); + } + latents.resize(static_cast(frames * config.latent_size)); + speaker_latent_ = std::move(latents); + speaker_frames_ = frames; + reference_cache_.put(identity, EchoPreparedSpeaker{speaker_latent_, speaker_frames_}); +} + +runtime::AudioBuffer EchoTtsSession::synthesize_chunk( + const std::string & text, + const EchoSamplerOptions & sampler) { + const auto & config = assets_->config; + auto tokens = tokenize_echo_text(text, config.max_text_length, true, false); + if (tokens.truncated) { + std::fprintf( + stderr, + "[echo_tts] warning: text truncated at %lld bytes; the tail will not be spoken\n", + static_cast(config.max_text_length)); + } + + EchoConditioning conditioning; + conditioning.text_input_ids = tokens.input_ids; + conditioning.text_mask = tokens.mask; + conditioning.text_length = static_cast(tokens.input_ids.size()); + conditioning.speaker_latent = speaker_latent_; + conditioning.speaker_mask.assign(static_cast(speaker_frames_), 1.0F); + conditioning.speaker_frames = speaker_frames_; + + dit_->prepare_conditioning(conditioning); + + // Adaptive window: OFF by default, enable with AUDIOCPP_ECHO_TTS_ADAPTIVE_WINDOW=1. + // + // An earlier revision defaulted this on, reasoning that an under-estimate + // "costs time, never fidelity" because a missing flattening point triggers a + // retry at full length on the same seed. That reasoning is wrong, and the + // seed is not what changes. Echo's generated self-attention is fully + // non-causal -- `self_mask = torch.ones((batch_size, seq_len))` at + // model.py:249 -- so every latent position attends over the whole window. + // Shrinking 640 to 128 therefore changes the computation at every retained + // position, not merely how many positions survive. The reference defaults to + // 640 (inference.py:353). + // + // The retry only fires when no flattening point is found, so a short window + // that happens to produce a plausible flat tail is never corrected and + // silently yields different audio. Keep the optimisation available -- the + // cost saving is real -- but it must not be the default until it has been + // A/B'd against the full window on a fixed seed. + EchoSamplerOptions attempt = sampler; + const bool adaptive = !sampler.window_pinned && echo_adaptive_window_enabled(); + if (adaptive) { + attempt.sequence_length = estimate_window_frames( + static_cast(tokens.input_ids.size()), config.max_sequence_length); + } + + std::vector latent; + int64_t frames = 0; + for (int pass = 0; pass < 2; ++pass) { + latent = dit_->sample(attempt); + frames = find_flattening_point(latent, attempt.sequence_length, config.latent_size); + const bool ran_out = frames >= attempt.sequence_length; + const bool can_retry = + adaptive && ran_out && attempt.sequence_length < config.max_sequence_length; + if (!can_retry) { + break; + } + if (echo_debug_enabled()) { + std::fprintf( + stderr, + "[echo_tts] window estimate of %lld frames was short for %lld bytes; " + "retrying at %lld\n", + static_cast(attempt.sequence_length), + static_cast(tokens.input_ids.size()), + static_cast(config.max_sequence_length)); + } + attempt.sequence_length = config.max_sequence_length; + } + report_stats("sampler.latent", latent); + + // The generated tail goes flat once the model finishes speaking; cropping + // there is what sets the output duration. + const double window_seconds = + static_cast(attempt.sequence_length * config.ae_downsample_factor) / + static_cast(config.sample_rate); + if (frames >= attempt.sequence_length) { + // No flat tail means the model was still speaking when the window ended, + // so the audio is cut mid-utterance. Almost always too much text for one + // chunk rather than a sampling problem. + std::fprintf( + stderr, + "[echo_tts] warning: no silence found within the %.2f s window for a " + "%lld-byte chunk; output is truncated mid-utterance. Try a smaller " + "text_chunk_size.\n", + window_seconds, static_cast(tokens.input_ids.size())); + } else if (echo_debug_enabled()) { + std::fprintf( + stderr, + "[echo_tts] chunk: %lld tokens -> %lld/%lld frames (%.2f s of %.2f s window)\n", + static_cast(tokens.input_ids.size()), + static_cast(frames), + static_cast(attempt.sequence_length), + static_cast(frames * config.ae_downsample_factor) / + static_cast(config.sample_rate), + window_seconds); + } + if (frames <= 0) { + return runtime::AudioBuffer{kSampleRate, 1, {}}; + } + if (echo_session_debug_enabled()) { + std::fprintf(stderr, " %-26s %lld of %lld frames (%.3f s)\n", + "flattening_point", static_cast(frames), + static_cast(attempt.sequence_length), + static_cast(frames * config.ae_downsample_factor) / + static_cast(config.sample_rate)); + } + latent.resize(static_cast(frames * config.latent_size)); + report_stats("latent.cropped", latent); + + auto z_q = pca_unproject(assets_->pca, config, latent, frames); + report_stats("decode.z_q", z_q); + auto audio = codec_->decode_zq(z_q, frames); + report_stats("decode.audio", audio.samples); + return audio; +} + +runtime::TaskResult EchoTtsSession::run(const runtime::TaskRequest & request) { + require_prepared("Echo-TTS run"); + runtime::validate_spec_backed_request_options(request.options, *contract_, "Echo-TTS"); + + if (!request.text_input.has_value() || request.text_input->text.empty()) { + throw std::runtime_error("Echo-TTS requires text input"); + } + if (!request.voice.has_value() || !request.voice->speaker.has_value() || + !request.voice->speaker->audio.has_value()) { + throw std::runtime_error( + "Echo-TTS requires speaker reference audio; pass --voice-ref " + "(--target-voice is for path-based voice conversion, not cloning)"); + } + + const auto sampler = parse_sampler_options(request.options); + reference_max_samples_ = resolve_reference_max_samples(request.options); + // Encoded once per request; the timbre is then identical across chunk seams + // by construction. + encode_speaker(*request.voice->speaker->audio); + report_stats("speaker.latent", speaker_latent_); + + const int64_t chunk_size = + engine::text::parse_text_chunk_size_override(request.options) + .value_or(kDefaultTextChunkSize); + const auto chunks = runtime::chunk_text_request(request, chunk_size); + if (echo_debug_enabled()) { + std::fprintf( + stderr, "[echo_tts] %zu chunk(s) at a %lld-codepoint budget\n", + chunks.size(), static_cast(chunk_size)); + } + + runtime::TaskResult result; + runtime::AudioBuffer output{kSampleRate, 1, {}}; + for (const auto & chunk : chunks) { + if (!chunk.text_input.has_value() || chunk.text_input->text.empty()) { + continue; + } + auto audio = synthesize_chunk(chunk.text_input->text, sampler); + runtime::append_audio_buffer(output, audio); + } + // Echo's output level is prosody-dependent and can exceed full scale on + // emphatic prompts; upstream's own loader carries a "should we target a + // specific energy level?" note. Divide by the peak only when it exceeds + // 1.0, so quiet output is left untouched and loud output is limited rather + // than clipped at the WAV writer. + engine::audio::normalize_peak_to_unit_range_and_clamp_in_place(output.samples); + result.audio_output = std::move(output); + return result; +} + +void EchoTtsSession::reset() { + speaker_latent_.clear(); + speaker_frames_ = 0; +} + +std::shared_ptr make_echo_tts_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = kFamily; + config.load_assets = load_echo_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::echo_tts diff --git a/src/community_models/echo_tts/tokenizer.cpp b/src/community_models/echo_tts/tokenizer.cpp new file mode 100644 index 00000000..54783691 --- /dev/null +++ b/src/community_models/echo_tts/tokenizer.cpp @@ -0,0 +1,88 @@ +#include "engine/community_models/echo_tts/tokenizer.h" + +#include +#include +#include + +namespace engine::models::echo_tts { +namespace { + +// UTF-8 spellings of the codepoints upstream rewrites. Searching for these as +// byte substrings is safe: UTF-8 is self-synchronising, so a valid multi-byte +// sequence can never match across a character boundary. +constexpr std::string_view kEllipsis = "\xE2\x80\xA6"; // U+2026 +constexpr std::string_view kRightSingleQuote = "\xE2\x80\x99"; // U+2019 +constexpr std::string_view kRightDoubleQuote = "\xE2\x80\x9D"; // U+201D +constexpr std::string_view kEmDash = "\xE2\x80\x94"; // U+2014 + +void replace_all(std::string & text, std::string_view needle, std::string_view replacement) { + if (needle.empty()) { + return; + } + size_t pos = 0; + while ((pos = text.find(needle, pos)) != std::string::npos) { + text.replace(pos, needle.size(), replacement); + pos += replacement.size(); + } +} + +} // namespace + +std::string normalize_echo_text(const std::string & text) { + std::string out = text; + + replace_all(out, kEllipsis, "..."); + replace_all(out, kRightSingleQuote, "'"); + // Upstream applies the right-double-quote rewrite twice and never rewrites + // the *left* double quote (U+201C). Reproduced verbatim so token streams + // match the reference implementation; see inference.py::tokenizer_encode. + replace_all(out, kRightDoubleQuote, "\""); + replace_all(out, kRightDoubleQuote, "\""); + replace_all(out, "\n", " "); + replace_all(out, ":", ","); + replace_all(out, ";", ","); + replace_all(out, kEmDash, ", "); + + const bool has_bracket = !out.empty() && (out.front() == '[' || out.front() == '('); + const bool has_speaker_tag = + out.find("S1") != std::string::npos || out.find("S2") != std::string::npos; + if (!has_bracket && !has_speaker_tag) { + out = "[S1] " + out; + } + return out; +} + +EchoTokenizedText tokenize_echo_text( + const std::string & text, + int64_t max_length, + bool normalize, + bool pad_to_max) { + if (max_length <= 0) { + throw std::runtime_error("Echo-TTS tokenizer requires a positive max_length"); + } + + EchoTokenizedText out; + out.normalized_text = normalize ? normalize_echo_text(text) : text; + + std::vector ids; + ids.reserve(out.normalized_text.size() + 1); + ids.push_back(0); // BOS + for (const char byte : out.normalized_text) { + ids.push_back(static_cast(static_cast(byte))); + } + + const auto encoded_length = static_cast(ids.size()); + const int64_t length = std::min(encoded_length, max_length); + out.truncated = encoded_length > max_length; + + const int64_t output_length = pad_to_max ? max_length : length; + out.input_ids.assign(static_cast(output_length), 0); + out.mask.assign(static_cast(output_length), 0.0F); + for (int64_t i = 0; i < length; ++i) { + out.input_ids[static_cast(i)] = ids[static_cast(i)]; + out.mask[static_cast(i)] = 1.0F; + } + return out; +} + +} // namespace engine::models::echo_tts diff --git a/src/framework/audio/wav_reader.cpp b/src/framework/audio/wav_reader.cpp index d55a9f84..a831eaac 100644 --- a/src/framework/audio/wav_reader.cpp +++ b/src/framework/audio/wav_reader.cpp @@ -1,6 +1,9 @@ #include "engine/framework/audio/wav_reader.h" +#include +#include #include +#include #include #include #include @@ -63,6 +66,89 @@ void skip_bytes(std::istream & input, std::streamoff count) { } } +// WAVE format tags. EXTENSIBLE is the one that matters in practice: many +// encoders emit it for ordinary PCM16 whenever there are more than two channels +// or a channel mask is set, and the real codec then lives in a SubFormat GUID +// rather than in the format tag itself. +constexpr uint16_t kFormatPcm = 0x0001; +constexpr uint16_t kFormatFloat = 0x0003; +constexpr uint16_t kFormatALaw = 0x0006; +constexpr uint16_t kFormatMuLaw = 0x0007; +constexpr uint16_t kFormatExtensible = 0xFFFE; + +// Names a container we can recognise but not decode, so the error can say what +// the file actually is instead of "invalid WAV RIFF header". +const char * identify_foreign_container(const std::array & header) { + const auto * bytes = reinterpret_cast(header.data()); + if (std::memcmp(header.data(), "fLaC", 4) == 0) { + return "FLAC"; + } + if (std::memcmp(header.data(), "OggS", 4) == 0) { + return "Ogg (Vorbis/Opus)"; + } + if (std::memcmp(header.data(), "ID3", 3) == 0) { + return "MP3"; + } + // MPEG audio frame sync: 11 set bits. + if (bytes[0] == 0xFF && (bytes[1] & 0xE0) == 0xE0) { + return "MP3"; + } + if (std::memcmp(header.data() + 4, "ftyp", 4) == 0) { + return "MP4/M4A (AAC or ALAC)"; + } + if (std::memcmp(header.data(), "FORM", 4) == 0) { + return "AIFF"; + } + if (std::memcmp(header.data(), "RF64", 4) == 0) { + return "RF64"; + } + if (std::memcmp(header.data(), "caff", 4) == 0) { + return "CAF"; + } + if (bytes[0] == 0x1A && bytes[1] == 0x45 && bytes[2] == 0xDF && bytes[3] == 0xA3) { + return "Matroska/WebM"; + } + return nullptr; +} + +// G.711 expansion. Both are 8-bit logarithmic codings still common in +// telephony recordings and in WAVs produced by conferencing tools. +float decode_mu_law(uint8_t value) { + value = static_cast(~value); + const int sign = (value & 0x80) != 0 ? -1 : 1; + const int exponent = (value >> 4) & 0x07; + const int mantissa = value & 0x0F; + const int magnitude = ((mantissa << 3) + 0x84) << exponent; + return static_cast(sign * (magnitude - 0x84)) / 32768.0F; +} + +float decode_a_law(uint8_t value) { + value ^= 0x55; + const int sign = (value & 0x80) != 0 ? -1 : 1; + const int exponent = (value >> 4) & 0x07; + const int mantissa = value & 0x0F; + int magnitude = 0; + if (exponent == 0) { + magnitude = (mantissa << 4) + 8; + } else { + magnitude = ((mantissa << 4) + 0x108) << (exponent - 1); + } + return static_cast(sign * magnitude) / 32768.0F; +} + +std::string describe_encoding(uint16_t format, uint16_t bits) { + std::string name; + switch (format) { + case kFormatPcm: name = "PCM"; break; + case kFormatFloat: name = "IEEE float"; break; + case kFormatALaw: name = "A-law"; break; + case kFormatMuLaw: name = "mu-law"; break; + case kFormatExtensible: name = "extensible"; break; + default: name = "format tag " + std::to_string(format); break; + } + return name + ", " + std::to_string(bits) + "-bit"; +} + } // namespace WavData read_wav_f32(std::istream & input) { @@ -70,17 +156,22 @@ WavData read_wav_f32(std::istream & input) { throw std::runtime_error("could not open WAV input"); } - char riff[4]; - input.read(riff, 4); - if (!input || std::string(riff, 4) != "RIFF") { + std::array header{}; + input.read(header.data(), static_cast(header.size())); + const auto header_read = static_cast(input.gcount()); + input.clear(); + input.seekg(static_cast(header_read), std::ios::beg); + + if (header_read < 12 || std::memcmp(header.data(), "RIFF", 4) != 0 || + std::memcmp(header.data() + 8, "WAVE", 4) != 0) { + if (const char * container = identify_foreign_container(header)) { + throw std::runtime_error( + std::string("input is ") + container + + ", not WAV; convert it first, e.g. " + "`ffmpeg -i input -ac 1 -ar 44100 -c:a pcm_s16le output.wav`"); + } throw std::runtime_error("invalid WAV RIFF header"); } - skip_bytes(input, 4); - char wave[4]; - input.read(wave, 4); - if (!input || std::string(wave, 4) != "WAVE") { - throw std::runtime_error("invalid WAV WAVE header"); - } uint16_t audio_format = 0; uint16_t channels = 0; @@ -102,8 +193,18 @@ WavData read_wav_f32(std::istream & input) { sample_rate = read_scalar(input); skip_bytes(input, 6); bits_per_sample = read_scalar(input); - if (chunk_size > 16) { - skip_bytes(input, static_cast(chunk_size - 16)); + std::streamoff consumed = 16; + if (audio_format == kFormatExtensible && chunk_size >= 40) { + skip_bytes(input, 2); // cbSize + skip_bytes(input, 2); // wValidBitsPerSample + skip_bytes(input, 4); // dwChannelMask + // The SubFormat GUID begins with the real format tag. + audio_format = read_scalar(input); + skip_bytes(input, 14); // remainder of the GUID + consumed = 40; + } + if (chunk_size > consumed) { + skip_bytes(input, static_cast(chunk_size) - consumed); } } else if (id == "data") { // chunk_size is a 32-bit field read straight from the file, so a @@ -143,6 +244,54 @@ WavData read_wav_f32(std::istream & input) { wav.sample_rate = static_cast(sample_rate); wav.channels = static_cast(channels); + if (audio_format == kFormatPcm && bits_per_sample == 8) { + // 8-bit PCM in WAV is unsigned, offset by 128. + wav.samples.resize(data.size()); + const auto * pcm = reinterpret_cast(data.data()); + for (size_t i = 0; i < data.size(); ++i) { + wav.samples[i] = (static_cast(pcm[i]) - 128.0F) / 128.0F; + } + return wav; + } + + if (audio_format == kFormatMuLaw && bits_per_sample == 8) { + wav.samples.resize(data.size()); + const auto * pcm = reinterpret_cast(data.data()); + for (size_t i = 0; i < data.size(); ++i) { + wav.samples[i] = decode_mu_law(pcm[i]); + } + return wav; + } + + if (audio_format == kFormatALaw && bits_per_sample == 8) { + wav.samples.resize(data.size()); + const auto * pcm = reinterpret_cast(data.data()); + for (size_t i = 0; i < data.size(); ++i) { + wav.samples[i] = decode_a_law(pcm[i]); + } + return wav; + } + + if (audio_format == kFormatPcm && bits_per_sample == 32) { + const size_t sample_count = data.size() / sizeof(int32_t); + wav.samples.resize(sample_count); + const auto * pcm = reinterpret_cast(data.data()); + for (size_t i = 0; i < sample_count; ++i) { + wav.samples[i] = static_cast(pcm[i]) / 2147483648.0F; + } + return wav; + } + + if (audio_format == kFormatFloat && bits_per_sample == 64) { + const size_t sample_count = data.size() / sizeof(double); + wav.samples.resize(sample_count); + const auto * pcm = reinterpret_cast(data.data()); + for (size_t i = 0; i < sample_count; ++i) { + wav.samples[i] = static_cast(pcm[i]); + } + return wav; + } + if (audio_format == 1 && bits_per_sample == 16) { const size_t sample_count = data.size() / sizeof(int16_t); wav.samples.resize(sample_count); @@ -184,7 +333,10 @@ WavData read_wav_f32(std::istream & input) { return wav; } - throw std::runtime_error("unsupported WAV encoding (need PCM16, PCM24, or float32)"); + throw std::runtime_error( + "unsupported WAV encoding (" + describe_encoding(audio_format, bits_per_sample) + + "); supported: PCM 8/16/24/32-bit, float 32/64-bit, A-law and mu-law. " + "Convert with `ffmpeg -i input -ac 1 -ar 44100 -c:a pcm_s16le output.wav`"); } WavData read_wav_f32(std::string_view input) { diff --git a/src/models/fish_audio/codec.cpp b/src/models/fish_audio/codec.cpp index e6705eb9..1f4f6d2f 100644 --- a/src/models/fish_audio/codec.cpp +++ b/src/models/fish_audio/codec.cpp @@ -549,9 +549,11 @@ core::TensorValue build_quantizer_out( return modules::Conv1dModule({8, kCodecDim, 1, 1, 0, 1, true}).build(ctx, emb_bdt, weights.out_proj); } -core::TensorValue build_decode_quantizer( +// Dequantises codes into the continuous z_q space, which is the boundary the +// Echo-TTS family enters at. Kept separate from the post_module/upsample tail so +// both a code-driven decode and a latent-driven decode can share that tail. +core::TensorValue build_zq_from_codes( 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); @@ -559,7 +561,17 @@ core::TensorValue build_decode_quantizer( 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); + return latent; +} + +// post_module -> upsample. This is exactly the body of DAC.decode_zq in the +// upstream Echo-TTS autoencoder, which is why it is factored out. +core::TensorValue build_decode_from_zq( + core::ModuleBuildContext & ctx, + core::ConstantTensorCache & constants, + const core::TensorValue & z_q, + const FishCodecWeights & weights) { + auto latent = build_window_transformer(ctx, constants, z_q, 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); @@ -567,13 +579,23 @@ core::TensorValue build_decode_quantizer( return latent; } +core::TensorValue build_decode_quantizer( + core::ModuleBuildContext & ctx, + core::ConstantTensorCache & constants, + const std::vector & code_inputs, + const FishCodecWeights & weights) { + return build_decode_from_zq( + ctx, constants, build_zq_from_codes(ctx, code_inputs, weights), weights); +} + 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) { + std::vector> & trace_outputs, + core::TensorValue * z_q_out) { auto x = encoder_latent; for (const auto & stage : weights.downsample) { x = causal_conv1d(ctx, x, stage.first, kCodecDim, kCodecDim, 2, 2, 1, true); @@ -607,6 +629,14 @@ core::TensorValue build_encode_quantizer( for (const auto & quantizer : weights.residual_quantizers) { quantize_one(quantizer, 1024); } + if (z_q_out != nullptr) { + // DAC.encode_zq sums the dequantised contribution of every codebook. + // Each quantize_one subtracts exactly that contribution from `residual`, + // which starts at `x`, so the sum is the difference between the two. + // This avoids a second accumulator and stays exact by construction. + *z_q_out = core::wrap_tensor( + ggml_sub(ctx.ggml, x.tensor, residual.tensor), x.shape, GGML_TYPE_F32); + } return x; } @@ -936,6 +966,104 @@ struct DecodeGraph { core::ConstantTensorCache constants_; }; +// Decodes continuous z_q latents rather than discrete codes. Shares the whole +// post_module -> upsample -> decoder tail with DecodeGraph via +// build_decode_from_zq; only the input differs. Added for Echo-TTS, whose DiT +// generates latents directly and never produces codebook indices. +struct LatentDecodeGraph { + LatentDecodeGraph( + std::shared_ptr assets, + std::shared_ptr weights, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + int64_t frames) + : assets_(std::move(assets)), + weights_(std::move(weights)), + backend_(execution_context.backend()), + backend_type_(execution_context.backend_type()), + threads_(std::max(1, execution_context.config().threads)), + frame_capacity_(frames), + constants_(backend_, threads_, "Fish Audio codec latent decode constants") { + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio codec latent decode graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.codec.latent_decode", backend_type_}; + constants_.begin_graph(); + // (batch, channels, time), matching what the quantizer tail expects. + latent_input_ = core::make_tensor( + ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, kCodecDim, frame_capacity_})); + ggml_set_input(latent_input_.tensor); + auto latent = build_decode_from_zq(ctx, constants_, latent_input_, *weights_); + auto waveform = build_decoder(ctx, latent, *weights_); + output_ = waveform.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 1048576, false); + ggml_build_forward_expand(graph_, output_); + constants_.finish_graph(); + constants_.ensure_uploaded(); + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_))); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_.get(), graph_)) { + throw std::runtime_error("failed to allocate Fish Audio codec latent decode graph"); + } + } + + ~LatentDecodeGraph() { + 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); + } + + // `latents` is (frames, kCodecDim) row-major, which is the layout the PCA + // inverse produces. It is transposed here into the (channels, time) order + // ggml holds, and zero-padded out to the graph's frame capacity. + runtime::AudioBuffer run(const std::vector & latents, int64_t frames) { + if (frames <= 0 || static_cast(latents.size()) != frames * kCodecDim) { + throw std::runtime_error("Fish Audio codec latent decode received a mis-shaped latent buffer"); + } + if (frames > frame_capacity_) { + throw std::runtime_error("Fish Audio codec latent decode request exceeds graph capacity"); + } + std::vector padded(static_cast(kCodecDim * frame_capacity_), 0.0F); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t channel = 0; channel < kCodecDim; ++channel) { + padded[static_cast(channel * frame_capacity_ + frame)] = + latents[static_cast(frame * kCodecDim + channel)]; + } + } + core::write_tensor_f32(latent_input_, padded); + core::set_backend_threads(backend_, threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio codec latent decode graph compute failed"); + } + auto values = core::read_tensor_f32(output_); + const int64_t expected_samples = 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_; + core::TensorValue latent_input_; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; + core::ConstantTensorCache constants_; +}; + struct EncodeGraph { EncodeGraph( std::shared_ptr assets, @@ -943,7 +1071,8 @@ struct EncodeGraph { core::ExecutionContext & execution_context, size_t graph_arena_bytes, int64_t samples, - int64_t frames) + int64_t frames, + bool want_z_q) : assets_(std::move(assets)), weights_(std::move(weights)), backend_(execution_context.backend()), @@ -951,6 +1080,7 @@ struct EncodeGraph { threads_(std::max(1, execution_context.config().threads)), sample_capacity_(samples), frame_capacity_(frames), + wants_z_q_(want_z_q), constants_(backend_, threads_, "Fish Audio codec encode constants") { ggml_init_params params{graph_arena_bytes, nullptr, true}; ctx_.reset(ggml_init(params)); @@ -963,8 +1093,25 @@ struct EncodeGraph { ggml_set_input(input_.tensor); auto encoded = build_encoder(ctx, constants_, input_, *weights_); trace_outputs_.push_back({"fish_audio.codec.encoder_latent", encoded}); - build_encode_quantizer(ctx, constants_, encoded, *weights_, code_outputs_, trace_outputs_); + // Continuous latents are what Echo-TTS conditions on; fish_audio itself + // only needs the codes. Building the extra node unconditionally would + // not change the arithmetic, but ggml_set_output pins that buffer and + // keeps both `x` and the quantiser residual live to the end of the + // graph, where otherwise the residual is free to be reused. That is a + // real allocation change for a family that gains nothing from it, so + // the output only exists when a caller has asked for it. + core::TensorValue z_q; + build_encode_quantizer( + ctx, constants_, encoded, *weights_, code_outputs_, trace_outputs_, + wants_z_q_ ? &z_q : nullptr); + if (wants_z_q_) { + z_q_output_ = core::ensure_backend_addressable_layout(ctx, z_q).tensor; + ggml_set_output(z_q_output_); + } graph_ = ggml_new_graph_custom(ctx_.get(), 1048576, false); + if (z_q_output_ != nullptr) { + ggml_build_forward_expand(graph_, z_q_output_); + } for (const auto & trace_output : trace_outputs_) { ggml_set_output(trace_output.second.tensor); ggml_build_forward_expand(graph_, trace_output.second.tensor); @@ -984,11 +1131,18 @@ struct EncodeGraph { engine::core::release_backend_graph_resources(backend_, graph_); } - bool matches(int64_t samples, int64_t frames, ggml_backend_t backend, int threads) const { + // A graph that also produces z_q can serve a codes-only request -- the + // codes are identical either way -- but not the reverse. In practice the + // two never mix on one codec instance: fish_audio only ever calls + // encode_reference and echo_tts only ever calls encode_zq, so fish_audio + // never gets the z_q-bearing graph at all. + bool matches(int64_t samples, int64_t frames, ggml_backend_t backend, int threads, + bool want_z_q) const { return sample_capacity_ >= samples && frame_capacity_ >= frames && backend_ == backend && - threads_ == std::max(1, threads); + threads_ == std::max(1, threads) && + (wants_z_q_ || !want_z_q); } FishAudioCodes run(const runtime::AudioBuffer & audio) { @@ -1032,14 +1186,39 @@ struct EncodeGraph { return out; } + // Continuous latents from the most recent encode, shaped + // (frames, kCodecDim) row-major. Valid only after a run() on a graph that + // was built with want_z_q, which is what encode_zq does. + std::vector read_z_q(int64_t frames) const { + if (z_q_output_ == nullptr) { + throw std::runtime_error("Fish Audio codec encode graph was not built with z_q output"); + } + auto values = core::read_tensor_f32(z_q_output_); + const size_t wanted = static_cast(frames * kCodecDim); + if (values.size() < wanted) { + throw std::runtime_error("Fish Audio codec z_q output is smaller than the frame count"); + } + // The graph is built for frame_capacity_; drop the padded tail. + std::vector out(wanted); + for (int64_t frame = 0; frame < frames; ++frame) { + for (int64_t channel = 0; channel < kCodecDim; ++channel) { + out[static_cast(frame * kCodecDim + channel)] = + values[static_cast(channel * frame_capacity_ + frame)]; + } + } + return out; + } + private: std::shared_ptr assets_; std::shared_ptr weights_; ggml_backend_t backend_ = nullptr; + ggml_tensor * z_q_output_ = nullptr; core::BackendType backend_type_ = core::BackendType::Cpu; int threads_ = 1; int64_t sample_capacity_ = 0; int64_t frame_capacity_ = 0; + bool wants_z_q_ = false; std::unique_ptr ctx_; core::TensorValue input_; std::vector code_outputs_; @@ -1079,8 +1258,10 @@ class FishAudioCodecRuntime::Impl { 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); + if (encode_graph_ == nullptr || + !encode_graph_->matches(samples, frames, execution_.backend(), threads_, false)) { + encode_graph_ = std::make_unique( + assets_, weights_, execution_, graph_arena_bytes_, samples, frames, false); } return encode_graph_->run(audio); } @@ -1092,6 +1273,37 @@ class FishAudioCodecRuntime::Impl { return decode_graph_->run(codes); } + // Continuous-latent counterparts of encode_reference/decode, matching + // DAC.encode_zq and DAC.decode_zq in the upstream Echo-TTS autoencoder. + FishAudioLatents encode_zq(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_, true)) { + encode_graph_ = std::make_unique( + assets_, weights_, execution_, graph_arena_bytes_, samples, frames, true); + } + // The codes are discarded; running the same graph keeps the quantiser + // path identical to encode_reference so the two cannot drift. + (void)encode_graph_->run(audio); + FishAudioLatents out; + out.frames = frames; + out.channels = kCodecDim; + out.values = encode_graph_->read_z_q(frames); + return out; + } + + runtime::AudioBuffer decode_zq(const std::vector & latents, int64_t frames) { + if (latent_decode_graph_ == nullptr || + !latent_decode_graph_->matches(frames, execution_.backend(), threads_)) { + latent_decode_graph_ = + std::make_unique(assets_, weights_, execution_, graph_arena_bytes_, frames); + } + return latent_decode_graph_->run(latents, frames); + } + void release_encode_graph() { encode_graph_.reset(); } @@ -1099,6 +1311,7 @@ class FishAudioCodecRuntime::Impl { void release_runtime_graphs() { encode_graph_.reset(); decode_graph_.reset(); + latent_decode_graph_.reset(); } private: @@ -1109,6 +1322,7 @@ class FishAudioCodecRuntime::Impl { std::shared_ptr weights_; std::unique_ptr encode_graph_; std::unique_ptr decode_graph_; + std::unique_ptr latent_decode_graph_; }; FishAudioCodecRuntime::FishAudioCodecRuntime( @@ -1138,6 +1352,14 @@ runtime::AudioBuffer FishAudioCodecRuntime::decode(const FishAudioCodes & codes) return impl_->decode(codes); } +FishAudioLatents FishAudioCodecRuntime::encode_zq(const runtime::AudioBuffer & audio) { + return impl_->encode_zq(audio); +} + +runtime::AudioBuffer FishAudioCodecRuntime::decode_zq(const std::vector & latents, int64_t frames) { + return impl_->decode_zq(latents, frames); +} + void FishAudioCodecRuntime::release_encode_graph() { impl_->release_encode_graph(); } diff --git a/tests/echo_tts/echo_tts_dit_parity.cpp b/tests/echo_tts/echo_tts_dit_parity.cpp new file mode 100644 index 00000000..9794a82c --- /dev/null +++ b/tests/echo_tts/echo_tts_dit_parity.cpp @@ -0,0 +1,377 @@ +// Numerical parity harness for the Echo-TTS DiT graph. +// +// Not registered with add_test: it needs the 5.5 GB GGUF and a reference dump +// from the upstream PyTorch implementation, so it is driven by hand, the same +// way dots_tts_vocoder_parity is. +// +// python3 tools/community_models/echo_tts_reference.py --speaker ref.wav +// --full-blocks -o echo_ref.npz +// python3 tools/community_models/echo_tts_pack_reference.py echo_ref.npz +// -o echo_ref.bin +// ./echo_tts_dit_parity --model /path/to/Echo-TTS-GGUF --reference echo_ref.bin +// +// Two checks, deliberately separate: +// +// denoiser feeds the reference's own x and t through one conditional +// forward, removing the sampler from the comparison entirely. +// +// It does NOT isolate the DiT blocks by themselves. The reference +// text ids and speaker latents are injected, but prepare_conditioning() +// then runs our own text encoder, speaker encoder and KV +// projections, so a difference here could originate in any of +// them. It is a combined conditioning-plus-denoiser comparison, +// which is still enough to catch a wrong block -- nothing is being +// compared against itself -- but not enough to localise one. +// +// sampler runs the full 40-step trajectory. Run it twice: once from the +// reference's own initial noise, which removes the RNG from the +// comparison, and once from our seeded draw. +// +// Neither is expected to be exact, and the dominant term is NOT the +// RNG. The reference defaults to bfloat16 while our GGUF is F16; +// the two round differently, and dual CFG at 3.0/8.0 amplifies the +// per-step difference every step. Dumping the reference with +// --force-dtype float16 moves the 40-step cosine from 0.905 to +// 0.977 and the denoiser probe from 0.999977 to 0.999999. Compare +// like dtypes or expect the gap. Cosine, never equality. + +#include "engine/community_models/echo_tts/config.h" +#include "engine/community_models/echo_tts/dit.h" +#include "engine/community_models/echo_tts/sampler.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/model_spec/package.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// --- reference bundle ------------------------------------------------------ + +struct Tensor { + bool is_int = false; + std::vector f32; + std::vector i32; + + int64_t size() const { return is_int ? static_cast(i32.size()) + : static_cast(f32.size()); } +}; + +class ReferenceBundle { +public: + // The largest bundle the packer emits with --blocks is 24 x 640 x 2048 + // floats in one entry; these caps sit well above that and well below + // anything that could exhaust memory. + static constexpr int32_t kMaxEntries = 4096; + static constexpr int32_t kMaxNameLength = 1024; + static constexpr int64_t kMaxElements = 1LL << 32; + + explicit ReferenceBundle(const std::filesystem::path & path) { + std::ifstream in(path, std::ios::binary); + if (!in) { + throw std::runtime_error("cannot open reference bundle: " + path.string()); + } + char magic[8] = {}; + in.read(magic, 8); + if (std::memcmp(magic, "ECHOPAR1", 8) != 0) { + throw std::runtime_error("not an ECHOPAR1 bundle: " + path.string()); + } + // Every length below is signed on the wire and comes from a file this + // process did not write. Validate before it reaches resize(), or a + // negative value becomes a huge size_t and a malformed header becomes + // an enormous allocation instead of a format error. + const int32_t count = read_i32(in); + if (count < 0 || count > kMaxEntries) { + throw std::runtime_error("reference bundle declares an implausible entry count"); + } + for (int32_t i = 0; i < count; ++i) { + const int32_t name_len = read_i32(in); + if (name_len < 0 || name_len > kMaxNameLength) { + throw std::runtime_error("reference bundle has an implausible tensor name length"); + } + std::string name(static_cast(name_len), '\0'); + in.read(name.data(), name_len); + const int32_t dtype = read_i32(in); + if (dtype != 0 && dtype != 1) { + throw std::runtime_error("reference bundle has an unknown dtype tag"); + } + const int64_t elements = read_i64(in); + if (elements < 0 || elements > kMaxElements) { + throw std::runtime_error("reference bundle declares an implausible element count"); + } + if (!in) { + throw std::runtime_error("truncated reference bundle header at entry " + name); + } + + Tensor tensor; + tensor.is_int = dtype == 1; + if (tensor.is_int) { + tensor.i32.resize(static_cast(elements)); + in.read(reinterpret_cast(tensor.i32.data()), elements * 4); + } else { + tensor.f32.resize(static_cast(elements)); + in.read(reinterpret_cast(tensor.f32.data()), elements * 4); + } + if (!in) { + throw std::runtime_error("truncated reference bundle at entry " + name); + } + entries_.emplace(std::move(name), std::move(tensor)); + } + } + + const Tensor & at(const std::string & name) const { + const auto it = entries_.find(name); + if (it == entries_.end()) { + throw std::runtime_error("reference bundle has no tensor named " + name); + } + return it->second; + } + +private: + // The packer emits little-endian and documents that audio.cpp targets only + // little-endian hosts, so these native reads are correct here. They would + // need byte-swapping on a big-endian build. + static int32_t read_i32(std::istream & in) { + int32_t value = 0; + in.read(reinterpret_cast(&value), 4); + return value; + } + static int64_t read_i64(std::istream & in) { + int64_t value = 0; + in.read(reinterpret_cast(&value), 8); + return value; + } + + std::map entries_; +}; + +// --- metrics --------------------------------------------------------------- + +struct Metrics { + double cosine = 0.0; + double max_abs_error = 0.0; + double rms_error = 0.0; +}; + +// Cosine over the flattened tensors, reported with max-absolute-error beside it. +// Cosine alone hides a uniform scale error; max-abs alone is dominated by one +// outlier. The pair is what the port's own gate is written against. +Metrics compare(const std::vector & actual, const std::vector & expected) { + if (actual.size() != expected.size()) { + throw std::runtime_error( + "size mismatch: actual=" + std::to_string(actual.size()) + + " expected=" + std::to_string(expected.size())); + } + double dot = 0.0; + double norm_a = 0.0; + double norm_b = 0.0; + double sq = 0.0; + Metrics metrics; + for (size_t i = 0; i < actual.size(); ++i) { + const double a = actual[i]; + const double b = expected[i]; + dot += a * b; + norm_a += a * a; + norm_b += b * b; + const double diff = std::fabs(a - b); + sq += diff * diff; + metrics.max_abs_error = std::max(metrics.max_abs_error, diff); + } + const double denom = std::sqrt(norm_a) * std::sqrt(norm_b); + metrics.cosine = denom > 0.0 ? dot / denom : 0.0; + metrics.rms_error = std::sqrt(sq / static_cast(actual.size())); + return metrics; +} + +// Cosine alone is not a gate: `actual = 1000 * expected` scores a perfect 1.0 +// while being catastrophically wrong in amplitude. The max-absolute error is +// what closes that hole, so both must hold for a PASS. +bool report(const std::string & label, const Metrics & m, double gate, double max_abs_gate) { + const bool pass = m.cosine >= gate && m.max_abs_error <= max_abs_gate; + std::cout << std::left << std::setw(10) << label + << " cosine=" << std::fixed << std::setprecision(9) << m.cosine + << " max_abs=" << std::setprecision(6) << m.max_abs_error + << " rms=" << m.rms_error + << " gate=" << std::setprecision(3) << gate + << "/" << max_abs_gate + << (pass ? " PASS" : " FAIL") << "\n"; + return pass; +} + +// --- arg parsing ----------------------------------------------------------- + +std::string arg_value(int argc, char ** argv, const std::string & name, + const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +engine::core::BackendType parse_backend(const std::string & value) { + if (value == "cuda") { + return engine::core::BackendType::Cuda; + } + if (value == "vulkan") { + return engine::core::BackendType::Vulkan; + } + if (value == "cpu") { + return engine::core::BackendType::Cpu; + } + if (value == "best") { + return engine::core::BackendType::BestAvailable; + } + throw std::runtime_error("echo_tts_dit_parity supports cuda, vulkan, cpu, or best"); +} + +bool has_flag(int argc, char ** argv, const std::string & name) { + for (int i = 1; i < argc; ++i) { + if (argv[i] == name) { + return true; + } + } + return false; +} + +std::vector to_float(const Tensor & tensor) { + if (!tensor.is_int) { + return tensor.f32; + } + std::vector out(tensor.i32.size()); + std::transform(tensor.i32.begin(), tensor.i32.end(), out.begin(), + [](int32_t v) { return static_cast(v); }); + return out; +} + +} // namespace + +int main(int argc, char ** argv) try { + const std::filesystem::path model_path = arg_value(argc, argv, "--model", ""); + const std::filesystem::path reference_path = arg_value(argc, argv, "--reference", ""); + const std::string backend_name = arg_value(argc, argv, "--backend", "cuda"); + const double denoiser_gate = std::stod(arg_value(argc, argv, "--denoiser-gate", "0.999")); + const double sampler_gate = std::stod(arg_value(argc, argv, "--sampler-gate", "0.999")); + // Amplitude gates, deliberately loose relative to the cosine gate: they + // exist to catch a scale error that cosine cannot see, not to re-litigate + // the rounding difference the cosine gate already bounds. + const double denoiser_max_abs = + std::stod(arg_value(argc, argv, "--denoiser-max-abs", "0.25")); + const double sampler_max_abs = + std::stod(arg_value(argc, argv, "--sampler-max-abs", "4.0")); + const bool skip_sampler = has_flag(argc, argv, "--skip-sampler"); + + if (model_path.empty() || reference_path.empty()) { + std::cerr << "usage: echo_tts_dit_parity --model --reference \n" + << " [--backend cuda|vulkan|cpu|best] [--denoiser-gate 0.999]\n" + << " [--sampler-gate 0.999] [--denoiser-max-abs 0.25]\n" + << " [--sampler-max-abs 4.0] [--skip-sampler]\n"; + return 2; + } + + const ReferenceBundle reference(reference_path); + + engine::core::BackendConfig backend_config; + backend_config.type = parse_backend(backend_name); + engine::core::ExecutionContext execution(backend_config); + + auto bundle = engine::model_spec::load_resource_bundle_for_family(model_path, "echo_tts"); + auto dit_weights = bundle.open_tensor_source("dit_weights"); + + engine::models::echo_tts::EchoTtsConfig config; + config.validate(); + + engine::models::echo_tts::EchoDitRuntime dit( + config, *dit_weights, "", execution, engine::assets::TensorStorageType::Native); + + // Inject the reference's own conditioning rather than recomputing it, so + // this measures the DiT and not the speaker encoder feeding it. + engine::models::echo_tts::EchoConditioning conditioning; + const auto & text_ids = reference.at("text.input_ids"); + conditioning.text_input_ids = text_ids.i32; + conditioning.text_mask = to_float(reference.at("text.mask")); + conditioning.text_length = text_ids.size(); + + const auto & speaker_latent = reference.at("speaker.latent"); + conditioning.speaker_latent = speaker_latent.f32; + conditioning.speaker_mask = to_float(reference.at("speaker.mask")); + conditioning.speaker_frames = speaker_latent.size() / config.latent_size; + + std::cout << "text_length=" << conditioning.text_length + << " speaker_frames=" << conditioning.speaker_frames << "\n"; + + dit.prepare_conditioning(conditioning); + + bool ok = true; + + // 1. Fixed-timestep denoiser probe. + { + const auto & x_input = reference.at("dit.x_input"); + const auto t = static_cast(reference.at("dit.t").f32.at(0)); + const auto predicted = dit.denoise_once(x_input.f32, t); + std::cout << "denoiser probe at t=" << std::fixed << std::setprecision(4) << t << "\n"; + ok &= report("denoiser", compare(predicted, reference.at("dit.v_pred").f32), denoiser_gate, + denoiser_max_abs); + } + + // 2. Sampler driven from the reference's OWN initial noise. + // + // This is the discriminator. Check 3 below runs the sampler from its own + // seeded draw, which cannot be bit-identical to CUDA's. If that one + // diverges while this one holds, the divergence is the RNG plus the + // trajectory's sensitivity to it, not a defect in the integration. If this + // one also diverges, the sampler itself is wrong. + if (!skip_sampler) { + engine::models::echo_tts::EchoSamplerOptions options; + options.num_steps = static_cast(reference.at("config.steps").i32.at(0)); + options.sequence_length = reference.at("config.sequence_length").i32.at(0); + options.window_pinned = true; + + auto denoise = [&dit](const std::vector & x, float t, int lanes) { + return dit.denoise_once(x, t, lanes); + }; + const auto latent = engine::models::echo_tts::run_euler_sampler( + options, + options.sequence_length, + config.latent_size, + reference.at("sampler.initial_noise").f32, + denoise); + std::cout << "sampler, reference initial noise injected\n"; + ok &= report("injected", compare(latent, reference.at("sampler.latent").f32), sampler_gate, + sampler_max_abs); + } + + // 3. Full sampler trajectory from our own seeded noise. Expected to be + // close, not exact: see header. + if (!skip_sampler) { + engine::models::echo_tts::EchoSamplerOptions options; + options.num_steps = static_cast(reference.at("config.steps").i32.at(0)); + options.sequence_length = reference.at("config.sequence_length").i32.at(0); + options.seed = reference.at("config.seed").i32.at(0); + options.window_pinned = true; + const auto latent = dit.sample(options); + std::cout << "sampler steps=" << options.num_steps + << " sequence_length=" << options.sequence_length + << " seed=" << options.seed << "\n"; + ok &= report("sampler", compare(latent, reference.at("sampler.latent").f32), sampler_gate, + sampler_max_abs); + } + + std::cout << (ok ? "echo_tts_dit_parity: ok\n" : "echo_tts_dit_parity: FAILED\n"); + return ok ? 0 : 1; +} catch (const std::exception & ex) { + std::cerr << "echo_tts_dit_parity: " << ex.what() << "\n"; + return 1; +} diff --git a/tests/echo_tts/echo_tts_host_units.cpp b/tests/echo_tts/echo_tts_host_units.cpp new file mode 100644 index 00000000..6a0ba09c --- /dev/null +++ b/tests/echo_tts/echo_tts_host_units.cpp @@ -0,0 +1,378 @@ +// Host-side unit tests for the Echo-TTS port. +// +// These cover the parts of the pipeline that run on the CPU and need neither a +// GPU nor the 5.5 GB checkpoint: the byte tokenizer and its WhisperD +// normalisation, the PCA forward/inverse pair, and the flattening-point crop +// that sets the output duration. +// +// Every expected value here was produced by executing the reference +// implementation at tts-bench/venvs/echo/src/inference.py -- `tokenizer_encode` +// for the token streams and `find_flattening_point` for the crop indices -- not +// by reasoning about what it ought to return. +// +// Two review findings shaped the fixtures, and both are worth stating so they +// are not "simplified" back out: +// +// * The PCA basis here is RECTANGULAR and non-symmetric on purpose. An +// identity basis is its own transpose, so a round trip over one cannot +// distinguish `components[c * features + k]` from the transposed indexing, +// and a mean or scale dropped on both legs cancels. Both projection and +// inversion are therefore pinned against independently computed values +// rather than against each other. +// +// * The flattening fixtures include a tail that is quiet but NOT zero and a +// tail that is flat but too loud. Only all-zero fixtures would let an +// implementation that merely searches for a zero window pass without ever +// evaluating the standard-deviation and mean thresholds. + +#include "engine/community_models/echo_tts/config.h" +#include "engine/community_models/echo_tts/latent_post.h" +#include "engine/community_models/echo_tts/tokenizer.h" + +#include "../unittests/test_assert.h" + +#include +#include +#include +#include +#include + +namespace { + +using engine::test::require; +using engine::test::require_eq; + +using namespace engine::models::echo_tts; + +constexpr int64_t kMaxLength = 768; // upstream's hard cap + +// The shared require_close compares `fabs(a - b) > tolerance`, which is FALSE +// for a NaN difference, so a NaN silently passes every float assertion. Reject +// non-finite values explicitly before deferring to it. +void require_close(float actual, float expected, float tolerance, const std::string & label) { + if (!std::isfinite(actual)) { + throw std::runtime_error(label + " is not finite"); + } + engine::test::require_close(actual, expected, tolerance, label); +} + +void require_ids(const std::vector & actual, + const std::vector & expected, + const std::string & label) { + require_eq(static_cast(actual.size()), static_cast(expected.size()), + label + " length"); + for (size_t i = 0; i < expected.size(); ++i) { + require_eq(actual[i], expected[i], label + " id " + std::to_string(i)); + } +} + +std::vector encode(const std::string & text) { + return tokenize_echo_text(text, kMaxLength).input_ids; +} + +// --- tokenizer ------------------------------------------------------------- + +void test_normalisation_matches_reference() { + require_eq(normalize_echo_text("Hello world."), std::string("[S1] Hello world."), + "bare text gets the [S1] tag"); + + require_eq(normalize_echo_text("[S1] Already tagged."), std::string("[S1] Already tagged."), + "an existing tag is not doubled"); + + require_eq(normalize_echo_text("(parenthesised start)"), std::string("(parenthesised start)"), + "a leading paren suppresses the tag"); + + // Upstream's tag check is a bare substring search for "S1"/"S2" anywhere in + // the string, not a prefix check, so prose containing those two characters + // loses the speaker tag. Faithful, and pinned so it stays faithful. + require_eq(normalize_echo_text("This is S1 talking"), std::string("This is S1 talking"), + "a bare S1 anywhere suppresses the tag, as upstream does"); + + // Colons and semicolons become commas, an em dash becomes ", ", an ellipsis + // becomes "...", a right single quote becomes an apostrophe, and a newline + // becomes a space. + // + // The asymmetric quote handling is deliberate and is reproduced from + // upstream: the RIGHT double quote is rewritten to ASCII, the LEFT one is + // not. Upstream applies the right-quote replacement twice, which is a + // no-op, and never touches U+201C. If someone "fixes" that asymmetry the + // token stream silently stops matching the reference, so it is pinned here. + require_eq( + normalize_echo_text("Time: 3; place \xE2\x80\x94 here\xE2\x80\xA6 he said " + "\xE2\x80\x9Cgo\xE2\x80\x9D and it\xE2\x80\x99s fine.\nNext line."), + std::string("[S1] Time, 3, place , here... he said \xE2\x80\x9Cgo\" and it's fine. Next line."), + "punctuation rewrites match the reference"); +} + +void test_tokenisation_matches_reference() { + // Full id vectors, not lengths. A length-preserving rewrite -- signed char + // sign-extension on the multibyte U+201C below being the obvious one -- + // passes a count check and fails these. + require_ids(encode("Hello world."), + {0, 91, 83, 49, 93, 32, 72, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100, 46}, + "hello"); + + require_ids(encode("[S1] Already tagged."), + {0, 91, 83, 49, 93, 32, 65, 108, 114, 101, 97, 100, 121, 32, 116, 97, 103, 103, + 101, 100, 46}, + "pre-tagged"); + + require_ids(encode("(parenthesised start)"), + {0, 40, 112, 97, 114, 101, 110, 116, 104, 101, 115, 105, 115, 101, 100, 32, 115, + 116, 97, 114, 116, 41}, + "paren"); + + require_ids(encode("This is S1 talking"), + {0, 84, 104, 105, 115, 32, 105, 115, 32, 83, 49, 32, 116, 97, 108, 107, 105, 110, + 103}, + "bare S1"); + + // 226/128/156 is the untouched left double quote: three unsigned bytes. + require_ids(encode("Time: 3; place \xE2\x80\x94 here\xE2\x80\xA6 he said " + "\xE2\x80\x9Cgo\xE2\x80\x9D and it\xE2\x80\x99s fine.\nNext line."), + {0, 91, 83, 49, 93, 32, 84, 105, 109, 101, 44, 32, 51, 44, 32, 112, 108, 97, 99, + 101, 32, 44, 32, 32, 104, 101, 114, 101, 46, 46, 46, 32, 104, 101, 32, 115, 97, + 105, 100, 32, 226, 128, 156, 103, 111, 34, 32, 97, 110, 100, 32, 105, 116, 39, + 115, 32, 102, 105, 110, 101, 46, 32, 78, 101, 120, 116, 32, 108, 105, 110, 101, + 46}, + "punctuation-heavy"); +} + +void test_tokeniser_truncates_at_max_length() { + const std::string long_text(4000, 'a'); + const auto tokens = tokenize_echo_text(long_text, kMaxLength); + require_eq(static_cast(tokens.input_ids.size()), kMaxLength, + "truncated length includes the BOS"); + require(tokens.truncated, "over-long input is reported as truncated"); + require_eq(tokens.input_ids.front(), static_cast(0), "BOS survives truncation"); + + // The surviving ids are the *leading* bytes, not zero padding: everything + // after the BOS is the tag "[S1] " and then 'a'. + const std::vector head{0, 91, 83, 49, 93, 32, 97, 97}; + for (size_t i = 0; i < head.size(); ++i) { + require_eq(tokens.input_ids[i], head[i], "truncated head id " + std::to_string(i)); + } + require_eq(tokens.input_ids.back(), static_cast('a'), "truncation keeps a real byte"); + for (const float m : tokens.mask) { + require_close(m, 1.0F, 1e-6F, "a fully truncated sequence has no padding"); + } + + const auto shortish = tokenize_echo_text("Hello world.", kMaxLength); + require(!shortish.truncated, "short input is not reported as truncated"); +} + +void test_mask_marks_real_tokens() { + const auto tokens = tokenize_echo_text("Hello world.", kMaxLength); + require_eq(tokens.mask.size(), tokens.input_ids.size(), "mask and ids are the same length"); + for (size_t i = 0; i < tokens.mask.size(); ++i) { + require_close(tokens.mask[i], 1.0F, 1e-6F, "unpadded mask entry " + std::to_string(i)); + } +} + +void test_pad_to_max_zeroes_the_tail() { + const auto padded = tokenize_echo_text("Hello world.", kMaxLength, true, true); + require_eq(static_cast(padded.input_ids.size()), kMaxLength, "padded id length"); + require_eq(static_cast(padded.mask.size()), kMaxLength, "padded mask length"); + + constexpr int64_t kReal = 18; // "[S1] Hello world." plus the BOS + for (int64_t i = 0; i < kMaxLength; ++i) { + const auto idx = static_cast(i); + if (i < kReal) { + require_close(padded.mask[idx], 1.0F, 1e-6F, "real mask " + std::to_string(i)); + } else { + require_close(padded.mask[idx], 0.0F, 1e-6F, "pad mask " + std::to_string(i)); + require_eq(padded.input_ids[idx], static_cast(0), + "pad id " + std::to_string(i)); + } + } +} + +// --- PCA ------------------------------------------------------------------- + +// A rectangular, non-symmetric orthonormal basis: 2 components over 4 features. +// Rectangular so a transposed read is a different computation rather than the +// same one; orthonormal so the projected coefficients are exact in binary +// floating point and can be written down by hand. +// +// b0 = [ 0.5, 0.5, 0.5, 0.5] +// b1 = [ 0.5, -0.5, 0.5, -0.5] +// mean = [1, 2, 3, 4], latent_scale = 0.5 +// +// Because 2 components cannot span 4 features, the round trip is deliberately +// LOSSY -- it returns the projection of the input onto span{b0, b1}, which is +// what the real (80, 1024) basis does too. An identity fixture hides that. +constexpr int64_t kFeatures = 4; +constexpr int64_t kComponents = 2; + +EchoPcaState rectangular_pca() { + EchoPcaState pca; + pca.components = {0.5F, 0.5F, 0.5F, 0.5F, + 0.5F, -0.5F, 0.5F, -0.5F}; + pca.mean = {1.0F, 2.0F, 3.0F, 4.0F}; + pca.latent_scale = 0.5F; + return pca; +} + +EchoTtsConfig rectangular_config() { + EchoTtsConfig config; + config.latent_size = kComponents; + config.ae_latent_dim = kFeatures; + return config; +} + +void test_pca_projection_matches_hand_computed_values() { + const auto pca = rectangular_pca(); + const auto config = rectangular_config(); + + // Frame 0: z_q - mean = [1, 2, 3, 4]; dot(b0) = 5, dot(b1) = -1; scaled by 0.5. + // Frame 1: z_q - mean = [-1, -2, -3, -4]; dot(b0) = -5, dot(b1) = 1. + const std::vector z_q{2.0F, 4.0F, 6.0F, 8.0F, + 0.0F, 0.0F, 0.0F, 0.0F}; + const std::vector expected{2.5F, -0.5F, + -2.5F, 0.5F}; + + const auto latents = pca_project(pca, config, z_q, 2); + require_eq(static_cast(latents.size()), static_cast(expected.size()), + "projected size"); + for (size_t i = 0; i < expected.size(); ++i) { + require_close(latents[i], expected[i], 1e-6F, "projection element " + std::to_string(i)); + } +} + +void test_pca_inversion_matches_hand_computed_values() { + const auto pca = rectangular_pca(); + const auto config = rectangular_config(); + + // Pinned independently of the forward pass so a mean or scale dropped on + // both legs cannot cancel: + // frame 0: coeffs [2.5, -0.5] / 0.5 = [5, -1] + // mean + 5*b0 - 1*b1 = [1,2,3,4] + [2.5]*4 + [-0.5, 0.5, -0.5, 0.5] + // = [3, 5, 5, 7] + // frame 1: coeffs [-5, 1] -> [1,2,3,4] + [-2.5]*4 + [0.5,-0.5,0.5,-0.5] + // = [-1, -1, 1, 1] + const std::vector latents{2.5F, -0.5F, + -2.5F, 0.5F}; + const std::vector expected{3.0F, 5.0F, 5.0F, 7.0F, + -1.0F, -1.0F, 1.0F, 1.0F}; + + const auto recovered = pca_unproject(pca, config, latents, 2); + require_eq(static_cast(recovered.size()), static_cast(expected.size()), + "unprojected size"); + for (size_t i = 0; i < expected.size(); ++i) { + require_close(recovered[i], expected[i], 1e-6F, "inversion element " + std::to_string(i)); + } +} + +void test_pca_round_trip_recovers_an_in_subspace_vector() { + const auto pca = rectangular_pca(); + const auto config = rectangular_config(); + + // Exactly on span{b0, b1} once the mean is removed, so the lossy projection + // is an identity here and the round trip must be exact -- 1e-6, not 1e-3. + const std::vector z_q{1.0F + 2.0F, 2.0F + 1.0F, 3.0F + 2.0F, 4.0F + 1.0F}; + + const auto latents = pca_project(pca, config, z_q, 1); + const auto recovered = pca_unproject(pca, config, latents, 1); + for (size_t i = 0; i < z_q.size(); ++i) { + require_close(recovered[i], z_q[i], 1e-6F, "round trip element " + std::to_string(i)); + } +} + +void test_pca_rejects_mis_shaped_buffers() { + const auto pca = rectangular_pca(); + const auto config = rectangular_config(); + + bool projected_threw = false; + try { + pca_project(pca, config, std::vector(7, 0.0F), 2); + } catch (const std::exception &) { + projected_threw = true; + } + require(projected_threw, "a mis-shaped z_q buffer is rejected rather than read out of bounds"); + + bool inverted_threw = false; + try { + pca_unproject(pca, config, std::vector(3, 0.0F), 2); + } catch (const std::exception &) { + inverted_threw = true; + } + require(inverted_threw, "a mis-shaped latent buffer is rejected"); + + bool zero_scale_threw = false; + try { + EchoPcaState broken = pca; + broken.latent_scale = 0.0F; + pca_unproject(broken, config, std::vector(2, 0.0F), 1); + } catch (const std::exception &) { + zero_scale_threw = true; + } + require(zero_scale_threw, "a zero latent_scale is rejected rather than dividing by zero"); +} + +// --- flattening point ------------------------------------------------------ + +constexpr int64_t kCropFrames = 60; +constexpr int64_t kCropLatent = 4; + +// Loud alternating +-1 for `active_frames`, then a constant `tail` value. +std::vector loud_then_tail(int64_t active_frames, float tail) { + std::vector out(static_cast(kCropFrames * kCropLatent), tail); + for (int64_t f = 0; f < active_frames; ++f) { + for (int64_t c = 0; c < kCropLatent; ++c) { + out[static_cast(f * kCropLatent + c)] = ((f + c) % 2 == 0) ? 1.0F : -1.0F; + } + } + return out; +} + +int64_t crop(const std::vector & latents) { + return find_flattening_point(latents, kCropFrames, kCropLatent); +} + +void test_flattening_point_matches_reference() { + // Active for 30 frames, then silent. Reference returns 30. + require_eq(crop(loud_then_tail(30, 0.0F)), static_cast(30), + "crop lands where the signal goes flat"); + + // Never flattens: the reference falls through to len(data). + require_eq(crop(loud_then_tail(kCropFrames, 0.0F)), kCropFrames, + "a latent that never flattens keeps every frame"); + + // Flat from the first frame. + require_eq(crop(loud_then_tail(0, 0.0F)), static_cast(0), + "an all-silent latent crops to nothing"); + + // Quiet but NOT zero: std is 0 and |mean - 0| = 0.02 < 0.1, so this must + // still crop at 30. An implementation that looks for an all-zero window + // rather than evaluating both thresholds fails here. + require_eq(crop(loud_then_tail(30, 0.02F)), static_cast(30), + "a quiet non-zero tail still counts as flat"); + + // Flat but too loud: std is 0, yet |mean - 0| = 0.5 exceeds 0.1, so no + // window qualifies and the crop falls through to every frame. This is the + // case that pins the mean threshold rather than just the std threshold. + require_eq(crop(loud_then_tail(30, 0.5F)), kCropFrames, + "a flat but loud tail is not a flattening point"); +} + +} // namespace + +int main() { + try { + test_normalisation_matches_reference(); + test_tokenisation_matches_reference(); + test_tokeniser_truncates_at_max_length(); + test_mask_marks_real_tokens(); + test_pad_to_max_zeroes_the_tail(); + test_pca_projection_matches_hand_computed_values(); + test_pca_inversion_matches_hand_computed_values(); + test_pca_round_trip_recovers_an_in_subspace_vector(); + test_pca_rejects_mis_shaped_buffers(); + test_flattening_point_matches_reference(); + std::cout << "echo_tts_host_units: ok\n"; + return 0; + } catch (const std::exception & ex) { + std::cerr << "echo_tts_host_units: " << ex.what() << "\n"; + return 1; + } +} diff --git a/tests/echo_tts/echo_tts_warm_bench_cases.json b/tests/echo_tts/echo_tts_warm_bench_cases.json new file mode 100644 index 00000000..3db37f5d --- /dev/null +++ b/tests/echo_tts/echo_tts_warm_bench_cases.json @@ -0,0 +1,12 @@ +{ + "default_clone": { + "requests": [ + { + "id": "chris_ref_p1", + "target_voice": "reference/chris_hemsworth_15s.wav", + "text": "[S1] The operations desk reviewed the morning brief and confirmed the relay stayed online through the storm.", + "seed": 0 + } + ] + } +} diff --git a/tools/community_models/convert_echo_tts.py b/tools/community_models/convert_echo_tts.py new file mode 100644 index 00000000..70e3556b --- /dev/null +++ b/tools/community_models/convert_echo_tts.py @@ -0,0 +1,667 @@ +#!/usr/bin/env python3 +"""Convert the published Echo-TTS checkpoint into the GGUF layout audio.cpp loads. + +Inputs (download them yourself; this tool does not fetch anything): + + jordand/echo-tts-base : pytorch_model.safetensors, pca_state.safetensors + +Usage: + + python3 convert_echo_tts.py \ + --model-dir /path/to/echo-tts-base \ + --outfile Echo-TTS-GGUF/model.gguf \ + --precision orig + +Tensors are emitted under three prefixes matching model_specs/echo_tts.json: + + dit_weights/* the EchoDiT, its text encoder and its speaker encoder + pca/* the PCA basis mapping 80-D latents to Fish z_q space + ae/* the Fish S1-DAC autoencoder, weight norm folded + +The Fish S1-DAC autoencoder IS packaged here, under the codec_weights prefix, +when --fish-dir is supplied. audio.cpp already implements this codec for the +fish_audio family and Echo reuses that implementation -- but not its weights: +fish_audio ships Fish Audio S2 Pro, while Echo is trained against the S1 DAC +(jordand/fish-s1-dac-min). Packaging the S1 weights alongside the DiT keeps the +two independent and guarantees Echo gets the exact autoencoder its PCA basis was +fitted to. + +Weight normalisation is folded during conversion. The checkpoint stores it in +two forms -- modern `conv.parametrizations.weight.original0/original1` and legacy +`weight_g`/`weight_v` -- and both reduce to `w = g * v / ||v||`, with the norm +taken over every axis except 0. Note that for ConvTranspose1d axis 0 is the +INPUT channel count, not the output, so g is sized differently there. + +The six registered buffers (three causal masks and three RoPE tables, 305 MB of +the 1.87 GB checkpoint) are regenerable and are dropped. + +Blockwise-continuation weights (latent_encoder, latent_norm, w{k,v}_latent) are +dropped by default, mirroring inference.py's delete_blockwise_modules=True. That +is 420M of the checkpoint's 2.80B parameters. Pass --keep-blockwise to retain +them; nothing in the port consumes them yet. +""" + +from __future__ import annotations + +import argparse +import json +import struct +import sys +from pathlib import Path +from typing import Dict, Iterable, List, Tuple + +import numpy as np + +try: + import gguf +except ImportError: # pragma: no cover - guidance path + gguf = None # reported in main(), so --help still works without it + + +# --- architecture ------------------------------------------------------- +# Not stored in the checkpoint; these are the EchoDiT constructor arguments in +# inference.py::load_model_from_hf. They are re-emitted as GGUF metadata so the +# C++ loader can cross-check rather than hardcode. +ARCH = "echo_tts" + +# GGML_MAX_NAME. gguf.cpp rejects any tensor name of this length or longer, and +# the failure surfaces only at load time as "tensor name N is too long". +GGML_MAX_NAME = 64 + +# The codec prefix is terse because it has to be. The longest name codec.cpp +# loads is 60 characters +# ("encoder.block.4.block.5.layers.3.attention_layer_scale.gamma"), leaving room +# for a three-character prefix and nothing more. "codec_weights." would push +# 157 of the 455 codec tensors past the limit. +# The namespace separator is "/", not ".". PrefixedTensorSourceView in +# src/framework/assets/tensor_source.cpp matches on `prefix + "/"`, so a +# dot-separated name is simply never routed and the namespace looks empty. +NAMESPACE_SEPARATOR = "/" +DIT_TENSOR_PREFIX = "dit_weights" + NAMESPACE_SEPARATOR +PCA_TENSOR_PREFIX = "pca" + NAMESPACE_SEPARATOR +CODEC_TENSOR_PREFIX = "ae" + NAMESPACE_SEPARATOR + +CONFIG: Dict[str, int] = { + "latent_size": 80, + "model_size": 2048, + "num_layers": 24, + "num_heads": 16, + "intermediate_size": 5888, + "text_vocab_size": 256, + "text_model_size": 1280, + "text_num_layers": 14, + "text_num_heads": 10, + "text_intermediate_size": 3328, + "speaker_patch_size": 4, + "speaker_model_size": 1280, + "speaker_num_layers": 14, + "speaker_num_heads": 10, + "speaker_intermediate_size": 3328, + "timestep_embed_size": 512, + "adaln_rank": 256, + "max_sequence_length": 640, + "max_text_length": 768, + "max_speaker_latent_length": 6400, + "ae_downsample_factor": 2048, + "ae_latent_dim": 1024, + "sample_rate": 44100, +} +NORM_EPS = 1.0e-5 + +# Upstream's own filter for the blockwise path, kept verbatim so the two stay +# in sync: inference.py::load_model_from_hf. +# Registered buffers, recomputed at graph build time rather than stored. +CODEC_BUFFER_SUFFIXES = ("causal_mask", "freqs_cis") + +BLOCKWISE_PREFIXES = ("latent_encoder.", "latent_norm") +BLOCKWISE_SUBSTRINGS = (".wk_latent", ".wv_latent") + +# Tensors that must stay F32 regardless of --precision: norm scales are tiny and +# quantising them costs accuracy for no meaningful saving, and biases likewise. +KEEP_F32_SUFFIXES = (".bias", ".alpha") +KEEP_F32_SUBSTRINGS = ( + "_norm.", + "norm.weight", + "q_norm", + "k_norm", + # LayerScale and ConvNeXt scales, and the Snake activation's alpha. These + # are small in count (0.08 MB total) and small in magnitude: LayerScale + # initialises around 1e-6, below F16's smallest normal of 6.1e-05, and the + # Snake activation uses alpha's reciprocal, which amplifies any error. + "gamma", + ".alpha", + # Codebook entries are summed to form z_q, which is exactly the latent + # Echo's PCA basis maps from, so quantising them perturbs the speaker + # conditioning directly. 0.43 MB to keep exact. + "codebook", +) + + +def is_blockwise(name: str) -> bool: + return name.startswith(BLOCKWISE_PREFIXES) or any( + token in name for token in BLOCKWISE_SUBSTRINGS + ) + + +def keep_f32(name: str) -> bool: + return name.endswith(KEEP_F32_SUFFIXES) or any( + token in name for token in KEEP_F32_SUBSTRINGS + ) + + +# --- safetensors reading ------------------------------------------------ + +_DTYPES = { + "F64": np.float64, + "F32": np.float32, + "F16": np.float16, + "I64": np.int64, + "I32": np.int32, + "I16": np.int16, + "I8": np.int8, + "U8": np.uint8, + "BOOL": np.bool_, +} + + +def read_safetensors(path: Path) -> Dict[str, np.ndarray]: + """Minimal safetensors reader with explicit bfloat16 handling. + + numpy has no bfloat16, and the Echo checkpoint is stored in it, so BF16 is + widened to float32 by placing the 16 stored bits in the high half of the + f32 mantissa/exponent. That is exact -- bf16 and f32 share an exponent + layout -- so nothing is lost on the way in. + """ + with path.open("rb") as handle: + (header_length,) = struct.unpack(" payload_bytes: + raise RuntimeError(f"{path.name}: tensor {name} runs past end of file") + shape = tuple(int(dim) for dim in entry["shape"]) + dtype_name = entry["dtype"] + handle.seek(payload_start + start) + raw = handle.read(end - start) + + if dtype_name == "BF16": + bits = np.frombuffer(raw, dtype=np.uint16).astype(np.uint32) << 16 + array = bits.view(np.float32).reshape(shape) + elif dtype_name in _DTYPES: + array = np.frombuffer(raw, dtype=_DTYPES[dtype_name]).reshape(shape) + else: + raise RuntimeError(f"{path.name}: unsupported dtype {dtype_name}") + tensors[name] = np.ascontiguousarray(array) + return tensors + + +# --- Fish S1-DAC codec --------------------------------------------------- + + +def fold_weight_norm(g: np.ndarray, v: np.ndarray) -> np.ndarray: + """Reconstruct a weight-normalised tensor: w = g * v / ||v||. + + torch's weight_norm(dim=0) normalises over every axis except the first, so + the reduction axes are the same regardless of layer type. What differs is + what axis 0 *means*: for Conv1d it is out_channels, for ConvTranspose1d it + is in_channels. Because the reduction is expressed relative to axis 0 rather + than to a named channel count, one implementation covers both -- and the + shapes of g and v carry the distinction for free. + """ + axes = tuple(range(1, v.ndim)) + norm = np.sqrt(np.sum(v.astype(np.float64) ** 2, axis=axes, keepdims=True)) + if not np.all(norm > 0): + raise RuntimeError("weight_norm folding hit a zero-norm direction vector") + return (g.astype(np.float64) * v.astype(np.float64) / norm).astype(np.float32) + + +def split_fused_qkv(name: str, array: np.ndarray) -> Dict[str, np.ndarray]: + """Split a fused wqkv projection into the q/k/v codec.cpp loads separately. + + autoencoder.py::Attention keeps one nn.Linear and splits its output into + three equal kv_size blocks (`wqkv(x).split([kv_size]*3, dim=-1)`), so the + weight rows partition in the same order. codec.cpp instead loads + attention.q_proj / k_proj / v_proj, one square matrix each. + """ + base = name[: -len(".wqkv.weight")] + rows = array.shape[0] + if rows % 3 != 0: + raise RuntimeError(f"{name}: fused qkv has {rows} rows, not divisible by 3") + size = rows // 3 + if array.shape[1] != size: + raise RuntimeError( + f"{name}: expected square projections, got {array.shape} -> {size}") + return { + f"{base}.q_proj.weight": np.ascontiguousarray(array[:size]), + f"{base}.k_proj.weight": np.ascontiguousarray(array[size:2 * size]), + f"{base}.v_proj.weight": np.ascontiguousarray(array[2 * size:]), + } + + +def resolve_codec_tensors(raw: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: + """Fold weight norm, split fused qkv, and drop buffers, leaving the names + codec.cpp expects.""" + out: Dict[str, np.ndarray] = {} + dropped_buffers = 0 + folded = 0 + split = 0 + for name, array in raw.items(): + if name.endswith(CODEC_BUFFER_SUFFIXES): + dropped_buffers += 1 + continue + if name.endswith(".parametrizations.weight.original1"): + base = name[: -len(".parametrizations.weight.original1")] + g = raw[base + ".parametrizations.weight.original0"] + out[base + ".weight"] = fold_weight_norm(g, array) + folded += 1 + continue + if name.endswith(".parametrizations.weight.original0"): + continue + if name.endswith(".weight_v"): + base = name[: -len(".weight_v")] + out[base + ".weight"] = fold_weight_norm(raw[base + ".weight_g"], array) + folded += 1 + continue + if name.endswith(".weight_g"): + continue + if name.endswith(".wqkv.weight"): + out.update(split_fused_qkv(name, array)) + split += 1 + continue + out[name] = array + print(f" folded {folded} weight-normalised tensors, split {split} fused qkv " + f"projections, dropped {dropped_buffers} buffers") + return out + + +# --- expected tensor manifest ------------------------------------------ +# PyTorch state_dict keys follow the nn.Module attribute path, so this manifest +# is derived directly from model.py. It is checked against the checkpoint at +# convert time: a mismatch aborts rather than silently dropping weights. + + +def encoder_block_tensors(prefix: str, dim: int, ff: int, heads: int) -> List[Tuple[str, Tuple[int, ...]]]: + head_dim = dim // heads + out: List[Tuple[str, Tuple[int, ...]]] = [] + for name in ("wq", "wk", "wv", "wo", "gate"): + out.append((f"{prefix}.attention.{name}.weight", (dim, dim))) + out.append((f"{prefix}.attention.q_norm.weight", (heads, head_dim))) + out.append((f"{prefix}.attention.k_norm.weight", (heads, head_dim))) + out.append((f"{prefix}.mlp.w1.weight", (ff, dim))) + out.append((f"{prefix}.mlp.w3.weight", (ff, dim))) + out.append((f"{prefix}.mlp.w2.weight", (dim, ff))) + out.append((f"{prefix}.attention_norm.weight", (dim,))) + out.append((f"{prefix}.mlp_norm.weight", (dim,))) + return out + + +def expected_tensors(keep_blockwise: bool) -> Dict[str, Tuple[int, ...]]: + c = CONFIG + D, TD, SD = c["model_size"], c["text_model_size"], c["speaker_model_size"] + L, RANK = c["latent_size"], c["adaln_rank"] + expected: Dict[str, Tuple[int, ...]] = {} + + expected["text_encoder.text_embedding.weight"] = (c["text_vocab_size"], TD) + for i in range(c["text_num_layers"]): + for name, shape in encoder_block_tensors( + f"text_encoder.blocks.{i}", TD, c["text_intermediate_size"], c["text_num_heads"] + ): + expected[name] = shape + + speaker_stacks = ["speaker_encoder"] + (["latent_encoder"] if keep_blockwise else []) + for stack in speaker_stacks: + expected[f"{stack}.in_proj.weight"] = (SD, L * c["speaker_patch_size"]) + expected[f"{stack}.in_proj.bias"] = (SD,) + for i in range(c["speaker_num_layers"]): + for name, shape in encoder_block_tensors( + f"{stack}.blocks.{i}", SD, c["speaker_intermediate_size"], c["speaker_num_heads"] + ): + expected[name] = shape + + expected["text_norm.weight"] = (TD,) + expected["speaker_norm.weight"] = (SD,) + if keep_blockwise: + expected["latent_norm.weight"] = (SD,) + + # nn.Sequential(Linear, SiLU, Linear, SiLU, Linear) -> indices 0, 2, 4. + expected["cond_module.0.weight"] = (D, c["timestep_embed_size"]) + expected["cond_module.2.weight"] = (D, D) + expected["cond_module.4.weight"] = (D * 3, D) + + expected["in_proj.weight"] = (D, L) + expected["in_proj.bias"] = (D,) + + head_dim = D // c["num_heads"] + for i in range(c["num_layers"]): + p = f"blocks.{i}.attention" + for name in ("wq", "wk", "wv", "gate", "wo"): + expected[f"{p}.{name}.weight"] = (D, D) + expected[f"{p}.wk_text.weight"] = (D, TD) + expected[f"{p}.wv_text.weight"] = (D, TD) + expected[f"{p}.wk_speaker.weight"] = (D, SD) + expected[f"{p}.wv_speaker.weight"] = (D, SD) + if keep_blockwise: + expected[f"{p}.wk_latent.weight"] = (D, SD) + expected[f"{p}.wv_latent.weight"] = (D, SD) + expected[f"{p}.q_norm.weight"] = (c["num_heads"], head_dim) + expected[f"{p}.k_norm.weight"] = (c["num_heads"], head_dim) + + expected[f"blocks.{i}.mlp.w1.weight"] = (c["intermediate_size"], D) + expected[f"blocks.{i}.mlp.w3.weight"] = (c["intermediate_size"], D) + expected[f"blocks.{i}.mlp.w2.weight"] = (D, c["intermediate_size"]) + + for adaln in ("attention_adaln", "mlp_adaln"): + a = f"blocks.{i}.{adaln}" + for field in ("shift", "scale", "gate"): + expected[f"{a}.{field}_down.weight"] = (RANK, D) + expected[f"{a}.{field}_up.weight"] = (D, RANK) + expected[f"{a}.{field}_up.bias"] = (D,) + + expected["out_norm.weight"] = (D,) + expected["out_proj.weight"] = (L, D) + expected["out_proj.bias"] = (L,) + return expected + + +def verify_manifest( + found: Dict[str, np.ndarray], keep_blockwise: bool, strict: bool +) -> None: + expected = expected_tensors(keep_blockwise) + kept = {k: v for k, v in found.items() if keep_blockwise or not is_blockwise(k)} + + missing = sorted(set(expected) - set(kept)) + unexpected = sorted(set(kept) - set(expected)) + mismatched = [ + (name, expected[name], kept[name].shape) + for name in sorted(set(expected) & set(kept)) + if tuple(kept[name].shape) != expected[name] + ] + + for name in missing: + print(f" MISSING {name} {expected[name]}", file=sys.stderr) + for name in unexpected: + print(f" UNEXPECTED {name} {tuple(kept[name].shape)}", file=sys.stderr) + for name, want, got in mismatched: + print(f" SHAPE {name}: expected {want}, got {got}", file=sys.stderr) + + if missing or mismatched or (unexpected and strict): + raise RuntimeError( + "checkpoint does not match the expected Echo-TTS manifest; " + "the architecture may have changed upstream" + ) + + +# --- conversion --------------------------------------------------------- + + +# Q8_0 stores 32 weights per block with one shared F16 scale, so a tensor is +# only quantisable when its fastest-varying axis is a multiple of 32. In GGUF +# that axis is the LAST logical dimension. +Q8_0_BLOCK = 32 + + +def q8_0_eligible(name: str, array: np.ndarray) -> bool: + if keep_f32(name) or array.ndim < 2: + return False + # 3-D tensors here are convolution kernels. ggml_conv_1d has no quantised + # path, which is why codec.cpp takes matmul and conv storage types + # separately; quantising them would fail at graph build, not at load. + if array.ndim > 2: + return False + return array.shape[-1] % Q8_0_BLOCK == 0 + + +def resolve_dtype(name: str, array: np.ndarray, precision: str): + if precision == "f32" or keep_f32(name) or array.ndim < 2: + return gguf.GGMLQuantizationType.F32, array.astype(np.float32) + if precision in ("orig", "f16"): + return gguf.GGMLQuantizationType.F16, array.astype(np.float16) + if precision == "q8_0": + if not q8_0_eligible(name, array): + return gguf.GGMLQuantizationType.F16, array.astype(np.float16) + quantised = gguf.quants.quantize( + np.ascontiguousarray(array, dtype=np.float32), + gguf.GGMLQuantizationType.Q8_0, + ) + return gguf.GGMLQuantizationType.Q8_0, quantised + raise RuntimeError(f"unknown precision {precision}") + + +def load_model_spec_json(explicit: str | None) -> str: + """Read the spec that will be embedded, and sanity-check it.""" + if explicit: + path = Path(explicit) + else: + # tools/community_models/convert_echo_tts.py -> model_specs/echo_tts.json + path = Path(__file__).resolve().parents[2] / "model_specs" / f"{ARCH}.json" + if not path.is_file(): + raise RuntimeError( + f"model spec not found at {path}; pass --model-spec explicitly") + text = path.read_text(encoding="utf-8") + spec = json.loads(text) + if spec.get("family") != ARCH: + raise RuntimeError( + f"{path} declares family {spec.get('family')!r}, expected {ARCH!r}") + if spec.get("schema_version") != 1: + raise RuntimeError(f"{path} is not a schema_version 1 spec") + return text + + +def convert(args: argparse.Namespace) -> int: + model_dir = Path(args.model_dir) + model_path = model_dir / "pytorch_model.safetensors" + pca_path = model_dir / "pca_state.safetensors" + for path in (model_path, pca_path): + if not path.is_file(): + print(f"missing required input: {path}", file=sys.stderr) + return 2 + + print(f"reading {model_path}") + weights = read_safetensors(model_path) + print(f" {len(weights)} tensors") + + print(f"reading {pca_path}") + pca = read_safetensors(pca_path) + + for required in ("pca_components", "pca_mean", "latent_scale"): + if required not in pca: + print(f"pca_state is missing {required}", file=sys.stderr) + return 2 + + components = pca["pca_components"].astype(np.float32) + mean = pca["pca_mean"].astype(np.float32) + scale = float(np.asarray(pca["latent_scale"]).reshape(-1)[0]) + + want = (CONFIG["latent_size"], CONFIG["ae_latent_dim"]) + if components.shape != want: + # ae_encode does `... @ pca_components.T` into 80-D, so the basis must be + # (80, 1024). Accept the transpose but say so loudly. + if components.shape == want[::-1]: + print( + f" note: pca_components stored as {components.shape}, transposing to {want}", + file=sys.stderr, + ) + components = np.ascontiguousarray(components.T) + else: + print( + f"pca_components has shape {components.shape}, expected {want}", + file=sys.stderr, + ) + return 2 + if mean.shape != (CONFIG["ae_latent_dim"],): + print(f"pca_mean has shape {mean.shape}, expected {(CONFIG['ae_latent_dim'],)}", + file=sys.stderr) + return 2 + + codec_tensors: Dict[str, np.ndarray] = {} + if not args.no_codec: + if not args.fish_dir: + print( + "--fish-dir is required (or pass --no-codec).\n" + "Echo decodes through the Fish S1 DAC and its PCA basis is fitted to " + "that codec's latent space; audio.cpp's fish_audio package ships S2 Pro, " + "which is a different model.", + file=sys.stderr, + ) + return 2 + fish_path = Path(args.fish_dir) / "pytorch_model.safetensors" + if not fish_path.is_file(): + print(f"missing required input: {fish_path}", file=sys.stderr) + return 2 + print(f"reading {fish_path}") + codec_tensors = resolve_codec_tensors(read_safetensors(fish_path)) + print(f" {len(codec_tensors)} codec tensors") + + try: + spec_json = load_model_spec_json(args.model_spec) + except RuntimeError as error: + print(str(error), file=sys.stderr) + return 2 + print(f"embedding model spec ({len(spec_json)} bytes)") + + print("verifying tensor manifest against model.py") + verify_manifest(weights, args.keep_blockwise, strict=not args.allow_extra) + print(" manifest OK") + + outfile = Path(args.outfile) + outfile.parent.mkdir(parents=True, exist_ok=True) + writer = gguf.GGUFWriter(str(outfile), ARCH) + + over_limit: List[str] = [] + # ggml_n_dims() ignores trailing dimensions of size 1, so a (1, C, 1) snake + # alpha reads back as (C, 1) and fails codec.cpp's {1, C, 1} shape check. + # audio.cpp preserves exact logical shapes through two parallel arrays, in + # tensor order: audiocpp.tensor_ranks (INT32) and the concatenated + # audiocpp.tensor_shapes (INT64). Both must be present or neither. + tensor_ranks: List[int] = [] + tensor_shapes: List[int] = [] + + def emit(name: str, data, dtype, logical_shape) -> None: + # Checked here rather than trusted, because ggml only reports this at + # load time and the message does not say which tensor is at fault. + if len(name) >= GGML_MAX_NAME: + over_limit.append(name) + # audiocpp.tensor_shapes records the LOGICAL shape; a quantised payload + # arrives packed, so it has to come from the source array. gguf's own + # raw_shape, by contrast, wants the packed byte shape and derives the + # logical one itself, so it is left to default. + tensor_ranks.append(len(logical_shape)) + tensor_shapes.extend(int(dim) for dim in logical_shape) + writer.add_tensor(name, data, raw_dtype=dtype) + + for key, value in CONFIG.items(): + writer.add_uint32(f"{ARCH}.{key}", int(value)) + writer.add_float32(f"{ARCH}.norm_eps", NORM_EPS) + writer.add_float32(f"{ARCH}.pca_latent_scale", scale) + writer.add_bool(f"{ARCH}.has_blockwise_modules", bool(args.keep_blockwise)) + writer.add_bool(f"{ARCH}.has_codec_weights", not args.no_codec) + writer.add_string(f"{ARCH}.source_precision", args.precision) + writer.add_string("general.license", "cc-by-nc-sa-4.0") + writer.add_string("general.name", "Echo-TTS") + + # A published GGUF must carry its own model spec: package.cpp refuses to load + # one that does not, so that a distributed file is self-describing and does + # not depend on a matching model_specs/ checkout. + writer.add_uint32("audiocpp.model_spec.version", 1) + writer.add_string("audiocpp.model_spec.family", ARCH) + writer.add_string("audiocpp.model_spec.json", spec_json) + + dropped = 0 + written = 0 + total_bytes = 0 + for name in sorted(weights): + if not args.keep_blockwise and is_blockwise(name): + dropped += 1 + continue + array = weights[name] + dtype, data = resolve_dtype(name, array, args.precision) + emit(f"{DIT_TENSOR_PREFIX}{name}", data, dtype, array.shape) + written += 1 + total_bytes += data.nbytes + + emit(f"{PCA_TENSOR_PREFIX}components", components, gguf.GGMLQuantizationType.F32, components.shape) + emit(f"{PCA_TENSOR_PREFIX}mean", mean, gguf.GGMLQuantizationType.F32, mean.shape) + + codec_bytes = 0 + for name in sorted(codec_tensors): + array = codec_tensors[name] + dtype, data = resolve_dtype(name, array, args.precision) + emit(f"{CODEC_TENSOR_PREFIX}{name}", data, dtype, array.shape) + codec_bytes += data.nbytes + + writer.add_key_value( + "audiocpp.tensor_ranks", + tensor_ranks, + gguf.GGUFValueType.ARRAY, + sub_type=gguf.GGUFValueType.INT32, + ) + writer.add_key_value( + "audiocpp.tensor_shapes", + tensor_shapes, + gguf.GGUFValueType.ARRAY, + sub_type=gguf.GGUFValueType.INT64, + ) + + if over_limit: + longest = max(over_limit, key=len) + print( + f"{len(over_limit)} tensor names reach or exceed GGML_MAX_NAME " + f"({GGML_MAX_NAME}); longest is {len(longest)} chars:\n {longest}\n" + "Shorten a tensor prefix; the GGUF would fail to load.", + file=sys.stderr, + ) + writer.close() + return 2 + + print(f"writing {outfile}") + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + print(f" {written} DiT tensors written, {dropped} blockwise tensors dropped") + if codec_tensors: + print(f" {len(codec_tensors)} codec tensors written ({codec_bytes / 1e9:.2f} GB)") + total_bytes += codec_bytes + print(f" approx tensor payload: {total_bytes / 1e9:.2f} GB") + print(f" pca latent_scale: {scale}") + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", required=True, + help="directory holding pytorch_model.safetensors and pca_state.safetensors") + parser.add_argument("--outfile", required=True) + parser.add_argument("--precision", default="orig", choices=("orig", "f16", "f32", "q8_0"), + help="orig and f16 both emit F16 matmul weights; the checkpoint " + "ships bf16, which has no GGUF matmul equivalent") + parser.add_argument("--fish-dir", + help="directory holding the Fish S1-DAC checkpoint " + "(jordand/fish-s1-dac-min/pytorch_model.safetensors). " + "Required unless --no-codec is passed.") + parser.add_argument("--no-codec", action="store_true", + help="omit the autoencoder; the resulting GGUF cannot synthesise") + parser.add_argument("--keep-blockwise", action="store_true", + help="retain latent_encoder / w{k,v}_latent (+840 MB, unused today)") + parser.add_argument("--model-spec", + help="path to model_specs/echo_tts.json to embed. Defaults to the " + "copy alongside this script's checkout.") + parser.add_argument("--allow-extra", action="store_true", + help="tolerate checkpoint tensors absent from the expected manifest") + args = parser.parse_args() + if gguf is None: + print("the 'gguf' package is required (pip install gguf)", file=sys.stderr) + return 2 + return convert(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/community_models/echo_tts_manifest.py b/tools/community_models/echo_tts_manifest.py new file mode 100644 index 00000000..95e8bb79 --- /dev/null +++ b/tools/community_models/echo_tts_manifest.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Dump the tensor manifest (names, shapes, dtypes) of the Echo-TTS checkpoints. + +Safetensors stores a JSON header at the front of the file: an 8-byte +little-endian length, then that many bytes of JSON. So the full tensor listing +can be read with a couple of HTTP range requests -- no weights are downloaded. + +Usage: + + # remote, no download (default: all three Echo-TTS checkpoints) + python3 echo_tts_manifest.py -o echo_manifest.json + + # a checkpoint you already have on disk + python3 echo_tts_manifest.py --local /path/to/pytorch_model.safetensors -o out.json + +Set HF_TOKEN in the environment if a repo needs auth. Requires only the standard +library. +""" + +from __future__ import annotations + +import argparse +import json +import os +import struct +import sys +import urllib.error +import urllib.request +from typing import Any, Dict, List, Tuple + +DEFAULT_TARGETS: List[Tuple[str, str, str]] = [ + ("echo_dit", "jordand/echo-tts-base", "pytorch_model.safetensors"), + ("echo_pca", "jordand/echo-tts-base", "pca_state.safetensors"), + ("fish_ae", "jordand/fish-s1-dac-min", "pytorch_model.safetensors"), +] + +# A safetensors header is JSON; 64 MiB is far more than any real one needs and +# still bounds a malformed-length read. +MAX_HEADER_BYTES = 64 * 1024 * 1024 + + +def _request(url: str, byte_range: Tuple[int, int]) -> bytes: + start, end = byte_range + headers = { + "Range": f"bytes={start}-{end}", + "User-Agent": "echo-tts-manifest/1.0", + } + token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + if token: + headers["Authorization"] = f"Bearer {token}" + + request = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(request, timeout=60) as response: + status = response.status + data = response.read(end - start + 1) + if status != 206: + raise RuntimeError( + f"server ignored the range request (HTTP {status}); refusing to " + f"download the whole file" + ) + return data + + +def read_remote_header(repo_id: str, filename: str, revision: str) -> Dict[str, Any]: + url = f"https://huggingface.co/{repo_id}/resolve/{revision}/{filename}" + prefix = _request(url, (0, 7)) + if len(prefix) != 8: + raise RuntimeError(f"short read on the length prefix of {filename}") + (header_length,) = struct.unpack(" Dict[str, Any]: + with open(path, "rb") as handle: + prefix = handle.read(8) + if len(prefix) != 8: + raise RuntimeError(f"short read on the length prefix of {path}") + (header_length,) = struct.unpack(" Dict[str, Any]: + """Reduce a safetensors header to name -> {dtype, shape, nbytes}.""" + tensors: Dict[str, Any] = {} + total_bytes = 0 + total_params = 0 + for name, entry in header.items(): + if name == "__metadata__": + continue + offsets = entry.get("data_offsets", [0, 0]) + nbytes = int(offsets[1]) - int(offsets[0]) + shape = [int(dim) for dim in entry.get("shape", [])] + count = 1 + for dim in shape: + count *= dim + tensors[name] = { + "dtype": entry.get("dtype"), + "shape": shape, + "nbytes": nbytes, + } + total_bytes += nbytes + total_params += count + return { + "metadata": header.get("__metadata__", {}), + "tensor_count": len(tensors), + "total_parameters": total_params, + "total_bytes": total_bytes, + "tensors": tensors, + } + + +def group_key(name: str) -> str: + """Collapse numeric path segments so repeated blocks fold into one entry.""" + parts = [] + for part in name.split("."): + parts.append("{N}" if part.isdigit() else part) + return ".".join(parts) + + +def summarize(label: str, manifest: Dict[str, Any]) -> str: + groups: Dict[str, Dict[str, Any]] = {} + for name, info in manifest["tensors"].items(): + key = group_key(name) + bucket = groups.setdefault( + key, {"count": 0, "shape": info["shape"], "dtype": info["dtype"]} + ) + bucket["count"] += 1 + + lines = [ + f"=== {label} ===", + f" tensors={manifest['tensor_count']} " + f"params={manifest['total_parameters']:,} " + f"bytes={manifest['total_bytes']:,}", + ] + if manifest["metadata"]: + lines.append(f" metadata={manifest['metadata']}") + for key in sorted(groups): + bucket = groups[key] + suffix = f" x{bucket['count']}" if bucket["count"] > 1 else "" + lines.append(f" {key:<62} {bucket['dtype']:<8} {bucket['shape']}{suffix}") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--local", + action="append", + default=[], + metavar="PATH", + help="read a local .safetensors file instead of querying the Hub " + "(repeatable)", + ) + parser.add_argument( + "--revision", default="main", help="git revision to resolve (default: main)" + ) + parser.add_argument( + "-o", + "--output", + default="echo_manifest.json", + help="where to write the full JSON manifest", + ) + args = parser.parse_args() + + results: Dict[str, Any] = {} + failures: List[str] = [] + + if args.local: + for path in args.local: + label = os.path.basename(path) + try: + results[label] = normalize(read_local_header(path)) + except Exception as error: # noqa: BLE001 - report and continue + failures.append(f"{label}: {error}") + else: + for label, repo_id, filename in DEFAULT_TARGETS: + try: + header = read_remote_header(repo_id, filename, args.revision) + manifest = normalize(header) + manifest["source"] = f"{repo_id}/{filename}@{args.revision}" + results[label] = manifest + except Exception as error: # noqa: BLE001 - report and continue + failures.append(f"{label} ({repo_id}/{filename}): {error}") + + for label in sorted(results): + print(summarize(label, results[label])) + print() + + for failure in failures: + print(f"FAILED {failure}", file=sys.stderr) + + if results: + with open(args.output, "w", encoding="utf-8") as handle: + json.dump(results, handle, indent=1, sort_keys=True) + print(f"wrote {args.output}") + + return 1 if failures and not results else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/community_models/echo_tts_pack_reference.py b/tools/community_models/echo_tts_pack_reference.py new file mode 100644 index 00000000..a4c7eb3a --- /dev/null +++ b/tools/community_models/echo_tts_pack_reference.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Pack an echo_ref.npz reference dump into a flat binary the C++ parity +harness can read without an npz parser. + + python3 echo_tts_pack_reference.py echo_ref.npz -o echo_ref.bin + +Only the tensors the harness actually consumes are packed; the per-block +activations are 24 x 640 x 2048 and are included only with --blocks, which +takes the archive from a few MB to a few hundred. + +Format, all little-endian, which is the only byte order audio.cpp targets: + + magic 8 bytes "ECHOPAR1" + count int32 number of entries + entry int32 name length + bytes name, not NUL-terminated + int32 dtype, 0 = float32, 1 = int32 + int64 element count + data element count * 4 bytes + +Entries appear in the order written here; the reader looks them up by name, so +order is not load-bearing. +""" + +from __future__ import annotations + +import argparse +import struct +import sys + +import numpy as np + +MAGIC = b"ECHOPAR1" +DTYPE_F32 = 0 +DTYPE_I32 = 1 + +# The minimum needed to drive the DiT at a fixed timestep and score the result. +REQUIRED = [ + "dit.x_input", + "dit.v_pred", + "dit.t", + "text.input_ids", + "text.mask", + "speaker.latent", + "speaker.mask", + "sampler.latent", + "sampler.initial_noise", + "config.sequence_length", + "config.steps", + "config.seed", +] + + +def pack_entry(name: str, array: np.ndarray) -> bytes: + flat = np.ascontiguousarray(array).reshape(-1) + if flat.dtype in (np.int32, np.int64): + dtype, payload = DTYPE_I32, flat.astype(" int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("npz", help="echo_ref.npz from echo_tts_reference.py") + parser.add_argument("-o", "--output", default="echo_ref.bin") + parser.add_argument( + "--blocks", + action="store_true", + help="also pack the 24 per-block DiT activations (large)", + ) + args = parser.parse_args() + + data = np.load(args.npz) + names = list(REQUIRED) + if args.blocks: + names += [f"dit.block.{i}" for i in range(24) if f"dit.block.{i}" in data.files] + + missing = [n for n in names if n not in data.files] + if missing: + print(f"missing from {args.npz}: {', '.join(missing)}", file=sys.stderr) + return 1 + + chunks = [pack_entry(n, data[n]) for n in names] + with open(args.output, "wb") as handle: + handle.write(MAGIC) + handle.write(struct.pack(".wav -o echo_ref.npz + +Everything is pinned to a fixed seed and a fixed text so the C++ port can be +compared stage by stage. Small tensors are dumped in full; the 24 DiT block +activations are dumped as statistics plus a value prefix unless --full-blocks is +passed, which keeps the archive to a few MB rather than a few hundred. + +The DiT forward pass is captured at a single fixed timestep with a fixed input +latent, deliberately *not* the sampler's own trajectory: that isolates a wrong +block from a wrong integration step. The sampler is then run separately. +""" + +from __future__ import annotations + +import argparse +import sys +from typing import Any, Dict + +import numpy as np +import torch + +try: + from inference import ( + get_speaker_latent_and_mask, + get_text_input_ids_and_mask, + ae_decode, + ae_encode, + load_audio, + load_fish_ae_from_hf, + load_model_from_hf, + load_pca_state_from_hf, + sample_euler_cfg_independent_guidances, + tokenizer_encode, + ) +except ImportError as error: # pragma: no cover - guidance path + print( + f"could not import the upstream inference module ({error}).\n" + "Run this script from inside a checkout of jordandare/echo-tts.", + file=sys.stderr, + ) + raise SystemExit(2) + +DEFAULT_TEXT = ( + "[S1] Alright, I'm going to demo this new model called Echo TTS. " + "Hopefully this works, I'm super excited to try this and see what it can do." +) + +SEED = 0 +FIXED_T = 0.7 +SEQUENCE_LENGTH = 640 + + +def to_numpy(tensor: torch.Tensor) -> np.ndarray: + return tensor.detach().float().cpu().numpy() + + +def summarize(name: str, tensor: torch.Tensor, out: Dict[str, Any], prefix: int = 64) -> None: + """Store shape, moments, and a value prefix -- enough to localise drift.""" + array = to_numpy(tensor) + flat = array.reshape(-1) + out[f"{name}.shape"] = np.array(array.shape, dtype=np.int64) + out[f"{name}.stats"] = np.array( + [flat.mean(), flat.std(), flat.min(), flat.max()], dtype=np.float64 + ) + out[f"{name}.prefix"] = flat[:prefix].astype(np.float32) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--speaker", required=True, help="reference wav for cloning") + parser.add_argument("--text", default=DEFAULT_TEXT) + parser.add_argument("-o", "--output", default="echo_ref.npz") + parser.add_argument( + "--full-blocks", + action="store_true", + help="dump every DiT block activation in full (large)", + ) + parser.add_argument( + "--steps", type=int, default=40, help="sampler steps (default: 40)" + ) + args = parser.parse_args() + + torch.manual_seed(SEED) + out: Dict[str, Any] = {} + + model = load_model_from_hf(delete_blockwise_modules=True) + fish_ae = load_fish_ae_from_hf() + pca_state = load_pca_state_from_hf() + device, dtype = model.device, model.dtype + + out["pca.components"] = to_numpy(pca_state.pca_components) + out["pca.mean"] = to_numpy(pca_state.pca_mean) + out["pca.latent_scale"] = np.array([pca_state.latent_scale], dtype=np.float64) + + # ---- tokenizer ------------------------------------------------------- + ids, normalized = tokenizer_encode(args.text, return_normalized_text=True) + out["tokenizer.input_ids"] = to_numpy(ids).astype(np.int32) + out["tokenizer.normalized_text"] = np.array([normalized.encode("utf-8")]) + + text_input_ids, text_mask = get_text_input_ids_and_mask( + [args.text], max_length=None, device=device + ) + out["text.input_ids"] = to_numpy(text_input_ids).astype(np.int32) + out["text.mask"] = to_numpy(text_mask).astype(np.int32) + + # ---- autoencoder round trip ----------------------------------------- + speaker_audio = load_audio(args.speaker).to(device) + out["audio.speaker_input"] = to_numpy(speaker_audio) + + z_q = fish_ae.encode_zq(speaker_audio.unsqueeze(0).to(fish_ae.dtype)) + summarize("ae.encode_zq", z_q, out) + + latent_from_audio = ae_encode(fish_ae, pca_state, speaker_audio.unsqueeze(0).to(fish_ae.dtype)) + summarize("ae.encode_pca", latent_from_audio, out) + + reconstructed = ae_decode(fish_ae, pca_state, latent_from_audio) + summarize("ae.decode_roundtrip", reconstructed, out) + + speaker_latent, speaker_mask = get_speaker_latent_and_mask( + fish_ae, pca_state, speaker_audio.to(fish_ae.dtype) + ) + out["speaker.latent"] = to_numpy(speaker_latent) + out["speaker.mask"] = to_numpy(speaker_mask).astype(np.int32) + + # ---- conditioning encoders ------------------------------------------ + with torch.inference_mode(): + text_state = model.text_encoder(text_input_ids, text_mask) + text_state = model.text_norm(text_state) + summarize("text_encoder.output", text_state, out) + + speaker_state = model.speaker_encoder(speaker_latent.to(dtype)) + speaker_state = model.speaker_norm(speaker_state) + summarize("speaker_encoder.output", speaker_state, out) + + kv_text = model.get_kv_cache_text(text_input_ids, text_mask) + kv_speaker = model.get_kv_cache_speaker(speaker_latent.to(dtype)) + for layer in (0, len(kv_text) // 2, len(kv_text) - 1): + summarize(f"kv_text.{layer}.k", kv_text[layer][0], out) + summarize(f"kv_text.{layer}.v", kv_text[layer][1], out) + summarize(f"kv_speaker.{layer}.k", kv_speaker[layer][0], out) + summarize(f"kv_speaker.{layer}.v", kv_speaker[layer][1], out) + + # ---- single fixed-timestep DiT forward, with per-block hooks ----- + block_outputs: Dict[int, torch.Tensor] = {} + + def make_hook(index: int): + def hook(_module, _inputs, output): + block_outputs[index] = output.detach() + + return hook + + handles = [ + block.register_forward_hook(make_hook(index)) + for index, block in enumerate(model.blocks) + ] + + generator = torch.Generator(device=device).manual_seed(SEED) + x_fixed = torch.randn( + (1, SEQUENCE_LENGTH, 80), + device=device, + dtype=torch.float32, + generator=generator, + ) + out["dit.x_input"] = to_numpy(x_fixed) + out["dit.t"] = np.array([FIXED_T], dtype=np.float64) + + t_tensor = (torch.ones((1,), device=device) * FIXED_T).to(dtype) + v_pred = model( + x=x_fixed.to(dtype), + t=t_tensor, + text_mask=text_mask, + speaker_mask=speaker_mask, + kv_cache_text=kv_text, + kv_cache_speaker=kv_speaker, + ) + for handle in handles: + handle.remove() + + out["dit.v_pred"] = to_numpy(v_pred) + for index, value in sorted(block_outputs.items()): + if args.full_blocks: + out[f"dit.block.{index}"] = to_numpy(value).astype(np.float16) + summarize(f"dit.block.{index}", value, out) + + # ---- full sampler + decode -------------------------------------- + latent_out = sample_euler_cfg_independent_guidances( + model=model, + speaker_latent=speaker_latent, + speaker_mask=speaker_mask, + text_input_ids=text_input_ids, + text_mask=text_mask, + rng_seed=SEED, + num_steps=args.steps, + cfg_scale_text=3.0, + cfg_scale_speaker=8.0, + cfg_min_t=0.5, + cfg_max_t=1.0, + truncation_factor=0.8, + rescale_k=None, + rescale_sigma=None, + speaker_kv_scale=None, + speaker_kv_max_layers=None, + speaker_kv_min_t=None, + sequence_length=SEQUENCE_LENGTH, + ) + out["sampler.latent"] = to_numpy(latent_out) + + # The initial noise, reproduced exactly as the sampler draws it, so the + # C++ Philox path can be checked independently of the model. + noise_generator = torch.Generator(device=device).manual_seed(SEED) + noise = torch.randn( + (1, SEQUENCE_LENGTH, 80), + device=device, + dtype=torch.float32, + generator=noise_generator, + ) + out["sampler.initial_noise"] = to_numpy(noise) + + audio_out = ae_decode(fish_ae, pca_state, latent_out) + summarize("decode.audio_full", audio_out, out) + out["decode.audio_prefix"] = to_numpy(audio_out).reshape(-1)[:44100] + + out["config.seed"] = np.array([SEED], dtype=np.int64) + out["config.steps"] = np.array([args.steps], dtype=np.int64) + out["config.sequence_length"] = np.array([SEQUENCE_LENGTH], dtype=np.int64) + out["config.model_dtype"] = np.array([str(dtype).encode("utf-8")]) + + np.savez_compressed(args.output, **out) + print(f"wrote {args.output} ({len(out)} arrays)") + for key in sorted(out): + if key.endswith(".stats"): + mean, std, low, high = out[key] + print(f" {key[:-6]:<40} mean={mean:+.6f} std={std:.6f} " + f"min={low:+.4f} max={high:+.4f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/community_models/verify_echo_gguf.py b/tools/community_models/verify_echo_gguf.py new file mode 100644 index 00000000..4e815c66 --- /dev/null +++ b/tools/community_models/verify_echo_gguf.py @@ -0,0 +1,350 @@ +#!/usr/bin/env python3 +"""Verify a converted Echo-TTS GGUF before trying to load it in audio.cpp. + +Checks the artifact rather than the conversion process: tensor coverage against +what the C++ loaders ask for, shapes against the architecture, and dtypes +against the precision policy. Catches a bad convert in a second instead of +after a model load and a wasted GPU run. + + python3 verify_echo_gguf.py /path/to/echo-tts.gguf + +Exits non-zero if anything is wrong. Requires the `gguf` package. +""" + +from __future__ import annotations + +import argparse +import sys +from typing import Dict, List, Set, Tuple + +try: + import gguf +except ImportError: # pragma: no cover - guidance path + gguf = None + +DIT_PREFIX = "dit_weights/" +PCA_PREFIX = "pca/" +CODEC_PREFIX = "ae/" + +# gguf.cpp rejects names at or beyond GGML_MAX_NAME, and only at load time. +GGML_MAX_NAME = 64 + +# Groups that must stay F32 regardless of --precision. Snake's alpha is used via +# its reciprocal, LayerScale/ConvNeXt gammas initialise below F16's smallest +# normal (6.1e-05), and codebook entries sum to form the z_q that the PCA basis +# maps from. +F32_REQUIRED_MARKERS = ("gamma", ".alpha", "codebook", "norm.weight", "_norm.") + + +def codec_probe_names() -> List[str]: + """Names src/models/fish_audio/codec.cpp loads, transcribed from its paths.""" + want: List[str] = ["encoder.block.0.conv.weight", "encoder.block.0.conv.bias"] + for b in range(1, 5): + for r in range(3): + want += [ + f"encoder.block.{b}.block.{r}.block.0.alpha", + f"encoder.block.{b}.block.{r}.block.1.conv.weight", + f"encoder.block.{b}.block.{r}.block.2.alpha", + f"encoder.block.{b}.block.{r}.block.3.conv.weight", + ] + want += [ + f"encoder.block.{b}.block.3.alpha", + f"encoder.block.{b}.block.4.conv.weight", + ] + for layer in range(4): + p = f"encoder.block.4.block.5.layers.{layer}" + want += [ + f"{p}.attention.q_proj.weight", + f"{p}.attention.k_proj.weight", + f"{p}.attention.v_proj.weight", + f"{p}.attention.wo.weight", + f"{p}.attention_norm.weight", + f"{p}.ffn_norm.weight", + f"{p}.feed_forward.w1.weight", + f"{p}.feed_forward.w2.weight", + f"{p}.feed_forward.w3.weight", + f"{p}.attention_layer_scale.gamma", + f"{p}.ffn_layer_scale.gamma", + ] + want += [ + "encoder.block.4.block.5.norm.weight", + "encoder.block.5.alpha", + "encoder.block.6.conv.weight", + ] + for stage in ("pre_module", "post_module"): + for layer in range(8): + want += [ + f"quantizer.{stage}.layers.{layer}.attention.q_proj.weight", + f"quantizer.{stage}.layers.{layer}.attention.k_proj.weight", + f"quantizer.{stage}.layers.{layer}.attention.v_proj.weight", + f"quantizer.{stage}.layers.{layer}.attention.wo.weight", + ] + want += [f"quantizer.{stage}.norm.weight"] + want += [ + "quantizer.semantic_quantizer.quantizers.0.codebook.weight", + "quantizer.semantic_quantizer.quantizers.0.in_proj.weight", + "quantizer.semantic_quantizer.quantizers.0.out_proj.weight", + ] + for q in range(9): + want += [ + f"quantizer.quantizer.quantizers.{q}.codebook.weight", + f"quantizer.quantizer.quantizers.{q}.in_proj.weight", + f"quantizer.quantizer.quantizers.{q}.out_proj.weight", + ] + for stage in ("downsample", "upsample"): + for i in range(2): + want += [ + f"quantizer.{stage}.{i}.0.conv.weight", + f"quantizer.{stage}.{i}.1.dwconv.conv.weight", + f"quantizer.{stage}.{i}.1.pwconv1.weight", + f"quantizer.{stage}.{i}.1.pwconv2.weight", + f"quantizer.{stage}.{i}.1.norm.weight", + f"quantizer.{stage}.{i}.1.gamma", + ] + want += ["decoder.model.0.conv.weight"] + for b in range(1, 5): + want += [ + f"decoder.model.{b}.block.0.alpha", + f"decoder.model.{b}.block.1.conv.weight", + ] + for r in range(3): + want += [ + f"decoder.model.{b}.block.{r + 2}.block.0.alpha", + f"decoder.model.{b}.block.{r + 2}.block.1.conv.weight", + ] + want += ["decoder.model.5.alpha", "decoder.model.6.conv.weight"] + return want + + +def dit_probe_names(config: Dict[str, int]) -> List[str]: + want = ["text_encoder.text_embedding.weight", "in_proj.weight", "in_proj.bias"] + want += ["cond_module.0.weight", "cond_module.2.weight", "cond_module.4.weight"] + want += ["text_norm.weight", "speaker_norm.weight", "out_norm.weight"] + want += ["out_proj.weight", "out_proj.bias"] + want += ["speaker_encoder.in_proj.weight", "speaker_encoder.in_proj.bias"] + for i in range(config["text_num_layers"]): + want += [ + f"text_encoder.blocks.{i}.attention.wq.weight", + f"text_encoder.blocks.{i}.attention.q_norm.weight", + f"text_encoder.blocks.{i}.mlp.w1.weight", + f"text_encoder.blocks.{i}.attention_norm.weight", + ] + for i in range(config["speaker_num_layers"]): + want += [f"speaker_encoder.blocks.{i}.attention.wq.weight"] + for i in range(config["num_layers"]): + want += [ + f"blocks.{i}.attention.wq.weight", + f"blocks.{i}.attention.wk_text.weight", + f"blocks.{i}.attention.wv_speaker.weight", + f"blocks.{i}.attention.q_norm.weight", + f"blocks.{i}.mlp.w2.weight", + f"blocks.{i}.attention_adaln.shift_down.weight", + f"blocks.{i}.attention_adaln.shift_up.bias", + f"blocks.{i}.mlp_adaln.gate_up.weight", + ] + return want + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("gguf", help="path to the converted Echo-TTS GGUF") + args = parser.parse_args() + + if gguf is None: + print("the 'gguf' package is required (pip install gguf)", file=sys.stderr) + return 2 + + try: + reader = gguf.GGUFReader(args.gguf) + tensors = {t.name: t for t in reader.tensors} + fields = reader.fields + except Exception as error: # noqa: BLE001 - any parse failure is a bad file + print(f"could not read {args.gguf} as a GGUF: {error}", file=sys.stderr) + print("the file may be truncated or still being written", file=sys.stderr) + return 1 + + def kv(name): + field = fields.get(name) + if field is None: + return None + try: + return field.parts[field.data[0]][0] + except Exception: # noqa: BLE001 - metadata shape varies by writer + return None + + problems: List[str] = [] + too_long = [n for n in tensors if len(n) >= GGML_MAX_NAME] + if too_long: + longest = max(too_long, key=len) + problems.append( + f"{len(too_long)} tensor names reach GGML_MAX_NAME ({GGML_MAX_NAME}); " + f"longest is {len(longest)} chars: {longest}") + + notes: List[str] = [] + + config = {} + for key in ("num_layers", "text_num_layers", "speaker_num_layers", + "model_size", "latent_size", "ae_latent_dim"): + value = kv(f"echo_tts.{key}") + if value is None: + problems.append(f"missing metadata echo_tts.{key}") + else: + config[key] = int(value) + + by_prefix: Dict[str, int] = {} + for name in tensors: + head = name.split("/")[0] + "/" if "/" in name else "(no namespace)" + by_prefix[head] = by_prefix.get(head, 0) + 1 + + print(f"file : {args.gguf}") + print(f"tensors : {len(tensors)}") + for head in sorted(by_prefix): + print(f" {head:16} {by_prefix[head]}") + scale = kv("echo_tts.pca_latent_scale") + has_codec = kv("echo_tts.has_codec_weights") + print(f"latent_scale : {scale}") + print(f"has_codec_weights : {has_codec}") + + # PrefixedTensorSourceView routes on `prefix + "/"`. A dot-separated name + # is never routed and the namespace reports as non-existent at load time, + # so check the separator explicitly rather than inferring it from coverage. + # package.cpp refuses to load a published GGUF that does not embed its spec. + spec_json = None + for key, kind in (("audiocpp.model_spec.version", "uint32"), + ("audiocpp.model_spec.family", "string"), + ("audiocpp.model_spec.json", "string")): + field = fields.get(key) + if field is None: + problems.append(f"missing embedded model spec key '{key}'") + continue + raw = field.parts[field.data[0]] + if kind == "string": + value = bytes(raw).decode("utf-8", "replace") + if key.endswith(".family") and value != "echo_tts": + problems.append(f"{key} is {value!r}, expected 'echo_tts'") + if key.endswith(".json"): + spec_json = value + elif int(raw[0]) != 1: + problems.append(f"{key} is {int(raw[0])}, expected 1") + if spec_json is not None: + try: + import json as _json + embedded = _json.loads(spec_json) + if embedded.get("schema_version") != 1: + problems.append("embedded model spec is not schema_version 1") + except Exception as error: # noqa: BLE001 + problems.append(f"embedded model spec is not valid JSON: {error}") + + unrouted = [n for n in tensors if "/" not in n] + if unrouted: + problems.append( + f"{len(unrouted)} tensors are outside any namespace (no '/' separator), " + f"e.g. {unrouted[:3]}; the loader will report the namespace as missing") + for namespace in ("dit_weights/", "pca/", "ae/" if has_codec else None): + if namespace and not any(n.startswith(namespace) for n in tensors): + problems.append(f"namespace '{namespace}' is empty; the loader will refuse to open it") + + if not config: + print("\nno echo_tts metadata found; is this an Echo-TTS GGUF?", file=sys.stderr) + return 1 + + # --- coverage --- + dit_names = {n[len(DIT_PREFIX):] for n in tensors if n.startswith(DIT_PREFIX)} + missing_dit = [n for n in dit_probe_names(config) if n not in dit_names] + if missing_dit: + problems.append(f"{len(missing_dit)} DiT tensors missing, e.g. {missing_dit[:3]}") + + for required in ("components", "mean"): + if PCA_PREFIX + required not in tensors: + problems.append(f"missing {PCA_PREFIX}{required}") + + codec_names = {n[len(CODEC_PREFIX):] for n in tensors if n.startswith(CODEC_PREFIX)} + if has_codec: + missing_codec = [n for n in codec_probe_names() if n not in codec_names] + if missing_codec: + problems.append( + f"{len(missing_codec)} codec tensors missing, e.g. {missing_codec[:3]}") + leftovers = [n for n in codec_names + if "parametrizations" in n or n.endswith(("weight_g", "weight_v"))] + if leftovers: + problems.append( + f"{len(leftovers)} codec tensors still weight-normalised, e.g. {leftovers[:2]}") + buffers = [n for n in codec_names if n.endswith(("causal_mask", "freqs_cis"))] + if buffers: + notes.append(f"{len(buffers)} regenerable buffers were packaged (harmless, wastes space)") + else: + problems.append("has_codec_weights is false; re-run the converter with --fish-dir") + + # --- shapes --- + # ggml_n_dims() ignores trailing 1s, so a (1, C, 1) tensor would read back as + # (C, 1). audio.cpp restores exact logical shapes from two parallel arrays in + # tensor order; without them the loader falls back to the lossy inference. + exact_shapes = {} + rank_field = fields.get("audiocpp.tensor_ranks") + shape_field = fields.get("audiocpp.tensor_shapes") + if rank_field is None or shape_field is None: + problems.append( + "missing audiocpp.tensor_ranks/tensor_shapes; tensors with trailing " + "size-1 dimensions (snake alphas) will fail their shape checks") + else: + ranks = [int(rank_field.parts[i][0]) for i in rank_field.data] + flat = [int(shape_field.parts[i][0]) for i in shape_field.data] + order = list(reader.tensors) + if len(ranks) != len(order): + problems.append( + f"audiocpp.tensor_ranks has {len(ranks)} entries for {len(order)} tensors") + elif sum(ranks) != len(flat): + problems.append( + f"audiocpp.tensor_shapes has {len(flat)} values, expected {sum(ranks)}") + else: + cursor = 0 + for tensor, rank in zip(order, ranks): + exact_shapes[tensor.name] = tuple(flat[cursor:cursor + rank]) + cursor += rank + + def shape(name): + if name in exact_shapes: + return exact_shapes[name] + t = tensors.get(name) + # GGUF stores dimensions reversed relative to the logical order. + return tuple(int(d) for d in reversed(t.shape)) if t is not None else None + + expect_shapes = { + PCA_PREFIX + "components": (config["latent_size"], config["ae_latent_dim"]), + PCA_PREFIX + "mean": (config["ae_latent_dim"],), + DIT_PREFIX + "out_proj.weight": (config["latent_size"], config["model_size"]), + } + for name, want in expect_shapes.items(): + got = shape(name) + if got is not None and got != want: + problems.append(f"{name}: shape {got}, expected {want}") + + # --- dtype policy --- + wrong_dtype = [] + for name, tensor in tensors.items(): + if any(marker in name for marker in F32_REQUIRED_MARKERS) or name.endswith(".bias"): + if tensor.tensor_type != gguf.GGMLQuantizationType.F32: + wrong_dtype.append((name, tensor.tensor_type.name)) + if wrong_dtype: + problems.append( + f"{len(wrong_dtype)} precision-sensitive tensors are not F32, " + f"e.g. {wrong_dtype[:3]} -- re-run with the current converter") + + counts: Dict[str, int] = {} + for tensor in tensors.values(): + counts[tensor.tensor_type.name] = counts.get(tensor.tensor_type.name, 0) + 1 + print(f"dtypes : {counts}") + + print() + for note in notes: + print(f" note: {note}") + if problems: + for problem in problems: + print(f" FAIL: {problem}") + return 1 + print(" GGUF looks good: tensor coverage, shapes, and precision policy all check out") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())