From a782a59cfcb2eb72259a129b99501da8f62f25b6 Mon Sep 17 00:00:00 2001 From: SearchSavior Date: Sat, 12 Sep 2026 15:06:17 -0400 Subject: [PATCH 1/2] moce code predictor and codec embedding onto gpu device --- config.yaml | 46 +++++++++++++++++----- src/engine/openvino/qwen3_tts/qwen3_tts.py | 4 +- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/config.yaml b/config.yaml index 7c16a87..88c2286 100644 --- a/config.yaml +++ b/config.yaml @@ -27,20 +27,48 @@ models: enable_prefix_caching: use_cache_eviction: use_sparse_attention: - samplers: + sampler_config: temperature: 0.7 top_k: 40 top_p: 0.95 repetition_penalty: 1.05 max_tokens: 1024 + qwen3-8b-spec: + load_config: + engine: ovgenai + model_type: vlm + model_path: /mnt/Ironwolf-4TB/Models/OpenVINO/Qwen/Qwen3-8B-ShiningValiant3-int4-asym-ov + device: GPU.0 + tool_call_parser: hermes + draft_model_path: "/mnt/Ironwolf-4TB/Models/OpenVINO/Qwen/Qwen3-pruned-6L-from-0.6B-int8-ov" + draft_device: "CPU" + num_assistant_tokens: 5 + runtime_config: + PERFORMANCE_HINT: LATENCY + scheduler_config: + max_num_batched_tokens: + num_kv_blocks: + cache_size: + num_linear_attention_blocks: + cache_interval_multiplier: + dynamic_split_fuse: + enable_prefix_caching: + use_cache_eviction: + use_sparse_attention: + sampler_config: + temperature: 0.7 + top_k: 40 + top_p: 0.95 + repetition_penalty: 1.05 + max_tokens: 1024 kokoro: engine: openvino model_type: kokoro model_path: device: CPU - kokoro_options: + kokoro_config: voice: af_sarah voice_blend: af_heart:0.7,af_nicole:0.3 lang_code: @@ -52,7 +80,7 @@ models: model_type: qwen3_asr model_path: device: CPU - qwen3_asr_options: + qwen3_asr_config: language: max_tokens: max_chunk_sec: @@ -65,11 +93,11 @@ models: model_path: engine: openvino device: GPU.0 - qwen3_tts_custom_voice_options: + qwen3_tts_custom_voice_config: language: speaker: instruct: - qwen3_tts_options: + qwen3_tts_config: max_new_tokens: do_sample: top_k: @@ -90,9 +118,9 @@ models: model_path: engine: openvino device: GPU.0 - qwen3_tts_voice_design_options: + qwen3_tts_voice_design_config: voice_description: - qwen3_tts_options: + qwen3_tts_config: max_new_tokens: do_sample: top_k: @@ -112,11 +140,11 @@ models: model_path: engine: openvino device: GPU.0 - qwen3_tts_voice_clone_options: + qwen3_tts_voice_clone_config: ref_text: x_vector_only: instruct: - qwen3_tts_options: + qwen3_tts_config: max_new_tokens: do_sample: top_k: diff --git a/src/engine/openvino/qwen3_tts/qwen3_tts.py b/src/engine/openvino/qwen3_tts/qwen3_tts.py index 7b0a751..0328b90 100644 --- a/src/engine/openvino/qwen3_tts/qwen3_tts.py +++ b/src/engine/openvino/qwen3_tts/qwen3_tts.py @@ -134,7 +134,7 @@ def load_model(self, load_config: ModelLoadConfig) -> None: self._text_model_c = core.compile_model(str(p / "text_model.xml"), device, _hint) self._codec_emb_c = core.compile_model(str(p / "codec_embedding.xml"), device, _hint) # Code predictor: many tiny inferences per frame; CPU avoids GPU launch/transfer overhead. - self._cp_codec_emb_c = core.compile_model(str(p / "cp_codec_embedding.xml"), "CPU", _hint) + self._cp_codec_emb_c = core.compile_model(str(p / "cp_codec_embedding.xml"), "GPU.0", _hint) # Speech decoder: single-shot vocoding; CPU fits typical sequence lengths without GPU overhead. self._decoder_c = core.compile_model( str(p / "speech_tokenizer" / "speech_decoder.xml"), "CPU", _hint, @@ -143,7 +143,7 @@ def load_model(self, load_config: ModelLoadConfig) -> None: talker_c = core.compile_model(str(p / "talker.xml"), device, _hint) self._talker_req = talker_c.create_infer_request() - cp_c = core.compile_model(str(p / "code_predictor.xml"), "CPU", _hint) + cp_c = core.compile_model(str(p / "code_predictor.xml"), "GPU.0", _hint) self._cp_req = cp_c.create_infer_request() if "GPU" in device: logger.info( From 8528ffcf8b7fa2ac930d58ab55349d8d5aeb4043 Mon Sep 17 00:00:00 2001 From: SearchSavior Date: Mon, 14 Sep 2026 22:10:17 -0400 Subject: [PATCH 2/2] fix kokoro chunking --- src/engine/openvino/kokoro.py | 115 ++++++++++++------ .../schemas/modeling/contract_kokoro.py | 2 +- src/server/worker_registry.py | 3 - tests/unit/test_ov_genai_kokoro_unit.py | 103 ++++++++++++++++ 4 files changed, 181 insertions(+), 42 deletions(-) diff --git a/src/engine/openvino/kokoro.py b/src/engine/openvino/kokoro.py index c905c98..cf080da 100644 --- a/src/engine/openvino/kokoro.py +++ b/src/engine/openvino/kokoro.py @@ -82,46 +82,76 @@ async def unload_model(self, registry: ModelRegistry, model_name: str) -> bool: return removed - def make_chunks(self, text: str, chunk_size: int) -> list[str]: - """ - Split text into chunks up to `chunk_size` characters, - preferring sentence boundaries. - """ - if len(text) <= chunk_size: - return [text] + # Sentence end: .!? (optionally followed by quotes/brackets), then whitespace. + _SENTENCE_RE = re.compile(r'(?<=[.!?…])["\')\]]*\s+') + # Clause punctuation worth pausing on when a mid-sentence split is needed. + _CLAUSE_PUNCT = ",;:\u2014\u2013-" + # Prefer clause splits at least this far into the chunk to avoid + # degenerate tiny heads (e.g. a comma at position 4). + _MIN_CLAUSE_FRACTION = 0.3 - chunks = [] - current_chunk = "" + @staticmethod + def _split_head(text: str, limit: int) -> tuple[str, str]: + """Split off a head of at most `limit` chars, preferring a clause + boundary, then a word boundary, then a hard cut. Keeps punctuation + on the head so the model still hears the pause.""" + head = text[:limit] + cut = -1 + best = max(head.rfind(p) for p in OV_Kokoro._CLAUSE_PUNCT) + if best >= int(limit * OV_Kokoro._MIN_CLAUSE_FRACTION): + cut = best + 1 + else: + cut = head.rfind(" ") + if cut <= 0: + cut = limit + return text[:cut].strip(), text[cut:].strip() - # Regex: split after ., !, ? followed by space - sentences = re.split(r'(?<=[.!?]) +', text) + def make_chunks(self, text: str, chunk_size: int) -> list[str]: + """ + Split text into chunks of at most `chunk_size` characters. - for sentence in sentences: - sentence = sentence.strip() - if not sentence: + Boundary preference: paragraph (blank line) -> sentence -> clause + (comma/semicolon/colon/dash) -> word -> hard cut. Guaranteed: every + returned chunk respects the size limit and no character is ever + dropped (concatenation reproduces the input modulo whitespace). + """ + if not text or not text.strip(): + return [] + if len(text.strip()) <= chunk_size: + return [text.strip()] + + # Paragraph breaks are hard boundaries; sentences within paragraphs. + segments: list[str] = [] + for paragraph in re.split(r'\n\s*\n+|\n', text): + paragraph = paragraph.strip() + if not paragraph: continue - - if len(current_chunk) + len(sentence) > chunk_size: - if current_chunk: - chunks.append(current_chunk.strip()) - current_chunk = sentence - else: - # sentence itself longer than chunk_size -> word splitting - words = sentence.split() - temp = "" - for word in words: - if len(temp) + len(word) + 1 > chunk_size: - if temp: - chunks.append(temp.strip()) - temp = word - else: - temp += (" " if temp else "") + word - current_chunk = temp + segments.extend(s for s in self._SENTENCE_RE.split(paragraph) if s.strip()) + + chunks: list[str] = [] + current = "" + + for seg in segments: + if len(seg) > chunk_size: + # Oversized segment: flush buffer, then carve off heads until + # the remainder fits. The tail becomes the new buffer so it + # can merge with the next sentence. + if current: + chunks.append(current) + current = "" + while len(seg) > chunk_size: + head, seg = self._split_head(seg, chunk_size) + chunks.append(head) + if not current: + current = seg + elif len(current) + 1 + len(seg) <= chunk_size: + current = f"{current} {seg}" else: - current_chunk += (" " if current_chunk else "") + sentence + chunks.append(current) + current = seg - if current_chunk: - chunks.append(current_chunk.strip()) + if current: + chunks.append(current) return chunks @@ -149,14 +179,23 @@ def infer_on_chunk(): """Blocking inference run in background thread.""" with torch.no_grad(): infer = pipeline(chunk_text, voice=voice_arg, speed=config.speed) - result = next(infer) if hasattr(infer, "__iter__") else infer - return result + if not hasattr(infer, "__iter__"): + return torch.as_tensor(infer.audio) + # KPipeline packs text into <=context_length phoneme + # buckets and yields one result per bucket. Consume ALL + # of them; taking only the first silently drops audio. + parts = [torch.as_tensor(r.audio) for r in infer] + if not parts: + return None + return parts[0] if len(parts) == 1 else torch.cat(parts, dim=-1) # Run blocking inference off the main loop - result = await asyncio.to_thread(infer_on_chunk) + audio = await asyncio.to_thread(infer_on_chunk) + if audio is None: + continue yield StreamChunk( - audio=result.audio, + audio=audio, chunk_text=chunk_text, chunk_index=idx, total_chunks=total_chunks, diff --git a/src/server/schemas/modeling/contract_kokoro.py b/src/server/schemas/modeling/contract_kokoro.py index 15bc46d..d657d3b 100644 --- a/src/server/schemas/modeling/contract_kokoro.py +++ b/src/server/schemas/modeling/contract_kokoro.py @@ -103,7 +103,7 @@ class OV_KokoroGenConfig(BaseModel): ) lang_code: KokoroLanguage = Field(KokoroLanguage.AMERICAN_ENGLISH, description="Language code for the voice") speed: float = Field(1.0, description="Speech speed multiplier") - character_count_chunk: int = Field(100, description="Max characters per chunk") + character_count_chunk: int = Field(400, description="Max characters per chunk") response_format: str = Field("wav", description="Output format") @field_validator("voice_blend") diff --git a/src/server/worker_registry.py b/src/server/worker_registry.py index 3a3da5a..df2ccb3 100644 --- a/src/server/worker_registry.py +++ b/src/server/worker_registry.py @@ -243,12 +243,10 @@ async def infer_kokoro(packet: WorkerPacket, kokoro_model: OV_Kokoro) -> WorkerP then converts to bytes for response. """ audio_chunks = [] - chunk_texts = [] try: async for chunk in kokoro_model.chunk_forward_pass(packet.gen_config): audio_chunks.append(chunk.audio) - chunk_texts.append(chunk.chunk_text) if audio_chunks: # Concatenate all audio chunks @@ -268,7 +266,6 @@ async def infer_kokoro(packet: WorkerPacket, kokoro_model: OV_Kokoro) -> WorkerP # Add some basic metrics packet.metrics = { "chunks_processed": len(audio_chunks), - "chunk_texts": chunk_texts, "total_samples": sum(len(chunk) for chunk in audio_chunks) if audio_chunks else 0 } except Exception as e: diff --git a/tests/unit/test_ov_genai_kokoro_unit.py b/tests/unit/test_ov_genai_kokoro_unit.py index e3c1803..bb9239f 100644 --- a/tests/unit/test_ov_genai_kokoro_unit.py +++ b/tests/unit/test_ov_genai_kokoro_unit.py @@ -151,6 +151,109 @@ async def _run_test(): assert pipeline_calls[1][0] == "call" +def test_make_chunks_splits_overlong_sentence_after_flush(load_config: ModelLoadConfig) -> None: + """Regression: an oversized sentence arriving right after a buffer flush + must still be split, not passed through whole.""" + kokoro = OV_Kokoro(load_config) + + short = "Short one." + long_sentence = "word " * 30 + "end" # 154 chars, no punctuation + text = f"{short} {long_sentence}" + chunks = kokoro.make_chunks(text, chunk_size=50) + + assert all(len(chunk) <= 50 for chunk in chunks) + joined = " ".join(chunks) + assert "end" in joined + assert joined.split().count("word") == 30 + + +def test_make_chunks_size_and_lossless_invariants(load_config: ModelLoadConfig) -> None: + """Every chunk respects the size limit and no word is ever dropped.""" + kokoro = OV_Kokoro(load_config) + + text = ( + "First sentence here! Followed by a question? And one more, with a " + "clause, to split on; plus a colon: like this. " * 8 + ) + chunks = kokoro.make_chunks(text, chunk_size=120) + + assert len(chunks) > 1 + assert all(len(chunk) <= 120 for chunk in chunks) + assert " ".join(chunks).split() == text.split() + + +def test_make_chunks_prefers_clause_boundary(load_config: ModelLoadConfig) -> None: + """Mid-sentence splits should land on clause punctuation when available.""" + kokoro = OV_Kokoro(load_config) + + text = "before the comma there are words, after it there are more words and yet more and more words" + chunks = kokoro.make_chunks(text, chunk_size=40) + + assert len(chunks) >= 2 + assert chunks[0].endswith(",") + assert all(len(chunk) <= 40 for chunk in chunks) + + +def test_make_chunks_handles_newlines_and_empty(load_config: ModelLoadConfig) -> None: + kokoro = OV_Kokoro(load_config) + + assert kokoro.make_chunks("", 100) == [] + assert kokoro.make_chunks(" \n ", 100) == [] + + chunks = kokoro.make_chunks("Para one is short.\n\nPara two is also short.", 100) + assert len(chunks) == 2 + + +def test_chunk_forward_pass_concatenates_multi_bucket_results( + monkeypatch: pytest.MonkeyPatch, load_config: ModelLoadConfig +) -> None: + """Regression: when KPipeline yields multiple buckets for one text chunk, + all audio must be concatenated — taking only the first drops speech.""" + import torch + + kokoro = OV_Kokoro(load_config) + kokoro.model = object() + kokoro.make_chunks = MagicMock(return_value=["Only chunk"]) # type: ignore[assignment] + + async def immediate_to_thread(func, *args, **kwargs): # type: ignore[override] + return func(*args, **kwargs) + + monkeypatch.setattr(kokoro_module.asyncio, "to_thread", immediate_to_thread) + + class DummyResult: + def __init__(self, audio) -> None: + self.audio = audio + + class DummyPipeline: + def __init__(self, model, lang_code): + pass + + def __call__(self, text, voice, speed): + yield DummyResult(torch.zeros(10)) + yield DummyResult(torch.ones(10)) + + monkeypatch.setattr("kokoro.pipeline.KPipeline", DummyPipeline) + + config = OV_KokoroGenConfig( + input="ignored", + voice=KokoroVoice.AF_SARAH, + lang_code=KokoroLanguage.AMERICAN_ENGLISH, + speed=1.0, + character_count_chunk=50, + response_format="wav", + ) + + async def _run_test(): + return [item async for item in kokoro.chunk_forward_pass(config)] + + chunks = asyncio.run(_run_test()) + + assert len(chunks) == 1 + assert chunks[0].audio.shape[0] == 20 + assert int(chunks[0].audio[:10].sum()) == 0 + assert int(chunks[0].audio[10:].sum()) == 10 + + def test_unload_model_resets_state(monkeypatch: pytest.MonkeyPatch, load_config: ModelLoadConfig) -> None: kokoro = OV_Kokoro(load_config) kokoro.model = object()