Spaces:
Runtime error
Runtime error
Claude Opus 5 (1M context)
Stream volumes by URL, cache checkpoints, reclaim archive space, quiet logs
f265752 | """ | |
| Shared driver for the upstream NV-Generate-CTMR scripts. | |
| Strategy: the upstream code is structured around argparse + global filesystem layout | |
| (reads relative-path configs, writes to a relative output_dir). Rather than refactor | |
| its internals, we treat it as an external tool: chdir into the upstream root, write | |
| modified config copies that override the user-controlled fields, ensure weights are | |
| downloaded, then call the upstream entry-point function. We then return the path of | |
| the most recently produced NIfTI from the configured output dir. | |
| """ | |
| from __future__ import annotations | |
| import contextlib | |
| import importlib | |
| from collections import OrderedDict | |
| import json | |
| import os | |
| import sys | |
| import time | |
| import uuid | |
| from pathlib import Path | |
| from typing import Iterable, Optional | |
| from utils.timing import Stopwatch, log, phase | |
| ROOT = Path(__file__).resolve().parent.parent | |
| UPSTREAM = ROOT / "repos" / "NV-Generate-CTMR" | |
| GENERATED_OUTPUT = UPSTREAM / "output" | |
| # Upstream `scripts.inference` reads MONAI_DATA_DIRECTORY and passes it to | |
| # download_model_data() as root_dir, which is where it expects the `datasets/` | |
| # payload to live. We prefetch to the same place so it finds the files already | |
| # there instead of downloading a second copy. | |
| DATA_DIR = UPSTREAM / "temp_work_dir" | |
| MASK_ARCHIVE = DATA_DIR / "datasets" / "all_masks_flexible_size_and_spacing_4000.zip" | |
| ALL_VERSIONS = ("rflow-ct", "rflow-mr", "rflow-mr-brain") | |
| # Once the mask archive is extracted, nothing reads the .zip again -- but | |
| # upstream's download_model_data() re-fetches any file whose destination is | |
| # missing, and that is an 11.6 GB round trip. Truncating it to a placeholder | |
| # keeps the path present for that existence check while giving back the space. | |
| # prefetch_weights() repairs the state if the extracted directory ever goes | |
| # away. Set NV_GENERATE_PRUNE_ARCHIVE=0 to keep the full archive on disk. | |
| PRUNE_ARCHIVE = os.environ.get("NV_GENERATE_PRUNE_ARCHIVE", "1") == "1" | |
| # Cap for the checkpoint cache below, in GiB of host RAM. | |
| CHECKPOINT_CACHE_GIB = float(os.environ.get("NV_GENERATE_CKPT_CACHE_GIB", "8")) | |
| class _LinkingShutil: | |
| """Drop-in for the `shutil` module that hard-links instead of copying. | |
| Upstream's fetch_to_hf_path_cmd() calls hf_hub_download() (which stores the | |
| blob in the HF cache) and then shutil.copy2(cached, dst) — so every file is | |
| written to disk twice. For CT that is 11.6 GB of anatomy masks duplicated. | |
| A hard link produces the same destination path for zero extra bytes. | |
| Moving the blob out of the cache instead is not safe: hf_hub_download keeps | |
| a blobs/ + snapshots/ layout and removing the blob both corrupts the cache | |
| and forces a full re-download on the next call. A link keeps one inode with | |
| two names, so either path can be removed independently. | |
| Falls back to a symlink, then to a real copy, if the filesystem refuses | |
| (e.g. the cache and the app are on different mounts). | |
| """ | |
| def __init__(self, real): | |
| self._real = real | |
| def __getattr__(self, name): | |
| return getattr(self._real, name) | |
| def copy2(self, src, dst, *args, **kwargs): | |
| # Resolve first. hf_hub_download returns the snapshot entry, which is a | |
| # *relative* symlink into blobs/ (../../blobs/<sha>). os.link() on a | |
| # symlink source copies the link itself, target string and all, so | |
| # linking it into <upstream>/models/ would produce a symlink resolving | |
| # to <upstream>/../blobs/<sha> — which does not exist. Link the blob. | |
| src_real = os.path.realpath(os.fspath(src)) | |
| dst_s = os.fspath(dst) | |
| # Clear any stale entry. A dangling symlink is invisible to exists() but | |
| # makes link()/symlink() fail EEXIST and makes copy2 write *through* it. | |
| if os.path.lexists(dst_s): | |
| try: | |
| os.unlink(dst_s) | |
| except OSError: | |
| pass | |
| for attempt in ( | |
| lambda: os.link(src_real, dst_s), | |
| lambda: os.symlink(src_real, dst_s), | |
| ): | |
| try: | |
| attempt() | |
| except OSError: | |
| continue | |
| if os.path.exists(dst_s): # must actually resolve to real bytes | |
| return dst_s | |
| try: | |
| os.unlink(dst_s) | |
| except OSError: | |
| pass | |
| return self._real.copy2(src_real, dst_s, *args, **kwargs) | |
| def _linking_downloads(download_mod): | |
| """Make upstream downloads link rather than copy, for the duration of the block. | |
| Swaps the module's own `shutil` binding rather than patching shutil.copy2 | |
| globally, so nothing outside the upstream download path is affected. | |
| """ | |
| real = download_mod.shutil | |
| download_mod.shutil = _LinkingShutil(real) | |
| try: | |
| yield | |
| finally: | |
| download_mod.shutil = real | |
| def _download_module(): | |
| return importlib.import_module("scripts.download_model_data") | |
| # key -> (payload, nbytes). Sizes are tracked here rather than as an attribute | |
| # on the payload, because state dicts are plain dicts and reject attributes. | |
| _ckpt_cache: "OrderedDict[tuple[str, int, int], tuple[object, int]]" = OrderedDict() | |
| _ckpt_bytes = 0 | |
| def _cached_torch_load(): | |
| """Serve repeat torch.load() calls for the same checkpoint from RAM. | |
| scripts/inference.py performs five torch.load() calls per CT run -- the | |
| autoencoder, diffusion UNet, controlnet, and both mask-generation nets, | |
| ~3.35 GB in total -- and re-reads every one of them from disk on each | |
| generation. The deserialised objects are only read from (load_state_dict | |
| copies out of them), so handing back the same object is safe. | |
| This caches host-side state dicts, not GPU memory: weights are ~3.35 GB for | |
| CT against a 96 GB xlarge slice, so VRAM was never the limit here. RAM is, | |
| hence the byte budget and LRU eviction. | |
| """ | |
| import torch | |
| original = torch.load | |
| def _load(f, *args, **kwargs): | |
| global _ckpt_bytes | |
| try: | |
| key_path = os.path.realpath(os.fspath(f)) | |
| st = os.stat(key_path) | |
| key = (key_path, st.st_size, int(st.st_mtime)) | |
| except (TypeError, OSError): | |
| return original(f, *args, **kwargs) # file-like object, not a path | |
| if key in _ckpt_cache: | |
| _ckpt_cache.move_to_end(key) | |
| log(f" ckpt cache HIT {os.path.basename(key_path)}") | |
| return _ckpt_cache[key][0] | |
| obj = original(f, *args, **kwargs) | |
| budget = int(CHECKPOINT_CACHE_GIB * 1024 ** 3) | |
| if st.st_size <= budget: | |
| while _ckpt_bytes + st.st_size > budget and _ckpt_cache: | |
| _, (_evicted, nbytes) = _ckpt_cache.popitem(last=False) | |
| _ckpt_bytes -= nbytes | |
| del _evicted | |
| _ckpt_cache[key] = (obj, st.st_size) | |
| _ckpt_bytes += st.st_size | |
| log(f" ckpt cache MISS {os.path.basename(key_path)}" | |
| f" (cached {_ckpt_bytes / 1e9:.2f} GB)") | |
| return obj | |
| torch.load = _load | |
| try: | |
| yield | |
| finally: | |
| torch.load = original | |
| def upstream_context(): | |
| """Temporarily add the upstream repo to sys.path and switch CWD to it.""" | |
| if not UPSTREAM.exists(): | |
| raise RuntimeError( | |
| f"Upstream repo not found at {UPSTREAM}. Run `bash pre-build.sh` first." | |
| ) | |
| prev_cwd = os.getcwd() | |
| added = False | |
| try: | |
| upstream_str = str(UPSTREAM) | |
| if upstream_str not in sys.path: | |
| sys.path.insert(0, upstream_str) | |
| added = True | |
| os.chdir(upstream_str) | |
| yield UPSTREAM | |
| finally: | |
| os.chdir(prev_cwd) | |
| if added and upstream_str in sys.path: | |
| sys.path.remove(upstream_str) | |
| def ensure_weights(version: str, root_dir: str | None = None) -> None: | |
| """Download model weights from HF Hub if not already on disk. | |
| `root_dir` only affects the `datasets/` payload; model checkpoints always | |
| land in <upstream>/models. Defaults to DATA_DIR so the CT datasets end up | |
| where upstream `scripts.inference` looks for them. | |
| """ | |
| if root_dir is None: | |
| root_dir = str(DATA_DIR) | |
| with upstream_context(): | |
| download_mod = _download_module() | |
| # download_model_data is idempotent — it skips files that already exist. | |
| model_only = version not in ("rflow-ct", "ddpm-ct") | |
| with _linking_downloads(download_mod): | |
| download_mod.download_model_data(version, root_dir, model_only=model_only) | |
| def extract_mask_archive() -> bool: | |
| """Unpack the CT anatomy-mask archive if it is not already unpacked. | |
| Upstream does this lazily in scripts/find_masks.py on the first CT run, | |
| which cost ~100s inside the GPU window. Doing it here means find_masks | |
| finds the directory present and skips straight past. Returns True if this | |
| call did the extraction. | |
| """ | |
| mask_dir = MASK_ARCHIVE.parent / MASK_ARCHIVE.stem | |
| if mask_dir.exists() or not MASK_ARCHIVE.is_file(): | |
| return False | |
| if MASK_ARCHIVE.stat().st_size == 0: | |
| return False # pruned placeholder; repair_mask_archive_state handles it | |
| from monai.apps.utils import extractall | |
| extractall(filepath=str(MASK_ARCHIVE), output_dir=str(MASK_ARCHIVE.parent), file_type="zip") | |
| return True | |
| def repair_mask_archive_state() -> None: | |
| """Undo a prune whose extracted directory has since disappeared. | |
| Removing the placeholder makes the destination look absent again, so the | |
| next download_model_data() refetches the real archive rather than trying to | |
| unzip an empty file. | |
| """ | |
| mask_dir = MASK_ARCHIVE.parent / MASK_ARCHIVE.stem | |
| if MASK_ARCHIVE.is_file() and MASK_ARCHIVE.stat().st_size == 0 and not mask_dir.exists(): | |
| log("mask archive placeholder without extracted data; forcing refetch") | |
| MASK_ARCHIVE.unlink() | |
| def prune_mask_archive() -> int: | |
| """Truncate the extracted archive to a placeholder. Returns bytes freed.""" | |
| mask_dir = MASK_ARCHIVE.parent / MASK_ARCHIVE.stem | |
| if not (PRUNE_ARCHIVE and mask_dir.is_dir() and MASK_ARCHIVE.is_file()): | |
| return 0 | |
| if not any(mask_dir.iterdir()): | |
| return 0 # extraction produced nothing; keep the archive | |
| freed = MASK_ARCHIVE.stat().st_size | |
| if freed == 0: | |
| return 0 | |
| MASK_ARCHIVE.unlink() # drop our link; the cache blob goes too if last | |
| MASK_ARCHIVE.touch() # placeholder keeps the existence check happy | |
| return freed | |
| def prefetch_weights(versions: Iterable[str] = ALL_VERSIONS) -> tuple[Stopwatch, list[str]]: | |
| """Warm the weight cache at startup, outside any @spaces.GPU function. | |
| ~15 GB for CT (11.6 GB of that the anatomy-mask archive) plus ~2.3 GB each | |
| for MR and MR-Brain. Fetching this lazily inside the decorated function | |
| burned the whole ZeroGPU allocation on I/O before inference began. Failures | |
| are logged and swallowed: the lazy path still works, it is just slow. | |
| """ | |
| os.environ.setdefault("MONAI_DATA_DIRECTORY", str(DATA_DIR)) | |
| DATA_DIR.mkdir(parents=True, exist_ok=True) | |
| watch = Stopwatch("prefetch") | |
| failures: list[str] = [] | |
| repair_mask_archive_state() | |
| for version in versions: | |
| with watch.split(f"fetch {version}"): | |
| try: | |
| log(f"-> fetching weights: {version}") | |
| ensure_weights(version) | |
| except Exception as e: # noqa: BLE001 - startup must not be fatal | |
| failures.append(version) | |
| log(f"xx fetch FAILED for {version}: {type(e).__name__}: {e}") | |
| with watch.split("extract mask archive"): | |
| try: | |
| if extract_mask_archive(): | |
| log("<- extracted CT anatomy-mask archive") | |
| else: | |
| log("<- mask archive already extracted, skipped") | |
| except Exception as e: # noqa: BLE001 | |
| failures.append("mask-archive") | |
| log(f"xx mask archive extraction FAILED: {type(e).__name__}: {e}") | |
| with watch.split("prune mask archive"): | |
| try: | |
| freed = prune_mask_archive() | |
| if freed: | |
| log(f"<- reclaimed {freed / 1e9:.1f} GB from the extracted archive") | |
| except Exception as e: # noqa: BLE001 | |
| log(f"xx archive prune skipped: {type(e).__name__}: {e}") | |
| return watch, failures | |
| def _list_outputs_before(output_dir: Path) -> set[str]: | |
| if not output_dir.exists(): | |
| return set() | |
| return {p.name for p in output_dir.glob("*.nii.gz")} | |
| def _newest_outputs(output_dir: Path, before: set[str]) -> list[Path]: | |
| if not output_dir.exists(): | |
| return [] | |
| new = [p for p in output_dir.glob("*.nii.gz") if p.name not in before] | |
| new.sort(key=lambda p: p.stat().st_mtime) | |
| return new | |
| def _write_temp_configs( | |
| base_env_config: Path, | |
| base_model_config: Path, | |
| overrides: dict, | |
| tag: str, | |
| ) -> tuple[Path, Path]: | |
| """ | |
| Write modified copies of the env + model configs into a per-call temp dir under | |
| UPSTREAM / configs / _temp /. Returns (env_path, model_path). | |
| """ | |
| temp_dir = UPSTREAM / "configs" / "_temp" | |
| temp_dir.mkdir(parents=True, exist_ok=True) | |
| env = json.loads(base_env_config.read_text()) | |
| model = json.loads(base_model_config.read_text()) | |
| if "env" in overrides: | |
| env.update(overrides["env"]) | |
| if "diffusion_unet_inference" in overrides: | |
| model.setdefault("diffusion_unet_inference", {}).update(overrides["diffusion_unet_inference"]) | |
| suffix = f"{tag}_{uuid.uuid4().hex[:8]}" | |
| env_path = temp_dir / f"env_{suffix}.json" | |
| model_path = temp_dir / f"model_{suffix}.json" | |
| env_path.write_text(json.dumps(env, indent=2)) | |
| model_path.write_text(json.dumps(model, indent=2)) | |
| return env_path, model_path | |
| def run_image_only( | |
| *, | |
| version: str, | |
| output_size: tuple[int, int, int], | |
| spacing: tuple[float, float, float], | |
| modality: int, | |
| seed: int, | |
| num_inference_steps: int = 30, | |
| cfg_guidance_scale: Optional[float] = None, | |
| ) -> Path: | |
| """ | |
| Run the image-only diffusion pipeline (`scripts.diff_model_infer`) for the given | |
| version (rflow-ct / rflow-mr / rflow-mr-brain). Returns path to generated NIfTI. | |
| """ | |
| ensure_weights(version) | |
| base_env = UPSTREAM / "configs" / f"environment_maisi_diff_model_{version}.json" | |
| base_model = UPSTREAM / "configs" / f"config_maisi_diff_model_{version}.json" | |
| network_def = UPSTREAM / "configs" / "config_network_rflow.json" | |
| inference_overrides = { | |
| "dim": list(output_size), | |
| "spacing": list(spacing), | |
| "modality": modality, | |
| "random_seed": seed, | |
| "num_inference_steps": num_inference_steps, | |
| } | |
| if cfg_guidance_scale is not None: | |
| inference_overrides["cfg_guidance_scale"] = cfg_guidance_scale | |
| with upstream_context(): | |
| env_path, model_path = _write_temp_configs( | |
| base_env_config=base_env, | |
| base_model_config=base_model, | |
| overrides={"diffusion_unet_inference": inference_overrides}, | |
| tag=version, | |
| ) | |
| # Read env to determine output_dir (relative to upstream root) | |
| env_data = json.loads(env_path.read_text()) | |
| output_dir = (UPSTREAM / env_data["output_dir"]).resolve() | |
| existing = _list_outputs_before(output_dir) | |
| diff_mod = importlib.import_module("scripts.diff_model_infer") | |
| t0 = time.time() | |
| with phase(f"{version}.diffusion"), _cached_torch_load(): | |
| diff_mod.diff_model_infer( | |
| env_config_path=str(env_path.relative_to(UPSTREAM)), | |
| model_config_path=str(model_path.relative_to(UPSTREAM)), | |
| model_def_path=str(network_def.relative_to(UPSTREAM)), | |
| num_gpus=1, | |
| ) | |
| runtime = time.time() - t0 | |
| new_files = _newest_outputs(output_dir, existing) | |
| if not new_files: | |
| raise RuntimeError(f"No new NIfTI produced in {output_dir}") | |
| latest = new_files[-1] | |
| # Cleanup temp configs (don't fail if cleanup errors) | |
| for p in (env_path, model_path): | |
| try: | |
| p.unlink() | |
| except OSError: | |
| pass | |
| return latest | |
| def run_paired_ct( | |
| *, | |
| output_size: tuple[int, int, int], | |
| spacing: tuple[float, float, float], | |
| body_region: list[str], | |
| anatomy_list: list[str], | |
| seed: int, | |
| num_inference_steps: int = 30, | |
| num_output_samples: int = 1, | |
| ) -> tuple[Path, Optional[Path]]: | |
| """ | |
| Run the paired CT image+mask pipeline (`scripts.inference`). Returns | |
| (image_path, mask_path). Mask is the corresponding label volume. | |
| """ | |
| version = "rflow-ct" | |
| # No ensure_weights() here: upstream scripts/inference.py calls | |
| # download_model_data(version, MONAI_DATA_DIRECTORY) itself. Calling it here | |
| # too — with a different root_dir — downloaded the 11.6 GB anatomy-mask | |
| # archive to a second location. prefetch_weights() warms the same paths at | |
| # startup, so this run finds everything already present. | |
| base_env = UPSTREAM / "configs" / f"environment_{version}.json" | |
| base_infer = UPSTREAM / "configs" / "config_infer.json" | |
| # Build a custom config_infer with overrides | |
| infer_data = json.loads(base_infer.read_text()) | |
| infer_data["output_size"] = list(output_size) | |
| infer_data["spacing"] = list(spacing) | |
| infer_data["body_region"] = list(body_region) | |
| infer_data["anatomy_list"] = list(anatomy_list) | |
| infer_data["num_inference_steps"] = num_inference_steps | |
| infer_data["num_output_samples"] = num_output_samples | |
| temp_dir = UPSTREAM / "configs" / "_temp" | |
| temp_dir.mkdir(parents=True, exist_ok=True) | |
| suffix = uuid.uuid4().hex[:8] | |
| infer_path = temp_dir / f"config_infer_{version}_{suffix}.json" | |
| infer_path.write_text(json.dumps(infer_data, indent=2)) | |
| env_data = json.loads(base_env.read_text()) | |
| output_dir = (UPSTREAM / env_data["output_dir"]).resolve() | |
| with upstream_context(): | |
| existing = _list_outputs_before(output_dir) | |
| inference_mod = importlib.import_module("scripts.inference") | |
| # The upstream `main()` parses argv directly. Patch sys.argv around the call. | |
| old_argv = sys.argv | |
| sys.argv = [ | |
| "scripts.inference", | |
| "-t", "./configs/config_network_rflow.json", | |
| "-i", str(infer_path.relative_to(UPSTREAM)), | |
| "-e", str(base_env.relative_to(UPSTREAM)), | |
| "--random-seed", str(seed), | |
| "--version", version, | |
| ] | |
| os.environ.setdefault("MONAI_DATA_DIRECTORY", str(DATA_DIR)) | |
| try: | |
| # Link-instead-of-copy also covers the download inference.py runs | |
| # internally, for the case where prefetch was skipped or failed. | |
| with phase("ct.paired-inference"), _cached_torch_load(), _linking_downloads(_download_module()): | |
| inference_mod.main() | |
| finally: | |
| sys.argv = old_argv | |
| new_files = _newest_outputs(output_dir, existing) | |
| try: | |
| infer_path.unlink() | |
| except OSError: | |
| pass | |
| # Paired pipeline writes both image and label NIfTIs. Convention: filenames | |
| # contain "image" / "label" or are emitted as adjacent files. | |
| image_path: Optional[Path] = None | |
| mask_path: Optional[Path] = None | |
| for p in new_files: | |
| name = p.name.lower() | |
| if "label" in name or "_mask" in name or "seg" in name: | |
| mask_path = p | |
| elif "image" in name or "img" in name: | |
| image_path = p | |
| # Fallback: if naming is ambiguous, treat the smaller-modality-time file as image | |
| if image_path is None and new_files: | |
| image_path = new_files[0] | |
| if mask_path is None and len(new_files) > 1: | |
| mask_path = new_files[-1] | |
| if image_path is None: | |
| raise RuntimeError(f"No NIfTI produced in {output_dir}") | |
| return image_path, mask_path | |
| def labels_present(mask_path: Path) -> set[int]: | |
| """Return the set of unique non-zero label IDs present in the mask volume.""" | |
| import nibabel as nib | |
| import numpy as np | |
| with phase("ct.label-scan"): | |
| img = nib.load(str(mask_path)) | |
| data = np.asarray(img.dataobj) | |
| uniq = np.unique(data).astype(int).tolist() | |
| return {int(u) for u in uniq if u != 0} | |