Scripts/OM_train_2modes.py

"""
OM_train_2modes.py — Dual-mode training (diffusion + registration) using OMorpher.

Drop-in replacement for OM_train_2modes.py. Uses the OMorpher object-oriented
wrapper instead of procedural DeformDDPM calls, while preserving the same
training logic, DDP support, loss functions, and checkpoint format.

Usage:
    # Single-GPU
    python Scripts/OM_train_2modes.py -C Config/config_om.yaml

    # Multi-GPU (DDP)
    CUDA_VISIBLE_DEVICES=0,1 python Scripts/OM_train_2modes.py -C Config/config_om.yaml
"""

import os
import sys

# Add project root to path so imports work from Scripts/
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT_DIR)

import gc
import glob
import random

import numpy as np
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import Adam
from torch.utils.data import DataLoader
from tqdm import tqdm
import argparse

from OMorpher import OMorpher
from Diffusion.networks import DefRec_MutAttnNet
from Diffusion.losses import Grad, LNCC, LMSE, MSLNCC
from Dataloader.dataLoader import OMDataset_indiv, OMDataset_pair
from Dataloader.dataloader_utils import thresh_img
import utils

# ========================== Constants ==========================

EPS = 1e-5
MSK_EPS = 0.01
TEXT_EMBED_PROB = 0.7
AUG_RESAMPLE_PROB = 0.6
LOSS_WEIGHTS_DIFF = [2.0, 1.0, 16]     # [ang, dist, reg]
LOSS_WEIGHTS_REGIST = [1.0, 0.05, 128]  # [imgsim, imgmse, ddf]
DIFF_REG_BATCH_RATIO = 2

use_distributed = True
# use_distributed = False

# ========================== Arguments ==========================

parser = argparse.ArgumentParser()
parser.add_argument(
    "--config", "-C",
    help="Path for the config file",
    type=str,
    default="Config/config_all.yaml",
    required=False,
)
args = parser.parse_args()

# ========================== DDP Setup ==========================


