#!/usr/bin/env python
"""Multi-view SD3 dynamics model.

    {I_{t-km..t}^1..V}, a[t:t+8], s_t  ->  {I_{t+8}^1..V}

Conditioning follows Ctrl-World (arXiv:2510.10125): context frames and the noisy target live in
ONE token sequence so they interact through full self-attention, instead of the target attending to
frame features through the text/cross-attention pathway. Here that is done by tiling latents into a
single canvas -- rows are timesteps (history..current, then the target), columns are camera views --
and letting the unmodified MMDiT attend over it. Loss and sampling touch the target row only.

An earlier design fed frames in as `encoder_hidden_states` only; it trained to a validation PSNR of
7.9 dB against 19.0 dB for simply copying the current frame, i.e. worse than doing nothing.

Env: policy_learning/lerobot/.venv (lerobot + diffusers + torchmetrics + wandb + peft).

    P=policy_learning/lerobot/.venv/bin/python
    $P sd3_dynamics.py --data <success> <failure> --inspect
    PYTHONUNBUFFERED=1 $P sd3_dynamics.py --data <success> <failure> --steps 20000
"""

import argparse
import json
import math
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel
from lerobot.datasets.lerobot_dataset import LeRobotDataset

HORIZON = 8  # action chunk length AND image lead time, per spec


# ── data ──────────────────────────────────────────────────────────────────────


def dataset_info(spec) -> tuple[str, Path | None, dict]:
    """Accepts a local dataset root or an HF repo id -> (repo_id, root|None, info.json)."""
    root = Path(spec)
    if (root / "meta" / "info.json").exists():
        return f"local/{root.name}", root, json.loads((root / "meta" / "info.json").read_text())
    from huggingface_hub import hf_hub_download

    return spec, None, json.loads(Path(hf_hub_download(spec, "meta/info.json", repo_type="dataset")).read_text())


def dataset_features(spec) -> dict:
    return dataset_info(spec)[2]["features"]


def dataset_revision(repo_id: str, root, info: dict) -> str | None:
    """LeRobotDataset resolves codebase_version as a git tag; some repos only have `main`."""
    if root is not None:
        return None
    from huggingface_hub import HfApi

    tag = info.get("codebase_version")
    try:
        tags = [t.name for t in HfApi().list_repo_refs(repo_id, repo_type="dataset").tags]
    except Exception:
        return None
    return None if tag in tags else "main"


def discover_cameras(features: dict) -> list[str]:
    """Every RGB camera feature, deterministically ordered. No hard-coded names."""
    cams = sorted(k for k, v in features.items() if v.get("dtype") in ("video", "image"))
    if not cams:
        raise ValueError(f"no image/video features found among {sorted(features)}")
    return cams


def object_state_keys(features: dict) -> list[str]:
    """Simulator/object-state features. Present in MimicGen data; must never reach the model."""
    return sorted(k for k in features if "object" in k.lower() or "sim_state" in k.lower())


def assert_object_free(features: dict, state_key: str, state_dim: int) -> str:
    """MimicGen packs object state INSIDE observation.state: [robot 9D | object ND].

    Slicing [:9] drops it, but a wrong --state-dim would feed object state to the model and
    violate the spec silently. Verified for coffee_d0 (66=9+57) and square_d1 (26=9+17).
    """
    assert state_key not in object_state_keys(features), f"{state_key} is itself object state"
    full = features[state_key]["shape"][0]
    obj = features.get("observation.object", {}).get("shape", [0])[0]
    if obj and state_dim + obj == full:
        assert state_dim <= full - obj, (
            f"--state-dim {state_dim} reaches into the object block of {state_key} "
            f"({full}D = {full - obj}D robot + {obj}D object)"
        )
        return f"{state_key}[:{state_dim}] of {full}D = robot only ({obj}D object tail dropped)"
    if full > state_dim:
        return (f"{state_key}[:{state_dim}] of {full}D — dropping {full - state_dim} tail dims that "
                f"this repo has no observation.object to cross-check against")
    return f"{state_key}[:{state_dim}] of {full}D (full state, nothing dropped)"


