LGTM: PyTorch + ONNX weights and inference code

README.md

@@ -0,0 +1,98 @@
+---
+language: [en, es, pt, fr, de, it, sv, vi, ja, ko, id]
+pipeline_tag: text-to-speech
+tags: [text-to-speech, tts, voice-cloning, onnx, multilingual]
+---
+
+# LGTM
+
+Multilingual text-to-speech at 44.1 kHz with 10 built-in voices and zero-shot voice cloning.
+Available in **PyTorch** and **ONNX** (ONNX Runtime needs no PyTorch).
+
+**Languages:** English `en`, Spanish `es`, Portuguese `pt`, French `fr`, German `de`, Italian `it`,
+Swedish `sv`, Vietnamese `vi`, Japanese `ja`, Korean `ko`, Indonesian `id`
+
+**Built-in voices:** `F1` `F2` `F3` `F4` `F5` (female), `M1` `M2` `M3` `M4` `M5` (male)
+
+## Setup
+
+```bash
+git clone https://huggingface.co/qnguyen3/LGTM
+cd LGTM
+pip install -r requirements.txt          # PyTorch backend
+pip install -r requirements-onnx.txt     # ONNX backend
+```
+
+## PyTorch
+
+```python
+from lgtm import LGTMTTS
+
+tts = LGTMTTS.from_pretrained(".")            # or "qnguyen3/LGTM" to download
+wav = tts.synthesize("Xin chào, hôm nay trời đẹp quá!", lang="vi", voice="F1")
+tts.save_wav(wav, "out.wav")                   # 44.1 kHz mono
+```
+
+## ONNX Runtime
+
+```python
+from lgtm import LGTMOnnx
+
+tts = LGTMOnnx.from_pretrained(".", use_gpu=False)   # use_gpu=True with onnxruntime-gpu
+wav = tts.synthesize("Bonjour à tous, comment allez-vous ?", lang="fr", voice="M1")
+tts.save_wav(wav, "out.wav")
+```
+
+## Voice cloning
+
+Give 5-15 seconds of clean speech; the voice can then speak any supported language.
+
+```python
+voice = tts.clone_voice("reference.wav")      # works with both backends
+wav = tts.synthesize("This is my cloned voice.", lang="en", voice=voice)
+
+from lgtm import save_voice_style              # ONNX: from lgtm.onnx_inference import save_voice_style
+save_voice_style("my_voice.json", voice)       # reuse later: voice="my_voice.json"
+```
+
+## Command line
+
+```bash
+python -m lgtm.cli --text "Hej! Hur mår du idag?" --lang sv --voice F2 --out out.wav
+python -m lgtm.cli --text "안녕하세요" --lang ko --ref reference.wav --save_voice my_voice.json --out out.wav
+python -m lgtm.cli --text "Selamat pagi" --lang id --voice M3 --backend onnx --out out.wav
+```
+
+## Options
+
+| argument | default | |
+|---|---|---|
+| `lang` | `"en"` | language code (see above) |
+| `voice` | `"F1"` | preset name, path to a voice `.json`, or `clone_voice()` output |
+| `steps` | `8` | denoising steps (fewer = faster, e.g. 5) |
+| `speed` | `1.05` | speaking rate (higher = faster) |
+| `silence` | `0.3` | seconds of silence between sentences (long text is split automatically) |
+
+## Files
+
+| path | contents |
+|---|---|
+| `pytorch/model.safetensors` | all weights (synthesis + voice cloning) |
+| `onnx/text_encoder.onnx`, `duration_predictor.onnx`, `vector_estimator.onnx`, `vocoder.onnx` | synthesis graphs |
+| `onnx/voice_encoder.onnx` | reference audio → voice style (cloning) |
+| `voice_styles/*.json` | built-in voices |
+| `config.json`, `unicode_indexer.json` | model config, text vocabulary |
+| `lgtm/` | inference code (`inference.py` PyTorch, `onnx_inference.py` ONNX, `cli.py`) |
+
+### ONNX graph I/O (for custom runtimes)
+
+| graph | inputs | outputs |
+|---|---|---|
+| `text_encoder` | `text_ids` int64 [B,T], `style_ttl` [B,50,256], `text_mask` [B,1,T] | `text_emb` [B,256,T] |
+| `duration_predictor` | `text_ids`, `style_dp` [B,8,16], `text_mask` | `duration` [B] (seconds) |
+| `vector_estimator` | `noisy_latent` [B,144,L], `text_emb`, `style_ttl`, `latent_mask` [B,1,L], `text_mask`, `current_step` [B], `total_step` [B] | `denoised_latent` [B,144,L] |
+| `vocoder` | `latent` [B,144,L] | `wav` [B, 3072·L] |
+| `voice_encoder` | `wav` [1,N] (44.1 kHz) | `style_ttl` [1,50,256], `style_dp` [1,8,16] |
+
+Sampling loop: `L = ceil(duration·44100 / 3072)`, start from Gaussian noise masked by `latent_mask`,
+call `vector_estimator` for `current_step = 0 … total_step-1`, then `vocoder`. See `lgtm/onnx_inference.py`.

config.json

@@ -0,0 +1,311 @@
+{
+  "tts_version": "v1.7.3",
+  "split": "opensource-multilingual",
+  "ttl": {
+    "latent_dim": 24,
+    "chunk_compress_factor": 6,
+    "batch_expander": {
+      "n_batch_expand": 6
+    },
+    "normalizer": {
+      "scale": 0.25
+    },
+    "text_encoder": {
+      "n_langs": 0,
+      "lang_emb_dim": 0,
+      "text_embedder": {
+        "char_emb_dim": 256
+      },
+      "convnext": {
+        "idim": 256,
+        "ksz": 5,
+        "intermediate_dim": 1024,
+        "num_layers": 6,
+        "dilation_lst": [
+          1,
+          1,
+          2,
+          2,
+          4,
+          4
+        ]
+      },
+      "attn_encoder": {
+        "hidden_channels": 256,
+        "filter_channels": 1024,
+        "n_heads": 4,
+        "n_layers": 4,
+        "p_dropout": 0.0
+      },
+      "proj_out": {
+        "idim": 256,
+        "odim": 256
+      }
+    },
+    "flow_matching": {
+      "sig_min": 1e-08
+    },
+    "style_encoder": {
+      "proj_in": {
+        "ldim": 24,
+        "chunk_compress_factor": 6,
+        "odim": 256
+      },
+      "convnext": {
+        "idim": 256,
+        "ksz": 5,
+        "intermediate_dim": 1024,
+        "num_layers": 6,
+        "dilation_lst": [
+          1,
+          1,
+          1,
+          1,
+          1,
+          1
+        ]
+      },
+      "style_token_layer": {
+        "input_dim": 256,
+        "n_style": 50,
+        "style_key_dim": 256,
+        "style_value_dim": 256,
+        "prototype_dim": 256,
+        "n_units": 256,
+        "n_heads": 2
+      }
+    },
+    "speech_prompted_text_encoder": {
+      "text_dim": 256,
+      "style_dim": 256,
+      "n_units": 256,
+      "n_heads": 2
+    },
+    "uncond_masker": {
+      "prob_both_uncond": 0.04,
+      "prob_text_uncond": 0.01,
+      "std": 0.1,
+      "text_dim": 256,
+      "n_style": 50,
+      "style_key_dim": 256,
+      "style_value_dim": 256
+    },
+    "vector_field": {
+      "n_langs": 0,
+      "lang_emb_dim": 0,
+      "proj_in": {
+        "ldim": 24,
+        "chunk_compress_factor": 6,
+        "odim": 512
+      },
+      "time_encoder": {
+        "time_dim": 64,
+        "hdim": 256
+      },
+      "main_blocks": {
+        "n_blocks": 4,
+        "time_cond_layer": {
+          "idim": 512,
+          "time_dim": 64
+        },
+        "style_cond_layer": {
+          "idim": 512,
+          "style_dim": 256
+        },
+        "text_cond_layer": {
+          "idim": 512,
+          "text_dim": 256,
+          "n_heads": 8,
+          "n_units": 512,
+          "use_residual": true,
+          "rotary_base": 10000,
+          "rotary_scale": 10
+        },
+        "convnext_0": {
+          "idim": 512,
+          "ksz": 5,
+          "intermediate_dim": 2048,
+          "num_layers": 4,
+          "dilation_lst": [
+            1,
+            2,
+            4,
+            8
+          ]
+        },
+        "convnext_1": {
+          "idim": 512,
+          "ksz": 5,
+          "intermediate_dim": 2048,
+          "num_layers": 1,
+          "dilation_lst": [
+            1
+          ]
+        },
+        "convnext_2": {
+          "idim": 512,
+          "ksz": 5,
+          "intermediate_dim": 2048,
+          "num_layers": 1,
+          "dilation_lst": [
+            1
+          ]
+        }
+      },
+      "last_convnext": {
+        "idim": 512,
+        "ksz": 5,
+        "intermediate_dim": 2048,
+        "num_layers": 4,
+        "dilation_lst": [
+          1,
+          1,
+          1,
+          1
+        ]
+      },
+      "proj_out": {
+        "idim": 512,
+        "chunk_compress_factor": 6,
+        "ldim": 24
+      }
+    }
+  },
+  "ae": {
+    "sample_rate": 44100,
+    "n_delay": 0,
+    "base_chunk_size": 512,
+    "chunk_compress_factor": 1,
+    "ldim": 24,
+    "encoder": {
+      "spec_processor": {
+        "n_fft": 2048,
+        "win_length": 2048,
+        "hop_length": 512,
+        "n_mels": 228,
+        "sample_rate": 44100,
+        "eps": 1e-05,
+        "norm_mean": 0.0,
+        "norm_std": 1.0
+      },
+      "ksz_init": 7,
+      "ksz": 7,
+      "num_layers": 10,
+      "dilation_lst": [
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1
+      ],
+      "intermediate_dim": 2048,
+      "idim": 1253,
+      "hdim": 512,
+      "odim": 24
+    },
+    "decoder": {
+      "ksz_init": 7,
+      "ksz": 7,
+      "num_layers": 10,
+      "dilation_lst": [
+        1,
+        2,
+        4,
+        1,
+        2,
+        4,
+        1,
+        1,
+        1,
+        1
+      ],
+      "intermediate_dim": 2048,
+      "idim": 24,
+      "hdim": 512,
+      "head": {
+        "idim": 512,
+        "hdim": 2048,
+        "odim": 512,
+        "ksz": 3
+      }
+    }
+  },
+  "dp": {
+    "latent_dim": 24,
+    "chunk_compress_factor": 6,
+    "normalizer": {
+      "scale": 1.0
+    },
+    "sentence_encoder": {
+      "char_emb_dim": 64,
+      "text_embedder": {
+        "char_emb_dim": 64
+      },
+      "convnext": {
+        "idim": 64,
+        "ksz": 5,
+        "intermediate_dim": 256,
+        "num_layers": 6,
+        "dilation_lst": [
+          1,
+          1,
+          1,
+          1,
+          1,
+          1
+        ]
+      },
+      "attn_encoder": {
+        "hidden_channels": 64,
+        "filter_channels": 256,
+        "n_heads": 2,
+        "n_layers": 2,
+        "p_dropout": 0.0
+      },
+      "proj_out": {
+        "idim": 64,
+        "odim": 64
+      }
+    },
+    "style_encoder": {
+      "proj_in": {
+        "ldim": 24,
+        "chunk_compress_factor": 6,
+        "odim": 64
+      },
+      "convnext": {
+        "idim": 64,
+        "ksz": 5,
+        "intermediate_dim": 256,
+        "num_layers": 4,
+        "dilation_lst": [
+          1,
+          1,
+          1,
+          1
+        ]
+      },
+      "style_token_layer": {
+        "input_dim": 64,
+        "n_style": 8,
+        "style_key_dim": 0,
+        "style_value_dim": 16,
+        "prototype_dim": 64,
+        "n_units": 64,
+        "n_heads": 2
+      }
+    },
+    "predictor": {
+      "sentence_dim": 64,
+      "n_style": 8,
+      "style_dim": 16,
+      "hdim": 128,
+      "n_layer": 2
+    }
+  }
+}

