Merge branch 'maisi' of https://github.com/Can-Zhao/tutorials into maisi

Author: Can-Zhao

generation/maisi/README.md

@@ -1,4 +1,4 @@
-# 🚨🚨🚨 THIS FOLDER IS DEPRECATED 🚨🚨🚨  
+# 🚨🚨🚨 THIS FOLDER IS DEPRECATED 🚨🚨🚨
 # 👉 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)

generation/maisi/maisi_inference_tutorial.ipynb

@@ -205,7 +205,7 @@
     "    os.makedirs(directory, exist_ok=True)\n",
     "root_dir = tempfile.mkdtemp() if directory is None else directory\n",
     "\n",
-    "download_model_data(maisi_version,root_dir)\n",
+    "download_model_data(maisi_version, root_dir)\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",

generation/maisi/scripts/download_model_data.py

@@ -17,12 +17,13 @@
 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
+    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
+    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"}
@@ -36,11 +37,11 @@ def fetch_to_hf_path_cmd(
     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
+    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_id = it["repo_id"]
         repo_file = it["filename"]
         dst = Path(it["path"])
         dst.parent.mkdir(parents=True, exist_ok=True)
@@ -51,11 +52,15 @@ def fetch_to_hf_path_cmd(
 
         # Build command (no shell=True; no quoting issues)
         cmd = [
-            "huggingface-cli", "download",
+            "huggingface-cli",
+            "download",
             repo_id,
-            "--include", repo_file,
-            "--revision", revision,
-            "--local-dir", str(root),
+            "--include",
+            repo_file,
+            "--revision",
+            revision,
+            "--local-dir",
+            str(root),
         ]
         # Run
         subprocess.run(cmd, check=True, env=env)
@@ -80,15 +85,14 @@ def fetch_to_hf_path_cmd(
     return saved
 
 
-
-def download_model_data(generate_version,root_dir, model_only=False):
+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",
+                "filename": "models/autoencoder_v1.pt",
             },
             {
                 "path": "models/mask_generation_autoencoder.pt",
@@ -99,7 +103,8 @@ def download_model_data(generate_version,root_dir, model_only=False):
                 "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 += [
                 {
@@ -124,10 +129,12 @@ def download_model_data(generate_version,root_dir, model_only=False):
                 "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}.")
+        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 += [
             {
@@ -139,7 +146,8 @@ def download_model_data(generate_version,root_dir, model_only=False):
                 "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 += [
                 {
@@ -159,7 +167,8 @@ def download_model_data(generate_version,root_dir, model_only=False):
                 "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 += [
                 {
@@ -168,11 +177,11 @@ def download_model_data(generate_version,root_dir, model_only=False):
                     "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")
+            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"])
@@ -191,7 +200,9 @@ def download_model_data(generate_version,root_dir, model_only=False):
         type=str,
         default="./",
     )
-    parser.add_argument("--model_only", dest="model_only", action="store_true", help="Download model only, not any dataset")
+    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)

generation/maisi/scripts/inference.py

@@ -89,7 +89,7 @@ def main():
     root_dir = tempfile.mkdtemp() if directory is None else directory
     print(root_dir)
 
-    download_model_data(maisi_version,root_dir)
+    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