class MultiViewWindows(torch.utils.data.Dataset):
    """(history+current frames, future frames, 8-action chunk, state) windows from one dataset."""

    def __init__(
        self,
        spec,
        cameras,
        state_key,
        state_dim,
        image_size,
        episodes=None,
        horizon=HORIZON,
        history=1,
        stride=4,
    ):
        self.cameras, self.state_key, self.state_dim = cameras, state_key, state_dim
        self.image_size, self.horizon = image_size, horizon
        self.history, self.stride = history, stride
        repo_id, root, info = dataset_info(spec)
        fps = info["fps"]
        # oldest history frame first, then the current frame, then the target
        self.offsets = [-(history - 1 - i) * stride for i in range(history)] + [horizon]
        dt = {"action": [i / fps for i in range(horizon)]}
        dt |= {c: [o / fps for o in self.offsets] for c in cameras}
        rev = dataset_revision(repo_id, root, info)
        self.ds = LeRobotDataset(repo_id, root=root, delta_timestamps=dt, revision=rev)

        # episode boundaries: every offset, past and future, must land in the same episode
        ep = np.asarray(self.ds.hf_dataset.select_columns("episode_index")["episode_index"])
        back, fwd = (history - 1) * stride, horizon
        idx = np.arange(back, len(ep) - fwd)
        valid = idx[(ep[idx - back] == ep[idx]) & (ep[idx + fwd] == ep[idx])]
        if episodes is not None:
            valid = valid[np.isin(ep[valid], list(episodes))]
        self.indices = valid

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

    def __getitem__(self, i):
        item = self.ds[int(self.indices[i])]
        frames = [
            # antialiased resize overshoots [0,1] slightly; LPIPS and the VAE both want it clamped
            F.interpolate(item[c], self.image_size, mode="bilinear", antialias=True, align_corners=False)
            .clamp(0, 1)
            for c in self.cameras
        ]  # each (history+1, 3, h, w)
        stacked = torch.stack(frames, 1)  # (history+1, V, 3, h, w)
        return {
            "context": stacked[:-1],  # (history, V, 3, h, w), oldest first
            "future": stacked[-1],  # (V, 3, h, w)
            "action": item["action"],
            # first state_dim entries only: eef_pos(3) + eef_quat(4) + gripper_qpos(2) on MimicGen.
            # Object/sim state lives in separate features and is never read.
            "state": item[self.state_key][: self.state_dim],
        }


def build_datasets(specs, image_size, state_key, state_dim, val_episodes, history=1, stride=4):
    """Concat all datasets; hold out the last `val_episodes` episodes of each for validation."""
    cameras = None
    train, val = [], []
    for spec in specs:
        _, _, info = dataset_info(spec)
        cams = discover_cameras(info["features"])
        if cameras is None:
            cameras = cams
        elif cams != cameras:
            raise ValueError(f"camera mismatch: {spec} has {cams}, expected {cameras}")
        n_ep = info["total_episodes"]
        holdout = set(range(max(0, n_ep - val_episodes), n_ep))
        args = (cams, state_key, state_dim, image_size)
        kw = dict(history=history, stride=stride)
        train.append(MultiViewWindows(spec, *args, set(range(n_ep)) - holdout, **kw))
        if holdout:
            val.append(MultiViewWindows(spec, *args, holdout, **kw))
    cat = torch.utils.data.ConcatDataset
    return cat(train), (cat(val) if val else None), cameras


# ── model ─────────────────────────────────────────────────────────────────────


