#!/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. """Wrapper for NV-Generate-CTMR VAE finetuning. The upstream repository documents VAE training in ``train_vae_tutorial.ipynb`` and provides reusable configs/transforms/utilities, but no dedicated ``scripts.train_vae`` entrypoint. This wrapper stages the same config and datalist contract into deterministic files, then runs the VAE training loop against existing upstream helper APIs. It does not execute the notebook. Engineering verification only. Outputs are not clinically meaningful. """ from __future__ import annotations import argparse import copy import json import os import subprocess import sys import time from pathlib import Path from typing import Any SKILL_NAME = "nv_generate_vae_finetune" MODEL_NAME = "maisi-vae" UPSTREAM_REPO = "https://github.com/NVIDIA-Medtech/NV-Generate-CTMR" UPSTREAM_ENTRYPOINT = "python skills/nv-generate-vae-finetune/scripts/run_vae_finetune.py" REPO_ROOT = Path(__file__).resolve().parents[3] DEFAULT_UPSTREAM = REPO_ROOT / ".workbench_data" / "upstreams" / "NV-Generate-CTMR" REQUIRED_UPSTREAM_FILES = ( "scripts/download_model_data.py", "scripts/transforms.py", "scripts/utils.py", "configs/config_network_rflow.json", "configs/environment_maisi_vae_train.json", "configs/config_maisi_vae_train.json", ) SUPPORTED_MODALITIES = ("ct", "mri") def _emit(payload: dict[str, Any]) -> None: sys.stdout.write(json.dumps(payload, indent=2)) sys.stdout.flush() def _tail(text: str, n_chars: int = 4000) -> str: return text if len(text) <= n_chars else "..." + text[-n_chars:] def _load_json(path: Path) -> Any: return json.loads(path.read_text()) def _write_json(path: Path, payload: dict[str, Any]) -> Path: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") return path def _parse_triplet(value: str, cast: type = int) -> list[Any]: parts = [cast(v.strip()) for v in value.split(",")] if len(parts) != 3: raise argparse.ArgumentTypeError("expected three comma-separated values") return parts def _normalize_modality(value: str) -> str: lower = value.lower() if lower == "ct": return "ct" if lower == "mri" or lower.startswith("mri_"): return "mri" raise ValueError(f"unsupported modality {value!r}; expected ct or mri") def _resolve_upstream_root(explicit: str | None = None) -> tuple[Path | None, list[str]]: candidates: list[Path] = [] if explicit: candidates.append(Path(explicit).expanduser()) env_root = os.environ.get("NV_GENERATE_ROOT") if env_root: candidates.append(Path(env_root).expanduser()) candidates.extend([DEFAULT_UPSTREAM, Path.home() / "NV-Generate-CTMR"]) checked: list[str] = [] seen: set[str] = set() for candidate in candidates: resolved = candidate.resolve() key = str(resolved) if key in seen: continue seen.add(key) checked.append(key) if all((resolved / rel).is_file() for rel in REQUIRED_UPSTREAM_FILES): return resolved, checked return None, checked def _resolve_data_path(data_base_dir: Path, image: str) -> Path: image_path = Path(image) if image_path.is_absolute(): raise ValueError("datalist image paths must be relative to --data-base-dir") return data_base_dir / image_path def _split_entries(raw: dict[str, Any]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: training = raw.get("training") validation = raw.get("validation", raw.get("testing", raw.get("val", []))) if not isinstance(training, list) or not training: raise ValueError("datalist.training must be a non-empty list") if not isinstance(validation, list) or not validation: raise ValueError("datalist must include non-empty validation[] or testing[] entries") return training, validation def _validate_datalist( data_base_dir: Path, datalist: Path, default_modality: str ) -> dict[str, Any]: default_class = _normalize_modality(default_modality) raw = _load_json(datalist) if not isinstance(raw, dict): raise ValueError("datalist must be a JSON object") training, validation = _split_entries(raw) missing: list[str] = [] classes: set[str] = set() for split_name, entries in (("training", training), ("validation", validation)): for i, item in enumerate(entries): if not isinstance(item, dict) or "image" not in item: raise ValueError(f"{split_name}[{i}] must contain image") modality = _normalize_modality( str(item.get("class", item.get("modality", default_class))) ) image_path = _resolve_data_path(data_base_dir, str(item["image"])) if not image_path.is_file(): missing.append(str(image_path)) classes.add(modality) if missing: raise FileNotFoundError(f"missing datalist image(s): {missing[:5]}") return { "data_base_dir": str(data_base_dir), "datalist": str(datalist), "training_cases": len(training), "validation_cases": len(validation), "modalities": sorted(classes), "default_modality": default_class, } def _stage_entries( data_base_dir: Path, entries: list[dict[str, Any]], default_modality: str, ) -> list[dict[str, Any]]: staged: list[dict[str, Any]] = [] for item in entries: next_item = dict(item) next_item["image"] = str( _resolve_data_path(data_base_dir, str(next_item["image"])).resolve() ) next_item["class"] = _normalize_modality( str(next_item.get("class", next_item.get("modality", default_modality))) ) next_item.pop("modality", None) staged.append(next_item) return staged def _stage_datalist( data_base_dir: Path, input_path: Path, output_path: Path, default_modality: str, ) -> tuple[Path, dict[str, Any]]: raw = _load_json(input_path) training, validation = _split_entries(raw) staged = { "training": _stage_entries(data_base_dir, training, default_modality), "validation": _stage_entries(data_base_dir, validation, default_modality), } return _write_json(output_path, staged), staged 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=10, ) except Exception: return "" return proc.stdout.strip() if proc.returncode == 0 else "" def _resolve_from_upstream(upstream_root: Path, value: str | None) -> str | None: if value in (None, ""): return value path = Path(str(value)).expanduser() if path.is_absolute(): return str(path) return str((upstream_root / path).resolve()) def _stage_configs(args: argparse.Namespace, upstream_root: Path) -> dict[str, Any]: model_def_src = upstream_root / "configs" / "config_network_rflow.json" env_src = upstream_root / "configs" / "environment_maisi_vae_train.json" train_src = upstream_root / "configs" / "config_maisi_vae_train.json" for path in (model_def_src, env_src, train_src): if not path.is_file(): raise FileNotFoundError(path) work_dir = args.output_dir.resolve() / "workflow" artifacts_dir = args.output_dir.resolve() / "artifacts" config_dir = work_dir / "configs" model_dir = artifacts_dir / "models" tfevent_path = artifacts_dir / "tfevent" staged_datalist_path, staged_datalist = _stage_datalist( args.data_base_dir.resolve(), args.datalist.resolve(), work_dir / "dataset.json", args.modality, ) model_def = copy.deepcopy(_load_json(model_def_src)) env_config = copy.deepcopy(_load_json(env_src)) train_config = copy.deepcopy(_load_json(train_src)) env_config["model_dir"] = str(model_dir) env_config["tfevent_path"] = str(tfevent_path) env_config["finetune"] = not args.train_from_scratch env_config["trained_autoencoder_path"] = ( str(args.trained_autoencoder_path.resolve()) if args.trained_autoencoder_path else _resolve_from_upstream(upstream_root, env_config.get("trained_autoencoder_path")) ) data_option = train_config.setdefault("data_option", {}) data_option["random_aug"] = args.random_aug data_option["spacing_type"] = args.spacing_type data_option["spacing"] = args.spacing data_option["select_channel"] = args.select_channel auto_train = train_config.setdefault("autoencoder_train", {}) auto_train["batch_size"] = args.batch_size auto_train["patch_size"] = args.patch_size auto_train["val_batch_size"] = args.val_batch_size auto_train["val_patch_size"] = args.val_patch_size auto_train["val_sliding_window_patch_size"] = args.val_sliding_window_patch_size auto_train["lr"] = args.lr auto_train["perceptual_weight"] = args.perceptual_weight auto_train["kl_weight"] = args.kl_weight auto_train["adv_weight"] = args.adv_weight auto_train["recon_loss"] = args.recon_loss auto_train["val_interval"] = args.val_interval auto_train["cache"] = args.cache_rate auto_train["amp"] = not args.no_amp auto_train["n_epochs"] = args.epochs if "autoencoder_def" in model_def and args.autoencoder_num_splits is not None: model_def["autoencoder_def"]["num_splits"] = args.autoencoder_num_splits artifacts_dir.mkdir(parents=True, exist_ok=True) return { "env_config": _write_json(config_dir / "environment_maisi_vae_train.json", env_config), "model_config": _write_json(config_dir / "config_maisi_vae_train.json", train_config), "model_def": _write_json(config_dir / "config_network_rflow.json", model_def), "datalist": _write_json(config_dir / "datalist_staged.json", staged_datalist), "staged_datalist": staged_datalist, "artifacts_dir": artifacts_dir, "model_dir": model_dir, "tfevent_path": tfevent_path, } def _namespace_from_staged(staged: dict[str, Any]) -> argparse.Namespace: ns = argparse.Namespace() for path_key in ("env_config", "model_def"): for key, value in _load_json(staged[path_key]).items(): setattr(ns, key, value) train_config = _load_json(staged["model_config"]) for section in ("data_option", "autoencoder_train"): for key, value in train_config.get(section, {}).items(): setattr(ns, key, value) return ns def _warmup_rule(epoch: int) -> float: if epoch < 10: return 0.01 if epoch < 20: return 0.1 return 1.0 def _loss_weighted_sum(args: argparse.Namespace, losses: dict[str, float]) -> float: return ( losses["recons_loss"] + args.kl_weight * losses["kl_loss"] + args.perceptual_weight * losses["p_loss"] ) def _run_training( args: argparse.Namespace, upstream_root: Path, staged: dict[str, Any] ) -> dict[str, Any]: if args.num_gpus != 1: raise ValueError("VAE finetuning runner currently supports exactly one CUDA GPU") sys.path.insert(0, str(upstream_root)) import torch from monai.data import CacheDataset, DataLoader from monai.inferers.inferer import SimpleInferer, SlidingWindowInferer from monai.losses.adversarial_loss import PatchAdversarialLoss from monai.losses.perceptual import PerceptualLoss from monai.networks.nets import PatchDiscriminator from monai.utils import set_determinism from scripts.download_model_data import download_model_data from scripts.transforms import VAE_Transform from scripts.utils import KL_loss, define_instance, dynamic_infer from torch.amp import GradScaler, autocast from torch.nn import L1Loss, MSELoss from torch.optim import lr_scheduler from torch.utils.tensorboard import SummaryWriter if args.download_model_data: previous_cwd = os.getcwd() try: os.chdir(upstream_root) download_model_data("rflow-ct", str(upstream_root), model_only=True) finally: os.chdir(previous_cwd) set_determinism(seed=args.random_seed) cfg = _namespace_from_staged(staged) device = torch.device("cuda") train_transform = VAE_Transform( is_train=True, random_aug=cfg.random_aug, k=4, patch_size=cfg.patch_size, val_patch_size=cfg.val_patch_size, output_dtype=torch.float16, spacing_type=cfg.spacing_type, spacing=cfg.spacing, image_keys=["image"], label_keys=[], additional_keys=[], select_channel=cfg.select_channel, ) val_transform = VAE_Transform( is_train=False, random_aug=False, k=4, val_patch_size=cfg.val_patch_size, output_dtype=torch.float16, image_keys=["image"], label_keys=[], additional_keys=[], select_channel=cfg.select_channel, ) staged_datalist = staged["staged_datalist"] dataset_train = CacheDataset( data=staged_datalist["training"], transform=train_transform, cache_rate=cfg.cache, num_workers=args.cache_num_workers, ) dataloader_train = DataLoader( dataset_train, batch_size=cfg.batch_size, num_workers=args.loader_num_workers, shuffle=True, drop_last=True, ) dataset_val = CacheDataset( data=staged_datalist["validation"], transform=val_transform, cache_rate=cfg.cache, num_workers=args.cache_num_workers, ) dataloader_val = DataLoader( dataset_val, batch_size=cfg.val_batch_size, num_workers=args.loader_num_workers, shuffle=False, ) if len(dataloader_train) == 0: raise ValueError("training dataloader is empty; add cases or reduce batch size") if len(dataloader_val) == 0: raise ValueError("validation dataloader is empty; add validation/testing cases") Path(cfg.model_dir).mkdir(parents=True, exist_ok=True) tensorboard_path = Path(cfg.tfevent_path) / "autoencoder" tensorboard_path.mkdir(parents=True, exist_ok=True) writer = SummaryWriter(str(tensorboard_path)) trained_g_path = Path(cfg.model_dir) / "autoencoder.pt" trained_d_path = Path(cfg.model_dir) / "discriminator.pt" autoencoder = define_instance(cfg, "autoencoder_def").to(device) discriminator = PatchDiscriminator( spatial_dims=cfg.spatial_dims, num_layers_d=3, channels=32, in_channels=1, out_channels=1, norm="INSTANCE", ).to(device) if cfg.finetune: checkpoint_autoencoder = torch.load(cfg.trained_autoencoder_path, map_location=device) if "unet_state_dict" in checkpoint_autoencoder: checkpoint_autoencoder = checkpoint_autoencoder["unet_state_dict"] autoencoder.load_state_dict(checkpoint_autoencoder) intensity_loss = MSELoss() if cfg.recon_loss == "l2" else L1Loss(reduction="mean") adv_loss = PatchAdversarialLoss(criterion="least_squares") loss_perceptual = ( PerceptualLoss(spatial_dims=3, network_type="squeeze", is_fake_3d=True, fake_3d_ratio=0.2) .eval() .to(device) ) optimizer_g = torch.optim.Adam( params=autoencoder.parameters(), lr=cfg.lr, eps=1e-6 if cfg.amp else 1e-8 ) optimizer_d = torch.optim.Adam( params=discriminator.parameters(), lr=cfg.lr, eps=1e-6 if cfg.amp else 1e-8 ) scheduler_g = lr_scheduler.LambdaLR(optimizer_g, lr_lambda=_warmup_rule) scheduler_d = lr_scheduler.LambdaLR(optimizer_d, lr_lambda=_warmup_rule) scaler_g = GradScaler("cuda", init_scale=2.0**8, growth_factor=1.5) if cfg.amp else None scaler_d = GradScaler("cuda", init_scale=2.0**8, growth_factor=1.5) if cfg.amp else None val_inferer = ( SlidingWindowInferer( roi_size=cfg.val_sliding_window_patch_size, sw_batch_size=1, progress=False, overlap=0.0, device=torch.device("cpu"), sw_device=device, ) if cfg.val_sliding_window_patch_size else SimpleInferer() ) history: list[dict[str, Any]] = [] best_val_loss = float("inf") best_paths: list[str] = [] total_step = 0 for epoch in range(cfg.n_epochs): autoencoder.train() discriminator.train() train_losses = {"recons_loss": 0.0, "kl_loss": 0.0, "p_loss": 0.0} adv_total = 0.0 for batch in dataloader_train: images = batch["image"].to(device).contiguous() optimizer_g.zero_grad(set_to_none=True) optimizer_d.zero_grad(set_to_none=True) with autocast("cuda", enabled=cfg.amp): reconstruction, z_mu, z_sigma = autoencoder(images) losses = { "recons_loss": intensity_loss(reconstruction, images), "kl_loss": KL_loss(z_mu, z_sigma), "p_loss": loss_perceptual(reconstruction.float(), images.float()), } logits_fake = discriminator(reconstruction.contiguous().float())[-1] generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False) loss_g = ( losses["recons_loss"] + cfg.kl_weight * losses["kl_loss"] + cfg.perceptual_weight * losses["p_loss"] + cfg.adv_weight * generator_loss ) if cfg.amp and scaler_g is not None: scaler_g.scale(loss_g).backward() scaler_g.unscale_(optimizer_g) scaler_g.step(optimizer_g) scaler_g.update() else: loss_g.backward() optimizer_g.step() logits_fake = discriminator(reconstruction.contiguous().detach())[-1] loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True) logits_real = discriminator(images.contiguous().detach())[-1] loss_d_real = adv_loss(logits_real, target_is_real=True, for_discriminator=True) loss_d = (loss_d_fake + loss_d_real) * 0.5 if cfg.amp and scaler_d is not None: scaler_d.scale(loss_d).backward() scaler_d.step(optimizer_d) scaler_d.update() else: loss_d.backward() optimizer_d.step() total_step += 1 for loss_name, loss_value in losses.items(): value = float(loss_value.item()) writer.add_scalar(f"train_{loss_name}_iter", value, total_step) train_losses[loss_name] += value adv_total += float(generator_loss.item()) writer.add_scalar("train_adv_loss_iter", float(generator_loss.item()), total_step) writer.add_scalar("train_fake_loss_iter", float(loss_d_fake.item()), total_step) writer.add_scalar("train_real_loss_iter", float(loss_d_real.item()), total_step) scheduler_g.step() scheduler_d.step() for key in train_losses: train_losses[key] /= len(dataloader_train) writer.add_scalar(f"train_{key}_epoch", train_losses[key], epoch) train_weighted = _loss_weighted_sum(cfg, train_losses) torch.save(autoencoder.state_dict(), trained_g_path) torch.save(discriminator.state_dict(), trained_d_path) epoch_record: dict[str, Any] = { "epoch": epoch, "train_losses": train_losses, "train_weighted_loss": train_weighted, "train_adv_loss": adv_total / len(dataloader_train), } if epoch % cfg.val_interval == 0: autoencoder.eval() val_losses = {"recons_loss": 0.0, "kl_loss": 0.0, "p_loss": 0.0} last_z_mu = None for batch in dataloader_val: with torch.no_grad(), autocast("cuda", enabled=cfg.amp): images = batch["image"].to(device).contiguous() reconstruction, z_mu, z_sigma = dynamic_infer(val_inferer, autoencoder, images) reconstruction = reconstruction.to(device) target = images val_losses["recons_loss"] += float( intensity_loss(reconstruction, target).item() ) val_losses["kl_loss"] += float(KL_loss(z_mu, z_sigma).item()) val_losses["p_loss"] += float(loss_perceptual(reconstruction, target).item()) last_z_mu = z_mu for key in val_losses: val_losses[key] /= len(dataloader_val) writer.add_scalar(key, val_losses[key], epoch) val_loss = _loss_weighted_sum(cfg, val_losses) epoch_record["val_losses"] = val_losses epoch_record["val_weighted_loss"] = val_loss if last_z_mu is not None: writer.add_scalar( "val_one_sample_scale_factor", float(1.0 / last_z_mu.flatten().std()), epoch ) if val_loss < best_val_loss: best_val_loss = val_loss best_path = Path(str(trained_g_path)[:-3] + f"_epoch{epoch}.pt") torch.save(autoencoder.state_dict(), best_path) best_paths.append(str(best_path)) epoch_record["best_autoencoder_checkpoint"] = str(best_path) history.append(epoch_record) writer.close() return { "autoencoder_checkpoint": str(trained_g_path), "discriminator_checkpoint": str(trained_d_path), "best_autoencoder_checkpoints": best_paths, "history": history, "tensorboard_dir": str(tensorboard_path), } def _write_workflow_summary(staged: dict[str, Any], result: dict[str, Any]) -> Path: summary = { "model": MODEL_NAME, "training_cases": len(staged["staged_datalist"].get("training", [])), "validation_cases": len(staged["staged_datalist"].get("validation", [])), "staged_configs": { "env_config": str(staged["env_config"]), "model_config": str(staged["model_config"]), "model_def": str(staged["model_def"]), "datalist": str(staged["datalist"]), }, **result, } return _write_json(staged["artifacts_dir"] / "workflow_summary.json", summary) def _build_command(args: argparse.Namespace) -> list[str]: command = [ sys.executable, str(Path(__file__).resolve()), str(args.datalist.resolve()), "--data-base-dir", str(args.data_base_dir.resolve()), "--output-dir", str(args.output_dir.resolve()), "--modality", args.modality, "--epochs", str(args.epochs), "--batch-size", str(args.batch_size), "--lr", str(args.lr), "--cache-rate", str(args.cache_rate), "--patch-size", ",".join(str(v) for v in args.patch_size), "--num-gpus", str(args.num_gpus), ] if args.download_model_data: command.append("--download-model-data") if args.train_from_scratch: command.append("--train-from-scratch") if args.no_amp: command.append("--no-amp") if args.preflight: command.append("--preflight") return command def _summarize_output(output_dir: Path) -> dict[str, Any]: artifacts_dir = output_dir / "artifacts" summary_path = artifacts_dir / "workflow_summary.json" summary = _load_json(summary_path) if summary_path.is_file() else {} autoencoder = Path(summary.get("autoencoder_checkpoint") or "") discriminator = Path(summary.get("discriminator_checkpoint") or "") best_paths = [Path(p) for p in summary.get("best_autoencoder_checkpoints", [])] return { "directory": str(output_dir), "artifacts_dir": str(artifacts_dir), "workflow_summary": str(summary_path) if summary_path.is_file() else None, "autoencoder_checkpoint": str(autoencoder) if str(autoencoder) else None, "autoencoder_checkpoint_present": autoencoder.is_file() if str(autoencoder) else False, "discriminator_checkpoint": str(discriminator) if str(discriminator) else None, "discriminator_checkpoint_present": ( discriminator.is_file() if str(discriminator) else False ), "best_autoencoder_checkpoints": [str(p) for p in best_paths], "num_best_autoencoder_checkpoints": len(best_paths), "loss_history": summary.get("history", []), "tensorboard_dir": summary.get("tensorboard_dir"), } def _empty_output(output_dir: Path) -> dict[str, Any]: return { "directory": str(output_dir), "artifacts_dir": str(output_dir / "artifacts"), "workflow_summary": None, "autoencoder_checkpoint": None, "autoencoder_checkpoint_present": False, "discriminator_checkpoint": None, "discriminator_checkpoint_present": False, "best_autoencoder_checkpoints": [], "num_best_autoencoder_checkpoints": 0, "loss_history": [], "tensorboard_dir": None, } def _payload( args: argparse.Namespace, dataset: dict[str, Any], upstream_root: Path | None, checked_roots: list[str], exit_code: int, elapsed: float, stdout: str = "", stderr: str = "", ) -> dict[str, Any]: output = ( _summarize_output(args.output_dir) if exit_code == 0 and not args.preflight else _empty_output(args.output_dir) ) return { "skill": SKILL_NAME, "model": MODEL_NAME, "model_repo": UPSTREAM_REPO, "license": "Apache-2.0", "input": { **dataset, "epochs": args.epochs, "batch_size": args.batch_size, "lr": args.lr, "cache_rate": args.cache_rate, "patch_size": args.patch_size, "val_patch_size": args.val_patch_size, "num_gpus": args.num_gpus, "finetune": not args.train_from_scratch, "random_seed": args.random_seed, }, "output": output, "invocation": { "official_entrypoint": UPSTREAM_ENTRYPOINT, "upstream_root": str(upstream_root) if upstream_root else None, "upstream_commit": _git_commit(upstream_root) if upstream_root else "", "checked_upstream_roots": checked_roots, "command": _build_command(args), "exit_code": exit_code, "subprocess_seconds": elapsed, }, "runtime": { "subprocess_seconds": elapsed, "device": "cuda" if args.num_gpus > 0 else "cpu", "preflight_only": bool(args.preflight), }, "logs": {"stdout_tail": _tail(stdout), "stderr_tail": _tail(stderr)}, "intended_use_disclaimer": ( "Engineering wrapper for synthetic-imaging VAE finetuning; not for clinical " "interpretation, regulatory use, or production training data approval." ), } def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("datalist", type=Path) parser.add_argument("--data-base-dir", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--upstream-root") parser.add_argument( "--modality", default="mri", help="Default modality/class for entries missing class/modality: ct or mri", ) parser.add_argument("--epochs", type=int, default=1) parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--val-batch-size", type=int, default=1) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--cache-rate", type=float, default=0.0) parser.add_argument("--patch-size", type=lambda s: _parse_triplet(s, int), default=[64, 64, 64]) parser.add_argument("--val-patch-size", type=lambda s: _parse_triplet(s, int)) parser.add_argument( "--val-sliding-window-patch-size", type=lambda s: _parse_triplet(s, int), default=[96, 96, 64], ) parser.add_argument("--autoencoder-num-splits", type=int, default=1) parser.add_argument("--num-gpus", type=int, default=1) parser.add_argument("--perceptual-weight", type=float, default=0.3) parser.add_argument("--kl-weight", type=float, default=1e-7) parser.add_argument("--adv-weight", type=float, default=0.1) parser.add_argument("--recon-loss", choices=("l1", "l2"), default="l1") parser.add_argument("--val-interval", type=int, default=1) parser.add_argument( "--spacing-type", choices=("original", "fixed", "rand_zoom"), default="original" ) parser.add_argument("--spacing", type=lambda s: _parse_triplet(s, float)) parser.add_argument("--select-channel", type=int, default=0) parser.add_argument("--cache-num-workers", type=int, default=0) parser.add_argument("--loader-num-workers", type=int, default=0) parser.add_argument("--random-seed", type=int, default=0) parser.add_argument("--trained-autoencoder-path", type=Path) parser.add_argument("--download-model-data", action="store_true") parser.add_argument("--train-from-scratch", action="store_true") parser.add_argument("--no-random-aug", dest="random_aug", action="store_false") parser.set_defaults(random_aug=True) parser.add_argument("--no-amp", action="store_true") parser.add_argument("--preflight", action="store_true") return parser def main() -> None: args = build_parser().parse_args() args.modality = _normalize_modality(args.modality) args.output_dir.mkdir(parents=True, exist_ok=True) dataset = _validate_datalist( args.data_base_dir.resolve(), args.datalist.resolve(), args.modality ) upstream_root, checked = _resolve_upstream_root(args.upstream_root) start = time.time() if upstream_root is None and args.preflight: payload = _payload(args, dataset, None, checked, 0, time.time() - start) payload["logs"]["stderr_tail"] = ( "Preflight did not find an NV-Generate-CTMR checkout containing VAE configs/helpers. " "A real training run requires NV_GENERATE_ROOT to point at a current checkout." ) _emit(payload) return if upstream_root is None: payload = _payload(args, dataset, None, checked, 2, time.time() - start) payload["logs"]["stderr_tail"] = ( "NV-Generate-CTMR checkout with VAE configs/helpers was not found. " "Set NV_GENERATE_ROOT or pass --upstream-root." ) _emit(payload) raise SystemExit(2) try: staged = _stage_configs(args, upstream_root) except Exception as exc: payload = _payload( args, dataset, upstream_root, checked, 2, time.time() - start, stderr=str(exc) ) _emit(payload) raise SystemExit(2) if args.preflight: payload = _payload(args, dataset, upstream_root, checked, 0, time.time() - start) payload["logs"][ "stderr_tail" ] = "Preflight staged VAE configs and validated datalist paths." _emit(payload) return stdout = "" try: result = _run_training(args, upstream_root, staged) _write_workflow_summary(staged, result) exit_code = 0 stderr = "" except Exception as exc: exit_code = 2 stderr = f"{type(exc).__name__}: {exc}" payload = _payload( args, dataset, upstream_root, checked, exit_code, time.time() - start, stdout, stderr ) _emit(payload) raise SystemExit(exit_code) if __name__ == "__main__": main()