Add FireANTs presets (Affine/SyN/IMPACT); optional-mask defaults; elastix skips non-restrictive masks

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

CBCT_CT_HeadNeck/app.json

@@ -1,7 +1,7 @@
 {
     "display_name": "CBCT/CT Head&Neck",
     "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.",
+    "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\u00d72\u00d73 mm, 2\u00d72\u00d73 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": 0,
     "mc_dropout": 0,
@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

CBCT_CT_HeadNeck/elastix_engine.py

@@ -44,6 +44,17 @@ from Model import _sorted_specs, generate_impact_parameter_map, load_models_regi
 ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
 
 
+def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
+    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
+    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
+    the whole image) instead of paying for a mask that restricts nothing."""
+    if mask is None:
+        return False
+    arr = sitk.GetArrayViewFromImage(mask)
+    return bool((arr > 0).any()) and bool((arr == 0).any())
+
+
 class ElastixEngine:
     """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
 
@@ -235,7 +246,7 @@ class ElastixEngine:
 
             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:
+                if _is_partial_mask(mask):
                     mask_path = work / name
                     sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
                     args += [flag, str(mask_path)]

CBCT_CT_MRSeg/app.json

@@ -1,7 +1,7 @@
 {
     "display_name": "CBCT/CT preset with MRSegmentator",
     "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.",
+    "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\u00d72\u00d73 mm, 2\u00d72\u00d73 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": 0,
     "mc_dropout": 0,
@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

CBCT_CT_MRSeg/elastix_engine.py

@@ -44,6 +44,17 @@ from Model import _sorted_specs, generate_impact_parameter_map, load_models_regi
 ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
 
 
+def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
+    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
+    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
+    the whole image) instead of paying for a mask that restricts nothing."""
+    if mask is None:
+        return False
+    arr = sitk.GetArrayViewFromImage(mask)
+    return bool((arr > 0).any()) and bool((arr == 0).any())
+
+
 class ElastixEngine:
     """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
 
@@ -235,7 +246,7 @@ class ElastixEngine:
 
             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:
+                if _is_partial_mask(mask):
                     mask_path = work / name
                     sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
                     args += [flag, str(mask_path)]

CBCT_CT_TS/app.json

@@ -1,7 +1,7 @@
 {
     "display_name": "CBCT/CT preset with TotalSegmentator",
     "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.",
+    "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\u00d72\u00d73 mm, 2\u00d72\u00d73 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": 0,
     "mc_dropout": 0,
@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

CBCT_CT_TS/elastix_engine.py

@@ -44,6 +44,17 @@ from Model import _sorted_specs, generate_impact_parameter_map, load_models_regi
 ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
 
 
+def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
+    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
+    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
+    the whole image) instead of paying for a mask that restricts nothing."""
+    if mask is None:
+        return False
+    arr = sitk.GetArrayViewFromImage(mask)
+    return bool((arr > 0).any()) and bool((arr == 0).any())
+
+
 class ElastixEngine:
     """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
 
@@ -235,7 +246,7 @@ class ElastixEngine:
 
             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:
+                if _is_partial_mask(mask):
                     mask_path = work / name
                     sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
                     args += [flag, str(mask_path)]

ConvexAdam_Coarse/app.json

@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

ConvexAdam_Composite/app.json

@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

ConvexAdam_Fine/app.json

