feat(loader): support published ONNX SAM composites

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

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

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

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

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,23 +1336,46 @@ 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})"
                         )
 
                         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,

src/winml/modelkit/commands/config.py

@@ -363,40 +363,70 @@ 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 "
                         "sub-components each have their own export inputs. Generate "
                         "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,27 +651,40 @@ 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(
             f"[dim]Generating config for component '{component_name}' "
             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)

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:

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