Hugging Face's logo

LiquidAI
/
LFM3-Audio-2B

Any-to-Any
Transformers
Safetensors
lfm3_audio
automatic-speech-recognition
liquid
lfm3
audio
speech
edge
conversational
multilingual
custom_code
Eval Results
Model card Files
xet
 Community
3
LFM3-Audio-2B / modeling_lfm3_audio.py
paulpak58's picture
paulpak58
orech's picture
orech
saniyakarwa's picture
saniyakarwa
haerski's picture
haerski
mlabonne's picture
mlabonne
CalderKat's picture
CalderKat
Paul Pak
LFM3-Audio-2B
9d49701
Raw History Blame Contribute Delete
91.2 kB
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import warnings
from collections import deque
from dataclasses import dataclass
from enum import Enum
import torch
import torch.nn.functional as F
from torch import nn
from transformers import initialization as init
from transformers.cache_utils import Cache, DynamicCache
from transformers.configuration_utils import PreTrainedConfig
from transformers.conversion_mapping import register_checkpoint_conversion_mapping
from transformers.core_model_loading import WeightRenaming
from transformers.generation import GenerationConfig, GenerationMixin
from transformers.generation.logits_process import (
MinPLogitsWarper,
RepetitionPenaltyLogitsProcessor,
TemperatureLogitsWarper,
TopKLogitsWarper,
TopPLogitsWarper,
)
from transformers.generation.streamers import BaseStreamer
from transformers.masking_utils import (
create_causal_mask,
create_recurrent_attention_mask,
create_sliding_window_causal_mask,
)
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.models.auto import AutoModel
from transformers.processing_utils import Unpack
from transformers.utils import (
ModelOutput,
TransformersKwargs,
auto_docstring,
can_return_tuple,
torch_compilable_check,
)
from transformers.utils.import_utils import is_tracing
from .configuration_lfm3_audio import Lfm3AudioConfig, Lfm3AudioEncoderConfig
from .modeling_lfm3 import Lfm3Model
class Lfm3AudioRMSNorm(nn.Module):
def __init__(self, hidden_size, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
# The heads were trained with a norm that scales in fp32 and rounds once, not weight * bf16(x) like the backbone.
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return (hidden_states * self.weight).to(input_dtype)
class Lfm3AudioDownsampler(nn.Module):
"""Three stride-2 convolutions over the (mel, time) plane: 128 mel bins and 8 frames become one vector."""
def __init__(self, config: Lfm3AudioEncoderConfig):
super().__init__()
channels = config.conv_channels
conv = {"kernel_size": 3, "stride": 2, "padding": 1}
self.conv_in = nn.Conv2d(1, channels, **conv)
self.conv_mid = nn.Conv2d(channels, channels, groups=channels, **conv)
self.lin_mid = nn.Conv2d(channels, channels, kernel_size=1)
self.conv_end = nn.Conv2d(channels, channels, groups=channels, **conv)
self.lin_end = nn.Conv2d(channels, channels, kernel_size=1)
self.lin_out = nn.Conv1d(config.num_mel_bins * channels // 8, config.hidden_size, kernel_size=1)
self.act = nn.SiLU()
def forward(self, input_features: torch.Tensor, feature_lengths: torch.Tensor | None = None) -> torch.Tensor:
def mask_padding(hidden_states, lengths):
if lengths is None:
return hidden_states
valid = torch.arange(hidden_states.shape[-1], device=hidden_states.device)[None] < lengths[:, None]
return hidden_states * valid[:, None, None].to(hidden_states.dtype)
hidden_states = mask_padding(input_features.unsqueeze(1), feature_lengths)
hidden_states = self.act(self.conv_in(hidden_states))
lengths = None if feature_lengths is None else (feature_lengths + 1) // 2
hidden_states = mask_padding(hidden_states, lengths)
hidden_states = self.act(self.lin_mid(self.conv_mid(hidden_states)))
lengths = None if lengths is None else (lengths + 1) // 2
hidden_states = mask_padding(hidden_states, lengths)
hidden_states = self.act(self.lin_end(self.conv_end(hidden_states)))
return self.lin_out(hidden_states.flatten(1, 2))
@staticmethod
def output_lengths(lengths: torch.Tensor) -> torch.Tensor:
return -(lengths // -8)
class Lfm3AudioEncoderStreamingCache:
"""Bounded downsampler context and causal encoder state for synchronized audio chunks."""
def __init__(self, config: Lfm3AudioEncoderConfig):
cache_config = copy.deepcopy(config.text_config)
cache_config.sliding_window = config.sliding_window
cache_config.layer_types = [
"sliding_attention" if kind == "full_attention" else kind for kind in cache_config.layer_types
]
self.past_key_values = DynamicCache(config=cache_config)
self.features: torch.Tensor | None = None
self.offset = 0
self.num_frames = 0
self.num_tokens = 0
self.finished = False
class Lfm3AudioEncoder(nn.Module):
def __init__(self, config: Lfm3AudioEncoderConfig):
super().__init__()
self.config = config
self.downsample = Lfm3AudioDownsampler(config)
self.lfm = AutoModel.from_config(config.text_config)
# the encoder only ever consumes spectrogram features
self.lfm.register_module("embed_tokens", None)
def forward(
self, input_features: torch.Tensor, feature_lengths: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
hidden_states = self.downsample(input_features, feature_lengths).transpose(1, 2)
# The mask bounds the window for mask-based backends; flash attention uses the kernel kwarg.
self.lfm.config.sliding_window = self.config.sliding_window
attention_mask = create_sliding_window_causal_mask(
config=self.lfm.config,
inputs_embeds=hidden_states,
attention_mask=None,
past_key_values=None,
position_ids=None,
)
outputs = self.lfm(
inputs_embeds=hidden_states,
attention_mask={"full_attention": attention_mask, "conv": None},
use_cache=False,
sliding_window=self.config.sliding_window,
)
return outputs.last_hidden_state, self.downsample.output_lengths(feature_lengths)
@torch.no_grad()
def stream(
self,
input_features: torch.Tensor,
cache: Lfm3AudioEncoderStreamingCache | None = None,
is_final: bool = False,
) -> tuple[torch.Tensor, Lfm3AudioEncoderStreamingCache]:
"""Encode complete stride-8 windows, with seven frames of lookahead and bounded left context."""
if self.training:
raise ValueError("Streaming audio encoding requires eval mode")
if cache is None:
cache = Lfm3AudioEncoderStreamingCache(self.config)
if cache.finished:
raise ValueError("This stream is finalized; start with a new cache")
features = input_features.to(self.downsample.conv_in.weight.dtype)
if features.ndim != 3 or features.shape[1] != self.config.num_mel_bins:
raise ValueError("Expected features of shape (batch_size, num_mel_bins, num_frames)")
if cache.features is not None:
features = torch.cat((cache.features, features), dim=-1)
num_frames = cache.num_frames + input_features.shape[-1]
num_tokens = (num_frames + (7 if is_final else 0)) // 8
count = num_tokens - cache.num_tokens
if count:
start = cache.num_tokens - cache.offset // 8
hidden_states = self.downsample(features).transpose(1, 2)[:, start : start + count]
positions = torch.arange(cache.num_tokens, num_tokens, device=features.device)[None]
self.lfm.config.sliding_window = self.config.sliding_window
mask = create_sliding_window_causal_mask(
config=self.lfm.config,
inputs_embeds=hidden_states,
attention_mask=None,
past_key_values=cache.past_key_values,
position_ids=positions,
layer_idx=self.lfm.config.layer_types.index("full_attention"),
)
try:
outputs = self.lfm(
inputs_embeds=hidden_states,
attention_mask={"full_attention": mask, "conv": None},
past_key_values=cache.past_key_values,
position_ids=positions,
use_cache=True,
sliding_window=self.config.sliding_window,
)
except Exception:
# the layers already run have advanced the key/value cache: the stream cannot be resumed
cache.finished = True
raise
encoded = outputs.last_hidden_state
else:
encoded = features.new_empty((features.shape[0], 0, self.config.hidden_size))
offset = max(0, (num_tokens - 1) * 8)
cache.features = features[..., offset - cache.offset :].clone() if not is_final else None
cache.num_frames, cache.offset, cache.num_tokens, cache.finished = num_frames, offset, num_tokens, is_final
return encoded, cache
class Lfm3AudioProjector(nn.Module):
def __init__(self, config: Lfm3AudioConfig):
super().__init__()
in_features = config.audio_encoder_config.hidden_size
out_features = config.text_config.hidden_size
self.layer_norm = nn.LayerNorm(in_features)
self.linear_1 = nn.Linear(in_features, out_features)
self.act = nn.GELU()
self.linear_2 = nn.Linear(out_features, out_features)
def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
return self.linear_2(self.act(self.linear_1(self.layer_norm(audio_features))))
class Lfm3AudioCodebookEmbedding(nn.Module):
"""Input embedding and tied output head of one codebook of the depth decoder."""
def __init__(self, config: Lfm3AudioConfig):
super().__init__()
self.embedding = nn.Embedding(config.codebook_size, config.depth_config.hidden_size)
self.embedding_norm = Lfm3AudioRMSNorm(config.depth_config.hidden_size, eps=config.depth_config.norm_eps)
def forward(self, codes: torch.Tensor) -> torch.Tensor:
return self.embedding(codes)
def logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
return F.linear(self.embedding_norm(hidden_states), self.embedding.weight)
class Lfm3AudioDepthDecoder(nn.Module):
"""
Predicts the codebooks of one speech frame from the language model's hidden state. Every codebook gets its own
slot of the projected hidden state plus either its known code embedding or a learned mask embedding; a
bidirectional LFM3 model over the slots then scores the masked codebooks.
"""
def __init__(self, config: Lfm3AudioConfig):
super().__init__()
self.config = config
depth_config = config.depth_config
self.num_codebooks = config.num_codebooks
self.hidden_size = depth_config.hidden_size
self.depth_linear = nn.Linear(config.text_config.hidden_size, self.num_codebooks * self.hidden_size)
self.mask_embedding = nn.Parameter(torch.empty(self.hidden_size))
self.model = Lfm3Model(depth_config)
self.model.register_module("embed_tokens", None)
self.depth_embeddings = nn.ModuleList(Lfm3AudioCodebookEmbedding(config) for _ in range(self.num_codebooks))
def embed_codes(self, codes: torch.Tensor) -> torch.Tensor:
return torch.stack(
[embedding(codes[..., index]) for index, embedding in enumerate(self.depth_embeddings)], dim=-2
)
def forward(self, hidden_states: torch.Tensor, codes: torch.Tensor, masked: torch.Tensor) -> torch.Tensor:
r"""
hidden_states (`torch.FloatTensor` of shape `(num_frames, hidden_size)`):
The language model's hidden state at the position preceding each frame.
codes (`torch.LongTensor` of shape `(num_frames, num_codebooks)`):
The codes of each frame.
masked (`torch.BoolTensor` of shape `(num_frames, num_codebooks)`):
The codebooks hidden from the decoder, which it is asked to predict.
Returns:
`torch.FloatTensor` of shape `(num_frames, num_codebooks, codebook_size)`: the logits of every codebook.
"""
depth_in = self.depth_linear(hidden_states).unflatten(-1, (self.num_codebooks, self.hidden_size))
code_embeds = self.embed_codes(codes)
depth_in = depth_in + torch.where(masked[..., None], self.mask_embedding.to(code_embeds.dtype), code_embeds)
depth_out = self.trunk(depth_in)
return torch.stack(
[embedding.logits(depth_out[:, index]) for index, embedding in enumerate(self.depth_embeddings)], dim=1
)
def trunk(self, hidden_states: torch.Tensor) -> torch.Tensor:
# all codebooks of a frame are predicted jointly, so the attention is bidirectional
return self.model(
inputs_embeds=hidden_states,
attention_mask={"full_attention": None},
is_causal=False,
use_cache=False,
).last_hidden_state
def create_masks_by_layer_type(
config: PreTrainedConfig,
inputs_embeds: torch.Tensor,
attention_mask: torch.Tensor | None,
past_key_values: Cache | None,
position_ids: torch.Tensor | None = None,
) -> dict:
"""The cache sizes each mask from the layer it is asked about, and the generic mask utilities ask about the first
non-sliding layer, which for a leading `conv` layer redirects to the first attention layer of any type: name the
first layer of each mask's own type instead."""
config = config.get_text_config()
mask_kwargs = {
"config": config,
"inputs_embeds": inputs_embeds,
"attention_mask": attention_mask,
"past_key_values": past_key_values,
"position_ids": position_ids,
}
layer_types, causal_mask_mapping = config.layer_types, {}
if "full_attention" in layer_types:
layer_idx = layer_types.index("full_attention")
causal_mask_mapping["full_attention"] = create_causal_mask(**mask_kwargs, layer_idx=layer_idx)
if "sliding_attention" in layer_types:
layer_idx = layer_types.index("sliding_attention")
causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
**mask_kwargs, layer_idx=layer_idx
)
if "conv" in layer_types:
causal_mask_mapping["conv"] = create_recurrent_attention_mask(**mask_kwargs)
return causal_mask_mapping
@auto_docstring
class Lfm3AudioPreTrainedModel(PreTrainedModel):
config: Lfm3AudioConfig
base_model_prefix = "model"
supports_gradient_checkpointing = True
_no_split_modules = ["Lfm3AudioDepthDecoder"]
_skip_keys_device_placement = ["past_key_values"]
_supports_flash_attn = True
_supports_sdpa = True
_supports_flex_attn = True
_can_compile_fullgraph = True
_supports_attention_backend = True
input_modalities = ("audio", "text")
# the training export stores the tied output head of every depth codebook embedding as well
_keys_to_ignore_on_load_unexpected = [r"depth_embeddings\.\d+\.to_logits\.weight"]
@torch.no_grad()
def _init_weights(self, module):
super()._init_weights(module)
if isinstance(module, Lfm3AudioDepthDecoder):
init.normal_(module.mask_embedding, std=0.02)
@staticmethod
def create_masks_for_generate(
config: PreTrainedConfig,
inputs_embeds: torch.Tensor,
attention_mask: torch.Tensor | None,
past_key_values: Cache | None,
position_ids: torch.Tensor | None = None,
**kwargs,
) -> dict:
return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
@auto_docstring(
custom_intro="""
Base class for Lfm3Audio outputs, with hidden states and attentions.
"""
)
@dataclass
class Lfm3AudioModelOutputWithPast(BaseModelOutputWithPast):
r"""
past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
`past_key_values` input) to speed up sequential decoding.
audio_hidden_states (`torch.FloatTensor`, *optional*):
A `torch.FloatTensor` of size `(num_audio_tokens, hidden_size)`.
audio_hidden_states of the model produced by the speech encoder and after projecting the last hidden state.
"""
audio_hidden_states: torch.FloatTensor | None = None
@auto_docstring(
custom_intro="""
Base class for Lfm3Audio causal language model (or autoregressive) outputs.
"""
)
@dataclass
class Lfm3AudioCausalLMOutputWithPast(CausalLMOutputWithPast):
r"""
audio_hidden_states (`torch.FloatTensor`, *optional*):
A `torch.FloatTensor` of size `(num_audio_tokens, hidden_size)`.
audio_hidden_states of the model produced by the speech encoder and after projecting the last hidden state.
text_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
The language modeling loss of the text tokens.
audio_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `audio_labels` is provided):
The depth decoder loss of the speech frames.
"""
audio_hidden_states: torch.FloatTensor | None = None
text_loss: torch.FloatTensor | None = None
audio_loss: torch.FloatTensor | None = None
@auto_docstring(
custom_intro="""
Outputs of [`Lfm3AudioForConditionalGeneration.generate`].
"""
)
@dataclass
class Lfm3AudioGenerateOutput(ModelOutput):
r"""
sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
The prompt followed by the generated tokens: text tokens as their ids, and one `audio_out_token_id` per
generated speech frame.
audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`):
The codes of the generated speech frames in generation order, without the end-of-audio frames unless
`return_audio_eos` was set. Rows are padded with -1; slice to `audio_lengths` before codec decoding.
sequence_lengths (`torch.LongTensor` of shape `(batch_size,)`):
The prompt width (including input padding) plus the number of returned generated slots for each row.
audio_lengths (`torch.LongTensor` of shape `(batch_size,)`):
The number of returned speech frames per row.
"""
sequences: torch.LongTensor = None
audio_codes: torch.LongTensor | None = None
sequence_lengths: torch.LongTensor | None = None
audio_lengths: torch.LongTensor | None = None
@auto_docstring(custom_intro="State returned by [`Lfm3AudioForConditionalGeneration.prefill_audio`].")
@dataclass
class Lfm3AudioStreamingOutput(ModelOutput):
r"""
past_key_values (`Cache`, *optional*):
Language-model cache including all completed audio tokens. Updated in place by subsequent calls.
encoder_cache (`Lfm3AudioEncoderStreamingCache`, *optional*):
Streaming encoder state for the current utterance. Start a new encoder cache for each utterance.
attention_mask (`torch.LongTensor` of shape `(batch_size, cached_sequence_length)`, *optional*):
Full language-model mask, including the text prefix and completed audio tokens.
last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_new_tokens, hidden_size)`, *optional*):
Language-model states for this chunk, or `None` when no audio token was completed.
"""
past_key_values: Cache | None = None
encoder_cache: Lfm3AudioEncoderStreamingCache | None = None
attention_mask: torch.LongTensor | None = None
last_hidden_state: torch.FloatTensor | None = None
class Lfm3AudioLiveTranscriptionState:
"""State of one live transcription: the encoder follows the whole stream, the language model one context."""
def __init__(self):
self.encoder_cache: Lfm3AudioEncoderStreamingCache | None = None
self.past_key_values: Cache | None = None
self.attention_mask: torch.LongTensor | None = None
# fed before the next chunk's audio: the prompt of a new context, or the `<|text_end|>` closing the last block
self.pending: torch.FloatTensor | None = None
self.finished = False
@auto_docstring(custom_intro="Output of [`Lfm3AudioForConditionalGeneration.transcribe_live`].")
@dataclass
class Lfm3AudioLiveTranscriptionOutput(ModelOutput):
r"""
sequences (`torch.LongTensor` of shape `(1, num_generated)`):
The text tokens of this chunk, without the `<|text_end|>` closing the block or the final stop token.
num_audio_tokens (`int`):
Audio tokens the chunk completed.
closed (`bool`):
Whether the model ended the block itself, with `<|text_end|>` or a stop token on the final chunk, instead of
exhausting `max_new_tokens`.
state ([`Lfm3AudioLiveTranscriptionState`]):
The stream state for the next call.
"""
sequences: torch.LongTensor | None = None
num_audio_tokens: int | None = None
closed: bool | None = None
state: Lfm3AudioLiveTranscriptionState | None = None
class Lfm3AudioGenerationMode(str, Enum):
TEXT = "text"
AUDIO = "audio"
INTERLEAVED = "interleaved"
@auto_docstring
class Lfm3AudioModel(Lfm3AudioPreTrainedModel):
def __init__(self, config: Lfm3AudioConfig):
super().__init__(config)
self.language_model = Lfm3Model(config.text_config)
self.audio_encoder = Lfm3AudioEncoder(config.audio_encoder_config)
self.audio_projector = Lfm3AudioProjector(config)
self.codebook_embedding = nn.Embedding(
config.num_codebooks * config.codebook_size, config.text_config.hidden_size
)
self.depth_decoder = Lfm3AudioDepthDecoder(config)
self.post_init()
def get_input_embeddings(self):
return self.language_model.get_input_embeddings()
def set_input_embeddings(self, value):
self.language_model.set_input_embeddings(value)
@auto_docstring
def get_audio_features(
self, input_features: torch.FloatTensor, feature_attention_mask: torch.Tensor
) -> torch.FloatTensor:
r"""
input_features (`torch.FloatTensor` of shape `(num_clips, num_mel_bins, num_frames)`):
The log-mel spectrograms of the input clips, right-padded to the longest clip.
feature_attention_mask (`torch.Tensor` of shape `(num_clips, num_frames)`):
The mask of the valid spectrogram frames of each clip.
Returns:
`torch.FloatTensor` of shape `(num_audio_tokens, hidden_size)`: the projected speech features of all clips
in order, one per eight spectrogram frames.
"""
# the streaming paths refuse these shapes too; offline, a few missing mel bins would encode silently
num_mel_bins = self.config.audio_encoder_config.num_mel_bins
if input_features.ndim != 3 or input_features.shape[1] != num_mel_bins:
raise ValueError(
f"Expected input_features of shape (num_clips, {num_mel_bins}, num_frames), got "
f"{tuple(input_features.shape)}"
)
if feature_attention_mask.shape != (input_features.shape[0], input_features.shape[-1]):
raise ValueError(
f"Expected feature_attention_mask of shape {(input_features.shape[0], input_features.shape[-1])}, "
f"got {tuple(feature_attention_mask.shape)}"
)
feature_lengths = feature_attention_mask.sum(-1)
input_features = input_features.to(self.audio_encoder.downsample.conv_in.weight.dtype)
encoded, encoded_lengths = self.audio_encoder(input_features, feature_lengths)
valid = torch.arange(encoded.shape[1], device=encoded.device)[None, :] < encoded_lengths[:, None]
return self.audio_projector(encoded[valid])
@torch.no_grad()
def get_audio_features_stream(
self,
input_features: torch.FloatTensor,
cache: Lfm3AudioEncoderStreamingCache | None = None,
is_final: bool = False,
) -> tuple[torch.Tensor, Lfm3AudioEncoderStreamingCache]:
"""Project newly completed audio tokens, returning `(batch, tokens, hidden)` and the updated cache."""
encoded, cache = self.audio_encoder.stream(input_features, cache, is_final)
return self.audio_projector(encoded), cache
def get_audio_out_embeds(self, audio_out_codes: torch.LongTensor) -> torch.FloatTensor:
r"""
audio_out_codes (`torch.LongTensor` of shape `(num_frames, num_codebooks)`):
The codes of the speech frames in the prompt.
Returns:
`torch.FloatTensor` of shape `(num_frames, hidden_size)`: the sum of the codebook embeddings of each frame.
"""
num_codebooks, codebook_size = self.config.num_codebooks, self.config.codebook_size
if audio_out_codes.ndim != 2 or audio_out_codes.shape[-1] != num_codebooks:
raise ValueError(
f"Expected audio_out_codes of shape (num_frames, {num_codebooks}), got {tuple(audio_out_codes.shape)}"
)
# the codebooks share one table: an out-of-range code would silently read a neighbouring codebook. The check
# is data-dependent, so a traced or captured forward skips it (`generate` prefills eagerly)
if not is_tracing(audio_out_codes):
torch_compilable_check(
(audio_out_codes >= 0) & (audio_out_codes < codebook_size),
f"audio_out_codes must be in [0, {codebook_size})",
)
offsets = torch.arange(num_codebooks, device=audio_out_codes.device) * codebook_size
return self.codebook_embedding(audio_out_codes + offsets).sum(-2)
def get_placeholder_positions(self, input_ids, inputs_embeds, token_id: int) -> torch.Tensor:
if input_ids is not None:
return input_ids == token_id
token_embed = self.get_input_embeddings()(
torch.full((), token_id, dtype=torch.long, device=inputs_embeds.device)
)
return (inputs_embeds == token_embed).all(-1)
def check_speech_placeholders(
self, input_ids, inputs_embeds, input_features=None, audio_out_codes=None, live: bool = False
) -> None:
"""A speech placeholder with nothing behind it would be read as text and switch `generate` into speech mode.
The check is data-dependent, so a compiled, exported or CUDA-graph-captured forward skips it; `generate` and
`transcribe_live` run eagerly."""
if is_tracing(input_ids if input_ids is not None else inputs_embeds):
return
for field, supplied, name in (
("audio_token_id", input_features is not None, "input_features"),
("audio_out_token_id", audio_out_codes is not None, "audio_out_codes"),
):
if not supplied:
torch_compilable_check(
~self.get_placeholder_positions(input_ids, inputs_embeds, getattr(self.config, field)),
"Live transcription prompts cannot hold speech placeholders: the audio arrives as `input_features`"
if live
else f"{field} placeholders require `{name}`",
)
def get_placeholder_mask(self, input_ids, inputs_embeds, token_id: int, num_features: int) -> torch.Tensor:
mask = self.get_placeholder_positions(input_ids, inputs_embeds, token_id)
# data-dependent, so a traced or captured forward skips it, like `check_speech_placeholders`
if not is_tracing(mask):
torch_compilable_check(
mask.sum() == num_features,
lambda: f"Placeholder tokens and features do not match: {int(mask.sum())} tokens {token_id}, "
f"{num_features} features",
)
return mask.unsqueeze(-1).to(inputs_embeds.device)
@can_return_tuple
@auto_docstring
def forward(
self,
input_ids: torch.LongTensor | None = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | None = None,
inputs_embeds: torch.FloatTensor | None = None,
input_features: torch.FloatTensor | None = None,
feature_attention_mask: torch.Tensor | None = None,
audio_out_codes: torch.LongTensor | None = None,
use_cache: bool | None = None,
**kwargs: Unpack[TransformersKwargs],
) -> tuple | Lfm3AudioModelOutputWithPast:
r"""
input_features (`torch.FloatTensor` of shape `(num_clips, num_mel_bins, num_frames)`, *optional*):
The log-mel spectrograms of the input clips, right-padded to the longest clip.
feature_attention_mask (`torch.Tensor` of shape `(num_clips, num_frames)`, *optional*):
The mask of the valid spectrogram frames of each clip.
audio_out_codes (`torch.LongTensor` of shape `(num_frames, num_codebooks)`, *optional*):
The codes of the speech frames in the prompt, one row per `audio_out_token_id` placeholder. A speech
placeholder without its features or codes is refused, except in a compiled or captured forward, which
skips the check.
"""
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
if inputs_embeds is None:
inputs_embeds = self.get_input_embeddings()(input_ids)
self.check_speech_placeholders(input_ids, inputs_embeds, input_features, audio_out_codes)
audio_features = None
if input_features is not None:
if feature_attention_mask is None:
feature_attention_mask = torch.ones(
input_features.shape[0], input_features.shape[-1], dtype=torch.bool, device=input_features.device
)
audio_features = self.get_audio_features(input_features, feature_attention_mask).to(
inputs_embeds.device, inputs_embeds.dtype
)
audio_mask = self.get_placeholder_mask(
input_ids, inputs_embeds, self.config.audio_token_id, audio_features.shape[0]
)
inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features)
if audio_out_codes is not None:
audio_out_embeds = self.get_audio_out_embeds(audio_out_codes).to(inputs_embeds.device, inputs_embeds.dtype)
audio_out_mask = self.get_placeholder_mask(
input_ids, inputs_embeds, self.config.audio_out_token_id, audio_out_embeds.shape[0]
)
inputs_embeds = inputs_embeds.masked_scatter(audio_out_mask, audio_out_embeds)
outputs = self.language_model(
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
**kwargs,
)
return Lfm3AudioModelOutputWithPast(
last_hidden_state=outputs.last_hidden_state,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
audio_hidden_states=audio_features,
)
@auto_docstring(
custom_intro="""
The LFM3-Audio model: an LFM3 language model that reads speech through the Liquid speech encoder and writes
speech as codebook frames through a depth decoder, alongside text.
"""
)
class Lfm3AudioForConditionalGeneration(Lfm3AudioPreTrainedModel, GenerationMixin):
# Trainer counts text labels only; the joint loss also supervises speech frames.
accepts_loss_kwargs = False
_tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
_checkpoint_conversion_mapping = {}
def __init__(self, config: Lfm3AudioConfig):
super().__init__(config)
self.model = Lfm3AudioModel(config)
self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
self.post_init()
def get_input_embeddings(self):
return self.model.get_input_embeddings()
def set_input_embeddings(self, value):
self.model.set_input_embeddings(value)
def get_output_embeddings(self) -> nn.Module:
return self.lm_head
@auto_docstring
def get_audio_features(self, input_features: torch.FloatTensor, feature_attention_mask: torch.Tensor):
r"""
input_features (`torch.FloatTensor` of shape `(num_clips, num_mel_bins, num_frames)`):
The log-mel spectrograms of the input clips, right-padded to the longest clip.
feature_attention_mask (`torch.Tensor` of shape `(num_clips, num_frames)`):
The mask of the valid spectrogram frames of each clip.
"""
return self.model.get_audio_features(input_features, feature_attention_mask)
@property
def language_model(self):
return self.model.language_model
@property
def _placeholder_token_ids(self) -> list[int]:
# the placeholders stand for media features, never for sampled text: one in `sequences` would break the
# one-placeholder-per-frame contract
return [self.config.audio_token_id, self.config.audio_out_token_id]
@property
def depth_decoder(self):
return self.model.depth_decoder
@can_return_tuple
@auto_docstring
def forward(
self,
input_ids: torch.LongTensor | None = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | None = None,
inputs_embeds: torch.FloatTensor | None = None,
input_features: torch.FloatTensor | None = None,
feature_attention_mask: torch.Tensor | None = None,
audio_out_codes: torch.LongTensor | None = None,
labels: torch.LongTensor | None = None,
audio_labels: torch.LongTensor | None = None,
use_cache: bool | None = None,
logits_to_keep: int | torch.Tensor = 0,
**kwargs: Unpack[TransformersKwargs],
) -> tuple | Lfm3AudioCausalLMOutputWithPast:
r"""
input_features (`torch.FloatTensor` of shape `(num_clips, num_mel_bins, num_frames)`, *optional*):
The log-mel spectrograms of the input clips, right-padded to the longest clip.
feature_attention_mask (`torch.Tensor` of shape `(num_clips, num_frames)`, *optional*):
The mask of the valid spectrogram frames of each clip.
audio_out_codes (`torch.LongTensor` of shape `(num_frames, num_codebooks)`, *optional*):
The codes of the speech frames in the sequence, one row per `audio_out_token_id` placeholder. A speech
placeholder without its features or codes is refused, except in a compiled or captured forward, which
skips the check.
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
audio_labels (`torch.LongTensor` of shape `(num_frames, num_codebooks)`, *optional*):
The codes of the speech frames to learn, one row per `audio_out_token_id` placeholder; rows of `-100` are
ignored. The depth decoder predicts a random subset of each frame's codebooks from the others and the
hidden state preceding the frame.
"""
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
input_features=input_features,
feature_attention_mask=feature_attention_mask,
audio_out_codes=audio_out_codes,
use_cache=use_cache,
**kwargs,
)
hidden_states = outputs[0]
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])
loss = text_loss = audio_loss = None
if labels is not None:
text_loss = (
self.loss_function(
logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
)
if (labels[..., 1:] != -100).any()
else logits.sum() * 0.0
)
if audio_labels is not None:
if input_ids is None:
raise ValueError("`audio_labels` requires `input_ids` to locate the speech frames")
audio_loss = self.audio_loss(input_ids, hidden_states, audio_labels)
if text_loss is not None or audio_loss is not None:
loss = self.combined_loss(input_ids, labels, audio_labels, text_loss, audio_loss)
return Lfm3AudioCausalLMOutputWithPast(
loss=loss,
text_loss=text_loss,
audio_loss=audio_loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
audio_hidden_states=outputs.audio_hidden_states,
)
def audio_loss(
self, input_ids: torch.LongTensor, hidden_states: torch.FloatTensor, audio_labels: torch.LongTensor
) -> torch.FloatTensor:
config = self.config
# under `device_map` sharding the last hidden states sit on the last layer's device, not the inputs'
input_ids, audio_labels = input_ids.to(hidden_states.device), audio_labels.to(hidden_states.device)
placeholders = input_ids == config.audio_out_token_id
if audio_labels.shape != (int(placeholders.sum()), config.num_codebooks):
raise ValueError(
f"Expected audio_labels of shape (num_frames, {config.num_codebooks}), one row per "
f"`audio_out_token_id` placeholder, got {tuple(audio_labels.shape)}"
)
# rows of -100 are ignored; any other code indexes the codebook embeddings and the loss
ignored = (audio_labels == -100).all(-1, keepdim=True)
torch_compilable_check(
ignored | ((audio_labels >= 0) & (audio_labels < config.codebook_size)),
f"audio_labels must be in [0, {config.codebook_size}) or rows of -100",
)
# a frame is predicted from the position before it, so the first position of a row has no frame
preceding = placeholders.clone()
preceding[:, 0] = False
supervised = (audio_labels != -100).all(-1)[preceding[placeholders]]
latents = hidden_states[preceding.roll(-1, dims=1)][supervised]
codes = audio_labels[preceding[placeholders]][supervised]
if codes.shape[0] == 0:
return hidden_states.sum() * 0.0
ratio = torch.rand(codes.shape[0], 1, device=codes.device)
masked = torch.rand(codes.shape, device=codes.device) < ratio
masked |= ~masked.any(-1, keepdim=True)
masked |= (codes == config.audio_eos_code).all(-1, keepdim=True)
logits = self.model.depth_decoder(latents, codes, masked)
# under `device_map` sharding the depth decoder answers on its own device
codes, masked = codes.to(logits.device), masked.to(logits.device)
losses = F.cross_entropy(logits.transpose(1, 2).float(), codes, reduction="none")
return (losses * masked).sum() / masked.sum().clamp_min(1)
def combined_loss(self, input_ids, labels, audio_labels, text_loss, audio_loss) -> torch.Tensor:
"""Pooled mean over the supervised text tokens and speech frames, each modality's mean weighted by its
count and its configured weight, as in training."""
# under `device_map` sharding the labels and the two losses can sit on three devices
device = (text_loss if text_loss is not None else audio_loss).device
text_count = 0 if labels is None else (labels[..., 1:] != -100).sum().to(device)
audio_count = 0
if audio_labels is not None:
placeholders = input_ids.to(audio_labels.device) == self.config.audio_out_token_id
preceding = placeholders.clone()
preceding[:, 0] = False
audio_count = (audio_labels != -100).all(-1)[preceding[placeholders]].sum().to(device)
text_weight = self.config.text_loss_weight * text_count
audio_weight = self.config.audio_loss_weight * audio_count
return (
text_weight * (text_loss if text_loss is not None else 0.0)
+ audio_weight * (audio_loss.to(device) if audio_loss is not None else 0.0)
) / (text_weight + audio_weight + 1e-6)
@staticmethod
def _build_warpers(
temperature: float | None, top_k: int | None, top_p: float | None = None, min_p: float | None = None
) -> list:
"""The logits warpers of one sampling setting, empty for greedy decoding; building them validates the values."""
# an explicit zero temperature is greedy; without a temperature, top_k alone samples at unit temperature
if temperature == 0 or (temperature is None and not top_k):
return []
warpers = []
if temperature:
warpers.append(TemperatureLogitsWarper(float(temperature)))
if top_k:
warpers.append(TopKLogitsWarper(int(top_k)))
if top_p is not None and top_p < 1.0:
warpers.append(TopPLogitsWarper(float(top_p)))
if min_p:
warpers.append(MinPLogitsWarper(float(min_p)))
return warpers
@classmethod
def _sample(
cls,
logits: torch.Tensor,
temperature: float | None,
top_k: int | None,
repetition_penalty: float | None = None,
history: torch.Tensor | None = None,
suppress: list[int] = (),
top_p: float | None = None,
min_p: float | None = None,
warpers: list | None = None,
) -> torch.Tensor:
scores = logits.float()[None].clone()
if repetition_penalty is not None and repetition_penalty != 1.0 and history is not None:
penalty = RepetitionPenaltyLogitsProcessor(float(repetition_penalty))
scores = penalty(history.reshape(1, -1).to(scores.device), scores)
if suppress:
scores[:, list(suppress)] = -float("inf")
if warpers is None:
warpers = cls._build_warpers(temperature, top_k, top_p, min_p)
if not warpers:
return scores.argmax(-1)
for warper in warpers:
scores = warper(None, scores)
return torch.multinomial(scores.softmax(-1), 1)[0]
def sample_audio_frame(
self,
hidden_state: torch.Tensor,
temperature: float | None = None,
top_k: int | None = None,
repetition_penalty: float | None = None,
history: torch.Tensor | list[torch.Tensor | None] | None = None,
warpers: list | None = None,
) -> torch.LongTensor:
r"""
Samples the codes of one speech frame from a language model hidden state of shape `(hidden_size,)` or `(batch_size, hidden_size)`: the
codebooks are revealed in a random order over the passes of `config.maskgit_schedule`, each pass scoring the
codebooks still hidden from the depth decoder. `history` holds the previous frames, `(num_frames,
num_codebooks)`, for the repetition penalty, or a list of per-row histories for a batch. `warpers` are the
logits warpers of `temperature` and `top_k`, prebuilt by `_build_warpers`.
"""
single = hidden_state.ndim == 1
if single:
hidden_state = hidden_state[None]
history = [history]
batch_size = hidden_state.shape[0]
if history is None:
history = [None] * batch_size
num_codebooks = self.config.num_codebooks
codes = hidden_state.new_zeros((batch_size, num_codebooks), dtype=torch.long)
revealed = torch.zeros_like(codes, dtype=torch.bool)
order = torch.stack([torch.randperm(num_codebooks, device=hidden_state.device) for _ in range(batch_size)])
for pick in order.split(self.config.maskgit_schedule, dim=1):
# the depth decoder's forward keeps every pass on its own device under `device_map` sharding or offload
logits = self.model.depth_decoder(hidden_state, codes, ~revealed)
for row in range(batch_size):
for index in pick[row].tolist():
codebook_history = None if history[row] is None else history[row][:, index]
codes[row, index] = self._sample(
logits[row, index], temperature, top_k, repetition_penalty, codebook_history, warpers=warpers
).to(codes.device)
revealed.scatter_(1, pick, True)
return codes[0] if single else codes
@torch.no_grad()
def prefill_audio(
self,
input_features: torch.FloatTensor,
past_key_values: Cache | None = None,
encoder_cache: Lfm3AudioEncoderStreamingCache | None = None,
attention_mask: torch.Tensor | None = None,
is_final: bool = False,
) -> Lfm3AudioStreamingOutput:
"""
Encode a mel chunk and append completed audio tokens to the language-model cache.
Args:
input_features (`torch.FloatTensor` of shape `(batch_size, num_mel_bins, num_frames)`):
New mel frames, without padding. Batched streams must have synchronized chunk lengths.
past_key_values (`Cache`, *optional*):
Cache from a text-prefix forward pass or the preceding `prefill_audio` call. Updated in place.
encoder_cache (`Lfm3AudioEncoderStreamingCache`, *optional*):
Encoder state from the preceding call. Omit at the start of each new utterance.
attention_mask (`torch.Tensor` of shape `(batch_size, cached_sequence_length)`, *optional*):
Mask of the tokens already in `past_key_values`, required for a nonempty cache. Pass the returned
mask to the next call. Audio tokens are appended as valid positions.
is_final (`bool`, *optional*, defaults to `False`):
Flush the encoder's trailing token and finalize this utterance, including for an empty final chunk.
Returns:
[`Lfm3AudioStreamingOutput`]: Updated encoder state, language-model cache and full attention mask.
Chunks that complete no token leave the language-model cache unchanged. After finalization, pass the
closing audio/user delimiters and assistant prefix to `generate` with this cache and a mask extended
for those new tokens. Returned generation sequences include only that suffix and the generated tokens.
"""
if self.training:
raise ValueError("Streaming audio prefill requires eval mode")
if input_features.ndim != 3 or input_features.shape[1] != self.config.audio_encoder_config.num_mel_bins:
raise ValueError("Expected features of shape (batch_size, num_mel_bins, num_frames)")
batch_size = input_features.shape[0]
past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
if attention_mask is None:
if past_length:
raise ValueError("Pass attention_mask covering the tokens already in past_key_values")
attention_mask = torch.ones((batch_size, 0), dtype=torch.long, device=input_features.device)
if attention_mask.shape != (batch_size, past_length):
raise ValueError("attention_mask must cover exactly the tokens already in past_key_values")
not_binary = ((attention_mask != 0) & (attention_mask != 1)).any()
not_monotonic = ((attention_mask[:, 1:].long() - attention_mask[:, :-1].long()) < 0).any()
if bool(not_binary | not_monotonic):
raise ValueError("attention_mask must describe contiguous left-padded prompts")
if past_length and not bool(attention_mask[:, -1].all()):
raise ValueError("Each cached prompt must end in a valid token")
if past_key_values is None:
past_key_values = DynamicCache(config=self.config.text_config)
embeddings, encoder_cache = self.model.get_audio_features_stream(input_features, encoder_cache, is_final)
num_tokens = embeddings.shape[1]
hidden_states = None
if num_tokens:
attention_mask = torch.cat((attention_mask, attention_mask.new_ones((batch_size, num_tokens))), dim=-1)
position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -num_tokens:]
outputs = self.model.language_model(
inputs_embeds=embeddings,
past_key_values=past_key_values,
attention_mask=attention_mask,
position_ids=position_ids,
use_cache=True,
)
hidden_states, past_key_values = outputs.last_hidden_state, outputs.past_key_values
return Lfm3AudioStreamingOutput(
past_key_values=past_key_values,
encoder_cache=encoder_cache,
attention_mask=attention_mask,
last_hidden_state=hidden_states,
)
def _live_forward(self, state: Lfm3AudioLiveTranscriptionState, inputs_embeds: torch.FloatTensor) -> torch.Tensor:
past_length = state.attention_mask.shape[1]
state.attention_mask = torch.cat(
(state.attention_mask, state.attention_mask.new_ones((1, inputs_embeds.shape[1]))), dim=-1
)
position_ids = torch.arange(past_length, state.attention_mask.shape[1], device=inputs_embeds.device)[None]
outputs = self.model.language_model(
inputs_embeds=inputs_embeds,
past_key_values=state.past_key_values,
attention_mask=state.attention_mask,
position_ids=position_ids,
use_cache=True,
)
return outputs.last_hidden_state[:, -1]
@torch.no_grad()
def transcribe_live(
self,
input_features: torch.FloatTensor,
state: Lfm3AudioLiveTranscriptionState | None = None,
input_ids: torch.LongTensor | None = None,
is_final: bool = False,
max_new_tokens: int = 48,
suppress_token_ids: list[int] | None = None,
stop_token_ids: list[int] | None = None,
) -> Lfm3AudioLiveTranscriptionOutput:
"""
Transcribe one chunk of a live stream with a checkpoint trained for live ASR.
The audio goes inside the assistant turn in chunks of `config.interleaved_n_audio` encoder tokens, each
followed by the text it completes:
<prompt> [A x 12] text<|text_end|> [A x 12] text<|text_end|> ... [A x r] <|mixed_end|> text<|text_end|><|im_end|>
Args:
input_features (`torch.FloatTensor` of shape `(1, num_mel_bins, num_frames)`):
The chunk's new mel frames, normally `8 * config.interleaved_n_audio`. A chunk that completes no
audio token decodes no text.
state ([`Lfm3AudioLiveTranscriptionState`], *optional*):
The state returned by the previous call. Omit for a new stream.
input_ids (`torch.LongTensor` of shape `(1, prompt_length)`, *optional*):
The prompt of a new language-model context (the system turn and the assistant header), required
for a new stream. With a `state`, it starts the context over while the encoder keeps following the
stream, the way training samples of at most 30 seconds are chained.
is_final (`bool`, *optional*, defaults to `False`):
Flush the encoder and decode the closing block, which ends at `<|text_end|>` or a stop token.
max_new_tokens (`int`, *optional*, defaults to 48):
Text budget of the block; a block that exhausts it is still closed with `<|text_end|>`.
suppress_token_ids (`list[int]`, *optional*):
Tokens never valid before the final chunk, on top of `<|audio_start|>`, `<|mixed_start|>` and the
end-of-turn token: pass `<|mixed_end|>` and `<|endoftext|>`.
stop_token_ids (`list[int]`, *optional*):
Tokens ending the final block on top of the end-of-turn token: pass `<|endoftext|>` and `<|im_start|>`.
Returns:
[`Lfm3AudioLiveTranscriptionOutput`]: the chunk's text tokens (the final block starts with
`<|mixed_end|>`), the number of audio tokens it completed, and the updated state.
"""
if self.training:
raise ValueError("Live transcription requires eval mode")
num_mel_bins = self.config.audio_encoder_config.num_mel_bins
if input_features.ndim != 3 or input_features.shape[0] != 1 or input_features.shape[1] != num_mel_bins:
raise ValueError(
f"Live transcription takes one stream: features of shape (1, {num_mel_bins}, num_frames), got "
f"{tuple(input_features.shape)}"
)
if input_ids is not None:
if input_ids.ndim != 2 or input_ids.shape[0] != 1:
raise ValueError("Live transcription takes one prompt of shape (1, prompt_length)")
self.model.check_speech_placeholders(input_ids, None, live=True)
if state is None:
if input_ids is None:
raise ValueError("A new stream needs the prompt `input_ids`")
state = Lfm3AudioLiveTranscriptionState()
if state.finished:
raise ValueError("This stream is finalized; start with a new state")
eos_token_ids = self.config.text_config.eos_token_id
eos_token_ids = (
[] if eos_token_ids is None else [eos_token_ids] if isinstance(eos_token_ids, int) else list(eos_token_ids)
)
text_end_token_id = self.config.text_end_token_id
suppress = [
token_id
for token_id in (
self.config.audio_start_token_id,
self.config.mixed_start_token_id,
*eos_token_ids,
*(suppress_token_ids or []),
)
# `<|text_end|>` closes every block, so it is never suppressed even when it is listed as an end-of-turn id
if token_id is not None and token_id != text_end_token_id
]
# the ids index the logits; refused before the encoder cache advances
vocab_size = self.config.text_config.vocab_size
if any(not 0 <= token_id < vocab_size for token_id in suppress):
raise ValueError(f"Suppressed token ids must lie in [0, {vocab_size}), got {suppress}")
stop = {*eos_token_ids, *(stop_token_ids or [])}
# the encoder runs before a restart replaces the context, so a rejected chunk leaves the context intact
embeddings, state.encoder_cache = self.model.get_audio_features_stream(
input_features, state.encoder_cache, is_final
)
embed_tokens = self.get_input_embeddings()
# an offloaded table sits on `meta`: the device comes from the embedding output it is joined with
embeddings = embeddings.to(dtype=embed_tokens.weight.dtype)
if input_ids is not None:
state.past_key_values = DynamicCache(config=self.config.text_config)
state.attention_mask = torch.zeros((1, 0), dtype=torch.long, device=input_features.device)
state.pending = embed_tokens(input_ids)
num_audio_tokens = embeddings.shape[1]
generated, closed = [], False
if num_audio_tokens or is_final:
# the pending text goes in with the audio it precedes, so a chunk completing no token leaves it pending
inputs_embeds = (
embeddings
if state.pending is None
else torch.cat((state.pending, embeddings.to(state.pending.device)), dim=1)
)
logits = self.lm_head(self._live_forward(state, inputs_embeds))
state.pending = None
# the placeholders are masked on every chunk, the final one included
placeholder_ids = torch.tensor(self._placeholder_token_ids, device=logits.device)
suppress = torch.tensor(suppress, dtype=torch.long, device=logits.device)
for _ in range(max_new_tokens):
logits[:, placeholder_ids] = float("-inf")
if not is_final:
logits[:, suppress] = float("-inf")
next_token = logits.argmax(-1, keepdim=True)
token = int(next_token)
if token == text_end_token_id or (is_final and token in stop):
closed = True
break
generated.append(token)
logits = self.lm_head(self._live_forward(state, embed_tokens(next_token)))
if is_final:
state.finished = True
else:
# the block closes with `<|text_end|>`, fed together with the next chunk's audio
state.pending = embed_tokens(torch.tensor([[text_end_token_id]], device=input_features.device))
return Lfm3AudioLiveTranscriptionOutput(
sequences=torch.tensor([generated], dtype=torch.long, device=input_features.device),
num_audio_tokens=num_audio_tokens,
closed=closed,
state=state,
)
# `GenerationConfig` fields the text/speech mode machine does not implement, with their inert values
_generation_fields_unsupported = {
"min_length": 0,
"min_new_tokens": None,
"max_time": None,
"stop_strings": None,
"early_stopping": False,
"penalty_alpha": None,
"dola_layers": None,
"use_cache": True,
"cache_implementation": None,
"cache_config": None,
"return_legacy_cache": None,
"prefill_chunk_size": None,
"typical_p": 1.0,
"epsilon_cutoff": 0.0,
"eta_cutoff": 0.0,
"diversity_penalty": 0.0,
"length_penalty": 1.0,
"no_repeat_ngram_size": 0,
"bad_words_ids": None,
"force_words_ids": None,
"renormalize_logits": False,
"remove_invalid_values": False,
"top_h": None,
"encoder_repetition_penalty": 1.0,
"use_mtp": None,
"constraints": None,
"forced_bos_token_id": None,
"forced_eos_token_id": None,
"exponential_decay_length_penalty": None,
"suppress_tokens": None,
"begin_suppress_tokens": None,
"sequence_bias": None,
"token_healing": False,
"guidance_scale": None,
"watermarking_config": None,
"output_attentions": False,
"output_hidden_states": False,
"output_scores": False,
"output_logits": None,
"encoder_no_repeat_ngram_size": 0,
"prompt_lookup_num_tokens": None,
"assistant_early_exit": None,
"compile_config": None,
}
@torch.no_grad()
def generate(
self,
input_ids: torch.LongTensor,
input_features: torch.FloatTensor | None = None,
feature_attention_mask: torch.Tensor | None = None,
audio_out_codes: torch.LongTensor | None = None,
generation_config: GenerationConfig | None = None,
generation_mode: Lfm3AudioGenerationMode | str | list[str] | None = None,
audio_temperature: float | None = None,
audio_top_k: int | None = None,
audio_repetition_penalty: float | None = None,
audio_repetition_window: int | None = None,
text_repetition_window: int | None = None,
suppress_mixed: bool = False,
suppress_audio_out: bool = False,
return_audio_eos: bool = False,
streamer: BaseStreamer | list[BaseStreamer | None] | None = None,
audio_streamer: BaseStreamer | list[BaseStreamer | None] | None = None,
past_key_values: Cache | None = None,
**kwargs,
) -> Lfm3AudioGenerateOutput | torch.LongTensor:
r"""
Generates text and speech for a batch of left-padded prompts.
Text is the default output. `<|audio_start|>` switches to speech until the depth decoder emits an
end-of-audio frame, `<|mixed_start|>` to blocks of `config.interleaved_n_text` text tokens and
`config.interleaved_n_audio` speech frames until `<|text_end|>` closes the text and the remaining speech
ends with an end-of-audio frame. A prompt ending on `<|audio_start|>`, a speech frame or `<|mixed_start|>`
starts in that mode. Generation stops at `<|im_end|>` emitted in text mode or after `max_new_tokens`
generated slots, a speech frame counting as one slot.
Args:
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
The uncached prompt tokens, with one `audio_token_id` per encoded frame of each input clip and one
`audio_out_token_id` per speech frame. With `past_key_values`, supply only the new suffix; returned
sequences contain this suffix and generated tokens, excluding the cached prefix.
input_features (`torch.FloatTensor` of shape `(num_clips, num_mel_bins, num_frames)`, *optional*):
The log-mel spectrograms of the input clips.
feature_attention_mask (`torch.Tensor` of shape `(num_clips, num_frames)`, *optional*):
The mask of the valid spectrogram frames of each clip.
audio_out_codes (`torch.LongTensor` of shape `(num_frames, num_codebooks)`, *optional*):
The codes of the speech frames in the prompt.
generation_config ([`GenerationConfig`], *optional*):
`max_new_tokens`, `do_sample`, `temperature`, `top_k`, `top_p`, `min_p` and `repetition_penalty`
apply to the text tokens; the penalty covers the last `text_repetition_window` generated tokens.
Without `max_new_tokens`, an explicit `max_length` bounds the prompt plus the generated slots and
must exceed the prompt; the default `max_length` counts generated slots alone, as in
[`GenerationMixin`].
Fields the text/speech mode machine does not implement (beam search, contrastive search,
`stop_strings`, `output_*`, `use_cache=False`, cache implementations, ...) raise instead of being
ignored. The bare call returns [`Lfm3AudioGenerateOutput`]. `return_dict_in_generate` is read from
the call, then a passed configuration, then the model's `generation_config`: `False` returns the
sequences alone, and so does a passed configuration when none of them sets it, like
[`GenerationMixin`], as pipelines expect.
generation_mode (`str` or `list[str]`, *optional*):
`"text"`, `"audio"` or `"interleaved"`, or one mode per row; inferred per row when omitted. An
explicit `"text"` keeps the reply in text by forbidding `<|mixed_start|>` and `<|audio_start|>`; an
inferred text mode lets the model switch to speech.
audio_temperature (`float`, *optional*):
The sampling temperature of the speech codes; `0` picks the most likely code. This and the four
arguments below are read from `generation_config` when the call omits them, so a checkpoint's
`generation_config.json` can carry the speech sampling settings.
audio_top_k (`int`, *optional*):
The top-k filtering of the speech codes.
audio_repetition_penalty (`float`, *optional*):
The penalty on the codes each codebook emitted in the last `audio_repetition_window` frames.
audio_repetition_window (`int`, *optional*):
The number of past frames the audio repetition penalty covers; 64 when nothing sets it.
text_repetition_window (`int`, *optional*):
The number of past tokens the text repetition penalty covers; 32 when nothing sets it.
suppress_mixed (`bool`, *optional*, defaults to `False`):
Whether to forbid `<|mixed_start|>`.
suppress_audio_out (`bool`, *optional*, defaults to `False`):
Whether to forbid `<|audio_start|>`.
return_audio_eos (`bool`, *optional*, defaults to `False`):
Whether to keep the end-of-audio frames in `audio_codes` and their placeholders in `sequences`, so
the output can be fed back as a prompt; they must be dropped before decoding with the codec.
streamer ([`~generation.BaseStreamer`], *optional*):
Receives the prompt and then every generated text token. For batches, pass one streamer per row.
audio_streamer ([`~generation.BaseStreamer`], *optional*):
Receives every generated speech frame, `(num_codebooks,)`, as soon as it is sampled. For batches,
pass one streamer per row. An entry may be `None`.
past_key_values (`Cache`, *optional*):
Cache from a prefix forward pass or `prefill_audio`, updated in place to cover the prompt and every
generated slot, end-of-audio frames included (`return_audio_eos=True` returns them); a batch whose rows
finished at different steps holds masked slots and cannot be continued. Pass `attention_mask`
covering both the cached prefix and `input_ids`. Positions are computed from that full mask unless
supplied.
"""
# a consumer blocked on a streamer must see the end whichever statement raises, argument checks included
streams = list(
{
id(stream): stream
for value in (streamer, audio_streamer)
for stream in (value if isinstance(value, (list, tuple)) else (value,))
if stream is not None
}.values()
)
try:
output = self._generate(
input_ids,
input_features=input_features,
feature_attention_mask=feature_attention_mask,
audio_out_codes=audio_out_codes,
generation_config=generation_config,
generation_mode=generation_mode,
audio_temperature=audio_temperature,
audio_top_k=audio_top_k,
audio_repetition_penalty=audio_repetition_penalty,
audio_repetition_window=audio_repetition_window,
text_repetition_window=text_repetition_window,
suppress_mixed=suppress_mixed,
suppress_audio_out=suppress_audio_out,
return_audio_eos=return_audio_eos,
streamer=streamer,
audio_streamer=audio_streamer,
past_key_values=past_key_values,
**kwargs,
)
except BaseException:
self._end_streamers(streams, raise_errors=False)
raise
self._end_streamers(streams)
return output
@staticmethod
def _end_streamers(streams: list, raise_errors: bool = True) -> None:
# every consumer is released before an `end()` failure surfaces
errors = []
for stream in streams:
try:
stream.end()
except Exception as error:
errors.append(error)
if errors and raise_errors:
raise errors[0]
def _generate(
self,
input_ids: torch.LongTensor,
input_features: torch.FloatTensor | None = None,
feature_attention_mask: torch.Tensor | None = None,
audio_out_codes: torch.LongTensor | None = None,
generation_config: GenerationConfig | None = None,
generation_mode: Lfm3AudioGenerationMode | str | list[str] | None = None,
audio_temperature: float | None = None,
audio_top_k: int | None = None,
audio_repetition_penalty: float | None = None,
audio_repetition_window: int | None = None,
text_repetition_window: int | None = None,
suppress_mixed: bool = False,
suppress_audio_out: bool = False,
return_audio_eos: bool = False,
streamer: BaseStreamer | list[BaseStreamer | None] | None = None,
audio_streamer: BaseStreamer | list[BaseStreamer | None] | None = None,
past_key_values: Cache | None = None,
**kwargs,
) -> Lfm3AudioGenerateOutput | torch.LongTensor:
if input_ids.ndim != 2 or not input_ids.shape[0] or not input_ids.shape[1]:
raise ValueError("input_ids must be a nonempty batch of prompts")
batch_size = input_ids.shape[0]
for name in (
"logits_processor",
"stopping_criteria",
"prefix_allowed_tokens_fn",
"assistant_model",
"negative_prompt_ids",
"negative_prompt_attention_mask",
"custom_generate",
):
if kwargs.get(name) is not None:
raise ValueError(f"{type(self).__name__}.generate does not support `{name}`")
# rows stop independently, so ZeRO-3 / FSDP ranks would leave the collective forwards at different steps
if kwargs.pop("synced_gpus", None):
raise ValueError(f"{type(self).__name__}.generate does not support `synced_gpus=True`")
# call over passed configuration over `self.generation_config`; only a bare call defaults to the output object
return_dict = kwargs.get("return_dict_in_generate")
if return_dict is None and generation_config is not None:
return_dict = generation_config.return_dict_in_generate
if return_dict is None:
return_dict = self.generation_config.return_dict_in_generate
plain_sequences = return_dict is False or (return_dict is None and generation_config is not None)
has_default_max_length = (
kwargs.get("max_length") is None
and (generation_config is None or generation_config.max_length is None)
and self.generation_config.max_length is None
)
generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs)
unsupported = sorted(
name
for name, inert in self._generation_fields_unsupported.items()
if getattr(generation_config, name, inert) not in (inert, None)
)
if unsupported:
raise ValueError(
f"{type(self).__name__}.generate does not implement {unsupported}; see its docstring for the "
"supported fields"
)
def speech_setting(name, value, default=None):
# the call, then the generation configuration, like the text sampling fields
if value is None:
value = getattr(generation_config, name, None)
return default if value is None else value
audio_temperature = speech_setting("audio_temperature", audio_temperature)
audio_top_k = speech_setting("audio_top_k", audio_top_k)
audio_repetition_penalty = speech_setting("audio_repetition_penalty", audio_repetition_penalty)
audio_repetition_window = speech_setting("audio_repetition_window", audio_repetition_window, 64)
text_repetition_window = speech_setting("text_repetition_window", text_repetition_window, 32)
attention_mask = model_kwargs.pop("attention_mask", None)
# only a caller that keeps the cache pays the decode step of the last sampled slot
caller_cache = past_key_values is not None
past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
if attention_mask is None and past_length:
raise ValueError("Cached generation requires attention_mask covering the cached prefix and input_ids")
if attention_mask is None:
pad_id = generation_config.pad_token_id
if pad_id is None:
pad_id = self.config.text_config.pad_token_id
eos_id = generation_config.eos_token_id
if eos_id is None:
eos_id = self.config.text_config.eos_token_id
eos_ids = list(eos_id) if isinstance(eos_id, (list, tuple)) else [eos_id]
attention_mask = (
input_ids.ne(pad_id).long()
if pad_id is not None and pad_id not in eos_ids
else torch.ones_like(input_ids)
)
position_ids = model_kwargs.pop("position_ids", None)
self._validate_model_kwargs(model_kwargs.copy())
if generation_config.num_beams != 1 or generation_config.num_return_sequences != 1:
raise ValueError(f"{type(self).__name__} supports one greedy or sampled continuation per input row")
if attention_mask.shape != (batch_size, past_length + input_ids.shape[1]) or not bool(
attention_mask[:, -1].all()
):
raise ValueError("Use left-padded prompts and an attention_mask covering the cached prefix and input_ids")
not_binary = ((attention_mask != 0) & (attention_mask != 1)).any()
not_monotonic = ((attention_mask[:, 1:].long() - attention_mask[:, :-1].long()) < 0).any()
if bool(not_binary | not_monotonic):
raise ValueError("attention_mask must describe contiguous left-padded prompts of zeros and ones")
if position_ids is None:
position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -input_ids.shape[1] :]
if position_ids.shape not in ((1, input_ids.shape[1]), input_ids.shape):
raise ValueError("position_ids must cover the uncached input_ids, with one row or one row per prompt")
if audio_repetition_window < 0 or text_repetition_window < 0:
raise ValueError("Repetition windows must be nonnegative")
text_temperature = generation_config.temperature if generation_config.do_sample else None
text_top_k = generation_config.top_k if generation_config.do_sample else None
text_top_p = generation_config.top_p if generation_config.do_sample else None
text_min_p = generation_config.min_p if generation_config.do_sample else None
config = self.config
pad_token_id = generation_config.pad_token_id
if pad_token_id is None:
pad_token_id = config.text_config.pad_token_id
if pad_token_id is None:
pad_token_id = config.text_config.eos_token_id
if isinstance(pad_token_id, (list, tuple)):
pad_token_id = pad_token_id[0]
if pad_token_id is None:
raise ValueError("Set pad_token_id to assemble generated sequences")
eos_ids = generation_config.eos_token_id
if eos_ids is None:
eos_ids = config.text_config.eos_token_id
eos_ids = list(eos_ids) if isinstance(eos_ids, (list, tuple)) else [eos_ids]
suppress = self._placeholder_token_ids + [
token_id
for flag, token_id in (
(suppress_mixed, config.mixed_start_token_id),
(suppress_audio_out, config.audio_start_token_id),
)
if flag
]
text_only_suppress = suppress + [config.mixed_start_token_id, config.audio_start_token_id]
text_only = [False] * batch_size
if generation_mode is None:
modes = [
Lfm3AudioGenerationMode.AUDIO
if token in (config.audio_start_token_id, config.audio_out_token_id)
else Lfm3AudioGenerationMode.INTERLEAVED
if token == config.mixed_start_token_id
else Lfm3AudioGenerationMode.TEXT
for token in input_ids[:, -1].tolist()
]
else:
modes = generation_mode if isinstance(generation_mode, list) else [generation_mode] * batch_size
if len(modes) != batch_size:
raise ValueError("generation_mode must have one entry per input row")
modes = [Lfm3AudioGenerationMode(mode) for mode in modes]
text_only = [mode is Lfm3AudioGenerationMode.TEXT for mode in modes]
# speech always follows its marker in training; append the one the requested mode needs
markers = {
Lfm3AudioGenerationMode.AUDIO: (config.audio_start_token_id, config.audio_out_token_id),
Lfm3AudioGenerationMode.INTERLEAVED: (config.mixed_start_token_id,),
}
needed = [
markers[mode][0] if mode in markers and token not in markers[mode] else None
for mode, token in zip(modes, input_ids[:, -1].tolist())
]
if len(set(needed)) > 1:
raise ValueError("Rows need different generation markers; end each prompt with its marker instead")
if needed[0] is not None:
marker = input_ids.new_full((batch_size, 1), needed[0])
input_ids = torch.cat((input_ids, marker), 1)
attention_mask = torch.cat((attention_mask, attention_mask.new_ones((batch_size, 1))), 1)
position_ids = torch.cat((position_ids, position_ids[:, -1:] + 1), 1)
def row_streamers(value):
if value is None:
return [None] * batch_size
if isinstance(value, (list, tuple)) and len(value) == batch_size:
return list(value)
if batch_size == 1 and not isinstance(value, (list, tuple)):
return [value]
raise ValueError("Pass one streamer per input row")
text_streams, audio_streams = row_streamers(streamer), row_streamers(audio_streamer)
# the appended marker is part of the prompt, so `max_length` counts it
max_new_tokens = generation_config.max_new_tokens
if max_new_tokens is None and has_default_max_length:
warnings.warn(
f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the "
"generation length. We recommend setting `max_new_tokens` to control the maximum length of the "
"generation.",
UserWarning,
)
max_new_tokens = generation_config.max_length
elif max_new_tokens is None:
prompt_length = past_length + input_ids.shape[1]
if prompt_length >= generation_config.max_length:
raise ValueError(
f"Input length of input_ids is {prompt_length}, but `max_length` is set to "
f"{generation_config.max_length}. Increase `max_length` or set `max_new_tokens`."
)
max_new_tokens = generation_config.max_length - prompt_length
# the warper constructors validate the sampling values: build them before the prefill mutates a passed
# cache
text_warpers = self._build_warpers(text_temperature, text_top_k, text_top_p, text_min_p)
audio_warpers = self._build_warpers(audio_temperature, audio_top_k)
for penalty in (generation_config.repetition_penalty, audio_repetition_penalty):
if penalty is not None and penalty != 1.0:
RepetitionPenaltyLogitsProcessor(float(penalty))
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
input_features=input_features,
feature_attention_mask=feature_attention_mask,
audio_out_codes=audio_out_codes,
use_cache=True,
past_key_values=past_key_values
if past_key_values is not None
else DynamicCache(config=config.text_config),
**model_kwargs,
)
# Prefill has validated the flattened frame/placeholder mapping. A replayed terminal frame resumes
# the text phase, just as it does when sampled inside this call; explicit mode overrides still win.
if generation_mode is None and audio_out_codes is not None:
last_frames = (input_ids == config.audio_out_token_id).sum(-1).cumsum(0) - 1
for row, last_frame in enumerate(last_frames.tolist()):
if input_ids[row, -1] == config.audio_out_token_id:
if bool((audio_out_codes[last_frame] == config.audio_eos_code).any()):
modes[row] = Lfm3AudioGenerationMode.TEXT
hidden, cache = outputs.last_hidden_state[:, -1], outputs.past_key_values
embed_tokens = self.get_input_embeddings()
for row, stream in enumerate(text_streams):
if stream is not None:
stream.put(input_ids[row : row + 1])
text_out = [[] for _ in range(batch_size)]
frames = [[] for _ in range(batch_size)]
text_history = [deque(maxlen=text_repetition_window) for _ in range(batch_size)]
audio_history = [deque(maxlen=audio_repetition_window) for _ in range(batch_size)]
slots_left = [config.interleaved_n_text] * batch_size
speaking, text_done, finished = [False] * batch_size, [False] * batch_size, [False] * batch_size
for step in range(max_new_tokens):
interleaving = [mode is Lfm3AudioGenerationMode.INTERLEAVED for mode in modes]
audio_rows = [
row
for row in range(batch_size)
if not finished[row]
and (modes[row] is Lfm3AudioGenerationMode.AUDIO or (interleaving[row] and speaking[row]))
]
sampled_frames = {}
if audio_rows:
histories = [
torch.stack(tuple(audio_history[row])) if audio_history[row] else None for row in audio_rows
]
if len(audio_rows) == 1:
sampled = self.sample_audio_frame(
hidden[audio_rows[0]],
audio_temperature,
audio_top_k,
audio_repetition_penalty,
histories[0],
audio_warpers,
)[None]
else:
sampled = self.sample_audio_frame(
hidden[audio_rows],
audio_temperature,
audio_top_k,
audio_repetition_penalty,
histories,
audio_warpers,
)
sampled_frames = dict(zip(audio_rows, sampled))
# the rows still running feed this step's slot, so a caller's cache covers every returned slot
active = input_ids.new_tensor([not done for done in finished])[:, None]
next_embeds = hidden.new_zeros((batch_size, 1, hidden.shape[-1]))
for row in range(batch_size):
if finished[row]:
continue
if row in sampled_frames:
codes = sampled_frames[row]
is_eos = bool((codes == config.audio_eos_code).any())
if is_eos:
codes = torch.full_like(codes, config.audio_eos_code)
modes[row] = Lfm3AudioGenerationMode.TEXT
else:
audio_history[row].append(codes)
if audio_streams[row] is not None:
audio_streams[row].put(codes)
if not is_eos or return_audio_eos:
frames[row].append(codes)
text_out[row].append(config.audio_out_token_id)
next_embeds[row] = self.model.get_audio_out_embeds(codes[None])
else:
history = torch.stack(tuple(text_history[row])) if text_history[row] else None
token = self._sample(
self.lm_head(hidden[row]),
text_temperature,
text_top_k,
generation_config.repetition_penalty,
history,
text_only_suppress if text_only[row] else suppress,
top_p=text_top_p,
min_p=text_min_p,
warpers=text_warpers,
)
token_id = int(token)
text_history[row].append(token)
text_out[row].append(token_id)
if text_streams[row] is not None:
text_streams[row].put(token.view(1, 1))
if token_id == config.text_end_token_id:
# inside an interleaved block it ends the text and starts the audio tail; otherwise it ends
# the turn, so listing it in `eos_token_id` for text-only serving changes nothing here
if interleaving[row]:
text_done[row], speaking[row] = True, True
else:
finished[row] = True
elif token_id in eos_ids:
# `<|im_end|>` ends the turn in any text slot, interleaved blocks included
finished[row] = True
elif interleaving[row]:
pass # the mode markers only open a block; inside one they are ordinary text
elif token_id == config.audio_start_token_id:
modes[row] = Lfm3AudioGenerationMode.AUDIO
elif token_id == config.mixed_start_token_id:
modes[row] = Lfm3AudioGenerationMode.INTERLEAVED
slots_left[row], speaking[row], text_done[row] = config.interleaved_n_text, False, False
next_embeds[row] = embed_tokens(token.view(1))
if interleaving[row] and not text_done[row]:
slots_left[row] -= 1
if slots_left[row] == 0:
speaking[row] = not speaking[row]
slots_left[row] = config.interleaved_n_audio if speaking[row] else config.interleaved_n_text
last_step = all(finished) or step == max_new_tokens - 1
if last_step and not caller_cache:
# the last sampled token needs no hidden state: skip the decode step nothing would consume
break
attention_mask = torch.cat((attention_mask, active), dim=1)
position_ids = position_ids[:, -1:] + active
outputs = self.model.language_model(
inputs_embeds=next_embeds,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=cache,
use_cache=True,
)
if last_step:
break
hidden, cache = outputs.last_hidden_state[:, -1], outputs.past_key_values
generated = input_ids.new_full((batch_size, max(map(len, text_out))), pad_token_id)
audio_codes = input_ids.new_full((batch_size, config.num_codebooks, max(map(len, frames))), -1)
for row in range(batch_size):
generated[row, : len(text_out[row])] = input_ids.new_tensor(text_out[row])
if frames[row]:
audio_codes[row, :, : len(frames[row])] = torch.stack(frames[row], dim=-1)
output = Lfm3AudioGenerateOutput(
sequences=torch.cat((input_ids, generated), dim=1),
audio_codes=audio_codes,
sequence_lengths=input_ids.new_tensor([input_ids.shape[1] + len(tokens) for tokens in text_out]),
audio_lengths=input_ids.new_tensor([len(row_frames) for row_frames in frames]),
)
return output.sequences if plain_sequences else output
__all__ = [
"Lfm3AudioEncoderStreamingCache",
"Lfm3AudioLiveTranscriptionState",
"Lfm3AudioLiveTranscriptionOutput",
"Lfm3AudioPreTrainedModel",
"Lfm3AudioModel",
"Lfm3AudioForConditionalGeneration",
"Lfm3AudioGenerateOutput",
"Lfm3AudioStreamingOutput",
]
register_checkpoint_conversion_mapping(
"lfm3_audio",
[
# the released checkpoints name the modules differently from the model
WeightRenaming(source_patterns=r"^lfm\.model\.", target_patterns="model.language_model."),
WeightRenaming(source_patterns=r"^lfm\.lm_head\.", target_patterns="lm_head."),
WeightRenaming(
source_patterns=r"^encoder\.downsample\.", target_patterns="model.audio_encoder.downsample."
),
WeightRenaming(source_patterns=r"^encoder\.lfm\.lfm\.", target_patterns="model.audio_encoder.lfm."),
WeightRenaming(source_patterns=r"^mlp\.model\.0\.", target_patterns="model.audio_projector.layer_norm."),
WeightRenaming(source_patterns=r"^mlp\.model\.1\.", target_patterns="model.audio_projector.linear_1."),
WeightRenaming(source_patterns=r"^mlp\.model\.3\.", target_patterns="model.audio_projector.linear_2."),
WeightRenaming(source_patterns=r"^codebook_embedding\.", target_patterns="model.codebook_embedding."),
WeightRenaming(source_patterns=r"^depth_linear\.", target_patterns="model.depth_decoder.depth_linear."),
WeightRenaming(source_patterns=r"^mask_embedding$", target_patterns="model.depth_decoder.mask_embedding"),
WeightRenaming(source_patterns=r"^depthformer\.model\.", target_patterns="model.depth_decoder.model."),
WeightRenaming(
source_patterns=r"^depth_embeddings\.", target_patterns="model.depth_decoder.depth_embeddings."
),
],
overwrite=True,
)