init

.gitattributes

@@ -1,35 +1,3 @@
-*.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
+*.pth filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,44 @@
+# Ignore non-release weights; release checkpoints are tracked with Git LFS.
+*.safetensors
+*.bin
+*.pt
+*.pth
+*.ckpt
+*.gguf
+*.onnx
+
+# Weight index — references .safetensors, meaningless without them
+model.safetensors.index.json
+
+!model.safetensors
+!layout_analysis/last.pth
+
+# Training artifacts (not part of the modeling spec)
+trainer_state.json
+training_args.bin
+zero_to_fp32.py
+optimizer.pt
+scheduler.pt
+rng_state.pth
+
+# Misc
+*.bak
+.DS_Store
+__pycache__/
+*.pyc
+*.egg-info/
+build/
+dist/
+.pytest_cache/
+.venv/
+
+# Local coding-assistant instructions are not model-release assets.
+[Aa][Gg][Ee][Nn][Tt][Ss].md
+[Aa][Gg][Ee][Nn][Tt][Ss].override.md
+[Cc][Ll][Aa][Uu][Dd][Ee].md
+[Gg][Ee][Mm][Ii][Nn][Ii].md
+copilot-instructions.md
+.claude/
+.codex/
+.agents/
+.cursor/

LICENSE

@@ -0,0 +1,140 @@
+# OXR License
+
+Copyright (c) 2026 SoMark Team
+
+OXR is licensed under the Apache License, Version 2.0, subject to the
+additional terms below. These additional terms take precedence in the event
+of a conflict; otherwise, Apache License 2.0 applies.
+
+## 1. Commercial License and Thresholds
+
+Commercial use of OXR and its Derivative Works is permitted without a
+separate commercial license only while you and your Affiliates, on a
+consolidated basis, do not exceed either of the following thresholds:
+
+(a) USD 2,000,000 in Total Funding; or
+(b) USD 2,000,000 in Annual Revenue.
+
+If either threshold is exceeded, you must obtain a separate commercial
+license from SoMark Team before commencing or continuing commercial use.
+Being below these thresholds does not waive any other term of this License,
+including the competing hosted API restriction in Section 3.
+
+## 2. Online Service Attribution
+
+When providing an online service to third parties using OXR or its Derivative
+Works, you must prominently identify SoMark OXR in the service's user interface or
+publicly accessible documentation.
+
+## 3. Competing Hosted API Services
+
+Unless expressly authorized in advance by a separate written license from
+SoMark Team, you may not use OXR or its Derivative Works to offer, operate,
+or make available to third parties a hosted API service that competes with
+a hosted API service offered by SoMark Team or its Affiliates.
+
+This restriction includes services offered through a reseller, a white-label
+arrangement, or an intermediary interface, whether paid or free. A competing
+hosted API service is one whose principal functionality provides document
+parsing, optical character recognition (OCR), document layout analysis, or
+document content extraction as a substitute for substantially similar
+functionality of a hosted API service offered by SoMark Team or its Affiliates.
+
+This restriction applies regardless of Total Funding or Annual Revenue.
+Use solely within your organization and its Affiliates, without offering such
+a service to third parties, is not prohibited by this Section; all other
+terms of this License still apply.
+
+## 4. Termination
+
+Your rights under this License terminate automatically, without notice from
+SoMark Team, if you fail to obtain a commercial license as required by
+Section 1, fail to comply with Section 2, or violate Section 3.
+
+## 5. Definitions and Scope
+
+"Affiliates" are entities that directly or indirectly control you, are
+controlled by you, or share common control with you. "Control" means the
+power to direct an entity's management or operations through ownership,
+voting rights, contractual rights, or other means.
+
+"Total Funding" means cumulative equity and debt financing raised from all
+sources by you and your Affiliates, on a consolidated basis, excluding
+intra-group transactions to avoid double counting.
+
+"Annual Revenue" means the gross revenue of you and your Affiliates for the
+most recently completed fiscal year, on a consolidated basis, excluding
+intra-group transactions to avoid double counting. Amounts in other
+currencies are measured at their US dollar equivalent.
+
+"Derivative Works" has the meaning given in Apache License 2.0. When
+redistributing OXR or its Derivative Works, you must include these additional
+terms together with Apache License 2.0. Third-party components separately
+identified as being under other licenses remain governed by those licenses.
+
+The full text of Apache License 2.0 is available at
+https://www.apache.org/licenses/LICENSE-2.0 .
+
+---
+
+# OXR 许可协议
+
+版权所有 (c) 2026 SoMark Team
+
+OXR 依据 Apache License 2.0 进行许可，同时受下列附加条款约束。若存在冲突，
+以本附加条款为准；其他事项适用 Apache License 2.0。
+
+## 1. 商业许可与门槛
+
+仅当您及您的关联方按合并口径计算，均未超过以下任一门槛时，方可无需另行
+取得商业许可而将 OXR 及其衍生作品用于商业用途：
+
+(a) 累计融资总额 2,000,000 美元；或
+(b) 年度总收入 2,000,000 美元。
+
+一旦超过任一门槛，您必须在开始或继续商业使用之前，从 SoMark Team 取得
+单独的商业许可。未超过上述门槛并不免除本许可中的其他义务，包括第 3 条
+关于竞争性托管 API 服务的限制。
+
+## 2. 在线服务标识义务
+
+如果您使用 OXR 或其衍生作品向第三方提供在线服务，必须在该服务的用户界面
+或公开可访问的文档中，显著标明使用了 SoMark OXR。
+
+## 3. 竞争性托管 API 服务限制
+
+除非事先通过与 SoMark Team 签订单独的书面许可取得明确授权，否则您不得
+使用 OXR 或其衍生作品，向第三方提供、运营或开放与 SoMark Team 或其关联方
+所提供的托管 API 服务构成竞争的托管 API 服务。
+
+本限制包括通过转售、白标或中间接口提供的服务，无论收费或免费。竞争性
+托管 API 服务，是指以文档解析、光学字符识别（OCR）、文档版面分析或文档
+内容提取为主要功能，并能够替代 SoMark Team 或其关联方提供的托管 API
+服务中实质相似功能的服务。
+
+本限制不受累计融资总额或年度总收入门槛影响。仅在您所属组织及其关联方内部
+使用，且不向第三方提供上述服务的，不属于本条禁止的行为；本许可的其他
+条款仍然适用。
+
+## 4. 许可终止
+
+如果您未按第 1 条要求取得商业许可、未遵守第 2 条标识义务，或违反第 3 条，
+您在本许可项下的权利将自动终止，SoMark Team 无须另行通知。
+
+## 5. 定义与适用范围
+
+“关联方”是指直接或间接控制您、受您控制，或与您处于共同控制之下的实体。
+“控制”是指通过所有权、表决权、合同权利或其他方式，决定某一实体管理或
+经营事项的能力。
+
+“累计融资总额”是指您及您的关联方按合并口径计算，从所有来源累计取得的
+股权融资和债务融资总额，剔除集团内部交易以避免重复计算。
+
+“年度总收入”是指您及您的关联方最近一个已结束财年的总收入，按合并口径
+计算，并剔除集团内部交易以避免重复计算。其他币种金额按其等值美元计算。
+
+“衍生作品”的含义与 Apache License 2.0 中的定义一致。再分发 OXR 或其
+衍生作品时，必须同时附上本附加条款和 Apache License 2.0。单独标明适用
+其他许可的第三方组件，仍受各自许可约束。
+
+Apache License 2.0 全文见：https://www.apache.org/licenses/LICENSE-2.0 。

README.md

@@ -1,3 +0,0 @@
----
-license: apache-2.0
----

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,18 @@
+{%- for message in messages -%}
+{{- '<|im_start|>' + message['role'] + '\n' -}}
+{%- if message['content'] is string -%}
+{{- message['content'] -}}
+{%- else -%}
+{%- for content in message['content'] -%}
+{%- if content['type'] == 'image_url' or content['type'] == 'image' -%}
+{{- '<image>' -}}
+{%- elif content['type'] == 'text' -%}
+{{- content['text'] -}}
+{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{{- '<|im_end|>\n' -}}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+{{- '<|im_start|>assistant\n' -}}
+{%- endif -%}

config.json

@@ -0,0 +1,109 @@
+{
+  "architectures": [
+    "OXR"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_oxr.OXRConfig",
+    "AutoModelForCausalLM": "modeling_oxr.OXR"
+  },
+  "conversation_formatter_class": "Qwen3ConversationFormatter",
+  "dtype": "bfloat16",
+  "eos_token_id": 151643,
+  "hidden_size": 1024,
+  "keys_to_ignore_at_inference": [
+    "past_key_values"
+  ],
+  "llm_config": {
+    "_name_or_path": "Qwen/Qwen3-0.6B-Base",
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "dtype": "bfloat16",
+    "eos_token_id": 151643,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 32768,
+    "max_window_layers": 28,
+    "model_type": "qwen3",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 8,
+    "pad_token_id": 151643,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000,
+    "sliding_window": null,
+    "tie_word_embeddings": true,
+    "use_cache": false,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+  },
+  "max_position_embeddings": 40960,
+  "model_type": "oxr",
+  "num_attention_heads": 32,
+  "pad_token_id": 151643,
+  "transformers_version": "5.6.2",
+  "use_cache": false,
+  "visual_vocab_size": 65536,
+  "vit_config": {
+    "_attn_implementation_autoset": true,
+    "_name_or_path": "google/siglip2-so400m-patch16-512",
+    "attention_dropout": 0.0,
+    "dtype": "bfloat16",
+    "fullatt_block_indexes": null,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "hidden_stride": 2,
+    "image_size": 512,
+    "intermediate_size": 4304,
+    "layer_norm_eps": 1e-06,
+    "model_type": "siglip2_navit",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 27,
+    "num_patches": -1,
+    "pad_token_id": 151643,
+    "patch_size": 16,
+    "preserve_original_pe": true,
+    "temporal_patch_size": 1,
+    "use_rope": true,
+    "window_size": 112
+  },
+  "vocab_size": 151936
+}

configuration_oxr.py

@@ -0,0 +1,100 @@
+from typing import Any, Optional, List, Union
+
+from transformers import AutoConfig, Qwen3Config
+from transformers.configuration_utils import PretrainedConfig
+
+__all__ = ["Siglip2NavitConfig", "OXRConfig"]
+
+
+class Siglip2NavitConfig(PretrainedConfig):
+    """This is the configuration class to store the configuration of an [`AIMv2Model`].
+
+    Instantiating a configuration with the defaults will yield a similar configuration
+    to that of the [apple/aimv2-large-patch14-224](https://huggingface.co/apple/aimv2-large-patch14-224).
+
+    Args:
+        hidden_size: Dimension of the hidden representations.
+        intermediate_size: Dimension of the SwiGLU representations.
+        num_hidden_layers: Number of hidden layers in the Transformer.
+        num_attention_heads: Number of attention heads for each attention layer
+            in the Transformer.
+        num_channels: Number of input channels.
+        image_size: Image size.
+        patch_size: Patch size.
+        rms_norm_eps: Epsilon value used for the RMS normalization layer.
+        attention_dropout: Dropout ratio for attention probabilities.
+        projection_dropout: Dropout ratio for the projection layer after the attention.
+        qkv_bias: Whether to add a bias to the queries, keys and values.
+        use_bias: Whether to add a bias in the feed-forward and projection layers.
+        kwargs: Keyword arguments for the [`PretrainedConfig`].
+    """
+
+    model_type: str = "siglip2_navit"
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        intermediate_size: int = 4096,
+        num_hidden_layers: int = 24,
+        num_attention_heads: int = 16,
+        num_channels: int = 3,
+        num_patches: int = -1,
+        image_size: int = 512,
+        patch_size: int = 16,
+        hidden_act: str="gelu_pytorch_tanh",
+        layer_norm_eps: float = 1e-6,
+        attention_dropout: float = 0.0,
+        hidden_stride: int = 2,
+        window_size: int = 112,
+        fullatt_block_indexes: Optional[list] = None,
+        temporal_patch_size: int = 1,
+        preserve_original_pe: bool = True,
+        use_rope: bool = True,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.hidden_act = hidden_act
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_stride = hidden_stride
+        self.window_size = window_size
+        self.fullatt_block_indexes = fullatt_block_indexes
+        self.temporal_patch_size = temporal_patch_size
+        self.preserve_original_pe = preserve_original_pe
+        self.use_rope = use_rope
+
+class OXRConfig(PretrainedConfig):
+    model_type = "oxr"
+    sub_configs = dict(llm_config=Qwen3Config, vit_config=Siglip2NavitConfig)
+
+    def __init__(self,
+        llm_config: Optional[Union[Qwen3Config, dict]] = None,
+        vit_config: Optional[Union[Siglip2NavitConfig, dict]] = None,
+        visual_vocab_size=65536,
+        hidden_size=None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if isinstance(llm_config, dict):
+            if "model_type" in llm_config:
+                config_kwargs = {k: v for k, v in llm_config.items() if k != "model_type"}
+                llm_config = AutoConfig.for_model(llm_config["model_type"], **config_kwargs)
+            else:
+                llm_config = Qwen3Config(**llm_config)
+        self.llm_config = llm_config
+        if isinstance(vit_config, dict):
+            vit_config = Siglip2NavitConfig(**vit_config)
+        self.vit_config = vit_config
+        self.visual_vocab_size = visual_vocab_size
+        self.hidden_size = hidden_size
+        if kwargs.get('attn_implementation'):
+            self.llm_config._attn_implementation = kwargs['attn_implementation']
+            self.vit_config._attn_implementation = kwargs['attn_implementation']

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": [
+    151643,
+    151645
+  ],
+  "transformers_version": "4.57.1"
+}