lgtm/__init__.py

@@ -0,0 +1,14 @@
+"""LGTM text-to-speech. `LGTMTTS` = PyTorch backend, `LGTMOnnx` = ONNX Runtime backend (no torch needed)."""
+
+
+def __getattr__(name):  # lazy imports: the ONNX backend must not require torch
+    if name in ("LGTMTTS", "load_voice_style", "save_voice_style"):
+        from . import inference
+        return getattr(inference, name)
+    if name == "LGTM":
+        from .model import LGTM
+        return LGTM
+    if name == "LGTMOnnx":
+        from .onnx_inference import LGTMOnnx
+        return LGTMOnnx
+    raise AttributeError(name)

lgtm/cli.py

@@ -0,0 +1,34 @@
+"""Command line:  python -m lgtm.cli --text "Xin chào" --lang vi --voice F1 --out out.wav [--ref ref.wav] [--backend onnx]"""
+import argparse
+
+
+def main():
+    ap = argparse.ArgumentParser(description="LGTM text-to-speech")
+    ap.add_argument("--text", required=True)
+    ap.add_argument("--lang", default="en", help="en es pt fr de it sv vi ja ko id")
+    ap.add_argument("--voice", default="F1", help="preset (F1-F5, M1-M5) or voice .json")
+    ap.add_argument("--ref", default=None, help="reference .wav to clone the voice from (overrides --voice)")
+    ap.add_argument("--save_voice", default=None, help="save the cloned voice to this .json")
+    ap.add_argument("--out", default="out.wav")
+    ap.add_argument("--model", default=".", help="local model dir or Hugging Face repo id")
+    ap.add_argument("--backend", choices=["torch", "onnx"], default="torch")
+    ap.add_argument("--steps", type=int, default=8)
+    ap.add_argument("--speed", type=float, default=1.05)
+    ap.add_argument("--gpu", action="store_true", help="(onnx) use CUDAExecutionProvider")
+    a = ap.parse_args()
+    if a.backend == "onnx":
+        from .onnx_inference import LGTMOnnx, save_voice_style
+        tts = LGTMOnnx.from_pretrained(a.model, use_gpu=a.gpu)
+    else:
+        from .inference import LGTMTTS, save_voice_style
+        tts = LGTMTTS.from_pretrained(a.model)
+    voice = tts.clone_voice(a.ref) if a.ref else a.voice
+    if a.ref and a.save_voice:
+        save_voice_style(a.save_voice, voice)
+    wav = tts.synthesize(a.text, lang=a.lang, voice=voice, steps=a.steps, speed=a.speed)
+    tts.save_wav(wav, a.out)
+    print(f"wrote {a.out} ({len(wav) / 44100:.2f} s)")
+
+
+if __name__ == "__main__":
+    main()

lgtm/export_onnx.py

