Add SACFlow source code

sacflow/__init__.py

@@ -0,0 +1,2 @@
+"""SACFlow-FM research package."""
+__version__ = "0.1.0"

sacflow/data/__init__.py

@@ -0,0 +1 @@
+

sacflow/data/loader.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+from torch.utils.data import DataLoader, DistributedSampler, Sampler
+from .nifti_dataset import NiftiSegDataset
+from sacflow.utils.distributed import get_world_size, get_rank
+
+
+class DistributedEvalSamplerNoPad(Sampler):
+    """Shard evaluation data across ranks without padding/duplication.
+
+    PyTorch's DistributedSampler pads samples so every rank has equal length.
+    That is useful for training but biases validation metrics because some cases
+    are duplicated. This sampler uses rank::world_size indices exactly once.
+    """
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.rank = get_rank()
+        self.world_size = get_world_size()
+        self.indices = list(range(self.rank, len(dataset), self.world_size))
+
+    def __iter__(self):
+        return iter(self.indices)
+
+    def __len__(self):
+        return len(self.indices)
+
+
+def build_loader(cfg, split: str, training: bool, require_label: bool = False, distributed: bool | None = None):
+    """Build a NIfTI segmentation loader.
+
+    Important DDP behavior:
+    - Training loaders use DistributedSampler when world_size > 1.
+    - Evaluation/validation loaders default to *no* DistributedSampler. This is deliberate:
+      training-time validation is run only on rank 0, and standalone eval usually uses one rank.
+      Using a DistributedSampler for validation without metric all-gather biases metrics to a
+      rank-local subset.
+    """
+    data_cfg = cfg["data"]
+    ds = NiftiSegDataset(data_cfg["manifest"], split=split, cfg=data_cfg, training=training, require_label=require_label)
+    if distributed is None:
+        distributed = bool(training and get_world_size() > 1)
+    if distributed:
+        sampler = DistributedSampler(ds, shuffle=True) if training else DistributedEvalSamplerNoPad(ds)
+    else:
+        sampler = None
+    loader = DataLoader(
+        ds,
+        batch_size=data_cfg.get("batch_size" if training else "val_batch_size", 1),
+        shuffle=(training and sampler is None),
+        sampler=sampler,
+        num_workers=cfg.get("num_workers", 4),
+        pin_memory=cfg.get("pin_memory", True),
+        persistent_workers=cfg.get("num_workers", 4) > 0,
+    )
+    return loader

