from __future__ import annotations import logging from threading import Event from time import perf_counter from typing import Any, Iterator import numpy as np from rich.console import Console 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 logger = logging.getLogger(__name__) console = Console() class PocketTTSHandler(BaseHandler[TTSIn, TTSOut]): """ Handler for Pocket TTS model from Kyutai Labs. Supports streaming audio generation with voice cloning. """ def setup( self, should_listen: Event, device: str = "cpu", voice: str = "alba", # Default voice from catalog sample_rate: int = 16000, # Match the pipeline's audio output (LocalAudioStreamer uses 16kHz) blocksize: int = 512, max_tokens: int = 50, gen_kwargs: dict[str, Any] | None = None, # For compatibility with pipeline, not used cancel_scope: CancelScope | None = None, speculative_turns: SpeculativeTurnTracker | None = None, ) -> None: """ Initialize Pocket TTS handler. Args: should_listen: Event to control when to start listening again device: Device to run model on ('cpu', 'cuda', 'mps') voice: Voice to use. Can be: - A preset name: 'alba', 'marius', 'javert', 'jean', 'fantine', 'cosette', 'eponine', 'azelma' - A local audio file path - A Hugging Face path like "hf://kyutai/tts-voices/..." sample_rate: Output sample rate (pocket-tts generates at 24kHz and will be resampled to this rate). Default 16kHz matches the pipeline's audio output. blocksize: Size of audio blocks to yield max_tokens: Maximum tokens to generate """ self.should_listen = should_listen self.cancel_scope = cancel_scope self.speculative_turns = speculative_turns self.device = device self.voice = voice self.sample_rate = sample_rate self.blocksize = blocksize self.max_tokens = max_tokens # Suppress verbose logging from pocket_tts library logging.getLogger("pocket_tts").setLevel(logging.WARNING) logging.getLogger("pocket_tts.models.tts_model").setLevel(logging.WARNING) logging.getLogger("pocket_tts.utils.utils").setLevel(logging.WARNING) # Import and load model from pocket_tts import TTSModel logger.info("Loading Pocket TTS model") self.model = TTSModel.load_model() # Move model to specified device if device == "cuda": self.model = self.model.cuda() elif device == "mps": self.model = self.model.to("mps") elif device != "cpu": self.model = self.model.to(device) logger.info(f"Pocket TTS model moved to {device}") # Load voice state logger.info(f"Loading voice: {voice}") self.voice_state = self.model.get_state_for_audio_prompt(voice) logger.info(f"Pocket TTS model sample rate: {self.model.sample_rate}") @property def min_time_to_debug(self) -> float: """ Override to suppress logging for individual audio chunks. Pocket TTS yields many small chunks (~10-20ms each), which would flood logs. Only log if a chunk takes unusually long (>100ms). """ return 0.1 # 100ms threshold 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 logger.debug(f"Received language code: {language_code}") console.print(f"[green]ASSISTANT: {text}") # Generate audio stream logger.debug(f"Generating audio for: {text[:50]}...") pipeline_start = perf_counter() first_chunk = True # Calculate target chunk size in original sample rate # We need enough samples so that after resampling we get blocksize samples needs_resampling = self.model.sample_rate != self.sample_rate if needs_resampling: from scipy.signal import resample_poly # Target chunk size in original sample rate (before resampling) resample_ratio = self.model.sample_rate / self.sample_rate target_chunk_size = int(self.blocksize * resample_ratio) # Calculate up/down factors for polyphase resampling # For 24kHz -> 16kHz: 16000/24000 = 2/3 from math import gcd g = gcd(self.sample_rate, self.model.sample_rate) self._resample_up = self.sample_rate // g self._resample_down = self.model.sample_rate // g else: target_chunk_size = self.blocksize # Buffer to accumulate audio until we have enough for one block audio_buffer = [] buffer_size = 0 for audio_chunk in self.model.generate_audio_stream( self.voice_state, text, max_tokens=self.max_tokens, copy_state=True, # Don't modify the original voice state ): 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 if first_chunk: logger.debug(f"Time to first audio: {perf_counter() - pipeline_start:.3f}s") first_chunk = False # Convert from torch tensor to numpy and add to buffer audio_np = audio_chunk.cpu().numpy() audio_buffer.append(audio_np) buffer_size += len(audio_np) # Process buffer when we have enough samples for at least one block while buffer_size >= target_chunk_size: # Concatenate buffer concatenated = np.concatenate(audio_buffer) # Extract enough samples for one block chunk_to_process = concatenated[:target_chunk_size] remainder = concatenated[target_chunk_size:] # Resample this chunk if needed if needs_resampling: chunk_resampled = resample_poly( chunk_to_process, up=self._resample_up, down=self._resample_down, ) else: chunk_resampled = chunk_to_process # Convert to int16 format expected by audio output audio_int16 = (chunk_resampled * 32768).astype(np.int16) # Ensure exact blocksize (pad or trim if resampling caused slight size differences) if len(audio_int16) < self.blocksize: audio_int16 = np.pad(audio_int16, (0, self.blocksize - len(audio_int16))) elif len(audio_int16) > self.blocksize: audio_int16 = audio_int16[: self.blocksize] yield audio_int16 # Update buffer with remainder audio_buffer = [remainder] if len(remainder) > 0 else [] buffer_size = len(remainder) # Process any remaining audio in buffer if audio_buffer and buffer_size > 0: concatenated = np.concatenate(audio_buffer) # Resample the remainder if needs_resampling: concatenated = resample_poly( concatenated, up=self._resample_up, down=self._resample_down, ) # Convert to int16 audio_int16 = (concatenated * 32768).astype(np.int16) # Yield in blocks for i in range(0, len(audio_int16), self.blocksize): chunk = audio_int16[i : i + self.blocksize] # Pad last chunk if needed if len(chunk) < self.blocksize: chunk = np.pad(chunk, (0, self.blocksize - len(chunk))) yield chunk