Sync Step3.7 remote code and processor config

config.json

@@ -4,6 +4,7 @@
   ],
   "auto_map": {
     "AutoConfig": "configuration_step3p7.Step3p7Config",
+    "AutoProcessor": "processing_step3.Step3VLProcessor",
     "AutoModelForCausalLM": "modeling_step3p7.Step3p7ForConditionalGeneration"
   },
   "dtype": "bfloat16",

configuration_step3p6.py

@@ -1,216 +0,0 @@
-from typing import Any, Optional, Sequence, Union
-
-from transformers.configuration_utils import PretrainedConfig
-
-class StepRoboticsVisionEncoderConfig(PretrainedConfig):
-    model_type = "perception_encoder"
-
-    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_cls_token: Optional[bool] = None,
-        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
-        if use_cls_token is None:
-            use_cls_token = ues_cls_token
-        self.ues_cls_token = use_cls_token
-        self.use_cls_token = use_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 Step3p6TextConfig(PretrainedConfig):
-    model_type = "step3p5"
-    architectures = ["Step3p5ForCausalLM"]
-
-    def __init__(
-        self,
-        hidden_size: int = 4096,
-        intermediate_size: int = 11264,
-        num_attention_heads: int = 64,
-        num_attention_groups: int = 8,
-        num_hidden_layers: int = 45,
-        max_seq_len: int = 128000,
-        vocab_size: int = 128815,
-        rms_norm_eps: float = 1e-5,
-        moe_intermediate_size: int = 1280,
-        moe_num_experts: int = 288,
-        moe_top_k: int = 8,
-        rope_theta: float = 10000,
-        rope_scaling: Optional[dict[str, Any]] = None,
-        max_position_embeddings: int = 128000,
-        share_expert_dims: int = 1280,
-        share_expert_dim: Optional[int] = None,
-        head_dim: int = 128,
-        norm_expert_weight: bool = True,
-        layer_types: list[str] = None,
-        sliding_window: Optional[int] = None,
-        pad_token_id: int = 1,
-        attention_dropout: float = 0.0,
-        use_head_wise_attn_gate: bool = False,
-        use_moe_router_bias: bool = False,
-        moe_router_activation: str = "softmax",
-        moe_router_scaling_factor: float = 1.0,
-        need_fp32_gate: bool = False,
-        attention_other_setting: Optional[dict[str, Any]] = None,
-        swiglu_limits: Optional[list[Optional[float]]] = None,
-        swiglu_limits_shared: Optional[list[Optional[float]]] = None,
-        use_rope_layers: Optional[list[bool]] = None,
-        yarn_only_types: Optional[list[str]] = None,
-        moe_layers_enum: tuple[int] = (3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
-                                       15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
-                                       25, 26, 27, 28, 29, 30, 31, 32, 33, 34,
-                                       35, 36, 37, 38, 39, 40, 41, 42, 43, 44),
-        **kwargs,
-    ) -> None:
-        torch_dtype = kwargs.get("torch_dtype")
-        layer_types = _normalize_per_layer_values(layer_types,
-                                                  num_hidden_layers)
-        swiglu_limits = _normalize_per_layer_values(swiglu_limits,
-                                                    num_hidden_layers)
-        swiglu_limits_shared = _normalize_per_layer_values(
-            swiglu_limits_shared, num_hidden_layers)
-        partial_rotary_factors = kwargs.get("partial_rotary_factors")
-        kwargs["partial_rotary_factors"] = _normalize_per_layer_values(
-            partial_rotary_factors, num_hidden_layers)
-        if isinstance(rope_theta, list):
-            rope_theta = _normalize_per_layer_values(rope_theta,
-                                                     num_hidden_layers)
-        if isinstance(rope_scaling, dict):
-            rope_scaling = dict(rope_scaling)
-        if use_rope_layers:
-            use_rope_layers = _normalize_per_layer_values(
-                use_rope_layers, num_hidden_layers)
-        if share_expert_dim is None:
-            share_expert_dim = share_expert_dims
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_attention_heads = num_attention_heads
-        self.num_attention_groups = num_attention_groups
-        self.num_hidden_layers = num_hidden_layers
-        self.max_seq_len = max_seq_len
-        self.vocab_size = vocab_size
-        self.rms_norm_eps = rms_norm_eps
-        self.moe_intermediate_size = moe_intermediate_size
-        self.moe_num_experts = moe_num_experts
-        self.moe_top_k = moe_top_k
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.max_position_embeddings = max_position_embeddings
-        self.share_expert_dim = share_expert_dim
-        self.head_dim = head_dim
-        self.norm_expert_weight = norm_expert_weight
-        self.moe_layers_enum = moe_layers_enum
-        self.layer_types = layer_types
-        self.sliding_window = sliding_window
-        self.pad_token_id = pad_token_id
-        self.attention_dropout = attention_dropout
-        self.use_head_wise_attn_gate = use_head_wise_attn_gate
-        self.use_moe_router_bias = use_moe_router_bias
-        self.moe_router_activation = moe_router_activation
-        self.moe_router_scaling_factor = moe_router_scaling_factor
-        self.need_fp32_gate = need_fp32_gate
-        self.attention_other_setting = attention_other_setting
-        self.swiglu_limits = swiglu_limits
-        self.swiglu_limits_shared = swiglu_limits_shared
-        self.use_rope_layers = use_rope_layers
-        self.yarn_only_types = yarn_only_types
-        super().__init__(**kwargs)
-        if torch_dtype is not None:
-            self.torch_dtype = torch_dtype
-
-    def to_dict(self):
-        output = super().to_dict()
-        torch_dtype = getattr(self, "torch_dtype", None)
-        if torch_dtype is not None:
-            output["torch_dtype"] = torch_dtype
-        return output
-
-
-def _normalize_per_layer_values(
-    values: Optional[Sequence[Any]],
-    num_hidden_layers: int,
-) -> Optional[list[Any]]:
-    if values is None:
-        return None
-    normalized = list(values)
-    if not normalized:
-        return normalized
-    if len(normalized) < num_hidden_layers:
-        normalized.extend([normalized[-1]] *
-                          (num_hidden_layers - len(normalized)))
-    return normalized
-
-class Step3p6Config(PretrainedConfig):
-    # This loader is a compatibility shim for original Step VL checkpoints
-    # whose top-level config model_type is `step3p5v`.
-    model_type = "step3p5v"
-
-    def __init__(
-        self,
-        vision_config: Optional[Union[dict, StepRoboticsVisionEncoderConfig]] = None,
-        text_config: Optional[Union[dict, Step3p6TextConfig]] = None,
-        understand_projector_stride: int = 2,
-        projector_bias: bool = False,
-        image_token_id: int = 151679,
-        **kwargs,
-    ) -> None:
-        shared_rope_scaling = kwargs.get("rope_scaling")
-        if isinstance(shared_rope_scaling, dict):
-            shared_rope_scaling = dict(shared_rope_scaling)
-
-        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 = Step3p6TextConfig(rope_scaling=shared_rope_scaling)
-        elif isinstance(text_config, dict):
-            text_config = dict(text_config)
-            if shared_rope_scaling is not None and "rope_scaling" not in text_config:
-                text_config["rope_scaling"] = shared_rope_scaling
-            text_config = Step3p6TextConfig(**text_config)
-        elif shared_rope_scaling is not None and text_config.rope_scaling is None:
-            text_config.rope_scaling = dict(shared_rope_scaling)
-        self.text_config = text_config
-
-        rope_scaling = kwargs.get("rope_scaling")
-        if isinstance(rope_scaling, dict):
-            kwargs["rope_scaling"] = dict(rope_scaling)
-
-        self.understand_projector_stride = understand_projector_stride
-        self.projector_bias = projector_bias
-        self.hidden_size = text_config.hidden_size
-        self.max_position_embeddings = text_config.max_position_embeddings
-        self.image_token_id = image_token_id
-        # Help Auto classes find the correct implementation when saving/loading.
-        super().__init__(**kwargs)

