diff --git a/src/engine/openvino/kokoro.py b/src/engine/openvino/kokoro.py index 59c86ce6..02e94e4b 100644 --- a/src/engine/openvino/kokoro.py +++ b/src/engine/openvino/kokoro.py @@ -5,9 +5,14 @@ """ import asyncio +import contextlib import gc +import hashlib import json +import logging +import math import re +import types from pathlib import Path from typing import AsyncIterator, NamedTuple @@ -15,6 +20,9 @@ import openvino as ov import soundfile as sf import torch +import torch.nn as nn +import torch.nn.functional as F +from kokoro import istftnet from kokoro.model import KModel @@ -22,6 +30,198 @@ from src.server.schemas.registration import ModelLoadConfig from src.server.schemas.modeling.contract_kokoro import OV_KokoroGenConfig +logger = logging.getLogger(__name__) + +# On GPU the OpenVINO plugin compiles kernels for every new input shape (5-11 s +# each), and the vocoder's length changes with every sentence. So the vocoder is +# compiled once per frame bucket with a static shape and inputs are padded up to +# the next bucket. Static shapes also +# matter for speed: with dynamic shapes the plugin falls back to reference kernels +# for most of the graph. Each bucket holds its own weights and buffers, so the +# list stops at 512 frames (~12.8 s of audio, ~2.4 GB of VRAM in total); longer +# vocoder inputs are cut in a pause and decoded piecewise. +GPU_DECODER_BUCKETS = (64, 96, 128, 192, 256, 384, 512) +# Audio samples the decoder produces per frame (24 kHz). +SAMPLES_PER_FRAME = 600 +# CPU threads for the OpenVINO front end on the GPU path. +FRONT_CPU_THREADS = 4 +# When streaming, the first text chunk is cut to about this many characters +# (a short sentence) so the first audio is ready quickly; later chunks use the +# request's character_count_chunk. +STREAM_FIRST_CHUNK_CHARS = 120 +# Bump when the converted graphs change, to invalidate cached IRs. +_IR_VERSION = 2 + + +def _pad_reflect(x: torch.Tensor, length: int) -> torch.Tensor: + """Pad the last axis to `length` by mirroring the signal. + + The decoder's AdaIN layers normalise over the whole time axis, so zero + padding shifts their statistics and audibly changes short clips. Mirrored + content keeps the statistics close to the unpadded input. + """ + out = x + while out.shape[-1] < length: + out = torch.cat([out, out.flip(-1)], dim=-1) + return out[..., :length] + + +def _f02sine_exact(self, f0_values): + """SineGen._f02sine with explicit interpolate sizes. + + Upstream uses scale_factor=1/upsample_scale. Exported to OpenVINO, that + float scale rounds the output length off by one above ~2^20 samples + (~1536 frames, ~38 s), and the graph then fails with "Argument shapes are + inconsistent". Integer sizes are exact at any length and identical below it. + """ + if self.flag_for_pulse: + return _ORIGINAL_F02SINE(self, f0_values) + rad_values = (f0_values / self.sampling_rate) % 1 + 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 + n = rad_values.shape[1] + rad_values = F.interpolate(rad_values.transpose(1, 2), size=n // 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, size=n, mode="linear").transpose(1, 2) + return torch.sin(phase) + + +_ORIGINAL_F02SINE = istftnet.SineGen._f02sine + + +@contextlib.contextmanager +def _exact_sinegen(): + istftnet.SineGen._f02sine = _f02sine_exact + try: + yield + finally: + istftnet.SineGen._f02sine = _ORIGINAL_F02SINE + + +class _PolyphaseConvTranspose1d(nn.Module): + """ConvTranspose1d with stride s, rewritten as one Conv1d that produces the + s output phases side by side, then interleaved. Mathematically identical. + + The GPU plugin runs ConvTranspose1d on a reference kernel; on Kokoro's two + generator upsamplers that was half the vocoder's time. Conv1d runs on oneDNN. + """ + + def __init__(self, ct: nn.ConvTranspose1d): + super().__init__() + if ct.dilation[0] != 1: + raise ValueError("dilated ConvTranspose1d is not supported") + w = ct.weight.detach() # (Cin, Cout/groups, K) + cin, cout_g, kernel = w.shape + stride, groups = ct.stride[0], ct.groups + taps = math.ceil(kernel / stride) + # phase r, tap j uses kernel index r + j * stride + phases = torch.zeros(cin, cout_g, stride, taps) + for r in range(stride): + for j in range(taps): + if r + j * stride < kernel: + phases[:, :, r, j] = w[:, :, r + j * stride] + # Conv1d weight (groups * Cout/g * s, Cin/g, taps); taps flipped because + # Conv1d is a correlation and the transposed conv is a true convolution. + phases = phases.view(groups, cin // groups, cout_g, stride, taps).permute(0, 2, 3, 1, 4) + self.weight = nn.Parameter(phases.reshape(groups * cout_g * stride, cin // groups, taps).flip(-1).contiguous(), requires_grad=False) + self.bias = None + if ct.bias is not None: + self.bias = nn.Parameter(ct.bias.detach().repeat_interleave(stride), requires_grad=False) + self.stride, self.padding, self.output_padding = stride, ct.padding[0], ct.output_padding[0] + self.kernel, self.taps, self.groups, self.out_channels = kernel, taps, groups, groups * cout_g + + def forward(self, x): + length = x.shape[-1] + y = F.conv1d(F.pad(x, (self.taps - 1, self.taps - 1)), self.weight, self.bias, groups=self.groups) + batch, _, steps = y.shape + y = y.view(batch, self.out_channels, self.stride, steps).transpose(2, 3).reshape(batch, self.out_channels, steps * self.stride) + out_len = (length - 1) * self.stride - 2 * self.padding + self.kernel + self.output_padding + return y[..., self.padding : self.padding + out_len] + + +def _remove_weight_norm(module: nn.Module) -> None: + for m in module.modules(): + for hook in list(m._forward_pre_hooks.values()): + if type(hook).__name__ == "WeightNorm": + nn.utils.remove_weight_norm(m) + + +def _resblock1_folded_forward(self, x, s): + """AdaINResBlock1.forward after _fold_snake_alpha. + + Snake is x + (1/a) * sin(a x)^2 = (u + sin(u)^2) / a with u = a x. The a is + folded into the AdaIN affine before it and the 1/a into the conv after it, + which removes the per-channel multiplies the GPU ran on reference kernels. + """ + for c1, c2, n1, n2 in zip(self.convs1, self.convs2, self.adain1, self.adain2): + xt = n1(x, s) + xt = xt + torch.sin(xt) ** 2 + xt = c1(xt) + xt = n2(xt, s) + xt = xt + torch.sin(xt) ** 2 + xt = c2(xt) + x = xt + x + return x + + +@torch.no_grad() +def _fold_snake_alpha(block: "istftnet.AdaINResBlock1") -> None: + pairs = list(zip(block.convs1, block.adain1, block.alpha1)) + list(zip(block.convs2, block.adain2, block.alpha2)) + for conv, adain, alpha in pairs: + a = alpha.detach().view(-1) + channels = a.numel() + weight, bias = adain.fc.weight, adain.fc.bias # rows [0, C) gamma, [C, 2C) beta + # a * ((1 + gamma) * n + beta) == (1 + gamma') * n + beta' + weight[:channels] *= a[:, None] + bias[:channels] = a * bias[:channels] + a - 1 + weight[channels:] *= a[:, None] + bias[channels:] *= a + conv.weight /= a.view(1, -1, 1) + block.forward = types.MethodType(_resblock1_folded_forward, block) + + +def _prepare_decoder_for_gpu(decoder: nn.Module) -> None: + """Rewrite the decoder in place into GPU-friendly, numerically equivalent ops.""" + _remove_weight_norm(decoder) + generator = decoder.generator + for i, up in enumerate(generator.ups): + generator.ups[i] = _PolyphaseConvTranspose1d(up) + for block in decoder.decode: + if isinstance(getattr(block, "pool", None), nn.ConvTranspose1d): + block.pool = _PolyphaseConvTranspose1d(block.pool) + for block in list(generator.resblocks) + list(generator.noise_res): + _fold_snake_alpha(block) + + +class _Front(nn.Module): + """KModel.forward_with_tokens (kokoro 0.9.4) up to the decoder call, as a + graph OpenVINO can run on CPU. Not registered on the model: it only exists + for conversion.""" + + def __init__(self, kmodel: KModel, style_dim: int): + super().__init__() + self.kmodel = kmodel + self.style_dim = style_dim + + def forward(self, input_ids, ref_s, speed): + km = self.kmodel + tokens = input_ids.shape[-1] + input_lengths = torch.full((1,), tokens, dtype=torch.long) + text_mask = torch.gt(torch.arange(tokens).unsqueeze(0) + 1, input_lengths.unsqueeze(1)) + d_en = km.bert_encoder(km.bert(input_ids, attention_mask=(~text_mask).int())).transpose(-1, -2) + s = ref_s[:, self.style_dim :] + d = km.predictor.text_encoder(d_en, s, input_lengths, text_mask) + x, _ = km.predictor.lstm(d) + duration = torch.sigmoid(km.predictor.duration_proj(x)).sum(axis=-1) / speed + pred_dur = torch.round(duration).clamp(min=1).long().squeeze(0) + indices = torch.repeat_interleave(torch.arange(tokens), pred_dur) + alignment = (indices.unsqueeze(0) == torch.arange(tokens).unsqueeze(1)).float().unsqueeze(0) + en = d.transpose(-1, -2) @ alignment + F0_pred, N_pred = km.predictor.F0Ntrain(en, s) + asr = km.text_encoder(input_ids, input_lengths, text_mask) @ alignment + return asr, F0_pred, N_pred, pred_dur + class StreamChunk(NamedTuple): audio: torch.Tensor @@ -33,12 +233,23 @@ class StreamChunk(NamedTuple): class OV_Kokoro(KModel): """ We subclass the KModel from Kokoro to use with OpenVINO inputs. + + CPU: the exported openvino_model.xml runs end to end. + GPU: the text front end (BERT, duration and F0 predictors, text encoder) + runs as an OpenVINO graph on CPU, and the vocoder runs on the GPU as one + static-shape graph per frame bucket, rewritten so its ops land on oneDNN + kernels. Nothing compiles per request. """ - + def __init__(self, load_config: ModelLoadConfig): super().__init__() self.model = None self._device = None + self._gpu_decoders = {} + self._front_cpu = None + self._buckets = GPU_DECODER_BUCKETS + self._space_id = None + self._clause_end_ids = set() self._pipelines = {} def load_model(self, load_config: ModelLoadConfig): @@ -56,9 +267,113 @@ def load_model(self, load_config: ModelLoadConfig): core.set_property({"CACHE_DIR": load_config.cache_dir}) if load_config.runtime_config: core.set_property(load_config.runtime_config) - self.model = core.compile_model(self.model_path / "openvino_model.xml", self._device) + + if str(self._device).upper().startswith("GPU"): + self._hidden_dim = model_config["hidden_dim"] + self._style_dim = model_config["style_dim"] + self._space_id = self.vocab.get(" ") + self._clause_end_ids = {self.vocab[c] for c in ".!?;:," if c in self.vocab} + front, decoder = self._gpu_irs(core, load_config.cache_dir, model_config) + self._front_cpu = core.compile_model(front, "CPU", {"INFERENCE_NUM_THREADS": FRONT_CPU_THREADS}) + for frames in self._buckets: + static = decoder.clone() + static.reshape({ + 0: [1, self._hidden_dim, frames], + 1: [1, 2 * frames], + 2: [1, 2 * frames], + 3: [1, self._style_dim], + }) + self._gpu_decoders[frames] = core.compile_model(static, self._device) + self._warm_buckets() + self.model = self._gpu_decoders[self._buckets[-1]] + # Build the English pipelines now; the first request in each + # language would otherwise pay ~1 s for G2P and spaCy loading. + for lang_code in ("a", "b"): + self._pipeline(lang_code) + else: + self.model = core.compile_model(self.model_path / "openvino_model.xml", self._device) return self.model + def _gpu_irs(self, core: ov.Core, cache_dir, model_config: dict) -> tuple[ov.Model, ov.Model]: + """Convert the front end and the rewritten decoder, reusing cached IRs.""" + key = hashlib.sha256(f"{_IR_VERSION}:{json.dumps(model_config, sort_keys=True)}".encode()).hexdigest()[:16] + cached = {name: Path(cache_dir) / f"kokoro_{name}_{key}.xml" for name in ("front", "decoder")} if cache_dir else {} + if cached and all(p.exists() for p in cached.values()): + return core.read_model(cached["front"]), core.read_model(cached["decoder"]) + + frames, tokens = 100, 32 + with torch.no_grad(): + front = ov.convert_model( + _Front(self, self._style_dim).eval(), + example_input=(torch.ones(1, tokens, dtype=torch.long), torch.zeros(1, 2 * self._style_dim), torch.tensor(1.0)), + ) + _prepare_decoder_for_gpu(self.decoder) + example = ( + torch.zeros(1, self._hidden_dim, frames), + torch.zeros(1, 2 * frames), + torch.zeros(1, 2 * frames), + torch.zeros(1, self._style_dim), + ) + with _exact_sinegen(), torch.no_grad(): + decoder = ov.convert_model(self.decoder, example_input=example) + for name, model in (("front", front), ("decoder", decoder)): + if name in cached: + cached[name].parent.mkdir(parents=True, exist_ok=True) + ov.save_model(model, cached[name], compress_to_fp16=False) + return front, decoder + + def _warm_buckets(self) -> None: + """Run every bucket once so no request pays for first-inference setup.""" + for frames, decoder in self._gpu_decoders.items(): + decoder([ + torch.zeros(1, self._hidden_dim, frames).numpy(), + torch.zeros(1, 2 * frames).numpy(), + torch.zeros(1, 2 * frames).numpy(), + torch.zeros(1, self._style_dim).numpy(), + ]) + + def _front(self, input_ids: torch.LongTensor, ref_s: torch.FloatTensor, speed: float): + asr, F0_pred, N_pred, pred_dur = self._front_cpu([input_ids.numpy(), ref_s.numpy(), torch.tensor(float(speed)).numpy()]).to_tuple() + return (torch.from_numpy(asr), torch.from_numpy(F0_pred), torch.from_numpy(N_pred), + ref_s[:, : self._style_dim], torch.from_numpy(pred_dur)) + + def _decode_bucketed(self, asr, F0_pred, N_pred, style) -> torch.FloatTensor: + frames = asr.shape[-1] + bucket = next(b for b in self._buckets if b >= frames) + audio = self._gpu_decoders[bucket]([ + _pad_reflect(asr, bucket).numpy(), + _pad_reflect(F0_pred, 2 * bucket).numpy(), + _pad_reflect(N_pred, 2 * bucket).numpy(), + style.numpy(), + ])[0] + return torch.from_numpy(audio).reshape(-1)[: frames * SAMPLES_PER_FRAME] + + def _plan_cuts(self, input_ids: torch.LongTensor, pred_dur: torch.LongTensor) -> list[int]: + """Frame indices at which to cut the vocoder input so every piece fits the + largest bucket. Cuts go in the middle of a space token's frames, i.e. in a + pause: preferably after sentence or clause punctuation, else at any word + gap, else (no space in reach) at the bucket limit.""" + limit = self._buckets[-1] + if int(pred_dur.sum()) <= limit: + return [] + tokens = input_ids[0].tolist() + ends = torch.cumsum(pred_dur, dim=0).tolist() + starts = [0] + ends[:-1] + total = ends[-1] + clause, word = [], [] + for i in range(1, len(tokens) - 1): + if tokens[i] == self._space_id: + mid = (starts[i] + ends[i]) // 2 + (clause if tokens[i - 1] in self._clause_end_ids else word).append(mid) + cuts, pos = [], 0 + while total - pos > limit: + window = lambda c: pos < c <= pos + limit + # a clause cut anywhere in the back 60% of the window beats a later word cut + options = [c for c in clause if window(c) and c > pos + 0.4 * limit] or [c for c in word if window(c)] + pos = max(options) if options else pos + limit + cuts.append(pos) + return cuts + def _pipeline(self, lang_code: str): """One KPipeline per language. Building one loads the G2P and spaCy models, which costs ~0.9 s, so it must not happen per request.""" @@ -77,9 +392,27 @@ def forward_with_tokens( ) -> tuple[torch.FloatTensor, torch.LongTensor]: """Run the compiled OpenVINO model. Without this override KModel's PyTorch version runs on CPU and the compiled model is never called.""" + if self._gpu_decoders: + asr, F0_pred, N_pred, style, pred_dur = self._front(input_ids, ref_s, speed) + bounds = [0, *self._plan_cuts(input_ids, pred_dur), asr.shape[-1]] + pieces = [ + self._decode_bucketed(asr[..., a:b], F0_pred[..., 2 * a : 2 * b], N_pred[..., 2 * a : 2 * b], style) + for a, b in zip(bounds, bounds[1:]) + ] + return torch.cat(pieces), pred_dur outputs = self.model([input_ids, ref_s, torch.tensor(speed)]) return torch.from_numpy(outputs[0]), torch.from_numpy(outputs[1]) + def _pipeline(self, lang_code: str): + """One KPipeline per language. Building one loads the G2P and spaCy + models, which costs ~0.9 s, so it must not happen per request.""" + pipeline = self._pipelines.get(lang_code) + if pipeline is None: + from kokoro.pipeline import KPipeline + pipeline = KPipeline(model=self, lang_code=lang_code) + self._pipelines[lang_code] = pipeline + return pipeline + async def unload_model(self, registry: ModelRegistry, model_name: str) -> bool: """Unregister model from registry and free memory resources. @@ -94,6 +427,8 @@ async def unload_model(self, registry: ModelRegistry, model_name: str) -> bool: if self.model is not None: del self.model self.model = None + self._gpu_decoders = {} + self._front_cpu = None self._pipelines.clear() # Unregister from registry @@ -178,6 +513,16 @@ def make_chunks(self, text: str, chunk_size: int) -> list[str]: return chunks + def _stream_chunks(self, text: str, chunk_size: int) -> list[str]: + """make_chunks, except the first chunk is cut down to a short sentence or + clause so a streaming client gets audio as soon as possible.""" + chunks = self.make_chunks(text, chunk_size) + if not chunks or len(chunks[0]) <= STREAM_FIRST_CHUNK_CHARS: + return chunks + head = self.make_chunks(chunks[0], STREAM_FIRST_CHUNK_CHARS) + rest = " ".join(head[1:]) + return head[:1] + (self.make_chunks(rest, chunk_size) if rest else []) + chunks[1:] + async def chunk_forward_pass( self, config: OV_KokoroGenConfig ) -> AsyncIterator[StreamChunk]: @@ -191,7 +536,10 @@ async def chunk_forward_pass( # blended FloatTensor; otherwise the plain voice name. voice_arg = self._resolve_voice(config, pipeline) - text_chunks = self.make_chunks(config.input, config.character_count_chunk) + if getattr(config, "stream", False): + text_chunks = self._stream_chunks(config.input, config.character_count_chunk) + else: + text_chunks = self.make_chunks(config.input, config.character_count_chunk) total_chunks = len(text_chunks) for idx, chunk_text in enumerate(text_chunks): diff --git a/src/server/routes/openai.py b/src/server/routes/openai.py index 34014731..01d5df36 100644 --- a/src/server/routes/openai.py +++ b/src/server/routes/openai.py @@ -727,6 +727,11 @@ async def openai_audio_speech(request: OpenAISpeechRequest): raise ValueError(f"Unknown Kokoro language code: '{request.language}'. See KokoroLanguage for valid values.") if "response_format" not in gen_config.model_fields_set and request.response_format is not None: gen_config.response_format = request.response_format + if gen_config.stream: + return StreamingResponse( + _workers.stream_generate_speech_kokoro(request.model, gen_config), + media_type="audio/L16;rate=24000;channels=1", + ) result = await _workers.generate_speech_kokoro(request.model, gen_config) metrics = result.get("metrics", {}) diff --git a/src/server/schemas/modeling/contract_kokoro.py b/src/server/schemas/modeling/contract_kokoro.py index d657d3b2..78719c21 100644 --- a/src/server/schemas/modeling/contract_kokoro.py +++ b/src/server/schemas/modeling/contract_kokoro.py @@ -105,6 +105,11 @@ class OV_KokoroGenConfig(BaseModel): speed: float = Field(1.0, description="Speech speed multiplier") character_count_chunk: int = Field(400, description="Max characters per chunk") response_format: str = Field("wav", description="Output format") + stream: bool = Field( + False, + description="Stream raw int16 LE mono PCM at 24 kHz (audio/L16) chunk by chunk instead of one WAV. " + "The first chunk is kept short so audio starts quickly.", + ) @field_validator("voice_blend") @classmethod diff --git a/src/server/worker_registry.py b/src/server/worker_registry.py index 2400dd11..37de820f 100644 --- a/src/server/worker_registry.py +++ b/src/server/worker_registry.py @@ -274,6 +274,28 @@ async def infer_kokoro(packet: WorkerPacket, kokoro_model: OV_Kokoro) -> WorkerP return packet + @staticmethod + async def infer_kokoro_stream(packet: WorkerPacket, kokoro_model: OV_Kokoro) -> WorkerPacket: + """Stream Kokoro PCM chunks (int16 LE bytes) onto packet.stream_queue as each + text chunk finishes; ends with None.""" + if packet.stream_queue is None: + raise RuntimeError("infer_kokoro_stream requires stream_queue") + chunks = samples = 0 + try: + async for chunk in kokoro_model.chunk_forward_pass(packet.gen_config): + audio = chunk.audio.numpy() if hasattr(chunk.audio, "numpy") else np.asarray(chunk.audio) + pcm = np.clip(audio * 32768.0, -32768.0, 32767.0).astype(np.int16).tobytes() + await packet.stream_queue.put(pcm) + chunks += 1 + samples += len(audio) + except Exception: + logger.error("Kokoro streaming inference failed!", exc_info=True) + finally: + await packet.stream_queue.put(None) + packet.response = "" + packet.metrics = {"chunks_processed": chunks, "total_samples": samples} + return packet + @staticmethod async def infer_qwen3_tts(packet: WorkerPacket, tts_model: OVQwen3TTS) -> WorkerPacket: """Generate speech audio for a single packet using the OVQwen3TTS engine.""" @@ -469,7 +491,10 @@ async def queue_worker_kokoro(model_name: str, model_queue: asyncio.Queue, kokor logger.info(f"[Kokoro Worker: {model_name}] Shutdown signal received.") break - completed_packet = await InferWorker.infer_kokoro(packet, kokoro_model) + if getattr(packet.gen_config, "stream", False) and packet.stream_queue is not None: + completed_packet = await InferWorker.infer_kokoro_stream(packet, kokoro_model) + else: + completed_packet = await InferWorker.infer_kokoro(packet, kokoro_model) if completed_packet.metrics: logger.info(f"[Kokoro Worker: {model_name}] Metrics: {completed_packet.metrics}") @@ -963,6 +988,28 @@ async def stream_generate_speech_qwen3_tts( break yield item + async def stream_generate_speech_kokoro( + self, model_name: str, gen_config: OV_KokoroGenConfig, + ) -> AsyncIterator[bytes]: + """Stream raw int16 LE mono PCM chunks at 24 kHz (RFC 4856 audio/L16 on the HTTP layer).""" + request_id = uuid.uuid4().hex + stream_queue: asyncio.Queue = asyncio.Queue() + result_future: asyncio.Future = asyncio.get_running_loop().create_future() + packet = WorkerPacket( + request_id=request_id, + id_model=model_name, + gen_config=gen_config, + stream_queue=stream_queue, + result_future=result_future, + ) + q = self._get_kokoro_queue(model_name) + await q.put(packet) + while True: + item = await stream_queue.get() + if item is None: + break + yield item + async def generate_speech_kokoro(self, model_name: str, gen_config: OV_KokoroGenConfig) -> Dict[str, Any]: """Generate speech using a loaded Kokoro model asynchronously via worker queue. diff --git a/tests/unit/test_ov_genai_kokoro_unit.py b/tests/unit/test_ov_genai_kokoro_unit.py index 55a9cdd7..5651a6e7 100644 --- a/tests/unit/test_ov_genai_kokoro_unit.py +++ b/tests/unit/test_ov_genai_kokoro_unit.py @@ -296,6 +296,122 @@ def test_unload_model_resets_state(monkeypatch: pytest.MonkeyPatch, load_config: gc_mock.assert_called_once() + +def test_pad_reflect_mirrors_and_keeps_prefix() -> None: + x = torch.arange(1, 5, dtype=torch.float32).reshape(1, 1, 4) # [1, 2, 3, 4] + + padded = kokoro_module._pad_reflect(x, 11) + + assert padded.shape[-1] == 11 + assert padded[..., :4].tolist() == x.tolist() + assert padded.flatten().tolist() == [1, 2, 3, 4, 4, 3, 2, 1, 1, 2, 3] + assert torch.equal(kokoro_module._pad_reflect(x, 4), x) + + +def test_gpu_path_pads_to_bucket_and_trims(load_config: ModelLoadConfig) -> None: + kokoro = OV_Kokoro(load_config) + kokoro._style_dim = 128 + kokoro._buckets = (64, 128) + frames = 70 + calls = [] + + def fake_decoder(inputs): + calls.append([a.shape for a in inputs]) + return (np.ones((1, 1, inputs[0].shape[-1] * kokoro_module.SAMPLES_PER_FRAME), dtype=np.float32),) + + kokoro._gpu_decoders = {64: None, 128: fake_decoder} + pred_dur = torch.tensor([3, 4]) + kokoro._front = MagicMock(return_value=( # type: ignore[assignment] + torch.randn(1, 512, frames), torch.randn(1, 2 * frames), torch.randn(1, 2 * frames), torch.zeros(1, 128), pred_dur, + )) + + audio, dur = kokoro.forward_with_tokens(torch.tensor([[0, 5, 0]]), torch.zeros(1, 256), 1.0) + + assert calls == [[(1, 512, 128), (1, 256), (1, 256), (1, 128)]] + assert audio.shape == (frames * kokoro_module.SAMPLES_PER_FRAME,) + assert dur is pred_dur + + +def test_plan_cuts_prefers_pause_after_punctuation(load_config: ModelLoadConfig) -> None: + kokoro = OV_Kokoro(load_config) + kokoro._buckets = (100,) + kokoro._space_id, kokoro._clause_end_ids = 16, {4} + # BOS a . sp b sp c EOS + ids = torch.tensor([[0, 1, 4, 16, 2, 16, 3, 0]]) + dur = torch.tensor([10, 40, 5, 10, 30, 10, 40, 5]) # ends 10 50 55 65 95 105 145 150 + + cuts = kokoro._plan_cuts(ids, dur) + + # the word gap at 100 is later, but the clause pause at 60 is in the back 60% of the window + assert cuts == [60] + + +def test_plan_cuts_falls_back_to_word_gap_then_hard_cut(load_config: ModelLoadConfig) -> None: + kokoro = OV_Kokoro(load_config) + kokoro._buckets = (100,) + kokoro._space_id, kokoro._clause_end_ids = 16, set() + ids = torch.tensor([[0, 1, 16, 2, 3, 0]]) + dur = torch.tensor([5, 60, 10, 100, 100, 5]) # ends 5 65 75 175 275 280 + + assert kokoro._plan_cuts(ids, dur) == [70, 170, 270] # word gap, then no space in reach: hard cuts + + +def test_gpu_path_decodes_long_inputs_in_pieces(load_config: ModelLoadConfig) -> None: + kokoro = OV_Kokoro(load_config) + kokoro._buckets = (64,) + kokoro._space_id, kokoro._clause_end_ids = 16, set() + seen = [] + + def fake_decoder(inputs): + seen.append(inputs[0].shape[-1]) + return (np.zeros((1, 1, inputs[0].shape[-1] * 600), np.float32),) + + kokoro._gpu_decoders = {64: fake_decoder} + dur = torch.tensor([10, 40, 10, 40, 10]) # 110 frames; the space's middle is frame 55 + kokoro._front = MagicMock(return_value=( # type: ignore[assignment] + torch.zeros(1, 512, 110), torch.zeros(1, 220), torch.zeros(1, 220), torch.zeros(1, 128), dur, + )) + + audio, out_dur = kokoro.forward_with_tokens(torch.tensor([[0, 1, 16, 2, 0]]), torch.zeros(1, 256), 1.0) + + assert seen == [64, 64] # 55 + 55 frames, each padded to 64 + assert audio.numel() == 110 * 600 # exact length, nothing inserted + assert out_dur is dur + + +def test_polyphase_conv_transpose_is_exact() -> None: + torch.manual_seed(0) + cases = [ + dict(in_channels=8, out_channels=6, kernel_size=20, stride=10, padding=5), + dict(in_channels=6, out_channels=4, kernel_size=12, stride=6, padding=3), + dict(in_channels=10, out_channels=10, kernel_size=3, stride=2, padding=1, output_padding=1, groups=10), + dict(in_channels=4, out_channels=4, kernel_size=5, stride=3, padding=0, bias=False), + ] + for case in cases: + ct = torch.nn.ConvTranspose1d(**case) + poly = kokoro_module._PolyphaseConvTranspose1d(ct) + x = torch.randn(2, case["in_channels"], 37) + with torch.no_grad(): + expected, actual = ct(x), poly(x) + assert actual.shape == expected.shape, case + assert torch.allclose(actual, expected, atol=1e-5), case + + +def test_fold_snake_alpha_is_exact() -> None: + torch.manual_seed(0) + block = kokoro_module.istftnet.AdaINResBlock1(16, 3, (1, 3, 5), style_dim=8) + with torch.no_grad(): + for p in list(block.alpha1) + list(block.alpha2): + p.uniform_(0.3, 2.0) + kokoro_module._remove_weight_norm(block) + x, s = torch.randn(1, 16, 50), torch.randn(1, 8) + with torch.no_grad(): + expected = block(x, s) + kokoro_module._fold_snake_alpha(block) + actual = block(x, s) + assert torch.allclose(actual, expected, atol=1e-4) + + def test_pipeline_is_built_once_per_language(monkeypatch: pytest.MonkeyPatch, load_config: ModelLoadConfig) -> None: built = [] @@ -311,3 +427,34 @@ def __init__(self, model, lang_code): kokoro._pipeline("b") assert built == ["a", "b"] + + +def test_exact_sinegen_matches_upstream_below_float_limit() -> None: + # Same seed, same input: the size-based interpolation must reproduce the + # upstream scale-factor version exactly where the latter is still exact. + gen = kokoro_module.istftnet.SineGen(24000, upsample_scale=300, harmonic_num=8) + f0 = torch.full((1, 300 * 40, 1), 180.0) + f0_values = f0 * torch.arange(1, 10).view(1, 1, 9) + + torch.manual_seed(0) + expected = kokoro_module._ORIGINAL_F02SINE(gen, f0_values.clone()) + torch.manual_seed(0) + actual = kokoro_module._f02sine_exact(gen, f0_values.clone()) + + assert torch.allclose(actual, expected, atol=1e-5) + assert kokoro_module.istftnet.SineGen._f02sine is kokoro_module._ORIGINAL_F02SINE + with kokoro_module._exact_sinegen(): + assert kokoro_module.istftnet.SineGen._f02sine is kokoro_module._f02sine_exact + assert kokoro_module.istftnet.SineGen._f02sine is kokoro_module._ORIGINAL_F02SINE + + +def test_stream_chunks_shortens_only_the_first_chunk(load_config: ModelLoadConfig) -> None: + kokoro = OV_Kokoro(load_config) + first = "This opening sentence is deliberately long, so that it runs well past the streaming limit, which is about one hundred and twenty characters." + text = first + " Second sentence. Third sentence here." + + chunks = kokoro._stream_chunks(text, 400) + + assert len(chunks[0]) <= kokoro_module.STREAM_FIRST_CHUNK_CHARS + assert " ".join(chunks).split() == text.split() + assert kokoro._stream_chunks("Short. Text.", 400) == kokoro.make_chunks("Short. Text.", 400) diff --git a/tests/unit/test_worker_registry_unit.py b/tests/unit/test_worker_registry_unit.py index 415f14ae..b4c874ba 100644 --- a/tests/unit/test_worker_registry_unit.py +++ b/tests/unit/test_worker_registry_unit.py @@ -402,3 +402,44 @@ def test_missing_model_queue(worker_registry: worker_module.WorkerRegistry) -> N with pytest.raises(ValueError): worker_registry._get_model_queue("missing") + + +def test_infer_kokoro_stream_puts_pcm_per_chunk_then_none() -> None: + import numpy as np + import torch + from src.engine.openvino.kokoro import StreamChunk + + class FakeKokoro: + async def chunk_forward_pass(self, config): + for i, audio in enumerate((torch.tensor([0.0, 0.5, -0.5]), torch.tensor([1.0, -1.0]))): + yield StreamChunk(audio=audio, chunk_text=f"c{i}", chunk_index=i, total_chunks=2) + + async def _run(): + queue: asyncio.Queue = asyncio.Queue() + packet = worker_module.WorkerPacket(request_id="r", id_model="k", gen_config=None, stream_queue=queue) + done = await worker_module.InferWorker.infer_kokoro_stream(packet, FakeKokoro()) # type: ignore[arg-type] + items = [] + while not queue.empty(): + items.append(queue.get_nowait()) + return done, items + + done, items = asyncio.run(_run()) + assert items[-1] is None and len(items) == 3 + assert np.frombuffer(items[0], dtype=np.int16).tolist() == [0, 16384, -16384] + assert np.frombuffer(items[1], dtype=np.int16).tolist() == [32767, -32768] + assert done.metrics == {"chunks_processed": 2, "total_samples": 5} + + +def test_infer_kokoro_stream_ends_stream_on_error() -> None: + class BrokenKokoro: + async def chunk_forward_pass(self, config): + raise RuntimeError("boom") + yield # pragma: no cover + + async def _run(): + queue: asyncio.Queue = asyncio.Queue() + packet = worker_module.WorkerPacket(request_id="r", id_model="k", gen_config=None, stream_queue=queue) + await worker_module.InferWorker.infer_kokoro_stream(packet, BrokenKokoro()) # type: ignore[arg-type] + return [queue.get_nowait() for _ in range(queue.qsize())] + + assert asyncio.run(_run()) == [None]