refactor: split elastix engine and type model parameters across presets

Split each elastix preset's Model.py into a parameter mapper plus an elastix_engine.py runtime, and type the constructor parameters (Literal / Annotated Range / Choices) so UIs derive their constraints. Reg presets set tta=0 and carry the real grid-spacing / spatial-samples in the config (source of truth). ConvexAdam presets carry typed Range parameters.

CBCT_CT_HeadNeck/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

CBCT_CT_HeadNeck/Prediction.yml

@@ -7,9 +7,9 @@ Predictor:
       - ParameterMap_CBCT_HN.txt
       outputs_criterions: None
       max_iterations: 0
-      final_grid_spacing: 0.0
+      final_grid_spacing: 10.0
       subset_features: 0
-      spatial_samples: 0
+      spatial_samples: 2000
       parameter_overrides: []
       resolutions:
         '0':
@@ -87,7 +87,6 @@ Predictor:
               subset_features: 64
               pca: 0
               distance: L1
-      models_registry: VBoussot/impact-torchscript-models:models.json
       mode: Static
   Dataset:
     groups_src:

CBCT_CT_HeadNeck/app.json

@@ -3,7 +3,7 @@
     "short_description": "Optimized preset for CBCT/CT registration on head & neck",
     "description": "A five-level recursive B-spline deformable registration optimized for CBCT/CT head-and-neck alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 200, 200, and 150 iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm) using L1 distances on selected internal layers of the network. A composite objective (IMPACT + mutual information + bending energy penalty, with increased MI weight) ensures robust cross-modality alignment in complex head-and-neck anatomy while enforcing smooth, physically plausible deformations.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

CBCT_CT_HeadNeck/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)

CBCT_CT_MRSeg/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

CBCT_CT_MRSeg/Prediction.yml

@@ -7,9 +7,9 @@ Predictor:
       - ParameterMap_CBCT_generic_MRSeg.txt
       outputs_criterions: None
       max_iterations: 0
-      final_grid_spacing: 0.0
+      final_grid_spacing: 14.0
       subset_features: 0
-      spatial_samples: 0
+      spatial_samples: 2000
       parameter_overrides: []
       resolutions:
         '0':
@@ -120,7 +120,6 @@ Predictor:
               subset_features: 64
               pca: 0
               distance: Dice
-      models_registry: VBoussot/impact-torchscript-models:models.json
       mode: Static
   Dataset:
     groups_src:

CBCT_CT_MRSeg/app.json

@@ -3,7 +3,7 @@
     "short_description": "Generic CBCT/CT deformable registration using MRSegmentator features",
     "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from the pretrained MRSegmentator model. The scheme follows a multi-resolution strategy with up to 300, 300, 250, and 200 ASGD iterations and 2000 stochastic spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), with a level-dependent combination of Dice-based segmentation overlap and L1 feature distances on selected internal layers of MRSegmentator. Early levels rely on pure Dice supervision, while finer stages progressively integrate feature-level alignment with increasing L1 contribution (0.3/0.7, 0.5/0.5) and a final purely feature-based stage. The optimization minimizes a composite objective (IMPACT + mutual information + bending energy penalty), enabling robust cross-modality alignment between CBCT and CT while enforcing smooth, physically plausible deformations.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

CBCT_CT_MRSeg/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)

CBCT_CT_TS/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

CBCT_CT_TS/Prediction.yml

@@ -7,9 +7,9 @@ Predictor:
       - ParameterMap_CBCT_generic_TS.txt
       outputs_criterions: None
       max_iterations: 0
-      final_grid_spacing: 0.0
+      final_grid_spacing: 14.0
       subset_features: 0
-      spatial_samples: 0
+      spatial_samples: 2000
       parameter_overrides: []
   Dataset:
     groups_src:

CBCT_CT_TS/app.json

@@ -3,7 +3,7 @@
     "short_description": "Generic CBCT/CT deformable registration using TotalSegmentator features",
     "description": "A four-level recursive B-spline deformable registration optimized for generic CBCT/CT alignment, driven by the IMPACT metric using semantic features extracted from pretrained TotalSegmentator TorchScript models. The optimization follows a multi-resolution ASGD scheme with up to 300, 300, 250, and 200 iterations using 2000 random spatial samples per level. Features are extracted at progressively finer voxel scales (3 mm, 3 mm, 2×2×3 mm, 2×2×3 mm), starting with Dice-based overlap on segmentation outputs and progressively integrating feature-level alignment via L1 distances on selected internal layers (0.3/0.7 then 0.5/0.5 L1/Dice), ending with a final purely feature-based stage. A composite objective (IMPACT + mutual information + bending energy penalty) ensures robust cross-modality alignment while enforcing smooth, physically plausible deformations.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

CBCT_CT_TS/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)

ConvexAdam_Coarse/Model.py

@@ -33,6 +33,8 @@ NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engi
 runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
 """
 
+from typing import Annotated
+
 import itk
 import numpy as np
 import SimpleITK as sitk
@@ -40,12 +42,17 @@ import torch
 import tqdm
 from huggingface_hub import hf_hub_download
 from konfai.network import network
+from konfai.utils.config import Range
 from konfai.utils.dataset import Attribute, data_to_image, image_to_data
 
 DIM = 3
 # The feature model's input channel count is an intrinsic property of the pretrained model (grayscale
 # medical images), not a tunable — so it's fixed here, never a config/signature parameter.
 NUM_CHANNELS = 1
+
+# A UI reads the tuning knobs straight from the TYPES on ``RegistrationNet.__init__``: ``Annotated[.., Range]``
+# gives numeric spin bounds; a flat list (``models`` / ``distance`` / ``voxel_size``) has no constraint and is
+# edited as a plain list. The engine downloads models straight from the ``models`` refs (a local path is OK).
 _IMAGE_F = itk.Image[itk.F, DIM]
 
 
@@ -417,22 +424,22 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         models: list[str] = [],
         voxel_size: list[float] = [3.0, 3.0, 3.0],
-        overlap: int = 2,
+        overlap: Annotated[int, Range(1, 128)] = 2,
         layers_mask: list[bool] = [True],
         mixed_precision: bool = False,
-        grid_spacing: int = 4,
-        displacement_half_width: int = 6,
-        iterations: int = 150,
-        learning_rate: float = 0.2,
-        regularization_weight: float = 1.0,
-        grid_shrink: int = 4,
+        grid_spacing: Annotated[int, Range(1, 512)] = 4,
+        displacement_half_width: Annotated[int, Range(1, 512)] = 6,
+        iterations: Annotated[int, Range(0, 100000)] = 150,
+        learning_rate: Annotated[float, Range(0.0, 100.0)] = 0.2,
+        regularization_weight: Annotated[float, Range(0.0, 1000.0)] = 1.0,
+        grid_shrink: Annotated[int, Range(1, 128)] = 4,
         distance: list[str] = ["L1"],
         layers_weight: list[float] = [1.0],
         subset_features: list[int] = [],  # feature-channel indices to keep (empty = all); NOT a count
         pca: list[int] = [0],
         stages: list[str] = ["coarse", "fine"],
         linear: bool = True,
-        linear_iterations: int = 200,
+        linear_iterations: Annotated[int, Range(0, 100000)] = 200,
         seed: int = 42,
     ) -> None:
         super().__init__(

ConvexAdam_Coarse/app.json

@@ -3,7 +3,7 @@
     "short_description": "Global coarse ConvexAdam initialization (whole-volume, IMPACT/MIND features).",
     "description": "First stage of the ConvexAdam pipeline: an optional moments+affine linear pre-align followed by the ConvexAdam coarse coupled-convex initialization on the whole volume, driven by the IMPACT metric on MIND features. Produces a robust low-resolution displacement field on the fixed grid, meant to warm-start (and pre-resample the moving for) the fine stage.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

ConvexAdam_Composite/Model.py

@@ -33,6 +33,8 @@ NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engi
 runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
 """
 
+from typing import Annotated
+
 import itk
 import numpy as np
 import SimpleITK as sitk
@@ -40,12 +42,17 @@ import torch
 import tqdm
 from huggingface_hub import hf_hub_download
 from konfai.network import network
