Copy the files in the new repo for public release

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
+weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/tts_b6369a24.safetensors
+
+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
+  #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,126 @@
+"""
+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 using Python's wave module."""
+    with wave.open(str(filepath), "rb") as wav_file:
+        sample_rate = wav_file.getframerate()
+
+        # Read all audio data as 16-bit signed integers
+        raw_data = wav_file.readframes(-1)
+        samples = np.frombuffer(raw_data, dtype=np.int16).astype(np.float32) / 32768.0
+
+        # Return as mono tensor (channels, samples)
+        wav = torch.from_numpy(samples.reshape(1, -1))
+        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 = "marius"
+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,254 @@
+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, truncate=True)
+        logging.warning("Using voice from URL: %s", voice_url)
+    elif voice_wav is not None:
+        # Use uploaded voice file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") 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)
+
+
+if __name__ == "__main__":
+    cli_app()

pocket_tts/models/__init__.py

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

pocket_tts/models/flow_lm.py

@@ -0,0 +1,209 @@
+import logging
+from functools import partial
+
+import torch
+import torch.export
+from beartype.typing import Callable
+from torch import nn
+from typing_extensions import Self
+
+from pocket_tts.conditioners.text import LUTConditioner
+from pocket_tts.modules.mimi_transformer import StreamingTransformer
+from pocket_tts.modules.mlp import SimpleMLPAdaLN
+from pocket_tts.utils.config import FlowLMConfig
+
+logger = logging.getLogger(__name__)
+
+FlowNet2 = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+
+
+def lsd_decode(v_t: FlowNet2, x_0: torch.Tensor, num_steps: int = 1) -> torch.Tensor:
+    """Rebuilds the data sample from starting point x_0.
+
+    Lagrangian Self Distillation (https://arxiv.org/pdf/2505.18825)
+
+    Args:
+        v_t: Function taking t and x_t as input and returning the flow.
+        x_0: Starting point from the known distribution.
+        num_steps: Number of steps to take.
+
+    Returns:
+        x_1_hat: (B, D) Reconstructed data sample.
+    """
+    current = x_0
+    for i in range(num_steps):
+        s = i / num_steps
+        t = (i + 1) / num_steps
+        flow_dir = v_t(
+            s * torch.ones_like(x_0[..., :1]), t * torch.ones_like(x_0[..., :1]), current
+        )
+        current += flow_dir / num_steps
+    return current
+
+
+class FlowLMModel(nn.Module):
+    """Transformer-based flow language model on multiple streams of latents.
+
+    Args:
+        conditioner (LUTConditioner): Text conditioner for processing text inputs.
+        flow: Flow module that defines the flow loss and sampling strategy.
+        flow_net: Trainable function (cond, t, x_t) -> u_t.
+        dim (int): Dimension of the transformer encoder.
+        norm (str): Normalization method.
+        attribute_dropouts (dict): Attribute dropout probabilities.
+        ldim (int): Latent dimension.
+        stats_ema_decay (float): Decay for the EMA of the latent statistics.
+        **kwargs: Additional parameters for the transformer encoder.
+    """
+
+    def __init__(
+        self,
+        conditioner: LUTConditioner,
+        flow_net: SimpleMLPAdaLN,
+        transformer: StreamingTransformer,
+        dim: int = 128,
+        ldim: int = 64,
+        stats_ema_decay: float = 0.999,
+        text_padding_weight: float = 1.0,
+        dtype=None,
+    ):
+        super().__init__()
+        self.conditioner = conditioner
+        self.ldim = ldim
+        self.stats_ema_decay = stats_ema_decay
+        self.dim = dim
+        self.text_padding_weight = text_padding_weight
+        self.dtype = dtype
+
+        self.flow_net = flow_net
+        self.register_buffer("emb_std", torch.ones(ldim, dtype=dtype))
+        self.register_buffer("emb_mean", torch.zeros(ldim, dtype=dtype))
+        self.bos_emb = torch.nn.Parameter(torch.randn(ldim, dtype=dtype))
+
+        self.input_linear = nn.Linear(self.ldim, dim, bias=False, dtype=dtype)
+        self.transformer = transformer
+        self.out_norm = nn.LayerNorm(dim, eps=1e-5)
+        self.out_eos = nn.Linear(dim, 1, dtype=dtype)
+
+    @property
+    def device(self) -> str:
+        return next(self.parameters()).device.type
+
+    def forward(
+        self,
+        sequence: torch.Tensor,
+        text_embeddings: torch.Tensor,
+        model_state: dict,
+        lsd_decode_steps: int,
+        temp: float,
+        noise_clamp: float | None,
+        eos_threshold: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply language model on sequence and conditions.
+        Given a tensor of sequence of shape [B, S, ldim], returns the loss in training mode
+        or the reconstructed latent in generation mode.
+
+        Args:
+            sequence (torch.Tensor): Latents to model.
+            text_embeddings (torch.Tensor): Pre-computed conditioning
+                tensor.
+            lsd_decode_steps (int): Number of steps to decode when generating audio.
+                If zero, the model computes the loss.
+        Returns:
+            (output, eos_output, metrics). If `lsd_decode_steps` is zero, `output` is the loss tensor of shape [B, S],
+            otherwise it is the reconstructed latent.
+        """
+        # NaN values signal a BOS position.
+        sequence = torch.where(torch.isnan(sequence), self.bos_emb, sequence)
+        input_ = self.input_linear(sequence)
+
+        transformer_out = self.backbone(input_, text_embeddings, sequence, model_state=model_state)
+        transformer_out = transformer_out.to(torch.float32)
+        assert lsd_decode_steps > 0
+
+        transformer_out = transformer_out[:, -1]
+        out_eos = self.out_eos(transformer_out) > eos_threshold
+
+        noise_shape = transformer_out.shape[:-1] + (self.ldim,)
+        std = temp**0.5
+        noise = torch.empty(noise_shape, dtype=transformer_out.dtype, device=transformer_out.device)
+        if noise_clamp is None:
+            torch.nn.init.normal_(noise, mean=0.0, std=std)
+        else:
+            torch.nn.init.trunc_normal_(noise, mean=0.0, std=std, a=-noise_clamp, b=noise_clamp)
+        conditioned_flow = partial(self.flow_net, transformer_out)
+        return lsd_decode(conditioned_flow, noise, lsd_decode_steps), out_eos
+
+    def backbone(
+        self, input_, text_embeddings: torch.Tensor, sequence, model_state: dict
+    ) -> torch.Tensor:
+        # Most of the time, one of those two tensors is empty, it allows us
+        # to input text or audio embeddings into the model without adding an
+        # if-else branch.
+        # print("text_embeddings shape:", text_embeddings.shape)
+        # if text_embeddings.numel() != 0:
+        #     torch.save(text_embeddings, "debug_flow_lm_text_embeddings.pt")
+        input_ = torch.cat([text_embeddings, input_], dim=1)
+        # transformer_out = self.transformer(input_, model_state=model_state)
+        transformer_out = self.transformer(input_, model_state)
+        if self.out_norm:
+            transformer_out = self.out_norm(transformer_out)
+        # remove the prefix from the model outputs (condition is prepended)
+        transformer_out = transformer_out[:, -sequence.shape[1] :]
+        return transformer_out
+
+    def _sample_next_latent(
+        self,
+        sequence: torch.Tensor,
+        text_embeddings: torch.Tensor,
+        model_state: dict,
+        lsd_decode_steps: int,
+        temp: float,
+        noise_clamp: float | None,
+        eos_threshold: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Sample next latent from the model given a sequence and a set of conditions.
+        Args:
+            sequence (torch.Tensor): Current sequence of shape [B, K, S]
+                with K corresponding to the number of codebooks and S the number of sequence steps.
+                S = 1 in streaming mode, except for the first step that contains a bigger prompt.
+            text_embeddings (torch.Tensor): Condition tensor.
+            n_steps (int): Number of flow steps to decode when generating audio.
+        Returns:
+            next_latent (torch.Tensor), is_eos (torch.Tensor): Next latent tensor of shape [B, 1, ldim]
+                and is_eos tensor of shape [B, 1] with 1 on EOS positions.
+        """
+        result = self(
+            sequence=sequence,
+            text_embeddings=text_embeddings,
+            lsd_decode_steps=lsd_decode_steps,
+            temp=temp,
+            noise_clamp=noise_clamp,
+            eos_threshold=eos_threshold,
+            model_state=model_state,
+        )
+
+        return result
+
+    @classmethod
+    def from_pydantic_config(cls, config: FlowLMConfig, latent_dim: int) -> Self:
+        d_model = config.transformer.d_model
+        flow_mlp = SimpleMLPAdaLN.from_pydantic_config(config, latent_dim, d_model)
+
+        conditioner = LUTConditioner(
+            n_bins=config.lookup_table.n_bins,
+            tokenizer_path=str(config.lookup_table.tokenizer_path),
+            dim=config.lookup_table.dim,
+            output_dim=d_model,
+        )
+
+        transformer = StreamingTransformer.from_pydantic_config(config.transformer)
+
+        return cls(
+            flow_net=flow_mlp,
+            transformer=transformer,
+            dim=d_model,
+            conditioner=conditioner,
+            ldim=latent_dim,
+            dtype=getattr(torch, config.dtype),
+        )

pocket_tts/models/mimi.py

@@ -0,0 +1,111 @@
+import logging
+
+import torch
+from torch import nn
+
+from pocket_tts.modules.conv import pad_for_conv1d
+from pocket_tts.modules.dummy_quantizer import DummyQuantizer
+from pocket_tts.modules.mimi_transformer import ProjectedTransformer
+from pocket_tts.modules.resample import ConvDownsample1d, ConvTrUpsample1d
+from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder
+
+logger = logging.getLogger()
+
+
+class MimiModel(nn.Module):
+    def __init__(
+        self,
+        encoder: SEANetEncoder,
+        decoder: SEANetDecoder,
+        quantizer: DummyQuantizer,
+        frame_rate: float,
+        encoder_frame_rate: float,
+        sample_rate: int,
+        channels: int,
+        encoder_transformer: ProjectedTransformer,
+        decoder_transformer: ProjectedTransformer,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.encoder_transformer = encoder_transformer
+        self.decoder_transformer = decoder_transformer
+        self.quantizer = quantizer
+        self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.encoder_frame_rate = encoder_frame_rate
+
+        # We will need the dimension for the resampling. In general the encoder will be a SeanetEncoder
+        # which exposes a `dimension` attribute.
+        dimension = encoder.dimension
+        assert isinstance(dimension, int), (
+            f"Dimension should be int, got {dimension} of type {type(dimension)}."
+        )
+        self.dimension = dimension
+
+        if encoder_frame_rate != frame_rate:
+            assert self.encoder_frame_rate > self.frame_rate, "Cannot upsample with conv."
+            downsample_stride = self.encoder_frame_rate / self.frame_rate
+            assert downsample_stride == int(downsample_stride), (
+                f"Only integer strides are supported, got {downsample_stride}"
+            )
+            self.downsample = ConvDownsample1d(int(downsample_stride), dimension=dimension)
+            self.upsample = ConvTrUpsample1d(int(downsample_stride), dimension=dimension)
+
+    @property
+    def frame_size(self) -> int:
+        return int(self.sample_rate / self.frame_rate)
+
+    def _to_framerate(self, x: torch.Tensor):
+        # Convert from the encoder frame rate to the overall framerate.
+        _, _, length = x.shape
+        frame_rate = self.encoder_frame_rate
+        new_frame_rate = self.frame_rate
+        if frame_rate == new_frame_rate:
+            return x
+        return self.downsample(x, model_state=None)
+
+    def _to_encoder_framerate(self, x: torch.Tensor, mimi_state) -> torch.Tensor:
+        # Convert from overall framerate to the encoder frame rate.
+        _, _, length = x.shape
+        frame_rate = self.encoder_frame_rate
+        new_frame_rate = self.frame_rate
+        if frame_rate == new_frame_rate:
+            return x
+        return self.upsample(x, mimi_state)
+
+    def forward(self, x: torch.Tensor):
+        raise NotImplementedError()
+
+    def decode_from_latent(self, latent: torch.Tensor, mimi_state) -> torch.Tensor:
+        emb = self._to_encoder_framerate(latent, mimi_state)
+        (emb,) = self.decoder_transformer(emb, mimi_state)
+        out = self.decoder(emb, mimi_state)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def encode_to_latent(self, x: torch.Tensor) -> torch.Tensor:
+        """Projects a batch of waveforms to unquantized latent space.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T].
+
+        Returns:
+            Unquantized embeddings.
+        """
+        assert x.dim() == 3, (
+            f"CompressionModel._encode_to_unquantized_latent expects audio of shape [B, C, T] but got {x.shape}"
+        )
+
+        frame_size = self.frame_size
+
+        # The underlying convolutions no longer accept partial inputs,
+        # `x` needs to be exactly a multiple of the frame size,
+        # reproducing the previous padding behavior here.
+        x = pad_for_conv1d(x, frame_size, frame_size)
+        emb = self.encoder(x, model_state=None)
+
+        (emb,) = self.encoder_transformer(emb, model_state=None)
+        emb = self._to_framerate(emb)
+        return emb

pocket_tts/models/tts_model.py

@@ -0,0 +1,719 @@
+import copy
+import logging
+import os
+import queue
+import statistics
+import threading
+import time
+from functools import lru_cache
+from pathlib import Path
+
+import safetensors
+import torch
+from torch import nn
+from torch.nn import functional as F
+from typing_extensions import Self
+
+from pocket_tts.conditioners.base import TokenizedText
+from pocket_tts.data.audio import audio_read
+from pocket_tts.data.audio_utils import convert_audio
+from pocket_tts.default_parameters import (
+    DEFAULT_EOS_THRESHOLD,
+    DEFAULT_LSD_DECODE_STEPS,
+    DEFAULT_NOISE_CLAMP,
+    DEFAULT_TEMPERATURE,
+    DEFAULT_VARIANT,
+)
+from pocket_tts.models.flow_lm import FlowLMModel
+from pocket_tts.models.mimi import MimiModel
+from pocket_tts.modules import mimi_transformer
+from pocket_tts.modules.dummy_quantizer import DummyQuantizer
+from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder
+from pocket_tts.modules.stateful_module import increment_steps, init_states
+from pocket_tts.utils.config import Config, load_config
+from pocket_tts.utils.utils import (
+    PREDEFINED_VOICES,
+    display_execution_time,
+    download_if_necessary,
+    load_predefined_voice,
+    size_of_dict,
+)
+from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_state_dict
+
+torch.set_num_threads(1)
+logger = logging.getLogger(__name__)
+
+
+KEYS_TO_MODIFY = [
+    "mimi.decoder.model.0.conv.conv.bias",
+    "mimi.decoder.model.0.conv.conv.weight",
+    "mimi.decoder.model.11.conv.conv.bias",
+    "mimi.decoder.model.11.conv.conv.weight",
+    "mimi.decoder.model.2.convtr.convtr.bias",
+    "mimi.decoder.model.2.convtr.convtr.weight",
+    "mimi.decoder.model.3.block.1.conv.conv.bias",
+    "mimi.decoder.model.3.block.1.conv.conv.weight",
+    "mimi.decoder.model.3.block.3.conv.conv.bias",
+    "mimi.decoder.model.3.block.3.conv.conv.weight",
+    "mimi.decoder.model.5.convtr.convtr.bias",
+    "mimi.decoder.model.5.convtr.convtr.weight",
+    "mimi.decoder.model.6.block.1.conv.conv.bias",
+    "mimi.decoder.model.6.block.1.conv.conv.weight",
+    "mimi.decoder.model.6.block.3.conv.conv.bias",
+    "mimi.decoder.model.6.block.3.conv.conv.weight",
+    "mimi.decoder.model.8.convtr.convtr.bias",
+    "mimi.decoder.model.8.convtr.convtr.weight",
+    "mimi.decoder.model.9.block.1.conv.conv.bias",
+    "mimi.decoder.model.9.block.1.conv.conv.weight",
+    "mimi.decoder.model.9.block.3.conv.conv.bias",
+    "mimi.decoder.model.9.block.3.conv.conv.weight",
+    "mimi.downsample.conv.conv.conv.weight",
+    "mimi.encoder.model.0.conv.conv.bias",
+    "mimi.encoder.model.0.conv.conv.weight",
+    "mimi.encoder.model.1.block.1.conv.conv.bias",
+    "mimi.encoder.model.1.block.1.conv.conv.weight",
+    "mimi.encoder.model.1.block.3.conv.conv.bias",
+    "mimi.encoder.model.1.block.3.conv.conv.weight",
+    "mimi.encoder.model.11.conv.conv.bias",
+    "mimi.encoder.model.11.conv.conv.weight",
+    "mimi.encoder.model.3.conv.conv.bias",
+    "mimi.encoder.model.3.conv.conv.weight",
+    "mimi.encoder.model.4.block.1.conv.conv.bias",
+    "mimi.encoder.model.4.block.1.conv.conv.weight",
+    "mimi.encoder.model.4.block.3.conv.conv.bias",
+    "mimi.encoder.model.4.block.3.conv.conv.weight",
+    "mimi.encoder.model.6.conv.conv.bias",
+    "mimi.encoder.model.6.conv.conv.weight",
+    "mimi.encoder.model.7.block.1.conv.conv.bias",
+    "mimi.encoder.model.7.block.1.conv.conv.weight",
+    "mimi.encoder.model.7.block.3.conv.conv.bias",
+    "mimi.encoder.model.7.block.3.conv.conv.weight",
+    "mimi.encoder.model.9.conv.conv.bias",
+    "mimi.encoder.model.9.conv.conv.weight",
+    "mimi.upsample.convtr.convtr.convtr.weight",
+]
+
+KEY_WITH_REMOVED_ZERO = [
+    "mimi.encoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
+    "mimi.encoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
+    "mimi.encoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
+    "mimi.encoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
+    "mimi.decoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
+    "mimi.decoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
+    "mimi.decoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
+    "mimi.decoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
+]
+
+
+def adapt_state_dict_of_TTSModel(state_dict: dict) -> dict:
+    # should be removed before release
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        if key == "mimi.quantizer.input_proj.weight":
+            continue
+        if key in KEYS_TO_MODIFY:
+            new_key = key.replace(".conv.conv.", ".conv.").replace(".convtr.convtr.", ".convtr.")
+            new_state_dict[new_key] = value
+        elif key in KEY_WITH_REMOVED_ZERO:
+            new_key = key.replace("s.0.weight", ".weight")
+            new_state_dict[new_key] = value
+        else:
+            new_state_dict[key] = value
+    return new_state_dict
+
+
+class TTSModel(nn.Module):
+    def __init__(
+        self,
+        flow_lm: FlowLMModel,
+        temp: float,
+        lsd_decode_steps: int,
+        noise_clamp: float | None,
+        eos_threshold,
+        config: Config,
+    ):
+        super().__init__()
+        self.flow_lm = flow_lm
+        self.temp = temp
+        self.lsd_decode_steps = lsd_decode_steps
+        self.noise_clamp = noise_clamp
+        self.eos_threshold = eos_threshold
+        self.config = config
+        self.has_voice_cloning = True
+
+    @property
+    def device(self) -> str:
+        return next(self.parameters()).device.type
+
+    @property
+    def sample_rate(self) -> int:
+        return self.config.mimi.sample_rate
+
+    @classmethod
+    def _from_pydantic_config(
+        cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold
+    ) -> Self:
+        flow_lm = FlowLMModel.from_pydantic_config(
+            config.flow_lm, latent_dim=config.mimi.quantizer.dimension
+        )
+        tts_model = cls(flow_lm, temp, lsd_decode_steps, noise_clamp, eos_threshold, config)
+        return tts_model
+
+    @classmethod
+    def _from_pydantic_config_with_weights(
+        cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold
+    ) -> Self:
+        tts_model = cls._from_pydantic_config(
+            config, temp, lsd_decode_steps, noise_clamp, eos_threshold
+        )
+        tts_model.flow_lm.speaker_proj_weight = torch.nn.Parameter(
+            torch.zeros((1024, 512), dtype=torch.float32)
+        )
+        if config.flow_lm.weights_path is not None:
+            if config.mimi.weights_path is None:
+                raise ValueError(
+                    "If you specify flow_lm.weights_path you should specify mimi.weights_path"
+                )
+            logger.info(f"Loading FlowLM weights from {config.flow_lm.weights_path}")
+            state_dict_flowlm = get_flow_lm_state_dict(
+                download_if_necessary(config.flow_lm.weights_path)
+            )
+            tts_model.flow_lm.load_state_dict(state_dict_flowlm, strict=True)
+
+        # safetensors.torch.save_file(tts_model.state_dict(), "7442637a.safetensors")
+        # Create mimi config directly from the provided config using model_dump
+        mimi_config = config.mimi.model_dump()
+
+        # Build mimi model from config
+        encoder = SEANetEncoder(**mimi_config["seanet"])
+        decoder = SEANetDecoder(**mimi_config["seanet"])
+
+        encoder_transformer = mimi_transformer.ProjectedTransformer(**mimi_config["transformer"])
+        decoder_transformer = mimi_transformer.ProjectedTransformer(**mimi_config["transformer"])
+        quantizer = DummyQuantizer(**mimi_config["quantizer"])
+
+        tts_model.mimi = MimiModel(
+            encoder,
+            decoder,
+            quantizer,
+            channels=mimi_config["channels"],
+            sample_rate=mimi_config["sample_rate"],
+            frame_rate=mimi_config["frame_rate"],
+            encoder_frame_rate=mimi_config["sample_rate"] / encoder.hop_length,
+            encoder_transformer=encoder_transformer,
+            decoder_transformer=decoder_transformer,
+        ).to(device="cpu")
+
+        # Load mimi weights from the config safetensors file with complete mapping for strict loading
+
+        if config.mimi.weights_path is not None:
+            if config.flow_lm.weights_path is None:
+                raise ValueError(
+                    "If you specify mimi.weights_path you should specify flow_lm.weights_path"
+                )
+            logger.info(f"Loading Mimi weights from {config.mimi.weights_path}")
+            mimi_state = get_mimi_state_dict(download_if_necessary(config.mimi.weights_path))
+            tts_model.mimi.load_state_dict(mimi_state, strict=True)
+
+        tts_model.mimi.eval()
+        # tts_model.to(dtype=torch.float32)
+
+        # uncomment to save the weights
+        # tts_model = tts_model.to(dtype=torch.bfloat16)
+        # safetensors.torch.save_file(tts_model.state_dict(), "tts_b6369a24.safetensors")
+        if config.weights_path is not None:
+            logger.info(f"Loading TTSModel weights from {config.weights_path}")
+            try:
+                weights_file = download_if_necessary(config.weights_path)
+            except Exception:
+                tts_model.has_voice_cloning = False
+                weights_file = download_if_necessary(config.weights_path_without_voice_cloning)
+
+            state_dict = safetensors.torch.load_file(weights_file)
+            state_dict = adapt_state_dict_of_TTSModel(state_dict)
+            tts_model.load_state_dict(state_dict, strict=True)
+
+        if config.flow_lm.weights_path is None and config.weights_path is None:
+            logger.warning(
+                "No weights_path specified for FlowLM or TTSModel, model is uninitialized!"
+            )
+        size_in_mb = size_of_dict(tts_model.state_dict()) // 1e6
+        logging.info(f"TTS Model loaded successfully. Its size is {size_in_mb} MB")
+
+        return tts_model
+
+    def load_model(
+        variant: str = DEFAULT_VARIANT,
+        temp: float | int = DEFAULT_TEMPERATURE,
+        lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS,
+        noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP,
+        eos_threshold: float = DEFAULT_EOS_THRESHOLD,
+    ) -> Self:
+        """Load a pre-trained TTS model with specified configuration.
+
+        This class method loads a complete TTS model including the flow language model
+        and Mimi compression model from pre-trained weights. The model is initialized
+        with the specified generation parameters and ready for inference.
+
+        Args:
+            variant: Model variant identifier corresponding to a config file name
+                (e.g., '610b0b2c'). Must match a YAML file in the config directory.
+            temp: Sampling temperature for generation. Higher values produce more
+                diverse but potentially lower quality output.
+            lsd_decode_steps: Number of steps for Lagrangian Self Distillation
+                decoding. More steps can improve quality but increase computation.
+            noise_clamp: Maximum value for noise sampling. If None, no clamping
+                is applied. Helps prevent extreme values in generation.
+            eos_threshold: Threshold for end-of-sequence detection. Higher values
+                make the model more likely to continue generating.
+
+        Returns:
+            TTSModel: Fully initialized model with loaded weights on cpu, ready for
+                text-to-speech generation.
+
+        Raises:
+            FileNotFoundError: If the specified config file or model weights
+                are not found.
+            ValueError: If the configuration is invalid or incompatible.
+        """
+        config = load_config(Path(__file__).parents[1] / f"config/{variant}.yaml")
+        tts_model = TTSModel._from_pydantic_config_with_weights(
+            config, temp, lsd_decode_steps, noise_clamp, eos_threshold
+        )
+        return tts_model
+
+    def _run_flow_lm_and_increment_step(
+        self,
+        model_state: dict,
+        text_tokens: torch.Tensor | None = None,
+        backbone_input_latents: torch.Tensor | None = None,
+        audio_conditioning: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """First one is the backbone output, second one is the audio decoding output."""
+        if text_tokens is None:
+            text_tokens = torch.zeros((1, 0), dtype=torch.int64, device=self.flow_lm.device)
+        if backbone_input_latents is None:
+            backbone_input_latents = torch.empty(
+                (1, 0, self.flow_lm.ldim), dtype=self.flow_lm.dtype, device=self.flow_lm.device
+            )
+        if audio_conditioning is None:
+            audio_conditioning = torch.empty(
+                (1, 0, self.flow_lm.dim), dtype=self.flow_lm.dtype, device=self.flow_lm.device
+            )
+
+        output = self._run_flow_lm(
+            text_tokens=text_tokens,
+            backbone_input_latents=backbone_input_latents,
+            model_state=model_state,
+            audio_conditioning=audio_conditioning,
+        )
+        increment_by = (
+            text_tokens.shape[1] + backbone_input_latents.shape[1] + audio_conditioning.shape[1]
+        )
+        increment_steps(self.flow_lm, model_state, increment=increment_by)
+        return output
+
+    def _run_flow_lm(
+        self,
+        model_state: dict,
+        text_tokens: torch.Tensor,
+        backbone_input_latents: torch.Tensor,
+        audio_conditioning: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        text_embeddings = self.flow_lm.conditioner(TokenizedText(text_tokens))
+        text_embeddings = torch.cat([text_embeddings, audio_conditioning], dim=1)
+
+        output_embeddings, is_eos = self.flow_lm._sample_next_latent(
+            backbone_input_latents,
+            text_embeddings,
+            model_state=model_state,
+            lsd_decode_steps=self.lsd_decode_steps,
+            temp=self.temp,
+            noise_clamp=self.noise_clamp,
+            eos_threshold=self.eos_threshold,
+        )
+        return output_embeddings[:, None, :], is_eos
+
+    def _encode_audio(self, audio: torch.Tensor) -> torch.Tensor:
+        encoded = self.mimi.encode_to_latent(audio)
+        latents = encoded.transpose(-1, -2).to(torch.float32)
+        conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight)
+        return conditioning
+
+    def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue):
+        """Worker thread function for decoding audio latents from queue with immediate streaming."""
+        try:
+            audio_chunks = []
+            mimi_state = init_states(self.mimi, batch_size=1, sequence_length=1000)
+            while True:
+                latent = latents_queue.get()
+                if latent is None:
+                    break
+                mimi_decoding_input = latent * self.flow_lm.emb_std + self.flow_lm.emb_mean
+                transposed = mimi_decoding_input.transpose(-1, -2)
+                quantized = self.mimi.quantizer(transposed)
+
+                t = time.monotonic()
+                audio_frame = self.mimi.decode_from_latent(quantized, mimi_state)
+                increment_steps(self.mimi, mimi_state, increment=16)
+                audio_frame_duration = audio_frame.shape[2] / self.config.mimi.sample_rate
+                # We could log the timings here.
+                logger.debug(
+                    " " * 30 + "Decoded %d ms of audio with mimi in %d ms",
+                    int(audio_frame_duration * 1000),
+                    int((time.monotonic() - t) * 1000),
+                )
+                audio_chunks.append(audio_frame)
+
+                result_queue.put(("chunk", audio_frame))
+
+                latents_queue.task_done()
+
+            # Signal completion
+            result_queue.put(("done", None))
+
+        except Exception as e:
+            # Put error in result queue
+            result_queue.put(("error", e))
+
+    @torch.no_grad
+    def generate_audio(
+        self,
+        model_state: dict,
+        text_to_generate: str,
+        frames_after_eos: int | None = None,
+        copy_state: bool = True,
+    ) -> torch.Tensor:
+        """Generate complete audio tensor from text input.
+
+        This method generates the full audio output for the given text prompt
+        and returns it as a single tensor. It internally uses the streaming
+        generation method but collects all chunks before returning.
+
+        This method is NOT thread-safe; separate model instances should be used
+        for concurrent generation.
+
+        Args:
+            model_state: Model state dictionary containing hidden states and
+                positional information. Can be obtained from get_state_for_audio_prompt()
+                or init_states(). The state may be modified during generation.
+            text_to_generate: Input text to convert to speech. The text will be
+                automatically formatted (capitalization, punctuation) for optimal
+                generation quality.
+            frames_after_eos: Number of additional frames to generate after
+                detecting end-of-sequence. If None, automatically determined
+                based on text length (1-3 frames).
+            copy_state: Whether to create a deep copy of the model state before
+                generation. If True, preserves the original state for reuse.
+                If False, modifies the input state in-place. Defaults to True.
+
+        Returns:
+            torch.Tensor: Generated audio tensor with shape [channels, samples]
+                at the model's sample rate (typically 24kHz). The audio is
+                normalized and ready for playback or saving.
+                You can get the sample rate from the `sample_rate` attribute.
+
+        Raises:
+            ValueError: If text_to_generate is empty or invalid.
+            RuntimeError: If generation fails due to model errors.
+        """
+        audio_chunks = []
+        for chunk in self.generate_audio_stream(
+            model_state=model_state,
+            text_to_generate=text_to_generate,
+            frames_after_eos=frames_after_eos,
+            copy_state=copy_state,
+        ):
+            audio_chunks.append(chunk)
+        return torch.cat(audio_chunks, dim=0)
+
+    @torch.no_grad
+    def generate_audio_stream(
+        self,
+        model_state: dict,
+        text_to_generate: str,
+        frames_after_eos: int | None = None,
+        copy_state: bool = True,
+    ):
+        """Generate audio streaming chunks from text input.
+
+        This method generates audio from text and yields chunks as they become
+        available, enabling real-time playback or processing. It uses multithreading
+        to parallelize generation and decoding for optimal performance.
+        This method is NOT thread-safe; separate model instances should be used
+        for concurrent generation.
+
+        Args:
+            model_state: Model state dictionary containing hidden states and
+                positional information. Can be obtained from get_state_for_audio_prompt()
+                or init_states(). The state may be modified during generation.
+            text_to_generate: Input text to convert to speech. The text will be
+                automatically formatted (capitalization, punctuation) for optimal
+                generation quality.
+            frames_after_eos: Number of additional frames to generate after
+                detecting end-of-sequence. If None, automatically determined
+                based on text length (1-3 frames). Defaults to None.
+            copy_state: Whether to create a deep copy of the model state before
+                generation. If True, preserves the original state for reuse.
+                If False, modifies the input state in-place. Defaults to True.
+
+        Yields:
+            torch.Tensor: Audio chunks with shape [samples] at the model's
+                sample rate (typically 24kHz). Chunks are yielded as soon as
+                they are decoded, enabling real-time streaming.
+
+        Raises:
+            ValueError: If text_to_generate is empty or invalid.
+            RuntimeError: If generation fails due to model errors or threading issues.
+
+        Note:
+            This method uses multithreading to parallelize latent generation
+            and audio decoding. Generation performance is logged including
+            real-time factor (RTF) metrics.
+        """
+        text_to_generate, frames_after_eos_guess = prepare_text_prompt(text_to_generate)
+        if frames_after_eos is None:
+            frames_after_eos = frames_after_eos_guess
+
+        if copy_state:
+            model_state = copy.deepcopy(model_state)
+
+        # Set up multithreaded generation and decoding
+        latents_queue = queue.Queue()
+        result_queue = queue.Queue()
+
+        # Start decoder worker thread
+        decoder_thread = threading.Thread(
+            target=self._decode_audio_worker, args=(latents_queue, result_queue), daemon=True
+        )
+        logger.info("starting timer now!")
+        t_generating = time.monotonic()
+        decoder_thread.start()
+
+        # Generate latents and add them to queue (decoder processes them in parallel)
+        self._generate(
+            model_state=model_state,
+            text_to_generate=text_to_generate,
+            frames_after_eos=frames_after_eos,
+            latents_queue=latents_queue,
+            result_queue=result_queue,
+        )
+
+        # Stream audio chunks as they become available
+        total_generated_samples = 0
+        while True:
+            result = result_queue.get()
+            if result[0] == "chunk":
+                # Audio chunk available immediately for streaming/playback
+                audio_chunk = result[1]
+                total_generated_samples += audio_chunk.shape[-1]
+                yield audio_chunk[0, 0]  # Remove batch, channel
+            elif result[0] == "done":
+                # Generation complete
+                break
+            elif result[0] == "error":
+                # Wait for decoder thread to finish cleanly before propagating error
+                with display_execution_time("Waiting for mimi decoder to finish"):
+                    decoder_thread.join()
+                # Propagate error
+                raise result[1]
+
+        # Wait for decoder thread to finish cleanly
+        with display_execution_time("Waiting for mimi decoder to finish"):
+            decoder_thread.join()
+
+        # Print timing information
+        duration_generated_audio = int(
+            total_generated_samples * 1000 / self.config.mimi.sample_rate
+        )
+        generation_time = int((time.monotonic() - t_generating) * 1000)
+        real_time_factor = duration_generated_audio / generation_time
+
+        logger.info(
+            "Generated: %d ms of audio in %d ms so %.2fx faster than real-time",
+            duration_generated_audio,
+            generation_time,
+            real_time_factor,
+        )
+
+    @torch.no_grad
+    def _generate(
+        self,
+        model_state: dict,
+        text_to_generate: str,
+        frames_after_eos: int,
+        latents_queue: queue.Queue,
+        result_queue: queue.Queue,
+    ):
+        gen_len_sec = len(text_to_generate.split()) * 1 + 2.0
+        max_gen_len = int(gen_len_sec * 12.5)
+        prepared = self.flow_lm.conditioner.prepare(text_to_generate)
+
+        with display_execution_time("Prompting text"):
+            self._run_flow_lm_and_increment_step(
+                model_state=model_state, text_tokens=prepared.tokens
+            )
+
+        def run_generation():
+            try:
+                self._autoregressive_generation(
+                    model_state, max_gen_len, frames_after_eos, latents_queue
+                )
+            except Exception as e:
+                logger.error(f"Error in autoregressive generation: {e}")
+                # Signal decoder to stop by putting None (completion sentinel)
+                if latents_queue is not None:
+                    latents_queue.put(None)
+                # Report error to main thread
+                if result_queue is not None:
+                    result_queue.put(("error", e))
+
+        generation_thread = threading.Thread(target=run_generation, daemon=True)
+        generation_thread.start()
+
+    def _autoregressive_generation(
+        self, model_state: dict, max_gen_len: int, frames_after_eos: int, latents_queue: queue.Queue
+    ):
+        backbone_input = torch.full(
+            (1, 1, self.flow_lm.ldim),
+            fill_value=float("NaN"),
+            device=next(iter(self.flow_lm.parameters())).device,
+            dtype=self.flow_lm.dtype,
+        )
+        steps_times = []
+        eos_step = None
+        for generation_step in range(max_gen_len):
+            with display_execution_time("Generating latent", print_output=False) as timer:
+                next_latent, is_eos = self._run_flow_lm_and_increment_step(
+                    model_state=model_state, backbone_input_latents=backbone_input
+                )
+                if is_eos.item() and eos_step is None:
+                    eos_step = generation_step
+                if eos_step is not None and generation_step >= eos_step + frames_after_eos:
+                    break
+
+                # Add generated latent to queue for immediate decoding
+                latents_queue.put(next_latent)
+                backbone_input = next_latent
+            steps_times.append(timer.elapsed_time_ms)
+        else:
+            if os.environ.get("KPOCKET_TTS_ERROR_WITHOUT_EOS", "0") == "1":
+                raise RuntimeError("Generation reached maximum length without EOS!")
+            logger.warning(
+                "Maximum generation length reached without EOS, this very often indicates an error."
+            )
+
+        # Add sentinel value to signal end of generation
+        latents_queue.put(None)
+        logger.info("Average generation step time: %d ms", int(statistics.mean(steps_times)))
+
+    @lru_cache(maxsize=2)
+    def _cached_get_state_for_audio_prompt(
+        self, audio_conditioning: Path | str | torch.Tensor, truncate: bool = False
+    ) -> dict:
+        return self.get_state_for_audio_prompt(audio_conditioning, truncate)
+
+    @torch.no_grad
+    def get_state_for_audio_prompt(
+        self, audio_conditioning: Path | str | torch.Tensor, truncate: bool = False
+    ) -> dict:
+        """Create model state conditioned on audio prompt for continuation.
+
+        This method processes an audio prompt and creates a model state that
+        captures the acoustic characteristics (speaker voice, style, prosody)
+        for use in subsequent text-to-speech generation. The resulting state
+        enables voice cloning and audio continuation with speaker consistency.
+
+        Args:
+            audio_conditioning: Audio prompt to condition on. Can be:
+                - Path: Local file path to audio file
+                - str: URL to download audio file from
+                - torch.Tensor: Pre-loaded audio tensor with shape [channels, samples]
+            truncate: Whether to truncate long audio prompts to 30 seconds.
+                Helps prevent memory issues with very long inputs. Defaults to False.
+
+        Returns:
+            dict: Model state dictionary containing hidden states and positional
+                information conditioned on the audio prompt. This state can be
+                passed to `generate_audio()` or `generate_audio_stream()` for
+                voice-consistent generation.
+
+        Raises:
+            FileNotFoundError: If audio file path doesn't exist.
+            ValueError: If audio tensor is invalid or empty.
+            RuntimeError: If audio processing or encoding fails.
+
+        Note:
+            - Audio is automatically resampled to the model's sample rate (24kHz)
+            - The audio is encoded using the Mimi compression model and projected
+              to the flow model's latent space
+            - Processing time is logged for performance monitoring
+            - The state preserves speaker characteristics for voice cloning
+        """
+        if isinstance(audio_conditioning, str) and audio_conditioning in PREDEFINED_VOICES:
+            # We get the audio conditioning directly from the safetensors file.
+            prompt = load_predefined_voice(audio_conditioning)
+        else:
+            if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
+                raise ValueError(
+                    f"We could not download the weights for the model with voice cloning, "
+                    f"but you're trying to use voice cloning. "
+                    f"Without voice cloning, you can use our catalog of voices {list(PREDEFINED_VOICES)}. "
+                    f"If you want access to the model with voice cloning, go to "
+                    f"https://huggingface.co/kyutai/pocket-tts and accept the terms, "
+                    f"then make sure you're logged in locally with `uvx hf auth login`."
+                )
+            if isinstance(audio_conditioning, str):
+                audio_conditioning = download_if_necessary(audio_conditioning)
+
+            if isinstance(audio_conditioning, Path):
+                audio, conditioning_sample_rate = audio_read(audio_conditioning)
+
+                if truncate:
+                    max_samples = int(30 * conditioning_sample_rate)  # 30 seconds of audio
+                    if audio.shape[-1] > max_samples:
+                        audio = audio[..., :max_samples]
+                        logger.info(f"Audio truncated to first 30 seconds ({max_samples} samples)")
+
+                audio_conditioning = convert_audio(
+                    audio, conditioning_sample_rate, self.config.mimi.sample_rate, 1
+                )
+
+            with display_execution_time("Encoding audio prompt"):
+                prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device))
+                # import safetensors.torch
+                # safetensors.torch.save_file({"audio_prompt": prompt}, "embeddings/eponine.safetensors")
+
+        model_state = init_states(self.flow_lm, batch_size=1, sequence_length=1000)
+
+        with display_execution_time("Prompting audio"):
+            self._run_flow_lm_and_increment_step(model_state=model_state, audio_conditioning=prompt)
+
+        return model_state
+
+
+def prepare_text_prompt(text: str) -> tuple[str, int]:
+    text = text.strip()
+    if text == "":
+        raise ValueError("Text prompt cannot be empty")
+    number_of_words = len(text.split())
+    if number_of_words <= 4:
+        frames_after_eos_guess = 3
+    else:
+        frames_after_eos_guess = 1
+
+    # Make sure it starts with an uppercase letter
+    if not text[0].isupper():
+        text = text[0].upper() + text[1:]
+
+    # Let's make sure it ends with some kind of punctuation
+    # If it ends with a letter or digit, we add a period.
+    if text[-1].isalnum():
+        text = text + "."
+
+    # The model does not perform well when there are very few tokens, so
+    # we can add empty spaces at the beginning to increase the token count.
+    if len(text.split()) < 5:
+        text = " " * 8 + text
+
+    return text, frames_after_eos_guess

