Merge branch 'Project-MONAI:main' into main

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: end-of-file-fixer
       - id: trailing-whitespace
@@ -23,7 +23,7 @@ repos:
       - id: detect-private-key
 
   - repo: https://github.com/psf/black
-    rev: "24.4.2"
+    rev: "24.10.0"
     hooks:
       - id: black
       - id: black-jupyter

generation/2d_diffusion_autoencoder/2d_diffusion_autoencoder_tutorial.ipynb

@@ -474,9 +474,8 @@
     }
    ],
    "source": [
-    "max_epochs = (\n",
-    "    1000  # training for longer (1e4 ~ 3h) helps a lot with reconstruction quality, even if the loss is already low\n",
-    ")\n",
+    "# training for longer (1e4 ~ 3h) helps a lot with reconstruction quality, even if the loss is already low\n",
+    "max_epochs = 1000\n",
     "val_interval = 100\n",
     "print_interval = 50\n",
     "iter_loss_list, val_iter_loss_list = [], []\n",

generation/README.md

@@ -43,7 +43,7 @@ Example shows the use cases of using MONAI to evaluate the performance of a gene
 ## [Training a 2D VQ-VAE + Autoregressive Transformers](./2d_vqvae_transformer/2d_vqvae_transformer_tutorial.ipynb):
 Example shows how to train a Vector-Quantized Variation Autoencoder + Transformers on the MedNIST dataset.
 
-## Training VQ-VAEs and VQ-GANs: [2D VAE](./2d_vqvae/2d_vqvae_tutorial.ipynb), [3D VAE](./3d_vqvae/3d_vqvae_tutorial.ipynb) and [2D GAN](./3d_autoencoderkl/2d_vqgan_tutorial.ipynb)
+## Training VQ-VAEs and VQ-GANs: [2D VAE](./2d_vqvae/2d_vqvae_tutorial.ipynb), [3D VAE](./3d_vqvae/3d_vqvae_tutorial.ipynb) and [2D GAN](./2d_vqgan/2d_vqgan_tutorial.ipynb)
 Examples show how to train Vector Quantized Variation Autoencoder on [2D](./2d_vqvae/2d_vqvae_tutorial.ipynb) and [3D](./3d_vqvae/3d_vqvae_tutorial.ipynb), and how to use the PatchDiscriminator class to train a [VQ-GAN](./2d_vqgan/2d_vqgan_tutorial.ipynb) and improve the quality of the generated images.
 
 ## [Training a 2D Denoising Diffusion Probabilistic Model](./2d_ddpm/2d_ddpm_tutorial.ipynb):

generation/anomaly_detection/2d_classifierfree_guidance_anomalydetection_tutorial.ipynb

@@ -426,7 +426,7 @@
    },
    "source": [
     "condition_dropout = 0.15\n",
-    "max_epochs = 2e4\n",
+    "max_epochs = 20000\n",
     "batch_size = 64\n",
     "val_interval = 100\n",
     "iter_loss_list = []\n",

generation/maisi/configs/config_infer_16g_256x256x128.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        256,
+        256,
+        128
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        1.5,
+        1.5,
+        4.0
+    ],
+    "autoencoder_sliding_window_infer_size": [96,96,96],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 2
+}

generation/maisi/configs/config_infer_16g_256x256x256.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        256,
+        256,
+        256
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        1.5,
+        1.5,
+        2.0
+    ],
+    "autoencoder_sliding_window_infer_size": [48,48,64],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 2
+}

generation/maisi/configs/config_infer_16g_512x512x128.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        512,
+        512,
+        128
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        0.75,
+        0.75,
+        4.0
+    ],
+    "autoencoder_sliding_window_infer_size": [64,64,32],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 1
+}

generation/maisi/configs/config_infer_24g_256x256x256.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        256,
+        256,
+        256
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        1.5,
+        1.5,
+        2.0
+    ],
+    "autoencoder_sliding_window_infer_size": [64,64,64],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 4
+}

generation/maisi/configs/config_infer_24g_512x512x128.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        512,
+        512,
+        128
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        0.75,
+        0.75,
+        4.0
+    ],
+    "autoencoder_sliding_window_infer_size": [80,80,32],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 1
+}

generation/maisi/configs/config_infer_24g_512x512x512.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        512,
+        512,
+        512
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        0.75,
+        0.75,
+        1.0
+    ],
+    "autoencoder_sliding_window_infer_size": [64,64,48],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 2
+}

