Skip to content

Commit 9f6d19f

Browse files
Can-Zhaopre-commit-ci[bot]KumoLiu
authored
Accelerated MAISI and make it compatible with previous DDPM (#1953)
Let MAISI tutorial work for more than one version name, and make it compatible with previous DDPM. Inference Tested with DDPM and RFlow. Training tested for DDPM and RFlow Some updates come from isort and black Fixes # . ### Description A few sentences describing the changes proposed in this pull request. ### Checks <!--- Put an `x` in all the boxes that apply, and remove the not applicable items --> - [ ] Avoid including large-size files in the PR. - [ ] Clean up long text outputs from code cells in the notebook. - [ ] For security purposes, please check the contents and remove any sensitive info such as user names and private key. - [ ] Ensure (1) hyperlinks and markdown anchors are working (2) use relative paths for tutorial repo files (3) put figure and graphs in the `./figure` folder - [x] Notebook runs automatically `./runner.sh -t <path to .ipynb file>` --------- Signed-off-by: Can-Zhao <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: YunLiu <[email protected]>
1 parent 77015e8 commit 9f6d19f

38 files changed

+1706
-1110
lines changed

‎generation/maisi/README.md

+69-29
Large diffs are not rendered by default.

‎generation/maisi/configs/config_infer.json

+6-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
256,
@@ -18,10 +18,11 @@
1818
2.0
1919
],
2020
"autoencoder_sliding_window_infer_size": [48,48,48],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
21+
"autoencoder_sliding_window_infer_overlap": 0.6666,
2222
"controlnet": "$@controlnet_def",
2323
"diffusion_unet": "$@diffusion_unet_def",
2424
"autoencoder": "$@autoencoder_def",
2525
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
26-
"mask_generation_diffusion": "$@mask_generation_diffusion_def"
26+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
27+
"modality": 1
2728
}

‎generation/maisi/configs/config_infer_16g_256x256x128.json

+10-4
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
256,
@@ -19,5 +19,11 @@
1919
],
2020
"autoencoder_sliding_window_infer_size": [96,96,96],
2121
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 2
22+
"autoencoder_tp_num_splits": 2,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_16g_256x256x256.json

+11-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
256,
@@ -18,6 +18,12 @@
1818
2.0
1919
],
2020
"autoencoder_sliding_window_infer_size": [48,48,64],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 2
21+
"autoencoder_sliding_window_infer_overlap": 0.6666,
22+
"autoencoder_tp_num_splits": 4,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_16g_512x512x128.json

+11-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
512,
@@ -18,6 +18,12 @@
1818
4.0
1919
],
2020
"autoencoder_sliding_window_infer_size": [64,64,32],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 1
21+
"autoencoder_sliding_window_infer_overlap": 0.5,
22+
"autoencoder_tp_num_splits": 2,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_24g_256x256x256.json

+10-4
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
256,
@@ -19,5 +19,11 @@
1919
],
2020
"autoencoder_sliding_window_infer_size": [64,64,64],
2121
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 4
22+
"autoencoder_tp_num_splits": 4,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_24g_512x512x128.json

+11-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
512,
@@ -18,6 +18,12 @@
1818
4.0
1919
],
2020
"autoencoder_sliding_window_infer_size": [80,80,32],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 1
21+
"autoencoder_sliding_window_infer_overlap": 0.4,
22+
"autoencoder_tp_num_splits": 2,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_24g_512x512x512.json

+11-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
512,
@@ -18,6 +18,12 @@
1818
1.0
1919
],
2020
"autoencoder_sliding_window_infer_size": [64,64,48],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 2
21+
"autoencoder_sliding_window_infer_overlap": 0.4,
22+
"autoencoder_tp_num_splits": 2,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_32g_512x512x512.json

+12-6
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
512,
@@ -17,7 +17,13 @@
1717
0.75,
1818
1.0
1919
],
20-
"autoencoder_sliding_window_infer_size": [64,64,64],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 2
20+
"autoencoder_sliding_window_infer_size": [80,80,48],
21+
"autoencoder_sliding_window_infer_overlap": 0.4,
22+
"autoencoder_tp_num_splits": 4,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_80g_512x512x128.json

-23
This file was deleted.

‎generation/maisi/configs/config_infer_80g_512x512x512.json

+11-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
3-
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
3+
"body_region": ["chest"],
4+
"anatomy_list": ["lung tumor"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
512,
@@ -18,6 +18,12 @@
1818
1.0
1919
],
2020
"autoencoder_sliding_window_infer_size": [80,80,80],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 2
21+
"autoencoder_sliding_window_infer_overlap": 0.4,
22+
"autoencoder_tp_num_splits": 2,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_infer_80g_512x512x768.json

+11-5
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
{
22
"num_output_samples": 1,
33
"body_region": ["abdomen"],
4-
"anatomy_list": ["liver","hepatic tumor"],
4+
"anatomy_list": ["liver"],
55
"controllable_anatomy_size": [],
6-
"num_inference_steps": 1000,
6+
"num_inference_steps": 30,
77
"mask_generation_num_inference_steps": 1000,
88
"output_size": [
99
512,
@@ -17,7 +17,13 @@
1717
0.75,
1818
0.66667
1919
],
20-
"autoencoder_sliding_window_infer_size": [80,80,112],
21-
"autoencoder_sliding_window_infer_overlap": 0.25,
22-
"autoencoder_tp_num_splits": 4
20+
"autoencoder_sliding_window_infer_size": [80,80,96],
21+
"autoencoder_sliding_window_infer_overlap": 0.4,
22+
"autoencoder_tp_num_splits": 4,
23+
"controlnet": "$@controlnet_def",
24+
"diffusion_unet": "$@diffusion_unet_def",
25+
"autoencoder": "$@autoencoder_def",
26+
"mask_generation_autoencoder": "$@mask_generation_autoencoder_def",
27+
"mask_generation_diffusion": "$@mask_generation_diffusion_def",
28+
"modality": 1
2329
}

‎generation/maisi/configs/config_maisi.json renamed to ‎generation/maisi/configs/config_maisi3d-ddpm.json

+3-2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
"spatial_dims": 3,
33
"image_channels": 1,
44
"latent_channels": 4,
5+
"include_body_region": true,
56
"mask_generation_latent_shape": [
67
4,
78
64,
@@ -60,8 +61,8 @@
6061
],
6162
"num_res_blocks": 2,
6263
"use_flash_attention": true,
63-
"include_top_region_index_input": true,
64-
"include_bottom_region_index_input": true,
64+
"include_top_region_index_input": "@include_body_region",
65+
"include_bottom_region_index_input": "@include_body_region",
6566
"include_spacing_input": true
6667
},
6768
"controlnet_def": {

0 commit comments

Comments
 (0)