+from konfai.utils.config import Range
 from konfai.utils.dataset import Attribute, data_to_image, image_to_data
 
 DIM = 3
 # The feature model's input channel count is an intrinsic property of the pretrained model (grayscale
 # medical images), not a tunable — so it's fixed here, never a config/signature parameter.
 NUM_CHANNELS = 1
+
+# A UI reads the tuning knobs straight from the TYPES on ``RegistrationNet.__init__``: ``Annotated[.., Range]``
+# gives numeric spin bounds; a flat list (``models`` / ``distance`` / ``voxel_size``) has no constraint and is
+# edited as a plain list. The engine downloads models straight from the ``models`` refs (a local path is OK).
 _IMAGE_F = itk.Image[itk.F, DIM]
 
 
@@ -417,22 +424,22 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         models: list[str] = [],
         voxel_size: list[float] = [3.0, 3.0, 3.0],
-        overlap: int = 2,
+        overlap: Annotated[int, Range(1, 128)] = 2,
         layers_mask: list[bool] = [True],
         mixed_precision: bool = False,
-        grid_spacing: int = 4,
-        displacement_half_width: int = 6,
-        iterations: int = 150,
-        learning_rate: float = 0.2,
-        regularization_weight: float = 1.0,
-        grid_shrink: int = 4,
+        grid_spacing: Annotated[int, Range(1, 512)] = 4,
+        displacement_half_width: Annotated[int, Range(1, 512)] = 6,
+        iterations: Annotated[int, Range(0, 100000)] = 150,
+        learning_rate: Annotated[float, Range(0.0, 100.0)] = 0.2,
+        regularization_weight: Annotated[float, Range(0.0, 1000.0)] = 1.0,
+        grid_shrink: Annotated[int, Range(1, 128)] = 4,
         distance: list[str] = ["L1"],
         layers_weight: list[float] = [1.0],
         subset_features: list[int] = [],  # feature-channel indices to keep (empty = all); NOT a count
         pca: list[int] = [0],
         stages: list[str] = ["coarse", "fine"],
         linear: bool = True,
