Merge branch 'dev' into GHSA-wg9g-w2j2-8pgr

Author: ericspod

monai/apps/nnunet/nnunetv2_runner.py

@@ -14,6 +14,7 @@
 
 import glob
 import os
+import re
 import shlex
 import subprocess
 from typing import Any
@@ -34,6 +35,8 @@
 
 __all__ = ["nnUNetV2Runner"]
 
+DATASET_ID_FORMAT = r"Dataset[0-9]{3}|[0-9]+"  # regex format for a valid nnUnet dataset name
+
 
 class nnUNetV2Runner:  # noqa: N801
     """
@@ -195,6 +198,13 @@ def __init__(
 
         # dataset_name_or_id has to be a string
         self.dataset_name_or_id = str(self.input_info.pop("dataset_name_or_id", 1))
+        self.dataset_name: str | None = None
+
+        # ensure the dataset name is a single identifier/number, this prevents code injection when composing commands
+        if re.fullmatch(DATASET_ID_FORMAT, self.dataset_name_or_id) is None:
+            raise ValueError(
+                f"Value for dataset_name_or_id `{self.dataset_name_or_id}` not a valid dataset name or ID."
+            )
 
         try:
             from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name
@@ -239,7 +249,7 @@ def convert_dataset(self):
 
             from nnunetv2.utilities.dataset_name_id_conversion import maybe_convert_to_dataset_name
 
-            self.dataset_name = maybe_convert_to_dataset_name(int(self.dataset_name_or_id))
+            self.dataset_name = maybe_convert_to_dataset_name(self.dataset_name_or_id)
 
             datalist_json = ConfigParser.load_config_file(self.input_info.pop("datalist"))
 
@@ -548,7 +558,7 @@ def train_single_model_command(
         Raises:
             ValueError: If gpu_id is an empty tuple or list.
         """
-        env = os.environ.copy()
+        env: dict[str, str] = os.environ.copy()
         device_setting: str = "0"
         num_gpus = 1
         if isinstance(gpu_id, str):