class SD3MultiViewDynamics(nn.Module):
    """SD3 MMDiT over a canvas of [history..current | target] frames × camera views.

    Camera identity is positional: view v always occupies column v of the canvas, in the sorted
    camera order, so it is consistent between training and evaluation. Because every view sits in
    one sequence, the predicted views also attend to each other -- they are denoised jointly, not
    in independent batch rows.
    """

    def __init__(
        self,
        vae,
        transformer,
        n_views,
        action_dim,
        state_dim,
        ctx_noise=0.0,
        aux_weight=0.0,
        aux_layer=None,
        action_dropout=0.2,
        horizon=HORIZON,
    ):
        super().__init__()
        self.vae = vae.eval().requires_grad_(False)
        self.transformer = transformer
        self.n_views, self.ctx_noise = n_views, ctx_noise
        self.action_dim, self.horizon = action_dim, horizon
        self.aux_weight, self.action_dropout = aux_weight, action_dropout
        d = transformer.config.joint_attention_dim
        p = transformer.config.pooled_projection_dim
        # actions and state are not spatial, so they stay on the cross-attention pathway
        self.act_tok = nn.Linear(action_dim, d)
        self.state_tok = nn.Linear(state_dim, d)
        self.pool = nn.Linear(d, p)
        # zero-init: step 0 is the pretrained MMDiT on empty context, not on random noise
        for m in (self.act_tok, self.state_tok, self.pool):
            nn.init.zeros_(m.weight)
            nn.init.zeros_(m.bias)

        self._aux_feat = None
        if aux_weight:
            inner = transformer.config.num_attention_heads * transformer.config.attention_head_dim
            self.aux_head = nn.Sequential(
                nn.LayerNorm(inner), nn.Linear(inner, inner), nn.GELU(), nn.Linear(inner, horizon * action_dim)
            )
            blocks = transformer.transformer_blocks
            self.aux_layer = len(blocks) // 2 if aux_layer is None else aux_layer
            # JointTransformerBlock returns (encoder_hidden_states, hidden_states); grab the image stream
            blocks[self.aux_layer].register_forward_hook(lambda m, i, o: setattr(self, "_aux_feat", o[-1]))

    @classmethod
    def from_pretrained(cls, sd3_id, n_views, action_dim, state_dim, dtype=torch.bfloat16, **kw):
        vae = AutoencoderKL.from_pretrained(sd3_id, subfolder="vae", torch_dtype=dtype)
        tr = SD3Transformer2DModel.from_pretrained(sd3_id, subfolder="transformer", torch_dtype=dtype)
        return cls(vae, tr, n_views, action_dim, state_dim, **kw)

    # -- latents --
    def encode(self, imgs, sample=False):
        """(N,3,H,W) in [0,1] -> scaled latents (N,C,h,w)."""
        dist = self.vae.encode(imgs.to(self.vae.dtype) * 2 - 1).latent_dist
        z = dist.sample() if sample else dist.mode()
        return (z - self.vae.config.shift_factor) * self.vae.config.scaling_factor

    def decode(self, z):
        z = z.to(self.vae.dtype) / self.vae.config.scaling_factor + self.vae.config.shift_factor
        return (self.vae.decode(z).sample / 2 + 0.5).clamp(0, 1)

    def encode_grid(self, imgs, **kw):
        """(B,R,V,3,H,W) -> (B,R,V,C,h,w)."""
        b, r, v = imgs.shape[:3]
        z = self.encode(imgs.flatten(0, 2), **kw)
        return z.reshape(b, r, v, *z.shape[1:])

    # -- canvas: rows are timesteps, columns are views --
    @staticmethod
    def to_canvas(z):
        """(B,R,V,C,h,w) -> (B,C,R*h,V*w)."""
        b, r, v, c, h, w = z.shape
        return z.permute(0, 3, 1, 4, 2, 5).reshape(b, c, r * h, v * w)

    @staticmethod
    def from_canvas(canvas, v):
        """(B,C,h,V*w) single row -> (B,V,C,h,w)."""
        b, c, h, vw = canvas.shape
        return canvas.reshape(b, c, h, v, vw // v).permute(0, 3, 1, 2, 4)

    def conditioning(self, action, state, drop=None):
        """Action chunk and robot state as cross-attention tokens -> (ctx, pooled).

        `drop` is a (B,) bool mask zeroing the action tokens, so the auxiliary head has to recover
        the action from the frames instead of reading it straight off the conditioning.
        """
        dt = self.act_tok.weight.dtype
        act = self.act_tok(action.to(dt))
        if drop is not None:
            act = act * (~drop).to(dt).view(-1, 1, 1)
        ctx = torch.cat([act, self.state_tok(state.to(dt))[:, None]], 1)
        return ctx, self.pool(ctx.mean(1))

    def denoise(self, canvas, timestep, ctx, pooled):
        dt = self.transformer.dtype  # conditioning heads stay fp32; cast at the MMDiT boundary
        return self.transformer(
            hidden_states=canvas.to(dt),
            encoder_hidden_states=ctx.to(dt),
            pooled_projections=pooled.to(dt),
            timestep=timestep,
            return_dict=False,
        )[0]

    def _context_rows(self, context):
        """Encoded history+current rows, optionally noise-perturbed for robustness (Ctrl-World §3)."""
        z = self.encode_grid(context)
        if self.ctx_noise:
            z = z + self.ctx_noise * torch.randn_like(z)
        return z

    # -- training / sampling --
    def loss(self, batch):
        """Rectified-flow matching on the target row of the canvas."""
        z_ctx = self._context_rows(batch["context"])
        z_tgt = self.encode_grid(batch["future"][:, None], sample=True)  # (B,1,V,C,h,w)
        b, _, v, c, h, w = z_tgt.shape
        t = torch.sigmoid(torch.randn(b, device=z_tgt.device, dtype=z_tgt.dtype))  # logit-normal
        noise = torch.randn_like(z_tgt)
        tt = t.view(-1, 1, 1, 1, 1, 1)  # (B,R,V,C,h,w) is 6-D; a 5-D view broadcasts into the row axis
        z_t = (1 - tt) * z_tgt + tt * noise
        canvas = self.to_canvas(torch.cat([z_ctx, z_t], 1))
        drop = torch.rand(b, device=canvas.device) < self.action_dropout if self.aux_weight else None
        ctx, pooled = self.conditioning(batch["action"], batch["state"], drop)
        pred = self.denoise(canvas, t * 1000, ctx, pooled)[:, :, -h:, :]  # target row only
        flow = F.mse_loss(pred.float(), self.to_canvas(noise - z_tgt).float())
        if not self.aux_weight or not drop.any():
            return flow
        # inverse dynamics on the action-dropped rows only: with the action still conditioned in,
        # the head would just read it back off the context tokens and learn nothing.
        feat = self._aux_feat[drop].float().mean(1)
        a_pred = self.aux_head(feat).view(-1, self.horizon, self.action_dim)
        return flow + self.aux_weight * F.mse_loss(a_pred, batch["action"][drop].float())

    @torch.no_grad()
    def predict(self, batch, scheduler, steps=20):
        """-> predicted future frames (B,V,3,H,W) in [0,1]."""
        z_ctx = self._context_rows(batch["context"])
        b, _, v, c, h, w = z_ctx.shape
        ctx, pooled = self.conditioning(batch["action"], batch["state"])
        z = torch.randn((b, 1, v, c, h, w), device=z_ctx.device, dtype=z_ctx.dtype)
        scheduler.set_timesteps(steps, device=z_ctx.device)
        for t in scheduler.timesteps:
            canvas = self.to_canvas(torch.cat([z_ctx, z], 1))
            vel = self.denoise(canvas, t.expand(b), ctx, pooled)[:, :, -h:, :]
            row = scheduler.step(vel.float(), t, self.to_canvas(z).float(), return_dict=False)[0]
            z = self.from_canvas(row.to(z_ctx.dtype), v)[:, None]
        imgs = self.decode(z.flatten(0, 2))
        return imgs.float().reshape(b, v, *imgs.shape[1:])


# ── evaluation ────────────────────────────────────────────────────────────────


def per_camera_metrics(pred, gt, cameras, lpips=None):
    """pred/gt: (B,V,3,H,W) in [0,1]. -> {'<cam>/mse': ..., 'mean/mse': ...}"""
    from torchmetrics.functional import structural_similarity_index_measure as ssim_fn

    out = {}
    for i, cam in enumerate(cameras):
        p, g = pred[:, i], gt[:, i]
        mse = F.mse_loss(p, g).item()
        out[f"{cam}/mse"] = mse
        out[f"{cam}/psnr"] = 10 * math.log10(1.0 / max(mse, 1e-12))  # data range 1.0
        out[f"{cam}/ssim"] = ssim_fn(p, g, data_range=1.0).item()
        if lpips is not None:
            out[f"{cam}/lpips"] = lpips(p, g).item()
    for m in ("mse", "psnr", "ssim", "lpips"):
        vals = [v for k, v in out.items() if k.endswith(f"/{m}")]
        if vals:
            out[f"mean/{m}"] = sum(vals) / len(vals)
    return out


def to_wandb_image(img):
    import wandb

    return wandb.Image((img.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))


def reconstruction_table(current, gt, pred, cameras, start_id=0):
    """One row per example, one Current/GT/Pred column group per camera."""
    import wandb

    cols = ["sample_id"] + [f"{c}/{k}" for c in cameras for k in ("current", "gt", "pred")]
    table = wandb.Table(columns=cols)
    for b in range(current.shape[0]):
        row = [start_id + b]
        for i in range(len(cameras)):
            row += [to_wandb_image(x[b, i]) for x in (current, gt, pred)]
        table.add_data(*row)
    return table


# ── dataset safety checks (spec §15 / §22) ────────────────────────────────────


def inspect(specs, state_key, state_dim, image_size, history=1, stride=4, horizon=HORIZON):
    ref_cams = None
    for spec in specs:
        _, _, info = dataset_info(spec)
        feats = info["features"]
        cams = discover_cameras(feats)
        print(f"\n=== {spec}")
        print(f"  fps={info['fps']} episodes={info['total_episodes']} frames={info['total_frames']}")
        print(f"  splits={info.get('splits')}")
        print(f"  all keys: {sorted(feats)}")
        print(f"  cameras (deterministic order): {cams}")
        for c in cams:
            f = feats[c]
            print(f"    {c}: shape={f['shape']} names={f.get('names')} dtype={f['dtype']}")
        print(f"  action dim: {feats['action']['shape']}")
        full = feats[state_key]["shape"][0]
        print(f"  state: {state_key}[:{state_dim}] of {full}D")
        assert full >= state_dim, f"{state_key} is only {full}D, need {state_dim}"
        objs = object_state_keys(feats)
        print(f"  object/sim-state keys present but NOT read by the model: {objs}")
        assert state_key not in objs, f"{state_key} is an object-state feature"
        if ref_cams is None:
            ref_cams = cams
        assert cams == ref_cams, f"camera mismatch vs first dataset: {cams} != {ref_cams}"

        ds = MultiViewWindows(spec, cams, state_key, state_dim, image_size, history=history, stride=stride)
        item = ds[0]
        print(f"  frame offsets (history..current, target): {ds.offsets}")
        print(f"  valid windows: {len(ds)} / {info['total_frames']} frames")
        for c, ctx, fut in zip(cams, item["context"].transpose(0, 1), item["future"]):
            print(f"    {c}: context {tuple(ctx.shape)}  future {tuple(fut.shape)}")
        print(f"    action {tuple(item['action'].shape)}  state {tuple(item['state'].shape)}")

        i0 = int(ds.indices[0])
        raw = ds.ds[i0]
        fi = int(raw["frame_index"])
        nxt = ds.ds.hf_dataset[i0 + horizon]
        print(f"    frame_index {fi} -> {int(nxt['frame_index'])} (must differ by exactly {horizon})")
        assert int(nxt["frame_index"]) - fi == horizon
        assert int(nxt["episode_index"]) == int(raw["episode_index"])
        for c in cams:
            assert not raw[f"{c}_is_pad"].any(), f"{c} padded at a supposedly valid window"
        assert not raw["action_is_pad"].any()
    print(f"\nOK: {len(specs)} dataset(s), cameras={ref_cams}")


# ── training ──────────────────────────────────────────────────────────────────


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument(
        "--data", nargs="+", required=True, help="local dataset roots or HF repo ids (success, failure, ...)"
    )
    ap.add_argument("--inspect", action="store_true", help="run dataset safety checks and exit")
    ap.add_argument("--sd3", default="stabilityai/stable-diffusion-3.5-medium")
    ap.add_argument("--state-key", default="observation.state")
    ap.add_argument("--state-dim", type=int, default=9, help="use the first N dims of --state-key (spec: 9)")
    ap.add_argument("--image-size", type=int, nargs=2, default=(224, 224))
    ap.add_argument("--history", type=int, default=1, help="context frames incl. current (1 = current only)")
    ap.add_argument("--history-stride", type=int, default=4, help="frames between history frames")
    ap.add_argument("--ctx-noise", type=float, default=0.0, help="noise on context latents (Ctrl-World)")
    ap.add_argument("--aux-weight", type=float, default=0.0, help="auxiliary action-reconstruction loss")
    ap.add_argument("--aux-layer", type=int, default=None, help="MMDiT block to read (default: middle)")
    ap.add_argument("--action-dropout", type=float, default=0.2, help="action-drop rate the aux loss uses")
    ap.add_argument("--batch-size", type=int, default=16)
    ap.add_argument("--grad-accum", type=int, default=1)
    ap.add_argument("--steps", type=int, default=20000)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--lora-rank", type=int, default=32, help="0 = full finetune of the MMDiT")
    ap.add_argument("--log-every", type=int, default=50)
    ap.add_argument(
        "--no-grad-ckpt",
        action="store_true",
        help="28%% faster per sample but far more activation memory; needs a small batch",
    )
    ap.add_argument("--val-episodes", type=int, default=20, help="held-out episodes per dataset")
    ap.add_argument("--val-every", type=int, default=1000)
    ap.add_argument("--val-batches", type=int, default=8)
    ap.add_argument("--log-images", type=int, default=4, help="examples logged to the W&B table")
    ap.add_argument("--sample-steps", type=int, default=20)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--out", default=str(Path(__file__).parent / "outputs"))
    ap.add_argument("--wandb-project", default="sd3-dynamics", help="empty string disables W&B")
    ap.add_argument("--run-name", default=None, help="W&B run name (default: auto-generated)")
    args = ap.parse_args()

    image_size = tuple(args.image_size)
    if args.inspect:
        inspect(args.data, args.state_key, args.state_dim, image_size, args.history, args.history_stride)
        return

    torch.manual_seed(0)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    train_ds, val_ds, cameras = build_datasets(
        args.data, image_size, args.state_key, args.state_dim, args.val_episodes,
        args.history, args.history_stride,
    )
    feats = dataset_features(args.data[0])
    state_dim = args.state_dim
    full = feats[args.state_key]["shape"][0]
    if full < state_dim:
        raise SystemExit(f"{args.state_key} is only {full}D, cannot take the first {state_dim}D")
    print(f"state: {assert_object_free(feats, args.state_key, state_dim)}")
    action_dim = feats["action"]["shape"][0]
    print(f"cameras={cameras} action_dim={action_dim}")
    print(f"excluded from the model: {object_state_keys(feats)}")
    print(f"canvas: {args.history + 1} rows x {len(cameras)} views, history stride {args.history_stride}")
    print(f"train windows={len(train_ds)} val windows={len(val_ds) if val_ds else 0}")

    model = SD3MultiViewDynamics.from_pretrained(
        args.sd3,
        len(cameras),
        action_dim,
        state_dim,
        ctx_noise=args.ctx_noise,
        aux_weight=args.aux_weight,
        aux_layer=args.aux_layer,
        action_dropout=args.action_dropout,
    ).to(dev)
    if not args.no_grad_ckpt:
        model.transformer.enable_gradient_checkpointing()
    if args.lora_rank:
        from diffusers.training_utils import cast_training_params
        from peft import LoraConfig

        model.transformer.requires_grad_(False)
        model.transformer.add_adapter(
            LoraConfig(
                r=args.lora_rank,
                lora_alpha=args.lora_rank,
                init_lora_weights="gaussian",
                target_modules=["to_q", "to_k", "to_v", "to_out.0"],
            )
        )
        cast_training_params(model.transformer, dtype=torch.float32)  # bf16 Adam states diverge
    for m in (model.act_tok, model.state_tok, model.pool):
        m.to(torch.float32)
    if args.aux_weight:
        model.aux_head.to(torch.float32)
        print(f"aux action head on block {model.aux_layer}/{len(model.transformer.transformer_blocks)}")
    trainable = {n for n, p in model.named_parameters() if p.requires_grad}
    params = [p for p in model.parameters() if p.requires_grad]
    print(f"trainable params: {sum(p.numel() for p in params) / 1e6:.1f}M")
    opt = torch.optim.AdamW(params, lr=args.lr)
    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.sd3, subfolder="scheduler")

    lpips = None
    if val_ds is not None:
        from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity

        lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).to(dev)

    run = None
    if args.wandb_project:
        import wandb

        run = wandb.init(
            project=args.wandb_project, name=args.run_name, config=vars(args) | {"cameras": cameras}
        )

    def loader(ds, bs, shuffle):
        return torch.utils.data.DataLoader(
            ds, batch_size=bs, shuffle=shuffle, num_workers=args.workers, drop_last=True, pin_memory=True
        )

    def to_dev(batch):
        return {k: v.to(dev, non_blocking=True) for k, v in batch.items()}

    def log(d, step):
        print(f"step {step}: " + " ".join(f"{k}={v:.4f}" for k, v in d.items() if isinstance(v, float)))
        if run:
            run.log(d, step=step)

    @torch.no_grad()
    def validate(step):
        model.eval()
        sums, n = {}, 0
        for i, batch in enumerate(loader(val_ds, args.batch_size, False)):
            if i >= args.val_batches:
                break
            batch = to_dev(batch)
            pred = model.predict(batch, scheduler, args.sample_steps)
            current = batch["context"][:, -1]  # newest context frame
            for k, v in per_camera_metrics(pred, batch["future"], cameras, lpips).items():
                sums[k] = sums.get(k, 0.0) + v
            # copy-the-current-frame baseline on the same batches: the number to beat
            for k, v in per_camera_metrics(current, batch["future"], cameras, lpips).items():
                sums[f"copy_{k}"] = sums.get(f"copy_{k}", 0.0) + v
            n += 1
            if i == 0 and run:
                k = min(args.log_images, pred.shape[0])
                run.log(
                    {
                        "val/reconstructions": reconstruction_table(
                            current[:k], batch["future"][:k], pred[:k], cameras
                        )
                    },
                    step=step,
                )
        model.train()
        log({f"val/{k}": v / max(n, 1) for k, v in sums.items()}, step)

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    step, running, t0 = 0, 0.0, time.time()
    model.train()
    while step < args.steps:
        for batch in loader(train_ds, args.batch_size, True):
            loss = model.loss(to_dev(batch)) / args.grad_accum
            loss.backward()
            running += loss.item() * args.grad_accum  # undo the accumulation scaling for reporting
            if (step + 1) % args.grad_accum == 0:
                torch.nn.utils.clip_grad_norm_(params, 1.0)
                opt.step()
                opt.zero_grad(set_to_none=True)
            step += 1
            if step % args.log_every == 0:
                peak = torch.cuda.max_memory_allocated() / 1e9 if dev == "cuda" else 0.0
                dt = (time.time() - t0) / args.log_every
                log(
                    {
                        "train/loss": running / args.log_every,
                        "train/peak_gb": peak,
                        "train/s_per_step": dt,
                        "train/samples_per_s": args.batch_size / dt,
                    },
                    step,
                )
                running, t0 = 0.0, time.time()
            if val_ds is not None and step % args.val_every == 0:
                validate(step)
            if step % 5000 == 0 or step == args.steps:
                # trainable tensors only: ~93MB of LoRA + conditioning heads, not 5GB of frozen MMDiT
                torch.save(
                    {k: v for k, v in model.state_dict().items() if k in trainable},
                    out / f"step_{step}.pt",
                )
            if step >= args.steps:
                break


if __name__ == "__main__":
    main()