Upload folder using huggingface_hub

.gitattributes

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

README.md

@@ -0,0 +1,226 @@
+---
+library_name: transformers
+base_model:
+- stepfun-ai/Step3-VL-10B
+---
+
+This tiny model is intended for debugging. It is randomly initialized using the configuration adapted from [stepfun-ai/Step3-VL-10B](https://huggingface.co/stepfun-ai/Step3-VL-10B).
+
+| File path | Size |
+|------|------|
+| model.safetensors | 6.0MB |
+
+
+### Example usage:
+
+- vLLM
+
+```bash
+vllm serve tiny-random/step3-vl \
+    --trust-remote-code \
+    --reasoning-parser deepseek_r1 \
+    --enable-auto-tool-choice \
+    --tool-call-parser hermes
+```
+
+- Transformers
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+model_id = "tiny-random/step3-vl"
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG"
+            },
+            {
+                "type": "text",
+                "text": "describe this image"
+            }
+        ],
+    }
+]
+processor = AutoProcessor.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="cuda",
+    trust_remote_code=True,
+    key_mapping={
+        "^vision_model": "model.vision_model",
+        r"^model(?!\.(language_model|vision_model))": "model.language_model",
+        "vit_large_projector": "model.vit_large_projector",
+    }
+)
+inputs = processor.apply_chat_template(
+    messages,
+    tokenize=True,
+    add_generation_prompt=True,
+    return_dict=True,
+    return_tensors="pt"
+).to(model.device)
+inputs.pop("token_type_ids", None)
+generated_ids = model.generate(**inputs, max_new_tokens=16)
+output_text = processor.decode(
+    generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=False)
+print(output_text)
+```
+
+### Codes to create this repo:
+
+<details>
+<summary>Python codes</summary>
+```python
+import json
+from pathlib import Path
+
+import accelerate
+import torch
+from huggingface_hub import file_exists, hf_hub_download, list_repo_files
+from safetensors.torch import save_file
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoProcessor,
+    AutoTokenizer,
+    GenerationConfig,
+    set_seed,
+)
+
+source_model_id = "stepfun-ai/Step3-VL-10B"
+save_folder = "/tmp/tiny-random/step3-vl"
+
+Path(save_folder).mkdir(parents=True, exist_ok=True)
+for f in list_repo_files(source_model_id, repo_type="model"):
+    if (f.endswith('.json') or f.endswith('.py') or f.endswith('.model') or f.endswith('.jinja')) and (
+        not f.endswith('.index.json')
+    ):
+        hf_hub_download(repo_id=source_model_id, filename=f,
+                        repo_type="model", local_dir=save_folder)
+
+def replace_file(filepath, old_string, new_string):
+    with open(filepath, 'r', encoding='utf-8') as f:
+        code = f.read()
+    code = code.replace(old_string, new_string)
+    with open(filepath, 'w', encoding='utf-8') as f:
+        f.write(code)
+
+with open(f'{save_folder}/config.json') as f:
+    config_json = json.load(f)
+
+config_json['text_config'].update({
+    'num_hidden_layers': 2,
+    'hidden_size': 8,
+    'head_dim': 32,
+    'intermediate_size': 64,
+    'num_attention_heads': 8,
+    "num_key_value_heads": 4,
+    'tie_word_embeddings': False,
+})
+config_json['vision_config'].update({
+    'width': 64,
+    'layers': 2,
+    'heads': 2,
+})
+with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
+    json.dump(config_json, f, indent=2)
+
+config = AutoConfig.from_pretrained(
+    save_folder,
+    trust_remote_code=True,
+)
+print(config)
+torch.set_default_dtype(torch.bfloat16)
+model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
+torch.set_default_dtype(torch.float32)
+# if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
+#     model.generation_config = GenerationConfig.from_pretrained(
+#         source_model_id, trust_remote_code=True,
+#     )
+set_seed(42)
+model = model.cpu()
+with torch.no_grad():
+    for name, p in sorted(model.named_parameters()):
+        torch.nn.init.normal_(p, 0, 0.1)
+        print(name, p.shape)
+model_new = torch.nn.Identity()
+model_new.model = model.model.language_model
+model_new.vision_model = model.model.vision_model
+model_new.lm_head = model.lm_head
+model_new.vit_large_projector = model.model.vit_large_projector
+state_dict = model_new.state_dict()
+save_file(state_dict, f"{save_folder}/model.safetensors")
+```
+</details>
+
+### Printing the model:
+
+```text
+Step3VL10BForCausalLM(
+  (model): StepRoboticsModel(
+    (vision_model): StepRoboticsVisionEncoder(
+      (conv1): Conv2d(3, 64, kernel_size=(14, 14), stride=(14, 14), bias=False)
+      (ln_pre): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
+      (ln_post): Identity()
+      (transformer): EncoderVisionTransformer(
+        (resblocks): ModuleList(
+          (0-1): 2 x EncoderVisionBlock(
+            (attn): EncoderVisionAttention(
+              (out_proj): Linear(in_features=64, out_features=64, bias=True)
+              (rope): EncoderRope2D()
+            )
+            (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
+            (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
+            (mlp): EncoderMLP(
+              (c_fc): Linear(in_features=64, out_features=373, bias=True)
+              (act_fn): QuickGELUActivation()
+              (c_proj): Linear(in_features=373, out_features=64, bias=True)
+            )
+            (ls_1): EncoderLayerScale()
+            (ls_2): EncoderLayerScale()
+          )
+        )
+      )
+      (vit_downsampler1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+      (vit_downsampler2): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
+    )
+    (language_model): Qwen3Model(
+      (embed_tokens): Embedding(151936, 8)
+      (layers): ModuleList(
+        (0-1): 2 x Qwen3DecoderLayer(
+          (self_attn): Qwen3Attention(
+            (q_proj): Linear(in_features=8, out_features=256, bias=False)
+            (k_proj): Linear(in_features=8, out_features=128, bias=False)
+            (v_proj): Linear(in_features=8, out_features=128, bias=False)
+            (o_proj): Linear(in_features=256, out_features=8, bias=False)
+            (q_norm): Qwen3RMSNorm((32,), eps=1e-06)
+            (k_norm): Qwen3RMSNorm((32,), eps=1e-06)
+          )
+          (mlp): Qwen3MLP(
+            (gate_proj): Linear(in_features=8, out_features=64, bias=False)
+            (up_proj): Linear(in_features=8, out_features=64, bias=False)
+            (down_proj): Linear(in_features=64, out_features=8, bias=False)
+            (act_fn): SiLUActivation()
+          )
+          (input_layernorm): Qwen3RMSNorm((8,), eps=1e-06)
+          (post_attention_layernorm): Qwen3RMSNorm((8,), eps=1e-06)
+        )
+      )
+      (norm): Qwen3RMSNorm((8,), eps=1e-06)
+      (rotary_emb): Qwen3RotaryEmbedding()
+    )
+    (vit_large_projector): Linear(in_features=256, out_features=8, bias=False)
+  )
+  (lm_head): Linear(in_features=8, out_features=151936, bias=False)
+)
+```

added_tokens.json

@@ -0,0 +1,51 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_calls>": 151670,
+  "</tool_response>": 151666,
+  "<dream>": 151682,
+  "<dream_end>": 151684,
+  "<dream_start>": 151683,
+  "<im_end>": 151681,
+  "<im_patch>": 151679,
+  "<im_start>": 151680,
+  "<patch_end>": 151690,
+  "<patch_newline>": 151691,
+  "<patch_start>": 151689,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_calls>": 151669,
+  "<tool_response>": 151665,
+  "<video_end>": 151688,
+  "<video_start>": 151687,
+  "<|BOT|>": 151672,
+  "<|CALL_END|>": 151674,
+  "<|CALL_START|>": 151673,
+  "<|EOT|>": 151671,
+  "<|IMG_END|>": 151678,
+  "<|IMG_START|>": 151677,
+  "<|MASK_1e69f|>": 151685,
+  "<|THINK_END|>": 151676,
+  "<|THINK_START|>": 151675,
+  "<|UNMASK_1e69f|>": 151686,
+  "<|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,81 @@
+{% macro render_content(content) %}{% if content is none %}{{- '' }}{% elif content is string %}{{- content }}{% elif content is mapping %}{{- content['value'] if 'value' in content else content['text'] }}{% elif content is iterable %}{% for item in content %}{% if item.type == 'text' %}{{- item['value'] if 'value' in item else item['text'] }}{% elif item.type == 'image' %}<im_patch>{% endif %}{% endfor %}{% endif %}{% endmacro %}
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- render_content(messages[0].content) + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nAlways adhere to this exact format for tool use:\n<tool_calls>\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>\n{additional_tool_calls}</tool_calls>\n\nNote:\n- For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags.\n- `<function-name>` must be an exact match to one of the available tools.\n- `<args-json-object>` must be valid JSON that strictly follows the tool's parameters schema.<|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + render_content(messages[0].content) + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and render_content(message.content) is string and not(render_content(message.content).startswith('<tool_response>') and render_content(message.content).endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content) %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {%- set role_name = 'observation' if (message.role == "system" and not loop.first and message.name == 'observation') else message.role %}
+        {{- '<|im_start|>' + role_name + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = render_content(message.reasoning_content) %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- else %}
+                {%- set reasoning_content = '' %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {{- '\n<tool_calls>' }}
+            {%- for tool_call in message.tool_calls %}
+                {{- '\n' }}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+            {{- '\n</tool_calls>' }}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>tool_response' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n<think>\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,67 @@
+{
+  "architectures": [
+    "StepVLForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_step_vl.StepRoboticsConfig",
+    "AutoModelForCausalLM": "modeling_step_vl.Step3VL10BForCausalLM"
+  },
+  "model_type": "step_robotics",
+  "im_end_token": "<im_end>",
+  "im_patch_token": "<im_patch>",
+  "im_start_token": "<im_start>",
+  "image_token_len": 169,
+  "patch_token_len": 81,
+  "image_token_id": 151679,
+  "understand_projector_stride": 2,
+  "use_im_start_end": "true",
+  "vision_select_layer": -1,
+  "projector_bias": false,
+  "vision_config": {
+    "image_size": 728,
+    "patch_size": 14,
+    "width": 64,
+    "layers": 2,
+    "heads": 2,
+    "pool_type": "none",
+    "output_dim": null,
+    "use_cls_token": false,
+    "ls_init_value": 0.1,
+    "use_ln_post": false,
+    "hidden_act": "quick_gelu"
+  },
+  "text_config": {
+    "architectures": [
+      "Qwen3ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 151643,
+    "eos_token_id": [
+      151643,
+      151645,
+      151679
+    ],
+    "head_dim": 32,
+    "hidden_act": "silu",
+    "hidden_size": 8,
+    "initializer_range": 0.02,
+    "intermediate_size": 64,
+    "max_position_embeddings": 65536,
+    "max_window_layers": 36,
+    "model_type": "qwen3",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 2,
+    "num_key_value_heads": 4,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000,
+    "sliding_window": null,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.51.0",
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+  }
+}

configuration_step_vl.py

@@ -0,0 +1,77 @@
+from typing import Any, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers import Qwen3Config
+
+
+class StepRoboticsVisionEncoderConfig(PretrainedConfig):
+
+    def __init__(
+        self,
+        width=1536,
+        layers=47,
+        heads=16,
+        num_channels=3,
+        image_size=728,
+        mlp_ratio = 8960/1536,
+        patch_size=14,
+        hidden_act="quick_gelu",
+        layer_norm_eps=1e-5,
+        ues_cls_token=False,
+        use_ln_pre=True,
+        use_ln_post=False,
+        use_abs_posemb=True,
+        use_rope2d=True,
+        ls_init_value=0.1,
+        **kwargs,
+    ):
+        self.width = width
+        self.layers = layers
+        self.heads = heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.mlp_ratio = mlp_ratio
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.ues_cls_token = ues_cls_token
+        self.use_ln_pre = use_ln_pre
+        self.ls_init_value = ls_init_value
+        self.use_ln_post = use_ln_post
+        self.use_abs_posemb = use_abs_posemb
+        self.use_rope2d = use_rope2d
+        super().__init__(**kwargs)
+
+
+
+class StepRoboticsConfig(PretrainedConfig):
+    model_type = "step_robotics"
+    architectures = ["StepVLForConditionalGeneration"]
+
+    def __init__(
+        self,
+        vision_config: Optional[Union[dict, StepRoboticsVisionEncoderConfig]] = None,
+        text_config: Optional[Union[dict, Qwen3Config]] = None,
+        understand_projector_stride: int = 2,
+        projector_bias: bool = False,
+        image_token_id: int = 151679,
+        **kwargs,
+    ) -> None:
+        if vision_config is None:
+            vision_config = StepRoboticsVisionEncoderConfig()
+        elif isinstance(vision_config, dict):
+            vision_config = StepRoboticsVisionEncoderConfig(**vision_config)
+        self.vision_config = vision_config
+
+        if text_config is None:
+            text_config = Qwen3Config()
+        elif isinstance(text_config, dict):
+            text_config = Qwen3Config(**text_config)
+        self.text_config = text_config
+
+        self.understand_projector_stride = understand_projector_stride
+        self.projector_bias = projector_bias
+        self.hidden_size = text_config.hidden_size
+        self.image_token_id = image_token_id
+        # Help Auto classes find the correct implementation when saving/loading.
+        super().__init__(**kwargs)

generation_config.json

@@ -0,0 +1,10 @@
+{
+    "temperature": 1.0,
+    "top_p": 1.0,
+    "top_k": 0,
+    "eos_token_id": [
+        151643,
+        151645,
+        151679
+    ]
+}

model.safetensors

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

modeling_step_vl.py

@@ -0,0 +1,557 @@
+# Copyright 2025 The LLAMA4 and HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple, Union
+from PIL import Image
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import Qwen3Model
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, can_return_tuple, logging
+
+from typing import Any, Literal, Optional, TypedDict, Union
+
+from .configuration_step_vl import StepRoboticsConfig
+from .vision_encoder import StepRoboticsVisionEncoder
+logger = logging.get_logger(__name__)
+
+class StepVLImagePixelInputs(TypedDict):
+    type: Literal["pixel_values"]
+    pixel_values: torch.Tensor
+    patch_pixel_values: Optional[torch.Tensor]
+    num_patches: list[int]
+
+
+class StepVLImageEmbeddingInputs(TypedDict):
+    type: Literal["image_embeds"]
+    image_embeds: torch.Tensor
+
+
+StepVLImageInputs = Union[StepVLImagePixelInputs,
+                           StepVLImageEmbeddingInputs]
+
+
+@dataclass
+class StepVLCausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[list[torch.FloatTensor]] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
+    image_hidden_states: Optional[torch.FloatTensor] = None
+
+def _flatten_embeddings(embeddings) -> torch.Tensor:
+    """
+    Recursively flattens and concatenates NestedTensors on all but the last
+    dimension.
+    """
+
+    if isinstance(embeddings, torch.Tensor):
+        # Flatten all but the last dimension.
+        return embeddings.flatten(0, -2)
+
+    return torch.cat(tuple(_flatten_embeddings(t) for t in embeddings))
+
+def _embedding_count_expression(embeddings) -> str:
+    """
+    Constructs a debugging representation of the number of embeddings in the
+    NestedTensors.
+    """
+
+    if isinstance(embeddings, torch.Tensor):
+        return " x ".join([str(dim) for dim in embeddings.shape[:-1]])
+
+    return " + ".join(
+        _embedding_count_expression(inner) for inner in embeddings)
+
+def _merge_multimodal_embeddings(
+    inputs_embeds: torch.Tensor,
+    is_multimodal: torch.Tensor,
+    multimodal_embeddings,
+) -> torch.Tensor:
+    """
+    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
+    positions in ``inputs_embeds`` corresponding to placeholder tokens in
+    ``input_ids``.
+    Note:
+        This updates ``inputs_embeds`` in place.
+    """
+    num_expected_tokens = is_multimodal.sum().item()
+    assert isinstance(num_expected_tokens, int)
+
+    flattened = _flatten_embeddings(multimodal_embeddings)
+    if flattened.shape[0] != num_expected_tokens:
+        expr = _embedding_count_expression(multimodal_embeddings)
+        raise ValueError(
+            f"Attempted to assign {expr} = {flattened.shape[0]} "
+            f"multimodal tokens to {num_expected_tokens} placeholders")
+
+    is_multimodal = is_multimodal.to(inputs_embeds.device)
+    flattened = flattened.to(inputs_embeds.device)
+    inputs_embeds[is_multimodal] = flattened
+    return inputs_embeds
+
+def merge_multimodal_embeddings(
+    input_ids: torch.Tensor,
+    inputs_embeds: torch.Tensor,
+    multimodal_embeddings,
+    placeholder_token_id: Union[int, list[int]],
+) -> torch.Tensor:
+    """
+    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
+    positions in ``inputs_embeds`` corresponding to placeholder tokens in
+    ``input_ids``.
+    
+    ``placeholder_token_id`` can be a list of token ids (e.g, token ids 
+    of img_start, img_break, and img_end tokens) when needed: This means 
+    the order of these tokens in the ``input_ids`` MUST MATCH the order of 
+    their embeddings in ``multimodal_embeddings`` since we need to 
+    slice-merge instead of individually scattering.
+    For example, if input_ids is "TTTTTSIIIBIIIBIIIETTT", where
+    - T is text token
+    - S is image start token
+    - I is image embedding token
+    - B is image break token
+    - E is image end token.
+    
+    Then the image embeddings (that correspond to I's) from vision encoder 
+    must be padded with embeddings of S, B, and E in the same order of 
+    input_ids for a correct embedding merge.
+    Note:
+        This updates ``inputs_embeds`` in place.
+    """
+    if isinstance(placeholder_token_id, list):
+        placeholder_token_id = torch.tensor(placeholder_token_id,
+                                            device=input_ids.device)
+        return _merge_multimodal_embeddings(
+            inputs_embeds,
+            torch.isin(input_ids, placeholder_token_id),
+            multimodal_embeddings,
+        )
+
+    return _merge_multimodal_embeddings(
+        inputs_embeds,
+        (input_ids == placeholder_token_id),
+        multimodal_embeddings,
+    )
+
+class StepRoboticsPreTrainedModel(PreTrainedModel):
+    # Link this model family to its configuration class so PreTrainedModel.from_pretrained
+    # can load the config instead of failing with a NoneType error.
+    config_class = StepRoboticsConfig
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = False
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+
+class StepRoboticsModel(StepRoboticsPreTrainedModel, GenerationMixin):
+    config: StepRoboticsConfig
+    base_model_prefix = ""
+    def __init__(self, config: StepRoboticsConfig):
+        super().__init__(config)
+        self.vision_model = StepRoboticsVisionEncoder(config.vision_config)
+        self.language_model = Qwen3Model(config.text_config)
+        self.vocab_size = config.text_config.vocab_size
+        self.vit_large_projector = nn.Linear(
+                config.vision_config.width * 4,
+                config.text_config.hidden_size,                
+                bias=config.projector_bias) 
+        self.image_placeholder_token_id = config.image_token_id
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings  = None,
+    ) -> torch.Tensor:
+        input_ids = input_ids.squeeze(0)
+        if multimodal_embeddings is None:
+            inputs_embeds = self.language_model.embed_tokens(input_ids)
+        else:
+            is_text = input_ids != self.config.image_token_id
+            text_ids = input_ids[is_text]
+            text_embeds = self.language_model.embed_tokens(text_ids)
+                     
+            inputs_embeds = torch.empty(input_ids.shape[0],
+                                        text_embeds.shape[-1],
+                                        dtype=text_embeds.dtype,
+                                        device=text_embeds.device)
+            inputs_embeds[is_text] = text_embeds
+            inputs_embeds = merge_multimodal_embeddings(
+                input_ids, inputs_embeds, multimodal_embeddings,
+                self.config.image_token_id)
+        inputs_embeds = inputs_embeds.unsqueeze(0)
+        return inputs_embeds
+       
+
+    def set_input_embeddings(self, value):
+        return self.language_model.set_input_embeddings(value)
+
+    def set_decoder(self, decoder):
+        self.language_model = decoder
+
+    def get_decoder(self):
+        return self.language_model
+    
+    def _parse_and_validate_image_input(
+            self, **kwargs: object) -> Optional[StepVLImageInputs]:
+        pixel_values = kwargs.pop("pixel_values", None)
+        patch_pixel_values = kwargs.pop("patch_pixel_values", None)
+        num_patches = kwargs.pop("num_patches", None)
+        image_embeds = kwargs.pop("image_embeds", None)
+
+        if pixel_values is None and image_embeds is None:
+            return None
+
+        if pixel_values is not None:
+            # pixel_values = flatten_bn(pixel_values, concat=True)
+            if pixel_values.dim() >= 3:
+                pixel_values = pixel_values.view(-1, *pixel_values.shape[-3:])
+            if patch_pixel_values is not None:
+                # patch_pixel_values = flatten_bn(patch_pixel_values,
+                #                                 concat=True)
+                patch_pixel_values = patch_pixel_values.view(
+                    -1, *patch_pixel_values.shape[-3:])
+                # Handle empty patch_pixel_values by setting to None
+                if patch_pixel_values.shape[0] == 0:
+                    patch_pixel_values = None
+
+            return StepVLImagePixelInputs(
+                type="pixel_values",
+                pixel_values=pixel_values.to(self.dtype).to(self.device),
+                patch_pixel_values=patch_pixel_values.to(self.dtype).to(
+                    self.device) if patch_pixel_values is not None else None,
+                num_patches=num_patches,
+            )
+
+        if image_embeds is not None:
+            if image_embeds.dim() == 2 or image_embeds.dim() >= 3:
+                image_embeds = image_embeds.view(-1, image_embeds.shape[-1])
+            else:
+                raise ValueError(
+                    f"Unexpected shape for image_embeds: {image_embeds.shape}")
+
+            return StepVLImageEmbeddingInputs(
+                type="image_embeds",
+                image_embeds=image_embeds.to(self.dtype).to(self.device),
+            )
+        return None
+    
+    def _process_image_features(self,
+                                image_features: torch.Tensor) -> torch.Tensor:
+        B, P = image_features.shape[:2]
+        HW = int(P ** 0.5)
+        image_features = image_features.permute(0, 2, 1).view(B, -1, HW, HW)
+        image_features = self.vision_model.vit_downsampler1(image_features)
+        image_features = self.vision_model.vit_downsampler2(image_features)
+
+        B, C, HW, HW = image_features.shape
+        image_features = image_features.view(B, -1, HW * HW).permute(0, 2, 1)
+        image_features = self.vit_large_projector(image_features)
+        return image_features
+
+    def _get_vision_model_output(self,
+                                 input_tensor: torch.Tensor) -> torch.Tensor:
+        return self.vision_model(input_tensor)
+
+    def _process_image_input(
+            self, image_input: StepVLImageInputs) -> tuple[torch.Tensor, ...]:
+
+        if image_input["type"] == "image_embeds":
+            image_features = image_input["image_embeds"]
+        else:
+            image_features = self._get_vision_model_output(
+                image_input["pixel_values"])
+            patch_image_features = self._get_vision_model_output(
+                image_input["patch_pixel_values"]
+            ) if image_input["patch_pixel_values"] is not None else None
+            num_patches = image_input["num_patches"]
+
+        image_features = self._process_image_features(image_features)
+        patch_image_features = self._process_image_features(
+            patch_image_features) if patch_image_features is not None else None
+
+        merged_image_features = []
+        cur_patch_idx = 0
+        for i, num_patch in enumerate(num_patches):
+            cur_feature = []
+            if num_patch > 0:
+                patch_slice = patch_image_features[
+                    cur_patch_idx:cur_patch_idx + num_patch]
+                cur_feature.append(patch_slice.view(-1, patch_slice.shape[-1]))
+            cur_feature.append(image_features[i].view(
+                -1, image_features.shape[-1]))
+            cur_patch_idx += num_patch
+            merged_image_features.append(
+                torch.cat(cur_feature) if len(cur_feature) >
+                1 else cur_feature[0])
+    
+        return merged_image_features
+    
+    def get_multimodal_embeddings(self, **kwargs):
+        image_input = self._parse_and_validate_image_input(**kwargs)
+        if image_input is None:
+            return None
+        vision_embeddings = self._process_image_input(image_input)
+        return vision_embeddings
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        images: Optional[list[Image.Image]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, StepVLCausalLMOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        Example:
+        ```python
+        >>> from transformers import AutoTokenizer, Llama4ForCausalLM
+        >>> model = Llama4ForCausalLM.from_pretrained("meta-llama4/Llama4-2-7b-hf")
+        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama4/Llama4-2-7b-hf")
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        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
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if inputs_embeds is None:
+            input_ids = input_ids
+            vision_embeddings = self.get_multimodal_embeddings(**kwargs)
+            inputs_embeds = self.get_input_embeddings(input_ids,
+                                                      vision_embeddings)
+            input_ids = None
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.language_model(
+            input_ids=None,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        output = StepVLCausalLMOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            attentions=outputs.attentions,
+            
+        )
+        return output if return_dict else output.to_tuple()
+
+
+
+class Step3VL10BForCausalLM(StepRoboticsPreTrainedModel, GenerationMixin):
+    _checkpoint_conversion_mapping = {
+        "^vision_model": "model.vision_model",
+        r"^model(?!\.(language_model|vision_model))": "model.language_model",
+        "^vit_large_projector": "model.vit_large_projector"
+        }
+    _tied_weights_keys = ["lm_head.weight"]
+    config: StepRoboticsConfig
+
+    def __init__(self, config: StepRoboticsConfig):
+        super().__init__(config)
+        self.model = StepRoboticsModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.text_config.vocab_size, bias=False)
+
+        self.post_init()
+    
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.get_decoder()
+    
+    @property
+    def language_model(self):
+        return self.model.language_model
+
+    @property
+    def visual(self):
+        return self.model.visual
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        num_patches = None,
+        patch_pixel_values = None,
+        patch_newline_mask = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> Union[tuple, StepVLCausalLMOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        Example:
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import AutoProcessor, LlavaForConditionalGeneration
+        >>> model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf")
+        >>> processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
+        >>> prompt = "USER: <image>\nWhat's the content of the image? ASSISTANT:"
+        >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+        >>> inputs = processor(images=image, text=prompt, return_tensors="pt")
+        >>> # Generate
+        >>> generate_ids = model.generate(**inputs, max_new_tokens=15)
+        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "USER:  \nWhat's the content of the image? ASSISTANT: The image features a busy city street with a stop sign prominently displayed"
+        ```"""
+
+        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
+        )
+
+        outputs = self.model(
+            input_ids=input_ids,
+            num_patches = num_patches,
+            patch_pixel_values = patch_pixel_values,
+            patch_newline_mask=patch_newline_mask,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        los = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+
+        return StepVLCausalLMOutputWithPast(
+            logits=logits,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        pixel_values=None,
+        attention_mask=None,
+        cache_position=None,
+        logits_to_keep=None,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+
+        if cache_position[0] == 0:
+            # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
+            # Otherwise we need pixel values to be passed to model
+            model_inputs["pixel_values"] = pixel_values
+
+        return model_inputs
+    
+    def _fix_state_dict_key_on_load(self, key: str) -> tuple[str, bool]:
+        if key.startswith("language_model."):
+            return key[len("language_model."):], True
+        
+        return key, False
+

processing_step3.py

@@ -0,0 +1,464 @@
+from transformers import BaseImageProcessor, ImageProcessingMixin
+from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs
+import math
+from typing import Iterable, Optional, Tuple, List, TypedDict, Literal, Union, overload
+
+from PIL import Image
+import torch
+import numpy as np
+import torchvision
+from torch import nn
+from torch.nn import functional as F, LayerNorm
+from torchvision.transforms.functional import InterpolationMode
+from transformers.activations import ACT2FN
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+from transformers.feature_extraction_utils import BatchFeature, TensorType
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from math import ceil
+from itertools import product
+
+
+
+MAX_IMAGE_SIZE: int = 3024
+
+class Step3VLImagePixelInputs(TypedDict):
+    type: Literal["pixel_values"]
+    pixel_values: torch.Tensor
+    patch_pixel_values: Optional[torch.Tensor]
+    num_patches: list[int]
+
+
+class Step3VLImageEmbeddingInputs(TypedDict):
+    type: Literal["image_embeds"]
+    image_embeds: torch.Tensor
+
+
+ImageWithPatches = tuple[Image.Image, list[Image.Image], list[int] | None]
+
+
+class GPUToTensor(torch.nn.Module):
+
+    def forward(self, raw_image: Union[np.ndarray,
+                                       Image.Image]) -> torch.Tensor:
+        if isinstance(raw_image, Image.Image):
+            return transforms.ToTensor()(raw_image)
+        if raw_image.ndim == 2:
+            raw_image = raw_image[:, :, None].repeat(3, -1)
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        else:
+            device = torch.device("cpu")
+        image_tensor = torch.from_numpy(raw_image).to(device)
+        image_tensor = torch.permute(image_tensor, (2, 0, 1)).contiguous()
+        if image_tensor.dtype == torch.uint8:
+            image_tensor = image_tensor.to(torch.float32).div(255)
+        return image_tensor
+
+class Step3VisionProcessor(BaseImageProcessor):
+
+    def __init__(self, size, interpolation_mode="bicubic", patch_size=None):
+        mean = [0.48145466, 0.4578275, 0.40821073]
+        std = [0.26862954, 0.26130258, 0.27577711]
+        patch_size = patch_size if patch_size is not None else size
+
+        self.transform = transforms.Compose([
+            GPUToTensor(),
+            transforms.Normalize(mean, std),
+            transforms.Resize(
+                (size, size),
+                interpolation=InterpolationMode.BICUBIC if interpolation_mode
+                == "bicubic" else InterpolationMode.BILINEAR,
+                antialias=True),
+        ])
+
+        self.patch_transform = transforms.Compose([
+            GPUToTensor(),
+            transforms.Normalize(mean, std),
+            transforms.Resize(
+                (patch_size, patch_size),
+                interpolation=InterpolationMode.BICUBIC if interpolation_mode
+                == "bicubic" else InterpolationMode.BILINEAR,
+                antialias=True),
+        ]) if patch_size is not None else None
+
+    def __call__(self, image, is_patch=False):
+        if is_patch:
+            return {"pixel_values": self.patch_transform(image).unsqueeze(0)}
+        else:
+            return {"pixel_values": self.transform(image).unsqueeze(0)}
+
+class ImagePatcher:
+    def determine_window_size(self, long: int, short: int) -> int:
+        if long <= 728:
+            return short if long / short > 1.5 else 0
+        return min(short, 504) if long / short > 4 else 504
+    def slide_window(
+        self,
+        width: int,
+        height: int,
+        sizes: list[tuple[int, int]],
+        steps: list[tuple[int, int]],
+        img_rate_thr: float = 0.6,
+    ) -> tuple[list[tuple[int, int, int, int]], tuple[int, int]]:
+        assert 1 >= img_rate_thr >= 0, "The `in_rate_thr` should lie in 0~1"
+        windows = []
+        # Sliding windows.
+        for size, step in zip(sizes, steps):
+            size_w, size_h = size
+            step_w, step_h = step
+
+            x_num = 1 if width <= size_w else ceil((width - size_w) / step_w +
+                                                   1)
+            x_start = [step_w * i for i in range(x_num)]
+            if len(x_start) > 1 and x_start[-1] + size_w > width:
+                x_start[-1] = width - size_w
+
+            y_num = 1 if height <= size_h else ceil((height - size_h) /
+                                                    step_h + 1)
+            y_start = [step_h * i for i in range(y_num)]
+            if len(y_start) > 1 and y_start[-1] + size_h > height:
+                y_start[-1] = height - size_h
+
+            start = np.array(list(product(y_start, x_start)), dtype=int)
+            start[:, [0, 1]] = start[:, [1, 0]]
+            windows.append(np.concatenate([start, start + size], axis=1))
+        windows = np.concatenate(windows, axis=0)
+
+        return [(int(box[0]), int(box[1]), int(box[2] - box[0]),
+                 int(box[3] - box[1])) for box in windows], (x_num, y_num)
+
+    def square_pad(self, img: Image.Image) -> Image.Image:
+        w, h = img.size
+        if w == h:
+            return img
+        size = max(w, h)
+        padded = Image.new(img.mode, (size, size), 0)
+        padded.paste(img, (0, 0))
+        return padded
+
+    def get_image_size_for_padding(self, img_width: int,
+                                   img_height: int) -> tuple[int, int]:
+        ratio = img_width / img_height
+        if min(img_height, img_width) < 32 and (ratio > 4 or ratio < 1 / 4):
+            new_size = max(img_height, img_width)
+            return new_size, new_size
+        return img_width, img_height
+
+    def get_image_size_for_preprocess(self, img_width: int,
+                                      img_height: int) -> tuple[int, int]:
+
+        if max(img_height, img_width) > MAX_IMAGE_SIZE:
+            scale_factor = MAX_IMAGE_SIZE / max(img_height, img_width)
+            img_width = int(img_width * scale_factor)
+            img_height = int(img_height * scale_factor)
+        return img_width, img_height
+
+    def get_image_size_for_crop(self, img_width: int, img_height: int,
+                                window_size: int):
+        w_ratio = img_width / window_size
+        h_ratio = img_height / window_size
+
+        if w_ratio < 1:
+            width_new = img_width
+        else:
+            decimal_w = w_ratio - img_width // window_size
+            w_ratio = int(w_ratio) + 1 if decimal_w > 0.2 else int(w_ratio)
+            width_new = window_size * w_ratio
+        if h_ratio < 1:
+            height_new = img_height
+        else:
+            decimal_h = h_ratio - img_height // window_size
+            h_ratio = int(h_ratio) + 1 if decimal_h > 0.2 else int(h_ratio)
+            height_new = window_size * h_ratio
+        return int(width_new), int(height_new)
+
+    def patch_crop(self, img: Image.Image, i: int, j: int, th: int, tw: int):
+        target = img.crop((j, i, j + tw, i + th))
+        return target
+
+    def get_num_patches(self, img_width: int,
+                        img_height: int) -> tuple[int, int]:
+        img_width, img_height = self.get_image_size_for_padding(
+            img_width, img_height)
+        img_width, img_height = self.get_image_size_for_preprocess(
+            img_width, img_height)
+        window_size = self.determine_window_size(max(img_height, img_width),
+                                                 min(img_height, img_width))
+        if window_size == 0:
+            return 0, 0
+        else:
+            img_width, img_height = self.get_image_size_for_crop(
+                img_width, img_height, window_size)
+            center_list, (x_num, y_num) = self.slide_window(
+                img_width, img_height, [(window_size, window_size)],
+                [(window_size, window_size)])
+            full_rows = (len(center_list) - 1) // x_num + 1
+            if len(center_list) > 0 and len(center_list) % x_num == 0:
+                full_rows -= 1
+            return len(center_list), full_rows
+
+    def __call__(
+        self, img: Image.Image
+    ) -> tuple[Image.Image, list[Image.Image], list[bool] | None]:
+        img_width, img_height = img.size
+        new_img_width, new_img_height = self.get_image_size_for_padding(
+            img_width, img_height)
+        if new_img_width != img_width or new_img_height != img_height:
+            img = self.square_pad(img)
+            img_width, img_height = img.size
+
+        new_img_width, new_img_height = self.get_image_size_for_preprocess(
+            img_width, img_height)
+        img = img.resize((new_img_width, new_img_height),
+                         Image.Resampling.BILINEAR)
+        window_size = self.determine_window_size(
+            max(new_img_height, new_img_width),
+            min(new_img_height, new_img_width))
+
+        if window_size == 0:
+            return img, [], None
+        else:
+            new_img_width, new_img_height = self.get_image_size_for_crop(
+                new_img_width, new_img_height, window_size)
+            if (new_img_width, new_img_height) != (img_width, img_height):
+                img_for_crop = img.resize((new_img_width, new_img_height),
+                                          Image.Resampling.BILINEAR)
+            else:
+                img_for_crop = img
+
+            patches = []
+            newlines = []
+            center_list, (x_num, y_num) = self.slide_window(
+                new_img_width, new_img_height, [(window_size, window_size)],
+                [(window_size, window_size)])
+            for patch_id, center_lf_point in enumerate(center_list):
+                x, y, patch_w, patch_h = center_lf_point
+                big_patch = self.patch_crop(img_for_crop, y, x, patch_h,
+                                            patch_w)
+                patches.append(big_patch)
+                if (patch_id + 1) % x_num == 0:
+                    newlines.append(patch_id)
+
+            if newlines and newlines[-1] == len(patches) - 1:
+                newlines.pop()
+
+            return img, patches, [i in newlines for i in range(len(patches))] if len(patches) > 0 else None
+
+
+
+
+class Step3VLProcessor(ProcessorMixin):
+    # Align ProcessorMixin with our custom components.
+    # We only have an image processor (not a feature extractor) plus a tokenizer.
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        tokenizer=None,
+        chat_template=None,
+        **kwargs
+    ) -> None:
+        self.image_size = 728
+        self.patch_size = 504
+
+        self.image_preprocessor = Step3VisionProcessor(self.image_size,
+                                                       "bilinear",
+                                                       self.patch_size)
+
+        self.num_image_feature_size = 169
+        self.num_patch_feature_size = 81
+        self.image_token = "<im_patch>"
+        self.image_feature_placeholder = (self.image_token *
+                                          self.num_image_feature_size)
+        self.patch_feature_placeholder = (self.image_token *
+                                          self.num_patch_feature_size)
+        super().__init__(tokenizer=tokenizer, chat_template=chat_template, **kwargs)
+        self.patcher = ImagePatcher()
+        
+    @property
+    def image_token_id(self) -> int:
+        return self.tokenizer.get_vocab()[self.image_token]
+
+    def get_num_image_tokens(self, img_width: int, img_height: int) -> int:
+        num_patches, num_newlines = self.patcher.get_num_patches(
+            img_width, img_height)
+
+        return num_patches * (
+            self.num_patch_feature_size +
+            2) + self.num_image_feature_size + 2 + num_newlines
+
+    def _split_images(self,
+                      images: list[Image.Image]) -> list[ImageWithPatches]:
+        result = []
+        for img in images:
+            result.append(self.patcher(img))
+        return result
+
+    def _convert_images_to_pixel_values(
+        self,
+        images: list[Image.Image],
+        is_patch: bool = False,
+    ) -> list[torch.Tensor]:
+        return [
+            self.image_preprocessor(img, is_patch=is_patch)["pixel_values"]
+            for img in images
+        ]
+
+    def _get_patch_repl(
+        self,
+        num_patches: int,
+        patch_newline_mask: list[bool] | None,
+    ) -> tuple[str, list[int]]:
+        text = ""
+        token_ids = []
+        for i in range(num_patches):
+            assert len(patch_newline_mask) == num_patches
+            text += f"<patch_start>{self.patch_feature_placeholder}<patch_end>"
+            token_ids.extend(
+                [self.tokenizer.convert_tokens_to_ids("<patch_start>")] +
+                [self.image_token_id] * self.num_patch_feature_size +
+                [self.tokenizer.convert_tokens_to_ids("<patch_end>")])
+            if patch_newline_mask and patch_newline_mask[i]:
+                text += "<patch_newline>"
+                token_ids.append(
+                    self.tokenizer.convert_tokens_to_ids("<patch_newline>"))
+        return text, token_ids
+
+    def _get_image_repl(
+        self,
+        num_images: int,
+    ) -> tuple[str, list[int]]:
+        text = f"<im_start>{self.image_feature_placeholder}<im_end>"
+        token_ids = [
+            self.tokenizer.convert_tokens_to_ids("<im_start>")
+        ] + [self.image_token_id] * self.num_image_feature_size + [
+            self.tokenizer.convert_tokens_to_ids("<im_end>")
+        ]
+        return text * num_images, token_ids * num_images
+
+    def _get_image_repl_features(
+        self,
+        num_images: int,
+        num_patches: int,
+        patch_new_line_idx: Optional[list[bool]],
+    ) -> tuple[str, list[int]]:
+        if num_patches > 0:
+            patch_repl, patch_repl_ids = self._get_patch_repl(
+                num_patches, patch_new_line_idx)
+        else:
+            patch_repl = ""
+            patch_repl_ids = []
+        image_repl, image_repl_ids = self._get_image_repl(num_images)
+        return patch_repl + image_repl, patch_repl_ids + image_repl_ids
+
+    def replace_placeholder(self, text: str, placeholder: str,
+                            repls: list[str]) -> str:
+        parts = text.split(placeholder)
+
+        if len(parts) - 1 != len(repls):
+            raise ValueError(
+                "The number of placeholders does not match the number of replacements."  # noqa: E501
+            )
+
+        result = [parts[0]]
+        for i, repl in enumerate(repls):
+            result.append(repl)
+            result.append(parts[i + 1])
+
+        return "".join(result)
+
+    def __call__(
+        self,
+        text: Optional[Union[str, list[str]]] = None,
+        images: ImageInput | None = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> BatchFeature:
+
+        if images is not None:
+            images = self.image_preprocessor.fetch_images(images)
+        if text is None:
+            text = []
+        if not isinstance(text, list):
+            text = [text]
+        if images is None:
+            images = []
+        elif not isinstance(images, list):
+            images = [images]
+        elif isinstance(images[0], list):
+            images = images[0]
+
+        if len(images) == 0:
+            image_inputs = {}
+            text_inputs = self.tokenizer(text)
+        else:
+            splitted_images_data = self._split_images(images)
+            pixel_values_lst = []
+            patch_pixel_values_lst = []
+            patch_newline_mask_lst = []
+            image_repl_str_lst = []
+            image_repl_ids_lst = []
+            num_patches = []
+            for raw_img, img_patches, patch_newline_mask in splitted_images_data:  # noqa: E501
+                pixel_values_lst.extend(
+                    self._convert_images_to_pixel_values([raw_img]))
+
+                if len(img_patches) > 0:
+                    patch_pixel_values_lst.extend(
+                        self._convert_images_to_pixel_values(img_patches,
+                                                             is_patch=True))
+                num_patches.append(len(img_patches))
+
+                image_repl_str, image_repl_ids = self._get_image_repl_features(
+                    1, len(img_patches), patch_newline_mask)
+                image_repl_str_lst.append(image_repl_str)
+                image_repl_ids_lst.extend(image_repl_ids)
+
+                if patch_newline_mask is not None:
+                    patch_newline_mask_lst.extend(patch_newline_mask)
+
+            image_inputs = {
+                "pixel_values": torch.cat(pixel_values_lst),
+                "num_patches": num_patches,
+            }
+            if patch_pixel_values_lst:
+                image_inputs["patch_pixel_values"] = torch.cat(
+                    patch_pixel_values_lst)
+            if patch_newline_mask_lst:
+                image_inputs["patch_newline_mask"] = torch.tensor(
+                    patch_newline_mask_lst, dtype=torch.bool)
+
+            text = [
+                self.replace_placeholder(t, self.image_token,
+                                         image_repl_str_lst) for t in text
+            ]
+            text_inputs = self.tokenizer(text)
+
+        return BatchFeature(
+            {
+                **text_inputs,
+                **image_inputs,
+            },
+            tensor_type=return_tensors,
+        )
+        
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Gemma
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Gemma
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+        
+__all__ = ["Step3VLProcessor"]

processor_config.json

@@ -0,0 +1,6 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_step3.Step3VLProcessor"
+  }
+}
+

special_tokens_map.json

@@ -0,0 +1,270 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<tool_calls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "</tool_calls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|EOT|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|BOT|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|CALL_START|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|CALL_END|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|THINK_START|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|THINK_END|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|IMG_START|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|IMG_END|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<im_patch>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<im_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<im_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<dream>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<dream_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<dream_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|MASK_1e69f|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|UNMASK_1e69f|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<video_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<video_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<patch_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<patch_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<patch_newline>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "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:439885437ceee7205faa72b808b543bb3fcd2f3df788a959aa3604d07a8f5c75
+size 11426994

tokenizer_config.json

@@ -0,0 +1,446 @@
+{
+  "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
+    },
+    "151669": {
+      "content": "<tool_calls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151670": {
+      "content": "</tool_calls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151671": {
+      "content": "<|EOT|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151672": {
+      "content": "<|BOT|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151673": {
+      "content": "<|CALL_START|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151674": {
+      "content": "<|CALL_END|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151675": {
+      "content": "<|THINK_START|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151676": {
+      "content": "<|THINK_END|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151677": {
+      "content": "<|IMG_START|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151678": {
+      "content": "<|IMG_END|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151679": {
+      "content": "<im_patch>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151680": {
+      "content": "<im_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151681": {
+      "content": "<im_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151682": {
+      "content": "<dream>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151683": {
+      "content": "<dream_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151684": {
+      "content": "<dream_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151685": {
+      "content": "<|MASK_1e69f|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151686": {
+      "content": "<|UNMASK_1e69f|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151687": {
+      "content": "<video_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151688": {
+      "content": "<video_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151689": {
+      "content": "<patch_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151690": {
+      "content": "<patch_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151691": {
+      "content": "<patch_newline>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "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|>",
+    "<tool_calls>",
+    "</tool_calls>",
+    "<|EOT|>",
+    "<|BOT|>",
+    "<|CALL_START|>",
+    "<|CALL_END|>",
+    "<|THINK_START|>",
+    "<|THINK_END|>",
+    "<|IMG_START|>",
+    "<|IMG_END|>",
+    "<im_patch>",
+    "<im_start>",
+    "<im_end>",
+    "<dream>",
+    "<dream_start>",
+    "<dream_end>",
+    "<|MASK_1e69f|>",
+    "<|UNMASK_1e69f|>",
+    "<video_start>",
+    "<video_end>",
+    "<patch_start>",
+    "<patch_end>",
+    "<patch_newline>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

vision_encoder.py

@@ -0,0 +1,451 @@
+from typing import Literal, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from transformers.activations import ACT2FN
+
+from .configuration_step_vl import StepRoboticsVisionEncoderConfig
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate last dimension halves (used by RoPE)."""
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+def apply_rotary_emb(freqs: torch.Tensor,
+                     t: torch.Tensor,
+                     start_index: int = 0,
+                     scale: float = 1.0,
+                     seq_dim: int = -2) -> torch.Tensor:
+    """Apply 2D rotary embeddings to queries / keys."""
+    dtype = t.dtype
+
+    if t.ndim == 3:
+        seq_len = t.shape[seq_dim]
+        freqs = freqs[-seq_len:]
+
+    rot_dim = freqs.shape[-1]
+    end_index = start_index + rot_dim
+    assert rot_dim <= t.shape[-1], (
+        f"feature dimension {t.shape[-1]} is too small for rot_dim {rot_dim}")
+
+    t_left, t, t_right = (
+        t[..., :start_index],
+        t[..., start_index:end_index],
+        t[..., end_index:],
+    )
+    t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
+    out = torch.cat((t_left, t, t_right), dim=-1)
+    return out.type(dtype)
+
+
+class EncoderRope2D(nn.Module):
+    """Cacheable 2D rotary positional embedding."""
+
+    def __init__(
+        self,
+        dim: int,
+        max_grid_height: int,
+        max_grid_width: int,
+        use_cls_token: bool = False,
+        theta: Union[int, float] = 10000,
+        max_freq: int = 10,
+        num_freqs: int = 1,
+        theta_rescale_factor: float = 1.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_grid_height = max_grid_height
+        self.max_grid_width = max_grid_width
+        self.use_cls_token = use_cls_token
+        self.theta = theta * theta_rescale_factor**(dim / (dim - 2))
+        self.max_freq = max_freq
+        self.num_freqs = num_freqs
+        cache = self._compute_2d_freqs()
+        self.register_buffer("freqs_cache", cache, persistent=False)
+
+    def _compute_inv_freq(self, base: Union[int, float],
+                          dim: int) -> torch.Tensor:
+
+        freqs = 1.0 / (base**(
+            torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        return freqs
+
+    def _compute_freqs(self, t: torch.Tensor, inv_freq: torch.Tensor):
+        freqs = torch.einsum("..., f -> ... f", t.type(inv_freq.dtype),
+                             inv_freq)
+        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
+        return freqs
+
+    def _compute_2d_freqs(self) -> torch.Tensor:
+        grid_h_range = torch.arange(self.max_grid_height, dtype=torch.float)
+        grid_w_range = torch.arange(self.max_grid_width, dtype=torch.float)
+        if self.use_cls_token:
+            grid_h_range += 1
+            grid_w_range += 1
+        inv_freq = self._compute_inv_freq(self.theta, self.dim // 2)
+        freqs_h = self._compute_freqs(grid_h_range, inv_freq)[:, None].expand(
+            self.max_grid_height, self.max_grid_width, -1)
+        freqs_w = self._compute_freqs(grid_w_range, inv_freq)[None, :].expand(
+            self.max_grid_height, self.max_grid_width, -1)
+        freqs = torch.cat([freqs_w, freqs_h], dim=-1).reshape(
+            self.max_grid_height * self.max_grid_width, -1)
+        if self.use_cls_token:
+            freqs = torch.cat([torch.zeros(1, freqs.shape[-1]), freqs], dim=0)
+        freqs = freqs[None, None, ...]
+        return freqs
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor,
+                grid_hw: tuple[int, int]):
+        # If grid matches cached shape we reuse directly to avoid recomputation.
+        if grid_hw[0] != self.max_grid_height or grid_hw[1] != self.max_grid_width:
+            rows = torch.arange(grid_hw[0], device=q.device).view(-1, 1)
+            cols = torch.arange(grid_hw[1], device=q.device).view(1, -1)
+            positions = (rows * self.max_grid_width + cols).reshape(-1).to(
+                torch.long)
+            if self.use_cls_token:
+                positions = torch.cat(
+                    [torch.zeros(1, device=q.device), positions + 1], dim=0)
+            freqs = self.freqs_cache.index_select(2, positions)
+        else:
+            freqs = self.freqs_cache
+        q = apply_rotary_emb(freqs, q)
+        k = apply_rotary_emb(freqs, k)
+        return q, k
+
+
+class EncoderLayerScale(nn.Module):
+    """Per-channel residual scaling used when ls_init_value is set."""
+
+    def __init__(self, dim: int, init_values: float):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.full((dim,), init_values))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:  # (B, L, D)
+        return hidden_states * self.gamma
+
+
+class EncoderMLP(nn.Module):
+    """Feed-forward network used inside each transformer block."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int,
+                 hidden_act: str):
+        super().__init__()
+        self.c_fc = nn.Linear(hidden_size, intermediate_size, bias=True)
+        self.act_fn = ACT2FN[hidden_act]
+        self.c_proj = nn.Linear(intermediate_size, hidden_size, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+
+        hidden_states = self.c_proj(self.act_fn(self.c_fc(hidden_states)))
+        return hidden_states
+
+
+class EncoderVisionAttention(nn.Module):
+    """Multi-head self attention with optional 2D RoPE."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        max_grid_height: int,
+        max_grid_width: int,
+        use_cls_token: bool = False,
+        use_rope2d: bool = True,
+        rope_theta: Union[int, float] = 10000,
+        rope_max_freq: int = 10,
+        rope_num_freqs: int = 1,
+        rope_theta_rescale_factor: float = 1.0,
+        rope_freqs_for: Literal["lang", "pixel", "constant"] = "lang",
+    ):
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(
+                f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})."
+            )
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scale = self.head_dim**-0.5
+        self.in_proj_weight = nn.Parameter(torch.zeros(hidden_size * 3, hidden_size))
+        self.in_proj_bias = nn.Parameter(torch.zeros(hidden_size * 3))
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+
+        self.rope = None
+        if use_rope2d:
+            self.rope = EncoderRope2D(
+                dim=self.head_dim,
+                max_grid_height=max_grid_height,
+                max_grid_width=max_grid_width,
+                use_cls_token=use_cls_token,
+                theta=rope_theta,
+                max_freq=rope_max_freq,
+                num_freqs=rope_num_freqs,
+                theta_rescale_factor=rope_theta_rescale_factor,
+            )
+
+    def forward(self, hidden_states: torch.Tensor, grid_hw: tuple[int, int]) -> torch.Tensor:
+        bsz, seq_len, _ = hidden_states.shape
+        qkv = F.linear(
+            hidden_states,
+            self.in_proj_weight,
+            self.in_proj_bias,
+        )
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(bsz, seq_len, self.num_heads,
+                   self.head_dim).transpose(1, 2)
+        k = k.view(bsz, seq_len, self.num_heads,
+                   self.head_dim).transpose(1, 2)
+        if self.rope is not None:
+            q, k = self.rope(q, k, grid_hw=grid_hw)
+        v = v.view(bsz, seq_len, self.num_heads,
+                   self.head_dim).transpose(1, 2)
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v, is_causal=False, scale=self.scale)
+        attn_output = attn_output.transpose(1, 2).reshape(
+            bsz, seq_len, self.num_heads * self.head_dim)
+        return self.out_proj(attn_output)
+
+
+class EncoderVisionBlock(nn.Module):
+    """A single Vision Transformer block (self-attention + MLP)."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float,
+        hidden_act: str,
+        layer_norm_eps: float,
+        ls_init_value: Optional[float] = None,
+        max_grid_height: Optional[int] = None,
+        max_grid_width: Optional[int] = None,
+        use_cls_token: bool = False,
+        use_rope2d: bool = True,
+        rope_kwargs: Optional[dict] = None,
+    ):
+        super().__init__()
+        rope_kwargs = rope_kwargs or {}
+        self.attn = EncoderVisionAttention(
+            hidden_size,
+            num_heads,
+            max_grid_height=max_grid_height,
+            max_grid_width=max_grid_width,
+            use_cls_token=use_cls_token,
+            use_rope2d=use_rope2d,
+            **rope_kwargs,
+        )
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+        intermediate = int(hidden_size * mlp_ratio)
+        self.mlp = EncoderMLP(hidden_size, intermediate, hidden_act)
+
+        self.ls_1 = EncoderLayerScale(hidden_size, ls_init_value)
+        self.ls_2 = EncoderLayerScale(hidden_size, ls_init_value)
+
+    def forward(self, hidden_states: torch.Tensor,
+                grid_hw: tuple[int, int]) -> torch.Tensor:
+        # breakpoint()
+        residual = hidden_states
+        hidden_states = self.ln_1(hidden_states)
+        hidden_states = self.attn(hidden_states, grid_hw=grid_hw)
+        hidden_states = residual + self.ls_1(hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.ln_2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + self.ls_2(hidden_states)
+        return hidden_states
+
+
+class EncoderVisionTransformer(nn.Module):
+    """Stack of encoder blocks parameterised by Step35VisionEncoderConfig."""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        depth: int,
+        num_heads: int,
+        mlp_ratio: float,
+        hidden_act: str,
+        layer_norm_eps: float,
+        ls_init_value: Optional[float] = None,
+        max_grid_height: Optional[int] = None,
+        max_grid_width: Optional[int] = None,
+        use_cls_token: bool = False,
+        use_rope2d: bool = True,
+        rope_kwargs: Optional[dict] = None,
+    ):
+        super().__init__()
+        self.layers = depth
+        rope_kwargs = rope_kwargs or {}
+        self.resblocks = nn.ModuleList([
+            EncoderVisionBlock(embed_dim, num_heads, mlp_ratio, hidden_act,
+                               layer_norm_eps,
+                               max_grid_height=max_grid_height,
+                               max_grid_width=max_grid_width,
+                               use_cls_token=use_cls_token,
+                               use_rope2d=use_rope2d,
+                               ls_init_value=ls_init_value,
+                               rope_kwargs=rope_kwargs)
+            for _ in range(depth)
+        ])
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                grid_hw: tuple[int, int]) -> torch.Tensor:
+        for block in self.resblocks:
+            hidden_states = block(hidden_states, grid_hw=grid_hw)
+        return hidden_states
+
+
+class StepRoboticsVisionEncoder(nn.Module):
+    """
+    Vision encoder built from StepRoboticsVisionEncoderConfig.
+
+    The encoder performs patch embedding followed by a stack of transformer
+    blocks. Only the config fields defined in StepRoboticsVisionEncoderConfig (and
+    StepRoboticVLConfig.vision_config) are expected.
+    """
+
+    def __init__(self, config: StepRoboticsVisionEncoderConfig):
+        super().__init__()
+        self.config = config
+
+        # Align commonly used attributes so downstream code (e.g. StepRoboticVL)
+        # can access them without extra renaming.
+        self.hidden_size = config.width
+        self.num_heads = config.heads
+        self.num_hidden_layers = config.layers
+        self.patch_size = config.patch_size
+        self.image_size = config.image_size
+        self.use_cls_token = getattr(config, "use_cls_token", False)
+        self.use_rope2d = getattr(config, "use_rope2d", True)
+        self.use_abs_posemb = getattr(config, "use_abs_posemb", True)
+        self.layer_norm_eps = config.layer_norm_eps
+        self.mlp_ratio = getattr(config, "mlp_ratio", 8960 / 1536)
+        self.ls_init_value = getattr(config, "ls_init_value", None)
+        self.hidden_act = config.hidden_act
+        self.use_ln_pre = getattr(config, "use_ln_pre", False)
+        self.use_ln_post = getattr(config, "use_ln_post", True)
+
+        # Patch embedding.
+        self.conv1 = nn.Conv2d(in_channels=config.num_channels,
+                               out_channels=self.hidden_size,
+                               kernel_size=self.patch_size,
+                               stride=self.patch_size,
+                               bias=False)
+
+        self.ln_pre = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps) if self.use_ln_pre else nn.Identity()
+        self.ln_post =  nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps) if self.use_ln_post else nn.Identity()
+
+        grid_size = self.image_size // self.patch_size
+        self.base_grid = (grid_size, grid_size)
+
+        if self.use_cls_token:
+            self.class_embedding = nn.Parameter(
+                torch.randn(self.hidden_size) * (self.hidden_size**-0.5))
+        else:
+            self.class_embedding = None
+        
+        if self.use_abs_posemb:
+            self.posemb_grid_size = self.image_size // self.patch_size
+            self.positional_embedding = nn.Parameter(
+                (self.hidden_size**-0.5) * torch.randn(
+                    int(self.use_cls_token) + self.posemb_grid_size**2,
+                    self.hidden_size,
+                ))
+
+        self.transformer = EncoderVisionTransformer(
+            embed_dim=self.hidden_size,
+            depth=self.num_hidden_layers,
+            num_heads=self.num_heads,
+            mlp_ratio=self.mlp_ratio,
+            hidden_act=self.hidden_act,
+            layer_norm_eps=self.layer_norm_eps,
+            ls_init_value=self.ls_init_value,
+            max_grid_height=self.base_grid[0],
+            max_grid_width=self.base_grid[1],
+            use_cls_token=self.use_cls_token,
+            use_rope2d=self.use_rope2d,
+            rope_kwargs={
+                "rope_theta": getattr(config, "rope_theta", 10000),
+                "rope_max_freq": getattr(config, "rope_max_freq", 10),
+                "rope_num_freqs": getattr(config, "rope_num_freqs", 1),
+                "rope_theta_rescale_factor":
+                getattr(config, "rope_theta_rescale_factor", 1.0),
+                "rope_freqs_for": getattr(config, "rope_freqs_for", "lang"),
+            },
+        )
+        self.vit_downsampler1 = nn.Conv2d(self.hidden_size,
+                                          self.hidden_size * 2,
+                                          kernel_size=3,
+                                          stride=2,
+                                          padding=1)
+        self.vit_downsampler2 = nn.Conv2d(self.hidden_size * 2,
+                                          self.hidden_size * 4,
+                                          kernel_size=3,
+                                          stride=2,
+                                          padding=1)
+
+
+    def sample_abs_posemb(self, grid_h: int, grid_w: int):
+        if self.posemb_grid_size == grid_h and self.posemb_grid_size == grid_w:
+            return self.positional_embedding[None, ...]
+
+        pos_embed = self.positional_embedding
+        if self.use_cls_token:
+            cls_token_embed, pos_embed = pos_embed[:1], pos_embed[1:]
+
+        pos_embed = (pos_embed.reshape(1, self.posemb_grid_size,
+                                       self.posemb_grid_size,
+                                       -1).permute(0, 3, 1, 2).contiguous())
+        pos_embed = F.interpolate(pos_embed,
+                                  size=(grid_h, grid_w),
+                                  mode="bilinear",
+                                  align_corners=False)
+        pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
+
+        if self.use_cls_token:
+            pos_embed = torch.cat([cls_token_embed, pos_embed], dim=0)
+
+        return pos_embed[None, ...]
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            pixel_values: Image tensor of shape (B, C, H, W).
+            layer_idx: Negative indices stop after a given block (e.g., -1 uses all blocks).
+            strip_cls_token: If True and cls token is used, remove it from output.
+        """
+        bsz, _, height, width = pixel_values.shape
+        grid_h, grid_w = height // self.patch_size, width // self.patch_size
+
+        hidden_state = self.conv1(pixel_values)  # (B, D, Gh, Gw)
+        hidden_state = hidden_state.flatten(2).transpose(1, 2)  # (B, Gh*Gw, D)
+
+        if self.use_cls_token:
+            cls_token = self.class_embedding.view(1, 1,
+                                                  -1).expand(bsz, -1, -1)
+            hidden_state = torch.cat([cls_token, hidden_state], dim=1)
+
+        if self.use_abs_posemb:
+            pos_emb = self.sample_abs_posemb(grid_h, grid_w)
+            hidden_state = hidden_state + pos_emb
+        hidden_state = self.ln_pre(hidden_state)
+        hidden_state = self.transformer(hidden_state, grid_hw=(grid_h, grid_w))
+
+        if self.use_ln_post:
+            hidden_state = self.ln_post(hidden_state)
+
+        if self.use_cls_token:
+            hidden_state = hidden_state[:, 1:, :]
+
+        return hidden_state