213 lines
8.0 KiB
Python
213 lines
8.0 KiB
Python
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")
|