configuration_step3p7.py

@@ -91,23 +91,10 @@ class Step3p7TextConfig(PretrainedConfig):
         **kwargs,
     ) -> None:
         torch_dtype = kwargs.get("torch_dtype")
-        layer_types = _normalize_per_layer_values(layer_types,
+        trim_layer_types = _normalize_per_layer_values(layer_types,
                                                   num_hidden_layers)
-        swiglu_limits = _normalize_per_layer_values(swiglu_limits,
-                                                    num_hidden_layers)
-        swiglu_limits_shared = _normalize_per_layer_values(
-            swiglu_limits_shared, num_hidden_layers)
-        partial_rotary_factors = kwargs.get("partial_rotary_factors")
-        kwargs["partial_rotary_factors"] = _normalize_per_layer_values(
-            partial_rotary_factors, num_hidden_layers)
-        if isinstance(rope_theta, list):
-            rope_theta = _normalize_per_layer_values(rope_theta,
-                                                     num_hidden_layers)
         if isinstance(rope_scaling, dict):
             rope_scaling = dict(rope_scaling)
-        if use_rope_layers:
-            use_rope_layers = _normalize_per_layer_values(
-                use_rope_layers, num_hidden_layers)
         if share_expert_dim is None:
             share_expert_dim = share_expert_dims
         self.hidden_size = hidden_size
@@ -128,7 +115,7 @@ class Step3p7TextConfig(PretrainedConfig):
         self.head_dim = head_dim
         self.norm_expert_weight = norm_expert_weight
         self.moe_layers_enum = moe_layers_enum
-        self.layer_types = layer_types
+        self.layer_types = trim_layer_types
         self.sliding_window = sliding_window
         self.pad_token_id = pad_token_id
         self.attention_dropout = attention_dropout
@@ -145,6 +132,7 @@ class Step3p7TextConfig(PretrainedConfig):
         super().__init__(**kwargs)
         if torch_dtype is not None:
             self.torch_dtype = torch_dtype
+        self.layer_types = layer_types
 
     def to_dict(self):
         output = super().to_dict()

modeling_step3p6.py

@@ -1,1324 +0,0 @@
-# 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.
-import copy
-import inspect
-from dataclasses import dataclass
-from typing import Callable, Optional, Tuple, Union, TypedDict, Literal
-from dataclasses import dataclass
-from PIL import Image
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache
-from transformers.generation import GenerationMixin
-from transformers.masking_utils import (create_causal_mask,
-                                        create_sliding_window_causal_mask)
-from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
-from transformers.modeling_layers import GradientCheckpointingLayer
-from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
-from transformers.modeling_rope_utils import (ROPE_INIT_FUNCTIONS,
-                                              dynamic_rope_update)
-from transformers.modeling_utils import (ALL_ATTENTION_FUNCTIONS,
-                                         PreTrainedModel)
-from transformers.processing_utils import Unpack
-from transformers.utils import TransformersKwargs, can_return_tuple, logging
-from .configuration_step3p6 import Step3p6Config, Step3p6TextConfig
-from .vision_encoder import StepRoboticsVisionEncoder
-
-
-logger = logging.get_logger(__name__)
-_MASK_INPUT_EMBEDS_ARG = ("inputs_embeds" if "inputs_embeds" in
-                          inspect.signature(create_causal_mask).parameters else
-                          "input_embeds")
-
-__all__ = [
-    "Step3p6Model",
-]
-
-
-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 Step3p6CausalLMOutputWithPast(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 Step3p6PreTrainedModel(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 = Step3p6Config
-    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
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *model_args,
-                        **kwargs):
-        key_mapping = getattr(cls, "_checkpoint_conversion_mapping", None)
-        if key_mapping is not None and kwargs.get("key_mapping") is None:
-            # Transformers only applies checkpoint renaming when key_mapping is
-            # passed explicitly; inheriting the class attribute alone is not enough.
-            kwargs["key_mapping"] = copy.deepcopy(key_mapping)
-        return super().from_pretrained(pretrained_model_name_or_path,
-                                       *model_args, **kwargs)
-
-
-
-class Step3p6RotaryEmbedding(nn.Module):
-
-    def __init__(self, config: Step3p6TextConfig, device=None, layer_idx=None):
-        super().__init__()
-        # BC: "rope_type" was originally "type"
-        self.layer_idx = layer_idx
-        self.original_rope_parameters = None
-        if config.rope_parameters is not None:
-            self.original_rope_parameters = config.rope_parameters
-            config.rope_parameters = dict(config.rope_parameters)
-            self.rope_type = config.rope_parameters.get(
-                "rope_type", config.rope_parameters.get("type"))
-        else:
-            self.rope_type = "default"
-        self.max_seq_len_cached = config.max_position_embeddings
-        self.original_max_seq_len = config.max_position_embeddings
-
-        partial_rotary_factors = getattr(config, "partial_rotary_factors",
-                                         None)
-        if partial_rotary_factors is not None:
-            config.partial_rotary_factor = partial_rotary_factors[
-                self.layer_idx]
-        else:
-            config.partial_rotary_factor = 1.0
-
-        self.rope_theta = config.rope_theta
-        if isinstance(config.rope_theta, list):
-            self.rope_theta = config.rope_theta.copy()
-            config.rope_theta = self.rope_theta[self.layer_idx]
-
-        self.config = copy.copy(config)
-        if config.rope_parameters is not None:
-            self.config.rope_parameters = dict(config.rope_parameters)
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-        inv_freq, self.attention_scaling = self.rope_init_fn(
-            self.config, device)
-
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-        config.rope_theta = self.rope_theta
-        config.rope_parameters = self.original_rope_parameters
-
-    @torch.no_grad()
-    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
-    def forward(self, x, position_ids):
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
-            position_ids.shape[0], -1, 1).to(x.device)
-        position_ids_expanded = position_ids[:, None, :].float().to(x.device)
-
-        device_type = x.device.type if isinstance(
-            x.device.type, str) and x.device.type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type,
-                            enabled=False):  # Force float32
-            freqs = (inv_freq_expanded.float()
-                     @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos() * self.attention_scaling
-            sin = emb.sin() * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., :x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2:]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    rotary_dim = cos.shape[-1]
-    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
-    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
-
-    # Apply rotary embeddings on the first half or full tensor
-    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
-    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
-
-    # Concatenate back to full shape
-    q_embed = torch.cat([q_embed, q_pass], dim=-1)
-    k_embed = torch.cat([k_embed, k_pass], dim=-1)
-    return q_embed, k_embed
-
-
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-    """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :,
-                                  None, :, :].expand(batch,
-                                                     num_key_value_heads,
-                                                     n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen,
-                                 head_dim)
-
-
-# Adapted from transformers.models.llama.modeling_llama.eager_attention_forward -> llama4 doesn't cast attn weights to fp32
-def eager_attention_forward(
-    module: nn.Module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: Optional[torch.Tensor],
-    scaling: float,
-    dropout: float = 0.0,
-    **kwargs,
-):
-    key_states = repeat_kv(key, module.num_key_value_groups)
-    value_states = repeat_kv(value, module.num_key_value_groups)
-    # breakpoint()
-    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
-    if attention_mask is not None:
-        causal_mask = attention_mask[:, :, :, :key_states.shape[-2]]
-        attn_weights = attn_weights + causal_mask
-
-    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
-    attn_weights = nn.functional.dropout(attn_weights,
-                                         p=dropout,
-                                         training=module.training)
-    attn_output = torch.matmul(attn_weights, value_states)
-    attn_output = attn_output.transpose(1, 2).contiguous()
-
-    return attn_output, attn_weights
-
-@dataclass
-class Step3p6CausalLMOutputWithPast(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
-
-
-class Step3p6MLP(nn.Module):
-
-    def __init__(self, config, intermediate_size=None, swiglu_limit=None):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size,
-                                   self.intermediate_size,
-                                   bias=False)
-        self.up_proj = nn.Linear(self.hidden_size,
-                                 self.intermediate_size,
-                                 bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size,
-                                   self.hidden_size,
-                                   bias=False)
-        self.act_fn = ACT2FN["silu"]
-        self.limit = swiglu_limit
-
-    def forward(self, x):
-        up = self.up_proj(x)
-        gate = self.act_fn(self.gate_proj(x))
-        if self.limit is not None:
-            gate = gate.clamp(min=None, max=self.limit)
-            up = up.clamp(min=-self.limit, max=self.limit)
-
-        return self.down_proj(gate * up)
-
-
-def sigmoid_routing_function(gating_output: torch.Tensor, topk: int,
-                             renormalize: bool):
-    gating_output = gating_output.float()
-    gate_prob = torch.sigmoid(gating_output)
-    gate_prob = gate_prob / gate_prob.sum(dim=-1, keepdim=True)
-    topk_prob, indices = torch.topk(gate_prob, k=topk, dim=1)
-    expert_topk_weight = topk_prob
-    if renormalize:
-        expert_topk_weight = expert_topk_weight / torch.sum(
-            expert_topk_weight, dim=-1, keepdim=True)
-    return expert_topk_weight, indices
-
-
-def softmax_routing_function(gating_output: torch.Tensor, top_k: int,
-                             renormalize: bool):
-    gating_output = gating_output.float()
-    gate_prob = torch.softmax(gating_output, dim=-1)
-    gate_prob = gate_prob / gate_prob.sum(dim=-1, keepdim=True)
-    topk_prob, indices = torch.topk(gate_prob, k=top_k, dim=1)
-    expert_topk_weight = topk_prob
-    if renormalize:
-        expert_topk_weight = expert_topk_weight / torch.sum(
-            expert_topk_weight, dim=-1, keepdim=True)
-    return expert_topk_weight, indices.to(torch.int32)
-
-
-class MoELinear(nn.Module):
-
-    def __init__(self, num_experts, in_features, out_features):
-        super().__init__()
-        self.num_experts = num_experts
-        self.in_features = in_features
-        self.out_features = out_features
-        self.weight = nn.Parameter(
-            torch.empty(num_experts, out_features, in_features))
-
-    def forward(self, x, expert_id):
-        x = F.linear(x.float(), self.weight[expert_id].float())
-        return x
-
-
-class Step3p6MoEMLP(nn.Module):
-
-    def __init__(self, config, swiglu_limit=None):
-        super().__init__()
-        self.num_experts = config.moe_num_experts
-        self.top_k = config.moe_top_k
-        self.hidden_size = config.hidden_size
-        self.moe_intermediate_size = config.moe_intermediate_size
-
-        self.use_moe_router_bias = config.use_moe_router_bias
-        if self.use_moe_router_bias:
-            self.router_bias = nn.Parameter(torch.zeros(config.moe_num_experts,
-                                                        dtype=torch.float32),
-                                            requires_grad=False)
-            self.custom_routing_function = self.router_bias_func
-        elif config.moe_router_activation == "sigmoid":
-            self.custom_routing_function = sigmoid_routing_function
-        else:
-            self.custom_routing_function = None
-        self.need_fp32_gate = config.need_fp32_gate
-        self.routed_scaling_factor = getattr(config,
-                                             "moe_router_scaling_factor", 1.0)
-        
-        # gating
-        self.gate = nn.Linear(self.hidden_size, self.num_experts, bias=False)
-            
-        self.act_fn = ACT2FN["silu"]
-        self.limit = swiglu_limit
-
-        self.up_proj = MoELinear(self.num_experts, self.hidden_size,
-                                 self.moe_intermediate_size)
-        self.gate_proj = MoELinear(self.num_experts, self.hidden_size,
-                                   self.moe_intermediate_size)
-        self.down_proj = MoELinear(self.num_experts,
-                                   self.moe_intermediate_size,
-                                   self.hidden_size)
-
-    def router_bias_func(self, gating_output: torch.Tensor, topk: int,
-                         renormalize: bool):
-        gate_prob = torch.sigmoid(gating_output.float())
-        gate_prob_with_bias = gate_prob + self.router_bias.unsqueeze(0)
-        _, indices = torch.topk(gate_prob_with_bias, k=topk, dim=1)
-        topk_prob = torch.gather(gate_prob, 1, indices)
-        expert_topk_weight = topk_prob
-        if renormalize:
-            expert_topk_weight = expert_topk_weight / (
-                torch.sum(expert_topk_weight, dim=-1, keepdim=True) + 1e-20)
-        return expert_topk_weight, indices
-
-    def get_expert_output(self, inputs: torch.Tensor, expert_id):
-        #if self.limit is None:
-        up = self.up_proj(inputs, expert_id)
-        gate = self.act_fn(self.gate_proj(inputs, expert_id))
-        if self.limit is not None:
-            gate = gate.clamp(min=None, max=self.limit)
-            up = up.clamp(min=-self.limit, max=self.limit)
-
-        return self.down_proj(gate * up, expert_id)
-
-    def forward(self, hidden_states):
-        """ """
-        batch_size, sequence_length, hidden_dim = hidden_states.shape
-        hidden_states = hidden_states.view(-1, hidden_dim)
-        if self.need_fp32_gate:
-            router_logits = torch.matmul(hidden_states.to(torch.float32), self.gate.weight.t().to(torch.float32))
-        else:
-            # router_logits: (batch * sequence_length, n_experts)
-            router_logits = self.gate(hidden_states)
-        
-        if self.custom_routing_function:
-            routing_weights, selected_experts = self.custom_routing_function(
-                router_logits, self.top_k, renormalize=True)
-        else:
-            routing_weights = F.softmax(router_logits,
-                                        dim=1,
-                                        dtype=torch.float)
-            routing_weights, selected_experts = torch.topk(routing_weights,
-                                                           self.top_k,
-                                                           dim=-1)
-
-        routing_weights = routing_weights * self.routed_scaling_factor
-
-        final_hidden_states = torch.zeros(
-            (batch_size * sequence_length, hidden_dim),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device)
-
-        # One hot encode the selected experts to create an expert mask
-        # this will be used to easily index which expert is going to be sollicitated
-        expert_mask = torch.nn.functional.one_hot(
-            selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
-
-        # Loop over all available experts in the model and perform the computation on each expert
-        for expert_idx in range(self.num_experts):
-            idx, top_x = torch.where(expert_mask[expert_idx])
-
-            # Index the correct hidden states and compute the expert hidden state for
-            # the current expert. We need to make sure to multiply the output hidden
-            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
-            current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
-            current_hidden_states = (
-                self.get_expert_output(current_state, expert_idx) *
-                routing_weights[top_x, idx, None])
-
-            # However `index_add_` only support torch tensors for indexing so we'll use
-            # the `top_x` tensor here.
-            final_hidden_states.index_add_(
-                0, top_x, current_hidden_states.to(hidden_states.dtype))
-        final_hidden_states = final_hidden_states.reshape(
-            batch_size, sequence_length, hidden_dim)
-        return final_hidden_states
-
-
-class Step3p6RMSNorm(nn.Module):
-
-    def __init__(
-        self,
-        hidden_size: int,
-        eps: float = 1e-5,
-    ) -> None:
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        dtype = x.dtype
-        x = x.float()
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-        normed = x * torch.rsqrt(variance + self.variance_epsilon)
-        normed = normed * (self.weight.float() + 1)
-        return normed.to(dtype)
-class Step3p6Attention(nn.Module):
-
-    def __init__(self, config: Step3p6TextConfig, layer_idx):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        self.num_attention_heads = config.num_attention_heads
-        self.num_key_value_heads = config.num_attention_groups
-
-        layer_types = getattr(config, "layer_types", [])
-        if layer_types:
-            enable_sliding_window = layer_types[
-                self.layer_idx] == "sliding_attention"
-        else:
-            enable_sliding_window = self.layer_idx % 2 == 0
-        
-        yarn_only_types = getattr(config, "yarn_only_types", None)
-        if yarn_only_types and layer_types[
-                self.layer_idx] not in yarn_only_types:
-            config.rope_parameters = None
-        else:
-            config.rope_parameters = getattr(config, "rope_scaling", None)
-
-        self.sliding_window = config.sliding_window
-        if enable_sliding_window:
-            self.num_attention_heads = config.attention_other_setting[
-                "num_attention_heads"]
-            self.num_key_value_heads = config.attention_other_setting[
-                "num_attention_groups"]
-
-        if self.sliding_window is not None and enable_sliding_window:
-            self.sliding_window = (self.sliding_window)
-        else:
-            self.sliding_window = None
-        self.head_dim = getattr(config, "head_dim",
-                        config.hidden_size // self.num_attention_heads)
-        self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
-
-        self.rotary_emb = Step3p6RotaryEmbedding(config, layer_idx=layer_idx)
-
-        self.q_size = self.num_attention_heads * self.head_dim
-        self.kv_size = self.num_key_value_heads * self.head_dim
-        self.scaling = self.head_dim**-0.5
-
-        self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=False)
-        self.k_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
-        self.v_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
-        self.o_proj = nn.Linear(self.q_size, config.hidden_size, bias=False)
-        self.attention_dropout = getattr(config, "attention_dropout", 0.0)
-        self.q_norm = Step3p6RMSNorm(self.head_dim,
-                                    eps=config.rms_norm_eps)
-        self.k_norm = Step3p6RMSNorm(self.head_dim,
-                                    eps=config.rms_norm_eps)
-
-        self.use_head_wise_attn_gate = config.use_head_wise_attn_gate
-        if self.use_head_wise_attn_gate:
-            self.g_proj = nn.Linear(config.hidden_size,
-                                    self.num_attention_heads,
-                                    bias=False)
-
-        self.use_rope = True
-        use_rope_layers = getattr(config, "use_rope_layers", None)
-        if use_rope_layers:
-            self.use_rope = use_rope_layers[self.layer_idx]
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor],
-        past_key_value: Optional[Cache] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
-               Optional[Tuple[torch.Tensor]]]:
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
-
-        query_states = self.q_norm(
-            self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        key_states = self.k_norm(
-            self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(
-            1, 2)
-        if self.use_head_wise_attn_gate:
-            gate_states = self.g_proj(hidden_states)
-        cos, sin = self.rotary_emb(hidden_states, position_ids)
-
-        # cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(
-            query_states, key_states, cos, sin)
-
-        # query_states, key_states = apply_rotary_pos_emb(query_norm_states, key_norm_states, cos, sin)
-        if past_key_value is not None:
-            # sin and cos are specific to RoPE models; position_ids needed for the static cache
-            cache_kwargs = {
-                "sin": sin,
-                "cos": cos,
-                "cache_position": cache_position
-            }
-            key_states, value_states = past_key_value.update(
-                key_states, value_states, self.layer_idx, cache_kwargs)
-
-        attention_interface: Callable = eager_attention_forward
-        # TODO: considering FP8；
-        # RuntimeError: Expected attn_mask dtype to be bool or float or to match query dtype,
-        # but got attn_mask.dtype: long int and  query.dtype: c10::BFloat16 instead.
-        if self.config._attn_implementation != "eager":
-            attention_interface = ALL_ATTENTION_FUNCTIONS[
-                self.config._attn_implementation]
-
-        attn_output, attn_weights = attention_interface(
-            self,
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            dropout=0.0 if not self.training else self.attention_dropout,
-            scaling=self.scaling,
-            sliding_window=self.sliding_window,  # main diff with Llama
-            **kwargs,
-        )
-        attn_output = attn_output.reshape(*input_shape, -1)
-        if self.use_head_wise_attn_gate:
-            output = attn_output.view(
-                *attn_output.shape[:-1], self.num_attention_heads,
-                self.head_dim) * gate_states.unsqueeze(-1).sigmoid()
-            attn_output = output.view(*attn_output.shape)
-        attn_output = self.o_proj(attn_output)
-
-        return attn_output, attn_weights
-
-
-class Step3p6DecoderLayer(GradientCheckpointingLayer):
-
-    def __init__(self, config, layer_idx):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.layer_idx = layer_idx
-        self.self_attn = Step3p6Attention(config, layer_idx)
-        layer_types = getattr(config, "layer_types", None) or []
-        if layer_types:
-            self.attention_type = layer_types[layer_idx]
-        else:
-            self.attention_type = (
-                "sliding_attention" if layer_idx % 2 == 0 else "full_attention"
-            )
-
-        moe_layers_enum = getattr(config, "moe_layers_enum", None)
-        if moe_layers_enum is not None:
-            if isinstance(moe_layers_enum, str):
-                moe_layers_idx = [
-                    int(i) for i in moe_layers_enum.split(',') if i.strip()
-                ]
-            else:
-                moe_layers_idx = [int(i) for i in moe_layers_enum]
-        else:
-            moe_layers_idx = [i for i in range(1, config.num_hidden_layers)]
-        self.is_moe_layer = layer_idx in moe_layers_idx
-        self.use_moe = False
-
-        if config.swiglu_limits_shared and config.swiglu_limits_shared[
-                layer_idx] is not None and config.swiglu_limits_shared[
-                    layer_idx] != 0:
-            swiglu_limit_shared = config.swiglu_limits_shared[layer_idx]
-        else:
-            swiglu_limit_shared = None
-        if config.swiglu_limits and config.swiglu_limits[
-                layer_idx] is not None and config.swiglu_limits[layer_idx] != 0:
-            swiglu_limit = config.swiglu_limits[layer_idx]
-        else:
-            swiglu_limit = None
-        if self.is_moe_layer:
-            self.moe = Step3p6MoEMLP(config, swiglu_limit=swiglu_limit)  #
-            self.share_expert = Step3p6MLP(
-                config,
-                intermediate_size=config.share_expert_dim,
-                swiglu_limit=swiglu_limit_shared)
-            self.use_moe = True
-        else:
-            self.mlp = Step3p6MLP(config,
-                                 intermediate_size=config.intermediate_size,
-                                 swiglu_limit=swiglu_limit_shared)
-
-        self.input_layernorm = Step3p6RMSNorm(
-            config.hidden_size,
-            eps=config.rms_norm_eps)
-        self.post_attention_layernorm = Step3p6RMSNorm(
-            config.hidden_size,
-            eps=config.rms_norm_eps)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[tuple[torch.Tensor]] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
-    ) -> torch.FloatTensor:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        hidden_states, _ = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_value=past_key_value,
-            cache_position=cache_position,
-            **kwargs,
-        )
-        hidden_states = residual + hidden_states
-        
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        if self.use_moe:
-            share_output = self.share_expert(hidden_states)
-            moe_output = self.moe(hidden_states)
-            ffn_output = moe_output + share_output
-        else:
-            ffn_output = self.mlp(hidden_states)
-        if isinstance(ffn_output, tuple):
-            hidden_states, _ = ffn_output
-        else:
-            hidden_states = ffn_output
-
-        hidden_states = residual + hidden_states
-        return hidden_states
-
-
-class Step3p6PreTrainedModel(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 = Step3p6TextConfig
-    supports_gradient_checkpointing = True
-    _skip_keys_device_placement = ["past_key_values"]
-    _keys_to_ignore_on_load_unexpected = [
-        r"model\.layers\.45\.*",
-        r"model\.layers\.46\.*",
-        r"model\.layers\.47\.*"
-    ]
-    _supports_flash_attn = False
-    _supports_sdpa = True
-    _supports_flex_attn = True
-    _supports_static_cache = True
-    _supports_attention_backend = True
-
-
-class Step3p6TextModel(Step3p6PreTrainedModel, GenerationMixin):
-    _no_split_modules = ["Step3p6DecoderLayer"]
-    base_model_prefix = "model"
-    _tied_weights_keys = ["lm_head.weight"]
-    config: Step3p6TextConfig
-    def __init__(self, config: Step3p6TextConfig):
-        super().__init__(config)
-        self.padding_idx = config.pad_token_id
-        self.vocab_size = config.vocab_size
-
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size,
-                                         self.padding_idx)
-        self.layers = nn.ModuleList([
-            Step3p6DecoderLayer(config, layer_idx)
-            for layer_idx in range(config.num_hidden_layers)
-        ])
-        self.norm = Step3p6RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.gradient_checkpointing = False
-        layer_types = self.config.layer_types or []
-        self.has_sliding_layers = (not layer_types or
-                                   "sliding_attention" in layer_types)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    
-    def get_input_embeddings(self, input_ids):
-        return self.embed_tokens(input_ids)
-
-    @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[Cache] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = 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, BaseModelOutputWithPast]:
-        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)
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        return_dict = return_dict if return_dict is not None else getattr(
-            self.config, "return_dict", True)
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You must specify exactly one of input_ids or inputs_embeds")
-
-        if self.gradient_checkpointing and self.training and use_cache:
-            logger.warning_once(
-                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
-            )
-            use_cache = False
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(
-                input_ids.to(self.embed_tokens.weight.device))
-
-        if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length(
-            ) if past_key_values is not None else 0
-            cache_position = torch.arange(past_seen_tokens,
-                                          past_seen_tokens +
-                                          inputs_embeds.shape[1],
-                                          device=inputs_embeds.device)
-
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        hidden_states = inputs_embeds
-
-        # It may already have been prepared by e.g. `generate`
-        if not isinstance(causal_mask_mapping := attention_mask, dict):
-            # Prepare mask arguments
-            mask_kwargs = {
-                "config": self.config,
-                "attention_mask": attention_mask,
-                "cache_position": cache_position,
-                "past_key_values": past_key_values,
-                "position_ids": position_ids,
-            }
-            mask_kwargs[_MASK_INPUT_EMBEDS_ARG] = inputs_embeds
-            # Create the masks
-            causal_mask_mapping = {
-                "full_attention": create_causal_mask(**mask_kwargs),
-            }
-
-            # The sliding window alternating layers are not always activated depending on the config
-            if self.has_sliding_layers:
-                causal_mask_mapping[
-                    "sliding_attention"] = create_sliding_window_causal_mask(
-                        **mask_kwargs)
-
-        # # create position embeddings to be shared across the decoder layers
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-        for decoder_layer in self.layers[:self.config.num_hidden_layers]:
-            if output_hidden_states:
-                all_hidden_states += (hidden_states, )
-
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=causal_mask_mapping[
-                    decoder_layer.attention_type],
-                position_ids=position_ids,
-                past_key_value=past_key_values,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                **kwargs,
-            )
-
-            hidden_states = layer_outputs
-
-        hidden_states = self.norm(hidden_states)
-
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=past_key_values if use_cache else None,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-
-
-class Step3p6Model(Step3p6PreTrainedModel, GenerationMixin):
-    config: Step3p6Config
-    _tied_weights_keys = ["lm_head.weight"]
-    base_model_prefix = ""
-
-    def __init__(self, config: Step3p6Config):
-        super().__init__(config)
-        self.vision_model = StepRoboticsVisionEncoder(config.vision_config)
-        self.language_model = Step3p6TextModel(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:
-        # breakpoint()
-        input_ids = input_ids.squeeze(0)
-        if multimodal_embeddings is None:
-            inputs_embeds = self.language_model.get_input_embeddings(input_ids)
-        else:
-            is_text = input_ids != self.config.image_token_id
-            text_ids = input_ids[is_text]
-            text_embeds = self.language_model.get_input_embeddings(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):
-        # breakpoint()
-        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, Step3p6CausalLMOutputWithPast]:
-        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 = Step3p6CausalLMOutputWithPast(
-            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 Step3p6ForConditionalGeneration(Step3p6PreTrainedModel, 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: Step3p6Config
-
-    def __init__(self, config: Step3p6Config):
-        super().__init__(config)
-        self.model = Step3p6Model(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.vision_model
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        pixel_values: Optional[torch.Tensor] = None,
-        num_patches=None,
-        patch_pixel_values=None,
-        patch_newline_mask=None,
-        image_embeds: Optional[torch.FloatTensor] = 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, Step3p6CausalLMOutputWithPast]:
-        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 Step3p6CausalLMOutputWithPast(
-            logits=logits,
-        )
-
-
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        past_key_values=None,
-        inputs_embeds=None,
-        pixel_values=None,
-        patch_pixel_values=None,
-        num_patches=None,
-        image_embeds=None,
-        attention_mask=None,
-        cache_position=None,
-        logits_to_keep=None,
-        **kwargs,
-    ):
-        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

modeling_step3p7.py

@@ -199,36 +199,40 @@ class Step3p7PreTrainedModel(PreTrainedModel):
 class Step3p7RotaryEmbedding(nn.Module):
     def __init__(self, config: Step3p7TextConfig, device=None, layer_idx=None):
         super().__init__()
-        # BC: "rope_type" was originally "type"
         self.layer_idx = layer_idx
-        self.original_rope_parameters = None
-        if config.rope_parameters is not None:
-            self.original_rope_parameters = config.rope_parameters
-            config.rope_parameters = dict(config.rope_parameters)
-            self.rope_type = config.rope_parameters.get(
-                "rope_type", config.rope_parameters.get("type")
-            )
-        else:
-            self.rope_type = "default"
         self.max_seq_len_cached = config.max_position_embeddings
         self.original_max_seq_len = config.max_position_embeddings
 
-        partial_rotary_factors = getattr(
-            config, "partial_rotary_factors", None
-        )
+        rope_theta = config.rope_theta
+        if isinstance(rope_theta, list):
+            rope_theta = rope_theta[0 if layer_idx is None else layer_idx]
+
+        partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+        partial_rotary_factors = getattr(config, "partial_rotary_factors", None)
         if partial_rotary_factors is not None:
-            config.partial_rotary_factor = partial_rotary_factors[self.layer_idx]
-        else:
-            config.partial_rotary_factor = 1.0
+            partial_rotary_factor = partial_rotary_factors[
+                0 if layer_idx is None else layer_idx
+            ]
 
-        self.rope_theta = config.rope_theta
-        if isinstance(config.rope_theta, list):
-            self.rope_theta = config.rope_theta.copy()
-            config.rope_theta = self.rope_theta[self.layer_idx]
+        self.rope_theta = rope_theta
+        self.partial_rotary_factor = partial_rotary_factor
 
         self.config = copy.copy(config)
+        self.config.rope_theta = rope_theta
+        self.config.partial_rotary_factor = partial_rotary_factor
+
         if config.rope_parameters is not None:
-            self.config.rope_parameters = dict(config.rope_parameters)
+            self.config.rope_parameters = copy.deepcopy(config.rope_parameters)
+            self.config.rope_parameters["rope_theta"] = rope_theta
+            self.config.rope_parameters["partial_rotary_factor"] = (
+                partial_rotary_factor
+            )
+            self.rope_type = self.config.rope_parameters.get(
+                "rope_type", self.config.rope_parameters.get("type")
+            )
+        else:
+            self.rope_type = "default"
+
         self.rope_init_fn = self.compute_default_rope_parameters
         if self.rope_type != "default":
             self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
@@ -238,8 +242,6 @@ class Step3p7RotaryEmbedding(nn.Module):
 
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.original_inv_freq = self.inv_freq
-        config.rope_theta = self.rope_theta
-        config.rope_parameters = self.original_rope_parameters
 
     @torch.no_grad()
     @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
@@ -288,10 +290,14 @@ class Step3p7RotaryEmbedding(nn.Module):
             post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
         """
         base = config.rope_theta