sacflow/data/nifti_dataset.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+import json
+import random
+import numpy as np
+import nibabel as nib
+import torch
+from torch.utils.data import Dataset
+
+
+def load_nifti(path: str | Path):
+    img = nib.load(str(path))
+    data = img.get_fdata(dtype=np.float32)
+    return data, img.affine, img.header
+
+
+def normalize_intensity(x: np.ndarray, mode: str = "zscore_nonzero", clip=None) -> np.ndarray:
+    x = x.astype(np.float32)
+    if clip is not None:
+        lo, hi = clip
+        x = np.clip(x, lo, hi)
+    if mode == "zscore_nonzero":
+        mask = np.abs(x) > 1e-6
+        if mask.sum() > 10:
+            mu, sd = x[mask].mean(), x[mask].std()
+        else:
+            mu, sd = x.mean(), x.std()
+        x = (x - mu) / (sd + 1e-6)
+    elif mode == "zscore":
+        x = (x - x.mean()) / (x.std() + 1e-6)
+    elif mode == "minmax":
+        x = (x - x.min()) / (x.max() - x.min() + 1e-6)
+    elif mode in ("none", None):
+        pass
+    else:
+        raise ValueError(f"Unknown intensity mode: {mode}")
+    return x.astype(np.float32)
+
+
+def _pad_to_shape(arr: np.ndarray, shape: Tuple[int, int, int], value=0):
+    pads = []
+    for dim, target in zip(arr.shape[-3:], shape):
+        total = max(0, target - dim)
+        before = total // 2
+        after = total - before
+        pads.append((before, after))
+    if arr.ndim == 4:
+        pad_width = [(0, 0)] + pads
+    else:
+        pad_width = pads
+    return np.pad(arr, pad_width, mode="constant", constant_values=value)
+
+
+def random_crop_pair(img: np.ndarray, lab: Optional[np.ndarray], patch_size: Tuple[int, int, int], foreground_prob: float = 0.5):
+    # img: [C,H,W,D], lab: [H,W,D]
+    img = _pad_to_shape(img, patch_size, 0)
+    if lab is not None:
+        lab = _pad_to_shape(lab, patch_size, 0)
+    H, W, D = img.shape[-3:]
+    ph, pw, pd = patch_size
+    if lab is not None and random.random() < foreground_prob and (lab > 0).sum() > 0:
+        coords = np.argwhere(lab > 0)
+        center = coords[random.randrange(len(coords))]
+        starts = []
+        for c, dim, p in zip(center, (H, W, D), patch_size):
+            s = int(c) - p // 2
+            s = max(0, min(s, dim - p))
+            starts.append(s)
+    else:
+        starts = [random.randint(0, max(0, dim - p)) for dim, p in zip((H, W, D), patch_size)]
+    sh, sw, sd = starts
+    img_c = img[:, sh:sh+ph, sw:sw+pw, sd:sd+pd]
+    lab_c = None if lab is None else lab[sh:sh+ph, sw:sw+pw, sd:sd+pd]
+    return img_c, lab_c
+
+
+def center_crop_pair(img: np.ndarray, lab: Optional[np.ndarray], patch_size: Tuple[int, int, int]):
+    img = _pad_to_shape(img, patch_size, 0)
+    if lab is not None:
+        lab = _pad_to_shape(lab, patch_size, 0)
+    H, W, D = img.shape[-3:]
+    ph, pw, pd = patch_size
+    sh, sw, sd = [(dim - p)//2 for dim, p in zip((H,W,D), patch_size)]
+    img_c = img[:, sh:sh+ph, sw:sw+pw, sd:sd+pd]
+    lab_c = None if lab is None else lab[sh:sh+ph, sw:sw+pw, sd:sd+pd]
+    return img_c, lab_c
+
+
+def augment(img: np.ndarray, lab: Optional[np.ndarray], cfg: Dict):
+    if cfg.get("random_flip", False):
+        for ax in range(3):
+            if random.random() < 0.5:
+                img = np.flip(img, axis=ax+1).copy()
+                if lab is not None:
+                    lab = np.flip(lab, axis=ax).copy()
+    shift = float(cfg.get("random_intensity_shift", 0.0) or 0.0)
+    scale = float(cfg.get("random_intensity_scale", 0.0) or 0.0)
+    if shift > 0:
+        img = img + np.random.uniform(-shift, shift)
+    if scale > 0:
+        img = img * np.random.uniform(1-scale, 1+scale)
+    return img, lab
+
+
+class NiftiSegDataset(Dataset):
+    def __init__(self, manifest: str | Path, split: str, cfg: Dict, training: bool = True, require_label: bool = False):
+        self.manifest_path = Path(manifest)
+        with open(self.manifest_path, "r") as f:
+            man = json.load(f)
+        if split not in man:
+            raise KeyError(f"Split {split} not found in {manifest}")
+        self.items = man[split]
+        self.split = split
+        self.cfg = cfg
+        self.training = training
+        self.require_label = require_label
+        self.patch_size = tuple(cfg.get("patch_size", [96,96,96]))
+        self.intensity_cfg = cfg.get("intensity", {})
+        self.aug_cfg = cfg.get("augmentation", {})
+
+    def __len__(self):
+        return len(self.items)
+
+    def __getitem__(self, idx):
+        item = self.items[idx]
+        img, affine, header = load_nifti(item["image"])
+        if img.ndim == 3:
+            img = img[None]
+        elif img.ndim == 4:
+            # assume channels last if last dim small
+            if img.shape[-1] <= 8:
+                img = np.moveaxis(img, -1, 0)
+            else:
+                img = img[None, ..., 0]
+        img = normalize_intensity(img, self.intensity_cfg.get("normalize", "zscore_nonzero"), self.intensity_cfg.get("clip"))
+        lab = None
+        if item.get("label"):
+            lab, _, _ = load_nifti(item["label"])
+            lab = lab.astype(np.int64)
+        elif self.require_label:
+            raise ValueError(f"Item {idx} has no label in split {self.split}")
+        if self.training:
+            img, lab = random_crop_pair(img, lab, self.patch_size)
+            img, lab = augment(img, lab, self.aug_cfg)
+        else:
+            # keep full volume for sliding window eval; no crop
+            pass
+        batch = {
+            "image": torch.from_numpy(img.copy()).float(),
+            "case_id": item.get("id", Path(item["image"]).stem),
+            "image_path": item["image"],
+        }
+        if lab is not None:
+            batch["label"] = torch.from_numpy(lab.copy()).long()
+        batch["spacing"] = torch.tensor(header.get_zooms()[:3] if header is not None else (1,1,1), dtype=torch.float32)
+        return batch
+
+
+def load_manifest(path: str | Path):
+    with open(path, "r") as f:
+        return json.load(f)

sacflow/engine/__init__.py

@@ -0,0 +1 @@
+

sacflow/engine/train_loop.py

@@ -0,0 +1,398 @@
+from __future__ import annotations
+from pathlib import Path
+import time
+import copy
+import re
+import torch
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+from tqdm import tqdm
+from monai.inferers import sliding_window_inference
+from sacflow.data.loader import build_loader
+from sacflow.models.unet3d import build_model, freeze_except_adapters
+from sacflow.models.velocity_field import VelocityField3D
+from sacflow.methods.sacflow_step import sacflow_forward_step, ce_loss_masked, dice_loss_masked
+from sacflow.methods.source_memory import load_source_memory, class_moments
+from sacflow.utils.metrics import torch_soft_dice_loss, entropy_loss, confidence_and_margin, dice_per_class, hd95_per_class
+from sacflow.utils.misc import ensure_dir, count_trainable, move_to_device, unwrap_model
+from sacflow.utils.distributed import is_main_process, get_world_size, get_rank, reduce_mean, barrier, is_dist_avail_and_initialized
+import torch.distributed as dist
+from sacflow.utils.wandb_utils import wandb_log
+
+
+def build_optimizer(params, cfg):
+    ocfg = cfg["optim"]
+    params = [p for p in params if p.requires_grad]
+    if ocfg.get("optimizer", "adamw").lower() == "sgd":
+        return torch.optim.SGD(params, lr=float(ocfg["lr"]), momentum=0.9, weight_decay=float(ocfg.get("weight_decay", 0)))
+    return torch.optim.AdamW(params, lr=float(ocfg["lr"]), weight_decay=float(ocfg.get("weight_decay", 0)), betas=tuple(ocfg.get("betas", [0.9, 0.999])))
+
+
+def update_ema(teacher, student, decay):
+    with torch.no_grad():
+        for pt, ps in zip(teacher.parameters(), student.parameters()):
+            pt.data.mul_(decay).add_(ps.data, alpha=1-decay)
+
+
+def load_checkpoint_into(model, path, strict=False):
+    ckpt = torch.load(path, map_location="cpu")
+    state = ckpt.get("model", ckpt)
+    missing, unexpected = model.load_state_dict(state, strict=strict)
+    return missing, unexpected
+
+
+def save_checkpoint(path, model, optimizer, epoch, step, best_metric=None, velocity_field=None, cfg=None, teacher=None, include_optimizer=True):
+    """Save a checkpoint on rank 0 only.
+
+    Disk policy:
+    - best.pt is intended for evaluation/inference and is saved without optimizer by default.
+    - last.pt is intended for resume and includes optimizer.
+    This avoids filling the disk with epoch_N.pt checkpoints.
+    """
+    if not is_main_process():
+        return
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    obj = {
+        "model": unwrap_model(model).state_dict(),
+        "epoch": epoch,
+        "step": step,
+        "best_metric": best_metric,
+        "cfg": cfg,
+    }
+    if include_optimizer and optimizer is not None:
+        obj["optimizer"] = optimizer.state_dict()
+    if velocity_field is not None:
+        obj["velocity_field"] = unwrap_model(velocity_field).state_dict()
+    if teacher is not None:
+        obj["teacher"] = unwrap_model(teacher).state_dict()
+    torch.save(obj, path)
+
+
+def _epoch_number(path: Path) -> int:
+    m = re.search(r"epoch_(\d+)\.pt$", path.name)
+    return int(m.group(1)) if m else -1
+
+
+def resolve_resume_checkpoint(ckpt_dir: Path, resume_value):
+    """Return a usable resume checkpoint path.
+
+    resume_value can be:
+    - None/False: do not resume
+    - "auto"/True: prefer last.pt, then newest epoch_*.pt, then best.pt
+    - explicit checkpoint path
+    Corrupted/incomplete checkpoints are skipped.
+    """
+    if not resume_value:
+        return None
+    if str(resume_value).lower() not in ("auto", "true", "1", "yes"):
+        return Path(resume_value)
+    candidates = []
+    last = ckpt_dir / "last.pt"
+    if last.exists():
+        candidates.append(last)
+    candidates.extend(sorted(ckpt_dir.glob("epoch_*.pt"), key=_epoch_number, reverse=True))
+    best = ckpt_dir / "best.pt"
+    if best.exists():
+        candidates.append(best)
+    for c in candidates:
+        try:
+            torch.load(c, map_location="cpu")
+            return c
+        except Exception as e:
+            if is_main_process():
+                print(f"Skipping unusable checkpoint {c}: {e}")
+    return None
+
+
+def load_training_checkpoint(path, model, optimizer=None, velocity_field=None, teacher=None):
+    ckpt = torch.load(path, map_location="cpu")
+    missing, unexpected = unwrap_model(model).load_state_dict(ckpt.get("model", ckpt), strict=False)
+    if is_main_process():
+        print(f"Loaded resume model from {path} missing={len(missing)} unexpected={len(unexpected)}")
+    if velocity_field is not None and "velocity_field" in ckpt:
+        unwrap_model(velocity_field).load_state_dict(ckpt["velocity_field"], strict=False)
+    if teacher is not None and "teacher" in ckpt:
+        unwrap_model(teacher).load_state_dict(ckpt["teacher"], strict=False)
+    elif teacher is not None:
+        unwrap_model(teacher).load_state_dict(unwrap_model(model).state_dict(), strict=False)
+    if optimizer is not None and ckpt.get("optimizer") is not None:
+        optimizer.load_state_dict(ckpt["optimizer"])
+    start_epoch = int(ckpt.get("epoch", 0))
+    global_step = int(ckpt.get("step", 0))
+    best = float(ckpt.get("best_metric", -1e9) if ckpt.get("best_metric", None) is not None else -1e9)
+    return start_epoch, global_step, best
+
+
+def supervised_step(model, batch, cfg):
+    x = batch["image"]
+    y = batch["label"]
+    logits = model(x)
+    ce = F.cross_entropy(logits, y.long())
+    dice = torch_soft_dice_loss(logits, y, cfg["data"]["num_classes"])
+    loss = cfg["train"].get("loss", {}).get("ce", 1.0)*ce + cfg["train"].get("loss", {}).get("dice", 1.0)*dice
+    return loss, {"loss_total": loss.detach(), "loss_ce": ce.detach(), "loss_dice": dice.detach()}
+
+
+
+def proto_align_step(model, teacher, batch, memory, cfg):
+    x = batch["image"]
+    logits, feats = model(x, return_features=True)
+    feat = feats["prelogit"]
+    with torch.no_grad():
+        tlogits = teacher(x)
+        tprobs = torch.softmax(tlogits, dim=1)
+        conf, margin, pseudo = confidence_and_margin(tprobs)
+        mask = conf > float(cfg["train"].get("pseudo_conf_threshold", 0.75))
+        probs_f = tprobs
+        if probs_f.shape[-3:] != feat.shape[-3:]:
+            probs_f = F.interpolate(probs_f, size=feat.shape[-3:], mode="trilinear", align_corners=False)
+    ce = ce_loss_masked(logits, pseudo, mask)
+    dice = dice_loss_masked(logits, pseudo, mask, cfg["data"]["num_classes"])
+    proto_loss = torch.tensor(0.0, device=x.device)
+    if memory is not None and "feature_mu" in memory:
+        mu = memory["feature_mu"].to(feat.device, feat.dtype)  # [C,d]
+        # expected source prototype at each voxel based on teacher probabilities
+        proto = torch.einsum("bchwz,cf->bfhwz", probs_f.detach(), mu)
+        proto_loss = ((feat - proto).pow(2) * probs_f.max(1, keepdim=True).values.detach()).mean()
+    ent = entropy_loss(logits)
+    loss_cfg = cfg["train"].get("loss", {})
+    loss = float(loss_cfg.get("ce", 1.0))*ce + float(loss_cfg.get("dice", 1.0))*dice + float(loss_cfg.get("prototype", 0.1))*proto_loss + float(loss_cfg.get("entropy", 0.01))*ent
+    return loss, {"loss_total": loss.detach(), "loss_pseudo_ce": ce.detach(), "loss_pseudo_dice": dice.detach(), "loss_proto_align": proto_loss.detach(), "loss_entropy": ent.detach(), "pseudo_conf_mean": conf.mean().detach(), "pseudo_accept_rate": mask.float().mean().detach()}
+
+
+def pseudo_step(model, teacher, batch, cfg):
+    x = batch["image"]
+    with torch.no_grad():
+        tlogits = teacher(x)
+        tprobs = torch.softmax(tlogits, dim=1)
+        conf, margin, pseudo = confidence_and_margin(tprobs)
+        mask = conf > float(cfg["train"].get("pseudo_conf_threshold", 0.75))
+    logits = model(x)
+    ce = ce_loss_masked(logits, pseudo, mask)
+    dice = dice_loss_masked(logits, pseudo, mask, cfg["data"]["num_classes"])
+    ent = entropy_loss(logits)
+    loss_cfg = cfg["train"].get("loss", {})
+    loss = float(loss_cfg.get("ce", 1.0))*ce + float(loss_cfg.get("dice", 1.0))*dice + float(loss_cfg.get("entropy", 0.01))*ent
+    return loss, {"loss_total": loss.detach(), "loss_pseudo_ce": ce.detach(), "loss_pseudo_dice": dice.detach(), "loss_entropy": ent.detach(), "pseudo_conf_mean": conf.mean().detach(), "pseudo_accept_rate": mask.float().mean().detach()}
+
+
+@torch.no_grad()
+def evaluate(model, loader, cfg, device, max_batches=None):
+    """Evaluate segmentation metrics.
+
+    In DDP this function is called on *all* ranks with a no-padding sharded
+    validation loader. It then all-reduces metric sums/counts so rank 0 gets
+    exact full-validation metrics without other ranks idling at a barrier.
+    """
+    model.eval()
+    all_metrics = []
+    num_classes = cfg["data"]["num_classes"]
+    roi_size = tuple(cfg.get("eval", {}).get("roi_size", cfg["data"].get("patch_size", [96,96,96])))
+    sw_batch_size = int(cfg.get("eval", {}).get("sw_batch_size", 1))
+    overlap = float(cfg.get("eval", {}).get("overlap", 0.5))
+    iterator = enumerate(loader)
+    if is_main_process():
+        iterator = tqdm(iterator, total=len(loader), desc="eval", leave=False)
+    for i, batch in iterator:
+        if max_batches is not None and i >= max_batches:
+            break
+        if "label" not in batch:
+            continue
+        x = batch["image"].to(device, non_blocking=True)
+        y = batch["label"].numpy()
+        if cfg.get("eval", {}).get("sliding_window", True):
+            logits = sliding_window_inference(x, roi_size=roi_size, sw_batch_size=sw_batch_size, predictor=model, overlap=overlap)
+        else:
+            logits = model(x)
+        pred = logits.argmax(1).cpu().numpy()
+        for b in range(pred.shape[0]):
+            m = {}
+            m.update(dice_per_class(pred[b], y[b], num_classes))
+            spacing = tuple(batch.get("spacing", torch.ones(1,3))[b].cpu().numpy().tolist()) if "spacing" in batch else (1,1,1)
+            m.update(hd95_per_class(pred[b], y[b], num_classes, spacing=spacing))
+            all_metrics.append(m)
+
+    metric_keys = [f"dice_c{c}" for c in range(1, num_classes)] + ["dice_mean"] + [f"hd95_c{c}" for c in range(1, num_classes)] + ["hd95_mean"]
+    sums = torch.zeros(len(metric_keys), device=device, dtype=torch.float64)
+    counts = torch.zeros(len(metric_keys), device=device, dtype=torch.float64)
+    for m in all_metrics:
+        for j, k in enumerate(metric_keys):
+            v = m.get(k, float("nan"))
+            if v == v:  # not NaN
+                sums[j] += float(v)
+                counts[j] += 1.0
+    if is_dist_avail_and_initialized():
+        dist.all_reduce(sums, op=dist.ReduceOp.SUM)
+        dist.all_reduce(counts, op=dist.ReduceOp.SUM)
+    out = {}
+    for j, k in enumerate(metric_keys):
+        if counts[j].item() > 0:
+            out[f"val/{k}"] = float((sums[j] / counts[j]).item())
+    if not out:
+        out["val/dice_mean"] = float("nan")
+    return out
+
+
+def run_training(cfg, device, wandb_run=None):
+    mode = cfg["train"]["mode"]
+    out_dir = ensure_dir(cfg["output_dir"])
+    ckpt_dir = ensure_dir(out_dir / "checkpoints")
+    require_label = mode in ("source_train", "oracle_train")
+    split = "source_train" if mode == "source_train" else ("target_train" if mode in ("oracle_train", "self_train", "peft", "sacflow_fm", "proto_align") else "target_train")
+    train_loader = build_loader(cfg, split=split, training=True, require_label=require_label)
+    val_split = "source_val" if mode == "source_train" else "target_val"
+    try:
+        val_loader = build_loader(cfg, split=val_split, training=False, require_label=True, distributed=(get_world_size() > 1))
+    except Exception:
+        val_loader = None
+    model = build_model(cfg).to(device)
+    if cfg["train"].get("source_checkpoint"):
+        missing, unexpected = load_checkpoint_into(model, cfg["train"]["source_checkpoint"], strict=False)
+        if is_main_process():
+            print("Loaded source checkpoint", cfg["train"]["source_checkpoint"], "missing", len(missing), "unexpected", len(unexpected))
+    if mode in ("peft", "sacflow_fm", "proto_align") and cfg.get("model", {}).get("adapter", {}).get("enabled", False):
+        freeze_except_adapters(model, train_norm_affine=True)
+    teacher = copy.deepcopy(model).to(device)
+    for p in teacher.parameters():
+        p.requires_grad = False
+    velocity_field = None
+    memory = None
+    if mode in ("sacflow_fm", "proto_align"):
+        if cfg["train"].get("memory_path"):
+            memory = load_source_memory(cfg["train"]["memory_path"], device=device)
+        feat_ch = model.prelogit_channels
+        vcfg = cfg.get("sacflow", {}).get("velocity", {})
+        if cfg.get("sacflow", {}).get("use_velocity_field", True):
+            velocity_field = VelocityField3D(
+                residual_channels=feat_ch,
+                num_classes=cfg["data"]["num_classes"],
+                hidden_ratio=float(vcfg.get("hidden_ratio", 0.25)),
+                depth=int(vcfg.get("depth", 2)),
+                tau_embedding_dim=int(vcfg.get("tau_embedding_dim", 32)),
+                organ_embedding_dim=int(vcfg.get("organ_embedding_dim", 16)),
+                include_teacher_probs=bool(vcfg.get("include_teacher_probs", True)),
+                include_confidence=bool(vcfg.get("include_confidence", True)),
+                include_boundary=bool(vcfg.get("include_boundary", True)),
+                use_depthwise=bool(vcfg.get("use_depthwise", True)),
+                use_group_norm=bool(vcfg.get("use_group_norm", True)),
+                use_film=bool(vcfg.get("use_film", True)),
+            ).to(device)
+    params = list(model.parameters()) + ([] if velocity_field is None else list(velocity_field.parameters()))
+    optimizer = build_optimizer(params, cfg)
+    if get_world_size() > 1:
+        # SACFlow uses a custom feature-path loss in addition to the ordinary forward pass.
+        # find_unused_parameters=True is safer for this mode because some classifier outputs
+        # from the first forward are not directly used in the loss, while PEFT adapters are
+        # used again for path-state classification.
+        find_unused = bool(cfg.get("distributed", {}).get("find_unused_parameters", False)) or mode == "sacflow_fm"
+        model = DDP(model, device_ids=[device.index] if device.type == "cuda" else None, find_unused_parameters=find_unused)
+        if velocity_field is not None:
+            velocity_field = DDP(velocity_field, device_ids=[device.index] if device.type == "cuda" else None, find_unused_parameters=True)
+    trainable, total = count_trainable(unwrap_model(model))
+    if velocity_field is not None:
+        vt, vtotal = count_trainable(unwrap_model(velocity_field))
+        trainable += vt
+        total += vtotal
+    if is_main_process():
+        print(f"Mode={mode} trainable={trainable:,} total={total:,} ({100*trainable/max(1,total):.2f}%)")
+    scaler = torch.cuda.amp.GradScaler(enabled=bool(cfg.get("amp", True)) and device.type == "cuda")
+    best = -1e9
+    global_step = 0
+    start_epoch = 0
+    resume_value = cfg.get("train", {}).get("resume_checkpoint")
+    resume_path = resolve_resume_checkpoint(ckpt_dir, resume_value)
+    if resume_path is not None:
+        start_epoch, global_step, best = load_training_checkpoint(
+            resume_path, model, optimizer=optimizer, velocity_field=velocity_field, teacher=teacher
+        )
+        if is_main_process():
+            print(f"Resuming from epoch={start_epoch}, step={global_step}, best={best:.6f}")
+    elif resume_value and is_main_process():
+        print(f"WARNING: requested resume={resume_value!r}, but no usable checkpoint was found in {ckpt_dir}")
+    epochs = int(cfg["train"].get("epochs", 100))
+    steps_per_epoch = int(cfg["train"].get("steps_per_epoch", len(train_loader)))
+    if start_epoch >= epochs and is_main_process():
+        print(f"Checkpoint epoch {start_epoch} is already >= configured epochs {epochs}; nothing to train.")
+    for epoch in range(start_epoch, epochs):
+        if hasattr(train_loader.sampler, "set_epoch"):
+            train_loader.sampler.set_epoch(epoch)
+        model.train()
+        if velocity_field is not None:
+            velocity_field.train()
+        iterator = iter(train_loader)
+        pbar = range(steps_per_epoch)
+        if is_main_process():
+            pbar = tqdm(pbar, desc=f"epoch {epoch+1}/{epochs}", dynamic_ncols=True)
+        for _ in pbar:
+            try:
+                batch = next(iterator)
+            except StopIteration:
+                iterator = iter(train_loader)
+                batch = next(iterator)
+            batch = move_to_device(batch, device)
+            optimizer.zero_grad(set_to_none=True)
+            with torch.cuda.amp.autocast(enabled=bool(cfg.get("amp", True)) and device.type == "cuda"):
+                if mode in ("source_train", "oracle_train"):
+                    loss, logs = supervised_step(model, batch, cfg)
+                elif mode in ("self_train", "peft"):
+                    loss, logs = pseudo_step(model, teacher, batch, cfg)
+                elif mode == "proto_align":
+                    loss, logs = proto_align_step(model, teacher, batch, memory, cfg)
+                elif mode == "sacflow_fm":
+                    loss, logs = sacflow_forward_step(model, teacher, velocity_field, batch, memory, cfg)
+                else:
+                    raise ValueError(f"Unknown train mode {mode}")
+            scaler.scale(loss).backward()
+            if float(cfg["optim"].get("grad_clip_norm", 0) or 0) > 0:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_([p for p in params if p.requires_grad], float(cfg["optim"].get("grad_clip_norm")))
+            scaler.step(optimizer)
+            scaler.update()
+            if mode in ("self_train", "peft", "sacflow_fm", "proto_align"):
+                update_ema(teacher, unwrap_model(model), float(cfg["train"].get("ema_decay", 0.995)))
+            global_step += 1
+            red_logs = {}
+            for k, v in logs.items():
+                if torch.is_tensor(v):
+                    red_logs[f"train/{k}"] = float(reduce_mean(v.float()).item())
+                else:
+                    red_logs[f"train/{k}"] = v
+            if is_main_process() and global_step % int(cfg["train"].get("log_every", 20)) == 0:
+                red_logs["train/epoch"] = epoch + 1
+                red_logs["train/lr"] = optimizer.param_groups[0]["lr"]
+                wandb_log(wandb_run, red_logs, step=global_step)
+                if hasattr(pbar, "set_postfix"):
+                    pbar.set_postfix({"loss": f"{red_logs.get('train/loss_total', 0):.4f}", "step": global_step})
+        if val_loader is not None and ((epoch + 1) % int(cfg["train"].get("val_every", 1)) == 0):
+            metrics = evaluate(unwrap_model(model), val_loader, cfg, device)
+            if is_main_process():
+                score = metrics.get("val/dice_mean", -1e9)
+                print(f"Epoch {epoch+1} validation: {metrics}")
+                wandb_log(wandb_run, metrics, step=global_step)
+                if score > best:
+                    best = score
+                    save_checkpoint(
+                        ckpt_dir / "best.pt", model, optimizer, epoch+1, global_step, best,
+                        velocity_field, cfg, teacher=teacher,
+                        include_optimizer=bool(cfg["train"].get("save_optimizer_in_best", False)),
+                    )
+            # Keep all ranks synchronized after rank0 best-checkpoint writing.
+            barrier()
+        # Always keep a resumable last.pt. Do not create epoch_N.pt files unless explicitly requested.
+        save_checkpoint(
+            ckpt_dir / "last.pt", model, optimizer, epoch+1, global_step, best,
+            velocity_field, cfg, teacher=teacher, include_optimizer=True,
+        )
+        # Ensure rank0 has finished writing last.pt before other ranks start the next epoch.
+        barrier()
+        if bool(cfg["train"].get("keep_epoch_checkpoints", False)) and int(cfg["train"].get("save_every", 0) or 0) > 0:
+            if (epoch + 1) % int(cfg["train"].get("save_every", 5)) == 0:
+                save_checkpoint(
+                    ckpt_dir / f"epoch_{epoch+1}.pt", model, optimizer, epoch+1, global_step, best,
+                    velocity_field, cfg, teacher=teacher, include_optimizer=True,
+                )
+                barrier()
+    # final last.pt is already saved after every epoch; save once more for completeness.
+    save_checkpoint(ckpt_dir / "last.pt", model, optimizer, epochs, global_step, best, velocity_field, cfg, teacher=teacher, include_optimizer=True)
+    return unwrap_model(model)

sacflow/methods/__init__.py

@@ -0,0 +1 @@
+

sacflow/methods/sacflow_step.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+from typing import Dict, Tuple
+import torch
+import torch.nn.functional as F
+from sacflow.utils.metrics import torch_soft_dice_loss, entropy_loss, confidence_and_margin, finite_difference_boundary
+from sacflow.methods.task_space import centered_classifier_basis, random_basis, project_task_and_residual, project_to_residual
+from sacflow.methods.source_memory import class_moments, moment_transport_residual
+from sacflow.utils.misc import unwrap_model
+
+
+def ce_loss_masked(logits, target, mask=None):
+    loss = F.cross_entropy(logits, target.long(), reduction="none")
+    if mask is not None:
+        loss = loss * mask.float()
+        return loss.sum() / (mask.float().sum() + 1e-6)
+    return loss.mean()
+
+
+
+
+def dice_loss_masked(logits, target, mask=None, num_classes=None, eps=1e-5):
+    if num_classes is None:
+        num_classes = logits.shape[1]
+    probs = torch.softmax(logits, dim=1)
+    target = target.clamp(0, num_classes - 1).long()
+    onehot = torch.nn.functional.one_hot(target, num_classes).permute(0,4,1,2,3).float()
+    if mask is not None:
+        m = mask.float().unsqueeze(1)
+        probs = probs * m
+        onehot = onehot * m
+    dims = tuple(range(2, logits.ndim))
+    inter = (probs * onehot).sum(dims)
+    denom = probs.sum(dims) + onehot.sum(dims)
+    dice = (2 * inter + eps) / (denom + eps)
+    return 1.0 - dice[:, 1:].mean()
+
+
+def kl_masked(p, q, mask=None, eps=1e-8):
+    # p,q probabilities [B,C,H,W,D]
+    kl = (p * ((p+eps).log() - (q+eps).log())).sum(dim=1)
+    if mask is not None:
+        kl = kl * mask.float()
+        return kl.sum() / (mask.float().sum() + 1e-6)
+    return kl.mean()
+
+
+def _stats_distance(mu_a, std_a, mu_b, std_b):
+    # simple differentiable class/channel statistic distance
+    return (mu_a - mu_b).abs().mean() + (std_a - std_b).abs().mean()
+
+
+def build_task_basis(model, cfg, feat_dim, device, dtype):
+    basis = cfg.get("sacflow", {}).get("basis", "centered_svd")
+    W = model.final_classifier_weight().detach().to(device=device, dtype=dtype)
+    if basis == "random":
+        rank = max(1, min(W.shape[0]-1, feat_dim))
+        return random_basis(feat_dim, rank, device, dtype)
+    return centered_classifier_basis(W, foreground_only=cfg.get("sacflow", {}).get("foreground_only_basis", False))
+
+
+def make_tau(B, device, dtype):
+    return torch.rand(B, device=device, dtype=dtype)
+
+
+def sacflow_forward_step(model, teacher, velocity_field, batch, memory, cfg):
+    x = batch["image"]
+    num_classes = cfg["data"]["num_classes"]
+    # Use the wrapped model for the main forward when DDP is active; unwrap only for helper methods.
+    base_model = unwrap_model(model)
+    logits, feats = model(x, return_features=True)
+    feat = feats["prelogit"]
+    B, d, H, W, D = feat.shape
+    with torch.no_grad():
+        tlogits = teacher(x)
+        tprobs = torch.softmax(tlogits, dim=1)
+        if tprobs.shape[-3:] != (H,W,D):
+            tprobs_f = F.interpolate(tprobs, size=(H,W,D), mode="trilinear", align_corners=False)
+        else:
+            tprobs_f = tprobs
+        conf, margin, pseudo = confidence_and_margin(tprobs)
+        mask = (conf > float(cfg["train"].get("pseudo_conf_threshold", 0.75))).float()
+    Q = build_task_basis(base_model, cfg, d, feat.device, feat.dtype)
+    F_task, Rt = project_task_and_residual(feat, Q)
+    if cfg.get("sacflow", {}).get("flow_space", "residual") == "whole_feature":
+        F_task = torch.zeros_like(feat)
+        Rt = feat
+        Q = torch.empty(d, 0, device=feat.device, dtype=feat.dtype)
+    use_mem = cfg.get("sacflow", {}).get("use_compact_memory", True) and memory is not None
+    whole_feature = cfg.get("sacflow", {}).get("flow_space", "residual") == "whole_feature"
+    if use_mem and whole_feature and "feature_mu" in memory:
+        src_mu = memory["feature_mu"].to(feat.device, feat.dtype)
+        src_std = memory["feature_std"].to(feat.device, feat.dtype)
+    elif use_mem and "residual_mu" in memory:
+        src_mu = memory["residual_mu"].to(feat.device, feat.dtype)
+        src_std = memory["residual_std"].to(feat.device, feat.dtype)
+    else:
+        # strict fallback: use target stats as weak source proxy; ablate this separately.
+        # This makes R0 close to Rt and is intentionally weaker than compact-memory SACFlow.
+        src_mu, src_std, _ = class_moments(Rt.detach(), tprobs_f.detach())
+    R0, tgt_mu, tgt_std = moment_transport_residual(Rt.detach(), tprobs_f.detach(), src_mu, src_std)
+    R1 = Rt
+    tau = make_tau(B, feat.device, feat.dtype)
+    tau_view = tau.view(B,1,1,1,1)
+    sigma_tau = float(cfg.get("sacflow", {}).get("sigma_tau", 0.0) or 0.0)
+    noise = torch.randn_like(Rt) if sigma_tau > 0 else torch.zeros_like(Rt)
+    R_interp = (1 - tau_view) * R0 + tau_view * R1 + sigma_tau * tau_view * (1 - tau_view) * noise
+    u = R1 - R0 + sigma_tau * (1 - 2*tau_view) * noise
+    anchors = {"probs": tprobs_f.detach(), "boundary": finite_difference_boundary(tprobs_f.detach())}
+    use_v = bool(cfg.get("sacflow", {}).get("use_velocity_field", True))
+    if use_v and velocity_field is not None:
+        v = velocity_field(R_interp, tau, anchors)
+        if cfg.get("sacflow", {}).get("project_velocity_to_nullspace", True) and Q.numel() > 0:
+            v = project_to_residual(v, Q)
+        fm_loss = F.mse_loss(v, u.detach())
+        path_mode = cfg.get("sacflow", {}).get("path_from_velocity", "one_step")
+        if path_mode == "one_step":
+            Rtau = R0 + tau_view * v
+        else:
+            Rtau = R_interp
+    else:
+        v = torch.zeros_like(Rt)
+        fm_loss = torch.tensor(0.0, device=feat.device)
+        Rtau = R_interp
+    Ftau = F_task + Rtau
+    # Classify path states through the wrapped model so DDP can track gradients.
+    path_logits = model(prelogit_features=Ftau)
+    # resize mask/pseudo if needed
+    pseudo_f = pseudo
+    mask_f = mask
+    if pseudo.shape[-3:] != path_logits.shape[-3:]:
+        pseudo_f = F.interpolate(pseudo[:, None].float(), size=path_logits.shape[-3:], mode="nearest")[:,0].long()
+        mask_f = F.interpolate(mask[:, None].float(), size=path_logits.shape[-3:], mode="nearest")[:,0]
+    path_probs = torch.softmax(path_logits, dim=1)
+    task_loss = ce_loss_masked(path_logits, pseudo_f, mask_f)
+    dice = dice_loss_masked(path_logits, pseudo_f, mask_f, num_classes)
+    ent = entropy_loss(path_logits)
+    if tprobs.shape[-3:] != path_logits.shape[-3:]:
+        tprobs_path = F.interpolate(tprobs.detach(), size=path_logits.shape[-3:], mode="trilinear", align_corners=False)
+    else:
+        tprobs_path = tprobs.detach()
+    task_kl = kl_masked(path_probs, tprobs_path, mask_f)
+    b_pred = finite_difference_boundary(path_probs)
+    b_ref = finite_difference_boundary(tprobs_path)
+    boundary_loss = F.l1_loss(b_pred, b_ref)
+    null_leak = torch.tensor(0.0, device=feat.device)
+    if Q.numel() > 0:
+        null_leak = (Rtau - R0 - project_to_residual(Rtau - R0, Q)).pow(2).mean()
+
+    # Domain-progress proxy in residual-statistic space. A valid path state should
+    # move monotonically from source-like residual stats toward target residual stats.
+    rtau_mu, rtau_std, _ = class_moments(Rtau, tprobs_f.detach())
+    dist_src = _stats_distance(rtau_mu, rtau_std, src_mu.detach(), src_std.detach())
+    dist_tgt = _stats_distance(rtau_mu, rtau_std, tgt_mu.detach(), tgt_std.detach())
+    rho = dist_src / (dist_src + dist_tgt + 1e-6)
+    domain_progress_loss = (rho - tau.mean()).abs()
+
+    losses_cfg = cfg.get("sacflow", {}).get("losses", {})
+    val_cfg = cfg.get("sacflow", {}).get("validation", {})
+    use_path_weights = bool(val_cfg.get("use_weights", True))
+    if use_path_weights:
+        alpha = float(val_cfg.get("task_alpha", 1.0))
+        beta = float(val_cfg.get("anatomy_beta", 1.0))
+        # Use detached scalar weights so the gate selects/weights path states but
+        # does not create degenerate gradients that simply lower the weight.
+        accepted_weight = torch.exp(-alpha * task_kl.detach() - beta * boundary_loss.detach()).clamp(
+            min=float(val_cfg.get("min_weight", 0.05)), max=1.0
+        )
+    else:
+        accepted_weight = torch.tensor(1.0, device=feat.device, dtype=feat.dtype)
+
+    path_loss = accepted_weight * (
+        float(losses_cfg.get("path_ce", 1.0)) * task_loss +
+        float(losses_cfg.get("path_dice", 1.0)) * dice
+    )
+    loss = (
+        float(losses_cfg.get("fm", 1.0)) * fm_loss +
+        path_loss +
+        float(losses_cfg.get("task_kl", 0.25)) * task_kl +
+        float(losses_cfg.get("boundary", 0.05)) * boundary_loss +
+        float(losses_cfg.get("null_leakage", 0.1)) * null_leak +
+        float(losses_cfg.get("domain_progress", 0.05)) * domain_progress_loss +
+        float(losses_cfg.get("entropy", 0.01)) * ent
+    )
+    logs = {
+        "loss_total": loss.detach(),
+        "loss_fm": fm_loss.detach(),
+        "loss_path_ce": task_loss.detach(),
+        "loss_path_dice": dice.detach(),
+        "loss_task_kl": task_kl.detach(),
+        "loss_boundary": boundary_loss.detach(),
+        "loss_null_leak": null_leak.detach(),
+        "loss_domain_progress": domain_progress_loss.detach(),
+        "loss_entropy": ent.detach(),
+        "sacflow_rho_mean": rho.detach(),
+        "sacflow_tau_mean": tau.mean().detach(),
+        "sacflow_velocity_mag": v.detach().abs().mean(),
+        "sacflow_residual_gap": (R1-R0).detach().abs().mean(),
+        "pseudo_conf_mean": conf.detach().mean(),
+        "pseudo_accept_rate": mask.detach().mean(),
+        "path_weight": accepted_weight.detach(),
+    }
+    return loss, logs

sacflow/methods/source_memory.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+from pathlib import Path
+from typing import Dict
+import torch
+import torch.nn.functional as F
+
+
+@torch.no_grad()
+def class_moments(residual: torch.Tensor, probs: torch.Tensor, eps: float = 1e-6):
+    """Class-conditional channel moments without materializing [B,C,d,N].
+
+    residual: [B,d,H,W,D]
+    probs:    [B,C,H,W,D], soft or one-hot class weights at the same resolution.
+
+    Returns:
+        mu:      [C,d]
+        std:     [C,d]
+        weights: [C]
+
+    This implementation uses first and second moments directly and is much more
+    memory efficient than computing (residual - mu)^2 for every class/voxel.
+    """
+    B, d, H, W, D = residual.shape
+    C = probs.shape[1]
+    r = residual.reshape(B, d, -1)
+    p = probs.reshape(B, C, -1).to(dtype=residual.dtype)
+    weights = p.sum(dim=(0, 2)).clamp_min(eps)  # [C]
+    mu = torch.einsum("bcn,bdn->cd", p, r) / weights[:, None]
+    second = torch.einsum("bcn,bdn->cd", p, r.pow(2)) / weights[:, None]
+    var = (second - mu.pow(2)).clamp_min(eps)
+    std = torch.sqrt(var)
+    return mu, std, weights
+
+
+def hard_onehot(label: torch.Tensor, num_classes: int):
+    return F.one_hot(label.long().clamp(0, num_classes-1), num_classes).permute(0,4,1,2,3).float()
+
+
+def save_source_memory(path: str | Path, memory: Dict):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(memory, path)
+
+
+def load_source_memory(path: str | Path, device=None):
+    mem = torch.load(path, map_location=device or "cpu")
+    return mem
+
+
+def moment_transport_residual(Rt: torch.Tensor, probs: torch.Tensor, source_mu: torch.Tensor, source_std: torch.Tensor,
+                              target_mu: torch.Tensor | None = None, target_std: torch.Tensor | None = None,
+                              eps: float = 1e-5):
+    """Class-gated diagonal moment transport from target residual stats to source residual stats.
+
+    Computes R0(u) = sum_c p_c(u) [mu_s_c + sigma_s_c / sigma_t_c * (Rt(u)-mu_t_c)].
+
+    The implementation intentionally avoids stacking all class-wise transported
+    residuals, because [B,C,d,H,W,D] can be very large for 3D volumes.
+    """
+    B, d, H, W, D = Rt.shape
+    C = probs.shape[1]
+    if target_mu is None or target_std is None:
+        target_mu, target_std, _ = class_moments(Rt, probs, eps=eps)
+    source_mu = source_mu.to(Rt.device, Rt.dtype)
+    source_std = source_std.to(Rt.device, Rt.dtype)
+    target_mu = target_mu.to(Rt.device, Rt.dtype)
+    target_std = target_std.to(Rt.device, Rt.dtype)
+    R0 = torch.zeros_like(Rt)
+    for c in range(C):
+        tr = source_mu[c].view(1,d,1,1,1) + (source_std[c].view(1,d,1,1,1) / (target_std[c].view(1,d,1,1,1) + eps)) * (Rt - target_mu[c].view(1,d,1,1,1))
+        R0 = R0 + probs[:, c:c+1].to(Rt.dtype) * tr
+    return R0, target_mu, target_std

sacflow/methods/task_space.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+import torch
+import torch.nn.functional as F
+
+
+def centered_classifier_basis(W: torch.Tensor, tol: float = 1e-5, foreground_only: bool = False):
+    # W: [C,d]
+    if foreground_only and W.shape[0] > 1:
+        W = W[1:]
+    Wc = W - W.mean(dim=0, keepdim=True)
+    # U basis in feature dimension for row space of Wc
+    U, S, Vh = torch.linalg.svd(Wc.T, full_matrices=False)
+    if S.numel() == 0:
+        return torch.empty(W.shape[1], 0, device=W.device, dtype=W.dtype)
+    r = int((S > tol * S.max()).sum().item())
+    return U[:, :max(1, r)]
+
+
+def random_basis(feature_dim: int, rank: int, device, dtype):
+    A = torch.randn(feature_dim, rank, device=device, dtype=dtype)
+    Q, _ = torch.linalg.qr(A, mode="reduced")
+    return Q
+
+
+def project_task_and_residual(Fmap: torch.Tensor, Q: torch.Tensor):
+    # Fmap [B,d,H,W,D], Q [d,r]
+    if Q.numel() == 0:
+        return torch.zeros_like(Fmap), Fmap
+    B, d, H, W, D = Fmap.shape
+    flat = Fmap.permute(0,2,3,4,1).reshape(-1, d)
+    task = (flat @ Q) @ Q.T
+    task = task.reshape(B,H,W,D,d).permute(0,4,1,2,3).contiguous()
+    residual = Fmap - task
+    return task, residual
+
+
+def project_to_residual(x: torch.Tensor, Q: torch.Tensor):
+    task, res = project_task_and_residual(x, Q)
+    return res

sacflow/models/__init__.py

@@ -0,0 +1 @@
+

sacflow/models/unet3d.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def norm_layer(kind: str, channels: int):
+    if kind == "batch":
+        return nn.BatchNorm3d(channels)
+    if kind == "group":
+        return nn.GroupNorm(num_groups=min(8, channels), num_channels=channels)
+    return nn.InstanceNorm3d(channels, affine=True)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, norm="instance", dropout=0.0):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv3d(in_ch, out_ch, 3, padding=1, bias=False),
+            norm_layer(norm, out_ch),
+            nn.LeakyReLU(0.01, inplace=True),
+            nn.Dropout3d(dropout) if dropout > 0 else nn.Identity(),
+            nn.Conv3d(out_ch, out_ch, 3, padding=1, bias=False),
+            norm_layer(norm, out_ch),
+            nn.LeakyReLU(0.01, inplace=True),
+        )
+    def forward(self, x):
+        return self.block(x)
+
+
+class BottleneckAdapter3D(nn.Module):
+    def __init__(self, channels: int, ratio: float = 0.25, norm: str = "group"):
+        super().__init__()
+        hidden = max(4, int(channels * ratio))
+        self.net = nn.Sequential(
+            nn.Conv3d(channels, hidden, 1, bias=False),
+            norm_layer(norm, hidden),
+            nn.GELU(),
+            nn.Conv3d(hidden, channels, 1, bias=True),
+        )
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+    def forward(self, x):
+        return x + self.net(x)
+
+
+class UNet3D(nn.Module):
+    def __init__(self, in_channels=1, num_classes=2, base_channels=32, levels=4, norm="instance", dropout=0.0, adapter_cfg=None):
+        super().__init__()
+        self.in_channels = in_channels
+        self.num_classes = num_classes
+        self.base_channels = base_channels
+        self.levels = levels
+        chs = [base_channels * (2 ** i) for i in range(levels)]
+        self.enc = nn.ModuleList()
+        self.down = nn.ModuleList()
+        prev = in_channels
+        for ch in chs:
+            self.enc.append(ConvBlock(prev, ch, norm, dropout))
+            self.down.append(nn.Conv3d(ch, ch, 3, stride=2, padding=1))
+            prev = ch
+        self.bottleneck = ConvBlock(chs[-1], chs[-1]*2, norm, dropout)
+        self.up = nn.ModuleList()
+        self.dec = nn.ModuleList()
+        prev = chs[-1]*2
+        for ch in reversed(chs):
+            self.up.append(nn.ConvTranspose3d(prev, ch, 2, stride=2))
+            self.dec.append(ConvBlock(ch*2, ch, norm, dropout))
+            prev = ch
+        self.prelogit_channels = base_channels
+        self.adapter_enabled = bool(adapter_cfg and adapter_cfg.get("enabled", False))
+        if self.adapter_enabled:
+            self.prelogit_adapter = BottleneckAdapter3D(base_channels, ratio=float(adapter_cfg.get("bottleneck_ratio", 0.25)))
+        else:
+            self.prelogit_adapter = nn.Identity()
+        self.out_conv = nn.Conv3d(base_channels, num_classes, 1)
+
+    def forward_features(self, x):
+        skips = []
+        out = x
+        for enc, down in zip(self.enc, self.down):
+            out = enc(out)
+            skips.append(out)
+            out = down(out)
+        out = self.bottleneck(out)
+        for up, dec, skip in zip(self.up, self.dec, reversed(skips)):
+            out = up(out)
+            if out.shape[-3:] != skip.shape[-3:]:
+                out = F.interpolate(out, size=skip.shape[-3:], mode="trilinear", align_corners=False)
+            out = torch.cat([out, skip], dim=1)
+            out = dec(out)
+        return out
+
+    def classify_from_features(self, feat):
+        feat = self.prelogit_adapter(feat)
+        return self.out_conv(feat)
+
+    def forward(self, x=None, return_features=False, prelogit_features=None):
+        # prelogit_features allows SACFlow to classify generated/intermediate
+        # feature states through the normal module/DDP call path. This is safer
+        # than bypassing DistributedDataParallel by calling module methods directly.
+        if prelogit_features is not None:
+            feat = prelogit_features
+        else:
+            if x is None:
+                raise ValueError("Either x or prelogit_features must be provided")
+            feat = self.forward_features(x)
+        logits = self.classify_from_features(feat)
+        if return_features:
+            return logits, {"prelogit": feat}
+        return logits
+
+    def final_classifier_weight(self):
+        return self.out_conv.weight.squeeze(-1).squeeze(-1).squeeze(-1)
+
+
+def build_model(cfg):
+    mcfg = cfg["model"]
+    return UNet3D(
+        in_channels=mcfg.get("in_channels", cfg["data"].get("in_channels", 1)),
+        num_classes=mcfg.get("num_classes", cfg["data"].get("num_classes", 2)),
+        base_channels=mcfg.get("base_channels", 32),
+        levels=mcfg.get("levels", 4),
+        norm=mcfg.get("norm", "instance"),
+        dropout=mcfg.get("dropout", 0.0),
+        adapter_cfg=mcfg.get("adapter", None),
+    )
+
+
+def freeze_except_adapters(model: nn.Module, train_norm_affine: bool = True):
+    """Freeze the backbone, train adapters and optionally norm affine params.
+
+    The previous name-based norm check missed InstanceNorm/GroupNorm modules
+    inside Sequential containers. This module-type based version reliably
+    unfreezes affine normalization parameters when requested.
+    """
+    for p in model.parameters():
+        p.requires_grad = False
+    for name, p in model.named_parameters():
+        if "adapter" in name:
+            p.requires_grad = True
+    if train_norm_affine:
+        norm_types = (nn.BatchNorm3d, nn.InstanceNorm3d, nn.GroupNorm)
+        for module in model.modules():
+            if isinstance(module, norm_types):
+                if getattr(module, "weight", None) is not None:
+                    module.weight.requires_grad = True
+                if getattr(module, "bias", None) is not None:
+                    module.bias.requires_grad = True

