pocket-tts 1.0.2__py3-none-any.whl

pocket_tts/__init__.py

@@ -0,0 +1,16 @@
+from beartype import BeartypeConf
+from beartype.claw import beartype_this_package
+
+beartype_this_package(conf=BeartypeConf(is_color=False))
+
+from pocket_tts.models.tts_model import TTSModel  # noqa: E402
+
+# Public methods:
+# TTSModel.device
+# TTSModel.sample_rate
+# TTSModel.load_model
+# TTSModel.generate_audio
+# TTSModel.generate_audio_stream
+# TTSModel.get_state_for_audio_prompt
+
+__all__ = ["TTSModel"]

pocket_tts/__main__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from pocket_tts.main import cli_app
+
+if __name__ == "__main__":
+    cli_app()

pocket_tts/conditioners/base.py

@@ -0,0 +1,38 @@
+import logging
+from typing import Generic, NamedTuple, TypeVar
+
+import torch
+from torch import nn
+
+logger = logging.getLogger(__name__)
+
+
+Prepared = TypeVar("Prepared")  # represents the prepared condition input type.
+
+
+class TokenizedText(NamedTuple):
+    tokens: torch.Tensor  # should be long tensor.
+
+
+class BaseConditioner(nn.Module, Generic[Prepared]):
+    """Base model for all conditioner modules.
+
+    Args:
+        dim (int): internal dim of the model.
+        output_dim (int): Output dim of the conditioner.
+        force_linear (bool, optional): Force linear projection even when `dim == output_dim`.
+        output_bias (bool): if True, the output projection will have a bias.
+        learn_padding (bool): if True, the padding value will be learnt, zero otherwise.
+    """
+
+    def __init__(
+        self, dim: int, output_dim: int, output_bias: bool = False, force_linear: bool = True
+    ):
+        super().__init__()
+        self.dim = dim
+        self.output_dim = output_dim
+        assert force_linear or dim != output_dim
+        assert not output_bias
+
+    def forward(self, inputs: TokenizedText) -> torch.Tensor:
+        return self._get_condition(inputs)

pocket_tts/conditioners/text.py

@@ -0,0 +1,61 @@
+import logging
+
+import sentencepiece
+import torch
+from torch import nn
+
+from pocket_tts.conditioners.base import BaseConditioner, TokenizedText
+from pocket_tts.utils.utils import download_if_necessary
+
+logger = logging.getLogger(__name__)
+
+
+class SentencePieceTokenizer:
+    """This tokenizer should be used for natural language descriptions.
+    For example:
+    ["he didn't, know he's going home.", 'shorter sentence'] =>
+    [[78, 62, 31,  4, 78, 25, 19, 34],
+    [59, 77, PAD, PAD, PAD, PAD, PAD, PAD]]
+
+    Args:
+        n_bins (int): should be equal to the number of elements in the sentencepiece tokenizer.
+        tokenizer_path (str): path to the sentencepiece tokenizer model.
+
+    """
+
+    def __init__(self, nbins: int, tokenizer_path: str) -> None:
+        logger.info("Loading sentencepiece tokenizer from %s", tokenizer_path)
+        tokenizer_path = download_if_necessary(tokenizer_path)
+        self.sp = sentencepiece.SentencePieceProcessor(str(tokenizer_path))
+        assert nbins == self.sp.vocab_size(), (
+            f"sentencepiece tokenizer has vocab size={self.sp.vocab_size()} but nbins={nbins} was specified"
+        )
+
+    def __call__(self, text: str) -> TokenizedText:
+        return TokenizedText(torch.tensor(self.sp.encode(text, out_type=int))[None, :])
+
+
+class LUTConditioner(BaseConditioner):
+    """Lookup table TextConditioner.
+
+    Args:
+        n_bins (int): Number of bins.
+        dim (int): Hidden dim of the model (text-encoder/LUT).
+        output_dim (int): Output dim of the conditioner.
+        tokenizer (str): Name of the tokenizer.
+        possible_values (list[str] or None): list of possible values for the tokenizer.
+    """
+
+    def __init__(self, n_bins: int, tokenizer_path: str, dim: int, output_dim: int):
+        super().__init__(dim=dim, output_dim=output_dim)
+        self.tokenizer = SentencePieceTokenizer(n_bins, tokenizer_path)
+        self.embed = nn.Embedding(n_bins + 1, self.dim)  # n_bins + 1 for padding.
+
+    def prepare(self, x: str) -> TokenizedText:
+        tokens = self.tokenizer(x)
+        tokens = tokens[0].to(self.embed.weight.device)
+        return TokenizedText(tokens)
+
+    def _get_condition(self, inputs: TokenizedText) -> torch.Tensor:
+        embeds = self.embed(inputs[0])
+        return embeds