pocket_tts/modules/__init__.py

@@ -0,0 +1 @@
+"""Modules used for building the models."""

pocket_tts/modules/conv.py

@@ -0,0 +1,161 @@
+import math
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from pocket_tts.modules.stateful_module import StatefulModule
+
+
+def get_extra_padding_for_conv1d(
+    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
+) -> int:
+    """See `pad_for_conv1d`."""
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
+    """Pad for a convolution to make sure that the last window is full.
+    Extra padding is added at the end. This is required to ensure that we can rebuild
+    an output of the same length, as otherwise, even with padding, some time steps
+    might get removed.
+    For instance, with total padding = 4, kernel size = 4, stride = 2:
+        0 0 1 2 3 4 5 0 0   # (0s are padding)
+        1   2   3           # (output frames of a convolution, last 0 is never used)
+        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)
+            1 2 3 4         # once you removed padding, we are missing one time step !
+    """
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+class StreamingConv1d(StatefulModule):
+    """Conv1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        pad_mode: str = "constant",
+    ):
+        super().__init__()
+        assert pad_mode in ["constant", "replicate"], pad_mode
+        self.pad_mode = pad_mode
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn(
+                "StreamingConv1d has been initialized with stride > 1 and dilation > 1"
+                f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
+            )
+        self.conv = nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+    @property
+    def _stride(self) -> int:
+        return self.conv.stride[0]
+
+    @property
+    def _kernel_size(self) -> int:
+        return self.conv.kernel_size[0]
+
+    @property
+    def _effective_kernel_size(self) -> int:
+        dilation = self.conv.dilation[0]
+        return (self._kernel_size - 1) * dilation + 1  # effective kernel size with dilations
+
+    def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
+        stride = self._stride
+        # Effective kernel size accounting for dilation.
+        kernel = self._effective_kernel_size
+        previous = torch.zeros(batch_size, self.conv.in_channels, kernel - stride)
+        first = torch.ones(batch_size, dtype=torch.bool)
+        return dict(previous=previous, first=first)
+
+    def forward(self, x, model_state: dict | None):
+        B, C, T = x.shape
+        S = self._stride
+        assert T > 0 and T % S == 0, "Steps must be multiple of stride"
+        if model_state is None:
+            state = self.init_state(B, 0)
+        else:
+            state = self.get_state(model_state)
+        TP = state["previous"].shape[-1]
+        if TP and self.pad_mode == "replicate":
+            assert T >= TP, "Not enough content to pad streaming."
+            init = x[..., :1]
+            state["previous"][:] = torch.where(
+                state["first"].view(-1, 1, 1), init, state["previous"]
+            )
+        if TP:
+            x = torch.cat([state["previous"], x], dim=-1)
+        y = self.conv(x)
+        if TP:
+            state["previous"][:] = x[..., -TP:]
+            if self.pad_mode == "replicate":
+                state["first"] = torch.zeros_like(state["first"])
+        return y
+
+
+class StreamingConvTranspose1d(StatefulModule):
+    """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+    ):
+        super().__init__()
+        self.convtr = nn.ConvTranspose1d(
+            in_channels, out_channels, kernel_size, stride, groups=groups, bias=bias
+        )
+
+    @property
+    def _stride(self) -> int:
+        return self.convtr.stride[0]
+
+    @property
+    def _kernel_size(self) -> int:
+        return self.convtr.kernel_size[0]
+
+    def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
+        K = self._kernel_size
+        S = self._stride
+        return dict(partial=torch.zeros(batch_size, self.convtr.out_channels, K - S))
+
+    def forward(self, x, mimi_state: dict):
+        layer_state = self.get_state(mimi_state)["partial"]
+        y = self.convtr(x)
+        PT = layer_state.shape[-1]
+        if PT > 0:
+            y[..., :PT] += layer_state
+            bias = self.convtr.bias
+            for_partial = y[..., -PT:]
+            if bias is not None:
+                for_partial -= bias[:, None]
+            layer_state[:] = for_partial
+            y = y[..., :-PT]
+        return y