sacflow/models/velocity_field.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def sinusoidal_embedding(t: torch.Tensor, dim: int):
+    # t: [B] or [B,1]
+    if t.ndim == 0:
+        t = t[None]
+    t = t.view(-1, 1)
+    half = dim // 2
+    freqs = torch.exp(torch.arange(half, device=t.device, dtype=t.dtype) * -(math.log(10000.0) / max(1, half - 1)))
+    args = t * freqs[None]
+    emb = torch.cat([torch.sin(args), torch.cos(args)], dim=1)
+    if dim % 2 == 1:
+        emb = F.pad(emb, (0,1))
+    return emb
+
+
+class FiLM(nn.Module):
+    def __init__(self, cond_dim: int, channels: int):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(cond_dim, channels*2), nn.SiLU(), nn.Linear(channels*2, channels*2))
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+    def forward(self, x, cond):
+        gb = self.net(cond)
+        gamma, beta = gb.chunk(2, dim=1)
+        shape = [x.shape[0], x.shape[1]] + [1]*(x.ndim-2)
+        gamma = gamma.view(*shape)
+        beta = beta.view(*shape)
+        return x * (1 + gamma) + beta
+
+
+class VelocityField3D(nn.Module):
+    def __init__(self, residual_channels: int, num_classes: int, hidden_ratio: float = 0.25, depth: int = 2,
+                 tau_embedding_dim: int = 32, organ_embedding_dim: int = 16,
+                 include_teacher_probs: bool = True, include_confidence: bool = True, include_boundary: bool = True,
+                 use_depthwise: bool = True, use_group_norm: bool = True, use_film: bool = True):
+        super().__init__()
+        self.residual_channels = residual_channels
+        self.num_classes = num_classes
+        self.include_teacher_probs = include_teacher_probs
+        self.include_confidence = include_confidence
+        self.include_boundary = include_boundary
+        self.use_film = use_film
+        self.tau_embedding_dim = tau_embedding_dim
+        self.organ_embedding_dim = organ_embedding_dim
+        self.organ_emb = nn.Embedding(num_classes, organ_embedding_dim)
+        anchor_ch = 0
+        if include_teacher_probs:
+            anchor_ch += num_classes
+        if include_confidence:
+            anchor_ch += 1
+        if include_boundary:
+            anchor_ch += 1
+        anchor_ch += organ_embedding_dim
+        in_ch = residual_channels + anchor_ch + 1  # tau map
+        hidden = max(8, int(residual_channels * hidden_ratio))
+        self.in_proj = nn.Conv3d(in_ch, hidden, 1)
+        blocks = []
+        for _ in range(depth):
+            if use_depthwise:
+                conv = nn.Sequential(
+                    nn.Conv3d(hidden, hidden, 3, padding=1, groups=hidden, bias=False),
+                    nn.Conv3d(hidden, hidden, 1, bias=False),
+                )
+            else:
+                conv = nn.Conv3d(hidden, hidden, 3, padding=1, bias=False)
+            norm = nn.GroupNorm(min(8, hidden), hidden) if use_group_norm else nn.InstanceNorm3d(hidden, affine=True)
+            blocks.append(nn.Sequential(conv, norm, nn.GELU()))
+        self.blocks = nn.ModuleList(blocks)
+        cond_dim = tau_embedding_dim
+        self.films = nn.ModuleList([FiLM(cond_dim, hidden) for _ in range(depth)]) if use_film else None
+        self.out_proj = nn.Conv3d(hidden, residual_channels, 1)
+        nn.init.zeros_(self.out_proj.weight)
+        nn.init.zeros_(self.out_proj.bias)
+
+    def organ_embedding_map(self, probs: torch.Tensor):
+        # probs [B,C,H,W,D]
+        emb = self.organ_emb.weight  # [C,E]
+        e = torch.einsum("bchwd,ce->behwd", probs, emb)
+        return e
+
+    def forward(self, residual: torch.Tensor, tau: torch.Tensor, anchors: dict):
+        B, _, H, W, D = residual.shape
+        xs = [residual]
+        tau_vec = tau.view(B).to(residual.dtype)
+        tau_map = tau_vec.view(B,1,1,1,1).expand(B,1,H,W,D)
+        xs.append(tau_map)
+        probs = anchors.get("probs")
+        if probs is None:
+            probs = torch.zeros(B, self.num_classes, H, W, D, device=residual.device, dtype=residual.dtype)
+            probs[:, 0] = 1.0
+        if probs.shape[-3:] != (H,W,D):
+            probs = F.interpolate(probs, size=(H,W,D), mode="trilinear", align_corners=False)
+        if self.include_teacher_probs:
+            xs.append(probs)
+        if self.include_confidence:
+            conf = probs.max(dim=1, keepdim=True).values
+            xs.append(conf)
+        if self.include_boundary:
+            bnd = anchors.get("boundary")
+            if bnd is None:
+                # simple finite diff boundary from probs
+                bnd = (torch.abs(probs[:, :, 1:] - probs[:, :, :-1]).mean(1, keepdim=True))
+                bnd = F.pad(bnd, (0,0,0,0,1,0))
+            if bnd.shape[-3:] != (H,W,D):
+                bnd = F.interpolate(bnd, size=(H,W,D), mode="trilinear", align_corners=False)
+            xs.append(bnd)
+        xs.append(self.organ_embedding_map(probs))
+        x = torch.cat(xs, dim=1)
+        h = self.in_proj(x)
+        cond = sinusoidal_embedding(tau_vec, self.tau_embedding_dim)
+        for i, block in enumerate(self.blocks):
+            h = block(h)
+            if self.use_film:
+                h = self.films[i](h, cond)
+        return self.out_proj(h)

