Update MAISI model URL

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.
 

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,14 +104,14 @@
     "\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",
     "from scripts.sample import LDMSampler, check_input\n",
     "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)"
    ]
   },
   {

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)