Copy the files in the new repo for public release

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