From 93a4dc522d849a104da53c5dc6e676a20be4a87b Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Wed, 28 Aug 2024 01:38:49 -0700 Subject: [PATCH 01/11] Added trt_compile configs for vista2d and vista3d Signed-off-by: Boris Fomitchev --- models/vista2d/configs/inference.json | 2 +- models/vista2d/configs/inference_trt.json | 10 ++++++++++ models/vista2d/docs/README.md | 8 +++++++- models/vista3d/configs/inference_trt.json | 9 +++++++++ models/vista3d/docs/README.md | 7 +++++++ 5 files changed, 34 insertions(+), 2 deletions(-) create mode 100644 models/vista2d/configs/inference_trt.json create mode 100644 models/vista3d/configs/inference_trt.json diff --git a/models/vista2d/configs/inference.json b/models/vista2d/configs/inference.json index f56a5c47..2439a4a5 100644 --- a/models/vista2d/configs/inference.json +++ b/models/vista2d/configs/inference.json @@ -11,7 +11,7 @@ 256, 256 ], - "input_dict": "${'image': '/home/venn/Desktop/data/medical/cellpose_dataset/test/001_img.png'}", + "input_dict": "${'image': '/cellpose_dataset/test/001_img.png'}", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", "sam_ckpt_path": "$@ckpt_dir + '/sam_vit_b_01ec64.pth'", "pretrained_ckpt_path": "$@ckpt_dir + '/model.pt'", diff --git a/models/vista2d/configs/inference_trt.json b/models/vista2d/configs/inference_trt.json new file mode 100644 index 00000000..17e9ca80 --- /dev/null +++ b/models/vista2d/configs/inference_trt.json @@ -0,0 +1,10 @@ +{ + "imports": [ + "$import numpy", + "$from monai.networks import trt_compile" + ], + "trt_args": { + "dynamic_batchsize": "$[1, @inferer#sw_batch_size, @inferer#sw_batch_size]" + }, + "network": "$trt_compile(@network_def.to(@device), @pretrained_ckpt_path, args=@trt_args)" +} diff --git a/models/vista2d/docs/README.md b/models/vista2d/docs/README.md index c2f221be..a3a902e4 100644 --- a/models/vista2d/docs/README.md +++ b/models/vista2d/docs/README.md @@ -66,7 +66,13 @@ torchrun --nproc_per_node=gpu -m monai.bundle run_workflow "scripts.workflow.Vis python -m monai.bundle run --config_file configs/inference.json ``` -Please note that the data used in the config file is: "/cellpose_dataset/test/001_img.png", if the dataset path is different or you want to do inference on another file, please modify in `configs/inference.json` accordingly. +Please note that the data used in this config file is: "/cellpose_dataset/test/001_img.png", if the dataset path is different or you want to do inference on another file, please modify in `configs/inference.json` accordingly. + +#### Execute inference with the TensorRT model: + +``` +python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']" +``` ### Execute multi-GPU inference ```bash diff --git a/models/vista3d/configs/inference_trt.json b/models/vista3d/configs/inference_trt.json new file mode 100644 index 00000000..9c3d52dc --- /dev/null +++ b/models/vista3d/configs/inference_trt.json @@ -0,0 +1,9 @@ +{ + "+imports": [ + "$from monai.networks import trt_compile" + ], + "trt_args": { + "dynamic_batchsize": "$[1, @inferer#sw_batch_size, @inferer#sw_batch_size]" + }, + "network": "$trt_compile(@network_def.to(@device), @bundle_root + '/models/model.pt', args=@trt_args, submodule=['image_encoder.encoder', 'class_head'])" +} diff --git a/models/vista3d/docs/README.md b/models/vista3d/docs/README.md index 966dba99..8fc53d7a 100644 --- a/models/vista3d/docs/README.md +++ b/models/vista3d/docs/README.md @@ -172,6 +172,13 @@ This default is overridable by changing the input folder `input_dir`, or the inp Set `"postprocessing#transforms#0#_disabled_": false` to move the postprocessing to cpu to reduce the GPU memory footprint. +#### Execute inference with the TensorRT model: + +``` +python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']" +``` + + ## Automatic segmentation label prompts : The mapping between organ name and label prompt is in the [json file](labels.json) From 89dfac19ea0627d00705b7273245e949e1151956 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Thu, 5 Sep 2024 17:04:35 -0700 Subject: [PATCH 02/11] Stash Signed-off-by: Boris Fomitchev --- models/maisi_ct_generative/configs/inference.json | 4 ++-- models/vista3d/configs/inference.json | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/models/maisi_ct_generative/configs/inference.json b/models/maisi_ct_generative/configs/inference.json index 2ad6aa08..2cb5bef1 100644 --- a/models/maisi_ct_generative/configs/inference.json +++ b/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", diff --git a/models/vista3d/configs/inference.json b/models/vista3d/configs/inference.json index e374ff57..52f1d701 100644 --- a/models/vista3d/configs/inference.json +++ b/models/vista3d/configs/inference.json @@ -15,7 +15,7 @@ "output_dtype": "$np.float32", "output_postfix": "trans", "separate_folder": true, - "input_dict": "${'image': '/home/venn/Desktop/data/medical/Task09_Spleen/imagesTr/spleen_10.nii.gz', 'label_prompt': [3]}", + "input_dict": "${'image': '/data/Task09_Spleen/imagesTr/spleen_10.nii.gz', 'label_prompt': ['spleen'], 'point_labels': [1], 'points': [(152, 170, 87)]}", "everything_labels": "$list(set([i+1 for i in range(132)]) - set([2,16,18,20,21,23,24,25,26,27,128,129,130,131,132]))", "metadata_path": "$@bundle_root + '/configs/metadata.json'", "metadata": "$json.loads(pathlib.Path(@metadata_path).read_text())", From b6628fc0e72532f91000f21ad4fba385bf48934e Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Fri, 11 Oct 2024 23:15:05 -0700 Subject: [PATCH 03/11] Working MAISI Signed-off-by: Boris Fomitchev --- models/maisi_ct_generative/configs/inference_trt.json | 9 +++++++++ models/maisi_ct_generative/scripts/sample.py | 4 ++-- 2 files changed, 11 insertions(+), 2 deletions(-) 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..5a9d4524 --- /dev/null +++ b/models/maisi_ct_generative/configs/inference_trt.json @@ -0,0 +1,9 @@ +{ + "+imports": [ + "$from monai.networks import trt_compile" + ], + "torch_trt_args" : {"method" : "torch_trt", "export_args": { "dynamo" : "$False", "verbose" : "$True" } }, + "trt_args" : {"export_args": { "dynamo" : "$False", "verbose" : "$True" } }, + "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path)", + "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path, args=@trt_args)" +} diff --git a/models/maisi_ct_generative/scripts/sample.py b/models/maisi_ct_generative/scripts/sample.py index d161597f..7f27cadd 100644 --- a/models/maisi_ct_generative/scripts/sample.py +++ b/models/maisi_ct_generative/scripts/sample.py @@ -241,10 +241,10 @@ 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( + controlnet_ret = controlnet( x=latents, timesteps=torch.Tensor((t,)).to(device), controlnet_cond=controlnet_cond_vis ) + down_block_res_samples, mid_block_res_sample = controlnet_ret[:-1], controlnet_ret[-1] latent_model_input = latents noise_pred = diffusion_unet( x=latent_model_input, From e91ad5721cf87763cd2f8fd49b4704988b91c390 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Tue, 15 Oct 2024 17:19:16 -0700 Subject: [PATCH 04/11] Adding TRT support Signed-off-by: Boris Fomitchev --- models/maisi_ct_generative/configs/inference_trt.json | 5 ++++- models/maisi_ct_generative/scripts/sample.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/models/maisi_ct_generative/configs/inference_trt.json b/models/maisi_ct_generative/configs/inference_trt.json index 5a9d4524..043ef1d9 100644 --- a/models/maisi_ct_generative/configs/inference_trt.json +++ b/models/maisi_ct_generative/configs/inference_trt.json @@ -3,7 +3,10 @@ "$from monai.networks import trt_compile" ], "torch_trt_args" : {"method" : "torch_trt", "export_args": { "dynamo" : "$False", "verbose" : "$True" } }, - "trt_args" : {"export_args": { "dynamo" : "$False", "verbose" : "$True" } }, + "trt_args" : { + "export_args": { "dynamo" : "$False", "verbose" : "$False" }, + "dynamic_batchsize": "$[1, 1, 4]" + }, "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path)", "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path, args=@trt_args)" } diff --git a/models/maisi_ct_generative/scripts/sample.py b/models/maisi_ct_generative/scripts/sample.py index 7f27cadd..eaa5aa5a 100644 --- a/models/maisi_ct_generative/scripts/sample.py +++ b/models/maisi_ct_generative/scripts/sample.py @@ -244,7 +244,10 @@ def ldm_conditional_sample_one_image( controlnet_ret = controlnet( x=latents, timesteps=torch.Tensor((t,)).to(device), controlnet_cond=controlnet_cond_vis ) - down_block_res_samples, mid_block_res_sample = controlnet_ret[:-1], controlnet_ret[-1] + if len(controlnet_ret==2): + down_block_res_samples, mid_block_res_sample = controlnet_ret + else: + down_block_res_samples, mid_block_res_sample = controlnet_ret[:-1], controlnet_ret[-1] latent_model_input = latents noise_pred = diffusion_unet( x=latent_model_input, From b336e069cdaf26d080e2a08a972e31b6f94c86c8 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Tue, 15 Oct 2024 17:29:57 -0700 Subject: [PATCH 05/11] cleanup Signed-off-by: Boris Fomitchev --- models/maisi_ct_generative/configs/inference_trt.json | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/models/maisi_ct_generative/configs/inference_trt.json b/models/maisi_ct_generative/configs/inference_trt.json index 043ef1d9..15eba2ff 100644 --- a/models/maisi_ct_generative/configs/inference_trt.json +++ b/models/maisi_ct_generative/configs/inference_trt.json @@ -2,11 +2,6 @@ "+imports": [ "$from monai.networks import trt_compile" ], - "torch_trt_args" : {"method" : "torch_trt", "export_args": { "dynamo" : "$False", "verbose" : "$True" } }, - "trt_args" : { - "export_args": { "dynamo" : "$False", "verbose" : "$False" }, - "dynamic_batchsize": "$[1, 1, 4]" - }, "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path)", - "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path, args=@trt_args)" + "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path)" } From 97d2b0e99653128c24966c9b6d06dc3590ce2d87 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Tue, 15 Oct 2024 19:18:07 -0700 Subject: [PATCH 06/11] fixing condition Signed-off-by: Boris Fomitchev --- models/maisi_ct_generative/scripts/sample.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/maisi_ct_generative/scripts/sample.py b/models/maisi_ct_generative/scripts/sample.py index eaa5aa5a..f2b054dd 100644 --- a/models/maisi_ct_generative/scripts/sample.py +++ b/models/maisi_ct_generative/scripts/sample.py @@ -244,7 +244,7 @@ def ldm_conditional_sample_one_image( controlnet_ret = controlnet( x=latents, timesteps=torch.Tensor((t,)).to(device), controlnet_cond=controlnet_cond_vis ) - if len(controlnet_ret==2): + if len(controlnet_ret) == 2: down_block_res_samples, mid_block_res_sample = controlnet_ret else: down_block_res_samples, mid_block_res_sample = controlnet_ret[:-1], controlnet_ret[-1] From 60000dffc239129fe1d0957bc380a11d6045f267 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Mon, 21 Oct 2024 15:15:12 -0700 Subject: [PATCH 07/11] Added output_lists option Signed-off-by: Boris Fomitchev --- models/maisi_ct_generative/configs/inference_trt.json | 3 ++- models/maisi_ct_generative/scripts/sample.py | 6 +----- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/models/maisi_ct_generative/configs/inference_trt.json b/models/maisi_ct_generative/configs/inference_trt.json index 15eba2ff..04ec7130 100644 --- a/models/maisi_ct_generative/configs/inference_trt.json +++ b/models/maisi_ct_generative/configs/inference_trt.json @@ -2,6 +2,7 @@ "+imports": [ "$from monai.networks import trt_compile" ], - "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path)", + "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)" } diff --git a/models/maisi_ct_generative/scripts/sample.py b/models/maisi_ct_generative/scripts/sample.py index f2b054dd..42913dd0 100644 --- a/models/maisi_ct_generative/scripts/sample.py +++ b/models/maisi_ct_generative/scripts/sample.py @@ -241,13 +241,9 @@ 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): - controlnet_ret = controlnet( + down_block_res_samples, mid_block_res_sample = controlnet( x=latents, timesteps=torch.Tensor((t,)).to(device), controlnet_cond=controlnet_cond_vis ) - if len(controlnet_ret) == 2: - down_block_res_samples, mid_block_res_sample = controlnet_ret - else: - down_block_res_samples, mid_block_res_sample = controlnet_ret[:-1], controlnet_ret[-1] latent_model_input = latents noise_pred = diffusion_unet( x=latent_model_input, From ebb50bcf0ffa03f59ac83eadd761f613eff9892a Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Tue, 22 Oct 2024 03:16:22 +0000 Subject: [PATCH 08/11] update pre commit config Signed-off-by: Yiheng Wang --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 145907a4..195a6192 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ default_language_version: - python: python3.8 + python: python ci: autofix_prs: true From 74a6ecc5daeeee0c5eee25c5818bed85f35ad744 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 22 Oct 2024 03:16:33 +0000 Subject: [PATCH 09/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../maisi_ct_generative/configs/inference_trt.json | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/models/maisi_ct_generative/configs/inference_trt.json b/models/maisi_ct_generative/configs/inference_trt.json index 04ec7130..10eacce4 100644 --- a/models/maisi_ct_generative/configs/inference_trt.json +++ b/models/maisi_ct_generative/configs/inference_trt.json @@ -1,8 +1,17 @@ { "+imports": [ - "$from monai.networks import trt_compile" + "$from monai.networks import trt_compile" ], - "c_trt_args" : { "output_lists" : [[-1],[0]] }, + "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)" } From a9a8b03e7bdc209d781b5e48a4b76ea41985b422 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Tue, 22 Oct 2024 03:18:18 +0000 Subject: [PATCH 10/11] update metadata Signed-off-by: Yiheng Wang --- .pre-commit-config.yaml | 2 +- models/maisi_ct_generative/configs/metadata.json | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 195a6192..21219e02 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ default_language_version: - python: python + python: python3 ci: autofix_prs: true 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", From e827f5cc3b6197bd184643353d3f186ad0eede20 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Tue, 22 Oct 2024 12:43:11 -0700 Subject: [PATCH 11/11] Addressing code review comments Signed-off-by: Boris Fomitchev --- models/maisi_ct_generative/configs/inference.json | 4 ++-- models/maisi_ct_generative/scripts/sample.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/models/maisi_ct_generative/configs/inference.json b/models/maisi_ct_generative/configs/inference.json index cd8d9cfd..283305ed 100644 --- a/models/maisi_ct_generative/configs/inference.json +++ b/models/maisi_ct_generative/configs/inference.json @@ -34,8 +34,8 @@ "spatial_dims": 3, "image_channels": 1, "latent_channels": 4, - "output_size_xy": 256, - "output_size_z": 256, + "output_size_xy": 512, + "output_size_z": 512, "output_size": [ "@output_size_xy", "@output_size_xy", diff --git a/models/maisi_ct_generative/scripts/sample.py b/models/maisi_ct_generative/scripts/sample.py index 42913dd0..d161597f 100644 --- a/models/maisi_ct_generative/scripts/sample.py +++ b/models/maisi_ct_generative/scripts/sample.py @@ -241,6 +241,7 @@ 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 )