kokoro-tts stream added

.gitignore

@@ -1,2 +1,3 @@
 .vscode
-.omo
+.omo
+__pycache__

modules/containers/kokoro-tts.nix

@@ -13,24 +13,18 @@ let
   # config revision: edit a file, `nixos-rebuild`, and the unit below rebuilds
   # and restarts. Reading the context off a checkout at runtime would leave the
   # running container untraceable back to any config.
-  source = pkgs.linkFarm "kokoro-tts-source" [
+  #
-    {
+  # runCommand rather than linkFarm: linkFarm entries are symlinks into other
-      name = "Dockerfile";
+  # store paths, and `podman build` only mounts the context root, so every COPY
-      path = toString ./kokoro-tts/Dockerfile;
+  # fails with "copier: get: lstat ...: no such file or directory". Copying the
-    }
+  # bytes in leaves the context with no symlinks that escape its root.
-    {
+  source = pkgs.runCommand "kokoro-tts-source" { } ''
-      name = "app.py";
+    mkdir -p "$out"
-      path = toString ./kokoro-tts/app.py;
+    cp -L ${./kokoro-tts/Dockerfile} "$out/Dockerfile"
-    }
+    cp -L ${./kokoro-tts/app.py} "$out/app.py"
-    {
+    cp -L ${./kokoro-tts/fetch_assets.py} "$out/fetch_assets.py"
-      name = "fetch_assets.py";
+    cp -L ${./kokoro-tts/requirements.txt} "$out/requirements.txt"
-      path = toString ./kokoro-tts/fetch_assets.py;
+  '';
-    }
-    {
-      name = "requirements.txt";
-      path = toString ./kokoro-tts/requirements.txt;
-    }
-  ];
 
   image = "localhost/kokoro-tts:latest";
 
@@ -64,11 +58,16 @@ in
           ];
 
           environment = {
-            # Inference is CPU-bound and already threaded inside torch; these
+            # Inference is CPU-bound and already threaded inside torch. Measured
-            # keep it from oversubscribing a small machine.
+            # on a 24-logical-core host: median end-to-end latency for a 5.6 s
-            KOKORO_THREADS = "4";
+            # utterance was 1.203 s at 4 threads, 0.979 s at 12, 0.980 s at 16
-            OMP_NUM_THREADS = "4";
+            # and 1.87 s at 24, so the useful ceiling is the physical core count
-            MKL_NUM_THREADS = "4";
+            # and oversubscribing it roughly doubles the wait. These three must
+            # stay equal to the Dockerfile ENV and the app.py default: whichever
+            # of the three is set wins over the others.
+            KOKORO_THREADS = "12";
+            OMP_NUM_THREADS = "12";
+            MKL_NUM_THREADS = "12";
             TZ = "Europe/Moscow";
           };
 
@@ -113,4 +112,4 @@ in
       };
     };
   };
-}
+}

modules/containers/kokoro-tts/Dockerfile

@@ -1,4 +1,7 @@
-FROM python:3.12-slim-bookworm
+# Fully qualified on purpose: NixOS ships a podman registries.conf without
+# unqualified-search-registries, so a bare "python:3.12-slim-bookworm" fails to
+# resolve before the build even starts.
+FROM docker.io/library/python:3.12-slim-bookworm
 
 # Pinned, not "main": a rebuild that only touched the Nix module must not
 # silently pick up different weights. Bump these deliberately.
@@ -8,6 +11,9 @@ ARG KOKORO_RU_REVISION=d649c57b239b18c4c384378127cbf01dba039bc1
 # a second 327 MB one.
 ARG KOKORO_RU_VOICES=sveta,masha,dima
 
+# Thread counts, not a guess: see app.py THREADS. 12 was the measured plateau on
+# a 24-logical-core host, and 24 was ~2x worse. Must stay equal to the Nix
+# module's environment.environment, which wins over this ENV.
 ENV PYTHONUNBUFFERED=1 \
     PIP_NO_CACHE_DIR=1 \
     PIP_DISABLE_PIP_VERSION_CHECK=1 \