@@ -574,22 +584,25 @@ def train_single_model_command(
 
         cmd = [
             "nnUNetv2_train",
-            f"{self.dataset_name_or_id}",
-            f"{config}",
-            f"{fold}",
+            self.dataset_name_or_id,
+            config,
+            fold,
             "-tr",
-            f"{self.trainer_class_name}",
+            self.trainer_class_name,
             "-num_gpus",
-            f"{num_gpus}",
+            num_gpus,
         ]
+
         if self.export_validation_probabilities:
             cmd.append("--npz")
+
         for _key, _value in kwargs.items():
-            if _key == "p" or _key == "pretrained_weights":
-                cmd.extend([f"-{_key}", f"{_value}"])
-            else:
-                cmd.extend([f"--{_key}", f"{_value}"])
-        return cmd, env
+            prefix = "-" if _key in {"p", "pretrained_weights"} else "--"
+            cmd += [f"{prefix}{_key}", str(_value)]
+
+        cmd_str: list[str] = [str(c) for c in cmd]
+
+        return cmd_str, env
 
     def train(
         self,
@@ -641,7 +654,14 @@ def train_parallel_cmd(
                 None (all available GPUs).
             kwargs: this optional parameter allows you to specify additional arguments defined in the
                 ``train_single_model`` method.
+
+        Raises:
+            ValueError: self.dataset_name must have a value, ie. when using an existing dataset or after creating one.
         """
+
+        if self.dataset_name is None:
+            raise ValueError(f"A valid dataset name must be given in {self.dataset_name=}.")
+
         # unpack compressed files
         folder_names = []
         for root, _, files in os.walk(os.path.join(self.nnunet_preprocessed, self.dataset_name)):
@@ -696,7 +716,14 @@ def train_parallel(
                 None (all available GPUs).
             kwargs: this optional parameter allows you to specify additional arguments defined in the
                 ``train_single_model`` method.
+
+        Raises:
+            ValueError: self.dataset_name must have a value, ie. when using an existing dataset or after creating one.
         """
+
+        if self.dataset_name is None:
+            raise ValueError(f"A valid dataset name must be given in {self.dataset_name=}.")
+
         all_cmds = self.train_parallel_cmd(configs=configs, gpu_id_for_all=gpu_id_for_all, **kwargs)
         for s, cmds in enumerate(all_cmds):
             for gpu_id, gpu_cmd in cmds.items():
@@ -908,7 +935,14 @@ def predict_ensemble_postprocessing(
             run_postprocessing: whether to conduct post-processing
             kwargs: this optional parameter allows you to specify additional arguments defined in the
                 ``predict`` method.
+
+        Raises:
+            ValueError: self.dataset_name must have a value, ie. when using an existing dataset or after creating one.
         """
+
+        if self.dataset_name is None:
+            raise ValueError(f"A valid dataset name must be given in {self.dataset_name=}.")
+
         from nnunetv2.ensembling.ensemble import ensemble_folders
         from nnunetv2.postprocessing.remove_connected_components import apply_postprocessing_to_folder
         from nnunetv2.utilities.file_path_utilities import get_output_folder

monai/losses/image_dissimilarity.py

@@ -233,9 +233,11 @@ def __init__(
         self.kernel_type = look_up_option(kernel_type, ["gaussian", "b-spline"])
         self.num_bins = num_bins
         self.kernel_type = kernel_type
+        self.bin_centers: torch.Tensor | None
+        self.register_buffer("bin_centers", None, persistent=False)
         if self.kernel_type == "gaussian":
             self.preterm = 1 / (2 * sigma**2)
-            self.bin_centers = bin_centers[None, None, ...]
+            self.register_buffer("bin_centers", bin_centers[None, None, ...], persistent=False)
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
 
@@ -314,6 +316,8 @@ def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, to
         """
         img = torch.clamp(img, 0, 1)
         img = img.reshape(img.shape[0], -1, 1)  # (batch, num_sample, 1)
+        if self.bin_centers is None:
+            raise ValueError("bin_centers must be defined for gaussian parzen windowing.")
         weight = torch.exp(
             -self.preterm.to(img) * (img - self.bin_centers.to(img)) ** 2
         )  # (batch, num_sample, num_bin)

monai/networks/blocks/crossattention.py

@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
 
@@ -139,7 +141,7 @@ def __init__(
         )
         self.input_size = input_size
 
-    def forward(self, x: torch.Tensor, context: torch.Tensor | None = None):
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None):  # noqa: UP045
         """
         Args:
             x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C

monai/networks/blocks/selfattention.py

@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -158,7 +160,7 @@ def __init__(
         )
         self.input_size = input_size
 
-    def forward(self, x, attn_mask: torch.Tensor | None = None):
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):  # noqa: UP045
         """
         Args:
             x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C

monai/networks/blocks/transformerblock.py

@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
 
@@ -23,6 +25,11 @@ class TransformerBlock(nn.Module):
     An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <https://arxiv.org/abs/2010.11929>"
     """
 
+    # Treat ``with_cross_attention`` as a TorchScript constant so the cross-attention branch in
+    # ``forward`` is pruned when it is False. Otherwise scripting tries to compile the
+    # ``self.cross_attn(..., context=context)`` call against ``nn.Identity`` and fails.
+    __constants__ = ["with_cross_attention"]
+
     def __init__(
         self,
         hidden_size: int,
@@ -46,11 +53,20 @@ def __init__(
             dropout_rate (float, optional): fraction of the input units to drop. Defaults to 0.0.
             qkv_bias(bool, optional): apply bias term for the qkv linear layer. Defaults to False.
             save_attn (bool, optional): to make accessible the attention matrix. Defaults to False.
+            causal (bool, optional): whether to apply causal masking in self-attention. Defaults to False.
+            sequence_length (int | None, optional): sequence length required for causal masking. Defaults to None.
+            with_cross_attention (bool, optional): whether to include cross-attention layers that attend to an
+                external context tensor. When False, cross_attn is set to nn.Identity() so that the attribute
+                always exists for typing and checkpoint compatibility. Defaults to False.
             use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism
                 (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html).
             include_fc: whether to include the final linear layer. Default to True.
             use_combined_linear: whether to use a single linear layer for qkv projection, default to True.
 
+        Raises:
+            ValueError: if dropout_rate is not in [0, 1].
+            ValueError: if hidden_size is not divisible by num_heads.
+
         """
 
         super().__init__()
@@ -79,17 +95,24 @@ def __init__(
         self.with_cross_attention = with_cross_attention
 
         self.norm_cross_attn = nn.LayerNorm(hidden_size)
-        self.cross_attn = CrossAttentionBlock(
-            hidden_size=hidden_size,
-            num_heads=num_heads,
-            dropout_rate=dropout_rate,
-            qkv_bias=qkv_bias,
-            causal=False,
-            use_flash_attention=use_flash_attention,
-        )
+        self.cross_attn: CrossAttentionBlock | nn.Identity
+        if with_cross_attention:
+            self.cross_attn = CrossAttentionBlock(
+                hidden_size=hidden_size,
+                num_heads=num_heads,
+                dropout_rate=dropout_rate,
+                qkv_bias=qkv_bias,
+                causal=False,
+                use_flash_attention=use_flash_attention,
+            )
+        else:
+            self.cross_attn = nn.Identity()
 
     def forward(
-        self, x: torch.Tensor, context: torch.Tensor | None = None, attn_mask: torch.Tensor | None = None
+        self,
+        x: torch.Tensor,
+        context: Optional[torch.Tensor] = None,  # noqa: UP045
+        attn_mask: Optional[torch.Tensor] = None,  # noqa: UP045
     ) -> torch.Tensor:
         x = x + self.attn(self.norm1(x), attn_mask=attn_mask)
         if self.with_cross_attention:

monai/networks/nets/masked_autoencoder_vit.py

@@ -209,3 +209,41 @@ def forward(self, x, masking_ratio: float | None = None):
 
         x = x[:, 1:, :]
         return x, mask
+
+    def load_old_state_dict(self, old_state_dict: dict, verbose: bool = False) -> None:
+        """
+        Load a state dict from a MaskedAutoEncoderViT model trained with an older version of MONAI
+        where ``CrossAttentionBlock`` was unconditionally instantiated in ``TransformerBlock``
+        even when ``with_cross_attention=False``. Old checkpoints contain stale
+        ``blocks.{i}.cross_attn.*`` and ``decoder_blocks.{i}.cross_attn.*`` keys that are not
+        present in the current model and are automatically dropped.
+
+        Args:
+            old_state_dict: state dict from the older MaskedAutoEncoderViT model.
+            verbose: if True, print keys that are missing or unmatched. Defaults to False.
+        """
+        new_state_dict = self.state_dict()
+        if all(k in new_state_dict for k in old_state_dict):
+            if verbose:
+                print("All keys match, loading state dict.")
+            self.load_state_dict(old_state_dict)
+            return
+
+        if verbose:
+            for k in new_state_dict:
+                if k not in old_state_dict:
+                    print(f"key {k} not found in old state dict")
+            print("----------------------------------------------")
+            for k in old_state_dict:
+                if k not in new_state_dict:
+                    print(f"key {k} not found in new state dict")
+
+        # copy over all matching keys; stale cross_attn.* keys in blocks and decoder_blocks
+        # are left as unmatched leftovers and are not inserted into new_state_dict
+        for k in new_state_dict:
+            if k in old_state_dict:
+                new_state_dict[k] = old_state_dict.pop(k)
+
+        if verbose:
+            print("remaining keys in old_state_dict:", old_state_dict.keys())
+        self.load_state_dict(new_state_dict)

monai/networks/nets/vit.py

@@ -131,3 +131,41 @@ def forward(self, x):
         if hasattr(self, "classification_head"):
             x = self.classification_head(x[:, 0])
         return x, hidden_states_out
+
+    def load_old_state_dict(self, old_state_dict: dict, verbose: bool = False) -> None:
+        """
+        Load a state dict from a ViT model trained with an older version of MONAI where
+        ``CrossAttentionBlock`` was unconditionally instantiated in ``TransformerBlock``
+        even when ``with_cross_attention=False``. Old checkpoints contain stale
+        ``blocks.{i}.cross_attn.*`` keys that are not present in the current model and
+        are automatically dropped.
+
+        Args:
+            old_state_dict: state dict from the older ViT model.
+            verbose: if True, print keys that are missing or unmatched. Defaults to False.
+        """
+        new_state_dict = self.state_dict()
+        if all(k in new_state_dict for k in old_state_dict):
+            if verbose:
+                print("All keys match, loading state dict.")
+            self.load_state_dict(old_state_dict)
+            return
+
+        if verbose:
+            for k in new_state_dict:
+                if k not in old_state_dict:
+                    print(f"key {k} not found in old state dict")
+            print("----------------------------------------------")
+            for k in old_state_dict:
+                if k not in new_state_dict:
+                    print(f"key {k} not found in new state dict")
+
+        # copy over all matching keys; stale cross_attn.* keys are left as unmatched
+        # leftovers and are not inserted into new_state_dict
+        for k in new_state_dict:
+            if k in old_state_dict:
+                new_state_dict[k] = old_state_dict.pop(k)
+
+        if verbose:
+            print("remaining keys in old_state_dict:", old_state_dict.keys())
+        self.load_state_dict(new_state_dict)

monai/networks/nets/vitautoenc.py

@@ -133,3 +133,41 @@ def forward(self, x):
         x = self.conv3d_transpose(x)
         x = self.conv3d_transpose_1(x)
         return x, hidden_states_out
+
+    def load_old_state_dict(self, old_state_dict: dict, verbose: bool = False) -> None:
+        """
+        Load a state dict from a ViTAutoEnc model trained with an older version of MONAI where
+        ``CrossAttentionBlock`` was unconditionally instantiated in ``TransformerBlock``
+        even when ``with_cross_attention=False``. Old checkpoints contain stale
+        ``blocks.{i}.cross_attn.*`` keys that are not present in the current model and
+        are automatically dropped.
+
+        Args:
+            old_state_dict: state dict from the older ViTAutoEnc model.
+            verbose: if True, print keys that are missing or unmatched. Defaults to False.
+        """
+        new_state_dict = self.state_dict()
+        if all(k in new_state_dict for k in old_state_dict):
+            if verbose:
+                print("All keys match, loading state dict.")
+            self.load_state_dict(old_state_dict)
+            return
+
+        if verbose:
+            for k in new_state_dict:
+                if k not in old_state_dict:
+                    print(f"key {k} not found in old state dict")
+            print("----------------------------------------------")
+            for k in old_state_dict:
+                if k not in new_state_dict:
+                    print(f"key {k} not found in new state dict")
+
+        # copy over all matching keys; stale cross_attn.* keys are left as unmatched
+        # leftovers and are not inserted into new_state_dict
+        for k in new_state_dict:
+            if k in old_state_dict:
+                new_state_dict[k] = old_state_dict.pop(k)
+
+        if verbose:
+            print("remaining keys in old_state_dict:", old_state_dict.keys())
+        self.load_state_dict(new_state_dict)