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

# Copyright 2025 SGLang Team
# 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.
# ==============================================================================

"""
Support for lightonai/LightOnOCR-2-1B.

LightOnOCR is a vision-language OCR model that combines:
- Pixtral vision encoder (24 layers, 1024 hidden dim)
- Spatial merge projection with RMSNorm + PatchMerger (2x2 = 4x token reduction)
- Qwen3 language decoder (28 layers, 1024 hidden dim)

Key differences from PixtralForConditionalGeneration:
- Uses Qwen3ForCausalLM instead of MistralLarge3ForCausalLM as the language model
- Has an RMSNorm applied to vision encoder output before patch merging
- Does not use image break/end tokens (single contiguous image token range)
- HuggingFace checkpoint uses a vision_projection namespace for norm, patch_merger,
  and adapter weights

References:
- https://huggingface.co/lightonai/LightOnOCR-2-1B
"""

from dataclasses import fields
from typing import Iterable, List, Tuple

import torch
import torch.nn as nn

from sglang.srt.layers.layernorm import RMSNorm
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.pixtral import (
    PATCH_MERGE,
    PatchMerger,
    PixtralHFVisionModel,
    VisionEncoderArgs,
    VisionLanguageAdapter,
)
from sglang.srt.models.qwen3 import Qwen3ForCausalLM


