Add MGP-STR (alibaba-damo/mgp-str-base) image-to-text task support

examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp16_config.json

@@ -0,0 +1,53 @@
+{
+  "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, 32, 128],
+        "value_range": [0, 1]
+      }
+    ],
+    "output_tensors": [
+      {"name": "char_logits"},
+      {"name": "bpe_logits"},
+      {"name": "wp_logits"}
+    ]
+  },
+  "optim": {},
+  "quant": {
+    "mode": "fp16",
+    "samples": 10,
+    "calibration_method": "minmax",
+    "weight_type": "uint8",
+    "activation_type": "uint8",
+    "per_channel": false,
+    "symmetric": false,
+    "weight_symmetric": null,
+    "activation_symmetric": null,
+    "save_calibration": false,
+    "distribution": "uniform",
+    "seed": null,
+    "calibration_load_path": null,
+    "calibration_save_path": null,
+    "op_types_to_quantize": null,
+    "nodes_to_exclude": null,
+    "fp16_keep_io_types": true,
+    "fp16_op_block_list": null
+  },
+  "compile": null,
+  "loader": {
+    "task": "image-to-text",
+    "model_class": "MgpstrForSceneTextRecognition",
+    "model_type": "mgp-str"
+  }
+}

examples/recipes/alibaba-damo_mgp-str-base/cpu/cpu/image-to-text_fp32_config.json

@@ -0,0 +1,34 @@
+{
+  "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, 32, 128],
+        "value_range": [0, 1]
+      }
+    ],
+    "output_tensors": [
+      {"name": "char_logits"},
+      {"name": "bpe_logits"},
+      {"name": "wp_logits"}
+    ]
+  },
+  "optim": {},
+  "quant": null,
+  "compile": null,
+  "loader": {
+    "task": "image-to-text",
+    "model_class": "MgpstrForSceneTextRecognition",
+    "model_type": "mgp-str"
+  }
+}

src/winml/modelkit/eval/image_to_text_evaluator.py

@@ -29,6 +29,7 @@
 
 if TYPE_CHECKING:
     from datasets import Dataset
+    from transformers.pipelines.base import Pipeline
 
     from ..models.winml.base import WinMLPreTrainedModel
     from .config import DatasetConfig, WinMLEvaluationConfig
@@ -56,6 +57,17 @@ def align_labels(self, dataset: Dataset, ds_config: DatasetConfig) -> Dataset:
         """No-op: free-text labels need no ClassLabel alignment."""
         return dataset
 
+    def prepare_pipeline(self) -> Pipeline:
+        """Create an image-to-text pipeline, including registered specializations."""
+        from typing import cast
+
+        from ..inference.pipeline import create_pipeline
+
+        return cast(
+            "Pipeline",
+            create_pipeline("image-to-text", self.model, self.config.model_id),
+        )
+
     def compute(self) -> dict[str, Any]:
         """Run the pipeline over each sample and return CER + CIDEr."""
         from tqdm.auto import tqdm

src/winml/modelkit/inference/pipeline.py

@@ -20,7 +20,7 @@
 
 import inspect
 import logging
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 
 if TYPE_CHECKING:
@@ -37,6 +37,10 @@
     "sentence-similarity": "feature-extraction",
 }
 
