Merge branch 'dev' into propose-fix-perceptual-loss-sqrt-nan

Author: ericspod

.github/workflows/pythonapp-min.yml

@@ -124,7 +124,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        pytorch-version: ['2.3.1', '2.4.1', '2.5.1', 'latest']
+        pytorch-version: ['2.4.1', '2.5.1', '2.6.0']  # FIXME: add 'latest' back once PyTorch 2.7 issues are resolved
     timeout-minutes: 40
     steps:
     - uses: actions/checkout@v4

.github/workflows/pythonapp.yml

@@ -155,7 +155,7 @@ jobs:
         # install the latest pytorch for testing
         # however, "pip install monai*.tar.gz" will build cpp/cuda with an isolated
         # fresh torch installation according to pyproject.toml
-        python -m pip install torch>=2.3.0 torchvision
+        python -m pip install torch>=2.4.1 torchvision
     - name: Check packages
       run: |
         pip uninstall monai

docs/requirements.txt

@@ -1,5 +1,5 @@
 -f https://download.pytorch.org/whl/cpu/torch-2.3.0%2Bcpu-cp39-cp39-linux_x86_64.whl
-torch>=2.3.0
+torch>=2.4.1, <2.7.0
 pytorch-ignite==0.4.11
 numpy>=1.20
 itk>=5.2

monai/bundle/scripts.py

@@ -312,21 +312,6 @@ def _get_ngc_token(api_key, retry=0):
         return token
 
 
-def _get_latest_bundle_version_monaihosting(name):
-    full_url = f"{MONAI_HOSTING_BASE_URL}/{name.lower()}"
-    if has_requests:
-        resp = requests.get(full_url)
-        try:
-            resp.raise_for_status()
-            model_info = json.loads(resp.text)
-            return model_info["model"]["latestVersionIdStr"]
-        except requests.exceptions.HTTPError:
-            # for monaihosting bundles, if cannot find the version, get from model zoo model_info.json
-            return get_bundle_versions(name)["latest_version"]
-
-    raise ValueError("NGC API requires requests package. Please install it.")
-
-
 def _examine_monai_version(monai_version: str) -> tuple[bool, str]:
     """Examine if the package version is compatible with the MONAI version in the metadata."""
     version_dict = get_versions()
