ema-lightning 1.0.0__py3-none-any.whl

ema_lightning/__init__.py

@@ -0,0 +1,5 @@
+"""EMA Lightning: tiny, fast and accurate Turkish text to speech."""
+from .api import EMA, Speech
+
+__all__ = ["EMA", "Speech"]
+__version__ = "1.0.0"

ema_lightning/api.py

@@ -0,0 +1,247 @@
+"""EMA Lightning: Turkish text to speech.
+
+    tts = EMA().lightning()
+    speech = tts.say("Merhaba, nasılsınız?", path="merhaba.wav")
+    speeches = tts.say(["Birinci cümle.", "İkinci cümle."])
+    for chunk in tts.stream("Uzun bir metin..."):
+        play(chunk)
+"""
+import json
+import os
+import random
+import statistics
+import time
+import warnings
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+
+from .audio import RATE, RATES, Resampler, write_wav
+from .chunker import chunk
+from .decoder import load_decoder
+from .engine import FIRST_WINDOW, Engine, windows
+from .frontend import Frontend
+from .graphs import Graphs
+from .model import load_acoustic
+from .scheduler import DONE, Playhead
+
+REPO = "canberkkkkkk/ema-lightning"
+PROBE = "Bugün hava çok güzel, yarın da yağmur yağacakmış; toplantı öğleden sonra başlayacak."
+CACHE = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "ema_lightning" / "batch_size_v2.json"
+
+
+@dataclass(frozen=True)
+class Speech:
+    """One spoken text: float32 audio in [-1, 1] at sample_rate, and the seed that made it."""
+
+    audio: np.ndarray = field(repr=False)
+    sample_rate: int
+    duration: float
+    seed: int
+
+
+class EMA:
+    def __init__(self, device="auto"):
+        device = torch.device(("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device)
+        model = load_acoustic(hf_hub_download(REPO, "ema.pt"), device)
+        decoder = load_decoder(hf_hub_download(REPO, "decoder.pt"), device)
+        self._setup(model, decoder, Frontend(model.vocab), device)
+
+    @classmethod
+    def _from_parts(cls, model, decoder, frontend, device):
+        self = cls.__new__(cls)
+        self._setup(model, decoder, frontend, device)
+        return self
+
+    def _setup(self, model, decoder, frontend, device):
+        self.device = torch.device(device)
+        self._engine = Engine(model, decoder, self.device)
+        self._frontend = frontend
+        self._batch_size = None
+        self._playhead = Playhead(self._engine, self._list_batch)
+
+    def lightning(self, batch_size=None):
+        """Compile and record every stage as CUDA graphs, check them against the plain path, report ready."""
+        if self.device.type != "cuda":
+            warnings.warn("lightning needs a CUDA GPU; EMA keeps working without it", stacklevel=2)
+            return self
+        if self._engine.graphs is not None:
+            return self
+        # cuDNN times each convolution shape once and keeps the fastest method. The decoder is all convolutions:
+        # on an RTX PRO 6000 this made it about twice as fast. It is process-wide and changes speed, not the audio.
+        torch.backends.cudnn.benchmark = True
+        start = time.perf_counter()
+        size = batch_size or self.best_batch_size()
+        try:
+            graphs = Graphs(self._engine.model, self._engine.decoder, size)
+        except Exception as error:  # e.g. no compiler on this machine: graphs alone still help
+            warnings.warn(f"compiling failed, recording graphs without it: {error}", stacklevel=2)
+            graphs = Graphs(self._engine.model, self._engine.decoder, size, compile=False)
+        built = time.perf_counter() - start
+        plain = self._probe()
+        self._engine.graphs = graphs
+        fast = self._probe()
+        if any((a - b).abs().max() > 1e-2 for a, b in zip(plain, fast, strict=True)):
+            self._engine.graphs = None
+            warnings.warn("recorded graphs disagree with the plain path; lightning stays off", stacklevel=2)
+            return self
+        self._batch_size = size
+        first = statistics.median(self._first_audio() for _ in range(5))
+        timings = [self._timed_probe() for _ in range(5)]
+        ms, seconds = statistics.median(t for t, _ in timings), timings[0][1]
+        print(f"EMA Lightning ready: {graphs.count} graphs built in {built:.1f} s, first audio in {first:.1f} ms, "
+              f"a {seconds:.1f} s sentence in {ms:.1f} ms, batch size {size}")
+        return self
+
+    def best_batch_size(self):
+        """The smallest batch that reaches 90% of this device's best throughput, measured once and cached."""
+        if self._batch_size:
+            return self._batch_size
+        key = torch.cuda.get_device_name(self.device) if self.device.type == "cuda" else f"cpu-{os.cpu_count()}"
+        cached = json.loads(CACHE.read_text()) if CACHE.exists() else {}
+        if key not in cached:
+            sizes = (1, 2, 4, 8, 16, 32, 64, 128) if self.device.type == "cuda" else (1, 2, 4, 8)
+            rates = {}
+            for size in sizes:
+                try:
+                    rates[size] = self._throughput(size)
+                except torch.cuda.OutOfMemoryError:
+                    torch.cuda.empty_cache()
+                    break
+            cached[key] = min(s for s, r in rates.items() if r >= 0.9 * max(rates.values()))
+            CACHE.parent.mkdir(parents=True, exist_ok=True)
+            CACHE.write_text(json.dumps(cached))
+        self._batch_size = cached[key]
+        return self._batch_size
+
+    def say(self, text, speed=1.0, seed=None, sample_rate=RATE, path=None):
+        """Speech for one text, or a list of Speech for a list of texts (batched)."""
+        _check(speed, seed, sample_rate)
+        if isinstance(text, str):
+            seed = _seed(seed)
+            request = self._submit(text, speed, seed)
+            speech = _speech(list(self._receive(request, sample_rate)), sample_rate, seed)
+            if path is not None:
+                write_wav(path, speech.audio, sample_rate)
+            return speech
+        texts = _texts(text)
+        seeds = [_seed(seed) for _ in texts]
+        requests = [self._submit(t, speed, s) for t, s in zip(texts, seeds, strict=True)]
+        out = [_speech(list(self._receive(r, sample_rate)), sample_rate, s)
+               for r, s in zip(requests, seeds, strict=True)]
+        if path is not None:
+            Path(path).mkdir(parents=True, exist_ok=True)
+            width = len(str(max(len(out) - 1, 0)))
+            for i, speech in enumerate(out):
+                write_wav(Path(path) / f"{i:0{width}d}.wav", speech.audio, sample_rate)
+        return out
+
+    def stream(self, text, speed=1.0, seed=None, sample_rate=RATE):
+        """float32 audio chunks as they are made: one second first, then four seconds at a time.
+
+        Call it from as many threads as you like; every stream shares the GPU through the same queues.
+        Stopping early (leaving the loop) drops the rest of that text's work.
+        """
+        _check(speed, seed, sample_rate)
+        if not isinstance(text, str):
+            raise TypeError("stream() takes one text; for many, call stream() once per text")
+        return self._stream(text, speed, _seed(seed), sample_rate)
+
+    def _stream(self, text, speed, seed, sample_rate):
+        """A stream's request, submitted when its first chunk is asked for."""
+        yield from self._receive(self._submit(text, speed, seed, first=FIRST_WINDOW), sample_rate)
+
+    def _submit(self, text, speed, seed, first=None):
+        pieces = self._pieces(text, speed, seed)
+        return self._playhead.submit(pieces, speed) if first is None else self._playhead.submit(pieces, speed, first)
+
+    def _receive(self, request, sample_rate):
+        """float32 chunks from a request's outbox, resampled. If the caller stops waiting, its work is dropped."""
+        resampler = Resampler(sample_rate)
+        try:
+            while (item := request.outbox.get()) is not DONE:
+                if isinstance(item, BaseException):
+                    raise item
+                out = resampler.push(item)
+                if out.numel():
+                    yield out.float().cpu().numpy()
+            tail = resampler.flush()
+            if tail is not None and tail.numel():
+                yield tail.float().cpu().numpy()
+        finally:
+            request.cancel()
+
+    def _pieces(self, text, speed, seed):
+        return [self._engine.piece(piece, pause, (seed * 1_000_003 + i) % 2**63)
+                for i, (piece, pause) in enumerate(chunk(self._frontend(text), speed))]
+
+    def _list_batch(self):
+        return self._batch_size or self.best_batch_size()
+
+    def _probe(self):
+        """Latents and both window sizes of a fixed sentence, for checking graphs against the plain path."""
+        (piece,) = self._pieces(PROBE, 1.0, 0)
+        with self._engine.lock:
+            self._engine.plan([piece], 1.0)
+            self._engine.think([piece])
+            short = self._engine.decode([(piece, windows(piece.frames, FIRST_WINDOW)[0])])[0]
+            full = self._engine.decode([(piece, windows(piece.frames)[0])])[0]
+            return piece.latents, short, full
+
+    def _first_audio(self):
+        """Milliseconds until a stream of the probe sentence hands over its first chunk."""
+        start = time.perf_counter()
+        next(iter(self._stream(PROBE, 1.0, 0, RATE)))
+        return 1000 * (time.perf_counter() - start)
+
+    def _timed_probe(self):
+        """Milliseconds for say() to return the probe sentence, and the seconds of audio it made."""
+        start = time.perf_counter()
+        speech = self.say(PROBE, seed=0)
+        return 1000 * (time.perf_counter() - start), speech.duration
+
+    def _throughput(self, size):
+        """Seconds of audio per second for a batch of `size` copies of the probe sentence."""
+        (piece,) = self._pieces(PROBE, 1.0, 0)
+        pieces = [self._engine.piece(piece.text, 0.0, i) for i in range(size)]
+        best = float("inf")
+        with self._engine.lock:
+            self._engine.plan(pieces, 1.0)
+            for _ in range(3):
+                if self.device.type == "cuda":
+                    torch.cuda.synchronize(self.device)
+                start = time.perf_counter()
+                self._engine.think(pieces)
+                for span in windows(pieces[0].frames):
+                    self._engine.decode([(p, span) for p in pieces])
+                if self.device.type == "cuda":
+                    torch.cuda.synchronize(self.device)
+                best = min(best, time.perf_counter() - start)
+        return size * pieces[0].frames / 25 / best
+
+
+def _check(speed, seed, sample_rate):
+    if isinstance(speed, bool) or not isinstance(speed, (int, float)) or not 0.25 <= speed <= 4:
+        raise ValueError("speed must be a number from 0.25 to 4")
+    if seed is not None and (isinstance(seed, bool) or not isinstance(seed, int) or seed < 0):
+        raise ValueError("seed must be a non-negative integer")
+    if sample_rate not in RATES:
+        raise ValueError(f"sample_rate must be one of {RATES}")
+
+
+def _texts(text):
+    if not isinstance(text, (list, tuple)) or not all(isinstance(t, str) for t in text):
+        raise TypeError("text must be a string or a list of strings")
+    return list(text)
+
+
+def _seed(seed):
+    return random.SystemRandom().randrange(2**31) if seed is None else seed
+
+
+def _speech(chunks, sample_rate, seed):
+    audio = np.concatenate(chunks).astype(np.float32) if chunks else np.zeros(0, np.float32)
+    return Speech(audio, sample_rate, len(audio) / sample_rate, seed)

ema_lightning/audio.py

@@ -0,0 +1,67 @@
+"""Audio helpers: sample-rate conversion from 48 kHz, window by window, and WAV writing.
+
+The resampler carries its state from window to window, so the output never depends on where the decoder's
+windows were cut.
+"""
+import wave
+
+import numpy as np
+import torch
+
+RATE = 48000
+RATES = (48000, 24000, 16000, 8000)
+
+
+class Resampler:
+    """48 kHz in, 48 kHz / factor out, through a linear-phase windowed-sinc low-pass."""
+
+    def __init__(self, rate):
+        self.factor = RATE // rate
+        self.half = 32 * self.factor
+        self.taps = None
+        self.buffer = None
+        self.start = -self.half  # absolute input index of buffer[0]; the signal is preceded by zeros
+        self.next = 0  # absolute input index at the centre of the next output sample
+        self.end = 0  # absolute input index one past the last real sample
+
+    def push(self, chunk):
+        if self.factor == 1:
+            return chunk
+        if self.buffer is None:
+            self.buffer = chunk.new_zeros(self.half)
+            n = torch.arange(-self.half, self.half + 1, dtype=torch.float64)
+            cutoff = 0.45 / self.factor
+            taps = 2 * cutoff * torch.sinc(2 * cutoff * n) * torch.blackman_window(2 * self.half + 1, False,
+                                                                                    dtype=torch.float64)
+            self.taps = (taps / taps.sum()).to(chunk)[None, None]
+        self.buffer = torch.cat([self.buffer, chunk])
+        self.end += chunk.numel()
+        return self._emit(self.end - 1 - self.half)
+
+    def flush(self):
+        if self.factor == 1 or self.buffer is None:
+            return None
+        self.buffer = torch.cat([self.buffer, self.buffer.new_zeros(self.half + self.factor)])
+        return self._emit(self.end - 1)
+
+    def _emit(self, last):
+        """Every output whose centre is at or before input index `last`."""
+        if last < self.next:
+            return self.buffer.new_zeros(0)
+        count = (last - self.next) // self.factor + 1
+        a = self.next - self.half - self.start
+        b = a + (count - 1) * self.factor + 2 * self.half + 1
+        out = torch.nn.functional.conv1d(self.buffer[a:b][None, None], self.taps, stride=self.factor)[0, 0]
+        self.next += count * self.factor
+        drop = self.next - self.half - self.start
+        self.buffer, self.start = self.buffer[drop:], self.start + drop
+        return out
+
+
+def write_wav(path, audio, rate):
+    pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).round().astype("<i2")
+    with wave.open(str(path), "wb") as f:
+        f.setnchannels(1)
+        f.setsampwidth(2)
+        f.setframerate(rate)
+        f.writeframes(pcm.tobytes())

ema_lightning/chunker.py

@@ -0,0 +1,44 @@
+"""Greedy chunking: spoken text in, pieces short enough for one pass of the model out.
+
+Text that fits in about ten seconds of speech stays one piece; the model reads its punctuation itself.
+Longer text is cut at the last good spot inside each ten-second window: a sentence end, then a clause
+mark, then a space, and only when there is none of those, exactly at the limit.
+"""
+import re
+
+LETTERS_PER_SECOND = 18.0
+MAX_SECONDS = 10.0
+MAX_LETTERS = 250
+SENTENCE_PAUSE = 0.25
+CLAUSE_PAUSE = 0.12
+CUTS = ((re.compile(r"[.!?]+[\"')]*(?= )"), SENTENCE_PAUSE), (re.compile(r"[,;:](?= )"), CLAUSE_PAUSE),
+        (re.compile(r"\S(?= )"), CLAUSE_PAUSE))
+LETTER = re.compile(r"[^\W\d_]")
+
+
+def chunk(text, speed):
+    """[(piece, seconds of silence after it)], each piece ending in terminal punctuation."""
+    limit = int(min(MAX_LETTERS, LETTERS_PER_SECOND * MAX_SECONDS * speed))
+    pieces, rest = [], text.strip()
+    while rest:
+        cut, pause = len(rest), 0.0
+        if len(rest) > limit:
+            cut = limit
+            for pattern, gap in CUTS:
+                ends = [m.end() for m in pattern.finditer(rest, 0, limit + 1)]
+                if ends:
+                    cut, pause = ends[-1], gap
+                    break
+        piece, rest = rest[:cut].strip(), rest[cut:].strip()
+        if LETTER.search(piece):
+            pieces.append((finish(piece), pause))
+    if pieces:
+        pieces[-1] = (pieces[-1][0], 0.0)
+    return pieces
+
+
+def finish(piece):
+    """End every piece the way the model was trained: on a sentence end."""
+    if piece.rstrip("\"')")[-1:] in (".", "!", "?"):
+        return piece
+    return piece.rstrip(",;:- ") + "."

ema_lightning/decoder.py

@@ -0,0 +1,65 @@
+"""EMA Lightning's decoder: 64-dim latents at 25 Hz in, 48 kHz audio out.
+
+With `lengths`, everything past each row's real length is held at zero after every layer, so a window
+padded to a fixed shape decodes exactly like the same window unpadded.
+"""
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _keep(x, length):
+    return x if length is None else x * (torch.arange(x.shape[-1], device=x.device) < length[:, None])[:, None]
+
+
+class ResBlock(nn.Module):
+    def __init__(self, ch, k, dilations):
+        super().__init__()
+        self.convs1 = nn.ModuleList(nn.Conv1d(ch, ch, k, dilation=d, padding=(k * d - d) // 2) for d in dilations)
+        self.convs2 = nn.ModuleList(nn.Conv1d(ch, ch, k, padding=(k - 1) // 2) for _ in dilations)
+
+    def forward(self, x, length=None):
+        for c1, c2 in zip(self.convs1, self.convs2, strict=True):
+            y = _keep(F.leaky_relu(c1(F.leaky_relu(x, 0.1)), 0.1), length)
+            x = _keep(x + c2(y), length)
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(self, latent_dim=64, ch=256, rates=(8, 6, 5, 2, 2, 2), kernels=(16, 12, 10, 4, 4, 4),
+                 rb_kernels=(3, 5, 9), rb_dilations=((1, 3, 5),) * 3):
+        super().__init__()
+        self.hop, self.nk = math.prod(rates), len(rb_kernels)
+        self.pre = nn.Conv1d(latent_dim, ch, 7, padding=3)
+        self.ups = nn.ModuleList(nn.ConvTranspose1d(ch >> i, ch >> (i + 1), k, r, padding=(k - r) // 2)
+                                 for i, (r, k) in enumerate(zip(rates, kernels, strict=True)))
+        self.blocks = nn.ModuleList(ResBlock(ch >> (i + 1), k, d)
+                                    for i in range(len(rates)) for k, d in zip(rb_kernels, rb_dilations, strict=True))
+        self.post = nn.Conv1d(ch >> len(rates), 1, 7, padding=3)
+
+    def forward(self, z, lengths=None):
+        """z: [B, latent_dim, T]; lengths: real frames per row, or None when nothing is padded."""
+        frames = z.shape[-1]
+        length = lengths
+        x = _keep(self.pre(_keep(z, length)), length)
+        for i, up in enumerate(self.ups):
+            x = up(F.leaky_relu(x, 0.1))
+            if length is not None:
+                length = (length - 1) * up.stride[0] - 2 * up.padding[0] + up.kernel_size[0]
+            x = _keep(x, length)
+            x = sum(b(x, length) for b in self.blocks[i * self.nk:(i + 1) * self.nk]) / self.nk
+        return torch.tanh(self.post(F.leaky_relu(x)))[:, 0, :frames * self.hop]
+
+
+def load_decoder(path, device):
+    ck = torch.load(path, map_location="cpu", weights_only=False)
+    cfg, sd = ck.get("cfg", {}), ck["G"]
+    for key in [k for k in sd if k.endswith("weight_g")]:
+        g, v = sd.pop(key), sd.pop(key[:-1] + "v")
+        sd[key[:-2]] = g * v / v.norm(dim=tuple(range(1, v.dim())), keepdim=True)
+    keys = ("latent_dim", "ch", "rates", "kernels", "rb_kernels", "rb_dilations")
+    model = Decoder(**{k: cfg[k] for k in keys if k in cfg})
+    model.load_state_dict(sd)
+    return model.to(device).float().eval().requires_grad_(False)

ema_lightning/engine.py

@@ -0,0 +1,158 @@
+"""The speech engine: pieces of spoken text in, 48 kHz audio out, one decoded window at a time.
+
+Every piece goes through three stages: plan (the text stage and its frame timeline), think (the
+aligner and the four steps) and decode (one window of latents to audio). Each stage takes any batch;
+with lightning on, the same stages replay from recorded CUDA graphs instead. Playhead (scheduler.py)
+decides which work runs in each batch.
+"""
+import threading
+from dataclasses import dataclass
+from itertools import pairwise
+
+import torch
+
+RATE = 48000
+FIRST_WINDOW = 25  # a stream's first window is one second, so its first audio comes fast
+WINDOW = 100  # four seconds for every later window
+CONTEXT = 8  # frames decoded on each side of a window; the decoder reaches 4
+# Every batch of windows is padded to one of these lengths before decoding. The GPU picks its convolution
+# method by input length, and on an RTX PRO 6000 the windows' own lengths (33, 108, 116 frames) ran about
+# twice as slow per frame as 48 and 120. The decoder ignores padding exactly, so the audio does not change.
+DECODE_SIZES = (48, 120)
+MAX_WORD_FRAMES = 250
+MAX_FRAMES = 3000
+
+
+@dataclass(eq=False)
+class Piece:
+    text: str
+    pause: float
+    seed: int
+    ids: torch.Tensor
+    cw: torch.Tensor  # word of each letter
+    wstart: torch.Tensor  # first letter of that word
+    h: torch.Tensor = None
+    dur: torch.Tensor = None
+    fw: torch.Tensor = None  # word of each frame
+    fp: torch.Tensor = None  # position of each frame inside its word
+    latents: torch.Tensor = None
+    spans: list = None  # its decoder windows, fixed when it is planned
+
+    @property
+    def letters(self):
+        return self.ids.numel()
+
+    @property
+    def frames(self):
+        return self.fw.numel()
+
+
+def windows(frames, first=WINDOW):
+    """(start, end) of every decoded window: `first` frames first, then four seconds each."""
+    spans, s = [], 0
+    while s < frames:
+        e = min(frames, s + (first if s == 0 else WINDOW))
+        spans.append((s, e))
+        s = e
+    return spans
+
+
+def batches(items, size):
+    for i in range(0, len(items), size):
+        yield items[i:i + size]
+
+
+class Engine:
+    def __init__(self, model, decoder, device):
+        self.model, self.decoder, self.device = model, decoder, torch.device(device)
+        self.hop = decoder.hop
+        self.graphs = None
+        self.lock = threading.RLock()
+        self.decode_sizes = DECODE_SIZES if self.device.type == "cuda" else ()  # slow lengths are a GPU problem
+
+    def piece(self, text, pause, seed):
+        ids = [self.model.stoi.get(ch, 1) for ch in text]
+        starts = [i for i, ch in enumerate(text) if ch != " " and (i == 0 or text[i - 1] == " ")] or [0]
+        bounds = [0] + starts[1:] + [len(text)]
+        cw, wstart = [], []
+        for w, (a, b) in enumerate(pairwise(bounds)):
+            cw += [w] * (b - a)
+            wstart += [a] * (b - a)
+        return Piece(text, pause, seed, torch.tensor(ids), torch.tensor(cw), torch.tensor(wstart))
+
+    @torch.no_grad()
+    def plan(self, pieces, speed):
+        """Durations for a batch, then each piece's exact frame timeline after a single wait."""
+        L = max(p.letters for p in pieces)
+        ids = self.stack([p.ids for p in pieces], L, 0)
+        mask = ids != 0
+        h, dur = self.run("text", ids=ids, mask=mask)
+        dur = dur / speed
+        word = self.stack([p.cw for p in pieces], L, 0)
+        counts = torch.zeros_like(dur).scatter_add_(1, word, dur).round().clamp(1, MAX_WORD_FRAMES).long().cpu()
+        for i, p in enumerate(pieces):
+            n = counts[i, :int(p.cw[-1]) + 1]
+            frames = min(int(n.sum()), MAX_FRAMES)
+            fw = torch.repeat_interleave(torch.arange(n.numel()), n)[:frames]
+            fp = ((torch.arange(frames) - (n.cumsum(0) - n)[fw]).double() / n[fw].double()).float()
+            p.h, p.dur, p.fw, p.fp = h[i, :p.letters].clone(), dur[i, :p.letters].clone(), fw, fp
+
+    @torch.no_grad()
+    def think(self, pieces):
+        """Latents for a batch, each piece from its own seeded noise."""
+        L, T = max(p.letters for p in pieces), max(p.frames for p in pieces)
+        latents = self.run(
+            "sound",
+            h=self.stack([p.h for p in pieces], L, 0.0),
+            dur=self.stack([p.dur for p in pieces], L, 0.0),
+            mask=self.stack([torch.ones(p.letters, dtype=torch.bool) for p in pieces], L, False),
+            cw=self.stack([p.cw for p in pieces], L, -1),
+            wstart=self.stack([p.wstart for p in pieces], L, 0),
+            fw=self.stack([p.fw for p in pieces], T, -1),
+            fp=self.stack([p.fp for p in pieces], T, 0.0),
+            fmask=self.stack([torch.ones(p.frames, dtype=torch.bool) for p in pieces], T, False),
+            noise=self.stack([self.noise(p) for p in pieces], T, 0.0, dim=1))
+        for i, p in enumerate(pieces):
+            p.latents = latents[i, :p.frames].clone()
+
+    @torch.no_grad()
+    def decode_size(self, piece, span):
+        """The length a window is decoded at: itself plus its margins, padded to a fast size on a GPU."""
+        s, e = span
+        n = min(piece.frames, e + CONTEXT) - max(0, s - CONTEXT)
+        return next((size for size in self.decode_sizes if size >= n), n)
+
+    def decode(self, items):
+        """Audio for a batch of (piece, window), each cut to its own window."""
+        spans = [(max(0, s - CONTEXT), min(p.frames, e + CONTEXT)) for p, (s, e) in items]
+        longest = max(b - a for a, b in spans)
+        P = next((size for size in self.decode_sizes if size >= longest), longest)
+        z = self.stack([p.latents[a:b] for (p, _), (a, b) in zip(items, spans, strict=True)], P, 0.0)
+        lengths = torch.tensor([b - a for a, b in spans], device=self.device)
+        audio = self.run("decode", z=z, lengths=lengths)
+        return [audio[i, (s - a) * self.hop:(e - a) * self.hop].clone()
+                for i, ((_, (s, e)), (a, _)) in enumerate(zip(items, spans, strict=True))]
+
+    def run(self, stage, **args):
+        out = self.graphs.run(stage, **args) if self.graphs is not None else None
+        if out is not None:
+            return out
+        if stage == "text":
+            return self.model.text_stage(**args)
+        if stage == "sound":
+            return self.model.sound_stage(**args)
+        full = bool((args["lengths"] == args["z"].shape[1]).all())  # nothing padded: skip the masks
+        return self.decoder(args["z"].transpose(1, 2), None if full else args["lengths"])
+
+    def noise(self, piece):
+        generator = torch.Generator(device=self.device).manual_seed(piece.seed)
+        shape = (len(self.model.times), piece.frames, self.model.latent_dim)
+        return torch.randn(shape, generator=generator, device=self.device)
+
+    def stack(self, rows, size, fill, dim=0):
+        shape = list(rows[0].shape)
+        shape[dim] = size
+        out = torch.full([len(rows)] + shape, fill, dtype=rows[0].dtype, device=self.device)
+        for i, row in enumerate(rows):
+            out[i].narrow(dim, 0, row.shape[dim]).copy_(row)
+        return out

ema_lightning/frontend.py

@@ -0,0 +1,59 @@
+"""Text frontend: any written Turkish in, text in the model's own alphabet out.
+
+normalizer-tr reads numbers, dates, times, money, units, symbols and abbreviations aloud. Its
+"fallback" policy leaves nothing unread: notation it cannot resolve is spoken literally instead of
+being kept as written. The alphabet step then lowercases the Turkish way and drops whatever the
+model cannot read. Nothing here ever raises on text.
+"""
+import re
+import unicodedata
+
+from normalizer_tr import Normalizer
+
+POLICY = "fallback"
+BLOCK_BYTES = 8 * 1024
+TURKISH = frozenset("çğıöşüÇĞİÖŞÜ")
+TYPOGRAPHY = str.maketrans({"’": "'", "‘": "'", "ʼ": "'", "´": "'", "`": "'", "“": '"', "”": '"', "„": '"',
+                            "«": '"', "»": '"', "–": "-", "—": "-", "−": "-", "…": "..."})
+UNSAFE = re.compile("[\x00-\x08\x0b-\x1f\x7f-\x9f\u061c\u200e\u200f\u202a-\u202e\u2066-\u2069]")
+
+
+class Frontend:
+    def __init__(self, vocab):
+        self.vocab = frozenset(vocab)
+        self.normalizer = Normalizer()
+
+    def __call__(self, text):
+        text = UNSAFE.sub(" ", text.encode("utf-8", "ignore").decode("utf-8"))
+        if not text.strip():
+            return ""
+        return self.alphabet(" ".join(self.spoken(block) for block in blocks(text)))
+
+    def spoken(self, text):
+        try:
+            return self.normalizer.normalize(text, ambiguity_policy=POLICY).normalized_text
+        except Exception:  # invalid input or a resource limit: keep the words rather than lose the sentence
+            return text
+
+    def alphabet(self, text):
+        text = text.translate(TYPOGRAPHY).replace("İ", "i").replace("I", "ı").lower()
+        out = []
+        for ch in text:
+            if ch not in TURKISH:
+                ch = "".join(c for c in unicodedata.normalize("NFKD", ch) if not unicodedata.combining(c))
+            out.append(ch if ch and all(c in self.vocab for c in ch) else " ")
+        return re.sub(r"\s+", " ", "".join(out)).strip()
+
+
+def blocks(text):
+    """Split at whitespace into pieces the normalizer accepts in one call."""
+    words, block, size = text.split(), [], 0
+    for word in words:
+        n = len(word.encode("utf-8")) + 1
+        if block and size + n > BLOCK_BYTES:
+            yield " ".join(block)
+            block, size = [], 0
+        block.append(word)
+        size += n
+    if block:
+        yield " ".join(block)