train.py

"""
Prism training loop — GAN + perceptual loss with temporal sequences.

Features:
  - FP8 training via torchao (RTX 4060/5060 Ti with compute 8.9+/12.0)
  - Mixed precision (AMP FP16/BF16) fallback
  - Apollo-Mini optimizer (SGD-level memory, AdamW-level quality)
  - Temporal training: sequences of N consecutive frames with hidden state
  - Multi-GPU: --device cuda:0 or cuda:1 (5060 Ti recommended)
  - Gradient accumulation for effective large batch sizes
  - wandb logging (optional)

Usage:
  # Fast on 5060 Ti with FP8 + Apollo-Mini:
  python train.py --data data/dataset --device cuda:1 --fp8 --optimizer apollo-mini --batch 8

  # Balanced on 4060:
  python train.py --data data/dataset --device cuda:0 --amp --batch 4

  # Quality with wandb logging:
  python train.py --data data/dataset --device cuda:1 --fp8 --model quality --wandb
"""

import argparse
import os
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.amp import autocast, GradScaler
from tqdm import tqdm

# V2 model: U-Net + Transformer (12.4ms at 1080p with cooperative vectors)
from model import PrismV2, ModelConfig, PRESETS as V2_PRESETS, warp
# V1 discriminator + losses (still used for GAN training)
from model import MultiScaleDiscriminator, PerceptualLoss, HingeLoss
# Game-style augmentation (applied to real video input so model learns color correction)
from augment_gpu import game_augment_gpu


# ============================================================================
# FP8 support via torchao
# ============================================================================

def enable_fp8_training(model: nn.Module) -> nn.Module:
    """Convert model's Linear layers to FP8 for faster training on Ada/Blackwell GPUs."""
    try:
        from torchao.float8 import Float8LinearConfig, convert_to_float8_training
        config = Float8LinearConfig()
        convert_to_float8_training(model, config=config)
        print("  FP8 training enabled via torchao")
        return model
    except Exception as e:
        print(f"  FP8 not available ({e}), falling back to FP16/BF16")
        return model


# ============================================================================
# Optimizer factory
# ============================================================================

def make_optimizer(name: str, params, lr: float) -> torch.optim.Optimizer:
    """Create optimizer by name."""
    if name == "apollo-mini":
        from apollo_torch import APOLLOAdamW
        param_groups = [{
            "params": list(params),
            "rank": 1,
            "proj": "random",
            "scale_type": "tensor",
            "scale": 128,
            "update_proj_gap": 200,
            "proj_type": "std",
        }]
        return APOLLOAdamW(param_groups, lr=lr, betas=(0.9, 0.999))

    elif name == "apollo":
        from apollo_torch import APOLLOAdamW
        param_groups = [{
            "params": list(params),
            "rank": 256,
            "proj": "random",
            "scale_type": "channel",
            "scale": 1,
            "update_proj_gap": 200,
            "proj_type": "std",
        }]
        return APOLLOAdamW(param_groups, lr=lr, betas=(0.9, 0.999))

    elif name == "adamw":
        return torch.optim.AdamW(params, lr=lr, betas=(0.0, 0.999))

    else:
        raise ValueError(f"Unknown optimizer: {name}")


# ============================================================================
# Dataset — supports single frames and temporal sequences
# ============================================================================

