diff --git a/models/maisi_ct_generative/configs/inference_trt.json b/models/maisi_ct_generative/configs/inference_trt.json index 90851711..f2a4a8f0 100644 --- 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)" }