oqyude/nixos
1
0
Fork 0
mirror of https://github.com/oqyude/nixos.git synced 2026-10-07 20:47:14 +03:00
Code Issues Packages Projects Releases Wiki Activity

kokoro-tts stream added

Browse Source
This commit is contained in:
oqyude
2026-10-03 15:20:33 +03:00
parent 509fd3dde0
commit 1c77ae658e
4 changed files with 213 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -1,2 +1,3 @@
.vscode
.omo
__pycache__
modules/containers/kokoro-tts.nix
+22 -23
View File
@@ -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" [
{
name = "Dockerfile";
path = toString ./kokoro-tts/Dockerfile;
}
{
name = "app.py";
path = toString ./kokoro-tts/app.py;
}
{
name = "fetch_assets.py";
path = toString ./kokoro-tts/fetch_assets.py;
}
{
name = "requirements.txt";
path = toString ./kokoro-tts/requirements.txt;
}
];
#
# runCommand rather than linkFarm: linkFarm entries are symlinks into other
# store paths, and `podman build` only mounts the context root, so every COPY
# 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" { } ''
mkdir -p "$out"
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"
cp -L ${./kokoro-tts/requirements.txt} "$out/requirements.txt"
'';
image = "localhost/kokoro-tts:latest";
@@ -64,11 +58,16 @@ in
];
environment = {
# Inference is CPU-bound and already threaded inside torch; these
# keep it from oversubscribing a small machine.
KOKORO_THREADS = "4";
OMP_NUM_THREADS = "4";
MKL_NUM_THREADS = "4";
# Inference is CPU-bound and already threaded inside torch. Measured
# on a 24-logical-core host: median end-to-end latency for a 5.6 s
# utterance was 1.203 s at 4 threads, 0.979 s at 12, 0.980 s at 16
# and 1.87 s at 24, so the useful ceiling is the physical core count
# 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";
};
modules/containers/kokoro-tts/Dockerfile
+17 -5
View File
@@ -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 \
OMP_NUM_THREADS=4 \
MKL_NUM_THREADS=4 \
KOKORO_THREADS=12 \
OMP_NUM_THREADS=12 \
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
+171 -16
View File
@@ -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
GET /v1/models OpenAI model list
GET /v1/voices voice inventory (extension, not part of OpenAI)
GET /healthz readiness, 503 until the model is loaded
POST /v1/audio/speech OpenAI text-to-speech
POST /v1/audio/speech/stream same, but mp3/opus emitted while synthesising
GET /v1/models OpenAI model list
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)
pieces: list[torch.Tensor] = []
gap = np.zeros(int(CHUNK_GAP_S * SAMPLE_RATE), dtype=np.float32)
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:
pieces.append(gap)
pieces.append(model(ps, style, speed, return_output=True).audio)
if index:
yield gap
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 {