first
This commit is contained in:
@@ -0,0 +1,212 @@
|
||||
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")
|
||||
Reference in New Issue
Block a user