Gabriel de Marmiesse
46bb28f448
Without a specific commit, huggingface cannot be sure that the file has not changed, so it needs to do a quick query to hf hub to check that there were no new commits. This adds some latency to the CLI and function calls, but it also prevents offline use, even if the files are already present locally. This PRs fixes this. Now with the commit hash, hf has the guarantee that the local files are the right ones and thus doesn't need to query the hub every time. Fix https://github.com/kyutai-labs/pocket-tts/pull/21
209 lines
8.0 KiB
Python
209 lines
8.0 KiB
Python
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),
|
|
)
|