Super-squash branch 'main' using huggingface_hub

Co-authored-by: nielsr <nielsr@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

Notice

@@ -0,0 +1,2 @@
+CASA-Qwen2_5-VL-3B is finetuned from Qwen2.5-VL-3B with additional CASA layers.
+Qwen is licensed under the Qwen LICENSE AGREEMENT, Copyright (c) Alibaba Cloud. All Rights Reserved.

README.md

@@ -0,0 +1,92 @@
+---
+base_model:
+- Qwen/Qwen2.5-VL-3B-Instruct
+datasets:
+- HuggingFaceM4/FineVision
+- mvp-lab/LLaVA-OneVision-1.5-Instruct-Data
+language:
+- en
+license: cc-by-nc-sa-4.0
+pipeline_tag: image-text-to-text
+library_name: transformers
+---
+
+# CASA-Qwen2_5-VL-3B
+
+This repository contains the model weights for **CASA-Qwen2_5-VL-3B**, introduced in the paper [CASA: Cross-Attention over Self-Attention for Efficient Vision-Language Fusion](https://huggingface.co/papers/2512.19535).
+
+This model is a [Qwen-2.5VL-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) model adapted from token insertion to a cross-attention-based architecture.
+
+- **Paper:** [CASA: Cross-Attention over Self-Attention for Efficient Vision-Language Fusion](https://arxiv.org/abs/2512.19535)
+- **Project Page:** [kyutai.org/casa](https://kyutai.org/casa)
+- **Code:** [github.com/kyutai-labs/casa](https://github.com/kyutai-labs/casa)
+
+## Sample Usage
+
+This model requires `trust_remote_code=True` to load the custom architecture. Below is a snippet to run inference using `transformers`.
+
+```python
+import torch
+from transformers.models.auto.modeling_auto import AutoModel
+from transformers.models.auto.processing_auto import AutoProcessor
+
+model_id = "kyutai/CASA-Qwen2_5-VL-3B"
+model = AutoModel.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    trust_remote_code=True,
+).cuda()
+
+processor = AutoProcessor.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+)
+
+conversation = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/ai2d-demo.png",
+            },
+            {
+                "type": "text",
+                "text": "Describe this image.",
+            },
+        ],
+    },
+]
+
+inputs = processor.tokenize_messages(messages=conversation)
+inputs = inputs.to(model.device)
+input_len = inputs["input_ids"].shape[1]
+
+output_ids = model.generate_from_image(
+  **inputs,
+  max_new_tokens=512,
+  pre_image_tokens=processor.pre_image_tokens,
+  post_image_tokens=processor.post_image_tokens,
+  eos_token_id=model.generation_config.eos_token_id,
+)[0, input_len:]
+
+response = processor.tokenizer.decode(output_ids, skip_special_tokens=True)
+print(response)
+```
+
+## Citation
+
+```bibtex
+@article{kyutai2025casa,
+  author = {Moritz B\"ohle and Am\'elie Royer and Juliette Marrie and Edouard Grave and Patrick P\'erez},
+  year = {2025},
+  title = {CASA: Cross-Attention over Self-Attention for Efficient Vision-Language Fusion},
+  journal = {ArXiv},
+  url = {https://arxiv.org/abs/2512.19535}
+}
+```
+
+## License
+
+The code in the official repository is provided under the **MIT license**. The weights for this model are released under the **CC-BY-NC-SA 4.0 license**. Additionally, as this model includes weights from Qwen2.5-VL-3B, it is subject to the [Qwen RESEARCH LICENSE AGREEMENT](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct/blob/main/LICENSE).

config.json

@@ -0,0 +1,80 @@
+{
+  "architectures": [
+    "V2Qwen2_5VL"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_qwen2_5vl_ca.Qwen2_5_VLCAConfig",
+    "AutoModel": "modeling_qwen2_5vl_ca.V2Qwen2_5VL"
+  },
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "cross_attention": true,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "image_token_id": 151655,
+  "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "max_position_embeddings": 128000,
+  "max_window_layers": 70,
+  "model_type": "qwen2_5vl_ca",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "mrope_section": [
+      16,
+      24,
+      24
+    ],
+    "rope_type": "default",
+    "type": "default"
+  },
+  "rope_theta": 1000000.0,
+  "sliding_window": 32768,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "video_token_id": 151656,
+  "vision_config": {
+    "depth": 32,
+    "fullatt_block_indexes": [
+      7,
+      15,
+      23,
+      31
+    ],
+    "hidden_act": "silu",
+    "hidden_size": 1280,
+    "image_mean": [
+      0.48145466,
+      0.4578275,
+      0.40821073
+    ],
+    "image_std": [
+      0.26862954,
+      0.26130258,
+      0.27577711
+    ],
+    "in_channels": 3,
+    "in_chans": 3,
+    "intermediate_size": 3420,
+    "model_type": "qwen2_5_vl",
+    "num_heads": 16,
+    "out_hidden_size": 2048,
+    "patch_size": 14,
+    "spatial_merge_size": 2,
+    "spatial_patch_size": 14,
+    "temporal_patch_size": 1,
+    "tokens_per_second": 2,
+    "window_size": 112
+  },
+  "vision_end_token_id": 151653,
+  "vision_start_token_id": 151652,
+  "vision_token_id": 151654,
+  "vocab_size": 151936,
+  "xa_layers": []
+}

configuration_qwen2_5vl_ca.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig
+
+
+class Qwen2_5_VLCAConfig(Qwen2_5_VLConfig):
+    """Qwen2.5-VL config augmented with cross-attention options"""
+
+    model_type = "qwen2_5vl_ca"
+
+    def __init__(
+        self,
+        *args: Any,
+        # Our fusion mechanism
+        xa_layers: None | tuple = None,
+        cross_attention: bool = False,
+        xa_share_qkvo: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(*args, **kwargs)
+        self.head_dim = self.hidden_size // self.num_attention_heads
+        # Our fusion mechanisms
+        self.xa_layers = xa_layers
+        self.cross_attention = cross_attention
+        self.xa_share_qkvo = xa_share_qkvo
+        # Derived from vision_start/end_token_id set by parent
+        self.pre_image_tokens = (
+            [self.vision_start_token_id]
+            if getattr(self, "vision_start_token_id", None) is not None
+            else []
+        )
+        self.post_image_tokens = (
+            [self.vision_end_token_id]
+            if getattr(self, "vision_end_token_id", None) is not None
+            else []
+        )

cross_attention.py

@@ -0,0 +1,396 @@
+"""Modular cross-attention module"""
+
+import bisect
+import copy
+from itertools import accumulate
+from typing import TYPE_CHECKING, Literal, TypedDict, TypeVar
+from typing import cast as type_cast
+
+import torch
+
+from .utils import StreamingModule, StreamingState
+
+if TYPE_CHECKING:
+    from transformers.configuration_utils import PretrainedConfig
+
+
+try:
+    from flash_attn import flash_attn_varlen_func
+except ImportError:
+    flash_attn_varlen_func = None  # type: ignore
+
+
+AttentionBaseT = TypeVar("AttentionBaseT", bound=StreamingModule)
+
+
+WindowsComputeKwargs = TypedDict(
+    "WindowsComputeKwargs",
+    {
+        "num_post_image_tokens": int,
+        "num_pre_image_tokens": int,
+    },
+    total=False,
+)
+
+
+def get_sample_lengths_for_xa(
+    image_embeds_insertion_points: list[torch.Tensor],
+    image_embeds: torch.Tensor | list[torch.Tensor] | None,
+    total_seq_len: int,
+    attention_mask: torch.Tensor | None = None,
+    **kwargs: WindowsComputeKwargs,
+) -> tuple[list[tuple[int, bool]], list[int], torch.Tensor | None]:
+    """Sample lengths for cross-attention. Compared to other functions in this file,
+    it also returns a mask on the text tokens to mark tokens which do not relate
+    to any images (e.g. squashed text-only-samples, or BoS in prefix_after_bos)
+    """
+    num_post_image_tokens = type_cast(int, kwargs.get("num_post_image_tokens", 0))
+    num_pre_image_tokens = type_cast(int, kwargs.get("num_pre_image_tokens", 0))
+    squashed_samples_lengths = type_cast(
+        list[list[int]] | None, kwargs.get("squashed_samples_lengths", None)
+    )
+    if squashed_samples_lengths is not None:
+        assert len(squashed_samples_lengths) == len(image_embeds_insertion_points)
+
+    def __insert_next_sample__(
+        batch_idx: int, insrt_pt: int, last_insrt_pt: int, end_of_batch_sample: bool = False
+    ) -> None:
+        nonlocal attention_mask, active_tokens
+        nonlocal text_sample_lengths, full_sample_lengths
+        nonlocal cum_samples_lengths, current_image_offset
+        # Add the sample between [last_insrt_pt, insrt_pt] with breaks in
+        # between any squashed samples we find on the way
+        nonlocal last_image_idx, current_image_idx, current_length
+        added_sample = False
+        start_pt = bisect.bisect_left(cum_samples_lengths, last_insrt_pt)
+        for end_of_sample in cum_samples_lengths[start_pt:]:
+            # we will break the loop at the end when end_of_sample = insrt_pt
+            end_of_sample = min(end_of_sample, insrt_pt)
+
+            # Add between [last_insrt_pt, end_of_sample]
+            current_length = end_of_sample - last_insrt_pt
+            num_padding_tokens = 0
+            if attention_mask is not None:
+                num_padding_tokens = int(
+                    torch.sum(~attention_mask[batch_idx, last_insrt_pt:end_of_sample]).item()
+                )
+                current_length -= num_padding_tokens
+            num_image_tokens = 0
+            if current_length > 0:
+                # add image tokens to current_length
+                added_sample = True
+                if current_image_idx > 0 and image_embeds is not None:
+                    images_in_sample = [
+                        img_idx
+                        for img_idx in range(last_image_idx, current_image_idx)
+                        if img_idx < len(image_embeds_insertion_points[batch_idx])
+                        and last_insrt_pt
+                        <= image_embeds_insertion_points[batch_idx][img_idx]
+                        < end_of_sample
+                    ]
+                    if len(images_in_sample) > 0:
+                        num_image_tokens = sum(
+                            _x.shape[0]
+                            for _x in image_embeds[
+                                current_image_offset + images_in_sample[0] : current_image_offset
+                                + images_in_sample[-1]
+                                + 1
+                            ]
+                        )
+                # If no image, we should not insert and instead make it as inactive
+                if num_image_tokens > 0:
+                    text_sample_lengths.append(
+                        (current_length, end_of_batch_sample and insrt_pt == end_of_sample)
+                    )
+                    full_sample_lengths.append(current_length + num_image_tokens)
+            # Active tokens
+            active_tokens += [int(num_image_tokens > 0)] * (current_length + num_padding_tokens)
+
+            # prepare for next loop
+            last_insrt_pt = end_of_sample
+            if end_of_sample == insrt_pt:
+                break
+        # End of loop: catching edge case where we end up on a span full of padding
+        if end_of_batch_sample:
+            assert added_sample, "Weird edge case. Don't do that, thank you"
+            text_sample_lengths[-1] = (text_sample_lengths[-1][0], True)
+
+    current_image_offset = 0
+    text_sample_lengths, full_sample_lengths = [], []
+    cum_samples_lengths: list[int] = []
+    active_tokens = []
+    current_length, last_insrt_pt, last_image_idx, current_image_idx = 0, 0, 0, 0
+    for batch_idx, pts in enumerate(image_embeds_insertion_points):
+        if squashed_samples_lengths is not None:
+            cum_samples_lengths = list(accumulate(squashed_samples_lengths[batch_idx]))
+        else:
+            cum_samples_lengths = [total_seq_len]
+        for current_image_idx, insrt_pt in enumerate(pts.cpu().tolist()):
+            # check if the images are consecutive in which way we want
+            # them to belong to the same window
+            if current_image_idx >= 1 and insrt_pt == (
+                image_embeds_insertion_points[batch_idx][current_image_idx - 1]
+                + num_pre_image_tokens
+                + num_post_image_tokens
+            ):
+                continue
+            # Otherwise, we found a new sample
+            # not very important but for completeness: the insertion points come *after*
+            # the pre-image tokens per design but for the document-id mask it is more consistent to
+            # have them correspond to the same image
+            insrt_pt -= num_pre_image_tokens
+
+            # Compute length between the two insertion points
+            current_length = insrt_pt - last_insrt_pt
+            if attention_mask is not None:
+                current_length -= int(
+                    torch.sum(~attention_mask[batch_idx, last_insrt_pt:insrt_pt]).item()
+                )
+
+            # Update text and full sample lengths
+            if insrt_pt > last_insrt_pt:
+                __insert_next_sample__(
+                    batch_idx, insrt_pt, last_insrt_pt, end_of_batch_sample=False
+                )
+            last_image_idx = current_image_idx
+            last_insrt_pt = insrt_pt
+
+        # End of batch: add sample in progress and reset
+        current_image_idx += 1
+        if cum_samples_lengths[-1] > last_insrt_pt:
+            __insert_next_sample__(
+                batch_idx, cum_samples_lengths[-1], last_insrt_pt, end_of_batch_sample=True
+            )
+        current_length, last_insrt_pt, last_image_idx, current_image_idx = 0, 0, 0, 0
+        current_image_offset += len(pts)
+
+    # Sample lengths
+    if image_embeds is None:
+        return text_sample_lengths, full_sample_lengths, None
+
+    return (
+        text_sample_lengths,
+        full_sample_lengths,
+        torch.tensor(active_tokens, dtype=torch.bool, device=image_embeds[0].device),
+    )
+
+
+class CrossAttentionHandler:
+    def __init__(
+        self,
+        inputs_embeds: torch.Tensor,
+        image_embeds: torch.Tensor | list[torch.Tensor],
+        image_embeds_insertion_points: list[torch.Tensor] | None,
+        # info for building text->image windows link
+        ca_windows_info: None | WindowsComputeKwargs = None,
+        training: bool = True,
+    ):
+        if image_embeds_insertion_points is None:
+            image_embeds_insertion_points = [
+                torch.tensor(
+                    [0] * len(image_embeds),  # type: ignore[arg-type]
+                    dtype=torch.long,
+                    device=image_embeds[0].device,  # type: ignore[index]
+                )
+            ]
+
+        # Create cu_seq_lens for queries (text tokens)
+        # Compute sample lengths based on image insertion points to get cu_seq_lens
+        text_sample_lengths, full_sample_lengths, self.active_tokens_mask = (
+            get_sample_lengths_for_xa(
+                image_embeds_insertion_points=image_embeds_insertion_points,
+                image_embeds=image_embeds,
+                total_seq_len=inputs_embeds.shape[1],
+                **(ca_windows_info or {}),  # pyright: ignore[reportArgumentType]
+            )
+        )
+
+        if self.active_tokens_mask is None:
+            self.active_tokens_mask = torch.zeros(
+                (inputs_embeds.shape[0] * inputs_embeds.shape[1],),
+                dtype=torch.bool,
+                device=inputs_embeds.device,
+            )
+        assert sum(_x[0] for _x in text_sample_lengths) == int(
+            torch.sum(self.active_tokens_mask).item()
+        ), "Sanity check"
+
+        self.cu_seqlens_q = torch.Tensor(
+            list(accumulate([_x[0] for _x in text_sample_lengths], initial=0))
+        ).to(dtype=torch.int32, device=inputs_embeds.device)
+        self.max_seqlen_q = max(_x[0] for _x in text_sample_lengths)
+
+        # Create cu_seq_lens for keys values (the image tokens) while grouping the
+        # images which are consecutive
+        image_lens = [(l2 - l1) for (l1, _), l2 in zip(text_sample_lengths, full_sample_lengths)]
+        self.cu_seqlens_kv = torch.Tensor(list(accumulate(image_lens, initial=0))).to(
+            dtype=torch.int32, device=inputs_embeds.device
+        )
+        self.max_seqlen_kv = max(image_lens)
+        self.image_embeds = torch.cat([_x for _x in image_embeds], dim=0)[None, :, :]
+
+    def get_active_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Return tokens to be used as queries while ignoring
+        tokens who have nothing to do with an image"""
+        channels = hidden_states.shape[-1]
+        src = hidden_states.flatten(0, 1)
+        if self.active_tokens_mask is None:
+            return src.reshape((1, -1, channels))
+        return torch.masked_select(src, self.active_tokens_mask[:, None]).reshape((1, -1, channels))
+
+    def replace_active_tokens(
+        self, token_updates: torch.Tensor, hidden_states_in: torch.Tensor
+    ) -> torch.Tensor:
+        if self.active_tokens_mask is None:
+            return token_updates
+        updates = torch.zeros_like(hidden_states_in.flatten(0, 1))
+        updates.masked_scatter_(source=token_updates, mask=self.active_tokens_mask[:, None])
+        return updates
+
+
+def tie_qkvo_projections(self_attn: torch.nn.Module, cross_attn: "CrossAttention") -> None:
+    """Alias the q/k/v/o projections of `cross_attn` onto those of `self_attn`.
+
+    Used in the `xa_share_qkvo` setting so the cross-attention reuses the host
+    self-attention's projection weights (a single set of parameters). The two attention
+    calls and their independent softmaxes are otherwise unchanged.
+
+    :param self_attn: the host self-attention module owning the q/k/v/o projections
+    :param cross_attn: the cross-attention module whose projections are aliased
+    """
+    for proj in ("q_proj", "k_proj", "v_proj", "o_proj"):
+        setattr(cross_attn, proj, getattr(self_attn, proj))
+
+
+class CrossAttention(StreamingModule[StreamingState]):
+    """Attention module between images and text tokens"""
+
+    def __init__(
+        self,
+        config: "PretrainedConfig",
+        layer_idx: int | None,
+        input_layernorm: torch.nn.Module | None = None,
+    ):
+        super().__init__(StreamingState)
+        self.head_dim = config.head_dim
+        self.config = config
+
+        self.is_first_ca_layer = layer_idx == (min(config.xa_layers) if config.xa_layers else 0)
+
+        # When weights are shared with the host self-attention, the q/k/v/o projections
+        # are not owned by this module; they are aliased onto the self-attn projections
+        # by the host model (see tie_qkvo_projections).
+        self.xa_share_qkvo: bool = getattr(config, "xa_share_qkvo", False)
+        if not self.xa_share_qkvo:
+            self.q_proj = self.init_from_config_proj("q", config)
+            self.k_proj = self.init_from_config_proj("k", config)
+            self.v_proj = self.init_from_config_proj("v", config)
+            self.o_proj = self.init_from_config_proj("o", config)
+
+        self.norm: torch.nn.Module | None = copy.deepcopy(input_layernorm)
+        self.cross_attention_handler = None
+        # (source image embeddings, projected keys, projected values); the projections
+        # are reused across decoding steps as long as the source tensor is unchanged
+        self._cached_image_kv: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
+
+    def init_from_config_proj(
+        self, key: Literal["q", "o", "k", "v"], config: "PretrainedConfig"
+    ) -> torch.nn.Linear:
+        """Initialize the Linear proj in this module"""
+        num_heads = config.num_key_value_heads if key in {"k", "v"} else config.num_attention_heads
+        return torch.nn.Linear(
+            config.hidden_size,
+            num_heads * config.head_dim,
+            bias=config.attention_bias if key != "o" else False,
+        )
+
+    def reset_streaming(self):
+        super().reset_streaming()
+        self._cached_image_kv = None
+
+    def forward(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self, hidden_states: torch.Tensor, cross_attention_handler: CrossAttentionHandler | None
+    ) -> torch.Tensor | None:
+        if self.is_streaming:
+            if self.cross_attention_handler is None:
+                self.cross_attention_handler = cross_attention_handler
+            else:
+                # extend the shared handler
+                cross_attention_handler = self.cross_attention_handler
+                if self.is_first_ca_layer:
+                    cross_attention_handler.active_tokens_mask = None
+                    cross_attention_handler.cu_seqlens_q = torch.tensor(
+                        range(0, hidden_states.shape[0] + 1),
+                        dtype=cross_attention_handler.cu_seqlens_q.dtype,
+                        device=cross_attention_handler.cu_seqlens_q.device,
+                    )
+
+        # Case of text-only samples (or inference when no handler was cached)
+        # in this case we just want to skip cross-attention
+        if cross_attention_handler is None:
+            return None
+
+        # Currently, we assume that images never change/get added on the fly at inference
+        if self.is_streaming and self.streaming_state.offset > 0:
+            assert hidden_states.shape[0] == (len(cross_attention_handler.cu_seqlens_kv) - 1)
+
+        og_dtype = hidden_states.dtype
+        og_shape = hidden_states.shape
+
+        # kv inputs: (1, num_total_image_tokens, dim)
+        q_inputs = cross_attention_handler.get_active_tokens(hidden_states)
+        kv_inputs = cross_attention_handler.image_embeds
+
+        if self.norm is not None:
+            q_inputs = self.norm(q_inputs)
+        assert q_inputs.shape[0] == kv_inputs.shape[0] == 1
+
+        # Compute QKV for the blockwise attention
+        bs = 1
+        hidden_shape_q = (bs, q_inputs.shape[1], -1, self.head_dim)
+        query_states = self.q_proj(q_inputs).view(*hidden_shape_q)
+
+        # The image keys/values are identical at every decoding step, so at inference
+        # they are computed on the first call and reused until the images change
+        if self._cached_image_kv is not None and self._cached_image_kv[0] is kv_inputs:
+            _, key_states, value_states = self._cached_image_kv
+        else:
+            normed_kv = self.norm(kv_inputs) if self.norm is not None else kv_inputs
+            hidden_shape_kv = (bs, kv_inputs.shape[1], -1, self.head_dim)
+            key_states = self.k_proj(normed_kv).view(*hidden_shape_kv)
+            value_states = self.v_proj(normed_kv).view(*hidden_shape_kv)
+            if self.is_streaming:
+                self._cached_image_kv = (kv_inputs, key_states, value_states)
+
+        assert flash_attn_varlen_func is not None, (
+            "flash_attention is not installed but required for block-wise attention"
+        )
+
+        assert cross_attention_handler.cu_seqlens_q[-1] == query_states.shape[1], (
+            f"{cross_attention_handler.cu_seqlens_q[-1]} != {query_states.shape[1]}"
+        )
+        attn_output: torch.Tensor = flash_attn_varlen_func(
+            query_states[0].to(torch.bfloat16),
+            key_states[0].to(torch.bfloat16),
+            value_states[0].to(torch.bfloat16),
+            cu_seqlens_q=cross_attention_handler.cu_seqlens_q,
+            cu_seqlens_k=cross_attention_handler.cu_seqlens_kv,
+            max_seqlen_q=cross_attention_handler.max_seqlen_q,
+            max_seqlen_k=cross_attention_handler.max_seqlen_kv,
+            dropout_p=0.0,
+            # No need for causality when cross-attending to image tokens since
+            # image tokens are never padded
+            causal=False,
+        ).to(og_dtype)
+
+        attn_output = attn_output.reshape(hidden_shape_q[1], -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        # Reshape from flattened to non-flattened
+        attn_output = cross_attention_handler.replace_active_tokens(attn_output, hidden_states)
+        attn_output = attn_output.reshape(og_shape)
+
+        if self.is_streaming:
+            self.streaming_state.offset += attn_output.shape[1]
+
+        return attn_output

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "pad_token_id": 151643,
+  "eos_token_id": [
+    151643,
+    151645
+  ],
+  "transformers_version": "4.51.3"
+}

image_encoder.py

@@ -0,0 +1,90 @@
+"""Qwen2.5VL encoder with delayed normalization"""
+
+import torch
+from einops import rearrange
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VisionTransformerPretrainedModel,
+)
+
+
+class _LinearPatchEmbedProj(torch.nn.Module):
+    """Replaces Conv3d in Qwen's patch embed after collapsing the temporal dim.
+
+    The training codebase replaces the Conv3d with a plain Linear for efficiency
+    when temporal_patch_size=1.  The checkpoint therefore stores the weight under
+    ``patch_embed.proj.linear.weight`` rather than ``patch_embed.proj.weight``,
+    so this wrapper is required for state-dict key compatibility.
+    """
+
+    def __init__(self, in_features: int, out_features: int) -> None:
+        super().__init__()
+        self.linear = torch.nn.Linear(in_features, out_features, bias=False)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.linear.weight
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        n = x.shape[0]
+        return self.linear(x.reshape(n, -1)).view(n, -1, 1, 1, 1)
+
+
+def _replace_patch_embed_proj(visual: "Qwen2_5_VisionTransformerPretrainedModel") -> None:
+    """Swap the Conv3d patch-embed proj to a Linear so parameter names match the checkpoint."""
+    proj = visual.patch_embed.proj
+    if not isinstance(proj, torch.nn.Conv3d):
+        return
+    _, in_ch, _t, p, _ = proj.weight.shape
+    visual.patch_embed.temporal_patch_size = 1
+    visual.patch_embed.proj = _LinearPatchEmbedProj(in_ch * p * p, proj.out_channels)  # type: ignore[assignment]
+
+
+def prepare_for_qwen_encoder(
+    x: torch.Tensor | list[torch.Tensor], mean: torch.Tensor, std: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Preprocessing for Qwen encoder
+    Image mean and std come from processor.image_processor.image_mean and image_std
+    """
+    grid_thw = torch.Tensor([[1, img.shape[0], img.shape[1]] for img in x]).to(x[0].device)
+    hws_flatten_shape = torch.prod(grid_thw, dim=-1)
+    x = torch.cat(
+        [img.reshape((int(hws_flatten_shape[idx].item()), -1)) for idx, img in enumerate(x)],
+        dim=0,
+    )
+    assert x.min() >= 0.0 and x.max() <= 1.0
+    og_shape = x.shape
+    x = rearrange(x, "L (c d) -> L c d", c=3)
+    x = (x - mean) / std
+    x = x.view(og_shape).to(torch.bfloat16)
+    return x, grid_thw
+
+
+class Qwen25VLEncoder(torch.nn.Module):
+    """Qwen2.5 VL encoder with pre/post processing compatible with our CA implementation"""
+
+    def __init__(
+        self,
+        visual: "Qwen2_5_VisionTransformerPretrainedModel",
+    ):
+        super().__init__()
+        self.visual = visual
+        # Match training checkpoint: Conv3d is replaced by a Linear for t=1 images
+        _replace_patch_embed_proj(self.visual)
+        self.image_mean = torch.tensor(self.visual.config.image_mean).view(1, 3, 1)
+        self.image_std = torch.tensor(self.visual.config.image_std).view(1, 3, 1)
+
+    def forward(
+        self, x: torch.Tensor | list[torch.Tensor]
+    ) -> dict[str, torch.Tensor | list[torch.Tensor]]:
+        x, grid_thw = prepare_for_qwen_encoder(
+            x, mean=self.image_mean.to(x[0].device), std=self.image_std.to(x[0].device)
+        )
+
+        grid_thw = grid_thw.type(torch.int)
+        assert len(x) == grid_thw.prod(dim=1).sum()
+        out = self.visual(x, grid_thw=grid_thw)
+
+        split_sizes = (grid_thw.prod(dim=-1) // self.visual.spatial_merge_size**2).tolist()
+        embeds = list(torch.split(out, split_sizes, dim=0))  # Ni * (seq, C)
+        return {"image_embeds": embeds, "grid_thw": grid_thw}

language_qwen2_5vl_ca.py

@@ -0,0 +1,128 @@
+from typing import Literal, Optional
+
+import torch
+from transformers.cache_utils import Cache
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VLDecoderLayer,
+    Qwen2_5_VLFlashAttention2,
+)
+
+from .cross_attention import (
+    CrossAttention,
+    CrossAttentionHandler,
+    tie_qkvo_projections,
+)
+from .configuration_qwen2_5vl_ca import Qwen2_5_VLCAConfig
+
+
+class QwenCrossAttention(CrossAttention):
+    """A CrossAttention layer compatible with Qwen's projection conventions"""
+
+    def __init__(
+        self,
+        config: Qwen2_5_VLCAConfig,
+        layer_idx: int | None,
+    ):
+        super().__init__(config, layer_idx)  # pyright: ignore[reportArgumentType]
+        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        assert config.rope_scaling is not None
+        self.mrope_section = config.rope_scaling["mrope_section"] * 2
+
+    def init_from_config_proj(
+        self, key: Literal["q", "o", "k", "v"], config: PretrainedConfig
+    ) -> torch.nn.Linear:
+        """Follows modeling_qwen2_5_vl.py initialization"""
+        head_dim = config.hidden_size // config.num_attention_heads
+        if key == "q":
+            return torch.nn.Linear(
+                config.hidden_size, config.num_attention_heads * head_dim, bias=True
+            )
+        if key in {"k", "v"}:
+            return torch.nn.Linear(
+                config.hidden_size, config.num_key_value_heads * head_dim, bias=True
+            )
+        if key == "o":
+            return torch.nn.Linear(
+                config.num_attention_heads * config.head_dim, config.hidden_size, bias=False
+            )
+        raise NotImplementedError(f"Unknown key {key}")
+
+
+class Qwen2_5_VLAttention_CrossAttention(Qwen2_5_VLFlashAttention2):
+    """
+    Qwen Attention with extra CrossAttention layer
+    """
+
+    def __init__(
+        self,
+        config: Qwen2_5_VLCAConfig,
+        layer_idx: Optional[int] = None,
+        input_layernorm: torch.nn.Module | None = None,
+    ):
+        super().__init__(config, layer_idx)  # pyright: ignore[reportArgumentType]
+        self.cross_attn = QwenCrossAttention(config, layer_idx=layer_idx)
+        self.cross_attention_handler: CrossAttentionHandler | None = None
+        if getattr(config, "xa_share_qkvo", False):
+            tie_qkvo_projections(self, self.cross_attn)
+
+    @classmethod
+    def from_qwen2_5_vl_attention(
+        cls, attention: Qwen2_5_VLFlashAttention2, input_layernorm: torch.nn.Module | None
+    ):
+        """Init this layer from an existing Qwen Attention layer"""
+        layer_idx = attention.layer_idx
+        assert layer_idx is not None
+        new_attention = cls(attention.config, layer_idx=layer_idx, input_layernorm=input_layernorm)  # pyright: ignore
+        new_attention.load_state_dict(attention.state_dict(), strict=False)
+        return new_attention
+
+    def forward(  # pyright: ignore[reportIncompatibleMethodOverride]
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+    ):
+        attn_output, attn_weights, past_key_values = super().forward(
+            hidden_states,
+            attention_mask,
+            position_ids,
+            past_key_value,
+            output_attentions,
+            use_cache,
+            cache_position,
+            position_embeddings,
+        )
+        if self.cross_attn is not None:
+            ca_out = self.cross_attn(
+                hidden_states=hidden_states,
+                cross_attention_handler=self.cross_attention_handler,
+            )
+            # ca_out is None when there is no handler (text-only or streaming non-first call)
+            if ca_out is not None:
+                attn_output = ca_out + attn_output
+        return attn_output, attn_weights, past_key_values
+
+
+def maybe_replace_with_cross_attention_layers(
+    m: torch.nn.Module, xa_layers: tuple[int, ...] | None, reindex: bool = False
+):
+    """Replace Attention layer by CrossAttention layer as needed"""
+    if isinstance(m, Qwen2_5_VLDecoderLayer):
+        layer_idx = m.self_attn.layer_idx
+        assert layer_idx is not None
+        if xa_layers is None or len(xa_layers) == 0 or layer_idx in xa_layers:
+            m.self_attn = Qwen2_5_VLAttention_CrossAttention.from_qwen2_5_vl_attention(
+                m.self_attn, input_layernorm=m.input_layernorm
+            )
+        elif reindex:
+            # shift left by number of cross-attention layers before this one
+            logical_idx = layer_idx - sum(j < layer_idx for j in (xa_layers or ()))
+            assert logical_idx >= 0
+            m.self_attn.layer_idx = logical_idx

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0831fdee11174fc55e3bf54d8b107d3efc299cb5bfb2c373f1230894f0be8721
+size 8187682152

modeling_qwen2_5vl_ca.py

@@ -0,0 +1,304 @@
+from functools import partial
+from typing import Any, Sequence
+from typing import cast as type_cast
+
+import torch
+from transformers.cache_utils import DynamicCache
+from transformers.generation.utils import GenerateOutput
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VLCausalLMOutputWithPast,
+    Qwen2_5_VLForConditionalGeneration,
+)
+
+from .cross_attention import CrossAttentionHandler, tie_qkvo_projections
+from .image_encoder import Qwen25VLEncoder
+from .configuration_qwen2_5vl_ca import Qwen2_5_VLCAConfig
+from .language_qwen2_5vl_ca import (
+    Qwen2_5_VLAttention_CrossAttention,
+    QwenCrossAttention,
+    maybe_replace_with_cross_attention_layers,
+)
+
+
+class V2Qwen2_5VL(Qwen2_5_VLForConditionalGeneration):  # pyright: ignore[reportIncompatibleMethodOverride]
+    config_class = Qwen2_5_VLCAConfig
+
+    def __init__(self, config: Qwen2_5_VLCAConfig, **kwargs: Any) -> None:
+        del kwargs
+        super().__init__(config)
+        # Wrap the Qwen visual encoder so its output matches our CA interface
+        self.image_prefix = Qwen25VLEncoder(self.visual)  # type: ignore[assignment]
+        self.visual = None
+        self.model.apply(
+            partial(maybe_replace_with_cross_attention_layers, xa_layers=self.config.xa_layers)
+        )
+        # The cross-attention layers are swapped in after the base post_init, so register
+        # the shared-weight alias keys and (re-)tie now that the CA modules exist.
+        if config.xa_share_qkvo:
+            shared_keys: list[str] = []
+            for i, layer in enumerate(self.model.layers):
+                if isinstance(layer.self_attn, Qwen2_5_VLAttention_CrossAttention):
+                    for proj, biased in (
+                        ("q_proj", True),
+                        ("k_proj", True),
+                        ("v_proj", True),
+                        ("o_proj", False),
+                    ):
+                        prefix = f"model.layers.{i}.self_attn.cross_attn.{proj}"
+                        shared_keys.append(f"{prefix}.weight")
+                        if biased:
+                            shared_keys.append(f"{prefix}.bias")
+            self._tied_weights_keys = list(self._tied_weights_keys or []) + shared_keys
+            self.tie_weights()
+
+    def _tie_weights(self) -> None:
+        if not getattr(self.config, "xa_share_qkvo", False):
+            return
+        for layer in self.model.layers:
+            if isinstance(layer.self_attn, Qwen2_5_VLAttention_CrossAttention):
+                tie_qkvo_projections(layer.self_attn, layer.self_attn.cross_attn)
+
+    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
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs: Any,
+        model_kwargs: dict[str, Any],
+        is_encoder_decoder: bool = False,
+        num_new_tokens: int = 1,
+    ):
+        """Override to handle multi-turn generation and propagate updated attention masks"""
+        if (am := outputs.get("updated_attention_mask", None)) is not None:
+            model_kwargs["attention_mask"] = am
+            if "updated_cache_position" in outputs:
+                model_kwargs["cache_position"] = outputs.get("updated_cache_position")
+            else:
+                start = 0
+                if (kv := model_kwargs.get("past_key_values", None)) is not None:
+                    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,
+                )
+
+        # Call parent to get default updates
+        model_kwargs = super()._update_model_kwargs_for_generation(
+            outputs, model_kwargs, is_encoder_decoder, num_new_tokens
+        )
+        # Used by prepare_inputs_for_generation
+        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,
+    ):
+        """Override to avoid Qwen erasing pixel_values on subsequent generation calls"""
+        backup = None
+        __is_first_gen_call__ = kwargs.get("__is_first_gen_call__", True)
+        if __is_first_gen_call__:
+            backup = kwargs.get("pixel_values", None)
+        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,
+            )
+        out = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            **kwargs,
+        )
+        if backup is not None:
+            out["pixel_values"] = backup
+        return out
+
+    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: dict = {
+            "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))
+            image_embeds = processed_inputs.get("image_embeds")
+            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(  # type: ignore[override] # pylint: disable=W0221
+        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,
+    ) -> tuple | Qwen2_5_VLCausalLMOutputWithPast:
+        """Multi-modal forward pass"""
+        if self.training:
+            assert return_loss is True, (
+                "Qwen2.5VL always computes its own labels/losses in train mode"
+            )
+
+        if inputs_embeds is None:
+            assert input_ids is not None
+            inputs_embeds = type_cast(torch.Tensor, self.model.embed_tokens(input_ids))
+
+        # Case 1: First generation call — compute image embeddings and set up CA handler
+        if kwargs.pop("__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,
+            )
+            image_embeds = processed_inputs.get("image_embeds", None)
+            inst_points = processed_inputs.get("image_embeds_insertion_points", None)
+
+            # Only build a handler when images are actually present
+            cross_attention_handler: CrossAttentionHandler | None = None
+            if image_embeds is not None and len(image_embeds) > 0:
+                cross_attention_handler = CrossAttentionHandler(
+                    inputs_embeds=torch.zeros_like(inputs_embeds),
+                    image_embeds=image_embeds,
+                    image_embeds_insertion_points=inst_points,
+                    ca_windows_info=processed_inputs.pop("ca_windows_info", None),
+                    training=self.training,
+                )
+            self.update_cross_attention_states(cross_attention_handler)
+
+        # Run Qwen with the attention layers replaced to use cross-attention
+        assert inputs_embeds is not None, "Could not compute input embeddings!"
+        out = super().forward(
+            inputs_embeds=inputs_embeds,  # type: ignore[arg-type]
+            attention_mask=attention_mask,
+            pixel_values=None,
+            **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]
+        self,
+        reset_streaming: bool = True,
+        temperature: float | None = 0.0,
+        eos_token_id: int | Sequence[int] | None = None,
+        **kwargs: Any,
+    ) -> GenerateOutput | torch.LongTensor:
+        """Custom generate function"""
+        # 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:
+            kwargs["pad_token_id"] = kwargs.get("eos_token_id", None)
+            if isinstance(kwargs["pad_token_id"], (list, tuple)):
+                kwargs["pad_token_id"] = kwargs["pad_token_id"][0]
+        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)
+
+        if not kwargs.get("do_sample", False):
+            temperature = None
+            kwargs.pop("top_p", None)
+            kwargs.pop("top_k", None)
+
+        # Generate
+        self.start_ca_streaming_states()
+        outputs = self.generate(
+            use_cache=True,
+            eos_token_id=eos_token_id,
+            temperature=temperature,
+            **kwargs,
+        )
+        if reset_streaming:
+            self.reset_ca_streaming_states()
+        return outputs
+
+    def update_cross_attention_states(self, handler: CrossAttentionHandler | None):
+        """Push the new handler into all CA attention layers"""
+
+        def __update__(m: torch.nn.Module):
+            nonlocal handler
+            if isinstance(m, Qwen2_5_VLAttention_CrossAttention):
+                m.cross_attention_handler = handler
+
+        self.apply(__update__)
+
+    def reset_ca_streaming_states(self) -> None:
+        def __reset__(m: torch.nn.Module):
+            if isinstance(m, QwenCrossAttention):
+                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, QwenCrossAttention):
+                m._set_streaming(True, ())
+
+        self.apply(__start__)