sacflow/utils/__init__.py

@@ -0,0 +1 @@
+

sacflow/utils/config.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+from pathlib import Path
+from typing import Any, Dict, List
+import copy
+import yaml
+
+
+def _deep_update(base: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]:
+    out = copy.deepcopy(base)
+    for k, v in update.items():
+        if isinstance(v, dict) and isinstance(out.get(k), dict):
+            out[k] = _deep_update(out[k], v)
+        else:
+            out[k] = copy.deepcopy(v)
+    return out
+
+
+def load_yaml(path: str | Path) -> Dict[str, Any]:
+    path = Path(path)
+    with open(path, "r") as f:
+        cfg = yaml.safe_load(f) or {}
+    base_files = cfg.pop("_base_", [])
+    if isinstance(base_files, str):
+        base_files = [base_files]
+    merged: Dict[str, Any] = {}
+    for base in base_files:
+        base_path = (path.parent / base).resolve()
+        merged = _deep_update(merged, load_yaml(base_path))
+    merged = _deep_update(merged, cfg)
+    return merged
+
+
+def save_yaml(cfg: Dict[str, Any], path: str | Path) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w") as f:
+        yaml.safe_dump(cfg, f, sort_keys=False)
+
+
+def get(cfg: Dict[str, Any], key: str, default=None):
+    cur = cfg
+    for part in key.split("."):
+        if not isinstance(cur, dict) or part not in cur:
+            return default
+        cur = cur[part]
+    return cur
+
+
+def set_by_path(cfg: Dict[str, Any], key: str, value: Any) -> None:
+    cur = cfg
+    parts = key.split(".")
+    for p in parts[:-1]:
+        cur = cur.setdefault(p, {})
+    cur[parts[-1]] = value

