# 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 math from dataclasses import dataclass import torch import torch.nn.functional as F from torch import nn from transformers import initialization as init from transformers.cache_utils import DynamicCache from transformers.conversion_mapping import register_checkpoint_conversion_mapping from transformers.core_model_loading import WeightRenaming from transformers.modeling_utils import PreTrainedAudioTokenizerBase from transformers.utils import ModelOutput, auto_docstring, can_return_tuple, torch_compilable_check from transformers.utils.import_utils import is_tracing from .configuration_liquid_codec import LiquidCodecConfig from .modeling_lfm3 import Lfm3Model @auto_docstring @dataclass class LiquidCodecEncoderOutput(ModelOutput): r""" audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`): The code of every codebook of every frame. """ audio_codes: torch.LongTensor = None @auto_docstring @dataclass class LiquidCodecDecoderOutput(ModelOutput): r""" audio_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`): The decoded 24 kHz audio in float32. """ audio_values: torch.FloatTensor = None @auto_docstring @dataclass class LiquidCodecOutput(ModelOutput): r""" audio_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`): The reconstructed audio, cut to the input length. audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`): The code of every codebook of every frame. """ audio_values: torch.FloatTensor = None audio_codes: torch.LongTensor = None class LiquidCodecStreamingCache: """Carries the decoder state between chunks: the LFM3 cache and the unfinished overlap-add tail.""" def __init__(self, config: LiquidCodecConfig): self.past_key_values = DynamicCache(config=config.lfm_config(config.decoder_window)) self.audio: torch.Tensor | None = None self.envelope: torch.Tensor | None = None self.started = False self.finished = False class LiquidCodecSTFT(nn.Module): def __init__(self, config: LiquidCodecConfig): super().__init__() self.n_fft = config.stft_n_fft self.hop_length = config.stft_hop_length self.win_length = config.stft_win_length self.window = nn.Buffer(torch.hann_window(self.win_length, dtype=torch.float32), persistent=False) self.num_features = 3 * (self.n_fft // 2 + 1) def forward(self, audio: torch.Tensor) -> torch.Tensor: total_pad = self.win_length - self.hop_length audio = F.pad(audio.float(), (total_pad // 2, total_pad - total_pad // 2), mode="reflect") spec = torch.stft( audio, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, window=self.window.float(), center=False, return_complex=True, ) magnitude = spec.abs() phase_denominator = magnitude.clamp_min(1e-5) return torch.cat((torch.log1p(magnitude), spec.real / phase_denominator, spec.imag / phase_denominator), dim=1) class LiquidCodecFrontEnd(nn.Module): def __init__(self, config: LiquidCodecConfig, num_features: int): super().__init__() self.frame_stack = config.stft_frame_stack self.in_proj = nn.Conv1d(num_features, config.encoder_conv_dim, kernel_size=1) self.conv = nn.Conv1d(config.encoder_conv_dim, config.encoder_conv_dim, kernel_size=3, padding=2) self.out = nn.Conv1d( config.encoder_conv_dim, config.encoder_dim, kernel_size=self.frame_stack, stride=self.frame_stack ) self.act = nn.SiLU() def forward(self, features: torch.Tensor) -> torch.Tensor: hidden_states = self.act(self.in_proj(features.to(self.in_proj.weight.dtype))) hidden_states = self.act(self.conv(hidden_states)[:, :, : hidden_states.shape[-1]]) return self.out(hidden_states) class LiquidCodecLfm(nn.Module): def __init__(self, config: LiquidCodecConfig, sliding_window: int | None = None): super().__init__() self.lfm = Lfm3Model(config.lfm_config(sliding_window)) self.lfm.register_module("embed_tokens", None) # the projections run in the model dtype but the rotary tables stay in float32; a named hook keeps the # module picklable self.lfm.rotary_emb.register_forward_pre_hook(self._rotary_in_float32) @staticmethod def _rotary_in_float32(module, args): return (args[0].float(), *args[1:]) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.lfm(inputs_embeds=hidden_states, use_cache=False).last_hidden_state class LiquidCodecQuantizer(nn.Module): def __init__(self, config: LiquidCodecConfig): super().__init__() levels = list(config.fsq_levels) self.num_codebooks = config.n_q self.codebook_dim = len(levels) self.levels = nn.Buffer(torch.tensor(levels, dtype=torch.int32), persistent=False) self.basis = nn.Buffer( torch.tensor([math.prod(levels[:index]) for index in range(self.codebook_dim)], dtype=torch.int32), persistent=False, ) effective_dim = self.codebook_dim * self.num_codebooks if config.quantizer_dim == effective_dim: self.project_in = nn.Identity() self.project_out = nn.Identity() else: self.project_in = nn.Linear(config.quantizer_dim, effective_dim) self.project_out = nn.Linear(effective_dim, config.quantizer_dim) def _bound(self, latents: torch.Tensor, eps: float = 1e-3) -> torch.Tensor: half_levels = (self.levels - 1) * (1 + eps) / 2 offset = torch.where(self.levels % 2 == 0, 0.5, 0.0) shift = torch.atanh(offset / half_levels) return torch.tanh(latents + shift) * half_levels - offset def encode(self, latents: torch.Tensor) -> torch.Tensor: # the quantizer runs outside `forward`, so no `device_map` hook moves its inputs (here and in `decode`) to the # device of its buffers latents = self.project_in(latents.movedim(1, -1).contiguous()).float().to(self.levels.device) if not is_tracing(latents): # finite audio far above full scale overflows the float32 STFT, and NaN latents round to INT64_MIN codes torch_compilable_check( torch.isfinite(latents), "The audio overflows the codec (non-finite latents); expected samples in [-1, 1]", ) latents = latents.reshape(*latents.shape[:-1], self.num_codebooks, self.codebook_dim) half_width = self.levels // 2 codes = self._bound(latents).round() / half_width indices = ((codes * half_width + half_width) * self.basis).sum(-1).round().long() return indices.transpose(1, 2) def decode(self, codes: torch.Tensor) -> torch.Tensor: indices = codes.to(self.basis.device).transpose(1, 2).unsqueeze(-1) half_width = self.levels // 2 latents = ((indices // self.basis) % self.levels - half_width) / half_width latents = latents.reshape(*latents.shape[:-2], -1) if isinstance(self.project_out, nn.Linear): latents = latents.to(self.project_out.weight.dtype) return self.project_out(latents).movedim(-1, 1) class LiquidCodecISTFT(nn.Module): def __init__(self, config: LiquidCodecConfig): super().__init__() self.n_fft = config.stft_n_fft self.hop_length = config.stft_hop_length self.win_length = config.stft_win_length self.window = nn.Buffer(torch.hann_window(self.win_length, dtype=torch.float32)) def forward(self, spec: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: num_frames = spec.shape[-1] frames = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") * self.window[None, :, None] output_size = (num_frames - 1) * self.hop_length + self.win_length fold = {"output_size": (1, output_size), "kernel_size": (1, self.win_length), "stride": (1, self.hop_length)} audio = F.fold(frames, **fold)[:, 0, 0] envelope = F.fold(self.window.square().expand(1, num_frames, -1).transpose(1, 2), **fold)[0, 0, 0] return audio, envelope class LiquidCodecDecoder(nn.Module): def __init__(self, config: LiquidCodecConfig): super().__init__() self.frame_stack = config.stft_frame_stack self.lfm = LiquidCodecLfm(config, sliding_window=config.decoder_window).lfm self.lin = nn.Linear(config.encoder_dim, 2 * (config.stft_n_fft // 2 + 1)) self.istft = LiquidCodecISTFT(config) def forward( self, embeddings: torch.Tensor, past_key_values: DynamicCache | None = None ) -> tuple[torch.Tensor, torch.Tensor]: hidden_states = self.lfm( inputs_embeds=F.interpolate( embeddings.mT, embeddings.shape[1] * self.frame_stack, mode="nearest-exact" ).mT, past_key_values=past_key_values, use_cache=past_key_values is not None, ).last_hidden_state log_magnitude, angle = torch.chunk(self.lin(hidden_states).mT.contiguous(), 2, 1) return self.istft(torch.polar(log_magnitude.float().exp(), angle.float())) @auto_docstring( custom_intro="The Liquid Codec neural audio codec: 24 kHz audio to 12 codebooks at 12.5 frames per second and back." ) class LiquidCodecModel(PreTrainedAudioTokenizerBase): config: LiquidCodecConfig base_model_prefix = "codec" main_input_name = "input_values" # the analysis window is a non-persistent buffer cast in the forward; only the synthesis window is a weight _keep_in_fp32_modules_strict = ["istft"] _supports_flash_attn = False _supports_sdpa = True def __init__(self, config: LiquidCodecConfig): super().__init__(config) self.stft = LiquidCodecSTFT(config) self.frontend = LiquidCodecFrontEnd(config, self.stft.num_features) self.encoder = LiquidCodecLfm(config) self.to_quantizer = nn.Linear(config.encoder_dim, config.quantizer_dim) self.quantizer = LiquidCodecQuantizer(config) self.decoder = LiquidCodecDecoder(config) self.from_quantizer = nn.Linear(config.quantizer_dim, config.encoder_dim) self.post_init() @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, (LiquidCodecSTFT, LiquidCodecISTFT)): init.copy_(module.window, torch.hann_window(module.win_length, dtype=torch.float32)) elif isinstance(module, LiquidCodecQuantizer): levels = list(self.config.fsq_levels) init.copy_(module.levels, torch.tensor(levels, dtype=torch.int32)) init.copy_( module.basis, torch.tensor([math.prod(levels[:index]) for index in range(len(levels))], dtype=torch.int32), ) @can_return_tuple def encode(self, input_values: torch.FloatTensor, *, return_dict: bool | None = None) -> LiquidCodecEncoderOutput: r""" Non-finite audio, or audio that overflows the codec, is refused eagerly; a compiled `encode` skips the check. input_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`): 24 kHz audio. The tail is zero-padded to a whole codec frame of 1920 samples. return_dict (`bool`, *optional*): Whether to return a [`~utils.ModelOutput`] instead of a plain tuple. """ if input_values.ndim != 2 or input_values.shape[-1] == 0: raise ValueError( f"Expected non-empty audio of shape (batch_size, num_samples), got {tuple(input_values.shape)}" ) if not is_tracing(input_values): # a non-finite sample would round to an INT64_MIN code that only fails far downstream; the tensor-valued # check would break a `fullgraph` compile or a CUDA-graph capture torch_compilable_check(torch.isfinite(input_values), "`input_values` must be finite") frame_size = self.config.frame_size num_samples = -(input_values.shape[-1] // -frame_size) * frame_size input_values = F.pad(input_values, (0, num_samples - input_values.shape[-1])) hidden_states = self.frontend(self.stft(input_values)).mT latents = self.to_quantizer(self.encoder(hidden_states)).mT return LiquidCodecEncoderOutput(audio_codes=self.quantizer.encode(latents)) @can_return_tuple def decode( self, audio_codes: torch.LongTensor, streaming_cache: LiquidCodecStreamingCache | None = None, finalize: bool = True, *, return_dict: bool | None = None, ) -> LiquidCodecDecoderOutput: r""" audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, num_frames)`): The codes to decode; every frame becomes 1920 samples. streaming_cache ([`LiquidCodecStreamingCache`], *optional*): The state carried between consecutive chunks of one stream. The chunks decode to the same audio as the concatenated codes. finalize (`bool`, *optional*, defaults to `True`): Whether this is the last chunk of the stream. Until then the trailing samples that later frames still contribute to are held back in the cache. return_dict (`bool`, *optional*): Whether to return a [`~utils.ModelOutput`] instead of a plain tuple. """ if audio_codes.ndim != 3 or audio_codes.shape[1] != self.config.n_q: raise ValueError( f"Expected codes of shape (batch_size, {self.config.n_q}, num_frames), got {tuple(audio_codes.shape)}" ) torch_compilable_check( (audio_codes >= 0) & (audio_codes < self.config.codebook_size), f"Codes must be in [0, {self.config.codebook_size})", ) if streaming_cache is not None and streaming_cache.finished: raise ValueError("This stream is finalized; start with a new cache") trim = (self.config.stft_win_length - self.config.stft_hop_length) // 2 tail = 2 * trim batch_size = audio_codes.shape[0] if not audio_codes.shape[-1] and not (finalize and streaming_cache is not None and streaming_cache.started): # no frames and no held-back tail to flush: the overlap-add state must stay untouched if streaming_cache is not None and finalize: streaming_cache.finished = True return LiquidCodecDecoderOutput(audio_values=audio_codes.new_zeros((batch_size, 0), dtype=torch.float32)) if audio_codes.shape[-1]: # `quantizer.decode` returns float32 latents when its projections are `nn.Identity` latents = self.quantizer.decode(audio_codes).mT.to(self.from_quantizer.weight.dtype) embeddings = self.from_quantizer(latents) past_key_values = None if streaming_cache is None else streaming_cache.past_key_values audio, envelope = self.decoder(embeddings, past_key_values) else: audio = self.decoder.istft.window.new_zeros(batch_size, tail) envelope = self.decoder.istft.window.new_zeros(tail) if streaming_cache is None: audio, envelope = audio[:, trim:-trim], envelope[trim:-trim] else: if streaming_cache.audio is not None: audio = torch.cat([audio[:, :tail] + streaming_cache.audio, audio[:, tail:]], dim=1) envelope = torch.cat([envelope[:tail] + streaming_cache.envelope, envelope[tail:]]) if not streaming_cache.started: audio, envelope = audio[:, trim:], envelope[trim:] streaming_cache.started = True if finalize: audio, envelope = audio[:, :-trim], envelope[:-trim] streaming_cache.audio, streaming_cache.envelope = None, None streaming_cache.finished = True else: streaming_cache.audio, streaming_cache.envelope = audio[:, -tail:], envelope[-tail:] audio, envelope = audio[:, :-tail], envelope[:-tail] return LiquidCodecDecoderOutput(audio_values=(audio / envelope).float()) @can_return_tuple @auto_docstring def forward(self, input_values: torch.FloatTensor) -> LiquidCodecOutput: r""" input_values (`torch.FloatTensor` of shape `(batch_size, num_samples)`): 24 kHz audio. """ audio_codes = self.encode(input_values, return_dict=True).audio_codes audio = self.decode(audio_codes, return_dict=True).audio_values return LiquidCodecOutput(audio_values=audio[..., : input_values.shape[-1]], audio_codes=audio_codes) __all__ = ["LiquidCodecModel", "LiquidCodecStreamingCache"] register_checkpoint_conversion_mapping( "liquid_codec", [ WeightRenaming(source_patterns=r"^_frontend\.", target_patterns="frontend."), WeightRenaming(source_patterns=r"^quantizer\.fsq\.", target_patterns="quantizer."), ], overwrite=True, )