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,
    Duration,
    SynthesisRequest,
    SynthesizedAudio,
)

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_voice_id(voice_id: str) -> 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.
    """
    prefix = str(voice_id or "").strip()[:1].lower()
    if prefix in _CODE_TO_LANGUAGE:
        return _CODE_TO_LANGUAGE[prefix]
    return Language.EN_US


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)

            audio_parts: list[np.ndarray] = []
            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_parts.append(np.asarray(audio, dtype="float32"))

            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) / _KOKORO_SAMPLE_RATE

            return SynthesizedAudio(
                data=audio_bytes,
                format=AudioFormat(mime="audio/wav", extension="wav"),
                duration=Duration(seconds=duration_seconds),
            )
        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 []