processing.py

@@ -0,0 +1,494 @@
+# pylint: disable=no-member  # avoid weird pylint warnings from SentencePieceProcessor
+"""Text and Image processor for CA models using Qwen2.5_VL image encoder"""
+
+from math import ceil
+from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast, overload
+from typing import cast as type_cast
+
+import torch
+import torchvision.transforms.v2 as T
+from einops import rearrange
+from PIL import Image
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms.functional import to_tensor as pil_to_tensor
+from torchvision.transforms.v2 import functional as F
+from transformers.image_processing_utils import BaseImageProcessor
+from transformers.processing_utils import ProcessorMixin
+
+if TYPE_CHECKING:
+    from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
+    from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+
+
+ImageMessage = TypedDict(
+    "ImageMessage",
+    {
+        "type": Literal["image"],
+        "image": str | Image.Image | None,
+    },
+)
+
+TextMessage = TypedDict(
+    "TextMessage",
+    {
+        "type": Literal["text"],
+        "text": str,
+    },
+)
+
+MessageContent = list[ImageMessage | TextMessage]
+
+Message = TypedDict(
+    "Message",
+    {
+        "role": Literal["system", "user", "assistant"],
+        "content": MessageContent,
+    },
+)
+
+ProcessorInput = list[list[Message]] | list[Message]
+
+__INTERP_NAME_TO_MODE__ = {
+    "nearest": InterpolationMode.NEAREST,
+    "bilinear": InterpolationMode.BILINEAR,
+    "bicubic": InterpolationMode.BICUBIC,
+    "lanczos": InterpolationMode.LANCZOS,
+}
+
+__INTERP_INT_TO_MODE__ = {
+    0: InterpolationMode.NEAREST,
+    2: InterpolationMode.BILINEAR,
+    3: InterpolationMode.BICUBIC,
+    4: InterpolationMode.BOX,
+    5: InterpolationMode.HAMMING,
+    1: InterpolationMode.LANCZOS,
+}
+
+
+@overload
+def universal_resize(
+    img: Image.Image,
+    size: tuple[int, int],
+    interpolation: str | InterpolationMode | int = "bilinear",
+    antialias: bool = True,
+) -> Image.Image: ...
+@overload
+def universal_resize(
+    img: torch.Tensor,
+    size: tuple[int, int],
+    interpolation: str | InterpolationMode | int = "bilinear",
+    antialias: bool = True,
+) -> torch.Tensor: ...
+def universal_resize(
+    img: Image.Image | torch.Tensor,
+    size: tuple[int, int],
+    interpolation: str | InterpolationMode | int = "bilinear",
+    antialias: bool = True,
+) -> Image.Image | torch.Tensor:
+    """Resize that works for PIL.Image, CHW tensor, or BCHW tensor"""
+    if isinstance(interpolation, str):
+        interpolation = __INTERP_NAME_TO_MODE__[interpolation]
+    elif isinstance(interpolation, int):
+        interpolation = __INTERP_INT_TO_MODE__[interpolation]
+
+    return F.resize(
+        img, size, interpolation=type_cast(InterpolationMode, interpolation), antialias=antialias
+    )
+
+
+@overload
+def convert_to_rgb(img: Image.Image) -> Image.Image: ...
+@overload
+def convert_to_rgb(img: torch.Tensor) -> torch.Tensor: ...
+def convert_to_rgb(img: Image.Image | torch.Tensor) -> Image.Image | torch.Tensor:
+    """Convert any image to RGB in a way that does not throw PIL warning"""
+    if isinstance(img, torch.Tensor):
+        return img
+    if img.mode == "RGB":  # no changes
+        return img
+    if img.mode == "P":  # palette images need to be converted to RGBA first
+        return img.convert("RGBA").convert("RGB")
+    return img.convert("RGB")
+
+
+class QwenImageProcessor(BaseImageProcessor):
+    """Resizing for the Qwen2.5VL encoder. Note that the normalization is
+    handled in the image_encoder in the model forward"""
+
+    def __init__(
+        self,
+        img_size: int = 448,
+        interpolation: Literal["bicubic", "bilinear", "nearest", "nearest_exact"] = "bicubic",
+        max_ratio: int = 10,
+        round_to_patch_size: int = 56,
+        use_fast: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        self._num_target_channels = 588
+        self._merge_size = 2
+        self._patch_size = 14
+        super().__init__(
+            use_fast=use_fast,
+            do_normalize=False,
+            **kwargs,
+        )
+        self.img_size = img_size
+        self.interpolation = interpolation
+        self.max_ratio = max_ratio
+        self.round_to_patch_size = round_to_patch_size
+
+    def resize_transform(
+        self, img: Image.Image | torch.Tensor, img_size: int | None = None
+    ) -> Image.Image | torch.Tensor:
+        if img_size is None:
+            img_size = self.img_size
+        max_area = img_size**2
+        if isinstance(img, Image.Image):
+            img = convert_to_rgb(img)
+            w_og, h_og = img.size
+        else:
+            h_og, w_og = img.shape[-2:]
+        w, h = w_og, h_og
+
+        # Qwen requires max ratio of 10 between max and min sizes
+        if self.max_ratio > 0:
+            w, h = max(w, h // self.max_ratio), max(h, w // self.max_ratio)
+
+        # resize to max area
+        current_area = w * h
+        if current_area > max_area:
+            scale = (max_area / current_area) ** 0.5
+            w, h = int(w * scale), int(h * scale)
+
+        # resize to patch size
+        if self.round_to_patch_size > 0:
+            w = ceil(w / self.round_to_patch_size) * self.round_to_patch_size
+            h = ceil((h / self.round_to_patch_size)) * self.round_to_patch_size
+
+        # resize
+        if w != w_og or h != h_og:
+            img = universal_resize(img, (h, w), self.interpolation)
+        if isinstance(img, torch.Tensor):
+            img = T.ToDtype(torch.float32, scale=True)(T.ToImage()(img))
+        return img
+
+    def __process_one__(
+        self, video_or_img: Image.Image | torch.Tensor, img_size: int | None = None
+    ) -> torch.Tensor:
+        """Same operation as __process_one_with_processor__ but without going through numpy"""
+        video_or_img = self.resize_transform(video_or_img, img_size)
+        if isinstance(video_or_img, Image.Image):
+            video_or_img = pil_to_tensor(video_or_img)
+        assert isinstance(video_or_img, torch.Tensor)
+        if video_or_img.ndim == 3:
+            video_or_img = video_or_img[None]
+        assert video_or_img.ndim == 4 and video_or_img.shape[1] == 3, (
+            f"Invalid shape {video_or_img.shape}."
+        )
+        t, c, h, w = video_or_img.shape
+        p = self._patch_size
+        m = self._merge_size
+
+        # Convert to RGB
+        if c == 1:
+            video_or_img = video_or_img.expand((-1, 3, -1, -1))
+        if c == 4:
+            video_or_img = video_or_img[:, :3]
+        c = video_or_img.shape[1]
+        assert c == 3, "Expecting RGB image in QwenNormalize"
+
+        # Reshape to t h w c' format
+        h, w = video_or_img.shape[2] // p, video_or_img.shape[3] // p
+        rearrange_dict = dict(p1=p, p2=p, m1=m, m2=m)
+
+        video_or_img = rearrange(
+            video_or_img,
+            "t c (h m1 p1) (w m2 p2) -> (t h w m1 m2) (c p1 p2)",
+            **rearrange_dict,
+        )
+        assert video_or_img.shape[-1] == self._num_target_channels, (
+            f"{video_or_img.shape[-1]} != {self._num_target_channels}"
+        )
+        video_or_img = video_or_img.view((-1, h, w, self._num_target_channels))
+
+        return video_or_img
+
+    @overload
+    def process_images(
+        self, image: Image.Image | torch.Tensor, img_size: int | None = None
+    ) -> torch.Tensor: ...
+    @overload
+    def process_images(
+        self, image: list[Image.Image] | list[torch.Tensor], img_size: int | None = None
+    ) -> list[torch.Tensor]: ...
+    def process_images(
+        self,
+        image: Image.Image | torch.Tensor | list[Image.Image] | list[torch.Tensor],
+        img_size: int | None = None,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        if isinstance(image, list):
+            return [self.__process_one__(_x, img_size) for _x in image]
+        return self.__process_one__(image, img_size)
+
+
+class ProcessorOutput(dict):
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    image_embeds_insertion_points: list[torch.Tensor] | None
+    pixel_values: torch.Tensor | list[torch.Tensor] | None
+
+    def to(
+        self, device: torch.device | str, dtype: torch.dtype = torch.bfloat16
+    ) -> "ProcessorOutput":
+        return ProcessorOutput(
+            {
+                "input_ids": self["input_ids"].to(device),
+                "attention_mask": self["attention_mask"].to(device),
+                "image_embeds_insertion_points": self["image_embeds_insertion_points"],
+                "pixel_values": (
+                    self["pixel_values"].to(dtype).to(device)
+                    if isinstance(self["pixel_values"], torch.Tensor)
+                    else [x.to(dtype).to(device) for x in self["pixel_values"]]
+                    if self["pixel_values"] is not None
+                    else None
+                ),
+            }
+        )
+
+
+class BaseProcessor(ProcessorMixin):
+    def __init__(
+        self,
+        tokenizer: "PreTrainedTokenizerFast | Qwen2Tokenizer",
+        pre_image_tokens: tuple[int, ...] = (),
+        post_image_tokens: tuple[int, ...] = (),
+        system_start_tokens: tuple[int, ...] = (),
+        system_end_tokens: tuple[int, ...] = (),
+        user_start_tokens: tuple[int, ...] = (),
+        user_end_tokens: tuple[int, ...] = (),
+        asst_start_tokens: tuple[int, ...] = (),
+        asst_end_tokens: tuple[int, ...] = (),
+        allow_system_prompt: bool = True,
+        pad_token: int = 0,
+        bos_token: int | None = None,
+    ) -> None:
+        self.pre_image_tokens = list(pre_image_tokens)
+        self.post_image_tokens = list(post_image_tokens)
+        self.system_start_tokens = list(system_start_tokens)
+        self.system_end_tokens = list(system_end_tokens)
+        self.user_start_tokens = list(user_start_tokens)
+        self.user_end_tokens = list(user_end_tokens)
+        self.asst_start_tokens = list(asst_start_tokens)
+        self.asst_end_tokens = list(asst_end_tokens)
+        self._allow_system_prompt = allow_system_prompt
+        self.tokenizer = tokenizer
+        self._image_processor = None
+        self._pad_token = pad_token
+        self.bos_token = bos_token
+
+    @property
+    def image_processor(self) -> QwenImageProcessor:
+        assert self._image_processor is not None
+        return self._image_processor
+
+    def _process_content(
+        self,
+        message_content: MessageContent,
+        role: Literal["system", "user", "assistant"],
+        tokenized_messages: list[torch.Tensor],
+        insertion_points: list[int],
+        image_list: list[torch.Tensor | None],
+        token_count: int,
+        img_size: int | None = None,
+        **kwargs: Any,
+    ) -> int:
+        mapping = {
+            "user": (self.user_start_tokens, self.user_end_tokens),
+            "assistant": (self.asst_start_tokens, self.asst_end_tokens),
+            "system": (self.system_start_tokens, self.system_end_tokens),
+        }
+        if role.lower() not in mapping:
+            raise ValueError(f"Unknown role '{role}' encountered in messages.")
+        start_tokens, end_tokens = mapping[role.lower()]
+        # 1) Add the start tokens
+        if start_tokens:
+            tokenized_messages.append(torch.Tensor(start_tokens).flatten().to(torch.long))
+            token_count += len(start_tokens)
+        # 2) Process the message content one by one (potentially interleaved image and text)
+        for part in message_content:
+            elt_type = part["type"]
+            if elt_type == "image":
+                part = cast(ImageMessage, part)
+                self._process_image_message(
+                    part,
+                    tokenized_messages,
+                    image_list,
+                    img_size=img_size,
+                )
+                token_count += len(self.pre_image_tokens)
+                insertion_points.append(token_count)
+                token_count += len(self.post_image_tokens)
+            else:
+                part = cast(TextMessage, part)
+                self._process_text_message(
+                    part["text"],
+                    role=role,
+                    token_list=tokenized_messages,
+                    **kwargs,
+                )
+                token_count += tokenized_messages[-1].size(0)
+        # 3) Add the end tokens
+        if end_tokens:
+            tokenized_messages.append(torch.Tensor(end_tokens).flatten().to(torch.long))
+            token_count += len(end_tokens)
+        return token_count
+
+    def _process_text_message(
+        self,
+        message: str,
+        role: Literal["system", "user", "assistant"],
+        token_list: list[torch.Tensor],
+        **kwargs: Any,
+    ) -> None:
+        if role.lower() == "system" and not self._allow_system_prompt:
+            raise ValueError("System prompts are not allowed in this tokenizer configuration.")
+        tokens = self.tokenizer.encode(
+            message, add_special_tokens=False, return_tensors="pt", **kwargs
+        )
+        tokens = cast(torch.Tensor, tokens)
+        token_list.append(tokens.flatten().to(torch.long))
+
+    def _process_image_message(
+        self,
+        message: ImageMessage,
+        token_list: list[torch.Tensor],
+        image_list: list[torch.Tensor | None],
+        img_size: int | None = None,
+    ) -> None:
+        img = message["image"]
+        if img is None:
+            image_list.append(None)
+        else:
+            image_list.append(
+                self.image_processor.process_images(
+                    self._load_image(img), img_size=img_size
+                ).squeeze(0)
+            )
+        if self.pre_image_tokens:
+            token_list.append(torch.Tensor(self.pre_image_tokens).flatten().to(torch.long))
+
+        if self.post_image_tokens:
+            token_list.append(torch.Tensor(self.post_image_tokens).flatten().to(torch.long))
+
+    def _load_image(self, image_path_or_image: str | Image.Image) -> Image.Image:
+        if isinstance(image_path_or_image, str):
+            return Image.open(image_path_or_image).convert("RGB")
+        return image_path_or_image
+
+    def _maybe_pad(self, tokens: torch.Tensor, pad_len: int, pad_value: int) -> torch.Tensor:
+        return torch.nn.functional.pad(
+            tokens,
+            (0, pad_len) if self.tokenizer.padding_side == "right" else (pad_len, 0),
+            value=pad_value,
+        )
+
+    def pad_tokenized_messages(
+        self,
+        tokenized_messages_batch: list[torch.Tensor],
+        image_insertion_points_batch: list[torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]:
+        max_len = max(len(x) for x in tokenized_messages_batch)
+        if image_insertion_points_batch is not None and self.tokenizer.padding_side == "left":
+            image_insertion_points_batch = [
+                x + max_len - len(tokenized_messages_batch[idx])
+                for idx, x in enumerate(image_insertion_points_batch)
+            ]
+        input_ids = torch.stack(
+            [
+                self._maybe_pad(s, max_len - s.size(0), self._pad_token)
+                for s in tokenized_messages_batch
+            ],
+            dim=0,
+        )
+        attention_mask = torch.stack(
+            [
+                self._maybe_pad(torch.ones_like(s), max_len - s.size(0), 0)
+                for s in tokenized_messages_batch
+            ],
+            dim=0,
+        )
+        return input_ids, attention_mask, image_insertion_points_batch
+
+    def tokenize_messages(
+        self,
+        messages: ProcessorInput,
+        suppress_bos_token: bool = False,
+        **kwargs: Any,
+    ) -> ProcessorOutput | None:
+        """Tokenize a batch of messages into token IDs suitable for the CA model.
+
+        Args:
+            messages: Batch of message lists (or single list of messages),
+              where each message is a dict with 'role' and 'content' keys.
+            suppress_bos_token: If True, the BOS token will not be added.
+            **kwargs: Additional keyword arguments passed to the underlying encode method.
+        """
+        if not messages:
+            return None
+        if isinstance(messages[0], dict):
+            messages = [messages]  # type: ignore[assignment]
+
+        messages = cast(list[list[Message]], messages)
+        image_insertion_points_batch = []
+        tokenized_messages_batch = []
+        image_list: list[torch.Tensor | None] = []
+        for msgs in messages:
+            tokenized_messages = []
+            if not suppress_bos_token and self.bos_token is not None:
+                tokenized_messages.append(torch.tensor([self.bos_token], dtype=torch.long))
+            insertion_points = []
+            token_count = 0
+            for msg in msgs:
+                token_count = self._process_content(
+                    msg["content"],
+                    role=msg["role"],
+                    tokenized_messages=tokenized_messages,
+                    insertion_points=insertion_points,
+                    image_list=image_list,
+                    token_count=token_count,
+                    **kwargs,
+                )
+            tokenized_messages_batch.append(torch.cat(tokenized_messages, dim=0).to(torch.long))
+            image_insertion_points_batch.append(torch.tensor(insertion_points, dtype=torch.long))
+
+            if msgs and self.asst_end_tokens and msgs[-1]["role"].lower() == "assistant":
+                # Remove the assistant end tokens from the final message
+                end_token_len = len(self.asst_end_tokens)
+                tokenized_messages_batch[-1] = tokenized_messages_batch[-1][:-end_token_len]
+            if msgs and self.asst_start_tokens and msgs[-1]["role"].lower() == "user":
+                tokenized_messages_batch[-1] = torch.cat(
+                    [
+                        tokenized_messages_batch[-1],
+                        torch.Tensor(self.asst_start_tokens).to(torch.long),
+                    ]
+                )
+
+        input_ids, attention_mask, image_embeds_insertion_points = self.pad_tokenized_messages(
+            tokenized_messages_batch, image_insertion_points_batch
+        )
+
+        if image_list:
+            assert sum(img is None for img in image_list) % len(image_list) == 0, (
+                "Either all or no image must be None."
+            )
+        pixel_values: None | torch.Tensor | list[torch.Tensor]
+        if image_list[0] is None:
+            pixel_values = None
+        else:
+            pixel_values = cast(list[torch.Tensor], image_list)
+        return ProcessorOutput(
+            input_ids=input_ids,
+            image_embeds_insertion_points=image_embeds_insertion_points,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+        )

processing_qwen2_5vl_ca.py

@@ -0,0 +1,39 @@
+from typing import Any
+
+from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
+
+from .processing import BaseProcessor, QwenImageProcessor
+
+
+class QwenCASAProcessor(BaseProcessor):
+    attributes = ["tokenizer"]
+    tokenizer_class = "Qwen2Tokenizer"
+
+    def __init__(
+        self,
+        tokenizer: Qwen2Tokenizer,
+        pre_image_tokens: tuple[int, ...] = (151652,),
+        post_image_tokens: tuple[int, ...] = (151653,),
+        system_start_tokens: tuple[int, ...] = (151644, 8948, 198),
+        system_end_tokens: tuple[int, ...] = (151645, 198),
+        user_start_tokens: tuple[int, ...] = (151644, 872, 198),
+        user_end_tokens: tuple[int, ...] = (151645, 198),
+        asst_start_tokens: tuple[int, ...] = (151644, 77091, 198),
+        asst_end_tokens: tuple[int, ...] = (151645, 198),
+        image_size: int = 448,
+        **kwargs: Any,
+    ):
+        del kwargs
+        super().__init__(
+            tokenizer=tokenizer,
+            pre_image_tokens=pre_image_tokens,
+            post_image_tokens=post_image_tokens,
+            system_start_tokens=system_start_tokens,
+            system_end_tokens=system_end_tokens,
+            user_start_tokens=user_start_tokens,
+            user_end_tokens=user_end_tokens,
+            asst_start_tokens=asst_start_tokens,
+            asst_end_tokens=asst_end_tokens,
+        )
+
+        self._image_processor = QwenImageProcessor(img_size=image_size)

processor_config.json

@@ -0,0 +1,13 @@
+{
+    "auto_map": {
+        "AutoProcessor": "processing_qwen2_5vl_ca.QwenCASAProcessor"
+    },
+    "image_size": 896,
+    "post_image_tokens": [
+        151653
+    ],
+    "pre_image_tokens": [
+        151652
+    ],
+    "processor_class": "QwenCASAProcessor"
+}

tokenizer_config.json

@@ -0,0 +1,208 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "add_bos_token": false
+}

utils.py

@@ -0,0 +1,337 @@
+# pylint: disable=protected-access
+"""Utils to handle CA layers construction"""
+
+from contextlib import contextmanager
+from dataclasses import dataclass, fields
+from typing import Any, Callable, Generic, Literal, Sequence, TypeVar, overload
+from typing import cast as type_cast
+
+import torch
+
+
+def __split_n_merge__(
+    x: torch.Tensor,
+    sample_lengths: list[int],
+    padding_side: Literal["left", "right"] = "right",
+    pad_value: int | float | bool = 0,
+) -> torch.Tensor:
+    max_sample_length = max(sample_lengths)
+    pad_tuple = tuple(0 for _ in range((x.ndim - 1) * 2))
+    return torch.stack(
+        [
+            torch.nn.functional.pad(
+                _x,
+                pad_tuple + (0, max_sample_length - _x.shape[0])
+                if padding_side == "right"
+                else pad_tuple + (max_sample_length - _x.shape[0], 0),
+                value=pad_value,
+            )
+            for _x in torch.split(x, sample_lengths, dim=0)
+        ],
+        dim=0,
+    )
+
+
+@overload
+def insert_image_tokens(
+    inputs_embeds: torch.Tensor,
+    image_embeds: torch.Tensor | Sequence[torch.Tensor],
+    image_embeds_insertion_points: list[torch.Tensor],
+    recover_batch_dim: Literal[True],
+    attention_mask: torch.Tensor | None = None,
+    padding_side: Literal["left", "right"] = "right",
+    keep_only_attended: bool = False,
+    pad_output: int | float | bool = 0.0,
+) -> tuple[
+    torch.Tensor,
+    None,
+    torch.Tensor | None,
+    torch.Tensor,
+]: ...
+@overload
+def insert_image_tokens(
+    inputs_embeds: torch.Tensor,
+    image_embeds: torch.Tensor | Sequence[torch.Tensor],
+    image_embeds_insertion_points: list[torch.Tensor],
+    recover_batch_dim: Literal[False],
+    attention_mask: torch.Tensor | None = None,
+    padding_side: Literal["left", "right"] = "right",
+    keep_only_attended: bool = False,
+    pad_output: int | float | bool = 0.0,
+) -> tuple[
+    torch.Tensor,
+    list[int],
+    torch.Tensor | None,
+    torch.Tensor,
+]: ...
+def insert_image_tokens(
+    inputs_embeds: torch.Tensor,
+    image_embeds: torch.Tensor | Sequence[torch.Tensor],
+    image_embeds_insertion_points: list[torch.Tensor],
+    recover_batch_dim: bool = True,
+    attention_mask: torch.Tensor | None = None,
+    padding_side: Literal["left", "right"] = "right",
+    keep_only_attended: bool = False,
+    pad_output: int | float | bool = 0.0,
+) -> tuple[
+    torch.Tensor | torch.Tensor,
+    list[int] | None,
+    torch.Tensor | torch.Tensor | None,
+    torch.Tensor | torch.Tensor,
+]:
+    """
+    Insert image embeddings into text embeddings
+
+    Args:
+        inputs_embeds (torch.Tensor): (B, S, D) input token embeddings.
+        image_embeds (torch.Tensor | list[torch.Tensor]): (N_images, Nt, D) |  List[(Nt, D)] image token embeddings.
+        image_embeds_insertion_points (list[torch.Tensor]): Insertion indices.
+        attention_mask (torch.Tensor, optional): (B, S) attention mask.
+        padding_side (Literal["left", "right"]): Padding scheme. Controls behavior for padded images.
+        return_indices (bool): Whether to return gather indices or the fused sequence directly.
+        keep_only_attended: This is only applicable when recover_batch_dim is False; whether to
+            remove any non-attended tokens in the whole array. In this case, the attention
+            mask returned is **still the original one**, so we can remember which indices have been
+            removed
+    Returns:
+        output (torch.Tensor): (B, S + Ni * Nt) gather indices or (B, S + Ni * Nt, D) fused sequence
+        image_embeds (torch.Tensor): (B, Ni * Nt) image embeds, padded and batch if input was a list
+        attention_mask (torch.Tensor): Same shape, 1 for real tokens, 0 for image and text padding.
+        image_tokens_mask (torch.Tensor): (B, S + Ni * Nt, 1), marks image token positions.
+    """
+    if isinstance(image_embeds, list) and len(image_embeds) == 0:
+        batch_size, text_seq_length, token_dim = inputs_embeds.shape
+        if recover_batch_dim:
+            return (
+                inputs_embeds,
+                None,
+                attention_mask,
+                torch.zeros((batch_size, text_seq_length, 1), dtype=torch.bool),
+            )
+        else:
+            flattened_seq_length = inputs_embeds.shape[0] * inputs_embeds.shape[1]
+            return (
+                torch.reshape(inputs_embeds, (flattened_seq_length, inputs_embeds.shape[2])),
+                [text_seq_length] * inputs_embeds.shape[0],
+                attention_mask.flatten() if attention_mask is not None else None,
+                torch.zeros((flattened_seq_length, 1), dtype=torch.bool),
+            )
+
+    # Sanity checks
+    if isinstance(image_embeds, torch.Tensor):
+        assert inputs_embeds.shape[-1] == image_embeds.shape[-1]
+    else:
+        assert all(inputs_embeds.shape[-1] == _x.shape[-1] for _x in image_embeds)
+
+    batch_size, text_seq_length, token_dim = inputs_embeds.shape
+    image_seq_length = [x.shape[0] for x in image_embeds]
+
+    # Flatten insertion points
+    insertion_offset = []
+    counter, offset_from_text, offset_from_image = 0, 0, 0
+    for sample in image_embeds_insertion_points:
+        for pt in sample:
+            insertion_offset.append(pt + offset_from_image + offset_from_text)
+            offset_from_image += image_seq_length[counter]
+            counter += 1
+        offset_from_text += text_seq_length
+    image_insert_positions = [
+        x for idx, pt in enumerate(insertion_offset) for x in range(pt, pt + image_seq_length[idx])
+    ]
+
+    # Flatten image embeds
+    if isinstance(image_embeds, list):
+        image_embeds = torch.cat(image_embeds, dim=0)
+    else:
+        image_embeds = type_cast(torch.Tensor, image_embeds)
+        image_embeds = torch.reshape(image_embeds, (-1, token_dim))
+
+    # Flatten text embeds across batch dim (B x S, D)
+    inputs_embeds = torch.reshape(inputs_embeds, (-1, token_dim))
+    flattened_seq_length = inputs_embeds.shape[0] + sum(image_seq_length)
+    text_insert_positions = sorted(
+        set(range(flattened_seq_length)).difference(set(image_insert_positions))
+    )
+
+    # Scatter image embeds in the flattened dict
+    # scatter text related stuff
+    output = torch.empty(
+        (flattened_seq_length, token_dim),
+        device=inputs_embeds.device,
+        dtype=inputs_embeds.dtype,
+    )
+    txt_positions_tensor = torch.Tensor(text_insert_positions).to(
+        dtype=torch.long, device=inputs_embeds.device
+    )
+    output.scatter_(0, txt_positions_tensor[:, None].expand(-1, token_dim), inputs_embeds)
+    attention_mask_new: torch.Tensor | None = None
+    if attention_mask is not None:
+        attention_mask_new = torch.ones(
+            (flattened_seq_length,), dtype=torch.bool, device=inputs_embeds.device
+        )
+        attention_mask_new.scatter_(
+            0, txt_positions_tensor, attention_mask.flatten().to(torch.bool)
+        )
+
+    # scatter image related stuff
+    image_tokens_mask = torch.zeros(
+        (flattened_seq_length,), dtype=torch.bool, device=inputs_embeds.device
+    )
+    img_positions_tensor = torch.Tensor(image_insert_positions).to(
+        device=inputs_embeds.device, dtype=torch.long
+    )
+    output.scatter_(0, img_positions_tensor[:, None].expand(-1, token_dim), image_embeds)
+    image_tokens_mask.scatter_(0, img_positions_tensor, True)
+
+    # Compute expected sample length, taking into account the real batch
+    # i.e. recover the batch dimension of image embeddings
+    sample_lengths = []
+    counter = 0
+    for sample_idx, pts in enumerate(image_embeds_insertion_points):
+        num_image_tokens = 0
+        for _ in pts:
+            num_image_tokens += image_seq_length[counter]
+            counter += 1
+        if keep_only_attended and attention_mask is not None:
+            attended_seq_length = torch.sum(attention_mask[sample_idx]).cpu().item()
+            sample_lengths.append(attended_seq_length + num_image_tokens)
+        else:
+            sample_lengths.append(text_seq_length + num_image_tokens)
+
+    # For CA attention, we can keep stuff flatten and return
+    # the sample_lengths for the blockwise attention
+    if not recover_batch_dim:
+        if keep_only_attended and attention_mask_new is not None:
+            output = output[attention_mask_new]
+            image_tokens_mask = image_tokens_mask[attention_mask_new]
+        return output, sample_lengths, attention_mask_new, image_tokens_mask[..., None]
+
+    # Otherwise, time to (pad) and reshape
+    # Easy case: everything has the same length
+    if all(x == sample_lengths[0] for x in sample_lengths):
+        output = torch.reshape(output, (batch_size, sample_lengths[0], token_dim))
+        image_tokens_mask = torch.reshape(image_tokens_mask, (batch_size, sample_lengths[0], 1))
+        if attention_mask_new is not None:
+            attention_mask_new = torch.reshape(attention_mask_new, (batch_size, sample_lengths[0]))
+    # if there is any size mismatch we break into a
+    # list and pad again
+    else:
+        # split and merge
+        output = __split_n_merge__(output, sample_lengths, padding_side, pad_value=pad_output)
+        # note that the extra padding tokens are also marked as image tokens to be removed later
+        image_tokens_mask = __split_n_merge__(
+            image_tokens_mask, sample_lengths, padding_side, True
+        )[:, :, None]
+        if attention_mask_new is not None:
+            attention_mask_new = __split_n_merge__(
+                attention_mask_new, sample_lengths, padding_side, 0
+            )
+    # Return
+    return output, sample_lengths, attention_mask_new, image_tokens_mask
+
+
+class SharedModuleType(type):
+    """Wrapper to build shared Pytorch modules. This can be used as a metaclass to build shared
+    modules; see an example in attention.py"""
+
+    _instances = {}
+
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+        if cls not in cls._instances:
+            cls._instances[cls] = super(SharedModuleType, cls).__call__(*args, **kwargs)
+        return cls._instances[cls]
+
+
+@dataclass
+class StreamingState:
+    """Streaming State used by CA layers at inference to save
+    e.g. the offset and other persistent states"""
+
+    offset: int = 0
+
+    def _is_valid_field(self, key: str) -> bool:
+        return key in {x.name for x in fields(self)}
+
+    def _init_field(self, key: str) -> None:
+        """Init function for non-argument dependent defaults"""
+        assert self._is_valid_field(key)
+        if key == "offset":
+            self.offset = 0
+        else:
+            # for fields which should be set explicitly and cannot be auto-initialized
+            setattr(self, key, None)
+
+    def init(self) -> None:
+        for key in [x.name for x in fields(self)]:
+            self._init_field(key)
+
+    def _reset_field(self, name: str) -> None:
+        """Resets the given field"""
+        self._init_field(name)
+
+    def reset(self) -> None:
+        for f in fields(self):
+            self._reset_field(f.name)
+
+    def _get_field(self, f: str) -> Any:
+        """Get field and init if not"""
+        assert self._is_valid_field(f)
+        if getattr(self, f) is None:
+            self._init_field(f)
+        return getattr(self, f)
+
+    def _set_field(self, f: str, value: Any) -> None:
+        assert self._is_valid_field(f)
+        setattr(self, f, value)
+
+
+StreamingStateT = TypeVar("StreamingStateT", bound=StreamingState)
+
+
+class StreamingModule(torch.nn.Module, Generic[StreamingStateT]):  # pylint: disable=abstract-method
+    """Streaming-aware module base class"""
+
+    def __init__(self, state_class: type) -> None:
+        torch.nn.Module.__init__(self)
+        self.is_streaming: bool = False
+        self.enable_viz: tuple[str, ...] = ()
+        self._streaming_state: StreamingStateT = state_class()
+
+    @property
+    def streaming_state(self) -> StreamingStateT:
+        return self._streaming_state
+
+    def _apply_named_streaming(self, fn: Callable):
+        """Apply function to all streaming modules"""
+        for name, module in self.named_modules():
+            if isinstance(module, StreamingModule):
+                fn(name, module)
+
+    def reset_streaming(self):
+        """Reset the streaming state."""
+
+        def _reset(_: str, module: StreamingModule):
+            module._streaming_state.reset()
+
+        self._apply_named_streaming(_reset)
+
+    def _set_streaming(self, streaming: bool, viz: tuple[str, ...] = ()):
+        """Set all streaming modules in streaming mode"""
+
+        def _set_streaming(_, module: StreamingModule) -> None:
+            module.is_streaming = streaming
+            module.enable_viz = viz
+            if streaming:
+                module.streaming_state.init()
+
+        self._apply_named_streaming(_set_streaming)
+
+    @contextmanager
+    def streaming(self, stream: bool = True, viz: tuple[str, ...] = ()):
+        """Context manager to enter streaming mode. Reset streaming state on exit."""
+        self._set_streaming(stream, viz)
+        try:
+            yield
+        finally:
+            self._set_streaming(False, ())
+            self.reset_streaming()