@@ -430,7 +415,7 @@ def _get_latest_bundle_version(
         name = _add_ngc_prefix(name)
         return _get_latest_bundle_version_ngc(name)
     elif source == "monaihosting":
-        return _get_latest_bundle_version_monaihosting(name)
+        return get_bundle_versions(name, repo="Project-MONAI/model-zoo", tag="dev")["latest_version"]
     elif source == "ngc_private":
         headers = kwargs.pop("headers", {})
         name = _add_ngc_prefix(name)

monai/networks/blocks/selfattention.py

@@ -101,7 +101,11 @@ def __init__(
 
         self.num_heads = num_heads
         self.hidden_input_size = hidden_input_size if hidden_input_size else hidden_size
-        self.out_proj = nn.Linear(self.inner_dim, self.hidden_input_size)
+        self.out_proj: Union[nn.Linear, nn.Identity]
+        if include_fc:
+            self.out_proj = nn.Linear(self.inner_dim, self.hidden_input_size)
+        else:
+            self.out_proj = nn.Identity()
 
         self.qkv: Union[nn.Linear, nn.Identity]
         self.to_q: Union[nn.Linear, nn.Identity]

monai/networks/nets/diffusion_model_unet.py

@@ -1847,9 +1847,9 @@ def load_old_state_dict(self, old_state_dict: dict, verbose=False) -> None:
             new_state_dict[f"{block}.attn.to_v.bias"] = old_state_dict.pop(f"{block}.to_v.bias")
 
             # projection
-            new_state_dict[f"{block}.attn.out_proj.weight"] = old_state_dict.pop(f"{block}.proj_attn.weight")
-            new_state_dict[f"{block}.attn.out_proj.bias"] = old_state_dict.pop(f"{block}.proj_attn.bias")
-
+            if f"{block}.attn.out_proj.weight" in new_state_dict and f"{block}.attn.out_proj.bias" in new_state_dict:
+                new_state_dict[f"{block}.attn.out_proj.weight"] = old_state_dict.pop(f"{block}.proj_attn.weight")
+                new_state_dict[f"{block}.attn.out_proj.bias"] = old_state_dict.pop(f"{block}.proj_attn.bias")
         # fix the cross attention blocks
         cross_attention_blocks = [
             k.replace(".out_proj.weight", "")

monai/networks/schedulers/ddpm.py

@@ -238,7 +238,7 @@ def step(
         pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample
 
         # 6. Add noise
-        variance = 0
+        variance: torch.Tensor = torch.tensor(0)
         if timestep > 0:
             noise = torch.randn(
                 model_output.size(),

monai/transforms/inverse.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import threading
 import warnings
 from collections.abc import Hashable, Mapping
 from contextlib import contextmanager
@@ -66,15 +67,41 @@ class TraceableTransform(Transform):
     The information in the stack of applied transforms must be compatible with the
     default collate, by only storing strings, numbers and arrays.
 
-    `tracing` could be enabled by `self.set_tracing` or setting
+    `tracing` could be enabled by assigning to `self.tracing` or setting
     `MONAI_TRACE_TRANSFORM` when initializing the class.
     """
 
-    tracing = MONAIEnvVars.trace_transform() != "0"
+    def _init_trace_threadlocal(self):
+        """Create a `_tracing` instance member to store the thread-local tracing state value."""
+        # needed since this class is meant to be a trait with no constructor
+        if not hasattr(self, "_tracing"):
+            self._tracing = threading.local()
+
+        # This is True while the above initialising _tracing is False when this is
+        # called from a different thread than the one initialising _tracing.
+        if not hasattr(self._tracing, "value"):
+            self._tracing.value = MONAIEnvVars.trace_transform() != "0"
+
+    def __getstate__(self):
+        """When pickling, remove the `_tracing` member from the output, if present, since it's not picklable."""
+        _dict = dict(getattr(self, "__dict__", {}))  # this makes __dict__ always present in the unpickled object
+        _slots = {k: getattr(self, k) for k in getattr(self, "__slots__", [])}
+        _dict.pop("_tracing", None)  # remove tracing
+        return _dict if len(_slots) == 0 else (_dict, _slots)
+
+    @property
+    def tracing(self) -> bool:
+        """
+        Returns the tracing state, which is thread-local and initialised to `MONAIEnvVars.trace_transform() != "0"`.
+        """
+        self._init_trace_threadlocal()
+        return bool(self._tracing.value)
 
-    def set_tracing(self, tracing: bool) -> None:
-        """Set whether to trace transforms."""
-        self.tracing = tracing
+    @tracing.setter
+    def tracing(self, val: bool):
+        """Sets the thread-local tracing state to `val`."""
+        self._init_trace_threadlocal()
+        self._tracing.value = val
 
     @staticmethod
     def trace_key(key: Hashable = None):
@@ -291,7 +318,7 @@ def check_transforms_match(self, transform: Mapping) -> None:
 
     def get_most_recent_transform(self, data, key: Hashable = None, check: bool = True, pop: bool = False):
         """
-        Get most recent transform for the stack.
+        Get most recent matching transform for the current class from the sequence of applied operations.
 
         Args:
             data: dictionary of data or `MetaTensor`.
@@ -316,9 +343,14 @@ def get_most_recent_transform(self, data, key: Hashable = None, check: bool = Tr
                 all_transforms = data.get(self.trace_key(key), MetaTensor.get_default_applied_operations())
         else:
             raise ValueError(f"`data` should be either `MetaTensor` or dictionary, got {type(data)}.")
+
+        if not all_transforms:
+            raise ValueError(f"Item of type {type(data)} (key: {key}, pop: {pop}) has empty 'applied_operations'")
+
         if check:
             self.check_transforms_match(all_transforms[-1])
-        return all_transforms.pop() if pop else all_transforms[-1]
+
+        return all_transforms.pop(-1) if pop else all_transforms[-1]
 
     def pop_transform(self, data, key: Hashable = None, check: bool = True):
         """

pyproject.toml

@@ -2,14 +2,14 @@
 requires = [
   "wheel",
   "setuptools",
-  "torch>=2.3.0",
+  "torch>=2.4.1, <2.7.0",
   "ninja",
   "packaging"
 ]
 
 [tool.black]
 line-length = 120
-target-version = ['py38', 'py39', 'py310']
+target-version = ['py39', 'py310', 'py311', 'py312']
 include = '\.pyi?$'
 exclude = '''
 (

requirements-dev.txt

@@ -17,7 +17,7 @@ mccabe
 pep8-naming
 pycodestyle
 pyflakes
-black>=22.12
+black>=25.1.0
 isort>=5.1, <6.0
 ruff
 pytype>=2020.6.1; platform_system != "Windows"
@@ -29,6 +29,7 @@ torchvision
 psutil
 cucim-cu12; platform_system == "Linux" and python_version >= "3.9" and python_version <= "3.10"
 openslide-python
+openslide-bin
 imagecodecs; platform_system == "Linux" or platform_system == "Darwin"
 tifffile; platform_system == "Linux" or platform_system == "Darwin"
 pandas