train_cvae.py

"""Entrypoint de treino do Conditional VAE (CVAE) para imagens CPRE.

Treina um beta-VAE condicionado pela classe. Depois usa-se
scripts/generate_cvae.py para gerar imagens sinteticas (uma classe ou todas).

Uso
---
    # treina com TODAS as classes (condicional) — recomendado:
    python train_cvae.py --config configs/cvae.yaml --data_root database/splits/seed1

    # treina so com uma classe (VAE dedicado a classe rara):
    python train_cvae.py --config configs/cvae.yaml --classes Stricture --run_name cvae_stricture

    # overrides rapidos:
    python train_cvae.py --config configs/cvae.yaml --epochs 50 --batch_size 24 --lr 1e-4
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

# Evita "OMP Error #15" (libiomp5md.dll vs libomp.dll) comum em Windows+conda
# quando torch (MKL) e matplotlib trazem copias diferentes do runtime OpenMP.
# Tem de ser definido ANTES de importar torch/numpy. Como os DataLoader workers
# no Windows usam spawn (re-importam este modulo), isto corre tambem em cada worker.
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")

import numpy as np
import torch

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

# NOTA: matplotlib e' importado de forma preguicosa dentro de save_sample_grid,
# para que os workers do DataLoader (spawn no Windows) nao o re-importem.

from src.data.loaders import build_dataloaders
from src.models.cvae import (
    AuxClassifier,
    ConditionalVAE,
    Discriminator,
    PerceptualLoss,
    d_hinge_loss,
    diff_augment,
    g_hinge_loss,
    vae_loss,
)
from src.utils.config import apply_overrides, load_config, save_config
from src.utils.logging_utils import setup_logger
from src.utils.seed import set_seed


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--config", type=str, required=True)
    p.add_argument("--run_name", type=str, default=None)
    p.add_argument("--epochs", type=int, default=None)
    p.add_argument("--batch_size", type=int, default=None)
    p.add_argument("--lr", type=float, default=None)
    p.add_argument("--img_size", type=int, default=None)
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--data_root", type=str, default=None)
    p.add_argument("--classes", nargs="+", default=None,
                   help="Subconjunto de classes a treinar (default: todas do config).")
    return p.parse_args()


def to_single_channel(x: torch.Tensor) -> torch.Tensor:
    """O dataset replica grayscale para 3 canais; o CVAE usa 1. Pega o canal 0."""
    if x.size(1) > 1:
        return x[:, :1]
    return x


def current_beta(epoch: int, beta_max: float, warmup: int) -> float:
    """KL annealing linear: 0 -> beta_max ao longo de `warmup` epochs."""
    if warmup <= 0:
        return beta_max
    return beta_max * min(1.0, (epoch + 1) / warmup)


@torch.no_grad()
def save_sample_grid(model, class_names, n_per_class, device, out_path, img_size):
    """Gera n_per_class imagens por classe e guarda um grid (linhas=classes)."""
    import matplotlib
    matplotlib.use("Agg")          # backend sem display (so escreve ficheiros)
    import matplotlib.pyplot as plt

    model.eval()
    n_cls = len(class_names)
    fig, axes = plt.subplots(n_cls, n_per_class,
                             figsize=(n_per_class * 1.4, n_cls * 1.4))
    axes = np.atleast_2d(axes)
    for ci in range(n_cls):
        imgs = model.sample(n_per_class, y=ci, device=device).cpu().numpy()
        for j in range(n_per_class):
            ax = axes[ci, j]
            ax.imshow(np.clip(imgs[j, 0], 0, 1), cmap="gray", vmin=0, vmax=1)
            ax.set_xticks([]); ax.set_yticks([])
            if j == 0:
                ax.set_ylabel(class_names[ci], fontsize=8, rotation=0,
                              ha="right", va="center")
    fig.suptitle(f"CVAE samples ({img_size}x{img_size})", fontsize=10)
    fig.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=120, bbox_inches="tight")
    plt.close(fig)


@torch.no_grad()
def evaluate(model, loader, device, beta, recon_loss, free_bits):
    model.eval()
    tot = {"total": 0.0, "recon": 0.0, "kl": 0.0}
    n = 0
    for x, y in loader:
        x = to_single_channel(x).to(device); y = y.to(device)
        recon, mu, logvar = model(x, y)
        L = vae_loss(recon.float(), x, mu.float(), logvar.float(),
                     beta=beta, recon_loss=recon_loss, free_bits=free_bits)
        bs = x.size(0); n += bs
        for k in tot:
            tot[k] += L[k].item() * bs
    return {k: v / max(n, 1) for k, v in tot.items()}


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    cfg = apply_overrides(cfg, {
        "run_name": args.run_name,
        "training.epochs": args.epochs,
        "training.batch_size": args.batch_size,
        "training.lr": args.lr,
        "data.img_size": args.img_size,
        "data.root": args.data_root,
        "seed": args.seed,
    })
    if args.classes is not None:
        cfg.data["classes"] = list(args.classes)

    set_seed(int(cfg.get("seed", 42)))
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Categoria do run -> isola VAE puro de VAE-GAN no disco. Auto-detectada
    # via adversarial.enabled (true -> "gan", false/missing -> "vae"). Tudo do
    # run vai para outputs/cvae/<categoria>/<run_name>/  (samples em samples/).
    category = "gan" if bool(cfg.get("adversarial", {}).get("enabled", False)) else "vae"
    run_dir = ROOT / "outputs" / "cvae" / category / str(cfg.run_name)
    run_dir.mkdir(parents=True, exist_ok=True)
    sample_dir = run_dir / "samples"
    logger = setup_logger("cvae", log_file=run_dir / "train.log")
    logger.info(f"Device: {device}  | run: {cfg.run_name}  | category: {category}")
    logger.info(f"Run dir: {run_dir}")
    save_config(cfg, run_dir / "config_used.yaml")

    classes = list(cfg.data.classes)
    img_size = int(cfg.data.img_size)
    tr = cfg.training

    # ---- dados (train + val)
    # Oversampling inverso a frequencia: o dataset e' ~5x desequilibrado
    # (Lithiasis 533 vs Biliary_Leaks 108). Sem isto a classe rara ve' ~5x
    # menos batches -> os seus parametros FiLM mal treinam -> e' a pior gerada
    # (exatamente a classe que mais precisamos de aumentar). default: True.
    use_weighted = bool(cfg.data.get("weighted_sampler", True))
    loaders = build_dataloaders(cfg, splits=("train", "val"),
                                use_weighted_sampler=use_weighted)
    logger.info(f"Classes: {classes}  | img {img_size}  | batch {int(tr.batch_size)} "
                f"| weighted_sampler={use_weighted}")

    # ---- modelo
    # latent_channels e' o que define o latente real (Cz*h*w); latent_dim no
    # config e' so' informativo (o modelo ignora-o, mantido por retrocompat).
    model = ConditionalVAE(
        img_size=img_size,
        in_channels=int(cfg.model.get("in_channels", 1)),
        num_classes=len(classes),
        latent_channels=int(cfg.model.get("latent_channels", 8)),
        latent_dim=int(cfg.model.get("latent_dim", 256)),     # retrocompat (ignorado)
        base_channels=int(cfg.model.base_channels),
        cond_embed_dim=int(cfg.model.cond_embed_dim),
        n_downsamples=cfg.model.get("n_downsamples", None),
    ).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    logger.info(f"CVAE params: {n_params/1e6:.2f}M  | latent_dim {model.latent_dim} "
                f"(Cz={model.latent_channels} x {model.feat_size}x{model.feat_size})")

    # ---- classificador auxiliar (forca as classes a ficarem distintas)
    # aux_start_epoch: ativa o aux SO a partir desse epoch (default 0 = sempre).
    # aux_warmup_epochs: rampa linear 0 -> aux_weight ao longo de N epochs apos
    # o start, em vez de ligar de chofre (que causaria "salto" na loss).
    #
    # Util quando o aux ativo desde epoch 1 empurra o decoder para texturas
    # discriminativas de alta frequencia ANTES de aprender a estrutura global
    # (FOV circular, fundo preto). Com start atrasado + warmup, o VAE aprende
    # primeiro a media + estrutura, depois o aux entra suavemente para refinar
    # a separacao de classes.
    aux_weight = float(tr.get("aux_weight", 0.0))
    aux_start_epoch = int(tr.get("aux_start_epoch", 0))
    aux_warmup_epochs = int(tr.get("aux_warmup_epochs", 0))
    use_aux = aux_weight > 0.0
    aux_clf = None
    if use_aux:
        aux_clf = AuxClassifier(
            in_channels=int(cfg.model.get("in_channels", 1)),
            num_classes=len(classes),
            base=int(cfg.model.get("aux_base_channels", 16)),
        ).to(device)
        when = (f"sempre" if aux_start_epoch <= 0
                else f"a partir do epoch {aux_start_epoch}"
                + (f" (warmup {aux_warmup_epochs}ep)" if aux_warmup_epochs > 0 else ""))
        logger.info(f"AuxClassifier ON | aux_weight={aux_weight} | "
                    f"ativo: {when} | "
                    f"params {sum(p.numel() for p in aux_clf.parameters())/1e6:.2f}M")

    def current_aux_w(ep: int) -> float:
        """Peso aux no epoch atual: 0 antes do start, rampa linear ate' aux_weight."""
        if not use_aux or ep < aux_start_epoch:
            return 0.0
        if aux_warmup_epochs <= 0:
            return aux_weight
        return aux_weight * min(1.0, (ep - aux_start_epoch + 1) / aux_warmup_epochs)

    # ---- perceptual loss (VGG16) — empurra para imagens nitidas (ductos finos)
    percep_weight = float(tr.get("perceptual_weight", 0.0))
    use_percep = percep_weight > 0.0
    percep_fn = None
    if use_percep:
        percep_resize = tr.get("perceptual_resize", 256)
        percep_fn = PerceptualLoss(resize_to=percep_resize).to(device).eval()
        for p in percep_fn.parameters():
            p.requires_grad_(False)
        logger.info(f"PerceptualLoss ON | perceptual_weight={percep_weight} | "
                    f"resize_to={percep_resize} (VGG16 congelada)")

    # ---- discriminador (VAE-GAN condicional) — empurra para alta frequencia real
    adv_cfg = cfg.get("adversarial", {})
    use_adv = bool(adv_cfg.get("enabled", False))
    adv_weight = float(adv_cfg.get("weight", 0.0))
    fm_weight = float(adv_cfg.get("fm_weight", 0.0))
    adv_start = int(adv_cfg.get("start_epoch", 0))
    adv_warmup = int(adv_cfg.get("warmup_epochs", 0))
    # DiffAugment (anti-memorizacao do D com poucos dados) + R1 (penaliza gradiente
    # do D no real; estabiliza). diffaug_policy="" desliga. r1_gamma=0 desliga o R1.
    diffaug_on = bool(adv_cfg.get("diffaugment", True))
    diffaug_policy = (str(adv_cfg.get("diffaug_policy", "color,translation,cutout"))
                      if diffaug_on else "")
    r1_gamma = float(adv_cfg.get("r1_gamma", 0.0))
    lr_d = float(adv_cfg.get("lr_d", tr.lr))
    disc = optD = None
    if use_adv:
        disc = Discriminator(
            in_channels=int(cfg.model.get("in_channels", 1)),
            num_classes=len(classes),
            base=int(adv_cfg.get("base_channels", 32)),
            img_size=img_size,
            fm_layer=int(adv_cfg.get("fm_layer", 3)),
        ).to(device)
        logger.info(
            f"Discriminator ON | adv_weight={adv_weight} fm_weight={fm_weight} | "
            f"start_epoch={adv_start} warmup={adv_warmup} | "
            f"lr_d={lr_d:.1e} diffaug='{diffaug_policy}' r1_gamma={r1_gamma} | "
            f"params {sum(p.numel() for p in disc.parameters())/1e6:.2f}M")

    # ---- optimizer / AMP
    opt_name = str(tr.get("optimizer", "adamw")).lower()
    OptCls = torch.optim.AdamW if opt_name == "adamw" else torch.optim.Adam
    params = list(model.parameters())            # gerador (G): VAE + aux
    if use_aux:
        params += list(aux_clf.parameters())
    optimizer = OptCls(params, lr=float(tr.lr),
                       weight_decay=float(tr.get("weight_decay", 0.0)))
    if use_adv:
        optD = torch.optim.Adam(disc.parameters(), lr=lr_d,
                                betas=(0.0, 0.9))     # betas tipicos de GAN (TTUR via lr_d)
    use_amp = bool(tr.get("use_amp", True)) and device.type == "cuda"
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)

    # ---- LR scheduler (cosine annealing) - opcional, mas estabiliza muito o
    # treino longo: o LR cai suavemente para eta_min, refinando sem oscilar nos
    # ultimos epochs. lr_scheduler: cosine | none. Aplica-se SO ao gerador.
    sched_name = str(tr.get("lr_scheduler", "none")).lower()
    sched = None
    if sched_name == "cosine":
        eta_min = float(tr.get("lr_min", float(tr.lr) * 0.1))
        sched = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=int(tr.epochs), eta_min=eta_min)
        logger.info(f"LR scheduler: cosine | T_max={int(tr.epochs)} "
                    f"| lr {float(tr.lr):.1e} -> {eta_min:.1e}")

    def adv_ramp(ep: int) -> float:
        """Peso adversarial: 0 antes de start, rampa linear 0->1 ao longo de warmup."""
        if not use_adv or ep < adv_start:
            return 0.0
        if adv_warmup <= 0:
            return 1.0
        return min(1.0, (ep - adv_start + 1) / adv_warmup)

    recon_loss = str(tr.get("recon_loss", "mse")).lower()
    beta_max = float(tr.get("beta", 1.0))
    warmup = int(tr.get("beta_warmup_epochs", 0))
    free_bits = float(tr.get("free_bits", 0.0))
    grad_clip = float(tr.get("grad_clip", 0.0))
    sample_every = int(tr.get("sample_every", 5))
    sample_n = int(tr.get("sample_n_per_class", 8))
    epochs = int(tr.epochs)

    es = cfg.get("early_stopping", {})
    es_enabled = bool(es.get("enabled", True))
    es_monitor = str(es.get("monitor", "val_total")).replace("val_", "")
    es_patience = int(es.get("patience", 20))
    es_min_delta = float(es.get("min_delta", 0.0))
    # Num VAE-GAN, recon (MSE) MAIOR pode ser MELHOR (imagens nitidas != media
    # borrada), por isso val_total deixa de ser um criterio fiavel de "best".
    # Desliga-se o early stopping e guarda-se best.pt = epoch mais recente; usa
    # as grelhas de amostras (sample_every) para escolher visualmente.
    if use_adv:
        es_enabled = False

    best = float("inf")
    best_epoch = -1
    no_improve = 0
    history = []

    import torch.nn.functional as F
    for epoch in range(epochs):
        model.train()
        # peso aux do epoch atual (0 antes do start, rampa ate' aux_weight)
        aux_w_now = current_aux_w(epoch)
        aux_active = aux_w_now > 0.0
        if use_aux:
            aux_clf.train()
        if use_adv:
            disc.train()
        beta = current_beta(epoch, beta_max, warmup)
        adv_w = adv_weight * adv_ramp(epoch)      # 0 ate' start_epoch, depois rampa
        run = {"total": 0.0, "recon": 0.0, "kl": 0.0, "aux": 0.0,
               "percep": 0.0, "d": 0.0, "adv": 0.0, "fm": 0.0, "r1": 0.0}
        aux_correct = 0
        seen = 0
        for x, y in loaders["train"]:
            x = to_single_channel(x).to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            bs = x.size(0)

            # ---- forward do gerador: recon + 1 amostra do prior (p/ adversarial/aux)
            with torch.amp.autocast("cuda", enabled=use_amp):
                recon, mu, logvar = model(x, y)
                samp = y_s = None
                if use_adv:
                    z = torch.randn(bs, model.latent_dim, device=device)
                    y_s = torch.randint(0, len(classes), (bs,), device=device)
                    samp = model.decode(z, y_s)
            recon_f = recon.float()
            samp_f = samp.float() if use_adv else None

            # ---- passo do DISCRIMINADOR (DiffAugment no real e nos fakes; hinge [+R1])
            d_val = r1_val = 0.0
            if use_adv and adv_w > 0.0:
                optD.zero_grad(set_to_none=True)
                # DiffAugment: MESMO tipo de augmentation no real e no fake -> o D
                # nao decora o conjunto real (chave com poucos dados).
                x_da = diff_augment(x, diffaug_policy)
                recon_da = diff_augment(recon_f.detach(), diffaug_policy)
                samp_da = diff_augment(samp_f.detach(), diffaug_policy)
                if r1_gamma > 0.0:
                    # R1 precisa de double-backward -> D-step em fp32 (sem AMP/scaler),
                    # com optD.step() direto. O scaler so' gere o passo do gerador.
                    x_da = x_da.detach().requires_grad_(True)
                    dl_real, _ = disc(x_da, y)
                    dl_recon, _ = disc(recon_da, y)
                    dl_samp, _ = disc(samp_da, y_s)
                    d_loss = d_hinge_loss(dl_real, [dl_recon, dl_samp])
                    grad_real = torch.autograd.grad(
                        dl_real.sum(), x_da, create_graph=True)[0]
                    r1 = grad_real.pow(2).flatten(1).sum(1).mean()
                    (d_loss + 0.5 * r1_gamma * r1).backward()
                    if grad_clip > 0:
                        torch.nn.utils.clip_grad_norm_(disc.parameters(), grad_clip)
                    optD.step()
                    d_val = d_loss.item(); r1_val = r1.item()
                else:
                    with torch.amp.autocast("cuda", enabled=use_amp):
                        dl_real, _ = disc(x_da, y)
                        dl_recon, _ = disc(recon_da, y)
                        dl_samp, _ = disc(samp_da, y_s)
                    d_loss = d_hinge_loss(dl_real.float(),
                                          [dl_recon.float(), dl_samp.float()])
                    scaler.scale(d_loss).backward()
                    if grad_clip > 0:
                        scaler.unscale_(optD)
                        torch.nn.utils.clip_grad_norm_(disc.parameters(), grad_clip)
                    scaler.step(optD)
                    d_val = d_loss.item()

            # ---- passo do GERADOR: VAE + perceptual + aux + adversarial + FM
            optimizer.zero_grad(set_to_none=True)
            L = vae_loss(recon_f, x, mu.float(), logvar.float(),
                         beta=beta, recon_loss=recon_loss, free_bits=free_bits)
            total = L["total"]
            aux_val = 0.0
            if aux_active:
                logits_real = aux_clf(x)
                logits_recon = aux_clf(recon_f)
                ce_real = F.cross_entropy(logits_real, y)
                ce_recon = F.cross_entropy(logits_recon, y)
                aux = ce_real + ce_recon
                if use_adv:
                    # CE tambem nas AMOSTRAS (z~prior) -> forca condicionamento
                    # na GERACAO (nao so na reconstrucao). Corrige o colapso visto.
                    ce_samp = F.cross_entropy(aux_clf(samp_f), y_s)
                    aux = aux + ce_samp
                # peso rampado (aux_w_now), nao o aux_weight fixo do config
                total = total + aux_w_now * aux
                aux_val = aux.item()
                aux_correct += (logits_real.argmax(1) == y).sum().item()
            percep_val = 0.0
            if use_percep:
                p = percep_fn(recon_f, x)
                total = total + percep_weight * p
                percep_val = p.item()
            adv_val = fm_val = 0.0
            if use_adv and adv_w > 0.0:
                # D nao se atualiza neste passo (so' propaga gradiente p/ o gerador)
                for p_ in disc.parameters():
                    p_.requires_grad_(False)
                # hinge: nas logits aplica-se o MESMO DiffAugment do D-step (consistencia)
                recon_da_g = diff_augment(recon_f, diffaug_policy)
                samp_da_g = diff_augment(samp_f, diffaug_policy)
                with torch.amp.autocast("cuda", enabled=use_amp):
                    gl_recon, _ = disc(recon_da_g, y)
                    gl_samp, _ = disc(samp_da_g, y_s)
                    # feature matching em imagens LIMPAS (alvo de features sem ruido
                    # de augmentation; o D nao treina aqui, logo nao ha leak)
                    _, feat_recon = disc(recon_f, y)
                    with torch.no_grad():
                        _, feat_real = disc(x, y)
                g_adv = g_hinge_loss([gl_recon.float(), gl_samp.float()])
                # Feature matching NORMALIZADO. As features do D (spectral norm +
                # LeakyReLU sobre input [0,1]) tem magnitude minima, por isso a L1
                # crua ficava ~0 e fm_weight nao tinha efeito (fm=0.00 durante as 150
                # epochs). Dividir pela magnitude media das features reais torna a fm
                # uma diferenca RELATIVA (~O(1)) -> fm_weight passa a estabilizar de
                # facto o treino adversarial (chave com poucos dados).
                feat_recon_f = feat_recon.float()
                feat_real_f = feat_real.float()
                fm = (F.l1_loss(feat_recon_f, feat_real_f)
                      / (feat_real_f.abs().mean() + 1e-6))
                total = total + adv_w * g_adv + fm_weight * fm
                for p_ in disc.parameters():
                    p_.requires_grad_(True)
                adv_val = g_adv.item()
                fm_val = fm.item()
            scaler.scale(total).backward()
            if grad_clip > 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(params, grad_clip)
            scaler.step(optimizer)
            scaler.update()

            seen += bs
            run["total"] += total.item() * bs
            run["recon"] += L["recon"].item() * bs
            run["kl"] += L["kl"].item() * bs
            run["aux"] += aux_val * bs
            run["percep"] += percep_val * bs
            run["d"] += d_val * bs
            run["adv"] += adv_val * bs
            run["fm"] += fm_val * bs
            run["r1"] += r1_val * bs
        tr_metrics = {k: v / max(seen, 1) for k, v in run.items()}
        aux_acc = aux_correct / max(seen, 1) if aux_active else 0.0
        val_metrics = evaluate(model, loaders["val"], device,
                               beta, recon_loss, free_bits)

        # mostra peso aux atual (rampado) quando ativo, OFF quando ainda nao iniciou
        if aux_active:
            aux_str = (f" | aux={tr_metrics['aux']:.2f} acc={aux_acc:.2f}"
                       f" aw={aux_w_now:.2f}")
        elif use_aux:
            aux_str = f" | aux=OFF (ate' epoch {aux_start_epoch})"
        else:
            aux_str = ""
        percep_str = (f" perc={tr_metrics['percep']:.3f}" if use_percep else "")
        r1_str = f" r1={tr_metrics['r1']:.3f}" if (use_adv and r1_gamma > 0) else ""
        adv_str = (f" | D={tr_metrics['d']:.3f} adv={tr_metrics['adv']:.3f} "
                   f"fm={tr_metrics['fm']:.3f}{r1_str} aw={adv_w:.2f}"
                   if (use_adv and adv_w > 0) else "")
        logger.info(
            f"[{epoch+1:3d}/{epochs}] beta={beta:.3f} | "
            f"train total={tr_metrics['total']:.2f} recon={tr_metrics['recon']:.2f} "
            f"kl={tr_metrics['kl']:.2f}{aux_str}{percep_str}{adv_str} | "
            f"val total={val_metrics['total']:.2f} recon={val_metrics['recon']:.2f} "
            f"kl={val_metrics['kl']:.2f}"
        )
        history.append({"epoch": epoch + 1, "beta": beta,
                        "train": tr_metrics, "val": val_metrics,
                        "lr": optimizer.param_groups[0]["lr"]})

        # LR scheduler step (apos calcular val, mas antes do checkpoint)
        if sched is not None:
            sched.step()

        # amostras periodicas
        if sample_every > 0 and ((epoch + 1) % sample_every == 0 or epoch == 0):
            save_sample_grid(model, classes, sample_n, device,
                             sample_dir / f"epoch_{epoch+1:03d}.png", img_size)

        # checkpoint + early stopping (monitor no val)
        score = val_metrics.get(es_monitor, val_metrics["total"])
        ckpt = {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "epoch": epoch + 1,
            "score": score,
            "classes": classes,
            "img_size": img_size,
            "model_kwargs": {
                "img_size": img_size,
                "in_channels": int(cfg.model.get("in_channels", 1)),
                "num_classes": len(classes),
                "latent_channels": int(model.latent_channels),       # define o latente real
                "latent_dim": int(model.latent_dim),                 # informativo (retrocompat)
                "base_channels": int(cfg.model.base_channels),
                "cond_embed_dim": int(cfg.model.cond_embed_dim),
                "n_downsamples": model.n_downsamples,
            },
        }
        if use_aux:
            ckpt["aux_state_dict"] = aux_clf.state_dict()
        if use_adv:
            ckpt["disc_state_dict"] = disc.state_dict()
            ckpt["optD_state_dict"] = optD.state_dict()
        torch.save(ckpt, run_dir / "last.pt")
        if use_adv:
            # GAN nao tem criterio de "best" fiavel no val (recon menor != melhor).
            # best.pt = epoch mais recente; escolhe visualmente pelas grelhas.
            best = score; best_epoch = epoch + 1
            torch.save(ckpt, run_dir / "best.pt")
        elif score < best - es_min_delta:
            best = score; best_epoch = epoch + 1; no_improve = 0
            torch.save(ckpt, run_dir / "best.pt")
            logger.info(f"  novo best ({es_monitor}={score:.2f}) -> best.pt")
        else:
            no_improve += 1
            if es_enabled and no_improve >= es_patience:
                logger.info(f"Early stopping (sem melhoria ha {no_improve} epochs).")
                break

    with (run_dir / "history.json").open("w", encoding="utf-8") as f:
        json.dump(history, f, indent=2)
    logger.info(f"Treino terminado. Best {es_monitor}={best:.2f} @ epoch {best_epoch}.")
    logger.info(f"Run dir: {run_dir}  | amostras: {sample_dir}")


if __name__ == "__main__":
    main()