sgl-project--sglang/python/sglang/srt/models/unlimited_ocr.py

"""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]