@@ -17,9 +23,9 @@ ENV PYTHONUNBUFFERED=1 \
     KOKORO_RU_REVISION=${KOKORO_RU_REVISION} \
     KOKORO_RU_VOICES=${KOKORO_RU_VOICES} \
     KOKORO_MODEL_DIR=/app/kokoro-ru \
-    KOKORO_THREADS=4 \
+    KOKORO_THREADS=12 \
-    OMP_NUM_THREADS=4 \
+    OMP_NUM_THREADS=12 \
-    MKL_NUM_THREADS=4 \
+    MKL_NUM_THREADS=12 \
     TZ=Europe/Moscow
 
 WORKDIR /app
@@ -39,13 +45,19 @@ RUN pip install --index-url https://download.pytorch.org/whl/cpu torch
 COPY requirements.txt ./
 RUN pip install -r requirements.txt
 
-COPY app.py fetch_assets.py ./
+# fetch_assets.py is copied on its own and app.py only after the snapshot, never
+# as one COPY. A single COPY would tie the 639 MB download to the application
+# source: any edit to app.py would invalidate this layer and refetch every
+# checkpoint as hundreds of anonymous, rate-limited requests.
+COPY fetch_assets.py ./
 
 # Bakes the checkpoints, the acute-aware espeak data and ruaccent's ONNX models
 # into the layer, which is what lets the container start with no network and no
 # writable volume.
 RUN python fetch_assets.py
 
+COPY app.py ./
+
 EXPOSE 8000
 
 HEALTHCHECK --interval=30s --timeout=5s --start-period=180s --retries=3 \

modules/containers/kokoro-tts/app.py

@@ -10,10 +10,11 @@ non-native (measured 27% vs 22% round-trip WER, per the model card).
 So: text -> RuG2P.phonemize -> KModel(ipa, voicepack[len(ipa) - 1]) -> waveform.
 
 Endpoints