-        linear_iterations: int = 200,
+        linear_iterations: Annotated[int, Range(0, 100000)] = 200,
         seed: int = 42,
     ) -> None:
         super().__init__(

ConvexAdam_Composite/app.json

@@ -3,7 +3,7 @@
     "short_description": "Coarse + fine ConvexAdam registration in one app (IMPACT/MIND features).",
     "description": "The full ConvexAdam pipeline chained in one app: optional moments+affine linear pre-align, then a ConvexAdam coarse coupled-convex initialization, then an Adam instance-optimisation refinement, all driven by the IMPACT metric on MIND features. Produces the moved image and the displacement field on the fixed grid.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

ConvexAdam_Fine/Model.py

@@ -33,6 +33,8 @@ NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engi
 runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
 """
 
+from typing import Annotated
+
 import itk
 import numpy as np
 import SimpleITK as sitk
@@ -40,12 +42,17 @@ import torch
 import tqdm
 from huggingface_hub import hf_hub_download
 from konfai.network import network
+from konfai.utils.config import Range
 from konfai.utils.dataset import Attribute, data_to_image, image_to_data
 
 DIM = 3
 # The feature model's input channel count is an intrinsic property of the pretrained model (grayscale
 # medical images), not a tunable — so it's fixed here, never a config/signature parameter.
 NUM_CHANNELS = 1
+
+# A UI reads the tuning knobs straight from the TYPES on ``RegistrationNet.__init__``: ``Annotated[.., Range]``
+# gives numeric spin bounds; a flat list (``models`` / ``distance`` / ``voxel_size``) has no constraint and is
+# edited as a plain list. The engine downloads models straight from the ``models`` refs (a local path is OK).
 _IMAGE_F = itk.Image[itk.F, DIM]
 
 
@@ -417,22 +424,22 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         models: list[str] = [],
         voxel_size: list[float] = [3.0, 3.0, 3.0],
-        overlap: int = 2,
+        overlap: Annotated[int, Range(1, 128)] = 2,
         layers_mask: list[bool] = [True],
         mixed_precision: bool = False,
-        grid_spacing: int = 4,
-        displacement_half_width: int = 6,
-        iterations: int = 150,
-        learning_rate: float = 0.2,
-        regularization_weight: float = 1.0,
-        grid_shrink: int = 4,
+        grid_spacing: Annotated[int, Range(1, 512)] = 4,
+        displacement_half_width: Annotated[int, Range(1, 512)] = 6,
+        iterations: Annotated[int, Range(0, 100000)] = 150,
+        learning_rate: Annotated[float, Range(0.0, 100.0)] = 0.2,
+        regularization_weight: Annotated[float, Range(0.0, 1000.0)] = 1.0,
+        grid_shrink: Annotated[int, Range(1, 128)] = 4,
         distance: list[str] = ["L1"],
         layers_weight: list[float] = [1.0],
         subset_features: list[int] = [],  # feature-channel indices to keep (empty = all); NOT a count
         pca: list[int] = [0],
         stages: list[str] = ["coarse", "fine"],
         linear: bool = True,
-        linear_iterations: int = 200,
+        linear_iterations: Annotated[int, Range(0, 100000)] = 200,
         seed: int = 42,
     ) -> None:
         super().__init__(

ConvexAdam_Fine/app.json

@@ -3,7 +3,7 @@
     "short_description": "Fine Adam refinement (tileable; expects a coarse/pre-aligned start).",
     "description": "Second stage of the ConvexAdam pipeline: the Adam instance-optimisation refinement driven by the IMPACT metric on MIND features, with no linear pre-align (it assumes the moving is already on the fixed grid, e.g. resampled by the coarse stage). Runs from a zero warm-start and tiles naturally for large images (set a patch size).",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "vram_plan": {
         "8": {"patch_size": [128, 128, 128], "batch_size": 1},

Generic_Rigid/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

Generic_Rigid/Prediction.yml

@@ -7,7 +7,7 @@ Predictor:
       - Parameters_Rigid.txt
       outputs_criterions: None
       max_iterations: 0
-      spatial_samples: 0
+      spatial_samples: 2048
       parameter_overrides: []
   Dataset:
     groups_src:

Generic_Rigid/app.json

@@ -3,7 +3,7 @@
     "short_description": "Rigid registration using mutual information and a multi-resolution pyramid.",
     "description": "This preset performs rigid alignment using an Euler transform optimized with Adaptive Stochastic Gradient Descent. It uses a 4-level multi-resolution strategy and Mattes mutual information as similarity metric. Initial alignment based on image centers are enabled to ensure robust convergence for multimodal images.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

Generic_Rigid/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)

Generic_Rigid_BSpline/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

Generic_Rigid_BSpline/Prediction.yml

@@ -8,8 +8,8 @@ Predictor:
       - Parameters_BSpline.txt
       outputs_criterions: None
       max_iterations: 0
-      final_grid_spacing: 0.0
-      spatial_samples: 0
+      final_grid_spacing: 16.0
+      spatial_samples: 2048
       parameter_overrides: []
   Dataset:
     groups_src:

Generic_Rigid_BSpline/app.json

@@ -3,7 +3,7 @@
     "short_description": "Two-stage registration: rigid alignment followed by BSpline refinement.",
     "description": "A combined registration strategy: first a rigid Euler transform corrects global misalignment, then a BSpline model captures localized anatomical deformations. Both stages use a multi-resolution pyramid, mutual information, and stochastic optimization for robust performance across a wide range of multimodal imaging scenarios.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

Generic_Rigid_BSpline/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)

MR_CT_HeadNeck/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

MR_CT_HeadNeck/Prediction.yml

@@ -7,9 +7,9 @@ Predictor:
       - ParameterMap_MRI_HN.txt
       outputs_criterions: None
       max_iterations: 0
-      final_grid_spacing: 0.0
+      final_grid_spacing: 12.0
       subset_features: 0
-      spatial_samples: 0
+      spatial_samples: 2000
       parameter_overrides: []
       resolutions:
         '0':
@@ -72,7 +72,6 @@ Predictor:
               subset_features: 32
               pca: 0
               distance: L1
-      models_registry: VBoussot/impact-torchscript-models:models.json
       mode: Static
   Dataset:
     groups_src:

MR_CT_HeadNeck/app.json

@@ -3,7 +3,7 @@
     "short_description": "Optimized preset for MR/CT registration on head & neck",
     "description": "A four-level recursive B-spline deformable registration optimized for MRI head-and-neck images, driven by the IMPACT metric and combining semantic features extracted from the pretrained MIND model at progressively finer voxel scales (6 mm, 3 mm, 2 mm, 2 mm). The optimization follows a multi-resolution scheme with up to 300, 300, 250, and 200 ASGD iterations, using stochastic sampling and a composite metric (IMPACT + mutual information + bending energy) to achieve robust semantic alignment in complex head-and-neck anatomy.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

MR_CT_HeadNeck/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)

MR_CT_MRSeg/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

MR_CT_MRSeg/Prediction.yml

@@ -7,9 +7,9 @@ Predictor:
       - ParameterMap_MRI_MRSeg.txt
       outputs_criterions: None
       max_iterations: 0
-      final_grid_spacing: 0.0
+      final_grid_spacing: 14.0
       subset_features: 0
-      spatial_samples: 0
+      spatial_samples: 2000
       parameter_overrides: []
       resolutions:
         '0':
@@ -120,7 +120,6 @@ Predictor:
               subset_features: 64
               pca: 0
               distance: Dice
-      models_registry: VBoussot/impact-torchscript-models:models.json
       mode: Static
   Dataset:
     groups_src:

MR_CT_MRSeg/app.json

@@ -3,7 +3,7 @@
     "short_description": "Generic MR/CT deformable registration using MIND + MRSegmentator features",
     "description": "A four-level recursive B-spline deformable registration optimized for generic MR/CT alignment, driven by the IMPACT metric and combining semantic features from two pretrained models: MIND (L1 distance on a subset of 32 features) and MRSegmentator (Dice overlap on segmentation outputs with 64 features). Features are extracted at progressively finer voxel scales with level-dependent weighting between MIND and MRSegmentator. The optimization follows a multi-resolution ASGD scheme with a composite objective (IMPACT + mutual information + bending energy penalty) to ensure robust cross-modality semantic alignment and smooth deformations.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

MR_CT_MRSeg/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)

MR_CT_TS/Model.py

@@ -14,115 +14,89 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-"""Registration as a KonfAI model, built the idiomatic way (an ``add_module`` graph).
+"""Registration as a KonfAI model: the config -> elastix parameter-map mapping + the ``add_module`` graph.
 
-``RegistrationNet`` wires a single custom module ``ElastixRegistration`` that takes the fixed image
-(input branch ``0``) and the moving image (input branch ``1``) and returns the moving image resampled
-onto the FIXED grid. The module opts into receiving the per-input geometry via ``accepts_attributes``
-(mirroring KonfAI's ``CriterionWithAttribute`` loss convention): the KonfAI graph then hands it the
-``Attribute`` (Origin/Spacing/Direction) of each branch alongside the tensors, which the elastix engine
-needs to register in physical space.
+``RegistrationNet`` wires ``ElastixRegistration`` (fixed = branch 0, moving = branch 1, fixed/moving masks =
+2/3) and splits its output into ``MovedImage`` / ``DisplacementField`` on the fixed grid. This module owns
+the MAPPING — the per-resolution model matrix (``resolutions``) turned into IMPACT parameter-map lines, and
+the config schema (``ModelSpec`` / ``ResolutionSpec``). The elastix RUNTIME (binary install, model download,
+subprocess, progress) lives in ``elastix_engine.py`` and is imported only when the graph is built.
 
-Runs through the standard ``konfai.predictor.predict`` path in whole-volume mode:
+A UI reads the tuning knobs straight from the TYPES below: ``Literal`` (a fixed set),
+``Annotated[.., Range]`` (numeric bounds), ``Annotated[str, Choices(...)]`` (a resolver the app owns).
 
-    Patch.patch_size: None        # one patch per case (registration is global)
-    batch_size: 1                 # one fixed/moving pair at a time
-
-NOTE: do NOT add ``from __future__ import annotations`` here — KonfAI's config engine relies on
-runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break arg resolution.
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine reads runtime annotations
+(``get_origin``); PEP 563 stringized annotations break arg resolution.
 """
 
 import json
 import os
 import re
-import shutil
-import subprocess  # nosec B404
-import tempfile
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Annotated, Literal
 
-import numpy as np
-import SimpleITK as sitk
 import torch
-import tqdm
 from huggingface_hub import hf_hub_download
-from install import get_elastix_bin, install_elastix_impact, try_elastix
 from konfai.network import network
-from konfai.utils.dataset import Attribute, data_to_image, image_to_data
-
-# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
-# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
-ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+from konfai.utils.config import Choices, Range
 
-# ---------------------------------------------------------------------------------------------------
-# Per-resolution model matrix (the config is the source of truth) -> generated IMPACT parameter map.
 # IMPACT field docs: https://github.com/vboussot/ImpactLoss/tree/main/ParameterMaps
-# The forced per-model props (dimension/channels/FOV formula) live in a registry (models.json on
-# VBoussot/impact-torchscript-models); the config carries the FREE knobs (which models per resolution,
-# feature voxel size, iterations, per-model layer weights/mask/subset/pca/distance) and the global
-# ``mode``. PatchSize follows ImpactMode: Static -> "0 0 0" (whole image); Jacobian -> the model FOV
-# evaluated from the registry formula (MIND 2*r*d+1, TS/MRSeg 2^l+3, SAM 29, DINOv2 14) as a cube.
-# ---------------------------------------------------------------------------------------------------
-
+# A model's FIXED props (dimension / channels / FOV formula) come from the registry (models.json on
+# VBoussot/impact-torchscript-models); the config carries the FREE knobs (models per resolution, voxel size,
+# iterations, per-model weights/mask/subset/pca/distance) and the global ``mode``.
 _IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
 
-# ``2^l+3`` grows with depth but the segmenters' receptive field plateaus: layers 7-8 share layer 6's
-# FOV (the "ramp max"). A config that deep should really run in Static (whole image) anyway; in Jacobian
-# we clamp ``l`` to this plateau so the patch stays finite and matches the real FOV.
+# ``2^l+3`` plateaus: segmenter layers 7-8 share layer 6's receptive field. Deeper configs should run
+# Static anyway; in Jacobian we clamp ``l`` to this plateau.
 _FOV_RAMP_MAX_LAYER = 6
 
 
+def registry_choices() -> list[str]:
+    """The ``ref`` picker's values — model refs (``repo:path``) from the registry the engine already fetches
+    (offline-first). A user may still point ``ref`` at a local model."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
 def _num(x: object) -> str:
-    """Format a number the elastix way: integers without a trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
+    """Format a number the elastix way: no trailing '.0' (6.0 -> '6', 0.2 -> '0.2')."""
     return "%g" % float(x)
 
 
+@dataclass
 class ModelSpec:
-    """One feature model at one resolution, with its OWN config (several models may share a resolution).
-
-    ``ref`` selects the model; ``voxel_size`` / ``layers_weight`` / ``subset_features`` / ``pca`` /
-    ``distance`` are its free per-(resolution, model) tuning knobs (the doc's per-model *tuning* fields).
-    The intrinsic per-model props — dimension, channels, ``layers_mask``, patch-size (FOV) — come from the
-    registry (read-only); ``layers_mask`` / ``distance`` left empty fall back to the registry default.
-    """
+    """One feature model at one resolution (several may share a resolution). ``ref`` picks the model; the
+    rest are its per-(resolution, model) knobs. Dimension / channels / FOV are intrinsic — from the registry
+    (``models.json``) keyed by ``ref`` — never tuned."""
 
-    def __init__(
-        self,
-        ref: str,
-        voxel_size: list[float] = [],
-        layers_weight: list[float] = [1.0],
-        subset_features: int = 0,
-        pca: int = 0,
-        distance: str = "",
-        layers_mask: str = "",
-    ) -> None:
-        self.ref = ref
-        self.voxel_size = voxel_size
-        self.layers_weight = layers_weight
-        self.subset_features = subset_features
-        self.pca = pca
-        self.distance = distance
-        self.layers_mask = layers_mask
+    ref: Annotated[str, Choices(registry_choices)]
+    voxel_size: list[float] = field(default_factory=list)
+    layers_weight: list[float] = field(default_factory=lambda: [1.0])
+    subset_features: Annotated[int, Range(0, 1000)] = 0
+    pca: Annotated[int, Range(0, 100)] = 0
+    distance: Literal["L1", "L2", "Dice", "Cosine", "NCC"] = "L1"
+    layers_mask: str = ""
 
 
+@dataclass
 class ResolutionSpec:
-    """One elastix resolution level: its iteration budget and the models compared there (each self-configured)."""
+    """One elastix resolution level: its iteration budget and the (self-configured) models compared there."""
 
-    def __init__(self, max_iterations: int, models: dict[str, ModelSpec]) -> None:
-        self.max_iterations = max_iterations
-        self.models = models
+    max_iterations: Annotated[int, Range(1, 100000)]
+    models: dict[str, ModelSpec]
 
 
 def _sorted_specs(mapping: dict) -> list:
-    """dict keyed by string indices ('0','1',...) -> values in numeric order (well-defined res/model order)."""
+    """dict keyed by string indices ('0','1',...) -> values in numeric order."""
     return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
 
 
 def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
-    """Load models.json (forced params per model) from the model repo on Hugging Face.
+    """Load models.json (the fixed params per model) from the model repo on Hugging Face.
 
-    The registry is NOT bundled with the preset — it lives on the models repo and is fetched from there.
-    Resolution: the ``KONFAI_IMPACT_MODELS_REGISTRY`` env path wins (dev/offline); otherwise ``ref`` must be
-    a ``repo:file`` Hugging Face reference.
+    The registry is NOT bundled with the preset. ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins for
+    dev/offline; otherwise ``ref`` must be a ``repo:file`` Hugging Face reference.
     """
     local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
     if local:
@@ -139,17 +113,16 @@ def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
 
 
 def _model_key(ref: str) -> str:
-    """Registry key / staged relative path = the model file within the models repo (strip a 'repo:' prefix)."""
+    """Registry key / staged relative path = the model file within the repo (strip a 'repo:' prefix)."""
     return ref.split(":", 1)[1] if ":" in ref else ref
 
 
 def _deepest_active_layer(layers_mask: str) -> int:
-    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index read left-to-right.
+    """Deepest (largest-FOV) layer selected by ``layers_mask``, as a 0-based index.
 
-    A model returns its feature layers shallow->deep (``[layer_0, layer_1, ...]``, see the model repo's
-    build scripts); ``layers_mask`` has one char per returned layer, position ``i`` == ``layer_i``, ``'1'``
-    = selected. In Jacobian the patch must cover the receptive field of the DEEPEST selected layer, so the
-    FOV is governed by the rightmost ``'1'``.
+    A model returns its layers shallow->deep; ``layers_mask`` has one char per returned layer, position ``i``
+    == ``layer_i``, ``'1'`` = selected. In Jacobian the patch must cover the DEEPEST selected layer's
+    receptive field, so the FOV is governed by the rightmost ``'1'``.
     """
     mask = layers_mask.strip().strip('"')
     active = [i for i, char in enumerate(mask) if char == "1"]
@@ -161,13 +134,13 @@ def _deepest_active_layer(layers_mask: str) -> int:
 def _fov_value(fov: dict, layers_mask: str) -> int:
     """Evaluate a model's field-of-view (in voxels) from its registry ``fov`` spec.
 
-    Supported formulas (from the model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
-      ``2*r*d+1``  MIND, from the handcrafted radius ``r`` / dilation ``d`` (e.g. R1D2 -> 5);
-      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = the deepest layer picked by ``layers_mask``,
-                   clamped to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
-      a bare int  a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
-      ``Global``   Anatomix — whole-image only (Static); has no finite Jacobian patch -> error.
-    An explicit ``value`` in the spec is honoured as a precomputed shortcut when the formula needs none.
+    Formulas (model repo, https://huggingface.co/VBoussot/impact-torchscript-models):
+      ``2*r*d+1``  MIND, from radius ``r`` / dilation ``d`` (R1D2 -> 5);
+      ``2^l+3``    TotalSegmentator / MRSegmentator, ``l`` = deepest layer picked by ``layers_mask``, clamped
+                   to the receptive-field plateau ``_FOV_RAMP_MAX_LAYER`` (layers 7-8 -> layer 6);
+      a bare int   a fixed FOV (SAM2.1 -> 29, DINOv2 -> 14);
+      ``Global``   Anatomix — whole-image only (Static); no finite Jacobian patch -> error.
+    An explicit ``value`` in the spec is honoured as a precomputed shortcut.
     """
     formula = str(fov.get("formula", "")).strip()
     key = re.sub(r"\s+", "", formula).lower()
@@ -185,9 +158,9 @@ def _fov_value(fov: dict, layers_mask: str) -> int:
 
 
 def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
-    """PatchSize from the model FOV, one token per model axis (2D model -> 2 tokens, 3D -> 3): Static ->
-    whole image (all zeros); Jacobian -> the evaluated FOV repeated over the axes. A 2D model mixed with a
-    3D one at a resolution concatenates as e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
+    """PatchSize from the model FOV, one token per model axis (2D -> 2 tokens, 3D -> 3): Static -> whole
+    image (all zeros); Jacobian -> the evaluated FOV per axis. A 2D+3D mix at a resolution concatenates,
+    e.g. ``29 29 11 11 11`` (SAM 2D + TS 3D), matching IMPACT."""
     dim = int(entry.get("dimension", 3))
     if mode.strip().strip('"').lower() != "jacobian":
         return " ".join(["0"] * dim)
@@ -195,16 +168,13 @@ def _patch_size(mode: str, entry: dict, layers_mask: str) -> str:
     return " ".join([str(fov)] * dim)
 
 
-def generate_impact_parameter_map(
-    template_text: str, resolutions: dict, registry: dict, mode: str = "Static"
-) -> str:
+def generate_impact_parameter_map(template_text: str, resolutions: dict, registry: dict, mode: str = "Static") -> str:
     """Rewrite the resolution-dependent lines of ``template_text`` from the model matrix ``resolutions``.
 
     Regenerated: MaximumNumberOfIterations, NumberOfResolutions, Fixed/MovingImagePyramidRescaleSchedule,
-    ImpactMode (from the config ``mode``), and the whole ImpactXxxK block; every other template line is
-    kept verbatim (optimizer, transform, metric weights, components...). N (number of resolutions) is
-    deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0`` (whole image); Jacobian -> the
-    per-model FOV evaluated from the registry formula and the cell's ``layers_mask``.
+    ImpactMode, and the whole ImpactXxxK block; every other line is kept verbatim. N (number of resolutions)
+    is deduced from the config. ``mode`` drives PatchSize: Static -> ``0 0 0``; Jacobian -> the per-model FOV
+    from the registry formula and the cell's ``layers_mask``.
     """
     res = _sorted_specs(resolutions)
     n = len(res)
@@ -218,9 +188,8 @@ def generate_impact_parameter_map(
         def row(stem: str, values: list[str]) -> None:
             impact.append(f"(Impact{stem}{k} " + " ".join(values) + ")")
 
-        # From the registry (models.json on the model repo) ONLY the 3 truly model-fixed props:
-        # Dimension, NumberOfChannels, PatchSize (the model FOV). Everything else is a per-model tuning knob
-        # taken straight from the cell: VoxelSize / LayersMask / SubsetFeatures / PCA / Distance / LayersWeight.
+        # From the registry ONLY the 3 truly model-fixed props (Dimension, NumberOfChannels, PatchSize = the
+        # model FOV); everything else is a per-model knob taken straight from the cell.
         row("ModelsPath", [f'"{_model_key(m.ref)}"' for m in models])
         row("Dimension", [e["dimension"] for e in entries])
         row("NumberOfChannels", [e["numberofchannels"] for e in entries])
@@ -234,8 +203,7 @@ def generate_impact_parameter_map(
         impact.append("")  # blank line between resolutions, mirroring the reference maps
 
     # The per-resolution block is the contiguous span from the first to the last ``Impact<name><k>`` line
-    # (the blank lines the reference maps put BETWEEN resolutions fall inside that span). Replace the whole
-    # span in one shot with the generated block, so the reference blanks are not kept on top of ours.
+    # (inner blanks fall inside it). Replace the whole span at its first line so reference blanks aren't kept.
     lines = template_text.splitlines()
     indexed = [(re.match(r"^\s*\((\S+?)\s+(.*?)\)\s*$", ln), ln) for ln in lines]
     block_rows = [i for i, (m, _) in enumerate(indexed) if m and re.match(r"^Impact[A-Za-z]+\d+$", m.group(1))]
@@ -260,352 +228,6 @@ def generate_impact_parameter_map(
     return "\n".join(out)
 
 
-class ElastixEngine:
-    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
-
-    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix
-    does NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
-    """
-
-    def __init__(
-        self,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        self._bundle_dir = Path(__file__).resolve().parent
-        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
-        self._max_iterations = max_iterations
-        self._final_grid_spacing = final_grid_spacing
-        self._subset_features = subset_features
-        self._spatial_samples = spatial_samples
-        self._parameter_overrides = list(parameter_overrides)
-        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
-        # samples random patches sized to the model FOV each iteration. Global knob: one mode per preset.
-        self._mode = mode
-        # Matrix mode: when `resolutions` is given the parameter map is GENERATED from it (the config is the
-        # source of truth). An empty `resolutions` = an intensity preset (no IMPACT feature models): the fixed
-        # parameter maps are staged with only the global knob overrides.
-        self._resolutions = resolutions
-        self._registry = load_models_registry(models_registry) if resolutions else {}
-        # The feature models are DERIVED — the unique refs across the matrix cells (no flat `models` param).
-        models: list[str] = []
-        for res in _sorted_specs(resolutions):
-            for model in _sorted_specs(res.models):
-                if model.ref not in models:
-                    models.append(model.ref)
-        self._models = models
-        # `iterations` (the progress-bar total) is NOT a config parameter — it is DERIVED: the sum of the
-        # per-resolution iteration budgets, read from the matrix (matrix mode) or the maps (legacy).
-        self._iterations = self._total_iterations()
-        self._elastix_bin = self._ensure_binary()
-        self._local_models = self._download_models()
-
-    def _total_iterations(self) -> int:
-        """Total iterations across all resolutions — the progress-bar budget, derived from the config."""
-        if self._resolutions:
-            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
-        total = 0
-        for src in self._parameter_maps:
-            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
-            if match:
-                total += sum(int(token) for token in match.group(1).split())
-        return total
-
-    def _ensure_binary(self) -> Path:
-        # Optional override: point at an existing elastix-IMPACT install (skips the download).
-        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
-        if override:
-            try_elastix(Path(override))
-            return get_elastix_bin(Path(override)).resolve()
-        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
-        try:
-            try_elastix(ELASTIX_CACHE)
-        except Exception:
-            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
-            try_elastix(ELASTIX_CACHE)
-        return get_elastix_bin(ELASTIX_CACHE).resolve()
-
-    def _download_models(self) -> list[tuple[str, Path]]:
-        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
-        models = []
-        for ref in self._models:
-            repo, filename = ref.split(":", 1)
-            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
-            models.append((filename, local))
-        return models
-
-    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
-        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
-
-        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value that replaces
-        **each** existing token, so per-resolution / per-model multiplicity is preserved (e.g.
-        ``(MaximumNumberOfIterations 500 250)`` -> ``(MaximumNumberOfIterations 300 300)``). ``exact``
-        entries (from ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win
-        over the named knobs. Overrides only REPLACE keys already present in a map — never inject new ones.
-        ``global_only`` (matrix mode) keeps just the map-wide knobs and drops ``max_iterations`` /
-        ``subset_features`` — the per-resolution matrix already sets those per cell.
-        """
-        per_token: dict[str, str] = {}
-        if not global_only and self._max_iterations > 0:
-            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
-        if self._final_grid_spacing > 0:
-            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
-        if not global_only and self._subset_features > 0:
-            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
-        if self._spatial_samples > 0:
-            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
-        exact: list[tuple[str, str]] = []
-        for entry in self._parameter_overrides:
-            key, sep, value = entry.partition("=")
-            if not sep or not key.strip():
-                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
-            exact.append((key.strip(), value.strip()))
-        return per_token, exact
-
-    @staticmethod
-    def _apply_map_overrides(
-        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
-    ) -> str:
-        """Patch a parameter map's text: set ImpactGPU to the device, apply exact key overrides, replace each
-        token of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
-        """
-        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
-        requested = set(per_token) | {key for key, _ in exact}
-        seen: set[str] = set()
-        lines = []
-        for line in text.splitlines():
-            match = entry_pattern.match(line)
-            if match:
-                indent, key, values = match.group(1), match.group(2), match.group(3)
-                if key == "ImpactGPU":
-                    line = f"{indent}(ImpactGPU {device_index})"
-                else:
-                    exact_value = next((value for k, value in exact if k == key), None)
-                    if exact_value is not None:
-                        seen.add(key)
-                        line = f"{indent}({key} {exact_value})"
-                    else:
-                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
-                        if token_key in per_token:
-                            seen.add(token_key)
-                            replaced = " ".join(per_token[token_key] for _ in values.split())
-                            line = f"{indent}({key} {replaced})"
-            lines.append(line)
-        # Overrides never inject keys, so a knob set for a key absent from every map would silently do
-        # nothing — surface it (e.g. final_grid_spacing on a rigid-only preset).
-        for key in sorted(requested - seen):
-            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
-        return "\n".join(lines)
-
-    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
-        """Stage the parameter maps into the work dir.
-
-        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
-        knobs (grid spacing, spatial samples, exact overrides) — the matrix already sets iterations and
-        features per cell. Legacy mode copies the preset's maps and applies every per-token / exact override.
-        Both set the ImpactGPU device.
-        """
-        staged = []
-        for src in self._parameter_maps:
-            if self._resolutions:
-                text = generate_impact_parameter_map(
-                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
-                )
-                per_token, exact = self._parameter_map_overrides(global_only=True)
-            else:
-                text = src.read_text(encoding="utf-8")
-                per_token, exact = self._parameter_map_overrides()
-            text = self._apply_map_overrides(text, per_token, exact, device_index)
-            dst = work / src.name
-            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
-            staged.append(dst)
-        return staged
-
-    def register(
-        self,
-        fixed: sitk.Image,
-        moving: sitk.Image,
-        device_index: int,
-        fixed_mask: sitk.Image | None = None,
-        moving_mask: sitk.Image | None = None,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
-
-        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region
-        (elastix ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
-        """
-        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
-        try:
-            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
-            sitk.WriteImage(fixed, str(fixed_path))
-            sitk.WriteImage(moving, str(moving_path))
-
-            # Stage the feature models at the relative path the parameter maps reference
-            # (e.g. ImpactModelsPath0 "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
-            for rel_name, model_path in self._local_models:
-                dst = work / rel_name
-                dst.parent.mkdir(parents=True, exist_ok=True)
-                if not dst.exists():
-                    dst.symlink_to(model_path)
-
-            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
-            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
-                if mask is not None:
-                    mask_path = work / name
-                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
-                    args += [flag, str(mask_path)]
-            args += ["-out", str(work)]
-            for pmap in self._stage_parameter_maps(work, device_index):
-                args += ["-p", str(pmap)]
-
-            # Stream elastix stdout and drive a tqdm bar over its iterations so SlicerKonfAI (which parses
-            # the "N% done/total" progress line) shows real progress during the long registration.
-            # Make the elastix binary's own libs (bundled libtorch under <install>/lib) and any extra
-            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
-            env = os.environ.copy()
-            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
-            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
-            proc = subprocess.Popen(  # nosec B603
-                args,
-                cwd=str(work),
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                text=True,
-                bufsize=1,
-                env=env,
-            )
-            captured: list[str] = []
-            iteration_line = re.compile(r"^\d+\s")
-            # ``iterations`` is the total iteration budget declared for the preset (summed over the
-            # chained parameter maps), so the bar spans the whole chain of registration stages. A tuned
-            # ``max_iterations`` makes that declared budget stale — fall back to an open-ended bar.
-            budget = None if self._max_iterations > 0 else (self._iterations or None)
-            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
-            assert proc.stdout is not None
-            resolution = 0
-            for line in proc.stdout:
-                captured.append(line)
-                stripped = line.strip()
-                if stripped.startswith("Resolution:"):
-                    try:
-                        resolution = int(stripped.split(":", 1)[1])
-                    except ValueError:
-                        pass
-                elif iteration_line.match(line):
-                    progress.update(1)
-                    # Mirror KonfAI's informative bars (which surface runtime state in the description):
-                    # show the elastix resolution level and the similarity metric being optimised so the
-                    # bar conveys convergence, not a bare iteration count. Column 2 of the iteration table
-                    # is the metric (header: "1:ItNr 2:Metric ...").
-                    columns = line.split()
-                    if len(columns) > 1:
-                        try:
-                            progress.set_description(
-                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
-                            )
-                        except ValueError:
-                            pass
-            progress.close()
-            returncode = proc.wait()
-            if returncode != 0:
-                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
-
-            transforms = sorted(
-                work.glob("TransformParameters.*-Composite.itk.txt"),
-                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
-            )
-            if not transforms:
-                raise FileNotFoundError("elastix produced no composite transform file.")
-            transform = sitk.ReadTransform(str(transforms[-1]))
-
-            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
-            dvf = sitk.TransformToDisplacementField(
-                transform,
-                sitk.sitkVectorFloat64,
-                fixed.GetSize(),
-                fixed.GetOrigin(),
-                fixed.GetSpacing(),
-                fixed.GetDirection(),
-            )
-            moved_np, _ = image_to_data(moved)
-            dvf_np, _ = image_to_data(dvf)
-            return moved_np, dvf_np
-        finally:
-            shutil.rmtree(work, ignore_errors=True)
-
-
-class ElastixRegistration(torch.nn.Module):
-    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
-
-    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
-    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix
-    needs the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
-    """
-
-    accepts_attributes = True
-
-    def __init__(
-        self,
-        engine: str,
-        parameter_maps: list[str],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
-        parameter_overrides: list[str] = [],
-        resolutions: dict = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
-    ) -> None:
-        super().__init__()
-        if engine != "elastix":
-            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
-        self._engine = ElastixEngine(
-            parameter_maps,
-            max_iterations,
-            final_grid_spacing,
-            subset_features,
-            spatial_samples,
-            parameter_overrides,
-            resolutions,
-            models_registry,
-            mode,
-        )
-
-    def forward(
-        self,
-        fixed: torch.Tensor,
-        moving: torch.Tensor,
-        fixed_mask: torch.Tensor,
-        moving_mask: torch.Tensor,
-        attributes: list[list[Attribute]],
-    ) -> torch.Tensor:
-        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each is a list[Attribute] over the batch.
-        # Returns, per sample, the moved image (1 channel) channel-stacked with the displacement field
-        # (dim channels), both on the fixed grid; downstream ChannelSelect modules split them. A mask covering
-        # the whole image (the auto-filled default when the user supplies none) restricts nothing.
-        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
-        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
-        combined = []
-        for b in range(fixed.shape[0]):
-            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
-            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
-            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
-            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
-            moved_np, dvf_np = self._engine.register(
-                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
-            )
-            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
-        return torch.stack(combined, dim=0).to(fixed.device)
-
-
 class ChannelSelect(torch.nn.Module):
     """Select a channel slice ``[start:stop]`` (splits the registration output into moved / DVF)."""
 
@@ -619,13 +241,13 @@ class ChannelSelect(torch.nn.Module):
 
 
 class RegistrationNet(network.Network):
-    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1,
-    fixed mask = branch 2, moving mask = branch 3; masks restrict the metric, whole-image = no restriction).
+    """Pairwise registration as an ``add_module`` graph (fixed = branch 0, moving = branch 1, fixed mask = 2,
+    moving mask = 3; masks restrict the metric, whole-image = no restriction).
 
-    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and
-    ``DisplacementField`` (the dim-component displacement field, in mm). ``ElastixRegistration`` produces
-    both channel-stacked; two ``ChannelSelect`` modules split them into the named outputs referenced by
-    ``Prediction.yml``. Output geometry is attached by the predictor via ``same_as_group: Volume_0:Fixed``.
+    Outputs (both on the fixed grid): ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the dim-component displacement field, mm). ``ElastixRegistration`` produces both channel-stacked; two
+    ``ChannelSelect`` modules split them. Output geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``.
     """
 
     def __init__(
@@ -637,23 +259,21 @@ class RegistrationNet(network.Network):
         outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
         engine: str = "elastix",
         parameter_maps: list[str] = [],
-        max_iterations: int = 0,
-        final_grid_spacing: float = 0.0,
-        subset_features: int = 0,
-        spatial_samples: int = 0,
+        max_iterations: Annotated[int, Range(0, 100000)] = 0,
+        final_grid_spacing: Annotated[float, Range(0.0, 100.0)] = 0.0,
+        subset_features: Annotated[int, Range(0, 1000)] = 0,
+        spatial_samples: Annotated[int, Range(0, 100000)] = 0,
         parameter_overrides: list[str] = [],
         resolutions: dict[str, ResolutionSpec] = {},
-        models_registry: str = _IMPACT_MODELS_REGISTRY,
-        mode: str = "Static",
+        mode: Literal["Static", "Jacobian"] = "Static",
     ) -> None:
-        # The registration is fully described by the per-resolution model matrix ``resolutions`` (config =
-        # source of truth): each resolution lists its models, each model self-configured (ref, voxel_size,
-        # layers_mask, layers_weight, subset_features, pca, distance); intrinsic per-model props come from
-        # ``models_registry``. The feature-model download list is DERIVED from the matrix (no flat ``models``).
-        # Global knobs override the generated map: final_grid_spacing -> FinalGridSpacingInPhysicalUnits (mm),
-        # spatial_samples -> NumberOfSpatialSamples, parameter_overrides ('Key=value') -> any other entry.
-        # An empty ``resolutions`` = an intensity-only preset (no IMPACT models): the fixed maps are staged
-        # with just the global overrides. The total iteration count is derived (sum of per-resolution budgets).
+        # The registration is fully described by ``resolutions`` (config = source of truth): each resolution
+        # lists its self-configured models; the download list is derived from the cells. Global knobs override
+        # the generated map (final_grid_spacing -> FinalGridSpacingInPhysicalUnits mm, spatial_samples ->
+        # NumberOfSpatialSamples, parameter_overrides 'Key=value'). Empty ``resolutions`` = an intensity-only
+        # preset (fixed maps + overrides). The elastix runtime is imported here (heavy: torch/sitk/subprocess).
+        from elastix_engine import ElastixRegistration
+
         super().__init__(
             in_channels=1,
             optimizer=optimizer,
@@ -672,7 +292,6 @@ class RegistrationNet(network.Network):
                 spatial_samples,
                 parameter_overrides,
                 resolutions,
-                models_registry,
                 mode,
             ),
             in_branch=[0, 1, 2, 3],

MR_CT_TS/Prediction.yml

@@ -7,9 +7,9 @@ Predictor:
       - ParameterMap_MRI_TS.txt
       outputs_criterions: None
       max_iterations: 0
-      final_grid_spacing: 0.0
+      final_grid_spacing: 14.0
       subset_features: 0
-      spatial_samples: 0
+      spatial_samples: 2000
       parameter_overrides: []
       resolutions:
         '0':
@@ -120,7 +120,6 @@ Predictor:
               subset_features: 64
               pca: 0
               distance: Dice
-      models_registry: VBoussot/impact-torchscript-models:models.json
       mode: Static
   Dataset:
     groups_src:

MR_CT_TS/app.json

@@ -3,7 +3,7 @@
     "short_description": "Generic MR/CT deformable registration using MIND + Totalsegmentator features",
     "description": "A four-level recursive B-spline deformable registration optimized for generic MR/CT alignment, driven by the IMPACT metric and combining semantic features from two pretrained models: MIND (L1 distance on a subset of 32 features) and Totalsegmentator (Dice overlap on segmentation outputs with 64 features). Features are extracted at progressively finer voxel scales (6 mm, 3 mm, 2 mm, 2 mm) with level-dependent weighting between MIND and MRSegmentator (0.2/0.8, 0.3/0.7, 0.6/0.4, 0.7/0.3). The optimization follows a multi-resolution ASGD scheme with up to 400, 300, 200, and 200 iterations using 2000 random spatial samples, and a composite objective (IMPACT + mutual information + bending energy penalty) to ensure robust cross-modality semantic alignment and smooth deformations.",
     "task": "registration",
-    "tta": 3,
+    "tta": 0,
     "mc_dropout": 0,
     "models": [
         "model.pt"

MR_CT_TS/elastix_engine.py

@@ -0,0 +1,375 @@
+# Copyright (c) 2025 Valentin Boussot
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Elastix-IMPACT runtime for the registration bundle.
+
+``ElastixEngine`` installs the elastix-IMPACT binary, downloads the TorchScript feature models, stages the
+parameter maps (generated from the model matrix or copied + overridden), runs the subprocess, and resamples.
+``ElastixRegistration`` is the graph module ``RegistrationNet`` wires — it bridges KonfAI tensors <-> SITK
+images. The config -> parameter-map MAPPING lives in ``Model.py`` and is imported here.
+"""
+
+import os
+import re
+import shutil
+import subprocess  # nosec B404
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+import tqdm
+from huggingface_hub import hf_hub_download
+from install import get_elastix_bin, install_elastix_impact, try_elastix
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+from Model import _sorted_specs, generate_impact_parameter_map, load_models_registry
+
+# Elastix + IMPACT binary is cached once here (heavy: binary + LibTorch) and reused across runs.
+# Set KONFAI_ELASTIX_DIR to point at an existing install and skip the download.
+ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
+
+
+class ElastixEngine:
+    """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
+
+    NOTE: the elastix-IMPACT metric lives only in the custom ``elastix-impact`` binary (SimpleElastix does
+    NOT ship it), so registration is a subprocess call, not ``sitk.ElastixImageFilter``.
+    """
+
+    def __init__(
+        self,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        self._bundle_dir = Path(__file__).resolve().parent
+        self._parameter_maps = [self._bundle_dir / Path(p).name for p in parameter_maps]
+        self._max_iterations = max_iterations
+        self._final_grid_spacing = final_grid_spacing
+        self._subset_features = subset_features
+        self._spatial_samples = spatial_samples
+        self._parameter_overrides = list(parameter_overrides)
+        # ImpactMode: Static computes features once per level (PatchSize 0 0 0 = whole image); Jacobian
+        # samples random FOV-sized patches each iteration. One mode per preset.
+        self._mode = mode
+        # Matrix mode: with ``resolutions`` the map is GENERATED from it. Empty ``resolutions`` = an
+        # intensity preset (no IMPACT models): the fixed maps are staged with only the global overrides.
+        self._resolutions = resolutions
+        self._registry = load_models_registry() if resolutions else {}
+        # Feature models are DERIVED — the unique refs across the matrix cells (no flat ``models`` param).
+        models: list[str] = []
+        for res in _sorted_specs(resolutions):
+            for model in _sorted_specs(res.models):
+                if model.ref not in models:
+                    models.append(model.ref)
+        self._models = models
+        # ``iterations`` (the progress-bar total) is DERIVED: the sum of per-resolution iteration budgets.
+        self._iterations = self._total_iterations()
+        self._elastix_bin = self._ensure_binary()
+        self._local_models = self._download_models()
+
+    def _total_iterations(self) -> int:
+        """Total iterations across resolutions — the progress-bar budget, from the config (or the maps)."""
+        if self._resolutions:
+            return sum(int(res.max_iterations) for res in _sorted_specs(self._resolutions))
+        total = 0
+        for src in self._parameter_maps:
+            match = re.search(r"\(MaximumNumberOfIterations\s+([^)]*)\)", src.read_text(encoding="utf-8"))
+            if match:
+                total += sum(int(token) for token in match.group(1).split())
+        return total
+
+    def _ensure_binary(self) -> Path:
+        # Optional override: point at an existing elastix-IMPACT install (skips the download).
+        override = os.environ.get("KONFAI_ELASTIX_DIR", "")
+        if override:
+            try_elastix(Path(override))
+            return get_elastix_bin(Path(override)).resolve()
+        ELASTIX_CACHE.mkdir(parents=True, exist_ok=True)
+        try:
+            try_elastix(ELASTIX_CACHE)
+        except Exception:
+            install_elastix_impact(ELASTIX_CACHE, force_cuda=False, force_cpu=False)
+            try_elastix(ELASTIX_CACHE)
+        return get_elastix_bin(ELASTIX_CACHE).resolve()
+
+    def _download_models(self) -> list[tuple[str, Path]]:
+        """Fetch the TorchScript feature models (``repo:filename``); keep (relative_name, local_path)."""
+        models = []
+        for ref in self._models:
+            repo, filename = ref.split(":", 1)
+            local = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+            models.append((filename, local))
+        return models
+
+    def _parameter_map_overrides(self, global_only: bool = False) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """The tuned knobs as parameter-map overrides: ``(per_token, exact)``.
+
+        ``per_token`` maps an elastix key (or the ``ImpactSubsetFeatures`` prefix) to a value replacing
+        **each** existing token, preserving per-resolution / per-model multiplicity. ``exact`` entries (from
+        ``parameter_overrides``, ``Key=value text``) replace the whole value verbatim and win over the named
+        knobs. Overrides only REPLACE keys already present — never inject. ``global_only`` (matrix mode) drops
+        ``max_iterations`` / ``subset_features`` (the matrix already sets those per cell).
+        """
+        per_token: dict[str, str] = {}
+        if not global_only and self._max_iterations > 0:
+            per_token["MaximumNumberOfIterations"] = str(int(self._max_iterations))
+        if self._final_grid_spacing > 0:
+            per_token["FinalGridSpacingInPhysicalUnits"] = str(float(self._final_grid_spacing))
+        if not global_only and self._subset_features > 0:
+            per_token["ImpactSubsetFeatures"] = str(int(self._subset_features))  # prefix: indexed per metric
+        if self._spatial_samples > 0:
+            per_token["NumberOfSpatialSamples"] = str(int(self._spatial_samples))
+        exact: list[tuple[str, str]] = []
+        for entry in self._parameter_overrides:
+            key, sep, value = entry.partition("=")
+            if not sep or not key.strip():
+                raise ValueError(f"Invalid parameter_overrides entry '{entry}': expected 'Key=value text'.")
+            exact.append((key.strip(), value.strip()))
+        return per_token, exact
+
+    @staticmethod
+    def _apply_map_overrides(
+        text: str, per_token: dict[str, str], exact: list[tuple[str, str]], device_index: int
+    ) -> str:
+        """Patch a parameter map: set ImpactGPU to the device, apply exact key overrides, replace each token
+        of a per-token knob (preserving multiplicity), and warn for a requested key absent from the map.
+        """
+        entry_pattern = re.compile(r"^(\s*)\((\S+)((?:\s+[^)]*)?)\)\s*$")
+        requested = set(per_token) | {key for key, _ in exact}
+        seen: set[str] = set()
+        lines = []
+        for line in text.splitlines():
+            match = entry_pattern.match(line)
+            if match:
+                indent, key, values = match.group(1), match.group(2), match.group(3)
+                if key == "ImpactGPU":
+                    line = f"{indent}(ImpactGPU {device_index})"
+                else:
+                    exact_value = next((value for k, value in exact if k == key), None)
+                    if exact_value is not None:
+                        seen.add(key)
+                        line = f"{indent}({key} {exact_value})"
+                    else:
+                        token_key = "ImpactSubsetFeatures" if key.startswith("ImpactSubsetFeatures") else key
+                        if token_key in per_token:
+                            seen.add(token_key)
+                            replaced = " ".join(per_token[token_key] for _ in values.split())
+                            line = f"{indent}({key} {replaced})"
+            lines.append(line)
+        # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing —
+        # surface it (e.g. final_grid_spacing on a rigid-only preset).
+        for key in sorted(requested - seen):
+            print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.")
+        return "\n".join(lines)
+
+    def _stage_parameter_maps(self, work: Path, device_index: int) -> list[Path]:
+        """Stage the parameter maps into ``work``.
+
+        Matrix mode GENERATES each map from ``resolutions`` + the registry, then applies only the map-wide
+        knobs (the matrix already sets iterations/features per cell). Legacy mode copies the preset's maps and
+        applies every per-token / exact override. Both set the ImpactGPU device.
+        """
+        staged = []
+        for src in self._parameter_maps:
+            if self._resolutions:
+                text = generate_impact_parameter_map(
+                    src.read_text(encoding="utf-8"), self._resolutions, self._registry, self._mode
+                )
+                per_token, exact = self._parameter_map_overrides(global_only=True)
+            else:
+                text = src.read_text(encoding="utf-8")
+                per_token, exact = self._parameter_map_overrides()
+            text = self._apply_map_overrides(text, per_token, exact, device_index)
+            dst = work / src.name
+            dst.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")
+            staged.append(dst)
+        return staged
+
+    def register(
+        self,
+        fixed: sitk.Image,
+        moving: sitk.Image,
+        device_index: int,
+        fixed_mask: sitk.Image | None = None,
+        moving_mask: sitk.Image | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Register ``moving`` onto ``fixed``; return (moved, dvf) as channel-first arrays on the fixed grid.
+
+        Optional ``fixed_mask`` / ``moving_mask`` restrict the similarity metric to a region (elastix
+        ``-fMask`` / ``-mMask``); a mask covering the whole image is equivalent to passing none.
+        """
+        work = Path(tempfile.mkdtemp(prefix="konfai_reg_"))
+        try:
+            fixed_path, moving_path = work / "Fixed.mha", work / "Moving.mha"
+            sitk.WriteImage(fixed, str(fixed_path))
+            sitk.WriteImage(moving, str(moving_path))
+
+            # Stage the feature models at the relative path the maps reference (e.g. ImpactModelsPath0
+            # "MIND/R1D2_3D.pt"), resolved from the elastix working directory.
+            for rel_name, model_path in self._local_models:
+                dst = work / rel_name
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                if not dst.exists():
+                    dst.symlink_to(model_path)
+
+            args = [str(self._elastix_bin), "-f", str(fixed_path), "-m", str(moving_path)]
+            for flag, mask, name in (("-fMask", fixed_mask, "FixedMask.mha"), ("-mMask", moving_mask, "MovingMask.mha")):
+                if mask is not None:
+                    mask_path = work / name
+                    sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
+                    args += [flag, str(mask_path)]
+            args += ["-out", str(work)]
+            for pmap in self._stage_parameter_maps(work, device_index):
+                args += ["-p", str(pmap)]
+
+            # Make the elastix binary's bundled libs (libtorch under <install>/lib) and any extra
+            # libtorch/CUDA dirs (KONFAI_ELASTIX_EXTRA_LIB) findable so the IMPACT metric plugin loads.
+            env = os.environ.copy()
+            extra_libs = [str(self._elastix_bin.parent.parent / "lib"), os.environ.get("KONFAI_ELASTIX_EXTRA_LIB", "")]
+            env["LD_LIBRARY_PATH"] = os.pathsep.join(p for p in [*extra_libs, env.get("LD_LIBRARY_PATH", "")] if p)
+            proc = subprocess.Popen(  # nosec B603
+                args,
+                cwd=str(work),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                bufsize=1,
+                env=env,
+            )
+            # Drive a tqdm bar over elastix's iteration lines so SlicerKonfAI (which parses the "N% done"
+            # progress line) shows real progress. A tuned max_iterations makes the declared budget stale ->
+            # open-ended bar. The description mirrors KonfAI's bars: resolution level + the metric value.
+            captured: list[str] = []
+            iteration_line = re.compile(r"^\d+\s")
+            budget = None if self._max_iterations > 0 else (self._iterations or None)
+            progress = tqdm.tqdm(total=budget, desc="Registration", ncols=0, leave=True)
+            assert proc.stdout is not None
+            resolution = 0
+            for line in proc.stdout:
+                captured.append(line)
+                stripped = line.strip()
+                if stripped.startswith("Resolution:"):
+                    try:
+                        resolution = int(stripped.split(":", 1)[1])
+                    except ValueError:
+                        pass
+                elif iteration_line.match(line):
+                    progress.update(1)
+                    columns = line.split()  # column 2 is the metric (header "1:ItNr 2:Metric ...")
+                    if len(columns) > 1:
+                        try:
+                            progress.set_description(
+                                f"Registration : res {resolution} | metric {float(columns[1]):.4f}"
+                            )
+                        except ValueError:
+                            pass
+            progress.close()
+            returncode = proc.wait()
+            if returncode != 0:
+                raise RuntimeError(f"elastix failed (code {returncode}):\n{''.join(captured[-40:])}")
+
+            transforms = sorted(
+                work.glob("TransformParameters.*-Composite.itk.txt"),
+                key=lambda p: int(p.name.split(".")[1].split("-")[0]),
+            )
+            if not transforms:
+                raise FileNotFoundError("elastix produced no composite transform file.")
+            transform = sitk.ReadTransform(str(transforms[-1]))
+
+            moved = sitk.Resample(moving, fixed, transform, sitk.sitkLinear, 0.0, moving.GetPixelID())
+            dvf = sitk.TransformToDisplacementField(
+                transform,
+                sitk.sitkVectorFloat64,
+                fixed.GetSize(),
+                fixed.GetOrigin(),
+                fixed.GetSpacing(),
+                fixed.GetDirection(),
+            )
+            moved_np, _ = image_to_data(moved)
+            dvf_np, _ = image_to_data(dvf)
+            return moved_np, dvf_np
+        finally:
+            shutil.rmtree(work, ignore_errors=True)
+
+
+class ElastixRegistration(torch.nn.Module):
+    """Custom graph module: (fixed, moving) tensors + their geometry -> moved image on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving, from the KonfAI graph, the per-branch
+    ``Attribute`` list alongside the tensors (same convention as ``CriterionWithAttribute``). elastix needs
+    the physical geometry (Origin/Spacing/Direction), which raw tensors do not carry.
+    """
+
+    accepts_attributes = True
+
+    def __init__(
+        self,
+        engine: str,
+        parameter_maps: list[str],
+        max_iterations: int = 0,
+        final_grid_spacing: float = 0.0,
+        subset_features: int = 0,
+        spatial_samples: int = 0,
+        parameter_overrides: list[str] = [],
+        resolutions: dict = {},
+        mode: str = "Static",
+    ) -> None:
+        super().__init__()
+        if engine != "elastix":
+            raise NotImplementedError(f"ElastixRegistration engine '{engine}' is not implemented yet.")
+        self._engine = ElastixEngine(
+            parameter_maps,
+            max_iterations,
+            final_grid_spacing,
+            subset_features,
+            spatial_samples,
+            parameter_overrides,
+            resolutions,
+            mode,
+        )
+
+    def forward(
+        self,
+        fixed: torch.Tensor,
+        moving: torch.Tensor,
+        fixed_mask: torch.Tensor,
+        moving_mask: torch.Tensor,
+        attributes: list[list[Attribute]],
+    ) -> torch.Tensor:
+        # attributes = [fixed, moving, fixed_mask, moving_mask] branch attrs; each a list[Attribute] over the
+        # batch. Returns, per sample, the moved image (1 channel) stacked with the DVF (dim channels), both on
+        # the fixed grid; downstream ChannelSelect splits them. A whole-image mask (the default) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        for b in range(fixed.shape[0]):
+            fixed_img = data_to_image(fixed[b].detach().cpu().numpy(), fixed_attrs[b])
+            moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b])
+            fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b])
+            moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b])
+            moved_np, dvf_np = self._engine.register(
+                fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img
+            )
+            combined.append(torch.from_numpy(np.concatenate([moved_np, dvf_np], axis=0)))
+        return torch.stack(combined, dim=0).to(fixed.device)