layout_analysis/README.md

@@ -0,0 +1,33 @@
+# Layout Analysis
+
+D-FINE Stage-2 B2 layout model with 12 classes. Detects page regions;
+text recognition and reading order are handled by OXR.
+
+## Files
+
+- `last.pth`: inference-only model weights (no optimizer or training state).
+- `metadata.json`: model identity, architecture, preprocessing, class IDs,
+  checkpoint size and SHA-256 (schema version 1).
+
+`model_version` is a public release version (`major.minor.patch`), not a
+training-run identifier.
+
+## Usage
+
+Use the LA inference service in the OXR code repository at
+`services/layout_analyze/`. Mount this directory into the service and set
+`MODEL_PATH` to the mounted `last.pth` path.
+
+The matching LA service reads `metadata.json` next to the checkpoint at startup
+and validates its supported architecture, preprocessing, class IDs and weight
+integrity. Both files must be mounted. `MODEL_METADATA_PATH` can override the
+metadata path. Unsupported metadata or a mismatched checkpoint fails startup.
+This is not a Transformers `AutoModel` configuration. Device, precision and
+CUDA Graph settings are selected by the service, not by the model metadata.
+
+Checkpoint verification:
+
+```bash
+shasum -a 256 layout_analysis/last.pth
+# dff78ccf18439955362ca44a69eb28fb0b40222a7d278860589682de0176dc17
+```

layout_analysis/last.pth

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

layout_analysis/metadata.json

