update trt convert files

Signed-off-by: binliu <binliu@nvidia.com>
Project-MONAI · Feb 19, 2025 · 18d7bb6 · 18d7bb6
1 parent 9c70853
commit 18d7bb6
Showing 1 changed file with 32 additions and 2 deletions.
diff --git a/models/maisi_ct_generative/configs/inference_trt.json b/models/maisi_ct_generative/configs/inference_trt.json
@@ -5,15 +5,45 @@
     "c_trt_args": {
         "export_args": {
             "dynamo": "$False",
-            "report": "$True"
+            "do_constant_folding": "$False"
         },
         "output_lists": [
             [
                 -1
             ],
             []
+        ],
+        "input_profiles": [
+            {
+                "x": [[1, 4, 64, 64, 32], [1, 4, 128, 128, 128], [1, 4, 128, 128, 128]],
+                "controlnet_cond": [[1, 8, 256, 256, 128], [1, 8, 512, 512, 512], [1, 8, 512, 512, 512]]
+            }
+        ]
+    },
+    "d_trt_args":{
+        "export_args": {
+            "dynamo": "$False",
+            "do_constant_folding": "$False"
+        },
+        "input_profiles": [
+            {
+                "x": [[1, 4, 64, 64, 32], [1, 4, 128, 128, 128], [1, 4, 128, 128, 128]],
+                "mid_block_additional_residual": [[1, 512, 8, 8, 4], [1, 512, 16, 16, 16], [1, 512, 16, 16, 16]],
+                "down_block_additional_residuals_0": [[1, 64, 64, 64, 32], [1, 64, 128, 128, 128], [1, 64, 128, 128, 128]],
+                "down_block_additional_residuals_1": [[1, 64, 64, 64, 32], [1, 64, 128, 128, 128], [1, 64, 128, 128, 128]],
+                "down_block_additional_residuals_2": [[1, 64, 64, 64, 32], [1, 64, 128, 128, 128], [1, 64, 128, 128, 128]],
+                "down_block_additional_residuals_3": [[1, 64, 32, 32, 16], [1, 64, 64, 64, 64], [1, 64, 64, 64, 64]],
+                "down_block_additional_residuals_4": [[1, 128, 32, 32, 16], [1, 128, 64, 64, 64], [1, 128, 64, 64, 64]],
+                "down_block_additional_residuals_5": [[1, 128, 32, 32, 16], [1, 128, 64, 64, 64], [1, 128, 64, 64, 64]],
+                "down_block_additional_residuals_6": [[1, 128, 16, 16, 8], [1, 128, 32, 32, 32], [1, 128, 32, 32, 32]],
+                "down_block_additional_residuals_7": [[1, 256, 16, 16, 8], [1, 256, 32, 32, 32], [1, 256, 32, 32, 32]],
+                "down_block_additional_residuals_8": [[1, 256, 16, 16, 8], [1, 256, 32, 32, 32], [1, 256, 32, 32, 32]],
+                "down_block_additional_residuals_9": [[1, 256, 8, 8, 4], [1, 256, 16, 16, 16], [1, 256, 16, 16, 16]],
+                "down_block_additional_residuals_10": [[1, 512, 8, 8, 4], [1, 512, 16, 16, 16], [1, 512, 16, 16, 16]],
+                "down_block_additional_residuals_11": [[1, 512, 8, 8, 4], [1, 512, 16, 16, 16], [1, 512, 16, 16, 16]]
+            }
         ]
     },
     "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)"
+    "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path, @d_trt_args)"
 }