class PrismDataset(Dataset):
    """
    Simple fast dataset. Loads pre-cropped .pt files directly.
    Splits samples into real/synthetic pools and serves balanced batches.
    """
    def __init__(self, data_dir: Path, crop_size: int = 256, seq_len: int = 1,
                 max_samples: int = 0, **kwargs):
        self.all_samples = sorted(data_dir.glob("sample_*.pt"))
        self.crop_size = crop_size
        self.seq_len = seq_len

        if not self.all_samples:
            raise RuntimeError(f"No samples in {data_dir}")
        if max_samples > 0:
            self.all_samples = self.all_samples[:max_samples]

        # Check if files are pre-cropped (small) or full-size (large)
        sample_size = self.all_samples[0].stat().st_size
        self.pre_cropped = sample_size < 2 * 1024 * 1024  # < 2MB = pre-cropped

        # Load or build real/synthetic split index
        import json
        import random
        self._rng = random.Random(42)

        index_path = data_dir / "split_index.json"
        if index_path.exists():
            with open(index_path) as f:
                index = json.load(f)
            self.real_indices = index["real"]
            self.synth_indices = index["synth"]
            print(f"Loaded split index: {len(self.real_indices)} real, {len(self.synth_indices)} synth")
        else:
            print(f"No split_index.json found — scanning {len(self.all_samples)} samples...")
            self.real_indices = []
            self.synth_indices = []
            for i, path in enumerate(self.all_samples):
                try:
                    data = torch.load(path, weights_only=True)
                    if data.get("is_real", torch.tensor(True)).item():
                        self.real_indices.append(i)
                    else:
                        self.synth_indices.append(i)
                except:
                    self.real_indices.append(i)
                if (i + 1) % 20000 == 0:
                    print(f"  Scanned {i+1}/{len(self.all_samples)}")
            # Save for next time
            with open(index_path, "w") as f:
                json.dump({"real": self.real_indices, "synth": self.synth_indices}, f)
            print(f"  Saved split index to {index_path}")

        print(f"Dataset: {len(self.all_samples)} samples ({len(self.real_indices)} real, "
              f"{len(self.synth_indices)} synthetic), seq_len={seq_len}, "
              f"crop={crop_size}, pre_cropped={self.pre_cropped}")

        self.balanced = len(self.real_indices) > 0 and len(self.synth_indices) > 0
        if self.balanced:
            print(f"  Balanced sampling: ~50/50 real/synthetic per batch")

    def __len__(self) -> int:
        return len(self.all_samples) - self.seq_len + 1

    def __getitem__(self, idx: int) -> list[dict]:
        seq = []

        # Pick which pool this sequence comes from (real or synthetic)
        # All frames in a sequence come from the same pool for consistency
        if self.balanced:
            use_synth = self._rng.random() < 0.5 and len(self.synth_indices) >= self.seq_len
            pool = self.synth_indices if use_synth else self.real_indices
            start = self._rng.randint(0, max(0, len(pool) - self.seq_len))
        else:
            pool = None
            start = idx

        for i in range(self.seq_len):
            if pool is not None:
                j = pool[min(start + i, len(pool) - 1)]
            else:
                j = (idx + i) % len(self.all_samples)
            for attempt in range(3):
                try:
                    data = torch.load(self.all_samples[j], weights_only=False)
                    break
                except Exception:
                    j = torch.randint(0, len(self.all_samples), (1,)).item()
            else:
                # Last resort: return zeros
                data = {
                    "color": torch.zeros(3, 128, 128, dtype=torch.float16),
                    "depth": torch.zeros(1, 128, 128, dtype=torch.float16),
                    "motion_vectors": torch.zeros(2, 128, 128, dtype=torch.float16),
                    "ground_truth": torch.zeros(3, 256, 256, dtype=torch.float16),
                    "is_real": torch.tensor(True),
                }

            # Convert uint8 -> float16 (V3 compressed format)
            if data["color"].dtype == torch.uint8:
                data["color"] = data["color"].half() / 255.0
            if data["ground_truth"].dtype == torch.uint8:
                data["ground_truth"] = data["ground_truth"].half() / 255.0

            # Normalize GT to 2x input size for batching (training loop handles scale)
            _, cH, cW = data["color"].shape
            target_h, target_w = cH * 2, cW * 2
            if data["ground_truth"].shape[1] != target_h or data["ground_truth"].shape[2] != target_w:
                data["ground_truth"] = F.interpolate(
                    data["ground_truth"].unsqueeze(0).float(),
                    (target_h, target_w), mode="bilinear", align_corners=False
                ).squeeze(0).half()

            if not self.pre_cropped:
                data = self._crop(data)

            seq.append(data)
        return seq

    def _crop(self, data: dict) -> dict:
        _, rH, rW = data["color"].shape
        _, dH, dW = data["ground_truth"].shape
        cr = min(self.crop_size, rH, rW)
        cd = cr * 2

        if rH > cr and rW > cr:
            y = torch.randint(0, rH - cr, (1,)).item()
            x = torch.randint(0, rW - cr, (1,)).item()
            dy, dx = int(y * dH / rH), int(x * dW / rW)
            dh, dw = int(cr * dH / rH), int(cr * dW / rW)
        else:
            y, x, dy, dx, dh, dw = 0, 0, 0, 0, dH, dW

        return {
            "color": data["color"][:, y:y+cr, x:x+cr],
            "depth": data["depth"][:, y:y+cr, x:x+cr],
            "motion_vectors": data["motion_vectors"][:, y:y+cr, x:x+cr],
            "ground_truth": F.interpolate(
                data["ground_truth"][:, dy:dy+dh, dx:dx+dw].unsqueeze(0),
                size=(cd, cd), mode="bilinear", align_corners=False
            ).squeeze(0),
            **{k: data[k] for k in ["is_real", "jitter"] if k in data},
        }


def collate_sequences(batch: list[list[dict]]) -> list[dict]:
    """Collate batch with mixed scales — resize gt to match a randomly chosen scale."""
    import random
    seq_len = len(batch[0])

    # Pick one scale for the entire batch: 2x (80%) or 3x (20%)
    scale = 3 if random.random() < 0.3 else 2

    result = []
    # Keys that all samples must have (skip extras like 'scale')
    common_keys = {"color", "depth", "motion_vectors", "ground_truth", "is_real"}
    for t in range(seq_len):
        frame = {}
        for key in common_keys:
            tensors = [batch[b][t][key] for b in range(len(batch))]

            if key == "ground_truth":
                # Resize all gt to the chosen scale
                rH = batch[0][t]["color"].shape[1]
                rW = batch[0][t]["color"].shape[2]
                target_h = rH * scale
                target_w = rW * scale
                resized = []
                for gt in tensors:
                    if gt.shape[1] != target_h or gt.shape[2] != target_w:
                        gt = F.interpolate(gt.unsqueeze(0), (target_h, target_w),
                                           mode="bilinear", align_corners=False).squeeze(0)
                    resized.append(gt)
                frame[key] = torch.stack(resized)
            else:
                frame[key] = torch.stack(tensors)

        frame["target_scale"] = torch.tensor(scale, dtype=torch.int32)
        result.append(frame)
    return result


