converse-framework 0.2.0__py3-none-any.whl

converse_framework/providers/faster_whisper.py

@@ -0,0 +1,290 @@
+"""faster-whisper ASR provider.
+
+The ``faster_whisper`` package is imported lazily inside
+:meth:`_ensure_model` and :meth:`check_status` so the base
+:mod:`converse_framework` package stays light. Install with::
+
+    pip install 'converse-framework[faster-whisper]'
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import time
+from collections.abc import AsyncIterator
+
+from converse_framework.audio_utils import pcm_s16le_to_float32
+from converse_framework.protocols import (
+    ASRProvider,
+    ProgressCallback,
+    ProviderCapabilities,
+    ProviderStatus,
+    TranscriptEvent,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FasterWhisperASRProvider(ASRProvider):
+    def __init__(self, config: dict):
+        self.model_name = str(config.get("model", "large-v3-turbo"))
+        self.device = str(config.get("device", "auto"))
+        self.compute_type = str(config.get("compute_type", "auto"))
+        self.language = config.get("language", "en")
+        self.beam_size = int(config.get("beam_size", 1))
+        self.vad_filter = bool(config.get("vad_filter", False))
+        self.initial_prompt = config.get("initial_prompt")
+        self.condition_on_previous_text = bool(
+            config.get("condition_on_previous_text", False)
+        )
+        self.temperature = config.get("temperature", 0)
+        self.compression_ratio_threshold = config.get(
+            "compression_ratio_threshold", 2.4
+        )
+        self.log_prob_threshold = config.get("log_prob_threshold", -0.5)
+        self.no_speech_threshold = config.get("no_speech_threshold", 0.2)
+        self.suppress_tokens = config.get("suppress_tokens")
+        self.timeout_s = float(config.get("timeout_s", 120))
+        self._model = config.get("_model")
+        self._load_error: str | None = None
+        self._auto_cuda_dll_dirs = bool(config.get("auto_cuda_dll_dirs", True))
+        self._cuda_dll_handles: list[object] = []
+
+    @property
+    def status(self) -> ProviderStatus:
+        if self._load_error:
+            return ProviderStatus(
+                name="faster-whisper",
+                kind="asr",
+                ready=False,
+                message=f"faster-whisper failed to load: {self._load_error}",
+                capabilities=ProviderCapabilities(),
+                status_level="error",
+            )
+        if self._model is not None:
+            message = f"Loaded {self.model_name} on {self.device}/{self.compute_type}."
+            status_level = "ready"
+        else:
+            message = (
+                f"Configured for {self.model_name} on {self.device}/{self.compute_type}. "
+                "Model loads on first voice transcription and may download if not cached."
+            )
+            status_level = "configured"
+        return ProviderStatus(
+            name="faster-whisper",
+            kind="asr",
+            ready=True,
+            message=message,
+            capabilities=ProviderCapabilities(
+                languages=(str(self.language),) if self.language else ("auto",)
+            ),
+            provider_id="faster-whisper",
+            loaded=self._model is not None,
+            active_model=self.model_name,
+            models=({"id": self.model_name, "label": self.model_name},),
+            status_level=status_level,
+        )
+
+    async def check_status(self) -> ProviderStatus:
+        return await self.probe_status()
+
+    async def probe_status(self) -> ProviderStatus:
+        """Cheap probe: check import availability, no model load."""
+        if self._model is None:
+            try:
+                import faster_whisper  # type: ignore[import-not-found]  # noqa: F401
+            except Exception as exc:  # pragma: no cover - import path
+                self._load_error = str(exc)
+        return self.status
+
+    async def load_status(self) -> ProviderStatus:
+        """May load heavy resources."""
+        return await self.load()
+
+    async def load(self) -> ProviderStatus:
+        if self._model is not None:
+            return self.status
+        try:
+            await asyncio.wait_for(
+                asyncio.to_thread(self._ensure_model), timeout=self.timeout_s
+            )
+        except asyncio.TimeoutError:
+            self._load_error = f"Model load timed out after {self.timeout_s}s"
+            raise
+        return self.status
+
+    async def transcribe_text_input(self, text: str) -> AsyncIterator[TranscriptEvent]:
+        stripped = text.strip()
+        if stripped:
+            yield TranscriptEvent(text=stripped, final=True)
+
+    async def transcribe_audio(
+        self,
+        pcm_s16le: bytes,
+        sample_rate: int,
+        progress: ProgressCallback | None = None,
+    ) -> AsyncIterator[TranscriptEvent]:
+        if sample_rate != 16000:
+            raise ValueError(
+                f"faster-whisper expects 16000 Hz audio, got {sample_rate}"
+            )
+        audio = pcm_s16le_to_float32(pcm_s16le)
+        if audio.size == 0:
+            return
+        if progress:
+            await progress(
+                "asr.progress",
+                {
+                    "stage": "queued",
+                    "message": (
+                        f"Queued {round(audio.size / sample_rate, 2)}s utterance "
+                        "for faster-whisper."
+                    ),
+                },
+            )
+        loop = asyncio.get_running_loop()
+        try:
+            segments_text = await asyncio.wait_for(
+                asyncio.to_thread(self._transcribe_blocking, audio, progress, loop),
+                timeout=self.timeout_s,
+            )
+        except asyncio.TimeoutError:
+            if self._model is None:
+                self._load_error = (
+                    f"Model load or transcription timed out after {self.timeout_s}s"
+                )
+            raise
+        text = " ".join(part for part in segments_text if part).strip()
+        if progress:
+            await progress(
+                "asr.progress",
+                {"stage": "complete", "message": "ASR transcription complete."},
+            )
+        if text:
+            yield TranscriptEvent(text=text, final=True)
+
+    def _transcribe_blocking(
+        self,
+        audio,
+        progress: ProgressCallback | None,
+        loop: asyncio.AbstractEventLoop,
+    ) -> list[str]:
+        started = time.perf_counter()
+        logger.info(
+            "[ASR] transcribe_blocking called, audio length=%d samples (%.2fs)",
+            audio.size,
+            audio.size / 16000,
+        )
+        self._emit_progress_threadsafe(
+            loop,
+            progress,
+            "loading",
+            f"Loading faster-whisper model {self.model_name}.",
+        )
+        with contextlib.suppress(Exception):
+            self._ensure_model()
+        if self._model is None:
+            raise RuntimeError(
+                f"faster-whisper model did not load: {self._load_error or 'unknown error'}"
+            )
+        self._emit_progress_threadsafe(
+            loop,
+            progress,
+            "loaded",
+            f"Model ready after {round(time.perf_counter() - started, 1)}s. "
+            "Running inference.",
+        )
+        logger.info(
+            "[ASR] model loaded in %.1fs, starting inference on %s/%s",
+            time.perf_counter() - started,
+            self.device,
+            self.compute_type,
+        )
+        transcribe_options = {
+            "language": self.language,
+            "beam_size": self.beam_size,
+            "vad_filter": self.vad_filter,
+            "initial_prompt": self.initial_prompt,
+            "condition_on_previous_text": self.condition_on_previous_text,
+            "temperature": self.temperature,
+            "compression_ratio_threshold": self.compression_ratio_threshold,
+            "log_prob_threshold": self.log_prob_threshold,
+            "no_speech_threshold": self.no_speech_threshold,
+        }
+        if self.suppress_tokens is not None:
+            transcribe_options["suppress_tokens"] = self.suppress_tokens
+        segments, _info = self._model.transcribe(audio, **transcribe_options)
+        logger.info("[ASR] inference call returned, iterating segments...")
+        texts: list[str] = []
+        for segment in segments:
+            text = segment.text.strip()
+            if text:
+                texts.append(text)
+                start = getattr(segment, "start", None)
+                end = getattr(segment, "end", None)
+                prefix = ""
+                if start is not None and end is not None:
+                    prefix = (
+                        f"Segment {round(float(start), 2)}-{round(float(end), 2)}s: "
+                    )
+                self._emit_progress_threadsafe(
+                    loop, progress, "segment", f"{prefix}{text}"
+                )
+        logger.info(
+            "[ASR] all segments collected in %.1fs, %d segments with text",
+            time.perf_counter() - started,
+            len(texts),
+        )
+        return texts
+
+    def _ensure_model(self) -> None:
+        if self._model is not None:
+            return
+        # Windows CUDA DLL directory discovery (conservative, best-effort).
+        if self._auto_cuda_dll_dirs:
+            self._add_cuda_dll_dirs()
+        try:
+            from faster_whisper import WhisperModel  # type: ignore[import-not-found]
+
+            self._model = WhisperModel(
+                self.model_name, device=self.device, compute_type=self.compute_type
+            )
+        except Exception as exc:  # pragma: no cover - import path
+            self._load_error = str(exc)
+            raise
+
+    def _add_cuda_dll_dirs(self) -> None:
+        """Register NVIDIA wheel DLL directories if on Windows."""
+        # Local import to keep cuda_utils optional.
+        try:
+            from converse_framework.cuda_utils import add_nvidia_dll_directories
+
+            self._cuda_dll_handles = add_nvidia_dll_directories()
+        except Exception:  # pragma: no cover
+            logger.debug("CUDA DLL discovery skipped or failed.")
+
+    async def unload(self) -> ProviderStatus:
+        if self._model is not None:
+            logger.info(
+                "[ASR] unloading faster-whisper model (%s/%s)",
+                self.device,
+                self.compute_type,
+            )
+            self._model = None
+        self._load_error = None
+        self._cuda_dll_handles.clear()
+        return self.status
+
+    def _emit_progress_threadsafe(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        progress: ProgressCallback | None,
+        stage: str,
+        message: str,
+    ) -> None:
+        if not progress:
+            return
+        coro = progress("asr.progress", {"stage": stage, "message": message})
+        asyncio.run_coroutine_threadsafe(coro, loop)  # type: ignore[arg-type]

converse_framework/providers/kokoro_onnx.py

@@ -0,0 +1,391 @@
+"""Kokoro ONNX TTS provider.
+
+The ``kokoro_onnx``, ``misaki``, and ``httpx`` packages are imported lazily
+inside :meth:`_ensure_model` and :meth:`_download_asset` so the base
+:mod:`converse_framework` package stays light. Install with::
+
+    pip install 'converse-framework[kokoro]'
+
+The default cache directory is platform-aware and does not depend on the
+harness ``PROJECT_ROOT``. The cache location can be overridden via the
+``cache_dir`` config key, or by setting the ``CONVERSE_FRAMEWORK_CACHE_DIR``
+environment variable (the provider appends ``/kokoro`` to that path).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import threading
+import time
+from collections.abc import AsyncIterator
+from dataclasses import replace
+from pathlib import Path
+
+from converse_framework.audio_utils import float_audio_to_pcm_s16le_bytes
+from converse_framework.protocols import (
+    AudioChunk,
+    ProgressCallback,
+    ProviderCapabilities,
+    ProviderStatus,
+    TTSProvider,
+)
+
+
+DEFAULT_KOKORO_MODEL_URL = (
+    "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/"
+    "kokoro-v1.0.int8.onnx"
+)
+DEFAULT_KOKORO_VOICES_URL = (
+    "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/"
+    "voices-v1.0.bin"
+)
+
+
+def _default_cache_dir() -> Path:
+    """Return a platform-appropriate cache directory for Kokoro assets.
+
+    Resolution order:
+      1. ``CONVERSE_FRAMEWORK_CACHE_DIR`` environment variable, with ``kokoro``
+         appended.
+      2. ``~/.cache/converse-framework/kokoro`` on POSIX-likes.
+      3. ``%LOCALAPPDATA%/converse-framework/kokoro`` (or
+         ``~/.cache/converse-framework/kokoro``) on Windows.
+    """
+    env_root = os.environ.get("CONVERSE_FRAMEWORK_CACHE_DIR")
+    if env_root:
+        return Path(env_root) / "kokoro"
+    if os.name == "nt":
+        local_app_data = os.environ.get("LOCALAPPDATA")
+        if local_app_data:
+            return Path(local_app_data) / "converse-framework" / "kokoro"
+    return Path.home() / ".cache" / "converse-framework" / "kokoro"
+
+
+class KokoroOnnxProvider(TTSProvider):
+    def __init__(self, config: dict):
+        self.voice = str(config.get("voice", "af_heart"))
+        self.lang = str(config.get("lang", "en-us"))
+        self.speed = float(config.get("speed", 1.0))
+        self.trim = bool(config.get("trim", True))
+        self.timeout_s = float(config.get("timeout_s", 300))
+        configured_cache = config.get("cache_dir")
+        self.cache_dir = (
+            Path(str(configured_cache)) if configured_cache else _default_cache_dir()
+        )
+        self.model_filename = str(config.get("model_filename", "kokoro-v1.0.int8.onnx"))
+        self.voices_filename = str(config.get("voices_filename", "voices-v1.0.bin"))
+        self.model_url = str(config.get("model_url", DEFAULT_KOKORO_MODEL_URL))
+        self.voices_url = str(config.get("voices_url", DEFAULT_KOKORO_VOICES_URL))
+        self.onnx_intra_op_num_threads = int(config.get("onnx_intra_op_num_threads", 4))
+        self.onnx_inter_op_num_threads = int(config.get("onnx_inter_op_num_threads", 1))
+        self.preload_g2p = bool(config.get("preload_g2p", True))
+        self._model = config.get("_model")
+        self._g2p = config.get("_g2p")
+        self._load_error: str | None = None
+        self._g2p_error: str | None = None
+        self._lock = threading.Lock()
+        self._generation_lock = asyncio.Lock()
+
+    @property
+    def status(self) -> ProviderStatus:
+        loaded = self._model is not None
+        if self._load_error:
+            return ProviderStatus(
+                name="kokoro-onnx",
+                kind="tts",
+                ready=False,
+                message=f"Kokoro ONNX failed to load: {self._load_error}",
+                capabilities=ProviderCapabilities(
+                    supports_streaming_tts=True, languages=("en",)
+                ),
+                provider_id="kokoro-onnx",
+                loaded=False,
+                supports_model_management=True,
+                supports_voice_selection=True,
+                active_voice=self.voice,
+                status_level="error",
+            )
+        if self._g2p_error:
+            return ProviderStatus(
+                name="kokoro-onnx",
+                kind="tts",
+                ready=False,
+                message=f"Kokoro English G2P failed: {self._g2p_error}",
+                capabilities=ProviderCapabilities(
+                    supports_streaming_tts=True, languages=("en",)
+                ),
+                provider_id="kokoro-onnx",
+                loaded=loaded,
+                supports_model_management=True,
+                supports_voice_selection=True,
+                active_voice=self.voice,
+                status_level="error",
+            )
+
+        if loaded:
+            message = f"Loaded Kokoro v1.0 ONNX voice '{self.voice}' ({self.lang})."
+            status_level = "ready"
+        else:
+            message = (
+                f"Configured for Kokoro v1.0 ONNX voice '{self.voice}' ({self.lang}). "
+                "Model loads on first TTS request."
+            )
+            status_level = "configured"
+        return ProviderStatus(
+            name="kokoro-onnx",
+            kind="tts",
+            ready=True,
+            message=message,
+            capabilities=ProviderCapabilities(
+                supports_streaming_tts=True, languages=("en",)
+            ),
+            provider_id="kokoro-onnx",
+            loaded=loaded,
+            supports_model_management=True,
+            supports_voice_selection=True,
+            active_voice=self.voice,
+            status_level=status_level,
+        )
+
+    async def check_status(self) -> ProviderStatus:
+        return await self.probe_status()
+
+    async def probe_status(self) -> ProviderStatus:
+        """Cheap probe: check import availability, no model load."""
+        try:
+            import kokoro_onnx  # type: ignore[import-not-found]  # noqa: F401
+
+            if self._should_use_misaki():
+                from misaki import en as _en  # type: ignore[import-not-found]  # noqa: F401
+                from misaki import espeak as _espeak  # type: ignore[import-not-found]  # noqa: F401
+        except Exception as exc:  # pragma: no cover - import path
+            if self._should_use_misaki():
+                self._g2p_error = str(exc)
+            else:
+                self._load_error = str(exc)
+        return self.status
+
+    async def load_status(self) -> ProviderStatus:
+        """May load heavy resources."""
+        return await self.load()
+
+    async def load(self) -> ProviderStatus:
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, self._ensure_model)
+        if self.preload_g2p and self._should_use_misaki():
+            await loop.run_in_executor(None, self._ensure_g2p)
+        return self.status
+
+    async def unload(self) -> ProviderStatus:
+        def release() -> None:
+            with self._lock:
+                self._model = None
+                self._load_error = None
+                self._g2p = None
+                self._g2p_error = None
+
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, release)
+        return self.status
+
+    async def stream_audio(self, text: str) -> AsyncIterator[AudioChunk]:
+        async for chunk in self.stream_audio_with_progress(text):
+            yield chunk
+
+    async def stream_audio_with_progress(
+        self,
+        text: str,
+        progress: ProgressCallback | None = None,
+    ) -> AsyncIterator[AudioChunk]:
+        loop = asyncio.get_running_loop()
+        started = time.perf_counter()
+        self._emit_progress(
+            loop,
+            progress,
+            "loading",
+            f"Loading Kokoro voice '{self.voice}'.",
+            started=started,
+        )
+        await loop.run_in_executor(None, self._ensure_model)
+        self._emit_progress(loop, progress, "loaded", "Kokoro ready.", started=started)
+
+        async with self._generation_lock:
+            stream_text = text
+            is_phonemes = False
+            if self._should_use_misaki():
+                self._emit_progress(
+                    loop,
+                    progress,
+                    "phonemizing",
+                    "Preparing English phonemes with Misaki.",
+                    started=started,
+                )
+                stream_text = await loop.run_in_executor(
+                    None, self._phonemize_english, text
+                )
+                is_phonemes = True
+
+            self._emit_progress(
+                loop,
+                progress,
+                "generating",
+                "Generating speech.",
+                started=started,
+            )
+            index = 0
+            previous_chunk: AudioChunk | None = None
+            assert self._model is not None
+            async for audio, sample_rate in self._model.create_stream(
+                stream_text,
+                voice=self.voice,
+                speed=self.speed,
+                lang=self.lang,
+                is_phonemes=is_phonemes,
+                trim=self.trim,
+            ):
+                pcm_bytes = float_audio_to_pcm_s16le_bytes(audio)
+                if not pcm_bytes:
+                    continue
+                index += 1
+                current_chunk = AudioChunk(
+                    pcm_bytes,
+                    sample_rate=sample_rate,
+                    channels=1,
+                    encoding="pcm_s16le",
+                    duration_ms=(
+                        int((len(pcm_bytes) // 2) * 1000 / sample_rate)
+                        if sample_rate
+                        else None
+                    ),
+                    final=False,
+                )
+                self._emit_progress(
+                    loop,
+                    progress,
+                    "chunk",
+                    f"Generated audio chunk {index}.",
+                    started=started,
+                    chunk_index=index,
+                    first_chunk=index == 1,
+                    duration_ms=current_chunk.duration_ms,
+                )
+                if previous_chunk is not None:
+                    yield previous_chunk
+                previous_chunk = current_chunk
+
+            if previous_chunk is not None:
+                yield replace(previous_chunk, final=True)
+            self._emit_progress(
+                loop,
+                progress,
+                "complete",
+                "TTS complete.",
+                started=started,
+                chunks=index,
+            )
+
+    def _ensure_model(self) -> None:
+        if self._model is not None:
+            return
+        model_path = self._download_asset(self.model_url, self.model_filename)
+        voices_path = self._download_asset(self.voices_url, self.voices_filename)
+        from kokoro_onnx import Kokoro  # type: ignore[import-not-found]
+
+        self._model = Kokoro(str(model_path), str(voices_path))
+        self._apply_onnx_session_options(str(model_path))
+        self._load_error = None
+        self._g2p_error = None
+
+    def _apply_onnx_session_options(self, model_path: str) -> None:
+        if self.onnx_intra_op_num_threads <= 0 and self.onnx_inter_op_num_threads <= 0:
+            return
+        if self._model is None or not hasattr(self._model, "sess"):
+            return
+        import onnxruntime as ort  # type: ignore[import-not-found]
+
+        options = ort.SessionOptions()
+        if self.onnx_intra_op_num_threads > 0:
+            options.intra_op_num_threads = self.onnx_intra_op_num_threads
+        if self.onnx_inter_op_num_threads > 0:
+            options.inter_op_num_threads = self.onnx_inter_op_num_threads
+        providers = ["CPUExecutionProvider"]
+        self._model.sess = ort.InferenceSession(
+            model_path, sess_options=options, providers=providers
+        )
+
+    def _ensure_g2p(self):
+        if self._g2p is not None:
+            return self._g2p
+        british = self._use_british_english()
+        from misaki import en, espeak  # type: ignore[import-not-found]
+
+        self._g2p = en.G2P(
+            trf=False,
+            british=british,
+            fallback=espeak.EspeakFallback(british=british),
+        )
+        self._g2p_error = None
+        return self._g2p
+
+    def _phonemize_english(self, text: str) -> str:
+        try:
+            g2p = self._ensure_g2p()
+            phonemes, _tokens = g2p(text)
+            return str(phonemes).strip()
+        except Exception as exc:  # pragma: no cover - exercised via tests
+            self._g2p_error = str(exc)
+            raise RuntimeError(f"Misaki English phonemization failed: {exc}") from exc
+
+    def _should_use_misaki(self) -> bool:
+        return self.lang.lower().startswith("en")
+
+    def _use_british_english(self) -> bool:
+        lang = self.lang.lower()
+        return lang.startswith("en-gb") or self.voice.lower().startswith("b")
+
+    def _download_asset(self, url: str, filename: str) -> Path:
+        try:
+            import httpx  # type: ignore[import-not-found]
+        except Exception as exc:  # pragma: no cover - import path
+            raise RuntimeError(
+                "kokoro-onnx provider requires httpx; install with "
+                "pip install 'converse-framework[kokoro]'."
+            ) from exc
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        target = self.cache_dir / filename
+        if target.exists():
+            return target
+        temp = target.with_suffix(target.suffix + ".part")
+        with (
+            httpx.Client(follow_redirects=True, timeout=self.timeout_s) as client,
+            temp.open("wb") as handle,
+        ):
+            with client.stream("GET", url) as response:
+                response.raise_for_status()
+                for chunk in response.iter_bytes():
+                    if chunk:
+                        handle.write(chunk)
+        temp.replace(target)
+        return target
+
+    def _emit_progress(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        progress: ProgressCallback | None,
+        stage: str,
+        message: str,
+        *,
+        started: float | None = None,
+        **extra,
+    ) -> None:
+        if not progress:
+            return
+        payload = {"stage": stage, "message": message, **extra}
+        if started is not None:
+            payload["elapsed_ms"] = int((time.perf_counter() - started) * 1000)
+
+        async def _fire() -> None:
+            await progress("tts.progress", payload)
+
+        loop.call_soon_threadsafe(asyncio.create_task, _fire())