From da7edb440b2870c3791f0264c1d3bfe2bcfd4b27 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 18:37:01 +0800 Subject: [PATCH 1/3] Add BiRefNet image segmentation support --- .../cpu/image-segmentation_fp16_config.json | 65 +++++ .../cpu/image-segmentation_fp32_config.json | 43 ++++ pyproject.toml | 2 + src/winml/modelkit/export/htp/exporter.py | 15 ++ src/winml/modelkit/loader/resolution.py | 82 ++++++- src/winml/modelkit/models/hf/__init__.py | 1 + src/winml/modelkit/models/hf/birefnet.py | 224 ++++++++++++++++++ .../test_htp_exporter_patcher_model_kwargs.py | 38 +++ .../loader/test_remote_auto_map_resolution.py | 95 ++++++++ .../unit/models/birefnet/test_onnx_config.py | 114 +++++++++ 10 files changed, 671 insertions(+), 8 deletions(-) create mode 100644 examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp16_config.json create mode 100644 examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp32_config.json create mode 100644 src/winml/modelkit/models/hf/birefnet.py create mode 100644 tests/unit/loader/test_remote_auto_map_resolution.py create mode 100644 tests/unit/models/birefnet/test_onnx_config.py diff --git a/examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp16_config.json b/examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp16_config.json new file mode 100644 index 000000000..1d322fc19 --- /dev/null +++ b/examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp16_config.json @@ -0,0 +1,65 @@ +{ + "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": "x", + "dtype": "float32", + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "value_range": [ + -2.1179039301310043, + 2.64 + ] + } + ], + "output_tensors": [ + { + "name": "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, + "task": "image-segmentation", + "model_id": "ZhengPeng7/BiRefNet", + "model_type": "SegformerForSemanticSegmentation", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "image-segmentation", + "model_class": "AutoModelForImageSegmentation", + "model_type": "SegformerForSemanticSegmentation", + "trust_remote_code": true + } +} diff --git a/examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp32_config.json b/examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp32_config.json new file mode 100644 index 000000000..627e03d02 --- /dev/null +++ b/examples/recipes/ZhengPeng7_BiRefNet/cpu/cpu/image-segmentation_fp32_config.json @@ -0,0 +1,43 @@ +{ + "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": "x", + "dtype": "float32", + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "value_range": [ + -2.1179039301310043, + 2.64 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "image-segmentation", + "model_class": "AutoModelForImageSegmentation", + "model_type": "SegformerForSemanticSegmentation", + "trust_remote_code": true + } +} diff --git a/pyproject.toml b/pyproject.toml index 4ffda3c89..e738c90bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,11 +35,13 @@ dependencies = [ # winml.modelkit.__init__.) "pyarrow==23.0.1", "diffusers>=0.36", + "einops>=0.8", "evaluate>=0.4.6", "fastapi>=0.135.3", "hf_xet>=1.1.10", "httpx>=0.24.0", "jsonschema>=4.23", + "kornia>=0.8", "mcp>=0.1.0", "ml-dtypes>=0.5.3", "numpy>=2.2,<2.3", diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index c61a162d3..c88aca4cb 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -512,6 +512,21 @@ def _get_optimum_patcher(model: nn.Module, task: str | None) -> Any: model_config = getattr(model, "config", None) model_type = getattr(model_config, "model_type", None) if model_config else None + if not model_type: + # Trusted custom-code models occasionally overwrite the + # PreTrainedModel config with an internal runtime settings object. + # Recover the Hugging Face config from the architecture-declared + # config_class so its registered export patcher still applies. + config_class = getattr(model.__class__, "config_class", None) + if callable(config_class): + try: + model_config = config_class() + model_type = getattr(model_config, "model_type", None) + except (TypeError, ValueError): + logger.debug( + "Could not instantiate %s.config_class for export patch resolution.", + model.__class__.__name__, + ) if not model_type: logger.debug("Model has no config.model_type; skipping Optimum patcher.") return contextlib.nullcontext() diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 548848175..57d7e6390 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -150,6 +150,64 @@ def _detect_task_from_model_class(model_class: type) -> str: return cast("str", TasksManager.infer_task_from_model(model_class)) +def _resolve_remote_auto_model_class( + config: PretrainedConfig, + task: str | None = None, +) -> tuple[str, type] | None: + """Resolve a trusted custom-code auto class from config metadata. + + A custom checkpoint may name an architecture that is intentionally absent + from the installed ``transformers`` package. In that case, ``auto_map`` is + the authoritative loader declaration. Accept an entry only when its remote + class basename is also declared in ``architectures``; this avoids selecting + unrelated tokenizer/config mappings and keeps resolution model-ID agnostic. + + Args: + config: Hugging Face config carrying ``architectures`` and ``auto_map``. + task: Optional Optimum-canonical task used to filter matching auto classes. + + Returns: + A unique ``(task, auto_model_class)`` match, or ``None`` when metadata + does not declare one. + + Raises: + ValueError: If metadata declares multiple matching model auto classes. + """ + architectures = set(getattr(config, "architectures", None) or []) + auto_map = getattr(config, "auto_map", None) + if not architectures or not isinstance(auto_map, dict): + return None + + import transformers + + matches: dict[tuple[str, type], None] = {} + for auto_class_name, remote_reference in auto_map.items(): + if not auto_class_name.startswith("AutoModel") or not isinstance(remote_reference, str): + continue + if remote_reference.rsplit(".", 1)[-1] not in architectures: + continue + auto_class = getattr(transformers, auto_class_name, None) + if not isinstance(auto_class, type): + continue + try: + inferred_task = _detect_task_from_model_class(auto_class) + except (KeyError, ValueError): + continue + if task is None or inferred_task == task: + matches[(inferred_task, auto_class)] = None + + if not matches: + return None + if len(matches) > 1: + names = sorted(f"{matched_task}:{cls.__name__}" for matched_task, cls in matches) + raise ValueError( + "Custom-code metadata declares multiple matching model auto classes: " + + ", ".join(names) + + ". Please specify task and model_class explicitly." + ) + return next(iter(matches)) + + def _upgrade_fill_mask_for_seq2seq(task: str, config: PretrainedConfig) -> str: """Correct Optimum's ``fill-mask`` mislabel for encoder-decoder generation heads. @@ -550,11 +608,14 @@ def resolve_task( # `text2text-generation`. So `--task summarization` tags the composite while # `--task text2text-generation` stays composite=None (single-decoder export). composite = resolve_composite(model_type_norm, original) if model_type_norm else None - resolved = None + remote_resolution = _resolve_remote_auto_model_class(config, normalized) + resolved = remote_resolution[1] if remote_resolution is not None else None if model_type_norm: - resolved = _get_custom_model_class( - model_type_norm, original - ) or _get_custom_model_class(model_type_norm, normalized) + resolved = ( + resolved + or _get_custom_model_class(model_type_norm, original) + or _get_custom_model_class(model_type_norm, normalized) + ) if resolved is None: try: resolved = TasksManager.get_model_class_for_task(normalized, framework="pt") @@ -606,11 +667,16 @@ def resolve_task( # 1c. TasksManager (reads config.architectures) if opt_task is None: - try: - opt_task = _infer_task_from_architecture(config) + remote_resolution = _resolve_remote_auto_model_class(config) + if remote_resolution is not None: + opt_task, resolved = remote_resolution source = TaskSource.TASKS_MANAGER - except ValueError: - opt_task = None + else: + try: + opt_task = _infer_task_from_architecture(config) + source = TaskSource.TASKS_MANAGER + except ValueError: + opt_task = None # 1d. Hub pipeline_tag fallback if opt_task is None and model_id and model_type: diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 3c5b0c12b..07cb1e9f5 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -37,6 +37,7 @@ BartSequenceClassificationIOConfig as _BartSequenceClassificationIOConfig, ) from .bert import BERT_CONFIG +from .birefnet import BiRefNetIOConfig as _BiRefNetIOConfig # triggers registration from .blip import BLIP_CONFIG from .blip import MODEL_CLASS_MAPPING as _BLIP_CLASS_MAPPING from .blip import BlipCaptioningIOConfig as _BlipCaptioningIOConfig # triggers registration diff --git a/src/winml/modelkit/models/hf/birefnet.py b/src/winml/modelkit/models/hf/birefnet.py new file mode 100644 index 000000000..9296607d6 --- /dev/null +++ b/src/winml/modelkit/models/hf/birefnet.py @@ -0,0 +1,224 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""BiRefNet Hugging Face model configuration.""" + +from __future__ import annotations + +import sys +from typing import Any, cast + +import torch +import torch.nn.functional as F +from optimum.exporters.onnx import OnnxConfig +from optimum.exporters.onnx.model_patcher import ModelPatcher, PatchingSpec +from optimum.utils import NormalizedConfig +from optimum.utils.input_generators import DummyVisionInputGenerator + +from ...export import register_onnx_overwrite + + +def _pair(value: int | tuple[int, int]) -> tuple[int, int]: + """Normalize a scalar or two-dimensional convolution parameter.""" + return value if isinstance(value, tuple) else (value, value) + + +def _exportable_deform_conv2d( + input: torch.Tensor, + offset: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + stride: int | tuple[int, int] = 1, + padding: int | tuple[int, int] = 0, + dilation: int | tuple[int, int] = 1, + mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Pure-PyTorch equivalent of ``torchvision.ops.deform_conv2d``. + + The legacy ONNX exporter cannot lower ``torchvision::deform_conv2d``. + This implementation expresses the same sampling operation with + ``grid_sample`` followed by grouped linear projection, both of which have + standard ONNX representations at the configured opset. + """ + stride_h, stride_w = _pair(stride) + pad_h, pad_w = _pair(padding) + dilation_h, dilation_w = _pair(dilation) + batch_size = input.shape[0] + in_channels = int(input.shape[1]) + out_channels = int(weight.shape[0]) + channels_per_group = int(weight.shape[1]) + kernel_h = int(weight.shape[2]) + kernel_w = int(weight.shape[3]) + output_h, output_w = offset.shape[-2:] + kernel_points = kernel_h * kernel_w + offset_groups = int(offset.shape[1]) // (2 * kernel_points) + if in_channels % offset_groups != 0: + raise ValueError("Input channels must be divisible by deformable offset groups") + conv_groups = in_channels // channels_per_group + if out_channels % conv_groups != 0: + raise ValueError("Output channels must be divisible by convolution groups") + + padded = F.pad(input, (pad_w, pad_w, pad_h, pad_h)) + padded_h, padded_w = padded.shape[-2:] + offsets = offset.reshape( + batch_size, + offset_groups, + kernel_points, + 2, + output_h, + output_w, + ) + modulation = ( + mask.reshape(batch_size, offset_groups, kernel_points, output_h, output_w) + if mask is not None + else None + ) + channels_per_offset_group = in_channels // offset_groups + base_y = torch.arange(output_h, device=input.device, dtype=input.dtype) * stride_h + base_x = torch.arange(output_w, device=input.device, dtype=input.dtype) * stride_w + grid_y, grid_x = torch.meshgrid(base_y, base_x, indexing="ij") + + sampled_groups: list[torch.Tensor] = [] + for offset_group in range(offset_groups): + channel_start = offset_group * channels_per_offset_group + channel_end = channel_start + channels_per_offset_group + point_samples: list[torch.Tensor] = [] + for point in range(kernel_points): + kernel_y, kernel_x = divmod(point, kernel_w) + sample_y = grid_y + kernel_y * dilation_h + offsets[:, offset_group, point, 0] + sample_x = grid_x + kernel_x * dilation_w + offsets[:, offset_group, point, 1] + normalized_y = sample_y * (2.0 / (padded_h - 1)) - 1.0 + normalized_x = sample_x * (2.0 / (padded_w - 1)) - 1.0 + grid = torch.stack((normalized_x, normalized_y), dim=-1) + sampled = F.grid_sample( + padded[:, channel_start:channel_end], + grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + if modulation is not None: + sampled = sampled * modulation[:, offset_group, point].unsqueeze(1) + point_samples.append(sampled) + sampled_groups.append(torch.stack(point_samples, dim=2)) + sampled_input = torch.cat(sampled_groups, dim=1) + + outputs: list[torch.Tensor] = [] + output_channels_per_group = out_channels // conv_groups + flattened_weight = weight.flatten(2) + for conv_group in range(conv_groups): + input_start = conv_group * channels_per_group + input_end = input_start + channels_per_group + output_start = conv_group * output_channels_per_group + output_end = output_start + output_channels_per_group + outputs.append( + torch.einsum( + "nckhw,ock->nohw", + sampled_input[:, input_start:input_end], + flattened_weight[output_start:output_end], + ) + ) + result = torch.cat(outputs, dim=1) + if bias is not None: + result = result + bias.reshape(1, -1, 1, 1) + return result + + +class _BiRefNetModelPatcher(ModelPatcher): # type: ignore[misc] + """Patch the remote module's imported deformable-convolution function.""" + + def __init__(self, config: OnnxConfig, model: torch.nn.Module, **kwargs: Any) -> None: + remote_module = sys.modules[model.__class__.__module__] + original = getattr(remote_module, "deform_conv2d", None) + if original is None: + raise ValueError( + "BiRefNet remote module does not expose the expected deform_conv2d function" + ) + config.PATCHING_SPECS = [ + PatchingSpec( + o=remote_module, + name="deform_conv2d", + custom_op=_exportable_deform_conv2d, + orig_op=original, + ) + ] + super().__init__(config, model, **kwargs) + + +class _BiRefNetVisionInputGenerator(DummyVisionInputGenerator): # type: ignore[misc] + """Generate the ``x`` image tensor consumed by remote BiRefNet classes.""" + + SUPPORTED_INPUT_NAMES = ("x",) + _NORMALIZED_MIN = -2.1179039301310043 + _NORMALIZED_MAX = 2.64 + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + batch_size: int = 1, + num_channels: int = 3, + width: int = 1024, + height: int = 1024, + **kwargs: Any, + ) -> None: + super().__init__( + task, + normalized_config, + batch_size=batch_size, + num_channels=num_channels, + width=width, + height=height, + **kwargs, + ) + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ) -> torch.Tensor: + """Generate an ImageNet-normalized RGB tensor for ``forward(x)``.""" + del input_name, int_dtype + return cast( + "torch.Tensor", + self.random_float_tensor( + shape=[self.batch_size, self.num_channels, self.height, self.width], + min_value=self._NORMALIZED_MIN, + max_value=self._NORMALIZED_MAX, + framework=framework, + dtype=float_dtype, + ), + ) + + +@register_onnx_overwrite( + "SegformerForSemanticSegmentation", + "image-segmentation", + library_name="transformers", +) +class BiRefNetIOConfig(OnnxConfig): # type: ignore[misc] + """ONNX contract for custom-code BiRefNet image segmentation models. + + BiRefNet's remote ``forward(x)`` accepts a normalized RGB image and, in + evaluation mode, returns a one-element list containing full-resolution, + single-channel foreground logits. The checkpoint's own inference example + uses 1024 by 1024 images. + """ + + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(allow_new=True) + DUMMY_INPUT_GENERATOR_CLASSES = (_BiRefNetVisionInputGenerator,) + _MODEL_PATCHER = _BiRefNetModelPatcher + + @property + def inputs(self) -> dict[str, dict[int, str]]: + """Return the remote model's exact forward input contract.""" + return {"x": {0: "batch_size", 2: "height", 3: "width"}} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + """Return the flattened foreground-logit output contract.""" + return {"logits": {0: "batch_size", 2: "height", 3: "width"}} diff --git a/tests/unit/export/test_htp_exporter_patcher_model_kwargs.py b/tests/unit/export/test_htp_exporter_patcher_model_kwargs.py index dd3ab8156..db54a94e1 100644 --- a/tests/unit/export/test_htp_exporter_patcher_model_kwargs.py +++ b/tests/unit/export/test_htp_exporter_patcher_model_kwargs.py @@ -36,6 +36,22 @@ def __init__(self) -> None: self.config = _FakeConfig() +class _ArchitectureConfig: + """HF config declared by a trusted custom architecture.""" + + model_type = "custom-architecture" + + +class _OverwritingConfigModel(nn.Module): + """Custom model that replaces its HF config with runtime settings.""" + + config_class = _ArchitectureConfig + + def __init__(self) -> None: + super().__init__() + self.config = object() + + class TestGetOptimumPatcherModelKwargs: """_get_optimum_patcher must pass an explicit mutable model_kwargs dict.""" @@ -71,3 +87,25 @@ def fake_ctor(*args: object, **kwargs: object): f"to patch_model_for_export, got {captured.get('model_kwargs')!r}. " "MoE patchers (e.g. ViTPose dataset_index) need a mutable dict." ) + + def test_uses_architecture_config_class_when_instance_config_is_overwritten(self) -> None: + """A custom model's declared HF config still drives patcher lookup.""" + captured: dict[str, object] = {} + fake_onnx_config = MagicMock() + fake_onnx_config.patch_model_for_export.return_value = MagicMock() + + def fake_ctor(config: object): + captured["config"] = config + return fake_onnx_config + + with patch( + "optimum.exporters.tasks.TasksManager.get_exporter_config_constructor", + return_value=fake_ctor, + ) as constructor: + HTPExporter._get_optimum_patcher( + _OverwritingConfigModel(), + task="image-segmentation", + ) + + assert isinstance(captured["config"], _ArchitectureConfig) + assert constructor.call_args.kwargs["model_type"] == "custom-architecture" diff --git a/tests/unit/loader/test_remote_auto_map_resolution.py b/tests/unit/loader/test_remote_auto_map_resolution.py new file mode 100644 index 000000000..910a92a1f --- /dev/null +++ b/tests/unit/loader/test_remote_auto_map_resolution.py @@ -0,0 +1,95 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from transformers import AutoModelForImageSegmentation, PretrainedConfig + +from winml.modelkit.loader import load_hf_model +from winml.modelkit.loader.config import resolve_loader_config +from winml.modelkit.loader.resolution import ( + TaskSource, + _resolve_remote_auto_model_class, + resolve_task, +) + + +class RemoteImageSegmentationConfig(PretrainedConfig): + model_type = "custom-image-segmentation" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.architectures = ["RemoteSegmentationModel"] + self.auto_map = { + "AutoConfig": "configuration.RemoteConfig", + "AutoModelForImageSegmentation": "modeling.RemoteSegmentationModel", + } + + +def test_resolve_task_uses_matching_remote_auto_map_metadata() -> None: + resolution = resolve_task(RemoteImageSegmentationConfig()) + + assert resolution.task == "image-segmentation" + assert resolution.model_class is AutoModelForImageSegmentation + assert resolution.source == TaskSource.TASKS_MANAGER + + +def test_explicit_task_uses_matching_remote_auto_map_metadata() -> None: + resolution = resolve_task( + RemoteImageSegmentationConfig(), + task="image-segmentation", + ) + + assert resolution.model_class is AutoModelForImageSegmentation + assert resolution.source == TaskSource.USER_TASK + + +def test_remote_auto_map_requires_architecture_match() -> None: + config = RemoteImageSegmentationConfig() + config.architectures = ["DifferentModel"] + + assert _resolve_remote_auto_model_class(config) is None + + +def test_loader_config_preserves_remote_auto_class_and_trust() -> None: + loader, _, resolved_class, resolution = resolve_loader_config( + "org/custom-model", + trust_remote_code=True, + hf_config=RemoteImageSegmentationConfig(), + ) + + assert loader.task == "image-segmentation" + assert loader.model_class == "AutoModelForImageSegmentation" + assert loader.model_type == "custom-image-segmentation" + assert loader.trust_remote_code is True + assert resolved_class is AutoModelForImageSegmentation + assert resolution.source == TaskSource.TASKS_MANAGER + + +def test_load_hf_model_calls_remote_auto_class_with_trust() -> None: + config = RemoteImageSegmentationConfig() + model = MagicMock() + model.parameters.return_value = [] + with ( + patch("winml.modelkit.loader.hf.AutoConfig.from_pretrained", return_value=config), + patch.object( + AutoModelForImageSegmentation, + "from_pretrained", + return_value=model, + ) as from_pretrained, + ): + loaded, _, task = load_hf_model( + "org/custom-model", + trust_remote_code=True, + ) + + assert loaded is model + assert task == "image-segmentation" + from_pretrained.assert_called_once_with( + "org/custom-model", + trust_remote_code=True, + ) diff --git a/tests/unit/models/birefnet/test_onnx_config.py b/tests/unit/models/birefnet/test_onnx_config.py new file mode 100644 index 000000000..237643bd3 --- /dev/null +++ b/tests/unit/models/birefnet/test_onnx_config.py @@ -0,0 +1,114 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +from torchvision.ops import deform_conv2d +from transformers import PretrainedConfig + +import winml.modelkit.models # noqa: F401 +from winml.modelkit.export import generate_dummy_inputs +from winml.modelkit.export.io import _get_onnx_config +from winml.modelkit.models.hf.birefnet import BiRefNetIOConfig, _exportable_deform_conv2d + + +class BiRefNetTestConfig(PretrainedConfig): + model_type = "SegformerForSemanticSegmentation" + + +def test_birefnet_onnx_config_is_registered() -> None: + config = _get_onnx_config( + "SegformerForSemanticSegmentation", + "image-segmentation", + BiRefNetTestConfig(), + ) + + assert isinstance(config, BiRefNetIOConfig) + assert list(config.inputs) == ["x"] + assert list(config.outputs) == ["logits"] + assert config.outputs["logits"] == {0: "batch_size", 2: "height", 3: "width"} + + +def test_birefnet_dummy_input_uses_remote_forward_name_and_size() -> None: + inputs = generate_dummy_inputs( + "SegformerForSemanticSegmentation", + "image-segmentation", + BiRefNetTestConfig(), + ) + + assert list(inputs) == ["x"] + assert inputs["x"].shape == torch.Size([1, 3, 1024, 1024]) + assert inputs["x"].dtype == torch.float32 + assert inputs["x"].min() >= -2.1179039301310043 + assert inputs["x"].max() < 2.64 + + +def test_birefnet_dummy_input_honors_explicit_shape() -> None: + inputs = generate_dummy_inputs( + "SegformerForSemanticSegmentation", + "image-segmentation", + BiRefNetTestConfig(), + height=256, + width=384, + ) + + assert inputs["x"].shape == torch.Size([1, 3, 256, 384]) + + +def test_exportable_deform_conv_matches_torchvision() -> None: + torch.manual_seed(7) + input_tensor = torch.randn(1, 4, 7, 8) + weight = torch.randn(6, 4, 3, 3) + bias = torch.randn(6) + offset = torch.randn(1, 18, 7, 8) * 0.2 + mask = torch.sigmoid(torch.randn(1, 9, 7, 8)) + + expected = deform_conv2d( + input_tensor, + offset, + weight, + bias=bias, + padding=1, + mask=mask, + ) + actual = _exportable_deform_conv2d( + input_tensor, + offset, + weight, + bias=bias, + padding=1, + mask=mask, + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_exportable_deform_conv_supports_convolution_and_offset_groups() -> None: + torch.manual_seed(11) + input_tensor = torch.randn(1, 4, 6, 5) + weight = torch.randn(6, 2, 3, 3) + bias = torch.randn(6) + offset = torch.randn(1, 36, 6, 5) * 0.2 + mask = torch.sigmoid(torch.randn(1, 18, 6, 5)) + + expected = deform_conv2d( + input_tensor, + offset, + weight, + bias=bias, + padding=1, + mask=mask, + ) + actual = _exportable_deform_conv2d( + input_tensor, + offset, + weight, + bias=bias, + padding=1, + mask=mask, + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) From 64fa85bf96b4aca47c1be943a33e67c8cd8517b0 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 19:48:30 +0800 Subject: [PATCH 2/3] Fix BiRefNet dependency locking --- pyproject.toml | 14 +++++++---- uv.lock | 65 ++++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 61 insertions(+), 18 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e738c90bf..857f4ba09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,20 +35,18 @@ dependencies = [ # winml.modelkit.__init__.) "pyarrow==23.0.1", "diffusers>=0.36", - "einops>=0.8", "evaluate>=0.4.6", "fastapi>=0.135.3", "hf_xet>=1.1.10", "httpx>=0.24.0", "jsonschema>=4.23", - "kornia>=0.8", "mcp>=0.1.0", "ml-dtypes>=0.5.3", "numpy>=2.2,<2.3", "onnx==1.18", "onnx-ir>=0.1", - "onnxruntime-genai-winml==0.14.1", - "onnxruntime-windowsml==1.24.5.202604171637", + "onnxruntime-genai-winml==0.14.1; sys_platform == 'win32'", + "onnxruntime-windowsml==1.24.5.202604171637; sys_platform == 'win32'", "onnxscript>=0.2", "opentelemetry-sdk>=1.39.1", "optimum>=2", @@ -79,7 +77,7 @@ dependencies = [ "tqdm>=4.67", "transformers>=4.57", "uvicorn[standard]>=0.42.0", - "windowsml==2.0.300", + "windowsml==2.0.300; sys_platform == 'win32'", ] optional-dependencies.dev = [ @@ -102,6 +100,12 @@ optional-dependencies.dev = [ "types-colorama>=0.4.15.20250801", ] optional-dependencies.audio = [ "soundfile>=0.13" ] +# Dependencies imported by the ZhengPeng7/BiRefNet trusted remote code during +# export. They are not required by winml itself; install with `winml-cli[birefnet]`. +optional-dependencies.birefnet = [ + "einops>=0.8", + "kornia>=0.8", +] optional-dependencies.openvino = [ "openvino>=2023" ] optional-dependencies.qnn = [ "onnxruntime-qnn>=1.24.1; python_version>='3.11'", diff --git a/uv.lock b/uv.lock index 3dec4c5c5..8e763f46f 100644 --- a/uv.lock +++ b/uv.lock @@ -523,6 +523,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, ] +[[package]] +name = "einops" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, +] + [[package]] name = "evaluate" version = "0.4.6" @@ -1212,6 +1221,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/dd/8050c947d435c8d4bc94e3252f4d8bb8a76cfb424f043a8680be637a57f1/kiwisolver-1.5.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:59cd8683f575d96df5bb48f6add94afc055012c29e28124fcae2b63661b9efb1", size = 73558, upload-time = "2026-03-09T13:15:52.112Z" }, ] +[[package]] +name = "kornia" +version = "0.8.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "kornia-rs", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "packaging", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "torch", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1a/80/460e9dad0c1b68f3e83f198f3983638d62f910a88542b492290010d25f50/kornia-0.8.3.tar.gz", hash = "sha256:c06887374eaf39fb614a77ca054383e6cc2092cb7225e45437111b4984d0f866", size = 727384, upload-time = "2026-05-19T20:23:35.048Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/05/0c1cc486e87d2a5ad6649eb96a8ccaa7c6002d8d73d676c13e2102f286c1/kornia-0.8.3-py3-none-any.whl", hash = "sha256:0b15f5d359aeafd7ff54ea631ed1943a3eb295c4a6dae3f745ddeada25e33289", size = 1189381, upload-time = "2026-05-19T20:23:33.209Z" }, +] + +[[package]] +name = "kornia-rs" +version = "0.1.14" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/09/0f/cd6985790031ad2f0d4fcdfdb9763a177bebb970edddc6fe29f9e6afd902/kornia_rs-0.1.14.tar.gz", hash = "sha256:7584f654a9db2b41bee05c9aaf865608b665e2f7195096372e001b6f220de1d2", size = 2556658, upload-time = "2026-05-19T07:45:06.366Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/da/18bbc67c2e54714540c011b13d6c96fd014dfb7d1588a04c0fabf3bdc228/kornia_rs-0.1.14-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:371ba151de638150554af9fad53a351d5c41ed80f50a73ae376b58622e0a3430", size = 3690285, upload-time = "2026-05-19T07:45:28.063Z" }, + { url = "https://files.pythonhosted.org/packages/c1/7f/95dbcdc340009efba12c159c8b8cce7edc2e2172ac628f9bbecc9ad3ca73/kornia_rs-0.1.14-cp311-cp311-win_amd64.whl", hash = "sha256:9175b704be9d2de5f1aefc6516eefa46835f71bb93605db67936996d2be42684", size = 3368862, upload-time = "2026-05-19T07:46:15.096Z" }, +] + [[package]] name = "lark" version = "1.3.1" @@ -1887,7 +1920,7 @@ name = "onnxruntime-genai-winml" version = "0.14.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "numpy", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/03/57/bb3f6bd5b2670a2343a381116f40aa23a297d0590d58e0ca66dd16e4ac83/onnxruntime_genai_winml-0.14.1-cp311-cp311-win_amd64.whl", hash = "sha256:79071ffa02edcb3f7aaee55f474c616807ffa8522329e09f3d358f9f3c8f8827", size = 30496923, upload-time = "2026-06-02T20:29:17.649Z" }, @@ -1916,11 +1949,11 @@ name = "onnxruntime-windowsml" version = "1.24.5.202604171637" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "flatbuffers", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, - { name = "numpy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, - { name = "packaging", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, - { name = "protobuf", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, - { name = "sympy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "flatbuffers", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "numpy", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "packaging", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "protobuf", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "sympy", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/0c/1c/a61254721626b1bca498664c1d9fe0309e4ae60f6e971ab6c1b0bad98efe/onnxruntime_windowsml-1.24.5.202604171637-cp311-cp311-win_amd64.whl", hash = "sha256:13c9cb07244c0e47561fbf8b551944904d0ba9eca47747a66e86ae95fc94b41d", size = 25563868, upload-time = "2026-05-04T22:05:44.546Z" }, @@ -3396,8 +3429,8 @@ dependencies = [ { name = "numpy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "onnx", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "onnx-ir", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, - { name = "onnxruntime-genai-winml", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, - { name = "onnxruntime-windowsml", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "onnxruntime-genai-winml", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, + { name = "onnxruntime-windowsml", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, { name = "onnxscript", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "opentelemetry-sdk", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "optimum", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, @@ -3425,13 +3458,17 @@ dependencies = [ { name = "tqdm", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "transformers", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "uvicorn", extra = ["standard"], marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, - { name = "windowsml", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "windowsml", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" }, ] [package.optional-dependencies] audio = [ { name = "soundfile", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, ] +birefnet = [ + { name = "einops", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "kornia", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, +] dev = [ { name = "ipykernel", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "ipywidgets", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, @@ -3485,6 +3522,7 @@ requires-dist = [ { name = "colorama", specifier = ">=0.4.6" }, { name = "datasets", specifier = ">=4.1.1" }, { name = "diffusers", specifier = ">=0.36" }, + { name = "einops", marker = "extra == 'birefnet'", specifier = ">=0.8" }, { name = "evaluate", specifier = ">=0.4.6" }, { name = "fastapi", specifier = ">=0.135.3" }, { name = "hf-xet", specifier = ">=1.1.10" }, @@ -3493,6 +3531,7 @@ requires-dist = [ { name = "ipywidgets", marker = "extra == 'dev'", specifier = ">=8.1.7" }, { name = "jsonschema", specifier = ">=4.23" }, { name = "jupyter", marker = "extra == 'dev'", specifier = ">=1.1.1" }, + { name = "kornia", marker = "extra == 'birefnet'", specifier = ">=0.8" }, { name = "markdown-it-py", marker = "extra == 'dev'", specifier = ">=3" }, { name = "matplotlib", marker = "extra == 'dev'", specifier = ">=3.10" }, { name = "mcp", specifier = ">=0.1.0" }, @@ -3504,9 +3543,9 @@ requires-dist = [ { name = "numpy", specifier = ">=2.2,<2.3" }, { name = "onnx", specifier = "==1.18" }, { name = "onnx-ir", specifier = ">=0.1" }, - { name = "onnxruntime-genai-winml", specifier = "==0.14.1" }, + { name = "onnxruntime-genai-winml", marker = "sys_platform == 'win32'", specifier = "==0.14.1" }, { name = "onnxruntime-qnn", marker = "python_full_version >= '3.11' and extra == 'qnn'", specifier = ">=1.24.1" }, - { name = "onnxruntime-windowsml", specifier = "==1.24.5.202604171637" }, + { name = "onnxruntime-windowsml", marker = "sys_platform == 'win32'", specifier = "==1.24.5.202604171637" }, { name = "onnxscript", specifier = ">=0.2" }, { name = "opentelemetry-sdk", specifier = ">=1.39.1" }, { name = "openvino", marker = "extra == 'openvino'", specifier = ">=2023" }, @@ -3544,9 +3583,9 @@ requires-dist = [ { name = "transformers", specifier = ">=4.57" }, { name = "types-colorama", marker = "extra == 'dev'", specifier = ">=0.4.15.20250801" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.42.0" }, - { name = "windowsml", specifier = "==2.0.300" }, + { name = "windowsml", marker = "sys_platform == 'win32'", specifier = "==2.0.300" }, ] -provides-extras = ["dev", "audio", "openvino", "qnn"] +provides-extras = ["dev", "audio", "birefnet", "openvino", "qnn"] [package.metadata.requires-dev] dev = [ From b174e95604039e25140e6ee39416cf374cecd248 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 20:44:49 +0800 Subject: [PATCH 3/3] Fix BiRefNet test registration import --- tests/unit/models/birefnet/test_onnx_config.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/models/birefnet/test_onnx_config.py b/tests/unit/models/birefnet/test_onnx_config.py index 237643bd3..7132ff780 100644 --- a/tests/unit/models/birefnet/test_onnx_config.py +++ b/tests/unit/models/birefnet/test_onnx_config.py @@ -9,7 +9,6 @@ from torchvision.ops import deform_conv2d from transformers import PretrainedConfig -import winml.modelkit.models # noqa: F401 from winml.modelkit.export import generate_dummy_inputs from winml.modelkit.export.io import _get_onnx_config from winml.modelkit.models.hf.birefnet import BiRefNetIOConfig, _exportable_deform_conv2d