sacflow/utils/distributed.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+import os
+from datetime import timedelta
+import torch
+import torch.distributed as dist
+
+
+def is_dist_avail_and_initialized() -> bool:
+    return dist.is_available() and dist.is_initialized()
+
+
+def get_rank() -> int:
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def get_world_size() -> int:
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def is_main_process() -> bool:
+    return get_rank() == 0
+
+
+def init_distributed(backend: str = "nccl") -> torch.device:
+    """Initialize distributed training.
+
+    Validation on 3D medical volumes can take longer than PyTorch's default
+    10 minute NCCL/RCCL watchdog timeout if some ranks are waiting at a
+    collective. We therefore set a longer timeout by default. The value can be
+    overridden with DIST_TIMEOUT_MINUTES.
+    """
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))
+        torch.cuda.set_device(local_rank)
+        timeout_min = int(os.environ.get("DIST_TIMEOUT_MINUTES", "180"))
+        dist.init_process_group(
+            backend=backend,
+            init_method="env://",
+            timeout=timedelta(minutes=timeout_min),
+        )
+        device = torch.device("cuda", local_rank)
+    else:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return device
+
+
+def barrier():
+    if is_dist_avail_and_initialized():
+        if torch.cuda.is_available():
+            dist.barrier(device_ids=[torch.cuda.current_device()])
+        else:
+            dist.barrier()
+
+
+def cleanup():
+    if is_dist_avail_and_initialized():
+        dist.destroy_process_group()
+
+
+def reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
+    if not is_dist_avail_and_initialized():
+        return tensor
+    rt = tensor.detach().clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    rt /= get_world_size()
+    return rt