def ddp_setup(rank, world_size):
    """
    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)


# ========================== Helpers ==========================


def reverse_diffuse_train(network, om, img_org, cond_imgs, T, text=None):
    """Registration reverse diffusion with selective gradient control.

    Mirrors DeformDDPM.diff_recover() with T=[None, T_regist].
    Only the last k=2 timesteps have gradients enabled for efficient training.

    Args:
        network: DDP-wrapped (or raw) network module.
        om: OMorpher instance (provides STN instances and device info).
        img_org: Source image [B, 1, S, S, S].
        cond_imgs: Processed conditioning image [B, 1, S, S, S].
        T: [T_init, T_schedule]. T_init=None means no forward diffusion.
            T_schedule is a list of batched timestep lists from the training loop.
        text: Optional text embedding [B, 1024].

    Returns:
        (ddf_comp, img_rec): Composed DDF and recovered image.
    """
    B = img_org.shape[0]
    S = om.img_size

    # T[0] = None → no forward diffusion, start from original image
    ddf_comp = torch.zeros(
        [B, om.ndims] + [S] * om.ndims,
        dtype=torch.float32, device=om.device,
    )
    img_rec = img_org.clone().detach()

    time_steps = T[1]
    k = 2
    trainable_iterations = time_steps[-1:-k - 1:-1]

    net_module = network.module if isinstance(network, DDP) else network

    for i in time_steps:
        t = torch.tensor(np.array([i])).to(om.device)

        if i in trainable_iterations:
            # Gradients enabled — call through DDP wrapper for gradient sync
            pre_dvf = network(x=img_rec, y=cond_imgs, t=t, rec_num=2, text=text)
        else:
            # No gradients — call underlying module directly
            with torch.no_grad():
                pre_dvf = net_module(x=img_rec, y=cond_imgs, t=t, rec_num=2, text=text)

        ddf_comp = om.stn_full(ddf_comp, pre_dvf) + pre_dvf
        img_rec = om.img_stn(img_org.clone().detach(), ddf_comp)

    return ddf_comp, img_rec


def ddp_load_checkpoint(gpu_id, network, optimizer, model_file,
                        use_dist=True, load_strict=False):
    """Load checkpoint with DDP-aware parameter broadcast."""
    if gpu_id == 0:
        utils.print_memory_usage("Before Loading Model")
        gc.collect()
        torch.cuda.empty_cache()

        checkpoint = torch.load(model_file)
        state_dict = checkpoint['model_state_dict']

        # Strip DDP 'module.' and DeformDDPM 'network.' prefixes
        cleaned = {}
        for k, v in state_dict.items():
            k = k.replace("module.", "")
            if k.startswith("network."):
                k = k[len("network."):]
            cleaned[k] = v

        net = network.module if use_dist else network
        net_keys = set(net.state_dict().keys())
        filtered = {k: v for k, v in cleaned.items() if k in net_keys}
        net.load_state_dict(filtered, strict=load_strict)

        if load_strict:
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        utils.print_memory_usage("After Loading Checkpoint on GPU")

    if use_dist:
        # Broadcast model weights from rank 0 to all other GPUs
        dist.barrier()
        for param in network.parameters():
            dist.broadcast(param.data, src=0)
        dist.barrier()
        for param_group in optimizer.param_groups:
            for param in param_group['params']:
                if param.grad is not None:
                    dist.broadcast(param.grad, src=0)

    initial_epoch = int(os.path.basename(model_file).split('.')[0][:6]) + 1
    return initial_epoch


def save_checkpoint(network, optimizer, epoch, save_path, use_dist=True):
    """Save checkpoint with 'network.' key prefix for backward compatibility."""
    net = network.module if use_dist and isinstance(network, DDP) else network
    state_dict = {f"network.{k}": v for k, v in net.state_dict().items()}
    torch.save({
        'model_state_dict': state_dict,
        'optimizer_state_dict': optimizer.state_dict(),
        'epoch': epoch,
    }, save_path)


# ========================== Main Training ==========================


def main_train(rank=0, world_size=1, train_mode_ratio=1, thresh_imgsim=0.01):
    if use_distributed:
        ddp_setup(rank, world_size)
        if torch.distributed.is_initialized():
            print(f"World size: {torch.distributed.get_world_size()}")
            print(f"Communication backend: {torch.distributed.get_backend()}")
    gpu_id = rank
    device = f"cuda:{rank}" if use_distributed else None

    # ---- OMorpher initialisation (config, network, STN, losses, auto-checkpoint) ----
    om = OMorpher(config=args.config, device=device)
    config = om.config
    if gpu_id == 0:
        print(config)

    epoch_per_save = config['epoch_per_save']
    suffix_pth = f"_{config['data_name']}_{config['net_name']}.pth"
    model_dir = os.path.join(ROOT_DIR, 'Models', f"{config['data_name']}_{config['net_name']}/")

    # ---- Additional loss functions for the two training modes ----
    # Diffusion losses reused from OMorpher: om._loss_dist (MRSE), om._loss_ang (NCC)
    loss_reg = Grad(
        penalty=['l1', 'negdetj', 'range'], ndims=om.ndims,
        outrange_thresh=0.2, outrange_weight=1e3,
    )
    loss_reg1 = Grad(
        penalty=['l1', 'negdetj', 'range'], ndims=om.ndims,
        outrange_thresh=0.6, outrange_weight=1e3,
    )
    # loss_imgsim = LNCC()
    loss_imgsim = MSLNCC()
    loss_imgmse = LMSE()

    # ---- DDP wrapping ----
    if use_distributed:
        om.network.to(rank)
        om.stn_full.to(rank)
        om.stn_ctl.to(rank)
        om.img_stn.to(rank)
        om.msk_stn.to(rank)
        network = DDP(om.network, device_ids=[rank])
    else:
        om.network.to(om.device)
        network = om.network

    # ---- Optimizer ----
    optimizer = Adam(network.parameters(), lr=config["lr"])

    # ---- Data loaders ----
    dataset = OMDataset_indiv(transform=None)
    train_loader = DataLoader(
        dataset, batch_size=config['batchsize'], shuffle=True, drop_last=True,
    )

    datasetp = OMDataset_pair(transform=None)
    train_loader_p = DataLoader(
        datasetp,
        batch_size=config['batchsize'] // DIFF_REG_BATCH_RATIO,
        shuffle=True, drop_last=True,
    )

    # ---- Auto-resume from checkpoint ----
    os.makedirs(model_dir, exist_ok=True)
    model_files = sorted(glob.glob(os.path.join(model_dir, "*.pth")))
    if model_files:
        if gpu_id == 0:
            print(model_files)
        initial_epoch = ddp_load_checkpoint(
            gpu_id, network, optimizer, model_files[-1], use_distributed,
        )
    else:
        initial_epoch = 0

    if gpu_id == 0:
        print('len_train_data: ', len(dataset))

    is_defrec = isinstance(om.network, DefRec_MutAttnNet)

    # ---- Training loop ----
    for epoch in range(initial_epoch, config["epoch"]):
        epoch_loss_tot = 0.0
        epoch_loss_gen_d = 0.0
        epoch_loss_gen_a = 0.0
        epoch_loss_reg_ = 0.0
        epoch_loss_regist = 0.0
        epoch_loss_imgsim_ = 0.0
        epoch_loss_imgmse_ = 0.0
        epoch_loss_ddfreg = 0.0

        network.train()
        loss_nan_step = 0
        total = min(len(train_loader), len(train_loader_p))

        for step, (batch, batch_p) in tqdm(
            enumerate(zip(train_loader, train_loader_p)), total=total,
        ):
            # ==========================================================
            # Mode 1: Diffusion training on single image
            # ==========================================================
            [x0, embd] = batch
            x0 = x0.to(om.device).type(torch.float32)
            if np.random.uniform(0, 1) < TEXT_EMBED_PROB:
                embd = embd.to(om.device).type(torch.float32)
            else:
                embd = None

            n = x0.size()[0]
            x0 = x0.to(om.device)

            blind_mask = utils.get_random_deformed_mask(
                x0.shape[2:], apply_possibility=0.6,
            ).to(om.device)

            # Data augmentation
            if om.ndims > 2:
                if np.random.uniform(0, 1) < AUG_RESAMPLE_PROB:
                    x0 = utils.random_resample(x0, deform_scale=0)
                else:
                    [x0] = utils.random_permute([x0], select_dims=[-1, -2, -3])

            if config['noise_scale'] > 0:
                if np.random.uniform(0, 1) < AUG_RESAMPLE_PROB:
                    x0 = thresh_img(x0, [0, 2 * config['noise_scale']])
                x0 = x0 * (np.random.normal(1, config['noise_scale'] * 1)) + np.random.normal(0, config['noise_scale'] * 1)

            t = torch.randint(0, om.timesteps, (n,)).to(om.device)

            proc_type = random.choice(
                ['adding', 'downsample', 'slice', 'slice1', 'none', 'uncon', 'uncon', 'uncon'],
            )
            cond_img, _, cond_ratio = om._proc_cond_img(x0, proc_type=proc_type)

            # Forward diffusion + network prediction
            noisy_img, dvf_gt, _ = om._get_random_ddf(x0, t)

            if is_defrec:
                pre_dvf_I = network(
                    x=noisy_img * blind_mask, y=cond_img, t=[t], rec_num=2, text=embd,
                )
            else:
                pre_dvf_I = network(
                    x=noisy_img * blind_mask, y=cond_img, t=t, rec_num=2, text=embd,
                )

            # Diffusion losses
            loss_tot = 0
            loss_ddf = loss_reg(pre_dvf_I, img=x0)
            trm_pred = om.stn_full(pre_dvf_I, dvf_gt)
            loss_gen_d = om._loss_dist(
                pred=trm_pred, inv_lab=dvf_gt, ddf_stn=None, mask=blind_mask,
            )
            loss_gen_a = om._loss_ang(
                pred=trm_pred, inv_lab=dvf_gt, ddf_stn=None, mask=blind_mask,
            )

            loss_tot += LOSS_WEIGHTS_DIFF[0] * loss_gen_a + LOSS_WEIGHTS_DIFF[1] * loss_gen_d
            loss_tot += LOSS_WEIGHTS_DIFF[2] * loss_ddf
            loss_tot = torch.sqrt(1. + MSK_EPS - cond_ratio) * loss_tot

            # NaN / divergence checks
            if torch.isnan(x0).any():
                print(f"*** Encountered NaN in input image x0 at epoch {epoch}, step {step}.")
            if loss_ddf > 0.001:
                print(f"*** High diffusion DDF loss at epoch {epoch}, step {step}: {loss_ddf.item()}.")
            if torch.isnan(loss_tot) or torch.isinf(loss_tot):
                print(f"*** Encountered NaN or Inf loss at epoch {epoch}, step {step}. Skipping this batch.")
                loss_nan_step += 1
                continue
            if loss_nan_step > 5:
                print(f"*** Too many NaN or Inf losses ({loss_nan_step} times) at epoch {epoch}, step {step}. Stopping training.")
                raise ValueError("Too many NaN losses detected in loss_tot. Code terminated.")

            optimizer.zero_grad()
            loss_tot.backward()
            optimizer.step()

            epoch_loss_tot += loss_tot.item() / total
            epoch_loss_gen_d += loss_gen_d.item() / total
            epoch_loss_gen_a += loss_gen_a.item() / total
            epoch_loss_reg_ += loss_ddf.item() / total

            # ==========================================================
            # Mode 2: Registration training on paired images
            # ==========================================================
            if step % train_mode_ratio == 0:
                [x1, y1, _, embd_y] = batch_p
                if np.random.uniform(0, 1) < TEXT_EMBED_PROB:
                    embd_y = embd_y.to(om.device).type(torch.float32)
                else:
                    embd_y = None

                x1 = x1.to(om.device).type(torch.float32)
                y1 = y1.to(om.device).type(torch.float32)
                n = x1.size()[0]

                # Augmentation
                [x1, y1] = utils.random_permute([x1, y1], select_dims=[-1, -2, -3])
                if config['noise_scale'] > 0:
                    [x1, y1] = thresh_img([x1, y1], [0, 2 * config['noise_scale']])
                    random_scale = np.random.normal(1, config['noise_scale'] * 1)
                    random_shift = np.random.normal(0, config['noise_scale'] * 1)
                    x1 = x1 * random_scale + random_shift
                    y1 = y1 * random_scale + random_shift

                # Timestep schedule for reverse diffusion
                scale_regist = np.random.uniform(0.0, 0.7)
                T_regist = sorted(
                    random.sample(
                        range(int(om.timesteps * scale_regist), om.timesteps), 16,
                    ),
                    reverse=True,
                )
                T_regist = [
                    [t_val for _ in range(config["batchsize"] // 2)]
                    for t_val in T_regist
                ]

                proc_type = random.choice(['downsample', 'slice', 'slice1', 'none', 'none'])
                y1_proc, msk_tgt, cond_ratio = om._proc_cond_img(
                    y1, proc_type=proc_type,
                )

                # Reverse diffusion for registration (via OMorpher's network & STN)
                ddf_comp, img_rec = reverse_diffuse_train(
                    network, om, x1, y1_proc, T=[None, T_regist], text=embd_y,
                )

                # Registration losses
                loss_sim = loss_imgsim(img_rec, y1, label=(y1 > thresh_imgsim))
                loss_mse = loss_imgmse(img_rec, y1, label=(y1 >= 0.0))
                loss_ddf1 = loss_reg1(ddf_comp, img=y1)

                loss_regist = 0
                loss_regist += LOSS_WEIGHTS_REGIST[0] * loss_sim
                loss_regist += LOSS_WEIGHTS_REGIST[1] * loss_mse
                loss_regist += LOSS_WEIGHTS_REGIST[2] * loss_ddf1

                # NaN / divergence checks
                if torch.isnan(x0).any():
                    print(f"*** Encountered NaN in input image x0 at epoch {epoch}, step {step}.")
                if loss_ddf1 > 0.002:
                    print(f"*** High registration DDF loss at epoch {epoch}, step {step}: {loss_ddf1.item()}.")

                loss_regist = torch.sqrt(cond_ratio + MSK_EPS) * loss_regist
                optimizer.zero_grad()
                loss_regist.backward()

                torch.nn.utils.clip_grad_norm_(network.parameters(), max_norm=0.4)
                optimizer.step()

                epoch_loss_regist += loss_regist.item() / total
                epoch_loss_imgsim_ += loss_sim.item() / total
                epoch_loss_imgmse_ += loss_mse.item() / total
                epoch_loss_ddfreg += loss_ddf1.item() / total

            if step % 10 == 0:
                print('step:', step, ':', loss_tot.item(), '=', loss_gen_a.item(), '+', loss_gen_d.item(), '+', loss_ddf.item())
                print(f'     loss_regist: {loss_regist} = {loss_sim} (imgsim) + {loss_mse} (imgmse) + {loss_ddf1} (ddf)')

        if 1:
            print('==================')
            print(epoch, ':', epoch_loss_tot, '=', epoch_loss_gen_a, '+', epoch_loss_gen_d, '+', epoch_loss_reg_, ' (ang+dist+regul)')
            print(f'     loss_regist: {epoch_loss_regist} = {epoch_loss_imgsim_} (imgsim) + {epoch_loss_imgmse_} (imgmse) + {epoch_loss_ddfreg} (ddf)')
            print('==================')

        if 0 == epoch % epoch_per_save:
            save_path = os.path.join(model_dir, str(epoch).rjust(6, '0') + suffix_pth)
            os.makedirs(model_dir, exist_ok=True)
            if not use_distributed:
                print(f"saved in {save_path}")
                save_checkpoint(network, optimizer, epoch, save_path, use_dist=False)
            elif gpu_id == 0:
                print(f"saved in {save_path}")
                save_checkpoint(network, optimizer, epoch, save_path, use_dist=True)

    # Resource cleanup
    torch.cuda.empty_cache()
    gc.collect()
    if use_distributed and dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    if use_distributed:
        world_size = torch.cuda.device_count()
        print(f"Distributed GPU number = {world_size}")
        mp.spawn(main_train, args=(world_size,), nprocs=world_size)
    else:
        main_train(0, 1)