# ============================================================================
# Trainer
# ============================================================================

class Trainer:
    def __init__(
        self,
        model_name: str = "balanced",
        model_config=None,
        optimizer_name: str = "adamw",
        lr_g: float = 1e-4,
        lr_d: float = 4e-4,
        device: str = "cuda:1",
        use_amp: bool = True,
        use_fp8: bool = False,
        grad_accum: int = 1,
        use_wandb: bool = False,
    ):
        self.device = torch.device(device)
        self.use_amp = use_amp and self.device.type == "cuda"
        self.use_fp8 = use_fp8
        self.grad_accum = grad_accum

        if self.device.type == "cuda":
            gpu_name = torch.cuda.get_device_name(self.device)
            gpu_mem = torch.cuda.get_device_properties(self.device).total_memory / 1024**3
            gpu_cc = torch.cuda.get_device_capability(self.device)
            print(f"Device: {gpu_name} ({gpu_mem:.0f}GB, compute {gpu_cc[0]}.{gpu_cc[1]})")

        # Models — V2 architecture (U-Net + Transformer, optionally with MoE)
        if model_config is not None:
            cfg = model_config
        elif model_name in V2_PRESETS:
            cfg = V2_PRESETS[model_name]
        else:
            raise ValueError(f"Unknown model: {model_name}. Available: {list(V2_PRESETS.keys())}")
        self.G = PrismV2(cfg).to(self.device)
        self.D = MultiScaleDiscriminator().to(self.device)
        self.perceptual = PerceptualLoss().to(self.device)

        # FP8 conversion (before optimizer creation)
        if use_fp8:
            self.G = enable_fp8_training(self.G)
            self.D = enable_fp8_training(self.D)

        # torch.compile — disabled for now (Triton issues on compute 12.0)
        # if hasattr(torch, "compile"):
        #     try:
        #         self.G = torch.compile(self.G, mode="reduce-overhead")
        #         self.D = torch.compile(self.D, mode="reduce-overhead")
        #         print("  torch.compile: enabled")
        #     except Exception:
        #         print("  torch.compile: not available")

        g_params = sum(p.numel() for p in self.G.parameters())
        d_params = sum(p.numel() for p in self.D.parameters())
        print(f"Generator: {g_params/1e6:.2f}M params ({model_name})")
        print(f"Discriminator: {d_params/1e6:.2f}M params")

        # Optimizers
        self.opt_G = make_optimizer(optimizer_name, self.G.parameters(), lr=lr_g)
        # Discriminator always uses AdamW (Apollo not needed — D is discarded after training)
        self.opt_D = torch.optim.AdamW(self.D.parameters(), lr=lr_d, betas=(0.0, 0.999))
        print(f"Optimizer G: {optimizer_name} | D: AdamW")

        # Determine autocast dtype
        if use_fp8:
            self.amp_dtype = torch.bfloat16
        elif self.device.type == "cuda":
            self.amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        else:
            self.amp_dtype = torch.float32

        # GradScaler for FP16 stability. For BF16: scaler is disabled but we still
        # use gradient clipping (critical for transformer stability)
        use_scaler = self.use_amp and self.amp_dtype == torch.float16 and not use_fp8
        # Always enable scaler wrapper (disabled scalers are no-ops but let us use same code path)
        self.scaler_G = GradScaler("cuda", enabled=use_scaler)
        self.scaler_D = GradScaler("cuda", enabled=use_scaler)

        print(f"Precision: {'FP8+BF16' if use_fp8 else self.amp_dtype}")
        print(f"Grad accumulation: {grad_accum}")

        self.gan_loss = HingeLoss()
        self.epoch = 0

        self.use_wandb = use_wandb
        if use_wandb:
            from dataclasses import asdict
            import wandb
            G_inner = self.G.module if hasattr(self.G, 'module') else self.G
            wandb.init(project="prism", config={
                "model": model_name, "optimizer": optimizer_name,
                "lr_g": lr_g, "lr_d": lr_d, "fp8": use_fp8, "amp": use_amp,
                "device": str(device),
                "g_params": sum(p.numel() for p in self.G.parameters()),
                **{f"model/{k}": v for k, v in asdict(G_inner.cfg).items()},
            })

    def train_epoch(self, loader: DataLoader, adv_weight: float = 0.1,
                    temporal_weight: float = 0.3, d_every: int = 1) -> dict:
        """
        Streaming temporal training — each sample is one step.
        Hidden state flows forward between consecutive samples (detached).
        No backprop through time = same speed as non-temporal.
        Exactly matches inference behavior (one frame at a time).
        """
        self.G.train()
        self.D.train()

        totals = {"g": 0, "d": 0, "l1": 0, "perc": 0, "adv": 0, "temp": 0, "moe": 0, "moe_bal": 0, "n": 0}

        # Persistent hidden state across steps (streaming)
        prev_output = None
        prev_hidden = None

        for step, sequence in enumerate(tqdm(loader, desc=f"Epoch {self.epoch}")):
            frame = sequence[0]  # seq_len=1 now, single frame per step
            color = frame["color"].to(self.device).float()
            depth = frame["depth"].to(self.device)
            mv = frame["motion_vectors"].to(self.device).float()
            gt = frame["ground_truth"].to(self.device).float()
            is_real = frame.get("is_real", torch.ones(color.shape[0], dtype=torch.bool))
            is_real = is_real.to(self.device)

            # Game-style augmentation on real video input
            # This teaches the model to do color correction (real video looks
            # different from game renders — flatten lighting, sharpen edges, etc.)
            if is_real.any():
                color[is_real] = game_augment_gpu(color[is_real], strength=0.5)

            # G forward with streaming hidden state, target scale from data
            target_h, target_w = gt.shape[2], gt.shape[3]
            with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp or self.use_fp8):
                fake, hidden = self.G(color, depth, mv,
                                      prev_output=prev_output,
                                      prev_hidden=prev_hidden,
                                      target_h=target_h, target_w=target_w)

                if gt.shape != fake.shape:
                    gt = F.interpolate(gt, fake.shape[2:], mode="bilinear", align_corners=False)

                # --- D step: uses fake.detach() so no graph shared with G ---
                fake_det = fake.detach()
                d_fake_preds = self.D(fake_det)

                if is_real.any():
                    real_gt = gt[is_real]
                    real_preds = self.D(real_gt.detach())
                else:
                    real_preds = None

                if real_preds is not None:
                    d_loss = self.gan_loss.d_loss(real_preds, d_fake_preds)
                else:
                    d_loss = F.relu(1 + d_fake_preds[0]).mean()

            # D backward + step (separate graph, no retain needed)
            self.opt_D.zero_grad()
            self.scaler_D.scale(d_loss).backward()
            self.scaler_D.unscale_(self.opt_D)
            torch.nn.utils.clip_grad_norm_(self.D.parameters(), 1.0)
            self.scaler_D.step(self.opt_D)
            self.scaler_D.update()

            # --- G step: own D forward so gradients flow through G ---
            with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp or self.use_fp8):
                g_fake_preds = self.D(fake)
                l1 = F.l1_loss(fake, gt)
                # Perceptual loss on subset to save VRAM (8 samples max)
                perc_n = min(8, fake.shape[0])
                perc = self.perceptual(fake[:perc_n].float(), gt[:perc_n].float())
                adv = self.gan_loss.g_loss(g_fake_preds)

                # Temporal consistency loss
                if prev_output is not None:
                    warped_prev = warp(prev_output, F.interpolate(mv, prev_output.shape[2:],
                                       mode="bilinear", align_corners=False))
                    if warped_prev.shape != fake.shape:
                        warped_prev = F.interpolate(warped_prev, fake.shape[2:],
                                                     mode="bilinear", align_corners=False)
                    temp_loss = F.l1_loss(fake, warped_prev)
                else:
                    temp_loss = torch.tensor(0.0, device=self.device)

                # MoE auxiliary losses (balance + z-loss for router stability)
                G_inner = self.G.module if hasattr(self.G, 'module') else self.G
                if G_inner.cfg.n_experts > 0:
                    moe_balance = G_inner._moe_balance_loss
                    moe_z = G_inner._moe_z_loss
                    moe_w = G_inner.cfg.moe_balance_weight
                    moe_loss = moe_w * moe_balance + 0.001 * moe_z
                else:
                    moe_loss = torch.tensor(0.0, device=self.device)
                    moe_balance = torch.tensor(0.0, device=self.device)

                g_loss = l1 + 0.5 * perc + adv_weight * adv + temporal_weight * temp_loss + moe_loss

            # Update hidden state (detached — no BPTT, just streaming)
            prev_output = fake.detach()
            prev_hidden = hidden.detach() if hidden is not None else None

            # Reset hidden state occasionally to prevent staleness
            if step % 100 == 0:
                prev_output = None
                prev_hidden = None

            # G backward + step
            self.opt_G.zero_grad()
            self.scaler_G.scale(g_loss).backward()
            self.scaler_G.unscale_(self.opt_G)
            torch.nn.utils.clip_grad_norm_(self.G.parameters(), 1.0)
            self.scaler_G.step(self.opt_G)
            self.scaler_G.update()

            B = color.shape[0]
            totals["g"] += g_loss.item() * B
            totals["d"] += d_loss.item() * B
            totals["l1"] += l1.item() * B
            totals["perc"] += perc.item() * B
            totals["adv"] += adv.item() * B
            totals["temp"] += temp_loss.item() * B
            totals["moe"] += moe_loss.item() * B
            totals["moe_bal"] += moe_balance.item() * B
            totals["n"] += B

        n = max(totals["n"], 1)
        self.epoch += 1
        metrics = {k: totals[k] / n for k in ["g", "d", "l1", "perc", "adv", "temp", "moe", "moe_bal"]}

        if self.use_wandb:
            import wandb
            log_data = dict(metrics)

            # MoE expert usage stats
            G_inner = self.G.module if hasattr(self.G, 'module') else self.G
            if G_inner.cfg.n_experts > 0 and hasattr(G_inner, '_moe_expert_usage') and G_inner._moe_expert_usage:
                import numpy as np
                # Average expert usage across all blocks
                all_usage = torch.stack(G_inner._moe_expert_usage)  # (n_blocks, n_experts)
                avg_usage = all_usage.mean(dim=0).cpu().numpy()  # (n_experts,)
                n_experts = len(avg_usage)
                uniform = 1.0 / n_experts

                # Key health metrics
                log_data["moe/expert_usage_std"] = float(avg_usage.std())
                log_data["moe/expert_usage_max"] = float(avg_usage.max())
                log_data["moe/expert_usage_min"] = float(avg_usage.min())
                log_data["moe/max_min_ratio"] = float(avg_usage.max() / max(avg_usage.min(), 1e-8))
                log_data["moe/dead_experts"] = int((avg_usage < uniform * 0.1).sum())
                log_data["moe/balance_loss"] = metrics["moe_bal"]

                # Per-expert usage histogram
                log_data["moe/expert_usage"] = wandb.Histogram(avg_usage)

                # Per-block heatmap (every 10 epochs to avoid log spam)
                if self.epoch % 10 == 0:
                    usage_grid = all_usage.cpu().numpy()
                    log_data["moe/block_expert_heatmap"] = wandb.Image(
                        self._make_usage_heatmap(usage_grid, n_experts),
                        caption=f"Expert usage by block (epoch {self.epoch})"
                    )

            wandb.log(log_data, step=self.epoch)

            # Log sample images every 5 epochs
            if self.epoch % 5 == 0:
                self._log_sample_images()

        return metrics

    @staticmethod
    def _make_usage_heatmap(usage_grid, n_experts):
        """Create a simple heatmap image of expert usage per block."""
        import numpy as np
        from PIL import Image as PILImage
        # usage_grid: (n_blocks, n_experts)
        n_blocks = usage_grid.shape[0]
        uniform = 1.0 / n_experts
        # Normalize: 0 = no usage, 1 = uniform, >1 = overloaded
        normalized = usage_grid / uniform
        # Clamp to [0, 3] for visualization
        normalized = np.clip(normalized, 0, 3) / 3.0
        # Scale up for visibility
        cell_h, cell_w = 12, max(4, 200 // n_experts)
        img = np.zeros((n_blocks * cell_h, n_experts * cell_w, 3), dtype=np.uint8)
        for b in range(n_blocks):
            for e in range(n_experts):
                v = normalized[b, e]
                # Green = balanced, Red = overloaded, Blue = underused
                if v < 0.33:
                    r, g, bl = int(50 * v * 3), int(50 * v * 3), int(200 * (1 - v * 3))
                elif v < 0.5:
                    r, g, bl = 0, int(200 * (v - 0.33) * 6), 0
                else:
                    r, g, bl = int(200 * (v - 0.5) * 2), int(200 * (1 - (v - 0.5) * 2)), 0
                img[b*cell_h:(b+1)*cell_h, e*cell_w:(e+1)*cell_w] = [r, g, bl]
        return PILImage.fromarray(img)

    def set_fixed_sample(self, loader):
        """Grab a fixed sample from the dataset for consistent wandb image logging."""
        for sequence in loader:
            frame = sequence[0]
            self._fixed_color = frame["color"][:1].to(self.device).float()
            self._fixed_depth = frame["depth"][:1].to(self.device).float()
            self._fixed_mv = frame["motion_vectors"][:1].to(self.device).float()
            self._fixed_gt = frame["ground_truth"][:1].to(self.device).float()
            self._fixed_is_real = frame.get("is_real", torch.ones(1, dtype=torch.bool))[0].item()
            break

    @torch.no_grad()
    def _log_sample_images(self):
        """Generate and log sample images to wandb."""
        import wandb
        import numpy as np

        if not hasattr(self, "_fixed_color"):
            return

        self.G.eval()

        try:
            c = self._fixed_color
            d = self._fixed_depth
            mv = self._fixed_mv
            gt = self._fixed_gt
            _, _, rH, rW = c.shape
            _, _, dH, dW = gt.shape

            def to_np(t):
                return (t[0].cpu().float().clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8)

            images = {}

            # Input visualizations
            images["inputs/color"] = wandb.Image(to_np(c), caption=f"Color {rW}x{rH}")
            depth_np = d[0, 0].cpu().numpy()
            depth_np = ((depth_np - depth_np.min()) / (depth_np.max() - depth_np.min() + 1e-7) * 255).astype(np.uint8)
            images["inputs/depth"] = wandb.Image(depth_np, caption="Depth")
            images["inputs/ground_truth"] = wandb.Image(to_np(gt), caption=f"Ground truth {dW}x{dH}")
            tag = "real" if self._fixed_is_real else "synthetic"
            images["inputs/type"] = wandb.Image(to_np(c), caption=f"Source: {tag}")

            # 2x upscale
            target_2x = (rH * 2, rW * 2)
            with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp):
                out_2x, _ = self.G(c, d, mv, target_h=target_2x[0], target_w=target_2x[1])
            nearest_2x = F.interpolate(c, target_2x, mode="nearest")
            bilinear_2x = F.interpolate(c, target_2x, mode="bilinear", align_corners=False)
            images["2x/nearest"] = wandb.Image(to_np(nearest_2x), caption="Nearest 2x")
            images["2x/bilinear"] = wandb.Image(to_np(bilinear_2x), caption="Bilinear 2x")
            images["2x/prism"] = wandb.Image(to_np(out_2x), caption=f"Prism 2x (ep{self.epoch})")

            # 3x upscale
            target_3x = (rH * 3, rW * 3)
            with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp):
                out_3x, _ = self.G(c, d, mv, target_h=target_3x[0], target_w=target_3x[1])
            nearest_3x = F.interpolate(c, target_3x, mode="nearest")
            bilinear_3x = F.interpolate(c, target_3x, mode="bilinear", align_corners=False)
            gt_3x = F.interpolate(gt, target_3x, mode="bilinear", align_corners=False)
            images["3x/nearest"] = wandb.Image(to_np(nearest_3x), caption="Nearest 3x")
            images["3x/bilinear"] = wandb.Image(to_np(bilinear_3x), caption="Bilinear 3x")
            images["3x/prism"] = wandb.Image(to_np(out_3x), caption=f"Prism 3x (ep{self.epoch})")
            images["3x/ground_truth"] = wandb.Image(to_np(gt_3x), caption="Ground Truth 3x")

            # Side-by-side comparisons
            gt_2x_resized = F.interpolate(gt, target_2x, mode="bilinear", align_corners=False)
            comp_2x = torch.cat([nearest_2x, bilinear_2x, out_2x, gt_2x_resized], dim=3)
            images["compare/2x_nearest_bilinear_prism_gt"] = wandb.Image(
                to_np(comp_2x), caption="Nearest | Bilinear | Prism | GT (2x)")

            comp_3x = torch.cat([nearest_3x, bilinear_3x, out_3x, gt_3x], dim=3)
            images["compare/3x_nearest_bilinear_prism_gt"] = wandb.Image(
                to_np(comp_3x), caption="Nearest | Bilinear | Prism | GT (3x)")

            wandb.log(images, step=self.epoch)

        except Exception as e:
            print(f"  [wandb image log failed: {e}]")
        finally:
            self.G.train()

    def save(self, path: Path):
        from dataclasses import asdict
        path.mkdir(parents=True, exist_ok=True)
        G_inner = self.G.module if hasattr(self.G, 'module') else self.G
        torch.save({
            "generator": G_inner.state_dict(),
            "discriminator": self.D.state_dict(),
            "opt_G": self.opt_G.state_dict(),
            "opt_D": self.opt_D.state_dict(),
            "epoch": self.epoch,
            "model_config": asdict(G_inner.cfg),
        }, path / f"checkpoint_ep{self.epoch}.pth")
        torch.save(G_inner.state_dict(), path / "prism_generator_latest.pth")
        print(f"Saved epoch {self.epoch}")

    @torch.no_grad()
    def _save_test_images(self, loader, epoch_num):
        """Save comparison images for browser viewing."""
        import numpy as np
        from PIL import Image as PILImage

        output_dir = Path(os.environ.get("PRISM_TEST_OUTPUT", "test_outputs"))
        output_dir.mkdir(exist_ok=True)

        self.G.eval()
        try:
            # Grab 5 samples from separate sequences (skip batches between)
            samples = []
            skip = 0
            for seq in loader:
                skip += 1
                if skip % 20 == 0:  # every 20th batch = different sequence
                    samples.append(seq[0])
                if len(samples) >= 5:
                    break

            def to_np(t):
                return (t[0].cpu().float().clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8)

            for i, frame in enumerate(samples):
                c = frame["color"][:1].to(self.device).float()
                d = frame["depth"][:1].to(self.device).float()
                mv = frame["motion_vectors"][:1].to(self.device).float()
                gt = frame["ground_truth"][:1].to(self.device).float()
                _, _, rH, rW = c.shape

                for scale, label in [(2, "2x"), (3, "3x")]:
                    tH, tW = rH * scale, rW * scale
                    with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp):
                        out, _ = self.G(c, d, mv, target_h=tH, target_w=tW)
                    gt_r = F.interpolate(gt, (tH, tW), mode="bilinear", align_corners=False)
                    nearest = F.interpolate(c, (tH, tW), mode="nearest")
                    bilinear = F.interpolate(c, (tH, tW), mode="bilinear", align_corners=False)
                    comp = torch.cat([nearest, bilinear, out, gt_r], dim=3)
                    img = PILImage.fromarray(to_np(comp))
                    img.save(output_dir / f"ep{epoch_num:03d}_sample{i}_{label}.png")

            # Update index.html
            html = f'<html><head><meta http-equiv="refresh" content="60"><style>body{{background:#111;color:#fff;font-family:monospace}}img{{max-width:100%;margin:5px 0}}</style></head><body>'
            html += f'<h1>Prism Training — Epoch {epoch_num}</h1><p>Nearest | Bilinear | Prism | Ground Truth</p>'
            for f in sorted(output_dir.glob("ep*.png"), reverse=True):
                html += f'<h3>{f.name}</h3><img src="{f.name}"><br>'
            html += "</body></html>"
            (output_dir / "index.html").write_text(html)
            print(f"  [test images saved to {output_dir}]")
        except Exception as e:
            print(f"  [test image save failed: {e}]")
        finally:
            self.G.train()

    def load(self, path: Path):
        ckpt = torch.load(path, map_location=self.device, weights_only=True)
        # Try strict load first, fall back to partial
        try:
            self.G.load_state_dict(ckpt["generator"])
            print("  Generator: loaded all weights (exact match)")
        except RuntimeError:
            # Architecture changed — load what we can
            self._load_partial(self.G, ckpt["generator"], "Generator")
        try:
            self.D.load_state_dict(ckpt["discriminator"])
            self.opt_D.load_state_dict(ckpt["opt_D"])
        except RuntimeError:
            print("  D weights incompatible (likely spectral norm change), using fresh D")
        try:
            self.opt_G.load_state_dict(ckpt["opt_G"])
        except Exception:
            print("  G optimizer state incompatible, using fresh optimizer")
        self.epoch = ckpt["epoch"]
        print(f"Resumed from epoch {self.epoch}")

    def load_conv_weights(self, path: Path):
        """Load encoder/decoder/output conv weights from a pretrained checkpoint.
        Transformer and MoE layers are initialized fresh. This lets you upgrade
        from a dense model to MoE while keeping trained conv layers."""
        ckpt = torch.load(path, map_location=self.device, weights_only=True)
        src = ckpt["generator"] if "generator" in ckpt else ckpt

        # Match conv layers: input_conv, enc*, dec*, to_rgb*, temporal
        conv_prefixes = ("input_conv.", "enc1.", "enc2.", "enc3.",
                         "dec1.", "dec2.", "dec3.",
                         "to_rgb_2x.", "to_rgb_3x.", "temporal.")
        loaded, skipped = 0, 0
        model_state = self.G.state_dict()
        for key, val in src.items():
            if any(key.startswith(p) for p in conv_prefixes):
                if key in model_state and model_state[key].shape == val.shape:
                    model_state[key] = val
                    loaded += 1
                else:
                    skipped += 1
                    if key in model_state:
                        print(f"  Shape mismatch {key}: ckpt={val.shape} vs model={model_state[key].shape}")
        self.G.load_state_dict(model_state)
        print(f"  Conv weights: loaded {loaded}, skipped {skipped}")
        print(f"  Transformer/MoE layers: initialized fresh")

    @staticmethod
    def _load_partial(model, state_dict, name="Model"):
        """Load matching weights, skip mismatched ones."""
        model_state = model.state_dict()
        loaded, skipped = 0, 0
        for key, val in state_dict.items():
            if key in model_state and model_state[key].shape == val.shape:
                model_state[key] = val
                loaded += 1
            else:
                skipped += 1
        model.load_state_dict(model_state)
        print(f"  {name}: loaded {loaded}/{loaded + skipped} weights ({skipped} skipped/mismatched)")



