diff --git a/examples/recipes/baidu_Unlimited-OCR/cpu/cpu/feature-extraction_config.json b/examples/recipes/baidu_Unlimited-OCR/cpu/cpu/feature-extraction_config.json new file mode 100644 index 000000000..7ec8e3ba4 --- /dev/null +++ b/examples/recipes/baidu_Unlimited-OCR/cpu/cpu/feature-extraction_config.json @@ -0,0 +1,41 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "image_embeds" + } + ] + }, + "optim": {}, + "quant": null, + "loader": { + "task": "feature-extraction", + "model_type": "unlimited-ocr", + "trust_remote_code": true + } +} diff --git a/examples/recipes/baidu_Unlimited-OCR/dml/gpu/feature-extraction_config.json b/examples/recipes/baidu_Unlimited-OCR/dml/gpu/feature-extraction_config.json new file mode 100644 index 000000000..7ec8e3ba4 --- /dev/null +++ b/examples/recipes/baidu_Unlimited-OCR/dml/gpu/feature-extraction_config.json @@ -0,0 +1,41 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "image_embeds" + } + ] + }, + "optim": {}, + "quant": null, + "loader": { + "task": "feature-extraction", + "model_type": "unlimited-ocr", + "trust_remote_code": true + } +} diff --git a/examples/recipes/baidu_Unlimited-OCR/openvino/cpu/feature-extraction_config.json b/examples/recipes/baidu_Unlimited-OCR/openvino/cpu/feature-extraction_config.json new file mode 100644 index 000000000..7ec8e3ba4 --- /dev/null +++ b/examples/recipes/baidu_Unlimited-OCR/openvino/cpu/feature-extraction_config.json @@ -0,0 +1,41 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "image_embeds" + } + ] + }, + "optim": {}, + "quant": null, + "loader": { + "task": "feature-extraction", + "model_type": "unlimited-ocr", + "trust_remote_code": true + } +} diff --git a/pyproject.toml b/pyproject.toml index 4ffda3c89..773092fe1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,6 +104,15 @@ optional-dependencies.openvino = [ "openvino>=2023" ] optional-dependencies.qnn = [ "onnxruntime-qnn>=1.24.1; python_version>='3.11'", ] +# Extra deps pulled in by the baidu/Unlimited-OCR trust_remote_code modeling +# code (DeepSeek-OCR family). Not required by winml itself; install with +# `pip install winml-modelkit[unlimited-ocr]` before building that model. +optional-dependencies.unlimited-ocr = [ + "addict>=2.4", + "einops>=0.8", + "easydict>=1.13", + "matplotlib>=3.10", +] urls."Bug Tracker" = "https://github.com/microsoft/winml-cli/issues" urls.Documentation = "https://github.com/microsoft/winml-cli/blob/main/README.md" urls.Homepage = "https://github.com/microsoft/winml-cli" diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 094714494..2e1c7b5ee 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -88,6 +88,10 @@ from .t5 import T5_CONFIG from .t5 import T5DecoderIOConfig as _T5DecoderIOConfig # triggers registration from .t5 import T5EncoderIOConfig as _T5EncoderIOConfig # triggers registration +from .unlimited_ocr import MODEL_CLASS_MAPPING as _UNLIMITED_OCR_CLASS_MAPPING +from .unlimited_ocr import ( + UnlimitedOCRVisionIOConfig as _UnlimitedOCRVisionIOConfig, # triggers registration +) from .vision_encoder_decoder import MODEL_CLASS_MAPPING as _VED_CLASS_MAPPING from .vision_encoder_decoder import VISION_ENCODER_DECODER_CONFIG from .vision_encoder_decoder import ( @@ -126,6 +130,7 @@ _SEGFORMER_CLASS_MAPPING, _SIGLIP_CLASS_MAPPING, _T5_CLASS_MAPPING, + _UNLIMITED_OCR_CLASS_MAPPING, _VED_CLASS_MAPPING, _VITPOSE_CLASS_MAPPING, ) diff --git a/src/winml/modelkit/models/hf/unlimited_ocr.py b/src/winml/modelkit/models/hf/unlimited_ocr.py new file mode 100644 index 000000000..7405a93ac --- /dev/null +++ b/src/winml/modelkit/models/hf/unlimited_ocr.py @@ -0,0 +1,154 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Unlimited-OCR (DeepSeek-OCR family) HuggingFace Model Configuration. + +``baidu/Unlimited-OCR`` (model_type ``unlimited-ocr``, architecture +``UnlimitedOCRForCausalLM``) is a ``trust_remote_code`` vision-language OCR +model. Its full ``forward`` is a generative pipeline: a dual vision encoder +(SAM ViT-B + CLIP-L-14) feeds an MLP projector, whose image embeddings are +spliced into a DeepSeek-V2 (MoE + MLA) causal LM via data-dependent tiling and +``masked_scatter_`` control flow. That generative decoder half is out of scope +for ONNX export (no vendor OnnxConfig exists; Optimum ships ``deepseek_v3``, +not ``deepseek_v2``). + +Optimum has NO ``unlimited-ocr`` OnnxConfig — this module writes the +vision-tower export contract from scratch. It exposes ONLY the pure-tensor +vision sub-graph (SAM -> CLIP -> projector) under the ``feature-extraction`` +task. That sub-graph is architecture-isolated by a thin wrapper whose +``forward`` bypasses the tiling/crop/scatter control flow, exporting just the +image-embedding computation the downstream LM consumes. + +This is an Effort-L1 contribution per the `adding-model-support` skill: a +from-scratch OnnxConfig plus a sub-graph wrapper, no changes to the export +engine itself. The generative decoder remains unexported by design. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from optimum.exporters.onnx import OnnxConfig +from optimum.utils import NormalizedConfig +from optimum.utils.input_generators import DummyVisionInputGenerator +from torch import nn +from transformers import AutoModel + +from ...export import register_onnx_overwrite + + +# Vision tower operates on fixed 1024x1024 RGB tiles (SAM ViT-B input geometry). +_VISION_NUM_CHANNELS = 3 +_VISION_IMAGE_SIZE = 1024 + + +# ============================================================================= +# Vision-tower sub-graph wrapper +# ============================================================================= + + +class UnlimitedOCRVisionTowerWrapper(nn.Module): + """Export-only wrapper exposing the SAM -> CLIP -> projector sub-graph. + + The full ``UnlimitedOCRModel.forward`` performs data-dependent tiling and + ``masked_scatter_`` splicing that is not ONNX-traceable. This wrapper + isolates the deterministic vision path that produces image embeddings: + + sam_feat = sam_model(pixel_values) + clip_feat = vision_model(pixel_values, sam_feat) + fused = cat(clip_feat[:, 1:], sam_feat.flatten(2).permute(0, 2, 1)) + return projector(fused) + """ + + def __init__(self, base: nn.Module) -> None: + super().__init__() + self.sam_model = base.sam_model + self.vision_model = base.vision_model + self.projector = base.projector + + @classmethod + def from_pretrained( + cls, model_name_or_path: str, **kwargs: Any + ) -> UnlimitedOCRVisionTowerWrapper: + """Load the full model via ``AutoModel``, then wrap its vision tower. + + ``trust_remote_code`` is forwarded by the loader through ``kwargs``; + the model's ``auto_map`` registers ``UnlimitedOCRForCausalLM`` under + ``AutoModel``. ``get_model()`` returns the base module owning the + ``sam_model`` / ``vision_model`` / ``projector`` attributes. + """ + full_model = AutoModel.from_pretrained(model_name_or_path, **kwargs) + base = full_model.get_model() if hasattr(full_model, "get_model") else full_model.model + wrapper = cls(base) + wrapper.eval() + return wrapper + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Return image embeddings ``[batch, 256, 1280]`` for the LM to consume.""" + sam_feat = self.sam_model(pixel_values) + clip_feat = self.vision_model(pixel_values, sam_feat) + fused = torch.cat( + (clip_feat[:, 1:], sam_feat.flatten(2).permute(0, 2, 1)), dim=-1 + ) + return self.projector(fused) + + +# ============================================================================= +# ONNX export config +# ============================================================================= + + +@register_onnx_overwrite("unlimited-ocr", "feature-extraction", library_name="transformers") +class UnlimitedOCRVisionIOConfig(OnnxConfig): + """From-scratch ONNX config for the Unlimited-OCR vision tower. + + Input geometry is pinned to 1024x1024 (the SAM encoder's fixed input + size); only the batch dimension is dynamic. The single output is the + projected image-embedding sequence consumed by the (unexported) LM. + """ + + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(allow_new=True) + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator,) + DEFAULT_ONNX_OPSET = 17 + + @property + def inputs(self) -> dict[str, dict[int, str]]: + """Single ``pixel_values`` input; only batch is dynamic (H,W pinned).""" + return { + "pixel_values": {0: "batch_size"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + """Projected image-embedding sequence consumed by the LM.""" + return { + "image_embeds": {0: "batch_size"}, + } + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs: Any): # type: ignore[override] + """Emit a fixed ``[1, 3, 1024, 1024]`` ``pixel_values`` tensor. + + The vision tower has no data-dependent control flow, so a zero tensor + of the correct geometry traces the full sub-graph. Geometry is pinned + here rather than derived from the nested ``vision_config`` to keep the + export contract explicit and architecture-agnostic. + """ + return { + "pixel_values": torch.zeros( + 1, _VISION_NUM_CHANNELS, _VISION_IMAGE_SIZE, _VISION_IMAGE_SIZE + ) + } + + +# ============================================================================= +# Model Class Mapping +# ============================================================================= + +# (model_type, task) -> export wrapper. Binds the ``feature-extraction`` task +# on ``unlimited-ocr`` to the vision-tower sub-graph wrapper so the loader +# exports image embeddings instead of attempting the generative full model. +MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = { + ("unlimited-ocr", "feature-extraction"): UnlimitedOCRVisionTowerWrapper, +} diff --git a/tests/unit/models/unlimited_ocr/test_onnx_config.py b/tests/unit/models/unlimited_ocr/test_onnx_config.py new file mode 100644 index 000000000..015ff4713 --- /dev/null +++ b/tests/unit/models/unlimited_ocr/test_onnx_config.py @@ -0,0 +1,89 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for the Unlimited-OCR (DeepSeek-OCR family) vision-tower ONNX config. + +Unlimited-OCR is a ``trust_remote_code`` VLM whose full ``forward`` is a +generative pipeline (out of scope for ONNX export). This module registers only +the pure-tensor vision sub-graph (SAM -> CLIP -> projector) under the +``feature-extraction`` task. These tests are network-free: they validate the +registry wiring and the pinned 1024x1024 dummy-input contract without loading +the 6.67 GB checkpoint. +""" + +from __future__ import annotations + +from transformers import PretrainedConfig + +# Import triggers ONNX config registration +import winml.modelkit.models # noqa: F401 +from winml.modelkit.export import generate_dummy_inputs +from winml.modelkit.models.hf.unlimited_ocr import ( + MODEL_CLASS_MAPPING, + UnlimitedOCRVisionIOConfig, + UnlimitedOCRVisionTowerWrapper, +) + + +# ============================================================================= +# Test Constants +# ============================================================================= + +VISION_NUM_CHANNELS = 3 +VISION_IMAGE_SIZE = 1024 + + +# ============================================================================= +# MODEL_CLASS_MAPPING — routing to the vision-tower wrapper +# ============================================================================= + + +class TestUnlimitedOCRModelClassMapping: + """The feature-extraction task routes to the vision-tower wrapper.""" + + def test_registered_in_local_mapping(self): + """Local sub-dict binds the wrapper to feature-extraction.""" + key = ("unlimited-ocr", "feature-extraction") + assert key in MODEL_CLASS_MAPPING + assert MODEL_CLASS_MAPPING[key] is UnlimitedOCRVisionTowerWrapper + + def test_aggregated_into_hf_mapping(self): + """Entry is merged into the package-level aggregated mapping.""" + from winml.modelkit.models import HF_MODEL_CLASS_MAPPING + + key = ("unlimited-ocr", "feature-extraction") + assert key in HF_MODEL_CLASS_MAPPING + assert HF_MODEL_CLASS_MAPPING[key] is UnlimitedOCRVisionTowerWrapper + + def test_wrapper_exposes_from_pretrained(self): + """Loader dispatches via a from_pretrained classmethod.""" + assert hasattr(UnlimitedOCRVisionTowerWrapper, "from_pretrained") + + +# ============================================================================= +# OnnxConfig — pinned 1024x1024 dummy input + IO contract +# ============================================================================= + + +class TestUnlimitedOCRVisionIOConfig: + """The from-scratch OnnxConfig pins geometry and declares a single output.""" + + def test_input_contract(self): + config = UnlimitedOCRVisionIOConfig(PretrainedConfig(), task="feature-extraction") + assert list(config.inputs.keys()) == ["pixel_values"] + assert config.inputs["pixel_values"] == {0: "batch_size"} + + def test_output_contract(self): + config = UnlimitedOCRVisionIOConfig(PretrainedConfig(), task="feature-extraction") + assert list(config.outputs.keys()) == ["image_embeds"] + assert config.outputs["image_embeds"] == {0: "batch_size"} + + def test_dummy_inputs_pinned_geometry(self): + """generate_dummy_inputs emits a fixed [1, 3, 1024, 1024] tensor.""" + inputs = generate_dummy_inputs( + "unlimited-ocr", "feature-extraction", PretrainedConfig() + ) + assert set(inputs.keys()) == {"pixel_values"} + pv = inputs["pixel_values"] + assert tuple(pv.shape) == (1, VISION_NUM_CHANNELS, VISION_IMAGE_SIZE, VISION_IMAGE_SIZE)