pocket_tts/modules/dummy_quantizer.py

@@ -0,0 +1,18 @@
+import torch
+from torch import nn
+
+
+class DummyQuantizer(nn.Module):
+    """Simplified quantizer that only provides output projection for TTS.
+
+    This removes all unnecessary quantization logic since we don't use actual quantization.
+    """
+
+    def __init__(self, dimension: int, output_dimension: int):
+        super().__init__()
+        self.dimension = dimension
+        self.output_dimension = output_dimension
+        self.output_proj = torch.nn.Conv1d(self.dimension, self.output_dimension, 1, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.output_proj(x)

pocket_tts/modules/layer_scale.py

@@ -0,0 +1,11 @@
+import torch
+import torch.nn as nn
+
+
+class LayerScale(nn.Module):
+    def __init__(self, channels: int, init: float):
+        super().__init__()
+        self.scale = nn.Parameter(torch.full((channels,), init))
+
+    def forward(self, x: torch.Tensor):
+        return self.scale * x

pocket_tts/modules/mimi_transformer.py

@@ -0,0 +1,285 @@
+from typing import NamedTuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+from typing_extensions import Self
+
+from pocket_tts.modules.layer_scale import LayerScale
+from pocket_tts.modules.rope import RotaryEmbedding
+from pocket_tts.modules.stateful_module import StatefulModule
+from pocket_tts.modules.transformer import StreamingMultiheadAttention
+from pocket_tts.utils.config import FlowLMTransformerConfig
+
+
+class KVCacheResult(NamedTuple):
+    keys: torch.Tensor
+    values: torch.Tensor
+    positions: torch.Tensor
+
+    @staticmethod
+    def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult":
+        B, H, T, D = keys.shape
+        assert tuple(values.shape[:-1]) == (B, H, T)
+        positions = torch.arange(T, device=keys.device, dtype=torch.long)
+        return KVCacheResult(keys, values, positions.expand(B, -1))
+
+
+def complete(
+    cache: torch.Tensor, end_offset: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+) -> KVCacheResult:
+    capacity = cache.shape[3]
+    assert k.shape[:-1] == v.shape[:-1], (k.shape, v.shape)
+    B, H, T, D = k.shape
+    assert T > 0
+    indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype)
+    indexes = indexes + end_offset.view(-1, 1)
+    indexes = indexes % capacity
+    # indexes is [B, T]
+    # k is [B, H, T, D]
+    # cache is [B, H, T', D]
+    this_indexes = indexes.view(B, 1, T, 1)
+    this_indexes = this_indexes.expand(-1, H, T, D)
+    cache[0].scatter_(2, this_indexes, k)
+    cache[1].scatter_(2, this_indexes, v)
+
+    keys = cache[0]
+    values = cache[1]
+
+    indexes = torch.arange(capacity, device=end_offset.device, dtype=torch.long)
+
+    # end_index correspond to the actual index where the last value was written.
+    last_offset = end_offset.view(-1, 1) + T - 1
+    end_index = last_offset % capacity
+    delta = indexes - end_index
+
+    positions = torch.where(delta <= 0, last_offset + delta, last_offset + delta - capacity)
+    end_offset[:] = end_offset + T
+    invalid = indexes >= end_offset.view(-1, 1)
+    positions = torch.where(invalid, torch.full_like(positions, -1), positions)
+
+    return KVCacheResult(keys, values, positions)
+
+
+class MimiStreamingMultiheadAttention(StatefulModule):
+    def __init__(self, embed_dim: int, num_heads: int, context: int, rope: RotaryEmbedding):
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.context = context
+        self.rope = rope
+        self.num_heads = num_heads
+        out_dim = 3 * embed_dim
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.in_proj = nn.Linear(embed_dim, out_dim, bias=False)
+
+    def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
+        dim_per_head = self.embed_dim // self.num_heads
+
+        state = {}
+        state["offset"] = torch.zeros(batch_size, dtype=torch.long)
+        state["cache"] = torch.zeros((2, batch_size, self.num_heads, sequence_length, dim_per_head))
+        state["end_offset"] = torch.zeros(batch_size, dtype=torch.long)
+        return state
+
+    def increment_step(self, state, increment: int = 1):
+        state["offset"] += increment
+
+    def _complete_kv(self, k, v, model_state: dict | None) -> KVCacheResult:
+        if model_state is None:
+            return KVCacheResult.from_kv(k, v)
+        else:
+            layer_state = self.get_state(model_state)
+            return complete(layer_state["cache"], layer_state["end_offset"], k, v)
+
+    def forward(self, query: torch.Tensor, model_state: dict | None) -> torch.Tensor:
+        B, T = query.shape[:2]
+
+        if model_state is None:
+            offset = torch.zeros(B, device=query.device, dtype=torch.long)
+        else:
+            offset = self.get_state(model_state)["offset"]
+
+        projected = self.in_proj(query)
+
+        q, k, v = rearrange(projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads)
+
+        # Permute from [b, h, t, d] to [b, t, h, d] for rope
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        q, k = self.rope(q, k, offset)
+        # Permute back from [b, t, h, d] to [b, h, t, d]
+        q = q.permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+
+        k, v, pos_k = self._complete_kv(k, v, model_state)
+        pos_k = pos_k[:, None]
+        pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(
+            -1, 1
+        )
+        delta = pos_q - pos_k
+        attn_bias = (pos_k >= 0) & (delta >= 0)
+        attn_bias = attn_bias & (delta < self.context)
+        attn_bias = attn_bias[:, None]
+
+        x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+
+        x = rearrange(x, "b h t d -> b t (h d)")
+        x = self.out_proj(x)
+        return x
+
+
+class StreamingTransformerLayer(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_feedforward: int,
+        context: int | None,
+        rope: RotaryEmbedding,
+        layer_scale: float | None = None,
+        attention_kind: str = "mimi",
+    ):
+        super().__init__()
+        # Redefine self_attn to our streaming multi-head attention
+        if attention_kind == "mimi":
+            # TODO: we should actually use StreamingMultiheadAttention here and add context window
+            # support. And we should then delete MimiStreamingMultiheadAttention.
+            # The implementation is really close.
+            self.self_attn = MimiStreamingMultiheadAttention(
+                context=context, rope=rope, embed_dim=d_model, num_heads=num_heads
+            )
+        else:
+            self.self_attn = StreamingMultiheadAttention(
+                rope=rope, embed_dim=d_model, num_heads=num_heads
+            )
+        self.norm1 = nn.LayerNorm(d_model, eps=1e-5)
+        self.norm2 = nn.LayerNorm(d_model, eps=1e-5)
+
+        self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False)
+        self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False)
+
+        if layer_scale is None:
+            self.layer_scale_1 = nn.Identity()
+            self.layer_scale_2 = nn.Identity()
+        else:
+            self.layer_scale_1 = LayerScale(d_model, layer_scale)
+            self.layer_scale_2 = LayerScale(d_model, layer_scale)
+
+    def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
+        x_orig = x
+        x = self.norm2(x)
+        update = self.linear2(F.gelu(self.linear1(x)))
+        return x_orig.to(update) + self.layer_scale_2(update)
+
+    def _sa_block(self, x: torch.Tensor, model_state: dict | None) -> torch.Tensor:
+        x_orig = x
+        x = self.norm1(x)
+        update = self.self_attn(x, model_state)
+        return x_orig.to(update) + self.layer_scale_1(update)
+
+    def forward(self, x: torch.Tensor, model_state: dict | None) -> torch.Tensor:
+        x = self._sa_block(x, model_state)
+        x = self._ff_block(x)
+        return x
+
+
+class StreamingTransformer(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_layers: int,
+        layer_scale: float | None = None,
+        dim_feedforward: int | list[int] = 2048,
+        context: int | None = None,
+        max_period: float = 10_000.0,
+        kind: str = "mimi",
+    ):
+        super().__init__()
+        assert d_model % num_heads == 0
+        self.max_period = max_period
+
+        self.rope = RotaryEmbedding(max_period=max_period)
+
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            self.layers.append(
+                StreamingTransformerLayer(
+                    d_model=d_model,
+                    num_heads=num_heads,
+                    dim_feedforward=dim_feedforward,
+                    context=context,
+                    rope=self.rope,
+                    layer_scale=layer_scale,
+                    attention_kind=kind,
+                )
+            )
+
+    @classmethod
+    def from_pydantic_config(cls, config: FlowLMTransformerConfig) -> Self:
+        dim_feedforward = int(config.d_model * config.hidden_scale)
+        return cls(
+            d_model=config.d_model,
+            num_heads=config.num_heads,
+            num_layers=config.num_layers,
+            dim_feedforward=dim_feedforward,
+            max_period=float(config.max_period),
+            kind="flow_lm",
+        )
+
+    def forward(self, x: torch.Tensor, model_state: dict | None):
+        for layer in self.layers:
+            x = layer(x, model_state)
+        return x
+
+
+class ProjectedTransformer(nn.Module):
+    def __init__(
+        self,
+        input_dimension: int,
+        output_dimensions: tuple[int, ...],
+        d_model: int,
+        num_heads: int,
+        num_layers: int,
+        layer_scale: float,
+        context: int,
+        max_period: float,
+        dim_feedforward: int,
+    ):
+        super().__init__()
+        self.transformer = StreamingTransformer(
+            d_model=d_model,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            layer_scale=layer_scale,
+            context=context,
+            max_period=max_period,
+            dim_feedforward=dim_feedforward,
+        )
+        self.input_dimension = input_dimension
+        self.output_dimensions = output_dimensions
+        self.input_proj = None
+        if d_model != input_dimension:
+            self.input_proj = nn.Linear(input_dimension, d_model, bias=False)
+
+        self.output_projs = nn.ModuleList()
+        for output_dimension in output_dimensions:
+            if d_model == output_dimension:
+                self.output_projs.append(nn.Identity())
+            else:
+                self.output_projs.append(nn.Linear(d_model, output_dimension, bias=False))
+
+    def forward(self, x, model_state: dict | None):
+        x = x.transpose(1, 2)
+        if self.input_proj is not None:
+            x = self.input_proj(x)
+        z = self.transformer(x, model_state)
+        ys = []
+        for output_proj in self.output_projs:
+            y = output_proj(z)
+            y = y.transpose(1, 2)
+            ys.append(y)
+        return ys

