chore: import upstream snapshot with attribution

python/sglang/srt/models/unlimited_ocr.py

@@ -0,0 +1,438 @@
+"""Standalone UNLIMITED-OCR model (SAM + CLIP vision encoders, Deepseek backbone)."""
+
+import logging
+from typing import Iterable, List, Optional, Set, Tuple, TypeAlias, Union
+
+import torch
+from torch import Tensor, nn
+
+from sglang.srt.configs.unlimited_ocr import UnlimitedVLConfig
+from sglang.srt.layers.quantization import QuantizationConfig
+from sglang.srt.managers.mm_utils import (
+    MultiModalityDataPaddingPatternMultimodalTokens,
+    general_mm_embed_routine,
+)
+from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models.deepseek import DeepseekForCausalLM
+from sglang.srt.models.deepseek_ocr import (
+    MlpProjector,
+    build_clip_l,
+    build_sam_vit_b,
+    merge_multimodal_embeddings,
+)
+from sglang.srt.models.transformers import maybe_prefix
+from sglang.srt.utils import cpu_has_amx_support, is_cpu
+
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = is_cpu()
+
+NestedTensors: TypeAlias = Union[
+    list["NestedTensors"],
+    list["torch.Tensor"],
+    "torch.Tensor",
+    tuple["torch.Tensor", ...],
+]
+
+MultiModalEmbeddings: TypeAlias = list[Tensor] | Tensor | tuple[Tensor, ...]
+
+logger = logging.getLogger(__name__)
+
+
+class UnlimitedOCRForCausalLM(nn.Module):
+    """Standalone UNLIMITED-OCR model (SAM + CLIP ViT) with prefill-aware SWA."""
+
+    def __init__(
+        self,
+        *,
+        config: UnlimitedVLConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        """Initialize UnlimitedOCRForCausalLM with vision encoders, projector, and LM."""
+        super().__init__()
+
+        self.config = config
+        self.vision_config = config.vision_config
+        self.projector_config = config.projector_config
+        self.text_config = config.text_config
+
+        n_embed = getattr(self.projector_config, "n_embed", 1280)
+
+        self.tile_tag = config.tile_tag
+        self.global_view_pos = config.global_view_pos
+
+        embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))
+        if self.tile_tag == "2D":
+            self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)
+            self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)
+        else:
+            raise ValueError(
+                f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
+            )
+
+        self.model = DeepseekForCausalLM(
+            config=config.text_config,
+            quant_config=quant_config,
+            prefix=maybe_prefix(prefix, "language"),
+        )
+
+        self.sam_model = build_sam_vit_b()
+        self.vision_model = build_clip_l()
+
+        self.projector = MlpProjector(
+            projector_type=self.projector_config.projector_type,
+            input_dim=self.projector_config.input_dim,
+            n_embed=n_embed,
+            depth=self.projector_config.depth,
+            mlp_ratio=self.projector_config.mlp_ratio,
+            downsample_ratio=self.projector_config.downsample_ratio,
+        )
+
+        self.image_token_id = None
+
+    def get_attention_sliding_window_size(self) -> Optional[int]:
+        """Return the sliding window size from the model config, or None."""
+        return getattr(self.config, "sliding_window_size", None)
+
+    def is_prefill_aware_swa(self) -> bool:
+        """Prefill tokens are always retained in KV cache during decode."""
+        return True
+
+    def _encode_ocr1_features(self, images: torch.Tensor) -> torch.Tensor:
+        """Encode images through SAM and CLIP encoders, then project features."""
+        features_1 = self.sam_model(images)
+        features_2 = self.vision_model(images, features_1)
+        features = torch.cat(
+            (
+                features_2[:, 1:],
+                features_1.flatten(2).permute(0, 2, 1),
+            ),
+            dim=-1,
+        )
+        return self.projector(features)
+
+    def _format_ocr1_global_features(self, features: torch.Tensor) -> torch.Tensor:
+        """Reshape global features into a flat sequence with newline tokens."""
+        _, hw, n_dim = features.shape
+        h = w = int(hw**0.5)
+        features = features.view(h, w, n_dim)
+        features = torch.cat(
+            [features, self.image_newline[None, None, :].expand(h, 1, n_dim)],
+            dim=1,
+        )
+        return features.view(-1, n_dim)
+
+    def _format_ocr1_local_features(
+        self, features: torch.Tensor, crop_shape: torch.Tensor
+    ) -> torch.Tensor:
+        """Reshape local crop features into a flat sequence with newline tokens."""
+        _, hw2, n_dim2 = features.shape
+        h2 = w2 = int(hw2**0.5)
+        width_crop_num, height_crop_num = int(crop_shape[0]), int(crop_shape[1])
+        features = (
+            features.view(height_crop_num, width_crop_num, h2, w2, n_dim2)
+            .permute(0, 2, 1, 3, 4)
+            .reshape(height_crop_num * h2, width_crop_num * w2, n_dim2)
+        )
+        features = torch.cat(
+            [
+                features,
+                self.image_newline[None, None, :].expand(
+                    height_crop_num * h2, 1, n_dim2
+                ),
+            ],
+            dim=1,
+        )
+        return features.view(-1, n_dim2)
+
+    @staticmethod
+    def _collect_mm_flag(
+        items: List[MultimodalDataItem], flag_name: str
+    ) -> Optional[List[bool]]:
+        """Collect a boolean multimodal flag from all data items."""
+        values = []
+        for item in items:
+            value = getattr(item, flag_name, None)
+            if value is None:
+                return None
+            if isinstance(value, list):
+                values.extend(value)
+            else:
+                values.append(bool(value))
+        return values
+
+    def _parse_and_validate_image_input(self, **kwargs: object):
+        """Parse and validate pixel values, spatial crops, and image crops."""
+        pixel_values = kwargs.pop("pixel_values", None)
+        images_spatial_crop = kwargs.pop("images_spatial_crop", None)
+        images_crop = kwargs.pop("images_crop", None)
+        has_images = kwargs.pop("has_images", None)
+
+        if pixel_values is None:
+            return None
+        if has_images is not None:
+            if not has_images:
+                return None
+        elif torch.sum(pixel_values).item() == 0:
+            return None
+
+        if pixel_values is not None:
+            if not isinstance(pixel_values, (torch.Tensor, list)):
+                raise ValueError(
+                    "Incorrect type of pixel values. " f"Got type: {type(pixel_values)}"
+                )
+            if not isinstance(images_spatial_crop, (torch.Tensor, list)):
+                raise ValueError(
+                    "Incorrect type of image sizes. "
+                    f"Got type: {type(images_spatial_crop)}"
+                )
+            if not isinstance(images_crop, (torch.Tensor, list)):
+                raise ValueError(
+                    "Incorrect type of image crop. " f"Got type: {type(images_crop)}"
+                )
+            return [pixel_values, images_crop, images_spatial_crop]
+
+        raise AssertionError("This line should be unreachable.")
+
+    def _pixel_values_to_embedding(
+        self,
+        pixel_values: torch.Tensor,
+        images_crop: torch.Tensor,
+        images_spatial_crop: torch.Tensor,
+        has_local_crops: Optional[List[bool]] = None,
+    ) -> NestedTensors:
+        """Encode pixel values into per-image embedding sequences."""
+        images_in_this_batch = []
+
+        with torch.no_grad():
+            for jdx in range(images_spatial_crop.size(0)):
+                patches = images_crop[jdx][0].to(torch.bfloat16)
+                image_ori = pixel_values[jdx]
+                crop_shape = images_spatial_crop[jdx][0]
+                use_local_crops = (
+                    has_local_crops[jdx]
+                    if has_local_crops is not None
+                    else torch.sum(patches).item() != 0
+                )
+
+                global_features = self._encode_ocr1_features(image_ori)
+                global_features = self._format_ocr1_global_features(global_features)
+
+                if use_local_crops:
+                    local_features = self._encode_ocr1_features(patches)
+                    local_features = self._format_ocr1_local_features(
+                        local_features, crop_shape
+                    )
+                    global_local_features = torch.cat(
+                        [
+                            local_features,
+                            global_features,
+                            self.view_seperator[None, :],
+                        ],
+                        dim=0,
+                    )
+                else:
+                    global_local_features = torch.cat(
+                        [global_features, self.view_seperator[None, :]], dim=0
+                    )
+
+                images_in_this_batch.append(global_local_features)
+
+        return images_in_this_batch
+
+    def _process_image_input(self, mm_items: List[MultimodalDataItem]) -> torch.Tensor:
+        """Process multimodal data items into concatenated vision features."""
+        target_dtype = self.vision_model.dtype
+        has_local_crops = self._collect_mm_flag(mm_items, "has_local_crops")
+        pixel_values = torch.stack([item.feature for item in mm_items], dim=0).type(
+            target_dtype
+        )
+
+        images_crop = (
+            torch.stack([item.images_crop for item in mm_items], dim=0)
+            .type(target_dtype)
+            .to(device=pixel_values.device)
+        )
+        images_spatial_crop = (
+            torch.cat([item.images_spatial_crop for item in mm_items], dim=0)
+            .type(torch.long)
+            .to(device=pixel_values.device)
+        )
+        pixel_values = pixel_values.view(
+            pixel_values.shape[0] * pixel_values.shape[1], 1, *pixel_values.shape[2:]
+        )
+        images_crop = images_crop.view(
+            images_crop.shape[0] * images_crop.shape[1], 1, *images_crop.shape[2:]
+        )
+        images_spatial_crop = images_spatial_crop.view(
+            images_spatial_crop.shape[0] * images_spatial_crop.shape[1],
+            1,
+            *images_spatial_crop.shape[2:],
+        )
+
+        assert images_crop.dim() == 6
+        assert images_spatial_crop.dim() == 3
+
+        vision_feature_lists = self._pixel_values_to_embedding(
+            pixel_values=pixel_values,
+            images_crop=images_crop,
+            images_spatial_crop=images_spatial_crop,
+            has_local_crops=has_local_crops,
+        )
+        vision_features = torch.cat(vision_feature_lists, dim=0).type(target_dtype)
+        return vision_features
+
+    def get_language_model(self) -> torch.nn.Module:
+        """Return the underlying language model."""
+        return self.model
+
+    def get_multimodal_embeddings(
+        self, **kwargs: object
+    ) -> Optional[MultiModalEmbeddings]:
+        """Compute multimodal embeddings from image inputs, if present."""
+        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
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
+    ) -> torch.Tensor:
+        """Get text embeddings and merge in multimodal embeddings if provided."""
+        inputs_embeds = self.model.get_input_embeddings(input_ids)
+        if multimodal_embeddings is not None:
+            inputs_embeds = merge_multimodal_embeddings(
+                input_ids, inputs_embeds, multimodal_embeddings, self.image_token_id
+            )
+        return inputs_embeds
+
+    def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
+        """Pad input token IDs with multimodal placeholder tokens."""
+        pattern = MultiModalityDataPaddingPatternMultimodalTokens()
+        return pattern.pad_input_tokens(input_ids, mm_inputs)
+
+    def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
+        """Extract vision features from multimodal data items."""
+        vision_embeddings = self._process_image_input(items)
+        return vision_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        **kwargs: object,
+    ):
+        """Run the full multimodal forward pass (embed, encode, decode)."""
+        hidden_states = general_mm_embed_routine(
+            input_ids=input_ids,
+            forward_batch=forward_batch,
+            language_model=self.model,
+            multimodal_model=self,
+            positions=positions,
+        )
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        """Load and remap checkpoint weights into the model parameters."""
+        stacked_params_mapping = [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if name == "lm_head.weight":
+                name = "model.lm_head.weight"
+            elif name.startswith("model."):
+                if (
+                    "image_newline" in name
+                    or ".projector" in name
+                    or "vision_model" in name
+                    or "sam_model" in name
+                    or "view_seperator" in name
+                ):
+                    name = name[len("model.") :]
+                elif not (
+                    ".projector" in name
+                    or "vision_model" in name
+                    or "sam_model" in name
+                    or "image_newline" in name
+                ):
+                    name = name.replace("model.", "model.model.")
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if (
+                    "mlp.experts." in name or "mlp.shared_experts." in name
+                ) and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if (
+                    "mlp.experts." in name or "mlp.shared_experts." in name
+                ) and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        unloaded_params = params_dict.keys() - loaded_params
+        if unloaded_params:
+            raise RuntimeError(
+                f"Some weights are not initialized from checkpoints: {unloaded_params}"
+            )
+        self.post_load_weights()
+
+    def post_load_weights(self):
+        """Apply post-loading weight transformations (e.g., AMX repacking on CPU)."""
+        if _is_cpu and _is_cpu_amx_available:
+            from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
+
+            layer_ids = int(self.config.num_hidden_layers)
+            first_k_dense_replace_id = (
+                self.config.first_k_dense_replace
+                if hasattr(self.config, "first_k_dense_replace")
+                else -1
+            )
+            moe_layer_freq_id = (
+                self.config.moe_layer_freq
+                if hasattr(self.config, "moe_layer_freq")
+                else 1
+            )
+            for layer_id in range(0, layer_ids):
+                if (
+                    layer_id >= first_k_dense_replace_id
+                    and layer_id % moe_layer_freq_id == 0
+                ):
+                    if (
+                        hasattr(self.model, "model")
+                        and hasattr(self.model.model, "layers")
+                        and hasattr(self.model.model.layers[layer_id], "mlp")
+                    ):
+                        self_moe = self.model.model.layers[layer_id].mlp
+                        if hasattr(self_moe, "w1") and hasattr(self_moe, "w2"):
+                            _amx_process_weight_after_loading(self_moe, ["w1", "w2"])
+
+
+EntryClass = [UnlimitedOCRForCausalLM]