pocket_tts/config/b6369a24.yaml

@@ -0,0 +1,57 @@
+# sig: b6369a24
+
+weights_path: hf://kyutai/pocket-tts/tts_b6369a24.safetensors@427e3d61b276ed69fdd03de0d185fa8a8d97fc5b
+weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/tts_b6369a24.safetensors@d4fdd22ae8c8e1cb3634e150ebeff1dab2d16df3
+
+flow_lm:
+  dtype: float32
+  flow:
+    depth: 6
+    dim: 512
+  transformer:
+    d_model: 1024
+    hidden_scale: 4
+    max_period: 10000
+    num_heads: 16
+    num_layers: 6
+  lookup_table:
+    dim: 1024
+    n_bins: 4000
+    tokenizer: sentencepiece
+    tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/tokenizer.model@d4fdd22ae8c8e1cb3634e150ebeff1dab2d16df3
+  #weights_path: flow_lm_b6369a24.safetensors
+
+mimi:
+  dtype: float32
+  sample_rate: 24000
+  channels: 1
+  frame_rate: 12.5
+  seanet:
+    dimension: 512
+    channels: 1
+    n_filters: 64
+    n_residual_layers: 1
+    ratios:
+    - 6
+    - 5
+    - 4
+    kernel_size: 7
+    residual_kernel_size: 3
+    last_kernel_size: 3
+    dilation_base: 2
+    pad_mode: constant
+    compress: 2
+  transformer:
+    d_model: 512
+    num_heads: 8
+    num_layers: 2
+    layer_scale: 0.01
+    context: 250
+    dim_feedforward: 2048
+    input_dimension: 512
+    output_dimensions:
+    - 512
+  quantizer:
+    dimension: 32
+    output_dimension: 512
+  #weights_path: mimi_b6369a24.safetensors

pocket_tts/data/__init__.py

@@ -0,0 +1,2 @@
+"""Audio loading and writing support. Datasets for raw audio
+or also including some metadata."""

pocket_tts/data/audio.py