+_SPECIALIZED_PIPELINE_REGISTRY: dict[tuple[str, str], str] = {
+    ("mgp-str", "image-to-text"): "MgpstrImageToTextPipeline",
+}
+
 
 def create_pipeline(
     task: str,
@@ -57,6 +61,23 @@ def create_pipeline(
     Returns:
         A configured ``transformers.Pipeline`` ready for inference.
     """
+    model_type = getattr(getattr(model, "config", None), "model_type", None)
+    if not isinstance(model_type, str):
+        model_type = ""
+    specialized = _SPECIALIZED_PIPELINE_REGISTRY.get((model_type, task))
+    if specialized is not None:
+        if model_id is None:
+            raise ValueError(f"{specialized} requires a model ID to load its processor.")
+        from ..models.winml.image_to_text import MgpstrImageToTextPipeline
+
+        specialized_classes = {"MgpstrImageToTextPipeline": MgpstrImageToTextPipeline}
+        pipe = specialized_classes[specialized](
+            cast("WinMLPreTrainedModel", model), model_id
+        )
+        _adapt_image_processor_size(pipe, task, model)
+        logger.info("Created specialized pipeline: task=%s model=%s", task, model_id)
+        return pipe
+
     from transformers import pipeline
 
     kwargs: dict[str, Any] = {

src/winml/modelkit/models/hf/__init__.py

@@ -52,6 +52,10 @@
 from .marian import MODEL_CLASS_MAPPING as _MARIAN_CLASS_MAPPING
 from .marian import MarianDecoderIOConfig as _MarianDecoderIOConfig  # triggers registration
 from .marian import MarianEncoderIOConfig as _MarianEncoderIOConfig  # triggers registration
+from .mgp_str import MODEL_CLASS_MAPPING as _MGPSTR_CLASS_MAPPING
+from .mgp_str import (
+    MgpstrImage2TextOnnxConfig as _MgpstrImage2TextOnnxConfig,  # triggers registration
+)
 from .mu2 import MODEL_CLASS_MAPPING as _MU2_CLASS_MAPPING
 from .mu2 import MU2_CONFIG
 from .mu2 import Mu2DecoderIOConfig as _Mu2DecoderIOConfig  # triggers registration
@@ -121,6 +125,7 @@
         _BLIP_CLASS_MAPPING,
         _CLIP_CLASS_MAPPING,
         _MARIAN_CLASS_MAPPING,
+        _MGPSTR_CLASS_MAPPING,
         _MU2_CLASS_MAPPING,
         _QWEN_CLASS_MAPPING,
         _QWEN_TO_CLASS_MAPPING,

src/winml/modelkit/models/hf/mgp_str.py

@@ -0,0 +1,59 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Configure MGP-STR for Hugging Face scene-text-recognition export.
+
+MGP-STR is a Vision Transformer-based scene text recognition (STR) model. The
+upstream ``MgpstrForSceneTextRecognition`` head produces three logit tensors —
+``char_logits``, ``bpe_logits``, ``wp_logits`` — at three granularities
+(character / byte-pair / word-piece), which the ``MgpstrProcessor`` combines
+into the final decoded string.
+
+The vendor ``MgpstrOnnxConfig`` (Optimum) already exposes the 3-head outputs
+correctly but is registered ONLY under the ``feature-extraction`` task. End
+users naturally reach for the ``image-to-text`` task label for STR work; this
+module registers the same export config under ``image-to-text`` so the
+user-facing task resolves cleanly.
+
+Registering the canonical task lets the shared resolver accept the published
+``image-to-text`` pipeline metadata even when checkpoint architecture metadata
+is stale, then bind the family to its concrete head-bearing class.
+"""
+
+from __future__ import annotations
+
+from optimum.exporters.onnx.model_configs import MgpstrOnnxConfig
+from transformers import MgpstrForSceneTextRecognition
+
+from ...export import register_onnx_overwrite
+
+
+# =============================================================================
+# Image-to-text alias for MGP-STR
+# =============================================================================
+
+
+@register_onnx_overwrite("mgp-str", "image-to-text", library_name="transformers")
+class MgpstrImage2TextOnnxConfig(MgpstrOnnxConfig):  # type: ignore[misc]
+    """MGP-STR ONNX config bound to the ``image-to-text`` task.
+
+    The 3-head ``(char_logits, bpe_logits, wp_logits)`` output contract and
+    the ``pixel_values`` input contract are inherited unchanged from
+    ``MgpstrOnnxConfig``. The only purpose of this subclass is to register
+    the same export semantics under the ``image-to-text`` task name so users
+    can build MGP-STR with the natural task label.
+    """
+
+
+# =============================================================================
+# Model Class Mapping
+# =============================================================================
+
+# (model_type, task) -> HF model class. Binds the ``image-to-text`` task on
+# MGP-STR to ``MgpstrForSceneTextRecognition`` (the head-bearing class with the
+# 3-granularity outputs), instead of letting the loader fall back to
+# ``AutoModelForVision2Seq`` — MGP-STR is NOT a Vision2Seq architecture.
+MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = {
+    ("mgp-str", "image-to-text"): MgpstrForSceneTextRecognition,
+}

src/winml/modelkit/models/winml/__init__.py

@@ -59,6 +59,7 @@
     # - Specific OPSET requirements
     # - Input remapping (non-standard tensor names)
     # - Custom pre/post-processing
+    ("mgp-str", "image-to-text"): "WinMLModelForMgpstrSceneTextRecognition",
 }
 
 
@@ -82,6 +83,7 @@ def _import_winml_class(class_name: str) -> type[WinMLPreTrainedModel]:
         WinMLModelForImageSegmentation,
         WinMLModelForSemanticSegmentation,
     )
+    from .image_to_text import WinMLModelForMgpstrSceneTextRecognition
     from .object_detection import WinMLModelForObjectDetection
     from .question_answering import WinMLModelForQuestionAnswering
     from .sequence_classification import WinMLModelForSequenceClassification
@@ -91,6 +93,7 @@ def _import_winml_class(class_name: str) -> type[WinMLPreTrainedModel]:
         "WinMLModelForDepthEstimation": WinMLModelForDepthEstimation,
         "WinMLModelForFeatureExtraction": WinMLModelForFeatureExtraction,
         "WinMLModelForImageClassification": WinMLModelForImageClassification,