@@ -0,0 +1,170 @@
+"""Export LGTM to ONNX (5 graphs) and verify each against PyTorch.
+
+    python -m lgtm.export_onnx --model_dir . --out onnx
+
+Graphs (all with dynamic batch / length):
+  text_encoder.onnx        text_ids[B,T] int64, style_ttl[B,50,256], text_mask[B,1,T]      -> text_emb[B,256,T]
+  duration_predictor.onnx  text_ids[B,T], style_dp[B,8,16], text_mask[B,1,T]                -> duration[B] (seconds)
+  vector_estimator.onnx    noisy_latent[B,144,L], text_emb, style_ttl, latent_mask[B,1,L],
+                           text_mask, current_step[B], total_step[B]                        -> denoised_latent[B,144,L]
+  vocoder.onnx             latent[B,144,L]                                                  -> wav[B, 3072*L] (44.1 kHz)
+  voice_encoder.onnx       wav[1,N] (44.1 kHz, prepared reference)                          -> style_ttl[1,50,256], style_dp[1,8,16]
+"""
+import argparse
+import json
+import os
+
+import numpy as np
+import onnxruntime as ort
+import torch
+from safetensors.torch import load_file
+
+from .model import LGTM
+from .strip_onnx import strip
+
+
+class TextEncoder(torch.nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.m = m
+
+    def forward(self, text_ids, style_ttl, text_mask):
+        return self.m.ttl.encode_text(text_ids, text_mask, style_ttl)
+
+
+class DurationPredictor(torch.nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.m = m
+
+    def forward(self, text_ids, style_dp, text_mask):
+        return self.m.dp(text_ids, text_mask, style_dp)
+
+
+class VectorEstimator(torch.nn.Module):
+    def __init__(self, m, cfg_scale=4.0):
+        super().__init__()
+        self.m, self.cfg_scale = m, cfg_scale
+
+    def forward(self, noisy_latent, text_emb, style_ttl, latent_mask, text_mask, current_step, total_step):
+        return self.m.ttl.euler_step(noisy_latent, current_step, total_step, text_emb, text_mask, style_ttl, latent_mask, self.cfg_scale)
+
+
+class Vocoder(torch.nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.m = m
+
+    def forward(self, latent):
+        return self.m.ae.decode_ttl(latent)
+
+
+class VoiceEncoder(torch.nn.Module):
+    def __init__(self, m):
+        super().__init__()
+        self.m = m
+
+    def forward(self, wav):
+        lat = self.m.ae.encode_ttl(wav)
+        mask = torch.ones_like(lat[:, :1])
+        return self.m.ttl.style_encoder(lat, mask), self.m.dp.style_encoder(lat, mask)
+
+
+def export(module, args, path, input_names, output_names, dynamic_shapes):
+    prog = torch.onnx.export(module, args, path, input_names=input_names, output_names=output_names,
+                             dynamic_shapes=dynamic_shapes, dynamo=True, opset_version=20, optimize=True, external_data=False)
+    return prog
+
+
+def check(path, module, feeds, names, atol_rel=1e-3):
+    sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
+    outs = sess.run(None, {k: v.numpy() for k, v in feeds.items()})
+    with torch.no_grad():
+        ref = module(*feeds.values())
+    ref = ref if isinstance(ref, (tuple, list)) else (ref,)
+    worst = 0.0
+    for o, r, n in zip(outs, ref, names):
+        r = r.numpy()
+        rel = np.abs(o - r).max() / (np.abs(r).max() + 1e-9)
+        worst = max(worst, rel)
+        print(f"    {n:16s} shape {tuple(o.shape)}  max rel diff {rel:.2e}")
+    assert worst < atol_rel, f"{path}: mismatch {worst}"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model_dir", default=".")
+    ap.add_argument("--out", default="onnx")
+    args = ap.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+    cfg = json.load(open(os.path.join(args.model_dir, "config.json")))
+    m = LGTM(cfg)
+    m.load_state_dict(load_file(os.path.join(args.model_dir, "pytorch", "model.safetensors")))
+    m.eval()
+    torch.manual_seed(0)
+    B, T, L = 3, 57, 63  # B must not equal any model dim (e.g. 2 attention heads) or the exporter fixes it
+    D = torch.export.Dim
+    b, t, l, n = D("batch", min=1, max=64), D("text_len", min=2, max=1000), D("latent_len", min=2, max=1000), D("samples", min=15000, max=44100 * 60)
+
+    def masks(T_, L_, B_=B):
+        tm = torch.ones(B_, 1, T_)
+        tm[-1, :, T_ - 9:] = 0
+        lm = torch.ones(B_, 1, L_)
+        lm[-1, :, L_ - 11:] = 0
+        return tm, lm
+
+    ids = torch.randint(0, 8000, (B, T))
+    tm, lm = masks(T, L)
+    s_ttl = torch.nn.functional.normalize(torch.randn(B, 50, 256), dim=-1)
+    s_dp = torch.nn.functional.normalize(torch.randn(B, 8, 16), dim=-1)
+    with torch.no_grad():
+        te = m.ttl.encode_text(ids, tm, s_ttl)
+    x = torch.randn(B, 144, L) * lm
+
+    jobs = [
+        ("text_encoder", TextEncoder(m), {"text_ids": ids, "style_ttl": s_ttl, "text_mask": tm}, ["text_emb"],
+         {"text_ids": {0: b, 1: t}, "style_ttl": {0: b}, "text_mask": {0: b, 2: t}}),
+        ("duration_predictor", DurationPredictor(m), {"text_ids": ids, "style_dp": s_dp, "text_mask": tm}, ["duration"],
+         {"text_ids": {0: b, 1: t}, "style_dp": {0: b}, "text_mask": {0: b, 2: t}}),
+        ("vector_estimator", VectorEstimator(m),
+         {"noisy_latent": x, "text_emb": te, "style_ttl": s_ttl, "latent_mask": lm, "text_mask": tm,
+          "current_step": torch.full((B,), 3.0), "total_step": torch.full((B,), 8.0)}, ["denoised_latent"],
+         {"noisy_latent": {0: b, 2: l}, "text_emb": {0: b, 2: t}, "style_ttl": {0: b}, "latent_mask": {0: b, 2: l},
+          "text_mask": {0: b, 2: t}, "current_step": {0: b}, "total_step": {0: b}}),
+        ("vocoder", Vocoder(m), {"latent": x}, ["wav"], {"latent": {0: b, 2: l}}),
+        ("voice_encoder", VoiceEncoder(m), {"wav": torch.randn(1, 44100 * 4) * 0.05}, ["style_ttl", "style_dp"],
+         {"wav": {1: n}}),
+    ]
+    for name, mod, feeds, outs, dyn in jobs:
+        path = os.path.join(args.out, f"{name}.onnx")
+        print(f"exporting {name} ...", flush=True)
+        export(mod, tuple(feeds.values()), path, list(feeds), outs, dyn)
+        print(f"  check (export shapes):")
+        check(path, mod, feeds, outs)
+        # different lengths than at export time, to prove the graph is really dynamic
+        B2 = 1 if name != "vocoder" else 5  # a different batch size than at export time
+        if name in ("text_encoder", "duration_predictor"):
+            T2 = 23
+            tm2, _ = masks(T2, L, B2)
+            second = list(feeds)[1]
+            f2 = {"text_ids": torch.randint(0, 8000, (B2, T2)), second: feeds[second][:B2], "text_mask": tm2}
+        elif name == "vector_estimator":
+            B2, T2, L2 = 5, 31, 140
+            tm2, lm2 = masks(T2, L2, B2)
+            st2 = torch.nn.functional.normalize(torch.randn(B2, 50, 256), dim=-1)
+            with torch.no_grad():
+                te2 = m.ttl.encode_text(torch.randint(0, 8000, (B2, T2)), tm2, st2)
+            f2 = {"noisy_latent": torch.randn(B2, 144, L2) * lm2, "text_emb": te2, "style_ttl": st2, "latent_mask": lm2,
+                  "text_mask": tm2, "current_step": torch.zeros(B2), "total_step": torch.full((B2,), 8.0)}
+        elif name == "vocoder":
+            f2 = {"latent": torch.randn(B2, 144, 17) * 0.27}
+        else:
+            f2 = {"wav": torch.randn(1, 44100 * 9) * 0.05}
+        print(f"  check (new shapes):")
+        check(path, mod, f2, outs)
+        strip(path)  # drop exporter debug metadata (stack traces with local paths)
+        print(f"  {os.path.getsize(path) / 1e6:.1f} MB", flush=True)
+
+
+if __name__ == "__main__":
+    main()

lgtm/inference.py

@@ -0,0 +1,151 @@
+"""PyTorch inference for LGTM.
+
+    from lgtm import LGTMTTS
+    tts = LGTMTTS.from_pretrained("path/or/hf-repo")      # downloads from the Hub if needed
+    wav = tts.synthesize("Xin chào!", lang="vi", voice="F1")
+    tts.save_wav(wav, "out.wav")
+    voice = tts.clone_voice("reference.wav")               # zero-shot voice from 5-15 s of audio
+    wav = tts.synthesize("Hello there.", lang="en", voice=voice)
+"""
+import json
+import os
+import re
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import torchaudio
+from safetensors.torch import load_file
+
+from .model import LGTM
+from .text import AVAILABLE_LANGS, TextProcessor
+
+SAMPLE_RATE = 44100
+
+
+def load_voice_style(path, device="cpu"):
+    d = json.load(open(path))
+    ttl = torch.tensor(np.array(d["style_ttl"]["data"], np.float32).reshape(d["style_ttl"]["dims"]))
+    dp = torch.tensor(np.array(d["style_dp"]["data"], np.float32).reshape(d["style_dp"]["dims"]))
+    return ttl.to(device), dp.to(device)
+
+
+def save_voice_style(path, voice):
+    ttl, dp = voice
+    d = {"style_ttl": {"data": ttl.detach().float().cpu().numpy().tolist(), "dims": list(ttl.shape), "type": "float32"},
+         "style_dp": {"data": dp.detach().float().cpu().numpy().tolist(), "dims": list(dp.shape), "type": "float32"}}
+    json.dump(d, open(path, "w"))
+
+
+def _prepare_reference(w, top_db=40.0, pad_s=0.1, target_rms_db=-23.0, peak=0.95):
+    """Trim leading/trailing silence and normalise loudness (same as the training data)."""
+    hop = SAMPLE_RATE // 100
+    n = len(w) // hop
+    if n >= 3:
+        db = 10 * torch.log10(w[: n * hop].view(n, hop).pow(2).mean(1) + 1e-10)
+        active = torch.nonzero(db > db.max() - top_db).flatten()
+        if len(active):
+            pad = int(pad_s * SAMPLE_RATE)
+            w = w[max(int(active[0]) * hop - pad, 0): min((int(active[-1]) + 1) * hop + pad, len(w))]
+    w = w * (10 ** (target_rms_db / 20) / (w.pow(2).mean().sqrt() + 1e-6))
+    m = w.abs().max()
+    return w * (peak / m) if m > peak else w
+
+
+def split_text(text, max_len=300):
+    """Split long text into sentence chunks of at most max_len characters."""
+    chunks = []
+    for para in [p.strip() for p in re.split(r"\n\s*\n+", text.strip()) if p.strip()]:
+        cur = ""
+        for sent in re.split(r"(?<=[.!?。！？])\s+", para):
+            if cur and len(cur) + len(sent) + 1 > max_len:
+                chunks.append(cur)
+                cur = sent
+            else:
+                cur = f"{cur} {sent}".strip()
+        if cur:
+            chunks.append(cur)
+    return chunks
+
+
+class LGTMTTS:
+    def __init__(self, model_dir, device=None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        cfg = json.load(open(os.path.join(model_dir, "config.json")))
+        self.model = LGTM(cfg)
+        self.model.load_state_dict(load_file(os.path.join(model_dir, "pytorch", "model.safetensors")))
+        self.model.to(self.device).eval()
+        self.tp = TextProcessor(os.path.join(model_dir, "unicode_indexer.json"))
+        self.voice_dir = os.path.join(model_dir, "voice_styles")
+
+    @classmethod
+    def from_pretrained(cls, repo_or_dir, device=None):
+        if not os.path.isdir(repo_or_dir):
+            from huggingface_hub import snapshot_download
+
+            repo_or_dir = snapshot_download(repo_or_dir, allow_patterns=["config.json", "unicode_indexer.json", "voice_styles/*", "pytorch/*"])
+        return cls(repo_or_dir, device)
+
+    @property
+    def voices(self):
+        return sorted(f[:-5] for f in os.listdir(self.voice_dir) if f.endswith(".json"))
+
+    def _voice(self, voice):
+        if isinstance(voice, str):
+            path = voice if voice.endswith(".json") else os.path.join(self.voice_dir, f"{voice}.json")
+            return load_voice_style(path, self.device)
+        return voice[0].to(self.device), voice[1].to(self.device)
+
+    @torch.no_grad()
+    def clone_voice(self, wav_path, max_seconds=15.0):
+        """Voice style from a reference recording (5-15 s of clean speech works best)."""
+        wav, sr = sf.read(wav_path, dtype="float32")
+        if wav.ndim == 2:
+            wav = wav.mean(1)
+        w = torch.from_numpy(wav)
+        if sr != SAMPLE_RATE:
+            w = torchaudio.functional.resample(w, sr, SAMPLE_RATE)
+        w = _prepare_reference(w)[: int(max_seconds * SAMPLE_RATE)]
+        lat = self.model.ae.encode_ttl(w[None].to(self.device))
+        mask = torch.ones(1, 1, lat.shape[-1], device=self.device)
+        return self.model.ttl.style_encoder(lat, mask), self.model.dp.style_encoder(lat, mask)
+
+    @torch.no_grad()
+    def _synth_batch(self, texts, langs, voice, steps, speed, cfg_scale):
+        s_ttl, s_dp = voice
+        b = len(texts)
+        ids, mask = self.tp.batch(texts, langs, device=self.device)
+        s_ttl, s_dp = s_ttl.expand(b, -1, -1), s_dp.expand(b, -1, -1)
+        m = self.model
+        dur = m.dp(ids, mask, s_dp) / speed
+        text_emb = m.ttl.encode_text(ids, mask, s_ttl)
+        chunk = m.ae.hop * m.ae.ccf
+        wav_len = (dur * SAMPLE_RATE).long()
+        lat_len = (wav_len + chunk - 1) // chunk
+        L = int(lat_len.max())
+        lmask = (torch.arange(L, device=self.device)[None] < lat_len[:, None]).float().unsqueeze(1)
+        x = torch.randn(b, m.ae.ldim * m.ae.ccf, L, device=self.device) * lmask
+        total = torch.full((b,), float(steps), device=self.device)
+        for i in range(steps):
+            x = m.ttl.euler_step(x, torch.full((b,), float(i), device=self.device), total, text_emb, mask, s_ttl, lmask, cfg_scale)
+        wav = m.ae.decode_ttl(x)
+        return [wav[k, : wav_len[k]].float().cpu().numpy() for k in range(b)]
+
+    def synthesize(self, text, lang="en", voice="F1", steps=8, speed=1.05, cfg_scale=4.0, silence=0.3):
+        """Text -> 44.1 kHz float32 waveform. Long text is split into sentence chunks.
+        voice: preset name (F1..F5, M1..M5), a .json path, or the output of clone_voice()."""
+        if lang not in AVAILABLE_LANGS:
+            raise ValueError(f"unsupported language {lang!r}")
+        v = self._voice(voice)
+        chunks = split_text(text, 120 if lang in ("ja", "ko") else 300)
+        wavs = self._synth_batch(chunks, [lang] * len(chunks), v, steps, speed, cfg_scale)
+        gap = np.zeros(int(silence * SAMPLE_RATE), np.float32)
+        out = []
+        for i, w in enumerate(wavs):
+            out += [w] + ([gap] if i < len(wavs) - 1 else [])
+        return np.concatenate(out)
+
+    @staticmethod
+    def save_wav(wav, path):
+        sf.write(path, wav, SAMPLE_RATE)

lgtm/model.py

@@ -0,0 +1,441 @@
+"""LGTM text-to-speech model (44.1 kHz, 11 languages, zero-shot voice cloning).
+
+    LGTM
+      ttl   text-to-latent: text encoder, style encoder, flow-matching vector field
+      ae    speech autoencoder: encoder (audio -> latent) and decoder (latent -> 44.1 kHz audio)
+      dp    utterance-level duration predictor (+ its style encoder)
+
+Latents: raw AE latent (B, 24, T_f) at 44100/512 Hz is normalised and 6 frames are stacked into
+channels -> (B, 144, T_f / 6); the flow-matching model works in that space.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .modules import (
+    AttnEncoder,
+    ChannelLayerNorm,
+    CharEmbedder,
+    ConvNeXtStack,
+    Linear,
+    PaddedConv1d,
+    RotaryCrossAttention,
+    StyleAttention,
+    TimeEncoder,
+)
+
+N_VOCAB = 8322
+
+
+# ==========================================================================
+# Text-to-latent
+# ==========================================================================
+class TextEncoder(nn.Module):
+    """char emb -> ConvNeXt x6 -> relpos transformer x4, with a long residual."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.text_embedder = CharEmbedder(N_VOCAB, cfg["text_embedder"]["char_emb_dim"])
+        self.convnext = ConvNeXtStack(**cfg["convnext"])
+        self.attn_encoder = AttnEncoder(**cfg["attn_encoder"])
+        # cfg["proj_out"] is an identity projection.
+
+    def forward(self, text_ids, text_mask):
+        x = self.text_embedder(text_ids, text_mask)
+        x = self.convnext(x, text_mask)
+        x = self.attn_encoder(x, text_mask) + x
+        return x * text_mask
+
+
+class SpeechPromptedTextEncoder(nn.Module):
+    """Two cross-attentions from text to the 50 style tokens; both residuals
+    are taken w.r.t. the *original* text features (as in the graph)."""
+
+    def __init__(self, text_dim, style_dim, n_units, n_heads):
+        super().__init__()
+        self.attention1 = StyleAttention(text_dim, style_dim, style_dim, n_units, n_heads, text_dim)
+        self.attention2 = StyleAttention(text_dim, style_dim, style_dim, n_units, n_heads, text_dim)
+        self.norm = ChannelLayerNorm(text_dim)
+
+    def forward(self, x, text_mask, style_key, style_value):
+        xt = x.transpose(1, 2)
+        m = text_mask.transpose(1, 2)
+        h = self.attention1(xt, style_key, style_value, m) * m + xt
+        h = self.attention2(h, style_key, style_value, m) * m + xt
+        return self.norm.norm(h).transpose(1, 2) * text_mask
+
+
+class StyleTokenLayer(nn.Module):
+    """Reference encoder pooling -> n_style tokens.
+
+    
+    """
+
+    def __init__(self, input_dim, n_style, style_key_dim, style_value_dim, prototype_dim, n_units, n_heads):
+        super().__init__()
+        if style_key_dim > 0:
+            self.style_key = nn.Parameter(torch.randn(1, n_style, style_key_dim) * 0.02)
+        else:
+            self.style_key = None
+        self.prototype = nn.Parameter(torch.randn(1, n_style, prototype_dim) * 0.02)
+        self.attention = StyleAttention(prototype_dim, input_dim, input_dim, n_units, n_heads, style_value_dim)
+        # Output = normalize(base + delta): `base` is the mean voice style, the network predicts the
+        # voice-specific deviation.
+        self.base = nn.Parameter(torch.zeros(1, n_style, style_value_dim))
+        nn.init.zeros_(self.attention.out_fc.linear.weight)
+        nn.init.zeros_(self.attention.out_fc.linear.bias)
+
+    def forward(self, h, mask):
+        """h: (B, C, T) encoded reference, mask: (B, 1, T) -> (B, n_style, style_value_dim)."""
+        b = h.shape[0]
+        ht = h.transpose(1, 2)
+        q = self.prototype.expand(b, -1, -1)
+        # mask padded reference frames by pushing their keys to a neutral value
+        # and removing their values; StyleAttention has no key mask of its own.
+        return self._masked_attention(q, ht, mask.transpose(1, 2))
+
+    def _masked_attention(self, q, kv, kv_mask):
+        att = self.attention
+        qh = att._heads(att.W_query(q))
+        kh = torch.tanh(att._heads(att.W_key(kv)))
+        vh = att._heads(att.W_value(kv))
+        scores = torch.matmul(qh, kh.transpose(-1, -2)) / (att.n_units ** 0.5)
+        scores = scores.masked_fill(kv_mask.transpose(1, 2).unsqueeze(0) == 0, float("-inf"))
+        out = torch.matmul(F.softmax(scores, dim=-1), vh)
+        out = self.base + att.out_fc(torch.cat(out.unbind(0), dim=-1))
+        # voice styles (ttl 50x256 and dp 8x16) are per-token unit vectors
+        return F.normalize(out, dim=-1)
+
+
+class StyleEncoder(nn.Module):
+    """Reference latent (B,144,T) -> style tokens."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        p = cfg["proj_in"]
+        self.proj_in = PaddedConv1d(p["ldim"] * p["chunk_compress_factor"], p["odim"], 1)
+        self.convnext = ConvNeXtStack(**cfg["convnext"])
+        self.style_token_layer = StyleTokenLayer(**cfg["style_token_layer"])
+
+    def forward(self, latent, mask):
+        h = self.proj_in(latent) * mask
+        h = self.convnext(h, mask)
+        return self.style_token_layer(h, mask)
+
+
+class UncondMasker(nn.Module):
+    """Learned null tokens for classifier-free guidance."""
+
+    def __init__(self, text_dim, n_style, style_key_dim, style_value_dim, **_):
+        super().__init__()
+        self.text_special_token = nn.Parameter(torch.zeros(1, text_dim, 1))
+        self.style_key_special_token = nn.Parameter(torch.zeros(1, n_style, style_key_dim))
+        self.style_value_special_token = nn.Parameter(torch.zeros(1, n_style, style_value_dim))
+
+
+class TimeCondBlock(nn.Module):
+    def __init__(self, idim, time_dim):
+        super().__init__()
+        self.linear = Linear(time_dim, idim)
+
+    def forward(self, x, mask, t_emb):
+        return (x + self.linear(t_emb).unsqueeze(-1)) * mask
+
+
+class TextCondBlock(nn.Module):
+    def __init__(self, idim, text_dim, n_heads, n_units, rotary_base, rotary_scale, **_):
+        super().__init__()
+        self.attn = RotaryCrossAttention(idim, text_dim, n_units, n_heads, rotary_base, rotary_scale)
+        self.norm = ChannelLayerNorm(idim)
+
+    def forward(self, x, mask, text_emb, text_mask):
+        x = x * mask
+        y = self.attn(x.transpose(1, 2), text_emb.transpose(1, 2), mask.transpose(1, 2), text_mask.transpose(1, 2))
+        x = x + y.transpose(1, 2) * mask
+        return self.norm(x) * mask
+
+
+class StyleCondBlock(nn.Module):
+    def __init__(self, idim, style_dim, n_units=256, n_heads=2):
+        super().__init__()
+        self.attention = StyleAttention(idim, style_dim, style_dim, n_units, n_heads, idim)
+        self.norm = ChannelLayerNorm(idim)
+
+    def forward(self, x, mask, style_key, style_value):
+        x = x * mask
+        m = mask.transpose(1, 2)
+        y = self.attention(x.transpose(1, 2), style_key, style_value, m) * m
+        x = x + y.transpose(1, 2)
+        return self.norm(x) * mask
+
+
+class VectorField(nn.Module):
+    """Flow-matching velocity estimator on compressed latents (B, 144, T)."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        p = cfg["proj_in"]
+        ldim = p["ldim"] * p["chunk_compress_factor"]
+        self.proj_in = PaddedConv1d(ldim, p["odim"], 1, bias=False)
+        self.time_encoder = TimeEncoder(cfg["time_encoder"]["time_dim"], cfg["time_encoder"]["hdim"])
+        mb = cfg["main_blocks"]
+        blocks = []
+        for _ in range(mb["n_blocks"]):
+            blocks += [
+                ConvNeXtStack(**mb["convnext_0"]),
+                TimeCondBlock(**mb["time_cond_layer"]),
+                ConvNeXtStack(**mb["convnext_1"]),
+                TextCondBlock(**mb["text_cond_layer"]),
+                ConvNeXtStack(**mb["convnext_2"]),
+                StyleCondBlock(**mb["style_cond_layer"]),
+            ]
+        self.main_blocks = nn.ModuleList(blocks)
+        self.last_convnext = ConvNeXtStack(**cfg["last_convnext"])
+        self.proj_out = PaddedConv1d(p["odim"], ldim, 1, bias=False)
+
+    def forward(self, x, t, text_emb, text_mask, style_key, style_value, latent_mask):
+        t_emb = self.time_encoder(t)
+        h = self.proj_in(x) * latent_mask
+        for blk in self.main_blocks:
+            if isinstance(blk, ConvNeXtStack):
+                h = blk(h, latent_mask)
+            elif isinstance(blk, TimeCondBlock):
+                h = blk(h, latent_mask, t_emb)
+            elif isinstance(blk, TextCondBlock):
+                h = blk(h, latent_mask, text_emb, text_mask)
+            else:
+                h = blk(h, latent_mask, style_key, style_value)
+        h = self.last_convnext(h, latent_mask)
+        return self.proj_out(h) * latent_mask
+
+
+class TextToLatent(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.text_encoder = TextEncoder(cfg["text_encoder"])
+        self.style_encoder = StyleEncoder(cfg["style_encoder"])
+        s = cfg["speech_prompted_text_encoder"]
+        self.speech_prompted_text_encoder = SpeechPromptedTextEncoder(s["text_dim"], s["style_dim"], s["n_units"], s["n_heads"])
+        self.uncond_masker = UncondMasker(**cfg["uncond_masker"])
+        self.vector_field = VectorField(cfg["vector_field"])
+        self.sig_min = cfg["flow_matching"]["sig_min"]
+
+    @property
+    def style_key(self):
+        return self.style_encoder.style_token_layer.style_key
+
+    def encode_text(self, text_ids, text_mask, style_ttl):
+        """== text_encoder.onnx. Returns text_emb (B, 256, T)."""
+        key = self.style_key.expand(text_ids.shape[0], -1, -1)
+        x = self.text_encoder(text_ids, text_mask)
+        return self.speech_prompted_text_encoder(x, text_mask, key, style_ttl)
+
+    def velocity(self, x, t, text_emb, text_mask, style_ttl, latent_mask, uncond=False):
+        """Single (conditional or unconditional) velocity prediction."""
+        b = x.shape[0]
+        if uncond:
+            um = self.uncond_masker
+            text_emb = um.text_special_token.expand(b, -1, text_emb.shape[-1])
+            key = um.style_key_special_token.expand(b, -1, -1)
+            style_ttl = um.style_value_special_token.expand(b, -1, -1)
+        else:
+            key = self.style_key.expand(b, -1, -1)
+        return self.vector_field(x, t, text_emb, text_mask, key, style_ttl, latent_mask)
+
+    def cfg_velocity(self, x, t, text_emb, text_mask, style_ttl, latent_mask, cfg_scale=4.0):
+        """CFG as baked into vector_estimator.onnx: v_u + 4 (v_c - v_u) = 4 v_c - 3 v_u."""
+        xx = torch.cat([x, x])
+        tt = torch.cat([t, t])
+        te = torch.cat([text_emb, self.uncond_masker.text_special_token.expand_as(text_emb)])
+        tm = torch.cat([text_mask, text_mask])
+        lm = torch.cat([latent_mask, latent_mask])
+        b = x.shape[0]
+        key = torch.cat([self.style_key.expand(b, -1, -1), self.uncond_masker.style_key_special_token.expand(b, -1, -1)])
+        val = torch.cat([style_ttl, self.uncond_masker.style_value_special_token.expand(b, -1, -1)])
+        v = self.vector_field(xx, tt, te, tm, key, val, lm)
+        v_c, v_u = v[:b], v[b:]  # (slicing instead of chunk keeps the batch dim dynamic in ONNX)
+        return v_u + cfg_scale * (v_c - v_u)
+
+    def euler_step(self, x, step, total_step, text_emb, text_mask, style_ttl, latent_mask, cfg_scale=4.0):
+        """== vector_estimator.onnx (one Euler step on t = step / total_step)."""
+        t = step / total_step
+        v = self.cfg_velocity(x, t, text_emb, text_mask, style_ttl, latent_mask, cfg_scale)
+        return (x + v / total_step.view(-1, 1, 1)) * latent_mask
+
+
+# ==========================================================================
+# Speech autoencoder
+# ==========================================================================
+class LatentDecoder(nn.Module):
+    """== vocoder.onnx decoder: causal ConvNeXt, head outputs 512 samples/frame."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        h = cfg["hdim"]
+        self.embed = PaddedConv1d(cfg["idim"], h, cfg["ksz_init"], causal=True)
+        self.convnext = ConvNeXtStack(
+            h, cfg["ksz"], cfg["intermediate_dim"], cfg["num_layers"], cfg["dilation_lst"], causal=True, wrapped_dwconv=True
+        ).convnext
+        self.final_norm = nn.Module()
+        self.final_norm.norm = nn.BatchNorm1d(h, eps=1e-5)
+        hd = cfg["head"]
+        self.head = nn.Module()
+        self.head.layer1 = PaddedConv1d(hd["idim"], hd["hdim"], hd["ksz"], causal=True)
+        self.head.act = nn.PReLU(1)
+        self.head.layer2 = nn.Conv1d(hd["hdim"], hd["odim"], 1, bias=False)
+
+    def forward(self, z):
+        """z: raw latent (B, 24, T_f) -> wav (B, T_f * 512)."""
+        x = self.embed(z)
+        for blk in self.convnext:
+            x = blk(x)
+        x = self.final_norm.norm(x)
+        x = self.head.layer2(self.head.act(self.head.layer1(x)))  # (B, 512, T_f)
+        return x.transpose(1, 2).reshape(x.shape[0], -1)
+
+
+class SpecProcessor(nn.Module):
+    """log |STFT| (1025) ++ log mel (228) = 1253 features at hop 512.
+
+    Frames are left-aligned to the decoder: frame i covers samples ending at
+    512*(i+1) (left pad n_fft - hop), so frame count == ceil(len / 512).
+    """
+
+    def __init__(self, n_fft, win_length, hop_length, n_mels, sample_rate, eps, **_):
+        super().__init__()
+        import torchaudio
+
+        self.n_fft, self.hop, self.win = n_fft, hop_length, win_length
+        self.eps = eps
+        self.register_buffer("window", torch.hann_window(win_length), persistent=False)
+        fb = torchaudio.functional.melscale_fbanks(n_fft // 2 + 1, 0.0, sample_rate / 2, n_mels, sample_rate, norm="slaney", mel_scale="slaney")
+        self.register_buffer("mel_fb", fb, persistent=False)
+
+    def forward(self, wav):
+        n = wav.shape[-1]
+        n_frames = (n + self.hop - 1) // self.hop
+        wav = F.pad(wav, (self.n_fft - self.hop, n_frames * self.hop - n))
+        spec = torch.stft(wav, self.n_fft, self.hop, self.win, self.window, center=False, return_complex=True).abs()
+        mel = torch.matmul(spec.transpose(1, 2), self.mel_fb).transpose(1, 2)
+        return torch.cat([torch.log(spec + self.eps), torch.log(mel + self.eps)], dim=1)
+
+
+class LatentEncoder(nn.Module):
+    """wav @ 44.1 kHz -> raw latent (B, 24, T_f)."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.spec_processor = SpecProcessor(**cfg["spec_processor"])
+        h = cfg["hdim"]
+        self.embed = PaddedConv1d(cfg["idim"], h, cfg["ksz_init"])
+        self.convnext = ConvNeXtStack(h, cfg["ksz"], cfg["intermediate_dim"], cfg["num_layers"], cfg["dilation_lst"], wrapped_dwconv=True).convnext
+        self.final_norm = ChannelLayerNorm(h)
+        self.proj_out = nn.Conv1d(h, cfg["odim"], 1)
+
+    def forward(self, wav):
+        x = self.embed(self.spec_processor(wav))
+        for blk in self.convnext:
+            x = blk(x)
+        return self.proj_out(self.final_norm(x))
+
+
+class SpeechAutoencoder(nn.Module):
+    def __init__(self, cfg, ttl_cfg):
+        super().__init__()
+        self.sample_rate = cfg["sample_rate"]
+        self.hop = cfg["base_chunk_size"]
+        self.ccf = ttl_cfg["chunk_compress_factor"]
+        self.ldim = cfg["ldim"]
+        self.register_buffer("latent_mean", torch.zeros(1, cfg["ldim"], 1))
+        self.register_buffer("latent_std", torch.ones(1, cfg["ldim"], 1))
+        self.register_buffer("normalizer_scale", torch.tensor(float(ttl_cfg["normalizer"]["scale"])))
+        self.encoder = LatentEncoder(cfg["encoder"])
+        self.decoder = LatentDecoder(cfg["decoder"])
+
+    # ---- latent <-> TTL representation -----------------------------------
+    def compress(self, z):
+        """raw (B, 24, T_f) -> TTL latent (B, 144, ceil(T_f/6)); pads with the latent mean."""
+        b, c, t = z.shape
+        zn = (z - self.latent_mean) / self.latent_std
+        pad = (-t) % self.ccf
+        if pad:
+            zn = F.pad(zn, (0, pad))
+        zn = zn.view(b, c, -1, self.ccf).permute(0, 1, 3, 2).reshape(b, c * self.ccf, -1)
+        return zn * self.normalizer_scale
+
+    def decompress(self, x):
+        """TTL latent (B, 144, T) -> raw (B, 24, 6T)."""
+        b, _, t = x.shape
+        x = x / self.normalizer_scale
+        z = x.view(b, self.ldim, self.ccf, t).permute(0, 1, 3, 2).reshape(b, self.ldim, t * self.ccf)
+        return z * self.latent_std + self.latent_mean
+
+    def decode_ttl(self, x):
+        """== vocoder.onnx: TTL latent (B, 144, T) -> wav (B, 3072 T)."""
+        return self.decoder(self.decompress(x))
+
+    def encode_ttl(self, wav):
+        """wav (B, N) @ 44.1 kHz -> TTL latent (B, 144, ceil(N/3072))."""
+        return self.compress(self.encoder(wav))
+
+
+# ==========================================================================
+# Duration predictor
+# ==========================================================================
+class SentenceEncoder(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        d = cfg["char_emb_dim"]
+        self.sentence_token = nn.Parameter(torch.randn(1, d, 1) * 0.02)
+        self.text_embedder = CharEmbedder(N_VOCAB, cfg["text_embedder"]["char_emb_dim"])
+        self.convnext = ConvNeXtStack(**cfg["convnext"])
+        self.attn_encoder = AttnEncoder(**cfg["attn_encoder"])
+        self.proj_out = PaddedConv1d(cfg["proj_out"]["idim"], cfg["proj_out"]["odim"], 1, bias=False)
+
+    def forward(self, text_ids, text_mask):
+        b = text_ids.shape[0]
+        x = self.text_embedder(text_ids, text_mask)
+        x = torch.cat([self.sentence_token.expand(b, -1, -1), x], dim=-1)
+        m = torch.cat([torch.ones_like(text_mask[:, :, :1]), text_mask], dim=-1)
+        x = self.convnext(x, m)
+        x = self.attn_encoder(x, m) + x
+        return (self.proj_out(x[:, :, :1]) * m[:, :, :1]).flatten(1)  # (B, 64)
+
+
+class DurationHead(nn.Module):
+    def __init__(self, sentence_dim, n_style, style_dim, hdim, n_layer):
+        super().__init__()
+        assert n_layer == 2
+        self.layers = nn.ModuleList([nn.Linear(sentence_dim + n_style * style_dim, hdim), nn.Linear(hdim, 1)])
+        self.activation = nn.PReLU(1)
+
+    def forward(self, s, style_dp):
+        h = torch.cat([s, style_dp.flatten(1)], dim=1)
+        return self.layers[1](self.activation(self.layers[0](h))).squeeze(1)  # log-seconds
+
+
+class DurationPredictor(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.sentence_encoder = SentenceEncoder(cfg["sentence_encoder"])
+        self.style_encoder = StyleEncoder(cfg["style_encoder"])  
+        self.predictor = DurationHead(**cfg["predictor"])
+
+    def log_duration(self, text_ids, text_mask, style_dp):
+        return self.predictor(self.sentence_encoder(text_ids, text_mask), style_dp)
+
+    def forward(self, text_ids, text_mask, style_dp):
+        """== duration_predictor.onnx: total utterance duration in seconds (B,)."""
+        return torch.exp(self.log_duration(text_ids, text_mask, style_dp))
+
+
+# ==========================================================================
+class LGTM(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.ttl = TextToLatent(cfg["ttl"])
+        self.ae = SpeechAutoencoder(cfg["ae"], cfg["ttl"])
+        self.dp = DurationPredictor(cfg["dp"])
+

lgtm/modules.py

@@ -0,0 +1,325 @@
+"""LGTM building blocks.
+
+Tensor layout: channels-first (B, C, T); masks are (B, 1, T) float tensors of 0/1.
+"""
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# --------------------------------------------------------------------------
+# Basic layers
+# --------------------------------------------------------------------------
+class ChannelLayerNorm(nn.Module):
+    """LayerNorm over channels of a (B, C, T) tensor. ONNX eps = 1e-6."""
+
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, eps=eps)
+
+    def forward(self, x):
+        return self.norm(x.transpose(1, 2)).transpose(1, 2)
+
+
+class Linear(nn.Module):
+    """nn.Linear wrapped as ``.linear`` to match original names (``W_query.linear.weight``)."""
+
+    def __init__(self, idim, odim, bias=True):
+        super().__init__()
+        self.linear = nn.Linear(idim, odim, bias=bias)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class PaddedConv1d(nn.Module):
+    """Conv1d with replicate ('edge') padding, wrapped as ``.net`` like the original.
+
+    causal=True pads (k-1)*d on the left only (autoencoder decoder),
+    otherwise (k-1)*d/2 on both sides (text encoder / vector field).
+    """
+
+    def __init__(self, idim, odim, ksz, dilation=1, groups=1, bias=True, causal=False):
+        super().__init__()
+        self.net = nn.Conv1d(idim, odim, ksz, dilation=dilation, groups=groups, bias=bias)
+        total = (ksz - 1) * dilation
+        self.pad = (total, 0) if causal else (total // 2, total - total // 2)
+
+    def forward(self, x):
+        if self.pad != (0, 0):
+            x = F.pad(x, self.pad, mode="replicate")
+        return self.net(x)
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt-1D block: dwconv -> LN -> pw(4x) -> GELU(erf) -> pw -> gamma, residual.
+
+    With a mask: input, dwconv output and block output are multiplied by the
+    mask (exactly as in the text encoder / vector-field graphs).  The decoder
+    uses no mask and causal padding.
+    """
+
+    def __init__(self, dim, intermediate_dim, ksz, dilation=1, causal=False, wrapped_dwconv=False):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.full((1, dim, 1), 1e-6))
+        dw = PaddedConv1d(dim, dim, ksz, dilation=dilation, groups=dim, causal=causal)
+        if wrapped_dwconv:
+            self.dwconv = dw  # params: dwconv.net.{weight,bias}
+        else:
+            # params: dwconv.{weight,bias}; keep padding logic in the block.
+            self.dwconv = dw.net
+            self._pad = dw.pad
+        self.wrapped = wrapped_dwconv
+        self.norm = ChannelLayerNorm(dim)
+        self.pwconv1 = nn.Conv1d(dim, intermediate_dim, 1)
+        self.act = nn.GELU()  # exact erf GELU, as in the graph
+        self.pwconv2 = nn.Conv1d(intermediate_dim, dim, 1)
+
+    def forward(self, x, mask=None):
+        if mask is not None:
+            x = x * mask
+        residual = x
+        if self.wrapped:
+            y = self.dwconv(x)
+        else:
+            y = self.dwconv(F.pad(x, self._pad, mode="replicate"))
+        if mask is not None:
+            y = y * mask
+        y = self.norm(y)
+        y = self.pwconv2(self.act(self.pwconv1(y)))
+        x = residual + self.gamma * y
+        if mask is not None:
+            x = x * mask
+        return x
+
+
+class ConvNeXtStack(nn.Module):
+    """Stack of ConvNeXt blocks, params at ``convnext.{i}.*``."""
+
+    def __init__(self, idim, ksz, intermediate_dim, num_layers, dilation_lst, causal=False, wrapped_dwconv=False, **_):
+        super().__init__()
+        assert len(dilation_lst) == num_layers
+        self.convnext = nn.ModuleList(
+            [
+                ConvNeXtBlock(idim, intermediate_dim, ksz, d, causal=causal, wrapped_dwconv=wrapped_dwconv)
+                for d in dilation_lst
+            ]
+        )
+
+    def forward(self, x, mask=None):
+        for blk in self.convnext:
+            x = blk(x, mask)
+        return x
+
+
+# --------------------------------------------------------------------------
+# VITS-style relative-position self-attention encoder (text / sentence enc.)
+# --------------------------------------------------------------------------
+class RelPosMultiHeadAttention(nn.Module):
+    """VITS MultiHeadAttention with shared relative position embeddings (window 4)."""
+
+    def __init__(self, channels, n_heads, window_size=4):
+        super().__init__()
+        assert channels % n_heads == 0
+        self.n_heads = n_heads
+        self.k_channels = channels // n_heads
+        self.window_size = window_size
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, channels, 1)
+        std = self.k_channels ** -0.5
+        self.emb_rel_k = nn.Parameter(torch.randn(1, 2 * window_size + 1, self.k_channels) * std)
+        self.emb_rel_v = nn.Parameter(torch.randn(1, 2 * window_size + 1, self.k_channels) * std)
+
+    def forward(self, x, attn_mask):
+        q, k, v = self.conv_q(x), self.conv_k(x), self.conv_v(x)
+        b, d, t = k.shape
+        h, kc = self.n_heads, self.k_channels
+        query = q.view(b, h, kc, t).transpose(2, 3) / math.sqrt(kc)
+        key = k.view(b, h, kc, t).transpose(2, 3)
+        value = v.view(b, h, kc, t).transpose(2, 3)
+
+        scores = torch.matmul(query, key.transpose(-2, -1))
+        key_rel = self._get_relative_embeddings(self.emb_rel_k, t)
+        rel_logits = torch.matmul(query, key_rel.unsqueeze(0).transpose(-2, -1))
+        scores = scores + self._relative_to_absolute(rel_logits)
+        scores = scores.masked_fill(attn_mask == 0, -1e4)
+        p = F.softmax(scores, dim=-1)
+        out = torch.matmul(p, value)
+        value_rel = self._get_relative_embeddings(self.emb_rel_v, t)
+        out = out + torch.matmul(self._absolute_to_relative(p), value_rel.unsqueeze(0))
+        out = out.transpose(2, 3).contiguous().view(b, d, t)
+        return self.conv_o(out)
+
+    def _get_relative_embeddings(self, emb, length):
+        w = self.window_size
+        pad_length = max(length - (w + 1), 0)
+        start = max((w + 1) - length, 0)
+        emb = F.pad(emb, (0, 0, pad_length, pad_length, 0, 0))  # pad of 0 is a no-op (keeps export branch-free)
+        return emb[:, start: start + 2 * length - 1]
+
+    @staticmethod
+    def _relative_to_absolute(x):
+        b, h, l, _ = x.shape
+        x = F.pad(x, (0, 1))
+        x = x.view(b, h, l * 2 * l)
+        x = F.pad(x, (0, l - 1))
+        return x.view(b, h, l + 1, 2 * l - 1)[:, :, :l, l - 1:]
+
+    @staticmethod
+    def _absolute_to_relative(x):
+        b, h, l, _ = x.shape
+        x = F.pad(x, (0, l - 1))
+        x = x.view(b, h, l * l + l * (l - 1))
+        x = F.pad(x, (l, 0))
+        return x.view(b, h, l, 2 * l)[:, :, :, 1:]
+
+
+class FFN(nn.Module):
+    def __init__(self, channels, filter_channels):
+        super().__init__()
+        self.conv_1 = nn.Conv1d(channels, filter_channels, 1)
+        self.conv_2 = nn.Conv1d(filter_channels, channels, 1)
+
+    def forward(self, x, mask):
+        x = torch.relu(self.conv_1(x * mask))
+        return self.conv_2(x * mask) * mask
+
+
+class AttnEncoder(nn.Module):
+    """VITS post-norm transformer encoder."""
+
+    def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, p_dropout=0.0):
+        super().__init__()
+        self.attn_layers = nn.ModuleList([RelPosMultiHeadAttention(hidden_channels, n_heads) for _ in range(n_layers)])
+        self.norm_layers_1 = nn.ModuleList([ChannelLayerNorm(hidden_channels) for _ in range(n_layers)])
+        self.ffn_layers = nn.ModuleList([FFN(hidden_channels, filter_channels) for _ in range(n_layers)])
+        self.norm_layers_2 = nn.ModuleList([ChannelLayerNorm(hidden_channels) for _ in range(n_layers)])
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x, mask):
+        attn_mask = mask.unsqueeze(2) * mask.unsqueeze(-1)
+        x = x * mask
+        for attn, n1, ffn, n2 in zip(self.attn_layers, self.norm_layers_1, self.ffn_layers, self.norm_layers_2):
+            x = n1(x + self.drop(attn(x, attn_mask)))
+            x = n2(x + self.drop(ffn(x, mask)))
+        return x * mask
+
+
+class CharEmbedder(nn.Module):
+    def __init__(self, n_vocab, dim):
+        super().__init__()
+        self.char_embedder = nn.Embedding(n_vocab, dim)
+
+    def forward(self, ids, mask):
+        return self.char_embedder(ids).transpose(1, 2) * mask
+
+
+# --------------------------------------------------------------------------
+# Style (GST-like) cross attention: keys go through tanh.
+# --------------------------------------------------------------------------
+class StyleAttention(nn.Module):
+    """Multi-head cross-attention to style tokens (used in the text encoder and
+    the vector field's style-conditioning layers).
+
+    Heads are formed by splitting the last dim and stacking on a new leading
+    axis; keys are passed through tanh; scores are divided by ``sqrt(n_units)``;
+    rows of padded queries are zeroed after softmax.
+    """
+
+    def __init__(self, q_dim, k_dim, v_dim, n_units, n_heads, out_dim):
+        super().__init__()
+        self.n_heads = n_heads
+        self.n_units = n_units
+        self.W_query = Linear(q_dim, n_units)
+        self.W_key = Linear(k_dim, n_units)
+        self.W_value = Linear(v_dim, n_units)
+        self.out_fc = Linear(n_units, out_dim)
+
+    def _heads(self, x):  # (B, T, U) -> (H, B, T, U/H)
+        return torch.stack(torch.chunk(x, self.n_heads, dim=-1), dim=0)
+
+    def forward(self, q, k, v, q_mask=None):
+        """q: (B, Tq, q_dim), k: (B, Tk, k_dim), v: (B, Tk, v_dim), q_mask: (B, Tq, 1)."""
+        q = self._heads(self.W_query(q))
+        k = torch.tanh(self._heads(self.W_key(k)))
+        v = self._heads(self.W_value(v))
+        p = F.softmax(torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.n_units), dim=-1)
+        if q_mask is not None:
+            p = p * q_mask.unsqueeze(0)
+        out = torch.matmul(p, v)  # (H, B, Tq, d)
+        out = torch.cat(out.unbind(0), dim=-1)
+        return self.out_fc(out)
+
+
+# --------------------------------------------------------------------------
+# Rotary cross attention (latent queries -> text keys), vector field.
+# --------------------------------------------------------------------------
+class RotaryCrossAttention(nn.Module):
+    """Text-conditioning attention in the vector field.
+
+    Rotary embedding uses *length-normalised* positions: pos = i / length,
+    angle = pos * theta, theta_j = rotary_scale * base^(-j/(d/2)).  Queries use
+    the latent length, keys the text length (so the attention learns a soft
+    monotonic alignment in relative position).  Non-interleaved rotation
+    (first half / second half).  Scores are divided by sqrt(n_units / 2).
+    """
+
+    def __init__(self, idim, text_dim, n_units, n_heads, rotary_base=10000, rotary_scale=10, **_):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = n_units // n_heads
+        self.scale = math.sqrt(n_units / 2)  # = 16 for n_units=512 (verified against ONNX)
+        self.W_query = Linear(idim, n_units)
+        self.W_key = Linear(text_dim, n_units)
+        self.W_value = Linear(text_dim, n_units)
+        self.out_fc = Linear(n_units, idim)
+        half = self.head_dim // 2
+        theta = rotary_scale * rotary_base ** (-torch.arange(half, dtype=torch.float32) / half)
+        self.register_buffer("theta", theta.view(1, 1, half))
+
+    def _heads(self, x):  # (B, T, U) -> (H, B, T, d)
+        b, t, _ = x.shape
+        return x.view(b, t, self.n_heads, self.head_dim).permute(2, 0, 1, 3)
+
+    def _rotate(self, x, mask):
+        # x: (H, B, T, d); mask: (B, T, 1)
+        t = x.shape[2]
+        length = mask.sum(dim=(1, 2)).view(-1, 1, 1)
+        pos = torch.arange(t, device=x.device, dtype=x.dtype).view(1, t, 1) / length
+        ang = pos * self.theta  # (B, T, d/2)
+        cos, sin = torch.cos(ang), torch.sin(ang)
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+    def forward(self, x, text, x_mask, text_mask):
+        """x: (B, Tq, idim), text: (B, Tk, text_dim), masks (B, T, 1)."""
+        q = self._rotate(self._heads(self.W_query(x)), x_mask)
+        k = self._rotate(self._heads(self.W_key(text)), text_mask)
+        v = self._heads(self.W_value(text))
+        scores = torch.matmul(q, k.transpose(-1, -2)) / self.scale
+        key_mask = text_mask.transpose(1, 2).unsqueeze(0)  # (1, B, 1, Tk)
+        scores = scores.masked_fill(key_mask == 0, float("-inf"))
+        p = F.softmax(scores, dim=-1) * x_mask.unsqueeze(0)
+        out = torch.matmul(p, v)  # (H, B, Tq, d)
+        b, tq = out.shape[1], out.shape[2]
+        out = out.permute(1, 2, 0, 3).reshape(b, tq, -1)
+        return self.out_fc(out)
+
+
+class TimeEncoder(nn.Module):
+    """sinusoidal(t * 1000) -> Linear -> Mish -> Linear."""
+
+    def __init__(self, time_dim, hdim):
+        super().__init__()
+        half = time_dim // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / (half - 1))
+        self.register_buffer("freqs", freqs.view(1, half))
+        self.mlp = nn.Sequential(Linear(time_dim, hdim), nn.Mish(), Linear(hdim, time_dim))
+
+    def forward(self, t):  # t: (B,) in [0, 1]
+        ang = t.view(-1, 1) * 1000.0 * self.freqs
+        return self.mlp(torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1))

lgtm/onnx_inference.py

@@ -0,0 +1,150 @@
+"""ONNX Runtime inference for LGTM (numpy + onnxruntime only).
+
+    from lgtm import LGTMOnnx
+    tts = LGTMOnnx.from_pretrained("path/or/hf-repo", use_gpu=False)
+    wav = tts.synthesize("Bonjour à tous.", lang="fr", voice="M1")
+    tts.save_wav(wav, "out.wav")
+    voice = tts.clone_voice("reference.wav")
+"""
+import json
+import os
+
+import numpy as np
+import onnxruntime as ort
+import soundfile as sf
+from scipy.signal import resample_poly
+
+from .text import AVAILABLE_LANGS, TextProcessor
+
+SAMPLE_RATE = 44100
+CHUNK = 512 * 6  # samples per latent frame
+
+
+def _split_text(text, max_len):
+    import re
+
+    chunks = []
+    for para in [p.strip() for p in re.split(r"\n\s*\n+", text.strip()) if p.strip()]:
+        cur = ""
+        for sent in re.split(r"(?<=[.!?。！？])\s+", para):
+            if cur and len(cur) + len(sent) + 1 > max_len:
+                chunks.append(cur)
+                cur = sent
+            else:
+                cur = f"{cur} {sent}".strip()
+        if cur:
+            chunks.append(cur)
+    return chunks
+
+
+def _prepare_reference(w, top_db=40.0, pad_s=0.1, target_rms_db=-23.0, peak=0.95):
+    hop = SAMPLE_RATE // 100
+    n = len(w) // hop
+    if n >= 3:
+        db = 10 * np.log10((w[: n * hop].reshape(n, hop) ** 2).mean(1) + 1e-10)
+        active = np.nonzero(db > db.max() - top_db)[0]
+        if len(active):
+            pad = int(pad_s * SAMPLE_RATE)
+            w = w[max(active[0] * hop - pad, 0): min((active[-1] + 1) * hop + pad, len(w))]
+    w = w * (10 ** (target_rms_db / 20) / (np.sqrt((w ** 2).mean()) + 1e-6))
+    m = np.abs(w).max()
+    return (w * (peak / m) if m > peak else w).astype(np.float32)
+
+
+def load_voice_style(path):
+    d = json.load(open(path))
+    ttl = np.array(d["style_ttl"]["data"], np.float32).reshape(d["style_ttl"]["dims"])
+    dp = np.array(d["style_dp"]["data"], np.float32).reshape(d["style_dp"]["dims"])
+    return ttl, dp
+
+
+def save_voice_style(path, voice):
+    ttl, dp = voice
+    json.dump({"style_ttl": {"data": ttl.tolist(), "dims": list(ttl.shape), "type": "float32"},
+               "style_dp": {"data": dp.tolist(), "dims": list(dp.shape), "type": "float32"}}, open(path, "w"))
+
+
+class LGTMOnnx:
+    def __init__(self, model_dir, use_gpu=False, threads=None):
+        providers = (["CUDAExecutionProvider"] if use_gpu else []) + ["CPUExecutionProvider"]
+        so = ort.SessionOptions()
+        if threads:
+            so.intra_op_num_threads = threads
+        load = lambda n: ort.InferenceSession(os.path.join(model_dir, "onnx", f"{n}.onnx"), so, providers=providers)  # noqa: E731
+        self.text_encoder = load("text_encoder")
+        self.duration_predictor = load("duration_predictor")
+        self.vector_estimator = load("vector_estimator")
+        self.vocoder = load("vocoder")
+        self._voice_encoder_path = os.path.join(model_dir, "onnx", "voice_encoder.onnx")
+        self._voice_encoder = None
+        self._so, self._providers = so, providers
+        self.tp = TextProcessor(os.path.join(model_dir, "unicode_indexer.json"))
+        self.voice_dir = os.path.join(model_dir, "voice_styles")
+
+    @classmethod
+    def from_pretrained(cls, repo_or_dir, **kw):
+        if not os.path.isdir(repo_or_dir):
+            from huggingface_hub import snapshot_download
+
+            repo_or_dir = snapshot_download(repo_or_dir, allow_patterns=["unicode_indexer.json", "voice_styles/*", "onnx/*"])
+        return cls(repo_or_dir, **kw)
+
+    @property
+    def voices(self):
+        return sorted(f[:-5] for f in os.listdir(self.voice_dir) if f.endswith(".json"))
+
+    def _voice(self, voice):
+        if isinstance(voice, str):
+            return load_voice_style(voice if voice.endswith(".json") else os.path.join(self.voice_dir, f"{voice}.json"))
+        return voice
+
+    def clone_voice(self, wav_path, max_seconds=15.0):
+        """Voice style from a reference recording (5-15 s of clean speech works best)."""
+        if self._voice_encoder is None:
+            self._voice_encoder = ort.InferenceSession(self._voice_encoder_path, self._so, providers=self._providers)
+        w, sr = sf.read(wav_path, dtype="float32")
+        if w.ndim == 2:
+            w = w.mean(1)
+        if sr != SAMPLE_RATE:
+            g = np.gcd(sr, SAMPLE_RATE)
+            w = resample_poly(w, SAMPLE_RATE // g, sr // g).astype(np.float32)
+        w = _prepare_reference(w)[: int(max_seconds * SAMPLE_RATE)]
+        ttl, dp = self._voice_encoder.run(None, {"wav": w[None]})
+        return ttl, dp
+
+    def _synth_batch(self, texts, langs, voice, steps, speed, rng):
+        s_ttl, s_dp = voice
+        b = len(texts)
+        ids, mask = self.tp.batch_numpy(texts, langs)
+        s_ttl = np.repeat(s_ttl, b, axis=0) if s_ttl.shape[0] == 1 else s_ttl
+        s_dp = np.repeat(s_dp, b, axis=0) if s_dp.shape[0] == 1 else s_dp
+        dur = self.duration_predictor.run(None, {"text_ids": ids, "style_dp": s_dp, "text_mask": mask})[0] / speed
+        text_emb = self.text_encoder.run(None, {"text_ids": ids, "style_ttl": s_ttl, "text_mask": mask})[0]
+        wav_len = (dur * SAMPLE_RATE).astype(np.int64)
+        lat_len = (wav_len + CHUNK - 1) // CHUNK
+        L = int(lat_len.max())
+        lmask = (np.arange(L)[None] < lat_len[:, None]).astype(np.float32)[:, None]
+        x = rng.standard_normal((b, 144, L)).astype(np.float32) * lmask
+        total = np.full(b, steps, np.float32)
+        for i in range(steps):
+            x = self.vector_estimator.run(None, {"noisy_latent": x, "text_emb": text_emb, "style_ttl": s_ttl, "latent_mask": lmask,
+                                                 "text_mask": mask, "current_step": np.full(b, i, np.float32), "total_step": total})[0]
+        wav = self.vocoder.run(None, {"latent": x})[0]
+        return [wav[k, : wav_len[k]] for k in range(b)]
+
+    def synthesize(self, text, lang="en", voice="F1", steps=8, speed=1.05, silence=0.3, seed=None):
+        """Text -> 44.1 kHz float32 waveform. voice: preset name, .json path, or clone_voice() output."""
+        if lang not in AVAILABLE_LANGS:
+            raise ValueError(f"unsupported language {lang!r}")
+        rng = np.random.default_rng(seed)
+        chunks = _split_text(text, 120 if lang in ("ja", "ko") else 300)
+        wavs = self._synth_batch(chunks, [lang] * len(chunks), self._voice(voice), steps, speed, rng)
+        gap = np.zeros(int(silence * SAMPLE_RATE), np.float32)
+        out = []
+        for i, w in enumerate(wavs):
+            out += [w] + ([gap] if i < len(wavs) - 1 else [])
+        return np.concatenate(out)
+
+    @staticmethod
+    def save_wav(wav, path):
+        sf.write(path, wav, SAMPLE_RATE)

lgtm/strip_onnx.py

@@ -0,0 +1,43 @@
+"""Remove exporter debug metadata (stack traces with local file paths, doc strings) from ONNX files."""
+import sys
+
+import onnx
+
+
+def _clean_graph(g):
+    g.doc_string = ""
+    del g.metadata_props[:]
+    for vi in list(g.input) + list(g.output) + list(g.value_info):
+        vi.doc_string = ""
+        del vi.metadata_props[:]
+    for init in g.initializer:
+        init.doc_string = ""
+        del init.metadata_props[:]
+    for n in g.node:
+        n.doc_string = ""
+        del n.metadata_props[:]
+        for a in n.attribute:
+            if a.g is not None and a.HasField("g"):
+                _clean_graph(a.g)
+            for sg in a.graphs:
+                _clean_graph(sg)
+
+
+def strip(path):
+    m = onnx.load(path)
+    m.doc_string = ""
+    del m.metadata_props[:]
+    _clean_graph(m.graph)
+    for f in m.functions:
+        f.doc_string = ""
+        del f.metadata_props[:]
+        for n in f.node:
+            n.doc_string = ""
+            del n.metadata_props[:]
+    onnx.save(m, path)
+
+
+if __name__ == "__main__":
+    for p in sys.argv[1:]:
+        strip(p)
+        print("stripped", p)

lgtm/text.py

@@ -0,0 +1,81 @@
+"""Text frontend: NFKD normalisation, light cleanup, "<lang>text</lang>" tagging, and a
+code-point -> id lookup (unicode_indexer.json, 8322 ids)."""
+import json
+import re
+from unicodedata import normalize
+
+
+AVAILABLE_LANGS = ["en", "ko", "ja", "ar", "bg", "cs", "da", "de", "el", "es", "et", "fi", "fr", "hi", "hr", "hu", "id", "it", "lt", "lv", "nl", "pl", "pt", "ro", "ru", "sk", "sl", "sv", "tr", "uk", "vi", "na"]
+
+_EMOJI = re.compile(
+    "[\U0001f600-\U0001f64f\U0001f300-\U0001f5ff\U0001f680-\U0001f6ff\U0001f700-\U0001f77f"
+    "\U0001f780-\U0001f7ff\U0001f800-\U0001f8ff\U0001f900-\U0001f9ff\U0001fa00-\U0001fa6f"
+    "\U0001fa70-\U0001faff☀-⛿✀-➿\U0001f1e6-\U0001f1ff]+",
+    flags=re.UNICODE,
+)
+_REPL = {"–": "-", "‑": "-", "—": "-", "_": " ", "“": '"', "”": '"', "‘": "'", "’": "'",
+         "´": "'", "`": "'", "[": " ", "]": " ", "|": " ", "/": " ", "#": " ", "→": " ", "←": " "}
+_EXPR = {"@": " at ", "e.g.,": "for example, ", "i.e.,": "that is, "}
+
+
+def preprocess_text(text: str, lang: str) -> str:
+    text = normalize("NFKD", text)
+    text = _EMOJI.sub("", text)
+    for k, v in _REPL.items():
+        text = text.replace(k, v)
+    text = re.sub(r"[♥☆♡©\\]", "", text)
+    for k, v in _EXPR.items():
+        text = text.replace(k, v)
+    for p in [",", r"\.", "!", r"\?", ";", ":", "'"]:
+        text = re.sub(r" " + p, p.replace("\\", ""), text)
+    while '""' in text:
+        text = text.replace('""', '"')
+    while "''" in text:
+        text = text.replace("''", "'")
+    while "``" in text:
+        text = text.replace("``", "`")
+    text = re.sub(r"\s+", " ", text).strip()
+    if not re.search(r"[.!?;:,'\"')\]}…。」』】〉》›»]$", text):
+        text += "."
+    if lang not in AVAILABLE_LANGS:
+        raise ValueError(f"Invalid language: {lang}")
+    return f"<{lang}>" + text + f"</{lang}>"
+
+
+class TextProcessor:
+    def __init__(self, unicode_indexer_path: str):
+        self.indexer_path = unicode_indexer_path
+        with open(unicode_indexer_path) as f:
+            self.indexer = json.load(f)  # list: code point -> id (-1 = unknown)
+
+    def encode(self, text: str, lang: str, preprocess: bool = True) -> list[int]:
+        s = preprocess_text(text, lang) if preprocess else text
+        return [self.indexer[ord(c)] if ord(c) < len(self.indexer) else -1 for c in s]
+
+    def unknown_chars(self, text: str, lang: str) -> set[str]:
+        s = preprocess_text(text, lang)
+        return {c for c in s if ord(c) >= len(self.indexer) or self.indexer[ord(c)] < 0}
+
+    def batch_numpy(self, texts, langs):
+        """ids (B, T) int64 and mask (B, 1, T) float32 as numpy arrays."""
+        import numpy as np
+
+        seqs = [self.encode(t, l) for t, l in zip(texts, langs)]
+        T = max(len(s) for s in seqs)
+        ids = np.zeros((len(seqs), T), np.int64)
+        mask = np.zeros((len(seqs), 1, T), np.float32)
+        for i, s in enumerate(seqs):
+            ids[i, : len(s)] = s
+            mask[i, :, : len(s)] = 1
+        return ids, mask
+
+    def batch(self, texts: list[str], langs: list[str], device=None):
+        import torch
+
+        seqs = [self.encode(t, l) for t, l in zip(texts, langs)]
+        lengths = torch.tensor([len(s) for s in seqs])
+        ids = torch.zeros(len(seqs), int(lengths.max()), dtype=torch.long)
+        for i, s in enumerate(seqs):
+            ids[i, : len(s)] = torch.tensor(s)
+        mask = (torch.arange(ids.shape[1])[None] < lengths[:, None]).float().unsqueeze(1)
+        return ids.to(device), mask.to(device)

onnx/duration_predictor.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6fbdf338c3768273d92ab1d14ed80f1db7214ac49e2377d1e81d6240bf8a8486
+size 3512757

onnx/text_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58e20a45e3a52d6755b35ffff1501fd990ae4cfd8f7431cf5ca561e422efe69b
+size 36088450

onnx/vector_estimator.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:829a06f43e524bf082658b5dd5d97f7e8e7082fa551fa80b829faf8127e64ba2
+size 256228310

onnx/vocoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96c75b381f26e701df5b9849304890310a9ebe7d8f585da8e5ec60bf3f8e8567
+size 101387918

onnx/voice_encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:594b3c8710a7e2f73ddbd1122500f6219a1aa7dd5394cb6f22dacc454f219a71
+size 117859547

pytorch/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcb429cd3bee3f37574bb47abaf0f9154de31adcc2f8570a248f011c59cf8c2e
+size 513767008

requirements-onnx.txt

@@ -0,0 +1,5 @@
+onnxruntime>=1.17   # or onnxruntime-gpu
+numpy
+soundfile
+scipy
+huggingface_hub

requirements.txt

@@ -0,0 +1,6 @@
+torch>=2.1
+torchaudio>=2.1
+numpy
+safetensors
+soundfile
+huggingface_hub