pocket_tts/modules/mlp.py

@@ -0,0 +1,215 @@
+"""
+Taken from
+https://github.com/LTH14/mar/blob/fe470ac24afbee924668d8c5c83e9fec60af3a73/models/diffloss.py
+
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+from typing_extensions import Self
+
+from pocket_tts.utils.config import FlowLMConfig
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale) + shift
+
+
+def _rms_norm(x: torch.Tensor, alpha: torch.Tensor, eps: float):
+    assert x.dim() >= alpha.dim()
+    x_dtype = x.dtype
+    var = eps + x.var(dim=-1, keepdim=True)
+    y = (x * (alpha.to(var) * torch.rsqrt(var))).to(x_dtype)
+    return y
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        alpha_shape = (dim,)
+        self.alpha = nn.Parameter(torch.full(alpha_shape, 1.0, requires_grad=True))
+
+    def forward(self, x: torch.Tensor):
+        return _rms_norm(x, self.alpha, self.eps)
+
+
+class LayerNorm(nn.Module):
+    """Reimplementation of LayerNorm because the default one doesn't support jvp."""
+
+    def __init__(self, channels, eps=1e-6, elementwise_affine=True):
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(channels))
+            self.bias = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        mean = x.mean(dim=-1, keepdim=True)
+        var = x.var(dim=-1, unbiased=False, keepdim=True)
+        x = (x - mean) / torch.sqrt(var + self.eps)
+        if hasattr(self, "weight"):
+            x = x * self.weight + self.bias
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    """Embeds scalar timesteps into vector representations."""
+
+    def __init__(
+        self, hidden_size: int, frequency_embedding_size: int = 256, max_period: int = 10000
+    ):
+        super().__init__()
+        blocks = [
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        ]
+        blocks.append(RMSNorm(hidden_size))
+        self.mlp = nn.Sequential(*blocks)
+        self.frequency_embedding_size = frequency_embedding_size
+        half = frequency_embedding_size // 2
+        self.register_buffer(
+            "freqs", torch.exp(-math.log(max_period) * torch.arange(start=0, end=half) / half)
+        )
+
+    def forward(self, t):
+        args = t * self.freqs.to(t.dtype)
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        assert not (self.frequency_embedding_size % 2)
+        t_emb = self.mlp(embedding)
+        return t_emb
+
+
+class ResBlock(nn.Module):
+    """
+    A residual block that can optionally change the number of channels.
+    :param channels: the number of input channels.
+    """
+
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+
+        self.in_ln = LayerNorm(channels, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(channels, channels, bias=True),
+        )
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True)
+        )
+
+    def forward(self, x, y):
+        shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
+        h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
+        h = self.mlp(h)
+        return x + gate_mlp * h
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer adopted from DiT.
+    """
+
+    def __init__(self, model_channels, out_channels):
+        super().__init__()
+        self.norm_final = LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(model_channels, out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class SimpleMLPAdaLN(nn.Module):
+    """Taken from https://arxiv.org/abs/2406.11838.
+
+    The MLP for Diffusion Loss.
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param cond_channels: channels in the condition.
+    :param num_res_blocks: number of residual blocks per downsample.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        out_channels,
+        cond_channels,
+        num_res_blocks,
+        num_time_conds=1,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.num_time_conds = num_time_conds
+
+        assert num_time_conds != 1
+        self.time_embed = nn.ModuleList(
+            [TimestepEmbedder(model_channels) for _ in range(num_time_conds)]
+        )
+        self.cond_embed = nn.Linear(cond_channels, model_channels)
+
+        self.input_proj = nn.Linear(in_channels, model_channels)
+
+        res_blocks = []
+        for i in range(num_res_blocks):
+            res_blocks.append(ResBlock(model_channels))
+
+        self.res_blocks = nn.ModuleList(res_blocks)
+        self.final_layer = FinalLayer(model_channels, out_channels)
+
+    @classmethod
+    def from_pydantic_config(cls, cfg: FlowLMConfig, latent_dim: int, cond_dim: int) -> Self:
+        config = cfg.flow
+
+        flow_dim = config.dim
+        flow_depth = config.depth
+        num_time_conds = 2
+        return SimpleMLPAdaLN(
+            latent_dim, flow_dim, latent_dim, cond_dim, flow_depth, num_time_conds=num_time_conds
+        )
+
+    def forward(
+        self, c: torch.Tensor, s: torch.Tensor, t: torch.Tensor, x: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Apply the model to an input batch.
+        :param c: conditioning from AR transformer.
+        :param s: start time tensor.
+        :param t: target time tensor.
+        :param x: an [N x C] Tensor of inputs.
+        :return: an [N x C] Tensor of outputs.
+        """
+        # Combine time conditions
+        ts = [s, t]
+        x = self.input_proj(x)
+        assert len(ts) == self.num_time_conds, (
+            f"Expected {self.num_time_conds} time conditions, got {len(ts)}"
+        )
+        assert self.num_time_conds != 1
+        t_combined = (
+            sum(self.time_embed[i](ts[i]) for i in range(self.num_time_conds)) / self.num_time_conds
+        )
+        c = self.cond_embed(c)
+        y = t_combined + c
+
+        for block in self.res_blocks:
+            x = block(x, y)
+
+        return self.final_layer(x, y)

pocket_tts/modules/resample.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn
+
+from pocket_tts.modules.conv import StreamingConv1d, StreamingConvTranspose1d
+
+
+class ConvDownsample1d(nn.Module):
+    """
+    Downsampling by some integer amount `stride` using convolutions
+    with a kernel size of twice the stride.
+    """
+
+    def __init__(self, stride: int, dimension: int):
+        super().__init__()
+        self.conv = StreamingConv1d(
+            dimension,
+            dimension,
+            kernel_size=2 * stride,
+            stride=stride,
+            groups=1,
+            bias=False,
+            pad_mode="replicate",
+        )
+
+    def forward(self, x: torch.Tensor, model_state: dict | None):
+        return self.conv(x, model_state)
+
+
+class ConvTrUpsample1d(nn.Module):
+    """
+    Upsample by some integer amount `stride` using transposed convolutions.
+    """
+
+    def __init__(self, stride: int, dimension: int):
+        super().__init__()
+        self.convtr = StreamingConvTranspose1d(
+            dimension,
+            dimension,
+            kernel_size=2 * stride,
+            stride=stride,
+            groups=dimension,
+            bias=False,
+        )
+
+    def forward(self, x: torch.Tensor, model_state: dict | None):
+        return self.convtr(x, model_state)

pocket_tts/modules/rope.py

@@ -0,0 +1,74 @@
+import math
+
+import torch
+from torch import nn
+
+
+def apply_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    offset: int | torch.Tensor = 0,
+    max_period: int | float = 10_000,
+):
+    """
+    Args:
+        q (torch.Tensor): Queries, shape `[B, T, H, D]`.
+        k (torch.Tensor): Keys, shape `[B, T, H, D]`.
+        offset (int): Current offset, e.g. when streaming.
+        max_period (float): Maximum period for the cos and sin.
+    """
+
+    B, T, H, D = q.shape
+    Bk, Tk, Hk, Dk = k.shape
+    assert (B, T, D) == (Bk, Tk, Dk)
+    assert D > 0
+    assert D % 2 == 0
+    assert max_period > 0
+
+    ds = torch.arange(D // 2, device=q.device, dtype=torch.float32)
+    freqs = torch.exp(ds * (-math.log(max_period) * 2 / D))
+
+    # could be optimized in one call
+    ts = torch.arange(T, device=q.device, dtype=torch.float32)
+    ts += offset
+    ts = ts.view(-1, 1, 1)
+
+    q = q.view(B, T, H, D // 2, 2)
+    k = k.view(B, T, Hk, D // 2, 2)
+
+    # convention is `r` suffix is real part, `i` is imaginary.
+    qr = q[..., 0].float()
+    qi = q[..., 1].float()
+
+    kr = k[..., 0].float()
+    ki = k[..., 1].float()
+
+    rotr = torch.cos(freqs * ts)
+    roti = torch.sin(freqs * ts)
+    qor = qr * rotr - qi * roti
+    qoi = qr * roti + qi * rotr
+
+    kor = kr * rotr - ki * roti
+    koi = kr * roti + ki * rotr
+
+    dtype = q.dtype
+    qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1)
+    ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1)
+
+    return qo.view(B, T, H, D), ko.view(B, T, Hk, D)
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864).
+
+    Args:
+        max_period (float): Maximum period of the rotation frequencies.
+    """
+
+    def __init__(self, max_period: float | int = 10000.0):
+        super().__init__()
+        self.max_period = max_period
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, offset: torch.Tensor | int):
+        """Apply rope rotation to query or key tensor."""
+        return apply_rope(q, k, offset, self.max_period)

