From b49648344819b199252fe918c3b0d5ec43c4d389 Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 10:32:52 -0400 Subject: [PATCH] fix(kokoro): COLA-normalize iSTFT deconv weights (FluidAudio #852) CoreMLCustomSTFT replaces torch.istft with ConvTranspose1d overlap-add but omitted the summed-squared-window (COLA) normalization torch.istft applies, leaving output exactly 1.5x too loud (periodic Hann, win 20, hop 5 -> interior envelope sum(w^2) = 1.5, a pure scalar; edge taps are sliced off by the center pad). Fold the envelope into the synthesis weights in both the laishere (en/ja) and v1.1-zh conversion scripts. This reconstructs the fix described in FluidAudio PR #699, which was built and measured (1.02x PyTorch raw level, jf_alpha peak 0.306 vs 0.299) but never committed; only the built artifact survived in build/tail-fix/. The reconstruction reproduces that artifact's deconv weights to within 1 ulp (max abs diff 7.5e-9, fp32 rounding order). Also commit convert-voices.py (laishere voice-pack extraction used for ANE-ja/, referenced by #699 but likewise never committed). Corrected tails are published as KokoroTail_v2.mlmodelc alongside the originals in ANE/, ANE-ja/, ANE-zh/ of FluidInference/kokoro-82m-coreml (commit acac8811); consumer adoption is tracked in FluidAudio #852. --- .../coreml/scripts/convert-coreml.py | 10 +++ .../kokoro/laishere-coreml/convert-coreml.py | 10 +++ .../kokoro/laishere-coreml/convert-voices.py | 81 +++++++++++++++++++ 3 files changed, 101 insertions(+) create mode 100644 models/tts/kokoro/laishere-coreml/convert-voices.py diff --git a/models/tts/kokoro-v1.1-zh/coreml/scripts/convert-coreml.py b/models/tts/kokoro-v1.1-zh/coreml/scripts/convert-coreml.py index c191f79..e9f35f3 100644 --- a/models/tts/kokoro-v1.1-zh/coreml/scripts/convert-coreml.py +++ b/models/tts/kokoro-v1.1-zh/coreml/scripts/convert-coreml.py @@ -157,6 +157,16 @@ def __init__(self, original): backward_imag = original.weight_backward_imag.clone() backward_real[1:-1] *= 2.0 backward_imag[1:-1] *= 2.0 + # COLA normalization torch.istft applies but the deconv path omits: + # analysis and synthesis each apply the window once, so the + # overlap-added output carries a summed-w^2 envelope (constant 1.5 in + # the interior for periodic Hann at hop = n_fft/4). Fold it into the + # synthesis weights; the edge taps are sliced off by the center pad. + window_sq = original.window.float() ** 2 + cola = window_sq.reshape(self.n_fft // self.hop_length, self.hop_length).sum(dim=0) + assert torch.allclose(cola, cola[:1].expand_as(cola), rtol=1e-5), cola + backward_real /= cola[0] + backward_imag /= cola[0] self.deconv_real.weight = nn.Parameter(backward_real, requires_grad=False) self.deconv_imag.weight = nn.Parameter(backward_imag, requires_grad=False) diff --git a/models/tts/kokoro/laishere-coreml/convert-coreml.py b/models/tts/kokoro/laishere-coreml/convert-coreml.py index b14f426..20324a5 100644 --- a/models/tts/kokoro/laishere-coreml/convert-coreml.py +++ b/models/tts/kokoro/laishere-coreml/convert-coreml.py @@ -152,6 +152,16 @@ def __init__(self, original): backward_imag = original.weight_backward_imag.clone() backward_real[1:-1] *= 2.0 backward_imag[1:-1] *= 2.0 + # COLA normalization torch.istft applies but the deconv path omits: + # analysis and synthesis each apply the window once, so the + # overlap-added output carries a summed-w^2 envelope (constant 1.5 in + # the interior for periodic Hann at hop = n_fft/4). Fold it into the + # synthesis weights; the edge taps are sliced off by the center pad. + window_sq = original.window.float() ** 2 + cola = window_sq.reshape(self.n_fft // self.hop_length, self.hop_length).sum(dim=0) + assert torch.allclose(cola, cola[:1].expand_as(cola), rtol=1e-5), cola + backward_real /= cola[0] + backward_imag /= cola[0] self.deconv_real.weight = nn.Parameter(backward_real, requires_grad=False) self.deconv_imag.weight = nn.Parameter(backward_imag, requires_grad=False) diff --git a/models/tts/kokoro/laishere-coreml/convert-voices.py b/models/tts/kokoro/laishere-coreml/convert-voices.py new file mode 100644 index 0000000..a3ea391 --- /dev/null +++ b/models/tts/kokoro/laishere-coreml/convert-voices.py @@ -0,0 +1,81 @@ +"""Extract Kokoro-82M voice packs (.pt -> .bin flat fp32 [510, 256]). + +The laishere 7-stage CoreML graphs + vocab.json are produced from the shared +base `hexgrad/Kokoro-82M` acoustic model and are language-agnostic, so a new +language variant (e.g. ANE-ja/) reuses those bundles unchanged and only needs +its voice packs in FluidAudio's `[510, 256]` flat float32 format. + +Unlike the v1.1-zh `convert-voices.py`, this loads each `.pt` tensor directly +(no `KPipeline`), so it needs no per-language G2P dependencies (misaki[ja], +fugashi/MeCab, ...). torch + huggingface_hub only. + +Usage: + # All Japanese voices into an ANE-ja staging dir: + python convert-voices.py --prefix jf jm --output-dir build/ANE-ja/voices + + # Specific voices: + python convert-voices.py --only jf_alpha jm_kumo --output-dir /tmp/voices +""" +from __future__ import annotations + +import argparse +import pathlib + +import numpy as np +import torch +from huggingface_hub import HfApi, hf_hub_download + +REPO_ID = "hexgrad/Kokoro-82M" +EXPECTED_SHAPE = (510, 1, 256) + + +def list_remote_voices(repo_id: str, prefixes: list[str] | None) -> list[str]: + files = HfApi().list_repo_files(repo_id=repo_id) + stems = sorted( + pathlib.Path(f).stem + for f in files + if f.startswith("voices/") and f.endswith(".pt") + ) + if prefixes: + stems = [s for s in stems if any(s.startswith(p) for p in prefixes)] + return stems + + +def main() -> None: + p = argparse.ArgumentParser(description="Extract Kokoro-82M voice packs to flat .bin") + p.add_argument("--output-dir", type=pathlib.Path, required=True) + p.add_argument("--repo-id", default=REPO_ID) + p.add_argument("--prefix", nargs="*", default=None, + help="Voice-id prefixes to keep (e.g. jf jm for Japanese)") + p.add_argument("--only", nargs="*", default=None, + help="Explicit voice ids (overrides remote enumeration)") + args = p.parse_args() + + args.output_dir.mkdir(parents=True, exist_ok=True) + + voices = sorted(args.only) if args.only else list_remote_voices(args.repo_id, args.prefix) + print(f"Converting {len(voices)} voice(s) from {args.repo_id}: {voices}") + + converted = 0 + for vid in voices: + out_path = args.output_dir / f"{vid}.bin" + try: + pt = hf_hub_download(args.repo_id, f"voices/{vid}.pt") + tensor = torch.load(pt, weights_only=True) + except Exception as e: # noqa: BLE001 - report and continue + print(f" [FAIL] {vid}: {type(e).__name__}: {e}") + continue + arr = tensor.cpu().numpy().astype(np.float32) + if arr.shape != EXPECTED_SHAPE: + print(f" [FAIL] {vid}: unexpected shape {arr.shape}, want {EXPECTED_SHAPE}") + continue + arr = arr.reshape(510, 256) + out_path.write_bytes(arr.tobytes()) + converted += 1 + print(f" [{converted}/{len(voices)}] {vid}.bin ({out_path.stat().st_size} bytes)") + + print(f"\nDone. converted={converted} -> {args.output_dir}/") + + +if __name__ == "__main__": + main()