| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only Unlimited-OCR model compatible with HuggingFace weights. |
| 4 | |
| 5 | Unlimited-OCR (``baidu/Unlimited-OCR``) shares |
| 6 | the exact DeepSeek-OCR (gundam, ``base_size=1024`` / ``image_size=640`` / crop) |
| 7 | vision stack: a DeepEncoder (SAM-ViT-B + CLIP-L) followed by a linear MLP |
| 8 | projector, with the same image-token tiling layout. The only difference is the |
| 9 | language backbone, which is a DeepSeek-V2 *MoE* (64 routed + 2 shared experts, |
| 10 | ``first_k_dense_replace=1``) that uses plain multi-head attention |
| 11 | (``use_mla=False``, ``qk_nope_head_dim == qk_rope_head_dim == 0``) instead of |
| 12 | the dense MLA decoder used by DeepSeek-OCR. |
| 13 | |
| 14 | vLLM's ``DeepseekV2DecoderLayer`` already dispatches to the plain-MHA |
| 15 | ``DeepseekAttention`` whenever ``qk_nope_head_dim == qk_rope_head_dim == 0`` and |
| 16 | builds the MoE blocks straight from the config, so the whole DeepSeek-OCR |
| 17 | multimodal wrapper can be reused verbatim. Model-specific config (language |
| 18 | backbone architecture, FlexAttention for R-SWA, vision encoder backend, and |
| 19 | ``rswa_window``) is applied in ``UnlimitedOCRForCausalLMConfig``. |
| 20 | |
| 21 | Attention backend: the reference applies Reference Sliding Window Attention |
| 22 | (R-SWA) -- the prompt/image tokens form a globally-visible prefix while the |
| 23 | *generated* tokens additionally attend only a fixed sliding window (128) of |
| 24 | recent tokens. We reproduce this with backend-specific custom masks: FA4 uses a |
| 25 | ``mask_mod``, FlexAttention uses a Triton block mask, and TritonAttention uses a |
| 26 | unified-attention decode mask. FlashInfer's paged decode exposes no custom mask. |
| 27 | The window size is published via ``model_config.rswa_window``, which the model |
| 28 | runner reads to plumb per-request prefix lengths into the attention metadata. |
| 29 | |
| 30 | The *vision encoder* (DeepEncoder's CLIP stage, head_dim 64) is unaffected and |
| 31 | does not use R-SWA: it runs a single full-attention prefill pass. FlashAttention, |
| 32 | Triton and torch SDPA all produce correct, equally fast results; only FlashInfer |
| 33 | is incompatible (its ViT path asserts on the varlen cu_seqlens metadata that this |
| 34 | CLIP encoder never builds). We default the encoder to FlashAttention and |
| 35 | transparently fall back to it if FlashInfer is requested. |
| 36 | |
| 37 | To suppress repetition on long documents, use ``NGramPerReqLogitsProcessor`` from |
| 38 | this module (same request-level processor as DeepSeek-OCR) with:: |
| 39 | |
| 40 | SamplingParams( |
| 41 | temperature=0.0, |
| 42 | max_tokens=8192, |
| 43 | extra_args={"ngram_size": 35, "window_size": 128}, |
| 44 | ) |
| 45 | |
| 46 | Image processing |
| 47 | ---------------- |
| 48 | Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR), i.e. |
| 49 | ``dynamic_preprocess`` runs with ``max_num=32``. |
| 50 | |
| 51 | Multi-image requests fall back to non-crop mode: crop ("gundam") mode is only |
| 52 | used for single-image input. DeepSeek-OCR does *not* have this restriction. |
| 53 | |
| 54 | Because that fallback makes the per-image processor output depend on *how many* |
| 55 | images are in the request, it breaks the assumption behind vLLM's per-item |
| 56 | multimodal processing cache (``MultiModalProcessorOnlyCache``). We handle this |
| 57 | the same way ``DeepseekVL2MultiModalProcessor`` does: only the single-image case |
| 58 | (which always crops) is cached, while multi-image requests bypass the cache and |
| 59 | are recomputed fresh -- see ``_cached_apply_hf_processor`` below. This keeps the |
| 60 | processing cache consistent (verified by ``test_processing_correctness``). |
| 61 | """ |
| 62 | |
| 63 | import math |
| 64 | from collections.abc import Mapping, Sequence |
| 65 | |
| 66 | from vllm.config import VllmConfig |
| 67 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 68 | from vllm.multimodal.inputs import MultiModalKwargsItems |
| 69 | from vllm.multimodal.parse import ( |
| 70 | ImageEmbeddingItems, |
| 71 | ImageProcessorItems, |
| 72 | ImageSize, |
| 73 | MultiModalDataItems, |
| 74 | ) |
| 75 | from vllm.multimodal.processing import PromptReplacement, PromptUpdate |
| 76 | from vllm.multimodal.processing.context import TimingContext |
| 77 | from vllm.multimodal.processing.inputs import ProcessorInputs |
| 78 | from vllm.multimodal.processing.processor import MultiModalProcessingResult |
| 79 | from vllm.transformers_utils.processors.deepseek_ocr import ( |
| 80 | BASE_SIZE, |
| 81 | CROP_MODE, |
| 82 | IMAGE_SIZE, |
| 83 | count_tiles, |
| 84 | ) |
| 85 | |
| 86 | from .deepseek_ocr import ( |
| 87 | DeepseekOCRDummyInputsBuilder, |
| 88 | DeepseekOCRForCausalLM, |
| 89 | DeepseekOCRMultiModalProcessor, |
| 90 | DeepseekOCRProcessingInfo, |
| 91 | NGramPerReqLogitsProcessor, |
| 92 | ) |
| 93 | |
| 94 | __all__ = [ |
| 95 | "NGramPerReqLogitsProcessor", |
| 96 | "UnlimitedOCRForCausalLM", |
| 97 | ] |
| 98 | |
| 99 | # Unlimited-OCR supports up to 32 local crops (vs 6 for DeepSeek-OCR). |
| 100 | _UNLIMITED_OCR_MAX_CROPS = 32 |
| 101 | |
| 102 | |
| 103 | class UnlimitedOCRProcessingInfo(DeepseekOCRProcessingInfo): |
| 104 | """ProcessingInfo for Unlimited-OCR: same as DeepSeek-OCR but with |
| 105 | max_crops=32 instead of 6. The higher crop count allows tiling very large |
| 106 | document pages into up to 32 640×640 patches (dynamic_preprocess max_num=32). |
| 107 | """ |
| 108 | |
| 109 | def get_hf_config(self): |
| 110 | from vllm.transformers_utils.configs.unlimited_ocr import UnlimitedOCRConfig |
| 111 | |
| 112 | return self.ctx.get_hf_config(UnlimitedOCRConfig) |
| 113 | |
| 114 | def get_hf_processor(self, **kwargs: object): |
| 115 | from vllm.transformers_utils.processors.unlimited_ocr import ( |
| 116 | UnlimitedOCRProcessor, |
| 117 | ) |
| 118 | |
| 119 | v1_processor_config = dict( |
| 120 | image_size=IMAGE_SIZE, |
| 121 | base_size=BASE_SIZE, |
| 122 | crop_mode=CROP_MODE, |
| 123 | strategy="v1", |
| 124 | max_crops=_UNLIMITED_OCR_MAX_CROPS, |
| 125 | ) |
| 126 | return self.ctx.get_hf_processor( |
| 127 | UnlimitedOCRProcessor, |
| 128 | **{**v1_processor_config, **kwargs}, |
| 129 | ) |
| 130 | |
| 131 | def get_num_image_tokens( |
| 132 | self, *, image_width: int, image_height: int, cropping: bool = True |
| 133 | ) -> int: |
| 134 | patch_size = 16 |
| 135 | downsample_ratio = 4 |
| 136 | |
| 137 | # Honour the caller-supplied `cropping` flag: multi-image callers pass |
| 138 | # cropping=False to match UnlimitedOCRProcessor.tokenize_with_images. |
| 139 | if cropping: |
| 140 | if image_width <= IMAGE_SIZE and image_height <= IMAGE_SIZE: |
| 141 | crop_ratio = [1, 1] |
| 142 | else: |
| 143 | crop_ratio = count_tiles( |
| 144 | image_width, |
| 145 | image_height, |
| 146 | max_num=_UNLIMITED_OCR_MAX_CROPS, |
| 147 | image_size=IMAGE_SIZE, |
| 148 | ) |
| 149 | num_width_tiles, num_height_tiles = crop_ratio |
| 150 | else: |
| 151 | num_width_tiles = num_height_tiles = 1 |
| 152 | |
| 153 | h = w = math.ceil((BASE_SIZE // patch_size) / downsample_ratio) |
| 154 | h2 = w2 = math.ceil((IMAGE_SIZE // patch_size) / downsample_ratio) |
| 155 | |
| 156 | global_views_tokens = h * (w + 1) |
| 157 | if num_width_tiles > 1 or num_height_tiles > 1: |
| 158 | local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1) |
| 159 | else: |
| 160 | local_views_tokens = 0 |
| 161 | |
| 162 | return global_views_tokens + local_views_tokens + 1 |
| 163 | |
| 164 | def get_image_size_with_most_features(self) -> ImageSize: |
| 165 | # With max_crops=32, the widest possible grid is 4×8 (aspect ratio 1:2). |
| 166 | # A 2560×5120 image (4×640 × 8×640) selects exactly 4×8=32 tiles and |
| 167 | # produces the maximum token count. |
| 168 | return ImageSize(width=640 * 4, height=640 * 8) |
| 169 | |
| 170 | |
| 171 | class UnlimitedOCRMultiModalProcessor(DeepseekOCRMultiModalProcessor): |
| 172 | """Multimodal processor for Unlimited-OCR. |
| 173 | |
| 174 | Disables crop mode for multi-image requests (to stay consistent with |
| 175 | ``UnlimitedOCRProcessor.tokenize_with_images``), and -- since that makes the |
| 176 | per-image output depend on the request's image count -- bypasses the |
| 177 | per-item processing cache for multi-image requests, exactly like |
| 178 | ``DeepseekVL2MultiModalProcessor``. |
| 179 | |
| 180 | DeepSeek-OCR does *not* apply either of these. |
| 181 | """ |
| 182 | |
| 183 | def _get_prompt_updates( |
| 184 | self, |
| 185 | mm_items: MultiModalDataItems, |
| 186 | hf_processor_mm_kwargs: Mapping[str, object], |
| 187 | out_mm_kwargs: MultiModalKwargsItems, |
| 188 | ) -> Sequence[PromptUpdate]: |
| 189 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 190 | |
| 191 | image_token_id = hf_processor.image_token_id |
| 192 | assert isinstance(image_token_id, int) |
| 193 | |
| 194 | def get_replacement_unlimited_ocr(item_idx: int): |
| 195 | images = mm_items.get_items( |
| 196 | "image", (ImageEmbeddingItems, ImageProcessorItems) |
| 197 | ) |
| 198 | |
| 199 | if isinstance(images, ImageEmbeddingItems): |
| 200 | num_image_tokens = images.get_feature_size(item_idx) |
| 201 | else: |
| 202 | assert isinstance(images, ImageProcessorItems) |
| 203 | size = images.get_image_size(item_idx) |
| 204 | |
| 205 | # Disable crop mode for multi-image input. |
| 206 | # UnlimitedOCRProcessor.tokenize_with_images applies the same |
| 207 | # fallback, so both paths must agree on the effective crop flag. |
| 208 | effective_cropping = CROP_MODE and len(images) == 1 |
| 209 | |
| 210 | num_image_tokens = self.info.get_num_image_tokens( |
| 211 | image_width=size.width, |
| 212 | image_height=size.height, |
| 213 | cropping=effective_cropping, |
| 214 | ) |
| 215 | return [image_token_id] * num_image_tokens |
| 216 | |
| 217 | return [ |
| 218 | PromptReplacement( |
| 219 | modality="image", |
| 220 | target=[image_token_id], |
| 221 | replacement=get_replacement_unlimited_ocr, |
| 222 | ) |
| 223 | ] |
| 224 | |
| 225 | def _cached_apply_hf_processor( |
| 226 | self, |
| 227 | inputs: ProcessorInputs, |
| 228 | timing_ctx: TimingContext, |
| 229 | ) -> MultiModalProcessingResult: |
| 230 | # The processor logic differs for single-image (crop) vs multi-image |
| 231 | # (no crop) requests. The processing cache assumes per-item output is |
| 232 | # invariant of how many images are passed per prompt, so we only cache |
| 233 | # the single-image case and recompute multi-image requests fresh. |
| 234 | if inputs.mm_data_items.get_count("image", strict=False) > 1: |
| 235 | return self._apply_hf_processor(inputs, timing_ctx) |
| 236 | |
| 237 | return super()._cached_apply_hf_processor(inputs, timing_ctx) |
| 238 | |
| 239 | |
| 240 | @MULTIMODAL_REGISTRY.register_processor( |
| 241 | UnlimitedOCRMultiModalProcessor, |
| 242 | info=UnlimitedOCRProcessingInfo, |
| 243 | dummy_inputs=DeepseekOCRDummyInputsBuilder, |
| 244 | ) |
| 245 | class UnlimitedOCRForCausalLM(DeepseekOCRForCausalLM): |
| 246 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 247 | super().__init__(vllm_config=vllm_config, prefix=prefix) |