pocket_tts/modules/seanet.py

@@ -0,0 +1,180 @@
+import numpy as np
+import torch.nn as nn
+
+from .conv import StreamingConv1d, StreamingConvTranspose1d
+
+
+class SEANetResnetBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        kernel_sizes: list[int] = [3, 1],
+        dilations: list[int] = [1, 1],
+        pad_mode: str = "reflect",
+        compress: int = 2,
+    ):
+        super().__init__()
+        assert len(kernel_sizes) == len(dilations), (
+            "Number of kernel sizes should match number of dilations"
+        )
+        hidden = dim // compress
+        block = nn.ModuleList([])
+        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            block += [
+                nn.ELU(alpha=1.0),
+                StreamingConv1d(
+                    in_chs, out_chs, kernel_size=kernel_size, dilation=dilation, pad_mode=pad_mode
+                ),
+            ]
+        self.block = block
+
+    def forward(self, x, model_state: dict | None):
+        v = x
+        for layer in self.block:
+            if isinstance(layer, StreamingConv1d):
+                v = layer(v, model_state)
+            else:
+                v = layer(v)
+        assert x.shape == v.shape, (x.shape, v.shape, x.shape)
+        return x + v
+
+
+class SEANetEncoder(nn.Module):
+    def __init__(
+        self,
+        channels: int = 1,
+        dimension: int = 128,
+        n_filters: int = 32,
+        n_residual_layers: int = 3,
+        ratios: list[int] = [8, 5, 4, 2],
+        kernel_size: int = 7,
+        last_kernel_size: int = 7,
+        residual_kernel_size: int = 3,
+        dilation_base: int = 2,
+        pad_mode: str = "reflect",
+        compress: int = 2,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.dimension = dimension
+        self.n_filters = n_filters
+        self.ratios = list(reversed(ratios))
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = int(np.prod(self.ratios))
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+
+        mult = 1
+        model = nn.ModuleList(
+            [StreamingConv1d(channels, mult * n_filters, kernel_size, pad_mode=pad_mode)]
+        )
+        # Downsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(
+                        mult * n_filters,
+                        kernel_sizes=[residual_kernel_size, 1],
+                        dilations=[dilation_base**j, 1],
+                        pad_mode=pad_mode,
+                        compress=compress,
+                    )
+                ]
+
+            # Add downsampling layers
+            model += [
+                nn.ELU(alpha=1.0),
+                StreamingConv1d(
+                    mult * n_filters,
+                    mult * n_filters * 2,
+                    kernel_size=ratio * 2,
+                    stride=ratio,
+                    pad_mode=pad_mode,
+                ),
+            ]
+            mult *= 2
+
+        model += [
+            nn.ELU(alpha=1.0),
+            StreamingConv1d(mult * n_filters, dimension, last_kernel_size, pad_mode=pad_mode),
+        ]
+
+        self.model = model
+
+    def forward(self, x, model_state: dict | None):
+        for layer in self.model:
+            if isinstance(layer, (StreamingConv1d, SEANetResnetBlock)):
+                x = layer(x, model_state)
+            else:
+                x = layer(x)
+        return x
+
+
+class SEANetDecoder(nn.Module):
+    def __init__(
+        self,
+        channels: int = 1,
+        dimension: int = 128,
+        n_filters: int = 32,
+        n_residual_layers: int = 3,
+        ratios: list[int] = [8, 5, 4, 2],
+        kernel_size: int = 7,
+        last_kernel_size: int = 7,
+        residual_kernel_size: int = 3,
+        dilation_base: int = 2,
+        pad_mode: str = "reflect",
+        compress: int = 2,
+    ):
+        super().__init__()
+        self.dimension = dimension
+        self.channels = channels
+        self.n_filters = n_filters
+        self.ratios = ratios
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = int(np.prod(self.ratios))
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        mult = int(2 ** len(self.ratios))
+        model = nn.ModuleList(
+            [StreamingConv1d(dimension, mult * n_filters, kernel_size, pad_mode=pad_mode)]
+        )
+        # Upsample to raw audio scale
+        for _, ratio in enumerate(self.ratios):
+            # Add upsampling layers
+            model += [
+                nn.ELU(alpha=1.0),
+                StreamingConvTranspose1d(
+                    mult * n_filters, mult * n_filters // 2, kernel_size=ratio * 2, stride=ratio
+                ),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(
+                        mult * n_filters // 2,
+                        kernel_sizes=[residual_kernel_size, 1],
+                        dilations=[dilation_base**j, 1],
+                        pad_mode=pad_mode,
+                        compress=compress,
+                    )
+                ]
+
+            mult //= 2
+
+        # Add final layers
+        model += [
+            nn.ELU(alpha=1.0),
+            StreamingConv1d(n_filters, channels, last_kernel_size, pad_mode=pad_mode),
+        ]
+        self.model = model
+
+    def forward(self, z, model_state: dict | None):
+        for layer in self.model:
+            if isinstance(layer, (StreamingConvTranspose1d, SEANetResnetBlock, StreamingConv1d)):
+                z = layer(z, model_state)
+            else:
+                z = layer(z)
+        return z

