diff --git a/api/src/inference/kokoro_v1.py b/api/src/inference/kokoro_v1.py index dc11ae3a..81c76a90 100644 --- a/api/src/inference/kokoro_v1.py +++ b/api/src/inference/kokoro_v1.py @@ -14,6 +14,69 @@ from ..structures.schemas import WordTimestamp from .base import AudioChunk, BaseModelBackend +_ESPEAK_TS_SCALE = 2.0 / 80.0 # pred_dur unit -> seconds (matches KPipeline.join_timestamps) + + +def _espeak_word_timestamps(graphemes, phonemes, pred_dur, g2p=None): + """Derive word timestamps for espeak-based (non-English) pipelines. + + KPipeline attaches timed tokens only for English (misaki G2P) voices, but + the model predicts a duration for every phoneme in every language — the + audio is rendered from exactly those durations. pred_dur[0] is BOS and + pred_dur[1 + i] covers phonemes[i], so summing per-character durations and + splitting at the phoneme string's spaces yields exact word times. + + Words espeak expands into several spoken words (numbers, some + abbreviations) are reconciled by re-phonemizing per word via `g2p`; if the + counts still disagree, returns None so the caller emits no timestamps and + clients can fall back to their own handling. + """ + try: + if pred_dur is None or not phonemes or len(pred_dur) < len(phonemes) + 1: + return None + words = graphemes.split() + if not words: + return None + groups = [] # [start, end) seconds per space-separated phoneme group + t = float(pred_dur[0]) * _ESPEAK_TS_SCALE + start = None + for i, ch in enumerate(phonemes): + d = float(pred_dur[1 + i]) * _ESPEAK_TS_SCALE + if ch.isspace(): + if start is not None: + groups.append((start, t)) + start = None + elif start is None: + start = t + t += d + if start is not None: + groups.append((start, t)) + + if len(groups) != len(words): + if g2p is None: + return None + counts = [] + for w in words: + ps = g2p(w) + if isinstance(ps, tuple): + ps = ps[0] + counts.append(max(len((ps or "").split()), 1)) + if sum(counts) != len(groups): + return None + merged, gi = [], 0 + for c in counts: + merged.append((groups[gi][0], groups[gi + c - 1][1])) + gi += c + groups = merged + + return [ + WordTimestamp(word=w, start_time=round(s, 3), end_time=round(e, 3)) + for w, (s, e) in zip(words, groups) + ] + except Exception as e: + logger.warning(f"espeak timestamp mapping failed: {e}") + return None + class KokoroV1(BaseModelBackend): """Kokoro backend with controlled resource management.""" @@ -326,6 +389,29 @@ async def generate( logger.error( f"Failed to process timestamps for chunk: {e}" ) + elif ( + return_timestamps + and result.phonemes + and type(getattr(pipeline, "g2p", None)).__name__ + == "EspeakG2P" + ): + # espeak pipelines (es/fr/it/hi/pt) yield no timed + # tokens; derive word times from the model's own + # phoneme durations instead. Espeak-only: other + # non-English G2Ps (ja/zh) have no space-separated + # word groups to map. + pred_dur = getattr(result, "pred_dur", None) + if ( + pred_dur is None + and getattr(result, "output", None) is not None + ): + pred_dur = getattr(result.output, "pred_dur", None) + word_timestamps = _espeak_word_timestamps( + result.graphemes, + result.phonemes, + pred_dur, + g2p=getattr(pipeline, "g2p", None), + ) yield AudioChunk( result.audio.numpy(), word_timestamps=word_timestamps diff --git a/api/tests/test_kokoro_v1.py b/api/tests/test_kokoro_v1.py index ff02d9a0..f43da344 100644 --- a/api/tests/test_kokoro_v1.py +++ b/api/tests/test_kokoro_v1.py @@ -158,3 +158,66 @@ async def test_generate_uses_correct_pipeline(kokoro_backend): voice_arg = Path(mock_pipeline.call_args[1]["voice"]) assert voice_arg.name == "temp_voice_ef_voice" assert voice_arg.parent == Path(tempfile.gettempdir()) + + +def test_espeak_word_timestamps_basic(): + """Word times derived from pred_dur for an espeak (non-English) result.""" + from api.src.inference.kokoro_v1 import _espeak_word_timestamps + + # phonemes "ab cd": BOS=8, a=4, b=4, space=8, c=4, d=4, EOS=0 + # scale is 2/80 s per unit -> BOS 0.2s, each char 0.1s, space 0.2s + pred_dur = torch.tensor([8, 4, 4, 8, 4, 4, 0]) + ts = _espeak_word_timestamps("hola mundo", "ab cd", pred_dur) + + assert [t.word for t in ts] == ["hola", "mundo"] + assert ts[0].start_time == pytest.approx(0.2) + assert ts[0].end_time == pytest.approx(0.4) + assert ts[1].start_time == pytest.approx(0.6) + assert ts[1].end_time == pytest.approx(0.8) + + +def test_espeak_word_timestamps_expansion(): + """Words espeak expands (numbers) are merged back via per-word g2p.""" + from api.src.inference.kokoro_v1 import _espeak_word_timestamps + + # Three text words but four phoneme groups ("1863" speaks as two words). + phonemes = "a bb cc d" + pred_dur = torch.tensor([0, 4] + [8, 4, 4] + [8, 4, 4] + [8, 4] + [0]) + + def g2p(word): + return ("bb cc" if word == "1863" else "x", None) + + ts = _espeak_word_timestamps("en 1863 el", phonemes, pred_dur, g2p=g2p) + + assert [t.word for t in ts] == ["en", "1863", "el"] + # "1863" spans both expanded groups. + assert ts[1].start_time == pytest.approx(0.3) + assert ts[1].end_time == pytest.approx(0.9) + # Whole-string timing stays monotonic and inside the clip. + assert ts[0].end_time <= ts[1].start_time <= ts[2].start_time + + +def test_espeak_word_timestamps_unreconcilable(): + """Group/word count mismatch that g2p can't explain yields None.""" + from api.src.inference.kokoro_v1 import _espeak_word_timestamps + + phonemes = "a b c" + pred_dur = torch.tensor([0, 4, 8, 4, 8, 4, 0]) + + # g2p says every word is a single group: 2 != 3 -> give up. + ts = _espeak_word_timestamps("uno dos", phonemes, pred_dur, g2p=lambda w: (w, None)) + assert ts is None + + # No g2p available -> also None. + assert _espeak_word_timestamps("uno dos", phonemes, pred_dur) is None + + +def test_espeak_word_timestamps_degenerate_inputs(): + """Missing durations or empty text return None instead of raising.""" + from api.src.inference.kokoro_v1 import _espeak_word_timestamps + + assert _espeak_word_timestamps("hola", "abc", None) is None + assert _espeak_word_timestamps("hola", "", torch.tensor([0, 0])) is None + assert _espeak_word_timestamps("", "abc", torch.tensor([0, 4, 4, 4, 0])) is None + # pred_dur too short for the phoneme string. + assert _espeak_word_timestamps("hola", "abcdef", torch.tensor([0, 4])) is None