@@ -0,0 +1,144 @@
+"""
+Audio IO methods are defined in this module (info, read, write),
+We rely on av library for faster read when possible, otherwise on torchaudio.
+"""
+
+import logging
+import os
+import sys
+import wave
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from beartype.typing import Iterator
+
+logger = logging.getLogger(__name__)
+
+FIRST_CHUNK_LENGTH_SECONDS = float(os.environ.get("FIRST_CHUNK_LENGTH_SECONDS", "0"))
+
+
+def audio_read(filepath: str | Path) -> tuple[torch.Tensor, int]:
+    """Read audio file. WAV uses built-in wave module; other formats require soundfile."""
+    filepath = Path(filepath)
+
+    if filepath.suffix.lower() == ".wav":
+        # Use built-in wave module for WAV files
+        with wave.open(str(filepath), "rb") as wav_file:
+            sample_rate = wav_file.getframerate()
+            n_channels = wav_file.getnchannels()
+            raw_data = wav_file.readframes(-1)
+            samples = np.frombuffer(raw_data, dtype=np.int16).astype(np.float32) / 32768.0
+            if n_channels > 1:
+                samples = samples.reshape(-1, n_channels).mean(axis=1)
+            return torch.from_numpy(samples).unsqueeze(0), sample_rate
+
+    # For non-WAV formats, use soundfile (optional dependency)
+    try:
+        import soundfile as sf
+    except ImportError as e:
+        raise ImportError(
+            "soundfile is required to read non-WAV audio files. "
+            "Install with: `pip install soundfile` or `uvx --with soundfile`"
+        ) from e
+
+    data, sample_rate = sf.read(str(filepath), dtype="float32")
+    if data.ndim == 1:
+        wav = torch.from_numpy(data).unsqueeze(0)
+    else:
+        wav = torch.from_numpy(data.mean(axis=1)).unsqueeze(0)
+    return wav, sample_rate
+
+
+class StreamingWAVWriter:
+    """WAV writer using Python's standard library wave module."""
+
+    def __init__(self, output_stream, sample_rate: int):
+        self.output_stream = output_stream
+        self.sample_rate = sample_rate
+        self.wave_writer = None
+        self.first_chunk_buffer = []
+
+    def write_header(self, sample_rate: int):
+        """Initialize WAV writer with header."""
+        # For stdout streaming, we need to handle the unseekable stream case
+        # The wave module supports unseekable streams since Python 3.4
+        self.wave_writer = wave.open(self.output_stream, "wb")
+        self.wave_writer.setnchannels(1)  # Mono
+        self.wave_writer.setsampwidth(2)  # 16-bit
+        self.wave_writer.setframerate(sample_rate)
+        self.wave_writer.setnframes(1_000_000_000)
+
+    def write_pcm_data(self, audio_chunk: torch.Tensor):
+        """Write PCM data using wave module."""
+        # Convert to int16 PCM bytes
+        chunk_int16 = (audio_chunk.clamp(-1, 1) * 32767).short()
+        chunk_bytes = chunk_int16.detach().cpu().numpy().tobytes()
+
+        if self.first_chunk_buffer is not None:
+            self.first_chunk_buffer.append(chunk_bytes)
+            total_length = sum(len(c) for c in self.first_chunk_buffer)
+            target_length = (
+                int(self.sample_rate * FIRST_CHUNK_LENGTH_SECONDS) * 2
+            )  # 2 bytes per sample
+            if total_length < target_length:
+                return
+            self._flush()
+            return
+
+        # Use writeframesraw to avoid frame count validation for streaming
+        self.wave_writer.writeframesraw(chunk_bytes)
+
+    def _flush(self):
+        if self.first_chunk_buffer is not None:
+            self.wave_writer.writeframesraw(b"".join(self.first_chunk_buffer))
+            self.first_chunk_buffer = None
+
+    def finalize(self):
+        """Close the wave writer."""
+        self._flush()
+
+        # Let's add 200ms of silence to ensure proper playback
+        silence_duration_sec = 0.2
+        num_silence_samples = int(self.sample_rate * silence_duration_sec)
+
+        self.wave_writer.writeframesraw(bytes(num_silence_samples * 2))
+
+        if self.wave_writer:
+            # do not update the header for unseekable streams
+            self.wave_writer._patchheader = lambda: None
+            self.wave_writer.close()
+
+
+def is_file_like(obj):
+    """Check if object has basic file-like methods."""
+    return all(hasattr(obj, attr) for attr in ["write", "close"])
+
+
+def stream_audio_chunks(
+    path: str | Path | None | Any, audio_chunks: Iterator[torch.Tensor], sample_rate: int
+):
+    """Stream audio chunks to a WAV file or stdout, optionally playing them."""
+    if path == "-":
+        f = sys.stdout.buffer
+    elif path is None:
+        f = nullcontext()
+    elif is_file_like(path):
+        f = path
+    else:
+        f = open(path, "wb")
+
+    with f:
+        if path is not None:
+            writer = StreamingWAVWriter(f, sample_rate)
+            writer.write_header(sample_rate)
+
+        for chunk in audio_chunks:
+            # Then write to file
+            if path is not None:
+                writer.write_pcm_data(chunk)
+
+        if path is not None:
+            writer.finalize()

pocket_tts/data/audio_utils.py

