Added TRT config for inference

generation/maisi/README.md

@@ -172,6 +172,16 @@ python -m scripts.inference -c ./configs/config_maisi.json -i ./configs/config_i
 
 Please refer to [maisi_inference_tutorial.ipynb](maisi_inference_tutorial.ipynb) for the tutorial for MAISI model inference.
 
+
+#### Accelerated Inference with TensorRT:
+To run the inference script with TensorRT acceleration, please run:
+```bash
+export MONAI_DATA_DIRECTORY=<dir_you_will_download_data>
+python -m scripts.inference -c ./configs/config_maisi.json -i ./configs/config_infer.json -e ./configs/environment.json -x ./configs/config_trt.json --random-seed 0
+```
+Extra config file,  [./configs/config_trt.json](./configs/config_trt.json) is using `trt_compile()` utility from MONAI to convert select modules to TensorRT by overriding their definitions from [./configs/config_infer.json](./configs/config_infer.json).
+
+
 #### Quality Check:
 We have implemented a quality check function for the generated CT images. The main idea behind this function is to ensure that the Hounsfield units (HU) intensity for each organ in the CT images remains within a defined range. For each training image used in the Diffusion network, we computed the median value for a few major organs. Then we summarize the statistics of these median values and save it to [./configs/image_median_statistics.json](./configs/image_median_statistics.json). During inference, for each generated image, we compute the median HU values for the major organs and check whether they fall within the normal range.
 

generation/maisi/configs/config_infer.json

@@ -18,5 +18,10 @@
         2.0
     ],
     "autoencoder_sliding_window_infer_size": [48,48,48],
-    "autoencoder_sliding_window_infer_overlap": 0.25
+    "autoencoder_sliding_window_infer_overlap": 0.25,
+    "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"
 }

generation/maisi/configs/config_trt.json

@@ -0,0 +1,24 @@
+{
+    "+imports": [
+        "$from monai.networks import trt_compile"
+    ],
+    "c_trt_args": {
+        "export_args": {
+            "dynamo": "$False",
+            "report": "$True"
+        },
+        "output_lists": [
+            [
+                -1
+            ],
+            [
+	    ]
+        ]
+    },
+    "device": "cuda",
+    "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path, @c_trt_args)",
+    "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path)",
+    "autoencoder": "$trt_compile(@autoencoder_def.to(@device), @trained_autoencoder_path, submodule='decoder')",
+    "mask_generation_autoencoder": "$trt_compile(@mask_generation_autoencoder_def.to(@device), @trained_mask_generation_autoencoder_path, submodule='decoder')",
+    "mask_generation_diffusion": "$trt_compile(@mask_generation_diffusion_def.to(@device), @trained_mask_generation_diffusion_path)"
+}

generation/maisi/scripts/inference.py

@@ -48,6 +48,12 @@ def main():
         default="./configs/config_infer.json",
         help="config json file that stores inference hyper-parameters",
     )
+    parser.add_argument(
+        "-x",
+        "--extra-config-file",
+        default=None,
+        help="config json file that stores inference extra parameters",
+    )
     parser.add_argument(
         "-s",
         "--random-seed",
@@ -140,6 +146,16 @@ def main():
         setattr(args, k, v)
         print(f"{k}: {v}")
 
+    #
+    # ## Read in optional extra configuration setting - typically acceleration options (TRT)
+    #
+    #
+    if args.extra_config_file is not None:
+        extra_config_dict = json.load(open(args.extra_config_file, "r"))
+        for k, v in extra_config_dict.items():
+            setattr(args, k, v)
+            print(f"{k}: {v}")
+
     check_input(
         args.body_region,
         args.anatomy_list,
@@ -158,25 +174,25 @@ def main():
 
     device = torch.device("cuda")
 
-    autoencoder = define_instance(args, "autoencoder_def").to(device)
+    autoencoder = define_instance(args, "autoencoder").to(device)
     checkpoint_autoencoder = torch.load(args.trained_autoencoder_path)
     autoencoder.load_state_dict(checkpoint_autoencoder)
 
-    diffusion_unet = define_instance(args, "diffusion_unet_def").to(device)
+    diffusion_unet = define_instance(args, "diffusion_unet").to(device)
     checkpoint_diffusion_unet = torch.load(args.trained_diffusion_path)
     diffusion_unet.load_state_dict(checkpoint_diffusion_unet["unet_state_dict"], strict=True)
     scale_factor = checkpoint_diffusion_unet["scale_factor"].to(device)
 
-    controlnet = define_instance(args, "controlnet_def").to(device)
+    controlnet = define_instance(args, "controlnet").to(device)
     checkpoint_controlnet = torch.load(args.trained_controlnet_path)
     monai.networks.utils.copy_model_state(controlnet, diffusion_unet.state_dict())
     controlnet.load_state_dict(checkpoint_controlnet["controlnet_state_dict"], strict=True)
 
-    mask_generation_autoencoder = define_instance(args, "mask_generation_autoencoder_def").to(device)
+    mask_generation_autoencoder = define_instance(args, "mask_generation_autoencoder").to(device)
     checkpoint_mask_generation_autoencoder = torch.load(args.trained_mask_generation_autoencoder_path)
     mask_generation_autoencoder.load_state_dict(checkpoint_mask_generation_autoencoder)
 
-    mask_generation_diffusion_unet = define_instance(args, "mask_generation_diffusion_def").to(device)
+    mask_generation_diffusion_unet = define_instance(args, "mask_generation_diffusion").to(device)
     checkpoint_mask_generation_diffusion_unet = torch.load(args.trained_mask_generation_diffusion_path)
     mask_generation_diffusion_unet.load_state_dict(checkpoint_mask_generation_diffusion_unet["unet_state_dict"])
     mask_generation_scale_factor = checkpoint_mask_generation_diffusion_unet["scale_factor"]