from __future__ import annotations import logging from threading import Event from typing import Any, Iterator import librosa import numpy as np import torch from rich.console import Console from transformers import AutoTokenizer, VitsModel from speech_to_speech.baseHandler import BaseHandler from speech_to_speech.pipeline.cancel_scope import CancelScope from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.DEBUG) logger = logging.getLogger(__name__) console = Console() WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE = { "en": "eng", # English "fr": "fra", # French "es": "spa", # Spanish "ko": "kor", # Korean "hi": "hin", # Hindi "ar": "ara", # Arabic "hy": "hyw", # Armenian "az": "azb", # Azerbaijani "bu": "bul", # Bulgarian "ca": "cat", # Catalan "nl": "nld", # Dutch "fi": "fin", # Finnish "de": "deu", # German "el": "ell", # Greek "he": "heb", # Hebrew "hu": "hun", # Hungarian "is": "isl", # Icelandic "id": "ind", # Indonesian "ka": "kan", # Kannada "kk": "kaz", # Kazakh "lv": "lav", # Latvian "zl": "zlm", # Malay "ma": "mar", # Marathi "fa": "fas", # Persian "po": "pol", # Polish "pt": "por", # Portuguese "ro": "ron", # Romanian "ru": "rus", # Russian "sw": "swh", # Swahili "sv": "swe", # Swedish "tg": "tgl", # Tagalog "ta": "tam", # Tamil "th": "tha", # Thai "tu": "tur", # Turkish "uk": "ukr", # Ukrainian "ur": "urd", # Urdu "vi": "vie", # Vietnamese "cy": "cym", # Welsh } class FacebookMMSTTSHandler(BaseHandler[TTSIn, TTSOut]): def setup( self, should_listen: Event, device: str = "cuda", torch_dtype: str = "float32", language: str = "en", stream: bool = True, chunk_size: int = 512, cancel_scope: CancelScope | None = None, speculative_turns: SpeculativeTurnTracker | None = None, **kwargs: Any, ) -> None: self.should_listen = should_listen self.cancel_scope = cancel_scope self.speculative_turns = speculative_turns self.device = device self.torch_dtype = getattr(torch, torch_dtype) self.stream = stream self.chunk_size = chunk_size self.language = language self._initial_language = self.language self.load_model(self.language) self.warmup() def load_model(self, language_code: str) -> None: try: model_name = f"facebook/mms-tts-{WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE[language_code]}" logger.info(f"Loading model: {model_name}") self.model = VitsModel.from_pretrained(model_name).to(self.device) # type: ignore[arg-type] self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.language = language_code except KeyError: logger.warning(f"Unsupported language: {language_code}. Falling back to English.") self.load_model("en") def warmup(self) -> None: logger.info(f"Warming up {self.__class__.__name__}") self.generate_audio("Hello, this is a test") def generate_audio(self, text: str) -> torch.Tensor | None: if not text: logger.warning("Received empty text input") return None try: logger.debug(f"Tokenizing text: {text}") logger.debug(f"Current language: {self.language}") logger.debug(f"Tokenizer: {self.tokenizer}") inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True) input_ids = inputs.input_ids.to(self.device).long() attention_mask = inputs.attention_mask.to(self.device) logger.debug(f"Input IDs shape: {input_ids.shape}, dtype: {input_ids.dtype}") logger.debug(f"Input IDs: {input_ids}") if input_ids.numel() == 0: logger.error("Input IDs tensor is empty") return None with torch.no_grad(): output = self.model(input_ids=input_ids, attention_mask=attention_mask) logger.debug(f"Output waveform shape: {output.waveform.shape}") return output.waveform except Exception as e: logger.error(f"Error in generate_audio: {str(e)}") logger.exception("Full traceback:") return None def process(self, tts_input: TTSIn) -> Iterator[TTSOut]: speculative_turns = getattr(self, "speculative_turns", None) if isinstance(tts_input, EndOfResponse): if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( tts_input.turn_id, tts_input.turn_revision, ): return yield AUDIO_RESPONSE_DONE return if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( tts_input.turn_id, tts_input.turn_revision, ): logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision) return if speculative_turns: speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision) gen = self.cancel_scope.generation if self.cancel_scope else None language_code = tts_input.language_code text = tts_input.text console.print(f"[green]ASSISTANT: {text}") logger.debug(f"Processing text: {text}") logger.debug(f"Language code: {language_code}") if language_code is not None and self.language != language_code: try: logger.info(f"Switching language from {self.language} to {language_code}") self.load_model(language_code) except KeyError: console.print( f"[red]Language {language_code} not supported by Facebook MMS. Using {self.language} instead." ) logger.warning(f"Unsupported language: {language_code}") audio_output = self.generate_audio(text) if audio_output is None or audio_output.numel() == 0: logger.warning("No audio output generated") return audio_numpy = audio_output.cpu().numpy().squeeze() logger.debug(f"Raw audio shape: {audio_numpy.shape}, dtype: {audio_numpy.dtype}") audio_resampled = librosa.resample(audio_numpy, orig_sr=self.model.config.sampling_rate, target_sr=16000) logger.debug(f"Resampled audio shape: {audio_resampled.shape}, dtype: {audio_resampled.dtype}") audio_int16 = (audio_resampled * 32768).astype(np.int16) logger.debug(f"Final audio shape: {audio_int16.shape}, dtype: {audio_int16.dtype}") if self.stream: for i in range(0, len(audio_int16), self.chunk_size): if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): logger.info("TTS generation cancelled (interruption)") return chunk = audio_int16[i : i + self.chunk_size] yield np.pad(chunk, (0, self.chunk_size - len(chunk))) else: for i in range(0, len(audio_int16), self.chunk_size): if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): logger.info("TTS generation cancelled (interruption)") return yield np.pad( audio_int16[i : i + self.chunk_size], (0, self.chunk_size - len(audio_int16[i : i + self.chunk_size])), ) def on_session_end(self) -> None: if self.language != self._initial_language: self.load_model(self._initial_language) logger.debug("Facebook MMS TTS session state reset")