@@ -0,0 +1,29 @@
+{
+  "schema_version": 1,
+  "model_type": "dfine",
+  "model_version": "1.0.0",
+  "architecture": "stage2_b2",
+  "checkpoint": "last.pth",
+  "checkpoint_size_bytes": 78791307,
+  "checkpoint_sha256": "dff78ccf18439955362ca44a69eb28fb0b40222a7d278860589682de0176dc17",
+  "num_classes": 12,
+  "image_size": 960,
+  "color_order": "rgb",
+  "resize_backend": "opencv_lanczos4",
+  "postprocess": "model_letterbox",
+  "provides_reading_order": false,
+  "id2label": {
+    "0": "Caption",
+    "1": "Footnote",
+    "2": "Formula",
+    "3": "Text",
+    "4": "Page-footer",
+    "5": "Page-header",
+    "6": "Picture",
+    "7": "Title",
+    "8": "Table",
+    "9": "Code Block",
+    "10": "Stamp",
+    "11": "Chemical Structure"
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:795bb2f1faac579f58829913f2c606202531a96c1203e19ef7a70b103fb7ea51
+size 3068857520

modeling_oxr.py

@@ -0,0 +1,1037 @@
+import math
+from typing import Dict, List, Optional, Tuple, Union
+
+import PIL.Image
+import numpy as np
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from transformers import (
+    AutoConfig,
+    AutoImageProcessor,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    __version__ as TRANSFORMERS_VERSION,
+)
+from transformers.activations import ACT2FN
+from transformers.generation.utils import GenerateOutput
+from transformers.modeling_outputs import BaseModelOutputWithNoAttention
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import is_flash_attn_2_available
+
+from .configuration_oxr import Siglip2NavitConfig, OXRConfig
+
+
+if is_flash_attn_2_available() :
+    from flash_attn import flash_attn_varlen_func
+    from flash_attn.layers.rotary import apply_rotary_emb
+
+
+IMAGE_PLACEHOLDER = "<image>"
+IMAGE_PLACEHOLDER_ID = -200
+VIDEO_PLACEHOLDER = "<video>"
+VIDEO_PLACEHOLDER_ID = -201
+
+VISUAL_ATOM_ID = -300
+INDICATOR_IDS = [-301, -302, -303, -304]
+
+
+# copied from qwen2.5-vl
+class VisionRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, seqlen: int) -> torch.Tensor:
+        inv_freq = 1.0 / (
+            self.theta ** (torch.arange(0, self.dim, 2, device=self.inv_freq.device, dtype=torch.float) / self.dim)
+        )
+        seq = torch.arange(seqlen, device=inv_freq.device, dtype=inv_freq.dtype)
+        freqs = torch.outer(seq, inv_freq)
+        return freqs
+
+
+class Siglip2VisionEmbeddings(nn.Module):
+    def __init__(self, config: Siglip2NavitConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.patch_size = config.patch_size
+        self.image_size = config.image_size
+        self.num_patches = config.num_patches
+        self.preserve_original_pe = config.preserve_original_pe
+        self.hidden_stride = config.hidden_stride
+
+
+        # siglip2 naflex
+        if self.num_patches > 0:
+            self.patch_embedding = nn.Linear(
+                    in_features=config.num_channels * self.patch_size * self.patch_size,
+                    out_features=self.embed_dim,
+                )
+            if self.preserve_original_pe:
+                self.position_embedding_size = int(self.num_patches**0.5)
+                self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
+
+        else:
+            self.patch_embedding = nn.Conv2d(
+                    in_channels=config.num_channels,
+                    out_channels=self.embed_dim,
+                    kernel_size=self.patch_size,
+                    stride=self.patch_size,
+                    padding="valid",
+                )
+            if self.preserve_original_pe:
+                self.num_patches = (self.image_size // self.patch_size) ** 2
+                self.position_embedding_size = self.image_size // self.patch_size
+                self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
+        
+    @staticmethod
+    def resize_positional_embeddings(
+        positional_embeddings: torch.Tensor,
+        spatial_shapes: torch.LongTensor,
+        max_length: int,
+    ) -> torch.Tensor:
+        """
+        Resize positional embeddings to image-specific size and pad to a fixed size.
+
+        Args:
+            positional_embeddings (`torch.Tensor`):
+                Position embeddings of shape (height, width, embed_dim)
+            spatial_shapes (`torch.LongTensor`):
+                Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
+            max_length (`int`):
+                Maximum length of the positional embeddings to pad resized positional embeddings to
+
+        Returns:
+            `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)
+        """
+        batch_size = spatial_shapes.shape[0]
+        embed_dim = positional_embeddings.shape[-1]
+        source_dtype = positional_embeddings.dtype
+
+        resulted_positional_embeddings = torch.empty(
+            (batch_size, max_length, embed_dim),
+            device=positional_embeddings.device,
+            dtype=source_dtype,
+        )
+
+        # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation
+        positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
+
+        # Upcast to float32 on CPU because antialias is not supported for bfloat16/float16 on CPU
+        if positional_embeddings.device.type == "cpu":
+            positional_embeddings = positional_embeddings.to(torch.float32)
+
+        for i in range(batch_size):
+            # (1, dim, height, width) -> (1, dim, target_height, target_width)
+            height, width = spatial_shapes[i]
+            resized_embeddings = F.interpolate(
+                positional_embeddings,
+                size=(height, width),
+                mode="bilinear",
+                align_corners=False,
+                antialias=True,
+            )
+
+            # (1, dim, target_height, target_width) -> (target_height * target_width, dim)
+            resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1)
+
+            # Cast to original dtype
+            resized_embeddings = resized_embeddings.to(source_dtype)
+
+            resulted_positional_embeddings[i, : height * width] = resized_embeddings
+            resulted_positional_embeddings[i, height * width :] = resized_embeddings[0]
+
+        return resulted_positional_embeddings
+
+    def forward(self, pixel_values: torch.FloatTensor, 
+                grid_thws: Optional[torch.LongTensor] = None) -> torch.Tensor:
+        """
+        Args:
+            pixel_values (`torch.FloatTensor`):
+                Pixel values of shape (num_patches, num_channels * temporal_patch_size * patch_size * patch_size)
+            grid_thws: (`torch.LongTensor`):
+                grid shape (num_patches, 3)
+        """
+
+        # Apply patch embeddings to already patchified pixel values
+        target_dtype = self.patch_embedding.weight.dtype
+        if isinstance(self.patch_embedding, nn.Linear):
+            patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
+        elif isinstance(self.patch_embedding, nn.Conv2d):
+            pixel_values = pixel_values.view(-1, self.config.num_channels * self.config.temporal_patch_size, self.patch_size,
+                   self.patch_size)
+            patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
+            patch_embeds = patch_embeds.reshape(-1, self.embed_dim)
+
+
+        if self.preserve_original_pe:
+            assert grid_thws is not None
+            pos_embed_new = torch.zeros_like(patch_embeds)
+            ori_h = ori_w = self.position_embedding_size
+            positional_embeddings = self.position_embedding.weight.reshape(
+                                self.position_embedding_size, self.position_embedding_size, -1
+                            ).unsqueeze(0).permute(0,3,1,2)
+            # pos_embed = self.pos_embed.reshape(1, ori_h, ori_w, -1).permute(0, 3, 1, 2)
+            cnt = 0
+            for t, h, w in grid_thws:
+                thw = t * h * w
+                pe = F.interpolate(positional_embeddings, size=(h, w), mode='bicubic', align_corners=False)
+                pe = pe.permute(0, 2, 3, 1).reshape(1, h * w, -1)
+                pe = pe[0].repeat(t, 1)
+                pe = pe.reshape(t, h // self.hidden_stride, self.hidden_stride, w // self.hidden_stride,
+                                self.hidden_stride, -1)
+                pe = pe.permute(0, 1, 3, 2, 4, 5).reshape(thw, -1)
+                pos_embed_new[cnt:cnt + thw] = pe
+                cnt += thw
+            patch_embeds = patch_embeds + pos_embed_new
+
+        return patch_embeds
+
+
+# copied from qwen2.5-vl
+def apply_rotary_pos_emb_flashatt(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.chunk(2, dim=-1)[0].contiguous()
+    sin = sin.chunk(2, dim=-1)[0].contiguous()
+    q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q)
+    k_embed = apply_rotary_emb(k.float(), cos.float(), sin.float()).type_as(k)
+    return q_embed, k_embed
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    orig_q_dtype = q.dtype
+    orig_k_dtype = k.dtype
+    q, k = q.float(), k.float()
+    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    q_embed = q_embed.to(orig_q_dtype)
+    k_embed = k_embed.to(orig_k_dtype)
+    return q_embed, k_embed
+
+
+class Siglip2Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+        self.is_causal = False
+
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+        self.use_rope = config.use_rope
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Input shape: Batch x Time x Channel"""
+
+        seq_length, embed_dim = hidden_states.shape
+
+        queries = self.q_proj(hidden_states)
+        keys = self.k_proj(hidden_states)
+        values = self.v_proj(hidden_states)
+
+        queries = queries.view(seq_length, self.num_heads, self.head_dim)
+        keys = keys.view(seq_length, self.num_heads, self.head_dim)
+        values = values.view(seq_length, self.num_heads, self.head_dim)
+
+        if self.use_rope:
+            cos, sin = position_embeddings
+            if is_flash_attn_2_available() :
+                queries, keys = apply_rotary_pos_emb_flashatt(queries.unsqueeze(0), keys.unsqueeze(0), cos, sin)
+            else:
+                queries, keys = apply_rotary_pos_emb_vision(queries.unsqueeze(0), keys.unsqueeze(0), cos, sin)
+            queries = queries.squeeze(0)
+            keys = keys.squeeze(0)
+
+        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+        if is_flash_attn_2_available() :
+            attn_output = flash_attn_varlen_func(queries, keys, values, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen).reshape(
+                                                seq_length, -1
+                                            )
+        else:
+            batch_size = cu_seqlens.shape[0] - 1
+            outputs = []
+            cu = cu_seqlens.tolist()
+            for i in range(batch_size):
+                start_idx = cu[i]
+                end_idx = cu[i + 1]
+                # Each sequence is processed independently.
+                q_i = queries[start_idx:end_idx].unsqueeze(0)
+                k_i = keys[start_idx:end_idx].unsqueeze(0)
+                v_i = values[start_idx:end_idx].unsqueeze(0)
+                # (1, seq_len, num_heads, head_dim) ->
+                # (1, num_heads, seq_len, head_dim)
+                q_i, k_i, v_i = [x.transpose(1, 2) for x in (q_i, k_i, v_i)]
+                output_i = F.scaled_dot_product_attention(q_i,
+                                                        k_i,
+                                                        v_i,
+                                                        dropout_p=0.0)
+                # (1, num_heads, seq_len, head_dim) -> (seq_len, embed_dim)
+                output_i = output_i.transpose(1, 2).reshape(-1, self.embed_dim)
+                outputs.append(output_i)
+            attn_output = torch.cat(outputs, dim=0)
+
+        attn_output = self.out_proj(attn_output)
+        return attn_output
+
+class Siglip2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class Siglip2EncoderLayer(nn.Module):
+    def __init__(self, config: Siglip2NavitConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.self_attn = Siglip2Attention(config)
+        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.mlp = Siglip2MLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor, 
+        position_embeddings: torch.Tensor
+    ) -> tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                Input to the layer of shape `(batch, seq_len, embed_dim)`.
+            attention_mask (`torch.FloatTensor`):
+                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        residual = hidden_states
+
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            cu_seqlens=cu_seqlens, 
+            position_embeddings=position_embeddings
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+class Siglip2Encoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`Siglip2EncoderLayer`].
+
+    Args:
+        config: Siglip2NavitConfig
+    """
+
+    def __init__(self, config: Siglip2NavitConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+        self.rotary_pos_emb = VisionRotaryEmbedding(config.hidden_size // config.num_attention_heads // 2)
+        self.patch_size = config.patch_size
+        self.hidden_stride = config.hidden_stride
+        self.window_size = config.window_size
+        self.spatial_merge_unit = config.hidden_stride * config.hidden_stride
+        self.fullatt_block_indexes = None if config.fullatt_block_indexes is None else [int(i) for i in config.fullatt_block_indexes.split('|')]
+        
+
+    # copied from qwen2.5_vl
+    def rot_pos_emb(self, grid_thw):
+        pos_ids = []
+        for t, h, w in grid_thw:
+            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+            hpos_ids = hpos_ids.reshape(
+                h // self.hidden_stride,
+                self.hidden_stride,
+                w // self.hidden_stride,
+                self.hidden_stride,
+            )
+            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+            hpos_ids = hpos_ids.flatten()
+
+            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+            wpos_ids = wpos_ids.reshape(
+                h // self.hidden_stride,
+                self.hidden_stride,
+                w // self.hidden_stride,
+                self.hidden_stride,
+            )
+            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+            wpos_ids = wpos_ids.flatten()
+            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+        pos_ids = torch.cat(pos_ids, dim=0)
+        max_grid_size = grid_thw[:, 1:].max()
+        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+        rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+        return rotary_pos_emb
+
+    def get_window_index(self, grid_thw):
+        window_index: list = []
+        cu_window_seqlens: list = [0]
+        window_index_id = 0
+        vit_merger_window_size = self.window_size // self.hidden_stride // self.patch_size  # patch (after merge) number in each window
+
+        for grid_t, grid_h, grid_w in grid_thw:
+            llm_grid_h, llm_grid_w = (
+                grid_h // self.hidden_stride,  # number of patch after merge
+                grid_w // self.hidden_stride,
+            )
+            index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)
+            pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
+            pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
+            num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
+            num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
+            index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100)
+            index_padded = index_padded.reshape(
+                grid_t,
+                num_windows_h,
+                vit_merger_window_size,
+                num_windows_w,
+                vit_merger_window_size,
+            )
+            index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(
+                grid_t,
+                num_windows_h * num_windows_w,
+                vit_merger_window_size,
+                vit_merger_window_size,
+            )
+            seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)
+            index_padded = index_padded.reshape(-1)
+            index_new = index_padded[index_padded != -100]
+            window_index.append(index_new + window_index_id)
+            cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1]
+            cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
+            window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
+        window_index = torch.cat(window_index, dim=0)
+
+        return window_index, cu_window_seqlens
+
+    # Ignore copy
+    def forward(
+        self,
+        inputs_embeds,
+        grid_thws: torch.Tensor,
+        output_hidden_states: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, ...]]]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+
+        rotary_pos_emb = self.rot_pos_emb(grid_thws)
+        window_index, cu_window_seqlens = self.get_window_index(grid_thws)
+        cu_window_seqlens = torch.tensor(
+            cu_window_seqlens,
+            device=inputs_embeds.device,
+            dtype=grid_thws.dtype if torch.jit.is_tracing() else torch.int32,
+        )
+        cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
+
+        seq_len, _ = inputs_embeds.size()
+        inputs_embeds = inputs_embeds.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+        inputs_embeds = inputs_embeds[window_index, :, :]
+        inputs_embeds = inputs_embeds.reshape(seq_len, -1)
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+        rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+        position_embeddings = (emb.cos(), emb.sin())
+
+        cu_seqlens = torch.repeat_interleave(grid_thws[:, 1] * grid_thws[:, 2], grid_thws[:, 0]).cumsum(
+            dim=0,
+            # Select dtype based on the following factors:
+            #  - FA2 requires that cu_seqlens_q must have dtype int32
+            #  - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+            # See https://github.com/huggingface/transformers/pull/34852 for more information
+            dtype=grid_thws.dtype if torch.jit.is_tracing() else torch.int32,
+        )
+        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+        reverse_indices = torch.argsort(window_index)
+        encoder_states = () if output_hidden_states else None
+
+        hidden_states = inputs_embeds
+        for index, block in enumerate(self.layers):
+            if self.fullatt_block_indexes is None or index in self.fullatt_block_indexes:
+                cu_seqlens_tmp = cu_seqlens
+            else:
+                cu_seqlens_tmp = cu_window_seqlens
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(block.__call__, hidden_states, cu_seqlens_tmp, position_embeddings)
+            else:
+                hidden_states = block(hidden_states, cu_seqlens_tmp, position_embeddings)
+            if output_hidden_states:
+                hidden_states_ = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+                encoder_states += (hidden_states_[reverse_indices, :].reshape(seq_len, -1),)
+        # tokens = self.post_trunk_norm(tokens)
+        hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+        hidden_states = hidden_states[reverse_indices, :].reshape(seq_len, -1)
+
+        return hidden_states, encoder_states
+
+class Siglip2VisionTransformer(nn.Module):
+    def __init__(self, config: Siglip2NavitConfig):
+        super().__init__()
+        self.config = config
+        embed_dim = config.hidden_size
+
+        self.embeddings = Siglip2VisionEmbeddings(config)
+        self.encoder = Siglip2Encoder(config)
+        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        grid_thws: torch.LongTensor,
+        output_hidden_states: Optional[bool] = True,
+        return_dict: Optional[bool] = True,
+    ) -> Union[
+        Tuple[torch.Tensor],
+        Tuple[torch.Tensor, Tuple[torch.Tensor, ...]],
+        BaseModelOutputWithNoAttention,
+    ]:
+        r"""
+        spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
+            Tensor containing the spatial dimensions (height, width) of the input images.
+        """
+        # output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        # output_hidden_states = (
+        #     output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        # )
+
+        hidden_states = self.embeddings(pixel_values, grid_thws)
+
+        last_hidden_state, hidden_states = self.encoder(hidden_states, grid_thws, output_hidden_states)
+        last_hidden_state = self.post_layernorm(last_hidden_state)
+
+        if not return_dict:
+            output = (last_hidden_state,)
+            output += (hidden_states,) if output_hidden_states else ()
+            return output
+        
+        return BaseModelOutputWithNoAttention(
+                last_hidden_state=last_hidden_state, 
+                hidden_states=hidden_states
+            )
+       
+class Siglip2PreTrainedModel(PreTrainedModel):
+    config_class = Siglip2NavitConfig
+    base_model_prefix = "siglip2_navit"
+    supports_gradient_checkpointing = True
+
+    _no_split_modules = [
+        "Siglip2VisionEmbeddings",
+        "Siglip2EncoderLayer",
+    ]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = False
+    _supports_flex_attn = False
+    _supports_attention_backend = True
+
+
+class Siglip2NavitModel(Siglip2PreTrainedModel):
+    config_class = Siglip2NavitConfig
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: Siglip2NavitConfig):
+        super().__init__(config)
+
+        self.vision_model = Siglip2VisionTransformer(config)
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.vision_model.embeddings.patch_embedding
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        grid_thws: torch.LongTensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[
+        Tuple[torch.Tensor],
+        Tuple[torch.Tensor, Tuple[torch.Tensor, ...]],
+        BaseModelOutputWithNoAttention,
+    ]:
+        
+        if output_hidden_states is None:
+            output_hidden_states = self.config.output_hidden_states
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+
+        return self.vision_model(
+            pixel_values=pixel_values,
+            grid_thws=grid_thws,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+class OXREmbedding(torch.nn.Embedding):
+    """
+    A visual embedding layer that can handle both discrete token IDs (long) and continuous
+    soft-token probabilities (float).
+    """
+
+    def forward(self, visual_tokens: Tensor) -> Tensor:
+        if visual_tokens.dtype in [torch.int8, torch.int16, torch.int32, torch.int64, torch.long]:
+            return super().forward(visual_tokens)
+        # Handle soft tokens (probabilities) by matrix multiplication with the embedding weight
+        return torch.matmul(visual_tokens, self.weight)
+
+
+class OXRVisualTokenizer(torch.nn.Module):
+    """
+    Tokenizes images or videos into a sequence of continuous visual tokens.
+    """
+
+    def __init__(self, vit, visual_vocab_size, image_processor_name_or_path, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.vit = vit
+        image_processor_kwargs = {"do_center_crop": False}
+        if int(TRANSFORMERS_VERSION.split(".", 1)[0]) >= 5:
+            image_processor_kwargs["backend"] = "pil"
+        else:
+            image_processor_kwargs["use_fast"] = False
+        self.image_processor = AutoImageProcessor.from_pretrained(
+            image_processor_name_or_path,
+            **image_processor_kwargs,
+        )
+        head_dim = visual_vocab_size - len(INDICATOR_IDS)
+        self.head = torch.nn.Sequential(
+            torch.nn.Linear(self.vit.config.hidden_size * self.vit.config.hidden_stride ** 2, head_dim, bias=False),
+            torch.nn.LayerNorm(head_dim)
+        )
+
+    def _encode(self, pixel_values, grid_thws):
+        output = self.vit(pixel_values, grid_thws, output_hidden_states=True, return_dict=True)
+        features = output.hidden_states[-1]
+        seq_len, _ = features.shape
+        features = features.reshape(seq_len // (self.vit.config.hidden_stride ** 2), -1)
+        return features
+
+    # Adapted from qwen2_vl
+    @staticmethod
+    def smart_resize(
+        height: int, width: int, factor: int = 28, min_pixels: int = 448 * 448, max_pixels: int = 1344 * 1792
+    ):
+        """Rescales the image so that the following conditions are met:
+        1. Both dimensions are divisible by 'factor'.
+        2. The total number of pixels is within ['min_pixels', 'max_pixels'].
+        3. The aspect ratio is maintained as closely as possible.
+        """
+        if height < factor or width < factor:
+            if height < width:
+                width = round(factor / height * width)
+                height = factor
+            else:
+                height = round(factor / width * height)
+                width = factor
+
+        elif max(height, width) / min(height, width) > 200:
+            if height > width:
+                height = 200 * width
+            else:
+                width = 200 * height
+
+        h_bar = round(height / factor) * factor
+        w_bar = round(width / factor) * factor
+        if h_bar * w_bar > max_pixels:
+            beta = math.sqrt((height * width) / max_pixels)
+            h_bar = math.floor(height / beta / factor) * factor
+            w_bar = math.floor(width / beta / factor) * factor
+        elif h_bar * w_bar < min_pixels:
+            beta = math.sqrt(min_pixels / (height * width))
+            h_bar = math.ceil(height * beta / factor) * factor
+            w_bar = math.ceil(width * beta / factor) * factor
+        return h_bar, w_bar
+
+    def preprocess(
+        self,
+        image: Optional[PIL.Image.Image] = None,
+        video: Optional[List[PIL.Image.Image]] = None,
+        min_pixels: Optional[int] = None,
+        max_pixels: Optional[int] = None
+    ):
+        patch_size = self.vit.config.patch_size
+        temporal_patch_size = self.vit.config.temporal_patch_size
+        hidden_stride = self.vit.config.hidden_stride
+        assert (image is None) ^ (video is None), "Invalid input: expect either image or video"
+        if image is not None:
+            images = [image]
+        else:
+            images = video
+        images = [image.convert("RGB") if image.mode != 'RGB' else image for image in images]
+        width, height = images[0].size
+        processed_images = []
+        for image in images:
+            resized_height, resized_width = self.smart_resize(
+                height,
+                width,
+                factor=patch_size * hidden_stride,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+            new_size = dict(height=resized_height, width=resized_width)
+            new_image = self.image_processor.preprocess(image, size=new_size, return_tensors="np")['pixel_values'][0]
+            processed_images.append(new_image)
+
+        patches = np.array(processed_images)
+        if patches.shape[0] % temporal_patch_size != 0:
+            repeats = np.repeat(patches[-1][np.newaxis], temporal_patch_size - 1, axis=0)
+            patches = np.concatenate([patches, repeats], axis=0)
+        channel = patches.shape[1]
+        grid_t = patches.shape[0] // temporal_patch_size
+        grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+        grid_thw = torch.tensor([[grid_t, grid_h, grid_w]])
+
+        patches = patches.reshape(
+            grid_t, temporal_patch_size, channel,
+            grid_h // hidden_stride, hidden_stride, patch_size,
+            grid_w // hidden_stride, hidden_stride, patch_size,
+        )
+        patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
+        flatten_patches = patches.reshape(
+            grid_t * grid_h * grid_w, channel * temporal_patch_size * patch_size * patch_size
+        )
+        flatten_patches = torch.tensor(flatten_patches)
+
+        return flatten_patches, grid_thw
+
+    def forward(
+        self, pixel_values, grid_thws
+    ) -> torch.Tensor:  # [BatchSize, ImageShape] -> [BatchSize, #Token, VocabSize]
+        features = self._encode(pixel_values, grid_thws)
+        logits = self.head(features)
+        tokens = torch.softmax(logits, dim=-1, dtype=torch.float32).to(logits.dtype)
+
+        token_len, _ = tokens.shape
+        padding_tensor = torch.zeros(size=(token_len, len(INDICATOR_IDS)),
+                                     dtype=tokens.dtype,
+                                     device=tokens.device,
+                                     layout=tokens.layout,
+                                     requires_grad=False)
+        tokens = torch.cat((tokens, padding_tensor), dim=1)
+        return tokens
+
+
+class OXRPreTrainedModel(PreTrainedModel):
+    config_class = OXRConfig
+    base_model_prefix = "oxr"
+
+
+class OXR(OXRPreTrainedModel):
+    _supports_flash_attn_2 = True
+
+    def __init__(self, config: OXRConfig, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+
+        self.llm = AutoModelForCausalLM.from_config(self.config.llm_config)
+        assert self.config.hidden_size == self.llm.config.hidden_size, "hidden size mismatch"
+        self.text_tokenizer = AutoTokenizer.from_pretrained(
+            self.config.name_or_path,
+            trust_remote_code=True,
+        )
+        self.visual_tokenizer = OXRVisualTokenizer(vit=AutoModel.from_config(self.config.vit_config),
+                                                visual_vocab_size=self.config.visual_vocab_size,
+                                                image_processor_name_or_path=self.config.name_or_path)
+
+        self.vte = OXREmbedding(self.config.visual_vocab_size, self.config.hidden_size,
+                                   device=self.visual_tokenizer.vit.device, dtype=self.visual_tokenizer.vit.dtype)
+        indicator_token_indices = torch.arange(
+            self.config.visual_vocab_size - len(INDICATOR_IDS),
+            self.config.visual_vocab_size,
+            dtype=torch.long
+        )
+        self.register_buffer("indicator_token_indices", indicator_token_indices, persistent=False)
+
+        def _merge_modules(modules_list: tuple):
+            merged_modules = []
+            for modules in modules_list:
+                merged_modules.extend(modules if modules else [])
+            return merged_modules
+
+        # Standard model configurations for parallelism and device placement
+        self._no_split_modules = _merge_modules(
+            (self.llm._no_split_modules, self.visual_tokenizer.vit._no_split_modules))
+        self._skip_keys_device_placement = self.llm._skip_keys_device_placement
+        self._keep_in_fp32_modules = _merge_modules(
+            (self.llm._keep_in_fp32_modules, self.visual_tokenizer.vit._keep_in_fp32_modules))
+        self.is_parallelizable = all((
+            getattr(self.llm, "is_parallelizable", False),
+            getattr(self.visual_tokenizer.vit, "is_parallelizable", False),
+        ))
+        self.supports_gradient_checkpointing = True
+        self.post_init()
+
+    def tie_weights(self, missing_keys: Optional[set[str]] = None, recompute_mapping: bool = True):
+        self.llm.tie_weights()
+        try:
+            return PreTrainedModel.tie_weights(
+                self,
+                missing_keys=missing_keys,
+                recompute_mapping=recompute_mapping,
+            )
+        except TypeError as exc:
+            if "unexpected keyword argument" not in str(exc):
+                raise
+            return PreTrainedModel.tie_weights(self)
+
+    def get_wte(self):
+        return self.llm.get_input_embeddings()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        pixel_values: Optional[torch.Tensor],
+        grid_thws: Optional[torch.Tensor],
+        labels: Optional[torch.Tensor] = None,
+        **kwargs
+    ):
+        inputs_embeds = self.merge_multimodal(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            grid_thws=grid_thws,
+        )
+        return self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels, **kwargs)
+
+    def merge_multimodal(
+        self,
+        input_ids: torch.Tensor,
+        pixel_values: Optional[torch.Tensor],
+        grid_thws: Optional[torch.Tensor],
+    ):
+        placeholder_token_mask = torch.lt(input_ids, 0)
+        multimodal_embeds = self.get_wte()(torch.masked_fill(input_ids, placeholder_token_mask, 0))
+
+        if pixel_values is not None:
+            visual_indicator_embeds = self.vte(self.indicator_token_indices).to(
+                dtype=multimodal_embeds.dtype, device=multimodal_embeds.device
+            )
+            visual_tokens = self.visual_tokenizer(pixel_values, grid_thws)
+            visual_embeds = self.vte(visual_tokens).to(dtype=multimodal_embeds.dtype, device=multimodal_embeds.device)
+
+            for i, indicator_id in enumerate(INDICATOR_IDS):
+                multimodal_embeds[input_ids == indicator_id] = visual_indicator_embeds[i]
+            multimodal_embeds[input_ids == VISUAL_ATOM_ID] = visual_embeds
+
+        return multimodal_embeds
+
+    def _merge_inputs(
+        self, raw_input_ids, placeholder_id, grid_thws, indicator_begin_id, indicator_end_id
+    ):
+        input_ids = []
+        prev_index = 0
+        placeholder_indexes = [i for i, v in enumerate(raw_input_ids) if v == placeholder_id]
+        for placeholder_index, grid_thw in zip(placeholder_indexes, grid_thws):
+            input_ids.extend(raw_input_ids[prev_index:placeholder_index])
+            num_image_atoms = grid_thw.prod().item()
+            num_image_atoms //= self.visual_tokenizer.vit.config.hidden_stride ** 2
+            num_image_atoms //= self.visual_tokenizer.vit.config.temporal_patch_size
+            input_ids.extend([indicator_begin_id] + [VISUAL_ATOM_ID] * num_image_atoms + [indicator_end_id])
+            prev_index = placeholder_index + 1
+        input_ids.extend(raw_input_ids[prev_index:])
+        return input_ids
+
+    def _tokenize_with_visual_placeholder(self, text):
+        placeholder = VIDEO_PLACEHOLDER if VIDEO_PLACEHOLDER in text else IMAGE_PLACEHOLDER
+        placeholder_id = VIDEO_PLACEHOLDER_ID if VIDEO_PLACEHOLDER in text else IMAGE_PLACEHOLDER_ID
+        chunks = [self.text_tokenizer(chunk, add_special_tokens=False).input_ids for chunk in text.split(placeholder)]
+        input_ids = chunks[0]
+        for chunk in chunks[1:]:
+            input_ids.append(placeholder_id)
+            input_ids.extend(chunk)
+        return input_ids
+
+    def preprocess_inputs(
+        self,
+        messages: List[Union[str, Dict]],
+        min_pixels=448 * 448,
+        max_pixels=1344 * 1792,
+        add_generation_prompt=True,
+        enable_thinking=False
+    ):
+        text = self.text_tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=add_generation_prompt,
+            enable_thinking=enable_thinking
+        )
+        input_ids = self._tokenize_with_visual_placeholder(text)
+        images = []
+        videos = []
+        for message in messages:
+            content = message["content"]
+            if isinstance(content, list):
+                images.extend([item["image"] for item in content if item.get("image") is not None])
+                videos.extend([item["video"] for item in content if item.get("video") is not None])
+        if images and videos:
+            raise ValueError(
+                "Multiple visual input data types detected (both image and video provided). "
+                "This model supports only one type of visual input data at a time. "
+                "Please provide either image or video, but not both."
+            )
+
+        pixel_values, grid_thws = None, None
+        if images:
+            pixel_values, grid_thws = zip(
+                *(self.visual_tokenizer.preprocess(image=image, min_pixels=min_pixels, max_pixels=max_pixels)
+                  for image in images)
+            )
+            input_ids = self._merge_inputs(
+                input_ids, IMAGE_PLACEHOLDER_ID, grid_thws, INDICATOR_IDS[0], INDICATOR_IDS[1]
+            )
+            pixel_values = torch.cat(pixel_values, dim=0)
+            grid_thws = torch.cat(grid_thws, dim=0)
+        elif videos:
+            assert len(videos) == 1, "only support single video"
+            pixel_values, grid_thws = self.visual_tokenizer.preprocess(
+                video=videos[0], min_pixels=min_pixels, max_pixels=max_pixels
+            )
+            input_ids = self._merge_inputs(
+                input_ids, VIDEO_PLACEHOLDER_ID, grid_thws, INDICATOR_IDS[2], INDICATOR_IDS[3]
+            )
+
+        input_ids = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0)
+
+        return input_ids, pixel_values, grid_thws
+
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        attention_mask = torch.ne(inputs, self.text_tokenizer.pad_token_id).to(device=inputs.device)
+        inputs_embeds = self.merge_multimodal(
+            input_ids=inputs,
+            pixel_values=kwargs.pop('pixel_values', None),
+            grid_thws=kwargs.pop('grid_thws', None)
+        )
+        enable_thinking = kwargs.pop('enable_thinking', False)
+        enable_thinking_budget = kwargs.pop('enable_thinking_budget', False)
+        thinking_budget = kwargs.pop('thinking_budget', 1024)
+        
+        if enable_thinking and enable_thinking_budget:
+            actual_max_new_tokens = kwargs['max_new_tokens']
+            kwargs['max_new_tokens'] = thinking_budget
+            generated_ids = self.llm.generate(inputs=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)
+            output_ids = generated_ids
+            output_ids_list = generated_ids[0]
+            
+            # check if the generation has already finished (151645 is <|im_end|>)
+            if 151645 not in output_ids_list:
+                # check if the thinking process has finished (151668 is </think>)
+                # and prepare the second model input
+                if 151668 not in output_ids_list:
+                    early_stopping_text = "\n\nConsidering the limited time by the user, I have to give the solution based on the thinking directly now.\n</think>\n\n"
+                    early_stopping_ids = self.text_tokenizer(early_stopping_text, return_tensors="pt", return_attention_mask=False).input_ids.to(inputs.device)
+                    input_ids_appendent = torch.cat([output_ids, early_stopping_ids], dim=-1)
+                    kwargs['streamer'].put(early_stopping_ids) if 'streamer' in kwargs else None
+                else:
+                    input_ids_appendent = output_ids
+                
+
+                # second generation
+                new_inputs = torch.cat([inputs, input_ids_appendent], dim=-1)
+                attention_mask = torch.ne(new_inputs, self.text_tokenizer.pad_token_id).to(device=inputs.device)
+                inputs_embeds_appendent = self.merge_multimodal(
+                    input_ids=input_ids_appendent,
+                    pixel_values=None,
+                    grid_thws=None
+                )
+                new_inputs_embeds = torch.cat([inputs_embeds, inputs_embeds_appendent], dim=-2)
+                
+                kwargs['max_new_tokens'] = inputs_embeds.size(-2) + actual_max_new_tokens - new_inputs_embeds.size(-2)
+                generated_ids2 = self.llm.generate(inputs=None, inputs_embeds=new_inputs_embeds, attention_mask=attention_mask, **kwargs)
+                kwargs['streamer'].manual_end() if 'streamer' in kwargs else None
+                return torch.cat([input_ids_appendent, generated_ids2], dim=-1)
+            
+            else:
+                kwargs['streamer'].manual_end() if 'streamer' in kwargs else None
+                return generated_ids
+            
+        else:
+            generated_ids = self.llm.generate(inputs=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)
+            kwargs['streamer'].manual_end() if 'streamer' in kwargs else None
+            return generated_ids
+
+
+AutoConfig.register('siglip2_navit', Siglip2NavitConfig)
+AutoModel.register(Siglip2NavitConfig, Siglip2NavitModel)
+try:
+    AutoModelForCausalLM.register(OXRConfig, OXR)
+except ValueError:
+    pass

plugins/vllm/README.md

@@ -0,0 +1,30 @@
+# OXR vLLM plugin
+
+Registers the OXR architecture through `vllm.general_plugins`. Requires
+Python 3.12+, vLLM 0.25.1 and Transformers 5.13.1. It does not modify the
+vLLM installation. The supported Docker runtime uses PyTorch 2.11.0+cu130.
+
+Install from the model repository root in a Linux CUDA environment:
+
+```bash
+python -m pip install ./plugins/vllm
+```
+
+Start with the model files and weights in the repository root:
+
+```bash
+VLLM_PLUGINS=oxr vllm serve . \
+  --trust-remote-code \
+  --served-model-name oxr \
+  --chat-template ./chat_template.jinja \
+  --dtype bfloat16 \
+  --attention-config '{"backend":"TRITON_ATTN"}' \
+  --max-model-len 16384
+```
+
+The LA model in `layout_analysis/` is served separately by the OXR layout
+analysis service, not by this plugin.
+
+To roll back to vLLM 0.19.1, use the previous model plugin together with the
+matching previous OXR launcher and Docker image. Do not install this plugin
+into the old runtime.

plugins/vllm/oxr_vllm/__init__.py

@@ -0,0 +1,11 @@
+"""Out-of-tree OXR model registration for vLLM."""
+
+
+def register() -> None:
+    """Register lazily so plugin discovery does not initialize CUDA."""
+    from vllm.model_executor.models.registry import ModelRegistry
+
+    ModelRegistry.register_model("OXR", "oxr_vllm.oxr_model:OXR")
+
+
+__all__ = ["register"]

plugins/vllm/oxr_vllm/oxr_model.py

@@ -0,0 +1,660 @@
+"""OXR model implementation for vLLM 0.25.1."""
+
+from collections.abc import Iterable, Mapping
+from functools import partial
+from typing import Annotated, Literal
+
+import torch
+import torch.nn as nn
+from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig
+
+from vllm.config import VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
+from vllm.model_executor.models.siglip2navit import Siglip2NavitModel
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    flatten_bn,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.parse import ImageSize, MultiModalDataItems
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+)
+from vllm.renderers import TokenizeParams
+from vllm.sequence import IntermediateTensors
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
+
+from .oxr_processor import OXRProcessor
+
+IMAGE_TOKEN = "<image>"
+VIDEO_TOKEN = "<video>"
+IMAGE_PLACEHOLDER_ID = -200
+VIDEO_PLACEHOLDER_ID = -201
+VISUAL_ATOM_ID = -300
+OXR_INDICATOR_IDS = [-301, -302, -303, -304]
+
+IMAGE_PAD_TOKEN_MAP = {
+    "gemma2": "<unused0>",
+    "llama": "<|reserved_special_token_0|>",
+    "qwen2": "<|image_pad|>",
+    "qwen3": "<|image_pad|>",
+}
+IMAGE_PAD_TOKEN_ID_MAP = {
+    "gemma2": 7,
+    "llama": 128002,
+    "qwen2": 151655,
+    "qwen3": 151655,
+}
+
+
+def get_llm_config(config: PretrainedConfig) -> PretrainedConfig:
+    llm_config = getattr(config, "llm_config", None)
+    if llm_config is not None:
+        return llm_config
+
+    text_config = getattr(config, "text_config", None)
+    if text_config is not None:
+        return text_config
+
+    get_text_config = getattr(config, "get_text_config", None)
+    if get_text_config is not None:
+        candidate = get_text_config()
+        if candidate is not config:
+            return candidate
+
+    raise ValueError("OXR config does not expose an LLM sub-config")
+
+
+def _oxr_field_config() -> Mapping[str, MultiModalFieldConfig]:
+    return dict(
+        pixel_values=MultiModalFieldConfig.batched("image"),
+        grids=MultiModalFieldConfig.batched("image"),
+        indicator_tokens=MultiModalFieldConfig.batched("image"),
+        video_pixel_values=MultiModalFieldConfig.batched("video"),
+        video_indicator_tokens=MultiModalFieldConfig.batched("video"),
+        video_grids=MultiModalFieldConfig.batched("video"),
+    )
+
+
+class OXRImagePatchInputs(TensorSchema):
+    type: Literal["image_patches"]
+    flat_data: Annotated[torch.Tensor, TensorShape("bnp", "patch_size")]
+    indicator_tokens: Annotated[torch.Tensor, TensorShape("indicators")]
+    patches_per_item: Annotated[list[int], TensorShape("bn")]
+    grids: Annotated[torch.Tensor, TensorShape("bn", 3)]
+
+
+class OXRVideoPatchInputs(TensorSchema):
+    type: Literal["video_patches"]
+    flat_data: Annotated[torch.Tensor, TensorShape("bnp", "patch_size")]
+    indicator_tokens: Annotated[torch.Tensor, TensorShape("indicators")]
+    patches_per_item: Annotated[list[int], TensorShape("bn")]
+    grids: Annotated[torch.Tensor, TensorShape("bn", 3)]
+
+
+class OXREmbedding(torch.nn.Embedding):
+    def forward(self, visual_tokens: torch.Tensor) -> torch.Tensor:
+        if visual_tokens.dtype in (
+            torch.int8,
+            torch.int16,
+            torch.int32,
+            torch.int64,
+            torch.long,
+        ):
+            return super().forward(visual_tokens)
+        return torch.matmul(visual_tokens, self.weight)
+
+    @property
+    def device(self) -> torch.device:
+        return self.weight.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.weight.dtype
+
+
+class OXRVisualTokenizer(torch.nn.Module):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        visual_vocab_size: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.config = config
+        self.vit = self._init_backbone(
+            config=config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.vit",
+        )
+        head_dim = visual_vocab_size - len(OXR_INDICATOR_IDS)
+        self.head = torch.nn.Sequential(
+            ReplicatedLinear(
+                self.config.hidden_size * self.config.hidden_stride**2,
+                head_dim,
+                bias=False,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.head.0",
+            ),
+            torch.nn.LayerNorm(head_dim),
+        )
+
+    def _init_backbone(
+        self,
+        config: PretrainedConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        if config.model_type == "siglip2_navit":
+            return Siglip2NavitModel(
+                config=config,
+                quant_config=quant_config,
+                prefix=prefix,
+            )
+        raise ValueError(f"Unsupported visual tokenizer model_type: {config.model_type}")
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.head.parameters()).dtype
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.head.parameters()).device
+
+    def tokenize(self, logits: torch.Tensor) -> torch.Tensor:
+        return torch.softmax(logits, dim=-1, dtype=torch.float32).to(logits.dtype)
+
+    def encode(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thws: torch.Tensor,
+    ) -> torch.Tensor:
+        features = self.vit(pixel_values, grid_thws)
+        seq_len, _ = features.shape
+        return features.reshape(seq_len // (self.config.hidden_stride**2), -1)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        grid_thws: torch.Tensor,
+    ) -> torch.Tensor:
+        features = self.encode(pixel_values, grid_thws)
+        logits = self.head(features)
+        tokens = self.tokenize(logits)
+        return torch.nn.functional.pad(
+            tokens,
+            (0, len(OXR_INDICATOR_IDS)),
+            mode="constant",
+            value=0,
+        )
+
+
+class OXRProcessingInfo(BaseProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config()
+
+    def get_hf_processor(self, **kwargs):
+        vit_config = self.get_hf_config().vit_config
+        return self.ctx.get_hf_processor(
+            OXRProcessor,
+            image_pad_token=self.get_image_pad_token(),
+            patch_size=vit_config.patch_size,
+            hidden_stride=vit_config.hidden_stride,
+            temporal_patch_size=vit_config.temporal_patch_size,
+            **kwargs,
+        )
+
+    def get_default_tok_params(self) -> TokenizeParams:
+        return super().get_default_tok_params().with_kwargs(add_special_tokens=False)
+
+    def get_image_pad_token(self) -> str:
+        text_model_type = get_llm_config(self.get_hf_config()).model_type
+        return IMAGE_PAD_TOKEN_MAP[text_model_type]
+
+    def get_image_processor(self) -> BaseImageProcessor:
+        return self.get_hf_processor().image_processor  # type: ignore[return-value]
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"image": None, "video": 1}
+
+    def get_image_size_with_most_features(self) -> ImageSize:
+        return ImageSize(width=1792, height=1792)
+
+    def get_num_image_tokens(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+        num_frames: int = 1,
+    ) -> int:
+        vit_config = self.get_hf_config().vit_config
+        padded_num_frames = num_frames + (-num_frames % vit_config.temporal_patch_size)
+        grid_t = max(padded_num_frames // vit_config.temporal_patch_size, 1)
+        grid_h = image_height // vit_config.patch_size
+        grid_w = image_width // vit_config.patch_size
+        return grid_t * grid_h * grid_w
+
+    def get_max_image_tokens(self) -> int:
+        target_width, target_height = self.get_image_size_with_most_features()
+        return self.get_num_image_tokens(
+            image_width=target_width,
+            image_height=target_height,
+        )
+
+    def _get_max_video_frames(self, max_tokens: int) -> int:
+        target_width, target_height = self.get_image_size_with_most_features()
+        num_frames = 0
+        while True:
+            next_num_frames = num_frames + 1
+            next_max_tokens = self.get_num_video_tokens(
+                image_width=target_width,
+                image_height=target_height,
+                num_frames=next_num_frames,
+            )
+            if next_max_tokens > max_tokens:
+                break
+            num_frames = next_num_frames
+        return num_frames
+
+    def get_num_frames_with_most_features(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> int:
+        max_images = mm_counts.get("image", 0)
+        max_videos = mm_counts.get("video", 0)
+        max_image_tokens = self.get_max_image_tokens() * max_images
+        max_total_frames = self._get_max_video_frames(seq_len - max_image_tokens)
+        max_frames_per_video = max_total_frames // max(max_videos, 1)
+        return max(max_frames_per_video, 1)
+
+    def get_num_video_tokens(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+        num_frames: int,
+    ) -> int:
+        return self.get_num_image_tokens(
+            image_width=image_width,
+            image_height=image_height,
+            num_frames=num_frames,
+        )
+
+    def get_max_video_tokens(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> int:
+        target_width, target_height = self.get_image_size_with_most_features()
+        return self.get_num_video_tokens(
+            image_width=target_width,
+            image_height=target_height,
+            num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
+        )
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        # Return a conservative upper bound so vLLM does not profile the
+        # visual tower with dozens of maximum-sized dummy images/videos.
+        return {
+            "image": self.get_max_image_tokens() + 2,
+            "video": self.get_max_video_tokens(seq_len, mm_counts) + 2,
+        }
+
+
+class OXRDummyInputsBuilder(BaseDummyInputsBuilder[OXRProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        return IMAGE_TOKEN * mm_counts.get("image", 0) + VIDEO_TOKEN * mm_counts.get(
+            "video", 0
+        )
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions],
+    ) -> MultiModalDataDict:
+        num_images = mm_counts.get("image", 0)
+        num_videos = mm_counts.get("video", 0)
+        target_width, target_height = self.info.get_image_size_with_most_features()
+        target_num_frames = self.info.get_num_frames_with_most_features(
+            seq_len, mm_counts
+        )
+
+        return {
+            "image": self._get_dummy_images(
+                width=target_width,
+                height=target_height,
+                num_images=num_images,
+                overrides=mm_options.get("image"),
+            ),
+            "video": self._get_dummy_videos(
+                width=target_width,
+                height=target_height,
+                num_frames=target_num_frames,
+                num_videos=num_videos,
+                overrides=mm_options.get("video"),
+            ),
+        }
+
+
+class OXRMultiModalProcessor(BaseMultiModalProcessor[OXRProcessingInfo]):
+    def visual_indicators_to_visual_tokens(
+        self,
+        visual_indicators: list[int],
+    ) -> list[int]:
+        vocab_size = self.info.get_hf_config().visual_vocab_size
+        indicator_by_id = {token_id: idx for idx, token_id in enumerate(OXR_INDICATOR_IDS)}
+        return [
+            vocab_size - len(OXR_INDICATOR_IDS) + indicator_by_id[token_id]
+            for token_id in visual_indicators
+            if token_id in indicator_by_id
+        ]
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        if not mm_data:
+            tokenizer = self.info.get_tokenizer()
+            prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
+            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
+
+        processed_outputs = super()._call_hf_processor(
+            prompt=prompt,
+            mm_data=mm_data,
+            mm_kwargs=mm_kwargs,
+            tok_kwargs=tok_kwargs,
+        )
+        hf_processor = self.info.get_hf_processor()
+
+        if "videos" in mm_data:
+            visual_indicators = [
+                hf_processor.construct_visual_indicators((1, 1, 1), True)
+                for _ in processed_outputs["video_grids"]
+            ]
+            processed_outputs["video_indicator_tokens"] = torch.tensor(
+                [
+                    self.visual_indicators_to_visual_tokens(indicators)
+                    for indicators in visual_indicators
+                ],
+                dtype=torch.long,
+            )
+
+        if "images" in mm_data:
+            visual_indicators = [
+                hf_processor.construct_visual_indicators((1, 1, 1), False)
+                for _ in processed_outputs["grids"]
+            ]
+            processed_outputs["indicator_tokens"] = torch.tensor(
+                [
+                    self.visual_indicators_to_visual_tokens(indicators)
+                    for indicators in visual_indicators
+                ],
+                dtype=torch.long,
+            )
+
+        return processed_outputs
+
+    def _apply_hf_processor_tokens_only(self, prompt_tokens: list[int]) -> list[int]:
+        return prompt_tokens
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: BatchFeature,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        return _oxr_field_config()
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> list[PromptReplacement]:
+        def get_replacement(item_idx: int, modality: str) -> list[int]:
+            if modality == "image":
+                grid = out_mm_kwargs["image"][item_idx]["grids"].data
+            elif modality == "video":
+                grid = out_mm_kwargs["video"][item_idx]["video_grids"].data
+            else:
+                raise ValueError(f"Unsupported modality: {modality}")
+
+            hf_processor = self.info.get_hf_processor()
+            return hf_processor.construct_visual_placeholders(
+                grid[0],
+                is_video=(modality == "video"),
+            )
+
+        return [
+            PromptReplacement(
+                modality=modality,
+                target=IMAGE_TOKEN if modality == "image" else VIDEO_TOKEN,
+                replacement=partial(get_replacement, modality=modality),
+            )
+            for modality in ("image", "video")
+        ]
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    OXRMultiModalProcessor,
+    info=OXRProcessingInfo,
+    dummy_inputs=OXRDummyInputsBuilder,
+)
+class OXR(nn.Module, SupportsMultiModal, SupportsPP):
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("image"):
+            return IMAGE_TOKEN
+        if modality.startswith("video"):
+            return VIDEO_TOKEN
+        raise ValueError("Only image or video modality is supported")
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.config: PretrainedConfig = config
+
+        with self._mark_language_model(vllm_config):
+            self.llm = init_vllm_registered_model(
+                vllm_config=vllm_config.with_hf_config(get_llm_config(config)),
+                prefix=maybe_prefix(prefix, "llm"),
+            )
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual_tokenizer = OXRVisualTokenizer(
+                config=config.vit_config,
+                visual_vocab_size=config.visual_vocab_size,
+                quant_config=quant_config,
+                prefix=f"{prefix}.visual_tokenizer",
+            )
+            self.vte = OXREmbedding(config.visual_vocab_size, config.hidden_size)
+
+        text_model_type = get_llm_config(config).model_type
+        self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]
+
+        self.make_empty_intermediate_tensors = (
+            self.get_language_model().make_empty_intermediate_tensors
+        )
+
+    def _parse_and_validate_image_input(
+        self,
+        **kwargs: object,
+    ) -> OXRImagePatchInputs | None:
+        pixel_values = kwargs.pop("pixel_values", None)
+        indicator_tokens = kwargs.pop("indicator_tokens", None)
+        grids = kwargs.pop("grids", None)
+        if pixel_values is None and indicator_tokens is None:
+            return None
+
+        if pixel_values is not None and indicator_tokens is not None:
+            if not isinstance(pixel_values, (torch.Tensor, list)):
+                raise ValueError(f"Incorrect type of pixel values: {type(pixel_values)}")
+            if not isinstance(indicator_tokens, (torch.Tensor, list)):
+                raise ValueError(
+                    f"Incorrect type of indicator_tokens: {type(indicator_tokens)}"
+                )
+            if grids is None:
+                raise ValueError("Missing grids for image input")
+
+            return OXRImagePatchInputs(
+                type="image_patches",
+                flat_data=flatten_bn(pixel_values, concat=True),
+                patches_per_item=[
+                    x.shape[0] // (self.config.vit_config.hidden_stride**2)
+                    for x in pixel_values
+                ],
+                indicator_tokens=flatten_bn(indicator_tokens, concat=True),
+                grids=flatten_bn(grids, concat=True),
+            )
+
+        raise AssertionError("This line should be unreachable.")
+
+    def _parse_and_validate_video_input(
+        self,
+        **kwargs: object,
+    ) -> OXRVideoPatchInputs | None:
+        pixel_values = kwargs.pop("video_pixel_values", None)
+        indicator_tokens = kwargs.pop("video_indicator_tokens", None)
+        grids = kwargs.pop("video_grids", None)
+        if pixel_values is None and indicator_tokens is None:
+            return None
+
+        if pixel_values is not None and indicator_tokens is not None:
+            if not isinstance(pixel_values, (torch.Tensor, list)):
+                raise ValueError(f"Incorrect type of pixel values: {type(pixel_values)}")
+            if not isinstance(indicator_tokens, (torch.Tensor, list)):
+                raise ValueError(
+                    f"Incorrect type of indicator_tokens: {type(indicator_tokens)}"
+                )
+            if grids is None:
+                raise ValueError("Missing grids for video input")
+
+            return OXRVideoPatchInputs(
+                type="video_patches",
+                flat_data=flatten_bn(pixel_values, concat=True),
+                patches_per_item=[
+                    x.shape[0] // (self.config.vit_config.hidden_stride**2)
+                    for x in pixel_values
+                ],
+                indicator_tokens=flatten_bn(indicator_tokens, concat=True),
+                grids=flatten_bn(grids, concat=True),
+            )
+
+        raise AssertionError("This line should be unreachable.")
+
+    def _process_visual_input(
+        self,
+        visual_input: OXRImagePatchInputs | OXRVideoPatchInputs,
+    ) -> MultiModalEmbeddings:
+        image_patches_flat = visual_input["flat_data"]
+        patches_per_item = visual_input["patches_per_item"]
+        indicator_tokens = visual_input["indicator_tokens"]
+        grid_thws = visual_input["grids"]
+
+        visual_tokens = self.visual_tokenizer(
+            image_patches_flat.to(self.visual_tokenizer.dtype),
+            grid_thws,
+        )
+        visual_embeds = self.vte(visual_tokens)
+        indicator_embeds = self.vte(indicator_tokens)
+
+        visual_embeds_per_item = visual_embeds.split(patches_per_item, dim=0)
+        indicator_embeds_per_item = indicator_embeds.split(
+            [2] * len(patches_per_item)
+        )
+
+        return tuple(
+            torch.cat([indicator[:1], visual, indicator[1:]], dim=0)
+            for indicator, visual in zip(
+                indicator_embeds_per_item,
+                visual_embeds_per_item,
+            )
+        )
+
+    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
+        modalities = {}
+        for input_key in kwargs:
+            if (
+                input_key in ("pixel_values", "indicator_tokens", "grids")
+                and "images" not in modalities
+            ):
+                modalities["images"] = self._parse_and_validate_image_input(**kwargs)
+            if (
+                input_key
+                in ("video_pixel_values", "video_indicator_tokens", "video_grids")
+                and "videos" not in modalities
+            ):
+                modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
+        return modalities
+
+    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+        modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
+        if not modalities:
+            return []
+
+        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
+        for modality in modalities:
+            if modality == "images":
+                image_input = modalities["images"]
+                multimodal_embeddings += tuple(self._process_visual_input(image_input))
+            if modality == "videos":
+                video_input = modalities["videos"]
+                multimodal_embeddings += tuple(self._process_visual_input(video_input))
+        return multimodal_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        return self.llm(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.llm.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)
+
+    def get_language_model(self) -> torch.nn.Module:
+        return self.llm

plugins/vllm/oxr_vllm/oxr_processor.py

@@ -0,0 +1,384 @@
+"""OXR processor for the vLLM 0.25.1 plugin."""
+
+import math
+from functools import cached_property
+from typing import Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from transformers import AutoProcessor, BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+
+IMAGE_TOKEN = "<image>"
+VIDEO_TOKEN = "<video>"
+MIN_PIXELS = 200704
+MAX_PIXELS = 921600
+
+
+class OXRProcessorKwargs(ProcessingKwargs, total=False):  # type: ignore[call-arg]
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+        "images_kwargs": {
+            "do_convert_rgb": True,
+        },
+        "videos_kwargs": {
+            "do_convert_rgb": True,
+        },
+    }
+
+
+class OXRProcessor(ProcessorMixin):
+    attributes = ["image_processor", "tokenizer"]
+    valid_kwargs = ["chat_template", "image_pad_token"]
+
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        chat_template=None,
+        image_pad_token=None,
+        patch_size=16,
+        hidden_stride=2,
+        temporal_patch_size=1,
+        **kwargs,
+    ):
+        self.image_token = IMAGE_TOKEN
+        self.video_token = VIDEO_TOKEN
+        self.image_pad_token = image_pad_token or "<|image_pad|>"
+        self.patch_size = patch_size
+        self.hidden_stride = hidden_stride
+        self.temporal_patch_size = temporal_patch_size
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    @cached_property
+    def extra_special_tokens(self):
+        image_pad_token_id = self.tokenizer.get_vocab()[self.image_pad_token]
+        return {
+            "image_token": -200,
+            "video_token": -201,
+            "visual_atom": -300,
+            "image_start": -301,
+            "image_end": -302,
+            "video_start": -303,
+            "video_end": -304,
+            "image_pad": image_pad_token_id,
+        }
+
+    def __call__(
+        self,
+        images: ImageInput = None,
+        videos: Union[np.ndarray, list[ImageInput]] = None,
+        text: Union[
+            TextInput,
+            PreTokenizedInput,
+            list[TextInput],
+            list[PreTokenizedInput],
+        ] = None,
+        **kwargs: Unpack[OXRProcessorKwargs],
+    ) -> BatchFeature:
+        output_kwargs = self._merge_kwargs(
+            OXRProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+        visual_features = {}
+        output = BatchFeature()
+
+        if images is not None:
+            processed_images = []
+            image_placeholders_list = []
+            grids = []
+            for image in images if isinstance(images, list) else [images]:
+                pixel_values, image_placeholders, grid = self.preprocess_multidata(
+                    images=image,
+                    **output_kwargs["images_kwargs"],
+                )
+                processed_images.append(pixel_values)
+                image_placeholders_list.append(image_placeholders)
+                grids.append(grid)
+
+            if processed_images:
+                visual_features["image_placeholders"] = image_placeholders_list
+            output["pixel_values"] = processed_images
+            output["grids"] = grids
+
+        if videos is not None:
+            processed_videos = []
+            video_placeholders_list = []
+            grids = []
+            for video in videos if isinstance(videos, list) else [videos]:
+                pixel_values, video_placeholders, grid = self.preprocess_multidata(
+                    video=video,
+                    **output_kwargs["videos_kwargs"],
+                )
+                processed_videos.append(pixel_values)
+                video_placeholders_list.append(video_placeholders)
+                grids.append(grid)
+
+            if processed_videos:
+                visual_features["video_placeholders"] = video_placeholders_list
+            output["video_pixel_values"] = processed_videos
+            output["video_grids"] = grids
+
+        if text is None:
+            return BatchFeature(data=visual_features)
+
+        if not isinstance(text, list):
+            text = [text]
+        tokenized_batched_text = self._tokenize_with_visual_symbol(text)
+        image_token_id = self.get_token_value("image_token")
+        video_token_id = self.get_token_value("video_token")
+        replaced_ids_list = []
+        image_idx = 0
+        video_idx = 0
+        for ids_tensor in tokenized_batched_text:
+            has_image_tokens = (
+                image_token_id in ids_tensor
+                and "image_placeholders" in visual_features
+                and image_idx < len(visual_features["image_placeholders"])
+            )
+            has_video_tokens = (
+                video_token_id in ids_tensor
+                and "video_placeholders" in visual_features
+                and video_idx < len(visual_features["video_placeholders"])
+            )
+            if has_image_tokens or has_video_tokens:
+                new_ids = []
+                for token_id in ids_tensor.tolist():
+                    if token_id == image_token_id:
+                        new_ids.extend(visual_features["image_placeholders"][image_idx])
+                        image_idx += 1
+                    elif token_id == video_token_id:
+                        new_ids.extend(visual_features["video_placeholders"][video_idx])
+                        video_idx += 1
+                    else:
+                        new_ids.append(token_id)
+                ids_tensor = torch.tensor(new_ids, dtype=torch.long)
+            replaced_ids_list.append(ids_tensor)
+
+        if replaced_ids_list:
+            output["input_ids"] = torch.stack(replaced_ids_list)
+        else:
+            output["input_ids"] = torch.tensor([], dtype=torch.long)
+        return output
+
+    def _tokenize_with_visual_symbol(self, text_list: list[str]) -> torch.LongTensor:
+        batch_token_ids = []
+        for text in text_list:
+            token_ids = []
+            video_token_id = self.get_token_value("video_token")
+            image_token_id = self.get_token_value("image_token")
+            video_split_texts = text.split(self.video_token)
+
+            for j, video_segment in enumerate(video_split_texts):
+                image_split_texts = video_segment.split(self.image_token)
+                text_chunks = [
+                    self.tokenizer(chunk, add_special_tokens=False).input_ids
+                    for chunk in image_split_texts
+                ]
+                segment_tokens = []
+                for i, chunk in enumerate(text_chunks):
+                    segment_tokens.extend(chunk)
+                    if i < len(text_chunks) - 1:
+                        segment_tokens.append(image_token_id)
+                token_ids.extend(segment_tokens)
+                if j < len(video_split_texts) - 1:
+                    token_ids.append(video_token_id)
+            batch_token_ids.append(token_ids)
+        return torch.tensor(batch_token_ids, dtype=torch.long)
+
+    def smart_resize(
+        self,
+        height: int,
+        width: int,
+        factor: int = 32,
+        min_pixels: int = MIN_PIXELS,
+        max_pixels: int = MAX_PIXELS,
+    ):
+        if height < factor or width < factor:
+            if height < width:
+                width = round(factor / height * width)
+                height = factor
+            else:
+                height = round(factor / width * height)
+                width = factor
+        elif max(height, width) / min(height, width) > 200:
+            if height > width:
+                height = 200 * width
+            else:
+                width = 200 * height
+
+        h_bar = round(height / factor) * factor
+        w_bar = round(width / factor) * factor
+        if h_bar * w_bar > max_pixels:
+            beta = math.sqrt((height * width) / max_pixels)
+            h_bar = math.floor(height / beta / factor) * factor
+            w_bar = math.floor(width / beta / factor) * factor
+        elif h_bar * w_bar < min_pixels:
+            beta = math.sqrt(min_pixels / (height * width))
+            h_bar = math.ceil(height * beta / factor) * factor
+            w_bar = math.ceil(width * beta / factor) * factor
+        return h_bar, w_bar
+
+    def get_token_value(self, token_name: str):
+        return self.extra_special_tokens[token_name]
+
+    def construct_visual_indicators(self, grid, is_video: bool = False):
+        if is_video:
+            start_token = self.get_token_value("video_start")
+            end_token = self.get_token_value("video_end")
+        else:
+            start_token = self.get_token_value("image_start")
+            end_token = self.get_token_value("image_end")
+
+        visual_indicators = [start_token, self.get_token_value("visual_atom")]
+        if grid[0] * grid[1] > 1:
+            for _ in range(grid[0]):
+                for _ in range(grid[1]):
+                    visual_indicators.append(self.get_token_value("visual_atom"))
+        visual_indicators.append(end_token)
+        return visual_indicators
+
+    def construct_visual_placeholders(self, grid, is_video: bool = False):
+        visual_placeholders = self.construct_visual_indicators((1, 1), is_video)
+        visual_atom_token_id = self.get_token_value("visual_atom")
+        image_padding_token_id = self.get_token_value("image_pad")
+
+        num_image_atoms = grid[0] * grid[1] * grid[2]
+        num_image_atoms //= self.hidden_stride**2
+        num_image_atoms //= self.temporal_patch_size
+
+        vision_start_token_id = self.tokenizer.get_vocab().get(
+            "<|vision_start|>",
+            image_padding_token_id,
+        )
+        vision_end_token_id = self.tokenizer.get_vocab().get(
+            "<|vision_end|>",
+            image_padding_token_id,
+        )
+
+        padded_placeholder_tokens = []
+        for token in visual_placeholders:
+            if token == visual_atom_token_id:
+                padded_placeholder_tokens.extend(
+                    [image_padding_token_id] * num_image_atoms
+                )
+            elif token in (
+                self.get_token_value("image_start"),
+                self.get_token_value("video_start"),
+            ):
+                padded_placeholder_tokens.append(vision_start_token_id)
+            elif token in (
+                self.get_token_value("image_end"),
+                self.get_token_value("video_end"),
+            ):
+                padded_placeholder_tokens.append(vision_end_token_id)
+            else:
+                padded_placeholder_tokens.append(image_padding_token_id)
+        return padded_placeholder_tokens
+
+    def preprocess_multidata(
+        self,
+        images: Optional[Union[PIL.Image.Image, list[PIL.Image.Image]]] = None,
+        video: Optional[Union[list[PIL.Image.Image], np.ndarray]] = None,
+        do_convert_rgb: Optional[bool] = True,
+        convert_to_rgb: Optional[bool] = None,
+        min_pixels: int = MIN_PIXELS,
+        max_pixels: int = MAX_PIXELS,
+        return_tensors: Optional[str] = "pt",
+    ):
+        if convert_to_rgb is not None:
+            do_convert_rgb = convert_to_rgb
+
+        is_video = False
+        if images is not None:
+            if not isinstance(images, list):
+                images = [images]
+        elif video is not None:
+            is_video = True
+            if isinstance(video, np.ndarray):
+                images = [
+                    PIL.Image.fromarray(video[i].astype(np.uint8))
+                    for i in range(video.shape[0])
+                ]
+            elif isinstance(video, list):
+                images = video
+        else:
+            raise ValueError("Expected either images or video")
+
+        min_pixels = min(
+            max_pixels if max_pixels is not None else MAX_PIXELS,
+            min_pixels if min_pixels is not None else MIN_PIXELS,
+        )
+        images = [
+            image.convert("RGB") if do_convert_rgb and image.mode != "RGB" else image
+            for image in images
+        ]
+
+        width, height = images[0].size
+        resized_height, resized_width = height, width
+        processed_images = []
+        for image in images:
+            resized_height, resized_width = self.smart_resize(
+                height,
+                width,
+                factor=self.patch_size * self.hidden_stride,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+            image_pt = self.image_processor.preprocess(
+                image,
+                size=dict(height=resized_height, width=resized_width),
+                return_tensors="np",
+            )["pixel_values"][0]
+            processed_images.append(image_pt)
+
+        patches = np.array(processed_images)
+        frame_remainder = patches.shape[0] % self.temporal_patch_size
+        if frame_remainder != 0:
+            num_to_pad = self.temporal_patch_size - frame_remainder
+            repeats = np.repeat(patches[-1][np.newaxis], num_to_pad, axis=0)
+            patches = np.concatenate([patches, repeats], axis=0)
+
+        channel = patches.shape[1]
+        grid_t = patches.shape[0] // self.temporal_patch_size
+        grid_h = resized_height // self.patch_size
+        grid_w = resized_width // self.patch_size
+
+        patches = patches.reshape(
+            grid_t,
+            self.temporal_patch_size,
+            channel,
+            grid_h // self.hidden_stride,
+            self.hidden_stride,
+            self.patch_size,
+            grid_w // self.hidden_stride,
+            self.hidden_stride,
+            self.patch_size,
+        )
+        patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
+        flatten_patches = patches.reshape(
+            grid_t * grid_h * grid_w,
+            channel * self.temporal_patch_size * self.patch_size * self.patch_size,
+        )
+
+        visual_placeholders = self.construct_visual_placeholders(
+            [grid_t, grid_h, grid_w],
+            is_video,
+        )
+        return (
+            torch.tensor(flatten_patches),
+            visual_placeholders,
+            torch.tensor([[grid_t, grid_h, grid_w]]),
+        )
+
+
+AutoProcessor.register("OXRProcessor", OXRProcessor)

plugins/vllm/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "oxr-vllm-plugin"
+version = "0.2.0"
+description = "Out-of-tree OXR model plugin for vLLM 0.25.1"
+requires-python = ">=3.12"
+dependencies = ["vllm==0.25.1", "transformers==5.13.1"]
+
+[project.entry-points."vllm.general_plugins"]
+oxr = "oxr_vllm:register"
+
+[tool.setuptools.packages.find]
+include = ["oxr_vllm", "oxr_vllm.*"]

preprocessor_config.json

@@ -0,0 +1,24 @@
+{
+  "do_convert_rgb": null,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "SiglipImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "processor_class": "SiglipProcessor",
+  "resample": 2,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 512,
+    "width": 512
+  }
+}

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|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|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,240 @@
+{
+  "add_bos_token": false,
+  "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
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|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,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "SiglipProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}