sacflow/utils/eta.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+import time
+
+class ETA:
+    def __init__(self, total: int, smoothing: float = 0.95):
+        self.total = max(1, total)
+        self.smoothing = smoothing
+        self.start = time.time()
+        self.last = self.start
+        self.avg = None
+        self.n = 0
+
+    def step(self, count: int = 1):
+        now = time.time()
+        dt = (now - self.last) / max(1, count)
+        self.last = now
+        self.n += count
+        if self.avg is None:
+            self.avg = dt
+        else:
+            self.avg = self.smoothing * self.avg + (1 - self.smoothing) * dt
+
+    def eta_seconds(self) -> float:
+        rem = max(0, self.total - self.n)
+        return rem * (self.avg or 0.0)
+
+    @staticmethod
+    def fmt(seconds: float) -> str:
+        seconds = int(seconds)
+        h = seconds // 3600
+        m = (seconds % 3600) // 60
+        s = seconds % 60
+        if h > 0:
+            return f"{h:d}h {m:02d}m {s:02d}s"
+        if m > 0:
+            return f"{m:d}m {s:02d}s"
+        return f"{s:d}s"

sacflow/utils/metrics.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+from typing import Dict, List
+import numpy as np
+import torch
+from scipy import ndimage
+
+
+def dice_per_class(pred: np.ndarray, target: np.ndarray, num_classes: int, include_bg: bool = False) -> Dict[str, float]:
+    out = {}
+    start = 0 if include_bg else 1
+    vals = []
+    for c in range(start, num_classes):
+        p = pred == c
+        t = target == c
+        denom = p.sum() + t.sum()
+        if denom == 0:
+            d = np.nan
+        else:
+            d = 2.0 * np.logical_and(p, t).sum() / denom
+        out[f"dice_c{c}"] = float(d) if not np.isnan(d) else np.nan
+        if not np.isnan(d):
+            vals.append(d)
+    out["dice_mean"] = float(np.mean(vals)) if vals else np.nan
+    return out
+
+
+def _surface(mask: np.ndarray) -> np.ndarray:
+    if mask.sum() == 0:
+        return mask.astype(bool)
+    eroded = ndimage.binary_erosion(mask, iterations=1, border_value=0)
+    return np.logical_xor(mask, eroded)
+
+
+def hd95_per_class(pred: np.ndarray, target: np.ndarray, num_classes: int, spacing=None, include_bg: bool = False, empty_penalty: str = "diagonal") -> Dict[str, float]:
+    """Compute per-class HD95.
+
+    If one mask is empty and the other is not, assigning NaN and dropping the
+    class makes failed segmentations look artificially good. We instead use the
+    physical image diagonal as a conservative finite penalty. If both masks are
+    empty, the class is not present and is excluded from the mean.
+    """
+    out = {}
+    start = 0 if include_bg else 1
+    vals = []
+    if spacing is None:
+        spacing = (1.0, 1.0, 1.0)
+    spacing = tuple(float(x) for x in spacing)
+    diag = float(np.sqrt(sum(((s * max(1, n - 1)) ** 2) for s, n in zip(spacing, target.shape))))
+    for c in range(start, num_classes):
+        p = pred == c
+        t = target == c
+        p_sum = int(p.sum())
+        t_sum = int(t.sum())
+        if p_sum == 0 and t_sum == 0:
+            hd = np.nan
+        elif p_sum == 0 or t_sum == 0:
+            hd = diag if empty_penalty == "diagonal" else np.nan
+        else:
+            ps = _surface(p)
+            ts = _surface(t)
+            if ps.sum() == 0 or ts.sum() == 0:
+                hd = diag if empty_penalty == "diagonal" else np.nan
+            else:
+                dt_t = ndimage.distance_transform_edt(~ts, sampling=spacing)
+                dt_p = ndimage.distance_transform_edt(~ps, sampling=spacing)
+                dists = np.concatenate([dt_t[ps], dt_p[ts]])
+                hd = np.percentile(dists, 95) if dists.size else np.nan
+        out[f"hd95_c{c}"] = float(hd) if not np.isnan(hd) else np.nan
+        if not np.isnan(hd):
+            vals.append(hd)
+    out["hd95_mean"] = float(np.mean(vals)) if vals else np.nan
+    return out
+
+
+def torch_soft_dice_loss(logits: torch.Tensor, target: torch.Tensor, num_classes: int, ignore_index: int | None = None, eps: float = 1e-5):
+    probs = torch.softmax(logits, dim=1)
+    if target.ndim == logits.ndim:
+        onehot = target.float()
+    else:
+        target_clamped = target.clamp(0, num_classes - 1).long()
+        onehot = torch.nn.functional.one_hot(target_clamped, num_classes).permute(0, 4, 1, 2, 3).float()
+    if ignore_index is not None:
+        mask = target != ignore_index
+        probs = probs * mask.unsqueeze(1)
+        onehot = onehot * mask.unsqueeze(1)
+    dims = tuple(range(2, logits.ndim))
+    inter = (probs * onehot).sum(dims)
+    denom = probs.sum(dims) + onehot.sum(dims)
+    dice = (2 * inter + eps) / (denom + eps)
+    return 1.0 - dice[:, 1:].mean()
+
+
+def entropy_loss(logits: torch.Tensor, eps: float = 1e-8):
+    p = torch.softmax(logits, dim=1)
+    return -(p * (p + eps).log()).sum(dim=1).mean()
+
+
+def confidence_and_margin(probs: torch.Tensor):
+    vals, inds = probs.topk(k=2, dim=1)
+    conf = vals[:, 0]
+    margin = vals[:, 0] - vals[:, 1]
+    pred = inds[:, 0]
+    return conf, margin, pred
+
+
+def finite_difference_boundary(x: torch.Tensor) -> torch.Tensor:
+    # x: [B,C,H,W,D] probabilities or one-hot maps
+    dx = torch.abs(x[:, :, 1:] - x[:, :, :-1]).mean(dim=1, keepdim=True)
+    dx = torch.nn.functional.pad(dx, (0,0,0,0,1,0))
+    dy = torch.abs(x[:, :, :, 1:] - x[:, :, :, :-1]).mean(dim=1, keepdim=True)
+    dy = torch.nn.functional.pad(dy, (0,0,1,0,0,0))
+    dz = torch.abs(x[:, :, :, :, 1:] - x[:, :, :, :, :-1]).mean(dim=1, keepdim=True)
+    dz = torch.nn.functional.pad(dz, (1,0,0,0,0,0))
+    return (dx + dy + dz) / 3.0