+        "WinMLModelForMgpstrSceneTextRecognition": WinMLModelForMgpstrSceneTextRecognition,
         "WinMLModelForImageSegmentation": WinMLModelForImageSegmentation,
         "WinMLModelForObjectDetection": WinMLModelForObjectDetection,
         "WinMLModelForQuestionAnswering": WinMLModelForQuestionAnswering,
@@ -204,6 +207,7 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None
     WinMLModelForImageSegmentation,
     WinMLModelForSemanticSegmentation,
 )
+from .image_to_text import WinMLModelForMgpstrSceneTextRecognition
 from .kv_cache import (
     WinMLCache,
     WinMLSlidingWindowCache,
@@ -233,6 +237,7 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None
     "WinMLModelForGenericTask",
     "WinMLModelForImageClassification",
     "WinMLModelForImageSegmentation",
+    "WinMLModelForMgpstrSceneTextRecognition",
     "WinMLModelForObjectDetection",
     "WinMLModelForSemanticSegmentation",
     "WinMLModelForSequenceClassification",

src/winml/modelkit/models/winml/image_to_text.py

@@ -0,0 +1,76 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Image-to-text inference wrappers and pipelines."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+from .base import WinMLPreTrainedModel
+
+
+class WinMLModelForMgpstrSceneTextRecognition(WinMLPreTrainedModel):
+    """Expose MGP-STR's three ONNX heads in the Transformers output contract."""
+
+    main_input_name = "pixel_values"
+
+    def forward(self, **kwargs: Any) -> Any:
+        """Run the ONNX graph and return its ordered three-head logits tuple."""
+        from transformers.models.mgp_str.modeling_mgp_str import MgpstrModelOutput
+
+        try:
+            pixel_values = kwargs.pop("pixel_values")
+        except KeyError as exc:
+            raise ValueError("MGP-STR inference requires 'pixel_values'.") from exc
+
+        outputs = self._run_inference(self._format_inputs(pixel_values=pixel_values))
+        return MgpstrModelOutput(
+            # Transformers types this heterogeneous three-head tuple as a
+            # one-item FloatTensor tuple even though that contradicts its
+            # runtime MGP-STR contract.
+            logits=cast(
+                "Any",
+                (
+                    outputs["char_logits"],
+                    outputs["bpe_logits"],
+                    outputs["wp_logits"],
+                ),
+            )
+        )
+
+
+class MgpstrImageToTextPipeline:
+    """Preprocess images and decode MGP-STR's character/BPE/WordPiece heads."""
+
+    def __init__(self, model: WinMLPreTrainedModel, model_id: str) -> None:
+        from transformers import AutoProcessor
+
+        self.model = model
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        self.image_processor = self.processor.image_processor
+        self.tokenizer = getattr(self.processor, "char_tokenizer", None)
+        self._preprocess_params: dict[str, Any] = {}
+
+    def _sanitize_parameters(self, **_kwargs: Any) -> tuple[dict, dict, dict]:
+        """Match the pipeline introspection contract used by inference clients."""
+        return {}, {}, {}
+
+    def __call__(self, images: Any, *, prompt: str | None = None, **_kwargs: Any) -> Any:
+        """Recognize text in one image or a batch of images."""
+        if prompt is not None:
+            raise ValueError("MGP-STR scene text recognition does not accept a text prompt.")
+
+        model_inputs = self.processor(images=images, return_tensors="pt")
+        outputs = self.model(**model_inputs)
+        decoded = self.processor.batch_decode(outputs.logits)
+
+        records = []
+        for index, text in enumerate(decoded["generated_text"]):
+            record: dict[str, Any] = {"generated_text": text}
+            if "scores" in decoded:
+                score = decoded["scores"][index]
+                record["score"] = float(score.item() if hasattr(score, "item") else score)
+            records.append(record)
+        return records

tests/unit/eval/test_image_to_text_evaluator.py

@@ -78,6 +78,29 @@ def test_registered(self):
         )
 
 
+class TestPreparePipeline:
+    def test_uses_shared_pipeline_factory(self):
+        from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig
+
+        evaluator = object.__new__(WinMLImageToTextEvaluator)
+        evaluator.model = MagicMock()
+        evaluator.config = WinMLEvaluationConfig(
+            model_id="local/mgp-str",
+            task="image-to-text",
+            dataset=DatasetConfig(path="local/dataset"),
+        )
+        expected = MagicMock()
+
+        with patch(
+            "winml.modelkit.inference.pipeline.create_pipeline",
+            return_value=expected,
+        ) as create:
+            result = evaluator.prepare_pipeline()
+
+        assert result is expected
+        create.assert_called_once_with("image-to-text", evaluator.model, "local/mgp-str")
+
+
 class TestCompute:
     """compute() iterates samples through the pipeline and aggregates metrics."""