first

2026-08-26 17:42:00 -04:00
commit ad95d530ad
134 changed files with 31750 additions and 0 deletions
@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Iterator, Optional
+
+import numpy as np
+import torch
+from lightning_whisper_mlx import LightningWhisperMLX
+from rich.console import Console
+
+from speech_to_speech.pipeline.handler_types import STTIn, STTOut
+from speech_to_speech.pipeline.messages import Transcription
+from speech_to_speech.STT.base_stt_handler import BaseSTTHandler
+
+logger = logging.getLogger(__name__)
+
+console = Console()
+
+SUPPORTED_LANGUAGES = [
+    "en",
+    "fr",
+    "es",
+    "zh",
+    "ja",
+    "ko",
+    "hi",
+    "de",
+    "pt",
+    "pl",
+    "it",
+    "nl",
+]
+
+
+class LightningWhisperSTTHandler(BaseSTTHandler):
+    """
+    Handles the Speech To Text generation using a Whisper model.
+    """
+
+    def setup(
+        self,
+        model_name: str = "distil-large-v3",
+        device: str = "mps",
+        torch_dtype: str = "float16",
+        compile_mode: Optional[str] = None,
+        language: Optional[str] = None,
+        gen_kwargs: dict[str, Any] = {},
+    ) -> None:
+        if len(model_name.split("/")) > 1:
+            model_name = model_name.split("/")[-1]
+        self.device = device
+        self.model = LightningWhisperMLX(model=model_name, batch_size=6, quant=None)
+        self.start_language = language
+        self.last_language = language
+
+        self.warmup()
+
+    def warmup(self) -> None:
+        logger.info(f"Warming up {self.__class__.__name__}")
+
+        # 2 warmup steps for no compile or compile mode with CUDA graphs capture
+        n_steps = 1
+        dummy_input = np.array([0] * 512)
+
+        for _ in range(n_steps):
+            _ = self.model.transcribe(dummy_input)["text"].strip()
+
+    def process(self, vad_audio: STTIn) -> Iterator[STTOut]:
+        logger.debug("infering whisper...")
+
+        audio = vad_audio.audio
+        if self.start_language != "auto":
+            transcription_dict = self.model.transcribe(audio, language=self.start_language)
+        else:
+            transcription_dict = self.model.transcribe(audio)
+            language_code = transcription_dict["language"]
+            if language_code not in SUPPORTED_LANGUAGES:
+                logger.warning(f"Whisper detected unsupported language: {language_code}")
+                if self.last_language in SUPPORTED_LANGUAGES:  # reprocess with the last language
+                    transcription_dict = self.model.transcribe(audio, language=self.last_language)
+                else:
+                    transcription_dict = {"text": "", "language": "en"}
+            else:
+                self.last_language = language_code
+
+        pred_text = transcription_dict["text"].strip()
+        language_code = transcription_dict["language"]
+        torch.mps.empty_cache()
+
+        logger.debug("finished whisper inference")
+        console.print(f"[yellow]USER: {pred_text}")
+        logger.debug(f"Language Code Whisper: {language_code}")
+
+        if self.start_language == "auto":
+            language_code += "-auto"
+
+        yield Transcription(
+            text=pred_text,
+            language_code=language_code,
+            turn_id=vad_audio.turn_id,
+            turn_revision=vad_audio.turn_revision,
+            speech_stopped_at_s=vad_audio.created_at_s,
+        )