-    POST /v1/audio/speech   OpenAI text-to-speech
+    POST /v1/audio/speech          OpenAI text-to-speech
-    GET  /v1/models         OpenAI model list
+    POST /v1/audio/speech/stream   same, but mp3/opus emitted while synthesising
-    GET  /v1/voices         voice inventory (extension, not part of OpenAI)
+    GET  /v1/models                OpenAI model list
-    GET  /healthz           readiness, 503 until the model is loaded
+    GET  /v1/voices                voice inventory (extension, not part of OpenAI)
+    GET  /healthz                  readiness, 503 until the model is loaded
 """
 
 from __future__ import annotations
@@ -21,6 +22,7 @@ from __future__ import annotations
 import io
 import logging
 import os
+import queue
 import re
 import subprocess
 import sys
@@ -28,11 +30,11 @@ import threading
 import wave
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Iterator, Literal
 
 import numpy as np
 from fastapi import FastAPI
-from fastapi.responses import JSONResponse, Response
+from fastapi.responses import JSONResponse, Response, StreamingResponse
 from pydantic import BaseModel, ConfigDict, Field
 
 if TYPE_CHECKING:  # torch is imported lazily so /healthz answers during boot
@@ -42,7 +44,12 @@ MODEL_ID = "kokoro-ru"
 SAMPLE_RATE = 24000
 MODEL_DIR = Path(os.environ.get("KOKORO_MODEL_DIR", "/app/kokoro-ru"))
 DEFAULT_VOICE = os.environ.get("KOKORO_DEFAULT_VOICE", "sveta")
-THREADS = int(os.environ.get("KOKORO_THREADS", os.cpu_count() or 4))
+# Measured on the host this was tuned for (Ryzen AI 9 HX 370, 24 logical cores):
+# median end-to-end latency for a 5.6 s utterance was 1.203 s @ 4 threads,
+# 1.066 s @ 8, 0.979 s @ 12, 0.980 s @ 16, then 1.87 s @ 24. The gain stops at
+# the physical core count and SMT oversubscription costs ~2x, so cap instead of
+# trusting os.cpu_count(), which reports logical CPUs. Override on other hosts.
+THREADS = int(os.environ.get("KOKORO_THREADS", min(12, os.cpu_count() or 4)))
 # 2026-07-29, when the kokoro-ru revision we pin was published. Clients that
 # cache on this treat any change as a new model, so it must stay stable.
 MODEL_CREATED = 1785353253
@@ -223,30 +230,39 @@ class KokoroRu:
             if ps:
                 yield from split_phonemes(ps)
 
-    def synthesize(self, text: str, voice: str, speed: float) -> np.ndarray:
+    def iter_audio_chunks(self, text: str, voice: str, speed: float):
+        """Yields float32 audio per phoneme chunk, silence gaps interleaved.
+
+        The engine lock is held for the whole iteration, so a caller that stops
+        consuming early releases synthesis for everyone else.
+        """
         torch = self._torch
         assert torch is not None, "synthesize() before load()"
         stem, _gender = VOICE_SPECS[voice]
         model = self._models[stem]
         pack = self._packs[voice]
 
-        gap = torch.zeros(int(CHUNK_GAP_S * SAMPLE_RATE), dtype=torch.float32)
+        gap = np.zeros(int(CHUNK_GAP_S * SAMPLE_RATE), dtype=np.float32)
-        pieces: list[torch.Tensor] = []
         with self._lock:
-            for ps in self.phonemes(text):
+            for index, ps in enumerate(self.phonemes(text)):
                 # The style vector is picked by phoneme-string length, which is
                 # why the model sounds deterministic for identical text.
                 style = pack[len(ps) - 1]
                 # The packs ship as [510, 256]; KModel wants a batch of one.
                 if style.dim() == 1:
                     style = style.unsqueeze(0)
-                if pieces:
+                if index:
-                    pieces.append(gap)
+                    yield gap
-                pieces.append(model(ps, style, speed, return_output=True).audio)
+                yield np.asarray(
+                    model(ps, style, speed, return_output=True).audio,
+                    dtype=np.float32,
+                ).reshape(-1)
 
-        if not pieces:
+    def synthesize(self, text: str, voice: str, speed: float) -> np.ndarray:
+        chunks = list(self.iter_audio_chunks(text, voice, speed))
+        if not chunks:
             return np.zeros(0, dtype=np.float32)
-        return torch.cat(pieces).numpy().astype(np.float32, copy=False)
+        return np.concatenate(chunks)
 
 
 def encode(audio: np.ndarray, fmt: str) -> bytes:
@@ -290,6 +306,86 @@ def encode(audio: np.ndarray, fmt: str) -> bytes:
     return done.stdout
 
 
+class StreamEncoder:
+    """One long-lived ffmpeg per request: raw PCM in, encoded bytes out.
+
+    A single process is what keeps the container valid. Handing it the audio in
+    pieces as they are synthesised avoids any byte-level concatenation, whereas
+    encoding the pieces separately and joining the results would emit chained
+    Ogg for opus, which plenty of players reject.
+    """
+
+    def __init__(self, fmt: str) -> None:
+        import imageio_ffmpeg
+
+        self._proc = subprocess.Popen(
+            [
+                imageio_ffmpeg.get_ffmpeg_exe(),
+                "-hide_banner",
+                "-loglevel",
+                "error",
+                "-f",
+                "s16le",
+                "-ar",
+                str(SAMPLE_RATE),
+                "-ac",
+                "1",
+                "-i",
+                "pipe:0",
+                *FFMPEG_ARGS[fmt],
+                "-f",
+                FFMPEG_CONTAINERS[fmt],
+                "pipe:1",
+            ],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        self._blocks: queue.Queue[bytes | None] = queue.Queue()
+        self._reader = threading.Thread(target=self._pump, daemon=True)
+        self._reader.start()
+
+    def _pump(self) -> None:
+        assert self._proc.stdout is not None
+        while True:
+            block = self._proc.stdout.read(8192)
+            if not block:
+                break
+            self._blocks.put(block)
+        self._blocks.put(None)
+
+    def push(self, audio: np.ndarray) -> None:
+        assert self._proc.stdin is not None
+        clipped = np.clip(audio, -1.0, 1.0)
+        self._proc.stdin.write((clipped * 32767.0).astype("<i2").tobytes())
+        self._proc.stdin.flush()
+
+    def drain(self) -> Iterator[bytes]:
+        """Yields whatever ffmpeg has already emitted, without waiting for more."""
+        while True:
+            try:
+                block = self._blocks.get_nowait()
+            except queue.Empty:
+                return
+            if block is None:
+                return
+            yield block
+
+    def finish(self) -> Iterator[bytes]:
+        assert self._proc.stdin is not None
+        self._proc.stdin.close()
+        self._reader.join(timeout=120)
+        code = self._proc.wait(timeout=30)
+        error = self._proc.stderr.read().decode("utf-8", "replace").strip()[-400:]
+        if code != 0:
+            raise RuntimeError(error or f"ffmpeg exited with {code}")
+        yield from self.drain()
+
+    def abort(self) -> None:
+        if self._proc.poll() is None:
+            self._proc.kill()
+
+
 engine = KokoroRu()
 state: dict[str, str | None] = {"status": "loading", "error": None}
 
@@ -331,6 +427,13 @@ class SpeechRequest(BaseModel):
     speed: float | None = Field(default=None, ge=0.25, le=4.0)
 
 
+class StreamSpeechRequest(SpeechRequest):
+    # Streaming needs a container that tolerates unknown length up front, so wav
+    # (whose header declares the final sizes) and the raw formats are out. mp3
+    # and opus emit bytes as they go, which is the whole point of the endpoint.
+    response_format: Literal["mp3", "opus"] = "mp3"
+
+
 def fail(status: int, message: str, param: str | None = None, code: str | None = None) -> JSONResponse:
     return JSONResponse(
         status_code=status,
@@ -395,6 +498,58 @@ def create_speech(request: SpeechRequest) -> Response | JSONResponse:
     )
 
 
+# response_model=None for the same reason as create_speech above.
+@app.post("/v1/audio/speech/stream", response_model=None)
+def stream_speech(request: StreamSpeechRequest) -> Response | JSONResponse:
+    if state["status"] != "ready":
+        return fail(503, f"model is not ready: {state['status']}", code="model_not_ready")
+
+    voice = resolve_voice(request.voice)
+    if voice is None:
+        available = ", ".join(engine.available_voices())
+        return fail(
+            400,
+            f"unknown voice {request.voice!r}; available: {available}",
+            param="voice",
+            code="unknown_voice",
+        )
+
+    chunks = engine.iter_audio_chunks(request.input, voice, request.speed or 1.0)
+    try:
+        # Pulled before responding: once the status line is sent it cannot become
+        # a 400, and input with no speakable text has to keep failing that way.
+        first = next(chunks)
+    except StopIteration:
+        return fail(
+            400,
+            "input contains no speakable text for the Russian G2P",
+            param="input",
+            code="no_phonemes",
+        )
+
+    def body() -> Iterator[bytes]:
+        encoder = StreamEncoder(request.response_format)
+        try:
+            encoder.push(first)
+            yield from encoder.drain()
+            for chunk in chunks:
+                encoder.push(chunk)
+                yield from encoder.drain()
+            yield from encoder.finish()
+        except Exception:
+            log.exception("streaming synthesis failed")
+            raise
+        finally:
+            chunks.close()
+            encoder.abort()
+
+    return StreamingResponse(
+        body(),
+        media_type=CONTENT_TYPES[request.response_format],
+        headers={"model-id": MODEL_ID, "voice-id": voice},
+    )
+
+
 @app.get("/v1/models")
 def list_models() -> dict:
     return {