diff --git a/generation/maisi/README.md b/generation/maisi/README.md index c9f4c117e..659b3c23a 100644 --- a/generation/maisi/README.md +++ b/generation/maisi/README.md @@ -1,3 +1,6 @@ +# 🚨🚨🚨 THIS FOLDER IS DEPRECATED (as of Oct 2025) 🚨🚨🚨 +# 👉 Please switch to: [https://github.com/NVIDIA-Medtech/NV-Generate-CTMR/tree/main](https://github.com/NVIDIA-Medtech/NV-Generate-CTMR/tree/main) + # Medical AI for Synthetic Imaging (MAISI) This example demonstrates the applications of training and validating NVIDIA MAISI, a 3D Latent Diffusion Model (LDM) capable of generating large CT images accompanied by corresponding segmentation masks. It supports variable volume size and voxel spacing and allows for the precise control of organ/tumor size. diff --git a/generation/maisi/maisi_inference_tutorial.ipynb b/generation/maisi/maisi_inference_tutorial.ipynb index 69ce91b9b..19283f2df 100644 --- a/generation/maisi/maisi_inference_tutorial.ipynb +++ b/generation/maisi/maisi_inference_tutorial.ipynb @@ -53,7 +53,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "67e2019e-1556-41a6-95e8-5d1a65f8b3a1", "metadata": { "scrolled": true @@ -104,7 +104,6 @@ "\n", "import monai\n", "import torch\n", - "from monai.apps import download_url\n", "from monai.config import print_config\n", "from monai.transforms import LoadImage, Orientation\n", "from monai.utils import set_determinism\n", @@ -112,6 +111,7 @@ "from scripts.utils import define_instance\n", "from scripts.utils_plot import find_label_center_loc, get_xyz_plot, show_image\n", "from scripts.diff_model_setting import setup_logging\n", + "from scripts.download_model_data import download_model_data\n", "\n", "print_config()\n", "\n", @@ -170,7 +170,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "e3c12dcc", "metadata": {}, "outputs": [ @@ -204,74 +204,7 @@ " os.makedirs(directory, exist_ok=True)\n", "root_dir = tempfile.mkdtemp() if directory is None else directory\n", "\n", - "# TODO: remove the `files` after the files are uploaded to the NGC\n", - "files = [\n", - " {\n", - " \"path\": \"models/autoencoder_epoch273.pt\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials\"\n", - " \"/model_zoo/model_maisi_autoencoder_epoch273_alternative.pt\",\n", - " },\n", - " {\n", - " \"path\": \"models/mask_generation_autoencoder.pt\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\" \"/tutorials/mask_generation_autoencoder.pt\",\n", - " },\n", - " {\n", - " \"path\": \"models/mask_generation_diffusion_unet.pt\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\"\n", - " \"/tutorials/model_zoo/model_maisi_mask_generation_diffusion_unet_v2.pt\",\n", - " },\n", - " {\n", - " \"path\": \"configs/all_anatomy_size_condtions.json\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/all_anatomy_size_condtions.json\",\n", - " },\n", - " {\n", - " \"path\": \"datasets/all_masks_flexible_size_and_spacing_4000.zip\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\"\n", - " \"/tutorials/all_masks_flexible_size_and_spacing_4000.zip\",\n", - " },\n", - "]\n", - "\n", - "if maisi_version == \"maisi3d-ddpm\":\n", - " files += [\n", - " {\n", - " \"path\": \"models/diff_unet_3d_ddpm.pt\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/model_zoo\"\n", - " \"/model_maisi_input_unet3d_data-all_steps1000size512ddpm_random_current_inputx_v1_alternative.pt\",\n", - " },\n", - " {\n", - " \"path\": \"models/controlnet_3d_ddpm.pt\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/model_zoo\"\n", - " \"/model_maisi_controlnet-20datasets-e20wl100fold0bc_noi_dia_fsize_current_alternative.pt\",\n", - " },\n", - " {\n", - " \"path\": \"configs/candidate_masks_flexible_size_and_spacing_3000.json\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\"\n", - " \"/tutorials/candidate_masks_flexible_size_and_spacing_3000.json\",\n", - " },\n", - " ]\n", - "elif maisi_version == \"maisi3d-rflow\":\n", - " files += [\n", - " {\n", - " \"path\": \"models/diff_unet_3d_rflow.pt\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/\"\n", - " \"diff_unet_ckpt_rflow_epoch19350.pt\",\n", - " },\n", - " {\n", - " \"path\": \"models/controlnet_3d_rflow.pt\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/controlnet_rflow_epoch60.pt\",\n", - " },\n", - " {\n", - " \"path\": \"configs/candidate_masks_flexible_size_and_spacing_4000.json\",\n", - " \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\"\n", - " \"/tutorials/candidate_masks_flexible_size_and_spacing_4000.json\",\n", - " },\n", - " ]\n", - "else:\n", - " raise ValueError(f\"maisi_version has to be chosen from ['maisi3d-ddpm', 'maisi3d-rflow'], yet got {maisi_version}.\")\n", - "\n", - "for file in files:\n", - " file[\"path\"] = file[\"path\"] if \"datasets/\" not in file[\"path\"] else os.path.join(root_dir, file[\"path\"])\n", - " download_url(url=file[\"url\"], filepath=file[\"path\"])" + "download_model_data(maisi_version, root_dir)" ] }, { diff --git a/generation/maisi/scripts/download_model_data.py b/generation/maisi/scripts/download_model_data.py new file mode 100644 index 000000000..bf39b4956 --- /dev/null +++ b/generation/maisi/scripts/download_model_data.py @@ -0,0 +1,208 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os, subprocess, shutil +import argparse +from tqdm.auto import tqdm +from monai.apps import download_url +from pathlib import Path +from huggingface_hub import snapshot_download +from typing import List, Dict, Optional + + +def fetch_to_hf_path_cmd( + items: List[Dict[str, str]], + root_dir: str = "./", # staging dir for CLI output + revision: str = "main", + overwrite: bool = False, + token: Optional[str] = None, # or rely on env HUGGINGFACE_HUB_TOKEN +) -> list[str]: + """ + items: list of {"repo_id": "...", "filename": "path/in/repo.ext", "path": "local/target.ext"} + Returns list of saved local paths (in the same order as items). + """ + saved = [] + root = Path(root_dir) + root.mkdir(parents=True, exist_ok=True) + + # Env for subprocess; keep Rust fast-path off to avoid notebook progress quirks + env = os.environ.copy() + if token: + env["HUGGINGFACE_HUB_TOKEN"] = token + env.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "0") # safer in Jupyter + env.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "0") # show CLI progress in terminal + + for it in items: + repo_id = it["repo_id"] + repo_file = it["filename"] + dst = Path(it["path"]) + dst.parent.mkdir(parents=True, exist_ok=True) + + if dst.exists() and not overwrite: + saved.append(str(dst)) + continue + + # Build command (no shell=True; no quoting issues) + cmd = [ + "huggingface-cli", + "download", + repo_id, + "--include", + repo_file, + "--revision", + revision, + "--local-dir", + str(root), + ] + # Run + subprocess.run(cmd, check=True, env=env) + + # Source path where CLI placed the file + src = root / repo_file + if not src.exists(): + raise FileNotFoundError( + f"Expected downloaded file missing: {src}\n" + f"Tip: authenticate (`huggingface-cli login` or pass token=...)," + f" and avoid shared-IP 429s." + ) + + # Move to desired target + if dst.exists() and overwrite: + dst.unlink() + if src.resolve() != dst.resolve(): + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(src), str(dst)) + saved.append(str(dst)) + + return saved + + +def download_model_data(generate_version, root_dir, model_only=False): + # TODO: remove the `files` after the files are uploaded to the NGC + if generate_version == "ddpm-ct" or generate_version == "rflow-ct": + files = [ + { + "path": "models/autoencoder_v1.pt", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "models/autoencoder_v1.pt", + }, + { + "path": "models/mask_generation_autoencoder.pt", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "models/mask_generation_autoencoder.pt", + }, + { + "path": "models/mask_generation_diffusion_unet.pt", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "models/mask_generation_diffusion_unet.pt", + }, + ] + if not model_only: + files += [ + { + "path": "datasets/all_anatomy_size_conditions.json", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "datasets/all_anatomy_size_conditions.json", + }, + { + "path": "datasets/all_masks_flexible_size_and_spacing_4000.zip", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "datasets/all_masks_flexible_size_and_spacing_4000.zip", + }, + ] + elif generate_version == "rflow-mr": + files = [ + { + "path": "models/autoencoder_v2.pt", + "repo_id": "nvidia/NV-Generate-MR", + "filename": "models/autoencoder_v2.pt", + }, + { + "path": "models/diff_unet_3d_rflow-mr.pt", + "repo_id": "nvidia/NV-Generate-MR", + "filename": "models/diff_unet_3d_rflow-mr.pt", + }, + ] + else: + raise ValueError( + f"generate_version has to be chosen from ['ddpm-ct', 'rflow-ct', 'rflow-mr'], yet got {generate_version}." + ) + if generate_version == "ddpm-ct": + files += [ + { + "path": "models/diff_unet_3d_ddpm-ct.pt", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "models/diff_unet_3d_ddpm-ct.pt", + }, + { + "path": "models/controlnet_3d_ddpm-ct.pt", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "models/controlnet_3d_ddpm-ct.pt", + }, + ] + if not model_only: + files += [ + { + "path": "datasets/candidate_masks_flexible_size_and_spacing_3000.json", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "datasets/candidate_masks_flexible_size_and_spacing_3000.json", + }, + ] + elif generate_version == "rflow-ct": + files += [ + { + "path": "models/diff_unet_3d_rflow-ct.pt", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "models/diff_unet_3d_rflow-ct.pt", + }, + { + "path": "models/controlnet_3d_rflow-ct.pt", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "models/controlnet_3d_rflow-ct.pt", + }, + ] + if not model_only: + files += [ + { + "path": "datasets/candidate_masks_flexible_size_and_spacing_4000.json", + "repo_id": "nvidia/NV-Generate-CT", + "filename": "datasets/candidate_masks_flexible_size_and_spacing_4000.json", + }, + ] + + for file in files: + file["path"] = file["path"] if "datasets/" not in file["path"] else os.path.join(root_dir, file["path"]) + if "repo_id" in file.keys(): + path = fetch_to_hf_path_cmd([file], root_dir=root_dir, revision="main") + print("saved to:", path) + else: + download_url(url=file["url"], filepath=file["path"]) + return + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Model downloading") + parser.add_argument( + "--version", + type=str, + default="rflow-ct", + ) + parser.add_argument( + "--root_dir", + type=str, + default="./", + ) + parser.add_argument( + "--model_only", dest="model_only", action="store_true", help="Download model only, not any dataset" + ) + + args = parser.parse_args() + download_model_data(args.version, args.root_dir, args.model_only) diff --git a/generation/maisi/scripts/inference.py b/generation/maisi/scripts/inference.py index 3f81f9c49..1eeef0c39 100644 --- a/generation/maisi/scripts/inference.py +++ b/generation/maisi/scripts/inference.py @@ -19,14 +19,12 @@ import monai import torch -from monai.apps import download_url -from monai.config import print_config -from monai.transforms import LoadImage, Orientation from monai.utils import set_determinism from scripts.sample import LDMSampler, check_input from scripts.utils import define_instance from scripts.utils_plot import find_label_center_loc, get_xyz_plot, show_image +from scripts.download_model_data import download_model_data def main(): @@ -88,78 +86,7 @@ def main(): root_dir = tempfile.mkdtemp() if directory is None else directory print(root_dir) - # TODO: remove the `files` after the files are uploaded to the NGC - files = [ - { - "path": "models/autoencoder_epoch273.pt", - "url": "https://developer.download.nvidia.com/assets/Clara/monai/tutorials" - "/model_zoo/model_maisi_autoencoder_epoch273_alternative.pt", - }, - { - "path": "models/mask_generation_autoencoder.pt", - "url": "https://developer.download.nvidia.com/assets/Clara/monai" - "/tutorials/mask_generation_autoencoder.pt", - }, - { - "path": "models/mask_generation_diffusion_unet.pt", - "url": "https://developer.download.nvidia.com/assets/Clara/monai" - "/tutorials/model_zoo/model_maisi_mask_generation_diffusion_unet_v2.pt", - }, - { - "path": "configs/all_anatomy_size_condtions.json", - "url": "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/all_anatomy_size_condtions.json", - }, - { - "path": "datasets/all_masks_flexible_size_and_spacing_4000.zip", - "url": "https://developer.download.nvidia.com/assets/Clara/monai" - "/tutorials/all_masks_flexible_size_and_spacing_4000.zip", - }, - ] - - if maisi_version == "maisi3d-ddpm": - files += [ - { - "path": "models/diff_unet_3d_ddpm.pt", - "url": "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/model_zoo" - "/model_maisi_input_unet3d_data-all_steps1000size512ddpm_random_current_inputx_v1_alternative.pt", - }, - { - "path": "models/controlnet_3d_ddpm.pt", - "url": "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/model_zoo" - "/model_maisi_controlnet-20datasets-e20wl100fold0bc_noi_dia_fsize_current_alternative.pt", - }, - { - "path": "configs/candidate_masks_flexible_size_and_spacing_3000.json", - "url": "https://developer.download.nvidia.com/assets/Clara/monai" - "/tutorials/candidate_masks_flexible_size_and_spacing_3000.json", - }, - ] - elif maisi_version == "maisi3d-rflow": - files += [ - { - "path": "models/diff_unet_3d_rflow.pt", - "url": "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/" - "diff_unet_ckpt_rflow_epoch19350.pt", - }, - { - "path": "models/controlnet_3d_rflow.pt", - "url": "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/" - "controlnet_rflow_epoch60.pt", - }, - { - "path": "configs/candidate_masks_flexible_size_and_spacing_4000.json", - "url": "https://developer.download.nvidia.com/assets/Clara/monai" - "/tutorials/candidate_masks_flexible_size_and_spacing_4000.json", - }, - ] - else: - raise ValueError( - f"maisi_version has to be chosen from ['maisi3d-ddpm', 'maisi3d-rflow'], yet got {maisi_version}." - ) - - for file in files: - file["path"] = file["path"] if "datasets/" not in file["path"] else os.path.join(root_dir, file["path"]) - download_url(url=file["url"], filepath=file["path"]) + download_model_data(maisi_version, root_dir) # ## Read in environment setting, including data directory, model directory, and output directory # The information for data directory, model directory, and output directory are saved in ./configs/environment.json