feat(tts/kokoro-v1.1-zh): CoreML conversion + atan2 phase noise fix

[Alex-Wengg]

Summary

• Adds the seven-stage CoreML conversion pipeline for Kokoro-82M-v1.1-zh (Mandarin) — ALBERT, PostAlbert, Alignment, Prosody, Noise, Vocoder, Tail.
• Fixes the audible high-frequency noise that survived every weight-quantisation, compute-precision, and dispatch-mode change tried in rounds 1–2 of the investigation. Root cause was a missing torch.atan2 PyTorch/ONNX-divergence correction in CoreMLForwardSTFT.
• Ships before/after/reference WAVs in docs/audio/ for direct A/B listening on GitHub.

Root cause

CoreMLForwardSTFT.transform (convert-coreml.py:432-456) was missing the same atan2 correction the upstream kokoro/custom_stft.py:135-138 already documents:

input	PyTorch atan2	CoreML MIL atan2
(imag = 0, real < 0)	+π	0

For real-valued STFT input the DC bin's imag is exactly 0 and the Nyquist bin's imag sits at the fp32 noise floor (~1e-15). Whenever real < 0 at those bins, the converted CoreML graph emitted phase = 0 instead of +π. The π/2π offset propagated through noise_convs[0] (a strided conv that mixes 11 frequency bins) into all 256 channels of x_source_0, surfacing as broad-spectrum noise above 10 kHz — present even in a pure-fp32 build on cpuOnly.

The fix

Two-stage clip-then-correct in CoreMLForwardSTFT.transform:

eps = 1e-5
imag_clipped = torch.where(imag_out.abs() < eps,
                           torch.zeros_like(imag_out), imag_out)
magnitude = torch.sqrt(real_out ** 2 + imag_clipped ** 2 + 1e-14)
phase = torch.atan2(imag_clipped, real_out)
correction_mask = (imag_clipped == 0) & (real_out < 0)
phase = torch.where(correction_mask,
                    torch.full_like(phase, math.pi), phase)

eps = 1e-5 catches both the exact-zero (DC) and computational-zero (Nyquist) cases without affecting any legitimate spectral imag value (≥ 1e-3 for typical audio). The where() patches the PyTorch branch.

Verification

KokoroNoise conversion fidelity (PyTorch trace vs CoreML, deterministic)

metric	before fix	after fix
x_source_0 rel-rms	0.397	0.057 (7×)
x_source_0 corr	0.911	0.998
x_source_1 rel-rms	0.070	0.002
x_source_1 corr	0.996	1.000
STFT phase max|diff|	6.283 (2π)	0.000

End-to-end audio HF (≥10 kHz) band power

zm_009	before fix	after fix
CoreML output	−81.31 dB	−83.91 dB
PyTorch reference	−83.81 dB	−83.90 dB
Δ(CoreML − PyTorch)	+2.50 dB	−0.01 dB
HF residual (coreml − pt)	−86.55 dB	−91.18 dB

zf_001 Δ(CoreML − PyTorch): +0.11 dB (no regression).

Reference audio (download from docs/audio/)

• before_fix_zm009.wav — broken Noise stage (audible HF noise, more obvious on the male voice)
• after_fix_zm009.wav — fixed Noise stage (HF matches PyTorch within 0.01 dB)
• pytorch_ref_zm009.wav — PyTorch teacher reference
• after_fix_zf001.wav — sanity check on the female voice

All 24 kHz mono float32. The zm_009 triplet is the same prompt (你好世界，今天天气很好。) so the residual is the noise-stage delta.

Investigation rounds

The full story is in report.md:

• Round 1 — ruled out weight precision (int8 → fp32), compute precision (FLOAT16 → FLOAT32), dispatch (ANE/GPU/CPU), voice timbre offset, and the cos-Snake1D identity rewrite.
• Round 2 — localised noise to KokoroNoise via a five-tier stage-swap chain (per_stage_diff.py); misdiagnosed the cause as sin-of-large-args inside CoreMLSineGenV2 because the substage probe conflated rand_ini randomness with graph drift.
• Round 3 — conversion-fidelity diff (probe_noise_fidelity.py) immediately exposed phase max|diff| = 2π, pointing at atan2. A wrap-before-sin attempt actually made x_source_0 worse (rel 0.025 vs 0.0009) before the right fix landed.

Diagnostic scripts (all self-contained, in models/tts/kokoro-v1.1-zh/coreml/)

script	what it does
tail_compare.py	three-way iSTFT comparison (torch.istft, CustomSTFT, CoreML Tail) on the same x_pre
vocoder_xpre_diff.py	Vocoder x_pre PT-vs-CoreML diff
per_stage_diff.py	five-tier stage-swap chain (full PT → … → full CoreML) with HF-band tracking
probe_noise.py	CU sweep + per-channel divergence on the noise stage
probe_noise_fidelity.py	conversion fidelity of CoreMLFullNoiseModel (the script that found the atan2 bug)
probe_sinegen_isolated.py	six standalone CoreML models that disprove the Round 2 sin-precision hypothesis

Test plan

• convert-coreml.py --stages noise --no-palettize noise produces a valid KokoroNoise.mlpackage.
• probe_noise_fidelity.py reports x_source_0 rel-rms ≤ 0.06, x_source_1 ≤ 0.003.
• per_stage_diff.py on zm_009 shows audio_D HF within 0.01 dB of audio_A.
• per_stage_diff.py on zf_001 shows ≤ 0.15 dB HF delta (no regression on the female voice).
• inference.py produces a 24 kHz WAV indistinguishable in HF band from the PyTorch reference.
• Listening test on a Mac/iPhone in the deployment app (FluidAudio) to confirm the noise is perceptually gone.

🤖 Generated with Claude Code

[Alex-Wengg]

@devin-ai-integration please take a look

[beta-devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.