generation/maisi/configs/config_infer_32g_512x512x512.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        512,
+        512,
+        512
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        0.75,
+        0.75,
+        1.0
+    ],
+    "autoencoder_sliding_window_infer_size": [64,64,64],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 2
+}

generation/maisi/configs/config_infer_80g_512x512x128.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        512,
+        512,
+        128
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        0.75,
+        0.75,
+        4.0
+    ],
+    "autoencoder_sliding_window_infer_size": [128,128,32],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 4
+}

generation/maisi/configs/config_infer_80g_512x512x512.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        512,
+        512,
+        512
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        0.75,
+        0.75,
+        1.0
+    ],
+    "autoencoder_sliding_window_infer_size": [80,80,80],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 2
+}

generation/maisi/configs/config_infer_80g_512x512x768.json

@@ -0,0 +1,23 @@
+{
+    "num_output_samples": 1,
+    "body_region": ["abdomen"],
+    "anatomy_list": ["liver","hepatic tumor"],
+    "controllable_anatomy_size": [],
+    "num_inference_steps": 1000,
+    "mask_generation_num_inference_steps": 1000,
+    "output_size": [
+        512,
+        512,
+        768
+    ],
+    "image_output_ext": ".nii.gz",
+    "label_output_ext": ".nii.gz",
+    "spacing": [
+        0.75,
+        0.75,
+        0.66667
+    ],
+    "autoencoder_sliding_window_infer_size": [80,80,112],
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_tp_num_splits": 4
+}

generation/maisi/scripts/inference.py

@@ -132,6 +132,10 @@ def main():
 
     # check the format of inference inputs
     config_infer_dict = json.load(open(args.inference_file, "r"))
+    # override num_split if asked
+    if "autoencoder_tp_num_splits" in config_infer_dict:
+        args.autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"]
+        args.mask_generation_autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"]
     for k, v in config_infer_dict.items():
         setattr(args, k, v)
         print(f"{k}: {v}")
@@ -225,4 +229,7 @@ def main():
         format="[%(asctime)s.%(msecs)03d][%(levelname)5s](%(name)s) - %(message)s",
         datefmt="%Y-%m-%d %H:%M:%S",
     )
+    torch.cuda.reset_peak_memory_stats()
     main()
+    peak_memory_gb = torch.cuda.max_memory_allocated() / (1024**3)  # Convert to GB
+    print(f"Peak GPU memory usage: {peak_memory_gb:.2f} GB")

generation/maisi/scripts/sample.py

@@ -132,7 +132,7 @@ def ldm_conditional_sample_one_mask(
             conditioning=anatomy_size.to(device),
         )
         # decode latents to synthesized masks
-        if math.prod(latent_shape[1:]) < math.prod(autoencoder_sliding_window_infer_size):
+        if math.prod(latent_shape[1:]) <= math.prod(autoencoder_sliding_window_infer_size):
             synthetic_mask = recon_model(latents).cpu().detach()
         else:
             synthetic_mask = (
@@ -274,15 +274,15 @@ def ldm_conditional_sample_one_image(
         # decode latents to synthesized images
         logging.info("---- Start decoding latent features into images... ----")
         start_time = time.time()
-        if math.prod(latent_shape[1:]) < math.prod(autoencoder_sliding_window_infer_size):
+        if math.prod(latent_shape[1:]) <= math.prod(autoencoder_sliding_window_infer_size):
             synthetic_images = recon_model(latents)
         else:
             synthetic_images = sliding_window_inference(
                 inputs=latents,
                 roi_size=(
-                    min(output_size[0] // 4 // 4 * 3, autoencoder_sliding_window_infer_size[0]),
-                    min(output_size[1] // 4 // 4 * 3, autoencoder_sliding_window_infer_size[1]),
-                    min(output_size[2] // 4 // 4 * 3, autoencoder_sliding_window_infer_size[2]),
+                    min(output_size[0] // 4, autoencoder_sliding_window_infer_size[0]),
+                    min(output_size[1] // 4, autoencoder_sliding_window_infer_size[1]),
+                    min(output_size[2] // 4, autoencoder_sliding_window_infer_size[2]),
                 ),
                 sw_batch_size=1,
                 predictor=recon_model,