class LightOnOCRForConditionalGeneration(nn.Module):
    """
    LightOnOCR model for SGLang inference.

    Architecture:
    - Pixtral-based vision encoder (PixtralHFVisionModel, 24 layers)
    - RMSNorm on vision encoder output
    - Spatial merge via PatchMerger (2x2 = 4x token reduction)
    - VisionLanguageAdapter projection to text hidden size
    - Qwen3-based decoder (28 layers) with QK norms
    """

    merge_by_field_config = True

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return None
        raise ValueError("Only image modality is supported")

    def __init__(self, *, config, prefix: str = "", **kwargs):
        super().__init__()
        self.config = config
        quant_config = kwargs.get("quant_config")

        # Build VisionEncoderArgs from config
        vision_config = config.vision_config
        config_dict = vision_config.to_dict()
        if config_dict.get("rope_parameters"):
            config_dict["rope_theta"] = config_dict["rope_parameters"].get("rope_theta")
        dataclass_fields = {field.name for field in fields(VisionEncoderArgs)}
        vision_args = {
            key: value for key, value in config_dict.items() if key in dataclass_fields
        }
        # LightOnOCR stores these at the top-level config
        if "image_token_id" not in vision_args:
            vision_args["image_token_id"] = getattr(config, "image_token_id", 151655)
        if "spatial_merge_size" not in vision_args:
            vision_args["spatial_merge_size"] = getattr(config, "spatial_merge_size", 2)
        if "adapter_bias" not in vision_args:
            vision_args["adapter_bias"] = getattr(
                config, "multimodal_projector_bias", True
            )
        # LightOnOCR uses patch merging for spatial merge
        vision_args["mm_projector_id"] = PATCH_MERGE
        self.vision_args = VisionEncoderArgs(**vision_args)

        # Vision encoder (Pixtral HF variant with SGLang parallel layers)
        self.vision_encoder = PixtralHFVisionModel(vision_config, quant_config=None)

        # RMSNorm applied to vision encoder output before patch merging
        self.vision_projection_norm = RMSNorm(self.vision_args.hidden_size, eps=1e-5)

        # Patch merger for spatial token reduction
        self.patch_merger = PatchMerger(
            vision_encoder_dim=self.vision_args.hidden_size,
            spatial_merge_size=self.vision_args.spatial_merge_size,
        )

        # Vision-to-language projection adapter
        self.vision_language_adapter = VisionLanguageAdapter(
            self.vision_args, dim=config.text_config.hidden_size
        )

        # Language model
        self.language_model = Qwen3ForCausalLM(
            config=config.text_config,
            quant_config=quant_config,
        )

    def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
        pattern = MultiModalityDataPaddingPatternMultimodalTokens()
        return pattern.pad_input_tokens(input_ids, mm_inputs)

    def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
        """Process images through vision encoder and projection pipeline."""
        images = [item.feature for item in items]

        # Extract image sizes from model-specific data or infer from tensor shape
        image_sizes_list = []
        for item in items:
            if item.model_specific_data and "image_sizes" in item.model_specific_data:
                sizes_tensor = item.model_specific_data["image_sizes"]
                for size in sizes_tensor:
                    image_sizes_list.append((int(size[0]), int(size[1])))
            else:
                img = item.feature
                for _ in range(img.shape[0]):
                    image_sizes_list.append((img.shape[-2], img.shape[-1]))

        # Stack pixel values
        if len(images) > 1:
            pixel_values = torch.cat(images, dim=0)
        else:
            pixel_values = images[0]

        # Vision encoder forward
        image_features = self.vision_encoder(pixel_values, image_sizes=image_sizes_list)
        image_features = image_features.view(-1, image_features.shape[-1])

        # Norm before patch merge (matches HF Mistral3MultiModalProjector order)
        image_features = self.vision_projection_norm(image_features)

        # Spatial merge via patch merger — use actual image sizes (not padded tensor
        # shape) because PixtralHFVisionModel crops embeddings to real dimensions.
        patch_size = self.vision_args.patch_size
        img_patch_dims = [
            (h // patch_size, w // patch_size) for (h, w) in image_sizes_list
        ]
        image_features = self.patch_merger(image_features, image_sizes=img_patch_dims)

        # Project to language model dimension
        image_embeds = self.vision_language_adapter(image_features)
        return image_embeds

    def get_language_model(self) -> torch.nn.Module:
        return self.language_model

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
    ):
        return general_mm_embed_routine(
            input_ids=input_ids,
            forward_batch=forward_batch,
            language_model=self.language_model,
            multimodal_model=self,
            positions=positions,
        )

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        return self.language_model.compute_logits(hidden_states)

    def get_embed_and_head(self):
        return self.language_model.get_embed_and_head()

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
        """Load weights from HuggingFace checkpoint.

        HF checkpoint weight layout (after stripping ``model.`` prefix):
        - ``vision_encoder.*`` -> self.vision_encoder
        - ``vision_projection.norm.*`` -> self.vision_projection_norm
        - ``vision_projection.patch_merger.*`` -> self.patch_merger
        - ``vision_projection.linear_1.*`` -> self.vision_language_adapter.w_in
        - ``vision_projection.linear_2.*`` -> self.vision_language_adapter.w_out
        - ``language_model.*`` -> self.language_model (Qwen3ForCausalLM)
        """
        vision_encoder_dict = dict(self.vision_encoder.named_parameters())
        patch_merger_dict = dict(self.patch_merger.named_parameters())
        norm_dict = dict(self.vision_projection_norm.named_parameters())
        adapter_dict = dict(self.vision_language_adapter.named_parameters())

        # PixtralHFVisionModel uses SGLang parallel layers with stacked params
        stacked_params_mapping = [
            (".attention.qkv_proj", ".attention.q_proj", "q"),
            (".attention.qkv_proj", ".attention.k_proj", "k"),
            (".attention.qkv_proj", ".attention.v_proj", "v"),
            (".feed_forward.gate_up_proj", ".feed_forward.gate_proj", 0),
            (".feed_forward.gate_up_proj", ".feed_forward.up_proj", 1),
        ]

        def llm_weights_generator():
            for name, w in weights:
                # HF checkpoint prefixes all weights with model.
                if name.startswith("model."):
                    name = name[len("model.") :]

                if name.startswith("vision_encoder."):
                    trimmed = name[len("vision_encoder.") :]

                    # Handle stacked params (QKV, gate/up)
                    loaded = False
                    for param_name, weight_name, shard_id in stacked_params_mapping:
                        if weight_name in trimmed:
                            transformed = trimmed.replace(weight_name, param_name)
                            if transformed in vision_encoder_dict:
                                param = vision_encoder_dict[transformed]
                                weight_loader = getattr(
                                    param, "weight_loader", default_weight_loader
                                )
                                with torch.no_grad():
                                    weight_loader(param, w, shard_id)
                                loaded = True
                                break

                    if not loaded:
                        # Handle o_proj -> proj rename
                        if ".attention.o_proj" in trimmed:
                            trimmed = trimmed.replace(
                                ".attention.o_proj", ".attention.proj"
                            )
                        if trimmed in vision_encoder_dict:
                            param = vision_encoder_dict[trimmed]
                            weight_loader = getattr(
                                param, "weight_loader", default_weight_loader
                            )
                            with torch.no_grad():
                                weight_loader(param, w)

                elif name.startswith("vision_projection."):
                    remaining = name[len("vision_projection.") :]

                    if remaining.startswith("patch_merger."):
                        trimmed = remaining[len("patch_merger.") :]
                        if trimmed in patch_merger_dict:
                            param = patch_merger_dict[trimmed]
                            with torch.no_grad():
                                default_weight_loader(param, w)

                    elif remaining.startswith("norm."):
                        trimmed = remaining[len("norm.") :]
                        if trimmed in norm_dict:
                            param = norm_dict[trimmed]
                            with torch.no_grad():
                                default_weight_loader(param, w)

                    else:
                        # linear_1 -> w_in, linear_2 -> w_out
                        trimmed = remaining.replace("linear_1.", "w_in.").replace(
                            "linear_2.", "w_out."
                        )
                        if trimmed in adapter_dict:
                            param = adapter_dict[trimmed]
                            with torch.no_grad():
                                default_weight_loader(param, w)

                else:
                    # Language model weights and any other weights
                    if name.startswith("language_model."):
                        # Qwen3ForCausalLM expects model.* prefix
                        name = "model." + name[len("language_model.") :]
                    yield (name, w)

        self.language_model.load_weights(llm_weights_generator())


EntryClass = LightOnOCRForConditionalGeneration