Swar 0.1.0: stage E release with default voice bijaya

LICENSE

@@ -0,0 +1,6 @@
+Swar weights (swar.safetensors, voices/) and Swar code: Creative Commons Attribution-NonCommercial
+4.0 International (CC BY-NC 4.0), https://creativecommons.org/licenses/by-nc/4.0/legalcode
+
+swar/kokoro/: vendored from Kokoro-82M (hexgrad), Apache License 2.0 -- see swar/kokoro/LICENSE.
+Style encoder (swar/model.py: LearnedDownSample, ResBlk, StyleEncoder): from StyleTTS2
+(yl4579/StyleTTS2), MIT License.

README.md

@@ -0,0 +1,152 @@
+---
+license: cc-by-nc-4.0
+language:
+  - ne
+  - en
+pipeline_tag: text-to-speech
+library_name: pytorch
+tags:
+  - text-to-speech
+  - tts
+  - nepali
+  - styletts2
+  - kokoro
+  - customer-service
+---
+
+# Swar
+
+**Swar** (स्वर, "voice") is a small, experimental text-to-speech model for **Nepali and
+English**, including sentences that mix the two, tuned toward a calm customer-service manner
+of speaking.
+
+- **38.5 M parameters**, 16 kHz output. The architecture is derived from
+  [Kokoro-82M](https://huggingface.co/hexgrad/Kokoro-82M) (StyleTTS2 family), built from its layer
+  classes: non-autoregressive, so a sentence is generated in one pass.
+- Reads **Devanagari and Latin letters directly**: no phonemiser, no language switch.
+  Numbers, currency ("Rs. 250", "रु ५००"), times and phone numbers are written out in words in
+  the language of their sentence.
+- **Default voice: `bijaya`** (`voices/bijaya.pt`). Any 3-10 s clip of one speaker gives
+  another voice.
+- Fast: about **15x real time on a desktop CPU** with all cores, about 6x on a single thread,
+  and about 30x on a consumer GPU (measured on a Ryzen 9 5900X and an RTX 3090).
+
+Try it in the Space: [Nepali Call Center Experimental - Powered by Swar](https://huggingface.co/spaces/elitexp/Nepali-Call-Center-Experimental).
+
+## Quick start
+
+```bash
+pip install torch torchaudio transformers safetensors huggingface_hub soundfile
+```
+
+```python
+import sys, soundfile as sf
+from huggingface_hub import snapshot_download
+
+folder = snapshot_download("elitexp/Swar")
+sys.path.insert(0, folder)
+from swar import Swar
+
+tts = Swar.from_pretrained(folder)                       # GPU if available, else CPU
+wav, sr = tts.synthesize("नमस्ते, ग्राहक सेवा केन्द्रमा यहाँलाई स्वागत छ।")   # voice "bijaya"
+sf.write("hello.wav", wav, sr)
+
+# another voice from a reference clip, a little slower, with clear pauses at punctuation
+wav, sr = tts.synthesize("Your balance is Rs. 250. धन्यवाद।",
+                         voice="my_reference.wav", speed=0.95, pause_floor_ms=300)
+```
+
+From the command line, inside the downloaded folder:
+
+```bash
+python -m swar.inference "नमस्ते, यहाँलाई स्वागत छ।" -o hello.wav
+```
+
+### Options
+
+| argument | meaning |
+|---|---|
+| `voice` | a name from `voices/` (default `bijaya`), a `.pt` of style vectors, or an audio clip (wav / flac / mp3 / ogg) |
+| `speed` | 0.5-2.0, scales the predicted durations (1.0 = as predicted) |
+| `pause_floor_ms` | every comma / full stop / question mark *inside* a sentence lasts at least this long (0 = as predicted) |
+| `return_info` | also return timing and the exact letters the model read |
+
+## What is inside
+
+| part | params | role |
+|---|---|---|
+| PL-BERT (ALBERT, letters) + projection | 3.5 M | contextual letter features |
+| prosody predictor | 9.7 M | duration per letter, pitch and energy at 80 Hz |
+| prosodic style encoder | 3.5 M | manner of speaking from a reference clip |
+| text encoder | 3.2 M | letter features for the decoder |
+| decoder (iSTFTNet with harmonic source) | 15.1 M | waveform at 16 kHz |
+| acoustic style encoder | 3.5 M | timbre from a reference clip |
+
+A voice is two 128-d vectors: acoustic (timbre, read by the decoder) and prosodic (manner,
+read by the predictor). `voices/bijaya.pt` stores them as `{"s_ac", "s_pr"}`.
+
+Files: `swar.safetensors` (inference weights only), `config.json`, `swar/` (model,
+text front-end, mel, inference; Kokoro's layer code is vendored under `swar/kokoro/` with its
+Apache-2.0 licence), `voices/`.
+
+## Training
+
+Trained from scratch in five stages, after the StyleTTS2 / Kokoro recipe adapted to letters at
+16 kHz:
+
+1. **PL-BERT** on Nepali and English text (masked letters).
+2. **Letter aligner** (CTC) on both languages, for durations.
+3. **Acoustic model**: text encoder, decoder and both style encoders, adversarial training with
+   multi-period and multi-resolution discriminators (300 k steps).
+4. **Prosody**: duration, pitch and energy predictors and the prosodic style encoder (100 k steps).
+5. **Register**: a short fine-tune (6 k steps) of the prosody half on about 8 hours of
+   customer-service style Nepali recordings, with general data replayed so the rest of the model
+   does not drift.
+
+Data (acoustic and prosody stages): about 800 hours of Nepali speech from public corpora
+(including OpenSLR 43 and OpenSLR 143) and other licensed and in-house recordings, part of it
+synthetic, and about 570 hours of English speech synthesised with Kokoro-82M (`af_heart`).
+
+## Evaluation
+
+- **Copy-synthesis intelligibility** (stage 3, durations and pitch taken from the real clip,
+  recognised by the same ASR as the real recordings): Nepali character error rate through a
+  G.711 telephone channel 0.199 against 0.220 for the real recordings (285 clips); English word
+  error rate 0.097 against 0.098 (200 clips).
+- **Customer-service register** (stage 5): median pitch stays within 15 Hz of the reference
+  voice. The pause rate (about 0.13 pauses per second, about 0.20 with `pause_floor_ms=300`)
+  is still below the 0.30-0.45 target band: pauses are only lengthened where the text has
+  punctuation.
+- Full text-to-speech intelligibility has not yet been measured with the same protocol.
+
+## Limitations
+
+- Experimental. Rare words, names and some English spellings can be mispronounced; Nepali
+  rhythm is less polished than English.
+- 16 kHz, telephone-oriented; not a studio-quality voice.
+- Voices from a reference clip copy that clip's recording conditions too; use clean audio.
+
+## Responsible use
+
+Swar can speak in a voice taken from a short recording. Only use voices you have permission to
+use, label synthetic speech as synthetic where people could mistake it for a person, and do not
+use the model to impersonate anyone or to deceive.
+
+## Acknowledgements
+
+Swar would not exist without **[Kokoro-82M](https://huggingface.co/hexgrad/Kokoro-82M)** by
+hexgrad. Its architecture is Swar's blueprint: the text encoder, the prosody predictor (duration,
+pitch and energy), the iSTFTNet decoder with its harmonic source and the split acoustic / prosodic
+style design all come from Kokoro, and Kokoro's own layer code is vendored in `swar/kokoro/`. Swar
+re-sizes that design for 16 kHz, reads letters instead of phonemes, adds trainable style encoders
+and is trained from scratch; Kokoro (voice `af_heart`) also generated Swar's English training
+speech. Kokoro in turn builds on [StyleTTS2](https://github.com/yl4579/StyleTTS2) (Li et al.), whose
+style encoder and PL-BERT recipe Swar also uses.
+
+## License
+
+The Swar weights (`swar.safetensors`, `voices/`) and Swar's own code are released under
+**[CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/)**: free for research and other
+non-commercial use, with attribution. The vendored Kokoro code in `swar/kokoro/` stays under its
+original Apache-2.0 licence (`swar/kokoro/LICENSE`); the StyleTTS2 style encoder code is MIT.
+Training data includes [OpenSLR](https://www.openslr.org/) corpora under their own licences.

config.json

@@ -0,0 +1,64 @@
+{
+  "sample_rate": 16000,
+  "n_token": 256,
+  "plbert": {
+    "hidden_size": 512,
+    "num_attention_heads": 8,
+    "intermediate_size": 2048,
+    "max_position_embeddings": 512,
+    "num_hidden_layers": 12,
+    "dropout": 0.1
+  },
+  "hidden_dim": 384,
+  "text_encoder_kernel_size": 5,
+  "n_layer": 3,
+  "style_dim": 128,
+  "max_dur": 50,
+  "dropout": 0.2,
+  "decoder": {
+    "hid": 512,
+    "out": 256,
+    "res": 64
+  },
+  "istftnet": {
+    "upsample_rates": [
+      10,
+      4
+    ],
+    "upsample_kernel_sizes": [
+      20,
+      8
+    ],
+    "upsample_initial_channel": 256,
+    "resblock_kernel_sizes": [
+      3,
+      7,
+      11
+    ],
+    "resblock_dilation_sizes": [
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ]
+    ],
+    "gen_istft_n_fft": 20,
+    "gen_istft_hop_size": 5
+  },
+  "style_encoder": {
+    "dim_in": 32,
+    "max_conv_dim": 256
+  },
+  "name": "swar",
+  "version": "0.1.0"
+}

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.1
+torchaudio>=2.1
+transformers>=4.34
+safetensors>=0.4
+huggingface_hub>=0.17
+soundfile>=0.12
+numpy

swar.safetensors

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

swar/__init__.py

@@ -0,0 +1,5 @@
+"""Swar: a small Nepali + English text-to-speech model (Kokoro / StyleTTS2 family)."""
+from .inference import Swar
+
+__all__ = ["Swar"]
+__version__ = "0.1.0"

swar/inference.py

@@ -0,0 +1,175 @@
+"""Swar inference: text (Nepali, English or mixed) -> 16 kHz speech.
+
+    from swar import Swar
+    tts = Swar.from_pretrained("elitexp/Swar")          # or a local folder
+    wav, sr = tts.synthesize("नमस्ते, यहाँलाई स्वागत छ।")  # default voice: bijaya
+    wav, sr = tts.synthesize("Your balance is Rs. 250.", voice="path/to/reference.wav",
+                             speed=1.0, pause_floor_ms=300)
+
+    python -m swar.inference "नमस्ते, यहाँलाई स्वागत छ।" -o hello.wav
+
+A voice is two 128-d style vectors: acoustic (timbre, read by the decoder) and
+prosodic (manner of speaking, read by the duration / pitch / energy predictor).
+voices/*.pt hold them ready-made; any clean 3-10 s clip of one speaker gives a new
+voice through style_from_audio(). pause_floor_ms makes every comma / full stop /
+question mark inside a sentence last at least that long.
+"""
+from __future__ import annotations
+
+import argparse
+import re
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .mel import SR, Mel
+from .model import build, conf
+from .text import STOI, normalize, to_ids
+
+PAUSE_CHARS = ",.?!;:…—"
+GAP_S = 0.15                                                       # silence between sentences
+MAX_CHARS = 220
+
+
+def _alignment(dur: torch.Tensor, T: int) -> torch.Tensor:
+    end = dur.cumsum(1)
+    start = end - dur
+    t = torch.arange(T, device=dur.device)[None, None]
+    return ((t >= start[..., None]) & (t < end[..., None])).float()
+
+
+def split_sentences(text: str) -> list[str]:
+    out = []
+    for s in re.split(r"(?<=[.?!])\s+", text):
+        s = s.strip()
+        while len(s) > MAX_CHARS:                                  # long run-on: cut at a comma or space
+            cut = max(s.rfind(",", 0, MAX_CHARS), s.rfind(" ", 0, MAX_CHARS))
+            cut = cut if cut > 0 else MAX_CHARS
+            out.append(s[: cut + 1].strip())
+            s = s[cut + 1:].strip()
+        if s:
+            out.append(s)
+    return out
+
+
+class Swar:
+    def __init__(self, folder: Path, device: str | None = None):
+        self.folder = Path(folder)
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.config = conf(self.folder / "config.json")
+        m = build("D", self.config)
+        from safetensors.torch import load_file
+        m.load_state_dict(load_file(str(self.folder / "swar.safetensors")), strict=True)
+        self.m = m.to(self.device).eval()
+        self.mel = Mel().to(self.device)
+        self._pause_ids = torch.tensor([STOI[c] for c in PAUSE_CHARS], device=self.device)
+        self._voices: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+
+    @classmethod
+    def from_pretrained(cls, repo_or_path: str | Path = "elitexp/Swar", device: str | None = None,
+                        revision: str | None = None) -> "Swar":
+        p = Path(repo_or_path)
+        if not p.is_dir():
+            from huggingface_hub import snapshot_download
+            p = Path(snapshot_download(str(repo_or_path), revision=revision))
+        return cls(p, device)
+
+    # ---- voices
+    def list_voices(self) -> list[str]:
+        return sorted(f.stem for f in (self.folder / "voices").glob("*.pt"))
+
+    @torch.no_grad()
+    def style_from_audio(self, path: str | Path) -> tuple[torch.Tensor, torch.Tensor]:
+        import soundfile as sf
+        import torchaudio.functional as AF
+        wav, sr = sf.read(str(path), dtype="float32", always_2d=True)
+        x = torch.from_numpy(wav.mean(1))
+        if sr != SR:
+            x = AF.resample(x, sr, SR)
+        mel = self.mel(x.to(self.device)[None])[:, None]
+        return self.m["style_encoder"](mel), self.m["predictor_encoder"](mel)
+
+    def voice(self, voice: str | Path | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """A name from voices/, a .pt of style vectors, or an audio clip."""
+        key = str(voice or "bijaya")
+        if key not in self._voices:
+            p = Path(key)
+            if p.suffix.lower() in (".wav", ".flac", ".mp3", ".ogg"):
+                self._voices[key] = self.style_from_audio(p)
+            else:
+                f = p if p.suffix == ".pt" else self.folder / "voices" / f"{key}.pt"
+                v = torch.load(f, map_location="cpu", weights_only=True)
+                self._voices[key] = (v["s_ac"].to(self.device), v["s_pr"].to(self.device))
+        return self._voices[key]
+
+    # ---- synthesis
+    @torch.no_grad()
+    def _tts(self, ids: list[int], s_ac, s_pr, speed: float, pause_floor: int) -> torch.Tensor:
+        m, dev = self.m, self.device
+        x_ids = torch.tensor([ids], device=dev)
+        n = torch.tensor([len(ids)], device=dev)
+        mask = torch.zeros(1, len(ids), dtype=torch.bool, device=dev)
+        bert = m["plbert"](x_ids, attention_mask=(~mask).int()).last_hidden_state
+        d_en = m["bert_encoder"](bert).transpose(1, 2)
+        d = m["predictor"].text_encoder(d_en, s_pr, n.cpu(), mask)
+        h, _ = m["predictor"].lstm(d)
+        dur = torch.round(torch.sigmoid(m["predictor"].duration_proj(h)).sum(-1) / speed).clamp(min=1).long()
+        if pause_floor:
+            L = len(ids)
+            pos = torch.arange(L, device=dev)
+            p = torch.isin(x_ids[0], self._pause_ids) & (pos > 0) & (pos < L - 2)   # inside the sentence
+            dur[0, p] = dur[0, p].clamp(min=pause_floor)
+        aln = _alignment(dur, int(dur.sum()))
+        F0, N = m["predictor"].F0Ntrain(d.transpose(-1, -2) @ aln, s_pr)
+        asr = m["text_encoder"](x_ids, n.cpu(), mask) @ aln
+        return m["decoder"](asr, F0, N, s_ac).squeeze().clamp(-1, 1)
+
+    def synthesize(self, text: str, voice: str | Path | None = None, speed: float = 1.0,
+                   pause_floor_ms: int = 0, return_info: bool = False):
+        """-> (float32 waveform, 16000), plus a timing dict if return_info."""
+        t0 = time.time()
+        s_ac, s_pr = self.voice(voice)
+        floor = int(round(pause_floor_ms / 25))
+        speed = float(min(2.0, max(0.5, speed)))
+        pieces, sentences, first = [], [], None
+        # split the RAW text first, so each sentence's numbers are read in its own
+        # language ("Rs. 250" stays together: no split before a digit)
+        raws = re.split(r"(?<=[।?!])\s+|(?<=\.)\s+(?![0-9०-९])", text)
+        for sent in (s for raw in raws for s in split_sentences(normalize(raw))):
+            sentences.append(sent)
+            y = self._tts(to_ids(sent), s_ac, s_pr, speed, floor).float().cpu().numpy()
+            first = first if first is not None else time.time() - t0
+            pieces += [y, np.zeros(int(GAP_S * SR), np.float32)]
+        if not pieces:
+            raise ValueError("nothing speakable left after normalising the text")
+        wav = np.concatenate(pieces[:-1])
+        if not return_info:
+            return wav, SR
+        return wav, SR, {"first_audio_sec": round(first, 3), "total_sec": round(time.time() - t0, 3),
+                         "duration_sec": round(len(wav) / SR, 2), "sentences": sentences}
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description="Swar text-to-speech")
+    ap.add_argument("text")
+    ap.add_argument("-o", "--out", default="swar.wav")
+    ap.add_argument("--model", default=str(Path(__file__).resolve().parents[1]))
+    ap.add_argument("--voice", default=None, help="voice name, .pt, or a reference audio clip")
+    ap.add_argument("--speed", type=float, default=1.0)
+    ap.add_argument("--pause-floor-ms", type=int, default=0)
+    ap.add_argument("--device", default=None)
+    a = ap.parse_args()
+    import soundfile as sf
+    tts = Swar.from_pretrained(a.model, device=a.device)
+    wav, sr, info = tts.synthesize(a.text, voice=a.voice, speed=a.speed,
+                                   pause_floor_ms=a.pause_floor_ms, return_info=True)
+    sf.write(a.out, wav, sr)
+    print(f"{a.out}: {info['duration_sec']} s audio in {info['total_sec']} s -- read as: "
+          f"{' | '.join(info['sentences'])}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

swar/kokoro/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

swar/kokoro/__init__.py

@@ -0,0 +1,2 @@
+"""Layer classes from Kokoro-82M (hexgrad/kokoro, Apache-2.0; see LICENSE here),
+vendored unchanged except for a relative import and one unused import removed."""

swar/kokoro/custom_stft.py

@@ -0,0 +1,196 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CustomSTFT(nn.Module):
+    """
+    STFT/iSTFT without unfold/complex ops, using conv1d and conv_transpose1d.
+
+    - forward STFT => Real-part conv1d + Imag-part conv1d
+    - inverse STFT => Real-part conv_transpose1d + Imag-part conv_transpose1d + sum
+    - avoids F.unfold, so easier to export to ONNX
+    - uses replicate or constant padding for 'center=True' to approximate 'reflect' 
+      (reflect is not supported for dynamic shapes in ONNX)
+    """
+
+    def __init__(
+        self,
+        filter_length=800,
+        hop_length=200,
+        win_length=800,
+        window="hann",
+        center=True,
+        pad_mode="replicate",  # or 'constant'
+    ):
+        super().__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_fft = filter_length
+        self.center = center
+        self.pad_mode = pad_mode
+
+        # Number of frequency bins for real-valued STFT with onesided=True
+        self.freq_bins = self.n_fft // 2 + 1
+
+        # Build window
+        assert window == 'hann', window
+        window_tensor = torch.hann_window(win_length, periodic=True, dtype=torch.float32)
+        if self.win_length < self.n_fft:
+            # Zero-pad up to n_fft
+            extra = self.n_fft - self.win_length
+            window_tensor = F.pad(window_tensor, (0, extra))
+        elif self.win_length > self.n_fft:
+            window_tensor = window_tensor[: self.n_fft]
+        self.register_buffer("window", window_tensor)
+
+        # Precompute forward DFT (real, imag)
+        # PyTorch stft uses e^{-j 2 pi k n / N} => real=cos(...), imag=-sin(...)
+        n = np.arange(self.n_fft)
+        k = np.arange(self.freq_bins)
+        angle = 2 * np.pi * np.outer(k, n) / self.n_fft  # shape (freq_bins, n_fft)
+        dft_real = np.cos(angle)
+        dft_imag = -np.sin(angle)  # note negative sign
+
+        # Combine window and dft => shape (freq_bins, filter_length)
+        # We'll make 2 conv weight tensors of shape (freq_bins, 1, filter_length).
+        forward_window = window_tensor.numpy()  # shape (n_fft,)
+        forward_real = dft_real * forward_window  # (freq_bins, n_fft)
+        forward_imag = dft_imag * forward_window
+
+        # Convert to PyTorch
+        forward_real_torch = torch.from_numpy(forward_real).float()
+        forward_imag_torch = torch.from_numpy(forward_imag).float()
+
+        # Register as Conv1d weight => (out_channels, in_channels, kernel_size)
+        # out_channels = freq_bins, in_channels=1, kernel_size=n_fft
+        self.register_buffer(
+            "weight_forward_real", forward_real_torch.unsqueeze(1)
+        )
+        self.register_buffer(
+            "weight_forward_imag", forward_imag_torch.unsqueeze(1)
+        )
+
+        # Precompute inverse DFT
+        # Real iFFT formula => scale = 1/n_fft, doubling for bins 1..freq_bins-2 if n_fft even, etc.
+        # For simplicity, we won't do the "DC/nyquist not doubled" approach here. 
+        # If you want perfect real iSTFT, you can add that logic. 
+        # This version just yields good approximate reconstruction with Hann + typical overlap.
+        inv_scale = 1.0 / self.n_fft
+        n = np.arange(self.n_fft)
+        angle_t = 2 * np.pi * np.outer(n, k) / self.n_fft  # shape (n_fft, freq_bins)
+        idft_cos = np.cos(angle_t).T  # => (freq_bins, n_fft)
+        idft_sin = np.sin(angle_t).T  # => (freq_bins, n_fft)
+
+        # Multiply by window again for typical overlap-add
+        # We also incorporate the scale factor 1/n_fft
+        inv_window = window_tensor.numpy() * inv_scale
+        backward_real = idft_cos * inv_window  # (freq_bins, n_fft)
+        backward_imag = idft_sin * inv_window
+
+        # We'll implement iSTFT as real+imag conv_transpose with stride=hop.
+        self.register_buffer(
+            "weight_backward_real", torch.from_numpy(backward_real).float().unsqueeze(1)
+        )
+        self.register_buffer(
+            "weight_backward_imag", torch.from_numpy(backward_imag).float().unsqueeze(1)
+        )
+        
+
+
+    def transform(self, waveform: torch.Tensor):
+        """
+        Forward STFT => returns magnitude, phase
+        Output shape => (batch, freq_bins, frames)
+        """
+        # waveform shape => (B, T).  conv1d expects (B, 1, T).
+        # Optional center pad
+        if self.center:
+            pad_len = self.n_fft // 2
+            waveform = F.pad(waveform, (pad_len, pad_len), mode=self.pad_mode)
+
+        x = waveform.unsqueeze(1)  # => (B, 1, T)
+        # Convolution to get real part => shape (B, freq_bins, frames)
+        real_out = F.conv1d(
+            x,
+            self.weight_forward_real,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        # Imag part
+        imag_out = F.conv1d(
+            x,
+            self.weight_forward_imag,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        # magnitude, phase
+        magnitude = torch.sqrt(real_out**2 + imag_out**2 + 1e-14)
+        phase = torch.atan2(imag_out, real_out)
+        # Handle the case where imag_out is 0 and real_out is negative to correct ONNX atan2 to match PyTorch
+        # In this case, PyTorch returns pi, ONNX returns -pi
+        correction_mask = (imag_out == 0) & (real_out < 0)
+        phase[correction_mask] = torch.pi
+        return magnitude, phase
+
+
+    def inverse(self, magnitude: torch.Tensor, phase: torch.Tensor, length=None):
+        """
+        Inverse STFT => returns waveform shape (B, T).
+        """
+        # magnitude, phase => (B, freq_bins, frames)
+        # Re-create real/imag => shape (B, freq_bins, frames)
+        real_part = magnitude * torch.cos(phase)
+        imag_part = magnitude * torch.sin(phase)
+
+        # conv_transpose wants shape (B, freq_bins, frames). We'll treat "frames" as time dimension
+        # so we do (B, freq_bins, frames) => (B, freq_bins, frames)
+        # But PyTorch conv_transpose1d expects (B, in_channels, input_length)
+        real_part = real_part  # (B, freq_bins, frames)
+        imag_part = imag_part
+
+        # real iSTFT => convolve with "backward_real", "backward_imag", and sum
+        # We'll do 2 conv_transpose calls, each giving (B, 1, time),
+        # then add them => (B, 1, time).
+        real_rec = F.conv_transpose1d(
+            real_part,
+            self.weight_backward_real,  # shape (freq_bins, 1, filter_length)
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        imag_rec = F.conv_transpose1d(
+            imag_part,
+            self.weight_backward_imag,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        # sum => (B, 1, time)
+        waveform = real_rec - imag_rec  # typical real iFFT has minus for imaginary part
+
+        # If we used "center=True" in forward, we should remove pad
+        if self.center:
+            pad_len = self.n_fft // 2
+            # Because of transposed convolution, total length might have extra samples
+            # We remove `pad_len` from start & end if possible
+            waveform = waveform[..., pad_len:-pad_len]
+
+        # If a specific length is desired, clamp
+        if length is not None:
+            waveform = waveform[..., :length]
+
+        # shape => (B, T)
+        return waveform
+
+    def forward(self, x: torch.Tensor):
+        """
+        Full STFT -> iSTFT pass: returns time-domain reconstruction.
+        Same interface as your original code.
+        """
+        mag, phase = self.transform(x)
+        return self.inverse(mag, phase, length=x.shape[-1])

swar/kokoro/istftnet.py

@@ -0,0 +1,421 @@
+# ADAPTED from https://github.com/yl4579/StyleTTS2/blob/main/Modules/istftnet.py
+from .custom_stft import CustomSTFT
+from torch.nn.utils import weight_norm
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# https://github.com/yl4579/StyleTTS2/blob/main/Modules/utils.py
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+class AdaIN1d(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        # affine should be False, however there's a bug in the old torch.onnx.export (not newer dynamo) that causes the channel dimension to be lost if affine=False. When affine is true, there's additional learnably parameters. This shouldn't really matter setting it to True, since we're in inference mode
+        self.norm = nn.InstanceNorm1d(num_features, affine=True)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+
+class AdaINResBlock1(nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
+        super(AdaINResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                                  padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                                  padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                                  padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+        self.convs2 = nn.ModuleList([
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+        self.adain1 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        self.adain2 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        self.alpha1 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
+        self.alpha2 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])
+
+    def forward(self, x, s):
+        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
+            xt = n1(x, s)
+            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
+            xt = c1(xt)
+            xt = n2(xt, s)
+            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+
+class TorchSTFT(nn.Module):
+    def __init__(self, filter_length=800, hop_length=200, win_length=800, window='hann'):
+        super().__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        assert window == 'hann', window
+        self.window = torch.hann_window(win_length, periodic=True, dtype=torch.float32)
+
+    def transform(self, input_data):
+        forward_transform = torch.stft(
+            input_data,
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(input_data.device),
+            return_complex=True)
+        return torch.abs(forward_transform), torch.angle(forward_transform)
+
+    def inverse(self, magnitude, phase):
+        inverse_transform = torch.istft(
+            magnitude * torch.exp(phase * 1j),
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device))
+        return inverse_transform.unsqueeze(-2)  # unsqueeze to stay consistent with conv_transpose1d implementation
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+
+
+class SineGen(nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(torch.pi) or cos(0)
+    """
+    def __init__(self, samp_rate, upsample_scale, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0,
+                 flag_for_pulse=False):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+        self.flag_for_pulse = flag_for_pulse
+        self.upsample_scale = upsample_scale
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = (f0 > self.voiced_threshold).type(torch.float32)
+        return uv
+
+    def _f02sine(self, f0_values):
+        """ f0_values: (batchsize, length, dim)
+            where dim indicates fundamental tone and overtones
+        """
+        # convert to F0 in rad. The interger part n can be ignored
+        # because 2 * torch.pi * n doesn't affect phase
+        rad_values = (f0_values / self.sampling_rate) % 1
+        # initial phase noise (no noise for fundamental component)
+        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device)
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
+        if not self.flag_for_pulse:
+            rad_values = F.interpolate(rad_values.transpose(1, 2), scale_factor=1/self.upsample_scale, mode="linear").transpose(1, 2)
+            phase = torch.cumsum(rad_values, dim=1) * 2 * torch.pi
+            phase = F.interpolate(phase.transpose(1, 2) * self.upsample_scale, scale_factor=self.upsample_scale, mode="linear").transpose(1, 2)
+            sines = torch.sin(phase)
+        else:
+            # If necessary, make sure that the first time step of every
+            # voiced segments is sin(pi) or cos(0)
+            # This is used for pulse-train generation
+            # identify the last time step in unvoiced segments
+            uv = self._f02uv(f0_values)
+            uv_1 = torch.roll(uv, shifts=-1, dims=1)
+            uv_1[:, -1, :] = 1
+            u_loc = (uv < 1) * (uv_1 > 0)
+            # get the instantanouse phase
+            tmp_cumsum = torch.cumsum(rad_values, dim=1)
+            # different batch needs to be processed differently
+            for idx in range(f0_values.shape[0]):
+                temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
+                temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
+                # stores the accumulation of i.phase within
+                # each voiced segments
+                tmp_cumsum[idx, :, :] = 0
+                tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
+            # rad_values - tmp_cumsum: remove the accumulation of i.phase
+            # within the previous voiced segment.
+            i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
+            # get the sines
+            sines = torch.cos(i_phase * 2 * torch.pi)
+        return sines
+
+    def forward(self, f0):
+        """ sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device)
+        # fundamental component
+        fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
+        # generate sine waveforms
+        sine_waves = self._f02sine(fn) * self.sine_amp
+        # generate uv signal
+        # uv = torch.ones(f0.shape)
+        # uv = uv * (f0 > self.voiced_threshold)
+        uv = self._f02uv(f0)
+        # noise: for unvoiced should be similar to sine_amp
+        #        std = self.sine_amp/3 -> max value ~ self.sine_amp
+        #        for voiced regions is self.noise_std
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+        # first: set the unvoiced part to 0 by uv
+        # then: additive noise
+        sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0):
+        super(SourceModuleHnNSF, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, upsample_scale, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshod)
+        # to merge source harmonics into a single excitation
+        self.l_linear = nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = nn.Tanh()
+
+    def forward(self, x):
+        """
+        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+        F0_sampled (batchsize, length, 1)
+        Sine_source (batchsize, length, 1)
+        noise_source (batchsize, length 1)
+        """
+        # source for harmonic branch
+        with torch.no_grad():
+            sine_wavs, uv, _ = self.l_sin_gen(x)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+        # source for noise branch, in the same shape as uv
+        noise = torch.randn_like(uv) * self.sine_amp / 3
+        return sine_merge, noise, uv
+
+
+class Generator(nn.Module):
+    def __init__(self, style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, disable_complex=False):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.m_source = SourceModuleHnNSF(
+                    sampling_rate=24000,
+                    upsample_scale=math.prod(upsample_rates) * gen_istft_hop_size,
+                    harmonic_num=8, voiced_threshod=10)
+        self.f0_upsamp = nn.Upsample(scale_factor=math.prod(upsample_rates) * gen_istft_hop_size)
+        self.noise_convs = nn.ModuleList()
+        self.noise_res = nn.ModuleList()
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(weight_norm(
+                nn.ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
+                                   k, u, padding=(k-u)//2)))
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel//(2**(i+1))
+            for j, (k, d) in enumerate(zip(resblock_kernel_sizes,resblock_dilation_sizes)):
+                self.resblocks.append(AdaINResBlock1(ch, k, d, style_dim))
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+            if i + 1 < len(upsample_rates):
+                stride_f0 = math.prod(upsample_rates[i + 1:])
+                self.noise_convs.append(nn.Conv1d(
+                    gen_istft_n_fft + 2, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
+                self.noise_res.append(AdaINResBlock1(c_cur, 7, [1,3,5], style_dim))
+            else:
+                self.noise_convs.append(nn.Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
+                self.noise_res.append(AdaINResBlock1(c_cur, 11, [1,3,5], style_dim))
+        self.post_n_fft = gen_istft_n_fft
+        self.conv_post = weight_norm(nn.Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.reflection_pad = nn.ReflectionPad1d((1, 0))
+        self.stft = (
+            CustomSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
+            if disable_complex
+            else TorchSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
+        )
+
+    def forward(self, x, s, f0):
+        with torch.no_grad():
+            f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+            har_source, noi_source, uv = self.m_source(f0)
+            har_source = har_source.transpose(1, 2).squeeze(1)
+            har_spec, har_phase = self.stft.transform(har_source)
+            har = torch.cat([har_spec, har_phase], dim=1)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, negative_slope=0.1) 
+            x_source = self.noise_convs[i](har)
+            x_source = self.noise_res[i](x_source, s)
+            x = self.ups[i](x)
+            if i == self.num_upsamples - 1:
+                x = self.reflection_pad(x)
+            x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x, s)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x, s)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
+        return self.stft.inverse(spec, phase)
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        else:
+            return F.interpolate(x, scale_factor=2, mode='nearest')
+
+
+class AdainResBlk1d(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2), upsample='none', dropout_p=0.0):
+        super().__init__()
+        self.actv = actv
+        self.upsample_type = upsample
+        self.upsample = UpSample1d(upsample)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out, style_dim)
+        self.dropout = nn.Dropout(dropout_p)
+        if upsample == 'none':
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
+
+    def _build_weights(self, dim_in, dim_out, style_dim):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
+        self.norm1 = AdaIN1d(style_dim, dim_in)
+        self.norm2 = AdaIN1d(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        x = self.upsample(x)
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        x = self.pool(x)
+        x = self.conv1(self.dropout(x))
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(self.dropout(x))
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) * torch.rsqrt(torch.tensor(2))
+        return out
+
+
+class Decoder(nn.Module):
+    def __init__(self, dim_in, style_dim, dim_out, 
+                 resblock_kernel_sizes,
+                 upsample_rates,
+                 upsample_initial_channel,
+                 resblock_dilation_sizes,
+                 upsample_kernel_sizes,
+                 gen_istft_n_fft, gen_istft_hop_size,
+                 disable_complex=False):
+        super().__init__()
+        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
+        self.decode = nn.ModuleList()
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 512, style_dim, upsample=True))
+        self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        self.asr_res = nn.Sequential(weight_norm(nn.Conv1d(512, 64, kernel_size=1)))
+        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates, 
+                                   upsample_initial_channel, resblock_dilation_sizes, 
+                                   upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, disable_complex=disable_complex)
+
+    def forward(self, asr, F0_curve, N, s):
+        F0 = self.F0_conv(F0_curve.unsqueeze(1))
+        N = self.N_conv(N.unsqueeze(1))
+        x = torch.cat([asr, F0, N], axis=1)
+        x = self.encode(x, s)
+        asr_res = self.asr_res(asr)
+        res = True
+        for block in self.decode:
+            if res:
+                x = torch.cat([x, asr_res, F0, N], axis=1)
+            x = block(x, s)
+            if block.upsample_type != "none":
+                res = False
+        x = self.generator(x, s, F0_curve)
+        return x

swar/kokoro/modules.py

@@ -0,0 +1,183 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/models.py
+from .istftnet import AdainResBlk1d
+from torch.nn.utils import weight_norm
+from transformers import AlbertModel
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LinearNorm(nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
+        super(LinearNorm, self).__init__()
+        self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias)
+        nn.init.xavier_uniform_(self.linear_layer.weight, gain=nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)):
+        super().__init__()
+        self.embedding = nn.Embedding(n_symbols, channels)
+        padding = (kernel_size - 1) // 2
+        self.cnn = nn.ModuleList()
+        for _ in range(depth):
+            self.cnn.append(nn.Sequential(
+                weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
+                LayerNorm(channels),
+                actv,
+                nn.Dropout(0.2),
+            ))
+        self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True)
+
+    def forward(self, x, input_lengths, m):
+        x = self.embedding(x)  # [B, T, emb]
+        x = x.transpose(1, 2)  # [B, emb, T]
+        m = m.unsqueeze(1)
+        x.masked_fill_(m, 0.0)
+        for c in self.cnn:
+            x = c(x)
+            x.masked_fill_(m, 0.0)
+        x = x.transpose(1, 2)  # [B, T, chn]
+        lengths = input_lengths if input_lengths.device == torch.device('cpu') else input_lengths.to('cpu')
+        x = nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+        x = x.transpose(-1, -2)
+        x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device)
+        x_pad[:, :, :x.shape[-1]] = x
+        x = x_pad
+        x.masked_fill_(m, 0.0)
+        return x
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, style_dim, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+        self.fc = nn.Linear(style_dim, channels*2)
+
+    def forward(self, x, s):
+        x = x.transpose(-1, -2)
+        x = x.transpose(1, -1)
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), eps=self.eps)
+        x = (1 + gamma) * x + beta
+        return x.transpose(1, -1).transpose(-1, -2)
+
+
+class ProsodyPredictor(nn.Module):
+    def __init__(self, style_dim, d_hid, nlayers, max_dur=50, dropout=0.1):
+        super().__init__()
+        self.text_encoder = DurationEncoder(sty_dim=style_dim, d_model=d_hid,nlayers=nlayers, dropout=dropout)
+        self.lstm = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
+        self.duration_proj = LinearNorm(d_hid, max_dur)
+        self.shared = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
+        self.F0 = nn.ModuleList()
+        self.F0.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.F0.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
+        self.F0.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
+        self.N = nn.ModuleList()
+        self.N.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.N.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
+        self.N.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
+        self.F0_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+        self.N_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+
+    def forward(self, texts, style, text_lengths, alignment, m):
+        d = self.text_encoder(texts, style, text_lengths, m)
+        m = m.unsqueeze(1)
+        lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu')
+        x = nn.utils.rnn.pack_padded_sequence(d, lengths, batch_first=True, enforce_sorted=False)
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+        x_pad = torch.zeros([x.shape[0], m.shape[-1], x.shape[-1]], device=x.device)
+        x_pad[:, :x.shape[1], :] = x
+        x = x_pad
+        duration = self.duration_proj(nn.functional.dropout(x, 0.5, training=False))
+        en = (d.transpose(-1, -2) @ alignment)
+        return duration.squeeze(-1), en
+
+    def F0Ntrain(self, x, s):
+        x, _ = self.shared(x.transpose(-1, -2))
+        F0 = x.transpose(-1, -2)
+        for block in self.F0:
+            F0 = block(F0, s)
+        F0 = self.F0_proj(F0)
+        N = x.transpose(-1, -2)
+        for block in self.N:
+            N = block(N, s)
+        N = self.N_proj(N)
+        return F0.squeeze(1), N.squeeze(1)
+
+
+class DurationEncoder(nn.Module):
+    def __init__(self, sty_dim, d_model, nlayers, dropout=0.1):
+        super().__init__()
+        self.lstms = nn.ModuleList()
+        for _ in range(nlayers):
+            self.lstms.append(nn.LSTM(d_model + sty_dim, d_model // 2, num_layers=1, batch_first=True, bidirectional=True, dropout=dropout))
+            self.lstms.append(AdaLayerNorm(sty_dim, d_model))
+        self.dropout = dropout
+        self.d_model = d_model
+        self.sty_dim = sty_dim
+
+    def forward(self, x, style, text_lengths, m):
+        masks = m
+        x = x.permute(2, 0, 1)
+        s = style.expand(x.shape[0], x.shape[1], -1)
+        x = torch.cat([x, s], axis=-1)
+        x.masked_fill_(masks.unsqueeze(-1).transpose(0, 1), 0.0)
+        x = x.transpose(0, 1)
+        x = x.transpose(-1, -2)
+        for block in self.lstms:
+            if isinstance(block, AdaLayerNorm):
+                x = block(x.transpose(-1, -2), style).transpose(-1, -2)
+                x = torch.cat([x, s.permute(1, 2, 0)], axis=1)
+                x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0)
+            else:
+                lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu')
+                x = x.transpose(-1, -2)
+                x = nn.utils.rnn.pack_padded_sequence(
+                    x, lengths, batch_first=True, enforce_sorted=False)
+                block.flatten_parameters()
+                x, _ = block(x)
+                x, _ = nn.utils.rnn.pad_packed_sequence(
+                    x, batch_first=True)
+                x = F.dropout(x, p=self.dropout, training=False)
+                x = x.transpose(-1, -2)
+                x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device)
+                x_pad[:, :, :x.shape[-1]] = x
+                x = x_pad
+
+        return x.transpose(-1, -2)
+
+
+# https://github.com/yl4579/StyleTTS2/blob/main/Utils/PLBERT/util.py
+class CustomAlbert(AlbertModel):
+    def forward(self, *args, **kwargs):
+        outputs = super().forward(*args, **kwargs)
+        return outputs.last_hidden_state

swar/mel.py

@@ -0,0 +1,28 @@
+"""16 kHz log-mel, exactly as Swar's style encoders were trained on it."""
+from __future__ import annotations
+
+import torch
+import torchaudio
+from torch import nn
+
+SR, N_FFT, WIN, HOP, N_MELS = 16000, 1024, 800, 200, 80
+MEAN, STD = -4.0, 4.0
+
+
+class Mel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mel = torchaudio.transforms.MelSpectrogram(
+            sample_rate=SR, n_fft=N_FFT, win_length=WIN, hop_length=HOP, n_mels=N_MELS,
+            f_min=0.0, f_max=8000.0, center=True, pad_mode="reflect", power=2.0)
+
+    def forward(self, wav: torch.Tensor) -> torch.Tensor:
+        if wav.dtype == torch.int16:
+            wav = wav.float() / 32768.0
+        if wav.dim() == 1:
+            wav = wav[None]
+        return (torch.log(1e-5 + self.mel(wav)) - MEAN) / STD
+
+    @staticmethod
+    def frames(n_samples: int) -> int:
+        return n_samples // HOP + 1

swar/model.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+"""The Swar network, built from Kokoro's own layer classes (swar/kokoro, Apache-2.0).
+
+    from swar.model import build
+    m = build("D", config)   # dict of nn.Modules, widths from config.json
+
+Stage C trains text_encoder + decoder + style_encoder (acoustic); stage D adds
+plbert + bert_encoder + predictor + predictor_encoder (prosodic style). At
+inference the two 128-d styles concatenate into Kokoro's 256-d ref_s layout:
+[:128] drives the decoder (timbre), [128:] drives the prosody predictor.
+
+Two deliberate departures from the Kokoro code, both because we run at 16 kHz:
+  * Kokoro's Generator hard-codes SourceModuleHnNSF(sampling_rate=24000). Its sine
+    source turns F0 into a waveform at that rate, so at 16 kHz every voice would
+    come out 1.5x too high. The source is rebuilt at our rate.
+  * Kokoro's Decoder hard-codes its widths (1024 / 512 / 64); they come from
+    config.json here (512 / 256 / 64).
+The style encoder is StyleTTS2's (Kokoro ships none -- it ships voice packs).
+"""
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.utils import spectral_norm, weight_norm
+
+from .kokoro.istftnet import AdainResBlk1d, Generator, SourceModuleHnNSF
+from .kokoro.modules import ProsodyPredictor, TextEncoder
+
+SR = 16000
+
+
+def conf(path: str | Path | None = None) -> dict:
+    p = Path(path) if path else Path(__file__).resolve().parents[1] / "config.json"
+    return json.loads(p.read_text())
+
+
+class Decoder(nn.Module):
+    """Kokoro's istftnet.Decoder: aligned text features + F0 + energy + style -> waveform."""
+
+    def __init__(self, dim_in, style_dim, hid, out, res, gen: dict, sr: int = SR):
+        super().__init__()
+        self.encode = AdainResBlk1d(dim_in + 2, hid, style_dim)
+        self.decode = nn.ModuleList([
+            AdainResBlk1d(hid + 2 + res, hid, style_dim),
+            AdainResBlk1d(hid + 2 + res, hid, style_dim),
+            AdainResBlk1d(hid + 2 + res, hid, style_dim),
+            AdainResBlk1d(hid + 2 + res, out, style_dim, upsample=True),
+        ])
+        self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        self.asr_res = nn.Sequential(weight_norm(nn.Conv1d(dim_in, res, kernel_size=1)))
+        self.generator = Generator(style_dim, **gen)
+        self.generator.m_source = SourceModuleHnNSF(
+            sampling_rate=sr,
+            upsample_scale=math.prod(gen["upsample_rates"]) * gen["gen_istft_hop_size"],
+            harmonic_num=8, voiced_threshod=10)
+
+    def forward(self, asr, F0_curve, N, s):
+        """asr [B, C, T] (40 Hz), F0/N [B, 2T] (80 Hz), s [B, style] -> [B, 1, 400 T]."""
+        F0 = self.F0_conv(F0_curve.unsqueeze(1))
+        N = self.N_conv(N.unsqueeze(1))
+        x = torch.cat([asr, F0, N], axis=1)
+        x = self.encode(x, s)
+        asr_res = self.asr_res(asr)
+        res = True
+        for block in self.decode:
+            if res:
+                x = torch.cat([x, asr_res, F0, N], axis=1)
+            x = block(x, s)
+            if block.upsample_type != "none":
+                res = False
+        return self.generator(x, s, F0_curve)
+
+
+# ---- StyleTTS2 style encoder (Modules from yl4579/StyleTTS2 models.py, MIT) --------
+class LearnedDownSample(nn.Module):
+    def __init__(self, dim_in):
+        super().__init__()
+        self.conv = spectral_norm(nn.Conv2d(dim_in, dim_in, kernel_size=(3, 3), stride=(2, 2),
+                                            groups=dim_in, padding=1))
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.actv = nn.LeakyReLU(0.2)
+        self.learned_sc = dim_in != dim_out
+        self.conv1 = spectral_norm(nn.Conv2d(dim_in, dim_in, 3, 1, 1))
+        self.conv2 = spectral_norm(nn.Conv2d(dim_in, dim_out, 3, 1, 1))
+        self.down_res = LearnedDownSample(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = spectral_norm(nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return F.avg_pool2d(x, 2)
+
+    def _residual(self, x):
+        x = self.conv1(self.actv(x))
+        x = self.down_res(x)
+        return self.conv2(self.actv(x))
+
+    def forward(self, x):
+        return (self._shortcut(x) + self._residual(x)) / math.sqrt(2)
+
+
+class StyleEncoder(nn.Module):
+    """mel [B, 1, 80, T] -> [B, style_dim]. T >= 80 frames (1 s) recommended."""
+
+    def __init__(self, dim_in=32, style_dim=128, max_conv_dim=256):
+        super().__init__()
+        blocks = [spectral_norm(nn.Conv2d(1, dim_in, 3, 1, 1))]
+        d = dim_in
+        for _ in range(4):
+            o = min(d * 2, max_conv_dim)
+            blocks.append(ResBlk(d, o))
+            d = o
+        blocks += [nn.LeakyReLU(0.2), spectral_norm(nn.Conv2d(d, d, 5, 1, 0)),
+                   nn.AdaptiveAvgPool2d(1), nn.LeakyReLU(0.2)]
+        self.shared = nn.Sequential(*blocks)
+        self.unshared = nn.Linear(d, style_dim)
+
+    def forward(self, x):
+        # 4 halvings: time must be a multiple of 16 or the pooled shortcut and the
+        # strided residual disagree by a frame (2 s of audio = 161 mel frames).
+        x = x[..., : (x.shape[-1] // 16) * 16]
+        return self.unshared(self.shared(x).flatten(1))
+
+
+def build(stage: str = "C", c: dict | None = None) -> nn.ModuleDict:
+    c = c or conf()
+    se = c["style_encoder"]
+    m = {
+        "text_encoder": TextEncoder(c["hidden_dim"], c["text_encoder_kernel_size"], c["n_layer"],
+                                    c["n_token"]),
+        "decoder": Decoder(c["hidden_dim"], c["style_dim"], c["decoder"]["hid"], c["decoder"]["out"],
+                           c["decoder"]["res"], c["istftnet"], sr=c["sample_rate"]),
+        "style_encoder": StyleEncoder(se["dim_in"], c["style_dim"], se["max_conv_dim"]),
+    }
+    if stage != "C":
+        from transformers import AlbertConfig, AlbertModel
+        m["plbert"] = AlbertModel(AlbertConfig(vocab_size=c["n_token"], hidden_act="gelu", **c["plbert"]),
+                                  add_pooling_layer=False)
+        m["bert_encoder"] = nn.Linear(c["plbert"]["hidden_size"], c["hidden_dim"])
+        m["predictor"] = ProsodyPredictor(c["style_dim"], c["hidden_dim"], c["n_layer"], c["max_dur"],
+                                          c["dropout"])
+        m["predictor_encoder"] = StyleEncoder(se["dim_in"], c["style_dim"], se["max_conv_dim"])
+    return nn.ModuleDict(m)
+
+
+def energy(mel_norm: torch.Tensor) -> torch.Tensor:
+    """StyleTTS2's N target: log of the L2 norm of the LINEAR mel frame, [B, 80, T] -> [B, T]."""
+    from .mel import MEAN, STD
+    return torch.log(torch.exp(mel_norm * STD + MEAN).norm(dim=1) + 1e-5)
+
+

swar/text.py

@@ -0,0 +1,296 @@
+#!/usr/bin/env python3
+"""The text front-end of Swar: raw text -> the letters the model reads.
+
+    from swar.text import normalize, to_ids, VOCAB
+    normalize("तपाईंको balance Rs. 250 छ।")  ->  "तपाईंको balance दुई सय पचास rupees छ."
+
+RULES, all enforced here so training and inference cannot drift:
+  * Output holds ONLY Devanagari letters/signs, Latin letters, space, apostrophe and
+    ordinary punctuation. No tags, no markers, no language switches -- the script
+    a letter is written in IS its language.
+  * Numbers become words in the language of the text around them (nearest letter
+    wins): Devanagari context -> Nepali words, Latin context -> English words.
+    Digit runs of 6+ are identifiers (phone, account, OTP) and are read digit by
+    digit with "double" (डबल in Nepali).
+  * Anything else is dropped (normalize) or makes the sentence unusable
+    (is_clean), so a stray symbol can never reach the model as an unknown id.
+
+This file is the vocabulary too: VOCAB is fixed here, and ids are positions in it,
+so a checkpoint is only valid with the VOCAB it was trained on (saved inside it).
+"""
+from __future__ import annotations
+
+import re
+import unicodedata
+
+# ---- vocabulary ----------------------------------------------------------------
+PAD, BOS, EOS = "<pad>", "<s>", "</s>"            # internal ids, never in text
+PUNCT = list(".,?!;:'\"-()…—")
+LATIN = [chr(c) for c in range(ord("a"), ord("z") + 1)] + \
+        [chr(c) for c in range(ord("A"), ord("Z") + 1)]
+# Every Devanagari code point Nepali text uses: U+0900..U+0963 (letters, signs,
+# vowel signs, virama, nukta) plus U+0970..U+097F; digits U+0966..U+096F are
+# excluded because numbers are always verbalised first; dandas fold to "."
+DEVANAGARI = [chr(c) for c in range(0x0900, 0x0964)] + [chr(c) for c in range(0x0970, 0x0980)]
+MASK = "<mask>"                                   # stage A (PL-BERT) only; appended LAST so
+VOCAB = [PAD, BOS, EOS, " "] + PUNCT + LATIN + DEVANAGARI + [MASK]   # no other id moves
+assert len(VOCAB) <= 256, len(VOCAB)
+STOI = {s: i for i, s in enumerate(VOCAB)}
+ALLOWED = set(VOCAB[3:-1])
+
+# ---- character folding -----------------------------------------------------------
+_FOLD = {
+    "।": ".", "॥": ".", "’": "'", "‘": "'", "`": "'", "“": '"', "”": '"',
+    "–": "-", "‐": "-", "‑": "-", "...": "…", " ": " ",
+}
+_DROP = dict.fromkeys(map(ord, "​‌‍‎‏"), None)
+
+# ---- English numbers ---------------------------------------------------------------
+_ONES = ("zero one two three four five six seven eight nine ten eleven twelve thirteen "
+         "fourteen fifteen sixteen seventeen eighteen nineteen").split()
+_TENS = "_ _ twenty thirty forty fifty sixty seventy eighty ninety".split()
+_SCALES = [(10 ** 9, "billion"), (10 ** 6, "million"), (1000, "thousand")]
+_ORD = {"one": "first", "two": "second", "three": "third", "five": "fifth",
+        "eight": "eighth", "nine": "ninth", "twelve": "twelfth"}
+
+
+def en_cardinal(n: int) -> str:
+    if n < 20:
+        return _ONES[n]
+    if n < 100:
+        return _TENS[n // 10] + ("" if n % 10 == 0 else "-" + _ONES[n % 10])
+    if n < 1000:
+        rest = n % 100
+        return _ONES[n // 100] + " hundred" + ("" if rest == 0 else " and " + en_cardinal(rest))
+    for v, w in _SCALES:
+        if n >= v:
+            rest = n % v
+            head = en_cardinal(n // v) + " " + w
+            if rest == 0:
+                return head
+            return head + (" and " if rest < 100 else " ") + en_cardinal(rest)
+    raise ValueError(n)
+
+
+def en_ordinal(n: int) -> str:
+    w = en_cardinal(n)
+    last = re.split(r"([ -])", w)
+    tail = last[-1]
+    if tail in _ORD:
+        tail = _ORD[tail]
+    elif tail.endswith("y"):
+        tail = tail[:-1] + "ieth"
+    else:
+        tail += "th"
+    return "".join(last[:-1]) + tail
+
+
+def en_digits(run: str) -> str:
+    """Identifier reading: 9851005996 -> nine eight five one double zero ..."""
+    out, i = [], 0
+    while i < len(run):
+        if i + 1 < len(run) and run[i] == run[i + 1]:
+            out.append("double " + _ONES[int(run[i])])
+            i += 2
+        else:
+            out.append(_ONES[int(run[i])])
+            i += 1
+    return " ".join(out)
+
+
+# ---- Nepali numbers ------------------------------------------------------------------
+_NE_ONES = ("शून्य एक दुई तीन चार पाँच छ सात आठ नौ दश एघार बाह्र तेह्र चौध पन्ध्र सोह्र सत्र "
+            "अठार उन्नाइस बीस एक्काइस बाइस तेइस चौबीस पच्चीस छब्बीस सत्ताइस अठ्ठाइस उनन्तीस "
+            "तीस एकतीस बत्तीस तेत्तीस चौँतीस पैँतीस छत्तीस सैँतीस अठतीस उनन्च���लीस चालीस "
+            "एकचालीस बयालीस त्रिचालीस चवालीस पैँतालीस छयालीस सच्चालीस अठचालीस उनन्चास पचास "
+            "एकाउन्न बाउन्न त्रिपन्न चवन्न पचपन्न छपन्न सन्ताउन्न अन्ठाउन्न उनन्साठी साठी "
+            "एकसट्ठी बयसट्ठी त्रिसट्ठी चौसट्ठी पैँसट्ठी छयसट्ठी सतसट्ठी अठसट्ठी उनन्सत्तरी सत्तरी "
+            "एकहत्तर बहत्तर त्रिहत्तर चौहत्तर पचहत्तर छयहत्तर सतहत्तर अठहत्तर उनासी असी "
+            "एकासी बयासी त्रियासी चौरासी पचासी छयासी सतासी अठासी उनान्नब्बे नब्बे "
+            "एकानब्बे बयानब्बे त्रियानब्बे चौरानब्बे पन्चानब्बे छयानब्बे सन्तानब्बे अन्ठानब्बे "
+            "उनान्सय").split()
+assert len(_NE_ONES) == 100, len(_NE_ONES)
+_NE_SCALES = [(10 ** 7, "करोड"), (10 ** 5, "लाख"), (1000, "हजार"), (100, "सय")]
+
+
+def ne_cardinal(n: int) -> str:
+    if n < 100:
+        return _NE_ONES[n]
+    for v, w in _NE_SCALES:
+        if n >= v:
+            rest = n % v
+            head = ne_cardinal(n // v) + " " + w
+            return head if rest == 0 else head + " " + ne_cardinal(rest)
+    raise ValueError(n)
+
+
+_NE_DIGIT = dict(zip("0123456789", _NE_ONES[:10]))
+
+
+def ne_digits(run: str) -> str:
+    """'9851005996' -> 'नौ आठ पाँच एक डबल शून्य पाँच डबल नौ छ'"""
+    d = run.translate(str.maketrans("०१२३४५६७८९", "0123456789"))
+    words, i = [], 0
+    while i < len(d):
+        if i + 1 < len(d) and d[i] == d[i + 1]:
+            words.append("डबल " + _NE_DIGIT[d[i]])
+            i += 2
+        else:
+            words.append(_NE_DIGIT[d[i]])
+            i += 1
+    return " ".join(words)
+
+
+# ---- script context -------------------------------------------------------------------
+def _is_dev(ch: str) -> bool:
+    return 0x0900 <= ord(ch) <= 0x097F and not ("०" <= ch <= "९")
+
+
+def _context_is_nepali(s: str, start: int, end: int) -> bool:
+    """The sentence's majority script decides; the nearest letter only breaks ties.
+
+    "तपाईंको balance Rs. 250 छ" is a Nepali sentence with one English loanword,
+    and a Nepali speaker reads its number in Nepali -- nearest-letter alone would
+    say English, because "balance" is closer than "छ".
+    """
+    dev = sum(1 for ch in s if _is_dev(ch))
+    lat = sum(1 for ch in s if ch.isascii() and ch.isalpha())
+    if dev != lat:
+        return dev > lat
+    for d in range(1, max(start, len(s) - end) + 1):
+        for j in (start - d, end - 1 + d):
+            if 0 <= j < len(s) and s[j].isalpha():
+                return _is_dev(s[j])
+    return False
+
+
+_CURRENCY = re.compile(r"(?:\b(?:Rs|NPR|NRs)\.?|रु\.?|\$)\s*(?=[0-9०-९])")
+_NUM = re.compile(r"[0-9०-९][0-9०-९,]*(?:\.[0-9०-९]+)*(?:st|nd|rd|th)?")
+_TIME = re.compile(r"(?<![0-9०-९:])([0-9०-९]{1,2}):([0-9०-९]{2})(?![0-9०-९:])")
+_DEV2ASCII = str.maketrans("०१२३४५६७८९", "0123456789")
+
+
+def _verbalise_numbers(s: str) -> str:
+    # currency prefix -> spoken unit after the number ("Rs 250" -> "250 rupees")
+    # Placeholders are control characters, NOT letters: a placeholder made of
+    # letters would itself be counted by the script-context check.
+    def cur(m):
+        return "\x01" if not m.group(0).startswith("$") else "\x02"
+    s = _CURRENCY.sub(cur, s)
+    s = _TIME.sub(lambda m: _time(s, m), s)
+
+    out, last = [], 0
+    for m in _NUM.finditer(s):
+        try:
+            word_or_none = _one_number(s, m, last, out)
+        except (ValueError, IndexError):
+            word_or_none = None
+        if word_or_none is None:
+            # Unreadable ("1,4th", "3.2.1"): leave the digits. They are not in
+            # VOCAB, so is_clean() rejects the sentence instead of the model
+            # hearing a made-up reading.
+            out.append(s[last:m.end()])
+        last = m.end()
+    out.append(s[last:])
+    return "".join(out).replace("\x01", "").replace("\x02", "")
+
+
+def _time(s, m):
+    """5:30 -> 'five thirty' / 'पाँच बजेर तीस मिनेट'; 9:00 -> "nine o'clock" / 'नौ बजे'."""
+    h, mi = (int(x.translate(_DEV2ASCII)) for x in m.groups())
+    if h > 24 or mi > 59:
+        return m.group(0)
+    if _context_is_nepali(s, m.start(), m.end()):
+        return ne_cardinal(h) + " बजे" + ("" if mi == 0 else "र " + ne_cardinal(mi) + " मिनेट")
+    if mi == 0:
+        return en_cardinal(h) + " o'clock"
+    return en_cardinal(h) + " " + ("oh " + _ONES[mi] if mi < 10 else en_cardinal(mi))
+
+
+def _one_number(s, m, last, out):
+    """Append the spoken form of number match m to out; return True, or None."""
+    if m.group(0).translate(_DEV2ASCII).count(".") > 1 or re.match(r"[.:/][0-9०-९]", s[m.end():m.end() + 2]):
+        return None                       # 3.2.1, 1/2, 12:30:15 -- no safe reading
+    tok = m.group(0)
+    ne = _context_is_nepali(s, m.start(), m.end())
+    raw = tok.translate(_DEV2ASCII)
+    suffix = ""
+    mo = re.match(r"^(.*?)(st|nd|rd|th)$", raw)
+    if mo and mo.group(1).isdigit():
+        raw, suffix = mo.group(1), "ord"
+    digits = raw.replace(",", "")
+    pre = s[last:m.start()]
+    money = pre.endswith(("\x01", "\x02"))
+    if money and "." in digits and len(digits.split(".", 1)[1]) == 2:
+        a, b = (int(x) for x in digits.split(".", 1))
+        usd = pre.endswith("\x02")
+        if ne:
+            big, small = ("डलर", "सेन्ट") if usd else ("रुपैयाँ", "पैसा")
+            word = f"{ne_cardinal(a)} {big}" + (f" {ne_cardinal(b)} {small}" if b else "")
+        else:
+            big, small = ("dollars", "cents") if usd else ("rupees", "paisa")
+            word = f"{en_cardinal(a)} {big}" + (f" and {en_cardinal(b)} {small}" if b else "")
+        out.append(pre[:-1] + word)
+        return True
+    if "." in digits:
+        a, b = digits.split(".", 1)
+        word = ((ne_cardinal(int(a)) + " दशमलव " + " ".join(_NE_ONES[int(c)] for c in b))
+                if ne else
+                (en_cardinal(int(a)) + " point " + " ".join(_ONES[int(c)] for c in b)))
+    elif len(digits) >= 6 or (digits.startswith("0") and len(digits) > 1):
+        word = ne_digits(digits) if ne else en_digits(digits)
+    elif suffix == "ord" and not ne:
+        word = en_ordinal(int(digits))
+    else:
+        word = ne_cardinal(int(digits)) if ne else en_cardinal(int(digits))
+    unit = ""
+    if pre.endswith("\x01"):
+        pre, unit = pre[:-1], (" रुपैयाँ" if ne else " rupees")
+    elif pre.endswith("\x02"):
+        pre, unit = pre[:-1], (" डलर" if ne else " dollars")
+    out.append(pre + word + unit)
+    return True
+
+
+# ---- public API -------------------------------------------------------------------------
+def normalize(text: str) -> str:
+    """Raw text -> model letters. Same function for training data and live requests."""
+    s = unicodedata.normalize("NFC", text or "").translate(_DROP)
+    for a, b in _FOLD.items():
+        s = s.replace(a, b)
+    s = _verbalise_numbers(s)
+    s = s.replace("%", " percent").replace("&", " and ")
+    s = "".join(ch if ch in ALLOWED else (" " if ch.isspace() else "") for ch in s)
+    s = re.sub(r"\s+", " ", s).strip()
+    s = re.sub(r"\s+([.,?!;:…])", r"\1", s)                 # no space before punctuation
+    s = re.sub(r"([.?!])\1+", r"\1", s)
+    return s
+
+
+def is_clean(raw: str, normalized: str) -> bool:
+    """True if normalising dropped nothing but whitespace and known foldings --
+    i.e. the sentence can be trained on without its text silently differing from
+    what the audio says."""
+    s = unicodedata.normalize("NFC", raw or "").translate(_DROP)
+    for a, b in _FOLD.items():
+        s = s.replace(a, b)
+    s = _verbalise_numbers(s).replace("%", " percent").replace("&", " and ")
+    return all(ch in ALLOWED or ch.isspace() for ch in s) and bool(normalized)
+
+
+def to_ids(normalized: str) -> list[int]:
+    return [STOI[BOS]] + [STOI[c] for c in normalized] + [STOI[EOS]]
+
+
+if __name__ == "__main__":
+    print(f"VOCAB {len(VOCAB)} symbols: {len(PUNCT)} punct, {len(LATIN)} Latin, "
+          f"{len(DEVANAGARI)} Devanagari, 4 internal/space")
+    for t in ["तपाईंको balance Rs. 250 छ।",
+              "Your balance is Rs 1,250.50, valid until the 15th.",
+              "Please call 9851005996 or 015574574.",
+              "मेरो नम्बर ९८५१००५९९६ हो, र म ५ बजे आउँछु।",
+              "रु ५०० को रिचार्ज गर्नुहोस्",
+              "A & B give 20% off… ‘really’?",
+              "emoji 😀 and © symbols are dropped"]:
+        n = normalize(t)
+        print(f"{'clean' if is_clean(t, n) else 'DIRTY'}  {t!r}\n       -> {n!r}")

voices/bijaya.pt

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