From 330edcbf99d185d0ab12e128ee1c8956a545e391 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Sun, 19 Jul 2026 19:40:01 +0800 Subject: [PATCH 1/5] feat(loader): support published ONNX SAM composites --- ...-generation_fp16_config_image-encoder.json | 13 ++ ...generation_fp16_config_prompt-decoder.json | 13 ++ ...-generation_fp32_config_image-encoder.json | 10 ++ ...generation_fp32_config_prompt-decoder.json | 10 ++ src/winml/modelkit/commands/build.py | 70 ++++++++-- src/winml/modelkit/commands/config.py | 114 ++++++++++----- src/winml/modelkit/eval/evaluate.py | 10 ++ .../eval/mask_generation_evaluator.py | 18 ++- src/winml/modelkit/loader/config.py | 4 + src/winml/modelkit/loader/onnx_hub.py | 130 +++++++++++++++++- src/winml/modelkit/loader/resolution.py | 64 +++++++++ src/winml/modelkit/loader/task.py | 1 + src/winml/modelkit/models/hf/sam.py | 43 +++++- .../modelkit/models/winml/composite_model.py | 29 ++++ src/winml/modelkit/pattern/base.py | 7 + .../test_pattern_matcher_robustness.py | 36 +++++ tests/unit/commands/test_config_cli.py | 101 ++++++++++++++ tests/unit/config/test_build_onnx.py | 16 +++ .../eval/test_mask_generation_evaluator.py | 27 ++++ tests/unit/loader/test_composite_tasks.py | 8 ++ tests/unit/loader/test_onnx_composite_hub.py | 82 +++++++++++ 21 files changed, 751 insertions(+), 55 deletions(-) create mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json create mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json create mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json create mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json create mode 100644 tests/unit/loader/test_onnx_composite_hub.py diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json new file mode 100644 index 000000000..2cb9e5235 --- /dev/null +++ b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json @@ -0,0 +1,13 @@ +{ + "export": null, + "optim": {}, + "quant": { + "mode": "fp16", + "fp16_keep_io_types": true + }, + "compile": null, + "loader": { + "task": "image-feature-extraction", + "component_name": "image-encoder" + } +} diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json new file mode 100644 index 000000000..6c70cf55a --- /dev/null +++ b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json @@ -0,0 +1,13 @@ +{ + "export": null, + "optim": {}, + "quant": { + "mode": "fp16", + "fp16_keep_io_types": true + }, + "compile": null, + "loader": { + "task": "mask-generation", + "component_name": "prompt-decoder" + } +} diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json new file mode 100644 index 000000000..05a2f7c52 --- /dev/null +++ b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json @@ -0,0 +1,10 @@ +{ + "export": null, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "image-feature-extraction", + "component_name": "image-encoder" + } +} diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json new file mode 100644 index 000000000..7a7a86f27 --- /dev/null +++ b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json @@ -0,0 +1,10 @@ +{ + "export": null, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "mask-generation", + "component_name": "prompt-decoder" + } +} diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index 3076b585a..1eb6b9511 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -1075,6 +1075,31 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: _configs_to_validate: list[WinMLBuildConfig] = ( config_or_configs if isinstance(config_or_configs, list) else [config_or_configs] ) + + # A portable composite-component recipe keeps the Hub model ID in + # ``-m`` and identifies its graph role in loader.component_name. Resolve + # that role to the published ONNX graph before deciding the build path. + if ( + model + and not isinstance(config_or_configs, list) + and isinstance(config_or_configs.loader.component_name, str) + and (model_input is None or model_input.kind is not ModelInputKind.ONNX_FILE) + ): + from ..loader.resolution import resolve_composite_onnx_sources + + sources = resolve_composite_onnx_sources( + model, + component_name=config_or_configs.loader.component_name, + precision=precision, + trust_remote_code=trust_remote_code, + ) + if sources is None: + raise click.UsageError( + f"Model {model!r} does not publish ONNX sources for composite " + f"component {config_or_configs.loader.component_name!r}." + ) + model = next(iter(sources.values())) + model_input = classify_model_input(model) try: for _cfg in _configs_to_validate: _cfg.validate() @@ -1309,6 +1334,14 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: completed: list[str] = [] try: + from ..loader.resolution import resolve_composite_onnx_sources + + component_sources = resolve_composite_onnx_sources( + model, + task=task_hint, + precision=precision, + trust_remote_code=trust_remote_code, + ) for name, component_task in components.items(): console.print( f"\n[bold blue]Sub-model:[/bold blue] {name} (task={component_task})" @@ -1316,16 +1349,31 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: from ..config import generate_build_config as gen_cfg - component_config = gen_cfg( - model, - task=component_task, - trust_remote_code=trust_remote_code, - device=device, - precision=precision, - ep=ep, - shape_config=shape_overrides, - override={"export": export_overrides} if export_overrides else None, + component_source = ( + component_sources.get(name) if component_sources else None ) + if component_source is not None: + from ..config import generate_onnx_build_config + + component_config = generate_onnx_build_config( + component_source, + task=component_task, + device=device, + precision=precision, + ep=ep, + ) + component_config.loader.component_name = name + else: + component_config = gen_cfg( + model, + task=component_task, + trust_remote_code=trust_remote_code, + device=device, + precision=precision, + ep=ep, + shape_config=shape_overrides, + override={"export": export_overrides} if export_overrides else None, + ) # Carry over quant/compile settings from the outer config # (already patched by CLI overrides like --no-quant / # --no-compile). Deep-copy to avoid sharing mutable state @@ -1378,8 +1426,8 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: _run_single_build( config=component_config, config_file=None, - model_id=model, - is_onnx=False, + model_id=component_source or model, + is_onnx=component_source is not None, resolved_dir=resolved_dir, rebuild=rebuild, cache_key=name, diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index 1f2bcade0..359fa5cca 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -363,13 +363,39 @@ def config( hf_model, model_type, task, trust_remote_code=trust_remote_code ) if pipeline_components: + from ..loader.resolution import ( + composite_requires_pretrained_onnx, + resolve_composite_onnx_sources, + ) + + published_only = composite_requires_pretrained_onnx(pipeline_components) + if not published_only and _export_flags_given: + raise click.UsageError( + "--input-specs, --export-config, and --dynamic-axes are not " + "supported for composite (multi-component) models, whose " + "sub-components each have their own export inputs. Generate " + "the per-component configs first, then edit their export " + "sections individually." + ) + component_sources = ( + resolve_composite_onnx_sources( + hf_model, + task=task, + precision=precision, + trust_remote_code=trust_remote_code, + ) + if hf_model + else None + ) + use_composite = component_sources is not None or not published_only + # Export controls target a single export graph; a composite model # has one export per sub-component with distinct inputs. Reject on # raw flag presence — before loading/validating the JSON — so the # composite-specific error wins (mirroring the ONNX path) instead # of a downstream "Invalid export configuration". config never fans # these overrides out across heterogeneous components. - if _export_flags_given: + if use_composite and _export_flags_given: raise click.UsageError( "--input-specs, --export-config, and --dynamic-axes are not " "supported for composite (multi-component) models, whose " @@ -377,26 +403,30 @@ def config( "the per-component configs first, then edit their export " "sections individually." ) - # composite model: generate one config per sub-component - _generate_pipeline_configs( - pipeline_components, - hf_model=hf_model, - model_class=model_class, - model_type=model_type, - override=override, - shape_config=shape_config, - library_name=library_name, - device=device, - precision=precision, - trust_remote_code=trust_remote_code, - ep=ep, - no_quant=not quant, - no_compile=no_compile, - output=output, - overwrite=overwrite, - console=console, - ) - return + # Composite model: generate one config per sub-component. An + # optional published-ONNX registration with no sources falls + # through to the pre-existing single PyTorch export path. + if use_composite: + _generate_pipeline_configs( + pipeline_components, + component_sources=component_sources, + hf_model=hf_model, + model_class=model_class, + model_type=model_type, + override=override, + shape_config=shape_config, + library_name=library_name, + device=device, + precision=precision, + trust_remote_code=trust_remote_code, + ep=ep, + no_quant=not quant, + no_compile=no_compile, + output=output, + overwrite=overwrite, + console=console, + ) + return # Load export CLI overrides now that the multi-graph paths (ONNX, # --module, composite) have all been rejected — the single HF config @@ -603,6 +633,7 @@ def _resolve_composite_model_components( def _generate_pipeline_configs( components: dict[str, str], *, + component_sources: dict[str, str] | None, hf_model: str | None, model_class: str | None, model_type: str | None, @@ -620,7 +651,7 @@ def _generate_pipeline_configs( console: Any, ) -> None: """Generate and save one config file per pipeline sub-component.""" - from ..config import generate_hf_build_config + from ..config import generate_hf_build_config, generate_onnx_build_config for component_name, component_task in components.items(): console.print( @@ -628,19 +659,32 @@ def _generate_pipeline_configs( f"(task={component_task})...[/dim]" ) - cfg = generate_hf_build_config( - model_id=hf_model, - task=component_task, - model_class=model_class, - model_type=model_type, - override=override, - shape_config=shape_config, - library_name=library_name, - device=device, - precision=precision, - trust_remote_code=trust_remote_code, - ep=ep, - ) + source = component_sources.get(component_name) if component_sources else None + if source is not None: + cfg = generate_onnx_build_config( + source, + task=component_task, + override=override, + device=device, + precision=precision, + ep=ep, + ) + cfg.loader.component_name = component_name + cfg.loader.model_type = model_type + else: + cfg = generate_hf_build_config( + model_id=hf_model, + task=component_task, + model_class=model_class, + model_type=model_type, + override=override, + shape_config=shape_config, + library_name=library_name, + device=device, + precision=precision, + trust_remote_code=trust_remote_code, + ep=ep, + ) _apply_stage_overrides(cfg, no_quant=no_quant, no_compile=no_compile) config_json = json.dumps(cfg.to_dict(), indent=2) diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 74863970f..0ba635e49 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -330,6 +330,16 @@ def evaluate(config: WinMLEvaluationConfig) -> EvalResult: config = replace( config, mode=mode, task=_resolve_task(config), dataset=deepcopy(config.dataset) ) + if config.task == "mask-generation" and config.model_path is None and config.model_id: + from ..loader.resolution import resolve_composite_onnx_sources + + sources = resolve_composite_onnx_sources( + config.model_id, + task=config.task, + precision=config.precision, + ) + if sources is not None: + config = replace(config, model_path=sources) if config.mode != "compare" and config.dataset.path is None: default = _DEFAULT_DATASETS.get(config.task) if config.task is not None else None if default is None: diff --git a/src/winml/modelkit/eval/mask_generation_evaluator.py b/src/winml/modelkit/eval/mask_generation_evaluator.py index 1e5064745..0514b4bd1 100644 --- a/src/winml/modelkit/eval/mask_generation_evaluator.py +++ b/src/winml/modelkit/eval/mask_generation_evaluator.py @@ -167,6 +167,13 @@ def __init__( "Use prompt_mode='bbox' or 'point'." ) self._enc_sess, self._dec_sess = self._load_sessions() + self._decoder_input_names = _node_names(self._dec_sess.get_inputs()) + if "input_boxes" not in self._decoder_input_names: + if "prompt_mode" in mapping and self._prompt_mode == "bbox": + raise ValueError( + "The decoder does not accept box prompts; use prompt_mode='point'." + ) + self._prompt_mode = "point" self._encoder_input_name = _resolve_encoder_input_name(self._enc_sess) self._encoder_output_names = _node_names(self._enc_sess.get_outputs()) self._embedding_input_names = _resolve_embedding_input_names( @@ -330,6 +337,7 @@ def _predict(self, sample: dict[str, Any]) -> np.ndarray: scale_y=scale_y, emb=emb, required_embed_names=self._embedding_input_names, + required_input_names=self._decoder_input_names, ) dec_out = self._dec_sess.run( list(self._decoder_output_names), @@ -427,7 +435,7 @@ def _resolve_embedding_input_names( ) -> tuple[str, ...]: """Match decoder embedding inputs to encoder output names.""" decoder_inputs = _node_names(dec_sess.get_inputs()) - required_prompt_inputs = ("input_points", "input_labels", "input_boxes") + required_prompt_inputs = ("input_points", "input_labels") missing_prompt = [name for name in required_prompt_inputs if name not in decoder_inputs] if missing_prompt: raise ValueError( @@ -599,6 +607,11 @@ def _build_decoder_inputs( "image_embeddings.1", "image_embeddings.2", ), + required_input_names: tuple[str, ...] = ( + "input_points", + "input_labels", + "input_boxes", + ), ) -> dict[str, np.ndarray]: """Assemble the decoder feed dict for bbox or point prompts. @@ -648,11 +661,12 @@ def _build_decoder_inputs( f"Encoder output missing required keys {missing_embeds}. Got: {list(emb.keys())}" ) - feed = { + prompt_feed = { "input_points": points, "input_labels": labels, "input_boxes": box, } + feed = {name: value for name, value in prompt_feed.items() if name in required_input_names} feed.update({name: emb[name] for name in required_embed_names}) return feed diff --git a/src/winml/modelkit/loader/config.py b/src/winml/modelkit/loader/config.py index 5e9ca0bff..9f3492f98 100644 --- a/src/winml/modelkit/loader/config.py +++ b/src/winml/modelkit/loader/config.py @@ -71,6 +71,7 @@ class WinMLLoaderConfig: model_class: str | None = None model_type: str | None = None module_path: str | None = None + component_name: str | None = None user_script: str | None = None trust_remote_code: bool = False @@ -89,6 +90,8 @@ def to_dict(self) -> dict[str, Any]: result["model_type"] = self.model_type if self.module_path is not None: result["module_path"] = self.module_path + if self.component_name is not None: + result["component_name"] = self.component_name if self.user_script is not None: result["user_script"] = self.user_script if self.trust_remote_code: @@ -110,6 +113,7 @@ def from_dict(cls, data: dict[str, Any]) -> WinMLLoaderConfig: model_class=data.get("model_class"), model_type=data.get("model_type"), module_path=data.get("module_path"), + component_name=data.get("component_name"), user_script=data.get("user_script"), trust_remote_code=data.get("trust_remote_code", False), ) diff --git a/src/winml/modelkit/loader/onnx_hub.py b/src/winml/modelkit/loader/onnx_hub.py index 5a9a5dc1b..cdec55b2f 100644 --- a/src/winml/modelkit/loader/onnx_hub.py +++ b/src/winml/modelkit/loader/onnx_hub.py @@ -19,12 +19,137 @@ from __future__ import annotations import logging +from collections import Counter +from dataclasses import dataclass from pathlib import Path +from typing import Any logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class _GraphContract: + path: Path + inputs: tuple[tuple[str, str, int], ...] + outputs: tuple[tuple[str, str, int], ...] + precision: str + + +def _inspect_runnable_graph(path: Path) -> _GraphContract | None: + """Return a runnable graph's I/O contract, or ``None`` when ORT rejects it.""" + import onnx + import onnxruntime as ort + + try: + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + session = ort.InferenceSession( + str(path), + sess_options=options, + providers=["CPUExecutionProvider"], + ) + model = onnx.load(path, load_external_data=False) + except Exception as error: + logger.warning("Ignoring unusable published ONNX graph %s: %s", path.name, error) + return None + + type_counts = Counter(initializer.data_type for initializer in model.graph.initializer) + fp16 = type_counts[onnx.TensorProto.FLOAT16] + fp32 = type_counts[onnx.TensorProto.FLOAT] + precision = "fp16" if fp16 > fp32 else "fp32" + + def contract(nodes: list[Any]) -> tuple[tuple[str, str, int], ...]: + return tuple((node.name, node.type, len(node.shape)) for node in nodes) + + return _GraphContract( + path=path, + inputs=contract(session.get_inputs()), + outputs=contract(session.get_outputs()), + precision=precision, + ) + + +def resolve_hf_onnx_encoder_decoder( + model_id: str, + *, + revision: str | None = None, + precision: str = "fp32", + cache_dir: str | Path | None = None, + token: str | bool | None = None, +) -> dict[str, Path]: + """Discover a runnable image-encoder/prompt-decoder pair in a Hub repo. + + Selection is graph-contract driven rather than filename driven. The encoder + has one rank-4 tensor input; the decoder consumes one or more encoder outputs + plus a rank-3 integer prompt tensor. Invalid published variants are ignored + after an ORT CPU session probe. An fp16 build may safely fall back to fp32 + source graphs because the normal build pipeline performs fp16 conversion. + """ + from huggingface_hub import list_repo_files + + files = sorted( + name + for name in list_repo_files(model_id, revision=revision, token=token) + if name.lower().endswith(".onnx") + ) + if not files: + raise FileNotFoundError(f"No published ONNX graphs found in Hub repo {model_id!r}.") + + graphs: list[_GraphContract] = [] + for filename in files: + path = resolve_hf_onnx_path( + f"{model_id}/{filename}", + revision=revision, + cache_dir=cache_dir, + token=token, + ) + graph = _inspect_runnable_graph(path) + if graph is not None: + graphs.append(graph) + + preferred_precisions = ("fp16", "fp32") if precision == "fp16" else ("fp32",) + candidates: list[tuple[int, int, str, _GraphContract, _GraphContract]] = [] + for encoder in graphs: + encoder_inputs = encoder.inputs + if len(encoder_inputs) != 1 or encoder_inputs[0][2] != 4: + continue + encoder_outputs = {name for name, _dtype, _rank in encoder.outputs} + for decoder in graphs: + if decoder.path == encoder.path or decoder.precision != encoder.precision: + continue + decoder_inputs = {name for name, _dtype, _rank in decoder.inputs} + shared = encoder_outputs & decoder_inputs + has_integer_prompt = any( + rank == 3 and "int" in dtype.lower() for _name, dtype, rank in decoder.inputs + ) + if not shared or not has_integer_prompt: + continue + try: + precision_rank = preferred_precisions.index(encoder.precision) + except ValueError: + continue + candidates.append((precision_rank, -len(shared), str(encoder.path), encoder, decoder)) + + if not candidates: + contracts = [ + { + "file": graph.path.name, + "precision": graph.precision, + "inputs": [name for name, _dtype, _rank in graph.inputs], + "outputs": [name for name, _dtype, _rank in graph.outputs], + } + for graph in graphs + ] + raise ValueError( + "Published ONNX graphs do not contain a runnable encoder/decoder pair " + f"for precision {precision!r}: {contracts}" + ) + + _rank, _shared, _path, encoder, decoder = min(candidates, key=lambda row: row[:3]) + return {"image-encoder": encoder.path, "prompt-decoder": decoder.path} + + def resolve_hf_onnx_path( model_id: str, *, @@ -79,9 +204,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 +294,6 @@ def _format_available_onnx_files( __all__ = [ + "resolve_hf_onnx_encoder_decoder", "resolve_hf_onnx_path", ] diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 548848175..2f1f1848b 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -350,6 +350,70 @@ def resolve_composite_components( return resolve_task(config).composite +def resolve_composite_onnx_sources( + hf_model: str, + *, + task: str | None = None, + component_name: str | None = None, + precision: str = "fp32", + trust_remote_code: bool = False, +) -> dict[str, str] | None: + """Resolve published ONNX sources for a registered composite model.""" + from transformers import AutoConfig + + config = AutoConfig.from_pretrained(hf_model, trust_remote_code=trust_remote_code) + model_type = config.model_type.lower().replace("_", "-") + registry = _composite_registry() + + if task is not None and (model_type, task) in registry: + composite_cls = registry[(model_type, task)] + else: + resolution = resolve_task(config, task=task) + matching = [ + cls + for (registered_type, _registered_task), cls in registry.items() + if registered_type == model_type + and resolution.composite == dict(cls._SUB_MODEL_CONFIG) + and (component_name is None or component_name in cls._SUB_MODEL_CONFIG) + ] + if not matching: + return None + if len(set(matching)) != 1: + raise ValueError( + f"{model_type!r} has multiple matching composite ONNX sources; pass --task." + ) + composite_cls = matching[0] + + sources = composite_cls.resolve_pretrained_onnx( + hf_model, + revision=getattr(config, "_commit_hash", None), + precision=precision, + ) + if sources is None: + return None + resolved = {name: str(path) for name, path in sources.items()} + if component_name is not None: + if component_name not in resolved: + raise ValueError( + f"Published composite has no component {component_name!r}; " + f"available: {sorted(resolved)}" + ) + return {component_name: resolved[component_name]} + return resolved + + +def composite_requires_pretrained_onnx(components: CompositeComponents) -> bool: + """Whether a registered composite exists only for published ONNX sources. + + Such registrations augment an architecture without replacing its existing + PyTorch export behavior when the checkpoint publishes no ONNX components. + """ + matching = { + cls for cls in _composite_registry().values() if components == dict(cls._SUB_MODEL_CONFIG) + } + return len(matching) == 1 and next(iter(matching))._PRETRAINED_ONNX_ONLY + + def composite_pipeline_tasks(model_type: str) -> list[str]: """Pipeline (composite) tasks a model_type can serve, sorted; ``[]`` for non-composites. diff --git a/src/winml/modelkit/loader/task.py b/src/winml/modelkit/loader/task.py index 0ae0ca6be..14a44c94b 100644 --- a/src/winml/modelkit/loader/task.py +++ b/src/winml/modelkit/loader/task.py @@ -129,6 +129,7 @@ COMPOSITE_TASKS: frozenset[str] = frozenset( { "image-to-text", + "mask-generation", "summarization", "table-question-answering", "text-generation", diff --git a/src/winml/modelkit/models/hf/sam.py b/src/winml/modelkit/models/hf/sam.py index 3e7cb2a6f..922677b8e 100644 --- a/src/winml/modelkit/models/hf/sam.py +++ b/src/winml/modelkit/models/hf/sam.py @@ -34,7 +34,7 @@ from __future__ import annotations import types -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast import torch import torch.nn.functional as F @@ -48,6 +48,7 @@ from transformers import Sam2Model, SamModel from ...export import register_onnx_overwrite +from ..winml.composite_model import WinMLCompositeModel, register_composite_model if TYPE_CHECKING: @@ -402,6 +403,41 @@ def forward( } +@register_composite_model("sam", "mask-generation") +class WinMLSAMModel(WinMLCompositeModel): + """SAM composite backed by published encoder and prompt-decoder ONNX graphs.""" + + _PRETRAINED_ONNX_ONLY: ClassVar[bool] = True + _SUB_MODEL_CONFIG: ClassVar[dict[str, str]] = { + "image-encoder": "image-feature-extraction", + "prompt-decoder": "mask-generation", + } + + @classmethod + def resolve_pretrained_onnx( + cls, + model_id: str, + *, + revision: str | None = None, + precision: str = "fp32", + ) -> dict[str, str] | None: + """Discover the checkpoint's runnable published ONNX graph pair.""" + from ...loader.onnx_hub import resolve_hf_onnx_encoder_decoder + + try: + sources = resolve_hf_onnx_encoder_decoder( + model_id, + revision=revision, + precision=precision, + ) + except FileNotFoundError: + # Published ONNX is an optional source for SAM. Ordinary PyTorch + # checkpoints must retain the existing export path. Malformed or + # ambiguous published graph sets still raise and fail closed. + return None + return {name: str(path) for name, path in sources.items()} + + # Note: No model-specific build config needed. The analyzer autoconf loop # discovers optimization flags automatically. See issue #232. @@ -577,9 +613,7 @@ def _patched_sam2_prompt_encoder_forward( # _original_forward is added dynamically by Sam2ModelPatcher (not on the class), # so it's reached via nn.Module.__getattr__ (typed Tensor | Module); type it callable. - orig_forward = cast( - "Callable[..., tuple[torch.Tensor, torch.Tensor]]", self._original_forward - ) + orig_forward = cast("Callable[..., tuple[torch.Tensor, torch.Tensor]]", self._original_forward) # If use_mask_input not provided, use original behavior if use_mask_input is None: @@ -1084,6 +1118,7 @@ def outputs(self) -> dict[str, dict[int, str]]: "Sam2ModelPatcher", "Sam2NormalizedVisionConfig", "SamMaskGenerationIOConfig", + "WinMLSAMModel", "_patched_sam2_multiscale_block_forward", "_patched_sam2_prompt_encoder_forward", ] diff --git a/src/winml/modelkit/models/winml/composite_model.py b/src/winml/modelkit/models/winml/composite_model.py index cc5d8e517..354c06b29 100644 --- a/src/winml/modelkit/models/winml/composite_model.py +++ b/src/winml/modelkit/models/winml/composite_model.py @@ -105,6 +105,18 @@ class WinMLCompositeModel(PreTrainedModel): """ _SUB_MODEL_CONFIG: ClassVar[dict[str, str]] = {} + _PRETRAINED_ONNX_ONLY: ClassVar[bool] = False + + @classmethod + def resolve_pretrained_onnx( + cls, + model_id: str, + *, + revision: str | None = None, + precision: str = "fp32", + ) -> Mapping[str, str | Path] | None: + """Resolve published ONNX components when this composite supports them.""" + return None def __init__( self, @@ -184,6 +196,23 @@ def from_pretrained( trust_remote_code=trust_remote_code, **kwargs, ) + + published = cls.resolve_pretrained_onnx( + model_id, + revision=getattr(hf_config, "_commit_hash", None), + precision=cast("str", kwargs.get("precision", "fp32")), + ) + if published is not None: + return cls.from_onnx( + published, + task=task, + hf_config=hf_config, + device=device, + precision=kwargs.pop("precision", "fp32"), + use_cache=use_cache, + force_rebuild=force_rebuild, + **kwargs, + ) from ..auto import WinMLAutoModel per_component = sub_model_kwargs or {} diff --git a/src/winml/modelkit/pattern/base.py b/src/winml/modelkit/pattern/base.py index 1ca29ab44..f52b55476 100644 --- a/src/winml/modelkit/pattern/base.py +++ b/src/winml/modelkit/pattern/base.py @@ -604,6 +604,13 @@ def _check_skeleton_result_impl( cast("tuple[int, ...]", info.shape) ).get_value(type_annotation) + # Pattern validation must fail closed when shape/value metadata is + # unavailable. Pattern implementations use these arrays to derive + # dtype-dependent constants; allowing a partial mapping through turns + # an incomplete shape-inference result into an analyzer-wide KeyError. + if any(name not in inputs for name in input_infos): + return None + # Build is_constant_map from input_infos is_constant_map = {name: info.is_constant for name, info in input_infos.items()} diff --git a/tests/unit/analyze/pattern/test_pattern_matcher_robustness.py b/tests/unit/analyze/pattern/test_pattern_matcher_robustness.py index 3285b9637..f90e192d4 100644 --- a/tests/unit/analyze/pattern/test_pattern_matcher_robustness.py +++ b/tests/unit/analyze/pattern/test_pattern_matcher_robustness.py @@ -13,6 +13,7 @@ from winml.modelkit.onnx import ONNXDomain from winml.modelkit.pattern import ( + LayerNormalizationPowPattern, Pattern, PatternMatcher, PatternSchema, @@ -185,3 +186,38 @@ def test_missing_edge_skips_match_gracefully(self): results = matcher.match_skeleton() # The pattern can't match because the edge is missing assert len(results) == 0 + + +class TestPatternMatcherIncompleteInputMetadata: + """Incomplete shape inference must reject a match, not abort analysis.""" + + def test_layernorm_without_input_shape_skips_match(self) -> None: + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, None) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, None) + scale = numpy_helper.from_array(np.ones([4], dtype=np.float32), name="Scale") + bias = numpy_helper.from_array(np.zeros([4], dtype=np.float32), name="B") + exponent = numpy_helper.from_array(np.array(2.0, dtype=np.float32), name="Exponent") + epsilon = numpy_helper.from_array(np.array(1e-5, dtype=np.float32), name="Epsilon") + nodes = [ + helper.make_node("ReduceMean", ["X"], ["mean"], axes=[-1], name="mean"), + helper.make_node("Sub", ["X", "mean"], ["centered"], name="sub"), + helper.make_node("Pow", ["centered", "Exponent"], ["squared"], name="pow"), + helper.make_node("ReduceMean", ["squared"], ["variance"], axes=[-1], name="variance"), + helper.make_node("Add", ["variance", "Epsilon"], ["stabilized"], name="epsilon"), + helper.make_node("Sqrt", ["stabilized"], ["stddev"], name="sqrt"), + helper.make_node("Div", ["centered", "stddev"], ["normalized"], name="div"), + helper.make_node("Mul", ["normalized", "Scale"], ["scaled"], name="scale"), + helper.make_node("Add", ["scaled", "B"], ["Y"], name="bias"), + ] + graph = helper.make_graph( + nodes, + "incomplete_layernorm", + [x], + [y], + [scale, bias, exponent, epsilon], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + matcher = PatternMatcher(model) + matcher.register_pattern(LayerNormalizationPowPattern()) + + assert matcher.match() == [] diff --git a/tests/unit/commands/test_config_cli.py b/tests/unit/commands/test_config_cli.py index 47c51bdc0..803a06ccd 100644 --- a/tests/unit/commands/test_config_cli.py +++ b/tests/unit/commands/test_config_cli.py @@ -211,6 +211,107 @@ def test_output_to_file( result = runner.invoke(config, ["-m", "test", "-o", str(output_file)]) assert result.exit_code == 0, f"Output to file should succeed: {result.output}" + +class TestSamConfigSourceSelection: + """SAM config routing distinguishes PyTorch and published-ONNX checkpoints.""" + + @staticmethod + def _single_sam_config() -> MagicMock: + cfg = MagicMock() + cfg.loader.task = "mask-generation" + cfg.loader.model_class = "SAMMaskGeneration" + cfg.export = None + cfg.quant = None + cfg.to_dict.return_value = { + "loader": { + "task": "mask-generation", + "model_class": "SAMMaskGeneration", + "model_type": "sam", + }, + "export": None, + "optim": {}, + "quant": None, + "compile": None, + } + return cfg + + def test_pytorch_sam_uses_existing_single_config_path( + self, runner: CliRunner, tmp_path: Path + ) -> None: + from transformers import SamConfig + + from winml.modelkit.commands.config import config + + output = tmp_path / "sam.json" + with ( + patch("transformers.AutoConfig.from_pretrained", return_value=SamConfig()), + patch( + "winml.modelkit.models.hf.sam.WinMLSAMModel.resolve_pretrained_onnx", + return_value=None, + ), + patch( + "winml.modelkit.config.generate_hf_build_config", + return_value=self._single_sam_config(), + ) as generate_hf, + patch("winml.modelkit.commands.config._generate_pipeline_configs") as pipeline, + ): + result = runner.invoke( + config, + [ + "-m", + "facebook/sam-vit-base", + "--task", + "mask-generation", + "--no-compile", + "-o", + str(output), + ], + ) + + assert result.exit_code == 0, result.output + assert json.loads(output.read_text())["loader"]["model_class"] == "SAMMaskGeneration" + generate_hf.assert_called_once() + pipeline.assert_not_called() + + def test_published_onnx_sam_uses_composite_config_path( + self, runner: CliRunner, tmp_path: Path + ) -> None: + from transformers import SamConfig + + from winml.modelkit.commands.config import config + + sources = { + "image-encoder": "encoder.onnx", + "prompt-decoder": "decoder.onnx", + } + with ( + patch("transformers.AutoConfig.from_pretrained", return_value=SamConfig()), + patch( + "winml.modelkit.models.hf.sam.WinMLSAMModel.resolve_pretrained_onnx", + return_value=sources, + ), + patch("winml.modelkit.commands.config._generate_pipeline_configs") as pipeline, + ): + result = runner.invoke( + config, + [ + "-m", + "Xenova/slimsam-77-uniform", + "--task", + "mask-generation", + "--no-compile", + "-o", + str(tmp_path / "slimsam.json"), + ], + ) + + assert result.exit_code == 0, result.output + assert pipeline.call_args.kwargs["component_sources"] == sources + + +class TestConfigCliInterfaceContinued: + """Remaining config CLI interface and override tests.""" + def test_model_type_without_model( self, runner: CliRunner, diff --git a/tests/unit/config/test_build_onnx.py b/tests/unit/config/test_build_onnx.py index 25180a255..72c43a02d 100644 --- a/tests/unit/config/test_build_onnx.py +++ b/tests/unit/config/test_build_onnx.py @@ -38,6 +38,22 @@ from winml.modelkit.utils.config_utils import merge_config +def test_loader_component_name_round_trips() -> None: + config = WinMLBuildConfig.from_dict( + { + "export": None, + "loader": { + "task": "mask-generation", + "model_type": "sam", + "component_name": "prompt-decoder", + }, + } + ) + + assert config.loader.component_name == "prompt-decoder" + assert config.to_dict()["loader"]["component_name"] == "prompt-decoder" + + # ============================================================================= # Fixtures # ============================================================================= diff --git a/tests/unit/eval/test_mask_generation_evaluator.py b/tests/unit/eval/test_mask_generation_evaluator.py index c0e8964a4..d8032e531 100644 --- a/tests/unit/eval/test_mask_generation_evaluator.py +++ b/tests/unit/eval/test_mask_generation_evaluator.py @@ -171,6 +171,33 @@ def test_point_shape_and_scale(self) -> None: # boxes empty assert feed["input_boxes"].shape == (1, 0, 4) + def test_point_only_decoder_omits_unsupported_box_tensor(self) -> None: + feed = _build_decoder_inputs( + prompt={"point": [15, 25], "label": 1}, + prompt_mode="point", + scale_x=2.0, + scale_y=3.0, + emb={ + "image_embeddings": np.zeros((1, 256, 64, 64), dtype=np.float32), + "image_positional_embeddings": np.zeros((1, 256, 64, 64), dtype=np.float32), + }, + required_embed_names=("image_embeddings", "image_positional_embeddings"), + required_input_names=( + "input_points", + "input_labels", + "image_embeddings", + "image_positional_embeddings", + ), + ) + + assert "input_boxes" not in feed + assert set(feed) == { + "input_points", + "input_labels", + "image_embeddings", + "image_positional_embeddings", + } + class TestBuildDecoderInputsInvalidMode: def test_text_mode_rejected_with_helpful_message(self) -> None: diff --git a/tests/unit/loader/test_composite_tasks.py b/tests/unit/loader/test_composite_tasks.py index b4f200bc4..0df725ad2 100644 --- a/tests/unit/loader/test_composite_tasks.py +++ b/tests/unit/loader/test_composite_tasks.py @@ -49,6 +49,14 @@ def test_matches_composite_registry(self) -> None: "Update COMPOSITE_TASKS in src/winml/modelkit/loader/task.py." ) + def test_sam_mask_generation_has_published_graph_roles(self) -> None: + from winml.modelkit.loader.resolution import resolve_composite + + assert resolve_composite("sam", "mask-generation") == { + "image-encoder": "image-feature-extraction", + "prompt-decoder": "mask-generation", + } + class TestCompositeTasksVsKnownTasks: def test_summarization_translation_excluded_from_known_tasks(self) -> None: diff --git a/tests/unit/loader/test_onnx_composite_hub.py b/tests/unit/loader/test_onnx_composite_hub.py new file mode 100644 index 000000000..6c39b038b --- /dev/null +++ b/tests/unit/loader/test_onnx_composite_hub.py @@ -0,0 +1,82 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for graph-contract-driven composite ONNX discovery.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from winml.modelkit.loader import onnx_hub +from winml.modelkit.models.hf.sam import WinMLSAMModel + + +def _graph( + name: str, + *, + inputs: tuple[tuple[str, str, int], ...], + outputs: tuple[tuple[str, str, int], ...], + precision: str, +) -> onnx_hub._GraphContract: + return onnx_hub._GraphContract(Path(name), inputs, outputs, precision) + + +def test_resolver_selects_matching_graph_contract_and_falls_back_to_fp32(monkeypatch) -> None: + graphs = { + "encoder_fp16.onnx": None, + "decoder_fp16.onnx": None, + "encoder.onnx": _graph( + "encoder.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(("embedding", "tensor(float)", 4), ("position", "tensor(float)", 4)), + precision="fp32", + ), + "decoder.onnx": _graph( + "decoder.onnx", + inputs=( + ("points", "tensor(float)", 4), + ("labels", "tensor(int64)", 3), + ("embedding", "tensor(float)", 4), + ("position", "tensor(float)", 4), + ), + outputs=(("scores", "tensor(float)", 3), ("masks", "tensor(float)", 5)), + precision="fp32", + ), + } + monkeypatch.setattr( + "huggingface_hub.list_repo_files", + lambda *args, **kwargs: list(graphs), + ) + monkeypatch.setattr( + onnx_hub, + "resolve_hf_onnx_path", + lambda model_id, **kwargs: Path(model_id.rsplit("/", 1)[-1]), + ) + monkeypatch.setattr(onnx_hub, "_inspect_runnable_graph", lambda path: graphs[path.name]) + + result = onnx_hub.resolve_hf_onnx_encoder_decoder("org/model", precision="fp16") + + assert result == { + "image-encoder": Path("encoder.onnx"), + "prompt-decoder": Path("decoder.onnx"), + } + + +def test_sam_published_onnx_absence_preserves_pytorch_fallback(monkeypatch) -> None: + monkeypatch.setattr("huggingface_hub.list_repo_files", lambda *args, **kwargs: []) + + assert WinMLSAMModel.resolve_pretrained_onnx("facebook/sam-vit-base") is None + + +def test_sam_malformed_published_onnx_still_fails_closed(monkeypatch) -> None: + monkeypatch.setattr( + onnx_hub, + "resolve_hf_onnx_encoder_decoder", + lambda *args, **kwargs: (_ for _ in ()).throw(ValueError("ambiguous pair")), + ) + + with pytest.raises(ValueError, match="ambiguous pair"): + WinMLSAMModel.resolve_pretrained_onnx("org/malformed-sam") From 760a97364bb42cb8898786e9d8fe12ff3e0bfdaa Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Sun, 19 Jul 2026 20:53:51 +0800 Subject: [PATCH 2/5] fix(loader): reject ambiguous ONNX composite pairs --- src/winml/modelkit/commands/build.py | 2 + src/winml/modelkit/loader/onnx_hub.py | 47 +++++++- tests/unit/loader/test_onnx_composite_hub.py | 109 ++++++++++++++++++- 3 files changed, 153 insertions(+), 5 deletions(-) diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index 1eb6b9511..abc2fc996 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -1322,6 +1322,8 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: components = {submodel: components[submodel]} if components: + if model is None: + raise RuntimeError("Composite components require a Hugging Face model ID.") if use_cache: raise click.UsageError( "--use-cache is not supported for composite models. " diff --git a/src/winml/modelkit/loader/onnx_hub.py b/src/winml/modelkit/loader/onnx_hub.py index cdec55b2f..2b72b39e1 100644 --- a/src/winml/modelkit/loader/onnx_hub.py +++ b/src/winml/modelkit/loader/onnx_hub.py @@ -34,6 +34,7 @@ class _GraphContract: inputs: tuple[tuple[str, str, int], ...] outputs: tuple[tuple[str, str, int], ...] precision: str + has_quantized_weights: bool = False def _inspect_runnable_graph(path: Path) -> _GraphContract | None: @@ -58,6 +59,9 @@ def _inspect_runnable_graph(path: Path) -> _GraphContract | None: fp16 = type_counts[onnx.TensorProto.FLOAT16] fp32 = type_counts[onnx.TensorProto.FLOAT] precision = "fp16" if fp16 > fp32 else "fp32" + has_quantized_weights = bool( + type_counts[onnx.TensorProto.INT8] or type_counts[onnx.TensorProto.UINT8] + ) def contract(nodes: list[Any]) -> tuple[tuple[str, str, int], ...]: return tuple((node.name, node.type, len(node.shape)) for node in nodes) @@ -67,6 +71,7 @@ def contract(nodes: list[Any]) -> tuple[tuple[str, str, int], ...]: inputs=contract(session.get_inputs()), outputs=contract(session.get_outputs()), precision=precision, + has_quantized_weights=has_quantized_weights, ) @@ -109,14 +114,18 @@ def resolve_hf_onnx_encoder_decoder( graphs.append(graph) preferred_precisions = ("fp16", "fp32") if precision == "fp16" else ("fp32",) - candidates: list[tuple[int, int, str, _GraphContract, _GraphContract]] = [] + candidates: list[tuple[int, int, _GraphContract, _GraphContract, tuple[str, ...]]] = [] for encoder in graphs: encoder_inputs = encoder.inputs if len(encoder_inputs) != 1 or encoder_inputs[0][2] != 4: continue encoder_outputs = {name for name, _dtype, _rank in encoder.outputs} for decoder in graphs: - if decoder.path == encoder.path or decoder.precision != encoder.precision: + if ( + decoder.path == encoder.path + or decoder.precision != encoder.precision + or decoder.has_quantized_weights != encoder.has_quantized_weights + ): continue decoder_inputs = {name for name, _dtype, _rank in decoder.inputs} shared = encoder_outputs & decoder_inputs @@ -129,7 +138,10 @@ def resolve_hf_onnx_encoder_decoder( precision_rank = preferred_precisions.index(encoder.precision) except ValueError: continue - candidates.append((precision_rank, -len(shared), str(encoder.path), encoder, decoder)) + quantization_rank = int(encoder.has_quantized_weights) + candidates.append( + (precision_rank, quantization_rank, encoder, decoder, tuple(sorted(shared))) + ) if not candidates: contracts = [ @@ -146,7 +158,34 @@ def resolve_hf_onnx_encoder_decoder( f"for precision {precision!r}: {contracts}" ) - _rank, _shared, _path, encoder, decoder = min(candidates, key=lambda row: row[:3]) + best_source_rank = min(candidate[:2] for candidate in candidates) + preferred_candidates = sorted( + (candidate for candidate in candidates if candidate[:2] == best_source_rank), + key=lambda candidate: (str(candidate[2].path), str(candidate[3].path)), + ) + if len(preferred_candidates) != 1: + pairs = [ + { + "image-encoder": encoder.path.name, + "prompt-decoder": decoder.path.name, + "precision": encoder.precision, + "shared_outputs": list(shared), + } + for ( + _precision_rank, + _quantization_rank, + encoder, + decoder, + shared, + ) in preferred_candidates + ] + raise ValueError( + "Published ONNX graphs contain multiple valid encoder/decoder pairs " + f"for precision {precision!r}; unable to select one unambiguously. " + f"Candidate pairs: {pairs}" + ) + + _precision_rank, _quantization_rank, encoder, decoder, _shared = preferred_candidates[0] return {"image-encoder": encoder.path, "prompt-decoder": decoder.path} diff --git a/tests/unit/loader/test_onnx_composite_hub.py b/tests/unit/loader/test_onnx_composite_hub.py index 6c39b038b..be75b465c 100644 --- a/tests/unit/loader/test_onnx_composite_hub.py +++ b/tests/unit/loader/test_onnx_composite_hub.py @@ -20,8 +20,9 @@ def _graph( inputs: tuple[tuple[str, str, int], ...], outputs: tuple[tuple[str, str, int], ...], precision: str, + has_quantized_weights: bool = False, ) -> onnx_hub._GraphContract: - return onnx_hub._GraphContract(Path(name), inputs, outputs, precision) + return onnx_hub._GraphContract(Path(name), inputs, outputs, precision, has_quantized_weights) def test_resolver_selects_matching_graph_contract_and_falls_back_to_fp32(monkeypatch) -> None: @@ -65,6 +66,112 @@ def test_resolver_selects_matching_graph_contract_and_falls_back_to_fp32(monkeyp } +def test_resolver_prefers_unquantized_graph_family(monkeypatch) -> None: + graphs = { + "encoder.onnx": _graph( + "encoder.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(("embedding", "tensor(float)", 4),), + precision="fp32", + ), + "decoder.onnx": _graph( + "decoder.onnx", + inputs=( + ("labels", "tensor(int64)", 3), + ("embedding", "tensor(float)", 4), + ), + outputs=(("masks", "tensor(float)", 5),), + precision="fp32", + ), + "encoder_quantized.onnx": _graph( + "encoder_quantized.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(("embedding", "tensor(float)", 4),), + precision="fp32", + has_quantized_weights=True, + ), + "decoder_quantized.onnx": _graph( + "decoder_quantized.onnx", + inputs=( + ("labels", "tensor(int64)", 3), + ("embedding", "tensor(float)", 4), + ), + outputs=(("masks", "tensor(float)", 5),), + precision="fp32", + has_quantized_weights=True, + ), + } + monkeypatch.setattr("huggingface_hub.list_repo_files", lambda *args, **kwargs: list(graphs)) + monkeypatch.setattr( + onnx_hub, + "resolve_hf_onnx_path", + lambda model_id, **kwargs: Path(model_id.rsplit("/", 1)[-1]), + ) + monkeypatch.setattr(onnx_hub, "_inspect_runnable_graph", lambda path: graphs[path.name]) + + result = onnx_hub.resolve_hf_onnx_encoder_decoder("org/model") + + assert result == { + "image-encoder": Path("encoder.onnx"), + "prompt-decoder": Path("decoder.onnx"), + } + + +def test_resolver_rejects_multiple_valid_graph_pairs_as_ambiguous(monkeypatch) -> None: + graphs = { + "encoder_a.onnx": _graph( + "encoder_a.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(("embedding_a", "tensor(float)", 4),), + precision="fp32", + ), + "decoder_a.onnx": _graph( + "decoder_a.onnx", + inputs=( + ("labels", "tensor(int64)", 3), + ("embedding_a", "tensor(float)", 4), + ), + outputs=(("masks", "tensor(float)", 5),), + precision="fp32", + ), + "encoder_b.onnx": _graph( + "encoder_b.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(("embedding_b", "tensor(float)", 4),), + precision="fp32", + ), + "decoder_b.onnx": _graph( + "decoder_b.onnx", + inputs=( + ("labels", "tensor(int64)", 3), + ("embedding_b", "tensor(float)", 4), + ), + outputs=(("masks", "tensor(float)", 5),), + precision="fp32", + ), + } + monkeypatch.setattr( + "huggingface_hub.list_repo_files", + lambda *args, **kwargs: list(reversed(graphs)), + ) + monkeypatch.setattr( + onnx_hub, + "resolve_hf_onnx_path", + lambda model_id, **kwargs: Path(model_id.rsplit("/", 1)[-1]), + ) + monkeypatch.setattr(onnx_hub, "_inspect_runnable_graph", lambda path: graphs[path.name]) + + with pytest.raises(ValueError) as error: + onnx_hub.resolve_hf_onnx_encoder_decoder("org/ambiguous-model") + + message = str(error.value) + assert "multiple valid encoder/decoder pairs" in message + assert "unable to select one unambiguously" in message + assert message.index("encoder_a.onnx") < message.index("encoder_b.onnx") + assert "decoder_a.onnx" in message + assert "decoder_b.onnx" in message + + def test_sam_published_onnx_absence_preserves_pytorch_fallback(monkeypatch) -> None: monkeypatch.setattr("huggingface_hub.list_repo_files", lambda *args, **kwargs: []) From cfa98a5c8efda9058d2540ac633346861db4d1f7 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Tue, 21 Jul 2026 23:12:35 +0800 Subject: [PATCH 3/5] feat(loader): support release-asset ONNX composites --- ...-generation_fp16_config_image-encoder.json | 13 - ...generation_fp16_config_prompt-decoder.json | 13 - ...-generation_fp32_config_image-encoder.json | 10 - ...generation_fp32_config_prompt-decoder.json | 10 - src/winml/modelkit/commands/build.py | 61 +++- src/winml/modelkit/commands/config.py | 8 +- .../eval/mask_generation_evaluator.py | 300 ++++++++++++++-- src/winml/modelkit/loader/onnx_hub.py | 208 +++++++---- src/winml/modelkit/loader/release_assets.py | 338 ++++++++++++++++++ src/winml/modelkit/loader/resolution.py | 66 +++- .../eval/test_mask_generation_evaluator.py | 148 ++++++++ tests/unit/loader/test_onnx_composite_hub.py | 87 ++++- tests/unit/loader/test_release_assets.py | 218 +++++++++++ 13 files changed, 1305 insertions(+), 175 deletions(-) delete mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json delete mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json delete mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json delete mode 100644 examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json create mode 100644 src/winml/modelkit/loader/release_assets.py create mode 100644 tests/unit/loader/test_release_assets.py diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json deleted file mode 100644 index 2cb9e5235..000000000 --- a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_image-encoder.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "export": null, - "optim": {}, - "quant": { - "mode": "fp16", - "fp16_keep_io_types": true - }, - "compile": null, - "loader": { - "task": "image-feature-extraction", - "component_name": "image-encoder" - } -} diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json deleted file mode 100644 index 6c70cf55a..000000000 --- a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp16_config_prompt-decoder.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "export": null, - "optim": {}, - "quant": { - "mode": "fp16", - "fp16_keep_io_types": true - }, - "compile": null, - "loader": { - "task": "mask-generation", - "component_name": "prompt-decoder" - } -} diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json deleted file mode 100644 index 05a2f7c52..000000000 --- a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_image-encoder.json +++ /dev/null @@ -1,10 +0,0 @@ -{ - "export": null, - "optim": {}, - "quant": null, - "compile": null, - "loader": { - "task": "image-feature-extraction", - "component_name": "image-encoder" - } -} diff --git a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json b/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json deleted file mode 100644 index 7a7a86f27..000000000 --- a/examples/recipes/Xenova_slimsam-77-uniform/cpu/cpu/mask-generation_fp32_config_prompt-decoder.json +++ /dev/null @@ -1,10 +0,0 @@ -{ - "export": null, - "optim": {}, - "quant": null, - "compile": null, - "loader": { - "task": "mask-generation", - "component_name": "prompt-decoder" - } -} diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index abc2fc996..7668dcbed 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -942,6 +942,8 @@ def build( if model_input.kind is ModelInputKind.INVALID: raise click.UsageError(model_input.error or f"Invalid model input: {model}") + published_composite_sources: dict[str, str] | None = None + # Load or auto-generate config if config_file is not None: config_or_configs = _load_config( @@ -1002,15 +1004,40 @@ def build( ep=ep, ) else: - config_or_configs = generate_build_config( - model, - trust_remote_code=trust_remote_code, - device=device, - precision=precision, - ep=ep, - shape_config=shape_overrides, - override={"export": export_overrides} if export_overrides else None, - ) + try: + config_or_configs = generate_build_config( + model, + trust_remote_code=trust_remote_code, + device=device, + precision=precision, + ep=ep, + shape_config=shape_overrides, + override={"export": export_overrides} if export_overrides else None, + ) + except (ValueError, OSError) as config_error: + from ..loader.resolution import resolve_composite_onnx_sources + + published_composite_sources = resolve_composite_onnx_sources( + model, + precision=precision, + trust_remote_code=trust_remote_code, + ) + if published_composite_sources is None: + raise + if export_overrides or shape_overrides: + raise click.UsageError( + "--input-specs, --export-config, --dynamic-axes, and " + "--shape-config are not supported for a pre-exported " + "release-asset composite." + ) from config_error + first_source = next(iter(published_composite_sources.values())) + config_or_configs = generate_build_config( + onnx_path=first_source, + task="mask-generation", + device=device, + precision=precision, + ep=ep, + ) if not quant: config_or_configs.quant = None # Auto-generated configs: compile disabled by default unless @@ -1106,7 +1133,9 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: except ValueError as e: raise click.UsageError(f"Config validation failed: {e}") from e - model_is_onnx = model_input is not None and model_input.kind is ModelInputKind.ONNX_FILE + model_is_onnx = ( + model_input is not None and model_input.kind is ModelInputKind.ONNX_FILE + ) or published_composite_sources is not None preloaded_hf_config = _validate_loader_tasks_for_model( model_id=model, @@ -1273,7 +1302,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: # export command). A composite fans out into one build per # sub-component; a plain model builds to the single output dir. components = None - if model and not model_is_onnx: + if model and (not model_is_onnx or published_composite_sources is not None): try: from ..loader.resolution import resolve_composite_components @@ -1290,6 +1319,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: model, task=task_hint, model_type=model_type_hint, + precision=precision, trust_remote_code=trust_remote_code, ) except click.ClickException: @@ -1430,7 +1460,11 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: config_file=None, model_id=component_source or model, is_onnx=component_source is not None, - resolved_dir=resolved_dir, + resolved_dir=( + resolved_dir / name + if component_source is not None + else resolved_dir + ), rebuild=rebuild, cache_key=name, ep=ep, @@ -1567,6 +1601,9 @@ def _run_single_build( device=device, extra_kwargs=extra_kwargs, ) + from ..loader.release_assets import copy_release_contract_files + + copy_release_contract_files(Path(model_id), resolved_dir) else: stage_timings = _build_hf_pipeline( config=config, diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index 359fa5cca..99a7d1445 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -360,7 +360,11 @@ def config( # Check composite model registry: (model_type, task) -> multi-config pipeline_components = _resolve_composite_model_components( - hf_model, model_type, task, trust_remote_code=trust_remote_code + hf_model, + model_type, + task, + precision=precision, + trust_remote_code=trust_remote_code, ) if pipeline_components: from ..loader.resolution import ( @@ -612,6 +616,7 @@ def _resolve_composite_model_components( hf_model: str | None, model_type: str | None, task: str | None, + precision: str = "fp32", trust_remote_code: bool = False, ) -> dict[str, str] | None: """Resolve the composite ``_SUB_MODEL_CONFIG`` for a build, else None. @@ -626,6 +631,7 @@ def _resolve_composite_model_components( hf_model, task=task, model_type=model_type, + precision=precision, trust_remote_code=trust_remote_code, ) diff --git a/src/winml/modelkit/eval/mask_generation_evaluator.py b/src/winml/modelkit/eval/mask_generation_evaluator.py index 0514b4bd1..8524599da 100644 --- a/src/winml/modelkit/eval/mask_generation_evaluator.py +++ b/src/winml/modelkit/eval/mask_generation_evaluator.py @@ -36,6 +36,7 @@ from __future__ import annotations +import json import logging import os from dataclasses import dataclass @@ -112,6 +113,17 @@ class _MaskGenProfile: ) +PROMPT_ENCODER_PROFILE = _MaskGenProfile( + name="prompt-encoder", + target_size=1024, + mean=(0.0, 0.0, 0.0), + std=(1.0, 1.0, 1.0), + resize_mode="direct", +) + +_RELEASE_METADATA_NAME = "winml_release_metadata.json" + + # Back-compat module-level SAM 3 constant (preserved so existing imports # from tests/scripts keep working unchanged). _TARGET_SIZE = SAM3_PROFILE.target_size @@ -167,20 +179,42 @@ def __init__( "Use prompt_mode='bbox' or 'point'." ) self._enc_sess, self._dec_sess = self._load_sessions() + self._encoder_prompt_inputs = _resolve_point_prompt_inputs(self._enc_sess) self._decoder_input_names = _node_names(self._dec_sess.get_inputs()) - if "input_boxes" not in self._decoder_input_names: + self._decoder_prompt_inputs = _resolve_point_prompt_inputs(self._dec_sess) + if self._encoder_prompt_inputs and self._decoder_prompt_inputs: + raise ValueError( + "Prompt inputs are present on both composite components; routing is ambiguous." + ) + if not self._encoder_prompt_inputs and not self._decoder_prompt_inputs: + raise ValueError("Neither composite component exposes a point-prompt contract.") + self._prompt_component = ( + self._ENCODER_ROLE if self._encoder_prompt_inputs else self._DECODER_ROLE + ) + if self._prompt_component == self._ENCODER_ROLE: + if "prompt_mode" in mapping and self._prompt_mode == "bbox": + raise ValueError( + "The encoder prompt contract accepts points only; use prompt_mode='point'." + ) + self._prompt_mode = "point" + elif "input_boxes" not in self._decoder_input_names: if "prompt_mode" in mapping and self._prompt_mode == "bbox": raise ValueError( "The decoder does not accept box prompts; use prompt_mode='point'." ) self._prompt_mode = "point" - self._encoder_input_name = _resolve_encoder_input_name(self._enc_sess) + self._encoder_input_name = _resolve_encoder_input_name( + self._enc_sess, + excluded=set(self._encoder_prompt_inputs or ()), + ) self._encoder_output_names = _node_names(self._enc_sess.get_outputs()) self._embedding_input_names = _resolve_embedding_input_names( self._dec_sess, self._encoder_output_names, ) - self._decoder_output_names = _resolve_decoder_output_names(self._dec_sess) + self._mask_output_name, self._score_output_name = _resolve_decoder_output_names( + self._dec_sess + ) # Pick the per-family preprocessing profile from the encoder's # static input shape (falling back to a model_id heuristic, then # SAM 3). Threaded through preprocess + postprocess in _predict. @@ -327,31 +361,45 @@ def _predict(self, sample: dict[str, Any]) -> np.ndarray: self._profile, image, ) - enc_out = self._enc_sess.run(None, {self._encoder_input_name: pixel_values}) + encoder_feed = {self._encoder_input_name: pixel_values} + if self._encoder_prompt_inputs is not None: + encoder_feed.update( + _build_encoder_prompt_inputs( + prompt=prompt, + prompt_mode=self._prompt_mode, + session=self._enc_sess, + coordinate_name=self._encoder_prompt_inputs[0], + label_name=self._encoder_prompt_inputs[1], + original_width=image.size[0], + original_height=image.size[1], + ) + ) + enc_out = self._enc_sess.run(None, encoder_feed) emb = dict(zip(self._encoder_output_names, enc_out, strict=True)) - dec_inputs = _build_decoder_inputs( - prompt=prompt, - prompt_mode=self._prompt_mode, - scale_x=scale_x, - scale_y=scale_y, - emb=emb, - required_embed_names=self._embedding_input_names, - required_input_names=self._decoder_input_names, - ) + if self._encoder_prompt_inputs is not None: + dec_inputs = {name: emb[name] for name in self._embedding_input_names} + else: + dec_inputs = _build_decoder_inputs( + prompt=prompt, + prompt_mode=self._prompt_mode, + scale_x=scale_x, + scale_y=scale_y, + emb=emb, + required_embed_names=self._embedding_input_names, + required_input_names=self._decoder_input_names, + ) dec_out = self._dec_sess.run( - list(self._decoder_output_names), + [self._score_output_name, self._mask_output_name], dec_inputs, ) - dec_by_name = dict(zip(self._decoder_output_names, dec_out, strict=True)) - iou_scores = dec_by_name["iou_scores"] - pred_masks = dec_by_name["pred_masks"] - - # pred_masks: (1, num_prompts, num_masks, H, W); pick the - # best-scoring of the candidate masks for the first prompt. - iou_preds = iou_scores[0, 0] # (num_masks,) - best_idx = int(iou_preds.argmax()) - best_low_res = pred_masks[0, 0, best_idx] + dec_by_name = dict( + zip((self._score_output_name, self._mask_output_name), dec_out, strict=True) + ) + best_low_res = _select_best_mask( + dec_by_name[self._mask_output_name], + dec_by_name[self._score_output_name], + ) return _postprocess_for_profile( self._profile, @@ -386,6 +434,9 @@ def _resolve_profile( ``sam3`` / ``sam-3`` -> SAM 3) when the encoder shape is dynamic. 3. **Default SAM 3** -- preserves the original evaluator behaviour. """ + if _resolve_point_prompt_inputs(enc_sess) is not None: + return _resolve_release_profile(config, enc_sess) + known = (SAM3_PROFILE, SAM2_PROFILE) try: @@ -406,6 +457,80 @@ def _resolve_profile( return SAM3_PROFILE +def _resolve_release_profile(config: WinMLEvaluationConfig, enc_sess: Any) -> _MaskGenProfile: + """Resolve image preprocessing from a release runtime contract, never an ID heuristic.""" + if not isinstance(config.model_path, dict): + raise TypeError("Release-backed mask generation requires composite model paths.") + encoder_path = Path(config.model_path[WinMLMaskGenerationEvaluator._ENCODER_ROLE]) + candidates = ( + encoder_path.parent / _RELEASE_METADATA_NAME, + encoder_path.parent / "metadata.json", + ) + metadata_path = next((path for path in candidates if path.is_file()), None) + if metadata_path is None: + raise ValueError( + "A prompt-bearing encoder requires persisted release runtime metadata; " + f"expected {_RELEASE_METADATA_NAME!r} beside {encoder_path.name!r}." + ) + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8-sig")) + model_files = metadata["model_files"] + except (OSError, json.JSONDecodeError, KeyError, TypeError) as error: + raise ValueError(f"Invalid release runtime metadata {metadata_path}: {error}") from error + if not isinstance(model_files, dict): + raise TypeError("Release runtime metadata 'model_files' must be an object.") + + session_inputs = {node.name: node for node in enc_sess.get_inputs()} + matching = [ + contract + for contract in model_files.values() + if isinstance(contract, dict) + and isinstance(contract.get("inputs"), dict) + and set(contract["inputs"]) == set(session_inputs) + ] + if len(matching) != 1: + raise ValueError( + "Release runtime metadata must contain exactly one graph contract matching " + f"encoder inputs {sorted(session_inputs)}; found {len(matching)}." + ) + inputs = matching[0]["inputs"] + image_names = [ + name + for name, spec in inputs.items() + if isinstance(spec, dict) and spec.get("io_type") == "image" + ] + if len(image_names) != 1: + raise ValueError( + "Release runtime metadata must identify exactly one encoder input as io_type=image." + ) + image_name = image_names[0] + image_spec = inputs[image_name] + shape = image_spec.get("shape") + value_range = image_spec.get("value_range") + session_shape = list(getattr(session_inputs[image_name], "shape", ())) + if ( + not isinstance(shape, list) + or len(shape) != 4 + or shape != session_shape + or shape[1] != 3 + or not isinstance(shape[-1], int) + or shape[-2] != shape[-1] + or image_spec.get("dtype") != "float32" + or value_range != [0.0, 1.0] + ): + raise ValueError( + "Unsupported or inconsistent release image contract; expected a matching " + "square NCHW RGB float32 input with value_range [0.0, 1.0]." + ) + return _MaskGenProfile( + name=PROMPT_ENCODER_PROFILE.name, + target_size=shape[-1], + mean=PROMPT_ENCODER_PROFILE.mean, + std=PROMPT_ENCODER_PROFILE.std, + resize_mode=PROMPT_ENCODER_PROFILE.resize_mode, + ) + + def _node_names(nodes: Any) -> tuple[str, ...]: """Return ORT input/output names, rejecting unnamed nodes.""" names = tuple(getattr(node, "name", "") for node in nodes) @@ -414,13 +539,18 @@ def _node_names(nodes: Any) -> tuple[str, ...]: return names -def _resolve_encoder_input_name(enc_sess: Any) -> str: +def _resolve_encoder_input_name(enc_sess: Any, *, excluded: set[str] | None = None) -> str: """Pick the encoder image input name from the actual ONNX session.""" - names = _node_names(enc_sess.get_inputs()) + excluded = excluded or set() + inputs = [node for node in enc_sess.get_inputs() if node.name not in excluded] + names = _node_names(inputs) if not names: raise ValueError("Encoder ONNX session has no inputs.") if "pixel_values" in names: return "pixel_values" + rank4 = [str(node.name) for node in inputs if len(getattr(node, "shape", ())) == 4] + if len(rank4) == 1: + return rank4[0] if len(names) == 1: return names[0] raise ValueError( @@ -435,14 +565,6 @@ def _resolve_embedding_input_names( ) -> tuple[str, ...]: """Match decoder embedding inputs to encoder output names.""" decoder_inputs = _node_names(dec_sess.get_inputs()) - required_prompt_inputs = ("input_points", "input_labels") - missing_prompt = [name for name in required_prompt_inputs if name not in decoder_inputs] - if missing_prompt: - raise ValueError( - f"Decoder ONNX session missing prompt input(s) {missing_prompt}. " - f"Got inputs: {list(decoder_inputs)}." - ) - embedding_inputs = tuple(name for name in decoder_inputs if name in encoder_output_names) if not embedding_inputs: raise ValueError( @@ -455,14 +577,114 @@ def _resolve_embedding_input_names( def _resolve_decoder_output_names(dec_sess: Any) -> tuple[str, str]: """Validate and return decoder outputs needed by the metric path.""" - outputs = _node_names(dec_sess.get_outputs()) - required = ("iou_scores", "pred_masks") - missing = [name for name in required if name not in outputs] - if missing: + nodes = dec_sess.get_outputs() + outputs = _node_names(nodes) + mask_candidates = [node.name for node in nodes if len(getattr(node, "shape", ())) >= 4] + score_candidates = [node.name for node in nodes if 1 <= len(getattr(node, "shape", ())) <= 3] + if "pred_masks" in outputs: + mask_candidates = ["pred_masks"] + elif "masks" in outputs: + mask_candidates = ["masks"] + if "iou_scores" in outputs: + score_candidates = ["iou_scores"] + elif "scores" in outputs: + score_candidates = ["scores"] + if len(mask_candidates) != 1 or len(score_candidates) != 1: + raise ValueError( + "Could not identify one mask output and one score output from decoder " + f"contract. Got outputs: {list(outputs)}." + ) + return mask_candidates[0], score_candidates[0] + + +def _resolve_point_prompt_inputs(session: Any) -> tuple[str, str] | None: + """Resolve one coordinate/label input pair by rank/shape, failing on ambiguity.""" + nodes = session.get_inputs() + coordinates = [] + for node in nodes: + shape = getattr(node, "shape", ()) + if isinstance(shape, (list, tuple)) and len(shape) in {3, 4} and shape[-1] == 2: + coordinates.append(node) + labels = [ + node + for node in nodes + if len(getattr(node, "shape", ())) in {2, 3} + and node not in coordinates + and ( + "int" in getattr(node, "type", "").lower() + or "label" in getattr(node, "name", "").lower() + ) + ] + if not coordinates and not labels: + return None + if len(coordinates) != 1 or len(labels) != 1: + raise ValueError( + "Point-prompt inputs are incomplete or ambiguous: " + f"coordinates={[node.name for node in coordinates]}, " + f"labels={[node.name for node in labels]}." + ) + return coordinates[0].name, labels[0].name + + +def _numpy_dtype(node_type: str) -> Any: + lowered = node_type.lower() + if "int64" in lowered: + return np.int64 + if "int32" in lowered: + return np.int32 + return np.float32 + + +def _build_encoder_prompt_inputs( + *, + prompt: dict[str, Any], + prompt_mode: str, + session: Any, + coordinate_name: str, + label_name: str, + original_width: int, + original_height: int, +) -> dict[str, np.ndarray]: + """Build normalized point tensors for a prompt-bearing encoder contract.""" + if prompt_mode != "point": + raise ValueError("Encoder-side prompt routing currently requires prompt_mode='point'.") + nodes = {node.name: node for node in session.get_inputs()} + coordinate = nodes[coordinate_name] + label = nodes[label_name] + coordinate_shape = tuple(coordinate.shape) + label_shape = tuple(label.shape) + if any(not isinstance(dim, int) or dim <= 0 for dim in coordinate_shape + label_shape): + raise ValueError( + "Encoder-side point prompts require positive static coordinate and label shapes." + ) + coords = np.zeros(coordinate_shape, dtype=_numpy_dtype(coordinate.type)) + labels = np.full(label_shape, -1, dtype=_numpy_dtype(label.type)) + px, py = prompt["point"] + normalized = (px / original_width, py / original_height) + if len(coordinate_shape) == 3: + coords[0, 0] = normalized + labels[0, 0] = 1 + else: + coords[0, 0, 0] = normalized + labels[0, 0, 0] = 1 + return {coordinate_name: coords, label_name: labels} + + +def _select_best_mask(masks: np.ndarray, scores: np.ndarray) -> np.ndarray: + """Select the highest-scoring low-resolution mask across common SAM layouts.""" + if masks.ndim == 4: + candidates = masks[0] + score_values = scores.reshape(-1) + elif masks.ndim == 5: + candidates = masks[0, 0] + score_values = scores[0, 0].reshape(-1) + else: + raise ValueError(f"Unsupported mask output rank {masks.ndim}; expected 4 or 5.") + if candidates.shape[0] != score_values.size: raise ValueError( - f"Decoder ONNX session missing output(s) {missing}. Got outputs: {list(outputs)}." + f"Mask/score candidate count mismatch: {candidates.shape[0]} vs {score_values.size}." ) - return required + return np.asarray(candidates[int(score_values.argmax())]) def _preprocess_for_profile( diff --git a/src/winml/modelkit/loader/onnx_hub.py b/src/winml/modelkit/loader/onnx_hub.py index 2b72b39e1..41d56082f 100644 --- a/src/winml/modelkit/loader/onnx_hub.py +++ b/src/winml/modelkit/loader/onnx_hub.py @@ -35,6 +35,16 @@ class _GraphContract: outputs: tuple[tuple[str, str, int], ...] precision: str has_quantized_weights: bool = False + input_shapes: tuple[tuple[Any, ...], ...] = () + output_shapes: tuple[tuple[Any, ...], ...] = () + + +@dataclass(frozen=True) +class _CompositeGraphPair: + encoder: _GraphContract + decoder: _GraphContract + shared_outputs: tuple[str, ...] + prompt_component: str def _inspect_runnable_graph(path: Path) -> _GraphContract | None: @@ -72,52 +82,45 @@ def contract(nodes: list[Any]) -> tuple[tuple[str, str, int], ...]: outputs=contract(session.get_outputs()), precision=precision, has_quantized_weights=has_quantized_weights, + input_shapes=tuple(tuple(node.shape) for node in session.get_inputs()), + output_shapes=tuple(tuple(node.shape) for node in session.get_outputs()), ) -def resolve_hf_onnx_encoder_decoder( - model_id: str, - *, - revision: str | None = None, - precision: str = "fp32", - cache_dir: str | Path | None = None, - token: str | bool | None = None, -) -> dict[str, Path]: - """Discover a runnable image-encoder/prompt-decoder pair in a Hub repo. +def _has_point_prompt(contract: _GraphContract) -> bool: + """Whether a graph contract exposes one coordinate/label prompt pair.""" + ports = contract.inputs + has_integer_labels = any( + rank in {2, 3} and "int" in dtype.lower() for _name, dtype, rank in ports + ) + has_point_tensor = any( + rank in {3, 4} and "float" in dtype.lower() for _name, dtype, rank in ports + ) + # Some deployment exports cast labels to float. Their unambiguous prompt + # shape is coordinates [B,N,2] plus labels [B,N]. Embedding-only decoders + # have no rank-2 input, so this does not mistake sparse embeddings for prompts. + has_float_pair = any(rank == 2 for _name, _dtype, rank in ports) and any( + rank == 3 for _name, _dtype, rank in ports + ) + return has_point_tensor and (has_integer_labels or has_float_pair) - Selection is graph-contract driven rather than filename driven. The encoder - has one rank-4 tensor input; the decoder consumes one or more encoder outputs - plus a rank-3 integer prompt tensor. Invalid published variants are ignored - after an ORT CPU session probe. An fp16 build may safely fall back to fp32 - source graphs because the normal build pipeline performs fp16 conversion. - """ - from huggingface_hub import list_repo_files - files = sorted( - name - for name in list_repo_files(model_id, revision=revision, token=token) - if name.lower().endswith(".onnx") - ) - if not files: - raise FileNotFoundError(f"No published ONNX graphs found in Hub repo {model_id!r}.") +def _has_mask_outputs(contract: _GraphContract) -> bool: + ranks = [rank for _name, _dtype, rank in contract.outputs] + return sum(rank >= 4 for rank in ranks) == 1 and sum(1 <= rank <= 3 for rank in ranks) == 1 - graphs: list[_GraphContract] = [] - for filename in files: - path = resolve_hf_onnx_path( - f"{model_id}/{filename}", - revision=revision, - cache_dir=cache_dir, - token=token, - ) - graph = _inspect_runnable_graph(path) - if graph is not None: - graphs.append(graph) +def _select_encoder_decoder_pair( + graphs: list[_GraphContract], + *, + precision: str, +) -> _CompositeGraphPair: + """Select one graph pair by connectivity, image, prompt, and mask contracts.""" preferred_precisions = ("fp16", "fp32") if precision == "fp16" else ("fp32",) - candidates: list[tuple[int, int, _GraphContract, _GraphContract, tuple[str, ...]]] = [] + candidates: list[tuple[int, int, _CompositeGraphPair]] = [] for encoder in graphs: - encoder_inputs = encoder.inputs - if len(encoder_inputs) != 1 or encoder_inputs[0][2] != 4: + image_inputs = [port for port in encoder.inputs if port[2] == 4] + if len(image_inputs) != 1: continue encoder_outputs = {name for name, _dtype, _rank in encoder.outputs} for decoder in graphs: @@ -125,23 +128,28 @@ def resolve_hf_onnx_encoder_decoder( decoder.path == encoder.path or decoder.precision != encoder.precision or decoder.has_quantized_weights != encoder.has_quantized_weights + or not _has_mask_outputs(decoder) ): continue decoder_inputs = {name for name, _dtype, _rank in decoder.inputs} - shared = encoder_outputs & decoder_inputs - has_integer_prompt = any( - rank == 3 and "int" in dtype.lower() for _name, dtype, rank in decoder.inputs - ) - if not shared or not has_integer_prompt: + shared = tuple(sorted(encoder_outputs & decoder_inputs)) + if not shared: + continue + encoder_prompt = _has_point_prompt(encoder) + decoder_prompt = _has_point_prompt(decoder) + if encoder_prompt == decoder_prompt: continue try: precision_rank = preferred_precisions.index(encoder.precision) except ValueError: continue - quantization_rank = int(encoder.has_quantized_weights) - candidates.append( - (precision_rank, quantization_rank, encoder, decoder, tuple(sorted(shared))) + pair = _CompositeGraphPair( + encoder=encoder, + decoder=decoder, + shared_outputs=shared, + prompt_component="image-encoder" if encoder_prompt else "prompt-decoder", ) + candidates.append((precision_rank, int(encoder.has_quantized_weights), pair)) if not candidates: contracts = [ @@ -154,39 +162,110 @@ def resolve_hf_onnx_encoder_decoder( for graph in graphs ] raise ValueError( - "Published ONNX graphs do not contain a runnable encoder/decoder pair " - f"for precision {precision!r}: {contracts}" + "Published ONNX graphs do not contain a runnable promptable " + f"encoder/decoder pair for precision {precision!r}: {contracts}" ) best_source_rank = min(candidate[:2] for candidate in candidates) - preferred_candidates = sorted( - (candidate for candidate in candidates if candidate[:2] == best_source_rank), - key=lambda candidate: (str(candidate[2].path), str(candidate[3].path)), + preferred = sorted( + (candidate[2] for candidate in candidates if candidate[:2] == best_source_rank), + key=lambda pair: (str(pair.encoder.path), str(pair.decoder.path)), ) - if len(preferred_candidates) != 1: + if len(preferred) != 1: pairs = [ { - "image-encoder": encoder.path.name, - "prompt-decoder": decoder.path.name, - "precision": encoder.precision, - "shared_outputs": list(shared), + "image-encoder": pair.encoder.path.name, + "prompt-decoder": pair.decoder.path.name, + "precision": pair.encoder.precision, + "shared_outputs": list(pair.shared_outputs), + "prompt_component": pair.prompt_component, } - for ( - _precision_rank, - _quantization_rank, - encoder, - decoder, - shared, - ) in preferred_candidates + for pair in preferred ] raise ValueError( "Published ONNX graphs contain multiple valid encoder/decoder pairs " f"for precision {precision!r}; unable to select one unambiguously. " f"Candidate pairs: {pairs}" ) + return preferred[0] + + +def resolve_hf_onnx_encoder_decoder( + model_id: str, + *, + revision: str | None = None, + precision: str = "fp32", + cache_dir: str | Path | None = None, + token: str | bool | None = None, +) -> dict[str, Path]: + """Discover a runnable image-encoder/prompt-decoder pair in a Hub repo. + + Selection is graph-contract driven rather than filename driven. The encoder + has one rank-4 tensor input; the decoder consumes one or more encoder outputs + plus a rank-3 integer prompt tensor. Invalid published variants are ignored + after an ORT CPU session probe. An fp16 build may safely fall back to fp32 + source graphs because the normal build pipeline performs fp16 conversion. + """ + from huggingface_hub import list_repo_files + + files = sorted( + name + for name in list_repo_files(model_id, revision=revision, token=token) + if name.lower().endswith(".onnx") + ) + if not files: + raise FileNotFoundError(f"No published ONNX graphs found in Hub repo {model_id!r}.") + + graphs: list[_GraphContract] = [] + for filename in files: + path = resolve_hf_onnx_path( + f"{model_id}/{filename}", + revision=revision, + cache_dir=cache_dir, + token=token, + ) + graph = _inspect_runnable_graph(path) + if graph is not None: + graphs.append(graph) + + pair = _select_encoder_decoder_pair(graphs, precision=precision) + return {"image-encoder": pair.encoder.path, "prompt-decoder": pair.decoder.path} + - _precision_rank, _quantization_rank, encoder, decoder, _shared = preferred_candidates[0] - return {"image-encoder": encoder.path, "prompt-decoder": decoder.path} +def resolve_hf_release_onnx_encoder_decoder( + model_id: str, + *, + revision: str | None = None, + precision: str = "fp32", + cache_dir: str | Path | None = None, + token: str | bool | None = None, +) -> dict[str, Path] | None: + """Resolve a config-less release archive into one promptable graph pair.""" + from .release_assets import acquire_hf_release_asset + + release = acquire_hf_release_asset( + model_id, + revision=revision, + precision=precision, + format_name="onnx", + cache_dir=cache_dir, + token=token, + ) + if release is None: + return None + pipeline_tag = release.provenance.get("pipeline_tag") + if pipeline_tag not in {"image-segmentation", "mask-generation"}: + raise ValueError( + f"Release archive graph pair requires an image-segmentation or " + f"mask-generation pipeline tag; got {pipeline_tag!r}." + ) + graphs = [ + graph + for path in sorted(release.root.rglob("*.onnx")) + if (graph := _inspect_runnable_graph(path)) is not None + ] + pair = _select_encoder_decoder_pair(graphs, precision=precision) + return {"image-encoder": pair.encoder.path, "prompt-decoder": pair.decoder.path} def resolve_hf_onnx_path( @@ -335,4 +414,5 @@ def _format_available_onnx_files( __all__ = [ "resolve_hf_onnx_encoder_decoder", "resolve_hf_onnx_path", + "resolve_hf_release_onnx_encoder_decoder", ] diff --git a/src/winml/modelkit/loader/release_assets.py b/src/winml/modelkit/loader/release_assets.py new file mode 100644 index 000000000..6c04427c1 --- /dev/null +++ b/src/winml/modelkit/loader/release_assets.py @@ -0,0 +1,338 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Acquire versioned model archives declared by Hub ``release_assets.json`` files. + +The resolver is intentionally metadata-driven. It does not know model IDs or +archive member names: callers select a precision/format tuple from the manifest, +then inspect the safely extracted artifacts by their graph contracts. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import shutil +import stat +import tempfile +import zipfile +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any +from urllib.parse import urlparse + + +logger = logging.getLogger(__name__) + +_MANIFEST_NAME = "release_assets.json" +_PROVENANCE_NAME = "winml_release_provenance.json" +_RUNTIME_METADATA_NAME = "winml_release_metadata.json" +_MAX_EXTRACTED_BYTES = 20 * 1024**3 + + +@dataclass(frozen=True) +class AcquiredReleaseAsset: + """A safely extracted immutable release asset and its provenance.""" + + root: Path + manifest_path: Path + metadata_path: Path | None + provenance_path: Path + provenance: dict[str, Any] + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _manifest_asset(manifest: dict[str, Any], precision: str, format_name: str) -> dict[str, Any]: + """Select exactly one release tuple, using float source for fp16 conversion.""" + source_precision = "float" if precision in {"fp32", "fp16"} else precision + try: + asset = manifest["precisions"][source_precision]["universal_assets"][format_name] + except (KeyError, TypeError) as error: + raise ValueError( + f"Release manifest has no {format_name!r} asset for requested precision " + f"{precision!r} (source precision {source_precision!r})." + ) from error + if not isinstance(asset, dict) or not isinstance(asset.get("download_url"), str): + raise TypeError( + f"Release manifest entry for {source_precision!r}/{format_name!r} must " + "contain one string download_url." + ) + return asset + + +def _cache_root(cache_dir: str | Path | None) -> Path: + if cache_dir is not None: + return Path(cache_dir) + from huggingface_hub.constants import HF_HUB_CACHE + + return Path(HF_HUB_CACHE) / "winml-release-assets" + + +def _download_archive(url: str, destination: Path) -> None: + """Stream an archive to a temporary file and atomically publish it.""" + import httpx + + parsed = urlparse(url) + if parsed.scheme != "https" or not parsed.netloc: + raise ValueError(f"Release asset download_url must be an absolute HTTPS URL: {url!r}") + + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(destination.suffix + ".part") + temporary.unlink(missing_ok=True) + try: + with httpx.stream("GET", url, follow_redirects=True, timeout=120.0) as response: + response.raise_for_status() + with temporary.open("wb") as stream: + for chunk in response.iter_bytes(chunk_size=1024 * 1024): + stream.write(chunk) + temporary.replace(destination) + except BaseException: + temporary.unlink(missing_ok=True) + raise + + +def _safe_member_path(info: zipfile.ZipInfo) -> PurePosixPath: + """Validate one ZIP member and return a normalized relative path.""" + raw = info.filename.replace("\\", "/") + member = PurePosixPath(raw) + mode = info.external_attr >> 16 + if ( + not raw + or member.is_absolute() + or any(part in {"", ".", ".."} for part in member.parts) + or (member.parts and ":" in member.parts[0]) + or stat.S_ISLNK(mode) + ): + raise ValueError(f"Unsafe ZIP member path or type: {info.filename!r}") + return member + + +def safe_extract_zip(archive: Path, destination: Path) -> None: + """Extract a ZIP without ``extractall`` and fail closed on unsafe members.""" + seen: set[PurePosixPath] = set() + total_size = 0 + with zipfile.ZipFile(archive) as bundle: + members = bundle.infolist() + for info in members: + member = _safe_member_path(info) + if member in seen: + raise ValueError(f"ZIP contains duplicate member {info.filename!r}.") + seen.add(member) + total_size += info.file_size + if total_size > _MAX_EXTRACTED_BYTES: + raise ValueError( + f"ZIP expands beyond the {_MAX_EXTRACTED_BYTES}-byte safety limit." + ) + + destination.mkdir(parents=True, exist_ok=False) + root = destination.resolve() + for info in members: + member = _safe_member_path(info) + target = destination.joinpath(*member.parts) + resolved = target.resolve() + if resolved != root and root not in resolved.parents: + raise ValueError(f"ZIP member escapes extraction root: {info.filename!r}") + if info.is_dir(): + target.mkdir(parents=True, exist_ok=True) + continue + target.parent.mkdir(parents=True, exist_ok=True) + with bundle.open(info) as source, target.open("xb") as output: + shutil.copyfileobj(source, output, length=1024 * 1024) + + +def validate_onnx_external_data(path: Path, root: Path) -> tuple[Path, ...]: + """Require every ONNX external-data location to be a present in-root file.""" + import onnx + + model = onnx.load(path, load_external_data=False) + sidecars: list[Path] = [] + root_resolved = root.resolve() + for initializer in model.graph.initializer: + if initializer.data_location != onnx.TensorProto.EXTERNAL: + continue + fields = {field.key: field.value for field in initializer.external_data} + location = fields.get("location") + if not location: + raise ValueError(f"External initializer in {path.name!r} has no location.") + relative = PurePosixPath(location.replace("\\", "/")) + if ( + relative.is_absolute() + or any(part in {"", ".", ".."} for part in relative.parts) + or (relative.parts and ":" in relative.parts[0]) + ): + raise ValueError( + f"ONNX graph {path.name!r} has unsafe external-data location {location!r}." + ) + sidecar = path.parent.joinpath(*relative.parts).resolve() + if root_resolved not in sidecar.parents or not sidecar.is_file(): + raise FileNotFoundError( + f"ONNX graph {path.name!r} requires missing or out-of-root external data " + f"{location!r}." + ) + sidecars.append(sidecar) + return tuple(dict.fromkeys(sidecars)) + + +def copy_release_contract_files(source_graph: Path, destination: Path) -> None: + """Copy immutable release provenance/runtime metadata beside a built graph.""" + release_root = next( + ( + parent + for parent in (source_graph.parent, *source_graph.parents) + if (parent / _PROVENANCE_NAME).is_file() + ), + None, + ) + if release_root is None: + return + metadata_files = sorted(release_root.rglob("metadata.json")) + if len(metadata_files) != 1: + raise ValueError( + "Release-backed ONNX graph requires exactly one metadata.json; " + f"found {[str(path.relative_to(release_root)) for path in metadata_files]}." + ) + destination.mkdir(parents=True, exist_ok=True) + shutil.copy2(release_root / _PROVENANCE_NAME, destination / _PROVENANCE_NAME) + shutil.copy2(metadata_files[0], destination / _RUNTIME_METADATA_NAME) + + +def acquire_hf_release_asset( + model_id: str, + *, + revision: str | None = None, + precision: str = "fp32", + format_name: str = "onnx", + cache_dir: str | Path | None = None, + token: str | bool | None = None, +) -> AcquiredReleaseAsset | None: + """Resolve and safely acquire a Hub repository's release archive. + + ``None`` means the repository does not publish ``release_assets.json``. + A present but malformed manifest/archive always raises; it never falls back + to another loader and thereby hides ambiguous or incomplete provenance. + """ + from huggingface_hub import HfApi, hf_hub_download + from huggingface_hub.errors import EntryNotFoundError + + info = HfApi().model_info(model_id, revision=revision, token=token) + resolved_revision = info.sha + try: + manifest_path = Path( + hf_hub_download( + repo_id=model_id, + filename=_MANIFEST_NAME, + revision=resolved_revision, + cache_dir=cache_dir, + token=token, + ) + ) + except EntryNotFoundError: + return None + + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8-sig")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"Invalid {_MANIFEST_NAME} for {model_id!r}: {error}") from error + if not isinstance(manifest, dict): + raise TypeError(f"{_MANIFEST_NAME} for {model_id!r} must contain a JSON object.") + + asset = _manifest_asset(manifest, precision, format_name) + url = asset["download_url"] + key_payload = json.dumps( + [model_id, resolved_revision, precision, format_name, url], separators=(",", ":") + ) + key = hashlib.sha256(key_payload.encode()).hexdigest() + asset_root = _cache_root(cache_dir) / key + archive_name = Path(urlparse(url).path).name or "release.zip" + archive_path = asset_root / archive_name + extracted = asset_root / "extracted" + provenance_path = extracted / _PROVENANCE_NAME + + if provenance_path.is_file(): + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + if provenance.get("archive_sha256") == _sha256(archive_path): + metadata_files = sorted(extracted.rglob("metadata.json")) + return AcquiredReleaseAsset( + root=extracted, + manifest_path=manifest_path, + metadata_path=metadata_files[0] if len(metadata_files) == 1 else None, + provenance_path=provenance_path, + provenance=provenance, + ) + + asset_root.mkdir(parents=True, exist_ok=True) + if not archive_path.is_file(): + _download_archive(url, archive_path) + + staging = Path(tempfile.mkdtemp(prefix="extract-", dir=asset_root)) + try: + shutil.rmtree(staging) + safe_extract_zip(archive_path, staging) + graph_paths = sorted(staging.rglob("*.onnx")) + if not graph_paths: + raise ValueError(f"Release archive {url!r} contains no ONNX graphs.") + external_data = { + str(path.relative_to(staging)): [ + str(sidecar.relative_to(staging)) + for sidecar in validate_onnx_external_data(path, staging) + ] + for path in graph_paths + } + provenance = { + "schema_version": 1, + "repo_id": model_id, + "requested_revision": revision, + "resolved_revision": resolved_revision, + "pipeline_tag": getattr(info, "pipeline_tag", None), + "manifest_sha256": _sha256(manifest_path), + "manifest_version": manifest.get("version"), + "requested_precision": precision, + "source_precision": "float" if precision in {"fp32", "fp16"} else precision, + "format": format_name, + "download_url": url, + "archive_sha256": _sha256(archive_path), + "tool_versions": asset.get("tool_versions", {}), + "graphs": {str(path.relative_to(staging)): _sha256(path) for path in graph_paths}, + "external_data": external_data, + } + (staging / _PROVENANCE_NAME).write_text( + json.dumps(provenance, indent=2, sort_keys=True), encoding="utf-8" + ) + if extracted.exists(): + shutil.rmtree(extracted) + staging.replace(extracted) + except BaseException: + shutil.rmtree(staging, ignore_errors=True) + raise + + metadata_files = sorted(extracted.rglob("metadata.json")) + if len(metadata_files) > 1: + raise ValueError( + f"Release archive contains multiple metadata.json files: " + f"{[str(path.relative_to(extracted)) for path in metadata_files]}" + ) + return AcquiredReleaseAsset( + root=extracted, + manifest_path=manifest_path, + metadata_path=metadata_files[0] if metadata_files else None, + provenance_path=provenance_path, + provenance=provenance, + ) + + +__all__ = [ + "AcquiredReleaseAsset", + "acquire_hf_release_asset", + "copy_release_contract_files", + "safe_extract_zip", + "validate_onnx_external_data", +] diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 2f1f1848b..de04f65e3 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -317,6 +317,7 @@ def resolve_composite_components( *, task: str | None = None, model_type: str | None = None, + precision: str = "fp32", trust_remote_code: bool = False, ) -> CompositeComponents | None: """Resolve a composite model's ``_SUB_MODEL_CONFIG`` (sub-name -> task), else None. @@ -331,21 +332,45 @@ def resolve_composite_components( """ from transformers import AutoConfig + config = None + config_error: ValueError | OSError | None = None + if hf_model is not None: + try: + config = AutoConfig.from_pretrained(hf_model, trust_remote_code=trust_remote_code) + except (ValueError, OSError) as error: + config_error = error + from .onnx_hub import resolve_hf_release_onnx_encoder_decoder + + release_sources = resolve_hf_release_onnx_encoder_decoder( + hf_model, + precision=precision, + ) + if release_sources is not None: + if task not in {None, "mask-generation"}: + raise ValueError( + f"Published promptable ONNX composite serves 'mask-generation', " + f"not {task!r}." + ) from error + return { + "image-encoder": "image-feature-extraction", + "prompt-decoder": "mask-generation", + } + if task is not None: resolved_type = model_type - if resolved_type is None and hf_model is not None: - resolved_type = AutoConfig.from_pretrained( - hf_model, trust_remote_code=trust_remote_code - ).model_type + if resolved_type is None and config is not None: + resolved_type = config.model_type + if resolved_type is None and config_error is not None: + raise config_error if resolved_type is None: return None return resolve_composite(resolved_type, task) - if hf_model is not None: - config = AutoConfig.from_pretrained(hf_model, trust_remote_code=trust_remote_code) - elif model_type is not None: + if config is None and model_type is not None: config = AutoConfig.for_model(model_type) - else: + elif config is None and config_error is not None: + raise config_error + elif config is None: return None return resolve_task(config).composite @@ -361,7 +386,30 @@ def resolve_composite_onnx_sources( """Resolve published ONNX sources for a registered composite model.""" from transformers import AutoConfig - config = AutoConfig.from_pretrained(hf_model, trust_remote_code=trust_remote_code) + try: + config = AutoConfig.from_pretrained(hf_model, trust_remote_code=trust_remote_code) + except (ValueError, OSError) as config_error: + from .onnx_hub import resolve_hf_release_onnx_encoder_decoder + + release_sources = resolve_hf_release_onnx_encoder_decoder( + hf_model, + precision=precision, + ) + if release_sources is None: + raise + if task not in {None, "mask-generation"}: + raise ValueError( + f"Published promptable ONNX composite serves 'mask-generation', not {task!r}." + ) from config_error + resolved_release = {name: str(path) for name, path in release_sources.items()} + if component_name is not None: + if component_name not in resolved_release: + raise ValueError( + f"Published composite has no component {component_name!r}; " + f"available: {sorted(resolved_release)}" + ) from config_error + return {component_name: resolved_release[component_name]} + return resolved_release model_type = config.model_type.lower().replace("_", "-") registry = _composite_registry() diff --git a/tests/unit/eval/test_mask_generation_evaluator.py b/tests/unit/eval/test_mask_generation_evaluator.py index d8032e531..fa27278f0 100644 --- a/tests/unit/eval/test_mask_generation_evaluator.py +++ b/tests/unit/eval/test_mask_generation_evaluator.py @@ -12,6 +12,8 @@ from __future__ import annotations +import json + import numpy as np import pytest from PIL import Image @@ -21,9 +23,12 @@ _TARGET_SIZE, WinMLMaskGenerationEvaluator, _build_decoder_inputs, + _build_encoder_prompt_inputs, _build_providers, _postprocess_mask, _preprocess_image, + _resolve_point_prompt_inputs, + _select_best_mask, ) @@ -363,6 +368,7 @@ def test_default_dataset_registered(self) -> None: _postprocess_for_profile, _preprocess_for_profile, _resolve_profile, + _resolve_release_profile, ) @@ -507,3 +513,145 @@ def test_default_is_sam3(self) -> None: sess = _StubSession([1, 3, "H", "W"]) prof = _resolve_profile(_cfg("unknown/family"), sess) assert prof is SAM3_PROFILE + + +class _PromptNode: + def __init__(self, name: str, shape: list[int], node_type: str = "tensor(float)") -> None: + self.name = name + self.shape = shape + self.type = node_type + + +class _PromptSession: + def __init__(self, inputs: list[_PromptNode]) -> None: + self._inputs = inputs + + def get_inputs(self): + return self._inputs + + +class TestEncoderPromptRouting: + def test_resolves_float_coordinate_and_label_pair(self) -> None: + session = _PromptSession( + [ + _PromptNode("image", [1, 3, 1024, 1024]), + _PromptNode("coordinates", [1, 2, 2]), + _PromptNode("labels", [1, 2]), + ] + ) + + assert _resolve_point_prompt_inputs(session) == ("coordinates", "labels") + + def test_builds_normalized_point_and_padding(self) -> None: + session = _PromptSession( + [ + _PromptNode("image", [1, 3, 1024, 1024]), + _PromptNode("coordinates", [1, 2, 2]), + _PromptNode("labels", [1, 2]), + ] + ) + + feed = _build_encoder_prompt_inputs( + prompt={"point": [160, 120]}, + prompt_mode="point", + session=session, + coordinate_name="coordinates", + label_name="labels", + original_width=640, + original_height=480, + ) + + np.testing.assert_allclose(feed["coordinates"][0, 0], [0.25, 0.25]) + np.testing.assert_allclose(feed["coordinates"][0, 1], [0.0, 0.0]) + np.testing.assert_allclose(feed["labels"], [[1.0, -1.0]]) + + def test_resolves_preprocessing_from_release_metadata(self, tmp_path) -> None: + encoder = tmp_path / "encoder" / "model.onnx" + encoder.parent.mkdir() + encoder.touch() + (encoder.parent / "winml_release_metadata.json").write_text( + json.dumps( + { + "model_files": { + "graph-a.onnx": { + "inputs": { + "image": { + "shape": [1, 3, 1024, 1024], + "dtype": "float32", + "io_type": "image", + "value_range": [0.0, 1.0], + }, + "coordinates": { + "shape": [1, 2, 2], + "dtype": "float32", + }, + "labels": {"shape": [1, 2], "dtype": "float32"}, + } + } + } + } + ), + encoding="utf-8", + ) + config = _cfg("unrelated/id") + config.model_path = { + "image-encoder": str(encoder), + "prompt-decoder": str(tmp_path / "decoder.onnx"), + } + session = _PromptSession( + [ + _PromptNode("image", [1, 3, 1024, 1024]), + _PromptNode("coordinates", [1, 2, 2]), + _PromptNode("labels", [1, 2]), + ] + ) + + profile = _resolve_release_profile(config, session) + + assert profile.name == "prompt-encoder" + assert profile.target_size == 1024 + assert profile.resize_mode == "direct" + + def test_missing_release_metadata_fails_closed(self, tmp_path) -> None: + config = _cfg("unrelated/id") + config.model_path = { + "image-encoder": str(tmp_path / "encoder.onnx"), + "prompt-decoder": str(tmp_path / "decoder.onnx"), + } + session = _PromptSession( + [ + _PromptNode("image", [1, 3, 1024, 1024]), + _PromptNode("coordinates", [1, 2, 2]), + _PromptNode("labels", [1, 2]), + ] + ) + + with pytest.raises(ValueError, match="requires persisted release runtime metadata"): + _resolve_release_profile(config, session) + + def test_ambiguous_coordinate_inputs_fail_closed(self) -> None: + session = _PromptSession( + [ + _PromptNode("coordinates_a", [1, 2, 2]), + _PromptNode("coordinates_b", [1, 2, 2]), + _PromptNode("labels", [1, 2]), + ] + ) + + with pytest.raises(ValueError, match="incomplete or ambiguous"): + _resolve_point_prompt_inputs(session) + + def test_selects_mask_from_rank4_release_layout(self) -> None: + masks = np.stack([np.zeros((4, 4), dtype=np.float32), np.ones((4, 4), dtype=np.float32)])[ + None, ... + ] + scores = np.array([[0.1, 0.9]], dtype=np.float32) + + np.testing.assert_array_equal(_select_best_mask(masks, scores), np.ones((4, 4))) + + def test_mask_score_candidate_mismatch_fails_closed(self) -> None: + with pytest.raises(ValueError, match="candidate count mismatch"): + _select_best_mask( + np.zeros((1, 2, 4, 4), dtype=np.float32), + np.zeros((1, 1), dtype=np.float32), + ) diff --git a/tests/unit/loader/test_onnx_composite_hub.py b/tests/unit/loader/test_onnx_composite_hub.py index be75b465c..99880d35f 100644 --- a/tests/unit/loader/test_onnx_composite_hub.py +++ b/tests/unit/loader/test_onnx_composite_hub.py @@ -66,6 +66,50 @@ def test_resolver_selects_matching_graph_contract_and_falls_back_to_fp32(monkeyp } +def test_resolver_discovers_prompt_bearing_encoder_and_embedding_only_decoder( + monkeypatch, +) -> None: + graphs = { + "first.onnx": _graph( + "first.onnx", + inputs=( + ("image", "tensor(float)", 4), + ("coordinates", "tensor(float)", 3), + ("labels", "tensor(float)", 2), + ), + outputs=( + ("embedding", "tensor(float)", 4), + ("sparse", "tensor(float)", 3), + ), + precision="fp32", + ), + "second.onnx": _graph( + "second.onnx", + inputs=( + ("embedding", "tensor(float)", 4), + ("sparse", "tensor(float)", 3), + ), + outputs=( + ("mask", "tensor(float)", 4), + ("quality", "tensor(float)", 2), + ), + precision="fp32", + ), + } + monkeypatch.setattr("huggingface_hub.list_repo_files", lambda *args, **kwargs: list(graphs)) + monkeypatch.setattr( + onnx_hub, + "resolve_hf_onnx_path", + lambda model_id, **kwargs: Path(model_id.rsplit("/", 1)[-1]), + ) + monkeypatch.setattr(onnx_hub, "_inspect_runnable_graph", lambda path: graphs[path.name]) + + assert onnx_hub.resolve_hf_onnx_encoder_decoder("org/model") == { + "image-encoder": Path("first.onnx"), + "prompt-decoder": Path("second.onnx"), + } + + def test_resolver_prefers_unquantized_graph_family(monkeypatch) -> None: graphs = { "encoder.onnx": _graph( @@ -77,10 +121,11 @@ def test_resolver_prefers_unquantized_graph_family(monkeypatch) -> None: "decoder.onnx": _graph( "decoder.onnx", inputs=( + ("points", "tensor(float)", 4), ("labels", "tensor(int64)", 3), ("embedding", "tensor(float)", 4), ), - outputs=(("masks", "tensor(float)", 5),), + outputs=(("scores", "tensor(float)", 3), ("masks", "tensor(float)", 5)), precision="fp32", ), "encoder_quantized.onnx": _graph( @@ -93,10 +138,11 @@ def test_resolver_prefers_unquantized_graph_family(monkeypatch) -> None: "decoder_quantized.onnx": _graph( "decoder_quantized.onnx", inputs=( + ("points", "tensor(float)", 4), ("labels", "tensor(int64)", 3), ("embedding", "tensor(float)", 4), ), - outputs=(("masks", "tensor(float)", 5),), + outputs=(("scores", "tensor(float)", 3), ("masks", "tensor(float)", 5)), precision="fp32", has_quantized_weights=True, ), @@ -128,10 +174,11 @@ def test_resolver_rejects_multiple_valid_graph_pairs_as_ambiguous(monkeypatch) - "decoder_a.onnx": _graph( "decoder_a.onnx", inputs=( + ("points", "tensor(float)", 4), ("labels", "tensor(int64)", 3), ("embedding_a", "tensor(float)", 4), ), - outputs=(("masks", "tensor(float)", 5),), + outputs=(("scores", "tensor(float)", 3), ("masks", "tensor(float)", 5)), precision="fp32", ), "encoder_b.onnx": _graph( @@ -143,10 +190,11 @@ def test_resolver_rejects_multiple_valid_graph_pairs_as_ambiguous(monkeypatch) - "decoder_b.onnx": _graph( "decoder_b.onnx", inputs=( + ("points", "tensor(float)", 4), ("labels", "tensor(int64)", 3), ("embedding_b", "tensor(float)", 4), ), - outputs=(("masks", "tensor(float)", 5),), + outputs=(("scores", "tensor(float)", 3), ("masks", "tensor(float)", 5)), precision="fp32", ), } @@ -172,6 +220,37 @@ def test_resolver_rejects_multiple_valid_graph_pairs_as_ambiguous(monkeypatch) - assert "decoder_b.onnx" in message +def test_resolver_rejects_decoder_without_score_output(monkeypatch) -> None: + graphs = { + "encoder.onnx": _graph( + "encoder.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(("embedding", "tensor(float)", 4),), + precision="fp32", + ), + "decoder.onnx": _graph( + "decoder.onnx", + inputs=( + ("points", "tensor(float)", 4), + ("labels", "tensor(int64)", 3), + ("embedding", "tensor(float)", 4), + ), + outputs=(("masks", "tensor(float)", 5),), + precision="fp32", + ), + } + monkeypatch.setattr("huggingface_hub.list_repo_files", lambda *args, **kwargs: list(graphs)) + monkeypatch.setattr( + onnx_hub, + "resolve_hf_onnx_path", + lambda model_id, **kwargs: Path(model_id.rsplit("/", 1)[-1]), + ) + monkeypatch.setattr(onnx_hub, "_inspect_runnable_graph", lambda path: graphs[path.name]) + + with pytest.raises(ValueError, match="do not contain a runnable promptable"): + onnx_hub.resolve_hf_onnx_encoder_decoder("org/model") + + def test_sam_published_onnx_absence_preserves_pytorch_fallback(monkeypatch) -> None: monkeypatch.setattr("huggingface_hub.list_repo_files", lambda *args, **kwargs: []) diff --git a/tests/unit/loader/test_release_assets.py b/tests/unit/loader/test_release_assets.py new file mode 100644 index 000000000..3448c25c1 --- /dev/null +++ b/tests/unit/loader/test_release_assets.py @@ -0,0 +1,218 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for safe, provenance-preserving release-asset acquisition.""" + +from __future__ import annotations + +import json +import shutil +import zipfile +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from winml.modelkit.loader import release_assets + + +if TYPE_CHECKING: + from pathlib import Path + + +def _identity_model(path: Path) -> None: + graph = helper.make_graph( + [helper.make_node("Identity", ["input"], ["output"])], + "identity", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])], + ) + onnx.save(helper.make_model(graph), path) + + +def test_safe_extract_zip_rejects_parent_traversal(tmp_path: Path) -> None: + archive = tmp_path / "unsafe.zip" + with zipfile.ZipFile(archive, "w") as bundle: + bundle.writestr("../escape.onnx", b"not a graph") + + with pytest.raises(ValueError, match="Unsafe ZIP member"): + release_assets.safe_extract_zip(archive, tmp_path / "out") + + assert not (tmp_path / "escape.onnx").exists() + + +def test_safe_extract_zip_rejects_duplicate_members(tmp_path: Path) -> None: + archive = tmp_path / "duplicate.zip" + with zipfile.ZipFile(archive, "w") as bundle: + bundle.writestr("model.onnx", b"first") + with pytest.warns(UserWarning): + bundle.writestr("model.onnx", b"second") + + with pytest.raises(ValueError, match="duplicate member"): + release_assets.safe_extract_zip(archive, tmp_path / "out") + + +def _external_model(path: Path) -> Path: + import numpy as np + + weight = numpy_helper.from_array(np.ones((2, 2), dtype=np.float32), name="weight") + graph = helper.make_graph( + [helper.make_node("MatMul", ["input", "weight"], ["output"])], + "external", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 2])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 2])], + [weight], + ) + model = helper.make_model(graph) + onnx.save_model( + model, + path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="weights.data", + size_threshold=0, + ) + return path.parent / "weights.data" + + +def test_external_data_must_be_present_and_colocated(tmp_path: Path) -> None: + model_path = tmp_path / "model.onnx" + sidecar = _external_model(model_path) + + assert release_assets.validate_onnx_external_data(model_path, tmp_path) == (sidecar,) + + sidecar.unlink() + with pytest.raises(FileNotFoundError, match="missing or out-of-root"): + release_assets.validate_onnx_external_data(model_path, tmp_path) + + +def test_external_data_rejects_traversal_location(tmp_path: Path) -> None: + model_path = tmp_path / "model.onnx" + _external_model(model_path) + model = onnx.load(model_path, load_external_data=False) + location = next( + field for field in model.graph.initializer[0].external_data if field.key == "location" + ) + location.value = "../weights.data" + model_path.write_bytes(model.SerializeToString()) + + with pytest.raises(ValueError, match="unsafe external-data location"): + release_assets.validate_onnx_external_data(model_path, tmp_path) + + +def test_copy_release_contract_files_persists_metadata_and_provenance(tmp_path: Path) -> None: + release_root = tmp_path / "extracted" + graph_root = release_root / "bundle" + graph_root.mkdir(parents=True) + graph = graph_root / "encoder.onnx" + graph.touch() + (release_root / "winml_release_provenance.json").write_text( + '{"resolved_revision":"abc"}', encoding="utf-8" + ) + (graph_root / "metadata.json").write_text('{"model_files":{}}', encoding="utf-8") + output = tmp_path / "built" + + release_assets.copy_release_contract_files(graph, output) + + assert json.loads((output / "winml_release_provenance.json").read_text()) == { + "resolved_revision": "abc" + } + assert json.loads((output / "winml_release_metadata.json").read_text()) == {"model_files": {}} + + +def test_acquisition_records_pinned_provenance_and_reuses_cache( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import huggingface_hub + + manifest = tmp_path / "release_assets.json" + manifest.write_text( + json.dumps( + { + "version": "1.2.3", + "precisions": { + "float": { + "universal_assets": { + "onnx": { + "download_url": "https://assets.example/model-onnx-float.zip", + "tool_versions": {"onnx_runtime": "1.2.3"}, + } + } + } + }, + } + ), + encoding="utf-8", + ) + source = tmp_path / "source" + source.mkdir() + _identity_model(source / "encoder.onnx") + _identity_model(source / "decoder.onnx") + (source / "metadata.json").write_text("{}", encoding="utf-8") + archive = tmp_path / "asset.zip" + with zipfile.ZipFile(archive, "w") as bundle: + for path in source.iterdir(): + bundle.write(path, f"bundle/{path.name}") + + info = SimpleNamespace(sha="a" * 40, pipeline_tag="image-segmentation") + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: SimpleNamespace(model_info=lambda *args, **kwargs: info), + ) + monkeypatch.setattr( + huggingface_hub, + "hf_hub_download", + lambda *args, **kwargs: str(manifest), + ) + downloads: list[str] = [] + + def copy_archive(url: str, destination: Path) -> None: + downloads.append(url) + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(archive, destination) + + monkeypatch.setattr(release_assets, "_download_archive", copy_archive) + + result = release_assets.acquire_hf_release_asset( + "org/model", revision="main", precision="fp16", cache_dir=tmp_path / "cache" + ) + assert result is not None + assert result.metadata_path is not None + assert result.provenance["requested_revision"] == "main" + assert result.provenance["resolved_revision"] == "a" * 40 + assert result.provenance["requested_precision"] == "fp16" + assert result.provenance["source_precision"] == "float" + assert result.provenance["archive_sha256"] + assert result.provenance["manifest_sha256"] + assert result.provenance["pipeline_tag"] == "image-segmentation" + assert downloads == ["https://assets.example/model-onnx-float.zip"] + + cached = release_assets.acquire_hf_release_asset( + "org/model", revision="main", precision="fp16", cache_dir=tmp_path / "cache" + ) + assert cached is not None + assert cached.root == result.root + assert len(downloads) == 1 + + +def test_present_malformed_manifest_fails_closed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import huggingface_hub + + manifest = tmp_path / "release_assets.json" + manifest.write_text("{}", encoding="utf-8") + info = SimpleNamespace(sha="b" * 40, pipeline_tag="image-segmentation") + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: SimpleNamespace(model_info=lambda *args, **kwargs: info), + ) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", lambda *args, **kwargs: str(manifest)) + + with pytest.raises(ValueError, match="has no 'onnx' asset"): + release_assets.acquire_hf_release_asset("org/model", cache_dir=tmp_path / "cache") From f97f40952482fbac409e4cd50f1993e223d5d685 Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Wed, 22 Jul 2026 02:17:35 +0800 Subject: [PATCH 4/5] recipe(sam2): add Qualcomm composite CPU recipes --- ...-generation_fp16_image-encoder_config.json | 29 +++++++++++++++++++ ...generation_fp16_prompt-decoder_config.json | 29 +++++++++++++++++++ ...-generation_fp32_image-encoder_config.json | 10 +++++++ ...generation_fp32_prompt-decoder_config.json | 10 +++++++ 4 files changed, 78 insertions(+) create mode 100644 examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_image-encoder_config.json create mode 100644 examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_prompt-decoder_config.json create mode 100644 examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_image-encoder_config.json create mode 100644 examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_prompt-decoder_config.json diff --git a/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_image-encoder_config.json b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_image-encoder_config.json new file mode 100644 index 000000000..6aba54fd1 --- /dev/null +++ b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_image-encoder_config.json @@ -0,0 +1,29 @@ +{ + "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": "image-feature-extraction", + "component_name": "image-encoder" + } +} \ No newline at end of file diff --git a/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_prompt-decoder_config.json b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_prompt-decoder_config.json new file mode 100644 index 000000000..de4bc5ef4 --- /dev/null +++ b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp16_prompt-decoder_config.json @@ -0,0 +1,29 @@ +{ + "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": "mask-generation", + "component_name": "prompt-decoder" + } +} \ No newline at end of file diff --git a/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_image-encoder_config.json b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_image-encoder_config.json new file mode 100644 index 000000000..cea40cbe4 --- /dev/null +++ b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_image-encoder_config.json @@ -0,0 +1,10 @@ +{ + "export": null, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "image-feature-extraction", + "component_name": "image-encoder" + } +} \ No newline at end of file diff --git a/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_prompt-decoder_config.json b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_prompt-decoder_config.json new file mode 100644 index 000000000..bd309641b --- /dev/null +++ b/examples/recipes/qualcomm_Segment-Anything-Model-2/cpu/cpu/mask-generation_fp32_prompt-decoder_config.json @@ -0,0 +1,10 @@ +{ + "export": null, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "mask-generation", + "component_name": "prompt-decoder" + } +} \ No newline at end of file From 448b3508740f1e3b7fc4bb0c004af6020cad9a8b Mon Sep 17 00:00:00 2001 From: "Shiyi Zheng (from Dev Box)" Date: Wed, 22 Jul 2026 03:33:06 +0800 Subject: [PATCH 5/5] fix(loader): harden release composite validation --- src/winml/modelkit/loader/onnx_hub.py | 88 ++++++++- src/winml/modelkit/loader/release_assets.py | 174 ++++++++++++++++-- src/winml/modelkit/loader/resolution.py | 33 +++- .../test_config_composite_resolution.py | 31 ++++ tests/unit/loader/test_onnx_composite_hub.py | 86 ++++++++- tests/unit/loader/test_release_assets.py | 128 ++++++++++++- 6 files changed, 508 insertions(+), 32 deletions(-) diff --git a/src/winml/modelkit/loader/onnx_hub.py b/src/winml/modelkit/loader/onnx_hub.py index 41d56082f..3d6eb82e4 100644 --- a/src/winml/modelkit/loader/onnx_hub.py +++ b/src/winml/modelkit/loader/onnx_hub.py @@ -110,6 +110,92 @@ def _has_mask_outputs(contract: _GraphContract) -> bool: return sum(rank >= 4 for rank in ranks) == 1 and sum(1 <= rank <= 3 for rank in ranks) == 1 +def _shared_ports_compatible( + encoder: _GraphContract, + decoder: _GraphContract, + shared: tuple[str, ...], +) -> bool: + """Check connected ports, unifying static and symbolic shape constraints.""" + encoder_ports = { + name: (dtype, rank, encoder.output_shapes[index] if encoder.output_shapes else None) + for index, (name, dtype, rank) in enumerate(encoder.outputs) + } + decoder_ports = { + name: (dtype, rank, decoder.input_shapes[index] if decoder.input_shapes else None) + for index, (name, dtype, rank) in enumerate(decoder.inputs) + } + if len(encoder_ports) != len(encoder.outputs) or len(decoder_ports) != len(decoder.inputs): + return False + + parents: dict[tuple[str, str], tuple[str, str]] = {} + static_values: dict[tuple[str, str], int] = {} + + def find(symbol: tuple[str, str]) -> tuple[str, str]: + parents.setdefault(symbol, symbol) + if parents[symbol] != symbol: + parents[symbol] = find(parents[symbol]) + return parents[symbol] + + def bind(symbol: tuple[str, str], value: int) -> bool: + root = find(symbol) + existing = static_values.get(root) + if existing is not None and existing != value: + return False + static_values[root] = value + return True + + def union(left: tuple[str, str], right: tuple[str, str]) -> bool: + left_root = find(left) + right_root = find(right) + if left_root == right_root: + return True + left_value = static_values.pop(left_root, None) + right_value = static_values.pop(right_root, None) + if left_value is not None and right_value is not None and left_value != right_value: + return False + parents[right_root] = left_root + merged_value = left_value if left_value is not None else right_value + if merged_value is not None: + static_values[left_root] = merged_value + return True + + for name in shared: + encoder_dtype, encoder_rank, encoder_shape = encoder_ports[name] + decoder_dtype, decoder_rank, decoder_shape = decoder_ports[name] + if encoder_dtype.lower() != decoder_dtype.lower() or encoder_rank != decoder_rank: + return False + if encoder_shape is None: + encoder_shape = (None,) * encoder_rank + if decoder_shape is None: + decoder_shape = (None,) * decoder_rank + if len(encoder_shape) != encoder_rank or len(decoder_shape) != decoder_rank: + return False + for encoder_dim, decoder_dim in zip(encoder_shape, decoder_shape, strict=True): + encoder_static = encoder_dim if isinstance(encoder_dim, int) else None + decoder_static = decoder_dim if isinstance(decoder_dim, int) else None + encoder_symbol = encoder_dim if isinstance(encoder_dim, str) and encoder_dim else None + decoder_symbol = decoder_dim if isinstance(decoder_dim, str) and decoder_dim else None + if ( + encoder_static is not None + and decoder_static is not None + and encoder_static != decoder_static + ): + return False + if encoder_symbol is not None and decoder_symbol is not None: + if not union(("encoder", encoder_symbol), ("decoder", decoder_symbol)): + return False + elif encoder_symbol is not None and decoder_static is not None: + if not bind(("encoder", encoder_symbol), decoder_static): + return False + elif ( + decoder_symbol is not None + and encoder_static is not None + and not bind(("decoder", decoder_symbol), encoder_static) + ): + return False + return True + + def _select_encoder_decoder_pair( graphs: list[_GraphContract], *, @@ -133,7 +219,7 @@ def _select_encoder_decoder_pair( continue decoder_inputs = {name for name, _dtype, _rank in decoder.inputs} shared = tuple(sorted(encoder_outputs & decoder_inputs)) - if not shared: + if not shared or not _shared_ports_compatible(encoder, decoder, shared): continue encoder_prompt = _has_point_prompt(encoder) decoder_prompt = _has_point_prompt(decoder) diff --git a/src/winml/modelkit/loader/release_assets.py b/src/winml/modelkit/loader/release_assets.py index 6c04427c1..bfaf032e7 100644 --- a/src/winml/modelkit/loader/release_assets.py +++ b/src/winml/modelkit/loader/release_assets.py @@ -13,18 +13,18 @@ import hashlib import json -import logging import shutil import stat import tempfile import zipfile from dataclasses import dataclass from pathlib import Path, PurePosixPath -from typing import Any +from typing import TYPE_CHECKING, Any from urllib.parse import urlparse -logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from collections.abc import Iterator _MANIFEST_NAME = "release_assets.json" _PROVENANCE_NAME = "winml_release_provenance.json" @@ -149,20 +149,68 @@ def safe_extract_zip(archive: Path, destination: Path) -> None: shutil.copyfileobj(source, output, length=1024 * 1024) +def _graph_tensors(graph: Any) -> Iterator[Any]: + """Yield every tensor reachable from a GraphProto, including nested graphs.""" + import onnx + + yield from graph.initializer + for sparse in graph.sparse_initializer: + yield sparse.values + yield sparse.indices + for node in graph.node: + yield from _node_attribute_tensors(node, onnx.AttributeProto) + + +def _node_attribute_tensors(node: Any, attribute_proto: Any) -> Iterator[Any]: + """Yield tensor and sparse-tensor attributes, recursing into graph attributes.""" + for attribute in node.attribute: + if attribute.type == attribute_proto.TENSOR: + yield attribute.t + elif attribute.type == attribute_proto.TENSORS: + yield from attribute.tensors + elif attribute.type == attribute_proto.SPARSE_TENSOR: + yield attribute.sparse_tensor.values + yield attribute.sparse_tensor.indices + elif attribute.type == attribute_proto.SPARSE_TENSORS: + for sparse in attribute.sparse_tensors: + yield sparse.values + yield sparse.indices + elif attribute.type == attribute_proto.GRAPH: + yield from _graph_tensors(attribute.g) + elif attribute.type == attribute_proto.GRAPHS: + for graph in attribute.graphs: + yield from _graph_tensors(graph) + + +def _model_tensors(model: Any) -> Iterator[Any]: + """Yield every TensorProto that may carry external data in a ModelProto.""" + import onnx + + yield from _graph_tensors(model.graph) + for function in model.functions: + for node in function.node: + yield from _node_attribute_tensors(node, onnx.AttributeProto) + + def validate_onnx_external_data(path: Path, root: Path) -> tuple[Path, ...]: - """Require every ONNX external-data location to be a present in-root file.""" + """Require every recursively referenced ONNX sidecar to be a present in-root file.""" import onnx + from onnx.external_data_helper import ExternalDataInfo, uses_external_data model = onnx.load(path, load_external_data=False) sidecars: list[Path] = [] root_resolved = root.resolve() - for initializer in model.graph.initializer: - if initializer.data_location != onnx.TensorProto.EXTERNAL: + for tensor in _model_tensors(model): + if not uses_external_data(tensor): continue - fields = {field.key: field.value for field in initializer.external_data} - location = fields.get("location") + try: + location = ExternalDataInfo(tensor).location + except ValueError as error: + raise ValueError( + f"External tensor in {path.name!r} has invalid external-data metadata." + ) from error if not location: - raise ValueError(f"External initializer in {path.name!r} has no location.") + raise ValueError(f"External tensor in {path.name!r} has no location.") relative = PurePosixPath(location.replace("\\", "/")) if ( relative.is_absolute() @@ -182,6 +230,79 @@ def validate_onnx_external_data(path: Path, root: Path) -> tuple[Path, ...]: return tuple(dict.fromkeys(sidecars)) +def _recorded_path(root: Path, relative_path: str) -> Path: + """Resolve one provenance path without allowing it to escape the cache root.""" + relative = PurePosixPath(relative_path.replace("\\", "/")) + if ( + relative.is_absolute() + or any(part in {"", ".", ".."} for part in relative.parts) + or (relative.parts and ":" in relative.parts[0]) + ): + raise ValueError(f"Unsafe path in release provenance: {relative_path!r}.") + resolved = root.joinpath(*relative.parts).resolve() + root_resolved = root.resolve() + if root_resolved not in resolved.parents: + raise ValueError(f"Release provenance path escapes extraction root: {relative_path!r}.") + return resolved + + +def _validate_cached_extraction( + extracted: Path, + archive_path: Path, + provenance: dict[str, Any], +) -> bool: + """Validate every cached graph and sidecar against its immutable provenance.""" + if provenance.get("schema_version") != 2 or not archive_path.is_file(): + return False + archive_sha256 = provenance.get("archive_sha256") + graphs = provenance.get("graphs") + external_data = provenance.get("external_data") + if ( + not isinstance(archive_sha256, str) + or not isinstance(graphs, dict) + or not graphs + or not isinstance(external_data, dict) + ): + return False + if _sha256(archive_path) != archive_sha256: + return False + + try: + actual_graphs = { + path.relative_to(extracted).as_posix() for path in extracted.rglob("*.onnx") + } + if actual_graphs != set(graphs): + return False + for graph_relative, graph_sha256 in graphs.items(): + if not isinstance(graph_relative, str) or not isinstance(graph_sha256, str): + return False + graph = _recorded_path(extracted, graph_relative) + if not graph.is_file() or _sha256(graph) != graph_sha256: + return False + recorded_sidecars = external_data.get(graph_relative) + if not isinstance(recorded_sidecars, dict): + return False + actual_sidecars = { + sidecar.relative_to(extracted).as_posix(): sidecar + for sidecar in validate_onnx_external_data(graph, extracted) + } + if set(actual_sidecars) != set(recorded_sidecars): + return False + for sidecar_relative, sidecar_sha256 in recorded_sidecars.items(): + if not isinstance(sidecar_relative, str) or not isinstance(sidecar_sha256, str): + return False + sidecar = _recorded_path(extracted, sidecar_relative) + if ( + actual_sidecars.get(sidecar_relative) != sidecar + or not sidecar.is_file() + or _sha256(sidecar) != sidecar_sha256 + ): + return False + return set(external_data) == set(graphs) + except (OSError, ValueError): + return False + + def copy_release_contract_files(source_graph: Path, destination: Path) -> None: """Copy immutable release provenance/runtime metadata beside a built graph.""" release_root = next( @@ -257,10 +378,27 @@ def acquire_hf_release_asset( extracted = asset_root / "extracted" provenance_path = extracted / _PROVENANCE_NAME + archive_is_trusted = False if provenance_path.is_file(): - provenance = json.loads(provenance_path.read_text(encoding="utf-8")) - if provenance.get("archive_sha256") == _sha256(archive_path): + try: + provenance = json.loads(provenance_path.read_text(encoding="utf-8")) + archive_is_trusted = ( + isinstance(provenance, dict) + and archive_path.is_file() + and isinstance(provenance.get("archive_sha256"), str) + and provenance["archive_sha256"] == _sha256(archive_path) + ) + except (OSError, json.JSONDecodeError): + provenance = {} + if isinstance(provenance, dict) and _validate_cached_extraction( + extracted, archive_path, provenance + ): metadata_files = sorted(extracted.rglob("metadata.json")) + if len(metadata_files) > 1: + raise ValueError( + f"Release cache contains multiple metadata.json files: " + f"{[str(path.relative_to(extracted)) for path in metadata_files]}" + ) return AcquiredReleaseAsset( root=extracted, manifest_path=manifest_path, @@ -270,6 +408,10 @@ def acquire_hf_release_asset( ) asset_root.mkdir(parents=True, exist_ok=True) + if extracted.exists(): + shutil.rmtree(extracted) + if archive_path.is_file() and not archive_is_trusted: + archive_path.unlink() if not archive_path.is_file(): _download_archive(url, archive_path) @@ -281,14 +423,14 @@ def acquire_hf_release_asset( if not graph_paths: raise ValueError(f"Release archive {url!r} contains no ONNX graphs.") external_data = { - str(path.relative_to(staging)): [ - str(sidecar.relative_to(staging)) + path.relative_to(staging).as_posix(): { + sidecar.relative_to(staging).as_posix(): _sha256(sidecar) for sidecar in validate_onnx_external_data(path, staging) - ] + } for path in graph_paths } provenance = { - "schema_version": 1, + "schema_version": 2, "repo_id": model_id, "requested_revision": revision, "resolved_revision": resolved_revision, @@ -301,7 +443,7 @@ def acquire_hf_release_asset( "download_url": url, "archive_sha256": _sha256(archive_path), "tool_versions": asset.get("tool_versions", {}), - "graphs": {str(path.relative_to(staging)): _sha256(path) for path in graph_paths}, + "graphs": {path.relative_to(staging).as_posix(): _sha256(path) for path in graph_paths}, "external_data": external_data, } (staging / _PROVENANCE_NAME).write_text( diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index de04f65e3..6acc906d6 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -312,6 +312,14 @@ def resolve_composite(model_type: str, task: str) -> CompositeComponents | None: return dict(cls._SUB_MODEL_CONFIG) if cls is not None else None +def _optional_hub_discovery_errors() -> tuple[type[Exception], ...]: + """Errors for which optional release discovery preserves the original path.""" + import httpx + from huggingface_hub.errors import HfHubHTTPError, OfflineModeIsEnabled + + return (HfHubHTTPError, OfflineModeIsEnabled, httpx.HTTPError) + + def resolve_composite_components( hf_model: str | None, *, @@ -332,6 +340,9 @@ def resolve_composite_components( """ from transformers import AutoConfig + if task is not None and model_type is not None: + return resolve_composite(model_type, task) + config = None config_error: ValueError | OSError | None = None if hf_model is not None: @@ -341,10 +352,13 @@ def resolve_composite_components( config_error = error from .onnx_hub import resolve_hf_release_onnx_encoder_decoder - release_sources = resolve_hf_release_onnx_encoder_decoder( - hf_model, - precision=precision, - ) + try: + release_sources = resolve_hf_release_onnx_encoder_decoder( + hf_model, + precision=precision, + ) + except _optional_hub_discovery_errors() as release_error: + raise error from release_error if release_sources is not None: if task not in {None, "mask-generation"}: raise ValueError( @@ -391,10 +405,13 @@ def resolve_composite_onnx_sources( except (ValueError, OSError) as config_error: from .onnx_hub import resolve_hf_release_onnx_encoder_decoder - release_sources = resolve_hf_release_onnx_encoder_decoder( - hf_model, - precision=precision, - ) + try: + release_sources = resolve_hf_release_onnx_encoder_decoder( + hf_model, + precision=precision, + ) + except _optional_hub_discovery_errors() as release_error: + raise config_error from release_error if release_sources is None: raise if task not in {None, "mask-generation"}: diff --git a/tests/unit/commands/test_config_composite_resolution.py b/tests/unit/commands/test_config_composite_resolution.py index a364af2d6..3a9175e67 100644 --- a/tests/unit/commands/test_config_composite_resolution.py +++ b/tests/unit/commands/test_config_composite_resolution.py @@ -22,6 +22,7 @@ from unittest.mock import patch +import httpx import pytest from transformers import BartConfig, Qwen3Config, T5Config @@ -145,3 +146,33 @@ def test_explicit_task_resolves_decoder_only_composite() -> None: components = _resolve(None, "qwen3", "text-generation") assert components is not None assert "decoder_prefill" in components and "decoder_gen" in components + + +def test_explicit_model_type_and_task_do_not_probe_hub() -> None: + """An explicit offline schema must not trigger optional release discovery.""" + with ( + patch("transformers.AutoConfig.from_pretrained") as auto_config, + patch( + "winml.modelkit.loader.onnx_hub.resolve_hf_release_onnx_encoder_decoder" + ) as release_resolver, + ): + assert _resolve("some-model", "bert", "fill-mask") is None + + auto_config.assert_not_called() + release_resolver.assert_not_called() + + +def test_optional_release_discovery_preserves_offline_config_fallback() -> None: + """An unavailable optional Hub probe must not replace the established error.""" + config_error = OSError("cached config unavailable") + with ( + patch("transformers.AutoConfig.from_pretrained", side_effect=config_error), + patch( + "winml.modelkit.loader.onnx_hub.resolve_hf_release_onnx_encoder_decoder", + side_effect=httpx.ConnectError("offline"), + ), + pytest.raises(OSError, match="cached config unavailable") as error, + ): + _resolve("some-model", None, None) + + assert error.value is config_error diff --git a/tests/unit/loader/test_onnx_composite_hub.py b/tests/unit/loader/test_onnx_composite_hub.py index 99880d35f..37a5bda78 100644 --- a/tests/unit/loader/test_onnx_composite_hub.py +++ b/tests/unit/loader/test_onnx_composite_hub.py @@ -21,8 +21,18 @@ def _graph( outputs: tuple[tuple[str, str, int], ...], precision: str, has_quantized_weights: bool = False, + input_shapes: tuple[tuple[object, ...], ...] = (), + output_shapes: tuple[tuple[object, ...], ...] = (), ) -> onnx_hub._GraphContract: - return onnx_hub._GraphContract(Path(name), inputs, outputs, precision, has_quantized_weights) + return onnx_hub._GraphContract( + Path(name), + inputs, + outputs, + precision, + has_quantized_weights, + input_shapes, + output_shapes, + ) def test_resolver_selects_matching_graph_contract_and_falls_back_to_fp32(monkeypatch) -> None: @@ -251,6 +261,80 @@ def test_resolver_rejects_decoder_without_score_output(monkeypatch) -> None: onnx_hub.resolve_hf_onnx_encoder_decoder("org/model") +@pytest.mark.parametrize( + ("encoder_port", "decoder_port", "encoder_shape", "decoder_shape"), + [ + (("embedding", "tensor(float)", 4), ("embedding", "tensor(float16)", 4), (), ()), + (("embedding", "tensor(float)", 4), ("embedding", "tensor(float)", 3), (), ()), + ( + ("embedding", "tensor(float)", 4), + ("embedding", "tensor(float)", 4), + (1, 256, 64, 64), + (1, 128, 32, 32), + ), + ( + ("embedding", "tensor(float)", 4), + ("embedding", "tensor(float)", 4), + (1, 256, "spatial", "spatial"), + (1, 256, 64, 32), + ), + ], +) +def test_pairing_rejects_incompatible_shared_port_contracts( + encoder_port, + decoder_port, + encoder_shape, + decoder_shape, +) -> None: + encoder = _graph( + "encoder.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(encoder_port,), + precision="fp32", + output_shapes=(encoder_shape,) if encoder_shape else (), + ) + decoder = _graph( + "decoder.onnx", + inputs=( + ("points", "tensor(float)", 4), + ("labels", "tensor(int64)", 3), + decoder_port, + ), + outputs=(("scores", "tensor(float)", 3), ("masks", "tensor(float)", 5)), + precision="fp32", + input_shapes=((None,) * 4, (None,) * 3, decoder_shape) if decoder_shape else (), + ) + + with pytest.raises(ValueError, match="do not contain a runnable promptable"): + onnx_hub._select_encoder_decoder_pair([encoder, decoder], precision="fp32") + + +def test_pairing_accepts_compatible_symbolic_and_static_shared_port_shapes() -> None: + encoder = _graph( + "encoder.onnx", + inputs=(("pixels", "tensor(float)", 4),), + outputs=(("embedding", "tensor(float)", 4),), + precision="fp32", + output_shapes=((1, 256, "height", "width"),), + ) + decoder = _graph( + "decoder.onnx", + inputs=( + ("points", "tensor(float)", 4), + ("labels", "tensor(int64)", 3), + ("embedding", "tensor(float)", 4), + ), + outputs=(("scores", "tensor(float)", 3), ("masks", "tensor(float)", 5)), + precision="fp32", + input_shapes=((None,) * 4, (None,) * 3, (1, 256, 64, 64)), + ) + + pair = onnx_hub._select_encoder_decoder_pair([encoder, decoder], precision="fp32") + + assert pair.encoder.path == Path("encoder.onnx") + assert pair.decoder.path == Path("decoder.onnx") + + def test_sam_published_onnx_absence_preserves_pytorch_fallback(monkeypatch) -> None: monkeypatch.setattr("huggingface_hub.list_repo_files", lambda *args, **kwargs: []) diff --git a/tests/unit/loader/test_release_assets.py b/tests/unit/loader/test_release_assets.py index 3448c25c1..1d14db421 100644 --- a/tests/unit/loader/test_release_assets.py +++ b/tests/unit/loader/test_release_assets.py @@ -12,9 +12,8 @@ from types import SimpleNamespace from typing import TYPE_CHECKING -import onnx import pytest -from onnx import TensorProto, helper, numpy_helper +from onnx import TensorProto, external_data_helper, helper, load, numpy_helper, save, save_model from winml.modelkit.loader import release_assets @@ -30,7 +29,7 @@ def _identity_model(path: Path) -> None: [helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])], [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])], ) - onnx.save(helper.make_model(graph), path) + save(helper.make_model(graph), path) def test_safe_extract_zip_rejects_parent_traversal(tmp_path: Path) -> None: @@ -67,7 +66,7 @@ def _external_model(path: Path) -> Path: [weight], ) model = helper.make_model(graph) - onnx.save_model( + save_model( model, path, save_as_external_data=True, @@ -92,7 +91,7 @@ def test_external_data_must_be_present_and_colocated(tmp_path: Path) -> None: def test_external_data_rejects_traversal_location(tmp_path: Path) -> None: model_path = tmp_path / "model.onnx" _external_model(model_path) - model = onnx.load(model_path, load_external_data=False) + model = load(model_path, load_external_data=False) location = next( field for field in model.graph.initializer[0].external_data if field.key == "location" ) @@ -103,6 +102,79 @@ def test_external_data_rejects_traversal_location(tmp_path: Path) -> None: release_assets.validate_onnx_external_data(model_path, tmp_path) +def _attribute_external_model(path: Path, location: str, *, nested: bool = False) -> Path: + import numpy as np + + tensor = numpy_helper.from_array(np.array([3.0], dtype=np.float32), name="attribute_value") + raw_data = tensor.raw_data + external_data_helper.set_external_data( + tensor, + location=location, + offset=0, + length=len(raw_data), + ) + tensor.ClearField("raw_data") + constant = helper.make_node("Constant", [], ["constant"], value=tensor) + if nested: + branch = helper.make_graph( + [constant], + "branch", + [], + [helper.make_tensor_value_info("constant", TensorProto.FLOAT, [1])], + ) + other_branch = helper.make_graph( + [helper.make_node("Identity", ["fallback"], ["constant"])], + "other_branch", + [helper.make_tensor_value_info("fallback", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("constant", TensorProto.FLOAT, [1])], + [helper.make_tensor("fallback", TensorProto.FLOAT, [1], [0.0])], + ) + nodes = [ + helper.make_node( + "If", + ["condition"], + ["output"], + then_branch=branch, + else_branch=other_branch, + ) + ] + inputs = [helper.make_tensor_value_info("condition", TensorProto.BOOL, [])] + else: + nodes = [constant, helper.make_node("Identity", ["constant"], ["output"])] + inputs = [] + graph = helper.make_graph( + nodes, + "attribute_external", + inputs, + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])], + ) + path.write_bytes(helper.make_model(graph).SerializeToString()) + sidecar = path.parent / location + if ".." not in location: + sidecar.parent.mkdir(parents=True, exist_ok=True) + sidecar.write_bytes(raw_data) + return sidecar + + +def test_external_data_recurses_into_nested_graph_tensor_attributes(tmp_path: Path) -> None: + model_path = tmp_path / "model.onnx" + sidecar = _attribute_external_model(model_path, "attribute.data", nested=True) + + assert release_assets.validate_onnx_external_data(model_path, tmp_path) == (sidecar,) + + sidecar.unlink() + with pytest.raises(FileNotFoundError, match="missing or out-of-root"): + release_assets.validate_onnx_external_data(model_path, tmp_path) + + +def test_external_data_rejects_malicious_tensor_attribute_location(tmp_path: Path) -> None: + model_path = tmp_path / "model.onnx" + _attribute_external_model(model_path, "../attribute.data") + + with pytest.raises(ValueError, match="unsafe external-data location"): + release_assets.validate_onnx_external_data(model_path, tmp_path) + + def test_copy_release_contract_files_persists_metadata_and_provenance(tmp_path: Path) -> None: release_root = tmp_path / "extracted" graph_root = release_root / "bundle" @@ -149,7 +221,7 @@ def test_acquisition_records_pinned_provenance_and_reuses_cache( ) source = tmp_path / "source" source.mkdir() - _identity_model(source / "encoder.onnx") + _external_model(source / "encoder.onnx") _identity_model(source / "decoder.onnx") (source / "metadata.json").write_text("{}", encoding="utf-8") archive = tmp_path / "asset.zip" @@ -189,6 +261,10 @@ def copy_archive(url: str, destination: Path) -> None: assert result.provenance["archive_sha256"] assert result.provenance["manifest_sha256"] assert result.provenance["pipeline_tag"] == "image-segmentation" + assert result.provenance["schema_version"] == 2 + assert result.provenance["external_data"]["bundle/encoder.onnx"] == { + "bundle/weights.data": release_assets._sha256(result.root / "bundle/weights.data") + } assert downloads == ["https://assets.example/model-onnx-float.zip"] cached = release_assets.acquire_hf_release_asset( @@ -198,6 +274,46 @@ def copy_archive(url: str, destination: Path) -> None: assert cached.root == result.root assert len(downloads) == 1 + graph = result.root / "bundle/encoder.onnx" + expected_graph_hash = result.provenance["graphs"]["bundle/encoder.onnx"] + graph.write_bytes(b"tampered graph") + repaired_graph = release_assets.acquire_hf_release_asset( + "org/model", revision="main", precision="fp16", cache_dir=tmp_path / "cache" + ) + assert repaired_graph is not None + assert release_assets._sha256(graph) == expected_graph_hash + assert len(downloads) == 1 + + sidecar = result.root / "bundle/weights.data" + expected_sidecar_hash = result.provenance["external_data"]["bundle/encoder.onnx"][ + "bundle/weights.data" + ] + sidecar.write_bytes(b"tampered sidecar") + repaired_sidecar = release_assets.acquire_hf_release_asset( + "org/model", revision="main", precision="fp16", cache_dir=tmp_path / "cache" + ) + assert repaired_sidecar is not None + assert release_assets._sha256(sidecar) == expected_sidecar_hash + assert len(downloads) == 1 + + unrecorded_graph = result.root / "bundle/unrecorded.onnx" + _identity_model(unrecorded_graph) + repaired_inventory = release_assets.acquire_hf_release_asset( + "org/model", revision="main", precision="fp16", cache_dir=tmp_path / "cache" + ) + assert repaired_inventory is not None + assert not unrecorded_graph.exists() + assert len(downloads) == 1 + + cached_archive = next((tmp_path / "cache").rglob("model-onnx-float.zip")) + cached_archive.write_bytes(b"tampered archive") + reacquired = release_assets.acquire_hf_release_asset( + "org/model", revision="main", precision="fp16", cache_dir=tmp_path / "cache" + ) + assert reacquired is not None + assert release_assets._sha256(reacquired.root / "bundle/encoder.onnx") == expected_graph_hash + assert len(downloads) == 2 + def test_present_malformed_manifest_fails_closed( tmp_path: Path, monkeypatch: pytest.MonkeyPatch