# ============================================================================
# Main
# ============================================================================

def main():
    # GPU optimizations: TF32 for convolutions (when available), cudnn autotuning
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.backends.cudnn.benchmark = True

    parser = argparse.ArgumentParser(description="Train Prism G-buffer decoder")
    parser.add_argument("--data", type=Path, default=Path("data/dataset"))
    parser.add_argument("--output", type=Path, default=Path("checkpoints"))
    parser.add_argument("--resume", type=Path, default=None)
    parser.add_argument("--model", choices=list(V2_PRESETS.keys()), default="balanced",
                        help="Model preset (dense: fast/balanced/quality, MoE: moe-16/moe-32/moe-64)")
    # Override individual model config fields (applied on top of preset)
    parser.add_argument("--n-blocks", type=int, default=None, help="Override transformer block count")
    parser.add_argument("--n-heads", type=int, default=None, help="Override attention head count")
    parser.add_argument("--t-dim", type=int, default=None, help="Override transformer dim (adds projection if != bottleneck)")
    parser.add_argument("--n-experts", type=int, default=None, help="Override expert count (0=dense)")
    parser.add_argument("--expert-hidden", type=int, default=None, help="Override per-expert FFN hidden dim")
    parser.add_argument("--top-k", type=int, default=None, help="Override MoE top-k routing")
    parser.add_argument("--window-size", type=int, default=None, help="Override attention window size (0=global)")
    parser.add_argument("--optimizer", choices=["adamw", "apollo-mini", "apollo"], default="apollo-mini")
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--batch", type=int, default=4)
    parser.add_argument("--crop", type=int, default=128)
    parser.add_argument("--seq-len", type=int, default=1)
    parser.add_argument("--progressive", action="store_true",
                        help="Progressive training: start small crops, grow over epochs")
    parser.add_argument("--lr-g", type=float, default=1e-4)
    parser.add_argument("--lr-d", type=float, default=4e-4)
    parser.add_argument("--amp", action="store_true", help="FP16/BF16 mixed precision")
    parser.add_argument("--fp8", action="store_true", help="FP8 training (RTX 4060+/5060 Ti)")
    parser.add_argument("--grad-accum", type=int, default=1)
    parser.add_argument("--adv-warmup", type=int, default=10)
    parser.add_argument("--save-every", type=int, default=10)
    parser.add_argument("--device", default="cuda:1", help="cuda:0=4060, cuda:1=5060Ti")
    parser.add_argument("--multi-gpu", action="store_true", help="Use DataParallel across all GPUs")
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--wandb", action="store_true")
    parser.add_argument("--d-every", type=int, default=1, help="Train D every N steps (1=every step)")
    parser.add_argument("--resume-conv", type=Path, default=None,
                        help="Load conv encoder/decoder weights from a pretrained checkpoint (for MoE upgrade)")
    args = parser.parse_args()

    # Progressive training schedule: crop size grows over epochs
    if args.progressive:
        schedule = [
            (64, int(args.epochs * 0.4)),    # 40% of epochs at 64x64 (fast)
            (128, int(args.epochs * 0.35)),   # 35% at 128x128 (medium)
            (192, int(args.epochs * 0.15)),   # 15% at 192x192 (detailed)
            (256, int(args.epochs * 0.10)),   # 10% at 256x256 (full context)
        ]
        print(f"Progressive training schedule:")
        for crop, epochs in schedule:
            print(f"  {crop}x{crop} for {epochs} epochs")
    else:
        schedule = [(args.crop, args.epochs)]

    total_epochs_done = 0

    for phase, (crop_size, phase_epochs) in enumerate(schedule):
        print(f"\n{'='*60}")
        print(f"Phase {phase}: crop={crop_size}x{crop_size}, {phase_epochs} epochs")
        print(f"{'='*60}")

        dataset = PrismDataset(args.data, crop_size=crop_size, seq_len=args.seq_len)
        loader = DataLoader(
            dataset, batch_size=args.batch, shuffle=True,
            num_workers=args.workers, pin_memory=True,
            drop_last=True, collate_fn=collate_sequences,
        )

        if phase == 0:
            # Build config from preset + CLI overrides
            from dataclasses import replace
            cfg = V2_PRESETS[args.model]
            overrides = {}
            if args.n_blocks is not None: overrides["n_transformer_blocks"] = args.n_blocks
            if args.n_heads is not None: overrides["n_heads"] = args.n_heads
            if args.t_dim is not None: overrides["transformer_dim"] = args.t_dim
            if args.n_experts is not None: overrides["n_experts"] = args.n_experts
            if args.expert_hidden is not None: overrides["expert_ffn_hidden"] = args.expert_hidden
            if args.top_k is not None: overrides["top_k"] = args.top_k
            if args.window_size is not None: overrides["window_size"] = args.window_size
            if overrides:
                cfg = replace(cfg, **overrides)
                print(f"Config overrides: {overrides}")

            trainer = Trainer(
                model_name=args.model, model_config=cfg,
                optimizer_name=args.optimizer,
                lr_g=args.lr_g, lr_d=args.lr_d,
                device=args.device, use_amp=args.amp, use_fp8=args.fp8,
                grad_accum=args.grad_accum, use_wandb=args.wandb,
            )

            # Multi-GPU with DataParallel
            if args.multi_gpu and torch.cuda.device_count() > 1:
                print(f"Using {torch.cuda.device_count()} GPUs with DataParallel")
                trainer.G = torch.nn.DataParallel(trainer.G)
                trainer.D = torch.nn.DataParallel(trainer.D)

            if args.resume:
                trainer.load(args.resume)
            elif args.resume_conv:
                trainer.load_conv_weights(args.resume_conv)

        # Grab a fixed sample for wandb image logging
        if args.wandb and phase == 0:
            trainer.set_fixed_sample(loader)

        print(f"\nTraining: {phase_epochs} epochs | batch={args.batch} | seq={args.seq_len} | "
              f"crop={crop_size} | {'FP8' if args.fp8 else 'AMP' if args.amp else 'FP32'} | "
              f"{args.optimizer} | d_every={args.d_every}\n")

        try:
            for epoch in range(phase_epochs):
                total_epoch = total_epochs_done + epoch
                adv_w = min(0.03, 0.03 * min(1.0, total_epoch / max(args.adv_warmup, 1)))

                t0 = time.time()
                m = trainer.train_epoch(loader, adv_weight=adv_w, d_every=args.d_every)
                dt = time.time() - t0

                moe_str = f" moe={m['moe']:.4f}" if m['moe'] > 0 else ""
                print(f"  [{total_epoch+1}] L1={m['l1']:.4f} adv={m['adv']:.4f}(w={adv_w:.2f}) "
                      f"temp={m['temp']:.4f} D={m['d']:.4f}{moe_str} [{dt:.1f}s]")

                if (total_epoch + 1) % args.save_every == 0:
                    trainer.save(args.output)

                # Generate test comparison images every epoch
                trainer._save_test_images(loader, total_epoch + 1)
        except KeyboardInterrupt:
            print("\nInterrupted")
            break

        total_epochs_done += phase_epochs

    trainer.save(args.output)
    print("Done!")


if __name__ == "__main__":
    main()