add time cost and gpu memory to readme asn add corresponding configs

Signed-off-by: Can-Zhao <[email protected]>

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,