Add tiny random LFM3-Omni model for optimum-intel tests

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

audio_tokenizer/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "LiquidCodecModel"
+  ],
+  "auto_map": {
+    "AutoModel": "modeling_liquid_codec.LiquidCodecModel"
+  },
+  "decoder_window": 30,
+  "dtype": "float32",
+  "encoder_conv_dim": 16,
+  "encoder_dim": 16,
+  "fsq_levels": [
+    8,
+    7,
+    6,
+    6
+  ],
+  "initializer_range": 0.02,
+  "model_type": "liquid_codec",
+  "n_q": 12,
+  "quantizer_dim": 16,
+  "sampling_rate": 24000,
+  "stft_frame_stack": 6,
+  "stft_hop_length": 16,
+  "stft_n_fft": 64,
+  "stft_win_length": 64,
+  "transformers_version": "5.20.0"
+}

audio_tokenizer/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d329bbdcf38f5c88a57ea6cddcc8f4cb96f2285387b6a13855f7ac54711f3b2
+size 7676072

chat_template.jinja

@@ -0,0 +1,296 @@
+{{- bos_token -}}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+
+{%- macro render_value(v) -%}
+    {%- if v is string -%}
+        {{- v -}}
+    {%- elif v is boolean -%}
+        {{- "true" if v else "false" -}}
+    {%- elif v is none -%}
+        {{- "null" -}}
+    {%- elif v is mapping or (v is sequence and v is not string) -%}
+        {{- v | tojson -}}
+    {%- else -%}
+        {{- v -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro parse_content(content) -%}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is mapping and content.get("type") == "image" -%}
+        {{- "<|vision_start|><|image|><|vision_end|>" -}}
+    {%- elif content is mapping and content.get("type") == "video" -%}
+        {{- "<|vision_start|><|video|><|vision_end|>" -}}
+    {%- elif content is mapping and content.get("type") in ["audio", "audio_url", "input_audio"] -%}
+        {{- "<|vision_start|><|audio|><|vision_end|>" -}}
+    {%- elif content is mapping and content.get("type") == "audio_out" -%}
+        {{- "<|audio_out|>" -}}
+    {%- elif content is mapping and content.get("type") == "text" -%}
+        {{- (content.get("text") or "") | string -}}
+    {%- elif content is mapping -%}
+        {{- content | tojson -}}
+    {%- elif content is iterable -%}
+        {%- set _ns = namespace(result="") -%}
+        {%- for item in content -%}
+            {%- if item is string -%}
+                {%- set _ns.result = _ns.result + item -%}
+            {%- elif item is mapping and item.get("type") == "image" -%}
+                {%- set _ns.result = _ns.result + "<|vision_start|><|image|><|vision_end|>" -%}
+            {%- elif item is mapping and item.get("type") == "video" -%}
+                {%- set _ns.result = _ns.result + "<|vision_start|><|video|><|vision_end|>" -%}
+            {%- elif item is mapping and item.get("type") in ["audio", "audio_url", "input_audio"] -%}
+                {%- set _ns.result = _ns.result + "<|vision_start|><|audio|><|vision_end|>" -%}
+            {%- elif item is mapping and item.get("type") == "audio_out" -%}
+                {%- set _ns.result = _ns.result + "<|audio_out|>" -%}
+            {%- elif item is mapping and item.get("type") == "text" -%}
+                {%- set _ns.result = _ns.result + ((item.get("text") or "") | string) -%}
+            {%- else -%}
+                {%- set _ns.result = _ns.result + (item | tojson) -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {{- _ns.result -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_calls(tool_calls) -%}
+    {%- for tool_call in tool_calls -%}
+        {%- if not loop.first -%}{{- "\n\n" -}}{%- endif -%}
+        {%- set func = tool_call["function"] if "function" in tool_call else tool_call -%}
+        {%- set func_args = func.get("arguments") -%}
+        {%- if func_args is string and (func_args | trim) not in ["", "{}", "null"] -%}
+            {{- raise_exception("Tool call arguments must be a mapping, got a JSON-encoded string: parse arguments with json.loads() before applying the chat template") -}}
+        {%- endif -%}
+        {{- "<tool_call>\n<function=" + func["name"] + ">\n" -}}
+        {%- if func_args is mapping -%}
+            {%- for arg_name, arg_value in func_args.items() -%}
+                {{- "<parameter=" + arg_name + ">\n" + render_value(arg_value) + "\n</parameter>\n" -}}
+            {%- endfor -%}
+        {%- endif -%}
+        {{- "</function>\n</tool_call>" -}}
+    {%- endfor -%}
+{%- endmacro -%}
+
+{%- macro tool_to_json(tool) -%}
+    {%- set fn = tool.function if tool.function is defined else tool -%}
+    {%- set ns_t = namespace(first=true) -%}
+    {{- "{" -}}
+    {%- for k, v in fn.items() -%}
+        {%- if k != "defer_loading" -%}
+            {%- if not ns_t.first -%}{{- ", " -}}{%- endif -%}
+            {%- set ns_t.first = false -%}
+            {{- k | tojson -}}{{- ": " -}}{{- v | tojson -}}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- "}" -}}
+{%- endmacro -%}
+
+{#- a tool_reference content part (vLLM, SGLang) names a tool from `tools`,
+    usually a deferred one; it renders as that tool's definition -#}
+{%- macro is_tool_reference(item) -%}
+    {{- "1" if item is mapping and item.get("type") == "tool_reference" else "" -}}
+{%- endmacro -%}
+
+{%- macro referenced_tool_to_json(name) -%}
+    {%- set _ns = namespace(found=none) -%}
+    {%- for tool in tools or [] -%}
+        {%- set fn = tool.function if tool.function is defined else tool -%}
+        {%- if fn.name == name -%}
+            {%- set _ns.found = fn -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if _ns.found is none -%}
+        {{- raise_exception("tool_reference names a tool that is not in tools: " + name | string) -}}
+    {%- endif -%}
+    {{- tool_to_json(_ns.found) -}}
+{%- endmacro -%}
+
+{%- macro render_tool_response(content) -%}
+    {#- content that already carries the tag (pre-formatted by the caller) is
+        passed through as-is instead of being wrapped a second time -#}
+    {%- if content is string and "<tool_response>" in content -%}
+        {{- content -}}
+    {%- else -%}
+    {{- "<tool_response>\n" -}}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is mapping -%}
+        {{- parse_content(content) -}}
+    {%- elif content is iterable -%}
+        {%- for item in content -%}
+            {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+            {%- if is_tool_reference(item) -%}
+                {#- consecutive references share one <additional_tools> block -#}
+                {%- if loop.first or not is_tool_reference(loop.previtem) -%}
+                    {{- "<additional_tools>\n" -}}
+                {%- endif -%}
+                {{- referenced_tool_to_json(item.name) -}}
+                {%- if loop.last or not is_tool_reference(loop.nextitem) -%}
+                    {{- "\n</additional_tools>" -}}
+                {%- endif -%}
+            {%- elif item is mapping and item.result is defined and item.result is mapping and item.result.tools is defined -%}
+                {{- "<additional_tools>\n" -}}
+                {%- set _tl = item.result.tools -%}
+                {%- if _tl is string -%}
+                    {{- _tl -}}
+                {%- elif _tl is mapping -%}
+                    {{- tool_to_json(_tl) -}}
+                {%- else -%}
+                    {%- for t in _tl -%}
+                        {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+                        {{- tool_to_json(t) -}}
+                    {%- endfor -%}
+                {%- endif -%}
+                {{- "\n</additional_tools>" -}}
+            {%- else -%}
+                {{- parse_content(item) -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {{- "\n</tool_response>" -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- set ns = namespace(system_prompt="", last_user_index=-1) -%}
+{%- if messages and messages[0]["role"] == "system" -%}
+    {%- if messages[0].get("content") -%}
+        {%- set ns.system_prompt = parse_content(messages[0]["content"]) -%}
+    {%- endif -%}
+    {%- set messages = messages[1:] -%}
+{%- endif -%}
+{%- if tools -%}
+    {%- set tools_ns = namespace(lines=[]) -%}
+    {%- for tool in tools -%}
+        {%- set fn = tool.function if tool.function is defined else tool -%}
+        {%- set _dyn = (fn.defer_loading is defined and fn.defer_loading) or (tool.defer_loading is defined and tool.defer_loading) -%}
+        {%- if not _dyn -%}
+            {%- set tools_ns.lines = tools_ns.lines + [tool_to_json(fn)] -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- if tools_ns.lines -%}
+        {%- set ns.system_prompt = ns.system_prompt + ("\n\n" if ns.system_prompt else "") + "<tools>\n" + (tools_ns.lines | join("\n")) + "\n</tools>" -%}
+    {%- endif -%}
+{%- endif -%}
+{%- if ns.system_prompt -%}
+    {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
+{%- endif -%}
+{%- for message in messages -%}
+    {%- if message["role"] == "user" -%}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif -%}
+{%- endfor -%}
+{#- vLLM's and SGLang's Anthropic /v1/messages adapters split one tool_result that carries
+    tool_reference blocks into consecutive tool messages with the same tool_call_id (text first,
+    references second). Merge such runs into one tool turn so they render exactly like a single
+    OpenAI-format tool message whose content holds the same items. -#}
+{%- set carry = namespace(items=[]) -%}
+{%- for message in messages -%}
+    {%- set _next = messages[loop.index0 + 1] if not loop.last else none -%}
+    {%- set _merge = message.role == "tool" and not message.get("additional_tools")
+        and _next is not none and _next.role == "tool" and not _next.get("additional_tools")
+        and message.get("tool_call_id") and _next.get("tool_call_id") == message.get("tool_call_id") -%}
+    {%- set _items = [] -%}
+    {%- if message.role == "tool" and not message.get("additional_tools") -%}
+        {%- set _c = message.get("content") -%}
+        {%- if _c is string -%}
+            {%- set _items = [{"type": "text", "text": _c}] if _c else [] -%}
+        {%- elif _c is mapping -%}
+            {%- set _items = [_c] -%}
+        {%- elif _c -%}
+            {%- set _items = _c | list -%}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if _merge -%}
+        {%- set carry.items = carry.items + _items -%}
+    {%- else -%}
+    {{- "<|im_start|>" + message.role + "\n" -}}
+    {%- if message.role == "assistant" -%}
+        {%- generation -%}
+        {%- set thinking = message.thinking if message.thinking is defined else message.reasoning if message.reasoning is defined else message.reasoning_content -%}
+        {%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
+        {%- if thinking is defined and keep_thinking -%}
+            {{- "<think>" + (thinking or "") + "</think>" -}}
+        {%- elif keep_thinking and "</think>" not in (message.get("content") if message.get("content") is string else "") -%}
+            {{- "<think></think>" -}}
+        {%- endif -%}
+        {%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
+        {%- set _has_cfm = false -%}
+        {%- set content = "" -%}
+        {%- if message.get("content") -%}
+            {%- set content = parse_content(message.content) -%}
+        {%- endif -%}
+        {%- if not keep_thinking -%}
+            {%- if "</think>" in content -%}
+                {%- set content = content.split("</think>")[-1] | trim -%}
+            {%- endif -%}
+            {%- set content = "<think></think>" + content -%}
+        {%- endif -%}
+        {%- if content.endswith(_cfm_tag) -%}
+            {%- set _has_cfm = true -%}
+            {%- set _trunc_len = (content | length) - (_cfm_tag | length) -%}
+            {%- set content = content[:_trunc_len] -%}
+        {%- endif -%}
+        {{- content -}}
+        {%- if message.tool_calls -%}
+            {{- render_tool_calls(message.tool_calls) -}}
+        {%- endif -%}
+        {%- if _has_cfm -%}
+            {{- _cfm_tag -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+        {%- endgeneration -%}
+    {%- elif message.role == "tool" -%}
+        {%- if message.get("additional_tools") -%}
+            {%- set _at = message["additional_tools"] -%}
+            {#- definitions may live in additional_tools itself, or in content
+                with additional_tools used as a plain truthy flag -#}
+            {%- if _at is mapping or (_at is sequence and _at is not string) -%}
+                {%- set _src = _at -%}
+            {%- else -%}
+                {%- set _src = message.get("content") -%}
+            {%- endif -%}
+            {#- nothing to reveal: skip the wrapper entirely rather than
+                emitting an empty <additional_tools> block -#}
+            {%- if _src is string and "<tool_response>" in _src -%}
+                {{- _src -}}
+            {%- elif _src is string and "<additional_tools>" in _src -%}
+                {{- "<tool_response>\n" + _src + "\n</tool_response>" -}}
+            {%- elif _src -%}
+                {{- "<tool_response>\n<additional_tools>\n" -}}
+                {%- if _src is string -%}
+                    {{- _src -}}
+                {%- elif _src is mapping -%}
+                    {{- tool_to_json(_src) -}}
+                {%- else -%}
+                    {%- for t in _src -%}
+                        {%- if not loop.first -%}{{- "\n" -}}{%- endif -%}
+                        {{- tool_to_json(t) -}}
+                    {%- endfor -%}
+                {%- endif -%}
+                {{- "\n</additional_tools>\n</tool_response>" -}}
+            {%- endif -%}
+        {%- elif carry.items -%}
+            {{- render_tool_response(carry.items + _items) -}}
+            {%- set carry.items = [] -%}
+        {%- else -%}
+            {#- an empty or missing result still gets the wrapper, so every tool turn has the same shape -#}
+            {{- render_tool_response(message.get("content") or "") -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+    {%- else %}
+        {%- if message.get("content") -%}
+            {{- parse_content(message["content"]) -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+    {%- endif %}
+    {%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{- "<|im_start|>assistant\n" -}}
+    {%- if enable_thinking -%}
+        {{- "<think>" -}}
+    {%- else -%}
+        {{- "<think></think>" -}}
+    {%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,267 @@
+{
+  "architectures": [
+    "Lfm3OmniForConditionalGeneration"
+  ],
+  "audio_encoder_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "chunk_size_feed_forward": 0,
+    "conv_channels": 8,
+    "dtype": null,
+    "hidden_size": 16,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "lfm3_omni_encoder",
+    "num_mel_bins": 128,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "problem_type": null,
+    "return_dict": true,
+    "sliding_window": 4,
+    "text_config": {
+      "_name_or_path": "",
+      "architectures": null,
+      "block_auto_adjust_ff_dim": false,
+      "block_ffn_dim_multiplier": 1.0,
+      "block_multiple_of": 256,
+      "bos_token_id": 1,
+      "chunk_size_feed_forward": 0,
+      "conv_L_cache": 3,
+      "conv_bias": false,
+      "dtype": null,
+      "eos_token_id": null,
+      "full_attn_idxs": null,
+      "hidden_size": 16,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "initializer_range": 0.02,
+      "intermediate_size": 32,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "layer_types": [
+        "conv",
+        "full_attention",
+        "conv",
+        "full_attention"
+      ],
+      "max_position_embeddings": 128000,
+      "model_type": "lfm2",
+      "norm_eps": 1e-05,
+      "num_attention_heads": 2,
+      "num_hidden_layers": 4,
+      "num_key_value_heads": 1,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "pad_token_id": 0,
+      "problem_type": null,
+      "return_dict": true,
+      "rope_parameters": {
+        "rope_theta": 1000000.0,
+        "rope_type": "default"
+      },
+      "tie_word_embeddings": true,
+      "use_cache": true,
+      "vocab_size": 2
+    }
+  },
+  "audio_eos_code": 2016,
+  "audio_loss_weight": 1.0,
+  "audio_out_token_id": 125032,
+  "audio_start_token_id": 125012,
+  "audio_token_id": 125031,
+  "auto_map": {
+    "AutoConfig": "configuration_lfm3_omni.Lfm3OmniConfig",
+    "AutoModel": "modeling_lfm3_omni.Lfm3OmniModel",
+    "AutoModelForImageTextToText": "modeling_lfm3_omni.Lfm3OmniForConditionalGeneration",
+    "AutoModelForMultimodalLM": "modeling_lfm3_omni.Lfm3OmniForConditionalGeneration",
+    "AutoModelForSeq2SeqLM": "modeling_lfm3_omni.Lfm3OmniForConditionalGeneration"
+  },
+  "codebook_size": 2017,
+  "depth_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "conv_L_cache": 7,
+    "conv_bias": false,
+    "dtype": null,
+    "eos_token_id": null,
+    "head_dim": 4,
+    "hidden_size": 16,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 32,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_types": [
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 12,
+    "model_type": "lfm3",
+    "norm_eps": 1e-05,
+    "num_attention_heads": 4,
+    "num_hidden_layers": 2,
+    "num_key_value_heads": 2,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": null,
+    "problem_type": null,
+    "return_dict": true,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 1.0,
+        "rope_theta": 1000000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 256,
+    "tie_word_embeddings": true,
+    "use_cache": false,
+    "use_sandwich_norm": false,
+    "vocab_size": 1
+  },
+  "downsample_factor": 2,
+  "dtype": "float32",
+  "image_token_id": 125030,
+  "initializer_range": 0.02,
+  "interleaved_n_audio": 12,
+  "interleaved_n_text": 6,
+  "maskgit_schedule": [
+    4,
+    4,
+    4
+  ],
+  "mixed_start_token_id": 125015,
+  "model_type": "lfm3_omni",
+  "num_codebooks": 12,
+  "projector_bias": true,
+  "projector_hidden_act": "gelu",
+  "projector_hidden_size": 32,
+  "projector_layer_norm_eps": 1e-06,
+  "text_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "bos_token_id": 124894,
+    "chunk_size_feed_forward": 0,
+    "conv_L_cache": 7,
+    "conv_bias": false,
+    "dtype": null,
+    "eos_token_id": 124900,
+    "head_dim": 8,
+    "hidden_size": 16,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 32,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_types": [
+      "conv",
+      "conv",
+      "full_attention",
+      "conv",
+      "sliding_attention",
+      "conv",
+      "sliding_attention"
+    ],
+    "max_position_embeddings": 128000,
+    "model_type": "lfm3",
+    "norm_eps": 1e-05,
+    "num_attention_heads": 4,
+    "num_hidden_layers": 7,
+    "num_key_value_heads": 2,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": 124893,
+    "problem_type": null,
+    "return_dict": true,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.5,
+        "rope_theta": 1000000.0,
+        "rope_type": "default"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 4,
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "use_sandwich_norm": true,
+    "vocab_size": 128000
+  },
+  "text_end_token_id": 125014,
+  "text_loss_weight": 1.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.20.0",
+  "video_token_id": 125029,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "chunk_size_feed_forward": 0,
+    "conv_kernel_size": 3,
+    "dtype": "float32",
+    "hidden_size": 16,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 32,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_types": [
+      "conv",
+      "full_attention",
+      "conv",
+      "full_attention"
+    ],
+    "model_type": "lfm3_vl_vision",
+    "norm_eps": 1e-05,
+    "num_attention_heads": 2,
+    "num_hidden_layers": 4,
+    "num_key_value_heads": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "patch_size": 16,
+    "problem_type": null,
+    "return_dict": true,
+    "rope_parameters": {
+      "rope_theta": 1000000.0,
+      "rope_type": "axial"
+    },
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  }
+}

configuration_lfm3.py

@@ -0,0 +1,132 @@
+# 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.
+from huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.utils import auto_docstring
+
+
+@auto_docstring(checkpoint="LiquidAI/LFM3-1.7B")
+@strict
+class Lfm3Config(PreTrainedConfig):
+    r"""
+    rope_parameters (`dict[str, RopeParameters]`, *optional*):
+        RoPE parameters keyed by attention layer type, for example `{"full_attention": {"rope_type": "default",
+        "rope_theta": 1000000.0, "partial_rotary_factor": 0.5}, "sliding_attention": {"rope_type": "default",
+        "rope_theta": 10000.0}}`. `conv` layers do not use RoPE and have no entry. Full-attention layers rotate the
+        first `partial_rotary_factor * head_dim` dimensions and leave the rest unrotated.
+    conv_bias (`bool`, *optional*, defaults to `False`):
+        Whether to use bias in the conv layers.
+    conv_L_cache (`int`, *optional*, defaults to 7):
+        L_cache dim in the conv layers.
+    layer_types (`list[str]`, *optional*):
+        Per-layer operator type, each one of `"full_attention"`, `"sliding_attention"` or `"conv"`. Defaults to the
+        LFM3 pattern of a full-attention layer every fourth layer starting from the third, a sliding-attention layer
+        in between, and gated short convolutions elsewhere.
+    sliding_window (`int`, *optional*, defaults to 256):
+        Sliding-window size of the `sliding_attention` layers.
+    use_sandwich_norm (`bool`, *optional*, defaults to `True`):
+        Whether to apply an RMSNorm to the operator and feed-forward outputs before they enter the residual stream.
+
+    ```python
+    >>> from transformers import Lfm3Model, Lfm3Config
+
+    >>> # Initializing a LFM3 model
+    >>> configuration = Lfm3Config()
+
+    >>> # Initializing a model from the LFM3-1.7B style configuration
+    >>> model = Lfm3Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "lfm3"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    default_theta = 1000000.0
+
+    vocab_size: int = 128_000
+    hidden_size: int = 2048
+    intermediate_size: int = 7680
+    num_hidden_layers: int = 24
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 4
+    max_position_embeddings: int = 128_000
+    initializer_range: float = 0.02
+    norm_eps: float = 0.00001
+    use_cache: bool = True
+    pad_token_id: int | None = 124_893
+    bos_token_id: int | None = 124_894
+    eos_token_id: int | list[int] | None = 124_900
+    tie_word_embeddings: bool = True
+    rope_parameters: dict | None = None
+    conv_bias: bool = False
+    conv_L_cache: int = 7
+    layer_types: list[str] | None = None
+    _attention_layer_types = ("full_attention", "sliding_attention")
+    head_dim: int | None = None
+    sliding_window: int | None = 256
+    use_sandwich_norm: bool = True
+
+    def __post_init__(self, **kwargs):
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                "full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 0.5},
+                "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0},
+            }
+        if self.layer_types is None:
+            self.layer_types = [
+                "full_attention" if i % 4 == 2 else "sliding_attention" if i % 4 == 0 and i > 0 else "conv"
+                for i in range(self.num_hidden_layers)
+            ]
+
+        self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings)
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of ``@strict``-powered validation."""
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
+            )
+        allowed_layer_types = (*self._attention_layer_types, "conv")
+        invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
+        if invalid_layer_types:
+            raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
+        if "sliding_attention" in self.layer_types and self.sliding_window is None:
+            raise ValueError("`layer_types` has `sliding_attention` layers but `sliding_window` is None.")
+        missing_rope = [
+            layer_type
+            for layer_type in self._attention_layer_types
+            if layer_type in self.layer_types and layer_type not in self.rope_parameters
+        ]
+        if missing_rope:
+            raise ValueError(f"`rope_parameters` has no entry for the {missing_rope} layers in `layer_types`.")
+        for layer_type, rope_params in self.rope_parameters.items():
+            if not isinstance(rope_params, dict):
+                continue
+            if (
+                rope_params.get("rope_type", "default") in ("dynamic", "longrope")
+                and rope_params.get("partial_rotary_factor", 1.0) < 1
+            ):
+                raise ValueError(
+                    f"`{rope_params['rope_type']}` RoPE recomputes its frequencies during the forward pass over the "
+                    f"rotated dimension only and cannot keep LFM3's partial rotation on `{layer_type}` layers."
+                )
+
+
+__all__ = ["Lfm3Config"]

configuration_lfm3_omni.py

@@ -0,0 +1,300 @@
+# 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.
+from huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.models.auto import CONFIG_MAPPING, AutoConfig
+from transformers.utils import auto_docstring
+
+from .configuration_lfm3 import Lfm3Config
+from .configuration_lfm3_vl import Lfm3VlVisionConfig
+
+
+def _config_class(model_type):
+    # the LFM3 configs ship with the checkpoint instead of the transformers registry
+    remote = {"lfm3": Lfm3Config, "lfm3_vl_vision": Lfm3VlVisionConfig, "lfm3_omni_encoder": Lfm3OmniEncoderConfig}
+    return remote.get(model_type) or CONFIG_MAPPING[model_type]
+
+
+# `dtype` is dropped so that the top-level dtype governs every component of a legacy export
+_LEGACY_SUBCONFIG_DROPPED_KEYS = (
+    "architectures",
+    "transformers_version",
+    "is_causal",
+    "auto_map",
+    "_name_or_path",
+    "dtype",
+)
+
+
+def _subconfig_dict(values: dict, default_model_type: str) -> dict:
+    values = {key: value for key, value in values.items() if key not in _LEGACY_SUBCONFIG_DROPPED_KEYS}
+    values.setdefault("model_type", default_model_type)
+    return values
+
+
+@auto_docstring(checkpoint="LiquidAI/lfm3_omni_unstable")
+@strict
+class Lfm3OmniEncoderConfig(PreTrainedConfig):
+    r"""
+    text_config (`Lfm2Config` or `dict`, *optional*):
+        The configuration of the LFM2 model run over the downsampled spectrogram. Defaults to the 17-layer,
+        768-wide model of LFM3-Audio.
+    num_mel_bins (`int`, *optional*, defaults to 128):
+        The number of mel bins of the input features.
+    conv_channels (`int`, *optional*, defaults to 256):
+        The number of channels of the downsampling convolutions.
+    hidden_size (`int`, *optional*, defaults to 768):
+        The width of the encoder output, equal to `text_config.hidden_size`.
+    sliding_window (`int`, *optional*, defaults to 32):
+        The number of past frames each frame attends to, itself included.
+    initializer_range (`float`, *optional*, defaults to 0.02):
+        The standard deviation of the truncated normal initializer of the convolutions.
+    """
+
+    model_type = "lfm3_omni_encoder"
+    sub_configs = {"text_config": AutoConfig}
+
+    text_config: dict | PreTrainedConfig | None = None
+    num_mel_bins: int = 128
+    conv_channels: int = 256
+    hidden_size: int = 768
+    sliding_window: int = 32
+    initializer_range: float = 0.02
+
+    def __post_init__(self, **kwargs):
+        if isinstance(self.text_config, dict):
+            # the encoder consumes embeddings only: no end-of-sequence id in its two-token vocabulary
+            self.text_config = {"eos_token_id": None, **_subconfig_dict(self.text_config, "lfm2")}
+            self.text_config = _config_class(self.text_config["model_type"])(**self.text_config)
+        elif self.text_config is None:
+            self.text_config = _config_class("lfm2")(
+                vocab_size=2,
+                hidden_size=self.hidden_size,
+                intermediate_size=2560,
+                num_hidden_layers=17,
+                num_attention_heads=8,
+                num_key_value_heads=4,
+                layer_types=(["conv", "conv", "full_attention"] * 5) + ["full_attention", "full_attention"],
+                conv_L_cache=3,
+                rope_parameters={"rope_type": "default", "rope_theta": 1_000_000.0},
+                pad_token_id=0,
+                bos_token_id=1,
+                eos_token_id=None,
+            )
+        super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="LiquidAI/LFM3-Omni-2.2B")
+@strict
+class Lfm3OmniConfig(PreTrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Lfm3OmniForConditionalGeneration`].
+
+    text_config (`Lfm3Config` or `dict`, *optional*):
+        The configuration of the LFM3 language model.
+    audio_encoder_config (`Lfm3OmniEncoderConfig` or `dict`, *optional*):
+        The configuration of the speech encoder.
+    depth_config (`Lfm3Config` or `dict`, *optional*):
+        The configuration of the depth decoder, a bidirectional LFM3 model over the codebook axis of one speech
+        frame. Defaults to the 3-layer, 1024-wide model of LFM3-Audio.
+    num_codebooks (`int`, *optional*, defaults to 12):
+        The number of codebooks per speech frame.
+    codebook_size (`int`, *optional*, defaults to 2017):
+        The size of each codebook, the codec's codes plus the end-of-audio code.
+    audio_eos_code (`int`, *optional*, defaults to 2016):
+        The end-of-audio code. A frame holding it in every codebook ends a speech segment.
+    maskgit_schedule (`list[int]`, *optional*, defaults to `[4, 4, 4]`):
+        The number of codebooks revealed by each pass of the depth decoder when sampling a frame.
+    interleaved_n_text (`int`, *optional*, defaults to 6):
+        The number of text tokens per block of the interleaved generation mode.
+    interleaved_n_audio (`int`, *optional*, defaults to 12):
+        The number of speech frames per block of the interleaved generation mode.
+    audio_token_id (`int`, *optional*, defaults to 125031):
+        The id of the placeholder token standing for one encoded input speech frame.
+    audio_out_token_id (`int`, *optional*, defaults to 125032):
+        The id of the placeholder token standing for one output speech frame.
+    audio_start_token_id (`int`, *optional*, defaults to 125012):
+        The id of `<|audio_start|>`, which switches generation to speech until an end-of-audio frame.
+    text_end_token_id (`int`, *optional*, defaults to 125014):
+        The id of `<|text_end|>`, which ends the text of an interleaved answer.
+    mixed_start_token_id (`int`, *optional*, defaults to 125015):
+        The id of `<|mixed_start|>`, which switches generation to interleaved text and speech.
+    tie_word_embeddings (`bool`, *optional*, defaults to `True`):
+        Whether to tie the language model head to the input embeddings.
+    text_loss_weight (`float`, *optional*, defaults to 1.0):
+        The weight of the text loss in the training loss.
+    audio_loss_weight (`float`, *optional*, defaults to 1.0):
+        The weight of the speech loss in the training loss.
+    initializer_range (`float`, *optional*, defaults to 0.02):
+        The standard deviation of the truncated normal initializer of the projector and depth decoder.
+
+    vision_config (`Lfm3VlVisionConfig` or `dict`, *optional*):
+        The configuration of the image encoder.
+    video_token_id (`int`, *optional*, defaults to 125029):
+        The id of the placeholder token standing for one encoded video patch group.
+    image_token_id (`int`, *optional*, defaults to 125030):
+        The placeholder for one projected image token.
+    projector_hidden_size (`int`, *optional*, defaults to 4096):
+        The intermediate dimension of the image projector.
+    projector_hidden_act (`str`, *optional*, defaults to `"gelu"`):
+        The image projector activation.
+    projector_bias (`bool`, *optional*, defaults to `True`):
+        Whether the image projector linear layers have biases.
+    projector_layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+        The epsilon of the image projector layer normalization.
+    downsample_factor (`int`, *optional*, defaults to 2):
+        The side length of each image patch block merged into one language token.
+    """
+
+    model_type = "lfm3_omni"
+    sub_configs = {
+        "text_config": AutoConfig,
+        "audio_encoder_config": Lfm3OmniEncoderConfig,
+        "depth_config": AutoConfig,
+        "vision_config": AutoConfig,
+    }
+
+    text_config: dict | PreTrainedConfig | None = None
+    audio_encoder_config: dict | PreTrainedConfig | None = None
+    depth_config: dict | PreTrainedConfig | None = None
+    num_codebooks: int = 12
+    codebook_size: int = 2017
+    audio_eos_code: int = 2016
+    maskgit_schedule: list[int] | None = None
+    interleaved_n_text: int = 6
+    interleaved_n_audio: int = 12
+    audio_token_id: int = 125_031
+    audio_out_token_id: int = 125_032
+    audio_start_token_id: int = 125_012
+    text_end_token_id: int = 125_014
+    mixed_start_token_id: int = 125_015
+    tie_word_embeddings: bool = True
+    text_loss_weight: float = 1.0
+    audio_loss_weight: float = 1.0
+    initializer_range: float = 0.02
+
+    # the ids `generate` and `transcribe_live` index the text logits with
+    _token_id_fields = (
+        "audio_token_id",
+        "audio_out_token_id",
+        "audio_start_token_id",
+        "text_end_token_id",
+        "mixed_start_token_id",
+    )
+    vision_config: dict | PreTrainedConfig | None = None
+    image_token_id: int = 125_030
+    video_token_id: int = 125_029
+    projector_hidden_size: int = 4096
+    projector_hidden_act: str = "gelu"
+    projector_bias: bool = True
+    projector_layer_norm_eps: float = 1e-6
+    downsample_factor: int = 2
+
+    def __post_init__(self, **kwargs):
+        if isinstance(self.vision_config, dict):
+            self.vision_config = _config_class("lfm3_vl_vision")(**self.vision_config)
+        elif self.vision_config is None:
+            self.vision_config = _config_class("lfm3_vl_vision")()
+        if "lfm" in kwargs or "format_version" in kwargs:
+            raise ValueError(
+                "LFM3-Omni reads the canonical checkpoint layout only; training exports (`lfm` / `format_version`) "
+                "are not supported, load a canonical Hugging Face checkpoint."
+            )
+        if isinstance(self.text_config, dict):
+            self.text_config = _subconfig_dict(self.text_config, "lfm3")
+            self.text_config = _config_class(self.text_config["model_type"])(**self.text_config)
+        elif self.text_config is None:
+            self.text_config = _config_class("lfm3")()
+        self.validate_vision_fields()
+        if isinstance(self.audio_encoder_config, dict):
+            self.audio_encoder_config = self.sub_configs["audio_encoder_config"](**self.audio_encoder_config)
+        elif self.audio_encoder_config is None:
+            self.audio_encoder_config = self.sub_configs["audio_encoder_config"]()
+        if isinstance(self.depth_config, dict):
+            # the depth decoder consumes embeddings only; its token ids are placeholders
+            self.depth_config = {
+                "pad_token_id": None,
+                "bos_token_id": None,
+                "eos_token_id": None,
+                **_subconfig_dict(self.depth_config, "lfm3"),
+            }
+            self.depth_config = _config_class(self.depth_config["model_type"])(**self.depth_config)
+        elif self.depth_config is None:
+            self.depth_config = _config_class("lfm3")(
+                vocab_size=1,
+                pad_token_id=None,
+                hidden_size=1024,
+                intermediate_size=2816,
+                num_hidden_layers=3,
+                num_attention_heads=32,
+                num_key_value_heads=8,
+                head_dim=32,
+                max_position_embeddings=self.num_codebooks,
+                layer_types=["full_attention"] * 3,
+                rope_parameters={
+                    "full_attention": {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 1.0}
+                },
+                sliding_window=None,
+                use_sandwich_norm=False,
+                use_cache=False,
+                bos_token_id=None,
+                eos_token_id=None,
+            )
+        steps = kwargs.pop("maskgit_steps", 3)
+        if self.maskgit_schedule is None:
+            if not 1 <= steps <= self.num_codebooks:
+                raise ValueError(
+                    f"maskgit_steps must be between 1 and num_codebooks={self.num_codebooks}, got {steps}"
+                )
+            self.maskgit_schedule = [self.num_codebooks // steps] * steps
+            self.maskgit_schedule[-1] += self.num_codebooks - sum(self.maskgit_schedule)
+        self.maskgit_schedule = list(self.maskgit_schedule)
+        self._validate_speech_fields()
+        super().__post_init__(**kwargs)
+
+    def validate_speech_fields(self):
+        """Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
+        self._validate_speech_fields()
+
+    def _validate_speech_fields(self):
+        if min(self.maskgit_schedule, default=0) < 1 or sum(self.maskgit_schedule) != self.num_codebooks:
+            raise ValueError(
+                f"maskgit_schedule {self.maskgit_schedule} must be positive counts summing to "
+                f"num_codebooks={self.num_codebooks}"
+            )
+        # the end-of-audio code indexes every codebook; a code outside them could never be sampled or embedded
+        if not 0 <= self.audio_eos_code < self.codebook_size:
+            raise ValueError(
+                f"audio_eos_code={self.audio_eos_code} must lie in [0, codebook_size={self.codebook_size})"
+            )
+        if self.interleaved_n_text < 1 or self.interleaved_n_audio < 1:
+            raise ValueError("interleaved_n_text and interleaved_n_audio must be positive")
+        vocab_size = self.text_config.vocab_size
+        for name in self._token_id_fields:
+            if not 0 <= getattr(self, name) < vocab_size:
+                raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
+
+    def validate_vision_fields(self):
+        """Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
+        if self.downsample_factor != self.vision_config.spatial_merge_size:
+            raise ValueError("downsample_factor must match vision_config.spatial_merge_size")
+        # an `inputs_embeds` forward embeds both placeholder ids to find them
+        vocab_size = self.text_config.vocab_size
+        for name in ("image_token_id", "video_token_id"):
+            if not 0 <= getattr(self, name) < vocab_size:
+                raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
+
+
+__all__ = ["Lfm3OmniConfig", "Lfm3OmniEncoderConfig"]

configuration_lfm3_vl.py

@@ -0,0 +1,211 @@
+# 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.
+from huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.models.auto import CONFIG_MAPPING, AutoConfig
+from transformers.utils import auto_docstring
+
+from .configuration_lfm3 import Lfm3Config
+
+
+def _config_class(model_type):
+    # the LFM3 configs ship with the checkpoint instead of the transformers registry
+    remote = {"lfm3": Lfm3Config, "lfm3_vl_vision": Lfm3VlVisionConfig}
+    return remote.get(model_type) or CONFIG_MAPPING[model_type]
+
+
+@auto_docstring(checkpoint="LiquidAI/LFM3-VL-1.9B")
+@strict
+class Lfm3VlVisionConfig(PreTrainedConfig):
+    r"""
+    layer_types (`list[str]`, *optional*):
+        Per-layer operator type, each one of `"full_attention"` or `"conv"`. Defaults to two gated convolution
+        layers followed by alternating attention and convolution layers.
+    patch_size (`int`, *optional*, defaults to 16):
+        The spatial size of each image patch.
+    temporal_patch_size (`int`, *optional*, defaults to 2):
+        The number of video frames packed into one patch.
+    spatial_merge_size (`int`, *optional*, defaults to 2):
+        The side of the square block of patches merged into one language token.
+    conv_kernel_size (`int`, *optional*, defaults to 3):
+        The kernel size of the depthwise 2D convolution layers.
+    norm_eps (`float`, *optional*, defaults to 1e-05):
+        The epsilon used by the RMS normalization layers.
+    """
+
+    model_type = "lfm3_vl_vision"
+    base_config_key = "vision_config"
+    default_theta = 1_000_000.0
+    default_rope_type = "axial"
+    _attention_layer_types = ("full_attention",)
+
+    hidden_size: int = 1024
+    intermediate_size: int = 2560
+    num_hidden_layers: int = 14
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    layer_types: list[str] | None = None
+    patch_size: int = 16
+    temporal_patch_size: int = 2
+    in_channels: int = 3
+    spatial_merge_size: int = 2
+    conv_kernel_size: int = 3
+    norm_eps: float = 1e-5
+    initializer_range: float = 0.02
+    rope_parameters: dict | None = None
+
+    def __post_init__(self, **kwargs):
+        if self.layer_types is None:
+            self.layer_types = ["conv" if i < 2 or i % 2 else "full_attention" for i in range(self.num_hidden_layers)]
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of ``@strict``-powered validation."""
+        if len(self.layer_types) != self.num_hidden_layers:
+            raise ValueError(
+                f"`layer_types` has {len(self.layer_types)} entries but `num_hidden_layers` is {self.num_hidden_layers}."
+            )
+        allowed_layer_types = (*self._attention_layer_types, "conv")
+        invalid_layer_types = [layer_type for layer_type in self.layer_types if layer_type not in allowed_layer_types]
+        if invalid_layer_types:
+            raise ValueError(f"`layer_types` entries must be one of {allowed_layer_types}; got {invalid_layer_types}.")
+        if self.hidden_size % self.num_attention_heads:
+            raise ValueError("`hidden_size` must be a multiple of `num_attention_heads`.")
+
+
+# The first checkpoints spell every text and `vision_*` field at the top level of the config
+_LEGACY_TEXT_KEYS = (
+    "vocab_size",
+    "hidden_size",
+    "intermediate_size",
+    "num_hidden_layers",
+    "num_attention_heads",
+    "num_key_value_heads",
+    "head_dim",
+    "layer_types",
+    "conv_bias",
+    "conv_L_cache",
+    "sliding_window",
+    "max_position_embeddings",
+    "use_sandwich_norm",
+    "use_cache",
+)
+_LEGACY_VISION_KEYS = {
+    "vision_hidden_size": "hidden_size",
+    "vision_intermediate_size": "intermediate_size",
+    "vision_num_hidden_layers": "num_hidden_layers",
+    "vision_num_attention_heads": "num_attention_heads",
+    "vision_num_key_value_heads": "num_key_value_heads",
+    "vision_layer_types": "layer_types",
+    "vision_patch_size": "patch_size",
+    "vision_temporal_patch_size": "temporal_patch_size",
+    "vision_num_channels": "in_channels",
+    "vision_spatial_merge_size": "spatial_merge_size",
+}
+_LEGACY_DROPPED_KEYS = ("vision_patch_dim", "vision_max_position_embeddings")
+
+
+@auto_docstring(checkpoint="LiquidAI/LFM3-VL-1.9B")
+@strict
+class Lfm3VlConfig(PreTrainedConfig):
+    r"""
+    video_token_id (`int`, *optional*, defaults to 125029):
+        The id of the placeholder token for video patches.
+    vision_start_token_id (`int`, *optional*, defaults to 125009):
+        The id of the token opening an image or a video frame.
+    vision_end_token_id (`int`, *optional*, defaults to 125010):
+        The id of the token closing an image or a video frame.
+    projector_layer_norm_eps (`float`, *optional*, defaults to 1e-06):
+        The epsilon of the layer normalization applied to the vision features before projection.
+    downsample_factor (`int`, *optional*, defaults to 2):
+        The side of the square block of vision tokens merged by the projector, equal to the vision encoder's
+        `spatial_merge_size`.
+    """
+
+    model_type = "lfm3_vl"
+    sub_configs = {"text_config": AutoConfig, "vision_config": Lfm3VlVisionConfig}
+
+    text_config: dict | PreTrainedConfig | None = None
+    vision_config: dict | PreTrainedConfig | None = None
+    image_token_id: int = 125_030
+    video_token_id: int = 125_029
+    vision_start_token_id: int = 125_009
+    vision_end_token_id: int = 125_010
+    projector_hidden_act: str = "gelu"
+    projector_hidden_size: int = 4096
+    projector_bias: bool = True
+    projector_layer_norm_eps: float = 1e-6
+    downsample_factor: int = 2
+    tie_word_embeddings: bool = True
+
+    def __post_init__(self, **kwargs):
+        if "vision_hidden_size" in kwargs:
+            kwargs = self._convert_legacy_flat_config(kwargs)
+
+        if isinstance(self.vision_config, dict):
+            self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
+        elif self.vision_config is None:
+            self.vision_config = self.sub_configs["vision_config"]()
+
+        if isinstance(self.text_config, dict):
+            self.text_config = {"model_type": "lfm3", **self.text_config}
+            self.text_config = _config_class(self.text_config["model_type"])(**self.text_config)
+        elif self.text_config is None:
+            self.text_config = _config_class("lfm3")()
+        self.validate_vision_fields()
+
+        super().__post_init__(**kwargs)
+
+    def validate_vision_fields(self):
+        """Part of `@strict`-powered validation, so `validate()` and `save_pretrained` refuse later edits too."""
+        if self.downsample_factor != self.vision_config.spatial_merge_size:
+            raise ValueError("downsample_factor must match vision_config.spatial_merge_size")
+        # an `inputs_embeds` forward embeds both placeholder ids to find them
+        vocab_size = self.text_config.vocab_size
+        for name in ("image_token_id", "video_token_id"):
+            if not 0 <= getattr(self, name) < vocab_size:
+                raise ValueError(f"{name}={getattr(self, name)} must lie in [0, text_config.vocab_size={vocab_size})")
+
+    def _convert_legacy_flat_config(self, kwargs: dict) -> dict:
+        text_config = {key: kwargs.pop(key) for key in _LEGACY_TEXT_KEYS if key in kwargs}
+        text_config["norm_eps"] = kwargs.pop("rms_norm_eps", 1e-5)
+        text_config["rope_parameters"] = {
+            "full_attention": {
+                "rope_type": "default",
+                "rope_theta": kwargs.pop("full_attention_rope_theta", 1_000_000.0),
+                "partial_rotary_factor": kwargs.pop("partial_rotary_factor", 0.5),
+            },
+            "sliding_attention": {
+                "rope_type": "default",
+                "rope_theta": kwargs.pop("sliding_attention_rope_theta", 10_000.0),
+            },
+        }
+        for key in ("pad_token_id", "bos_token_id", "eos_token_id"):
+            if kwargs.get(key) is not None:
+                text_config[key] = kwargs[key]
+        text_config["tie_word_embeddings"] = self.tie_word_embeddings
+
+        vision_config = {new: kwargs.pop(old) for old, new in _LEGACY_VISION_KEYS.items() if old in kwargs}
+        vision_config["norm_eps"] = text_config["norm_eps"]
+        vision_config["rope_parameters"] = {"rope_type": "axial", "rope_theta": kwargs.pop("vision_rope_theta", 1e6)}
+        for key in _LEGACY_DROPPED_KEYS:
+            kwargs.pop(key, None)
+
+        self.text_config = text_config
+        self.vision_config = vision_config
+        return kwargs
+
+
+__all__ = ["Lfm3VlConfig", "Lfm3VlVisionConfig"]

configuration_liquid_codec.py

@@ -0,0 +1,131 @@
+# 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 huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.models.auto import CONFIG_MAPPING
+from transformers.utils import auto_docstring
+
+from .configuration_lfm3 import Lfm3Config
+
+
+def _config_class(model_type):
+    # the LFM3 configs ship with the checkpoint instead of the transformers registry
+    remote = {"lfm3": Lfm3Config}
+    return remote.get(model_type) or CONFIG_MAPPING[model_type]
+
+
+@auto_docstring(checkpoint="LiquidAI/Liquid-Codec")
+@strict
+class LiquidCodecConfig(PreTrainedConfig):
+    r"""
+    encoder_dim (`int`, *optional*, defaults to 512):
+        The width of the encoder and decoder LFM3 models.
+    encoder_conv_dim (`int`, *optional*, defaults to 256):
+        The number of channels of the STFT front end.
+    stft_n_fft (`int`, *optional*, defaults to 1280):
+        The FFT size of the analysis and synthesis transforms.
+    stft_hop_length (`int`, *optional*, defaults to 320):
+        The hop of the analysis and synthesis transforms.
+    stft_win_length (`int`, *optional*, defaults to 1280):
+        The window length of the analysis and synthesis transforms.
+    stft_frame_stack (`int`, *optional*, defaults to 6):
+        The number of STFT frames per codec frame.
+    quantizer_dim (`int`, *optional*, defaults to 256):
+        The width of the latent handed to the quantizer.
+    fsq_levels (`list[int]`, *optional*, defaults to `[8, 7, 6, 6]`):
+        The number of levels of each finite-scalar-quantization dimension; their product is the codebook size.
+    n_q (`int`, *optional*, defaults to 12):
+        The number of codebooks.
+    decoder_window (`int`, *optional*, defaults to 30):
+        The sliding attention window of the decoder LFM3 model, in 75 Hz frames.
+    sampling_rate (`int`, *optional*, defaults to 24000):
+        The sampling rate of the audio.
+    initializer_range (`float`, *optional*, defaults to 0.02):
+        The standard deviation of the truncated normal initializer of the weights.
+    """
+
+    model_type = "liquid_codec"
+
+    encoder_dim: int = 512
+    encoder_conv_dim: int = 256
+    stft_n_fft: int = 1280
+    stft_hop_length: int = 320
+    stft_win_length: int = 1280
+    stft_frame_stack: int = 6
+    quantizer_dim: int = 256
+    fsq_levels: list[int] | None = None
+    n_q: int = 12
+    decoder_window: int = 30
+    sampling_rate: int = 24_000
+    initializer_range: float = 0.02
+
+    def __post_init__(self, **kwargs):
+        # the released configs nest the fields under `model` and dump the two LFM3 sub-configs
+        dropped = ("encoder_lfm_config", "decoder_lfm_config", "n_output_samples")
+        for key, value in kwargs.pop("model", {}).items():
+            if key in dropped:
+                continue
+            if key not in self.__dataclass_fields__:
+                raise ValueError(f"Unknown Liquid codec setting `{key}`")
+            setattr(self, key, value)
+        for key in dropped:
+            kwargs.pop(key, None)
+        if self.fsq_levels is None:
+            self.fsq_levels = [8, 7, 6, 6]
+        self.fsq_levels = list(self.fsq_levels)
+        super().__post_init__(**kwargs)
+
+    @property
+    def codebook_size(self) -> int:
+        return math.prod(self.fsq_levels)
+
+    @property
+    def frame_size(self) -> int:
+        return self.stft_hop_length * self.stft_frame_stack
+
+    def lfm_config(self, sliding_window: int | None = None) -> PreTrainedConfig:
+        attention_type = "full_attention" if sliding_window is None else "sliding_attention"
+        return _config_class("lfm3")(
+            # the backbones are not sub-configs, so the requested attention implementation is handed over here
+            attn_implementation=self._attn_implementation,
+            vocab_size=1,
+            pad_token_id=None,
+            bos_token_id=None,
+            eos_token_id=None,
+            hidden_size=self.encoder_dim,
+            intermediate_size=2304,
+            num_hidden_layers=8,
+            num_attention_heads=8,
+            num_key_value_heads=2,
+            head_dim=64,
+            max_position_embeddings=131_072,
+            norm_eps=1e-5,
+            conv_L_cache=7,
+            conv_bias=False,
+            use_sandwich_norm=True,
+            layer_types=[attention_type if i in (2, 4, 6) else "conv" for i in range(8)],
+            sliding_window=sliding_window,
+            rope_parameters={
+                attention_type: {"rope_type": "default", "rope_theta": 1_000_000.0, "partial_rotary_factor": 1.0}
+            },
+            use_cache=False,
+            tie_word_embeddings=False,
+        )
+
+
+__all__ = ["LiquidCodecConfig"]

feature_extraction_lfm3_omni.py

@@ -0,0 +1,202 @@
+# 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 numpy as np
+
+from transformers.audio_utils import mel_filter_bank
+from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.utils import TensorType, is_torch_available, logging
+
+from .modeling_lfm3_omni import Lfm3OmniFeatureStreamingCache
+
+
+if is_torch_available():
+    import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+class Lfm3OmniFeatureExtractor(SequenceFeatureExtractor):
+    """The LFM3-Audio log-mel front end; canonical checkpoints bundle its configuration files."""
+
+    model_input_names = ["input_features", "feature_attention_mask"]
+
+    def __init__(
+        self,
+        feature_size: int = 128,
+        sampling_rate: int = 16000,
+        n_fft: int = 400,
+        hop_length: int = 160,
+        win_length: int = 400,
+        log_floor: float = -12.0,
+        normalization_offset: float = 4.0,
+        normalization_divisor: float = 4.0,
+        encoder_stride: int = 8,
+        padding_value: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.log_floor = log_floor
+        self.normalization_offset = normalization_offset
+        self.normalization_divisor = normalization_divisor
+        self.encoder_stride = encoder_stride
+        self.mel_filters = mel_filter_bank(
+            num_frequency_bins=n_fft // 2 + 1,
+            num_mel_filters=feature_size,
+            min_frequency=0.0,
+            max_frequency=sampling_rate / 2,
+            sampling_rate=sampling_rate,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+
+    def num_frames(self, num_samples: int) -> int:
+        padded = self.pad_length(num_samples)
+        return 1 + (padded - self.n_fft) // self.hop_length
+
+    def pad_length(self, num_samples: int) -> int:
+        if num_samples < self.n_fft:
+            return self.n_fft
+        return num_samples + (-(num_samples - self.n_fft) % self.hop_length)
+
+    def num_encoder_tokens(self, num_frames: int) -> int:
+        return -(num_frames // -self.encoder_stride)
+
+    def _log_mel(self, waveform: torch.Tensor) -> torch.Tensor:
+        window = torch.hann_window(self.win_length, device=waveform.device)
+        spectrogram = torch.stft(
+            waveform,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=window,
+            center=False,
+            return_complex=True,
+        )
+        power = spectrogram.abs().square()
+        mel_filters = torch.from_numpy(self.mel_filters).to(device=waveform.device, dtype=power.dtype)
+        log_mel = (mel_filters.T @ power).log10().clamp(min=self.log_floor)
+        return (log_mel + self.normalization_offset) / self.normalization_divisor
+
+    @torch.no_grad()
+    def stream(
+        self,
+        raw_speech: np.ndarray | torch.Tensor | list[float],
+        cache: Lfm3OmniFeatureStreamingCache | None = None,
+        is_final: bool = False,
+        sampling_rate: int | None = None,
+    ) -> tuple[BatchFeature, Lfm3OmniFeatureStreamingCache]:
+        """Extract completed windows from a mono chunk, retaining overlap. Finalization pads the last window."""
+        if sampling_rate is not None and sampling_rate != self.sampling_rate:
+            raise ValueError(f"Expected {self.sampling_rate} Hz, got {sampling_rate} Hz")
+        if cache is None:
+            cache = Lfm3OmniFeatureStreamingCache()
+        if cache.finished:
+            raise ValueError("This stream is finalized; start with a new cache")
+        chunk = torch.as_tensor(raw_speech, dtype=torch.float32)
+        if chunk.ndim != 1:
+            raise ValueError("stream expects one mono waveform")
+        samples = chunk if cache.samples is None else torch.cat((cache.samples, chunk))
+        num_samples = cache.num_samples + chunk.numel()
+        if is_final:
+            count = self.num_frames(num_samples) - cache.num_frames
+            if count:
+                size = (count - 1) * self.hop_length + self.n_fft
+                samples = torch.nn.functional.pad(samples, (0, size - samples.numel()))
+        else:
+            count = max(0, 1 + (samples.numel() - self.n_fft) // self.hop_length)
+        features = self._log_mel(samples) if count else samples.new_empty((self.feature_size, 0))
+        # the cache changes only once the chunk's features exist, so a failed call can be retried
+        cache.samples = samples[count * self.hop_length :].clone() if not is_final else None
+        cache.num_samples = num_samples
+        cache.num_frames += count
+        cache.finished = is_final
+        return BatchFeature(
+            {
+                "input_features": features[None],
+                "feature_attention_mask": torch.ones((1, count), device=chunk.device, dtype=torch.long),
+            }
+        ), cache
+
+    def __call__(
+        self,
+        raw_speech: np.ndarray | torch.Tensor | list[np.ndarray] | list[torch.Tensor] | list[float],
+        sampling_rate: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ) -> BatchFeature:
+        r"""
+        raw_speech (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`):
+            One mono clip or a list of mono clips, as float samples at `sampling_rate`.
+        sampling_rate (`int`, *optional*):
+            The sampling rate of `raw_speech`, checked against the extractor's.
+        return_tensors (`str` or [`~utils.TensorType`], *optional*):
+            If set, will return tensors instead of numpy arrays.
+        """
+        if kwargs:
+            raise TypeError(
+                f"{self.__class__.__name__} does not implement {sorted(kwargs)}: clips are always right-padded to "
+                "whole frames and `feature_attention_mask` is always returned"
+            )
+        if sampling_rate is not None and sampling_rate != self.sampling_rate:
+            raise ValueError(
+                f"{self.__class__.__name__} expects audio sampled at {self.sampling_rate} Hz, got {sampling_rate} Hz"
+            )
+        if sampling_rate is None:
+            logger.warning_once(
+                f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()` "
+                "to avoid silent errors."
+            )
+        if isinstance(raw_speech, (list, tuple)) and len(raw_speech) == 0:
+            raise ValueError(f"{self.__class__.__name__} needs at least one clip")
+        is_batched = isinstance(raw_speech, (list, tuple)) and hasattr(raw_speech[0], "__len__")
+        is_batched = is_batched or (isinstance(raw_speech, (np.ndarray, torch.Tensor)) and raw_speech.ndim == 2)
+        clips = list(raw_speech) if is_batched else [raw_speech]
+        clips = [
+            clip.detach().cpu().float()
+            if isinstance(clip, torch.Tensor)
+            else torch.as_tensor(np.asarray(clip, dtype=np.float32))
+            for clip in clips
+        ]
+        for clip in clips:
+            if clip.ndim != 1:
+                raise ValueError(
+                    f"{self.__class__.__name__} expects mono clips, got a clip of shape {tuple(clip.shape)}"
+                )
+
+        features = []
+        for clip in clips:
+            clip = torch.nn.functional.pad(clip, (0, self.pad_length(clip.shape[0]) - clip.shape[0]))
+            features.append(self._log_mel(clip))
+        lengths = [feature.shape[-1] for feature in features]
+        max_length = max(lengths)
+        input_features = torch.full((len(features), self.feature_size, max_length), float(self.padding_value))
+        feature_attention_mask = torch.zeros((len(features), max_length), dtype=torch.long)
+        for index, feature in enumerate(features):
+            input_features[index, :, : feature.shape[-1]] = feature
+            feature_attention_mask[index, : feature.shape[-1]] = 1
+
+        data = {
+            "input_features": input_features.numpy(),
+            "feature_attention_mask": feature_attention_mask.numpy(),
+        }
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["Lfm3OmniFeatureExtractor"]

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 124894,
+  "eos_token_id": 124900,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 124893,
+  "transformers_version": "5.20.0",
+  "use_cache": true
+}

image_processing_lfm3_vl.py

@@ -0,0 +1,293 @@
+# 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
+import warnings
+from collections.abc import Iterable
+
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from transformers.image_processing_backends import TorchvisionBackend
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_transforms import group_images_by_shape, reorder_images
+from transformers.image_utils import ImageInput, PILImageResampling, SizeDict
+from transformers.processing_utils import ImagesKwargs, Unpack
+from transformers.utils import TensorType, auto_docstring
+
+
+class Lfm3VlImageProcessorKwargs(ImagesKwargs, total=False):
+    r"""
+    min_pixels (`int`, *optional*, defaults to `56 * 56`):
+        The min pixels of the image to resize the image.
+    max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
+        The max pixels of the image to resize the image.
+    patch_size (`int`, *optional*, defaults to 14):
+        The spatial patch size of the vision encoder.
+    temporal_patch_size (`int`, *optional*, defaults to 2):
+        The temporal patch size of the vision encoder.
+    merge_size (`int`, *optional*, defaults to 2):
+        The merge size of the vision encoder to llm encoder.
+    """
+
+    min_pixels: int
+    max_pixels: int
+    patch_size: int
+    temporal_patch_size: int
+    merge_size: int
+
+
+def smart_resize(
+    height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
+):
+    """Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+
+    """
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = round(height / factor) * factor
+    w_bar = round(width / factor) * factor
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = max(factor, math.floor(height / beta / factor) * factor)
+        w_bar = max(factor, math.floor(width / beta / factor) * factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+    return h_bar, w_bar
+
+
+@auto_docstring
+class Lfm3VlImageProcessor(TorchvisionBackend):
+    do_resize = True
+    resample = PILImageResampling.BICUBIC
+    size = {"shortest_edge": 65_536, "longest_edge": 16_777_216}
+    default_to_square = False
+    do_rescale = True
+    do_normalize = True
+    image_mean = [0.5, 0.5, 0.5]
+    image_std = [0.5, 0.5, 0.5]
+    do_convert_rgb = True
+    patch_size = 16
+    temporal_patch_size = 2
+    merge_size = 2
+    valid_kwargs = Lfm3VlImageProcessorKwargs
+    model_input_names = ["pixel_values", "image_grid_thw"]
+
+    def __init__(self, **kwargs: Unpack[Lfm3VlImageProcessorKwargs]):
+        # backward compatibility: override size with min_pixels and max_pixels if they are provided
+        size = kwargs.pop("size", None)
+        size = dict(self.size) if size is None else size
+        if (min_pixels := kwargs.pop("min_pixels", None)) is not None:
+            size["shortest_edge"] = min_pixels
+            size.pop("min_pixels", None)
+        if (max_pixels := kwargs.pop("max_pixels", None)) is not None:
+            size["longest_edge"] = max_pixels
+            size.pop("max_pixels", None)
+        super().__init__(size=size, **kwargs)
+
+    def _standardize_kwargs(
+        self,
+        size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+        min_pixels: int | None = None,
+        max_pixels: int | None = None,
+        **kwargs,
+    ) -> dict:
+        if min_pixels is not None or max_pixels is not None:
+            warnings.warn(
+                "Passing `min_pixels` and `max_pixels` to a processor call is deprecated and will be removed in v5.23. "
+                "Pass in `size={'longest_edge': xxx, 'shortest_edge': xxx} to override the target size.`",
+                FutureWarning,
+            )
+
+            size_dict = dict(size) if isinstance(size, (dict, SizeDict)) else {}
+            if min_pixels is not None:
+                size_dict["shortest_edge"] = min_pixels
+            if max_pixels is not None:
+                size_dict["longest_edge"] = max_pixels
+            size = SizeDict(**size_dict)
+        return super()._standardize_kwargs(size=size, **kwargs)
+
+    @auto_docstring
+    def preprocess(
+        self,
+        images: ImageInput,
+        **kwargs: Unpack[Lfm3VlImageProcessorKwargs],
+    ) -> BatchFeature:
+        return super().preprocess(images, **kwargs)
+
+    def resize(
+        self,
+        images: "torch.Tensor",
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        factor: int,
+        **kwargs,
+    ) -> "torch.Tensor":
+        """Resize dynamically based on input image aspect ratio."""
+        if not size.shortest_edge or not size.longest_edge:
+            raise ValueError(f"`size` dict must contain 'shortest_edge' and 'longest_edge' keys but got {size}.")
+
+        height, width = images.shape[-2:]
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=factor,
+            min_pixels=size.shortest_edge,
+            max_pixels=size.longest_edge,
+        )
+        return super().resize(
+            image=images,
+            size=SizeDict(height=resized_height, width=resized_width),
+            resample=resample,
+        )
+
+    def patchify(
+        self,
+        images: "torch.Tensor",
+        patch_size: int,
+        merge_size: int,
+        temporal_patch_size: int,
+    ) -> tuple["torch.Tensor", int, int]:
+        """Patchifies each image into flat layout of shape (`seq_len`, `patch_dim`) so we can concat dynamically shaped pixels."""
+        batch_size, channel, resized_height, resized_width = images.shape
+        grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+        patches = images.reshape(
+            batch_size,
+            channel,
+            grid_h // merge_size,
+            merge_size,
+            patch_size,
+            grid_w // merge_size,
+            merge_size,
+            patch_size,
+        )
+        # Reorder dimensions to group grid and patch information for subsequent flattening.
+        # [batch, grid_h/merge, grid_w/merge, merge, merge, channel, patch, patch]
+        patches = patches.permute(0, 2, 5, 3, 6, 1, 4, 7)
+
+        flatten_patches = (
+            patches.unsqueeze(6)
+            .expand(-1, -1, -1, -1, -1, -1, temporal_patch_size, -1, -1)
+            .reshape(
+                batch_size,
+                grid_h * grid_w,
+                channel * temporal_patch_size * patch_size * patch_size,
+            )
+        )
+        return flatten_patches, grid_h, grid_w
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int,
+        temporal_patch_size: int,
+        merge_size: int,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        **kwargs,
+    ) -> BatchFeature:
+        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            if do_resize:
+                stacked_images = self.resize(
+                    images=stacked_images,
+                    size=size,
+                    resample=resample,
+                    factor=patch_size * merge_size,
+                )
+            resized_images_grouped[shape] = stacked_images
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+        processed_images_grouped = {}
+        processed_grids = {}
+        for shape, stacked_images in grouped_images.items():
+            stacked_images = self.rescale_and_normalize(
+                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+            )
+            patches, grid_h, grid_w = self.patchify(
+                stacked_images,
+                patch_size=patch_size,
+                merge_size=merge_size,
+                temporal_patch_size=temporal_patch_size,
+            )
+
+            processed_images_grouped[shape] = patches
+            processed_grids[shape] = [[1, grid_h, grid_w]] * len(stacked_images)
+
+        processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+        processed_grids_ordered = reorder_images(processed_grids, grouped_images_index)
+        pixel_values = processed_images[0] if len(processed_images) == 1 else torch.cat(processed_images, dim=0)
+        image_grid_thw = torch.tensor(processed_grids_ordered, dtype=torch.long)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, tensor_type=return_tensors
+        )
+
+    def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None) -> int:
+        """
+        A utility that returns number of image patches for a given image size.
+
+        Note: Do not remove this method! It is used by vLLM to infer the number of patches and placeholders
+        without an image input.
+
+        Args:
+            height (`int`):
+                Height of the input image.
+            width (`int`):
+                Width of the input image.
+            images_kwargs (`dict`, *optional*)
+                Any kwargs to override defaults of the image processor.
+        Returns:
+            `int`: Number of image patches per image.
+        """
+        # `size` is the supported override (`min_pixels` / `max_pixels` are deprecated): count with the budget the
+        # preprocessing applies
+        images_kwargs = dict(images_kwargs or {})
+        size = images_kwargs.pop("size", None) or self.size
+        images_kwargs.setdefault("min_pixels", size["shortest_edge"])
+        images_kwargs.setdefault("max_pixels", size["longest_edge"])
+        images_kwargs = images_kwargs or {}
+        min_pixels = images_kwargs["min_pixels"] if "min_pixels" in images_kwargs else self.size["shortest_edge"]
+        max_pixels = images_kwargs["max_pixels"] if "max_pixels" in images_kwargs else self.size["longest_edge"]
+        patch_size = images_kwargs.get("patch_size", self.patch_size)
+        merge_size = images_kwargs.get("merge_size", self.merge_size)
+
+        factor = patch_size * merge_size
+        resized_height, resized_width = smart_resize(
+            height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels
+        )
+        grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+        return grid_h * grid_w
+
+
+__all__ = ["Lfm3VlImageProcessor"]

lfm3_conv.py

@@ -0,0 +1,202 @@
+# 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.
+"""Inference kernels for LFM3's width-7 depthwise convolution."""
+
+import torch
+
+from transformers.utils import is_triton_available
+
+
+if is_triton_available():
+    import triton
+    import triton.language as tl
+
+    # the length and strides are runtime arguments, or every prompt length would compile its own kernel
+    @triton.jit(do_not_specialize=["T", "X0", "X1", "X2", "S0", "S1", "ST"])
+    def _prefill(
+        X,
+        W,
+        BIAS,
+        SEQ,
+        Y,
+        T,
+        D: tl.constexpr,
+        X0,
+        X1,
+        X2,
+        S0,
+        S1,
+        ST,
+        HAS_BIAS: tl.constexpr,
+        PACKED: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        # int64 offsets: batch * channels * length can exceed 2**31 elements
+        b, d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+        t = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+        acc = tl.full((BLOCK,), 0, tl.float32)
+        if PACKED:
+            doc = tl.load(SEQ + b * S0 + (t - T + ST) * S1, (t < T) & (t >= T - ST), other=-1)
+        for k in tl.static_range(7):
+            pos = t + k - 6
+            valid = (t < T) & (pos >= 0)
+            if PACKED:
+                prev = tl.load(SEQ + b * S0 + (pos - T + ST) * S1, valid & (pos >= T - ST), other=-1)
+                valid = valid & (doc == prev)
+            x = tl.load(X + b * X0 + d * X1 + pos * X2, valid, other=0).to(tl.float32)
+            w = tl.load(W + d * 7 + k).to(tl.float32)
+            acc += x * w
+        if HAS_BIAS:
+            acc += tl.load(BIAS + d).to(tl.float32)
+        tl.store(Y + (b * D + d) * T + t, acc, t < T)
+
+    @triton.jit
+    def _decode(
+        X,
+        W,
+        BIAS,
+        STATE,
+        Y,
+        D: tl.constexpr,
+        K: tl.constexpr,
+        X0: tl.constexpr,
+        X1: tl.constexpr,
+        S0: tl.constexpr,
+        S1: tl.constexpr,
+        S2: tl.constexpr,
+        HAS_BIAS: tl.constexpr,
+        BLOCK: tl.constexpr,
+    ):
+        b = tl.program_id(0).to(tl.int64)
+        d = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+        k = tl.arange(0, 8)
+        # Load the complete state before writing the shifted window.
+        old = tl.load(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, (d[:, None] < D) & (k[None, :] < K), other=0)
+        x = tl.load(X + b * X0 + d * X1, d < D, other=0)
+        shifted = tl.gather(old, tl.broadcast_to(tl.minimum(k + 1, 7)[None, :], (BLOCK, 8)), 1)
+        shifted = tl.where(k[None, :] == K - 1, x[:, None], shifted)
+        # A cache may store either K=6 previous inputs or K=7 including the oldest discarded input.
+        positions = k + K - 6
+        values = tl.gather(old, tl.broadcast_to(tl.minimum(positions, 7)[None, :], (BLOCK, 8)), 1)
+        values = tl.where(k[None, :] == 6, x[:, None], values)
+        weights = tl.load(W + d[:, None] * 7 + k[None, :], (d[:, None] < D) & (k[None, :] < 7), other=0)
+        acc = tl.sum(values.to(tl.float32) * weights.to(tl.float32), 1)
+        if HAS_BIAS:
+            acc += tl.load(BIAS + d, d < D, other=0).to(tl.float32)
+        tl.store(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, shifted, (d[:, None] < D) & (k[None, :] < K))
+        tl.store(Y + b * D + d, acc, d < D)
+
+
+def can_use_lfm3_conv(hidden_states, weight):
+    return (
+        is_triton_available()
+        and hasattr(tl, "gather")  # the decode kernel needs Triton >= 3.3
+        and hidden_states.is_cuda
+        and not torch.is_grad_enabled()
+        and hidden_states.dtype == weight.dtype
+        and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32)
+        and weight.shape[-1] == 7
+    )
+
+
+def lfm3_causal_conv1d(hidden_states, weight, bias=None, seq_idx=None):
+    batch, channels, length = hidden_states.shape
+    # the kernels index every tensor by the batch and channel of `hidden_states` and cannot bound the others
+    if weight.shape != (channels, 7) or (bias is not None and bias.shape != (channels,)):
+        raise ValueError(
+            f"`weight` must have shape ({channels}, 7) and `bias` ({channels},), got {tuple(weight.shape)}"
+        )
+    if seq_idx is not None and (
+        seq_idx.ndim != 2 or seq_idx.shape[0] not in (1, batch) or seq_idx.dtype not in (torch.int32, torch.int64)
+    ):
+        raise ValueError(
+            f"`seq_idx` must be an int32 or int64 tensor of shape (1, length) or ({batch}, length), got "
+            f"{seq_idx.dtype} {tuple(seq_idx.shape)}"
+        )
+    if any(t is not None and t.device != hidden_states.device for t in (weight, bias, seq_idx)):
+        raise ValueError(
+            f"`weight`, `bias` and `seq_idx` must be on the device of `hidden_states` ({hidden_states.device})"
+        )
+    output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
+    # Triton launches on the current device; the model may live on another one (`device_map`, `.to("cuda:1")`).
+    with torch.cuda.device(hidden_states.device):
+        _prefill[(batch, channels, triton.cdiv(length, 128))](
+            hidden_states,
+            # the kernels index `bias` and `weight` densely
+            weight.contiguous(),
+            bias.contiguous() if bias is not None else None,
+            seq_idx,
+            output,
+            length,
+            channels,
+            *hidden_states.stride(),
+            *(
+                (0 if seq_idx.shape[0] == 1 else seq_idx.stride(0), seq_idx.stride(1))
+                if seq_idx is not None
+                else (0, 0)
+            ),
+            seq_idx.shape[-1] if seq_idx is not None else 0,
+            bias is not None,
+            seq_idx is not None,
+            128,
+            enable_fp_fusion=False,
+        )
+    return output
+
+
+def lfm3_causal_conv1d_update(hidden_states, conv_state, weight, bias=None):
+    batch, channels, length = hidden_states.shape
+    # the kernel addresses one position per row and bounds nothing by the other tensors' shapes
+    if length != 1:
+        raise ValueError(f"`hidden_states` must hold one position per row, got {tuple(hidden_states.shape)}")
+    # the 8-wide tile holds a 6- or 7-wide state (see `positions` in `_decode`)
+    if (
+        conv_state.ndim != 3
+        or conv_state.shape[:2] != (batch, channels)
+        or conv_state.shape[-1] not in (6, 7)
+        or weight.shape != (channels, 7)
+    ):
+        raise ValueError(
+            f"`conv_state` must have shape ({batch}, {channels}, 6 or 7) and `weight` ({channels}, 7), got "
+            f"{tuple(conv_state.shape)} and {tuple(weight.shape)}"
+        )
+    if bias is not None and bias.shape != (channels,):
+        raise ValueError(f"`bias` must have shape ({channels},), got {tuple(bias.shape)}")
+    # the kernel writes the state in place: rows or channels sharing memory (an `expand`ed cache) would race
+    extent = 0
+    for stride, size in sorted(zip(conv_state.stride(), conv_state.shape)):
+        if size > 1 and stride <= extent:
+            raise ValueError(f"`conv_state` must not overlap itself, got strides {conv_state.stride()}")
+        extent += stride * (size - 1)
+    if any(t is not None and t.device != hidden_states.device for t in (conv_state, weight, bias)):
+        raise ValueError(
+            f"`conv_state`, `weight` and `bias` must be on the device of `hidden_states` ({hidden_states.device})"
+        )
+    output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
+    with torch.cuda.device(hidden_states.device):
+        _decode[(batch, triton.cdiv(channels, 32))](
+            hidden_states,
+            weight.contiguous(),
+            bias.contiguous() if bias is not None else None,
+            conv_state,
+            output,
+            channels,
+            conv_state.shape[-1],
+            *hidden_states.stride()[:2],
+            *conv_state.stride(),
+            bias is not None,
+            32,
+            enable_fp_fusion=False,
+        )
+    return output

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:510de08fc53fed5cafed593ee8c058cedf2339546679980fc5de033fce5dd995
+size 11635872

modeling_lfm3.py

@@ -0,0 +1,882 @@
+# 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
+from collections.abc import Callable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers import initialization as init
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.generation import GenerationMixin
+from transformers.integrations import (
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub_with_fallback,
+    use_kernelized_func,
+)
+from transformers.integrations.accelerate import force_accelerate_hooks
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_recurrent_attention_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.import_utils import is_tracing
+from transformers.utils.output_capturing import capture_outputs
+
+from .configuration_lfm3 import Lfm3Config
+from .lfm3_conv import can_use_lfm3_conv, lfm3_causal_conv1d, lfm3_causal_conv1d_update
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Lfm3RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        Lfm3RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        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 self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Lfm3RotaryEmbedding(nn.Module):
+    """
+    Rotary embedding per attention layer type, like Gemma 3; `conv` layers register no buffers. Partial RoPE keeps the
+    frequency progression of the full head dimension and zeroes the tail instead of shrinking the rotated dimension.
+    """
+
+    def __init__(self, config: Lfm3Config, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        self.layer_types = sorted(set(config.layer_types) & set(config.rope_parameters))
+        self.rope_type = {}
+        for layer_type in self.layer_types:
+            rope_params = self.config.rope_parameters[layer_type]
+            self.rope_type[layer_type] = rope_params["rope_type"]
+            curr_inv_freq, curr_attention_scaling = self.compute_rope_parameters(config, device, layer_type=layer_type)
+            setattr(self, f"{layer_type}_inv_freq", nn.Buffer(curr_inv_freq, persistent=False))
+            setattr(self, f"{layer_type}_original_inv_freq", nn.Buffer(curr_inv_freq.clone(), persistent=False))
+            setattr(self, f"{layer_type}_attention_scaling", curr_attention_scaling)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: Lfm3Config, device=None, layer_type: str | None = None, **kwargs
+    ) -> tuple[torch.Tensor, float]:
+        """
+        Inverse frequencies of the default RoPE with LFM3's zero-tail partial rotation: the first
+        `partial_rotary_factor * head_dim // 2` frequencies follow the progression of the full head dimension and the
+        rest are zero, leaving those dimensions unrotated.
+
+        Returns:
+            Tuple of (`torch.Tensor`, `float`): the inverse frequencies and the attention scaling factor (always 1).
+        """
+        base = config.rope_parameters[layer_type]["rope_theta"]
+        partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        rope_angles = int(partial_rotary_factor * dim) // 2
+        inv_freq = 1.0 / (base ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float) / dim))
+        inv_freq = torch.cat((inv_freq, torch.zeros(dim // 2 - rope_angles, dtype=torch.float)))
+        return inv_freq.to(device), 1.0
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids, layer_type):
+        inv_freq = getattr(self, f"{layer_type}_inv_freq")
+        attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
+
+        freqs = position_ids[..., None].float() * inv_freq.to(device=x.device, dtype=torch.float)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos() * attention_scaling
+        sin = emb.sin() * attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+    @classmethod
+    def compute_rope_parameters(
+        cls, config: Lfm3Config, device=None, layer_type: str | None = None
+    ) -> tuple[torch.Tensor, float]:
+        """Inverse frequencies of any RoPE type in the zero-tail layout: scaled types are evaluated over the full head
+        dimension and truncated to the rotated prefix."""
+        rope_params = config.rope_parameters[layer_type]
+        if rope_params["rope_type"] == "default":
+            return cls.compute_default_rope_parameters(config, device, layer_type=layer_type)
+        full_config = copy.copy(config)
+        full_config.rope_parameters = {
+            **config.rope_parameters,
+            layer_type: {**rope_params, "partial_rotary_factor": 1.0},
+        }
+        inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_params["rope_type"]](
+            full_config, device, layer_type=layer_type
+        )
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        rope_angles = int(rope_params.get("partial_rotary_factor", 1.0) * dim) // 2
+        inv_freq = torch.cat((inv_freq[:rope_angles], inv_freq.new_zeros(dim // 2 - rope_angles)))
+        return inv_freq, attention_factor
+
+
+class Lfm3MLP(nn.Module):
+    def __init__(self, config: Lfm3Config):
+        super().__init__()
+        self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies RoPE to the query and key tensors in float32, like the reference implementation."""
+    cos = cos.unsqueeze(unsqueeze_dim).float()
+    sin = sin.unsqueeze(unsqueeze_dim).float()
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class Lfm3Attention(nn.Module):
+    """Llama attention with QK RMSNorm, a sigmoid output gate per head, a sliding window on `sliding_attention`
+    layers and RoPE in float32."""
+
+    def __init__(self, config: Lfm3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = True
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
+        self.k_norm = Lfm3RMSNorm(self.head_dim, eps=config.norm_eps)
+        self.gate_proj = nn.Linear(config.hidden_size, config.num_attention_heads, bias=False)
+        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(*hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(*hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        gate = torch.sigmoid(self.gate_proj(hidden_states)).unsqueeze(-1)
+        attn_output = (attn_output * gate).reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+    """
+    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
+    """
+    # NOTE: attention mask is a 2D boolean tensor
+    if attention_mask is not None:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+    return hidden_states
+
+
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_update", "causal_conv1d")
+def causal_conv1d_update(
+    hidden_states: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    state_len = conv_state.shape[-1]
+
+    hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
+    conv_state.copy_(hidden_states_new[:, :, -state_len:])
+    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
+    out = out[:, :, -seq_len:]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernel_func_from_hub_with_fallback("causal_conv1d_fn", "causal_conv1d")
+def causal_conv1d_fn(
+    hidden_states: torch.Tensor,
+    weight: nn.Parameter,
+    bias: nn.Parameter | None = None,
+    activation: str | None = None,
+    **kwargs,
+):
+    _, hidden_size, seq_len = hidden_states.shape
+    padding = weight.shape[-1] - 1
+
+    out = F.conv1d(
+        hidden_states.to(weight.dtype),
+        weight=weight.unsqueeze(1),
+        bias=bias,
+        padding=padding,
+        groups=hidden_size,
+    )[:, :, :seq_len]
+    if activation is not None:
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+@use_kernelized_func([causal_conv1d_update, causal_conv1d_fn])
+class Lfm3ShortConv(nn.Module):
+    """LFM2 gated short convolution with packed-document boundaries (`seq_idx`), the padding mask re-applied after a
+    biased `in_proj`, and width-7 Triton inference kernels with a PyTorch fallback."""
+
+    def __init__(
+        self,
+        config: Lfm3Config,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.conv_kernel_size = config.conv_L_cache
+        self.bias = config.conv_bias
+
+        self.conv = nn.Conv1d(
+            in_channels=config.hidden_size,
+            out_channels=config.hidden_size,
+            kernel_size=self.conv_kernel_size,
+            groups=config.hidden_size,
+            bias=self.bias,
+            padding=self.conv_kernel_size - 1,
+        )
+        self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=self.bias)
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=self.bias)
+
+        self.layer_type = config.layer_types[layer_idx]
+
+    @force_accelerate_hooks("conv")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        past_key_values: Cache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        seq_idx: torch.IntTensor | None = None,
+    ):
+        seq_len = hidden_states.shape[1]
+
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+        BCx = self.in_proj(hidden_states).transpose(-1, -2)
+        B, C, x = BCx.chunk(3, dim=-2)
+        hidden_states = B * x
+        if attention_mask is not None:
+            # with `conv_bias` the projection makes the padded positions non-zero again
+            hidden_states = hidden_states * attention_mask[:, None, :].to(hidden_states.dtype)
+
+        use_precomputed_states = past_key_values is not None and past_key_values.has_previous_state(self.layer_idx)
+        use_kernels = 2 <= self.conv_kernel_size <= 4
+        use_wide_kernel = can_use_lfm3_conv(hidden_states, self.conv.weight)
+
+        if use_precomputed_states and seq_len == 1 and not past_key_values.layers[self.layer_idx].record_past:
+            conv_state = past_key_values.layers[self.layer_idx].conv_states[0]
+            if use_wide_kernel and conv_state.shape[-1] in (6, 7):
+                hidden_states = lfm3_causal_conv1d_update(
+                    hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
+                )
+            elif use_kernels:
+                hidden_states = causal_conv1d_update(
+                    hidden_states, conv_state, self.conv.weight.squeeze(1), self.conv.bias
+                )
+            else:
+                hidden_states = torch.cat([conv_state, hidden_states], dim=-1)
+                conv_state.copy_(hidden_states[:, :, -conv_state.shape[-1] :])
+                hidden_states = self._conv(hidden_states)[:, :, -seq_len:]
+        else:
+            if past_key_values is not None and seq_idx is not None:
+                # `seq_idx` is only set on an empty cache, and the decode steps read the cached window without it:
+                # only the last document (after any positional left padding) may enter the window
+                last_document = (seq_idx == seq_idx[:, -1:])[:, None].to(hidden_states.dtype)
+                past_key_values.update_conv_state(
+                    hidden_states * last_document, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                )
+            elif past_key_values is not None:
+                hidden_states = past_key_values.update_conv_state(
+                    hidden_states, self.layer_idx, conv_kernel_size=self.conv_kernel_size
+                )
+
+            if use_wide_kernel:
+                hidden_states = lfm3_causal_conv1d(hidden_states, self.conv.weight.squeeze(1), self.conv.bias, seq_idx)
+            elif seq_idx is not None:
+                hidden_states = self._segmented_conv(hidden_states, seq_idx)
+            elif use_kernels:
+                hidden_states = causal_conv1d_fn(hidden_states, self.conv.weight.squeeze(1), self.conv.bias)
+            else:
+                hidden_states = self._conv(hidden_states)
+
+            # Drop the additional previous states
+            if past_key_values is not None:
+                hidden_states = hidden_states[:, :, -seq_len:]
+
+        y = C * hidden_states
+        y = y.transpose(-1, -2).contiguous()
+        y = self.out_proj(y)
+        return y
+
+    def _conv(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Causal depthwise conv in torch: the `causal_conv1d` kernels only accept widths 2 to 4."""
+        conv_out = self.conv(hidden_states.to(self.conv.weight.dtype))[..., : hidden_states.shape[-1]]
+        return conv_out.to(hidden_states.dtype)
+
+    def _segmented_conv(self, hidden_states: torch.Tensor, seq_idx: torch.Tensor) -> torch.Tensor:
+        """Causal depthwise conv that only reads positions sharing the current token's document index."""
+        kernel_size = self.conv_kernel_size
+        seq_idx = F.pad(seq_idx, (hidden_states.shape[-1] - seq_idx.shape[-1] + kernel_size - 1, 0), value=-1)
+        windows = F.pad(hidden_states, (kernel_size - 1, 0)).unfold(-1, kernel_size, 1)
+        seq_idx_windows = seq_idx.unfold(-1, kernel_size, 1)
+        same_document = (seq_idx_windows == seq_idx_windows[..., -1:]).unsqueeze(1)
+        conv_out = torch.einsum("bdtk,dk->bdt", windows * same_document.to(windows.dtype), self.conv.weight.squeeze(1))
+        if self.bias:
+            conv_out = conv_out + self.conv.bias[None, :, None]
+        return conv_out.to(hidden_states.dtype)
+
+
+class Lfm3DecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Lfm3Config, layer_idx: int):
+        super().__init__()
+        self.is_attention_layer = config.layer_types[layer_idx] in config._attention_layer_types
+        self.use_sandwich_norm = config.use_sandwich_norm
+
+        if self.is_attention_layer:
+            self.self_attn = Lfm3Attention(config, layer_idx)
+        else:
+            self.conv = Lfm3ShortConv(config, layer_idx)
+        self.feed_forward = Lfm3MLP(config)
+        self.operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+        if self.use_sandwich_norm:
+            self.post_operator_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+            self.post_ffn_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        seq_idx: torch.IntTensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        if self.is_attention_layer:
+            hidden_states, _ = self.self_attn(
+                hidden_states=self.operator_norm(hidden_states),
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+        else:
+            hidden_states = self.conv(
+                hidden_states=self.operator_norm(hidden_states),
+                past_key_values=past_key_values,
+                attention_mask=attention_mask,
+                seq_idx=seq_idx,
+            )
+        if self.use_sandwich_norm:
+            hidden_states = self.post_operator_norm(hidden_states)
+        hidden_states = hidden_states + residual
+
+        ffn_out = self.feed_forward(self.ffn_norm(hidden_states))
+        if self.use_sandwich_norm:
+            ffn_out = self.post_ffn_norm(ffn_out)
+        hidden_states = hidden_states + ffn_out
+
+        return hidden_states
+
+
+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 Lfm3PreTrainedModel(PreTrainedModel):
+    config: Lfm3Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Lfm3DecoderLayer"]
+    _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
+    _can_record_outputs = {
+        "hidden_states": Lfm3DecoderLayer,
+        "attentions": Lfm3Attention,
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, Lfm3RotaryEmbedding):
+            for layer_type in module.layer_types:
+                curr_inv_freq, _ = module.compute_rope_parameters(module.config, layer_type=layer_type)
+                init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
+                init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
+
+    @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)
+
+
+def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
+    """Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
+    cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
+    restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
+    valid = padding_mask[:, -restarts.shape[-1] :].bool()
+    packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
+    if filled:
+        # a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
+        # against; a mask covering only the chunk keeps the strict rule
+        cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
+        cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
+        packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
+    return packed
+
+
+@auto_docstring
+class Lfm3Model(Lfm3PreTrainedModel):
+    """LFM2 model with sliding-attention layers as a third layer type: the mask and the rotary embedding are built per
+    layer type, and packed documents (`position_ids` restarting) reset the short-conv window through `seq_idx`."""
+
+    def __init__(self, config: Lfm3Config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Lfm3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.rotary_emb = Lfm3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.embedding_norm = Lfm3RMSNorm(config.hidden_size, eps=config.norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @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,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)` or `dict`, *optional*):
+            Mask to avoid performing attention on padding token indices (1 keeps a token, 0 masks it), or the mask
+            dict of `create_masks_for_generate`. A mask dict is used as given: `generate` refuses packed
+            `position_ids` before its prefill.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        # `seq_idx` is a shared kwarg (`DataCollatorWithFlattening(return_seq_idx=True)`, read by Lfm2); LFM3 derives
+        # the documents from the restarting `position_ids` that come with it
+        if kwargs.pop("seq_idx", None) is not None and position_ids is None:
+            raise ValueError("Pass packed documents as `position_ids` restarting at 0; LFM3 does not read `seq_idx`")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+        filled = past_key_values is not None and (
+            past_key_values.get_seq_length() > 0
+            if any(t != "conv" for t in self.config.layer_types)
+            else past_key_values.has_previous_state()
+        )
+        packed_cache = getattr(past_key_values, "lfm3_packed", None)
+        if isinstance(packed_cache, torch.Tensor) or packed_cache:
+            # Reset layers have no history to mix. Ordinary caches carry a Python `False` and skip this, so the decode
+            # steps do no tensor work here; only a compiled packed prefill leaves a tensor flag, and `filled` stays
+            # tensor-valued for static caches.
+            packed_message = (
+                "A cache filled by a packed-document forward cannot be continued: its attention spans the documents"
+            )
+            packed_cache = packed_cache & filled if isinstance(packed_cache, torch.Tensor) else filled
+            if isinstance(packed_cache, torch.Tensor):
+                torch_compilable_check(~packed_cache, packed_message)
+                if not is_tracing(packed_cache):
+                    past_key_values.lfm3_packed = bool(packed_cache)
+            elif packed_cache:
+                raise ValueError(packed_message)
+            else:
+                past_key_values.lfm3_packed = False
+
+        # The mask utilities detect packing only with `attention_mask=None`, so a multi-token forward with explicit
+        # positions drops the all-ones mask tokenizers emit (a data-dependent check, hence not when traced).
+        explicit_positions = position_ids is not None and inputs_embeds.shape[1] > 1
+        if (
+            explicit_positions
+            and isinstance(attention_mask, torch.Tensor)
+            and attention_mask.ndim == 2
+            and not is_tracing(position_ids)
+            and bool(attention_mask.all())
+        ):
+            attention_mask = None
+        seq_idx = None
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+        elif explicit_positions and not isinstance(attention_mask, dict):
+            # a mask dict comes from `generate`, which refused packing before its prefill
+            if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 2:
+                # data-dependent: a traced forward skips it (`generate` refuses packing eagerly before its prefill)
+                if not is_tracing(position_ids):
+                    torch_compilable_check(
+                        ~packed_under_padding(position_ids, attention_mask, filled),
+                        "Packed documents (`position_ids` restarting) cannot be combined with an attention mask: "
+                        "attention would not isolate the documents. Pass `attention_mask=None` without a static cache.",
+                    )
+            else:
+                restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
+                document_ids = restarts.cumsum(-1)
+                if is_tracing(position_ids) or bool((document_ids[..., -1] > 0).any()):
+                    seq_idx = document_ids.to(torch.int32)
+                # Packed documents are prefilled in one forward on an empty cache (cached tokens would join the first
+                # document), and their cache is never continued
+                if filled:
+                    torch_compilable_check(
+                        (~restarts[:, 1:]).all() & (position_ids[:, 0] != 0).all(),
+                        "Packed documents must be prefilled in one forward; the cache already holds tokens",
+                    )
+                    # Ordinary continuation must retain the convolution window from the cached document.
+                    seq_idx = None
+                if attention_mask is None:
+                    # without a mask only the cacheless packed-sequence mask isolates repeated zero positions (left
+                    # padding): a cache-sized mask and every later decode step would attend them, so they count too
+                    packed = (document_ids[..., -1] > 0).any()
+                else:
+                    # a prepared 4D mask is the caller's and carries the padding (repeated zero positions), but may
+                    # or may not isolate documents, so the cache records the packing and is never continued
+                    steps = torch.diff(position_ids, dim=-1)
+                    packed = ((steps != 1) & ~((steps == 0) & (position_ids[:, :-1] == 0))).any()
+                if past_key_values is not None:
+                    # resolved here, where the prefill already syncs, so the decode steps carry a Python flag: no
+                    # per-step host sync, data-dependent op or CUDA-graph output on the cache
+                    past_key_values.lfm3_packed = packed if is_tracing(position_ids) else bool(packed)
+
+        if not isinstance(causal_mask_mapping := attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "inputs_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            # The mask utilities only isolate packed documents without a cache, and an empty cache adds nothing.
+            layer_types = self.config.layer_types
+            if seq_idx is not None and past_key_values is not None:
+                attention_layer = next((i for i, t in enumerate(layer_types) if t != "conv"), None)
+                if attention_layer is not None and past_key_values.get_seq_length(attention_layer) == 0:
+                    if past_key_values.is_compileable:
+                        # static layers always return their full length, so the masks are sized by the cache, which
+                        # isolates neither documents nor positional-only left padding
+                        torch_compilable_check(
+                            ~packed,
+                            "Packed documents, or left padding without an attention mask, cannot be prefilled into a "
+                            "static cache: its masks are sized by the cache. Pass the padding as `attention_mask`, or "
+                            "use a `DynamicCache` for packed documents",
+                        )
+                    else:
+                        mask_kwargs["past_key_values"] = None
+            causal_mask_mapping = create_masks_by_layer_type(**mask_kwargs)
+
+        hidden_states = inputs_embeds
+        position_embeddings = {
+            layer_type: self.rotary_emb(hidden_states, position_ids=position_ids, layer_type=layer_type)
+            for layer_type in self.rotary_emb.layer_types
+        }
+
+        for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+            layer_type = self.config.layer_types[i]
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[layer_type],
+                position_embeddings=position_embeddings.get(layer_type),
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                seq_idx=seq_idx,
+                **kwargs,
+            )
+
+        hidden_states = self.embedding_norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
+    """`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
+    (`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
+    position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
+    cache = model_kwargs.get("past_key_values")
+    if cache is None or position_ids is None or position_ids.ndim != 2:
+        return
+    if attention_mask is not None and attention_mask.ndim != 2:
+        return
+    has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
+    filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
+    inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
+    length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
+    if generation_config.prefill_chunk_size is not None:
+        # the chunks run from the start of the prompt
+        position_ids = position_ids[:, :length]
+    else:
+        if filled and has_attention:
+            # the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
+            mask_length = getattr(generation_config, "_mask_length", -1)
+            if attention_mask is not None:
+                mask_length = attention_mask.shape[1]
+            if inputs_embeds is not None or mask_length == length:
+                length -= int(cache.get_seq_length())
+        position_ids = position_ids[:, position_ids.shape[1] - length :]
+    if attention_mask is None:
+        attention_mask = torch.ones_like(position_ids)
+    if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
+        raise ValueError(
+            "`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
+            "holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
+        )
+
+
+@auto_docstring
+class Lfm3ForCausalLM(Lfm3PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Lfm3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @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,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Lfm3ForCausalLM
+
+        >>> model = Lfm3ForCausalLM.from_pretrained("LiquidAI/LFM3-1.7B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("LiquidAI/LFM3-1.7B")
+
+        >>> inputs = tokenizer.apply_chat_template(
+        ...     [{"role": "user", "content": "What is C. elegans?"}], add_generation_prompt=True, return_tensors="pt"
+        ... )
+        >>> generate_ids = model.generate(**inputs, max_new_tokens=64)
+        >>> tokenizer.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        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 = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
+        # every prefill passes here once before its first forward: a check per chunk would leave the earlier chunks
+        # in the cache, and a dynamic cache would only refuse packing at the first decode step
+        reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
+        return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
+
+
+__all__ = ["Lfm3ForCausalLM", "Lfm3Model", "Lfm3PreTrainedModel"]

modeling_lfm3_omni.py

@@ -0,0 +1,1863 @@
+# 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
+from typing import TypedDict
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers import initialization as init
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.configuration_utils import PreTrainedConfig
+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,
+    TensorType,
+    TransformersKwargs,
+    auto_docstring,
+    can_return_tuple,
+    torch_compilable_check,
+)
+from transformers.utils.import_utils import is_tracing
+
+from .configuration_lfm3_omni import Lfm3OmniConfig, Lfm3OmniEncoderConfig
+from .modeling_lfm3 import Lfm3Model
+from .modeling_lfm3_vl import Lfm3VlVisionModel
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Lfm3Omni outputs, with hidden states and attentions.
+    """
+)
+@dataclass
+class Lfm3OmniModelOutputWithPast(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 Lfm3Omni causal language model (or autoregressive) outputs.
+    """
+)
+@dataclass
+class Lfm3OmniCausalLMOutputWithPast(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
+
+
+class Lfm3OmniEncoderStreamingCache:
+    """Bounded downsampler context and causal encoder state for synchronized audio chunks."""
+
+    def __init__(self, config: Lfm3OmniEncoderConfig):
+        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
+
+
+@auto_docstring(
+    custom_intro="""
+    Outputs of [`Lfm3OmniForConditionalGeneration.generate`].
+    """
+)
+@dataclass
+class Lfm3OmniGenerateOutput(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 [`Lfm3OmniForConditionalGeneration.prefill_audio`].")
+@dataclass
+class Lfm3OmniStreamingOutput(ModelOutput):
+    r"""
+    past_key_values (`Cache`, *optional*):
+        Language-model cache including all completed audio tokens. Updated in place by subsequent calls.
+    encoder_cache (`Lfm3OmniEncoderStreamingCache`, *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: Lfm3OmniEncoderStreamingCache | None = None
+    attention_mask: torch.LongTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+
+
+class Lfm3OmniLiveTranscriptionState:
+    """State of one live transcription: the encoder follows the whole stream, the language model one context."""
+
+    def __init__(self):
+        self.encoder_cache: Lfm3OmniEncoderStreamingCache | 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 [`Lfm3OmniForConditionalGeneration.transcribe_live`].")
+@dataclass
+class Lfm3OmniLiveTranscriptionOutput(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 ([`Lfm3OmniLiveTranscriptionState`]):
+        The stream state for the next call.
+    """
+
+    sequences: torch.LongTensor | None = None
+    num_audio_tokens: int | None = None
+    closed: bool | None = None
+    state: Lfm3OmniLiveTranscriptionState | None = None
+
+
+class Lfm3OmniRMSNorm(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 trained heads scale in fp32 and round once (liquid_lfm `RMSNorm`), 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 Lfm3OmniCodebookEmbedding(nn.Module):
+    """Input embedding and tied output head of one codebook of the depth decoder."""
+
+    def __init__(self, config: Lfm3OmniConfig):
+        super().__init__()
+        self.embedding = nn.Embedding(config.codebook_size, config.depth_config.hidden_size)
+        self.embedding_norm = Lfm3OmniRMSNorm(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 Lfm3OmniDepthDecoder(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: Lfm3OmniConfig):
+        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(Lfm3OmniCodebookEmbedding(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 Lfm3OmniPreTrainedModel(PreTrainedModel):
+    config: Lfm3OmniConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Lfm3OmniDepthDecoder"]
+    _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 = ("image", "video", "audio", "text")
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, Lfm3OmniDepthDecoder):
+            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)
+
+
+class Lfm3OmniMultiModalProjector(nn.Module):
+    """Merges each block of `downsample_factor**2` vision tokens into one language token."""
+
+    def __init__(self, config: Lfm3OmniConfig):
+        super().__init__()
+        self.downsample_factor = config.downsample_factor
+        self.layer_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=config.projector_layer_norm_eps)
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size * config.downsample_factor**2,
+            config.projector_hidden_size,
+            bias=config.projector_bias,
+        )
+        self.act = ACT2FN[config.projector_hidden_act]
+        self.linear_2 = nn.Linear(
+            config.projector_hidden_size, config.text_config.hidden_size, bias=config.projector_bias
+        )
+
+    def forward(self, image_features: torch.Tensor):
+        hidden_states = self.layer_norm(image_features).reshape(-1, self.linear_1.in_features)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class Lfm3OmniDownsampler(nn.Module):
+    """Three stride-2 convolutions over the (mel, time) plane: 128 mel bins and 8 frames become one vector."""
+
+    def __init__(self, config: Lfm3OmniEncoderConfig):
+        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 Lfm3OmniEncoder(nn.Module):
+    def __init__(self, config: Lfm3OmniEncoderConfig):
+        super().__init__()
+        self.config = config
+        self.downsample = Lfm3OmniDownsampler(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: Lfm3OmniEncoderStreamingCache | None = None,
+        is_final: bool = False,
+    ) -> tuple[torch.Tensor, Lfm3OmniEncoderStreamingCache]:
+        """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 = Lfm3OmniEncoderStreamingCache(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 Lfm3OmniProjector(nn.Module):
+    def __init__(self, config: Lfm3OmniConfig):
+        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))))
+
+
+@auto_docstring
+class Lfm3OmniModel(Lfm3OmniPreTrainedModel):
+    def __init__(self, config: Lfm3OmniConfig):
+        super().__init__(config)
+        self.language_model = Lfm3Model(config.text_config)
+        self.audio_encoder = Lfm3OmniEncoder(config.audio_encoder_config)
+        self.audio_projector = Lfm3OmniProjector(config)
+        self.codebook_embedding = nn.Embedding(
+            config.num_codebooks * config.codebook_size, config.text_config.hidden_size
+        )
+        self.depth_decoder = Lfm3OmniDepthDecoder(config)
+        self.vision_tower = Lfm3VlVisionModel(config.vision_config)
+        self.multi_modal_projector = Lfm3OmniMultiModalProjector(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: Lfm3OmniEncoderStreamingCache | None = None,
+        is_final: bool = False,
+    ) -> tuple[torch.Tensor, Lfm3OmniEncoderStreamingCache]:
+        """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,
+        pixel_values: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Lfm3OmniModelOutputWithPast:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened image patches. `patch_dim` is channels times temporal patch size times squared spatial patch size.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in image placeholder order.
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened video patches; the tower encodes every temporal patch as its own spatial sequence.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in video placeholder order.
+        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)
+        # vision placeholders follow the Qwen3-VL convention and are left to the count check
+        self.check_speech_placeholders(input_ids, inputs_embeds, input_features, audio_out_codes)
+
+        if (pixel_values is None) != (image_grid_thw is None):
+            raise ValueError("pixel_values and image_grid_thw must be supplied together")
+        if pixel_values is not None:
+            features = self.get_image_features(pixel_values, image_grid_thw).to(inputs_embeds)
+            mask = self.get_placeholder_mask(input_ids, inputs_embeds, self.config.image_token_id, features.shape[0])
+            inputs_embeds = inputs_embeds.masked_scatter(mask, features)
+        if (pixel_values_videos is None) != (video_grid_thw is None):
+            raise ValueError("pixel_values_videos and video_grid_thw must be supplied together")
+        if pixel_values_videos is not None:
+            features = self.get_video_features(pixel_values_videos, video_grid_thw).to(inputs_embeds)
+            mask = self.get_placeholder_mask(input_ids, inputs_embeds, self.config.video_token_id, features.shape[0])
+            inputs_embeds = inputs_embeds.masked_scatter(mask, features)
+
+        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 Lfm3OmniModelOutputWithPast(
+            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
+    def get_image_features(self, pixel_values: torch.Tensor, image_grid_thw: torch.LongTensor) -> torch.Tensor:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened image patches. `patch_dim` is channels times temporal patch size times squared spatial patch size.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in image placeholder order.
+
+        Returns:
+            `torch.Tensor` of shape `(num_image_tokens, hidden_size)`: Projected image features in placeholder order.
+        """
+        vision = self.vision_tower(pixel_values, grid_thw=image_grid_thw, return_dict=True)
+        return self.multi_modal_projector(vision.last_hidden_state)
+
+    @auto_docstring
+    def get_video_features(self, pixel_values_videos: torch.Tensor, video_grid_thw: torch.LongTensor) -> torch.Tensor:
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened video patches; the tower encodes every temporal patch as its own spatial sequence.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in video placeholder order.
+
+        Returns:
+            `torch.Tensor` of shape `(num_video_tokens, hidden_size)`: Projected video features in placeholder order.
+        """
+        return self.get_image_features(pixel_values_videos, video_grid_thw)
+
+
+class Lfm3OmniGenerationMode(str, Enum):
+    TEXT = "text"
+    AUDIO = "audio"
+    INTERLEAVED = "interleaved"
+
+
+@auto_docstring(
+    custom_intro="""
+    The LFM3-Omni model: an LFM3 language model that reads speech through the Liquid speech encoder, reads images
+    through the LFM3-VL vision tower and writes speech as codebook frames through a depth decoder, alongside text.
+    """
+)
+class Lfm3OmniForConditionalGeneration(Lfm3OmniPreTrainedModel, 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 = {}
+
+    # `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,
+    }
+
+    def __init__(self, config: Lfm3OmniConfig):
+        super().__init__(config)
+        self.model = Lfm3OmniModel(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]:
+        return [
+            self.config.audio_token_id,
+            self.config.audio_out_token_id,
+            self.config.image_token_id,
+            self.config.video_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,
+        pixel_values: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Lfm3OmniCausalLMOutputWithPast:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened image patches. `patch_dim` is channels times temporal patch size times squared spatial patch size.
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened video patches; the tower encodes every temporal patch as its own spatial sequence.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in video placeholder order.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in image placeholder order.
+        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,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+            pixel_values_videos=pixel_values_videos,
+            video_grid_thw=video_grid_thw,
+            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 Lfm3OmniCausalLMOutputWithPast(
+            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
+        # like the training reference
+        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: Lfm3OmniEncoderStreamingCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        is_final: bool = False,
+    ) -> Lfm3OmniStreamingOutput:
+        """
+        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 (`Lfm3OmniEncoderStreamingCache`, *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:
+            [`Lfm3OmniStreamingOutput`]: 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 Lfm3OmniStreamingOutput(
+            past_key_values=past_key_values,
+            encoder_cache=encoder_cache,
+            attention_mask=attention_mask,
+            last_hidden_state=hidden_states,
+        )
+
+    def _live_forward(self, state: Lfm3OmniLiveTranscriptionState, 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: Lfm3OmniLiveTranscriptionState | 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,
+    ) -> Lfm3OmniLiveTranscriptionOutput:
+        """
+        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 ([`Lfm3OmniLiveTranscriptionState`], *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:
+            [`Lfm3OmniLiveTranscriptionOutput`]: 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 = Lfm3OmniLiveTranscriptionState()
+        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 Lfm3OmniLiveTranscriptionOutput(
+            sequences=torch.tensor([generated], dtype=torch.long, device=input_features.device),
+            num_audio_tokens=num_audio_tokens,
+            closed=closed,
+            state=state,
+        )
+
+    @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: Lfm3OmniGenerationMode | 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 = 64,
+        text_repetition_window: int = 32,
+        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,
+    ) -> Lfm3OmniGenerateOutput | 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 [`Lfm3OmniGenerateOutput`]. `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.
+            audio_temperature (`float`, *optional*):
+                The sampling temperature of the speech codes; `None` picks the most likely code.
+            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*, defaults to 64):
+                The number of past frames the audio repetition penalty covers.
+            text_repetition_window (`int`, *optional*, defaults to 32):
+                The number of past tokens the text repetition penalty covers.
+            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: Lfm3OmniGenerationMode | 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 = 64,
+        text_repetition_window: int = 32,
+        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,
+    ) -> Lfm3OmniGenerateOutput | 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"
+            )
+        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
+        ]
+        if generation_mode is None:
+            modes = [
+                Lfm3OmniGenerationMode.AUDIO
+                if token in (config.audio_start_token_id, config.audio_out_token_id)
+                else Lfm3OmniGenerationMode.INTERLEAVED
+                if token == config.mixed_start_token_id
+                else Lfm3OmniGenerationMode.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 = [Lfm3OmniGenerationMode(mode) for mode in modes]
+            # speech always follows its marker in training; append the one the requested mode needs
+            markers = {
+                Lfm3OmniGenerationMode.AUDIO: (config.audio_start_token_id, config.audio_out_token_id),
+                Lfm3OmniGenerationMode.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] = Lfm3OmniGenerationMode.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 Lfm3OmniGenerationMode.INTERLEAVED for mode in modes]
+            audio_rows = [
+                row
+                for row in range(batch_size)
+                if not finished[row]
+                and (modes[row] is Lfm3OmniGenerationMode.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] = Lfm3OmniGenerationMode.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,
+                        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] = Lfm3OmniGenerationMode.AUDIO
+                    elif token_id == config.mixed_start_token_id:
+                        modes[row] = Lfm3OmniGenerationMode.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 = Lfm3OmniGenerateOutput(
+            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
+
+    @auto_docstring
+    def get_image_features(self, pixel_values: torch.Tensor, image_grid_thw: torch.LongTensor):
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened image patches. `patch_dim` is channels times temporal patch size times squared spatial patch size.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in image placeholder order.
+
+        Returns:
+            `torch.Tensor` of shape `(num_image_tokens, hidden_size)`: Projected image features in placeholder order.
+        """
+        return self.model.get_image_features(pixel_values, image_grid_thw)
+
+    @auto_docstring
+    def get_video_features(self, pixel_values_videos: torch.Tensor, video_grid_thw: torch.LongTensor):
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, patch_dim)`, *optional*):
+            Flattened video patches; the tower encodes every temporal patch as its own spatial sequence.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            Temporal, height, and width patch counts before spatial merging, in video placeholder order.
+
+        Returns:
+            `torch.Tensor` of shape `(num_video_tokens, hidden_size)`: Projected video features in placeholder order.
+        """
+        return self.model.get_video_features(pixel_values_videos, video_grid_thw)
+
+
+@dataclass
+class Lfm3OmniFeatureStreamingCache:
+    """Unconsumed mono samples and frame count for one input stream."""
+
+    samples: torch.Tensor | None = None
+    num_samples: int = 0
+    num_frames: int = 0
+    finished: bool = False
+
+
+class Lfm3OmniFeatureExtractorKwargs(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
+
+
+__all__ = [
+    "Lfm3OmniFeatureStreamingCache",
+    "Lfm3OmniEncoderStreamingCache",
+    "Lfm3OmniLiveTranscriptionState",
+    "Lfm3OmniLiveTranscriptionOutput",
+    "Lfm3OmniModel",
+    "Lfm3OmniPreTrainedModel",
+    "Lfm3OmniForConditionalGeneration",
+    "Lfm3OmniGenerateOutput",
+    "Lfm3OmniStreamingOutput",
+]

modeling_lfm3_vl.py

@@ -0,0 +1,965 @@
+# 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.
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_recurrent_attention_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    TransformersKwargs,
+    auto_docstring,
+    can_return_tuple,
+    torch_compilable_check,
+)
+from transformers.utils.generic import (
+    accepts_precomputed_kwargs,
+    is_flash_attention_requested,
+    maybe_autocast,
+    merge_with_config_defaults,
+)
+from transformers.utils.output_capturing import OutputRecorder, capture_outputs
+from transformers.vision_utils import get_max_seqlen, get_vision_attention_seqlens, get_vision_position_ids
+
+from .configuration_lfm3_vl import Lfm3VlConfig, Lfm3VlVisionConfig
+from .modeling_lfm3 import Lfm3Model
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Lfm3VlRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        Lfm3VlRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        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 self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Lfm3VlMLP(nn.Module):
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__()
+        self.w1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w3 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.w2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class Lfm3VlVisionRotaryEmbedding(nn.Module):
+    """
+    Simple axial 2D rope with same freqs used for H and W grids. The freqs are
+    pre-computed using `head-dim//4` which is later used to concat H and W positions.
+    The final angles rotate over the whole head dim, no partial rotation involved.
+    """
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__()
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_axial_rope_parameters
+        if self.rope_type != "axial":
+            raise ValueError(f"{self.__class__.__name__} supports only axial rope, but requested {self.rope_type}")
+        inv_freq, self.attention_scaling = rope_init_fn(self.config)
+
+        self.inv_freq = nn.Buffer(inv_freq, persistent=False)
+        self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_axial_rope_parameters(config: Lfm3VlVisionConfig, **kwargs) -> tuple[torch.Tensor, float]:
+        """Inverse frequencies of the axial RoPE over half the head dimension per axis, with an attention scaling of 1."""
+        base = config.rope_parameters["rope_theta"]
+        spatial_dim = config.hidden_size // config.num_attention_heads // 2
+        inv_freq = 1.0 / (base ** (torch.arange(0, spatial_dim, 2, dtype=torch.float) / spatial_dim))
+        return inv_freq, 1.0
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        # position_ids: (2, N) — row 0 = h coords, row 1 = w coords
+        position_ids_expanded = position_ids[..., None].float()
+        device_type = x.device.type if isinstance(x.device.type, str) else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):
+            freqs = position_ids_expanded * self.inv_freq.float()
+            cos = freqs.cos() * self.attention_scaling
+            sin = freqs.sin() * self.attention_scaling
+
+        cos = self.recomposition_frequencies(cos)
+        sin = self.recomposition_frequencies(sin)
+        return cos, sin
+
+    def recomposition_frequencies(self, freq):
+        """
+        Recompose the frequencies into the final spatial layout used per each grid.
+        """
+        freq_h, freq_w = freq[:, 0], freq[:, 1]
+        freq_hw = torch.cat([freq_h, freq_w], dim=-1)
+        return torch.cat([freq_hw, freq_hw], dim=-1)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    orig_q_dtype = q.dtype
+    orig_k_dtype = k.dtype
+    q, k = q.float(), k.float()
+    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    q_embed = q_embed.to(orig_q_dtype)
+    k_embed = k_embed.to(orig_k_dtype)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class Lfm3VlVisionAttention(nn.Module):
+    """Bidirectional grouped-query attention over the patches of one frame, with RMSNorm on the query and key heads."""
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__()
+        self.config = config
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = 0.0
+        self.is_causal = False
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.out_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
+        self.k_layernorm = Lfm3VlRMSNorm(self.head_dim, eps=config.norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        max_seqlen: int | None = None,
+        frame_shapes: list[list[int]] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        seq_length = hidden_states.shape[0]
+        query_states = self.q_layernorm(self.q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim))
+        key_states = self.k_layernorm(
+            self.k_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
+        )
+        value_states = self.v_proj(hidden_states).view(seq_length, self.num_key_value_heads, self.head_dim)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+        query_states = query_states.transpose(0, 1).unsqueeze(0)
+        key_states = key_states.transpose(0, 1).unsqueeze(0)
+        value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        if is_flash_attention_requested(self.config):
+            max_seqlen = get_max_seqlen(cu_seqlens, self.config, kwargs={"max_seqlen": max_seqlen})
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask=None,
+                scaling=self.scaling,
+                dropout=0.0,
+                cu_seq_lens_q=cu_seqlens,
+                cu_seq_lens_k=cu_seqlens,
+                max_length_q=max_seqlen,
+                max_length_k=max_seqlen,
+                is_causal=False,
+                **kwargs,
+            )
+        else:
+            # Other implementations: frames of one shape (every video, equal-size images) attend as one batch, mixed
+            # shapes frame by frame. `frame_shapes` gives the lengths without a device sync.
+            if frame_shapes is not None:
+                lengths = [height * width for height, width in frame_shapes]
+            else:
+                lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+            if len(set(lengths)) == 1:
+                query_states, key_states, value_states = (
+                    tensor.reshape(-1, len(lengths), lengths[0], self.head_dim).transpose(0, 1)
+                    for tensor in (query_states, key_states, value_states)
+                )
+                attn_output, attn_weights = attention_interface(
+                    self,
+                    query_states,
+                    key_states,
+                    value_states,
+                    attention_mask=None,
+                    scaling=self.scaling,
+                    dropout=0.0,
+                    is_causal=False,
+                    **kwargs,
+                )
+            else:
+                splits = [torch.split(tensor, lengths, dim=2) for tensor in (query_states, key_states, value_states)]
+                attn_outputs, attn_weights = [], []
+                for q, k, v in zip(*splits):
+                    output, weights = attention_interface(
+                        self,
+                        q,
+                        k,
+                        v,
+                        attention_mask=None,
+                        scaling=self.scaling,
+                        dropout=0.0,
+                        is_causal=False,
+                        **kwargs,
+                    )
+                    attn_outputs.append(output)
+                    attn_weights.append(weights)
+                attn_output = torch.cat(attn_outputs, dim=1)
+                # frames of different shapes have differently sized maps: one entry per frame, or None under sdpa
+                attn_weights = None if attn_weights[0] is None else tuple(attn_weights)
+
+        attn_output = attn_output.reshape(seq_length, -1).contiguous()
+        attn_output = self.out_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Lfm3VlVisionConv2d(nn.Module):
+    """Gated depthwise 2D convolution over the patch grid of each frame."""
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__()
+        self.spatial_merge_size = config.spatial_merge_size
+        self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
+        self.spatial_conv = nn.Conv2d(
+            config.hidden_size,
+            config.hidden_size,
+            kernel_size=config.conv_kernel_size,
+            padding=config.conv_kernel_size // 2,
+            groups=config.hidden_size,
+            bias=False,
+        )
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor, frame_shapes: list[list[int]]):
+        B, C, x = self.in_proj(hidden_states).chunk(3, dim=-1)
+        gated = B * x
+        merge = self.spatial_merge_size
+        num_frames = len(frame_shapes)
+        if len(set(map(tuple, frame_shapes))) == 1:
+            # frames of one shape (every video, equal-size images) convolve as one batch: merge-block order -> raster
+            # grid, convolve, back
+            height, width = frame_shapes[0]
+            frames = gated.view(num_frames, height // merge, width // merge, merge, merge, -1)
+            frames = frames.permute(0, 5, 1, 3, 2, 4).reshape(num_frames, -1, height, width)
+            frames = self.spatial_conv(frames)
+            frames = frames.view(num_frames, -1, height // merge, merge, width // merge, merge).permute(
+                0, 2, 4, 3, 5, 1
+            )
+            conv_output = frames.reshape(gated.shape[0], -1)
+        else:
+            conv_outputs = []
+            lengths = [height * width for height, width in frame_shapes]
+            for frame, (height, width) in zip(torch.split(gated, lengths, dim=0), frame_shapes):
+                frame = frame.view(height // merge, width // merge, merge, merge, -1)
+                frame = frame.permute(4, 0, 2, 1, 3).reshape(1, -1, height, width)
+                frame = self.spatial_conv(frame)
+                frame = frame.view(-1, height // merge, merge, width // merge, merge).permute(1, 3, 2, 4, 0)
+                conv_outputs.append(frame.reshape(height * width, -1))
+            conv_output = torch.cat(conv_outputs, dim=0)
+        return self.out_proj(C * conv_output)
+
+
+class Lfm3VlVisionLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Lfm3VlVisionConfig, layer_idx: int):
+        super().__init__()
+        self.is_attention_layer = config.layer_types[layer_idx] == "full_attention"
+        if self.is_attention_layer:
+            self.self_attn = Lfm3VlVisionAttention(config)
+        else:
+            self.conv = Lfm3VlVisionConv2d(config)
+        self.feed_forward = Lfm3VlMLP(config)
+        self.operator_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.ffn_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        frame_shapes: list[list[int]],
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if self.is_attention_layer:
+            operator_output, _ = self.self_attn(
+                self.operator_norm(hidden_states),
+                cu_seqlens=cu_seqlens,
+                position_embeddings=position_embeddings,
+                frame_shapes=frame_shapes,
+                **kwargs,
+            )
+        else:
+            operator_output = self.conv(self.operator_norm(hidden_states), frame_shapes)
+        hidden_states = hidden_states + operator_output
+        hidden_states = hidden_states + self.feed_forward(self.ffn_norm(hidden_states))
+        return hidden_states
+
+
+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 Lfm3VlPreTrainedModel(PreTrainedModel):
+    config: Lfm3VlConfig
+    base_model_prefix = "model"
+    input_modalities = ("image", "video", "text")
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = ["past_key_values"]
+
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    # the cached decode step traces without graph breaks, like the text model
+    _can_compile_fullgraph = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _no_split_modules = ["Lfm3VlVisionLayer"]
+
+    @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:
+        # the static-cache masks of the text model are sized by layer type, not by the generic first attention layer
+        return create_masks_by_layer_type(config, inputs_embeds, attention_mask, past_key_values, position_ids)
+
+
+def get_vision_frame_shapes(grid_thw: torch.Tensor) -> list[list[int]]:
+    return grid_thw[:, 1:].repeat_interleave(grid_thw[:, 0], dim=0).tolist()
+
+
+@auto_docstring(
+    custom_intro="""
+    The LFM3-VL vision encoder: a stack of gated convolution and bidirectional attention layers over the patches of
+    each image or video frame, embedded by a linear patchifier.
+    """
+)
+class Lfm3VlVisionModel(Lfm3VlPreTrainedModel):
+    config: Lfm3VlVisionConfig
+    input_modalities = ("image", "video")
+    _input_embed_layer = "patchifier"
+    _can_record_outputs = {
+        "hidden_states": Lfm3VlVisionLayer,
+        "attentions": OutputRecorder(Lfm3VlVisionAttention, index=1),
+    }
+
+    def __init__(self, config: Lfm3VlVisionConfig):
+        super().__init__(config)
+        self.spatial_merge_size = config.spatial_merge_size
+        patch_dim = config.in_channels * config.temporal_patch_size * config.patch_size**2
+        self.patchifier = nn.Linear(patch_dim, config.hidden_size, bias=False)
+        self.embedding_norm = Lfm3VlRMSNorm(config.hidden_size, eps=config.norm_eps)
+        self.rotary_pos_emb = Lfm3VlVisionRotaryEmbedding(config)
+        self.layers = nn.ModuleList(
+            [Lfm3VlVisionLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+    ) -> BaseModelOutputWithPooling:
+        r"""
+        pixel_values (`torch.Tensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every image or video, grouped by merge block.
+        grid_thw (`torch.LongTensor` of shape `(num_images_or_videos, 3)`):
+            The temporal, height and width dimensions of the patch grid of each image or video.
+        """
+        position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size, kwargs=kwargs)
+        cu_seqlens, max_seqlen = get_vision_attention_seqlens(grid_thw, self.config, kwargs=kwargs)
+        frame_shapes = get_vision_frame_shapes(grid_thw)
+
+        hidden_states = self.patchifier(pixel_values.to(self.patchifier.weight.dtype))
+        hidden_states = self.embedding_norm(hidden_states)
+        position_embeddings = self.rotary_pos_emb(hidden_states, position_ids)
+
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                frame_shapes=frame_shapes,
+                max_seqlen=max_seqlen,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+
+class Lfm3VlMultiModalProjector(nn.Module):
+    """Merges each block of `downsample_factor**2` vision tokens into one language token."""
+
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__()
+        self.downsample_factor = config.downsample_factor
+        self.layer_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=config.projector_layer_norm_eps)
+        self.linear_1 = nn.Linear(
+            config.vision_config.hidden_size * config.downsample_factor**2,
+            config.projector_hidden_size,
+            bias=config.projector_bias,
+        )
+        self.act = ACT2FN[config.projector_hidden_act]
+        self.linear_2 = nn.Linear(
+            config.projector_hidden_size, config.text_config.hidden_size, bias=config.projector_bias
+        )
+
+    def forward(self, image_features: torch.Tensor):
+        hidden_states = self.layer_norm(image_features).reshape(-1, self.linear_1.in_features)
+        hidden_states = self.linear_1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Lfm3Vl outputs, with hidden states and attentions.
+    """
+)
+@dataclass
+class Lfm3VlModelOutputWithPast(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.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    video_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
+        video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    image_hidden_states: torch.FloatTensor | None = None
+    video_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Lfm3Vl causal language model (or autoregressive) outputs.
+    """
+)
+@dataclass
+class Lfm3VlCausalLMOutputWithPast(CausalLMOutputWithPast):
+    r"""
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_image_tokens, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    video_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(num_video_tokens, hidden_size)`.
+        video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    image_hidden_states: torch.FloatTensor | None = None
+    video_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    The Lfm3Vl model which consists of a vision backbone and a language model, without a language modeling head.
+    """
+)
+class Lfm3VlModel(Lfm3VlPreTrainedModel):
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__(config)
+        self.vision_tower = Lfm3VlVisionModel(config.vision_config)
+        self.multi_modal_projector = Lfm3VlMultiModalProjector(config)
+        self.language_model = Lfm3Model(config.text_config)
+        self.post_init()
+
+    @accepts_precomputed_kwargs(modality="image")
+    @can_return_tuple
+    @auto_docstring(
+        custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
+    )
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        image_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
+            The temporal, height and width dimensions of the patch grid of each image.
+        """
+        image_outputs = self.vision_tower(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs)
+        image_features = self.multi_modal_projector(image_outputs.last_hidden_state)
+        split_sizes = (image_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2).tolist()
+        image_outputs.pooler_output = torch.split(image_features, split_sizes)
+        return image_outputs
+
+    def get_placeholder_mask(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: torch.FloatTensor,
+        image_features: torch.FloatTensor | None = None,
+        video_features: torch.FloatTensor | None = None,
+    ):
+        """Placeholder masks from `input_ids` or `inputs_embeds`, checked against the feature counts."""
+        if input_ids is None:
+            special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_image_mask = special_image_mask.all(-1)
+            special_video_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.full((), self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_video_mask = special_video_mask.all(-1)
+        else:
+            special_image_mask = input_ids == self.config.image_token_id
+            special_video_mask = input_ids == self.config.video_token_id
+
+        n_image_tokens = special_image_mask.sum()
+        special_image_mask = special_image_mask.unsqueeze(-1).to(inputs_embeds.device)
+        if image_features is not None:
+            torch_compilable_check(
+                n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(),
+                f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
+            )
+
+        n_video_tokens = special_video_mask.sum()
+        special_video_mask = special_video_mask.unsqueeze(-1).to(inputs_embeds.device)
+        if video_features is not None:
+            torch_compilable_check(
+                n_video_tokens * inputs_embeds.shape[-1] == video_features.numel(),
+                f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
+            )
+        return special_image_mask, special_video_mask
+
+    @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,
+        pixel_values: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Lfm3VlModelOutputWithPast:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each image.
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each video.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if (pixel_values is not None or pixel_values_videos is not None) and mm_encoder_outputs is not None:
+            raise ValueError(
+                "You cannot specify both pixel_values/pixel_values_videos and mm_encoder_outputs at the same time"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        mm_encoder_outputs = mm_encoder_outputs if mm_encoder_outputs is not None else {}
+        if (pixel_values is None) != (image_grid_thw is None):
+            raise ValueError("pixel_values and image_grid_thw must be supplied together")
+        if (pixel_values_videos is None) != (video_grid_thw is None):
+            raise ValueError("pixel_values_videos and video_grid_thw must be supplied together")
+        if mm_encoder_outputs.get("image") is None and pixel_values is not None:
+            mm_encoder_outputs["image"] = self.get_image_features(
+                pixel_values, image_grid_thw, return_dict=True, **kwargs
+            )
+        if mm_encoder_outputs.get("video") is None and pixel_values_videos is not None:
+            mm_encoder_outputs["video"] = self.get_video_features(
+                pixel_values_videos, video_grid_thw, return_dict=True, **kwargs
+            )
+
+        image_features = None
+        if mm_encoder_outputs.get("image") is not None:
+            image_features = torch.cat(mm_encoder_outputs["image"].pooler_output, dim=0).to(
+                inputs_embeds.device, inputs_embeds.dtype
+            )
+            image_mask, _ = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features)
+
+        video_features = None
+        if mm_encoder_outputs.get("video") is not None:
+            video_features = torch.cat(mm_encoder_outputs["video"].pooler_output, dim=0).to(
+                inputs_embeds.device, inputs_embeds.dtype
+            )
+            _, video_mask = self.get_placeholder_mask(
+                input_ids, inputs_embeds=inputs_embeds, video_features=video_features
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_features)
+
+        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 Lfm3VlModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_features,
+            video_hidden_states=video_features,
+        )
+
+    @accepts_precomputed_kwargs(modality="video")
+    @can_return_tuple
+    @auto_docstring(
+        custom_intro="Obtains video last hidden states from the vision tower and apply multimodal projection."
+    )
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
+            The temporal, height and width dimensions of the patch grid of each video.
+        """
+        return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+
+def reject_packed_generation(config: PreTrainedConfig, input_ids, generation_config, model_kwargs, is_first_iteration):
+    """`generate` decodes from the cache its prefill fills, and no mask it builds isolates packed documents
+    (`position_ids` restarting): the prompt is refused before the prefill, one forward or every chunk, writes to it."""
+    position_ids, attention_mask = model_kwargs.get("position_ids"), model_kwargs.get("attention_mask")
+    cache = model_kwargs.get("past_key_values")
+    if cache is None or position_ids is None or position_ids.ndim != 2:
+        return
+    if attention_mask is not None and attention_mask.ndim != 2:
+        return
+    has_attention = any(t != "conv" for t in config.get_text_config().layer_types)
+    filled = bool(cache.get_seq_length() > 0 if has_attention else cache.has_previous_state())
+    inputs_embeds = model_kwargs.get("inputs_embeds") if is_first_iteration else None
+    length = input_ids.shape[1] if inputs_embeds is None else inputs_embeds.shape[1]
+    if generation_config.prefill_chunk_size is not None:
+        # the chunks run from the start of the prompt
+        position_ids = position_ids[:, :length]
+    else:
+        if filled and has_attention:
+            # the slice `_prefill` feeds: a prompt that spans the cache is cut to its uncached part
+            mask_length = getattr(generation_config, "_mask_length", -1)
+            if attention_mask is not None:
+                mask_length = attention_mask.shape[1]
+            if inputs_embeds is not None or mask_length == length:
+                length -= int(cache.get_seq_length())
+        position_ids = position_ids[:, position_ids.shape[1] - length :]
+    if attention_mask is None:
+        attention_mask = torch.ones_like(position_ids)
+    if bool(packed_under_padding(position_ids, attention_mask, filled).any()):
+        raise ValueError(
+            "`generate` cannot decode packed documents (`position_ids` restarting, or starting at 0 on a cache that "
+            "holds tokens): attention would span them. Generate each document in its own row of a left-padded batch."
+        )
+
+
+def packed_under_padding(position_ids: torch.Tensor, padding_mask: torch.Tensor, filled) -> torch.Tensor:
+    """Padding repeats position 0, so packing is a restart between two kept tokens, or a kept position 0 on a filled
+    cache whose row holds a kept token; attention under a padding mask cannot isolate it."""
+    restarts = torch.diff(position_ids, prepend=position_ids[:, :1] - 1, dim=-1) != 1
+    valid = padding_mask[:, -restarts.shape[-1] :].bool()
+    packed = restarts[:, 1:] & valid[:, 1:] & valid[:, :-1]
+    if filled:
+        # a row whose cached tokens are all padding (chunked prefill of a left-padded batch) has no document to pack
+        # against; a mask covering only the chunk keeps the strict rule
+        cached = padding_mask[:, : padding_mask.shape[-1] - restarts.shape[-1]]
+        cached_valid = cached.bool().any(-1, keepdim=True) if cached.shape[-1] else torch.ones_like(valid[:, :1])
+        packed = packed | (valid[:, :1] & (position_ids[:, :1] == 0) & cached_valid)
+    return packed
+
+
+@auto_docstring(
+    custom_intro="""
+    The LFM3_VL model which consists of a vision backbone and a language model.
+    """
+)
+class Lfm3VlForConditionalGeneration(Lfm3VlPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+
+    def __init__(self, config: Lfm3VlConfig):
+        super().__init__(config)
+        self.model = Lfm3VlModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    @auto_docstring
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        image_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
+            The temporal, height and width dimensions of the patch grid of each image.
+        """
+        return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
+
+    @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,
+        pixel_values: torch.FloatTensor | None = None,
+        image_grid_thw: torch.LongTensor | None = None,
+        pixel_values_videos: torch.FloatTensor | None = None,
+        video_grid_thw: torch.LongTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        mm_encoder_outputs: dict[str, BaseModelOutputWithPooling] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | Lfm3VlCausalLMOutputWithPast:
+        r"""
+        pixel_values (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every image, grouped by merge block.
+        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each image.
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`, *optional*):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+            The temporal, height and width dimensions of the patch grid of each video.
+        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]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoProcessor, AutoModelForImageTextToText
+
+        >>> model = AutoModelForImageTextToText.from_pretrained("LiquidAI/LFM3-VL-1.9B")
+        >>> processor = AutoProcessor.from_pretrained("LiquidAI/LFM3-VL-1.9B")
+
+        >>> conversation = [
+        ...     {
+        ...         "role": "user",
+        ...         "content": [
+        ...             {"type": "image", "url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+        ...             {"type": "text", "text": "What is in this image?"},
+        ...         ],
+        ...     },
+        ... ]
+
+        >>> inputs = processor.apply_chat_template(
+        ...     conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
+        ... )
+
+        >>> outputs = model.generate(**inputs, max_new_tokens=45)
+        >>> processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        ```"""
+        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,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+            pixel_values_videos=pixel_values_videos,
+            video_grid_thw=video_grid_thw,
+            use_cache=use_cache,
+            mm_encoder_outputs=mm_encoder_outputs,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        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 = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+            )
+
+        return Lfm3VlCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=outputs.image_hidden_states,
+            video_hidden_states=outputs.video_hidden_states,
+        )
+
+    def _prefill(self, input_ids, generation_config, model_kwargs, is_first_iteration=True):
+        # the text model's refusal, before the prefill writes to the cache
+        reject_packed_generation(self.config, input_ids, generation_config, model_kwargs, is_first_iteration)
+        return super()._prefill(input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration)
+
+    @auto_docstring
+    def get_video_features(
+        self,
+        pixel_values_videos: torch.FloatTensor,
+        video_grid_thw: torch.LongTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPooling:
+        r"""
+        pixel_values_videos (`torch.FloatTensor` of shape `(num_patches, in_channels * temporal_patch_size * patch_size**2)`):
+            The flattened patches of every video, grouped by merge block.
+        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`):
+            The temporal, height and width dimensions of the patch grid of each video.
+        """
+        return self.model.get_video_features(pixel_values_videos, video_grid_thw, **kwargs)
+
+
+__all__ = ["Lfm3VlForConditionalGeneration", "Lfm3VlModel", "Lfm3VlPreTrainedModel", "Lfm3VlVisionModel"]

modeling_liquid_codec.py

@@ -0,0 +1,381 @@
+# 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,
+)

processing_lfm3_omni.py

@@ -0,0 +1,769 @@
+# 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 math
+import os
+import re
+
+import numpy as np
+import torch
+
+from transformers.audio_utils import AudioInput, make_list_of_audio
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import is_pil_image, is_valid_image, make_flat_list_of_images
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, VideosKwargs
+from transformers.utils import auto_docstring, is_numpy_array, is_torch_tensor
+from transformers.video_utils import is_valid_video, make_batched_metadata, make_batched_videos
+
+from .modeling_lfm3_omni import Lfm3OmniFeatureExtractorKwargs
+from .modeling_liquid_codec import LiquidCodecModel
+
+
+class Lfm3OmniVideosKwargs(VideosKwargs, total=False):
+    """
+    use_audio_in_video (`bool`, *optional*, defaults to `False`):
+        Whether each video comes with its soundtrack: the first audio clip after the video in the text (the one
+        `apply_chat_template(load_audio_from_video=True)` puts right behind the video) is interleaved after the
+        video's frame groups instead of standing in its own media span.
+    """
+
+    use_audio_in_video: bool
+
+
+class Lfm3OmniProcessorKwargs(ProcessingKwargs, total=False):
+    audio_kwargs: Lfm3OmniFeatureExtractorKwargs
+    videos_kwargs: Lfm3OmniVideosKwargs
+    _defaults = {"videos_kwargs": {"return_metadata": True, "use_audio_in_video": False}}
+
+
+# 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 -%}"
+)
+
+
+def _media_per_prompt(media, is_item) -> list[int] | None:
+    # rows name the prompt that owns them while the expansion consumes media in flat order; a row is a list of whole
+    # items or one stacked batch, and a list of single frames is one video, not a row
+    if not isinstance(media, (list, tuple)):
+        return None
+    counts = []
+    for row in media:
+        if isinstance(row, (list, tuple)) and all(is_item(item) for item in row):
+            counts.append(len(row))
+        elif (is_numpy_array(row) or is_torch_tensor(row)) and row.ndim == 4 and (len(row) == 0 or is_item(row[0])):
+            counts.append(len(row))
+        else:
+            return None
+    return counts
+
+
+def _check_vision_placeholder_counts(processor, prompts, images, videos) -> None:
+    for name, token, media, flatten, is_item in (
+        (
+            "images",
+            processor.image_token,
+            images,
+            make_flat_list_of_images,
+            lambda item: is_valid_image(item) and (is_pil_image(item) or item.ndim == 3),
+        ),
+        (
+            "videos",
+            processor.video_token,
+            videos,
+            make_batched_videos,
+            lambda item: isinstance(item, str) or (not is_pil_image(item) and is_valid_video(item)),
+        ),
+    ):
+        num_placeholders = sum(prompt.count(token) for prompt in prompts)
+        num_media = 0 if media is None else len(flatten(media))
+        if num_placeholders != num_media:
+            raise ValueError(f"Found {num_placeholders} {token} tokens in the text but received {num_media} {name}")
+        # one prompt cannot take another's media
+        per_prompt = _media_per_prompt(media, is_item) if len(prompts) > 1 else None
+        in_text = [prompt.count(token) for prompt in prompts]
+        if per_prompt is not None and per_prompt != in_text:
+            raise ValueError(f"Found {in_text} {token} tokens per prompt but received {per_prompt} {name} per prompt")
+
+
+def _check_video_metadata_count(videos, video_metadata) -> None:
+    # the shared sampler pairs clips with metadata by `zip`, which drops the clips past the last entry
+    if video_metadata is None:
+        return
+    videos = make_batched_videos(videos)
+    entries = make_batched_metadata(videos, video_metadata) if video_metadata != [] else []
+    if len(entries) != len(videos):
+        raise ValueError(
+            f"Received {len(videos)} videos but {len(entries)} `video_metadata` entries; pass one entry per video"
+        )
+
+
+DEFAULT_OMNI_CHAT_TEMPLATE = DEFAULT_CHAT_TEMPLATE.replace(
+    "{%- if content['type'] == 'audio' -%}",
+    "{%- if content['type'] == 'image' -%}{{- '<|image|>' -}}"
+    "{%- elif content['type'] == 'video' -%}{{- '<|video|>' -}}"
+    "{%- elif content['type'] == 'audio' -%}",
+)
+
+
+@auto_docstring
+class Lfm3OmniProcessor(ProcessorMixin):
+    """The LFM3-Audio processor extended with the LFM3-VL image and video processors: every image, video frame group
+    and input clip is expanded in place, wrapped in the media span tokens the Omni SFT data uses."""
+
+    attributes = ["image_processor", "feature_extractor", "tokenizer", "video_processor"]
+    tokenizer_class = "AutoTokenizer"
+    audio_tokenizer_class = "LiquidCodecModel"
+    valid_processor_kwargs = Lfm3OmniProcessorKwargs
+
+    def __init__(
+        self,
+        image_processor=None,
+        feature_extractor=None,
+        tokenizer=None,
+        video_processor=None,
+        audio_tokenizer=None,
+        chat_template=None,
+        audio_token="<|audio|>",
+        audio_out_token="<|audio_out|>",
+        image_token="<|image|>",
+        video_token="<|video|>",
+        vision_start_token="<|media_start|>",
+        vision_end_token="<|media_end|>",
+        vision_markers=True,
+        audio_markers=True,
+        **kwargs,
+    ):
+        r"""
+        image_processor ([`Lfm3VlImageProcessor`]):
+            The image patch processor.
+        feature_extractor ([`Lfm3OmniFeatureExtractor`]):
+            The log-mel feature extractor.
+        tokenizer ([`PreTrainedTokenizerFast`]):
+            The LFM3 tokenizer.
+        video_processor ([`Lfm3VlVideoProcessor`], *optional*):
+            The video frame sampler and patch processor.
+        audio_tokenizer ([`LiquidCodecModel`], *optional*):
+            The codec for 24 kHz output audio.
+        chat_template (`str`, *optional*):
+            The conversation template.
+        audio_token (`str`, *optional*, defaults to `"<|audio|>"`):
+            The input audio placeholder.
+        audio_out_token (`str`, *optional*, defaults to `"<|audio_out|>"`):
+            The output audio placeholder.
+        image_token (`str`, *optional*, defaults to `"<|image|>"`):
+            The image placeholder.
+        video_token (`str`, *optional*, defaults to `"<|video|>"`):
+            The video placeholder.
+        vision_start_token (`str`, *optional*, defaults to `"<|media_start|>"`):
+            The token opening an image, a video and each of its timestamped frame groups.
+        vision_end_token (`str`, *optional*, defaults to `"<|media_end|>"`):
+            The token closing them.
+        vision_markers (`bool`, *optional*, defaults to `True`):
+            Whether images and videos are wrapped in the vision span tokens, as the Omni SFT data renders them.
+        audio_markers (`bool`, *optional*, defaults to `True`):
+            Whether input audio clips are wrapped in the same span tokens, as the Omni SFT data renders them.
+        """
+        self.audio_token = audio_token
+        self.audio_out_token = audio_out_token
+        self.image_token = image_token
+        self.video_token = video_token
+        self.vision_start_token = vision_start_token
+        self.vision_end_token = vision_end_token
+        self.vision_markers = vision_markers
+        self.audio_markers = audio_markers
+        for token in (audio_token, audio_out_token, image_token, video_token):
+            if tokenizer.convert_tokens_to_ids(token) in (None, tokenizer.unk_token_id):
+                tokenizer.add_tokens([token], special_tokens=True)
+        # the span tokens exist in every LFM3 tokenizer under one of two names (`<|vision_*|>` in liquid7.1,
+        # `<|media_*|>` in liquid7.1_media) at the same ids; never mint new ones
+        for attribute, token, fallback in (
+            ("vision_start_token", vision_start_token, "<|vision_start|>"),
+            ("vision_end_token", vision_end_token, "<|vision_end|>"),
+        ):
+            if tokenizer.convert_tokens_to_ids(token) in (None, tokenizer.unk_token_id):
+                if tokenizer.convert_tokens_to_ids(fallback) not in (None, tokenizer.unk_token_id):
+                    setattr(self, attribute, fallback)
+                else:
+                    tokenizer.add_tokens([token], special_tokens=True)
+        self.audio_token_id = tokenizer.convert_tokens_to_ids(audio_token)
+        self.audio_out_token_id = tokenizer.convert_tokens_to_ids(audio_out_token)
+        self.image_token_id = tokenizer.convert_tokens_to_ids(image_token)
+        self.video_token_id = tokenizer.convert_tokens_to_ids(video_token)
+        super().__init__(
+            image_processor,
+            feature_extractor,
+            tokenizer,
+            video_processor,
+            chat_template=chat_template if chat_template is not None else DEFAULT_OMNI_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)
+        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,
+        images=None,
+        videos=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.
+        videos (`VideoInput`, *optional*):
+            One clip per `video_token`. With `use_audio_in_video=True` each video takes the first audio clip after it
+            as its soundtrack, interleaved after each timestamped frame group like the Omni SFT data; otherwise every
+            audio clip is its own media span.
+        """
+        for name, media in (("audio", audio), ("audio_out", audio_out), ("images", images), ("videos", videos)):
+            if isinstance(media, (list, tuple)) and len(media) == 0:
+                raise ValueError(f"`{name}` must hold at least one item; pass `{name}=None` for prompts without one")
+        if text is None:
+            # media-only calls return the features alone, like the audio processor
+            if audio_out is not None:
+                raise ValueError("`audio_out` needs the text holding its placeholders")
+            return super().__call__(audio=audio, images=images, videos=videos, **kwargs)
+        single = isinstance(text, str)
+        caller = [text] if single else list(text)
+        stripped = [self._strip_spans(prompt) for prompt in caller]
+        prompts, offsets = [prompt for prompt, _ in stripped], [row for _, row in stripped]
+        audio_out_codes = None
+        if audio_out is not None:
+            prompts, audio_out_codes, audio_out_offsets = self._expand_audio_out(prompts, audio_out)
+            offsets = self._compose_replacement_offsets(caller, offsets, audio_out_offsets)
+            merged = self._merge_kwargs(
+                self.valid_processor_kwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs
+            )
+            return_tensors = merged["text_kwargs"].get("return_tensors")
+        if videos is not None:
+            outputs = self._call_with_videos(prompts, audio, images, videos, single, **kwargs)
+        else:
+            outputs = super().__call__(
+                text=prompts[0] if single else prompts, audio=audio, images=images, videos=videos, **kwargs
+            )
+        if "text_replacement_offsets" in outputs:
+            # the expansion saw the prompts after the span pairs were stripped and the speech placeholders repeated
+            outputs["text_replacement_offsets"] = self._compose_replacement_offsets(
+                caller, offsets, outputs["text_replacement_offsets"]
+            )
+        if audio_out_codes is not None:
+            self._check_audio_out_placeholders(outputs["input_ids"], audio_out_codes.shape[0])
+            outputs["audio_out_codes"] = self._as_return_tensors(audio_out_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"
+            )
+        if text is not None:
+            _check_vision_placeholder_counts(self, [text] if isinstance(text, str) else text, images, videos)
+
+    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._wrap(self.audio_token * audio_inputs["num_audio_tokens"][audio_idx], self.audio_markers)
+
+    @property
+    def unused_input_names(self) -> list[str]:
+        return ["num_audio_tokens"]
+
+    @property
+    def model_input_names(self):
+        return list(
+            dict.fromkeys(
+                self.tokenizer.model_input_names
+                + self.feature_extractor.model_input_names
+                + self.image_processor.model_input_names
+                + (self.video_processor.model_input_names if self.video_processor is not None else [])
+                + ["audio_out_codes"]
+            )
+        )
+
+    def _wrap(self, placeholders: str, markers: bool) -> str:
+        return self.vision_start_token + placeholders + self.vision_end_token if markers else placeholders
+
+    def _strip_spans(self, prompt: str) -> tuple[str, list[dict]]:
+        # a chat template may already wrap a placeholder in the span pair; the expansion adds exactly one pair
+        wrapped = {
+            self.vision_start_token + token + self.vision_end_token: token
+            for token in (self.image_token, self.video_token, self.audio_token)
+        }
+        return self._rewrite_with_offsets(
+            prompt, "|".join(map(re.escape, wrapped)), lambda match: wrapped[match.group()], "span"
+        )
+
+    def apply_chat_template(
+        self,
+        conversation,
+        chat_template=None,
+        tools=None,
+        documents=None,
+        add_generation_prompt=False,
+        continue_final_message=False,
+        return_assistant_tokens_mask=False,
+        tokenize=False,
+        return_tensors=None,
+        return_dict=False,
+        load_audio_from_video=False,
+        processor_kwargs=None,
+        **kwargs,
+    ):
+        """With `load_audio_from_video=True`, each video's soundtrack becomes the audio part right behind it, which
+        `use_audio_in_video` pairs with the video; standalone clips are kept (the shared implementation drops them)."""
+        if load_audio_from_video and tokenize:
+            load_audio_from_video = False
+            conversation = self._with_soundtrack_parts(conversation)
+        return super().apply_chat_template(
+            conversation,
+            chat_template,
+            tools=tools,
+            documents=documents,
+            add_generation_prompt=add_generation_prompt,
+            continue_final_message=continue_final_message,
+            return_assistant_tokens_mask=return_assistant_tokens_mask,
+            tokenize=tokenize,
+            return_tensors=return_tensors,
+            return_dict=return_dict,
+            load_audio_from_video=load_audio_from_video,
+            processor_kwargs=processor_kwargs,
+            **kwargs,
+        )
+
+    @staticmethod
+    def _with_soundtrack_parts(conversation):
+        if not conversation:
+            return conversation
+        batched = isinstance(conversation[0], (list, tuple))
+        rewritten = []
+        for messages in conversation if batched else [conversation]:
+            rewritten.append([])
+            for message in messages:
+                content = message.get("content")
+                if isinstance(content, list):
+                    parts = []
+                    for part in content:
+                        parts.append(part)
+                        source = next((part[key] for key in ("video", "url", "path") if key in part), None)
+                        if part.get("type") == "video" and source is not None:
+                            parts.append({"type": "audio", "audio": source})
+                    message = {**message, "content": parts}
+                rewritten[-1].append(message)
+        return rewritten if batched else rewritten[0]
+
+    def _attach_soundtracks(self, prompt: str) -> tuple[str, list[dict]]:
+        """`use_audio_in_video`: each video takes the first audio placeholder after it as its soundtrack, moved right
+        behind the video; returns the prompt and the offset records of the moves."""
+        row, taken = [], set()
+        audio = [match.start() for match in re.finditer(re.escape(self.audio_token), prompt)]
+        for video in re.finditer(re.escape(self.video_token), prompt):
+            position = next((p for p in audio if p >= video.end() and p not in taken), None)
+            if position is None:
+                raise ValueError("`use_audio_in_video` needs an audio clip after each video")
+            taken.add(position)
+            row.append(
+                {"type": "audio", "span": (video.end(), video.end()), "text": "", "replacement": self.audio_token}
+            )
+            row.append(
+                {
+                    "type": "audio",
+                    "span": (position, position + len(self.audio_token)),
+                    "text": self.audio_token,
+                    "replacement": "",
+                }
+            )
+        row.sort(key=lambda r: r["span"])
+        pieces, last, offset = [], 0, 0
+        for entry in row:
+            start, end = entry["span"]
+            pieces.extend((prompt[last:start], entry["replacement"]))
+            entry["new_span"] = (start + offset, start + offset + len(entry["replacement"]))
+            offset += len(entry["replacement"]) - (end - start)
+            last = end
+        pieces.append(prompt[last:])
+        return "".join(pieces), row
+
+    def _call_with_videos(self, prompts, audio, images, videos, single, **kwargs):
+        """Expands the placeholders itself: a soundtrack's encoder tokens are split across its video's frame groups,
+        which the generic one-clip-one-run expansion cannot do."""
+        images, prompts, videos, audio = self.prepare_inputs_layout(
+            images=images, text=prompts, videos=videos, audio=audio, **kwargs
+        )
+        self.validate_inputs(images=images, text=prompts, videos=videos, audio=audio, **kwargs)
+        merged = self._merge_kwargs(
+            self.valid_processor_kwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs
+        )
+        return_tensors = merged["text_kwargs"].pop("return_tensors", None)
+        return_mm_token_type_ids = merged["text_kwargs"].pop("return_mm_token_type_ids", False)
+        return_text_replacement_offsets = merged["text_kwargs"].pop("return_text_replacement_offsets", False)
+        # Timestamps always need metadata internally, independently of whether the caller wants it returned.
+        return_metadata = kwargs.get("return_metadata", kwargs.get("videos_kwargs", {}).get("return_metadata", False))
+        merged["videos_kwargs"]["return_metadata"] = True
+        use_audio_in_video = merged["videos_kwargs"].pop("use_audio_in_video")
+        unattached, attach_offsets = prompts, [[] for _ in prompts]
+        if use_audio_in_video:
+            attached = [self._attach_soundtracks(prompt) for prompt in prompts]
+            prompts, attach_offsets = [prompt for prompt, _ in attached], [row for _, row in attached]
+        outputs = {}
+        clips = make_list_of_audio(audio) if audio is not None else []
+        videos = make_batched_videos(videos)
+        self._check_video_clock(videos, merged["videos_kwargs"].get("video_metadata"))
+        image_inputs = self.image_processor(images=images, **merged["images_kwargs"]) if images is not None else None
+        if image_inputs is not None:
+            outputs.update(image_inputs)
+        if clips:
+            # the decoding backend is a loader option, not a feature-extractor one (the shared path drops it too)
+            merged["audio_kwargs"].pop("load_audio_backend", None)
+            audio_inputs = self.feature_extractor(clips, **merged["audio_kwargs"])
+            lengths = np.asarray(audio_inputs["feature_attention_mask"]).sum(-1).tolist()
+            audio_tokens = [self.feature_extractor.num_encoder_tokens(int(length)) for length in lengths]
+            outputs.update(audio_inputs)
+        video_inputs = self.video_processor(videos=videos, **merged["videos_kwargs"])
+        outputs.update(
+            {key: value for key, value in video_inputs.items() if key != "video_metadata" or return_metadata}
+        )
+        video = self.video_token + (self.audio_token if use_audio_in_video else "")
+        pattern = re.compile("|".join(map(re.escape, (video, self.image_token, self.audio_token))))
+        counters = {"image": 0, "video": 0, "audio": 0}
+
+        def expand(match):
+            token = match.group(0)
+            if token == self.image_token:
+                text = self.replace_image_token(image_inputs, counters["image"])
+                counters["image"] += 1
+                return text
+            if token == self.audio_token:
+                text = self._wrap(self.audio_token * audio_tokens[counters["audio"]], self.audio_markers)
+                counters["audio"] += 1
+                return text
+            soundtrack = None
+            if token != self.video_token:  # the pair `_attach_soundtracks` made
+                soundtrack = audio_tokens[counters["audio"]]
+                counters["audio"] += 1
+            text = self._video_expansion(video_inputs, counters["video"], soundtrack)
+            counters["video"] += 1
+            return text
+
+        expanded, batch_offsets = [], []
+        for prompt in prompts:
+            pieces, offsets = [], []
+            last, offset = 0, 0
+            for match in pattern.finditer(prompt):
+                start, end = match.span()
+                replacement = expand(match)
+                pieces.extend((prompt[last:start], replacement))
+                token = match.group(0)
+                offsets.append(
+                    {
+                        "type": "video"
+                        if token.startswith(self.video_token)
+                        else "image"
+                        if token == self.image_token
+                        else "audio",
+                        "span": (start, end),
+                        "new_span": (start + offset, start + offset + len(replacement)),
+                        "text": token,
+                        "replacement": replacement,
+                    }
+                )
+                offset += len(replacement) - (end - start)
+                last = end
+            pieces.append(prompt[last:])
+            expanded.append("".join(pieces))
+            batch_offsets.append(offsets)
+        text_inputs = self.tokenizer(expanded, **merged["text_kwargs"])
+        self._check_special_mm_tokens(expanded, text_inputs, modalities=["image", "video", "audio"])
+        if return_text_replacement_offsets:
+            text_inputs["text_replacement_offsets"] = self._compose_replacement_offsets(
+                unattached, attach_offsets, batch_offsets
+            )
+        if return_mm_token_type_ids:
+            text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"])
+        outputs.update(text_inputs)
+        return BatchFeature(outputs, tensor_type=return_tensors, skip_tensor_conversion=self.skip_tensor_conversion)
+
+    def _video_expansion(self, video_inputs: dict, video_idx: int, soundtrack_tokens: int | None = None) -> str:
+        grid = video_inputs["video_grid_thw"][video_idx]
+        patch_tokens = int(grid[1:].prod() // self.video_processor.merge_size**2)
+        metadata = video_inputs["video_metadata"][video_idx]
+        if metadata.frames_indices is None or not (
+            metadata.fps is not None and math.isfinite(metadata.fps) and metadata.fps > 0
+        ):
+            raise ValueError(
+                "LFM3-Omni prompts carry frame timestamps: pass videos with `video_metadata` holding the sampled frame "
+                "indices and a positive fps"
+            )
+        # a released transformers reads torchcodec's scanned average fps, which carries float noise
+        # (30.000000000000004) where the fork reads the container header
+        metadata.fps = round(metadata.fps, 6)
+        temporal = self.video_processor.temporal_patch_size
+        timestamps = self._calculate_timestamps(metadata.frames_indices, metadata.fps, temporal)[: int(grid[0])]
+        bounds = None
+        if soundtrack_tokens is not None:
+            # the soundtrack's encoder tokens are split at each frame group's first frame
+            starts = [index / metadata.fps for index in list(metadata.frames_indices)[::temporal]][: int(grid[0])]
+            seconds_per_token = (
+                self.feature_extractor.hop_length
+                * self.feature_extractor.encoder_stride
+                / self.feature_extractor.sampling_rate
+            )
+            bounds = (
+                [0]
+                + [min(soundtrack_tokens, round(start / seconds_per_token)) for start in starts[1:]]
+                + [soundtrack_tokens]
+            )
+            bounds = [max(bounds[: i + 1]) for i in range(len(bounds))]
+        groups = ""
+        for index, timestamp in enumerate(timestamps):
+            groups += f"<{timestamp:.1f} seconds>" + self._wrap(self.video_token * patch_tokens, self.vision_markers)
+            if bounds is not None:
+                groups += self.audio_token * (bounds[index + 1] - bounds[index])
+        return self._wrap(groups, self.vision_markers)
+
+    def _check_video_clock(self, videos, video_metadata) -> None:
+        _check_video_metadata_count(videos, video_metadata)
+        if video_metadata is None:
+            if any(not isinstance(video, str) for video in make_batched_videos(videos)):
+                # decoded frames carry no clock: without metadata the timestamps would silently assume 24 fps
+                raise ValueError(
+                    "LFM3-Omni prompts carry frame timestamps: pass decoded frames with `video_metadata` holding the "
+                    "clip's fps and total_num_frames (a path or URL is decoded with its own metadata)"
+                )
+        # the shared sampler writes a 24 fps default into metadata lacking a clock before the expansion runs
+        elif any(metadata.fps is None for metadata in make_batched_metadata(videos, video_metadata)):
+            raise ValueError(
+                "LFM3-Omni prompts carry frame timestamps: `video_metadata` must hold the clip's fps "
+                "(a path or URL is decoded with its own metadata)"
+            )
+
+    def _process_videos(self, videos, **kwargs):
+        # only media-only calls reach the shared path (prompts with videos take `_call_with_videos`): the pairing
+        # flag is not a video-processor argument and no placeholder text is needed
+        kwargs.pop("use_audio_in_video", None)
+        self._check_video_clock(videos, kwargs.get("video_metadata"))
+        return self.video_processor(videos, **kwargs), []
+
+    def replace_image_token(self, image_inputs, image_idx, **kwargs):
+        count = image_inputs["image_grid_thw"][image_idx].prod() // self.image_processor.merge_size**2
+        return self._wrap(self.image_token * int(count), self.vision_markers)
+
+    @staticmethod
+    def _calculate_timestamps(indices, video_fps: float, temporal_patch_size: int) -> list[float]:
+        indices = list(indices)
+        if pad := -len(indices) % temporal_patch_size:
+            indices.extend([indices[-1]] * pad)
+        return [
+            (indices[i] + indices[i + temporal_patch_size - 1]) / (2 * video_fps)
+            for i in range(0, len(indices), temporal_patch_size)
+        ]
+
+
+__all__ = ["Lfm3OmniProcessor"]

processor_config.json

@@ -0,0 +1,101 @@
+{
+  "audio_markers": true,
+  "audio_out_token": "<|audio_out|>",
+  "audio_token": "<|audio|>",
+  "audio_tokenizer_subfolder": "audio_tokenizer",
+  "auto_map": {
+    "AutoProcessor": "processing_lfm3_omni.Lfm3OmniProcessor"
+  },
+  "feature_extractor": {
+    "auto_map": {
+      "AutoFeatureExtractor": "feature_extraction_lfm3_omni.Lfm3OmniFeatureExtractor",
+      "AutoProcessor": "processing_lfm3_omni.Lfm3OmniProcessor"
+    },
+    "encoder_stride": 8,
+    "feature_extractor_type": "Lfm3OmniFeatureExtractor",
+    "feature_size": 128,
+    "hop_length": 160,
+    "log_floor": -12.0,
+    "n_fft": 400,
+    "normalization_divisor": 4.0,
+    "normalization_offset": 4.0,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "return_attention_mask": true,
+    "sampling_rate": 16000,
+    "win_length": 400
+  },
+  "image_processor": {
+    "auto_map": {
+      "AutoImageProcessor": "image_processing_lfm3_vl.Lfm3VlImageProcessor",
+      "AutoProcessor": "processing_lfm3_omni.Lfm3OmniProcessor"
+    },
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_processor_type": "Lfm3VlImageProcessor",
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "merge_size": 2,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+      "longest_edge": 2097152,
+      "shortest_edge": 65536
+    },
+    "temporal_patch_size": 2
+  },
+  "image_token": "<|image|>",
+  "processor_class": "Lfm3OmniProcessor",
+  "video_processor": {
+    "auto_map": {
+      "AutoProcessor": "processing_lfm3_omni.Lfm3OmniProcessor",
+      "AutoVideoProcessor": "video_processing_lfm3_vl.Lfm3VlVideoProcessor"
+    },
+    "cap_pixels_per_frame": false,
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "fps": 2.0,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "max_frames": 384,
+    "max_video_tokens": 768,
+    "merge_size": 2,
+    "min_frames": 4,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "size": {
+      "longest_edge": 25165824,
+      "shortest_edge": 4096
+    },
+    "temporal_patch_size": 2,
+    "video_processor_type": "Lfm3VlVideoProcessor"
+  },
+  "video_token": "<|video|>",
+  "vision_end_token": "<|vision_end|>",
+  "vision_markers": true,
+  "vision_start_token": "<|vision_start|>"
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41927bd314d17d7d7fadb44e10cd0919f7b71528d7741596d8cbd15c6945ec44
+size 17908634

tokenizer_config.json

@@ -0,0 +1,25 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_lfm3_omni.Lfm3OmniProcessor"
+  },
+  "backend": "tokenizers",
+  "bos_token": "<|startoftext|>",
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": [],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "image_token": "<|image|>",
+    "video_token": "<|video|>",
+    "vision_end_token": "<|vision_end|>",
+    "vision_start_token": "<|vision_start|>"
+  },
+  "pad_token": "<|pad|>",
+  "processor_class": "Lfm3OmniProcessor",
+  "tokenizer_class": "TokenizersBackend",
+  "video_token": "<|video|>",
+  "vision_end_token": "<|vision_end|>",
+  "vision_start_token": "<|vision_start|>"
+}

video_processing_lfm3_vl.py

@@ -0,0 +1,396 @@
+# 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
+
+import numpy as np
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_utils import PILImageResampling, SizeDict
+from transformers.processing_utils import Unpack, VideosKwargs
+from transformers.utils import TensorType, auto_docstring, logging
+from transformers.video_processing_utils import BaseVideoProcessor
+from transformers.video_utils import (
+    VideoMetadata,
+    group_videos_by_shape,
+    make_batched_metadata,
+    make_batched_videos,
+    reorder_videos,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+class Lfm3VlVideoProcessorInitKwargs(VideosKwargs, total=False):
+    r"""
+    patch_size (`int`, *optional*, defaults to 16):
+        The spatial patch size of the vision encoder.
+    temporal_patch_size (`int`, *optional*, defaults to 2):
+        The temporal patch size of the vision encoder.
+    merge_size (`int`, *optional*, defaults to 2):
+        The merge size of the vision encoder to llm encoder.
+    min_frames (`int`, *optional*, defaults to 4):
+        The minimum number of frames to sample from video.
+    max_frames (`int`, *optional*, defaults to 768):
+        The maximum number of frames to sample from video.
+    cap_pixels_per_frame (`bool`, *optional*):
+        Whether to cap each frame's pixel cost the way the reference implementation (qwen-vl-utils) does:
+        per-frame pixels are limited to `min(max_video_tokens * factor**2, size["longest_edge"] / num_frames)`
+        and floored at `1.05 * size["shortest_edge"]`, so token cost scales with clip duration. Without the
+        cap, videos that sample few frames spend the whole `size["longest_edge"]` budget on those frames and
+        keep near-native per-frame resolution, so a short clip can cost almost as many tokens as a long
+        video. If unset, the current uncapped behavior is kept and a warning is emitted: the default will
+        change to `True` in v5.22, after which the argument will be removed.
+    max_video_tokens (`int`, *optional*, defaults to 768):
+        The per-frame token ceiling applied by `cap_pixels_per_frame`, in vision tokens per frame
+        (qwen-vl-utils' `VIDEO_MAX_TOKEN_NUM`).
+    """
+
+    patch_size: int
+    temporal_patch_size: int
+    merge_size: int
+    min_frames: int
+    max_frames: int
+    cap_pixels_per_frame: bool
+    max_video_tokens: int
+
+
+def smart_resize(
+    num_frames: int,
+    height: int,
+    width: int,
+    temporal_factor: int = 2,
+    factor: int = 28,
+    min_pixels: int = 112 * 112,
+    max_pixels: int = 14 * 14 * 2 * 2 * 2 * 6144,
+):
+    if num_frames < temporal_factor:
+        raise ValueError(f"t:{num_frames} must be larger than temporal_factor:{temporal_factor}")
+    if height < factor or width < factor:
+        scale = max(factor / height, factor / width)
+        height = int(height * scale)
+        width = int(width * scale)
+
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = round(height / factor) * factor
+    w_bar = round(width / factor) * factor
+    t_bar = round(num_frames / temporal_factor) * temporal_factor
+
+    if t_bar * h_bar * w_bar > max_pixels:
+        beta = math.sqrt((num_frames * height * width) / max_pixels)
+        h_bar = max(factor, math.floor(height / beta / factor) * factor)
+        w_bar = max(factor, math.floor(width / beta / factor) * factor)
+    elif t_bar * h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (num_frames * height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+
+    return h_bar, w_bar
+
+
+def _check_video_fps(metadata: VideoMetadata) -> None:
+    if metadata.fps is not None and not (math.isfinite(metadata.fps) and metadata.fps > 0):
+        raise ValueError(f"`video_metadata.fps` must be a positive number, got {metadata.fps}")
+
+
+def _check_video_metadata_count(videos, video_metadata) -> None:
+    # the shared sampler pairs clips with metadata by `zip`, which drops the clips past the last entry
+    if video_metadata is None:
+        return
+    videos = make_batched_videos(videos)
+    entries = make_batched_metadata(videos, video_metadata) if video_metadata != [] else []
+    if len(entries) != len(videos):
+        raise ValueError(
+            f"Received {len(videos)} videos but {len(entries)} `video_metadata` entries; pass one entry per video"
+        )
+
+
+@auto_docstring
+class Lfm3VlVideoProcessor(BaseVideoProcessor):
+    resample = PILImageResampling.BICUBIC
+    size = {"shortest_edge": 4_096, "longest_edge": 25_165_824}
+    image_mean = [0.5, 0.5, 0.5]
+    image_std = [0.5, 0.5, 0.5]
+    do_resize = True
+    do_rescale = True
+    do_normalize = True
+    do_convert_rgb = True
+    patch_size = 16
+    temporal_patch_size = 2
+    merge_size = 2
+    min_frames = 4
+    max_frames = 768
+    do_sample_frames = True
+    cap_pixels_per_frame = False
+    max_video_tokens = 768
+    valid_kwargs = Lfm3VlVideoProcessorInitKwargs
+    model_input_names = ["pixel_values_videos", "video_grid_thw"]
+    max_image_size = {"longest_edge": 28 * 28 * 2 * 30000}
+    max_duration = None
+    num_frames = None
+    fps = 2.0
+
+    def __init__(self, **kwargs: Unpack[Lfm3VlVideoProcessorInitKwargs]):
+        super().__init__(**kwargs)
+
+    def sample_frames(
+        self, metadata: VideoMetadata, num_frames: int | None = None, fps: int | float | None = None, **kwargs
+    ):
+        """Decoded frames whose metadata carries no fps are sampled and timestamped at the 24 fps the shared sampler
+        assumes (with a warning); pass `video_metadata` with the clip's fps for accurate timestamps."""
+        _check_video_fps(metadata)
+        if fps is not None and not (math.isfinite(fps) and fps > 0):
+            raise ValueError(f"`fps` must be a positive number, got {fps}")
+        if fps is not None and num_frames is not None:
+            raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
+
+        total_num_frames = metadata.total_num_frames
+        fps = fps if fps is not None else self.fps
+
+        # If num_frames is not given but fps is, calculate num_frames from fps
+        if num_frames is None and fps is not None:
+            if metadata.fps is None:
+                metadata.fps = 24
+                logger.warning_once(
+                    "Asked to sample `fps` frames per second but no video metadata was provided which is required when sampling with `fps`. "
+                    "Defaulting to `fps=24`. Please provide `video_metadata` for more accurate results."
+                )
+            num_frames = int(total_num_frames / metadata.fps * fps)
+            num_frames = min(max(num_frames, self.min_frames), self.max_frames, total_num_frames)
+
+        if num_frames is None:
+            num_frames = min(max(total_num_frames, self.min_frames), self.max_frames)
+
+        indices = np.linspace(0, total_num_frames - 1, num_frames).round().astype(int)
+
+        return indices
+
+    def resize(
+        self,
+        videos: "torch.Tensor",
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        factor: int,
+        temporal_factor: int,
+        cap_pixels_per_frame: bool | None = None,
+        **kwargs,
+    ) -> "torch.Tensor":
+        """Resize dynamically based on input video aspect ratio."""
+        if not size.shortest_edge or not size.longest_edge:
+            raise ValueError(f"`size` dict must contain 'shortest_edge' and 'longest_edge' keys but got {size}.")
+
+        num_frames = videos.shape[1]
+        max_pixels = size.longest_edge
+        if cap_pixels_per_frame:
+            # per-frame pixels are capped at `max_video_tokens` patches or the budget's even share per frame
+            frame_cap = self.max_video_tokens * factor * factor
+            pixels_per_frame = max(min(frame_cap, size.longest_edge // num_frames), int(size.shortest_edge * 1.05))
+            max_pixels = pixels_per_frame * num_frames
+
+        height, width = videos.shape[-2:]
+        resized_height, resized_width = smart_resize(
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            factor=factor,
+            temporal_factor=temporal_factor,
+            min_pixels=size.shortest_edge,
+            max_pixels=max_pixels,
+        )
+        return super().resize(
+            image=videos,
+            size=SizeDict(height=resized_height, width=resized_width),
+            resample=resample,
+        )
+
+    def patchify(
+        self,
+        videos: "torch.Tensor",
+        patch_size: int,
+        merge_size: int,
+        temporal_patch_size: int,
+    ) -> tuple["torch.Tensor", int, int]:
+        "Patchifies each video into flat layout of shape (`seq_len`, `patch_dim`) so we can concat dynamically shaped pixels."
+        batch_size, num_frames, channel, resized_height, resized_width = videos.shape
+
+        # Check that videos have `num_frames` divisible by `temporal_patch_size`
+        if pad := -num_frames % temporal_patch_size:
+            repeats = videos[:, -1:].expand(-1, pad, -1, -1, -1)
+            videos = torch.cat((videos, repeats), dim=1)
+            num_frames += pad
+
+        grid_t = num_frames // temporal_patch_size
+        grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+        patches = videos.view(
+            batch_size,
+            grid_t,
+            temporal_patch_size,
+            channel,
+            grid_h // merge_size,
+            merge_size,
+            patch_size,
+            grid_w // merge_size,
+            merge_size,
+            patch_size,
+        )
+        patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
+        flatten_patches = patches.reshape(
+            batch_size,
+            grid_t * grid_h * grid_w,
+            channel * temporal_patch_size * patch_size * patch_size,
+        )
+
+        return flatten_patches, grid_t, grid_h, grid_w
+
+    def _preprocess(
+        self,
+        videos: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int | None = None,
+        temporal_patch_size: int | None = None,
+        merge_size: int | None = None,
+        cap_pixels_per_frame: bool | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ):
+        if cap_pixels_per_frame is None:
+            logger.warning_once(
+                "Lfm3Vl video processing does not apply the per-frame pixel cap the reference "
+                "implementation (qwen-vl-utils) applies, so some videos cost far more tokens than they "
+                "would there. In v5.22 the capped behavior will become the default and "
+                "`cap_pixels_per_frame` will be removed. Pass `cap_pixels_per_frame=True` to adopt the "
+                "reference behavior now, or `False` to keep the current behavior and silence this "
+                "warning."
+            )
+            cap_pixels_per_frame = False
+        # Group videos by size for batched resizing
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+        for shape, stacked_videos in grouped_videos.items():
+            if do_resize:
+                stacked_videos = self.resize(
+                    videos=stacked_videos,
+                    size=size,
+                    resample=resample,
+                    factor=patch_size * merge_size,
+                    temporal_factor=temporal_patch_size,
+                    cap_pixels_per_frame=cap_pixels_per_frame,
+                )
+            resized_videos_grouped[shape] = stacked_videos
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+
+        # Group videos by size for further processing
+        # Needed in case do_resize is False, or resize returns videos with different sizes
+        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
+        processed_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            # Fused rescale and normalize
+            stacked_videos = self.rescale_and_normalize(
+                stacked_videos, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+            )
+            patches, grid_t, grid_h, grid_w = self.patchify(
+                stacked_videos,
+                patch_size=patch_size,
+                merge_size=merge_size,
+                temporal_patch_size=temporal_patch_size,
+            )
+
+            processed_videos_grouped[shape] = patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * len(stacked_videos)
+
+        processed_videos = reorder_videos(processed_videos_grouped, grouped_videos_index)
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(processed_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values_videos": pixel_values_videos, "video_grid_thw": video_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    def get_num_of_video_patches(self, num_frames: int, height: int, width: int, videos_kwargs=None):
+        """
+        A utility that returns number of video patches a given video size.
+
+        Args:
+            num_frames (`int`):
+                Number of frames in the input video.
+            height (`int`):
+                Height of the input video.
+            width (`int`):
+                Width of the input video.
+            videos_kwargs (`dict`, *optional*)
+                Any kwargs to override defaults of the video processor.
+        Returns:
+            `int`: Number of video patches per video.
+        """
+        videos_kwargs = videos_kwargs if videos_kwargs is not None else {}
+        size = videos_kwargs.get("size", None) or self.size
+        min_pixels = size["shortest_edge"]
+        max_pixels = size["longest_edge"]
+        patch_size = videos_kwargs.get("patch_size", None) or self.patch_size
+        merge_size = videos_kwargs.get("merge_size", None) or self.merge_size
+        temporal_patch_size = videos_kwargs.get("temporal_patch_size", None) or self.temporal_patch_size
+        cap_pixels_per_frame = videos_kwargs.get("cap_pixels_per_frame", None) or self.cap_pixels_per_frame
+
+        factor = patch_size * merge_size
+        if cap_pixels_per_frame:
+            # Keep the count in sync with `resize` when the per-frame cap is active.
+            frame_cap = self.max_video_tokens * factor * factor
+            pixels_per_frame = max(min(frame_cap, max_pixels // num_frames), int(min_pixels * 1.05))
+            max_pixels = pixels_per_frame * num_frames
+
+        resized_height, resized_width = smart_resize(
+            num_frames,
+            height,
+            width,
+            temporal_factor=temporal_patch_size,
+            factor=factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+        grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+        grid_t = (num_frames + -num_frames % temporal_patch_size) // temporal_patch_size
+        return grid_t * grid_h * grid_w
+
+    def preprocess(self, videos, **kwargs):
+        _check_video_metadata_count(videos, kwargs.get("video_metadata"))
+        # the shared `preprocess` fills every unset option from the instance before `sample_frames` runs, so the
+        # caller's options are told apart from the instance ones here: a caller option beats an instance one, and a
+        # configured `num_frames` beats the default fps
+        num_frames, fps = kwargs.get("num_frames"), kwargs.get("fps")
+        if num_frames is not None and fps is not None:
+            raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!")
+        if num_frames is not None or (fps is None and self.num_frames is not None):
+            kwargs["fps"] = None
+        elif fps is not None:
+            kwargs["num_frames"] = None
+        return super().preprocess(videos, **kwargs)
+
+
+__all__ = ["Lfm3VlVideoProcessor"]