Added TRT config for inference (#1907)

### Description
Added TRT config for MAISI and the extension of the inference script to
handle extra config file.

Note: autoencoder.decoder currently cannot be exported to TRT (crashes
during engine generation).
It does not seem to take a big part of the whole run anyway.

---------

Signed-off-by: Boris Fomitchev <[email protected]>
Co-authored-by: Yiheng Wang <[email protected]>

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"]