Added TRT config for inference

Signed-off-by: Boris Fomitchev <[email protected]>

Author: borisfom

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,22 @@
+{
+    "+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)",
+    "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",
@@ -122,7 +128,7 @@ def main():
         print(f"{k}: {val}")
     print("Global config variables have been loaded.")
 
-    # ## Read in configuration setting, including network definition, body region and anatomy to generate, etc.
+     # ## Read in configuration setting, including network definition, body region and anatomy to generate, etc.
     #
     # The information for the inference input, like body region and anatomy to generate, is stored in "./configs/config_infer.json".
     # Please refer to README.md for the details.
@@ -135,11 +141,21 @@ def main():
     # 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"]
+        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}")
 
+    #
+    # ## 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"]