pocket_tts/modules/stateful_module.py

@@ -0,0 +1,45 @@
+from abc import ABC, abstractmethod
+
+import torch
+from torch import nn
+
+
+def init_states(
+    model: nn.Module, batch_size: int, sequence_length: int
+) -> dict[str, dict[str, torch.Tensor]]:
+    result = {}
+    for module_name, module in model.named_modules():
+        if not isinstance(module, StatefulModule):
+            continue
+        module._module_absolute_name = module_name
+        module_state = module.init_state(batch_size, sequence_length=sequence_length)
+        result[module_name] = module_state
+    return result
+
+
+def increment_steps(
+    module: nn.Module, model_state: dict[str, dict[str, torch.Tensor]], increment: int = 1
+):
+    # print("incrementing steps by", increment)
+    for module_name, module in module.named_modules():
+        if not isinstance(module, StatefulModule):
+            continue
+        module.increment_step(model_state[module_name], increment)
+
+
+class StatefulModule(ABC, nn.Module):
+    def __init__(self, *args, **kwds):
+        self._module_absolute_name = None
+        return super().__init__(*args, **kwds)
+
+    @abstractmethod
+    def init_state(self, batch_size: int, sequence_length: int):
+        """Initialize the state."""
+        raise NotImplementedError
+
+    def increment_step(self, state: dict, increment: int = 1):
+        pass
+
+    def get_state(self, model_state: dict[str, dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+        """Get the state for this module from the model state."""
+        return model_state[self._module_absolute_name]

pocket_tts/modules/transformer.py

@@ -0,0 +1,126 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from pocket_tts.modules.rope import RotaryEmbedding
+from pocket_tts.modules.stateful_module import StatefulModule
+
+
+def complete_kv(
+    cache: torch.Tensor, current_end: torch.Tensor, k: torch.Tensor, v: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    current_end = current_end.shape[0]
+
+    cache[0, :, current_end : current_end + k.shape[1]] = k
+    cache[1, :, current_end : current_end + v.shape[1]] = v
+    valid = cache[:, :, : current_end + k.shape[1]]
+    return valid[0], valid[1]
+
+
+def _materialize_causal_mask(
+    shape: tuple[int, ...], shift: int, device: str | torch.device = "cpu"
+) -> torch.Tensor:
+    dtype = torch.float32
+
+    num_queries, num_keys = shape[-2:]
+    shift = num_keys - num_queries
+
+    tensor = torch.full(shape, dtype=dtype, fill_value=1, device=device)
+    mask = torch.tril(tensor, diagonal=shift).to(dtype)
+    mask = torch.log(mask)
+    return mask.to(dtype)
+
+
+class StreamingMultiheadAttention(StatefulModule):
+    """Similar to `nn.MultiheadAttention` but with support for streaming.
+
+    Args:
+        embed_dim (int): Dimension to project to.
+        num_heads (int): Number of heads.
+        context (int, optional): Number of time steps the attention can access to.
+            Can access `context` time steps into the past.
+        rope (`RotaryEmbedding`, optional): Rope embedding to use.
+        device (torch.device, optional): Device on which to initialize.
+        dtype (torch.dtype, optional): dtype to use.
+    """
+
+    def __init__(self, embed_dim: int, num_heads: int, rope: RotaryEmbedding):
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.rope = rope
+        self.num_heads = num_heads
+
+        out_dim = embed_dim
+        num_kv = num_heads
+        kv_dim = (embed_dim // num_heads) * num_kv
+        out_dim += 2 * kv_dim
+        mult = 1
+        in_proj = nn.Linear(embed_dim, mult * out_dim, bias=False)
+        # We try to follow the default PyTorch MHA convention, to easily compare results.
+        self.in_proj_weight = in_proj.weight
+        self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False)
+
+    def _get_mask(self, shape: tuple[int, int], shift: int, device: torch.device) -> torch.Tensor:
+        return _materialize_causal_mask(shape, shift=shift, device=device)
+
+    def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
+        dim_per_head = self.embed_dim // self.num_heads
+        initial_current_end = torch.zeros((0,)).to(self.in_proj_weight.device)
+        return dict(
+            current_end=initial_current_end,
+            cache=torch.full(
+                (2, batch_size, sequence_length, self.num_heads, dim_per_head),
+                float("NaN"),
+                device=self.in_proj_weight.device,
+                dtype=self.in_proj_weight.dtype,
+            ),
+        )
+
+    def increment_step(self, state: dict, increment: int = 1):
+        new_size = state["current_end"].shape[0] + increment
+        state["current_end"] = torch.zeros((new_size,)).to(state["current_end"].device)
+
+    def _complete_kv(self, k, v, state: dict | None):
+        k, v = complete_kv(state["cache"], state["current_end"], k, v)
+        return k, v
+
+    def _apply_rope(self, query: torch.Tensor, key: torch.Tensor, state: dict | None):
+        # Apply rope embeddings to query and key tensors.
+        streaming_offset = self._streaming_offset(state)
+        return self.rope(query, key, offset=streaming_offset)
+
+    def _streaming_offset(self, state: dict | None) -> torch.Tensor | int:
+        return state["current_end"].shape[0]
+
+    def check_model_state(self, model_state: dict):
+        if model_state is None:
+            raise ValueError("model_state must be provided")
+        return self.get_state(model_state)
+
+    def forward(self, query: torch.Tensor, model_state: dict | None):
+        state = self.check_model_state(model_state)
+
+        projected = nn.functional.linear(query, self.in_proj_weight)
+        # Reshape from (b, t, p*h*d) to (b, t, p, h, d) where p=3, h=num_heads
+        b, t, _ = projected.shape
+        d = self.embed_dim // self.num_heads
+        packed = projected.view(b, t, 3, self.num_heads, d)
+        q, k, v = torch.unbind(packed, dim=2)
+        q, k = self._apply_rope(q, k, state)
+        k, v = self._complete_kv(k, v, state)
+
+        mask_shape = (query.shape[1], query.shape[1] + state["current_end"].shape[0])
+        shift = state["current_end"].shape[0]
+
+        attn_mask = self._get_mask(mask_shape, shift=shift, device=q.device)
+
+        q, k, v = [x.transpose(1, 2) for x in (q, k, v)]
+        x = F.scaled_dot_product_attention(q, k, v, attn_mask)
+        x = x.transpose(1, 2)
+        # Reshape from (b, t, h, d) to (b, t, h*d)
+        b, t, h, d = x.shape
+        x = x.reshape(b, t, h * d)
+        x = self.out_proj(x)
+
+        return x

pocket_tts/static/index.html

@@ -0,0 +1,374 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Pocket TTS - Streaming</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <style>
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+        }
+        .spinner {
+            animation: spin 1s linear infinite;
+        }
+        @keyframes spin {
+            from { transform: rotate(0deg); }
+            to { transform: rotate(360deg); }
+        }
+    </style>
+</head>
+<body class="bg-gray-50 min-h-screen">
+    <div class="max-w-xl mx-auto p-4 space-y-4">
+        <h1 class="text-xl font-bold text-center">Pocket TTS</h1>
+        
+        <div class="bg-white p-4 rounded-lg shadow">
+            <textarea
+                id="text-input"
+                class="w-full border rounded p-2"
+                placeholder="Enter your text here..."
+                rows="4"
+            >Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.</textarea>
+        </div>
+
+        <div class="bg-white p-4 rounded-lg shadow">
+            <label for="voice-url-input" class="block text-sm font-medium text-gray-700 mb-2">
+                Optional voice URL (leave empty to use default voice):
+            </label>
+            <input
+                type="text"
+                id="voice-url-input"
+                class="w-full border rounded p-2"
+                placeholder="hf://kyutai/tts-voices/alba-mackenna/casual.wav"
+                value="marius"
+            />
+            <p class="text-xs text-gray-500 mt-1">
+                Supports: http://, https://, or hf:// URLs.<br>
+                You can also use predefined voices:<br>
+                "marius", "javert", "jean", "fantine", "cosette", "eponine", "azelma".
+            </p>
+        </div>
+
+        <div class="bg-white p-4 rounded-lg shadow">
+            <label for="voice-wav-input" class="block text-sm font-medium text-gray-700 mb-2">
+                Or upload a WAV file for voice cloning:
+            </label>
+            <input
+                type="file"
+                id="voice-wav-input"
+                class="w-full border rounded p-2"
+                accept=".wav,audio/wav"
+            />
+            <p class="text-xs text-gray-500 mt-1">
+                Upload a WAV file to use as voice reference. Takes precedence over voice URL.
+            </p>
+        </div>
+
+        
+        <button
+            id="generate-btn"
+            class="w-full px-4 py-2 rounded text-white bg-blue-600 hover:bg-blue-700 disabled:bg-gray-400 disabled:cursor-not-allowed"
+        >
+            <span id="generate-text">Generate audio</span>
+        </button>
+
+        <div id="status" class="hidden">
+            <div class="flex items-center space-x-2">
+                <div id="status-spinner" class="spinner w-4 h-4 border-2 border-blue-600 border-t-transparent rounded-full"></div>
+                <span id="status-text"></span>
+            </div>
+        </div>
+
+        <div id="audio-section" class="hidden">
+            <div class="bg-white p-4 rounded-lg shadow">
+                <h3 class="text-lg font-semibold mb-2">Generated Audio</h3>
+                <audio id="audio" controls class="w-full"></audio>
+                <div class="mt-2">
+                    <a id="download-link" class="text-blue-600 text-sm hover:underline">Download</a>
+                </div>
+            </div>
+        </div>
+    </div>
+
+    <script>
+        class StreamingWavPlayer {
+            constructor() {
+                this.audioContext = new (window.AudioContext || window.webkitAudioContext)();
+                this.sampleRate = 0;
+                this.numChannels = 0;
+                this.headerParsed = false;
+                this.headerBuffer = new Uint8Array(44);
+                this.headerBytesReceived = 0;
+                this.nextStartTime = 0;
+                this.isPlaying = false;
+                this.minBufferSize = 16384;
+                this.pcmData = new Uint8Array(0);
+            }
+
+            parseWavHeader(header) {
+                const view = new DataView(header.buffer);
+
+                const riff = String.fromCharCode.apply(null, Array.from(header.slice(0, 4)));
+                const wave = String.fromCharCode.apply(null, Array.from(header.slice(8, 12)));
+
+                if (riff !== 'RIFF' || wave !== 'WAVE') {
+                    throw new Error('Invalid WAV file');
+                }
+
+                this.numChannels = view.getUint16(22, true);
+                this.sampleRate = view.getUint32(24, true);
+                const bitsPerSample = view.getUint16(34, true);
+
+                console.log(`WAV Format: ${this.sampleRate}Hz, ${this.numChannels} channels, ${bitsPerSample} bits`);
+
+                this.headerParsed = true;
+            }
+
+            appendPcmData(newData) {
+                const newBuffer = new Uint8Array(this.pcmData.length + newData.length);
+                newBuffer.set(this.pcmData);
+                newBuffer.set(newData, this.pcmData.length);
+                this.pcmData = newBuffer;
+            }
+
+            async tryPlayBuffer() {
+                if (!this.headerParsed || this.pcmData.length < this.minBufferSize) {
+                    return;
+                }
+
+                const bytesPerSample = this.numChannels * 2;
+                const samplesToPlay = Math.floor(this.pcmData.length / bytesPerSample);
+                const bytesToPlay = samplesToPlay * bytesPerSample;
+
+                if (bytesToPlay === 0) return;
+
+                const dataToPlay = this.pcmData.slice(0, bytesToPlay);
+                this.pcmData = this.pcmData.slice(bytesToPlay);
+
+                const audioBuffer = this.audioContext.createBuffer(
+                    this.numChannels,
+                    samplesToPlay,
+                    this.sampleRate
+                );
+
+                const int16Data = new Int16Array(dataToPlay.buffer, dataToPlay.byteOffset, samplesToPlay * this.numChannels);
+
+                for (let channel = 0; channel < this.numChannels; channel++) {
+                    const channelData = audioBuffer.getChannelData(channel);
+                    for (let i = 0; i < samplesToPlay; i++) {
+                        channelData[i] = int16Data[i * this.numChannels + channel] / 32768;
+                    }
+                }
+
+                const source = this.audioContext.createBufferSource();
+                source.buffer = audioBuffer;
+                source.connect(this.audioContext.destination);
+
+                const currentTime = this.audioContext.currentTime;
+                const startTime = Math.max(currentTime, this.nextStartTime);
+
+                source.start(startTime);
+
+                // Track first audio playback
+                if (!this.firstAudioPlayed && window.firstAudioCallback) {
+                    this.firstAudioPlayed = true;
+                    window.firstAudioCallback();
+                }
+
+                this.nextStartTime = startTime + audioBuffer.duration;
+                this.isPlaying = true;
+
+                if (this.pcmData.length >= this.minBufferSize) {
+                    setTimeout(() => this.tryPlayBuffer(), 10);
+                }
+            }
+
+            addChunk(chunk) {
+                if (!this.headerParsed) {
+                    const headerBytesNeeded = 44 - this.headerBytesReceived;
+                    const bytesToCopy = Math.min(headerBytesNeeded, chunk.length);
+
+                    this.headerBuffer.set(
+                        chunk.slice(0, bytesToCopy),
+                        this.headerBytesReceived
+                    );
+
+                    this.headerBytesReceived += bytesToCopy;
+
+                    if (this.headerBytesReceived >= 44) {
+                        this.parseWavHeader(this.headerBuffer);
+
+                        if (chunk.length > bytesToCopy) {
+                            this.appendPcmData(chunk.slice(bytesToCopy));
+                        }
+                    }
+                } else {
+                    this.appendPcmData(chunk);
+                }
+
+                this.tryPlayBuffer();
+            }
+
+            stop() {
+                this.audioContext.close();
+                this.isPlaying = false;
+            }
+        }
+
+        // Application state
+        let streamingPlayer = null;
+        let currentAudioBlob = null;
+
+        // DOM elements
+        const textInput = document.getElementById('text-input');
+        const generateBtn = document.getElementById('generate-btn');
+        const generateText = document.getElementById('generate-text');
+        const status = document.getElementById('status');
+        const statusText = document.getElementById('status-text');
+        const statusSpinner = document.getElementById('status-spinner');
+        const audioSection = document.getElementById('audio-section');
+        const audio = document.getElementById('audio');
+        const downloadLink = document.getElementById('download-link');
+
+        // Event listeners
+        generateBtn.addEventListener('click', generateAudio);
+
+        function showStatus(message, isLoading = false) {
+            statusText.textContent = message;
+            statusSpinner.style.display = isLoading ? 'block' : 'none';
+            status.classList.remove('hidden');
+        }
+
+        function hideStatus() {
+            status.classList.add('hidden');
+        }
+
+        async function generateAudio() {
+            const text = textInput.value.trim();
+
+            if (!text) {
+                showStatus('Please enter some text to generate speech.', false);
+                setTimeout(hideStatus, 3000);
+                return;
+            }
+
+            // Stop any currently playing audio
+            if (streamingPlayer) {
+                streamingPlayer.stop();
+                streamingPlayer = null;
+            }
+
+            // Track timing
+            const startTime = performance.now();
+            let firstAudioTime = null;
+
+            // Set callback for first audio
+            window.firstAudioCallback = () => {
+                if (!firstAudioTime) {
+                    firstAudioTime = performance.now();
+                    const timeToFirstAudio = ((firstAudioTime - startTime) / 1000).toFixed(2);
+                    showStatus(`First audio in ${timeToFirstAudio}s...`, true);
+                }
+            };
+
+            // Update UI
+            generateBtn.disabled = true;
+            generateText.textContent = 'Generating...';
+            showStatus('Generating speech...', true);
+            audioSection.classList.add('hidden');
+
+            try {
+                const formData = new FormData();
+                formData.append('text', text);
+
+                // Add voice URL if provided
+                const voiceUrl = document.getElementById('voice-url-input').value.trim();
+                if (voiceUrl) {
+                    formData.append('voice_url', voiceUrl);
+                }
+
+                // Add voice WAV file if uploaded
+                const voiceWavFile = document.getElementById('voice-wav-input').files[0];
+                if (voiceWavFile) {
+                    formData.append('voice_wav', voiceWavFile);
+                }
+
+                const response = await fetch('/tts', {
+                    method: 'POST',
+                    body: formData
+                });
+
+                if (!response.ok) {
+                    throw new Error(`Server error: ${response.status}`);
+                }
+
+                // Clone the response for both streaming and blob collection
+                const responseForPlayback = response.clone();
+                const responseForHistory = response.clone();
+
+                // Start streaming playback
+                const reader = responseForPlayback.body.getReader();
+                streamingPlayer = new StreamingWavPlayer();
+
+                const processStream = async () => {
+                    try {
+                        while (true) {
+                            const { done, value } = await reader.read();
+                            if (done) break;
+
+                            if (value) {
+                                streamingPlayer.addChunk(value);
+                            }
+                        }
+                    } catch (e) {
+                        console.error('Error processing stream:', e);
+                    }
+                };
+
+                // Start processing the stream
+                processStream();
+
+                // Also collect the full blob for download/audio element
+                const blob = await responseForHistory.blob();
+                const totalTime = ((performance.now() - startTime) / 1000).toFixed(2);
+                currentAudioBlob = blob;
+
+                // Update UI with audio
+                const audioUrl = URL.createObjectURL(blob);
+                audio.src = audioUrl;
+
+                // Wait for audio metadata to load to get duration
+                audio.addEventListener('loadedmetadata', () => {
+                    const audioDuration = audio.duration;
+                    const speedRatio = (audioDuration / parseFloat(totalTime)).toFixed(1);
+
+                    const timeToFirst = firstAudioTime ? ((firstAudioTime - startTime) / 1000).toFixed(2) : 'N/A';
+
+                    showStatus(
+                        `✨ First audio: ${timeToFirst}s | Total: ${totalTime}s | ${speedRatio}x faster than real-time`,
+                        false
+                    );
+                });
+                downloadLink.href = audioUrl;
+                downloadLink.download = `tts-audio.wav`;
+                
+                audioSection.classList.remove('hidden');
+
+            } catch (error) {
+                console.error('Error generating audio:', error);
+                showStatus(`Error: ${error.message}`, false);
+                setTimeout(hideStatus, 3000);
+            } finally {
+                // Re-enable button
+                generateBtn.disabled = false;
+                generateText.textContent = 'Generate Audio';
+            }
+        }
+
+        // Focus on text input when page loads
+        window.addEventListener('load', () => {
+            textInput.focus();
+        });
+    </script>
+</body>
+</html>

pocket_tts/utils/__init__.py

@@ -0,0 +1 @@
+"""Utilities."""

pocket_tts/utils/config.py

@@ -0,0 +1,122 @@
+"""Configuration models for loading YAML config files."""
+
+from pathlib import Path
+
+import yaml
+from pydantic import BaseModel, ConfigDict
+
+
+class StrictModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+# Flow configuration
+class FlowConfig(StrictModel):
+    dim: int
+    depth: int
+
+
+# Transformer configuration for FlowLM
+class FlowLMTransformerConfig(StrictModel):
+    hidden_scale: int
+    max_period: int
+    d_model: int
+    num_heads: int
+    num_layers: int
+
+
+class LookupTable(StrictModel):
+    dim: int
+    n_bins: int
+    tokenizer: str
+    tokenizer_path: str
+
+
+# Root configuration
+class FlowLMConfig(StrictModel):
+    """Root configuration model for YAML config files."""
+
+    dtype: str
+
+    # Nested configurations
+    flow: FlowConfig
+    transformer: FlowLMTransformerConfig
+
+    # conditioning
+    lookup_table: LookupTable
+    weights_path: str | None = None
+
+
+# SEANet configuration
+class SEANetConfig(StrictModel):
+    dimension: int
+    channels: int
+    n_filters: int
+    n_residual_layers: int
+    ratios: list[int]
+    kernel_size: int
+    residual_kernel_size: int
+    last_kernel_size: int
+    dilation_base: int
+    pad_mode: str
+    compress: int
+
+
+# Transformer configuration for Mimi
+class MimiTransformerConfig(StrictModel):
+    d_model: int
+    input_dimension: int
+    output_dimensions: tuple[int, ...]
+    num_heads: int
+    num_layers: int
+    layer_scale: float
+    context: int
+    max_period: float = 10000.0
+    dim_feedforward: int
+
+
+# Quantizer configuration
+class QuantizerConfig(StrictModel):
+    dimension: int
+    output_dimension: int
+
+
+# Root configuration
+class MimiConfig(StrictModel):
+    """Root configuration model for Mimi YAML config files."""
+
+    dtype: str
+
+    # Sample rate and channels
+    sample_rate: int
+    channels: int
+    frame_rate: float
+
+    # SEANet configurations
+    seanet: SEANetConfig
+
+    # Transformer
+    transformer: MimiTransformerConfig
+
+    # Quantizer
+    quantizer: QuantizerConfig
+    weights_path: str | None = None
+
+
+class Config(StrictModel):
+    flow_lm: FlowLMConfig
+    mimi: MimiConfig
+    weights_path: str | None = None
+    weights_path_without_voice_cloning: str | None = None
+
+
+def load_config(yaml_path: str | Path) -> Config:
+    yaml_path = Path(yaml_path)
+
+    if not yaml_path.exists():
+        raise FileNotFoundError(f"Config file not found: {yaml_path}")
+
+    with open(yaml_path, "r") as f:
+        config_dict = yaml.safe_load(f)
+
+    return Config(**config_dict)

pocket_tts/utils/debugging.py

@@ -0,0 +1,26 @@
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+
+def to_str(obj):
+    if isinstance(obj, (torch.Tensor, torch.nn.Parameter)):
+        return f"T(s={list(obj.shape)})"
+    elif isinstance(obj, (list, tuple)):
+        return "[" + ", ".join(to_str(o) for o in obj) + "]"
+    elif isinstance(obj, dict):
+        return "{" + ", ".join(f"{to_str(k)}: {to_str(v)}" for k, v in obj.items()) + "}"
+    else:
+        return str(obj)
+
+
+class LoggingMode(TorchDispatchMode):
+    """Useful to check implementation differences."""
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        output = func(*args, **kwargs or {})
+        print(
+            f"Aten function called: {func}, args: "
+            f"{to_str(args)}, kwargs: {to_str(kwargs)} -> "
+            f"output: {to_str(output)}"
+        )
+        return output

pocket_tts/utils/logging_utils.py

@@ -0,0 +1,41 @@
+import logging
+from contextlib import contextmanager
+
+
+class PocketTTSFilter(logging.Filter):
+    def filter(self, record):
+        return record.name.startswith("pocket_tts")
+
+
+@contextmanager
+def enable_logging(library_name, level):
+    # Get the specific logger and its parent
+    logger = logging.getLogger(library_name)
+    parent_logger = logging.getLogger("pocket_tts")
+
+    # Store original configuration
+    old_level = logger.level
+    old_parent_level = parent_logger.level
+    old_handlers = parent_logger.handlers.copy()
+
+    # Configure logging format for pocket_tts logger
+    parent_logger.setLevel(level)
+
+    # Clear existing handlers and add our custom formatter with filter
+    parent_logger.handlers.clear()
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter("%(levelname)s: %(message)s")
+    handler.setFormatter(formatter)
+    handler.addFilter(PocketTTSFilter())
+    parent_logger.addHandler(handler)
+    parent_logger.propagate = False
+
+    try:
+        yield logger
+    finally:
+        # Restore original configuration
+        logger.setLevel(old_level)
+        parent_logger.setLevel(old_parent_level)
+        parent_logger.handlers.clear()
+        for h in old_handlers:
+            parent_logger.addHandler(h)

pocket_tts/utils/utils.py

@@ -0,0 +1,102 @@
+import hashlib
+import logging
+import time
+from pathlib import Path
+
+import requests
+import safetensors.torch
+import torch
+from huggingface_hub import hf_hub_download
+from torch import nn
+
+PROJECT_ROOT = Path(__file__).parent.parent.parent
+
+PREDEFINED_VOICES = {
+    "marius": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/marius.safetensors",
+    "javert": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/javert.safetensors",
+    "jean": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/jean.safetensors",
+    "fantine": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/fantine.safetensors",
+    "cosette": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/cosette.safetensors",
+    "eponine": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/eponine.safetensors",
+    "azelma": "hf://kyutai/pocket-tts-without-voice-cloning/embeddings/azelma.safetensors",
+}
+
+
+def make_cache_directory() -> Path:
+    cache_dir = Path.home() / ".cache" / "pocket_tts"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def print_nb_parameters(model: nn.Module, model_name: str):
+    logger = logging.getLogger(__name__)
+    state_dict = model.state_dict()
+    total = 0
+    for key, value in state_dict.items():
+        logger.info("%s: %,d", key, value.numel())
+        total += value.numel()
+    logger.info("Total number of parameters in %s: %,d", model_name, total)
+
+
+def size_of_dict(state_dict: dict) -> int:
+    total_size = 0
+    for value in state_dict.values():
+        if isinstance(value, torch.Tensor):
+            total_size += value.numel() * value.element_size()
+        elif isinstance(value, dict):
+            total_size += size_of_dict(value)
+    return total_size
+
+
+class display_execution_time:
+    def __init__(self, task_name: str, print_output: bool = True):
+        self.task_name = task_name
+        self.print_output = print_output
+        self.start_time = None
+        self.elapsed_time_ms = None
+        self.logger = logging.getLogger(__name__)
+
+    def __enter__(self):
+        self.start_time = time.monotonic()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        end_time = time.monotonic()
+        self.elapsed_time_ms = int((end_time - self.start_time) * 1000)
+        if self.print_output:
+            self.logger.info("%s took %d ms", self.task_name, self.elapsed_time_ms)
+        return False  # Don't suppress exceptions
+
+
+def download_if_necessary(file_path: str) -> Path:
+    if file_path.startswith("http://") or file_path.startswith("https://"):
+        cache_dir = make_cache_directory()
+        cached_file = cache_dir / (
+            hashlib.sha256(file_path.encode()).hexdigest() + "." + file_path.split(".")[-1]
+        )
+        if not cached_file.exists():
+            response = requests.get(file_path)
+            response.raise_for_status()
+            with open(cached_file, "wb") as f:
+                f.write(response.content)
+        return cached_file
+    elif file_path.startswith("hf://"):
+        file_path = file_path.removeprefix("hf://")
+        splitted = file_path.split("/")
+        repo_id = "/".join(splitted[:2])
+        filename = "/".join(splitted[2:])
+        cached_file = hf_hub_download(repo_id=repo_id, filename=filename)
+        return Path(cached_file)
+    else:
+        return Path(file_path)
+
+
+def load_predefined_voice(voice_name: str) -> torch.Tensor:
+    if voice_name not in PREDEFINED_VOICES:
+        raise ValueError(
+            f"Predefined voice '{voice_name}' not found"
+            f", available voices are {list(PREDEFINED_VOICES)}."
+        )
+    voice_file = download_if_necessary(PREDEFINED_VOICES[voice_name])
+    # There is only one tensor in the file.
+    return safetensors.torch.load_file(voice_file)["audio_prompt"]

pocket_tts/utils/weights_loading.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+
+import safetensors
+
+
+def get_flow_lm_state_dict(path: Path) -> dict:
+    state_dict = {}
+    with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            if (
+                key.startswith("flow.w_s_t.")
+                or key == "condition_provider.conditioners.transcript_in_segment.learnt_padding"
+                or key == "condition_provider.conditioners.speaker_wavs.learnt_padding"
+            ):
+                # skip lookup table weights
+                continue
+            new_name = key
+            if key == "condition_provider.conditioners.transcript_in_segment.embed.weight":
+                new_name = "conditioner.embed.weight"
+            if key == "condition_provider.conditioners.speaker_wavs.output_proj.weight":
+                new_name = "speaker_proj_weight"
+            state_dict[new_name] = f.get_tensor(key)
+    return state_dict
+
+
+def get_mimi_state_dict(path: Path) -> dict:
+    state_dict = {}
+    with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            if key.startswith("model.quantizer.vq.") or key == "model.quantizer.logvar_proj.weight":
+                # skip vq weights
+                continue
+
+            state_dict[key.removeprefix("model.")] = f.get_tensor(key)
+    return state_dict