@@ -0,0 +1,28 @@
+"""Various utilities for audio convertion (pcm format, sample rate and channels),
+and volume normalization."""
+
+import torch
+from scipy.signal import resample_poly
+
+
+def convert_audio(
+    wav: torch.Tensor, from_rate: int | float, to_rate: int | float, to_channels: int
+) -> torch.Tensor:
+    """Convert audio to new sample rate and number of audio channels."""
+    if from_rate != to_rate:
+        # Convert to numpy for scipy resampling
+        wav_np = wav.detach().cpu().numpy()
+
+        # Calculate resampling parameters
+        gcd = int(torch.gcd(torch.tensor(from_rate), torch.tensor(to_rate)).item())
+        up = int(to_rate // gcd)
+        down = int(from_rate // gcd)
+
+        # Resample using scipy
+        resampled_np = resample_poly(wav_np, up, down, axis=-1)
+
+        # Convert back to torch tensor
+        wav = torch.from_numpy(resampled_np).to(wav.device).to(wav.dtype)
+
+    assert wav.shape[-2] == to_channels
+    return wav

pocket_tts/default_parameters.py

@@ -0,0 +1,7 @@
+DEFAULT_AUDIO_PROMPT = "alba"
+DEFAULT_VARIANT = "b6369a24"
+DEFAULT_TEMPERATURE = 0.7
+DEFAULT_LSD_DECODE_STEPS = 1
+DEFAULT_NOISE_CLAMP = None
+DEFAULT_EOS_THRESHOLD = -4.0
+DEFAULT_FRAMES_AFTER_EOS = None

pocket_tts/main.py

@@ -0,0 +1,262 @@
+import io
+import logging
+import os
+import tempfile
+import threading
+from pathlib import Path
+from queue import Queue
+
+import typer
+import uvicorn
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, StreamingResponse
+from typing_extensions import Annotated
+
+from pocket_tts.data.audio import stream_audio_chunks
+from pocket_tts.default_parameters import (
+    DEFAULT_AUDIO_PROMPT,
+    DEFAULT_EOS_THRESHOLD,
+    DEFAULT_FRAMES_AFTER_EOS,
+    DEFAULT_LSD_DECODE_STEPS,
+    DEFAULT_NOISE_CLAMP,
+    DEFAULT_TEMPERATURE,
+    DEFAULT_VARIANT,
+)
+from pocket_tts.models.tts_model import TTSModel
+from pocket_tts.utils.logging_utils import enable_logging
+from pocket_tts.utils.utils import PREDEFINED_VOICES, size_of_dict
+
+logger = logging.getLogger(__name__)
+
+cli_app = typer.Typer(
+    help="Kyutai Pocket TTS - Text-to-Speech generation tool", pretty_exceptions_show_locals=False
+)
+
+
+# ------------------------------------------------------
+# The pocket-tts server implementation
+# ------------------------------------------------------
+
+# Global model instance
+tts_model = None
+global_model_state = None
+
+web_app = FastAPI(
+    title="Kyutai Pocket TTS API", description="Text-to-Speech generation API", version="1.0.0"
+)
+web_app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[
+        "http://localhost:3000",
+        "https://pod1-10007.internal.kyutai.org",
+        "https://kyutai.org",
+    ],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@web_app.get("/")
+async def root():
+    """Serve the frontend."""
+    static_path = Path(__file__).parent / "static" / "index.html"
+    return FileResponse(static_path)
+
+
+@web_app.get("/health")
+async def health():
+    return {"status": "healthy"}
+
+
+def write_to_queue(queue, text_to_generate, model_state):
+    """Allows writing to the StreamingResponse as if it were a file."""
+
+    class FileLikeToQueue(io.IOBase):
+        def __init__(self, queue):
+            self.queue = queue
+
+        def write(self, data):
+            self.queue.put(data)
+
+        def flush(self):
+            pass
+
+        def close(self):
+            self.queue.put(None)
+
+    audio_chunks = tts_model.generate_audio_stream(
+        model_state=model_state, text_to_generate=text_to_generate
+    )
+    stream_audio_chunks(FileLikeToQueue(queue), audio_chunks, tts_model.config.mimi.sample_rate)
+
+
+def generate_data_with_state(text_to_generate: str, model_state: dict):
+    queue = Queue()
+
+    # Run your function in a thread
+    thread = threading.Thread(target=write_to_queue, args=(queue, text_to_generate, model_state))
+    thread.start()
+
+    # Yield data as it becomes available
+    i = 0
+    while True:
+        data = queue.get()
+        if data is None:
+            break
+        i += 1
+        yield data
+
+    thread.join()
+
+
+@web_app.post("/tts")
+def text_to_speech(
+    text: str = Form(...),
+    voice_url: str | None = Form(None),
+    voice_wav: UploadFile | None = File(None),
+):
+    """
+    Generate speech from text using the pre-loaded voice prompt or a custom voice.
+
+    Args:
+        text: Text to convert to speech
+        voice_url: Optional voice URL (http://, https://, or hf://)
+        voice_wav: Optional uploaded voice file (mutually exclusive with voice_url)
+    """
+    if not text.strip():
+        raise HTTPException(status_code=400, detail="Text cannot be empty")
+
+    if voice_url is not None and voice_wav is not None:
+        raise HTTPException(status_code=400, detail="Cannot provide both voice_url and voice_wav")
+
+    # Use the appropriate model state
+    if voice_url is not None:
+        if not (
+            voice_url.startswith("http://")
+            or voice_url.startswith("https://")
+            or voice_url.startswith("hf://")
+            or voice_url in PREDEFINED_VOICES
+        ):
+            raise HTTPException(
+                status_code=400, detail="voice_url must start with http://, https://, or hf://"
+            )
+        model_state = tts_model._cached_get_state_for_audio_prompt(voice_url)
+        logging.warning("Using voice from URL: %s", voice_url)
+    elif voice_wav is not None:
+        # Use uploaded voice file - preserve extension for format detection
+        suffix = Path(voice_wav.filename).suffix if voice_wav.filename else ".wav"
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
+            content = voice_wav.file.read()
+            temp_file.write(content)
+            temp_file.flush()
+
+            try:
+                model_state = tts_model.get_state_for_audio_prompt(
+                    Path(temp_file.name), truncate=True
+                )
+            finally:
+                os.unlink(temp_file.name)
+    else:
+        # Use default global model state
+        model_state = global_model_state
+
+    return StreamingResponse(
+        generate_data_with_state(text, model_state),
+        media_type="audio/wav",
+        headers={
+            "Content-Disposition": "attachment; filename=generated_speech.wav",
+            "Transfer-Encoding": "chunked",
+        },
+    )
+
+
+@cli_app.command()
+def serve(
+    voice: Annotated[
+        str, typer.Option(help="Path to voice prompt audio file (voice to clone)")
+    ] = DEFAULT_AUDIO_PROMPT,
+    host: Annotated[str, typer.Option(help="Host to bind to")] = "localhost",
+    port: Annotated[int, typer.Option(help="Port to bind to")] = 8000,
+    reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False,
+):
+    """Start the FastAPI server."""
+
+    global tts_model, global_model_state
+    tts_model = TTSModel.load_model(DEFAULT_VARIANT)
+
+    # Pre-load the voice prompt
+    global_model_state = tts_model.get_state_for_audio_prompt(voice)
+    logger.info(f"The size of the model state is {size_of_dict(global_model_state) // 1e6} MB")
+
+    uvicorn.run("pocket_tts.main:web_app", host=host, port=port, reload=reload)
+
+
+# ------------------------------------------------------
+# The pocket-tts single generation CLI implementation
+# ------------------------------------------------------
+
+
+@cli_app.command()
+def generate(
+    text: Annotated[
+        str, typer.Option(help="Text to generate")
+    ] = "Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.",
+    voice: Annotated[
+        str, typer.Option(help="Path to audio conditioning file (voice to clone)")
+    ] = DEFAULT_AUDIO_PROMPT,
+    quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False,
+    variant: Annotated[str, typer.Option(help="Model signature")] = DEFAULT_VARIANT,
+    lsd_decode_steps: Annotated[
+        int, typer.Option(help="Number of generation steps")
+    ] = DEFAULT_LSD_DECODE_STEPS,
+    temperature: Annotated[
+        float, typer.Option(help="Temperature for generation")
+    ] = DEFAULT_TEMPERATURE,
+    noise_clamp: Annotated[float, typer.Option(help="Noise clamp value")] = DEFAULT_NOISE_CLAMP,
+    eos_threshold: Annotated[float, typer.Option(help="EOS threshold")] = DEFAULT_EOS_THRESHOLD,
+    frames_after_eos: Annotated[
+        int, typer.Option(help="Number of frames to generate after EOS")
+    ] = DEFAULT_FRAMES_AFTER_EOS,
+    output_path: Annotated[
+        str, typer.Option(help="Output path for generated audio")
+    ] = "./tts_output.wav",
+    device: Annotated[str, typer.Option(help="Device to use")] = "cpu",
+):
+    """Generate speech using Kyutai Pocket TTS."""
+    if "cuda" in device:
+        # Cuda graphs capturing does not play nice with multithreading.
+        os.environ["NO_CUDA_GRAPH"] = "1"
+
+    log_level = logging.ERROR if quiet else logging.INFO
+    with enable_logging("pocket_tts", log_level):
+        tts_model = TTSModel.load_model(
+            variant, temperature, lsd_decode_steps, noise_clamp, eos_threshold
+        )
+        tts_model.to(device)
+
+        model_state_for_voice = tts_model.get_state_for_audio_prompt(voice)
+        # Stream audio generation directly to file or stdout
+        audio_chunks = tts_model.generate_audio_stream(
+            model_state=model_state_for_voice,
+            text_to_generate=text,
+            frames_after_eos=frames_after_eos,
+        )
+
+        stream_audio_chunks(output_path, audio_chunks, tts_model.config.mimi.sample_rate)
+
+        # Only print the result message if not writing to stdout
+        if output_path != "-":
+            logger.info("Results written in %s", output_path)
+        logger.info("-" * 20)
+        logger.info(
+            "If you want to try multiple voices and prompts quickly, try the `serve` command."
+        )
+        logger.info(
+            "If you like Kyutai projects, comment, like, subscribe at https://x.com/kyutai_labs"
+        )
+
+
+if __name__ == "__main__":
+    cli_app()

pocket_tts/models/__init__.py

@@ -0,0 +1,3 @@
+"""
+Models for EnCodec (and Mimi), and FlowLMModel.
+"""