TRT support for MAISI

.pre-commit-config.yaml

@@ -1,5 +1,5 @@
 default_language_version:
-  python: python3.8
+  python: python3
 
 ci:
   autofix_prs: true

models/maisi_ct_generative/configs/inference.json

@@ -34,8 +34,8 @@
     "spatial_dims": 3,
     "image_channels": 1,
     "latent_channels": 4,
-    "output_size_xy": 512,
-    "output_size_z": 512,
+    "output_size_xy": 256,
+    "output_size_z": 256,
     "output_size": [
         "@output_size_xy",
         "@output_size_xy",

models/maisi_ct_generative/configs/inference_trt.json

@@ -0,0 +1,17 @@
+{
+    "+imports": [
+        "$from monai.networks import trt_compile"
+    ],
+    "c_trt_args": {
+        "output_lists": [
+            [
+                -1
+            ],
+            [
+                0
+            ]
+        ]
+    },
+    "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path, args=@c_trt_args)",
+    "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path)"
+}

models/maisi_ct_generative/configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20240318.json",
-    "version": "0.4.5",
+    "version": "0.4.6",
     "changelog": {
+        "0.4.6": "add TensorRT support",
         "0.4.5": "update README",
         "0.4.4": "update issue for IgniteInfo",
         "0.4.3": "remove download large files, add weights_only when loading weights and add label_dict to large files",

models/maisi_ct_generative/scripts/sample.py

@@ -241,7 +241,6 @@ def ldm_conditional_sample_one_image(
         # synthesize latents
         noise_scheduler.set_timesteps(num_inference_steps=num_inference_steps)
         for t in tqdm(noise_scheduler.timesteps, ncols=110):
-            # Get controlnet output
             down_block_res_samples, mid_block_res_sample = controlnet(
                 x=latents, timesteps=torch.Tensor((t,)).to(device), controlnet_cond=controlnet_cond_vis
             )