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