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..abc2fc996 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() @@ -1297,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. " @@ -1309,6 +1336,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 +1351,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 +1428,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..2b72b39e1 100644 --- a/src/winml/modelkit/loader/onnx_hub.py +++ b/src/winml/modelkit/loader/onnx_hub.py @@ -19,12 +19,176 @@ 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 + has_quantized_weights: bool = False + + +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" + 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) + + return _GraphContract( + path=path, + inputs=contract(session.get_inputs()), + outputs=contract(session.get_outputs()), + precision=precision, + has_quantized_weights=has_quantized_weights, + ) + + +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, _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 + 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 + 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 + quantization_rank = int(encoder.has_quantized_weights) + candidates.append( + (precision_rank, quantization_rank, encoder, decoder, tuple(sorted(shared))) + ) + + 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}" + ) + + 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} + + def resolve_hf_onnx_path( model_id: str, *, @@ -79,9 +243,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 +333,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..be75b465c --- /dev/null +++ b/tests/unit/loader/test_onnx_composite_hub.py @@ -0,0 +1,189 @@ +# ------------------------------------------------------------------------- +# 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, + has_quantized_weights: bool = False, +) -> onnx_hub._GraphContract: + 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: + 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_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: []) + + 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")