@@ -6,10 +6,38 @@
     "tta": 0,
     "mc_dropout": 0,
     "vram_plan": {
-        "8": {"patch_size": [128, 128, 128], "batch_size": 1},
-        "16": {"patch_size": [192, 192, 192], "batch_size": 1},
-        "24": {"patch_size": [256, 256, 256], "batch_size": 1},
-        "40": {"patch_size": [320, 320, 320], "batch_size": 1}
+        "8": {
+            "patch_size": [
+                128,
+                128,
+                128
+            ],
+            "batch_size": 1
+        },
+        "16": {
+            "patch_size": [
+                192,
+                192,
+                192
+            ],
+            "batch_size": 1
+        },
+        "24": {
+            "patch_size": [
+                256,
+                256,
+                256
+            ],
+            "batch_size": 1
+        },
+        "40": {
+            "patch_size": [
+                320,
+                320,
+                320
+            ],
+            "batch_size": 1
+        }
     },
     "models": [
         "model.pt"
@@ -28,12 +56,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

FireANTs_Affine/Evaluation_with_fid.yml

@@ -0,0 +1,22 @@
+Evaluator:
+  metrics:
+    FixedFid:
+      targets_criterions:
+        MovingFid:
+          criterions_loader:
+            TRE: {}
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_Affine/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_Affine/Evaluation_with_seg.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    FixedSeg:
+      targets_criterions:
+        MovingSeg:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+      Reference_0:
+        groups_dest:
+          MovingSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_Affine/Model.py

@@ -0,0 +1,534 @@
+# 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
+#
+# This wrapper does NOT copy any FireANTs source: it only calls the public FireANTs API of the
+# separately-installed ``fireants`` wheel (PyPI). FireANTs is distributed under the FireANTs License
+# v1.0 and must be cited — see the NOTICE file in this directory for the license, copyright and
+# bibliography that ship with this app.
+
+"""FireANTs registration as a self-contained KonfAI model (shared by the FireANTs presets).
+
+Same idiomatic ``add_module`` graph and the same output contract as the ConvexAdam preset
+(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), so the
+orchestrator / app.json / ensemble / uncertainty are unchanged. The engine chains FireANTs' own
+composable stages (GPU, Riemannian Adam), each seeding the next like ANTs' ``-t`` stages:
+
+    Rigid (MI, centre-of-mass init) -> Affine (MI, seeded by the rigid) -> deformable
+
+The deformable stage is selected by ``deformable_method`` — the ONE knob that specialises this shared
+Model.py into the different presets (exactly as ConvexAdam's shared Model.py is specialised by
+``stages``):
+
+    "syn"    symmetric diffeomorphic SyN (CC)   — invertible, higher quality, averages cleanly for ensembling
+    "greedy" greedy diffeomorphic (CC)          — one-directional, faster / lower VRAM
+    "none"   linear only                        — Rigid+Affine, no deformable (the FireANTs_Affine preset)
+
+Masks: the optional Fixed/Moving masks restrict the metric to a region. FireANTs implements this by
+carrying the mask as the last image channel and prefixing the metric with ``masked_``; a mask is only
+honoured when it actually restricts (some voxels in, some out), so the common mask-free path is
+unchanged (an absent optional mask arrives as a whole-image default and is treated as no mask).
+
+The deformable stages produce the single TOTAL displacement field on the fixed grid (the linear
+pre-align is baked in via ``init_affine``, ANTs convention); ``none`` uses the affine matrix directly.
+``MovedImage`` and the emitted ``DisplacementField`` are rebuilt from that transform with SimpleITK —
+the same output path as the ConvexAdam engine — so all presets/engines are interchangeable in an
+ensemble. FireANTs' output-transform writer only serialises to a file, so the deformable field is
+round-tripped through a temporary NIfTI (no FireANTs internals are reimplemented here).
+
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
+"""
+
+import contextlib
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Literal
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+from konfai.metric.measure import IMPACTReg
+from konfai.network import network
+from konfai.utils.config import Choices, Range
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+DIM = 3
+
+# Feature-model registry (models.json): the available IMPACT feature models, fetched from HF (NOT bundled).
+# Only consulted by the "impact" deformable metric; ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins
+# for dev/offline. Mirrors the ConvexAdam preset so the same 30-model catalogue and picker are shared.
+_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
+
+_DISTANCES: dict[str, type[torch.nn.Module]] = {"L1": torch.nn.L1Loss, "L2": torch.nn.MSELoss}
+
+
+def registry_choices() -> list[str]:
+    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
+def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
+    """Load ``models.json`` (available feature models). ``KONFAI_IMPACT_MODELS_REGISTRY`` (local path) wins
+    for dev/offline; otherwise ``ref`` is a ``repo:file`` Hugging Face reference (fetched, not bundled)."""
+    from huggingface_hub import hf_hub_download
+
+    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
+    if local:
+        path = Path(local)
+    elif ":" in ref:
+        repo, filename = ref.split(":", 1)
+        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+    else:
+        raise ValueError(
+            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference — or set "
+            "KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
+        )
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _sorted_specs(mapping: dict) -> list:
+    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
+    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
+
+
+@dataclass
+class ModelSpec:
+    """One IMPACT feature model in the deformable metric (several are fused). ``ref`` picks the model; the
+    rest are its per-model knobs — the same as the ConvexAdam / elastix ``ModelSpec`` except ``voxel_size``
+    (an itk-impact resampling knob) has no meaning for FireANTs' geometry-free torch ``custom_loss`` and is
+    intentionally absent."""
+
+    ref: Annotated[str, Choices(registry_choices)]
+    layers_mask: str = "01"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
+    layers_weight: float = 1.0  # this model's weight in the multi-model fusion
+    pca: Annotated[int, Range(0, 100)] = 0  # keep the top-K principal components of the features (0 = no PCA)
+    distance: Literal["L1", "L2"] = "L1"
+
+
+@contextlib.contextmanager
+def _no_texpr_fuser():
+    """Disable the TensorExpr JIT fuser while IMPACT's TorchScript feature model runs under autograd.
+
+    The IMPACT feature models are TorchScript; run under FireANTs' gradient optimisation the TensorExpr
+    fuser trips on shape ops (``aten::size`` INTERNAL ASSERT). Scoped and restored so no other torch/JIT
+    user is affected; the modern profiling executor stays on (this is NOT the legacy executor).
+    """
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    try:
+        yield
+    finally:
+        torch._C._jit_set_texpr_fuser_enabled(True)
+
+
+class _ImpactCore(IMPACTReg):
+    """One IMPACT feature model, exposed as a FireANTs ``forward(moved, fixed)``.
+
+    Reuses ``IMPACTReg._compute`` / ``preprocessing`` verbatim — the stats-normalised feature extraction
+    (the model wants per-image ``[min, mean, max, std]``) and the per-layer weighted distance — so the
+    metric is exactly KonfAI's, not a re-derivation. Only KonfAI's config-binding ``__init__`` and its
+    ``Attribute``-based geometry are replaced: FireANTs passes raw tensors at the current pyramid scale, so
+    the intensity statistics are computed from those tensors directly. ``pca`` (absent from KonfAI's torch
+    ``IMPACTReg``) is added here as a per-layer feature-space reduction matching itk-impact.
+    """
+
+    def __init__(self, ref: str, in_channels: int, weights: list[float], distance: str, pca: int) -> None:
+        from huggingface_hub import hf_hub_download
+
+        torch.nn.Module.__init__(self)  # bypass IMPACTReg.__init__ (KONFAI_CONFIG_PATH / apply_config binding)
+        self.name = "Reg"
+        self.in_channels = int(in_channels)
+        self.weights = [float(w) for w in weights]
+        self.nb_layer = len(self.weights)
+        self.loss = _DISTANCES[distance]()
+        self.pca = int(pca)  # PCA lives in KonfAI's IMPACTReg._compute (same behaviour as itk-impact)
+        self.dim = DIM
+        self.shape = None  # score the whole (downsampled) tensor — no ModelPatch tiling
+        if ":" in ref:  # a "repo:path" HF reference; otherwise a local model file
+            repo, filename = ref.split(":", 1)
+            self.model_path = hf_hub_download(repo, filename, repo_type="model")  # nosec B615
+        else:
+            self.model_path = ref
+        self.model = None  # lazy-loaded on the first forward, like IMPACTReg
+
+    @staticmethod
+    def _stats(tensor: torch.Tensor) -> dict:
+        detached = tensor.detach()
+        return {
+            "ImageMin": float(detached.min()),
+            "ImageMean": float(detached.mean()),
+            "ImageMax": float(detached.max()),
+            "ImageStd": float(detached.std()),
+        }
+
+    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        if self.model is None:
+            self.model = torch.jit.load(self.model_path)  # nosec B614
+        self.model.to(moved.device).eval()
+        with _no_texpr_fuser():
+            loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], None)
+        return loss / max(true_nb, 1)
+
+
+class ImpactFeatureLoss(torch.nn.Module):
+    """FireANTs ``custom_loss`` = the KonfAI IMPACT metric fused over several feature models.
+
+    ``forward(moved, fixed)`` sums each model's ``layers_weight * IMPACT(model)``. A model's per-layer
+    weights come from its ``layers_mask`` bitmask; its input channel count is read from the registry
+    (``models.json`` ``numberofchannels``) so it never has to be configured by hand.
+    """
+
+    def __init__(self, specs: list["ModelSpec"]) -> None:
+        super().__init__()
+        registry = load_models_registry()
+        self._cores = torch.nn.ModuleList()
+        self._model_weights: list[float] = []
+        for spec in specs:
+            in_channels = int(registry.get(spec.ref.split(":", 1)[-1], {}).get("numberofchannels", 1))
+            weights = [1.0 if char == "1" else 0.0 for char in spec.layers_mask]
+            self._cores.append(_ImpactCore(spec.ref, in_channels, weights, spec.distance, spec.pca))
+            self._model_weights.append(float(spec.layers_weight))
+
+    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:
+        total: torch.Tensor | None = None
+        for weight, core in zip(self._model_weights, self._cores, strict=True):
+            term = weight * core(moved, fixed)
+            total = term if total is None else total + term
+        return total
+
+
+class FireANTsEngine:
+    """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return
+    (moved, dvf) on the fixed grid.
+
+    ``fireants`` is imported lazily inside :meth:`register` so this module can be imported for config
+    /signature introspection (SlicerImpactReg reads the tuning knobs off the ``RegistrationNet``
+    annotations) on a machine without a GPU or without FireANTs installed.
+    """
+
+    def __init__(
+        self,
+        scales: list[int],
+        affine_iterations: list[int],
+        deformable_iterations: list[int],
+        cc_kernel: int,
+        affine_metric: str,
+        affine_lr: float,
+        deformable_method: str,
+        deformable_metric: str,
+        deformable_lr: float,
+        integrator_n: int,
+        smooth_warp_sigma: float,
+        smooth_grad_sigma: float,
+        seed: int,
+        impact_specs: list["ModelSpec"],
+    ) -> None:
+        self._scales = [int(s) for s in scales]
+        self._affine_iterations = [int(i) for i in affine_iterations]
+        self._deformable_iterations = [int(i) for i in deformable_iterations]
+        self._cc_kernel = int(cc_kernel)
+        self._affine_metric = affine_metric
+        self._affine_lr = float(affine_lr)
+        self._deformable_method = deformable_method
+        self._deformable_metric = deformable_metric
+        self._deformable_lr = float(deformable_lr)
+        self._integrator_n = int(integrator_n)
+        self._smooth_warp_sigma = float(smooth_warp_sigma)
+        self._smooth_grad_sigma = float(smooth_grad_sigma)
+        self._seed = int(seed)
+        # IMPACT deformable metric (only used when deformable_metric == "impact"): KonfAI IMPACT feature
+        # models drive the SyN/greedy stage instead of the analytic CC/MI/MSE.
+        self._impact_specs = impact_specs
+
+    @staticmethod
+    def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+        """True only for a mask that actually restricts the region — some voxels in, some out. An absent
+        optional mask arrives as a whole-image (all-ones) default and an all-zero mask is degenerate; both
+        are treated as no mask so the plain (non-masked) metric path is used."""
+        if mask is None:
+            return False
+        arr = sitk.GetArrayViewFromImage(mask)
+        return bool((arr > 0).any()) and bool((arr == 0).any())
+
+    @staticmethod
+    def _affine_to_sitk(affine_matrix: "torch.Tensor") -> sitk.AffineTransform:
+        """FireANTs' physical (LPS) linear matrix -> SimpleITK AffineTransform (fixed -> moving points),
+        the same convention FireANTs writes into an ANTs ``0GenericAffine.mat``."""
+        matrix = affine_matrix.float().cpu().numpy()[0]
+        affine = sitk.AffineTransform(DIM)
+        affine.SetMatrix(matrix[:DIM, :DIM].flatten().astype(np.float64))
+        affine.SetTranslation(matrix[:DIM, DIM].astype(np.float64))
+        return affine
+
+    def _total_field_transform(self, reg) -> sitk.Transform:
+        """Optimise a deformable stage and return its TOTAL displacement (affine baked in) as a
+        SimpleITK ``DisplacementFieldTransform`` on the fixed grid.
+
+        FireANTs serialises the total field (ANTs convention, fixed grid) only to a file, so it is
+        round-tripped through a temporary NIfTI — its public API, no internals reimplemented."""
+        reg.optimize()
+        with tempfile.TemporaryDirectory() as tmp:
+            warp_path = os.path.join(tmp, "total_warp.nii.gz")
+            reg.save_as_ants_transforms(warp_path)
+            total_field = sitk.ReadImage(warp_path, sitk.sitkVectorFloat64)
+        return sitk.DisplacementFieldTransform(total_field)  # consumes total_field
+
+    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."""
+        from fireants.io import BatchedImages, Image
+        from fireants.io.imagemask import apply_mask_to_image, generate_image_mask_allones
+        from fireants.registration.affine import AffineRegistration
+        from fireants.registration.rigid import RigidRegistration
+
+        torch.manual_seed(self._seed)
+        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
+        # FireANTs' Image ctor accepts a SimpleITK image directly, so the fixed/moving cross into
+        # FireANTs in-memory (no file load) with their geometry preserved.
+        fixed_img = Image(fixed, device=device)
+        moving_img = Image(moving, device=device)
+
+        # Masked metric only when a mask genuinely restricts the region. FireANTs' masked mode wants the
+        # mask as the last channel of BOTH images (all-ones where one side has none) and a ``masked_``
+        # metric prefix; the plain path is untouched when no real mask is present.
+        use_fixed_mask = self._is_partial_mask(fixed_mask)
+        use_moving_mask = self._is_partial_mask(moving_mask)
+        masked = use_fixed_mask or use_moving_mask
+        if masked:
+            fmask = Image(fixed_mask, device=device) if use_fixed_mask else generate_image_mask_allones(fixed_img)
+            mmask = Image(moving_mask, device=device) if use_moving_mask else generate_image_mask_allones(moving_img)
+            fixed_img = apply_mask_to_image(fixed_img, fmask)
+            moving_img = apply_mask_to_image(moving_img, mmask)
+
+        bf = BatchedImages([fixed_img])
+        bm = BatchedImages([moving_img])
+        affine_loss = f"masked_{self._affine_metric}" if masked else self._affine_metric
+        deformable_loss = f"masked_{self._deformable_metric}" if masked else self._deformable_metric
+
+        # Linear: Rigid(MI, COM init) -> Affine(MI, seeded by the rigid), mirroring ANTs. The affine
+        # seeds the deformable stage (or is the whole transform when deformable_method == "none").
+        rigid = RigidRegistration(
+            scales=self._scales,
+            iterations=self._affine_iterations,
+            fixed_images=bf,
+            moving_images=bm,
+            loss_type=affine_loss,
+            optimizer="Adam",
+            optimizer_lr=self._affine_lr,
+            cc_kernel_size=self._cc_kernel,
+            init_translation="cof",
+        )
+        rigid.optimize()
+        rigid_matrix = rigid.get_rigid_matrix().detach()
+
+        affine = AffineRegistration(
+            scales=self._scales,
+            iterations=self._affine_iterations,
+            fixed_images=bf,
+            moving_images=bm,
+            loss_type=affine_loss,
+            optimizer="Adam",
+            optimizer_lr=self._affine_lr,
+            cc_kernel_size=self._cc_kernel,
+            init_rigid=rigid_matrix,
+        )
+        affine.optimize()
+        affine_matrix = affine.get_affine_matrix().detach()
+
+        # Deformable stage (or none). SyN and Greedy share the same constructor surface; both warm-start
+        # from the affine so their TOTAL transform already bakes in the linear pre-align.
+        if self._deformable_method == "none":
+            transform: sitk.Transform = self._affine_to_sitk(affine_matrix)
+        else:
+            if self._deformable_method == "syn":
+                from fireants.registration.syn import SyNRegistration as Deformable
+            elif self._deformable_method == "greedy":
+                from fireants.registration.greedy import GreedyRegistration as Deformable
+            else:
+                raise ValueError(
+                    f"Unknown deformable_method '{self._deformable_method}' (expected 'syn', 'greedy' or 'none')."
+                )
+            # "impact" swaps the analytic metric for a KonfAI IMPACT feature loss on the deformable stage
+            # (the linear pre-align keeps its own affine_metric); masks do not restrict the IMPACT metric.
+            if self._deformable_metric == "impact":
+                loss_type: str = "custom"
+                custom_loss: torch.nn.Module | None = ImpactFeatureLoss(self._impact_specs)
+            else:
+                loss_type, custom_loss = deformable_loss, None
+            reg = Deformable(
+                scales=self._scales,
+                iterations=self._deformable_iterations,
+                fixed_images=bf,
+                moving_images=bm,
+                loss_type=loss_type,
+                custom_loss=custom_loss,
+                cc_kernel_size=self._cc_kernel,
+                deformation_type="compositive",
+                integrator_n=self._integrator_n,
+                smooth_warp_sigma=self._smooth_warp_sigma,
+                smooth_grad_sigma=self._smooth_grad_sigma,
+                optimizer="Adam",
+                optimizer_lr=self._deformable_lr,
+                init_affine=affine_matrix,
+            )
+            transform = self._total_field_transform(reg)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        # Rebuild moved + DVF from the single transform on the fixed grid — the ConvexAdam output path,
+        # so every FireANTs preset emits identical-shaped results.
+        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
+
+
+class FireANTsRegistration(torch.nn.Module):
+    """Graph module: (fixed, moving) tensors + their geometry -> moved image + DVF on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving the per-branch ``Attribute`` list
+    alongside the tensors (same convention as the ConvexAdam / elastix engines); registration needs the
+    physical geometry, and the mask branches restrict the metric.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: FireANTsEngine) -> None:
+        super().__init__()
+        self._engine = engine
+
+    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) channel-stacked with the
+        # displacement field (DIM channels); downstream ChannelSelect modules split them. A whole-image
+        # mask (the default when none is supplied) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        # FireANTs runs a gradient-based instance optimisation (Riemannian Adam over the warp); the
+        # predictor calls forward under torch.inference_mode(), which forbids autograd. The image tensors
+        # have already crossed to numpy/SimpleITK here, so re-enable grad for the optimisation.
+        with torch.inference_mode(False), torch.enable_grad():
+            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)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise FireANTs 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 on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the DIM-component displacement field, in mm). Geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``. The knobs below are read straight from these annotations by the
+    UI: ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` a dropdown. ``deformable_method``
+    is the knob that specialises this shared model into each FireANTs preset.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        scales: list[int] = [4, 2, 1],
+        affine_iterations: list[int] = [200, 100, 50],
+        deformable_iterations: list[int] = [200, 100, 50],
+        cc_kernel: Annotated[int, Range(1, 21)] = 5,
+        affine_metric: Literal["mi", "cc", "mse"] = "mi",
+        affine_lr: Annotated[float, Range(0.0, 10.0)] = 0.003,
+        deformable_method: Literal["none", "syn", "greedy"] = "syn",
+        deformable_metric: Literal["cc", "mi", "mse", "impact"] = "cc",
+        deformable_lr: Annotated[float, Range(0.0, 10.0)] = 0.25,
+        integrator_n: Annotated[int, Range(1, 100)] = 10,
+        smooth_warp_sigma: Annotated[float, Range(0.0, 100.0)] = 0.5,
+        smooth_grad_sigma: Annotated[float, Range(0.0, 100.0)] = 1.0,
+        seed: int = 42,
+        models: dict[str, ModelSpec] = {},
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        engine = FireANTsEngine(
+            scales,
+            affine_iterations,
+            deformable_iterations,
+            cc_kernel,
+            affine_metric,
+            affine_lr,
+            deformable_method,
+            deformable_metric,
+            deformable_lr,
+            integrator_n,
+            smooth_warp_sigma,
+            smooth_grad_sigma,
+            seed,
+            _sorted_specs(models),
+        )
+        self.add_module(
+            "Registration", FireANTsRegistration(engine), in_branch=[0, 1, 2, 3], out_branch=["registration"]
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

FireANTs_Affine/NOTICE

@@ -0,0 +1,40 @@
+FireANTs_SyN — third-party attribution
+======================================
+
+This KonfAI app drives FireANTs, an external registration library. FireANTs is NOT
+redistributed here as source: the app depends on the official `fireants` wheel
+(https://pypi.org/project/fireants/), which is installed at resolve time (see
+requirements.txt) and called through its public Python API only. No FireANTs source
+code (functions, classes, or modules) is copied into this app.
+
+FireANTs is distributed under the **FireANTs License, Version 1.0 (July 2025)**, a
+custom license modified from Apache 2.0. Its redistribution clause requires that, when
+FireANTs is incorporated as a dependency in other projects, all license terms —
+including attribution and the bibliography below — be maintained.
+
+  Project : FireANTs
+  Source  : https://github.com/rohitrango/FireANTs
+  License : https://github.com/rohitrango/FireANTs/blob/main/LICENSE
+  Copyright (c) 2026 Rohit Jena. All rights reserved.
+
+Bibliography (as required by the FireANTs License — cite if you use this app)
+----------------------------------------------------------------------------
+
+  @article{jena2024fireants,
+    title={FireANTs: Adaptive Riemannian Optimization for Multi-Scale Diffeomorphic Registration},
+    author={Jena, Rohit and Chaudhari, Pratik and Gee, James C},
+    journal={Nature Communications},
+    year={2024}
+  }
+
+  @inproceedings{jena2025scalable,
+    title={A Scalable Distributed Framework for Multimodal {GigaVoxel} Image Registration},
+    author={Jena, Rohit and Zope, Vedant and Chaudhari, Pratik and Gee, James C},
+    booktitle={The Fourteenth International Conference on Learning Representations},
+    year={2026},
+    url={https://openreview.net/forum?id=8dLexnao2h}
+  }
+
+The wrapper code in this directory (Model.py, the KonfAI app configuration) is original
+work © 2025 Valentin Boussot, licensed under Apache-2.0, and is a separate work that
+merely links to the FireANTs interfaces.

FireANTs_Affine/Prediction.yml

@@ -0,0 +1,132 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      scales:
+      - 4
+      - 2
+      - 1
+      affine_iterations:
+      - 200
+      - 100
+      - 50
+      deformable_iterations:
+      - 200
+      - 100
+      - 50
+      cc_kernel: 5
+      affine_metric: mi
+      affine_lr: 0.003
+      deformable_method: none
+      deformable_metric: cc
+      deformable_lr: 0.25
+      integrator_n: 10
+      smooth_warp_sigma: 0.5
+      smooth_grad_sigma: 1.0
+      seed: 42
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations:
+      DataAugmentation_0:
+        data_augmentations:
+          Flip:
+            f_prob:
+            - 0
+            - 0.5
+            - 0.5
+            vector_field: true
+            prob: 1
+        nb: 2
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-FireANTs-Affine
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

FireANTs_Affine/Uncertainty.yml

@@ -0,0 +1,24 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean:
+              name: Uncertainty
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Norm: {}
+              StandardDeviation: {}
+              Save:
+                dataset: ./Uncertainties/ImpactReg/Output:mha
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_Affine/app.json

@@ -0,0 +1,96 @@
+{
+    "display_name": "FireANTs (Affine)",
+    "short_description": "Rigid + Affine linear alignment on GPU (FireANTs), no deformable stage.",
+    "description": "GPU linear registration with FireANTs: a Rigid (MI, centre-of-mass init) then Affine (MI) alignment optimised with Riemannian Adam, with no deformable stage. Produces the moved image and the affine displacement field on the fixed grid \u2014 use it on its own, or as a linear pre-align feeding a deformable preset. FireANTs is an external dependency under the FireANTs License v1.0 (see NOTICE).",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false,
+            "default": "ones"
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false,
+            "default": "ones"
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

FireANTs_Affine/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de99fbc36331ce674639acc774f52b4a2d0027f2f312d9d28669e831a0c4fd7e
+size 1249

FireANTs_Affine/requirements.txt

@@ -0,0 +1 @@
+fireants

FireANTs_IMPACT/Evaluation_with_fid.yml

@@ -0,0 +1,22 @@
+Evaluator:
+  metrics:
+    FixedFid:
+      targets_criterions:
+        MovingFid:
+          criterions_loader:
+            TRE: {}
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_IMPACT/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_IMPACT/Evaluation_with_seg.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    FixedSeg:
+      targets_criterions:
+        MovingSeg:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+      Reference_0:
+        groups_dest:
+          MovingSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_IMPACT/Model.py

@@ -0,0 +1,534 @@
+# 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
+#
+# This wrapper does NOT copy any FireANTs source: it only calls the public FireANTs API of the
+# separately-installed ``fireants`` wheel (PyPI). FireANTs is distributed under the FireANTs License
+# v1.0 and must be cited — see the NOTICE file in this directory for the license, copyright and
+# bibliography that ship with this app.
+
+"""FireANTs registration as a self-contained KonfAI model (shared by the FireANTs presets).
+
+Same idiomatic ``add_module`` graph and the same output contract as the ConvexAdam preset
+(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), so the
+orchestrator / app.json / ensemble / uncertainty are unchanged. The engine chains FireANTs' own
+composable stages (GPU, Riemannian Adam), each seeding the next like ANTs' ``-t`` stages:
+
+    Rigid (MI, centre-of-mass init) -> Affine (MI, seeded by the rigid) -> deformable
+
+The deformable stage is selected by ``deformable_method`` — the ONE knob that specialises this shared
+Model.py into the different presets (exactly as ConvexAdam's shared Model.py is specialised by
+``stages``):
+
+    "syn"    symmetric diffeomorphic SyN (CC)   — invertible, higher quality, averages cleanly for ensembling
+    "greedy" greedy diffeomorphic (CC)          — one-directional, faster / lower VRAM
+    "none"   linear only                        — Rigid+Affine, no deformable (the FireANTs_Affine preset)
+
+Masks: the optional Fixed/Moving masks restrict the metric to a region. FireANTs implements this by
+carrying the mask as the last image channel and prefixing the metric with ``masked_``; a mask is only
+honoured when it actually restricts (some voxels in, some out), so the common mask-free path is
+unchanged (an absent optional mask arrives as a whole-image default and is treated as no mask).
+
+The deformable stages produce the single TOTAL displacement field on the fixed grid (the linear
+pre-align is baked in via ``init_affine``, ANTs convention); ``none`` uses the affine matrix directly.
+``MovedImage`` and the emitted ``DisplacementField`` are rebuilt from that transform with SimpleITK —
+the same output path as the ConvexAdam engine — so all presets/engines are interchangeable in an
+ensemble. FireANTs' output-transform writer only serialises to a file, so the deformable field is
+round-tripped through a temporary NIfTI (no FireANTs internals are reimplemented here).
+
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
+"""
+
+import contextlib
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Literal
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+from konfai.metric.measure import IMPACTReg
+from konfai.network import network
+from konfai.utils.config import Choices, Range
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+DIM = 3
+
+# Feature-model registry (models.json): the available IMPACT feature models, fetched from HF (NOT bundled).
+# Only consulted by the "impact" deformable metric; ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins
+# for dev/offline. Mirrors the ConvexAdam preset so the same 30-model catalogue and picker are shared.
+_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
+
+_DISTANCES: dict[str, type[torch.nn.Module]] = {"L1": torch.nn.L1Loss, "L2": torch.nn.MSELoss}
+
+
+def registry_choices() -> list[str]:
+    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
+def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
+    """Load ``models.json`` (available feature models). ``KONFAI_IMPACT_MODELS_REGISTRY`` (local path) wins
+    for dev/offline; otherwise ``ref`` is a ``repo:file`` Hugging Face reference (fetched, not bundled)."""
+    from huggingface_hub import hf_hub_download
+
+    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
+    if local:
+        path = Path(local)
+    elif ":" in ref:
+        repo, filename = ref.split(":", 1)
+        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+    else:
+        raise ValueError(
+            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference — or set "
+            "KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
+        )
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _sorted_specs(mapping: dict) -> list:
+    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
+    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
+
+
+@dataclass
+class ModelSpec:
+    """One IMPACT feature model in the deformable metric (several are fused). ``ref`` picks the model; the
+    rest are its per-model knobs — the same as the ConvexAdam / elastix ``ModelSpec`` except ``voxel_size``
+    (an itk-impact resampling knob) has no meaning for FireANTs' geometry-free torch ``custom_loss`` and is
+    intentionally absent."""
+
+    ref: Annotated[str, Choices(registry_choices)]
+    layers_mask: str = "01"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
+    layers_weight: float = 1.0  # this model's weight in the multi-model fusion
+    pca: Annotated[int, Range(0, 100)] = 0  # keep the top-K principal components of the features (0 = no PCA)
+    distance: Literal["L1", "L2"] = "L1"
+
+
+@contextlib.contextmanager
+def _no_texpr_fuser():
+    """Disable the TensorExpr JIT fuser while IMPACT's TorchScript feature model runs under autograd.
+
+    The IMPACT feature models are TorchScript; run under FireANTs' gradient optimisation the TensorExpr
+    fuser trips on shape ops (``aten::size`` INTERNAL ASSERT). Scoped and restored so no other torch/JIT
+    user is affected; the modern profiling executor stays on (this is NOT the legacy executor).
+    """
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    try:
+        yield
+    finally:
+        torch._C._jit_set_texpr_fuser_enabled(True)
+
+
+class _ImpactCore(IMPACTReg):
+    """One IMPACT feature model, exposed as a FireANTs ``forward(moved, fixed)``.
+
+    Reuses ``IMPACTReg._compute`` / ``preprocessing`` verbatim — the stats-normalised feature extraction
+    (the model wants per-image ``[min, mean, max, std]``) and the per-layer weighted distance — so the
+    metric is exactly KonfAI's, not a re-derivation. Only KonfAI's config-binding ``__init__`` and its
+    ``Attribute``-based geometry are replaced: FireANTs passes raw tensors at the current pyramid scale, so
+    the intensity statistics are computed from those tensors directly. ``pca`` (absent from KonfAI's torch
+    ``IMPACTReg``) is added here as a per-layer feature-space reduction matching itk-impact.
+    """
+
+    def __init__(self, ref: str, in_channels: int, weights: list[float], distance: str, pca: int) -> None:
+        from huggingface_hub import hf_hub_download
+
+        torch.nn.Module.__init__(self)  # bypass IMPACTReg.__init__ (KONFAI_CONFIG_PATH / apply_config binding)
+        self.name = "Reg"
+        self.in_channels = int(in_channels)
+        self.weights = [float(w) for w in weights]
+        self.nb_layer = len(self.weights)
+        self.loss = _DISTANCES[distance]()
+        self.pca = int(pca)  # PCA lives in KonfAI's IMPACTReg._compute (same behaviour as itk-impact)
+        self.dim = DIM
+        self.shape = None  # score the whole (downsampled) tensor — no ModelPatch tiling
+        if ":" in ref:  # a "repo:path" HF reference; otherwise a local model file
+            repo, filename = ref.split(":", 1)
+            self.model_path = hf_hub_download(repo, filename, repo_type="model")  # nosec B615
+        else:
+            self.model_path = ref
+        self.model = None  # lazy-loaded on the first forward, like IMPACTReg
+
+    @staticmethod
+    def _stats(tensor: torch.Tensor) -> dict:
+        detached = tensor.detach()
+        return {
+            "ImageMin": float(detached.min()),
+            "ImageMean": float(detached.mean()),
+            "ImageMax": float(detached.max()),
+            "ImageStd": float(detached.std()),
+        }
+
+    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        if self.model is None:
+            self.model = torch.jit.load(self.model_path)  # nosec B614
+        self.model.to(moved.device).eval()
+        with _no_texpr_fuser():
+            loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], None)
+        return loss / max(true_nb, 1)
+
+
+class ImpactFeatureLoss(torch.nn.Module):
+    """FireANTs ``custom_loss`` = the KonfAI IMPACT metric fused over several feature models.
+
+    ``forward(moved, fixed)`` sums each model's ``layers_weight * IMPACT(model)``. A model's per-layer
+    weights come from its ``layers_mask`` bitmask; its input channel count is read from the registry
+    (``models.json`` ``numberofchannels``) so it never has to be configured by hand.
+    """
+
+    def __init__(self, specs: list["ModelSpec"]) -> None:
+        super().__init__()
+        registry = load_models_registry()
+        self._cores = torch.nn.ModuleList()
+        self._model_weights: list[float] = []
+        for spec in specs:
+            in_channels = int(registry.get(spec.ref.split(":", 1)[-1], {}).get("numberofchannels", 1))
+            weights = [1.0 if char == "1" else 0.0 for char in spec.layers_mask]
+            self._cores.append(_ImpactCore(spec.ref, in_channels, weights, spec.distance, spec.pca))
+            self._model_weights.append(float(spec.layers_weight))
+
+    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:
+        total: torch.Tensor | None = None
+        for weight, core in zip(self._model_weights, self._cores, strict=True):
+            term = weight * core(moved, fixed)
+            total = term if total is None else total + term
+        return total
+
+
+class FireANTsEngine:
+    """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return
+    (moved, dvf) on the fixed grid.
+
+    ``fireants`` is imported lazily inside :meth:`register` so this module can be imported for config
+    /signature introspection (SlicerImpactReg reads the tuning knobs off the ``RegistrationNet``
+    annotations) on a machine without a GPU or without FireANTs installed.
+    """
+
+    def __init__(
+        self,
+        scales: list[int],
+        affine_iterations: list[int],
+        deformable_iterations: list[int],
+        cc_kernel: int,
+        affine_metric: str,
+        affine_lr: float,
+        deformable_method: str,
+        deformable_metric: str,
+        deformable_lr: float,
+        integrator_n: int,
+        smooth_warp_sigma: float,
+        smooth_grad_sigma: float,
+        seed: int,
+        impact_specs: list["ModelSpec"],
+    ) -> None:
+        self._scales = [int(s) for s in scales]
+        self._affine_iterations = [int(i) for i in affine_iterations]
+        self._deformable_iterations = [int(i) for i in deformable_iterations]
+        self._cc_kernel = int(cc_kernel)
+        self._affine_metric = affine_metric
+        self._affine_lr = float(affine_lr)
+        self._deformable_method = deformable_method
+        self._deformable_metric = deformable_metric
+        self._deformable_lr = float(deformable_lr)
+        self._integrator_n = int(integrator_n)
+        self._smooth_warp_sigma = float(smooth_warp_sigma)
+        self._smooth_grad_sigma = float(smooth_grad_sigma)
+        self._seed = int(seed)
+        # IMPACT deformable metric (only used when deformable_metric == "impact"): KonfAI IMPACT feature
+        # models drive the SyN/greedy stage instead of the analytic CC/MI/MSE.
+        self._impact_specs = impact_specs
+
+    @staticmethod
+    def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+        """True only for a mask that actually restricts the region — some voxels in, some out. An absent
+        optional mask arrives as a whole-image (all-ones) default and an all-zero mask is degenerate; both
+        are treated as no mask so the plain (non-masked) metric path is used."""
+        if mask is None:
+            return False
+        arr = sitk.GetArrayViewFromImage(mask)
+        return bool((arr > 0).any()) and bool((arr == 0).any())
+
+    @staticmethod
+    def _affine_to_sitk(affine_matrix: "torch.Tensor") -> sitk.AffineTransform:
+        """FireANTs' physical (LPS) linear matrix -> SimpleITK AffineTransform (fixed -> moving points),
+        the same convention FireANTs writes into an ANTs ``0GenericAffine.mat``."""
+        matrix = affine_matrix.float().cpu().numpy()[0]
+        affine = sitk.AffineTransform(DIM)
+        affine.SetMatrix(matrix[:DIM, :DIM].flatten().astype(np.float64))
+        affine.SetTranslation(matrix[:DIM, DIM].astype(np.float64))
+        return affine
+
+    def _total_field_transform(self, reg) -> sitk.Transform:
+        """Optimise a deformable stage and return its TOTAL displacement (affine baked in) as a
+        SimpleITK ``DisplacementFieldTransform`` on the fixed grid.
+
+        FireANTs serialises the total field (ANTs convention, fixed grid) only to a file, so it is
+        round-tripped through a temporary NIfTI — its public API, no internals reimplemented."""
+        reg.optimize()
+        with tempfile.TemporaryDirectory() as tmp:
+            warp_path = os.path.join(tmp, "total_warp.nii.gz")
+            reg.save_as_ants_transforms(warp_path)
+            total_field = sitk.ReadImage(warp_path, sitk.sitkVectorFloat64)
+        return sitk.DisplacementFieldTransform(total_field)  # consumes total_field
+
+    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."""
+        from fireants.io import BatchedImages, Image
+        from fireants.io.imagemask import apply_mask_to_image, generate_image_mask_allones
+        from fireants.registration.affine import AffineRegistration
+        from fireants.registration.rigid import RigidRegistration
+
+        torch.manual_seed(self._seed)
+        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
+        # FireANTs' Image ctor accepts a SimpleITK image directly, so the fixed/moving cross into
+        # FireANTs in-memory (no file load) with their geometry preserved.
+        fixed_img = Image(fixed, device=device)
+        moving_img = Image(moving, device=device)
+
+        # Masked metric only when a mask genuinely restricts the region. FireANTs' masked mode wants the
+        # mask as the last channel of BOTH images (all-ones where one side has none) and a ``masked_``
+        # metric prefix; the plain path is untouched when no real mask is present.
+        use_fixed_mask = self._is_partial_mask(fixed_mask)
+        use_moving_mask = self._is_partial_mask(moving_mask)
+        masked = use_fixed_mask or use_moving_mask
+        if masked:
+            fmask = Image(fixed_mask, device=device) if use_fixed_mask else generate_image_mask_allones(fixed_img)
+            mmask = Image(moving_mask, device=device) if use_moving_mask else generate_image_mask_allones(moving_img)
+            fixed_img = apply_mask_to_image(fixed_img, fmask)
+            moving_img = apply_mask_to_image(moving_img, mmask)
+
+        bf = BatchedImages([fixed_img])
+        bm = BatchedImages([moving_img])
+        affine_loss = f"masked_{self._affine_metric}" if masked else self._affine_metric
+        deformable_loss = f"masked_{self._deformable_metric}" if masked else self._deformable_metric
+
+        # Linear: Rigid(MI, COM init) -> Affine(MI, seeded by the rigid), mirroring ANTs. The affine
+        # seeds the deformable stage (or is the whole transform when deformable_method == "none").
+        rigid = RigidRegistration(
+            scales=self._scales,
+            iterations=self._affine_iterations,
+            fixed_images=bf,
+            moving_images=bm,
+            loss_type=affine_loss,
+            optimizer="Adam",
+            optimizer_lr=self._affine_lr,
+            cc_kernel_size=self._cc_kernel,
+            init_translation="cof",
+        )
+        rigid.optimize()
+        rigid_matrix = rigid.get_rigid_matrix().detach()
+
+        affine = AffineRegistration(
+            scales=self._scales,
+            iterations=self._affine_iterations,
+            fixed_images=bf,
+            moving_images=bm,
+            loss_type=affine_loss,
+            optimizer="Adam",
+            optimizer_lr=self._affine_lr,
+            cc_kernel_size=self._cc_kernel,
+            init_rigid=rigid_matrix,
+        )
+        affine.optimize()
+        affine_matrix = affine.get_affine_matrix().detach()
+
+        # Deformable stage (or none). SyN and Greedy share the same constructor surface; both warm-start
+        # from the affine so their TOTAL transform already bakes in the linear pre-align.
+        if self._deformable_method == "none":
+            transform: sitk.Transform = self._affine_to_sitk(affine_matrix)
+        else:
+            if self._deformable_method == "syn":
+                from fireants.registration.syn import SyNRegistration as Deformable
+            elif self._deformable_method == "greedy":
+                from fireants.registration.greedy import GreedyRegistration as Deformable
+            else:
+                raise ValueError(
+                    f"Unknown deformable_method '{self._deformable_method}' (expected 'syn', 'greedy' or 'none')."
+                )
+            # "impact" swaps the analytic metric for a KonfAI IMPACT feature loss on the deformable stage
+            # (the linear pre-align keeps its own affine_metric); masks do not restrict the IMPACT metric.
+            if self._deformable_metric == "impact":
+                loss_type: str = "custom"
+                custom_loss: torch.nn.Module | None = ImpactFeatureLoss(self._impact_specs)
+            else:
+                loss_type, custom_loss = deformable_loss, None
+            reg = Deformable(
+                scales=self._scales,
+                iterations=self._deformable_iterations,
+                fixed_images=bf,
+                moving_images=bm,
+                loss_type=loss_type,
+                custom_loss=custom_loss,
+                cc_kernel_size=self._cc_kernel,
+                deformation_type="compositive",
+                integrator_n=self._integrator_n,
+                smooth_warp_sigma=self._smooth_warp_sigma,
+                smooth_grad_sigma=self._smooth_grad_sigma,
+                optimizer="Adam",
+                optimizer_lr=self._deformable_lr,
+                init_affine=affine_matrix,
+            )
+            transform = self._total_field_transform(reg)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        # Rebuild moved + DVF from the single transform on the fixed grid — the ConvexAdam output path,
+        # so every FireANTs preset emits identical-shaped results.
+        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
+
+
+class FireANTsRegistration(torch.nn.Module):
+    """Graph module: (fixed, moving) tensors + their geometry -> moved image + DVF on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving the per-branch ``Attribute`` list
+    alongside the tensors (same convention as the ConvexAdam / elastix engines); registration needs the
+    physical geometry, and the mask branches restrict the metric.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: FireANTsEngine) -> None:
+        super().__init__()
+        self._engine = engine
+
+    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) channel-stacked with the
+        # displacement field (DIM channels); downstream ChannelSelect modules split them. A whole-image
+        # mask (the default when none is supplied) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        # FireANTs runs a gradient-based instance optimisation (Riemannian Adam over the warp); the
+        # predictor calls forward under torch.inference_mode(), which forbids autograd. The image tensors
+        # have already crossed to numpy/SimpleITK here, so re-enable grad for the optimisation.
+        with torch.inference_mode(False), torch.enable_grad():
+            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)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise FireANTs 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 on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the DIM-component displacement field, in mm). Geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``. The knobs below are read straight from these annotations by the
+    UI: ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` a dropdown. ``deformable_method``
+    is the knob that specialises this shared model into each FireANTs preset.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        scales: list[int] = [4, 2, 1],
+        affine_iterations: list[int] = [200, 100, 50],
+        deformable_iterations: list[int] = [200, 100, 50],
+        cc_kernel: Annotated[int, Range(1, 21)] = 5,
+        affine_metric: Literal["mi", "cc", "mse"] = "mi",
+        affine_lr: Annotated[float, Range(0.0, 10.0)] = 0.003,
+        deformable_method: Literal["none", "syn", "greedy"] = "syn",
+        deformable_metric: Literal["cc", "mi", "mse", "impact"] = "cc",
+        deformable_lr: Annotated[float, Range(0.0, 10.0)] = 0.25,
+        integrator_n: Annotated[int, Range(1, 100)] = 10,
+        smooth_warp_sigma: Annotated[float, Range(0.0, 100.0)] = 0.5,
+        smooth_grad_sigma: Annotated[float, Range(0.0, 100.0)] = 1.0,
+        seed: int = 42,
+        models: dict[str, ModelSpec] = {},
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        engine = FireANTsEngine(
+            scales,
+            affine_iterations,
+            deformable_iterations,
+            cc_kernel,
+            affine_metric,
+            affine_lr,
+            deformable_method,
+            deformable_metric,
+            deformable_lr,
+            integrator_n,
+            smooth_warp_sigma,
+            smooth_grad_sigma,
+            seed,
+            _sorted_specs(models),
+        )
+        self.add_module(
+            "Registration", FireANTsRegistration(engine), in_branch=[0, 1, 2, 3], out_branch=["registration"]
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

FireANTs_IMPACT/NOTICE

@@ -0,0 +1,40 @@
+FireANTs_SyN — third-party attribution
+======================================
+
+This KonfAI app drives FireANTs, an external registration library. FireANTs is NOT
+redistributed here as source: the app depends on the official `fireants` wheel
+(https://pypi.org/project/fireants/), which is installed at resolve time (see
+requirements.txt) and called through its public Python API only. No FireANTs source
+code (functions, classes, or modules) is copied into this app.
+
+FireANTs is distributed under the **FireANTs License, Version 1.0 (July 2025)**, a
+custom license modified from Apache 2.0. Its redistribution clause requires that, when
+FireANTs is incorporated as a dependency in other projects, all license terms —
+including attribution and the bibliography below — be maintained.
+
+  Project : FireANTs
+  Source  : https://github.com/rohitrango/FireANTs
+  License : https://github.com/rohitrango/FireANTs/blob/main/LICENSE
+  Copyright (c) 2026 Rohit Jena. All rights reserved.
+
+Bibliography (as required by the FireANTs License — cite if you use this app)
+----------------------------------------------------------------------------
+
+  @article{jena2024fireants,
+    title={FireANTs: Adaptive Riemannian Optimization for Multi-Scale Diffeomorphic Registration},
+    author={Jena, Rohit and Chaudhari, Pratik and Gee, James C},
+    journal={Nature Communications},
+    year={2024}
+  }
+
+  @inproceedings{jena2025scalable,
+    title={A Scalable Distributed Framework for Multimodal {GigaVoxel} Image Registration},
+    author={Jena, Rohit and Zope, Vedant and Chaudhari, Pratik and Gee, James C},
+    booktitle={The Fourteenth International Conference on Learning Representations},
+    year={2026},
+    url={https://openreview.net/forum?id=8dLexnao2h}
+  }
+
+The wrapper code in this directory (Model.py, the KonfAI app configuration) is original
+work © 2025 Valentin Boussot, licensed under Apache-2.0, and is a separate work that
+merely links to the FireANTs interfaces.

FireANTs_IMPACT/Prediction.yml

@@ -0,0 +1,139 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      scales:
+      - 4
+      - 2
+      - 1
+      affine_iterations:
+      - 200
+      - 100
+      - 50
+      deformable_iterations:
+      - 200
+      - 100
+      - 50
+      cc_kernel: 5
+      affine_metric: mi
+      affine_lr: 0.003
+      deformable_method: syn
+      deformable_metric: impact
+      deformable_lr: 0.25
+      integrator_n: 10
+      smooth_warp_sigma: 0.5
+      smooth_grad_sigma: 1.0
+      seed: 42
+      models:
+        '0':
+          ref: VBoussot/impact-torchscript-models:TS/M291.pt
+          layers_mask: '01'
+          layers_weight: 1.0
+          pca: 0
+          distance: L1
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations:
+      DataAugmentation_0:
+        data_augmentations:
+          Flip:
+            f_prob:
+            - 0
+            - 0.5
+            - 0.5
+            vector_field: true
+            prob: 1
+        nb: 2
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-FireANTs-IMPACT
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

FireANTs_IMPACT/Uncertainty.yml

@@ -0,0 +1,24 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean:
+              name: Uncertainty
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Norm: {}
+              StandardDeviation: {}
+              Save:
+                dataset: ./Uncertainties/ImpactReg/Output:mha
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_IMPACT/app.json

@@ -0,0 +1,96 @@
+{
+    "display_name": "FireANTs (IMPACT)",
+    "short_description": "Rigid + Affine + SyN driven by the IMPACT deep-feature metric on GPU (FireANTs).",
+    "description": "GPU deformable registration with FireANTs where the SyN stage is driven by the KonfAI IMPACT metric — deep features from one or more pretrained TorchScript models — instead of an intensity metric. A Rigid (MI) + Affine (MI) linear pre-align seeds a symmetric diffeomorphic SyN optimised with Riemannian Adam against a weighted multi-model feature distance. Produces the moved image and the total displacement field on the fixed grid. FireANTs is an external dependency under the FireANTs License v1.0 (see NOTICE); the feature models are fetched from VBoussot/impact-torchscript-models.",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false,
+            "default": "ones"
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false,
+            "default": "ones"
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

FireANTs_IMPACT/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de99fbc36331ce674639acc774f52b4a2d0027f2f312d9d28669e831a0c4fd7e
+size 1249

FireANTs_IMPACT/requirements.txt

@@ -0,0 +1 @@
+fireants

FireANTs_SyN/Evaluation_with_fid.yml

@@ -0,0 +1,22 @@
+Evaluator:
+  metrics:
+    FixedFid:
+      targets_criterions:
+        MovingFid:
+          criterions_loader:
+            TRE: {}
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedFid:
+            transforms: None
+      Reference_0:
+        groups_dest:
+          MovingFid:
+            transforms: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_SyN/Evaluation_with_images.yml

@@ -0,0 +1,35 @@
+Evaluator:
+  metrics:
+    FixedImage:
+      targets_criterions:
+        MovingImage;Mask:
+          criterions_loader:
+            MAESaveMap:
+              reduction: mean
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Reference_0:
+        groups_dest:
+          MovingImage:
+            transforms: 
+              TensorCast:
+                dtype: float32
+      Mask_0:
+        groups_dest:
+          Mask:
+            transforms:
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_SyN/Evaluation_with_seg.yml

@@ -0,0 +1,29 @@
+Evaluator:
+  metrics:
+    FixedSeg:
+      targets_criterions:
+        MovingSeg:
+          criterions_loader:
+            DiceSaveMap:
+              labels: None
+              dataset: ./Evaluations/ImpactReg/Output:mha
+              group: Seg_MAE_map
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          FixedSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+      Reference_0:
+        groups_dest:
+          MovingSeg:
+            transforms: 
+              TensorCast:
+                dtype: uint8
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_SyN/Model.py

@@ -0,0 +1,534 @@
+# 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
+#
+# This wrapper does NOT copy any FireANTs source: it only calls the public FireANTs API of the
+# separately-installed ``fireants`` wheel (PyPI). FireANTs is distributed under the FireANTs License
+# v1.0 and must be cited — see the NOTICE file in this directory for the license, copyright and
+# bibliography that ship with this app.
+
+"""FireANTs registration as a self-contained KonfAI model (shared by the FireANTs presets).
+
+Same idiomatic ``add_module`` graph and the same output contract as the ConvexAdam preset
+(``MovedImage`` + ``DisplacementField`` on the FIXED grid, split by two ``ChannelSelect``), so the
+orchestrator / app.json / ensemble / uncertainty are unchanged. The engine chains FireANTs' own
+composable stages (GPU, Riemannian Adam), each seeding the next like ANTs' ``-t`` stages:
+
+    Rigid (MI, centre-of-mass init) -> Affine (MI, seeded by the rigid) -> deformable
+
+The deformable stage is selected by ``deformable_method`` — the ONE knob that specialises this shared
+Model.py into the different presets (exactly as ConvexAdam's shared Model.py is specialised by
+``stages``):
+
+    "syn"    symmetric diffeomorphic SyN (CC)   — invertible, higher quality, averages cleanly for ensembling
+    "greedy" greedy diffeomorphic (CC)          — one-directional, faster / lower VRAM
+    "none"   linear only                        — Rigid+Affine, no deformable (the FireANTs_Affine preset)
+
+Masks: the optional Fixed/Moving masks restrict the metric to a region. FireANTs implements this by
+carrying the mask as the last image channel and prefixing the metric with ``masked_``; a mask is only
+honoured when it actually restricts (some voxels in, some out), so the common mask-free path is
+unchanged (an absent optional mask arrives as a whole-image default and is treated as no mask).
+
+The deformable stages produce the single TOTAL displacement field on the fixed grid (the linear
+pre-align is baked in via ``init_affine``, ANTs convention); ``none`` uses the affine matrix directly.
+``MovedImage`` and the emitted ``DisplacementField`` are rebuilt from that transform with SimpleITK —
+the same output path as the ConvexAdam engine — so all presets/engines are interchangeable in an
+ensemble. FireANTs' output-transform writer only serialises to a file, so the deformable field is
+round-tripped through a temporary NIfTI (no FireANTs internals are reimplemented here).
+
+NOTE: do NOT add ``from __future__ import annotations`` — KonfAI's config engine relies on
+runtime-evaluated annotations (``get_origin``); PEP 563 stringized annotations break binding.
+"""
+
+import contextlib
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Literal
+
+import numpy as np
+import SimpleITK as sitk
+import torch
+from konfai.metric.measure import IMPACTReg
+from konfai.network import network
+from konfai.utils.config import Choices, Range
+from konfai.utils.dataset import Attribute, data_to_image, image_to_data
+
+DIM = 3
+
+# Feature-model registry (models.json): the available IMPACT feature models, fetched from HF (NOT bundled).
+# Only consulted by the "impact" deformable metric; ``KONFAI_IMPACT_MODELS_REGISTRY`` (a local path) wins
+# for dev/offline. Mirrors the ConvexAdam preset so the same 30-model catalogue and picker are shared.
+_IMPACT_MODELS_REGISTRY = "VBoussot/impact-torchscript-models:models.json"
+
+_DISTANCES: dict[str, type[torch.nn.Module]] = {"L1": torch.nn.L1Loss, "L2": torch.nn.MSELoss}
+
+
+def registry_choices() -> list[str]:
+    """The per-model ``ref`` picker's values — model refs (``repo:path``) from the feature-model registry."""
+    repo = _IMPACT_MODELS_REGISTRY.split(":", 1)[0]
+    return [f"{repo}:{key}" for key in load_models_registry()]
+
+
+def load_models_registry(ref: str = _IMPACT_MODELS_REGISTRY) -> dict:
+    """Load ``models.json`` (available feature models). ``KONFAI_IMPACT_MODELS_REGISTRY`` (local path) wins
+    for dev/offline; otherwise ``ref`` is a ``repo:file`` Hugging Face reference (fetched, not bundled)."""
+    from huggingface_hub import hf_hub_download
+
+    local = os.environ.get("KONFAI_IMPACT_MODELS_REGISTRY", "")
+    if local:
+        path = Path(local)
+    elif ":" in ref:
+        repo, filename = ref.split(":", 1)
+        path = Path(hf_hub_download(repo_id=repo, filename=filename, repo_type="model"))  # nosec B615
+    else:
+        raise ValueError(
+            f"models_registry '{ref}' must be a 'repo:file' Hugging Face reference — or set "
+            "KONFAI_IMPACT_MODELS_REGISTRY to a local file for offline use."
+        )
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _sorted_specs(mapping: dict) -> list:
+    """A dict keyed by string indices ('0','1',...) -> its values in numeric order."""
+    return [mapping[k] for k in sorted(mapping, key=lambda key: int(key))]
+
+
+@dataclass
+class ModelSpec:
+    """One IMPACT feature model in the deformable metric (several are fused). ``ref`` picks the model; the
+    rest are its per-model knobs — the same as the ConvexAdam / elastix ``ModelSpec`` except ``voxel_size``
+    (an itk-impact resampling knob) has no meaning for FireANTs' geometry-free torch ``custom_loss`` and is
+    intentionally absent."""
+
+    ref: Annotated[str, Choices(registry_choices)]
+    layers_mask: str = "01"  # per-layer bitmask, one char per model layer ('1' = use, '0' = skip), like elastix
+    layers_weight: float = 1.0  # this model's weight in the multi-model fusion
+    pca: Annotated[int, Range(0, 100)] = 0  # keep the top-K principal components of the features (0 = no PCA)
+    distance: Literal["L1", "L2"] = "L1"
+
+
+@contextlib.contextmanager
+def _no_texpr_fuser():
+    """Disable the TensorExpr JIT fuser while IMPACT's TorchScript feature model runs under autograd.
+
+    The IMPACT feature models are TorchScript; run under FireANTs' gradient optimisation the TensorExpr
+    fuser trips on shape ops (``aten::size`` INTERNAL ASSERT). Scoped and restored so no other torch/JIT
+    user is affected; the modern profiling executor stays on (this is NOT the legacy executor).
+    """
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    try:
+        yield
+    finally:
+        torch._C._jit_set_texpr_fuser_enabled(True)
+
+
+class _ImpactCore(IMPACTReg):
+    """One IMPACT feature model, exposed as a FireANTs ``forward(moved, fixed)``.
+
+    Reuses ``IMPACTReg._compute`` / ``preprocessing`` verbatim — the stats-normalised feature extraction
+    (the model wants per-image ``[min, mean, max, std]``) and the per-layer weighted distance — so the
+    metric is exactly KonfAI's, not a re-derivation. Only KonfAI's config-binding ``__init__`` and its
+    ``Attribute``-based geometry are replaced: FireANTs passes raw tensors at the current pyramid scale, so
+    the intensity statistics are computed from those tensors directly. ``pca`` (absent from KonfAI's torch
+    ``IMPACTReg``) is added here as a per-layer feature-space reduction matching itk-impact.
+    """
+
+    def __init__(self, ref: str, in_channels: int, weights: list[float], distance: str, pca: int) -> None:
+        from huggingface_hub import hf_hub_download
+
+        torch.nn.Module.__init__(self)  # bypass IMPACTReg.__init__ (KONFAI_CONFIG_PATH / apply_config binding)
+        self.name = "Reg"
+        self.in_channels = int(in_channels)
+        self.weights = [float(w) for w in weights]
+        self.nb_layer = len(self.weights)
+        self.loss = _DISTANCES[distance]()
+        self.pca = int(pca)  # PCA lives in KonfAI's IMPACTReg._compute (same behaviour as itk-impact)
+        self.dim = DIM
+        self.shape = None  # score the whole (downsampled) tensor — no ModelPatch tiling
+        if ":" in ref:  # a "repo:path" HF reference; otherwise a local model file
+            repo, filename = ref.split(":", 1)
+            self.model_path = hf_hub_download(repo, filename, repo_type="model")  # nosec B615
+        else:
+            self.model_path = ref
+        self.model = None  # lazy-loaded on the first forward, like IMPACTReg
+
+    @staticmethod
+    def _stats(tensor: torch.Tensor) -> dict:
+        detached = tensor.detach()
+        return {
+            "ImageMin": float(detached.min()),
+            "ImageMean": float(detached.mean()),
+            "ImageMax": float(detached.max()),
+            "ImageStd": float(detached.std()),
+        }
+
+    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        if self.model is None:
+            self.model = torch.jit.load(self.model_path)  # nosec B614
+        self.model.to(moved.device).eval()
+        with _no_texpr_fuser():
+            loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], None)
+        return loss / max(true_nb, 1)
+
+
+class ImpactFeatureLoss(torch.nn.Module):
+    """FireANTs ``custom_loss`` = the KonfAI IMPACT metric fused over several feature models.
+
+    ``forward(moved, fixed)`` sums each model's ``layers_weight * IMPACT(model)``. A model's per-layer
+    weights come from its ``layers_mask`` bitmask; its input channel count is read from the registry
+    (``models.json`` ``numberofchannels``) so it never has to be configured by hand.
+    """
+
+    def __init__(self, specs: list["ModelSpec"]) -> None:
+        super().__init__()
+        registry = load_models_registry()
+        self._cores = torch.nn.ModuleList()
+        self._model_weights: list[float] = []
+        for spec in specs:
+            in_channels = int(registry.get(spec.ref.split(":", 1)[-1], {}).get("numberofchannels", 1))
+            weights = [1.0 if char == "1" else 0.0 for char in spec.layers_mask]
+            self._cores.append(_ImpactCore(spec.ref, in_channels, weights, spec.distance, spec.pca))
+            self._model_weights.append(float(spec.layers_weight))
+
+    def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor:
+        total: torch.Tensor | None = None
+        for weight, core in zip(self._model_weights, self._cores, strict=True):
+            term = weight * core(moved, fixed)
+            total = term if total is None else total + term
+        return total
+
+
+class FireANTsEngine:
+    """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return
+    (moved, dvf) on the fixed grid.
+
+    ``fireants`` is imported lazily inside :meth:`register` so this module can be imported for config
+    /signature introspection (SlicerImpactReg reads the tuning knobs off the ``RegistrationNet``
+    annotations) on a machine without a GPU or without FireANTs installed.
+    """
+
+    def __init__(
+        self,
+        scales: list[int],
+        affine_iterations: list[int],
+        deformable_iterations: list[int],
+        cc_kernel: int,
+        affine_metric: str,
+        affine_lr: float,
+        deformable_method: str,
+        deformable_metric: str,
+        deformable_lr: float,
+        integrator_n: int,
+        smooth_warp_sigma: float,
+        smooth_grad_sigma: float,
+        seed: int,
+        impact_specs: list["ModelSpec"],
+    ) -> None:
+        self._scales = [int(s) for s in scales]
+        self._affine_iterations = [int(i) for i in affine_iterations]
+        self._deformable_iterations = [int(i) for i in deformable_iterations]
+        self._cc_kernel = int(cc_kernel)
+        self._affine_metric = affine_metric
+        self._affine_lr = float(affine_lr)
+        self._deformable_method = deformable_method
+        self._deformable_metric = deformable_metric
+        self._deformable_lr = float(deformable_lr)
+        self._integrator_n = int(integrator_n)
+        self._smooth_warp_sigma = float(smooth_warp_sigma)
+        self._smooth_grad_sigma = float(smooth_grad_sigma)
+        self._seed = int(seed)
+        # IMPACT deformable metric (only used when deformable_metric == "impact"): KonfAI IMPACT feature
+        # models drive the SyN/greedy stage instead of the analytic CC/MI/MSE.
+        self._impact_specs = impact_specs
+
+    @staticmethod
+    def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+        """True only for a mask that actually restricts the region — some voxels in, some out. An absent
+        optional mask arrives as a whole-image (all-ones) default and an all-zero mask is degenerate; both
+        are treated as no mask so the plain (non-masked) metric path is used."""
+        if mask is None:
+            return False
+        arr = sitk.GetArrayViewFromImage(mask)
+        return bool((arr > 0).any()) and bool((arr == 0).any())
+
+    @staticmethod
+    def _affine_to_sitk(affine_matrix: "torch.Tensor") -> sitk.AffineTransform:
+        """FireANTs' physical (LPS) linear matrix -> SimpleITK AffineTransform (fixed -> moving points),
+        the same convention FireANTs writes into an ANTs ``0GenericAffine.mat``."""
+        matrix = affine_matrix.float().cpu().numpy()[0]
+        affine = sitk.AffineTransform(DIM)
+        affine.SetMatrix(matrix[:DIM, :DIM].flatten().astype(np.float64))
+        affine.SetTranslation(matrix[:DIM, DIM].astype(np.float64))
+        return affine
+
+    def _total_field_transform(self, reg) -> sitk.Transform:
+        """Optimise a deformable stage and return its TOTAL displacement (affine baked in) as a
+        SimpleITK ``DisplacementFieldTransform`` on the fixed grid.
+
+        FireANTs serialises the total field (ANTs convention, fixed grid) only to a file, so it is
+        round-tripped through a temporary NIfTI — its public API, no internals reimplemented."""
+        reg.optimize()
+        with tempfile.TemporaryDirectory() as tmp:
+            warp_path = os.path.join(tmp, "total_warp.nii.gz")
+            reg.save_as_ants_transforms(warp_path)
+            total_field = sitk.ReadImage(warp_path, sitk.sitkVectorFloat64)
+        return sitk.DisplacementFieldTransform(total_field)  # consumes total_field
+
+    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."""
+        from fireants.io import BatchedImages, Image
+        from fireants.io.imagemask import apply_mask_to_image, generate_image_mask_allones
+        from fireants.registration.affine import AffineRegistration
+        from fireants.registration.rigid import RigidRegistration
+
+        torch.manual_seed(self._seed)
+        device = f"cuda:{device_index}" if device_index >= 0 else "cpu"
+        # FireANTs' Image ctor accepts a SimpleITK image directly, so the fixed/moving cross into
+        # FireANTs in-memory (no file load) with their geometry preserved.
+        fixed_img = Image(fixed, device=device)
+        moving_img = Image(moving, device=device)
+
+        # Masked metric only when a mask genuinely restricts the region. FireANTs' masked mode wants the
+        # mask as the last channel of BOTH images (all-ones where one side has none) and a ``masked_``
+        # metric prefix; the plain path is untouched when no real mask is present.
+        use_fixed_mask = self._is_partial_mask(fixed_mask)
+        use_moving_mask = self._is_partial_mask(moving_mask)
+        masked = use_fixed_mask or use_moving_mask
+        if masked:
+            fmask = Image(fixed_mask, device=device) if use_fixed_mask else generate_image_mask_allones(fixed_img)
+            mmask = Image(moving_mask, device=device) if use_moving_mask else generate_image_mask_allones(moving_img)
+            fixed_img = apply_mask_to_image(fixed_img, fmask)
+            moving_img = apply_mask_to_image(moving_img, mmask)
+
+        bf = BatchedImages([fixed_img])
+        bm = BatchedImages([moving_img])
+        affine_loss = f"masked_{self._affine_metric}" if masked else self._affine_metric
+        deformable_loss = f"masked_{self._deformable_metric}" if masked else self._deformable_metric
+
+        # Linear: Rigid(MI, COM init) -> Affine(MI, seeded by the rigid), mirroring ANTs. The affine
+        # seeds the deformable stage (or is the whole transform when deformable_method == "none").
+        rigid = RigidRegistration(
+            scales=self._scales,
+            iterations=self._affine_iterations,
+            fixed_images=bf,
+            moving_images=bm,
+            loss_type=affine_loss,
+            optimizer="Adam",
+            optimizer_lr=self._affine_lr,
+            cc_kernel_size=self._cc_kernel,
+            init_translation="cof",
+        )
+        rigid.optimize()
+        rigid_matrix = rigid.get_rigid_matrix().detach()
+
+        affine = AffineRegistration(
+            scales=self._scales,
+            iterations=self._affine_iterations,
+            fixed_images=bf,
+            moving_images=bm,
+            loss_type=affine_loss,
+            optimizer="Adam",
+            optimizer_lr=self._affine_lr,
+            cc_kernel_size=self._cc_kernel,
+            init_rigid=rigid_matrix,
+        )
+        affine.optimize()
+        affine_matrix = affine.get_affine_matrix().detach()
+
+        # Deformable stage (or none). SyN and Greedy share the same constructor surface; both warm-start
+        # from the affine so their TOTAL transform already bakes in the linear pre-align.
+        if self._deformable_method == "none":
+            transform: sitk.Transform = self._affine_to_sitk(affine_matrix)
+        else:
+            if self._deformable_method == "syn":
+                from fireants.registration.syn import SyNRegistration as Deformable
+            elif self._deformable_method == "greedy":
+                from fireants.registration.greedy import GreedyRegistration as Deformable
+            else:
+                raise ValueError(
+                    f"Unknown deformable_method '{self._deformable_method}' (expected 'syn', 'greedy' or 'none')."
+                )
+            # "impact" swaps the analytic metric for a KonfAI IMPACT feature loss on the deformable stage
+            # (the linear pre-align keeps its own affine_metric); masks do not restrict the IMPACT metric.
+            if self._deformable_metric == "impact":
+                loss_type: str = "custom"
+                custom_loss: torch.nn.Module | None = ImpactFeatureLoss(self._impact_specs)
+            else:
+                loss_type, custom_loss = deformable_loss, None
+            reg = Deformable(
+                scales=self._scales,
+                iterations=self._deformable_iterations,
+                fixed_images=bf,
+                moving_images=bm,
+                loss_type=loss_type,
+                custom_loss=custom_loss,
+                cc_kernel_size=self._cc_kernel,
+                deformation_type="compositive",
+                integrator_n=self._integrator_n,
+                smooth_warp_sigma=self._smooth_warp_sigma,
+                smooth_grad_sigma=self._smooth_grad_sigma,
+                optimizer="Adam",
+                optimizer_lr=self._deformable_lr,
+                init_affine=affine_matrix,
+            )
+            transform = self._total_field_transform(reg)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+        # Rebuild moved + DVF from the single transform on the fixed grid — the ConvexAdam output path,
+        # so every FireANTs preset emits identical-shaped results.
+        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
+
+
+class FireANTsRegistration(torch.nn.Module):
+    """Graph module: (fixed, moving) tensors + their geometry -> moved image + DVF on the fixed grid.
+
+    ``accepts_attributes = True`` opts this module into receiving the per-branch ``Attribute`` list
+    alongside the tensors (same convention as the ConvexAdam / elastix engines); registration needs the
+    physical geometry, and the mask branches restrict the metric.
+    """
+
+    accepts_attributes = True
+
+    def __init__(self, engine: FireANTsEngine) -> None:
+        super().__init__()
+        self._engine = engine
+
+    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) channel-stacked with the
+        # displacement field (DIM channels); downstream ChannelSelect modules split them. A whole-image
+        # mask (the default when none is supplied) restricts nothing.
+        fixed_attrs, moving_attrs, fmask_attrs, mmask_attrs = attributes
+        device_index = fixed.device.index if fixed.device.type == "cuda" else -1
+        combined = []
+        # FireANTs runs a gradient-based instance optimisation (Riemannian Adam over the warp); the
+        # predictor calls forward under torch.inference_mode(), which forbids autograd. The image tensors
+        # have already crossed to numpy/SimpleITK here, so re-enable grad for the optimisation.
+        with torch.inference_mode(False), torch.enable_grad():
+            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)."""
+
+    def __init__(self, start: int, stop: int) -> None:
+        super().__init__()
+        self._start = start
+        self._stop = stop
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor[:, self._start : self._stop]
+
+
+class RegistrationNet(network.Network):
+    """Pairwise FireANTs 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 on the fixed grid: ``MovedImage`` (moving resampled onto fixed) and ``DisplacementField``
+    (the DIM-component displacement field, in mm). Geometry is attached by the predictor via
+    ``same_as_group: Volume_0:Fixed``. The knobs below are read straight from these annotations by the
+    UI: ``Annotated[.., Range]`` gives numeric spin bounds; ``Literal`` a dropdown. ``deformable_method``
+    is the knob that specialises this shared model into each FireANTs preset.
+    """
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        scales: list[int] = [4, 2, 1],
+        affine_iterations: list[int] = [200, 100, 50],
+        deformable_iterations: list[int] = [200, 100, 50],
+        cc_kernel: Annotated[int, Range(1, 21)] = 5,
+        affine_metric: Literal["mi", "cc", "mse"] = "mi",
+        affine_lr: Annotated[float, Range(0.0, 10.0)] = 0.003,
+        deformable_method: Literal["none", "syn", "greedy"] = "syn",
+        deformable_metric: Literal["cc", "mi", "mse", "impact"] = "cc",
+        deformable_lr: Annotated[float, Range(0.0, 10.0)] = 0.25,
+        integrator_n: Annotated[int, Range(1, 100)] = 10,
+        smooth_warp_sigma: Annotated[float, Range(0.0, 100.0)] = 0.5,
+        smooth_grad_sigma: Annotated[float, Range(0.0, 100.0)] = 1.0,
+        seed: int = 42,
+        models: dict[str, ModelSpec] = {},
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=3,
+        )
+        engine = FireANTsEngine(
+            scales,
+            affine_iterations,
+            deformable_iterations,
+            cc_kernel,
+            affine_metric,
+            affine_lr,
+            deformable_method,
+            deformable_metric,
+            deformable_lr,
+            integrator_n,
+            smooth_warp_sigma,
+            smooth_grad_sigma,
+            seed,
+            _sorted_specs(models),
+        )
+        self.add_module(
+            "Registration", FireANTsRegistration(engine), in_branch=[0, 1, 2, 3], out_branch=["registration"]
+        )
+        self.add_module("MovedImage", ChannelSelect(0, 1), in_branch=["registration"], out_branch=["moved"])
+        self.add_module("DisplacementField", ChannelSelect(1, 4), in_branch=["registration"], out_branch=["dvf"])

