#!/usr/bin/env python3 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # 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. """NVIDIA-Medtech NV-Generate-CTMR rflow-mr skill. Thin wrapper around the upstream `scripts.diff_model_infer` entry point from https://github.com/NVIDIA-Medtech/NV-Generate-CTMR. The wrapper does NOT implement diffusion sampling or autoencoder decoding. It stages config overrides, shells out to the upstream command, and summarizes generated MR NIfTI outputs. Engineering verification only. Output is NOT clinically meaningful. """ from __future__ import annotations import json import os import subprocess import sys import time from pathlib import Path from typing import Any import nibabel as nib import numpy as np import typer SKILL_NAME = "nv_generate_mr" MODEL_REPO = "https://github.com/NVIDIA-Medtech/NV-Generate-CTMR" MODEL_WEIGHTS_REPO = "https://huggingface.co/nvidia/NV-Generate-MR" VERSION = "rflow-mr" REPO_ROOT = Path(__file__).resolve().parents[int("3")] UPSTREAM_NETWORK_CONFIG = "configs/config_network_rflow.json" UPSTREAM_MODEL_CONFIG = "configs/config_maisi_diff_model_rflow-mr.json" UPSTREAM_ENV_CONFIG = "configs/environment_maisi_diff_model_rflow-mr.json" UPSTREAM_MODALITY_MAPPING = "configs/modality_mapping.json" UPSTREAM_MODEL_FILES = ( "models/autoencoder_v2.pt", "models/diff_unet_3d_rflow-mr.pt", ) SUPPORTED_MODALITIES = ("mri", "mri_t1", "mri_t2", "mri_flair") OVERRIDE_KEYS = ( "dim", "spacing", "top_region_index", "bottom_region_index", "random_seed", "num_inference_steps", "modality", "cfg_guidance_scale", "output_prefix", ) MAX_VOXELS = int("512") * int("512") * int("128") app = typer.Typer(add_completion=False) def emit(payload: dict[str, Any]) -> None: sys.stdout.write(json.dumps(payload, indent=2)) sys.stdout.flush() def tail(s: str, n_chars: int = int("4000")) -> str: if len(s) <= n_chars: return s return "..." + s[-n_chars:] def sha256_file(path: Path, chunk: int = 1 << int("20")) -> str: import hashlib h = hashlib.sha256() with path.open("rb") as f: while True: buf = f.read(chunk) if not buf: break h.update(buf) return h.hexdigest() def file_sha256_safe(path: Path) -> str: if not path.is_file(): return "" try: return sha256_file(path) except Exception: return "" def git_commit(root: Path) -> str: try: proc = subprocess.run( ["git", "rev-parse", "HEAD"], cwd=str(root), check=False, capture_output=True, text=True, timeout=int("10"), ) except Exception: return "" if proc.returncode == 0: return proc.stdout.strip() return "" def _round(values: Any, ndigits: int = int("6")) -> Any: if isinstance(values, (list, tuple, np.ndarray)): return [round(float(v), ndigits) for v in values] return round(float(values), ndigits) def _load_json(path: Path) -> dict[str, Any]: return json.loads(path.read_text()) def _valid_upstream_root(path: Path) -> bool: return (path / UPSTREAM_NETWORK_CONFIG).is_file() def _candidate_upstream_roots(env_value: str) -> list[Path]: candidates: list[Path] = [] if env_value: candidates.append(Path(env_value).expanduser()) candidates.extend( [ REPO_ROOT / ".workbench_data/upstreams/NV-Generate-CTMR", Path.home() / "NV-Generate-CTMR", Path.home() / "nv-generate-ctmr", ] ) deduped: list[Path] = [] seen: set[str] = set() for candidate in candidates: key = str(candidate) if key not in seen: seen.add(key) deduped.append(candidate) return deduped def _resolve_upstream_root(env_value: str) -> tuple[Path | None, list[str]]: checked: list[str] = [] for candidate in _candidate_upstream_roots(env_value): resolved = candidate.resolve() checked.append(str(resolved)) if _valid_upstream_root(resolved): return resolved, checked return None, checked def _load_modality_mapping(upstream_root: Path) -> dict[str, int]: mapping = _load_json(upstream_root / UPSTREAM_MODALITY_MAPPING) return {str(k): int(v) for k, v in mapping.items()} def _modality_to_code(modality: str, mapping: dict[str, int]) -> int: if modality not in SUPPORTED_MODALITIES: raise typer.BadParameter( f"--modality must be one of {list(SUPPORTED_MODALITIES)}, got {modality!r}" ) if modality not in mapping: raise typer.BadParameter(f"modality {modality!r} not found in upstream modality mapping") return int(mapping[modality]) def _load_config_override(fixture_arg: str) -> tuple[dict[str, Any], str | None]: if fixture_arg == "default": return {}, None fixture_path = Path(fixture_arg).expanduser().resolve() if not fixture_path.is_file(): raise typer.BadParameter(f"model config override not found: {fixture_arg}") raw = json.loads(fixture_path.read_text()) cleaned = {k: v for k, v in raw.items() if not k.startswith("_")} if "diffusion_unet_inference" in cleaned: nested = cleaned.pop("diffusion_unet_inference") if not isinstance(nested, dict): raise typer.BadParameter("diffusion_unet_inference must be a JSON object") cleaned.update(nested) unknown = sorted(k for k in cleaned if k not in OVERRIDE_KEYS) if unknown: raise typer.BadParameter( f"MR override contains unknown key(s): {unknown}. Allowed: {sorted(OVERRIDE_KEYS)}" ) return cleaned, str(fixture_path) def _validate_inference_config(rendered_inference: dict[str, Any]) -> list[str]: errors: list[str] = [] dim = rendered_inference.get("dim") if not (isinstance(dim, (list, tuple)) and len(dim) == int("3")): errors.append(f"dim must be a 3-tuple, got {dim!r}") else: voxels = 1 for i, v in enumerate(dim): if not isinstance(v, int): errors.append(f"dim[{i}] must be int, got {v!r}") elif v < int("64") or v > int("512"): errors.append(f"dim[{i}]={v} outside rflow-mr range [64, 512]") elif v % int("32") != 0: errors.append(f"dim[{i}]={v} must be a multiple of 32") if isinstance(v, int): voxels *= int(v) if voxels > MAX_VOXELS: errors.append( f"dim product {voxels} exceeds rflow-mr max volume {MAX_VOXELS} " "(512x512x128 per upstream README)" ) spacing = rendered_inference.get("spacing") if not (isinstance(spacing, (list, tuple)) and len(spacing) == int("3")): errors.append(f"spacing must be a 3-tuple, got {spacing!r}") else: for i, v in enumerate(spacing): if not isinstance(v, (int, float)) or v <= 0: errors.append(f"spacing[{i}] must be a positive float, got {v!r}") for key in ("top_region_index", "bottom_region_index"): value = rendered_inference.get(key) if not (isinstance(value, (list, tuple)) and len(value) == int("4")): errors.append(f"{key} must be a 4-tuple, got {value!r}") elif not all(isinstance(v, (int, float)) for v in value): errors.append(f"{key} values must be numeric, got {value!r}") n_steps = rendered_inference.get("num_inference_steps") if not isinstance(n_steps, int) or n_steps < 1 or n_steps > int("2000"): errors.append(f"num_inference_steps must be int in [1, 2000], got {n_steps!r}") seed = rendered_inference.get("random_seed") if not isinstance(seed, int): errors.append(f"random_seed must be int, got {seed!r}") cfg = rendered_inference.get("cfg_guidance_scale") if not isinstance(cfg, (int, float)): errors.append(f"cfg_guidance_scale must be numeric, got {cfg!r}") modality = rendered_inference.get("modality") if not isinstance(modality, int) or modality < 0: errors.append(f"modality must be a non-negative int code, got {modality!r}") return errors def _stage_config( upstream_root: Path, stage_dir: Path, override: dict[str, Any], output_dir: Path, modality_code: int, modality_name: str, seed: int, ) -> tuple[dict[str, Any], dict[str, Any], Path, Path]: stage_dir.mkdir(parents=True, exist_ok=True) base_model = _load_json(upstream_root / UPSTREAM_MODEL_CONFIG) rendered_model = dict(base_model) inference = dict(rendered_model.get("diffusion_unet_inference") or {}) inference.update(override) inference["modality"] = modality_code inference["random_seed"] = seed rendered_model["diffusion_unet_inference"] = inference staged_model_path = stage_dir / "config_maisi_diff_model_rflow-mr.json" staged_model_path.write_text(json.dumps(rendered_model, indent=2)) base_env = _load_json(upstream_root / UPSTREAM_ENV_CONFIG) rendered_env = dict(base_env) rendered_env["output_dir"] = str(output_dir) if "output_prefix" in override: rendered_env["output_prefix"] = str(override["output_prefix"]) else: rendered_env["output_prefix"] = f"mr_{modality_name}" staged_env_path = stage_dir / "environment_maisi_diff_model_rflow-mr.json" staged_env_path.write_text(json.dumps(rendered_env, indent=2)) return rendered_model, rendered_env, staged_model_path, staged_env_path def _estimate_cost(rendered_inference: dict[str, Any]) -> dict[str, Any]: dim = rendered_inference.get("dim") or [int("128"), int("256"), int("256")] n_steps = int(rendered_inference.get("num_inference_steps") or int("30")) voxels = int(dim[0]) * int(dim[1]) * int(dim[2]) ref_voxels = int("128") * int("256") * int("256") ref_steps = int("30") ref_seconds = float("60.0") seconds = ref_seconds * (voxels / ref_voxels) * (n_steps / ref_steps) vram = float("16.0") if voxels <= ref_voxels else float("32.0") disk_mb = (voxels * 2.0) / (int("1024") * int("1024")) return { "version": VERSION, "voxels_per_sample": voxels, "num_inference_steps": n_steps, "estimated_wall_seconds": round(seconds, 1), "estimated_peak_vram_gb": round(vram, 1), "estimated_disk_mb": round(disk_mb, 1), } def _detect_cuda() -> dict[str, Any]: info: dict[str, Any] = {"available": False, "device_name": None, "total_memory_gb": None} try: import torch # noqa: PLC0415 info["torch_version"] = torch.__version__ info["available"] = bool(torch.cuda.is_available()) if info["available"]: props = torch.cuda.get_device_properties(0) info["device_name"] = props.name info["total_memory_gb"] = round(props.total_memory / (int("1024") ** int("3")), 1) info["cuda_version"] = torch.version.cuda except Exception as e: info["import_error"] = repr(e) return info def _preflight(rendered_inference: dict[str, Any]) -> tuple[list[str], list[str], dict[str, Any]]: errors = _validate_inference_config(rendered_inference) warnings: list[str] = [] cuda = _detect_cuda() cost = _estimate_cost(rendered_inference) if not cuda["available"]: errors.append( "CUDA not available. rflow-mr synthesis needs an NVIDIA GPU; " "there is no CPU fallback in the upstream code path." ) elif cuda["total_memory_gb"] is not None: usable = cuda["total_memory_gb"] * float("0.85") if cost["estimated_peak_vram_gb"] > usable: warnings.append( f"estimated peak VRAM {cost['estimated_peak_vram_gb']} GB exceeds " f"85% of detected GPU memory ({cuda['total_memory_gb']} GB on " f"{cuda['device_name']}). Risk of OOM; reduce dim or use a larger GPU." ) return errors, warnings, {"cuda": cuda, "estimated_cost": cost} def _model_inventory(upstream_root: Path) -> dict[str, Any]: files: list[dict[str, Any]] = [] all_present = True for rel in UPSTREAM_MODEL_FILES: path = upstream_root / rel present = path.is_file() files.append( { "path": rel, "present": present, "bytes": path.stat().st_size if present else None, "sha256": file_sha256_safe(path) if present else "", } ) all_present = all_present and present return {"all_present": all_present, "files": files} def _build_command(staged_model_path: Path, staged_env_path: Path, num_gpus: int) -> list[str]: cmd = [ sys.executable, "-m", "scripts.diff_model_infer", "-t", f"./{UPSTREAM_NETWORK_CONFIG}", "-e", str(staged_env_path), "-c", str(staged_model_path), ] if num_gpus != 1: cmd.extend(["-g", str(num_gpus)]) return cmd def _scan_outputs(output_dir: Path, run_started: float) -> list[Path]: if not output_dir.is_dir(): return [] paths: list[Path] = [] for path in output_dir.rglob("*.nii*"): if not path.is_file(): continue try: if path.stat().st_size > 0 and path.stat().st_mtime >= run_started - 1: paths.append(path) except OSError: continue return sorted(paths) def _summarize_image( image_path: Path, requested_dim: list[int], requested_spacing: list[float], ) -> dict[str, Any]: record: dict[str, Any] = { "image_path": str(image_path), "image_bytes": image_path.stat().st_size if image_path.exists() else None, "image_sha256": file_sha256_safe(image_path) if image_path.exists() else "", "image_readable": False, } try: img = nib.load(str(image_path)) arr = np.asarray(img.get_fdata(), dtype=np.float32) finite = arr[np.isfinite(arr)] record["image_readable"] = True record["image_shape"] = [int(v) for v in arr.shape] record["requested_shape"] = [int(v) for v in requested_dim] record["shape_match_requested"] = record["image_shape"] == record["requested_shape"] record["image_spacing"] = _round(img.header.get_zooms()[: int("3")]) record["requested_spacing"] = _round(requested_spacing) record["spacing_match_requested"] = record["image_spacing"] == record["requested_spacing"] record["image_affine"] = [list(map(float, row)) for row in img.affine.tolist()] record["finite_fraction"] = ( round(float(finite.size) / float(arr.size), int("6")) if arr.size else 0.0 ) record["all_finite"] = bool(finite.size == arr.size) if finite.size: record["intensity_min"] = _round(float(finite.min()), int("3")) record["intensity_max"] = _round(float(finite.max()), int("3")) record["intensity_mean"] = _round(float(finite.mean()), int("3")) record["intensity_std"] = _round(float(finite.std()), int("3")) record["image_nonconstant"] = bool(finite.max() - finite.min() > 1.0) record["image_nonnegative"] = bool(finite.min() >= 0) else: record["image_nonconstant"] = False record["image_nonnegative"] = False except Exception as e: record["image_error"] = repr(e) return record def _aggregate(samples: list[dict[str, Any]]) -> dict[str, Any]: n = len(samples) return { "num_samples": n, "all_images_readable": bool(n) and all(s.get("image_readable") for s in samples), "all_shapes_match_requested": bool(n) and all(s.get("shape_match_requested") for s in samples), "all_spacing_match_requested": bool(n) and all(s.get("spacing_match_requested") for s in samples), "all_images_finite": bool(n) and all(s.get("all_finite") for s in samples), "all_images_nonconstant": bool(n) and all(s.get("image_nonconstant") for s in samples), "all_images_nonnegative": bool(n) and all(s.get("image_nonnegative") for s in samples), } @app.command() def main( model_config: str = typer.Argument( ..., help='Path to a model-config override JSON, or "default" for upstream defaults.', ), output_dir: Path | None = typer.Option( None, "--output-dir", "-o", help="Absolute directory for generated NIfTI volumes." ), modality: str | None = typer.Option(None, "--modality", help="MR modality name."), seed: int = typer.Option(0, "--random-seed", "-s"), num_gpus: int = typer.Option(1, "--num-gpus", min=1), timeout_seconds: float = typer.Option(float("3600.0"), "--timeout-seconds"), preflight_only: bool = typer.Option( False, "--preflight-only", help="Validate config, CUDA, cost estimate, and model inventory without inference.", ), yes: bool = typer.Option( False, "--yes", "-y", help="Skip the cost-preview confirmation gate for large runs.", ), ) -> None: """Generate synthetic 3D MRI volumes via NV-Generate-CTMR rflow-mr.""" upstream_root_env = os.environ.get("NV_GENERATE_ROOT", "").strip() upstream_root, checked_roots = _resolve_upstream_root(upstream_root_env) if upstream_root is None and not upstream_root_env: emit( { "skill": SKILL_NAME, "error": "NV_GENERATE_ROOT is unset", "detail": "Clone https://github.com/NVIDIA-Medtech/NV-Generate-CTMR and export " "NV_GENERATE_ROOT to its path, or place the clone at " ".workbench_data/upstreams/NV-Generate-CTMR.", "checked_roots": checked_roots, } ) raise typer.Exit(2) if upstream_root is None: emit( { "skill": SKILL_NAME, "error": "NV_GENERATE_ROOT layout invalid", "detail": f"{UPSTREAM_NETWORK_CONFIG} not found in any checked root", "checked_roots": checked_roots, } ) raise typer.Exit(2) if output_dir is None: output_dir = upstream_root / "output" output_dir = output_dir.expanduser().resolve() output_dir.mkdir(parents=True, exist_ok=True) mapping = _load_modality_mapping(upstream_root) override, override_source = _load_config_override(model_config) override_modality = str(override.pop("modality")) if "modality" in override else None modality_name = modality or override_modality or "mri_t1" modality_code = _modality_to_code(modality_name, mapping) stage_dir = output_dir / "_staged_configs" rendered_model, rendered_env, staged_model_path, staged_env_path = _stage_config( upstream_root, stage_dir, override, output_dir, modality_code, modality_name, seed, ) rendered_inference = rendered_model["diffusion_unet_inference"] errors, warnings, context = _preflight(rendered_inference) inventory = _model_inventory(upstream_root) if not inventory["all_present"]: errors.append( "missing rflow-mr model weights. Run `python -m scripts.download_model_data " "--version rflow-mr --root_dir ./ --model_only` from $NV_GENERATE_ROOT." ) cost = context["estimated_cost"] cuda = context["cuda"] print( f"[nv_generate_mr] preflight: dim={rendered_inference.get('dim')} " f"spacing={rendered_inference.get('spacing')} modality={modality_name}({modality_code}) " f"steps={rendered_inference.get('num_inference_steps')}", file=sys.stderr, ) print( f"[nv_generate_mr] cost estimate: ~{cost['estimated_wall_seconds']}s wall, " f"~{cost['estimated_peak_vram_gb']} GB VRAM peak, ~{cost['estimated_disk_mb']} MB disk. " f"GPU: {cuda.get('device_name','?')} ({cuda.get('total_memory_gb','?')} GB)", file=sys.stderr, ) for warning in warnings: print(f"[nv_generate_mr] warning: {warning}", file=sys.stderr) if errors: for error in errors: print(f"[nv_generate_mr] error: {error}", file=sys.stderr) emit( { "skill": SKILL_NAME, "error": "preflight validation failed", "preflight_errors": errors, "preflight_warnings": warnings, "estimated_cost": cost, "cuda": cuda, "invocation": {"model_inventory": inventory}, } ) raise typer.Exit(2) if preflight_only: emit( { "skill": SKILL_NAME, "preflight": "ok", "preflight_warnings": warnings, "estimated_cost": cost, "cuda": cuda, "model_inventory": inventory, "rendered_model_config": rendered_model, "rendered_env_config": rendered_env, } ) raise typer.Exit(0) if not yes and ( cost["estimated_wall_seconds"] > float("300.0") or cost["estimated_peak_vram_gb"] > float("30.0") ): emit( { "skill": SKILL_NAME, "error": "cost gate: run would be expensive; re-run with --yes to proceed", "estimated_cost": cost, "cuda": cuda, } ) raise typer.Exit(2) cmd = _build_command(staged_model_path, staged_env_path, num_gpus) run_env = os.environ.copy() run_env.setdefault("MONAI_DATA_DIRECTORY", str(upstream_root / "temp_work_dir")) run_env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128,expandable_segments:True") run_started = time.time() t0 = time.monotonic() try: proc = subprocess.run( cmd, cwd=str(upstream_root), env=run_env, capture_output=True, text=True, timeout=timeout_seconds, check=False, ) rc = proc.returncode stdout = proc.stdout stderr = proc.stderr except subprocess.TimeoutExpired as e: rc = int("124") stdout = e.stdout.decode() if isinstance(e.stdout, bytes) else (e.stdout or "") stderr_raw = e.stderr.decode() if isinstance(e.stderr, bytes) else (e.stderr or "") stderr = stderr_raw + f"\n[TIMEOUT after {timeout_seconds}s]" elapsed = time.monotonic() - t0 requested_dim = [int(v) for v in rendered_inference["dim"]] requested_spacing = [float(v) for v in rendered_inference["spacing"]] output_paths = _scan_outputs(output_dir, run_started) samples = [_summarize_image(p, requested_dim, requested_spacing) for p in output_paths] aggregate = _aggregate(samples) payload: dict[str, Any] = { "skill": SKILL_NAME, "model": "NVIDIA-Medtech/NV-Generate-CTMR (rflow-mr)", "model_repo": MODEL_REPO, "model_weights_repo": MODEL_WEIGHTS_REPO, "license": "Wrapper Apache-2.0; NV-Generate-MR weights use NVIDIA Non-Commercial License.", "input": { "model_config_override_path": override_source, "model_config_override": override, "modality_name": modality_name, "modality_code": modality_code, "dim_requested": requested_dim, "spacing_requested": requested_spacing, "num_inference_steps_requested": rendered_inference.get("num_inference_steps"), "cfg_guidance_scale_requested": rendered_inference.get("cfg_guidance_scale"), "random_seed": seed, "version": VERSION, }, "output": { "directory": str(output_dir), "samples": samples, **aggregate, }, "invocation": { "official_entrypoint": "python -m scripts.diff_model_infer", "upstream_root": str(upstream_root), "upstream_commit": git_commit(upstream_root), "command": cmd, "exit_code": rc, "subprocess_seconds": round(elapsed, int("3")), "model_inventory": inventory, "rendered_model_config": rendered_model, "rendered_env_output_dir": rendered_env.get("output_dir"), "rendered_env_output_prefix": rendered_env.get("output_prefix"), }, "runtime": { "subprocess_seconds": round(elapsed, int("3")), "device": "cuda", }, "logs": { "stdout_tail": tail(stdout), "stderr_tail": tail(stderr), }, "preflight": { "warnings": warnings, "estimated_cost": cost, "cuda": cuda, }, "intended_use_disclaimer": ( "Engineering verification only. Output is synthetic and NOT clinically meaningful. " "This wrapper invokes the upstream scripts.diff_model_infer entry point from the " "NV-Generate-CTMR README; it does not modify diffusion sampling or autoencoder decoding." ), } emit(payload) raise typer.Exit(0) if __name__ == "__main__": app()