# Copyright 2026 The HuggingFace Team. All rights reserved. # Licensed under the Apache License, Version 2.0 """Simple PyTorch LFM3-VL with a variable-shape Liquid vision encoder.""" from dataclasses import dataclass import torch import torch.nn.functional as F from torch import nn from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from .configuration_lfm3_vl import Lfm3VlConfig class Lfm3VlRMSNorm(nn.Module): def __init__(self, hidden_size, eps): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps def forward(self, hidden_states): dtype = hidden_states.dtype normalized = hidden_states.float() normalized = normalized * torch.rsqrt(normalized.square().mean(-1, keepdim=True) + self.eps) return (normalized.to(dtype) * self.weight).to(dtype) def _rope_1d(hidden_states, position_ids, theta, rotary_factor): head_dim = hidden_states.shape[-1] pairs = int(head_dim * rotary_factor) // 2 frequencies = theta ** (-(2.0 * torch.arange(pairs, device=hidden_states.device, dtype=torch.float32) / head_dim)) angles = position_ids.float()[..., None] * frequencies cos = angles.cos().to(hidden_states.dtype)[:, None, :, :] sin = angles.sin().to(hidden_states.dtype)[:, None, :, :] first = hidden_states[..., :pairs] second = hidden_states[..., head_dim // 2 : head_dim // 2 + pairs] output = hidden_states.clone() output[..., :pairs] = first * cos - second * sin output[..., head_dim // 2 : head_dim // 2 + pairs] = first * sin + second * cos return output def _rope_2d(hidden_states, height, width, theta): head_dim = hidden_states.shape[-1] quarter = head_dim // 4 frequencies = theta ** ( -(2.0 * torch.arange(quarter, device=hidden_states.device, dtype=torch.float32) / (head_dim // 2)) ) rows = torch.arange(height, device=hidden_states.device).repeat_interleave(width) cols = torch.arange(width, device=hidden_states.device).repeat(height) angles = torch.cat((rows[:, None] * frequencies, cols[:, None] * frequencies), dim=-1) cos = angles.cos().to(hidden_states.dtype)[:, None, :] sin = angles.sin().to(hidden_states.dtype)[:, None, :] first, second = hidden_states[..., : head_dim // 2], hidden_states[..., head_dim // 2 :] return torch.cat((first * cos - second * sin, first * sin + second * cos), dim=-1) class Lfm3VlCache: def __init__(self, config, batch_size, dtype, device): self.key_cache = [None] * config.num_hidden_layers self.value_cache = [None] * config.num_hidden_layers self.conv_cache = [ torch.zeros(batch_size, config.hidden_size, config.conv_L_cache - 1, dtype=dtype, device=device) for _ in range(config.num_hidden_layers) ] self.layer_types = config.layer_types self.is_decode = False def get_seq_length(self): for layer_type, key in zip(self.layer_types, self.key_cache): if layer_type != "conv" and key is not None: return key.shape[-2] return 0 def reorder_cache(self, beam_idx): for index in range(len(self.layer_types)): if self.key_cache[index] is not None: self.key_cache[index] = self.key_cache[index].index_select(0, beam_idx) self.value_cache[index] = self.value_cache[index].index_select(0, beam_idx) self.conv_cache[index] = self.conv_cache[index].index_select(0, beam_idx) class Lfm3VlMLP(nn.Module): def __init__(self, hidden_size, intermediate_size): super().__init__() self.gate_up = nn.Linear(hidden_size, 2 * intermediate_size, bias=False) self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False) def forward(self, hidden_states): gate, up = self.gate_up(hidden_states).chunk(2, dim=-1) return self.w2(F.silu(gate) * up) class Lfm3VlAttention(nn.Module): def __init__(self, config, layer_idx, layer_type): super().__init__() self.layer_idx = layer_idx self.layer_type = layer_type self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.head_dim = config.head_dim or config.hidden_size // config.num_attention_heads qkv_size = (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim self.qkv = nn.Linear(config.hidden_size, qkv_size, bias=False) self.gate_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) self.q_norm = Lfm3VlRMSNorm(self.head_dim, config.rms_norm_eps) self.k_norm = Lfm3VlRMSNorm(self.head_dim, config.rms_norm_eps) self.theta = ( config.full_attention_rope_theta if layer_type == "full_attention" else config.sliding_attention_rope_theta ) self.rotary_factor = config.partial_rotary_factor if layer_type == "full_attention" else 1.0 self.sliding_window = config.sliding_window def forward(self, hidden_states, position_ids, attention_mask, past_key_values): batch_size, query_length, _ = hidden_states.shape q_size = self.num_heads * self.head_dim kv_size = self.num_key_value_heads * self.head_dim q, k, v = self.qkv(hidden_states).split((q_size, kv_size, kv_size), dim=-1) q = self.q_norm(q.view(batch_size, query_length, self.num_heads, self.head_dim)).transpose(1, 2) k = self.k_norm(k.view(batch_size, query_length, self.num_key_value_heads, self.head_dim)).transpose(1, 2) v = v.view(batch_size, query_length, self.num_key_value_heads, self.head_dim).transpose(1, 2) q = _rope_1d(q, position_ids, self.theta, self.rotary_factor) k = _rope_1d(k, position_ids, self.theta, self.rotary_factor) if past_key_values is not None: old_k = past_key_values.key_cache[self.layer_idx] old_v = past_key_values.value_cache[self.layer_idx] if old_k is not None: k = torch.cat((old_k, k), dim=-2) v = torch.cat((old_v, v), dim=-2) past_key_values.key_cache[self.layer_idx] = k past_key_values.value_cache[self.layer_idx] = v groups = self.num_heads // self.num_key_value_heads k = k.repeat_interleave(groups, dim=1) v = v.repeat_interleave(groups, dim=1) key_length = k.shape[-2] query_positions = torch.arange(key_length - query_length, key_length, device=hidden_states.device) key_positions = torch.arange(key_length, device=hidden_states.device) allowed = key_positions[None, :] <= query_positions[:, None] if self.layer_type == "sliding_attention": allowed &= key_positions[None, :] > query_positions[:, None] - self.sliding_window if attention_mask is not None: allowed = allowed[None, None] & attention_mask[:, None, None, :key_length].bool() else: allowed = allowed[None, None] additive_mask = torch.zeros(allowed.shape, dtype=q.dtype, device=q.device) additive_mask.masked_fill_(~allowed, torch.finfo(q.dtype).min) output = F.scaled_dot_product_attention(q, k, v, attn_mask=additive_mask) gate = torch.sigmoid(self.gate_proj(hidden_states)).transpose(1, 2).unsqueeze(-1) output = (output * gate).transpose(1, 2).reshape(batch_size, query_length, -1) return self.o_proj(output) class Lfm3VlShortConv(nn.Module): def __init__(self, config, layer_idx): super().__init__() self.layer_idx = layer_idx self.kernel_size = config.conv_L_cache self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False) self.conv = nn.Parameter(torch.empty(config.hidden_size, 1, self.kernel_size)) self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) def forward(self, hidden_states, attention_mask, past_key_values): first, middle, third = self.in_proj(hidden_states).chunk(3, dim=-1) gated = first * third if attention_mask is not None: gated = gated * attention_mask[:, -hidden_states.shape[1] :, None] gated = gated.transpose(1, 2) if past_key_values is not None and past_key_values.is_decode: state = past_key_values.conv_cache[self.layer_idx] full = torch.cat((state, gated), dim=-1) output = (full * self.conv[:, 0, :][None]).sum(-1, keepdim=True) past_key_values.conv_cache[self.layer_idx] = full[..., -self.kernel_size + 1 :] else: output = F.conv1d(gated, self.conv, padding=self.kernel_size - 1, groups=gated.shape[1]) output = output[..., : hidden_states.shape[1]] if past_key_values is not None: past_key_values.conv_cache[self.layer_idx] = F.pad( gated, (max(0, self.kernel_size - 1 - gated.shape[-1]), 0) )[..., -self.kernel_size + 1 :] return self.out_proj((middle.transpose(1, 2) * output).transpose(1, 2)) class Lfm3VlLanguageLayer(nn.Module): def __init__(self, config, layer_idx, layer_type): super().__init__() self.layer_type = layer_type self.operator_norm = Lfm3VlRMSNorm(config.hidden_size, config.rms_norm_eps) self.post_operator_norm = Lfm3VlRMSNorm(config.hidden_size, config.rms_norm_eps) self.ffn_norm = Lfm3VlRMSNorm(config.hidden_size, config.rms_norm_eps) self.post_ffn_norm = Lfm3VlRMSNorm(config.hidden_size, config.rms_norm_eps) if layer_type == "conv": self.conv = Lfm3VlShortConv(config, layer_idx) else: self.self_attn = Lfm3VlAttention(config, layer_idx, layer_type) self.feed_forward = Lfm3VlMLP(config.hidden_size, config.intermediate_size) def forward(self, hidden_states, position_ids, attention_mask, past_key_values): operator_input = self.operator_norm(hidden_states) if self.layer_type == "conv": operator_output = self.conv(operator_input, attention_mask, past_key_values) else: operator_output = self.self_attn(operator_input, position_ids, attention_mask, past_key_values) hidden_states = hidden_states + self.post_operator_norm(operator_output) return hidden_states + self.post_ffn_norm(self.feed_forward(self.ffn_norm(hidden_states))) class Lfm3VlBackbone(nn.Module): def __init__(self, config): super().__init__() self.config = config self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) self.layers = nn.ModuleList( [Lfm3VlLanguageLayer(config, index, layer_type) for index, layer_type in enumerate(config.layer_types)] ) self.embedding_norm = Lfm3VlRMSNorm(config.hidden_size, config.rms_norm_eps) def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_values=None): if past_key_values is not None: past_key_values.is_decode = past_key_values.get_seq_length() > 0 if position_ids is None: past_length = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = ( torch.arange(past_length, past_length + hidden_states.shape[1], device=hidden_states.device) .unsqueeze(0) .expand(hidden_states.shape[0], -1) ) for layer in self.layers: hidden_states = layer(hidden_states, position_ids, attention_mask, past_key_values) return self.embedding_norm(hidden_states) class Lfm3VlVisionAttention(nn.Module): def __init__(self, config): super().__init__() self.num_heads = config.vision_num_attention_heads self.num_key_value_heads = config.vision_num_key_value_heads self.head_dim = config.vision_hidden_size // self.num_heads self.theta = config.vision_rope_theta self.q_proj = nn.Linear(config.vision_hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(config.vision_hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(config.vision_hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.out_proj = nn.Linear(config.vision_hidden_size, config.vision_hidden_size, bias=False) self.q_layernorm = Lfm3VlRMSNorm(self.head_dim, config.rms_norm_eps) self.k_layernorm = Lfm3VlRMSNorm(self.head_dim, config.rms_norm_eps) def forward(self, hidden_states, height, width): tokens = hidden_states.shape[0] q = self.q_layernorm(self.q_proj(hidden_states).view(tokens, self.num_heads, self.head_dim)) k = self.k_layernorm(self.k_proj(hidden_states).view(tokens, self.num_key_value_heads, self.head_dim)) v = self.v_proj(hidden_states).view(tokens, self.num_key_value_heads, self.head_dim) q = _rope_2d(q, height, width, self.theta) k = _rope_2d(k, height, width, self.theta) groups = self.num_heads // self.num_key_value_heads output = F.scaled_dot_product_attention( q.transpose(0, 1).unsqueeze(0), k.repeat_interleave(groups, dim=1).transpose(0, 1).unsqueeze(0), v.repeat_interleave(groups, dim=1).transpose(0, 1).unsqueeze(0), ) return self.out_proj(output.squeeze(0).transpose(0, 1).reshape(tokens, -1)) class Lfm3VlVisionConv2d(nn.Module): def __init__(self, hidden_size): super().__init__() self.in_proj = nn.Linear(hidden_size, 3 * hidden_size, bias=False) self.spatial_conv = nn.Parameter(torch.empty(hidden_size, 1, 3, 3)) self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False) def forward(self, hidden_states, height, width): first, middle, third = self.in_proj(hidden_states).chunk(3, dim=-1) gated = (first * third).transpose(0, 1).reshape(1, -1, height, width) output = F.conv2d(gated, self.spatial_conv, padding=1, groups=gated.shape[1]) return self.out_proj(middle * output.reshape(-1, height * width).transpose(0, 1)) class Lfm3VlVisionLayer(nn.Module): def __init__(self, config, layer_type): super().__init__() self.layer_type = layer_type self.operator_norm = Lfm3VlRMSNorm(config.vision_hidden_size, config.rms_norm_eps) self.ffn_norm = Lfm3VlRMSNorm(config.vision_hidden_size, config.rms_norm_eps) if layer_type == "conv": self.conv = Lfm3VlVisionConv2d(config.vision_hidden_size) else: self.self_attn = Lfm3VlVisionAttention(config) self.feed_forward = Lfm3VlMLP(config.vision_hidden_size, config.vision_intermediate_size) def forward(self, hidden_states, height, width): operator_input = self.operator_norm(hidden_states) operator_output = ( self.conv(operator_input, height, width) if self.layer_type == "conv" else self.self_attn(operator_input, height, width) ) hidden_states = hidden_states + operator_output return hidden_states + self.feed_forward(self.ffn_norm(hidden_states)) class Lfm3VlVisionEncoder(nn.Module): def __init__(self, config): super().__init__() self.spatial_merge_size = config.vision_spatial_merge_size self.patchifier = nn.Linear(config.vision_patch_dim, config.vision_hidden_size, bias=False) self.embedding_norm = Lfm3VlRMSNorm(config.vision_hidden_size, config.rms_norm_eps) self.layers = nn.ModuleList( [Lfm3VlVisionLayer(config, layer_type) for layer_type in config.vision_layer_types] ) @staticmethod def _block_major_to_raster(hidden_states, temporal, height, width, merge_size): return ( hidden_states.view(temporal, height // merge_size, width // merge_size, merge_size, merge_size, -1) .permute(0, 1, 3, 2, 4, 5) .reshape(temporal, height, width, -1) ) def forward(self, pixel_values, grid_thw): """Encode Qwen-preprocessed patches, one independent spatial sequence per temporal patch.""" outputs = [] offset = 0 for temporal, height, width in grid_thw.tolist(): temporal, height, width = int(temporal), int(height), int(width) tokens = temporal * height * width hidden_states = pixel_values[offset : offset + tokens] offset += tokens hidden_states = self._block_major_to_raster( hidden_states, temporal, height, width, self.spatial_merge_size ) for temporal_hidden_states in hidden_states: temporal_hidden_states = self.embedding_norm( self.patchifier(temporal_hidden_states.reshape(-1, hidden_states.shape[-1])) ) for layer in self.layers: temporal_hidden_states = layer(temporal_hidden_states, height, width) outputs.append((temporal_hidden_states, height, width)) if offset != pixel_values.shape[0]: raise ValueError("pixel_values and grid_thw disagree") return outputs class Lfm3VlProjector(nn.Module): def __init__(self, config): super().__init__() self.spatial_merge_size = config.vision_spatial_merge_size self.layer_norm = nn.LayerNorm(config.vision_hidden_size, eps=config.projector_layer_norm_eps) self.linear_1 = nn.Linear( self.spatial_merge_size**2 * config.vision_hidden_size, config.projector_hidden_size, bias=True, ) self.linear_2 = nn.Linear(config.projector_hidden_size, config.hidden_size, bias=True) def forward(self, hidden_states, height, width): if height % self.spatial_merge_size or width % self.spatial_merge_size: raise ValueError("projector requires even patch-grid dimensions") hidden_states = self.layer_norm(hidden_states).view(height, width, -1) merge = self.spatial_merge_size hidden_states = hidden_states.view(height // merge, merge, width // merge, merge, -1) hidden_states = hidden_states.permute(0, 2, 1, 3, 4).reshape(-1, merge**2 * hidden_states.shape[-1]) return self.linear_2(F.gelu(self.linear_1(hidden_states), approximate="none")) @dataclass class Lfm3VlCausalLMOutputWithPast(CausalLMOutputWithPast): image_hidden_states: torch.FloatTensor | None = None video_hidden_states: torch.FloatTensor | None = None class Lfm3VlPreTrainedModel(PreTrainedModel): config_class = Lfm3VlConfig base_model_prefix = "" input_modalities = ("image", "video", "text") _no_split_modules = ["Lfm3VlLanguageLayer", "Lfm3VlVisionLayer"] class Lfm3VlForConditionalGeneration(Lfm3VlPreTrainedModel, GenerationMixin): def __init__(self, config): super().__init__(config) self.language_model = Lfm3VlBackbone(config) self.vision_encoder = Lfm3VlVisionEncoder(config) self.vision_projector = Lfm3VlProjector(config) self.post_init() def get_input_embeddings(self): return self.language_model.embed_tokens @classmethod def _supports_default_dynamic_cache(cls): return False def set_input_embeddings(self, value): self.language_model.embed_tokens = value def get_output_embeddings(self): return self.language_model.embed_tokens def get_visual_features(self, pixel_values, grid_thw): target_dtype = self.vision_encoder.patchifier.weight.dtype encoded = self.vision_encoder(pixel_values.to(dtype=target_dtype), grid_thw) return [self.vision_projector(hidden_states, height, width) for hidden_states, height, width in encoded] def get_image_features(self, pixel_values, image_grid_thw): return self.get_visual_features(pixel_values, image_grid_thw) def get_video_features(self, pixel_values_videos, video_grid_thw): return self.get_visual_features(pixel_values_videos, video_grid_thw) def forward( self, input_ids=None, attention_mask=None, position_ids=None, pixel_values=None, image_grid_thw=None, pixel_values_videos=None, video_grid_thw=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, logits_to_keep=0, **kwargs, ): del kwargs if inputs_embeds is None: inputs_embeds = self.language_model.embed_tokens(input_ids) image_hidden_states = None if pixel_values is not None: if image_grid_thw is None: raise ValueError("image_grid_thw is required with pixel_values") image_features = self.get_image_features(pixel_values, image_grid_thw) image_hidden_states = torch.cat(image_features, dim=0).to(inputs_embeds.dtype) image_mask = input_ids == self.config.image_token_id if int(image_mask.sum()) != image_hidden_states.shape[0]: raise ValueError("image token count does not match projected image features") inputs_embeds = inputs_embeds.clone() inputs_embeds[image_mask] = image_hidden_states video_hidden_states = None if pixel_values_videos is not None: if video_grid_thw is None: raise ValueError("video_grid_thw is required with pixel_values_videos") video_features = self.get_video_features(pixel_values_videos, video_grid_thw) video_hidden_states = torch.cat(video_features, dim=0).to(inputs_embeds.dtype) video_mask = input_ids == self.config.video_token_id if int(video_mask.sum()) != video_hidden_states.shape[0]: raise ValueError("video token count does not match projected video features") if image_hidden_states is None: inputs_embeds = inputs_embeds.clone() inputs_embeds[video_mask] = video_hidden_states use_cache = self.config.use_cache if use_cache is None else use_cache if use_cache and past_key_values is None: past_key_values = Lfm3VlCache( self.config, inputs_embeds.shape[0], inputs_embeds.dtype, inputs_embeds.device ) hidden_states = self.language_model(inputs_embeds, attention_mask, position_ids, past_key_values) hidden_for_logits = ( hidden_states[:, -logits_to_keep:] if isinstance(logits_to_keep, int) and logits_to_keep else hidden_states ) logits = F.linear(hidden_for_logits, self.language_model.embed_tokens.weight) loss = None if labels is not None: loss = F.cross_entropy( logits[:, :-1].reshape(-1, logits.shape[-1]), labels[:, 1:].reshape(-1), ignore_index=-100 ) return Lfm3VlCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=past_key_values, image_hidden_states=image_hidden_states, video_hidden_states=video_hidden_states, ) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, position_ids=None, **kwargs ): if past_key_values is not None and past_key_values.get_seq_length() > 0: input_ids = input_ids[:, -1:] if position_ids is not None: position_ids = position_ids[:, -1:] kwargs["pixel_values"] = None kwargs["image_grid_thw"] = None kwargs["pixel_values_videos"] = None kwargs["video_grid_thw"] = None return { "input_ids": input_ids, "past_key_values": past_key_values, "attention_mask": attention_mask, "position_ids": position_ids, "use_cache": kwargs.get("use_cache", True), "pixel_values": kwargs.get("pixel_values"), "image_grid_thw": kwargs.get("image_grid_thw"), "pixel_values_videos": kwargs.get("pixel_values_videos"), "video_grid_thw": kwargs.get("video_grid_thw"), } __all__ = ["Lfm3VlForConditionalGeneration", "Lfm3VlPreTrainedModel"]