FireANTs_SyN/NOTICE

@@ -0,0 +1,40 @@
+FireANTs_SyN — third-party attribution
+======================================
+
+This KonfAI app drives FireANTs, an external registration library. FireANTs is NOT
+redistributed here as source: the app depends on the official `fireants` wheel
+(https://pypi.org/project/fireants/), which is installed at resolve time (see
+requirements.txt) and called through its public Python API only. No FireANTs source
+code (functions, classes, or modules) is copied into this app.
+
+FireANTs is distributed under the **FireANTs License, Version 1.0 (July 2025)**, a
+custom license modified from Apache 2.0. Its redistribution clause requires that, when
+FireANTs is incorporated as a dependency in other projects, all license terms —
+including attribution and the bibliography below — be maintained.
+
+  Project : FireANTs
+  Source  : https://github.com/rohitrango/FireANTs
+  License : https://github.com/rohitrango/FireANTs/blob/main/LICENSE
+  Copyright (c) 2026 Rohit Jena. All rights reserved.
+
+Bibliography (as required by the FireANTs License — cite if you use this app)
+----------------------------------------------------------------------------
+
+  @article{jena2024fireants,
+    title={FireANTs: Adaptive Riemannian Optimization for Multi-Scale Diffeomorphic Registration},
+    author={Jena, Rohit and Chaudhari, Pratik and Gee, James C},
+    journal={Nature Communications},
+    year={2024}
+  }
+
+  @inproceedings{jena2025scalable,
+    title={A Scalable Distributed Framework for Multimodal {GigaVoxel} Image Registration},
+    author={Jena, Rohit and Zope, Vedant and Chaudhari, Pratik and Gee, James C},
+    booktitle={The Fourteenth International Conference on Learning Representations},
+    year={2026},
+    url={https://openreview.net/forum?id=8dLexnao2h}
+  }
+
+The wrapper code in this directory (Model.py, the KonfAI app configuration) is original
+work © 2025 Valentin Boussot, licensed under Apache-2.0, and is a separate work that
+merely links to the FireANTs interfaces.

FireANTs_SyN/Prediction.yml

@@ -0,0 +1,132 @@
+Predictor:
+  Model:
+    classpath: Model:RegistrationNet
+    RegistrationNet:
+      scales:
+      - 4
+      - 2
+      - 1
+      affine_iterations:
+      - 200
+      - 100
+      - 50
+      deformable_iterations:
+      - 200
+      - 100
+      - 50
+      cc_kernel: 5
+      affine_metric: mi
+      affine_lr: 0.003
+      deformable_method: syn
+      deformable_metric: cc
+      deformable_lr: 0.25
+      integrator_n: 10
+      smooth_warp_sigma: 0.5
+      smooth_grad_sigma: 1.0
+      seed: 42
+      outputs_criterions: None
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Fixed:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_1:
+        groups_dest:
+          Moving:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_2:
+        groups_dest:
+          FixedMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+      Volume_3:
+        groups_dest:
+          MovingMask:
+            transforms:
+              TensorCast:
+                dtype: float32
+                inverse: false
+            patch_transforms: None
+            is_input: true
+    augmentations:
+      DataAugmentation_0:
+        data_augmentations:
+          Flip:
+            f_prob:
+            - 0
+            - 0.5
+            - 0.5
+            vector_field: true
+            prob: 1
+        nb: 2
+    Patch:
+      patch_size: None
+      overlap: None
+      mask: None
+      pad_value: None
+      extend_slice: 0
+    subset: None
+    filter: None
+    dataset_filenames:
+    - ./Dataset/:mha
+    use_cache: false
+    batch_size: 1
+    num_workers: None
+    pin_memory: false
+    prefetch_factor: None
+    persistent_workers: None
+  outputs_dataset:
+    MovedImage:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: Moved:mha
+        group: Moved
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+    DisplacementField:
+      OutputDataset:
+        name_class: OutSameAsGroupDataset
+        before_reduction_transforms: None
+        after_reduction_transforms: None
+        final_transforms:
+          TensorCast:
+            dtype: float32
+            inverse: false
+        dataset_filename: DVF:mha
+        group: DVF
+        same_as_group: Volume_0:Fixed
+        patch_combine: None
+        inverse_transform: false
+        reduction: Mean
+        Mean: {}
+  train_name: ImpactReg-FireANTs-SyN
+  manual_seed: 42
+  gpu_checkpoints: None
+  images_log: None
+  combine: Mean
+  autocast: false
+  data_log: None

FireANTs_SyN/Uncertainty.yml

@@ -0,0 +1,24 @@
+Evaluator:
+  metrics:
+    Uncertainty:
+      targets_criterions:
+        None:
+          criterions_loader:
+            Mean:
+              name: Uncertainty
+  Dataset:
+    groups_src:
+      Volume_0:
+        groups_dest:
+          Uncertainty:
+            transforms:
+              Norm: {}
+              StandardDeviation: {}
+              Save:
+                dataset: ./Uncertainties/ImpactReg/Output:mha
+                group: None
+    subset: None
+    dataset_filenames:
+    - ./Dataset:mha
+    validation: None
+  train_name: ImpactReg

FireANTs_SyN/app.json

@@ -0,0 +1,96 @@
+{
+    "display_name": "FireANTs (SyN)",
+    "short_description": "Rigid + Affine + SyN diffeomorphic registration on GPU (FireANTs).",
+    "description": "GPU deformable registration with FireANTs: a Rigid (MI) + Affine (MI) linear pre-align followed by a symmetric diffeomorphic SyN stage (CC), all optimised with Riemannian Adam. Produces the moved image and the total displacement field on the fixed grid. FireANTs is an external dependency under the FireANTs License v1.0 (see NOTICE).",
+    "task": "registration",
+    "tta": 0,
+    "mc_dropout": 0,
+    "models": [
+        "model.pt"
+    ],
+    "inputs": {
+        "Fixed": {
+            "display_name": "Fixed image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "Moving": {
+            "display_name": "Moving image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "FixedMask": {
+            "display_name": "Fixed mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false,
+            "default": "ones"
+        },
+        "MovingMask": {
+            "display_name": "Moving mask (optional)",
+            "volume_type": "SEGMENTATION",
+            "required": false,
+            "default": "ones"
+        }
+    },
+    "outputs": {
+        "MovedImage": {
+            "display_name": "Moved image",
+            "volume_type": "VOLUME",
+            "required": true
+        },
+        "DisplacementField": {
+            "display_name": "Displacement field",
+            "volume_type": "VOLUME",
+            "required": false
+        }
+    },
+    "inputs_evaluations": {
+        "Image": {
+            "Evaluation_with_images.yml": {
+                "FixedImage": {
+                    "display_name": "Fixed image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "MovingImage": {
+                    "display_name": "Moving image",
+                    "volume_type": "VOLUME",
+                    "required": true
+                },
+                "Mask": {
+                    "display_name": "Evaluation mask",
+                    "volume_type": "SEGMENTATION",
+                    "required": false
+                }
+            }
+        },
+        "Segmentation": {
+            "Evaluation_with_seg.yml": {
+                "FixedSeg": {
+                    "display_name": "Fixed segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                },
+                "MovingSeg": {
+                    "display_name": "Moving segmentation",
+                    "volume_type": "SEGMENTATION",
+                    "required": true
+                }
+            }
+        },
+        "Landmarks": {
+            "Evaluation_with_fid.yml": {
+                "FixedFid": {
+                    "display_name": "Fixed landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                },
+                "MovingFid": {
+                    "display_name": "Moving landmarks",
+                    "volume_type": "FIDUCIALS",
+                    "required": true
+                }
+            }
+        }
+    }
+}

FireANTs_SyN/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de99fbc36331ce674639acc774f52b4a2d0027f2f312d9d28669e831a0c4fd7e
+size 1249

FireANTs_SyN/requirements.txt

@@ -0,0 +1 @@
+fireants

Generic_Rigid/app.json

@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

Generic_Rigid/elastix_engine.py

@@ -44,6 +44,17 @@ from Model import _sorted_specs, generate_impact_parameter_map, load_models_regi
 ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
 
 
+def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
+    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
+    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
+    the whole image) instead of paying for a mask that restricts nothing."""
+    if mask is None:
+        return False
+    arr = sitk.GetArrayViewFromImage(mask)
+    return bool((arr > 0).any()) and bool((arr == 0).any())
+
+
 class ElastixEngine:
     """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
 
@@ -235,7 +246,7 @@ class ElastixEngine:
 
             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:
+                if _is_partial_mask(mask):
                     mask_path = work / name
                     sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
                     args += [flag, str(mask_path)]

Generic_Rigid_BSpline/app.json

@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {

Generic_Rigid_BSpline/elastix_engine.py

@@ -44,6 +44,17 @@ from Model import _sorted_specs, generate_impact_parameter_map, load_models_regi
 ELASTIX_CACHE = Path.home() / ".cache" / "konfai" / "elastix-impact"
 
 
+def _is_partial_mask(mask: "sitk.Image | None") -> bool:
+    """True only for a mask that actually restricts the metric region — some voxels in, some out. An
+    absent optional mask arrives as a whole-image (all-ones) default from KonfAI, and an all-zero mask
+    is degenerate; both are treated as no mask, so elastix runs without ``-fMask`` / ``-mMask`` (i.e.
+    the whole image) instead of paying for a mask that restricts nothing."""
+    if mask is None:
+        return False
+    arr = sitk.GetArrayViewFromImage(mask)
+    return bool((arr > 0).any()) and bool((arr == 0).any())
+
+
 class ElastixEngine:
     """Run the elastix-IMPACT binary on a fixed/moving pair; return (moved, dvf) on the fixed grid.
 
@@ -235,7 +246,7 @@ class ElastixEngine:
 
             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:
+                if _is_partial_mask(mask):
                     mask_path = work / name
                     sitk.WriteImage(sitk.Cast(mask, sitk.sitkUInt8), str(mask_path))
                     args += [flag, str(mask_path)]

MR_CT_HeadNeck/app.json

@@ -22,12 +22,14 @@
         "FixedMask": {
             "display_name": "Fixed mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         },
         "MovingMask": {
             "display_name": "Moving mask (optional)",
             "volume_type": "SEGMENTATION",
-            "required": false
+            "required": false,
+            "default": "ones"
         }
     },
     "outputs": {