abogen/plugins/kokoro/engine.py

"""Kokoro Engine adapter for the TTS Plugin Architecture.

This module adapts the existing Kokoro backend to the new Engine/EngineSession
protocol. It wraps the KokoroBackend without modifying it.

Language mapping: this is the engine's responsibility. The engine knows
which languages it supports and converts Language enum → internal format.
Callers outside this module never see engine-specific codes.
"""

from __future__ import annotations

import logging
from typing import Any

import numpy as np

from abogen.domain.enums import Language
from abogen.tts_plugin.capabilities import VoiceLister
from abogen.tts_plugin.engine import Engine, EngineSession
from abogen.tts_plugin.errors import EngineError
from abogen.tts_plugin.manifest import VoiceManifest
from abogen.tts_plugin.types import (
    AudioFormat,
    AudioSegment,
    Duration,
    SynthesisRequest,
    SynthesizedAudio,
    TokenTiming,
)

logger = logging.getLogger(__name__)

# Sample rate for Kokoro audio
_KOKORO_SAMPLE_RATE = 24000

# Engine-internal language mapping: Language enum → kokoro code.
# ONLY visible inside this module — callers never see kokoro codes.
_KOKORO_LANG_MAP: dict[Language, str] = {
    Language.EN_US: "a",
    Language.EN_GB: "b",
    Language.ES: "e",
    Language.FR: "f",
    Language.HI: "h",
    Language.IT: "i",
    Language.JA: "j",
    Language.PT_BR: "p",
    Language.ZH: "z",
}

# Reverse mapping: engine-internal code → Language enum.
# Used by voice catalog and other places that need to convert
# engine codes back to Language enum (e.g. voice ID prefix extraction).
_CODE_TO_LANGUAGE: dict[str, Language] = {v: k for k, v in _KOKORO_LANG_MAP.items()}


def supported_languages() -> list[Language]:
    """Return the list of Language enum values this engine supports.

    This is the engine's responsibility — the engine knows which
    languages it supports and exposes them as Language enum values.
    UI layers query this to populate language selectors.
    """
    return list(_KOKORO_LANG_MAP.keys())


def engine_language(lang: Language) -> str:
    """Map a Language enum to the engine's internal code.

    This is the engine's responsibility — the engine owns the mapping
    between Language enum and its internal format. Callers pass Language
    enum; the engine converts internally. The returned string is ONLY
    used inside the engine implementation.
    """
    return _KOKORO_LANG_MAP.get(lang, "a")


def language_for_code(code: str | None) -> Language:
    """Map a kokoro engine language code (single letter) to a Language enum.

    Used to resolve legacy data such as old profile files that stored
    kokoro letter codes. This is kokoro-specific knowledge that stays
    inside the engine. Unparseable values fall back to EN_US.
    """
    letter = str(code or "").strip()[:1].lower()
    if letter in _CODE_TO_LANGUAGE:
        return _CODE_TO_LANGUAGE[letter]
    return Language.EN_US


def language_for_voice_id(voice_id: str | None) -> Language:
    """Determine which Language a voice belongs to from its voice ID.

    Kokoro voice IDs encode language as a prefix (e.g. "af_heart" → "a" → EN_US).
    This is kokoro-specific knowledge that stays inside the engine.
    Callers pass a voice ID string; the engine returns a Language enum.
    """
    return language_for_code(voice_id)


class KokoroSession:
    """EngineSession implementation for Kokoro.

    Owns mutable execution state for synthesis.
    NOT thread-safe.
    """

    def __init__(self, pipeline: Any) -> None:
        self._pipeline = pipeline
        self._disposed = False

    def synthesize(self, request: SynthesisRequest) -> SynthesizedAudio:
        """Synthesize audio from text using Kokoro."""
        if self._disposed:
            raise EngineError("Session disposed")

        try:
            voice = request.voice.key
            speed = request.parameters.values.get("speed", 1.0)
            split_pattern = request.parameters.values.get("split_pattern", None)

            sample_rate = _KOKORO_SAMPLE_RATE
            audio_parts: list[np.ndarray] = []
            segments: list[AudioSegment] = []
            for segment in self._pipeline(
                request.text,
                voice=voice,
                speed=speed,
                split_pattern=split_pattern,
            ):
                audio = segment.audio
                if hasattr(audio, "numpy"):
                    audio = audio.numpy()
                audio = np.asarray(audio, dtype="float32")
                if audio.size == 0:
                    continue
                audio_parts.append(audio)

                tokens = tuple(
                    TokenTiming(
                        text=str(tok.text),
                        whitespace=str(tok.whitespace or ""),
                        start=float(tok.start_ts or 0.0),
                        end=float(tok.end_ts or 0.0),
                    )
                    for tok in (getattr(segment, "tokens", None) or [])
                )
                segments.append(
                    AudioSegment(
                        graphemes=str(getattr(segment, "graphemes", "") or ""),
                        audio=audio.tobytes(),
                        sample_rate=sample_rate,
                        tokens=tokens,
                    )
                )

            if not audio_parts:
                return SynthesizedAudio(
                    data=b"",
                    format=AudioFormat(mime="audio/wav", extension="wav"),
                    duration=Duration(seconds=0.0),
                )

            combined = np.concatenate(audio_parts).astype("float32", copy=False)
            audio_bytes = combined.tobytes()
            duration_seconds = len(combined) / sample_rate

            return SynthesizedAudio(
                data=audio_bytes,
                format=AudioFormat(mime="audio/wav", extension="wav"),
                duration=Duration(seconds=duration_seconds),
                segments=tuple(segments),
            )
        except EngineError:
            raise
        except Exception as e:
            raise EngineError(f"Synthesis failed: {e}") from e

    def dispose(self) -> None:
        """Release session resources. Idempotent."""
        self._disposed = True


class KokoroEngine:
    """Engine implementation for Kokoro.

    Factory for KokoroSession instances. Stateless and thread-safe.
    """

    def __init__(self, pipeline: Any) -> None:
        self._pipeline = pipeline
        self._disposed = False

    def createSession(self) -> KokoroSession:
        """Create a new KokoroSession."""
        if self._disposed:
            raise EngineError("Engine disposed")
        return KokoroSession(self._pipeline)

    def dispose(self) -> None:
        """Release engine resources. Idempotent."""
        self._disposed = True

    def listVoices(self, sourceId: str) -> list[VoiceManifest]:
        """List available Kokoro voices. Implements VoiceLister capability.

        Note: Static voices are declared in the plugin manifest.
        This method is a fallback for dynamic plugins.
        """
        if self._disposed:
            raise EngineError("Engine disposed")
        return []