diff --git a/examples/recipes/caidas_swin2SR-classical-sr-x2-64/cpu/cpu/image-to-image_fp16_config.json b/examples/recipes/caidas_swin2SR-classical-sr-x2-64/cpu/cpu/image-to-image_fp16_config.json new file mode 100644 index 000000000..1a88a5375 --- /dev/null +++ b/examples/recipes/caidas_swin2SR-classical-sr-x2-64/cpu/cpu/image-to-image_fp16_config.json @@ -0,0 +1,64 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 64, + 64 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "reconstruction" + } + ] + }, + "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-to-image", + "model_id": "caidas/swin2SR-classical-sr-x2-64", + "model_type": "swin2sr", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "image-to-image", + "model_class": "AutoModelForImageToImage", + "model_type": "swin2sr" + } +} diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index b11e7b416..83be7d0d1 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -91,6 +91,8 @@ from .siglip import SIGLIP_CONFIG from .siglip import SiglipTextModelIOConfig as _SiglipTextModelIOConfig # triggers registration from .siglip import SiglipVisionModelIOConfig as _SiglipVisionModelIOConfig # triggers registration +from .swin2sr import MODEL_CLASS_MAPPING as _SWIN2SR_CLASS_MAPPING +from .swin2sr import Swin2SRIOConfig as _Swin2SRIOConfig # triggers registration from .t5 import MODEL_CLASS_MAPPING as _T5_CLASS_MAPPING from .t5 import T5_CONFIG from .t5 import T5DecoderIOConfig as _T5DecoderIOConfig # triggers registration @@ -132,6 +134,7 @@ _SAM2_CLASS_MAPPING, _SEGFORMER_CLASS_MAPPING, _SIGLIP_CLASS_MAPPING, + _SWIN2SR_CLASS_MAPPING, _T5_CLASS_MAPPING, _VED_CLASS_MAPPING, _VITPOSE_CLASS_MAPPING, diff --git a/src/winml/modelkit/models/hf/swin2sr.py b/src/winml/modelkit/models/hf/swin2sr.py index 88f9a411b..a07338c67 100644 --- a/src/winml/modelkit/models/hf/swin2sr.py +++ b/src/winml/modelkit/models/hf/swin2sr.py @@ -2,12 +2,35 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Swin2SR (Swin Transformer V2 for Super-Resolution) HuggingFace Model Patches. +"""Swin2SR HuggingFace model registration. -This module will provide Swin2SR-specific patches for ONNX export compatibility. -Currently empty — patches will be added as needed (see issue #236). - -Note: - No model-specific build config is needed. The analyzer autoconf loop in the - build pipeline discovers optimization flags automatically. See issue #232. +Swin2SR checkpoints default to image super-resolution (image-to-image). Optimum +ships a built-in ``Swin2srOnnxConfig``, but registering it here ensures WinML's +``get_supported_tasks("swin2sr")`` is populated as soon as +``winml.modelkit.models`` is imported, even if +``optimum.exporters.onnx.model_configs`` has not been imported yet. """ + +from __future__ import annotations + +from optimum.exporters.onnx.model_configs import Swin2srOnnxConfig +from transformers import AutoModelForImageToImage + +from ...export import register_onnx_overwrite + + +# (model_type, task) -> HuggingFace model class +# +# The (swin2sr, None) sentinel declares image-to-image as the default task for +# task auto-detection when --task is omitted. +MODEL_CLASS_MAPPING: dict[tuple[str, str | None], type] = { + ("swin2sr", "image-to-image"): AutoModelForImageToImage, + ("swin2sr", None): AutoModelForImageToImage, +} + + +@register_onnx_overwrite("swin2sr", "feature-extraction", library_name="transformers") +@register_onnx_overwrite("swin2sr", "image-to-image", library_name="transformers") +class Swin2SRIOConfig(Swin2srOnnxConfig): # type: ignore[misc] # optimum base is untyped + """Local registration shim for Swin2SR ONNX export tasks.""" + diff --git a/tests/unit/loader/test_get_supported_tasks.py b/tests/unit/loader/test_get_supported_tasks.py index 4ff146401..bf0b7ff5b 100644 --- a/tests/unit/loader/test_get_supported_tasks.py +++ b/tests/unit/loader/test_get_supported_tasks.py @@ -29,6 +29,7 @@ class TestGetSupportedTasks: ("whisper", "automatic-speech-recognition"), ("clip", "zero-shot-image-classification"), ("llama", "text-generation"), + ("swin2sr", "image-to-image"), ], ids=[ "bert", @@ -41,6 +42,7 @@ class TestGetSupportedTasks: "whisper", "clip", "llama", + "swin2sr", ], ) def test_known_model_type_contains_expected_task(self, model_type, expected_task): diff --git a/tests/unit/models/test_swin2sr_support.py b/tests/unit/models/test_swin2sr_support.py new file mode 100644 index 000000000..fad226670 --- /dev/null +++ b/tests/unit/models/test_swin2sr_support.py @@ -0,0 +1,58 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for Swin2SR model support registration. + +Swin2SR should resolve to image-to-image by default and expose ONNX export +registrations through WinML's local registry import path. +""" + +from __future__ import annotations + +from transformers import Swin2SRConfig + +import winml.modelkit.models # noqa: F401 # trigger registrations +from winml.modelkit.export.io import _get_onnx_config +from winml.modelkit.loader import get_supported_tasks, resolve_task +from winml.modelkit.models.hf import MODEL_CLASS_MAPPING +from winml.modelkit.models.hf.swin2sr import ( + MODEL_CLASS_MAPPING as SWIN2SR_MAPPING, +) +from winml.modelkit.models.hf.swin2sr import ( + Swin2SRIOConfig, +) + + +class TestSwin2SRSupport: + """Swin2SR is discoverable and resolvable as image-to-image.""" + + def test_get_supported_tasks_includes_image_to_image(self): + """swin2sr task list includes image-to-image after local registrations.""" + tasks = get_supported_tasks("swin2sr") + assert "image-to-image" in tasks + + def test_default_resolution_is_image_to_image(self): + """Task auto-detection defaults Swin2SR to image-to-image.""" + config = Swin2SRConfig() + config.architectures = ["Swin2SRForImageSuperResolution"] + + resolution = resolve_task(config) + + assert resolution.task == "image-to-image" + assert resolution.model_class.__name__ == "AutoModelForImageToImage" + + def test_onnx_config_registration(self): + """ONNX config lookup for swin2sr/image-to-image resolves to local shim.""" + onnx_config = _get_onnx_config("swin2sr", "image-to-image", Swin2SRConfig()) + assert isinstance(onnx_config, Swin2SRIOConfig) + + def test_mapping_is_merged_and_has_default_sentinel(self): + """Model-class mapping includes swin2sr task key and default sentinel.""" + assert SWIN2SR_MAPPING.items() <= MODEL_CLASS_MAPPING.items() + assert ("swin2sr", "image-to-image") in MODEL_CLASS_MAPPING + assert ("swin2sr", None) in MODEL_CLASS_MAPPING + assert ( + MODEL_CLASS_MAPPING[("swin2sr", None)] + is MODEL_CLASS_MAPPING[("swin2sr", "image-to-image")] + )