From 397f35c3ca1a5f5f603f90d9865f33f2d469aaae Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Wed, 11 Dec 2024 20:13:43 -0800 Subject: [PATCH] TRT support for MAISI (#701) ### Description TRT optimization support for MAISI. Depends on https://github.com/Project-MONAI/MONAI/pull/8153 To successfully export MAISI, either 24.11 PyTorch container (or latest public Torch nightly) is needed, or this patch needs to be applied to 24.10-based container: ``` --- /usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset14.bak 2024-10-09 01:38:04.920316673 +0000 +++ /usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset14.py 2024-10-09 01:38:25.228053951 +0000 @@ -148,7 +148,6 @@ is_causal and symbolic_helper._is_none(attn_mask) ), "is_causal and attn_mask cannot be set at the same time" - scale = symbolic_helper._maybe_get_const(scale, "f") if symbolic_helper._is_none(scale): scale = _attention_scale(g, query) --------- Signed-off-by: Boris Fomitchev Signed-off-by: Yiheng Wang Co-authored-by: Yiheng Wang Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../configs/inference_trt.json | 19 +++++++++++++++++++ .../maisi_ct_generative/configs/metadata.json | 3 ++- 2 files changed, 21 insertions(+), 1 deletion(-) create mode 100644 models/maisi_ct_generative/configs/inference_trt.json diff --git a/models/maisi_ct_generative/configs/inference_trt.json b/models/maisi_ct_generative/configs/inference_trt.json new file mode 100644 index 00000000..90851711 --- /dev/null +++ b/models/maisi_ct_generative/configs/inference_trt.json @@ -0,0 +1,19 @@ +{ + "+imports": [ + "$from monai.networks import trt_compile" + ], + "c_trt_args": { + "export_args": { + "dynamo": "$False", + "report": "$True" + }, + "output_lists": [ + [ + -1 + ], + [] + ] + }, + "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)" +} diff --git a/models/maisi_ct_generative/configs/metadata.json b/models/maisi_ct_generative/configs/metadata.json index fd2a306c..010a70d1 100644 --- a/models/maisi_ct_generative/configs/metadata.json +++ b/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",