Add ZhengPeng7/BiRefNet image segmentation support

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
+  }
+}

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
+  }
+}

pyproject.toml

@@ -45,8 +45,8 @@ dependencies = [
   "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",
@@ -77,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 = [
@@ -100,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'",

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()

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:

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