# 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. """Processor class for LFM3-Audio.""" import copy import math import os import re from typing import TypedDict import numpy as np import torch from transformers.audio_utils import AudioInput, make_list_of_audio from transformers.configuration_utils import PreTrainedConfig from transformers.feature_extraction_utils import BatchFeature from transformers.processing_utils import ProcessingKwargs, ProcessorMixin from transformers.utils import TensorType, auto_docstring from .modeling_liquid_codec import LiquidCodecModel # The prompt format of the training data: ChatML turns, one placeholder per clip, no thinking block DEFAULT_CHAT_TEMPLATE = ( "{{- bos_token -}}" "{%- for message in messages -%}" "{{- '<|im_start|>' + message['role'] + '\\n' -}}" "{%- if message['content'] is string -%}{{- message['content'] -}}" "{%- else -%}{%- for content in message['content'] -%}" "{%- if content['type'] == 'audio' -%}{{- '<|audio|>' -}}" "{%- elif content['type'] == 'audio_out' -%}{{- '<|audio_out|>' -}}" "{%- elif content['type'] == 'text' -%}{{- content['text'] -}}" "{%- endif -%}{%- endfor -%}{%- endif -%}" "{{- '<|im_end|>\\n' -}}" "{%- endfor -%}" "{%- if add_generation_prompt -%}{{- '<|im_start|>assistant\\n' -}}{%- endif -%}" ) class Lfm3AudioFeatureExtractorKwargs(TypedDict, total=False): """ sampling_rate (`int`, *optional*): The sampling rate of the clips, which must be the extractor's own; it is checked, never resampled. return_tensors (`str` or [`~utils.TensorType`], *optional*): The tensor type of the returned features. load_audio_backend (`str`, *optional*): The decoding backend `apply_chat_template` loads audio with; the extractor itself ignores it. """ # the extractor always right-pads to whole frames and returns the mask, so the padding, truncation and # attention-mask options of `AudioKwargs` are not offered rather than silently dropped sampling_rate: int | None return_tensors: str | TensorType | None load_audio_backend: str | None class Lfm3AudioProcessorKwargs(ProcessingKwargs, total=False): audio_kwargs: Lfm3AudioFeatureExtractorKwargs _defaults = {} @auto_docstring class Lfm3AudioProcessor(ProcessorMixin): r""" Wraps the LFM3-Audio feature extractor, the LFM3 tokenizer and optionally the Liquid codec. The text carries one `audio_token` per input clip and one `audio_out_token` per output clip or speech frame of the prompt; the processor expands the clips to their encoder tokens (one per 8 spectrogram frames) and codec frames. The two placeholder tokens are added to the tokenizer when it does not define them: the model replaces their embeddings with the speech features, so they never reach the model as text. """ attributes = ["feature_extractor", "tokenizer"] tokenizer_class = "AutoTokenizer" audio_tokenizer_class = "LiquidCodecModel" valid_processor_kwargs = Lfm3AudioProcessorKwargs def __init__( self, feature_extractor=None, tokenizer=None, audio_tokenizer=None, chat_template=None, audio_token: str = "<|audio|>", audio_out_token: str = "<|audio_out|>", **kwargs, ): r""" feature_extractor ([`Lfm3AudioFeatureExtractor`]): The feature extractor computing the log-mel spectrograms. tokenizer ([`PreTrainedTokenizerFast`]): The LFM3 tokenizer. audio_tokenizer ([`LiquidCodecModel`], *optional*): The codec turning 24 kHz output clips into speech frames. chat_template (`str`, *optional*): A Jinja template to turn conversations into prompts. audio_token (`str`, *optional*, defaults to `"<|audio|>"`): The placeholder token of one encoded frame of an input clip. audio_out_token (`str`, *optional*, defaults to `"<|audio_out|>"`): The placeholder token of one speech frame of the prompt. """ self.audio_token = getattr(tokenizer, "audio_token", audio_token) self.audio_out_token = getattr(tokenizer, "audio_out_token", audio_out_token) missing = [ token for token in (self.audio_token, self.audio_out_token) if tokenizer.convert_tokens_to_ids(token) in (None, tokenizer.unk_token_id) ] if missing: tokenizer.add_tokens(missing, special_tokens=True) self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token) self.audio_out_token_id = tokenizer.convert_tokens_to_ids(self.audio_out_token) if chat_template is None: chat_template = DEFAULT_CHAT_TEMPLATE super().__init__(feature_extractor, tokenizer, chat_template=chat_template) # a released transformers only accepts its own codec classes as `audio_tokenizer` self.audio_tokenizer = audio_tokenizer @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *args, audio_tokenizer=None, **kwargs): # `ProcessorMixin` forwards unknown kwargs to the tokenizer loader, which would keep the codec object in the # tokenizer config and break `save_pretrained` processor = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) if audio_tokenizer is not None: processor.audio_tokenizer = audio_tokenizer return processor def to_dict(self): # `ProcessorMixin.to_dict` deep-copies the attributes; the codec weights are saved on their own. A shallow # clone drops the codec, so concurrent calls on `self` keep it. clone = copy.copy(self) codec = clone.__dict__.pop("audio_tokenizer", None) output = ProcessorMixin.to_dict(clone) if codec is not None: output["audio_tokenizer_subfolder"] = "audio_tokenizer" return output def save_pretrained(self, save_directory, push_to_hub=False, **kwargs): files = super().save_pretrained(save_directory, push_to_hub=False, **kwargs) if getattr(self, "audio_tokenizer", None) is not None: self.audio_tokenizer.save_pretrained(os.path.join(save_directory, "audio_tokenizer")) if push_to_hub: repo_id = kwargs.pop("repo_id", os.path.basename(os.path.abspath(save_directory))) self.push_to_hub(repo_id, **kwargs) return files @classmethod def get_processor_dict(cls, pretrained_model_name_or_path, **kwargs): codec_kwargs = { key: value for key, value in kwargs.items() if key in ("cache_dir", "force_download", "local_files_only", "token", "revision") } parent = kwargs.get("subfolder") or "" processor_dict, kwargs = super().get_processor_dict(pretrained_model_name_or_path, **kwargs) if "processor_class" not in processor_dict: config, _ = PreTrainedConfig.get_config_dict( pretrained_model_name_or_path, subfolder=parent, **codec_kwargs ) if config.get("model_type") == "lfm3_audio" and config.get("format_version") == 1: # Legacy exports carry the text-only tokenizer template, not the trained audio prompt format. processor_dict["chat_template"] = DEFAULT_CHAT_TEMPLATE subfolder = processor_dict.pop("audio_tokenizer_subfolder", None) if subfolder is not None: if subfolder != "audio_tokenizer": raise ValueError("Expected the bundled codec in the audio_tokenizer subfolder") codec_class = LiquidCodecModel processor_dict["audio_tokenizer"] = codec_class.from_pretrained( pretrained_model_name_or_path, subfolder=os.path.join(parent, subfolder), **codec_kwargs ) return processor_dict, kwargs @auto_docstring def __call__(self, text=None, audio: AudioInput | None = None, audio_out: AudioInput | None = None, **kwargs): r""" audio_out (`np.ndarray`, `torch.Tensor` or a list of them, *optional*): 24 kHz clips, one per `audio_out_token` in the text, encoded with the codec into `audio_out_codes`; the placeholder is repeated once per frame plus once for the end-of-speech frame that closes every clip. """ # the shared audio loader indexes the first clip before `validate_inputs` runs for name, clips in (("audio", audio), ("audio_out", audio_out)): if isinstance(clips, (list, tuple)) and len(clips) == 0: raise ValueError(f"`{name}` must hold at least one clip; pass `{name}=None` for prompts without one") if audio_out is None: return super().__call__(text=text, audio=audio, **kwargs) if text is None: raise ValueError("`audio_out` needs the text holding its placeholders") single = isinstance(text, str) prompts = [text] if single else list(text) expanded, codes, audio_out_offsets = self._expand_audio_out(prompts, audio_out) merged = self._merge_kwargs( self.valid_processor_kwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs ) return_tensors = merged["text_kwargs"].get("return_tensors") outputs = super().__call__(text=expanded[0] if single else expanded, audio=audio, **kwargs) self._check_audio_out_placeholders(outputs["input_ids"], codes.shape[0]) if "text_replacement_offsets" in outputs: # the shared expansion saw the text after the speech placeholders were repeated outputs["text_replacement_offsets"] = self._compose_replacement_offsets( prompts, audio_out_offsets, outputs["text_replacement_offsets"] ) outputs["audio_out_codes"] = self._as_return_tensors(codes, return_tensors) return outputs def _check_audio_out_placeholders(self, input_ids, num_frames: int) -> None: # truncation may have dropped placeholders the codes still cover rows = input_ids.tolist() if hasattr(input_ids, "tolist") else input_ids rows = [rows] if rows and isinstance(rows[0], int) else rows num_placeholders = sum(row.count(self.audio_out_token_id) for row in rows) if num_placeholders != num_frames: raise ValueError( f"Found {num_placeholders} {self.audio_out_token} tokens after tokenization for {num_frames} " "speech frames; disable truncation or shorten the prompt" ) @staticmethod def _as_return_tensors(codes: torch.Tensor, return_tensors): return BatchFeature({"codes": codes}, tensor_type=return_tensors)["codes"] @staticmethod def _rewrite_with_offsets(prompt: str, pattern: str, replace, kind: str) -> tuple[str, list[dict]]: """Rewrites `prompt` like `ProcessorMixin.get_text_with_replacements`, with the same offset records.""" pieces, row, last, offset = [], [], 0, 0 for match in re.finditer(pattern, prompt): start, end = match.span() replacement = replace(match) pieces.extend((prompt[last:start], replacement)) row.append( { "type": kind, "span": (start, end), "new_span": (start + offset, start + offset + len(replacement)), "text": match.group(), "replacement": replacement, } ) offset += len(replacement) - (end - start) last = end pieces.append(prompt[last:]) return "".join(pieces), row @staticmethod def _compose_replacement_offsets(prompts: list[str], first: list[list[dict]], second) -> list[list[dict]]: """The records of two successive rewrites (the caller's text A to B, then B to C) as one rewrite of A: an entry of the second that covers entries of the first absorbs them, so every entry spans the caller's text.""" second = second or [[] for _ in prompts] composed = [] for prompt, row_first, row_second in zip(prompts, first, second): def to_first(position, row_first=row_first): return position - sum( len(r["replacement"]) - (r["span"][1] - r["span"][0]) for r in row_first if r["new_span"][1] <= position ) mapped = [] for entry in row_second: start, end = to_first(entry["span"][0]), to_first(entry["span"][1]) mapped.append({**entry, "span": (start, end), "text": prompt[start:end]}) kept = [ r for r in row_first if not any(m["span"][0] <= r["span"][0] and r["span"][1] <= m["span"][1] for m in mapped) ] row, offset = [], 0 for entry in sorted(kept + mapped, key=lambda r: r["span"]): start, end = entry["span"] row.append({**entry, "new_span": (start + offset, start + offset + len(entry["replacement"]))}) offset += len(entry["replacement"]) - (end - start) composed.append(row) return composed def _expand_audio_out( self, prompts: list[str], audio_out: AudioInput ) -> tuple[list[str], torch.Tensor, list[list[dict]]]: if getattr(self, "audio_tokenizer", None) is None: raise ValueError("`audio_out` requires the processor to hold the codec as `audio_tokenizer`") clips = make_list_of_audio(audio_out) if sum(prompt.count(self.audio_out_token) for prompt in prompts) != len(clips): raise ValueError(f"The number of {self.audio_out_token} tokens in the text must match the number of clips") codes = [self.encode_audio(clip) for clip in clips] # the model ends speech with a frame of end-of-speech codes before the next text token; the prompt and the # training targets carry it the same way `generate(return_audio_eos=True)` does eos_frame = torch.full((1, codes[0].shape[1]), self.audio_eos_code, dtype=codes[0].dtype) codes = [torch.cat((frames, eos_frame)) for frames in codes] frame_counts = iter(code.shape[0] for code in codes) expanded, offsets = [], [] for prompt in prompts: prompt, row = self._rewrite_with_offsets( prompt, re.escape(self.audio_out_token), lambda match: self.audio_out_token * next(frame_counts), "audio_out", ) expanded.append(prompt) offsets.append(row) return expanded, torch.cat(codes), offsets @property def audio_eos_code(self) -> int: """The end-of-speech code: the first id after the codec's vocabulary.""" return math.prod(self.audio_tokenizer.config.fsq_levels) @torch.no_grad() def encode_audio(self, audio) -> torch.Tensor: codec = self.audio_tokenizer waveform = torch.as_tensor(audio if torch.is_tensor(audio) else np.asarray(audio), dtype=torch.float32) # the quantizer rounds non-finite values to INT64_MIN, which would surface as codes if waveform.ndim != 1 or not torch.isfinite(waveform).all(): raise ValueError("`audio_out` clips must be mono waveforms of finite samples") # an offloaded codec keeps some weights on `meta`: send the clip to a device that holds weights, and the hooks # move it on from there device = next( (param.device for param in codec.parameters() if param.device.type != "meta"), torch.device("cpu") ) # the codec's return type follows its config's `return_dict` codes = codec.encode(waveform.to(device)[None], return_dict=True).audio_codes[0].T.cpu() # a compiled codec skips its latent check, and a non-finite latent rounds to INT64_MIN if codes.numel() and not 0 <= int(codes.min()) <= int(codes.max()) < self.audio_eos_code: raise ValueError("The audio overflows the codec; expected `audio_out` samples in [-1, 1]") return codes def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs): super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs) if text is None: return prompts = [text] if isinstance(text, str) else text num_placeholders = sum(prompt.count(self.audio_token) for prompt in prompts) num_clips = 0 if audio is None else len(make_list_of_audio(audio)) if num_placeholders != num_clips: raise ValueError( f"Found {num_placeholders} {self.audio_token} tokens in the text but received {num_clips} clips" ) def _process_audio(self, audio: AudioInput, **kwargs): # the backend only serves `apply_chat_template`'s audio loading kwargs.pop("load_audio_backend", None) audio_inputs = self.feature_extractor(audio, **kwargs) lengths = np.asarray(audio_inputs["feature_attention_mask"]).sum(-1).tolist() audio_inputs["num_audio_tokens"] = [ self.feature_extractor.num_encoder_tokens(int(length)) for length in lengths ] audio_replacements = [self.replace_audio_token(audio_inputs, audio_idx=idx) for idx in range(len(lengths))] return audio_inputs, audio_replacements def replace_audio_token(self, audio_inputs: dict, audio_idx: int, **kwargs) -> str: return self.audio_token * audio_inputs["num_audio_tokens"][audio_idx] @property def unused_input_names(self) -> list[str]: return ["num_audio_tokens"] @property def model_input_names(self): tokenizer_input_names = self.tokenizer.model_input_names feature_extractor_input_names = self.feature_extractor.model_input_names return list(dict.fromkeys(tokenizer_input_names + feature_extractor_input_names + ["audio_out_codes"])) __all__ = ["Lfm3AudioProcessor"]