Skip to content

Commit 397f35c

Browse files
borisfomyiheng-wang-nvpre-commit-ci[bot]
authored
TRT support for MAISI (#701)
### Description TRT optimization support for MAISI. Depends on Project-MONAI/MONAI#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 <bfomitchev@nvidia.com> Signed-off-by: Yiheng Wang <vennw@nvidia.com> Co-authored-by: Yiheng Wang <vennw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent d486306 commit 397f35c

File tree

2 files changed

+21
-1
lines changed

2 files changed

+21
-1
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
{
2+
"+imports": [
3+
"$from monai.networks import trt_compile"
4+
],
5+
"c_trt_args": {
6+
"export_args": {
7+
"dynamo": "$False",
8+
"report": "$True"
9+
},
10+
"output_lists": [
11+
[
12+
-1
13+
],
14+
[]
15+
]
16+
},
17+
"controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path, @c_trt_args)",
18+
"diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path)"
19+
}

models/maisi_ct_generative/configs/metadata.json

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
{
22
"schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20240318.json",
3-
"version": "0.4.5",
3+
"version": "0.4.6",
44
"changelog": {
5+
"0.4.6": "add TensorRT support",
56
"0.4.5": "update README",
67
"0.4.4": "update issue for IgniteInfo",
78
"0.4.3": "remove download large files, add weights_only when loading weights and add label_dict to large files",

0 commit comments

Comments
 (0)