-        dim = (
+        partial_rotary_factor = getattr(
+            config, "partial_rotary_factor", 1.0
+        )
+        head_dim = (
             getattr(config, "head_dim", None)
             or config.hidden_size // config.num_attention_heads
         )
+        dim = int(head_dim * partial_rotary_factor)
 
         attention_factor = 1.0  # Unused in this type of RoPE
 
@@ -968,7 +974,6 @@ class Step3p7TextModel(Step3p7TextPreTrainedModel, GenerationMixin):
             mask_kwargs = {
                 "config": self.config,
                 "attention_mask": attention_mask,
-                "cache_position": cache_position,
                 "past_key_values": past_key_values,
                 "position_ids": position_ids,
             }
@@ -1381,7 +1386,12 @@ class Step3p7ForConditionalGeneration(Step3p7PreTrainedModel, GenerationMixin):
             **kwargs,
         )
 
-        if cache_position[0] == 0:
+        generation_cache_position = model_inputs.get("cache_position", cache_position)
+        is_prefill = past_key_values is None
+        if generation_cache_position is not None and generation_cache_position.numel() > 0:
+            is_prefill = generation_cache_position[0].item() == 0
+
+        if is_prefill:
             # During cached decoding, input ids no longer contain image tokens,
             # so pixel values should only be passed at the first step.
             model_inputs["pixel_values"] = pixel_values
@@ -1392,4 +1402,4 @@ class Step3p7ForConditionalGeneration(Step3p7PreTrainedModel, GenerationMixin):
         if key.startswith("language_model."):
             return key[len("language_model.") :], True
         
-        return key, False
+        return key, False

processing_step3.py

@@ -0,0 +1,475 @@
+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 transformers.tokenization_utils_tokenizers import TokenizersBackend
+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))
+        # return img, [], None
+        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"
+
+    @classmethod
+    def _load_tokenizer_from_pretrained(
+        cls, sub_processor_type, pretrained_model_name_or_path, subfolder="", **kwargs
+    ):
+        return TokenizersBackend.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder=subfolder,
+            **kwargs,
+        )
+
+    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"]