Fix missing word timestamps for non-English (espeak) voices

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

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