sacflow/utils/misc.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+import os
+import random
+from pathlib import Path
+import numpy as np
+import torch
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+
+def ensure_dir(path: str | Path) -> Path:
+    p = Path(path)
+    p.mkdir(parents=True, exist_ok=True)
+    return p
+
+
+def count_trainable(model: torch.nn.Module) -> tuple[int, int]:
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return trainable, total
+
+
+def move_to_device(batch, device):
+    if isinstance(batch, torch.Tensor):
+        return batch.to(device, non_blocking=True)
+    if isinstance(batch, dict):
+        return {k: move_to_device(v, device) for k, v in batch.items()}
+    if isinstance(batch, (list, tuple)):
+        return type(batch)(move_to_device(v, device) for v in batch)
+    return batch
+
+
+def unwrap_model(model):
+    return model.module if hasattr(model, "module") else model

sacflow/utils/wandb_utils.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+from typing import Any, Dict
+import os
+from pathlib import Path
+try:
+    import wandb
+except Exception:  # pragma: no cover
+    wandb = None
+from .distributed import is_main_process
+
+
+def _read_text(path: Path) -> str | None:
+    try:
+        if path.exists():
+            text = path.read_text().strip()
+            return text or None
+    except Exception:
+        pass
+    return None
+
+
+def _write_text(path: Path, text: str) -> None:
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(text.strip() + "\n")
+    except Exception:
+        # W&B should never crash training just because the local ID file could not be written.
+        pass
+
+
+def _env_int(name: str, default: int) -> int:
+    try:
+        return int(os.environ.get(name, default))
+    except Exception:
+        return default
+
+
+def init_wandb(cfg: Dict[str, Any], run_name: str | None = None):
+    """Initialize W&B on rank 0 only and never crash training if W&B is slow/down.
+
+    Resume behavior:
+    - If wandb.id or WANDB_RUN_ID is provided, use it.
+    - Otherwise, if <output_dir>/wandb_run_id.txt exists, reuse that ID.
+    - If a run ID is used, default resume mode is "allow" unless overridden by
+      wandb.resume or WANDB_RESUME.
+
+    Reliability behavior:
+    - WANDB_INIT_TIMEOUT controls wandb.init timeout seconds (default 300).
+    - WANDB_SERVICE_WAIT controls W&B service startup wait seconds (default 300).
+    - If wandb.init raises any exception, training continues with W&B disabled.
+    """
+    wb_cfg = cfg.get("wandb", {})
+    if not wb_cfg.get("enabled", False) or not is_main_process() or wandb is None:
+        return None
+
+    mode = os.environ.get("WANDB_MODE", wb_cfg.get("mode", "online"))
+    entity = wb_cfg.get("entity") or os.environ.get("WANDB_ENTITY") or None
+    project = os.environ.get("WANDB_PROJECT", wb_cfg.get("project", "SACFlow-FM"))
+
+    out_dir = Path(cfg.get("output_dir", "."))
+    run_id_file = out_dir / "wandb_run_id.txt"
+
+    run_id = (
+        wb_cfg.get("id")
+        or os.environ.get("WANDB_RUN_ID")
+        or _read_text(run_id_file)
+    )
+
+    resume_mode = wb_cfg.get("resume") or os.environ.get("WANDB_RESUME")
+    if run_id and resume_mode is None:
+        resume_mode = "allow"
+
+    init_timeout = _env_int("WANDB_INIT_TIMEOUT", int(wb_cfg.get("init_timeout", 300)))
+    service_wait = _env_int("WANDB_SERVICE_WAIT", int(wb_cfg.get("service_wait", 300)))
+
+    # Keep W&B files away from home/cache quotas if an output directory exists.
+    os.environ.setdefault("WANDB_DIR", str(out_dir / "wandb"))
+    os.environ.setdefault("WANDB_CACHE_DIR", str(out_dir / "wandb_cache"))
+    os.environ.setdefault("WANDB_CONFIG_DIR", str(out_dir / "wandb_config"))
+
+    try:
+        settings = wandb.Settings(init_timeout=init_timeout, _service_wait=service_wait)
+        run = wandb.init(
+            project=project,
+            entity=entity,
+            mode=mode,
+            name=run_name,
+            id=run_id,
+            resume=resume_mode,
+            tags=wb_cfg.get("tags", []),
+            config=cfg,
+            settings=settings,
+        )
+    except Exception as e:  # pragma: no cover - depends on external W&B service/network
+        print(
+            f"[W&B WARNING] wandb.init failed ({type(e).__name__}: {e}). "
+            "Continuing training with W&B disabled for this run. "
+            "To avoid this, use WANDB_MODE=offline or increase WANDB_INIT_TIMEOUT.",
+            flush=True,
+        )
+        return None
+
+    # Persist the generated or reused run ID so future --resume auto attaches to
+    # the same W&B run from this output directory.
+    if getattr(run, "id", None):
+        _write_text(run_id_file, run.id)
+
+    return run
+
+
+def wandb_log(run, data: Dict[str, Any], step: int | None = None):
+    if run is not None:
+        try:
+            run.log(data, step=step)
+        except Exception as e:
+            print(f"[W&B WARNING] wandb.log failed: {e}", flush=True)
+
+
+def wandb_finish(run):
+    if run is not None:
+        try:
+            run.finish()
+        except Exception as e:
+            print(f"[W&B WARNING] wandb.finish failed: {e}", flush=True)