From 00c8f27dc524434ac65a37f27a15de4f09f98738 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 12:43:08 +0800 Subject: [PATCH 1/4] feat(inference): support direct ONNX inpainting models --- .../cpu/cpu/inpainting_fp16_config.json | 29 ++++ .../cpu/cpu/inpainting_fp32_config.json | 10 ++ src/winml/modelkit/build/common.py | 21 ++- src/winml/modelkit/build/hf.py | 16 +- src/winml/modelkit/build/onnx.py | 16 +- src/winml/modelkit/commands/build.py | 64 ++++++-- src/winml/modelkit/config/build.py | 9 +- src/winml/modelkit/inference/engine.py | 8 +- src/winml/modelkit/inference/inpainting.py | 146 ++++++++++++++++++ src/winml/modelkit/inference/pipeline.py | 18 +++ src/winml/modelkit/inference/tasks.py | 23 +++ src/winml/modelkit/loader/onnx_hub.py | 89 ++++++++++- src/winml/modelkit/models/auto.py | 2 +- src/winml/modelkit/utils/cli.py | 2 +- src/winml/modelkit/utils/model_input.py | 33 +++- tests/unit/build/test_hf.py | 19 +++ tests/unit/build/test_onnx.py | 19 +++ tests/unit/commands/test_build.py | 40 +++++ tests/unit/config/test_build_onnx.py | 24 +++ .../inference/test_inpainting_pipeline.py | 89 +++++++++++ tests/unit/inference/test_tasks.py | 12 ++ tests/unit/loader/test_onnx_hub.py | 86 ++++++++++- tests/unit/utils/test_model_input.py | 23 ++- 23 files changed, 750 insertions(+), 48 deletions(-) create mode 100644 examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json create mode 100644 examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json create mode 100644 src/winml/modelkit/inference/inpainting.py create mode 100644 tests/unit/inference/test_inpainting_pipeline.py diff --git a/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json new file mode 100644 index 000000000..ff7203f5d --- /dev/null +++ b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json @@ -0,0 +1,29 @@ +{ + "skip_optimize": true, + "export": null, + "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": "inpainting" + } +} diff --git a/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json new file mode 100644 index 000000000..309f0d8dc --- /dev/null +++ b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json @@ -0,0 +1,10 @@ +{ + "skip_optimize": true, + "export": null, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "inpainting" + } +} diff --git a/src/winml/modelkit/build/common.py b/src/winml/modelkit/build/common.py index 7e221d2ee..f6fe4d714 100644 --- a/src/winml/modelkit/build/common.py +++ b/src/winml/modelkit/build/common.py @@ -31,10 +31,12 @@ def ensure_pre_quantized_stamped( config: WinMLBuildConfig, onnx_path: Path, *, force: bool = False ) -> None: - """Stamp ``config.skip_optimize`` (and clear ``config.quant``) once. + """Stamp ``config.skip_optimize`` and suppress incompatible quantization. - Sets ``config.skip_optimize = True`` and clears ``config.quant`` if the - input ONNX is already quantized. + Sets ``config.skip_optimize = True`` when the input ONNX is already + quantized. QDQ/QOperator quantization is suppressed, but an explicit FP16 + conversion is preserved because it converts the graph's remaining float + tensors rather than re-quantizing integer weights. This is the **single defensive detection point** for the library entry points (``build_onnx_model``, ``build_hf_model``). When @@ -50,20 +52,25 @@ def ensure_pre_quantized_stamped( ``is_quantized_onnx`` (used to honor the legacy ``skip_optimize=True`` kwarg from direct callers). """ + + def _clear_incompatible_quant() -> None: + if config.quant is not None and config.quant.mode != "fp16": + config.quant = None + if config.skip_optimize: - config.quant = None + _clear_incompatible_quant() return if force: config.skip_optimize = True - config.quant = None + _clear_incompatible_quant() return if is_quantized_onnx(onnx_path): config.skip_optimize = True - config.quant = None + _clear_incompatible_quant() logger.info( "Pre-quantized model detected (QDQ or QOperator nodes present). " - "Skipping optimize + quantize stages." + "Skipping graph optimization and incompatible quantization." ) diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index 888c65df8..c3d77933f 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -254,14 +254,15 @@ def _name(base: str) -> str: is_pre_quantized = config.skip_optimize if is_pre_quantized: - logger.info("Skipping optimize + quantize stages (config.skip_optimize=True)") + logger.info("Skipping optimize stage (config.skip_optimize=True)") stages_skipped.append("optimize") # Skip the ORT-based graph optimization (no kernel for QOperator # ops like ConvInteger on the host EP). The autoconf re-optim/ # analyze loop is disabled too -- ``run_optimize_analyze_loop`` # forces ``max_optim_iterations=0`` when ``skip_optimize=True``, # so ``_run_analyze_loop`` is not invoked. The model still flows - # through later stages (quantize-skip + compile) for validation. + # through later stages. In particular, an explicit FP16 conversion + # still converts the graph's remaining float tensors before compile. current_path, _, analyze_iterations, analyze_unsupported_nodes, analyze_details = ( run_optimize_analyze_loop( model_path=current_path, @@ -304,15 +305,14 @@ def _name(base: str) -> str: # ========================================================================= # [4] QUANTIZE (optional — config.quant=None means skip) # ========================================================================= - # No defensive ``is_quantized_onnx`` re-check here: when the model is - # pre-quantized, ``ensure_pre_quantized_stamped`` has already set - # ``config.quant = None`` at stage [3], so this branch naturally - # falls through to the ``quant is None`` skip path. + # No defensive ``is_quantized_onnx`` re-check here. The stage-[3] + # stamping suppresses incompatible integer quantization while preserving + # an explicit FP16 conversion of the graph's remaining float tensors. quant_result = None - if is_pre_quantized: + if is_pre_quantized and (config.quant is None or config.quant.mode != "fp16"): if "quantize" not in stages_skipped: stages_skipped.append("quantize") - logger.info("Quantize skipped (pre-quantized model)") + logger.info("Quantize skipped (incompatible with pre-quantized model)") elif config.quant is not None: logger.info("Quantizing model...") t0 = time.monotonic() diff --git a/src/winml/modelkit/build/onnx.py b/src/winml/modelkit/build/onnx.py index 8e0095e5e..7bf225b74 100644 --- a/src/winml/modelkit/build/onnx.py +++ b/src/winml/modelkit/build/onnx.py @@ -165,14 +165,15 @@ def _name(base: str) -> str: is_pre_quantized = config.skip_optimize if is_pre_quantized: - logger.info("Skipping optimize + quantize stages (config.skip_optimize=True)") + logger.info("Skipping optimize stage (config.skip_optimize=True)") stages_skipped.append("optimize") # Skip the ORT-based graph optimization (no kernel for QOperator # ops like ConvInteger on the host EP). The autoconf re-optim/ # analyze loop is disabled too -- ``run_optimize_analyze_loop`` # forces ``max_optim_iterations=0`` when ``skip_optimize=True``, # so ``_run_analyze_loop`` is not invoked. The model still flows - # through later stages (quantize-skip + compile) for validation. + # through later stages. In particular, an explicit FP16 conversion + # still converts the graph's remaining float tensors before compile. current_path, _, analyze_iters, analyze_unsupported, analyze_details = ( run_optimize_analyze_loop( model_path=current_path, @@ -210,16 +211,15 @@ def _name(base: str) -> str: # ========================================================================= # [2] QUANTIZE (optional — config.quant=None means skip) # ========================================================================= - # No defensive ``is_quantized_onnx`` re-check here: when the model is - # pre-quantized, ``ensure_pre_quantized_stamped`` has already set - # ``config.quant = None`` at stage [1], so this branch naturally - # falls through to the ``quant is None`` skip path. + # No defensive ``is_quantized_onnx`` re-check here. The stage-[1] + # stamping suppresses incompatible integer quantization while preserving + # an explicit FP16 conversion of the graph's remaining float tensors. quant_result = None - if is_pre_quantized: + if is_pre_quantized and (config.quant is None or config.quant.mode != "fp16"): # Already handled above -- skip quantize for pre-quantized models if "quantize" not in stages_skipped: stages_skipped.append("quantize") - logger.info("Quantize skipped (pre-quantized model)") + logger.info("Quantize skipped (incompatible with pre-quantized model)") elif config.quant is not None: logger.info("Quantizing model...") t0 = time.monotonic() diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index 3076b585a..6165d524f 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -38,7 +38,7 @@ print_stages_header, ) from ..utils.logging import configure_logging -from ..utils.model_input import ModelInputKind, classify_model_input +from ..utils.model_input import ModelInputKind, classify_model_input, resolve_model_input if TYPE_CHECKING: @@ -935,12 +935,7 @@ def build( # Hub-hosted ONNX (e.g. ``onnx-community/sam3-tracker-ONNX/onnx/...``) # is downloaded once and treated as a local .onnx file thereafter. model_input = None - if model: - model = cli_utils.normalize_model_arg(model) - if model: - model_input = classify_model_input(model) - if model_input.kind is ModelInputKind.INVALID: - raise click.UsageError(model_input.error or f"Invalid model input: {model}") + source_model_input = None # Load or auto-generate config if config_file is not None: @@ -949,6 +944,19 @@ def build( no_quant=not quant, no_compile=no_compile, ) + if model: + is_single_config = not isinstance(config_or_configs, list) + direct_onnx = is_single_config and config_or_configs.export is None + source_model_input = resolve_model_input( + model, + discover_repo_onnx=direct_onnx, + ) + if source_model_input.kind is ModelInputKind.INVALID: + raise click.UsageError( + source_model_input.error or f"Invalid model input: {model}" + ) + model = source_model_input.local_path or model + model_input = classify_model_input(model) if export_overrides: from ..config import merge_export_overrides @@ -976,6 +984,12 @@ def build( else: if not model: raise click.UsageError("-m/--model is required when -c is not provided.") + source_model_input = resolve_model_input(model, discover_repo_onnx=True) + if source_model_input.kind is ModelInputKind.INVALID: + invalid_message = source_model_input.error or f"Invalid model input: {model}" + raise click.UsageError(invalid_message) + model = source_model_input.local_path or model + model_input = classify_model_input(model) from ..config import generate_build_config # When ``model`` resolves to an .onnx file (either a local path or @@ -1031,7 +1045,11 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: resolved_quant, _ = resolve_quant_compile_config( device=device, precision=precision, ep=ep ) - if cfg.skip_optimize or not quant or resolved_quant is None: + if ( + not quant + or resolved_quant is None + or (cfg.skip_optimize and resolved_quant.mode != "fp16") + ): cfg.quant = None elif cfg.quant is None: # Populate calibration identifiers from the loader/model @@ -1103,6 +1121,12 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: # on the key being present, matching the module-mode path which passes # allow_unsupported_nodes explicitly regardless of its value. extra_kwargs["allow_unsupported_nodes"] = allow_unsupported_nodes + source_kind = source_model_input.kind if source_model_input is not None else None + if source_kind is ModelInputKind.HUB_ONNX: + assert source_model_input is not None + extra_kwargs["source_model_id"] = source_model_input.hf_id + extra_kwargs["source_onnx_file"] = source_model_input.artifact_path + extra_kwargs["source_revision"] = source_model_input.revision # ---- OPTIMIZED (GENAI BUNDLE) EXPORT ---- # ``--export-type optimized`` builds the family's registered @@ -1779,7 +1803,7 @@ def _run_quantize_stage( from ..quant import quantize_onnx from ..utils.console import StageLive - if config.skip_optimize: + if config.skip_optimize and (config.quant is None or config.quant.mode != "fp16"): config.quant = None return current_path @@ -2116,6 +2140,9 @@ def _build_onnx_pipeline( max_iters: int = extra_kwargs.pop("hack_max_optim_iterations", 3) allow_unsupported_nodes: bool = extra_kwargs.pop("allow_unsupported_nodes", False) + source_model_id: str | None = extra_kwargs.pop("source_model_id", None) + source_onnx_file: str | None = extra_kwargs.pop("source_onnx_file", None) + source_revision: str | None = extra_kwargs.pop("source_revision", None) # ── Validate + setup ───────────────────────────────────────── if not onnx_path.exists(): @@ -2204,4 +2231,23 @@ def _build_onnx_pipeline( if current_path != final_path: copy_onnx_model(current_path, final_path) + from ..utils import MANIFEST_FILENAME, WinMLManifest + + manifest = WinMLManifest( + source="onnx", + model_id=source_model_id, + task=config.loader.task, + input_onnx=str(onnx_path), + final_artifact=final_path.name, + extras={ + key: value + for key, value in { + "hub_onnx_file": source_onnx_file, + "hub_revision": source_revision, + }.items() + if value is not None + }, + ) + manifest.save(output_dir / MANIFEST_FILENAME) + return stage_timings diff --git a/src/winml/modelkit/config/build.py b/src/winml/modelkit/config/build.py index 07f3b2091..495804af6 100644 --- a/src/winml/modelkit/config/build.py +++ b/src/winml/modelkit/config/build.py @@ -459,10 +459,15 @@ def generate_onnx_build_config( ) if is_quantized_onnx(onnx_path_resolved): - # Skip optimize+quantize, compile with resolved policy. + # Skip graph optimization and incompatible re-quantization, but + # preserve explicit FP16 conversion of the graph's float tensors. # ``skip_optimize`` is the single source of truth — downstream # pipelines must read this flag and not re-detect. - config.quant = None + config.quant = ( + resolved_quant + if resolved_quant is not None and resolved_quant.mode == "fp16" + else None + ) config.skip_optimize = True config.compile = resolved_compile logger.info("Quantized model (QDQ) detected") diff --git a/src/winml/modelkit/inference/engine.py b/src/winml/modelkit/inference/engine.py index c38ad9c02..dbc9546d7 100644 --- a/src/winml/modelkit/inference/engine.py +++ b/src/winml/modelkit/inference/engine.py @@ -327,7 +327,9 @@ def load( # is downloaded once and treated as a local .onnx path thereafter. from ..utils.model_input import resolve_model_input - model_path = resolve_model_input(str(model_path)).local_path or str(model_path) + model_path = resolve_model_input( + str(model_path), discover_repo_onnx=True + ).local_path or str(model_path) self._model_path = str(model_path) self._ep = ep @@ -408,7 +410,9 @@ def load_schema_only( # is downloaded once and treated as a local .onnx path thereafter. from ..utils.model_input import resolve_model_input - model_path = resolve_model_input(str(model_path)).local_path or str(model_path) + model_path = resolve_model_input( + str(model_path), discover_repo_onnx=True + ).local_path or str(model_path) self._model_path = str(model_path) self._device = device diff --git a/src/winml/modelkit/inference/inpainting.py b/src/winml/modelkit/inference/inpainting.py new file mode 100644 index 000000000..46a40d0b7 --- /dev/null +++ b/src/winml/modelkit/inference/inpainting.py @@ -0,0 +1,146 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Pipeline adapter for direct-ONNX image-inpainting models.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from PIL import Image + + +if TYPE_CHECKING: + from ..models.winml.base import WinMLPreTrainedModel + + +class WinMLInpaintingPipeline: + """Prepare an image/mask pair and postprocess a direct ONNX result. + + The adapter derives spatial sizes and tensor roles from ONNX metadata. It + implements the OpenCV inpainting interchange contract: float32 NCHW BGR + image values in ``[0, 1]`` and a float32 NCHW binary mask. The first + three-channel output is converted back to an RGB PIL image. + """ + + def __init__(self, model: WinMLPreTrainedModel) -> None: + self.model = model + self._inputs, self._output_name = self._resolve_contract(model.io_config) + + def _sanitize_parameters(self, **_kwargs: Any) -> tuple[dict, dict, dict]: + """Expose the standard pipeline parameter-discovery contract.""" + return {}, {}, {} + + @staticmethod + def _resolve_contract( + io_config: dict[str, Any], + ) -> tuple[dict[str, tuple[str, list[Any]]], str]: + names = list(io_config.get("input_names", [])) + shapes = list(io_config.get("input_shapes", [])) + output_names = list(io_config.get("output_names", [])) + output_shapes = list(io_config.get("output_shapes", [])) + try: + entries = list(zip(names, shapes, strict=True)) + outputs = list(zip(output_names, output_shapes, strict=True)) + except ValueError as exc: + raise ValueError("Inpainting ONNX I/O metadata is incomplete.") from exc + + def _pick(preferred_name: str, channels: int) -> tuple[str, list[Any]] | None: + by_name = next((entry for entry in entries if entry[0].lower() == preferred_name), None) + if by_name is not None and len(by_name[1]) == 4 and by_name[1][1] == channels: + return by_name + return next( + (entry for entry in entries if len(entry[1]) == 4 and entry[1][1] == channels), + None, + ) + + image = _pick("image", 3) + mask = _pick("mask", 1) + if image is None or mask is None or image[0] == mask[0]: + raise ValueError( + "Inpainting requires distinct 4-D image (3-channel) and mask " + "(1-channel) ONNX inputs." + ) + image_size = WinMLInpaintingPipeline._spatial_size(image[1], "image") + mask_size = WinMLInpaintingPipeline._spatial_size(mask[1], "mask") + if image_size != mask_size: + raise ValueError("Inpainting image and mask inputs must have the same spatial size.") + + output = next( + ( + entry + for entry in outputs + if entry[0].lower() in {"output", "image", "images"} + and len(entry[1]) == 4 + and entry[1][1] == 3 + ), + None, + ) + if output is None: + output = next( + (entry for entry in outputs if len(entry[1]) == 4 and entry[1][1] == 3), + None, + ) + if output is None: + raise ValueError("Inpainting requires a 4-D 3-channel ONNX image output.") + output_size = WinMLInpaintingPipeline._spatial_size(output[1], "output") + if output_size != image_size: + raise ValueError("Inpainting output must match the image input spatial size.") + + return {"image": image, "mask": mask}, output[0] + + @staticmethod + def _spatial_size(shape: list[Any], role: str) -> tuple[int, int]: + if len(shape) != 4 or not isinstance(shape[2], int) or not isinstance(shape[3], int): + raise ValueError(f"Inpainting {role} input requires fixed NCHW height and width.") + return shape[3], shape[2] + + @staticmethod + def _prepare_image(image: Image.Image, size: tuple[int, int]) -> np.ndarray: + rgb = np.asarray(image.convert("RGB").resize(size, Image.Resampling.BILINEAR)) + bgr = rgb[..., ::-1].astype(np.float32) / 255.0 + return np.ascontiguousarray(bgr.transpose(2, 0, 1)[None, ...]) + + @staticmethod + def _prepare_mask(mask: Image.Image, size: tuple[int, int]) -> np.ndarray: + grayscale = np.asarray(mask.convert("L").resize(size, Image.Resampling.NEAREST)) + binary = (grayscale > 0).astype(np.float32) + return binary[None, None, ...] + + @staticmethod + def _to_image(output: Any) -> Image.Image: + if isinstance(output, torch.Tensor): + array = output.detach().cpu().numpy() + else: + array = np.asarray(output) + if array.ndim != 4 or array.shape[0] != 1 or array.shape[1] != 3: + raise ValueError( + f"Inpainting output must have shape [1, 3, height, width], got {list(array.shape)}." + ) + image = array[0].transpose(1, 2, 0) + if image.size and float(np.nanmax(image)) <= 1.0: + image = image * 255.0 + rgb = np.clip(image[..., ::-1], 0, 255).astype(np.uint8) + return Image.fromarray(rgb, mode="RGB") + + def __call__(self, inputs: dict[str, Image.Image], **_kwargs: Any) -> Image.Image: + """Run inpainting for a decoded ``{"image": ..., "mask": ...}`` input.""" + image_name, image_shape = self._inputs["image"] + mask_name, mask_shape = self._inputs["mask"] + model_inputs = { + image_name: self._prepare_image( + inputs["image"], self._spatial_size(image_shape, "image") + ), + mask_name: self._prepare_mask(inputs["mask"], self._spatial_size(mask_shape, "mask")), + } + outputs = self.model(**model_inputs) + if not isinstance(outputs, dict) or not outputs: + raise ValueError("Inpainting model returned no named ONNX outputs.") + if self._output_name not in outputs: + raise ValueError( + f"Inpainting model did not return expected ONNX output '{self._output_name}'." + ) + return self._to_image(outputs[self._output_name]) diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index fc8bdb66c..3e219b312 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -38,6 +38,18 @@ } +def _create_inpainting_pipeline(model: WinMLPreTrainedModel) -> Any: + """Create the built-in direct-ONNX inpainting adapter.""" + from .inpainting import WinMLInpaintingPipeline + + return WinMLInpaintingPipeline(model) + + +_CUSTOM_PIPELINE_FACTORIES = { + "inpainting": _create_inpainting_pipeline, +} + + def create_pipeline( task: str, model: WinMLPreTrainedModel | WinMLCompositeModel, @@ -57,6 +69,12 @@ def create_pipeline( Returns: A configured ``transformers.Pipeline`` ready for inference. """ + custom_factory = _CUSTOM_PIPELINE_FACTORIES.get(task) + if custom_factory is not None: + pipe = custom_factory(model) + logger.info("Created WinML pipeline: task=%s model=%s", task, model_id) + return pipe + from transformers import pipeline kwargs: dict[str, Any] = { diff --git a/src/winml/modelkit/inference/tasks.py b/src/winml/modelkit/inference/tasks.py index e0602ff65..3669e79de 100644 --- a/src/winml/modelkit/inference/tasks.py +++ b/src/winml/modelkit/inference/tasks.py @@ -150,6 +150,16 @@ def _encode_mask_png(mask: Any) -> str: return base64.b64encode(buf.getvalue()).decode("ascii") +def _postprocess_inpainting(raw: Any, **_kwargs: Any) -> dict[str, Any]: + """Encode an inpainted PIL image as a JSON-safe PNG payload.""" + return { + "image": _encode_mask_png(raw), + "format": "png", + "width": raw.width, + "height": raw.height, + } + + def _postprocess_segmentation(raw: Any, **_kwargs: Any) -> Any: """Convert segmentation masks to predictions with coverage score and mask. @@ -287,6 +297,19 @@ def _postprocess_sentence_similarity( ], mapping=PipelineMapping(pipe_input="image"), ), + "inpainting": TaskInputSpec( + user_inputs=[ + InputField(name="image", type="image", required=True, description="Input image"), + InputField( + name="mask", + type="image", + required=True, + description="Binary mask; non-zero pixels are replaced", + ), + ], + mapping=PipelineMapping(pipe_input=["image", "mask"]), + postprocess=_postprocess_inpainting, + ), # -- Single text -------------------------------------------------------- "text-classification": TaskInputSpec( user_inputs=[ diff --git a/src/winml/modelkit/loader/onnx_hub.py b/src/winml/modelkit/loader/onnx_hub.py index 5a9a5dc1b..6e9324299 100644 --- a/src/winml/modelkit/loader/onnx_hub.py +++ b/src/winml/modelkit/loader/onnx_hub.py @@ -19,12 +19,95 @@ from __future__ import annotations import logging +from dataclasses import dataclass from pathlib import Path logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class ResolvedHubOnnx: + """A downloaded Hub ONNX artifact and its immutable provenance.""" + + local_path: Path + repo_id: str + filename: str + revision: str + + +def resolve_hf_repo_onnx( + repo_id: str, + *, + revision: str | None = None, + cache_dir: str | Path | None = None, + token: str | bool | None = None, +) -> ResolvedHubOnnx | None: + """Resolve a non-Transformers Hub repository containing one ONNX graph. + + Repositories with a usable Transformers ``model_type`` remain normal HF + model IDs. A repository without that metadata is routed to direct ONNX only + when exactly one ``.onnx`` sibling exists. Multiple graphs are deliberately + rejected because selecting one by filename or ordering would be a hidden, + model-specific policy. + + The repository listing resolves the requested branch/tag to a commit SHA, + and that SHA is used for the artifact download. This prevents a moving Hub + revision from changing between discovery and download. + """ + from huggingface_hub import HfApi + from huggingface_hub.errors import HfHubHTTPError + + try: + info = HfApi().model_info( + repo_id=repo_id, + revision=revision, + files_metadata=False, + token=token, + ) + except HfHubHTTPError as exc: + # Discovery is an optional preflight before the existing Transformers + # path. Auth, gated-repo, missing-repo, and transient Hub HTTP errors + # must therefore fall through so the authoritative downstream loader + # can preserve its established diagnostics and mocked/offline behavior. + logger.debug("Could not inspect %s for a direct ONNX artifact: %s", repo_id, exc) + return None + config = info.config if isinstance(info.config, dict) else {} + if config.get("model_type"): + return None + + onnx_files = sorted( + sibling.rfilename + for sibling in (info.siblings or []) + if sibling.rfilename.lower().endswith(".onnx") + ) + if not onnx_files: + return None + if len(onnx_files) > 1: + listing = "\n".join(f" - {repo_id}/{filename}" for filename in onnx_files) + raise ValueError( + f"Hub repo '{repo_id}' has no usable Transformers model_type and contains " + f"multiple ONNX files. Pass an explicit artifact path:\n{listing}" + ) + + resolved_revision = info.sha or revision + if not resolved_revision: + raise ValueError(f"Could not resolve an immutable revision for Hub repo '{repo_id}'.") + filename = onnx_files[0] + local_path = resolve_hf_onnx_path( + f"{repo_id}/{filename}", + revision=resolved_revision, + cache_dir=cache_dir, + token=token, + ) + return ResolvedHubOnnx( + local_path=local_path, + repo_id=repo_id, + filename=filename, + revision=resolved_revision, + ) + + def resolve_hf_onnx_path( model_id: str, *, @@ -79,9 +162,7 @@ def resolve_hf_onnx_path( # ``.onnx`` files so the user can pick the right one without leaving # the terminal. Re-raise as ``FileNotFoundError`` so callers that # already handle local-file-missing errors get a consistent type. - hint = _format_available_onnx_files( - repo_id, revision=revision, token=token - ) + hint = _format_available_onnx_files(repo_id, revision=revision, token=token) raise FileNotFoundError( f"ONNX file '{filename}' not found in Hub repo '{repo_id}'.\n{hint}" ) from e @@ -171,5 +252,7 @@ def _format_available_onnx_files( __all__ = [ + "ResolvedHubOnnx", "resolve_hf_onnx_path", + "resolve_hf_repo_onnx", ] diff --git a/src/winml/modelkit/models/auto.py b/src/winml/modelkit/models/auto.py index 54a36ffa4..7a071de89 100644 --- a/src/winml/modelkit/models/auto.py +++ b/src/winml/modelkit/models/auto.py @@ -336,7 +336,7 @@ def from_pretrained( # is downloaded once and treated as a local .onnx path thereafter. from ..utils.model_input import resolve_model_input - model_id = resolve_model_input(model_id).local_path or model_id + model_id = resolve_model_input(model_id, discover_repo_onnx=True).local_path or model_id # ===================================================================== # ONNX FAST PATH -- skip HF loading and export when given an .onnx file diff --git a/src/winml/modelkit/utils/cli.py b/src/winml/modelkit/utils/cli.py index 9a48d1c42..5f2583787 100644 --- a/src/winml/modelkit/utils/cli.py +++ b/src/winml/modelkit/utils/cli.py @@ -1215,7 +1215,7 @@ def normalize_model_arg(value: str | None) -> str | None: return None from .model_input import resolve_model_input - return resolve_model_input(value).local_path or value + return resolve_model_input(value, discover_repo_onnx=True).local_path or value def is_cli_provided(ctx: click.Context, param_name: str) -> bool: diff --git a/src/winml/modelkit/utils/model_input.py b/src/winml/modelkit/utils/model_input.py index bf48ad8b5..fe99b4ba6 100644 --- a/src/winml/modelkit/utils/model_input.py +++ b/src/winml/modelkit/utils/model_input.py @@ -76,6 +76,9 @@ class ModelInput: ``HF_ID``. hf_id: HuggingFace repo id (``org/name``) for ``HF_ID`` / ``HUB_ONNX``. ``None`` otherwise. + artifact_path: Path inside the Hub repository for ``HUB_ONNX``. + revision: Immutable Hub commit used to download a discovered repository + artifact, or the caller-provided revision for an explicit artifact. error: Human-readable reason the value is ``INVALID``. ``None`` for every valid kind. """ @@ -84,6 +87,8 @@ class ModelInput: raw: str local_path: str | None = None hf_id: str | None = None + artifact_path: str | None = None + revision: str | None = None error: str | None = None @@ -194,6 +199,7 @@ def resolve_model_input( revision: str | None = None, cache_dir: str | Path | None = None, token: str | bool | None = None, + discover_repo_onnx: bool = False, ) -> ModelInput: """Classify + download Hub-hosted ONNX refs in one call. @@ -213,6 +219,25 @@ def resolve_model_input( kind implies a filesystem path. """ mi = classify_model_input(value) + if mi.kind is ModelInputKind.HF_ID and discover_repo_onnx: + from ..loader.onnx_hub import resolve_hf_repo_onnx + + discovered = resolve_hf_repo_onnx( + value, + revision=revision, + cache_dir=cache_dir, + token=token, + ) + if discovered is not None: + return ModelInput( + kind=ModelInputKind.HUB_ONNX, + raw=value, + local_path=str(discovered.local_path), + hf_id=discovered.repo_id, + artifact_path=discovered.filename, + revision=discovered.revision, + ) + return mi if mi.kind is not ModelInputKind.HUB_ONNX: return mi @@ -222,7 +247,13 @@ def resolve_model_input( from ..loader.onnx_hub import resolve_hf_onnx_path local = resolve_hf_onnx_path(value, revision=revision, cache_dir=cache_dir, token=token) - return replace(mi, local_path=str(local)) + parts = [part for part in value.split("/") if part] + return replace( + mi, + local_path=str(local), + artifact_path="/".join(parts[2:]), + revision=revision, + ) __all__ = [ diff --git a/tests/unit/build/test_hf.py b/tests/unit/build/test_hf.py index cca015319..07df15284 100644 --- a/tests/unit/build/test_hf.py +++ b/tests/unit/build/test_hf.py @@ -836,6 +836,25 @@ def test_post_export_qdq_skips_optimize_and_quantize( mock_pipeline["optimize"].assert_not_called() mock_pipeline["quantize"].assert_not_called() + def test_post_export_qdq_preserves_explicit_fp16_conversion( + self, tmp_path: Path, sample_config, mock_pipeline + ) -> None: + """QDQ/QOperator exports still convert their remaining float tensors to FP16.""" + mock_pipeline["is_quantized_onnx"].return_value = True + sample_config.quant.mode = "fp16" + + result = build_hf_model( + config=sample_config, + output_dir=tmp_path / "output", + pytorch_model=mock_pipeline["model"], + ) + + assert "optimize" in result.stages_skipped + assert "quantize" in result.stages_completed + assert "quantize" not in result.stages_skipped + mock_pipeline["optimize"].assert_not_called() + mock_pipeline["quantize"].assert_called_once() + def test_post_export_qdq_still_exports( self, tmp_path: Path, sample_config, mock_pipeline ) -> None: diff --git a/tests/unit/build/test_onnx.py b/tests/unit/build/test_onnx.py index e26907287..1254bd2d8 100644 --- a/tests/unit/build/test_onnx.py +++ b/tests/unit/build/test_onnx.py @@ -363,6 +363,25 @@ def test_build_onnx_non_quantized_proceeds( assert "quantize" in result.stages_completed mock_onnx_pipeline["quantize"].assert_called_once() + def test_pre_quantized_preserves_explicit_fp16_conversion( + self, tmp_path: Path, fake_onnx: Path, sample_onnx_config, mock_onnx_pipeline + ) -> None: + """QDQ/QOperator graphs still convert their remaining float tensors to FP16.""" + mock_onnx_pipeline["is_quantized_onnx"].return_value = True + sample_onnx_config.quant.mode = "fp16" + + result = build_onnx_model( + fake_onnx, + config=sample_onnx_config, + output_dir=tmp_path / "output", + ) + + assert "optimize" in result.stages_skipped + assert "quantize" in result.stages_completed + assert "quantize" not in result.stages_skipped + mock_onnx_pipeline["optimize"].assert_not_called() + mock_onnx_pipeline["quantize"].assert_called_once() + def test_pre_quantized_skips_optimize_and_quantize( self, tmp_path: Path, fake_onnx: Path, sample_onnx_config, mock_onnx_pipeline ) -> None: diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 9d1a35cf3..1e386ae7c 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -623,6 +623,46 @@ def test_precision_fp16_sets_fp16_algorithm( assert quant is not None assert quant.mode == "fp16" + def test_precision_fp16_preserved_for_prequantized_onnx( + self, tmp_path: Path, mock_run_single_build: MagicMock + ): + """Pre-quantized ONNX still converts remaining float tensors to FP16.""" + onnx_path = tmp_path / "prequantized.onnx" + onnx_path.write_bytes(b"fake") + config_path = tmp_path / "prequantized_fp16.json" + config_path.write_text( + json.dumps( + { + "skip_optimize": True, + "loader": {"task": "inpainting"}, + "export": None, + "optim": {}, + "quant": {"mode": "fp16"}, + "compile": None, + } + ) + ) + + result = _invoke( + [ + "-c", + str(config_path), + "-m", + str(onnx_path), + "-o", + str(tmp_path / "out"), + "--device", + "cpu", + "--precision", + "fp16", + ] + ) + + assert result.exit_code == 0, result.output + quant = mock_run_single_build.call_args.kwargs["config"].quant + assert quant is not None + assert quant.mode == "fp16" + def test_precision_alone_triggers_quant_patch( self, tmp_path: Path, mock_run_single_build: MagicMock ): diff --git a/tests/unit/config/test_build_onnx.py b/tests/unit/config/test_build_onnx.py index 25180a255..6aff78471 100644 --- a/tests/unit/config/test_build_onnx.py +++ b/tests/unit/config/test_build_onnx.py @@ -255,6 +255,30 @@ def test_quantized_onnx_cpu(self, tmp_path) -> None: assert config.quant is None assert config.compile is None + def test_quantized_onnx_preserves_explicit_fp16_conversion(self, tmp_path) -> None: + """Pre-quantized integer weights do not suppress requested FP16 float conversion.""" + onnx_file = tmp_path / "quantized.onnx" + onnx_file.write_bytes(b"fake") + + with ( + patch("winml.modelkit.onnx.is_compiled_onnx", return_value=False), + patch("winml.modelkit.onnx.is_quantized_onnx", return_value=True), + patch( + "winml.modelkit.sysinfo.resolve_check_device_ep", + return_value=("cpu", ["cpu"], ["CPUExecutionProvider"]), + ), + ): + config = generate_onnx_build_config( + str(onnx_file), + device="cpu", + precision="fp16", + ) + + assert config.skip_optimize is True + assert config.quant is not None + assert config.quant.mode == "fp16" + assert config.compile is None + def test_compiled_onnx_skips_all(self, tmp_path) -> None: """Compiled ONNX (EPContext) sets quant=None and compile=None.""" onnx_file = tmp_path / "compiled.onnx" diff --git a/tests/unit/inference/test_inpainting_pipeline.py b/tests/unit/inference/test_inpainting_pipeline.py new file mode 100644 index 000000000..e1ef202e6 --- /dev/null +++ b/tests/unit/inference/test_inpainting_pipeline.py @@ -0,0 +1,89 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for the generic direct-ONNX inpainting pipeline.""" + +from __future__ import annotations + +from typing import ClassVar + +import numpy as np +import pytest +import torch +from PIL import Image + +from winml.modelkit.inference.inpainting import WinMLInpaintingPipeline + + +class _EchoInpaintingModel: + io_config: ClassVar[dict] = { + "input_names": ["source_pixels", "edit_region"], + "input_shapes": [[1, 3, 2, 2], [1, 1, 2, 2]], + "output_names": ["completed_image"], + "output_shapes": [[1, 3, 2, 2]], + } + + def __init__(self) -> None: + self.inputs: dict[str, np.ndarray] = {} + + def __call__(self, **kwargs): + self.inputs = kwargs + return {"completed_image": torch.from_numpy(kwargs["source_pixels"] * 255.0)} + + +class TestWinMLInpaintingPipeline: + def test_prepares_bgr_image_and_binary_mask(self) -> None: + model = _EchoInpaintingModel() + pipeline = WinMLInpaintingPipeline(model) # type: ignore[arg-type] + pixels = np.full((2, 2, 3), [10, 20, 30], dtype=np.uint8) + mask = np.array([[0, 1], [127, 255]], dtype=np.uint8) + + output = pipeline( + { + "image": Image.fromarray(pixels, mode="RGB"), + "mask": Image.fromarray(mask, mode="L"), + } + ) + + assert model.inputs["source_pixels"].shape == (1, 3, 2, 2) + np.testing.assert_allclose( + model.inputs["source_pixels"][0, :, 0, 0], + np.array([30, 20, 10], dtype=np.float32) / 255.0, + ) + np.testing.assert_array_equal( + model.inputs["edit_region"], + np.array([[[[0, 1], [1, 1]]]], dtype=np.float32), + ) + np.testing.assert_array_equal(np.asarray(output), pixels) + + def test_requires_image_and_mask_tensor_roles(self) -> None: + model = _EchoInpaintingModel() + model.io_config = { + "input_names": ["pixels"], + "input_shapes": [[1, 3, 2, 2]], + "output_names": ["output"], + "output_shapes": [[1, 3, 2, 2]], + } + with pytest.raises(ValueError, match=r"image .* and mask"): + WinMLInpaintingPipeline(model) # type: ignore[arg-type] + + def test_requires_matching_image_output(self) -> None: + model = _EchoInpaintingModel() + model.io_config = { + **model.io_config, + "output_names": ["embedding"], + "output_shapes": [[1, 128]], + } + with pytest.raises(ValueError, match="3-channel ONNX image output"): + WinMLInpaintingPipeline(model) # type: ignore[arg-type] + + def test_selects_named_image_output(self) -> None: + model = _EchoInpaintingModel() + model.io_config = { + **model.io_config, + "output_names": ["scores", "completed_image"], + "output_shapes": [[1, 10], [1, 3, 2, 2]], + } + pipeline = WinMLInpaintingPipeline(model) # type: ignore[arg-type] + assert pipeline._output_name == "completed_image" diff --git a/tests/unit/inference/test_tasks.py b/tests/unit/inference/test_tasks.py index e36dd9d00..37b44190f 100644 --- a/tests/unit/inference/test_tasks.py +++ b/tests/unit/inference/test_tasks.py @@ -22,6 +22,7 @@ # Private helpers under test from winml.modelkit.inference.tasks import ( _masked_mean_pool, + _postprocess_inpainting, _postprocess_segmentation, _postprocess_sentence_similarity, ) @@ -177,3 +178,14 @@ def test_has_postprocess_callback(self) -> None: def test_semantic_alias(self) -> None: assert TASK_REGISTRY["semantic-segmentation"].postprocess is _postprocess_segmentation + + +class TestInpaintingRegistry: + def test_requires_image_and_mask(self) -> None: + spec = TASK_REGISTRY["inpainting"] + assert [field.name for field in spec.user_inputs] == ["image", "mask"] + assert all(field.required for field in spec.user_inputs) + assert spec.mapping.pipe_input == ["image", "mask"] + + def test_has_png_postprocess(self) -> None: + assert TASK_REGISTRY["inpainting"].postprocess is _postprocess_inpainting diff --git a/tests/unit/loader/test_onnx_hub.py b/tests/unit/loader/test_onnx_hub.py index aae4c8bb5..9a25955b2 100644 --- a/tests/unit/loader/test_onnx_hub.py +++ b/tests/unit/loader/test_onnx_hub.py @@ -22,6 +22,7 @@ from winml.modelkit.loader.onnx_hub import ( _split_hf_onnx_path, resolve_hf_onnx_path, + resolve_hf_repo_onnx, ) @@ -183,9 +184,7 @@ def _fake_download(*, repo_id, filename, revision, cache_dir, token): ) as mock_list, pytest.raises(FileNotFoundError) as exc_info, ): - resolve_hf_onnx_path( - "onnx-community/sam3-tracker-ONNX/onnx/vision_encoder.onnx" - ) + resolve_hf_onnx_path("onnx-community/sam3-tracker-ONNX/onnx/vision_encoder.onnx") msg = str(exc_info.value) # Names the bad path and the repo @@ -202,8 +201,7 @@ def _fake_download(*, repo_id, filename, revision, cache_dir, token): mock_list.assert_called_once() assert ( mock_list.call_args.args[0] == "onnx-community/sam3-tracker-ONNX" - or mock_list.call_args.kwargs.get("repo_id") - == "onnx-community/sam3-tracker-ONNX" + or mock_list.call_args.kwargs.get("repo_id") == "onnx-community/sam3-tracker-ONNX" ) def test_missing_file_listing_failure_falls_back_gracefully(self) -> None: @@ -254,3 +252,81 @@ def _fake_download(*, repo_id, filename, revision, cache_dir, token): msg = str(exc_info.value) assert "No .onnx files were found" in msg assert "org/pytorch-only" in msg + + +class TestResolveHfRepoOnnx: + """Bare Hub repositories resolve only for an unambiguous ONNX sibling.""" + + def test_single_onnx_uses_resolved_commit(self, tmp_path: Path) -> None: + downloaded = tmp_path / "model.onnx" + downloaded.write_bytes(b"onnx") + info = type( + "Info", + (), + { + "config": {}, + "sha": "abc123", + "siblings": [type("Sibling", (), {"rfilename": "weights/model.onnx"})()], + }, + )() + + with ( + patch("huggingface_hub.HfApi.model_info", return_value=info), + patch( + "winml.modelkit.loader.onnx_hub.resolve_hf_onnx_path", + return_value=downloaded, + ) as mock_download, + ): + result = resolve_hf_repo_onnx("org/repo", revision="release") + + assert result is not None + assert result.local_path == downloaded + assert result.filename == "weights/model.onnx" + assert result.revision == "abc123" + mock_download.assert_called_once_with( + "org/repo/weights/model.onnx", + revision="abc123", + cache_dir=None, + token=None, + ) + + def test_transformers_repository_is_not_reclassified(self) -> None: + info = type( + "Info", + (), + { + "config": {"model_type": "bert"}, + "sha": "abc123", + "siblings": [type("Sibling", (), {"rfilename": "model.onnx"})()], + }, + )() + with patch("huggingface_hub.HfApi.model_info", return_value=info): + assert resolve_hf_repo_onnx("org/bert") is None + + def test_http_error_falls_back_to_existing_hf_loader(self) -> None: + """Optional discovery must not replace established HF loader errors.""" + from huggingface_hub.errors import HfHubHTTPError + from requests import Response + + response = Response() + response.status_code = 401 + response.url = "https://huggingface.co/api/models/org/private" + with patch( + "huggingface_hub.HfApi.model_info", + side_effect=HfHubHTTPError("unauthorized", response=response), + ): + assert resolve_hf_repo_onnx("org/private") is None + + def test_multiple_onnx_requires_explicit_path(self) -> None: + siblings = [ + type("Sibling", (), {"rfilename": name})() for name in ("a.onnx", "nested/b.onnx") + ] + info = type("Info", (), {"config": {}, "sha": "abc123", "siblings": siblings})() + with ( + patch("huggingface_hub.HfApi.model_info", return_value=info), + pytest.raises(ValueError, match="multiple ONNX files") as exc_info, + ): + resolve_hf_repo_onnx("org/ambiguous") + + assert "org/ambiguous/a.onnx" in str(exc_info.value) + assert "org/ambiguous/nested/b.onnx" in str(exc_info.value) diff --git a/tests/unit/utils/test_model_input.py b/tests/unit/utils/test_model_input.py index 8dfb35342..abdd6940f 100644 --- a/tests/unit/utils/test_model_input.py +++ b/tests/unit/utils/test_model_input.py @@ -13,7 +13,7 @@ from __future__ import annotations from typing import TYPE_CHECKING -from unittest.mock import patch +from unittest.mock import MagicMock, patch from winml.modelkit.utils.model_input import ( ModelInputKind, @@ -294,3 +294,24 @@ def _fake_download(*, repo_id, filename, revision, cache_dir, token): assert mi.kind == "hub_onnx" assert mi.local_path == str(downloaded) assert mi.hf_id == "onnx-community/sam3-tracker-ONNX" + + def test_bare_repo_discovery_preserves_provenance(self, tmp_path: Path) -> None: + downloaded = tmp_path / "model.onnx" + downloaded.write_bytes(b"") + resolved = MagicMock( + local_path=downloaded, + repo_id="org/repo", + filename="nested/model.onnx", + revision="abc123", + ) + with patch( + "winml.modelkit.loader.onnx_hub.resolve_hf_repo_onnx", + return_value=resolved, + ): + mi = resolve_model_input("org/repo", discover_repo_onnx=True) + + assert mi.kind is ModelInputKind.HUB_ONNX + assert mi.local_path == str(downloaded) + assert mi.hf_id == "org/repo" + assert mi.artifact_path == "nested/model.onnx" + assert mi.revision == "abc123" From 28815d61b548f3b255fa6be169c69e62020676a0 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 13:42:55 +0800 Subject: [PATCH 2/4] fix(inference): persist inpainting runtime contracts --- .../cpu/cpu/inpainting_fp16_config.json | 13 ++ .../cpu/cpu/inpainting_fp32_config.json | 13 ++ src/winml/modelkit/build/hf.py | 1 + src/winml/modelkit/build/onnx.py | 2 + src/winml/modelkit/commands/build.py | 6 +- src/winml/modelkit/config/__init__.py | 2 + src/winml/modelkit/config/build.py | 46 +++++ src/winml/modelkit/inference/engine.py | 20 ++- src/winml/modelkit/inference/inpainting.py | 170 +++++++++++++----- src/winml/modelkit/inference/pipeline.py | 10 +- src/winml/modelkit/loader/onnx_hub.py | 6 +- src/winml/modelkit/models/auto.py | 11 +- src/winml/modelkit/models/winml/base.py | 6 + src/winml/modelkit/utils/manifest.py | 1 + tests/unit/config/test_build_onnx.py | 17 ++ tests/unit/inference/test_engine.py | 53 ++++-- .../inference/test_inpainting_pipeline.py | 86 ++++++++- tests/unit/loader/test_onnx_hub.py | 20 +++ tests/unit/models/auto/test_auto_onnx.py | 33 ++++ tests/unit/utils/test_manifest.py | 2 + tests/unit/utils/test_model_input.py | 14 ++ 21 files changed, 442 insertions(+), 90 deletions(-) diff --git a/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json index ff7203f5d..cf944853f 100644 --- a/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json +++ b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp16_config.json @@ -25,5 +25,18 @@ "compile": null, "loader": { "task": "inpainting" + }, + "runtime": { + "pipeline": "inpainting", + "options": { + "image_input_name": "image", + "mask_input_name": "mask", + "output_name": "output", + "image_color_order": "bgr", + "image_value_range": [0, 1], + "mask_semantics": "nonzero-is-hole", + "output_color_order": "bgr", + "output_value_range": [0, 255] + } } } diff --git a/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json index 309f0d8dc..8a1fb9cae 100644 --- a/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json +++ b/examples/recipes/opencv_inpainting_lama/cpu/cpu/inpainting_fp32_config.json @@ -6,5 +6,18 @@ "compile": null, "loader": { "task": "inpainting" + }, + "runtime": { + "pipeline": "inpainting", + "options": { + "image_input_name": "image", + "mask_input_name": "mask", + "output_name": "output", + "image_color_order": "bgr", + "image_value_range": [0, 1], + "mask_semantics": "nonzero-is-hole", + "output_color_order": "bgr", + "output_value_range": [0, 255] + } } } diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index c3d77933f..58064a5e3 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -399,6 +399,7 @@ def _name(base: str) -> str: source="hf", model_id=model_label, task=task, + runtime=config.runtime.to_dict() if config.runtime is not None else None, cache_key=cache_key, config_hash=cache_key.rsplit("_", 1)[-1] if cache_key and "_" in cache_key else None, timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), diff --git a/src/winml/modelkit/build/onnx.py b/src/winml/modelkit/build/onnx.py index 7bf225b74..76570c1ad 100644 --- a/src/winml/modelkit/build/onnx.py +++ b/src/winml/modelkit/build/onnx.py @@ -307,6 +307,8 @@ def _name(base: str) -> str: timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), elapsed_seconds=round(elapsed, 3), final_artifact=final_path.name, + task=config.loader.task, + runtime=config.runtime.to_dict() if config.runtime is not None else None, stages=manifest_stages, analyze_iterations=analyze_iters, analyze_unsupported_node_count=analyze_unsupported, diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index 6165d524f..897e42f99 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -945,8 +945,9 @@ def build( no_compile=no_compile, ) if model: - is_single_config = not isinstance(config_or_configs, list) - direct_onnx = is_single_config and config_or_configs.export is None + direct_onnx = ( + not isinstance(config_or_configs, list) and config_or_configs.export is None + ) source_model_input = resolve_model_input( model, discover_repo_onnx=direct_onnx, @@ -2237,6 +2238,7 @@ def _build_onnx_pipeline( source="onnx", model_id=source_model_id, task=config.loader.task, + runtime=config.runtime.to_dict() if config.runtime is not None else None, input_onnx=str(onnx_path), final_artifact=final_path.name, extras={ diff --git a/src/winml/modelkit/config/__init__.py b/src/winml/modelkit/config/__init__.py index cfa3b7b6a..1e1e430cb 100644 --- a/src/winml/modelkit/config/__init__.py +++ b/src/winml/modelkit/config/__init__.py @@ -26,6 +26,7 @@ from .build import ( SubmoduleClassNotFoundError, WinMLBuildConfig, + WinMLRuntimeConfig, generate_build_config, generate_hf_build_config, generate_onnx_build_config, @@ -48,6 +49,7 @@ "PrecisionPolicy", "SubmoduleClassNotFoundError", "WinMLBuildConfig", + "WinMLRuntimeConfig", "expand_precision", "extract_activation_bits", "extract_weight_bits", diff --git a/src/winml/modelkit/config/build.py b/src/winml/modelkit/config/build.py index 495804af6..873c41490 100644 --- a/src/winml/modelkit/config/build.py +++ b/src/winml/modelkit/config/build.py @@ -81,6 +81,7 @@ __all__ = [ "WinMLBuildConfig", + "WinMLRuntimeConfig", "generate_build_config", "generate_hf_build_config", "generate_onnx_build_config", @@ -95,6 +96,38 @@ # ============================================================================= +@dataclass +class WinMLRuntimeConfig: + """Runtime adapter selection and task-specific options. + + The build system persists this data unchanged into the artifact manifest. + Runtime adapters own validation of their option schema, which keeps this + plumbing reusable without introducing model- or task-specific branches. + """ + + pipeline: str + options: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> WinMLRuntimeConfig: + """Create runtime configuration from a recipe dictionary.""" + return cls( + pipeline=data.get("pipeline", ""), + options=dict(data.get("options", {})), + ) + + def to_dict(self) -> dict[str, Any]: + """Convert runtime configuration to a JSON-safe dictionary.""" + return {"pipeline": self.pipeline, "options": copy.deepcopy(self.options)} + + def validate(self) -> None: + """Validate the generic runtime envelope.""" + if not self.pipeline.strip(): + raise ValueError("runtime.pipeline must be a non-empty string") + if not isinstance(self.options, dict): + raise TypeError("runtime.options must be an object") + + @dataclass class WinMLBuildConfig: """Combined configuration for WinML model pipeline. @@ -137,6 +170,7 @@ class WinMLBuildConfig: quant: WinMLQuantizationConfig | None = field(default_factory=WinMLQuantizationConfig) compile: WinMLCompileConfig | None = field(default_factory=WinMLCompileConfig) eval: WinMLEvaluationConfig | None = None + runtime: WinMLRuntimeConfig | None = None auto: bool = True # Stamped True by generate_*_build_config (or by the build_*_model # entry-point defensive fallback) when the input ONNX is already @@ -164,6 +198,7 @@ def from_dict(cls, config_dict: dict) -> WinMLBuildConfig: quant_data = config_dict.get("quant") compile_data = config_dict.get("compile") eval_data = config_dict.get("eval") + runtime_data = config_dict.get("runtime") eval_cfg = None if eval_data is not None: eval_cfg = WinMLEvaluationConfig.from_dict(eval_data) @@ -178,6 +213,9 @@ def from_dict(cls, config_dict: dict) -> WinMLBuildConfig: WinMLCompileConfig.from_dict(compile_data) if compile_data is not None else None ), eval=eval_cfg, + runtime=( + WinMLRuntimeConfig.from_dict(runtime_data) if runtime_data is not None else None + ), auto=config_dict.get("auto", True), skip_optimize=config_dict.get("skip_optimize", False), ) @@ -203,6 +241,8 @@ def to_dict(self) -> dict: result["loader"] = loader_dict if self.eval is not None: result["eval"] = self.eval.to_dict() + if self.runtime is not None: + result["runtime"] = self.runtime.to_dict() return result def validate(self) -> None: @@ -254,6 +294,12 @@ def validate(self) -> None: ): errors.append("compile.ep_config.provider is required when compile is enabled") + if self.runtime is not None: + try: + self.runtime.validate() + except (TypeError, ValueError) as exc: + errors.append(str(exc)) + if errors: raise ValueError("Invalid WinMLBuildConfig:\n" + "\n".join(f" - {e}" for e in errors)) diff --git a/src/winml/modelkit/inference/engine.py b/src/winml/modelkit/inference/engine.py index dbc9546d7..161440fa5 100644 --- a/src/winml/modelkit/inference/engine.py +++ b/src/winml/modelkit/inference/engine.py @@ -323,8 +323,7 @@ def load( no effect on raw .onnx files or pre-built build directories (no build/analyze step runs in those paths). """ - # Hub-hosted ONNX (e.g. ``onnx-community/sam3-tracker-ONNX/onnx/...``) - # is downloaded once and treated as a local .onnx path thereafter. + # Hub-hosted ONNX is downloaded once and treated as a local path. from ..utils.model_input import resolve_model_input model_path = resolve_model_input( @@ -406,13 +405,15 @@ def load_schema_only( Falls back to ``load()`` only when the task cannot be determined without a full model load. """ - # Hub-hosted ONNX (e.g. ``onnx-community/sam3-tracker-ONNX/onnx/...``) - # is downloaded once and treated as a local .onnx path thereafter. - from ..utils.model_input import resolve_model_input + # An explicit task is sufficient to display its schema. Do not run + # optional bare-repository discovery here: a single-ONNX repository + # discovery downloads model weights, violating schema-only's contract. + if task is None: + from ..utils.model_input import resolve_model_input - model_path = resolve_model_input( - str(model_path), discover_repo_onnx=True - ).local_path or str(model_path) + model_path = resolve_model_input( + str(model_path), discover_repo_onnx=True + ).local_path or str(model_path) self._model_path = str(model_path) self._device = device @@ -971,7 +972,8 @@ def _load_from_build_dir( from ..models.winml import get_winml_class winml_class = get_winml_class(None, task or "") - self._model = winml_class(onnx_path=onnx_path, config=None, device=device) + self._model = winml_class(onnx_path=onnx_path, config=None, device=device, ep=ep) + self._model._runtime_config = manifest.get("runtime") if manifest is not None else None self._task = task or getattr(self._model, "task", None) if model_id: diff --git a/src/winml/modelkit/inference/inpainting.py b/src/winml/modelkit/inference/inpainting.py index 46a40d0b7..7dc5dd327 100644 --- a/src/winml/modelkit/inference/inpainting.py +++ b/src/winml/modelkit/inference/inpainting.py @@ -2,10 +2,11 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Pipeline adapter for direct-ONNX image-inpainting models.""" +"""Data-driven pipeline adapter for direct-ONNX image-inpainting models.""" from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, Any import numpy as np @@ -17,26 +18,113 @@ from ..models.winml.base import WinMLPreTrainedModel +@dataclass(frozen=True) +class _InpaintingContract: + image_input_name: str + mask_input_name: str + output_name: str + image_color_order: str + image_value_range: tuple[float, float] + mask_semantics: str + output_color_order: str + output_value_range: tuple[float, float] + + @classmethod + def from_runtime_config(cls, runtime_config: dict[str, Any] | None) -> _InpaintingContract: + """Validate and materialize the explicit runtime contract.""" + if not isinstance(runtime_config, dict) or runtime_config.get("pipeline") != "inpainting": + raise ValueError( + "Inpainting requires runtime.pipeline='inpainting' and explicit runtime.options." + ) + options = runtime_config.get("options") + if not isinstance(options, dict): + raise TypeError("Inpainting runtime.options must be an object.") + + required = { + "image_input_name", + "mask_input_name", + "output_name", + "image_color_order", + "image_value_range", + "mask_semantics", + "output_color_order", + "output_value_range", + } + missing = sorted(required - options.keys()) + unknown = sorted(options.keys() - required) + if missing or unknown: + details = [] + if missing: + details.append(f"missing: {', '.join(missing)}") + if unknown: + details.append(f"unknown: {', '.join(unknown)}") + raise ValueError(f"Invalid inpainting runtime options ({'; '.join(details)}).") + + def _color_order(name: str) -> str: + value = options[name] + if value not in {"rgb", "bgr"}: + raise ValueError(f"Inpainting {name} must be 'rgb' or 'bgr'.") + return str(value) + + def _value_range(name: str) -> tuple[float, float]: + value = options[name] + if not isinstance(value, list) or len(value) != 2: + raise ValueError(f"Inpainting {name} must be a two-number array.") + result = (float(value[0]), float(value[1])) + if result not in {(0.0, 1.0), (-1.0, 1.0), (0.0, 255.0)}: + raise ValueError( + f"Unsupported inpainting {name} {value}; supported ranges are " + "[0, 1], [-1, 1], and [0, 255]." + ) + return result + + mask_semantics = options["mask_semantics"] + if mask_semantics not in {"nonzero-is-hole", "zero-is-hole"}: + raise ValueError( + "Inpainting mask_semantics must be 'nonzero-is-hole' or 'zero-is-hole'." + ) + for name in ("image_input_name", "mask_input_name", "output_name"): + if not isinstance(options[name], str) or not options[name].strip(): + raise ValueError(f"Inpainting {name} must be a non-empty string.") + + return cls( + image_input_name=options["image_input_name"], + mask_input_name=options["mask_input_name"], + output_name=options["output_name"], + image_color_order=_color_order("image_color_order"), + image_value_range=_value_range("image_value_range"), + mask_semantics=mask_semantics, + output_color_order=_color_order("output_color_order"), + output_value_range=_value_range("output_value_range"), + ) + + class WinMLInpaintingPipeline: """Prepare an image/mask pair and postprocess a direct ONNX result. - The adapter derives spatial sizes and tensor roles from ONNX metadata. It - implements the OpenCV inpainting interchange contract: float32 NCHW BGR - image values in ``[0, 1]`` and a float32 NCHW binary mask. The first - three-channel output is converted back to an RGB PIL image. + Tensor roles, color order, value ranges, mask polarity, and output semantics + come from a checked build/runtime contract. ONNX shape metadata is used only + to validate that explicit contract and determine spatial dimensions. """ - def __init__(self, model: WinMLPreTrainedModel) -> None: + def __init__( + self, + model: WinMLPreTrainedModel, + *, + runtime_config: dict[str, Any] | None, + ) -> None: self.model = model - self._inputs, self._output_name = self._resolve_contract(model.io_config) + self._contract = _InpaintingContract.from_runtime_config(runtime_config) + self._inputs, self._output_name = self._resolve_io(model.io_config, self._contract) def _sanitize_parameters(self, **_kwargs: Any) -> tuple[dict, dict, dict]: """Expose the standard pipeline parameter-discovery contract.""" return {}, {}, {} @staticmethod - def _resolve_contract( + def _resolve_io( io_config: dict[str, Any], + contract: _InpaintingContract, ) -> tuple[dict[str, tuple[str, list[Any]]], str]: names = list(io_config.get("input_names", [])) shapes = list(io_config.get("input_shapes", [])) @@ -48,44 +136,27 @@ def _resolve_contract( except ValueError as exc: raise ValueError("Inpainting ONNX I/O metadata is incomplete.") from exc - def _pick(preferred_name: str, channels: int) -> tuple[str, list[Any]] | None: - by_name = next((entry for entry in entries if entry[0].lower() == preferred_name), None) - if by_name is not None and len(by_name[1]) == 4 and by_name[1][1] == channels: - return by_name - return next( - (entry for entry in entries if len(entry[1]) == 4 and entry[1][1] == channels), - None, - ) - - image = _pick("image", 3) - mask = _pick("mask", 1) + image = next((entry for entry in entries if entry[0] == contract.image_input_name), None) + mask = next((entry for entry in entries if entry[0] == contract.mask_input_name), None) if image is None or mask is None or image[0] == mask[0]: raise ValueError( - "Inpainting requires distinct 4-D image (3-channel) and mask " - "(1-channel) ONNX inputs." + "The inpainting runtime contract names inputs that are absent or not distinct." + ) + if len(image[1]) != 4 or image[1][1] != 3 or len(mask[1]) != 4 or mask[1][1] != 1: + raise ValueError( + "Inpainting contract inputs must identify a 4-D 3-channel image and " + "a 4-D 1-channel mask." ) image_size = WinMLInpaintingPipeline._spatial_size(image[1], "image") mask_size = WinMLInpaintingPipeline._spatial_size(mask[1], "mask") if image_size != mask_size: raise ValueError("Inpainting image and mask inputs must have the same spatial size.") - output = next( - ( - entry - for entry in outputs - if entry[0].lower() in {"output", "image", "images"} - and len(entry[1]) == 4 - and entry[1][1] == 3 - ), - None, - ) - if output is None: - output = next( - (entry for entry in outputs if len(entry[1]) == 4 and entry[1][1] == 3), - None, + output = next((entry for entry in outputs if entry[0] == contract.output_name), None) + if output is None or len(output[1]) != 4 or output[1][1] != 3: + raise ValueError( + "The inpainting runtime contract must identify a 4-D 3-channel ONNX output." ) - if output is None: - raise ValueError("Inpainting requires a 4-D 3-channel ONNX image output.") output_size = WinMLInpaintingPipeline._spatial_size(output[1], "output") if output_size != image_size: raise ValueError("Inpainting output must match the image input spatial size.") @@ -98,20 +169,21 @@ def _spatial_size(shape: list[Any], role: str) -> tuple[int, int]: raise ValueError(f"Inpainting {role} input requires fixed NCHW height and width.") return shape[3], shape[2] - @staticmethod - def _prepare_image(image: Image.Image, size: tuple[int, int]) -> np.ndarray: + def _prepare_image(self, image: Image.Image, size: tuple[int, int]) -> np.ndarray: rgb = np.asarray(image.convert("RGB").resize(size, Image.Resampling.BILINEAR)) - bgr = rgb[..., ::-1].astype(np.float32) / 255.0 - return np.ascontiguousarray(bgr.transpose(2, 0, 1)[None, ...]) + pixels = rgb[..., ::-1] if self._contract.image_color_order == "bgr" else rgb + low, high = self._contract.image_value_range + pixels = pixels.astype(np.float32) / 255.0 * (high - low) + low + return np.ascontiguousarray(pixels.transpose(2, 0, 1)[None, ...]) - @staticmethod - def _prepare_mask(mask: Image.Image, size: tuple[int, int]) -> np.ndarray: + def _prepare_mask(self, mask: Image.Image, size: tuple[int, int]) -> np.ndarray: grayscale = np.asarray(mask.convert("L").resize(size, Image.Resampling.NEAREST)) binary = (grayscale > 0).astype(np.float32) + if self._contract.mask_semantics == "zero-is-hole": + binary = 1.0 - binary return binary[None, None, ...] - @staticmethod - def _to_image(output: Any) -> Image.Image: + def _to_image(self, output: Any) -> Image.Image: if isinstance(output, torch.Tensor): array = output.detach().cpu().numpy() else: @@ -121,9 +193,11 @@ def _to_image(output: Any) -> Image.Image: f"Inpainting output must have shape [1, 3, height, width], got {list(array.shape)}." ) image = array[0].transpose(1, 2, 0) - if image.size and float(np.nanmax(image)) <= 1.0: - image = image * 255.0 - rgb = np.clip(image[..., ::-1], 0, 255).astype(np.uint8) + low, high = self._contract.output_value_range + image = (image - low) / (high - low) * 255.0 + if self._contract.output_color_order == "bgr": + image = image[..., ::-1] + rgb = np.clip(image, 0, 255).astype(np.uint8) return Image.fromarray(rgb, mode="RGB") def __call__(self, inputs: dict[str, Image.Image], **_kwargs: Any) -> Image.Image: diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index 3e219b312..41e46c99f 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -38,11 +38,17 @@ } -def _create_inpainting_pipeline(model: WinMLPreTrainedModel) -> Any: +def _create_inpainting_pipeline( + model: WinMLPreTrainedModel | WinMLCompositeModel, +) -> Any: """Create the built-in direct-ONNX inpainting adapter.""" + from ..models.winml.base import WinMLPreTrainedModel from .inpainting import WinMLInpaintingPipeline - return WinMLInpaintingPipeline(model) + if not isinstance(model, WinMLPreTrainedModel): + raise TypeError("The inpainting runtime pipeline requires a single ONNX model.") + + return WinMLInpaintingPipeline(model, runtime_config=model.runtime_config) _CUSTOM_PIPELINE_FACTORIES = { diff --git a/src/winml/modelkit/loader/onnx_hub.py b/src/winml/modelkit/loader/onnx_hub.py index 6e9324299..cfd0dca0d 100644 --- a/src/winml/modelkit/loader/onnx_hub.py +++ b/src/winml/modelkit/loader/onnx_hub.py @@ -56,7 +56,9 @@ def resolve_hf_repo_onnx( revision from changing between discovery and download. """ from huggingface_hub import HfApi - from huggingface_hub.errors import HfHubHTTPError + from huggingface_hub.errors import HfHubHTTPError, OfflineModeIsEnabled + from requests.exceptions import ConnectionError as RequestsConnectionError + from requests.exceptions import Timeout as RequestsTimeout try: info = HfApi().model_info( @@ -65,7 +67,7 @@ def resolve_hf_repo_onnx( files_metadata=False, token=token, ) - except HfHubHTTPError as exc: + except (HfHubHTTPError, OfflineModeIsEnabled, RequestsConnectionError, RequestsTimeout) as exc: # Discovery is an optional preflight before the existing Transformers # path. Auth, gated-repo, missing-repo, and transient Hub HTTP errors # must therefore fall through so the authoritative downstream loader diff --git a/src/winml/modelkit/models/auto.py b/src/winml/modelkit/models/auto.py index 7a071de89..c85f3504e 100644 --- a/src/winml/modelkit/models/auto.py +++ b/src/winml/modelkit/models/auto.py @@ -202,7 +202,7 @@ def from_onnx( logger.info("Skipping build (compiled model or explicit skip). Using original ONNX.") # TODO: run analyze_onnx for validation/lint winml_class = get_winml_class(None, resolved_task) - return winml_class( + model = winml_class( onnx_path=onnx_path, config=None, device=device, @@ -210,6 +210,9 @@ def from_onnx( ep=ep, provider_options=provider_options, ) + model._build_config = config + model._runtime_config = config.runtime.to_dict() if config.runtime is not None else None + return model # Resolve output directory and cache key task_abbrev = get_task_abbrev(resolved_task) if resolved_task else "onnx" @@ -247,7 +250,7 @@ def from_onnx( winml_class = get_winml_class(None, resolved_task) logger.info("Creating inference wrapper: %s", winml_class.__name__) - return winml_class( + model = winml_class( onnx_path=result.final_onnx_path, config=None, # No HF PretrainedConfig for bare ONNX builds device=device, @@ -255,6 +258,9 @@ def from_onnx( ep=ep, provider_options=provider_options, ) + model._build_config = config + model._runtime_config = config.runtime.to_dict() if config.runtime is not None else None + return model @classmethod def from_pretrained( @@ -528,6 +534,7 @@ def from_pretrained( provider_options=provider_options, ) model._build_config = config # resolved build config (task, quant, compile) + model._runtime_config = config.runtime.to_dict() if config.runtime is not None else None return model @classmethod diff --git a/src/winml/modelkit/models/winml/base.py b/src/winml/modelkit/models/winml/base.py index ea9172032..8924cafe1 100644 --- a/src/winml/modelkit/models/winml/base.py +++ b/src/winml/modelkit/models/winml/base.py @@ -89,6 +89,7 @@ def __init__( # Set by WinMLAutoModel.from_pretrained() after construction self._build_config: Any = None + self._runtime_config: dict[str, Any] | None = None # Create WinMLSession (delegates ORT operations) self._session = WinMLSession( @@ -258,6 +259,11 @@ def precision(self) -> str | None: """ return None + @property + def runtime_config(self) -> dict[str, Any] | None: + """Data-driven runtime adapter contract attached by build/load plumbing.""" + return self._runtime_config + @property def dtype(self) -> torch.dtype: """Model dtype (for HF compatibility).""" diff --git a/src/winml/modelkit/utils/manifest.py b/src/winml/modelkit/utils/manifest.py index c4dc4b475..71c76eaf0 100644 --- a/src/winml/modelkit/utils/manifest.py +++ b/src/winml/modelkit/utils/manifest.py @@ -118,6 +118,7 @@ class WinMLManifest: model_id: str | None = None task: str | None = None + runtime: dict[str, Any] | None = None cache_key: str | None = None config_hash: str | None = None input_onnx: str | None = None diff --git a/tests/unit/config/test_build_onnx.py b/tests/unit/config/test_build_onnx.py index 6aff78471..5bd23d6d9 100644 --- a/tests/unit/config/test_build_onnx.py +++ b/tests/unit/config/test_build_onnx.py @@ -315,6 +315,23 @@ def test_compiled_onnx_with_device_npu(self, tmp_path) -> None: # Config structure invariants # ----------------------------------------------------------------- + def test_runtime_contract_round_trips_in_onnx_override(self, tmp_path) -> None: + onnx_file = tmp_path / "model.onnx" + onnx_file.write_bytes(b"fake") + runtime = {"pipeline": "inpainting", "options": {"contract": "test"}} + + with ( + patch("winml.modelkit.onnx.is_compiled_onnx", return_value=True), + patch("winml.modelkit.onnx.is_quantized_onnx", return_value=False), + ): + config = generate_onnx_build_config( + onnx_file, + task="inpainting", + override={"runtime": runtime}, + ) + + assert config.to_dict()["runtime"] == runtime + def test_export_always_none(self, tmp_path) -> None: """All ONNX model states produce export=None.""" onnx_file = tmp_path / "model.onnx" diff --git a/tests/unit/inference/test_engine.py b/tests/unit/inference/test_engine.py index 68172e5d2..c1fd202ee 100644 --- a/tests/unit/inference/test_engine.py +++ b/tests/unit/inference/test_engine.py @@ -366,15 +366,8 @@ def test_task_param_overrides_manifest(self, tmp_path: Any) -> None: engine.load_schema_only(tmp_path, task="image-classification") assert engine._task == "image-classification" - def test_hub_onnx_ref_is_resolved_before_routing(self, tmp_path: Any) -> None: - """A Hub-style ONNX ref (``//.onnx``) must be - resolved to a local path BEFORE the .onnx-suffix-and-exists check, - otherwise it falls through to the HF model-id branch and tries to - load a Hub-ONNX path string as if it were a transformers config. - - Regression test for ``winml run`` and ``winml serve`` on Hub refs - like ``onnx-community/sam3-tracker-ONNX/onnx/...``. - """ + def test_explicit_task_does_not_download_hub_onnx_ref(self, tmp_path: Any) -> None: + """Schema-only display must not download an explicit Hub ONNX artifact.""" from unittest.mock import patch local = tmp_path / "vision_encoder_int8.onnx" @@ -382,23 +375,45 @@ def test_hub_onnx_ref_is_resolved_before_routing(self, tmp_path: Any) -> None: hub_ref = "onnx-community/sam3-tracker-ONNX/onnx/vision_encoder_int8.onnx" engine = InferenceEngine() - # ``WinMLSession.load_schema_only`` now routes Hub-ONNX resolution - # through the unified ``resolve_model_input`` (in - # ``winml.modelkit.utils.model_input``). Patch the underlying - # downloader so the lazy ``from ..loader.onnx_hub import - # resolve_hf_onnx_path`` picks up the mock at call time. with patch( "winml.modelkit.loader.onnx_hub.resolve_hf_onnx_path", return_value=local, ) as mock_resolve: engine.load_schema_only(hub_ref, task="mask-generation") - mock_resolve.assert_called_once() - # After resolution the engine should treat the input as a local - # ONNX file (not as an HF model id), which means _model_id is the - # resolved local path string, not the original Hub ref. - assert engine._model_id == str(local) + mock_resolve.assert_not_called() + assert engine._model_id == hub_ref assert engine._task == "mask-generation" + def test_explicit_task_does_not_download_bare_repo_onnx(self) -> None: + """Schema-only needs only the explicit task, never repository weights.""" + engine = InferenceEngine() + with patch("winml.modelkit.loader.onnx_hub.resolve_hf_onnx_path") as mock_download: + engine.load_schema_only("opencv/inpainting_lama", task="inpainting") + + mock_download.assert_not_called() + assert engine._model_id == "opencv/inpainting_lama" + assert engine._task == "inpainting" + + def test_build_dir_preserves_runtime_contract(self, tmp_path: Any) -> None: + """Build-directory loading attaches the manifest's runtime contract.""" + import json + + runtime = {"pipeline": "inpainting", "options": {"contract": "test"}} + (tmp_path / "build_manifest.json").write_text( + json.dumps({"task": "inpainting", "runtime": runtime}) + ) + (tmp_path / "model.onnx").write_bytes(b"fake") + fake_model = MagicMock() + + with patch( + "winml.modelkit.models.winml.get_winml_class", + return_value=lambda **_kwargs: fake_model, + ): + engine = InferenceEngine() + engine._load_from_build_dir(tmp_path, task=None, device="cpu", ep=None) + + assert fake_model._runtime_config == runtime + # --------------------------------------------------------------------------- # _sanitize_numpy diff --git a/tests/unit/inference/test_inpainting_pipeline.py b/tests/unit/inference/test_inpainting_pipeline.py index e1ef202e6..6193c45cf 100644 --- a/tests/unit/inference/test_inpainting_pipeline.py +++ b/tests/unit/inference/test_inpainting_pipeline.py @@ -16,6 +16,25 @@ from winml.modelkit.inference.inpainting import WinMLInpaintingPipeline +_LAMA_RUNTIME = { + "pipeline": "inpainting", + "options": { + "image_input_name": "source_pixels", + "mask_input_name": "edit_region", + "output_name": "completed_image", + "image_color_order": "bgr", + "image_value_range": [0, 1], + "mask_semantics": "nonzero-is-hole", + "output_color_order": "bgr", + "output_value_range": [0, 255], + }, +} + + +def _make_pipeline(model: _EchoInpaintingModel) -> WinMLInpaintingPipeline: + return WinMLInpaintingPipeline(model, runtime_config=_LAMA_RUNTIME) # type: ignore[arg-type] + + class _EchoInpaintingModel: io_config: ClassVar[dict] = { "input_names": ["source_pixels", "edit_region"], @@ -35,7 +54,7 @@ def __call__(self, **kwargs): class TestWinMLInpaintingPipeline: def test_prepares_bgr_image_and_binary_mask(self) -> None: model = _EchoInpaintingModel() - pipeline = WinMLInpaintingPipeline(model) # type: ignore[arg-type] + pipeline = _make_pipeline(model) pixels = np.full((2, 2, 3), [10, 20, 30], dtype=np.uint8) mask = np.array([[0, 1], [127, 255]], dtype=np.uint8) @@ -65,8 +84,8 @@ def test_requires_image_and_mask_tensor_roles(self) -> None: "output_names": ["output"], "output_shapes": [[1, 3, 2, 2]], } - with pytest.raises(ValueError, match=r"image .* and mask"): - WinMLInpaintingPipeline(model) # type: ignore[arg-type] + with pytest.raises(ValueError, match="absent or not distinct"): + _make_pipeline(model) def test_requires_matching_image_output(self) -> None: model = _EchoInpaintingModel() @@ -75,8 +94,8 @@ def test_requires_matching_image_output(self) -> None: "output_names": ["embedding"], "output_shapes": [[1, 128]], } - with pytest.raises(ValueError, match="3-channel ONNX image output"): - WinMLInpaintingPipeline(model) # type: ignore[arg-type] + with pytest.raises(ValueError, match="3-channel ONNX output"): + _make_pipeline(model) def test_selects_named_image_output(self) -> None: model = _EchoInpaintingModel() @@ -85,5 +104,60 @@ def test_selects_named_image_output(self) -> None: "output_names": ["scores", "completed_image"], "output_shapes": [[1, 10], [1, 3, 2, 2]], } - pipeline = WinMLInpaintingPipeline(model) # type: ignore[arg-type] + pipeline = _make_pipeline(model) assert pipeline._output_name == "completed_image" + + def test_requires_explicit_runtime_contract(self) -> None: + with pytest.raises(ValueError, match=r"explicit runtime\.options"): + WinMLInpaintingPipeline( # type: ignore[arg-type] + _EchoInpaintingModel(), runtime_config=None + ) + + def test_rejects_unsupported_contract_instead_of_guessing(self) -> None: + runtime = { + **_LAMA_RUNTIME, + "options": {**_LAMA_RUNTIME["options"], "image_value_range": [0, 2]}, + } + with pytest.raises(ValueError, match="Unsupported inpainting image_value_range"): + WinMLInpaintingPipeline( # type: ignore[arg-type] + _EchoInpaintingModel(), runtime_config=runtime + ) + + def test_applies_declared_non_lama_rgb_range_and_mask_contract(self) -> None: + class _AlternateContractModel(_EchoInpaintingModel): + def __call__(self, **kwargs): + self.inputs = kwargs + return {"completed_image": torch.from_numpy(kwargs["source_pixels"])} + + runtime = { + "pipeline": "inpainting", + "options": { + **_LAMA_RUNTIME["options"], + "image_color_order": "rgb", + "image_value_range": [-1, 1], + "mask_semantics": "zero-is-hole", + "output_color_order": "rgb", + "output_value_range": [-1, 1], + }, + } + model = _AlternateContractModel() + pipeline = WinMLInpaintingPipeline(model, runtime_config=runtime) # type: ignore[arg-type] + pixels = np.full((2, 2, 3), [10, 20, 30], dtype=np.uint8) + mask = np.array([[0, 255], [255, 0]], dtype=np.uint8) + + output = pipeline( + { + "image": Image.fromarray(pixels, mode="RGB"), + "mask": Image.fromarray(mask, mode="L"), + } + ) + + np.testing.assert_allclose( + model.inputs["source_pixels"][0, :, 0, 0], + np.array([10, 20, 30], dtype=np.float32) / 255.0 * 2.0 - 1.0, + ) + np.testing.assert_array_equal( + model.inputs["edit_region"], + np.array([[[[1, 0], [0, 1]]]], dtype=np.float32), + ) + np.testing.assert_allclose(np.asarray(output), pixels, atol=1) diff --git a/tests/unit/loader/test_onnx_hub.py b/tests/unit/loader/test_onnx_hub.py index 9a25955b2..605416048 100644 --- a/tests/unit/loader/test_onnx_hub.py +++ b/tests/unit/loader/test_onnx_hub.py @@ -317,6 +317,26 @@ def test_http_error_falls_back_to_existing_hf_loader(self) -> None: ): assert resolve_hf_repo_onnx("org/private") is None + def test_offline_mode_falls_back_to_existing_hf_loader(self) -> None: + """Optional discovery preserves the established cached/offline HF path.""" + from huggingface_hub.errors import OfflineModeIsEnabled + + with patch( + "huggingface_hub.HfApi.model_info", + side_effect=OfflineModeIsEnabled("offline"), + ): + assert resolve_hf_repo_onnx("microsoft/resnet-50") is None + + def test_connection_failure_falls_back_to_existing_hf_loader(self) -> None: + """A connectivity-only preflight failure is deferred to the normal loader.""" + from requests.exceptions import ConnectionError + + with patch( + "huggingface_hub.HfApi.model_info", + side_effect=ConnectionError("network unavailable"), + ): + assert resolve_hf_repo_onnx("microsoft/resnet-50") is None + def test_multiple_onnx_requires_explicit_path(self) -> None: siblings = [ type("Sibling", (), {"rfilename": name})() for name in ("a.onnx", "nested/b.onnx") diff --git a/tests/unit/models/auto/test_auto_onnx.py b/tests/unit/models/auto/test_auto_onnx.py index ed360fe4d..7041a6c8b 100644 --- a/tests/unit/models/auto/test_auto_onnx.py +++ b/tests/unit/models/auto/test_auto_onnx.py @@ -198,6 +198,39 @@ def test_returns_winml_pretrained_model(self, fake_onnx: Path, tmp_path: Path): assert result is mock_instance + def test_programmatic_onnx_preserves_runtime_contract( + self, fake_onnx: Path, tmp_path: Path + ) -> None: + """An explicit config reaches the returned runtime wrapper unchanged.""" + from winml.modelkit.config import WinMLBuildConfig + + runtime = {"pipeline": "inpainting", "options": {"contract": "test"}} + explicit_config = WinMLBuildConfig.from_dict( + { + "loader": {"task": "inpainting"}, + "export": None, + "quant": None, + "compile": None, + "runtime": runtime, + } + ) + with ( + patch("winml.modelkit.onnx.is_compiled_onnx", return_value=True), + patch("winml.modelkit.onnx.is_quantized_onnx", return_value=False), + patch("winml.modelkit.models.auto.get_winml_class") as mock_get_class, + ): + model = MagicMock() + mock_get_class.return_value = lambda **_kwargs: model + result = WinMLAutoModel.from_onnx( + fake_onnx, + task="inpainting", + config=explicit_config, + skip_build=True, + ) + + assert result is model + assert model._runtime_config == runtime + class TestFromPretrainedDelegatesToFromOnnx: """Test that from_pretrained delegates .onnx files to from_onnx.""" diff --git a/tests/unit/utils/test_manifest.py b/tests/unit/utils/test_manifest.py index fc5d443c0..02aeaacd9 100644 --- a/tests/unit/utils/test_manifest.py +++ b/tests/unit/utils/test_manifest.py @@ -74,6 +74,7 @@ def test_round_trip(self) -> None: source="hf", model_id="microsoft/resnet-50", task="image-classification", + runtime={"pipeline": "custom", "options": {"threshold": 0.5}}, final_artifact="model.onnx", elapsed_seconds=12.345, stages=[ @@ -91,6 +92,7 @@ def test_round_trip(self) -> None: assert restored.source == original.source assert restored.model_id == original.model_id assert restored.task == original.task + assert restored.runtime == original.runtime assert restored.elapsed_seconds == original.elapsed_seconds assert len(restored.stages) == 2 assert restored.stages[0].name == "export" diff --git a/tests/unit/utils/test_model_input.py b/tests/unit/utils/test_model_input.py index abdd6940f..6c66d2890 100644 --- a/tests/unit/utils/test_model_input.py +++ b/tests/unit/utils/test_model_input.py @@ -315,3 +315,17 @@ def test_bare_repo_discovery_preserves_provenance(self, tmp_path: Path) -> None: assert mi.hf_id == "org/repo" assert mi.artifact_path == "nested/model.onnx" assert mi.revision == "abc123" + + def test_optional_bare_repo_discovery_falls_back_offline(self) -> None: + """Offline discovery leaves a normal HF ID available to cached loaders.""" + from huggingface_hub.errors import OfflineModeIsEnabled + + with patch( + "huggingface_hub.HfApi.model_info", + side_effect=OfflineModeIsEnabled("offline"), + ): + mi = resolve_model_input("microsoft/resnet-50", discover_repo_onnx=True) + + assert mi.kind is ModelInputKind.HF_ID + assert mi.hf_id == "microsoft/resnet-50" + assert mi.local_path is None From eca1e3b5896adedc2e93d7ed6016364399febe2f Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 15:14:21 +0800 Subject: [PATCH 3/4] fix(inference): clarify inpainting mask schema --- src/winml/modelkit/inference/tasks.py | 2 +- tests/unit/inference/test_tasks.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/winml/modelkit/inference/tasks.py b/src/winml/modelkit/inference/tasks.py index 3669e79de..cbc935dee 100644 --- a/src/winml/modelkit/inference/tasks.py +++ b/src/winml/modelkit/inference/tasks.py @@ -304,7 +304,7 @@ def _postprocess_sentence_similarity( name="mask", type="image", required=True, - description="Binary mask; non-zero pixels are replaced", + description="Binary mask; hole polarity is defined by the runtime contract", ), ], mapping=PipelineMapping(pipe_input=["image", "mask"]), diff --git a/tests/unit/inference/test_tasks.py b/tests/unit/inference/test_tasks.py index 37b44190f..6956258f2 100644 --- a/tests/unit/inference/test_tasks.py +++ b/tests/unit/inference/test_tasks.py @@ -187,5 +187,14 @@ def test_requires_image_and_mask(self) -> None: assert all(field.required for field in spec.user_inputs) assert spec.mapping.pipe_input == ["image", "mask"] + def test_mask_description_defers_polarity_to_runtime_contract(self) -> None: + mask = next( + field for field in TASK_REGISTRY["inpainting"].user_inputs if field.name == "mask" + ) + + assert "runtime contract" in mask.description + assert "non-zero pixels are replaced" not in mask.description + assert "zero pixels are replaced" not in mask.description + def test_has_png_postprocess(self) -> None: assert TASK_REGISTRY["inpainting"].postprocess is _postprocess_inpainting From f03add820f384ce725e83400d53d108ff43aad0a Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 16:12:18 +0800 Subject: [PATCH 4/4] fix(inference): define canonical inpainting masks --- src/winml/modelkit/inference/inpainting.py | 19 ++++++++-- src/winml/modelkit/inference/tasks.py | 2 +- .../inference/test_inpainting_pipeline.py | 37 +++++++++++++++++++ tests/unit/inference/test_tasks.py | 7 ++-- 4 files changed, 56 insertions(+), 9 deletions(-) diff --git a/src/winml/modelkit/inference/inpainting.py b/src/winml/modelkit/inference/inpainting.py index 7dc5dd327..a94e9465a 100644 --- a/src/winml/modelkit/inference/inpainting.py +++ b/src/winml/modelkit/inference/inpainting.py @@ -20,6 +20,13 @@ @dataclass(frozen=True) class _InpaintingContract: + """Model-side ONNX tensor contract for the inpainting adapter. + + Caller masks use the public canonical convention where nonzero pixels are + holes. ``mask_semantics`` declares how those same holes are encoded in the + model's ONNX mask input. + """ + image_input_name: str mask_input_name: str output_name: str @@ -81,7 +88,8 @@ def _value_range(name: str) -> tuple[float, float]: mask_semantics = options["mask_semantics"] if mask_semantics not in {"nonzero-is-hole", "zero-is-hole"}: raise ValueError( - "Inpainting mask_semantics must be 'nonzero-is-hole' or 'zero-is-hole'." + "Inpainting mask_semantics describes the model-side ONNX mask tensor and " + "must be 'nonzero-is-hole' or 'zero-is-hole'." ) for name in ("image_input_name", "mask_input_name", "output_name"): if not isinstance(options[name], str) or not options[name].strip(): @@ -102,9 +110,11 @@ def _value_range(name: str) -> tuple[float, float]: class WinMLInpaintingPipeline: """Prepare an image/mask pair and postprocess a direct ONNX result. - Tensor roles, color order, value ranges, mask polarity, and output semantics - come from a checked build/runtime contract. ONNX shape metadata is used only - to validate that explicit contract and determine spatial dimensions. + The public mask input always uses nonzero pixels for holes. Tensor roles, + color order, value ranges, model-side ONNX mask polarity, and output + semantics come from a checked build/runtime contract. ONNX shape metadata + is used only to validate that explicit contract and determine spatial + dimensions. """ def __init__( @@ -177,6 +187,7 @@ def _prepare_image(self, image: Image.Image, size: tuple[int, int]) -> np.ndarra return np.ascontiguousarray(pixels.transpose(2, 0, 1)[None, ...]) def _prepare_mask(self, mask: Image.Image, size: tuple[int, int]) -> np.ndarray: + """Convert a canonical nonzero-is-hole caller mask to the ONNX contract.""" grayscale = np.asarray(mask.convert("L").resize(size, Image.Resampling.NEAREST)) binary = (grayscale > 0).astype(np.float32) if self._contract.mask_semantics == "zero-is-hole": diff --git a/src/winml/modelkit/inference/tasks.py b/src/winml/modelkit/inference/tasks.py index cbc935dee..517fb845c 100644 --- a/src/winml/modelkit/inference/tasks.py +++ b/src/winml/modelkit/inference/tasks.py @@ -304,7 +304,7 @@ def _postprocess_sentence_similarity( name="mask", type="image", required=True, - description="Binary mask; hole polarity is defined by the runtime contract", + description="Binary mask; nonzero pixels identify holes (the replaced region)", ), ], mapping=PipelineMapping(pipe_input=["image", "mask"]), diff --git a/tests/unit/inference/test_inpainting_pipeline.py b/tests/unit/inference/test_inpainting_pipeline.py index 6193c45cf..bb3b1170e 100644 --- a/tests/unit/inference/test_inpainting_pipeline.py +++ b/tests/unit/inference/test_inpainting_pipeline.py @@ -76,6 +76,43 @@ def test_prepares_bgr_image_and_binary_mask(self) -> None: ) np.testing.assert_array_equal(np.asarray(output), pixels) + @pytest.mark.parametrize( + ("mask_semantics", "model_hole_value", "model_background_value"), + [ + pytest.param("nonzero-is-hole", 1.0, 0.0, id="model-nonzero-is-hole"), + pytest.param("zero-is-hole", 0.0, 1.0, id="model-zero-is-hole"), + ], + ) + def test_maps_canonical_caller_hole_to_model_mask_semantics( + self, + mask_semantics: str, + model_hole_value: float, + model_background_value: float, + ) -> None: + runtime = { + **_LAMA_RUNTIME, + "options": {**_LAMA_RUNTIME["options"], "mask_semantics": mask_semantics}, + } + model = _EchoInpaintingModel() + pipeline = WinMLInpaintingPipeline(model, runtime_config=runtime) # type: ignore[arg-type] + caller_mask = np.array([[0, 255], [0, 0]], dtype=np.uint8) + + pipeline( + { + "image": Image.new("RGB", (2, 2)), + "mask": Image.fromarray(caller_mask, mode="L"), + } + ) + + model_mask = model.inputs["edit_region"][0, 0] + assert model_mask[0, 1] == model_hole_value + np.testing.assert_array_equal( + model_mask[[0, 1, 1], [0, 0, 1]], + np.full(3, model_background_value, dtype=np.float32), + ) + model_holes = model_mask > 0 if mask_semantics == "nonzero-is-hole" else model_mask == 0 + np.testing.assert_array_equal(model_holes, caller_mask > 0) + def test_requires_image_and_mask_tensor_roles(self) -> None: model = _EchoInpaintingModel() model.io_config = { diff --git a/tests/unit/inference/test_tasks.py b/tests/unit/inference/test_tasks.py index 6956258f2..80f9eab22 100644 --- a/tests/unit/inference/test_tasks.py +++ b/tests/unit/inference/test_tasks.py @@ -187,14 +187,13 @@ def test_requires_image_and_mask(self) -> None: assert all(field.required for field in spec.user_inputs) assert spec.mapping.pipe_input == ["image", "mask"] - def test_mask_description_defers_polarity_to_runtime_contract(self) -> None: + def test_mask_description_defines_canonical_caller_polarity(self) -> None: mask = next( field for field in TASK_REGISTRY["inpainting"].user_inputs if field.name == "mask" ) - assert "runtime contract" in mask.description - assert "non-zero pixels are replaced" not in mask.description - assert "zero pixels are replaced" not in mask.description + assert "nonzero pixels identify holes" in mask.description + assert "replaced region" in mask.description def test_has_png_postprocess(self) -> None: assert TASK_REGISTRY["inpainting"].postprocess is _postprocess_inpainting