fix(kokoro, kitten_tts): match PyTorch reference semantics in five decoder paths (-2.5 dB level error, F0-path misalignment, window mismatch)

[mchen04]

Summary

Five places where the MLX iSTFTNet ports (kokoro, kitten_tts) diverge from the PyTorch reference (hexgrad/kokoro). Each was found by instrumenting both stacks layer-by-layer on identical phoneme/style inputs and comparing every intermediate tensor. Two are clearly audible; the rest are smaller semantic mismatches found on the way. Value-pinned regression tests included for the three numerically-checkable fixes.

1. Constant −2.5 dB output attenuation (MLXSTFT.inverse, kokoro)

dsp.istft defaults to normalized=False (overlap-add division by Σw). torch.istft always divides by Σw² (least-squares inversion — unrelated to torch.istft's own normalized argument, which is FFT scaling). For the periodic-hann, win = 4×hop configuration this is a constant amplitude factor of Σw²/Σw = 1.5/2.0 = 0.75 (−2.50 dB).

Measured: MLX/torch waveform RMS ratio 0.738 (−2.64 dB) before, −0.16 dB after (residual = stochastic noise paths). This also explains why downstream users compensate with hardcoded gain hacks.

⚠️ User-facing change: all Kokoro output becomes ~2.5 dB louder (×4/3). This restores the reference level the network was trained against (no clipping observed across a 55-utterance eval set; post-fix mean level ≈ −26 dB RMS on af_heart), but anyone with downstream loudness calibration or golden outputs will notice. kitten_tts already passed normalized=True.

2. Symmetric/periodic window mismatch (MLXSTFT, kokoro + kitten_tts)

The torch reference uses hann_window(win, periodic=True) for both analysis and synthesis. The MLX string path resolves "hann" to a symmetric window in dsp.stft while dsp.istft builds a periodic one. The mismatch breaks exact COLA inversion (~3% reconstruction ripple — caught by the new round-trip test) and skews the harmonic-source STFT features fed to the generator. MLXSTFT now materializes the periodic window once and passes it to both directions; dsp.stft's string behavior is untouched (no blast radius onto other callers).

3. One-frame misalignment in AdainResBlk1d upsample path (kokoro + kitten_tts)

torch uses ConvTranspose1d(k=3, stride=2, groups=dim_in, padding=1, output_padding=1), which maps T→2T by trimming one sample from the left of the unpadded (2T+1) transpose-conv output. The kokoro port ran the transpose conv with padding=1 (→ 2T−1) and left-zero-padded — shifting the residual branch one frame against the shortcut and replacing a computed tail sample with 0. kitten_tts right-padded (aligned) but still zeroed the computed tail sample.

Both now build the pool conv with padding=0 (construction-time, no runtime state mutation) and slice [:, 1:, :] — exact by derivation and pinned against torch constants in the new test.

This block sits in predictor.F0[1], predictor.N[1], and decoder.decode[3] — directly in the pitch/energy contour path. Measured on identical inputs in fp32 (kokoro): F0_pred relRMSE 0.134 (corr 0.975) before → 0.0000 (corr 1.0000) after.

4. SineGen initial harmonic phase distribution (kokoro + kitten_tts)

Reference draws initial phase offsets with torch.rand (uniform [0,1)); the ports used mx.random.normal. Uniform is the correct full-circle random phase.

5. interpolate1d negative source coordinates (shared tts/models/interpolate.py)

With align_corners=False, the source coordinate for the first outputs is negative (e.g. −0.498 at scale 300, as used by the harmonic-source upsampler). torch clamps to 0; here floor(x) = −1 and the gather read an out-of-range index (observed as the last frame), so the first half-frame interpolated against the end of the signal. Clamped to 0.

Note: this is a shared utility — the clamp corrects output for every linear/align_corners=False caller (kokoro, kitten_tts, soprano), not just kokoro.

Tests

mlx_audio/tts/tests/test_istftnet_fidelity.py (new, parametrized over kokoro + kitten_tts):

• transpose-conv emulation vs pinned torch.nn.ConvTranspose1d(k=3, s=2, p=1, op=1) constants
• MLXSTFT transform→inverse round trip at unity gain (old code reconstructs at 0.75×)

test_interpolate.py: the existing align_corners=False case asserted only the output shape (which is how the wrap bug survived); it now pins values from torch.nn.functional.interpolate — first element must be 1.0, old code produced 2.5.

Full mlx_audio/tts/tests suite passes (586 passed; TestSparkTTSModel::test_init fails identically on clean main — pre-existing, unrelated).

Validation beyond unit tests

• Layer probe vs PyTorch kokoro (fp32 weights prince-canuma/Kokoro-82M, torch 2.12.1, kokoro 0.9.4, misaki 0.9.4): all decoder-upstream tensors match at relRMSE ≤ 2e-3 (bf16) / 0.0 (fp32); predicted durations identical.
• 55-utterance eval battery against frozen fp32 reference audio (paired log-mel L1, DTW-MCD, multi-res STFT, F0/VUV via pyworld, WavLM-SV speaker cosine), thresholds calibrated against the reference model's own render-to-render noise floor (the decoder is stochastic): mel_l1 0.601 → 0.187 (floor 0.077), MCD 7.29 → 4.09 dB (floor 1.86), F0 RMSE 11.5 → 5.1 Hz (floor 3.7), speaker cosine 0.996 → 0.999 (floor 0.9997).
• Whisper-large-v3-turbo WER on the fixed stack matches the PyTorch reference within 0.04 pp on the same texts.
• Output durations are unchanged by all five fixes.

Scope notes

• Other dsp.istft callers (soprano, vocos, deepfilternet, lfm_audio, …) still use the plain-Σw default and were not audited here; if they were trained against torch.istft, they may carry the same 0.75× attenuation. Flagging rather than changing, since each model has its own training convention.
• Happy to split this into separate PRs if preferred.

[lucasnewman]

@mchen04 Thanks! This looks reasonable, but please see the comments inline.

[lucasnewman]

Note you'll also need to run the formatter with pre-commit run --all and make sure your commits are signed.

[mchen04]

this was my first open-source contribution, so thank for the guidance :) I'm glad to be able to help out with some issues i found

[lucasnewman]

Thanks!