Accelerated MAISI and make it compatible with previous DDPM

generation/maisi/configs/config_infer.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         256,
@@ -18,10 +18,11 @@
         2.0
     ],
     "autoencoder_sliding_window_infer_size": [48,48,48],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "autoencoder_sliding_window_infer_overlap": 0.6666,
     "controlnet": "$@controlnet_def",
     "diffusion_unet": "$@diffusion_unet_def",
     "autoencoder": "$@autoencoder_def",
     "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
-    "mask_generation_diffusion": "$@mask_generation_diffusion_def"
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_16g_256x256x128.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         256,
@@ -19,5 +19,11 @@
     ],
     "autoencoder_sliding_window_infer_size": [96,96,96],
     "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 2
+    "autoencoder_tp_num_splits": 2,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_16g_256x256x256.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         256,
@@ -18,6 +18,12 @@
         2.0
     ],
     "autoencoder_sliding_window_infer_size": [48,48,64],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 2
+    "autoencoder_sliding_window_infer_overlap": 0.6666,
+    "autoencoder_tp_num_splits": 4,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_16g_512x512x128.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         512,
@@ -18,6 +18,12 @@
         4.0
     ],
     "autoencoder_sliding_window_infer_size": [64,64,32],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 1
+    "autoencoder_sliding_window_infer_overlap": 0.5,
+    "autoencoder_tp_num_splits": 2,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_24g_256x256x256.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         256,
@@ -19,5 +19,11 @@
     ],
     "autoencoder_sliding_window_infer_size": [64,64,64],
     "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 4
+    "autoencoder_tp_num_splits": 4,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_24g_512x512x128.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         512,
@@ -18,6 +18,12 @@
         4.0
     ],
     "autoencoder_sliding_window_infer_size": [80,80,32],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 1
+    "autoencoder_sliding_window_infer_overlap": 0.4,
+    "autoencoder_tp_num_splits": 2,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_24g_512x512x512.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         512,
@@ -18,6 +18,12 @@
         1.0
     ],
     "autoencoder_sliding_window_infer_size": [64,64,48],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 2
+    "autoencoder_sliding_window_infer_overlap": 0.4,
+    "autoencoder_tp_num_splits": 2,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_32g_512x512x512.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         512,
@@ -17,7 +17,13 @@
         0.75,
         1.0
     ],
-    "autoencoder_sliding_window_infer_size": [64,64,64],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 2
+    "autoencoder_sliding_window_infer_size": [80,80,48],
+    "autoencoder_sliding_window_infer_overlap": 0.4,
+    "autoencoder_tp_num_splits": 4,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_80g_512x512x512.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         512,
@@ -18,6 +18,12 @@
         1.0
     ],
     "autoencoder_sliding_window_infer_size": [80,80,80],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 2
+    "autoencoder_sliding_window_infer_overlap": 0.4,
+    "autoencoder_tp_num_splits": 2,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_infer_80g_512x512x768.json

@@ -1,9 +1,9 @@
 {
     "num_output_samples": 1,
     "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "anatomy_list": ["liver"],
     "controllable_anatomy_size": [],
-    "num_inference_steps": 1000,
+    "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
     "output_size": [
         512,
@@ -17,7 +17,13 @@
         0.75,
         0.66667
     ],
-    "autoencoder_sliding_window_infer_size": [80,80,112],
-    "autoencoder_sliding_window_infer_overlap": 0.25,
-    "autoencoder_tp_num_splits": 4
+    "autoencoder_sliding_window_infer_size": [80,80,96],
+    "autoencoder_sliding_window_infer_overlap": 0.4,
+    "autoencoder_tp_num_splits": 4,
+    "controlnet": "$@controlnet_def",
+    "diffusion_unet": "$@diffusion_unet_def",
+    "autoencoder": "$@autoencoder_def",
+    "mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
+    "mask_generation_diffusion": "$@mask_generation_diffusion_def",
+    "modality": 1
 }

generation/maisi/configs/config_maisi.json → generation/maisi/configs/config_maisi3d-ddpm.json

@@ -2,6 +2,7 @@
     "spatial_dims": 3,
     "image_channels": 1,
     "latent_channels": 4,
+    "include_body_region": true,
     "mask_generation_latent_shape": [
         4,
         64,
@@ -60,8 +61,8 @@
         ],
         "num_res_blocks": 2,
         "use_flash_attention": true,
-        "include_top_region_index_input": true,
-        "include_bottom_region_index_input": true,
+        "include_top_region_index_input": "@include_body_region",
+        "include_bottom_region_index_input": "@include_body_region",
         "include_spacing_input": true
     },
     "controlnet_def": {