ALDM2.py

# train_ct_spade_only.py
import os, sys, copy, math, time, argparse
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler
from torch.utils.data import DataLoader, random_split
from omegaconf import OmegaConf
from tqdm import tqdm
import neptune

# ---------------- ALDM repo 경로 ----------------
REPO = Path("/pscratch/sd/s/seojw/seojw/ALDM")
sys.path.append(str(REPO/"VQ-GAN"))
sys.path.append(str(REPO/"LDM"))

# VQ-GAN: taming 버전 사용 (질문에 준 시그니처)
from taming.models.vqgan import VQModel

# ---------------- Dataset ----------------
#from dataset_ct_t1_pairs import CTMRIPairDataset

import os
import torch
from torch.utils.data import Dataset
from typing import List
import torchio as tio

class CTMRIPairDataset(Dataset):
    """
    각 subject 폴더에 존재하는
      - ct_brain_crop_pad_(0,1).pt
      - mri_brain_crop_pad_(0,1).pt
    를 로드하여 (CT, T1) 쌍을 [-1,1]로 변환하여 반환.
    TorchIO Subject로 묶어 동일한 rigid augmentation 적용.
    """
    def __init__(self, base_dir: str, subjects: List[str], augment: bool = True):
        self.base_dir = base_dir
        self.subjects = []
        for s in subjects:
            ct = f"{base_dir}/{s}/ct_brain_crop_pad_(0,1).pt"
            mr = f"{base_dir}/{s}/mri_brain_crop_pad_(0,1).pt"
            if os.path.exists(ct) and os.path.exists(mr):
                self.subjects.append({"id": s, "ct": ct, "mr": mr})
        assert len(self.subjects) > 0, "No subject pairs found."

        self.augment = augment
        if augment:
            # Rigid 계열: 동일 랜덤 변환을 CT/T1에 공통 적용
            self.tx = tio.Compose([
                tio.RandomFlip(axes=(0,1,2), flip_probability=0.5),
                tio.RandomAffine(
                    scales=(0.95, 1.05),        # ±5% scale
                    degrees=10,                 # ±10° 회전
                    translation=5,              # ±5 vox 평행이동
                    isotropic=False,
                    center='image',
                ),
            ])
        else:
            self.tx = None

    def __len__(self):
        return len(self.subjects)

    @staticmethod
    def to_minus1_1(x: torch.Tensor) -> torch.Tensor:
        # (0,1) -> (-1,1)
        return x*2.0 - 1.0

    def __getitem__(self, idx):
        it = self.subjects[idx]
        # [D,H,W] 텐서 로드 (bfloat16 저장되어 있어도 float32로 변환)
        ct = torch.load(it["ct"]).to(torch.float32)
        t1 = torch.load(it["mr"]).to(torch.float32)

        # 채널 추가: [1,D,H,W]
        ct = ct.unsqueeze(0)
        t1 = t1.unsqueeze(0)

        # (0,1)->(-1,1)
        ct = self.to_minus1_1(ct)
        t1 = self.to_minus1_1(t1)

        # TorchIO Subject로 묶어서 동일 변환
        if self.tx:
            subj = tio.Subject(
                ct=tio.ScalarImage(tensor=ct),
                t1=tio.ScalarImage(tensor=t1)
            )
            subj = self.tx(subj)
            ct = subj.ct.tensor
            t1 = subj.t1.tensor

        return {"id": it["id"], "ct": ct, "t1": t1}



# --------- SPADE 분기 확장 유틸 ---------
def expand_spade_classes(model: VQModel, new_num_classes: int, init_from_idx: int = 0):
    """
    model.spade.* 내부의 SPADE_Multimodal(spades: ModuleList)을 찾아
    spades의 길이를 new_num_classes로 확장. 새로 추가되는 분기는 init_from_idx의 weight로 복제 초기화.
    """
    assert hasattr(model, "spade"), "This VQModel has no SPADE module."
    blocks = model.spade.block   # ModuleList of SPADEResnetBlock
    for b in blocks:
        for norm_name in ["norm_0", "norm_1"]:
            normm = getattr(b, norm_name)               # SPADE_Multimodal
            spades = normm.spades                       # ModuleList of SPADE (per class)
            old_n = len(spades)
            if new_num_classes <= old_n:
                continue
            # 기준 분기 weight 복제
            base = spades[init_from_idx]
            for _ in range(old_n, new_num_classes):
                new_branch = copy.deepcopy(base)
                spades.append(new_branch)
    return model

