from typing import Any, Callable, Sequence from typing import cast as type_cast import torch from transformers.cache_utils import DynamicCache from transformers.configuration_utils import PretrainedConfig from transformers.generation.utils import GenerateOutput from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionTransformerPretrainedModel, ) from .cross_attention import CrossAttention from .image_encoder import Qwen25VLEncoder from .configuration_helium1_ca import Helium1CAConfig from .language_helium1_ca import ( CausalHeliumOutput, Helium1ForCausalLM, Helium1RMSNorm, ) def meta_project( logits: torch.Tensor | list[torch.Tensor], projector: torch.nn.Module, norm: torch.nn.Module | None = None, ) -> torch.Tensor | list[torch.Tensor]: """Projection operation that handles both tensors and list of tensors Outputs either a (N, S, D) tensor (same resolution images) or a list of N (S, D) tensors (where S can be a different sequence length per image) """ split_sizes: list[int] | None = None if not isinstance(logits, torch.Tensor): split_sizes = [_x.shape[0] for _x in logits] logits = torch.cat(logits, dim=0)[None, :, :] logits = type_cast(torch.Tensor, logits) logits = projector(logits) assert isinstance(logits, torch.Tensor) if norm is not None: logits = norm(logits) if split_sizes is not None: return list(torch.split(type_cast(torch.Tensor, logits[0]), split_sizes, dim=0)) return logits class ImageProjection(torch.nn.Module): """Takes in a batch or sequence of images and returns embeddings which are then fed to the LM. :param config: Helium1CAConfig object :param lm_model_dim: Output dimension (number of channels) for this module """ def __init__(self, config: PretrainedConfig, lm_model_dim: int) -> None: super().__init__() self.config = config self.out_dim = lm_model_dim # Inherit the parent model's attention implementation (default flash_attention_2). # Without this the standalone ViT silently falls back to "sdpa", whose dense # block-diagonal mask fails to isolate images in a multi-image batch, leaking # attention across images and corrupting per-image embeddings at batch_size > 1. visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( config.vision_config, attn_implementation=getattr(config, "_attn_implementation", None) or "flash_attention_2", ) self.enc = Qwen25VLEncoder(visual=visual) # Projection layer self.proj_xa = self.init_proj_module() # Output normalization — only when training checkpoint has this norm self.norm_xa: Helium1RMSNorm | None = ( Helium1RMSNorm(self.out_dim) if config.image_proj_norm else None ) def init_proj_module(self) -> torch.nn.Module: """Init the projection module for the cross-attended image tokens""" if self.config.vision_config.out_dim == self.out_dim: return torch.nn.Identity() return torch.nn.Linear(self.config.vision_config.out_dim, self.out_dim) def forward( self, x: torch.Tensor | list[torch.Tensor] ) -> dict[ str, torch.Tensor | list[torch.Tensor], ]: """Image embedding mapping :param x: Either a tensor with shape (Bi, C, H, W) or a list of Bi tensors with shape (C, H, W) (or (H, W, C) in the case of Qwen) :return: Either a tensor with shape (num_total_image, S, D) or, if images can have different seq length, a list of `num_total_images` Tensors with shape (S, D) """ # Apply image encoder og_dtype = x[0].dtype encoded = self.enc(x)["image_embeds"] encoded = [_x.to(og_dtype) for _x in encoded] if all(x.shape[0] == encoded[0].shape[0] for x in encoded): encoded = torch.stack(encoded, dim=0) # Extra projection image_embeds = meta_project(encoded, self.proj_xa, self.norm_xa) return {"cross_attention_src": image_embeds} class V2Helium1(Helium1ForCausalLM): # pyright: ignore[reportIncompatibleMethodOverride] config_class = Helium1CAConfig def __init__(self, config: Helium1CAConfig, **kwargs: Any) -> None: del kwargs super().__init__(config) self.image_prefix = ImageProjection(config=config, lm_model_dim=self.token_dim) def get_device(self) -> str: """Return the device type of the model""" return next(self.parameters()).device.type @property def token_dim(self) -> int: """Returns the number of dimensions for the token representation""" return self.config.hidden_size @property def rotary_embed(self) -> Callable: """Returns the rotary embedding function of the underlying model""" return self.model.rotary_emb def _update_model_kwargs_for_generation( self, outputs: Any, model_kwargs: dict[str, Any], is_encoder_decoder: bool = False, num_new_tokens: int = 1, ): # After token insertion on the first call, the KV cache holds more tokens than # model_kwargs["cache_position"] tracks. Resync both attention_mask and cache_position # so subsequent decode steps use the correct positions. if (am := outputs.get("updated_attention_mask", None)) is not None: model_kwargs["attention_mask"] = am start = 0 if (kv := model_kwargs.get("past_key_values", None)) is not None and hasattr( kv, "_seen_tokens" ): start = kv._seen_tokens - am.shape[1] model_kwargs["cache_position"] = torch.arange( start, start + am.shape[1], dtype=model_kwargs["cache_position"].dtype, device=model_kwargs["cache_position"].device, ) model_kwargs = super()._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder, num_new_tokens ) model_kwargs["__is_first_gen_call__"] = False return model_kwargs def prepare_inputs_for_generation( # pyright: ignore[reportIncompatibleMethodOverride] self, input_ids: torch.Tensor, past_key_values: DynamicCache | None = None, **kwargs: Any, ): __is_first_gen_call__ = kwargs.get("__is_first_gen_call__", True) if past_key_values is not None and ( kwargs.get("cache_position") is None or type_cast(torch.Tensor, kwargs.get("cache_position")).shape[0] == 0 ): # We're continuing from a cached state past_length = past_key_values._seen_tokens kwargs["cache_position"] = torch.arange( past_length, past_length + (input_ids.shape[1] if __is_first_gen_call__ else 1), dtype=torch.long, device=input_ids.device, ) return super().prepare_inputs_for_generation( type_cast(torch.LongTensor, input_ids), past_key_values=past_key_values, **kwargs, ) def prepare_multimodal_inputs( self, input_ids: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, image_embeds_insertion_points: list[torch.Tensor] | None = None, labels: torch.Tensor | None = None, pixel_values: torch.Tensor | list[torch.Tensor] | None = None, pre_image_tokens: list[int] | None = None, post_image_tokens: list[int] | None = None, **_kwargs: Any, ) -> dict: """Get a batch data mixing text and image data""" del _kwargs processed_inputs = { "input_ids": input_ids, "inputs_embeds": inputs_embeds, "labels": labels, "attention_mask": attention_mask, "image_embeds_insertion_points": image_embeds_insertion_points, } if pixel_values is not None: processed_inputs.update(self.image_prefix(pixel_values)) # In insertion mode the forward expects "image_embeds", not "cross_attention_src" if not self.config.cross_attention and "cross_attention_src" in processed_inputs: processed_inputs["image_embeds"] = processed_inputs.pop("cross_attention_src") image_embeds = processed_inputs.get( "image_embeds", processed_inputs.get("cross_attention_src") ) assert image_embeds is not None assert (isinstance(image_embeds, torch.Tensor) and image_embeds.ndim == 3) or ( isinstance(image_embeds, list) and all(_x.ndim == 2 for _x in image_embeds) ) # Add kwargs necessary to compute cu_seqlens windows for CA processed_inputs["ca_windows_info"] = { "num_post_image_tokens": 0 if post_image_tokens is None else len(post_image_tokens), "num_pre_image_tokens": 0 if pre_image_tokens is None else len(pre_image_tokens), } return processed_inputs def forward( # pyright: ignore[reportIncompatibleMethodOverride] self, input_ids: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, pixel_values: torch.Tensor | list[torch.Tensor] | None = None, return_loss: bool = True, labels: torch.Tensor | None = None, image_embeds_insertion_points: list[torch.Tensor] | None = None, pre_image_tokens: list[int] | None = None, post_image_tokens: list[int] | None = None, **kwargs: Any, ) -> CausalHeliumOutput: """Multi modal forward pass""" assert input_ids is not None or inputs_embeds is not None if self.training: assert return_loss is True, ( "Helium models always compute its own labels/losses in train mode" ) # Case 1: First generation call — compute pixel values and set up CA handler if kwargs.get("__is_first_gen_call__", True): processed_inputs = self.prepare_multimodal_inputs( input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, image_embeds_insertion_points=image_embeds_insertion_points, pixel_values=pixel_values, labels=labels, pre_image_tokens=pre_image_tokens, post_image_tokens=post_image_tokens, ) else: if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) processed_inputs = { "inputs_embeds": inputs_embeds, "attention_mask": attention_mask, } # For token insertion mode: update cache position to account for inserted image tokens _insertion_image_embeds = processed_inputs.get("image_embeds") if ( not self.config.cross_attention and (cp := kwargs.get("cache_position", None)) is not None and _insertion_image_embeds is not None ): # Compute the max per-sample image token count (insert_image_tokens pads to this length) image_seq_lengths = [_x.shape[0] for _x in _insertion_image_embeds] _insertion_points = processed_inputs.get("image_embeds_insertion_points") or [] counter = 0 max_img_tokens = 0 for pts in _insertion_points: sample_total = sum(image_seq_lengths[counter + j] for j in range(len(pts))) max_img_tokens = max(max_img_tokens, sample_total) counter += len(pts) num_image_tokens = max_img_tokens start = kwargs["cache_position"][0].item() num_tokens = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] # type: ignore kwargs["cache_position"] = torch.arange( start, start + num_tokens + num_image_tokens, dtype=cp.dtype, device=cp.device, ) kwargs.pop("__is_first_gen_call__", True) out = super().forward( **processed_inputs, # type: ignore **kwargs, ) return out @property def default_generation_eos_token_id(self) -> int | Sequence[int] | None: return self.generation_config.eos_token_id if self.generation_config is not None else None @torch.no_grad() def generate_from_image( # pyright: ignore[reportInconsistentOverload,reportIncompatibleMethodOverride] self, reset_streaming: bool = True, eos_token_id: int | Sequence[int] | None = None, **kwargs: Any, ) -> "GenerateOutput | torch.LongTensor": # init self-attention KVCache if kwargs.get("past_key_values", None) is None: kwargs["past_key_values"] = DynamicCache() if eos_token_id is None: eos_token_id = self.default_generation_eos_token_id # To avoid generate warning if kwargs.get("pad_token_id", None) is None: pad = eos_token_id kwargs["pad_token_id"] = pad[0] if isinstance(pad, (list, tuple)) else pad if "pre_image_tokens" not in kwargs: kwargs["pre_image_tokens"] = list(self.config.pre_image_tokens) if "post_image_tokens" not in kwargs: kwargs["post_image_tokens"] = list(self.config.post_image_tokens) self.start_ca_streaming_states() outputs = self.generate( use_cache=True, eos_token_id=eos_token_id, **kwargs, ) if reset_streaming: self.reset_ca_streaming_states() return outputs def reset_ca_streaming_states(self) -> None: def __reset__(m: torch.nn.Module): if isinstance(m, CrossAttention): m._set_streaming(False, ()) m.reset_streaming() if hasattr(m, "cross_attention_handler"): del m.cross_attention_handler m.cross_attention_handler = None self.apply(__reset__) def start_ca_streaming_states(self) -> None: def __start__(m: torch.nn.Module): if isinstance(m, CrossAttention): m._set_streaming(True, ()) self.apply(__start__)