pocket-tts 1.0.2__py3-none-any.whl

pocket_tts/models/flow_lm.py

@@ -0,0 +1,208 @@
+import logging
+from functools import partial
+
+import torch
+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