def freeze_all_but_ct_spade(model: VQModel, ct_class_index: int):
    # 전체 freeze
    for p in model.parameters():
        p.requires_grad_(False)
    # SPADE 중 ct 분기만 학습
    for b in model.spade.block:
        for norm_name in ["norm_0", "norm_1"]:
            spades = getattr(b, norm_name).spades
            for idx, sp in enumerate(spades):
                if idx == ct_class_index:
                    for p in sp.parameters():
                        p.requires_grad_(True)
    return model

# 2) freeze: T1 분기만 학습
def freeze_all_but_t1_spade(model: VQModel, t1_idx: int):
    for p in model.parameters():
        p.requires_grad_(False)
    for b in model.spade.block:
        for norm_name in ["norm_0", "norm_1"]:
            spades = getattr(b, norm_name).spades
            for i, sp in enumerate(spades):
                if i == t1_idx:
                    for p in sp.parameters():
                        p.requires_grad_(True)
    return model

def count_trainable_params(module: nn.Module) -> int:
    return sum(p.numel() for p in module.parameters() if p.requires_grad)

# --------- 학습 루프 ---------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base_dir", type=str, default="/pscratch/sd/s/seojw/CT_to_MRI/Task1-2/brain")
    parser.add_argument("--cfg_path", type=str, default=str(REPO/"VQ-GAN/configs/brats_vqgan_stage2.yaml"))
    parser.add_argument("--ckpt_path", type=str, default="/pscratch/sd/s/seojw/CT_to_MRI/stage2.ckpt")
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--num_workers", type=int, default=8)
    parser.add_argument("--max_epochs", type=int, default=20000)
    parser.add_argument("--lr", type=float, default=3e-5)
    parser.add_argument("--beta1", type=float, default=0.5)
    parser.add_argument("--beta2", type=float, default=0.9)
    parser.add_argument("--use_bf16", action="store_true", default=True)
    parser.add_argument("--log_interval", type=int, default=50)
    parser.add_argument("--ckpt_interval", type=int, default=5)
    parser.add_argument("--experiment_name", type=str, default="ct_spade_only_2")
    parser.add_argument("--grad_accum_steps", type=int, default=4)
    parser.add_argument("--neptune_project", type=str, default="ejswjawnj/CT-to-MRI")
    parser.add_argument("--use_discriminator", action="store_true")
    # 주의: 토큰은 환경변수로!
    args = parser.parse_args()

    # 실험/체크포인트 디렉토리
    save_root = Path("/pscratch/sd/s/seojw/CT_to_MRI/checkpoints")
    exp_dir = save_root / args.experiment_name
    exp_dir.mkdir(parents=True, exist_ok=True)

    # Neptune (토큰은 환경변수 NEPTUNE_API_TOKEN 권장)
    run = neptune.init_run(
        project=args.neptune_project,
        api_token="eyJhcGlfYWRkcmVzcyI6Imh0dHBzOi8vYXBwLm5lcHR1bmUuYWkiLCJhcGlfdXJsIjoiaHR0cHM6Ly9hcHAubmVwdHVuZS5haSIsImFwaV9rZXkiOiIzM2NmNzI4MC0yYmVkLTQ3YzEtYWU0Ni0yOTU4MmUwZmJmMTYifQ==",
        name=args.experiment_name,
        tags=["ALDM", "VQGAN-stage2", "SPADE-CT-only"]
    )

    # ----------------- 데이터 -----------------
    subjects = [s for s in os.listdir(args.base_dir) if "1B" in s]
    full = CTMRIPairDataset(args.base_dir, subjects, augment=True)
    n_train = int(len(full)*0.9)
    train_set, val_set = random_split(full, [n_train, len(full)-n_train])
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True,
                              num_workers=args.num_workers, pin_memory=True, persistent_workers=True)
    val_loader   = DataLoader(val_set, batch_size=args.batch_size, shuffle=False,
                              num_workers=args.num_workers, pin_memory=True, persistent_workers=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # ----------------- 모델 로드 -----------------
    cfg = OmegaConf.load(args.cfg_path)
    params = cfg.model.params

    # 기존 5 클래스(t1,t1ce,t2,flair,fake) + 1(CT)
    orig_n_classes = 5
    ct_class_index = 5
    t1_class_index = 0          # 레포에서 T1이 0번이었다면. 실제 맵에 맞게 설정.
    new_n_classes = ct_class_index + 1

    # (중요) SPADE 클래스 수를 new_n_classes로 덮어쓰기
    if isinstance(params, dict):
        params["num_classes"] = new_n_classes
    else:
        params.num_classes = new_n_classes

    model = VQModel(**params).to(device)
    ckpt = torch.load(args.ckpt_path, map_location="cpu", weights_only=False)
    sd = ckpt.get("state_dict", ckpt)

    # 먼저 로드 (spade 5개 기준으로 저장된 가중치). 새 분기는 아직 없음 → strict=False
    missing, unexpected = model.load_state_dict(sd, strict=False)
    print("First load => missing:", missing, "\nunexpected:", unexpected)

    # 실제 spade 분기 확장(6개) – 새 분기 weight는 T1 분기(=0) 복제
    model = expand_spade_classes(model, new_num_classes=new_n_classes, init_from_idx=0)

    # 나머지 모두 freeze, SPADE-CT만 학습
    #model = freeze_all_but_ct_spade(model, ct_class_index=ct_class_index)
    model = freeze_all_but_t1_spade(model, t1_class_index)      # 나머지 모두 freeze, SPADE-CT만 학습
    model.train()

    # --- Discriminator는 loss 내부에 있음: model.loss.discriminator ---
    # SPADE-CT만 학습 + Discriminator는 업데이트
    for p in model.loss.discriminator.parameters():
        p.requires_grad_(True)

    # Optimizers
    #spade_ct_params = []
    #for b in model.spade.block:
    #    for norm_name in ["norm_0", "norm_1"]:
    #        sp = getattr(b, norm_name).spades[ct_class_index]
    #        spade_ct_params += list(sp.parameters())
    #opt_ae   = torch.optim.Adam(spade_ct_params, lr=args.lr, betas=(args.beta1, args.beta2))
    
    
    spade_t1_params = []
    for b in model.spade.block:
        for norm_name in ["norm_0", "norm_1"]:
            sp = getattr(b, norm_name).spades[t1_class_index]
            spade_t1_params += list(sp.parameters())
    opt_ae   = torch.optim.Adam(spade_t1_params, lr=args.lr, betas=(args.beta1, args.beta2))
    
    
    if args.use_discriminator:
        for p in model.loss.discriminator.parameters():
            p.requires_grad_(True)
        # 더 약하게 시작하고 싶으면 lr를 줄이세요(예: 0.5x)
        opt_disc = torch.optim.Adam(model.loss.discriminator.parameters(),
                                    lr=args.lr * 0.5, betas=(args.beta1, args.beta2))
    else:
        for p in model.loss.discriminator.parameters():
            p.requires_grad_(False)
        opt_disc = None

    scaler = GradScaler(enabled=False)

    print(f"Trainable params (SPADE-CT only): {count_trainable_params(model):,d}")
    run["model/trainable_params"] = count_trainable_params(model)

    # ----------------- 학습 루프 -----------------
    global_step = 0
    accum_steps = max(1, args.grad_accum_steps)

    for epoch in range(1, args.max_epochs+1):
        pbar = tqdm(train_loader, desc=f"Epoch {epoch}")
        # ====== Train (AE step) ======
        model.train()
        opt_ae.zero_grad(set_to_none=True)
        #opt_disc.zero_grad(set_to_none=True)
        do_disc = args.use_discriminator and (epoch > args.disc_warmup_epochs)

        for it, batch in enumerate(pbar):
            global_step += 1
            ct = batch["ct"].to(device, non_blocking=True)  # [B,1,D,H,W], (-1,1)
            t1 = batch["t1"].to(device, non_blocking=True)

            # 1) Encode source(CT) & target(T1) → latent 4ch
            with torch.no_grad():
                z_ct, qloss_ct, _ = model.encode(ct)   # (B,4,D/4,H/4,W/4)
                z_t1, qloss_t1, _ = model.encode(t1)

            # 2) SPADE(CT 분기)로 z_ct → z_hat (target-like)
            #y_idx = torch.full((z_ct.size(0),), ct_class_index,
            #                   device=z_ct.device, dtype=torch.long)   # (B,)
            y_idx = torch.full((z_ct.size(0),), t1_class_index, device=z_ct.device, dtype=torch.long) # forward: y는 T1

            # ---------------- AE step (optimizer_idx=0) ----------------
            with autocast(enabled=args.use_bf16, dtype=torch.bfloat16):
                z_hat = model.spade(z_ct, y_idx)   # target-modulated latent
                # taming VQ loss: 입력=정답(latent), xrec=예측(latent)
                aeloss, log_dict_ae = model.loss(
                    qloss_t1,          # qloss (target의 양자화 손실 사용: 원본 구현 stage2와 동일 패턴)
                    z_t1,              # inputs (정답 latent)
                    z_hat,             # xrec   (예측 latent)
                    optimizer_idx=0,
                    global_step=global_step,
                    last_layer=model.get_last_layer(),  # decoder.conv_out.weight
                    label=y_idx,
                    split="train"
                )
                loss_ae = aeloss / accum_steps

            scaler.scale(loss_ae).backward()

            # Gradient Accumulation: step 조건
            if (it + 1) % accum_steps == 0:
                scaler.step(opt_ae)
                scaler.update()
                opt_ae.zero_grad(set_to_none=True)

            # ---------------- Disc step (optimizer_idx=1) ----------------
            do_disc = args.use_discriminator and (epoch > args.disc_warmup_epochs)
            if do_disc:
                # (기존 코드와 동일) discloss 역전파/스텝
                with autocast(enabled=args.use_bf16, dtype=torch.bfloat16), torch.no_grad():
                    z_hat = model.spade(z_ct, y_idx)
                with autocast(enabled=args.use_bf16, dtype=torch.bfloat16):
                    discloss, log_dict_disc = model.loss(
                        qloss_t1, z_t1, z_hat.detach(),
                        optimizer_idx=1, global_step=global_step,
                        last_layer=model.get_last_layer(), label=y_idx, split="train"
                    )
                    loss_disc = discloss / accum_steps

                scaler.scale(loss_disc).backward()
                if (it + 1) % accum_steps == 0:
                    scaler.step(opt_disc)
                    scaler.update()
                    opt_disc.zero_grad(set_to_none=True)

            # ---------------- Logging ----------------
            if global_step % args.log_interval == 0:
                run["train/aeloss"].append(aeloss.item())
                if do_disc:
                    run["train/discloss"].append(discloss.item())
                # 시각화: 중앙 슬라이스 (이미지 공간 아님 주의 — latent 시각화 대신 디코드해서 보여주기)
                with torch.no_grad(), autocast(enabled=args.use_bf16, dtype=torch.bfloat16):
                    x_hat = model.decode(z_hat.to(torch.float32))  # 디코더는 float32가 안전
                zc = x_hat.shape[2] // 2
                
                x_hat_log_img = (x_hat[0,0,zc].detach().cpu().float().numpy())
                run["train/images/x_hat"].append(
                    neptune.types.File.as_image( (x_hat_log_img - x_hat_log_img.min()) / (x_hat_log_img.max() - x_hat_log_img.min()) ))
                run["train/images/t1"].append(
                    neptune.types.File.as_image( (t1[0,0,zc].detach().cpu().float().numpy() + 1)/2 ))
                run["train/images/ct"].append(
                    neptune.types.File.as_image( (ct[0,0,zc].detach().cpu().float().numpy() + 1)/2 ))
                #pbar.set_postfix({"ae": f"{aeloss.item():.4f}", "disc": f"{discloss.item():.4f}"})
                pbar.set_postfix({"ae": f"{aeloss.item():.4f}"})

        # ====== Validation ======
        model.eval()
        val_losses = []
        with torch.no_grad():
            for batch in val_loader:
                ct = batch["ct"].to(device)
                t1 = batch["t1"].to(device)
                z_ct, _, _ = model.encode(ct)
                z_t1, qloss_t1, _ = model.encode(t1)
                #y_idx = torch.full((z_ct.size(0),), ct_class_index, device=z_ct.device, dtype=torch.long)
                y_idx = torch.full((z_ct.size(0),), t1_class_index, device=z_ct.device, dtype=torch.long)
                z_hat = model.spade(z_ct, y_idx)
                # AE쪽 기준으로만 간단히 평가 (latent L1)
                val_losses.append(F.l1_loss(z_hat.float(), z_t1.float()).item())
        val_loss = sum(val_losses)/max(1,len(val_losses))
        run["val/l1_latent"].append(val_loss)
        print(f"[val] epoch {epoch} latent L1={val_loss:.4f}")

        # ====== Checkpoint ======
        if epoch % args.ckpt_interval == 0:
            ckpt_out = exp_dir / f"{args.experiment_name}_epoch{epoch:03d}.pt"
            torch.save({
                "epoch": epoch,
                "model": model.state_dict(),
                "opt_ae": opt_ae.state_dict(),
                #"opt_disc": opt_disc.state_dict(),
                "ct_class_index": t1_class_index,
                "config": params,
            }, ckpt_out)
            run["checkpoints"].upload(str(ckpt_out))

    run.stop()

if __name__ == "__main__":
    main()