Add Orbis 2 hierarchical world model demo (app2.py)

Vendors the Orbis 2 inference code into orbis2/ and adds app2.py, a Gradio
demo for the hierarchical L1-L2 model. Orbis 1 (app.py, orbis/) is untouched;
the Space still boots app.py until app_file is switched.

Unlike Orbis 1, the Orbis 2 rollout takes an mp4 directly and samples its own
L1 (high-rate) and L2 (low-rate, further back) context windows, so app2.py
re-encodes the upload to a constant 10 fps instead of extracting frames and
writing a val config.

- orbis2/evaluate/rollout_demo_v2.py: add --num_videos to roll out N futures
from one context as a single minibatch (tiling context/steering before
building condition_kwargs), and save each to fake_images/sequence_XXXX/,
matching the layout app2.py reads.
- orbis2/models/second_stage/fm_model_v2.py: label the rollout tqdm bars so
the app can scrape progress from the child process.
- app2.py: downloads the three models (L1 detail predictor, L2 abstract
predictor, tokenizer) from sud0301/orbis2_test into the HF cache and points
$ORBIS2_MODELS_DIR at the snapshot, which the L1/L2 configs expand to locate
each other. Uses the cache path directly rather than local_dir= to avoid
duplicating ~20 GB on disk.
- .gitignore: exclude the root models/ staging dir (weights + configs live on
the Hub). Anchored so it does not match orbis2/models/ source.
- requirements.txt: decord, h5py, requests for the video-in rollout path.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

.gitignore

@@ -1,3 +1,11 @@
 __pycache__/
 *.pyc
 orbis/imgs/Rollout.png
+
+# Checkpoints are never committed (orbis1's logs_*/checkpoints, and anything else).
+*.ckpt
+
+# Orbis 2 models (L1/L2/tok — ~20 GB of weights plus their configs) live on the Hub
+# at sud0301/orbis2_test; app2.py downloads them at startup. Leading slash anchors
+# this to the repo root so it does NOT match orbis2/models/, which is source code.
+/models/

app2.py

@@ -0,0 +1,460 @@
+"""
+Orbis 2 World Model — interactive rollout demo for Hugging Face Spaces.
+
+This is the Orbis 2 counterpart to app.py (which drives the single-level Orbis 1
+model). Orbis 2 is a hierarchical L1-L2 world model whose rollout script takes an
+mp4 directly, so the flow differs from app.py:
+
+  1. User uploads a short driving clip.
+  2. We re-encode it to a constant CONTEXT_FPS (= the L1 frame rate) so the
+     rollout script's exact integer-multiple frame-rate check passes.
+  3. We call evaluate/rollout_demo_v2.py, which samples the L1 (high-rate) and L2
+     (low-rate, further back) context windows from the tail of that video.
+  4. The context is tiled NUM_VIDEOS times so one minibatch rolls out that many
+     independently sampled futures under a single seed.
+  5. Generated frames (fake_images/sequence_XXXX/*.jpg) are encoded back into mp4s.
+
+Checkpoint / config resolution (see resolve_exp_dir / download_checkpoints):
+  * If ORBIS2_EXP_DIR is set (and contains the config + checkpoint), it is used
+    as-is — this is the path that works today against a local training run.
+  * Otherwise, if ORBIS2_HF_REPO is set, the config + checkpoint are pulled from
+    that Hub repo into a local exp dir. Orbis 2 is not on the Hub yet; set these
+    env vars once the checkpoint is uploaded.
+"""
+
+import os
+import random
+import re
+import subprocess
+import tempfile
+import traceback
+import uuid
+from pathlib import Path
+
+import cv2
+import gradio as gr
+import spaces
+from huggingface_hub import snapshot_download
+
+# ----------------------------------------------------------------------------
+# Workaround for gradio 5.9.1 / gradio_client bug (gradio-app/gradio#11722):
+# get_api_info() walks the app's JSON schema and crashes with
+#   "TypeError: argument of type 'bool' is not iterable"
+# when a schema node is a boolean (e.g. "additionalProperties": true), because
+# get_type()/_json_schema_to_python_type() assume every schema is a dict. The
+# main page route calls api_info() on every load, so without this the endpoint
+# 500s continuously. Make the walker tolerate boolean schemas.
+# ----------------------------------------------------------------------------
+import gradio_client.utils as _gc_utils
+
+_orig_get_type = _gc_utils.get_type
+_orig_json_to_py = _gc_utils._json_schema_to_python_type
+
+
+def _safe_get_type(schema):
+    if isinstance(schema, bool):
+        return "bool"
+    return _orig_get_type(schema)
+
+
+def _safe_json_to_py(schema, defs=None):
+    if isinstance(schema, bool):
+        return "Any"
+    return _orig_json_to_py(schema, defs)
+
+
+_gc_utils.get_type = _safe_get_type
+_gc_utils._json_schema_to_python_type = _safe_json_to_py
+
+# ----------------------------------------------------------------------------
+# Configuration — adjust these to your setup
+# ----------------------------------------------------------------------------
+REPO_DIR = Path(__file__).parent / "orbis2"         # copied silviogalesso/orbis2_release
+
+# HF cache location. Prefer the Space's persistent /data volume so large
+# checkpoints download once and survive restarts; fall back to the default
+# ephemeral cache if /data isn't writable (e.g. persistent storage not enabled).
+def _pick_cache_dir() -> str:
+    for candidate in ("/data/.cache/huggingface",
+                      os.path.expanduser("~/.cache/huggingface")):
+        try:
+            probe = Path(candidate)
+            probe.mkdir(parents=True, exist_ok=True)
+            test = probe / ".write_probe"
+            test.touch()
+            test.unlink()
+            return candidate
+        except OSError:
+            continue
+    return os.path.expanduser("~/.cache/huggingface")
+
+
+HF_CACHE_DIR = os.environ.get("HF_HOME") or _pick_cache_dir()
+os.environ["HF_HOME"] = HF_CACHE_DIR
+HF_HUB_CACHE_DIR = str(Path(HF_CACHE_DIR) / "hub")
+print(f"[startup] HF cache dir: {HF_CACHE_DIR}"
+      f"{'' if HF_CACHE_DIR.startswith('/data') else '  (ephemeral — /data not writable)'}")
+
+# --- Orbis 2 models (configs + checkpoints) ----------------------------------
+# Three models are needed, laid out as <models>/{L1,L2,tok}/{config,checkpoints/last.ckpt}:
+#   L1  — detail predictor (the world model rollout_demo_v2.py loads directly)
+#   L2  — abstract predictor (frozen; referenced by L1's condition_preprocessor)
+#   tok — tokenizer (referenced by BOTH L1's and L2's tokenizer_config)
+#
+# L1's and L2's configs locate their dependencies via $ORBIS2_MODELS_DIR, which we
+# export below so OmegaConf's expandvars resolves them wherever the tree lands.
+CONFIG_NAME = "config_distill.yaml"     # L1's config, relative to EXP_DIR
+CKPT_NAME = "checkpoints/last.ckpt"     # every model stores its weights here
+
+#: (sub-dir, config filename) for each of the three required models.
+MODEL_SPECS = [
+    ("L1", CONFIG_NAME),
+    ("L2", "config.yaml"),
+    ("tok", "config.yaml"),
+]
+
+# Hub repo holding the three models (root contains L1/, L2/, tok/).
+HF_CKPT_REPO = os.environ.get("ORBIS2_HF_REPO", "sud0301/orbis2_test")
+
+
+def _missing_models(root: Path) -> list[str]:
+    """Return the models whose config or checkpoint is absent under `root`."""
+    return [sub for sub, cfg in MODEL_SPECS
+            if not (root / sub / cfg).exists() or not (root / sub / CKPT_NAME).exists()]
+
+
+def resolve_models_dir() -> Path:
+    """Return a directory containing L1/, L2/ and tok/, downloading it if needed.
+
+    A local tree wins if it's complete (a dev checkout, or ORBIS2_MODELS_DIR pointing
+    at one). Otherwise we snapshot the Hub repo and use the *cache* path directly
+    rather than copying into the Space — `local_dir=` would duplicate ~20 GB on disk.
+    The cache lives on /data when persistent storage is enabled, so the download
+    survives restarts.
+    """
+    local = Path(os.environ.get("ORBIS2_MODELS_DIR") or (Path(__file__).parent / "models"))
+    if not _missing_models(local):
+        print(f"[startup] Using local Orbis 2 models: {local}")
+        return local
+
+    print(f"[startup] Fetching Orbis 2 models from {HF_CKPT_REPO} (~20 GB, first run only)…")
+    snapshot = Path(snapshot_download(
+        repo_id=HF_CKPT_REPO,
+        allow_patterns=[f"{sub}/**" for sub, _ in MODEL_SPECS],
+        cache_dir=HF_HUB_CACHE_DIR,
+    ))
+
+    still_missing = _missing_models(snapshot)
+    if still_missing:
+        raise RuntimeError(
+            f"Downloaded {HF_CKPT_REPO} but these models are incomplete: "
+            f"{', '.join(still_missing)}. Each of L1/, L2/, tok/ must contain its "
+            f"config and '{CKPT_NAME}'."
+        )
+    print(f"[startup] Orbis 2 models ready: {snapshot}")
+    return snapshot
+
+
+MODELS_DIR = resolve_models_dir()
+os.environ["ORBIS2_MODELS_DIR"] = str(MODELS_DIR)   # L1/L2 configs expandvars this
+EXP_DIR = MODELS_DIR / "L1"                         # rollout_demo_v2.py --exp_dir
+
+FRAME_H, FRAME_W = 288, 512
+L1_FRAME_RATE = 10          # Hz to sample L1 context/rollout frames (and output fps)
+CONTEXT_FPS = L1_FRAME_RATE # we re-encode uploads to this exact fps for the rollout script
+DEFAULT_GEN_STEPS = 10      # rollout steps; each yields model.num_pred_frames frames
+DEFAULT_STEPS = 30          # flow-matching solver steps (matches the reference command)
+OUTPUT_FPS = L1_FRAME_RATE
+NUM_VIDEOS = 3              # futures rolled out per run, batched into one minibatch
+DEFAULT_ETA = 0.0           # 0 = deterministic ODE; >0 injects noise each solver step
+
+DESCRIPTION = f"""
+# Orbis 2: A Hierarchical World Model for Driving
+
+Sudhanshu Mittal\\*, Arian Mousakhan\\*, Silvio Galesso\\*, Karim Farid,
+Johannes Dienert, Rajat Sahay, Thomas Brox
+— University of Freiburg, Germany (\\*main contributors)
+
+[Project page](https://lmb-freiburg.github.io/orbis2.github.io/) ·
+[Code](https://github.com/lmb-freiburg/orbis) ·
+[Orbis 1](https://lmb-freiburg.github.io/orbis.github.io/)
+
+Upload a short driving clip (a few seconds is enough). The hierarchical model takes
+the tail of the clip as L1 (high-rate) and L2 (low-rate, further back) context and
+autoregressively predicts future frames.
+
+Each run samples **{NUM_VIDEOS} rollouts** from the same context in a single
+minibatch, so you see several plausible futures diverge from identical starting
+conditions.
+"""
+
+PAPER_MD = """
+### Abstract
+
+> Current world models typically operate at a single abstraction level, favoring
+> perceptual fidelity but lacking the spatial and semantic reasoning needed for
+> downstream driving tasks. We propose a hierarchical driving world model that
+> separates prediction into a high-level long-horizon scene forecaster and a
+> low-level detail generator conditioned on it. This design improves both visual
+> fidelity and spatial-semantic representation quality. We also introduce a
+> two-stage training strategy: diffusion-forcing pretraining for richer
+> representations, followed by teacher-forcing fine-tuning for stable
+> autoregressive rollouts. Our method achieves state-of-the-art results on
+> standard driving world model benchmarks, including long-horizon fidelity,
+> counterfactual steering responsiveness, and internal representation quality.
+
+### Method
+
+Orbis 2 splits future prediction across two levels of abstraction:
+
+- **Abstract predictor (L2)** — operates over a long temporal context to forecast a
+  future state in latent space, capturing abstract scene dynamics over long
+  horizons and enabling steering control.
+- **Detail predictor (L1)** — conditioned on that abstract prediction, generates
+  fine-grained short-horizon frames, so high-fidelity local prediction stays
+  grounded in long-range temporal context.
+
+Training runs in two stages: **diffusion-forcing pretraining** for richer
+representations, followed by **teacher-forcing fine-tuning** for stable
+autoregressive rollouts.
+
+### Links & citation
+
+```bibtex
+@article{orbis2_2026,
+  author  = {Mittal, Sudhanshu and Mousakhan, Arian and Galesso, Silvio and
+             Farid, Karim and Dienert, Johannes and Sahay, Rajat and Brox, Thomas},
+  title   = {Orbis 2: A Hierarchical World Model for Driving},
+  journal = {arXiv preprint arXiv:XXXX.XXXXX},
+  year    = {2026},
+}
+```
+"""
+
+DEMO_MD = f"""
+**Pipeline.** Your clip is re-encoded to a constant {CONTEXT_FPS} fps. The rollout
+script samples two context windows from the tail of that video: an **L1** window at
+{L1_FRAME_RATE} Hz (fine, recent frames) and an **L2** window sampled further back at
+the frozen L2 predictor's own trained rate (coarse, long-horizon context). Each frame
+is resized to {FRAME_H}×{FRAME_W}. The **L2 abstract predictor** forecasts a latent
+future, and the **L1 detail predictor**, conditioned on it, generates the next frames
+autoregressively.
+
+**Rollout length.** *Rollout steps* sets how many autoregressive steps to run; each
+step emits several image frames, so the final clip is longer than the step count.
+
+**Sampling.** Each step integrates the flow-matching ODE with *Diffusion steps*
+solver steps. *Eta* (under Advanced) injects noise per step: 0 is the deterministic
+ODE, higher values turn it into an SDE.
+
+**Randomness.** The {NUM_VIDEOS} rollouts share one context and one seed, but the
+solver draws independent initial noise per batch element, so the futures diverge.
+Leave the seed at -1 under Advanced for a fresh draw each run, or set it explicitly
+to reproduce a run exactly.
+
+**Note.** The clip must be long enough for the L2 look-back window; very short clips
+are rejected with an explicit error. Frame and step counts are kept modest to fit the
+ZeroGPU time budget.
+"""
+
+
+# ----------------------------------------------------------------------------
+# Video helpers (CPU)
+# ----------------------------------------------------------------------------
+def reencode_to_fps(src_path: str, out_path: Path, fps: int = CONTEXT_FPS):
+    """Re-encode a clip to a constant `fps` H.264 mp4.
+
+    The rollout script requires the video's native frame rate to be an exact integer
+    multiple of the L1/L2 sampling rates. Forcing a constant fps here (via -vsync cfr)
+    makes an arbitrary upload satisfy that check.
+    """
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    subprocess.run(
+        ["ffmpeg", "-y", "-i", str(src_path),
+         "-vsync", "cfr", "-r", str(fps),
+         "-an", "-c:v", "libx264", "-pix_fmt", "yuv420p",
+         "-movflags", "+faststart", str(out_path)],
+        check=True, capture_output=True,
+    )
+
+
+def frames_to_mp4(frame_paths: list[Path], out_path: Path, fps: int = OUTPUT_FPS):
+    first = cv2.imread(str(frame_paths[0]))
+    h, w = first.shape[:2]
+    # mp4v then re-encode with ffmpeg for browser-compatible H.264
+    tmp = out_path.with_suffix(".raw.mp4")
+    writer = cv2.VideoWriter(str(tmp), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
+    for p in frame_paths:
+        writer.write(cv2.imread(str(p)))
+    writer.release()
+    subprocess.run(
+        ["ffmpeg", "-y", "-i", str(tmp), "-c:v", "libx264",
+         "-pix_fmt", "yuv420p", "-movflags", "+faststart", str(out_path)],
+        check=True, capture_output=True,
+    )
+    tmp.unlink(missing_ok=True)
+
+
+# ----------------------------------------------------------------------------
+# GPU inference
+# ----------------------------------------------------------------------------
+#: tqdm renders the rollout bar as "Rolling out frames:  30%|### | 3/10 [00:05<...]"
+_ROLLOUT_RE = re.compile(r"Rolling out frames.*?(\d+)/(\d+)")
+
+
+def _stream_rollout(cmd: list[str], say) -> str:
+    """Run rollout_demo_v2.py, forwarding its output to the log and driving progress.
+
+    The heavy lifting happens in a child process, so gr.Progress(track_tqdm=True)
+    cannot see its tqdm bars. We scrape the child's own bar instead. tqdm redraws in
+    place with a carriage return rather than a newline, so read whatever bytes are
+    available and split on both.
+    """
+    proc = subprocess.Popen(
+        cmd, cwd=REPO_DIR, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+    )
+    tail: list[str] = []
+    buf = ""
+    while True:
+        chunk = proc.stdout.read1(4096)          # returns as soon as bytes are ready
+        if not chunk:
+            break
+        buf += chunk.decode("utf-8", errors="replace")
+        *lines, buf = re.split(r"[\r\n]", buf)
+        for line in lines:
+            line = line.rstrip()
+            if not line:
+                continue
+            print(line)                          # full trace stays in the Space logs
+            tail.append(line)
+            del tail[:-200]
+            m = _ROLLOUT_RE.search(line)
+            if m:
+                done, total = int(m.group(1)), max(int(m.group(2)), 1)
+                say(0.15 + 0.70 * min(done / total, 1.0),
+                    f"Rolling out step {done}/{total}…")
+
+    proc.wait()
+    joined = "\n".join(tail)
+    if proc.returncode != 0:
+        raise gr.Error(f"Rollout failed:\n{joined[-800:]}")
+    return joined
+
+
+@spaces.GPU(duration=300)
+def run_rollout(video_path: str, num_gen_steps: int, num_steps: int,
+                seed: int, eta: float, progress=gr.Progress()):
+    if video_path is None:
+        raise gr.Error("Please upload a short mp4 clip first.")
+
+    def say(frac, msg):
+        progress(frac, desc=msg)
+
+    try:
+        # The rollout script only seeds when seed > 0; within the batch each sample
+        # still draws its own noise, so one seed yields NUM_VIDEOS distinct rollouts.
+        # seed < 0 means "surprise me" — draw a fresh positive seed each run.
+        seed = random.randrange(1, 2**31 - 1) if int(seed) < 0 else max(1, int(seed))
+
+        job = Path(tempfile.gettempdir()) / f"orbis2_{uuid.uuid4().hex[:8]}"
+
+        say(0.02, "Preparing context video…")
+        ctx_video = job / "context.mp4"
+        reencode_to_fps(video_path, ctx_video, CONTEXT_FPS)
+        print(f"context @ {CONTEXT_FPS} fps | seed {seed} | "
+              f"{int(num_gen_steps)} rollout steps | {int(num_steps)} solver steps | eta {eta}")
+
+        output_dir = job / "rollout"
+
+        say(0.10, "Loading model and rolling out…")
+        cmd = [
+            "python", "evaluate/rollout_demo_v2.py",
+            "--exp_dir", str(EXP_DIR),
+            "--config", CONFIG_NAME,
+            "--ckpt", CKPT_NAME,
+            "--video", str(ctx_video),
+            "--l1_frame_rate", str(L1_FRAME_RATE),
+            # Given explicitly so resolve_context_image_size never falls back to the
+            # config's `data:` block (which points at training-only dataset paths).
+            "--height", str(FRAME_H),
+            "--width", str(FRAME_W),
+            "--num_gen_frames", str(int(num_gen_steps)),
+            "--num_steps", str(int(num_steps)),
+            "--eta", str(float(eta)),
+            "--seed", str(seed),
+            "--num_videos", str(NUM_VIDEOS),
+            "--output_dir", str(output_dir),
+            "--device", "cuda",
+        ]
+        _stream_rollout(cmd, say)
+
+        say(0.88, "Encoding videos…")
+        out_mp4s = []
+        for i in range(NUM_VIDEOS):
+            seq_dir = output_dir / "fake_images" / f"sequence_{i:04d}"
+            gen_frames = sorted(seq_dir.glob("*.jpg"))
+            if not gen_frames:
+                raise gr.Error(f"No generated frames found in {seq_dir} — "
+                               "check rollout output path.")
+            out_mp4 = job / f"rollout_{i}.mp4"
+            frames_to_mp4(gen_frames, out_mp4)
+            out_mp4s.append(str(out_mp4))
+
+        say(1.0, f"Done — {NUM_VIDEOS} rollouts of {int(num_gen_steps)} steps.")
+        return tuple(out_mp4s)
+
+    except gr.Error:
+        raise
+    except Exception as e:
+        tb = traceback.format_exc()
+        print(tb)
+        raise gr.Error(f"Rollout failed: {e}\n\n{tb[-1500:]}")
+
+
+# ----------------------------------------------------------------------------
+# UI
+# ----------------------------------------------------------------------------
+def build_demo():
+    with gr.Blocks(title="Orbis 2: A Hierarchical World Model for Driving") as demo:
+        gr.Markdown(DESCRIPTION)
+
+        with gr.Accordion("About the paper", open=False):
+            gr.Markdown(PAPER_MD)
+        with gr.Accordion("About this demo", open=False):
+            gr.Markdown(DEMO_MD)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                inp = gr.Video(label="Upload a short mp4 clip", sources=["upload"])
+                n_steps_gen = gr.Slider(4, 30, value=DEFAULT_GEN_STEPS, step=2,
+                                        label="Rollout steps",
+                                        info="Each step emits several frames")
+                n_steps = gr.Slider(10, 40, value=DEFAULT_STEPS, step=5,
+                                    label="Diffusion steps",
+                                    info="Flow-matching solver steps per frame")
+
+                with gr.Accordion("Advanced", open=False):
+                    seed_in = gr.Number(value=-1, precision=0, label="Seed",
+                                        info="-1 draws a fresh seed each run")
+                    eta = gr.Slider(0.0, 1.0, value=DEFAULT_ETA, step=0.05,
+                                    label="Eta (stochasticity)",
+                                    info="0 = deterministic ODE; >0 adds noise per step")
+
+                btn = gr.Button("Generate rollouts", variant="primary")
+
+            with gr.Column(scale=1):
+                outs = [gr.Video(label=f"Rollout {i + 1}", autoplay=True, loop=True)
+                        for i in range(NUM_VIDEOS)]
+
+        btn.click(
+            run_rollout,
+            inputs=[inp, n_steps_gen, n_steps, seed_in, eta],
+            outputs=outs,
+            concurrency_limit=1,
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    build_demo().queue(max_size=8).launch()

orbis2/data/datamodule.py

@@ -0,0 +1,306 @@
+import math
+import os
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader, ConcatDataset
+from torch.utils.data.distributed import DistributedSampler
+from torch.utils.data.dataloader import default_collate
+
+import pytorch_lightning as pl
+
+from omegaconf import ListConfig, DictConfig
+
+import logging
+
+from util import instantiate_from_config
+
+logger = logging.getLogger(__name__)
+
+
+def _collate_pad_missing(batch):
+    """Collate dicts that may have different keys across datasets.
+
+    Missing keys are filled with zero tensors matching the shape of the first
+    sample in the batch that has that key. Non-tensor values are filled with
+    None. Allows heterogeneous datasets (e.g. with/without steering) to be
+    mixed in the same batch.
+    """
+    if not isinstance(batch[0], dict):
+        return default_collate(batch)
+
+    all_keys = set().union(*[item.keys() for item in batch])
+    filled = []
+    for item in batch:
+        item = dict(item)
+        for key in all_keys:
+            if key not in item:
+                ref = next((b[key] for b in batch if key in b), None)
+                if isinstance(ref, torch.Tensor):
+                    item[key] = torch.full_like(ref, float("nan"))
+                else:
+                    item[key] = ref
+        filled.append(item)
+    return default_collate(filled)
+
+
+def _env_bool(name, default):
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    value = raw.strip().lower()
+    if value in {"1", "true", "yes", "y", "on"}:
+        return True
+    if value in {"0", "false", "no", "n", "off"}:
+        return False
+    logger.warning("Invalid boolean %s=%r; using default %s", name, raw, default)
+    return default
+
+
+def _env_int(name, default):
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        logger.warning("Invalid integer %s=%r; using default %s", name, raw, default)
+        return default
+
+
+class DataModuleFromConfig(pl.LightningDataModule):
+    def __init__(self, batch_size, val_batch_size=None, train=None, validation=None, test=None,
+                 wrap=False, num_workers=None, dbg=False, train_weights=None):
+        super().__init__()
+        self.batch_size = batch_size
+        self.val_batch_size = val_batch_size if val_batch_size is not None else batch_size
+        self.dataset_configs = dict()
+        self.num_workers = num_workers if num_workers is not None else batch_size*2
+        if train is not None:
+            self.dataset_configs["train"] = train
+            self.train_dataloader = self._train_dataloader
+        if validation is not None:
+            self.dataset_configs["validation"] = validation
+            self.val_dataloader = self._val_dataloader
+        if test is not None:
+            self.dataset_configs["test"] = test
+            self.test_dataloader = self._test_dataloader
+        self.wrap = wrap
+        self.dbg = dbg
+
+        if train_weights is not None:
+            if not isinstance(train, (list, ListConfig)):
+                raise ValueError("train_weights requires train to be a list of dataset configs")
+            if len(train_weights) != len(train):
+                raise ValueError(
+                    f"train_weights has {len(train_weights)} entries but train has {len(train)} datasets"
+                )
+        self.train_weights = train_weights
+
+        if self.wrap:
+            raise NotImplementedError("Wrapped datasets not implemented")
+
+        self._resume_epoch = None
+        self._resume_batches_completed = 0
+
+    def state_dict(self):
+        trainer = getattr(self, "trainer", None)
+        if trainer is None:
+            return {}
+
+        epoch = int(getattr(trainer, "current_epoch", 0))
+        batches_completed = 0
+        try:
+            # Lightning tracks completed batches for the current epoch in this counter.
+            batches_completed = int(trainer.fit_loop.epoch_loop.batch_progress.current.completed)
+        except Exception:
+            batches_completed = 0
+
+        return {
+            "resume_epoch": epoch,
+            "resume_batches_completed": max(batches_completed, 0),
+        }
+
+    def load_state_dict(self, state_dict):
+        if not isinstance(state_dict, dict):
+            return
+
+        self._resume_epoch = state_dict.get("resume_epoch")
+        self._resume_batches_completed = int(state_dict.get("resume_batches_completed", 0) or 0)
+        if self._resume_batches_completed < 0:
+            self._resume_batches_completed = 0
+
+    def setup(self, stage=None):
+        self.datasets = dict()
+        for k, cfg in self.dataset_configs.items():
+            logger.info("Loading dataset: %s", k)
+            if isinstance(cfg, (list, ListConfig)):
+                datasets = [instantiate_from_config(c) for c in cfg]
+                self.datasets[k] = ConcatDataset(datasets)
+                [logger.info(d) for d in datasets]
+            elif isinstance(cfg, DictConfig):
+                ds = instantiate_from_config(cfg)
+                self.datasets[k] = ds
+                logger.info(ds)
+            else:
+                raise ValueError(f"Invalid dataset config: {cfg}")
+
+    def _train_dataloader(self):
+        is_distributed = dist.is_available() and dist.is_initialized()
+        use_distributed_sampler = _env_bool("ORBIS_USE_DISTRIBUTED_SAMPLER", True)
+        sampler = None
+        if self.train_weights is not None:
+            train_ds = self.datasets["train"]
+            sub_datasets = getattr(train_ds, "datasets", [train_ds])
+            dataset_sizes = [len(ds) for ds in sub_datasets]
+            # Anchor epoch length to min(size/weight) across datasets.
+            # This fully covers the most "weight-adjusted-constrained" dataset (typically
+            # the highest-weight one) with exactly one pass, while letting smaller/lower-weight
+            # datasets cycle. Avoids the 2x repetition that len(ConcatDataset) causes when
+            # one dataset is large and dominates with high weight.
+            num_samples = math.ceil(min(s / w for s, w in zip(dataset_sizes, self.train_weights)))
+            sampler = _WeightedDistributedSampler(
+                self.train_weights, dataset_sizes, num_samples,
+                num_replicas=dist.get_world_size() if (is_distributed and use_distributed_sampler) else 1,
+                rank=dist.get_rank() if (is_distributed and use_distributed_sampler) else 0,
+            )
+            logger.info(
+                "Train DataLoader: weighted sampling fractions=%s dataset sizes=%s num_samples/epoch=%s",
+                list(self.train_weights),
+                dataset_sizes,
+                num_samples,
+            )
+        elif is_distributed and use_distributed_sampler:
+            sampler = _ResumableDistributedSampler(
+                self.datasets["train"],
+                shuffle=True,
+                drop_last=True,
+                resume_epoch=self._resume_epoch,
+                resume_batches_completed=self._resume_batches_completed,
+                batch_size=self.batch_size,
+            )
+
+        timeout_s = _env_int("ORBIS_DATALOADER_TIMEOUT_S", 0)
+        prefetch_factor = _env_int("ORBIS_DATALOADER_PREFETCH_FACTOR", 1)
+        persistent_workers = _env_bool(
+            "ORBIS_DATALOADER_PERSISTENT_WORKERS",
+            self.num_workers > 0,
+        )
+        mp_context = os.environ.get("ORBIS_DATALOADER_MP_CONTEXT", "").strip()
+
+        loader_kwargs = dict(
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            shuffle=sampler is None,
+            pin_memory=True,
+            drop_last=True,
+            sampler=sampler,
+            timeout=max(timeout_s, 0),
+            collate_fn=_collate_pad_missing if self.train_weights is not None else None,
+        )
+        if self.num_workers > 0:
+            loader_kwargs["persistent_workers"] = persistent_workers
+            if prefetch_factor > 0:
+                loader_kwargs["prefetch_factor"] = prefetch_factor
+            if mp_context:
+                loader_kwargs["multiprocessing_context"] = mp_context
+
+        logger.info(
+            "Train DataLoader: distributed=%s sampler=%s workers=%s timeout_s=%s mp_context=%s",
+            is_distributed,
+            sampler.__class__.__name__ if sampler is not None else "None",
+            self.num_workers,
+            loader_kwargs["timeout"],
+            mp_context or "<default>",
+        )
+
+        if self.dbg:
+            dbg_sampler = DistributedSampler(self.datasets["train"], shuffle=True)
+            loader_kwargs["sampler"] = dbg_sampler
+            loader_kwargs["shuffle"] = False
+
+        return DataLoader(self.datasets["train"], **loader_kwargs)
+
+    def _val_dataloader(self):
+        return DataLoader(self.datasets["validation"],
+                          batch_size=self.val_batch_size,
+                          num_workers=self.num_workers, pin_memory=True)
+
+    def _test_dataloader(self):
+        return DataLoader(self.datasets["test"], batch_size=self.val_batch_size,
+                          num_workers=self.num_workers)
+
+
+class _ResumableDistributedSampler(DistributedSampler):
+    def __init__(
+        self,
+        dataset,
+        *,
+        resume_epoch=None,
+        resume_batches_completed=0,
+        batch_size=1,
+        **kwargs,
+    ):
+        super().__init__(dataset, **kwargs)
+        self.resume_epoch = None if resume_epoch is None else int(resume_epoch)
+        self.resume_batches_completed = int(resume_batches_completed or 0)
+        self.batch_size = max(int(batch_size), 1)
+
+    def __iter__(self):
+        indices = list(super().__iter__())
+        if (
+            self.resume_epoch is not None
+            and int(self.epoch) == self.resume_epoch
+            and self.resume_batches_completed > 0
+        ):
+            skip = self.resume_batches_completed * self.batch_size
+            if skip > 0:
+                logger.info(
+                    "Resuming dataloader at epoch=%s after %s batches (%s samples/rank).",
+                    self.resume_epoch,
+                    self.resume_batches_completed,
+                    skip,
+                )
+                indices = indices[skip:]
+        return iter(indices)
+
+
+class _WeightedDistributedSampler(torch.utils.data.Sampler):
+    """Weighted sampler for dataset balancing, with optional distributed sharding.
+
+    Two-stage sampling: (1) pick a dataset by weight, (2) pick a uniform index
+    within that dataset. Avoids torch.multinomial's 2^24 category limit since
+    only num_datasets categories are ever sampled, not num_total_samples.
+    """
+
+    def __init__(self, dataset_weights, dataset_sizes, num_samples, num_replicas=1, rank=0):
+        self.dataset_weights = torch.as_tensor(dataset_weights, dtype=torch.float64)
+        self.dataset_sizes = torch.tensor(dataset_sizes, dtype=torch.long)
+        offsets = [0]
+        for s in dataset_sizes[:-1]:
+            offsets.append(offsets[-1] + s)
+        self.dataset_offsets = torch.tensor(offsets, dtype=torch.long)
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.num_samples_per_replica = math.ceil(num_samples / num_replicas)
+        self.epoch = 0
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def __iter__(self):
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+        total = self.num_samples_per_replica * self.num_replicas
+
+        # Stage 1: pick dataset for each draw (num_datasets categories, well within 2^24)
+        ds_idx = torch.multinomial(self.dataset_weights, total, replacement=True, generator=g)
+
+        # Stage 2: pick a uniform index within each chosen dataset
+        local = (torch.rand(total, generator=g) * self.dataset_sizes[ds_idx].float()).long()
+        indices = (self.dataset_offsets[ds_idx] + local).tolist()
+
+        return iter(indices[self.rank::self.num_replicas])
+
+    def __len__(self):
+        return self.num_samples_per_replica

orbis2/data/helper_types.py

@@ -0,0 +1,49 @@
+from typing import Dict, Tuple, Optional, NamedTuple, Union
+from PIL.Image import Image as pil_image
+from torch import Tensor
+
+try:
+  from typing import Literal
+except ImportError:
+  from typing_extensions import Literal
+
+Image = Union[Tensor, pil_image]
+BoundingBox = Tuple[float, float, float, float]  # x0, y0, w, h
+CropMethodType = Literal['none', 'random', 'center', 'random-2d']
+SplitType = Literal['train', 'validation', 'test']
+
+
+class ImageDescription(NamedTuple):
+    id: int
+    file_name: str
+    original_size: Tuple[int, int]  # w, h
+    url: Optional[str] = None
+    license: Optional[int] = None
+    coco_url: Optional[str] = None
+    date_captured: Optional[str] = None
+    flickr_url: Optional[str] = None
+    flickr_id: Optional[str] = None
+    coco_id: Optional[str] = None
+
+
+class Category(NamedTuple):
+    id: str
+    super_category: Optional[str]
+    name: str
+
+
+class Annotation(NamedTuple):
+    area: float
+    image_id: str
+    bbox: BoundingBox
+    category_no: int
+    category_id: str
+    id: Optional[int] = None
+    source: Optional[str] = None
+    confidence: Optional[float] = None
+    is_group_of: Optional[bool] = None
+    is_truncated: Optional[bool] = None
+    is_occluded: Optional[bool] = None
+    is_depiction: Optional[bool] = None
+    is_inside: Optional[bool] = None
+    segmentation: Optional[Dict] = None

orbis2/data/l2_context.py

@@ -0,0 +1,83 @@
+class L2ContextMixin:
+    """Reusable helper for hierarchical L1/L2 frame alignment."""
+
+    def _init_l2_context(
+        self,
+        *,
+        num_l2_context=0,
+        l2_frame_rate=1.0,
+        l1_context_frames=1,
+        require_l2_context=False,
+    ):
+        self.num_l2_context = int(num_l2_context)
+        self.l1_context_frames = int(l1_context_frames)
+
+        if self.num_l2_context < 0:
+            raise ValueError(f"num_l2_context must be non-negative, got {self.num_l2_context}.")
+        if self.l1_context_frames <= 0:
+            raise ValueError(f"l1_context_frames must be positive, got {self.l1_context_frames}.")
+        if require_l2_context and self.num_l2_context <= 0:
+            raise ValueError("num_l2_context must be positive for hierarchical L1/L2 loaders.")
+
+        self.l2_frame_interval = None
+        self._l2_anchor_offset = (self.l1_context_frames - 1) * self.frame_interval
+        if self.num_l2_context == 0:
+            return
+
+        ratio = float(self.stored_data_frame_rate) / float(l2_frame_rate)
+        rounded_ratio = round(ratio)
+        if abs(ratio - rounded_ratio) > 1e-8:
+            raise ValueError(
+                "stored_data_frame_rate must be an integer multiple of l2_frame_rate, "
+                f"got stored_data_frame_rate={self.stored_data_frame_rate}, "
+                f"l2_frame_rate={l2_frame_rate}."
+            )
+        self.l2_frame_interval = int(rounded_ratio)
+
+    @property
+    def l2_context_enabled(self):
+        return self.num_l2_context > 0
+
+    def get_required_l1_start_offset(self):
+        if not self.l2_context_enabled:
+            return 0
+        return max(
+            0,
+            (self.num_l2_context - 1) * self.l2_frame_interval - self._l2_anchor_offset,
+        )
+
+    def has_l2_context_for_start(self, start_frame):
+        if not self.l2_context_enabled:
+            return True
+        l2_anchor = start_frame + self._l2_anchor_offset
+        oldest_l2 = l2_anchor - self.l2_frame_interval * (self.num_l2_context - 1)
+        return oldest_l2 >= 0
+
+    def filter_index_map_with_l2_headroom(self, index_map, start_frame_idx=-1):
+        if not self.l2_context_enabled:
+            return list(index_map)
+
+        filtered = []
+        for item in index_map:
+            start_frame = item[start_frame_idx]
+            if self.has_l2_context_for_start(start_frame):
+                filtered.append(item)
+        return filtered
+
+    def get_l1_indices(self, start_frame, num_frames):
+        return list(range(
+            start_frame,
+            start_frame + num_frames * self.frame_interval,
+            self.frame_interval,
+        ))
+
+    def get_l2_indices(self, start_frame):
+        if not self.l2_context_enabled:
+            return []
+
+        l2_end = start_frame + self._l2_anchor_offset
+        l2_start = l2_end - (self.num_l2_context - 1) * self.l2_frame_interval
+        return list(range(l2_start, l2_end + 1, self.l2_frame_interval))
+
+    def get_l1_and_l2_indices(self, start_frame, num_frames):
+        return self.get_l1_indices(start_frame, num_frames), self.get_l2_indices(start_frame)

orbis2/data/steering_loaders.py

@@ -0,0 +1,948 @@
+import os
+
+import h5py
+import numpy as np
+import torch
+
+from data.utils import get_trajectory_from_speeds_and_yaw_rates
+from data.l2_context import L2ContextMixin
+from data.video_loaders import (
+    MultiHDF5DatasetMultiFrameIdxMapping,
+    MultiMP4DatasetMultiFrameIdxMapping,
+    MultiMP4DatasetMultiFrameIdxMappingWithL2Context,
+)
+from data.vista_style import (
+    VistaStyleNuScenesLoader,
+    VistaStyleNuScenesLoaderWithL2Context,
+    extract_pose_table,
+    get_pose,
+)
+from util import instantiate_from_config
+
+
+class OdometryHorizonMixin:
+    @staticmethod
+    def _validate_frame_rate_ratio(stored_rate, sampled_rate, rate_name):
+        ratio = stored_rate / sampled_rate
+        if not np.isclose(ratio, round(ratio), atol=1e-8):
+            raise ValueError(
+                f"stored_data_frame_rate={stored_rate} must be an integer multiple of "
+                f"{rate_name}={sampled_rate}"
+            )
+
+    @staticmethod
+    def _infer_num_frames_odo_for_matched_horizon(num_frames, frame_rate, odometry_frame_rate):
+        inferred = (num_frames - 1) * odometry_frame_rate / frame_rate + 1
+        if not np.isclose(inferred, round(inferred), atol=1e-8):
+            raise ValueError(
+                "Matched video/odometry horizons require an integer odometry length, got "
+                f"num_frames={num_frames}, frame_rate={frame_rate}, "
+                f"odometry_frame_rate={odometry_frame_rate}"
+            )
+        return int(round(inferred))
+
+    @classmethod
+    def _resolve_odometry_horizon(
+        cls,
+        num_frames,
+        frame_rate,
+        num_frames_odo,
+        odometry_frame_rate,
+        odometry_horizon,
+        inferred_num_frames_odo,
+    ):
+        valid_horizons = {"auto", "match_video", "explicit"}
+        if odometry_horizon not in valid_horizons:
+            raise ValueError(
+                f"Unknown odometry_horizon={odometry_horizon!r}. Expected one of {sorted(valid_horizons)}"
+            )
+
+        if odometry_horizon == "auto":
+            if num_frames_odo is None or num_frames_odo == inferred_num_frames_odo:
+                odometry_horizon = "match_video"
+            else:
+                odometry_horizon = "explicit"
+
+        if odometry_horizon == "match_video":
+            resolved_num_frames_odo = inferred_num_frames_odo
+            if num_frames_odo is not None and num_frames_odo != resolved_num_frames_odo:
+                raise ValueError(
+                    f"num_frames_odo={num_frames_odo} does not match the video horizon; "
+                    f"expected {resolved_num_frames_odo} for num_frames={num_frames}, "
+                    f"frame_rate={frame_rate}, odometry_frame_rate={odometry_frame_rate}"
+                )
+            return resolved_num_frames_odo, odometry_horizon
+
+        if num_frames_odo is None:
+            raise ValueError("odometry_horizon='explicit' requires num_frames_odo to be provided")
+        return int(num_frames_odo), odometry_horizon
+
+    @classmethod
+    def reconfigure_params_for_required_odometry_horizon(cls, params, required_odo_steps):
+        if required_odo_steps is None:
+            raise ValueError("`required_odo_steps` must be provided")
+        required_odo_steps = int(required_odo_steps)
+        if required_odo_steps <= 0:
+            raise ValueError(f"`required_odo_steps` must be positive, got {required_odo_steps}.")
+
+        if "odometry_frame_rate" not in params or params.odometry_frame_rate is None:
+            params.odometry_frame_rate = params.frame_rate
+
+        frame_rate = float(params.frame_rate)
+        odometry_frame_rate = float(params.odometry_frame_rate)
+
+        if odometry_frame_rate != frame_rate:
+            required_num_frames = int(np.ceil((required_odo_steps - 1) * frame_rate / odometry_frame_rate)) + 1
+            params.num_frames = max(int(params.num_frames), required_num_frames)
+            params.odometry_horizon = "match_video"
+            if "num_frames_odo" in params:
+                del params["num_frames_odo"]
+            return params
+
+        if "num_frames_odo" in params and params.num_frames_odo is not None:
+            params.num_frames_odo = max(int(params.num_frames_odo), required_odo_steps)
+        else:
+            params.num_frames_odo = required_odo_steps
+        params.odometry_horizon = "explicit"
+        return params
+
+
+class PlaceholderSteeringMixin:
+    """Mixin for datasets that have no real steering data.
+
+    Injects a NaN steering tensor of the correct shape so that rollout_steering_v2
+    can work with --steering_file or --no_steering without any changes to downstream
+    code.  Using such a dataset without either of those flags raises an explicit error
+    at rollout validation time (see validate_steering_source in rollout_steering_v2.py).
+    """
+
+    def _init_placeholder_steering(self, num_frames_odo, steering_dim=2, steering_format="speed_yawrate"):
+        self._placeholder_num_frames_odo = int(num_frames_odo)
+        self._placeholder_steering_dim = int(steering_dim)
+        self._placeholder_steering_format = steering_format
+
+    def _add_placeholder_steering(self, item: dict) -> dict:
+        item["steering"] = torch.full(
+            (self._placeholder_num_frames_odo, self._placeholder_steering_dim), float("nan")
+        )
+        item["steering_format"] = self._placeholder_steering_format
+        item["_steering_placeholder"] = True
+        return item
+
+    @classmethod
+    def reconfigure_params_for_required_odometry_horizon(cls, params, required_odo_steps):
+        if required_odo_steps is None:
+            raise ValueError("`required_odo_steps` must be provided")
+        required_odo_steps = int(required_odo_steps)
+        if required_odo_steps <= 0:
+            raise ValueError(f"`required_odo_steps` must be positive, got {required_odo_steps}.")
+        current = int(getattr(params, "num_frames_odo", 0) or 0)
+        params.num_frames_odo = max(current, required_odo_steps)
+        return params
+
+
+class OdometryLoaderConti:
+    steering_format = "speed_yawrate"
+
+    def __call__(self, odo_data):
+        speeds = np.array([odo_frame[0] for odo_frame in odo_data])
+        yaw_rates = np.array([odo_frame[6] for odo_frame in odo_data])
+        ret_odo = np.stack([speeds, yaw_rates], axis=1)
+        assert ret_odo.shape[1] == 2 and ret_odo.shape[0] == len(odo_data), f"Unexpected odometry shape {ret_odo.shape}, expected ({len(odo_data)}, 2)"
+        return ret_odo
+    
+
+class OdometryLoaderNuPlan:
+    steering_format = "speed_yawrate"
+
+    def __call__(self, odo_data):
+        """
+        odo_data: list of dicts, each dict contains IMU data for a frame
+        """
+        ret = np.stack([np.array([odo_frame["vx"], odo_frame["angular_rate_z"]]) for odo_frame in odo_data], axis=0)
+        assert ret.shape[1] == 2 and ret.shape[0] == len(odo_data), f"Unexpected odometry shape {ret.shape}, expected ({len(odo_data)}, 2)"
+        return ret
+
+
+class OdometryLoaderNVIDIAPhysAI:
+    steering_format = "speed_yawrate"
+
+    def __init__(self, speed_key="vx", curvature_key="curvature"):
+        self.speed_key = speed_key
+        self.curvature_key = curvature_key
+        
+    def __call__(self, odo_data):
+        """
+        odo_data: list of dicts, each dict contains vx (speed), and curvature (which can be used to compute yaw rate as curvature * speed)
+        
+        We return the speed and yaw rate as the odometry.
+        
+        """
+        speeds = np.array([odo_frame[self.speed_key] for odo_frame in odo_data])
+        yaw_rates = np.array([odo_frame[self.curvature_key] * odo_frame[self.speed_key] for odo_frame in odo_data])
+        ret_odo = np.stack([speeds, yaw_rates], axis=1)
+        assert ret_odo.shape[1] == 2 and ret_odo.shape[0] == len(odo_data), f"Unexpected odometry shape {ret_odo.shape}, expected ({len(odo_data)}, 2)"
+        return ret_odo
+
+
+class TrajectoryLoaderNuPlanFromSpeedYawRate:
+    steering_format = "trajectory_with_heading"
+
+    def __init__(self, frame_rate, speed_key="vx", yaw_rate_key="angular_rate_z"):
+        self.frame_rate = frame_rate
+        self.speed_key = speed_key
+        self.yaw_rate_key = yaw_rate_key
+
+    def __call__(self, odo_data):
+        """
+        odo_data: list of dicts, each dict contains IMU data for a frame
+        """
+        traj, headings = get_trajectory_from_speeds_and_yaw_rates(
+            speeds=np.array([odo_frame[self.speed_key] for odo_frame in odo_data]),
+            yaw_rates=np.array([odo_frame[self.yaw_rate_key] for odo_frame in odo_data]),
+            dt=1.0 / self.frame_rate,
+        )
+
+        assert traj.shape[1] == 2 and traj.shape[0] == len(odo_data), f"Unexpected odometry shape {traj.shape}, expected ({len(odo_data)}, 2)"
+        steering = np.concatenate([traj, headings[:, None]], axis=-1)  # (num_frames, 3)
+        return steering
+
+
+class TrajectoryLoaderNVIDIAPhysAIFromSpeedCurvature:
+    def __init__(self, frame_rate, speed_key="speed", curvature_key="curvature", return_headings=False):
+        self.frame_rate = frame_rate
+        self.speed_key = speed_key
+        self.curvature_key = curvature_key
+        self.return_headings = return_headings
+        self.steering_format = "trajectory_with_heading" if return_headings else "trajectory"
+
+    def __call__(self, odo_data):
+        """
+        odo_data: list of dicts, each dict contains vx (speed), and curvature (which can be used to compute yaw rate as curvature * speed)
+        
+        We then compute the trajectory by integrating the speeds and yaw rates over time.
+        
+        """
+        speeds = np.array([odo_frame[self.speed_key] for odo_frame in odo_data])
+        yaw_rates = np.array([odo_frame[self.curvature_key] * odo_frame[self.speed_key] for odo_frame in odo_data])
+        traj, headings = get_trajectory_from_speeds_and_yaw_rates(
+            speeds=speeds,
+            yaw_rates=yaw_rates,
+            dt=1.0 / self.frame_rate,
+        )
+        assert traj.shape[1] == 2 and traj.shape[0] == len(odo_data), f"Unexpected odometry shape {traj.shape}, expected ({len(odo_data)}, 2)"
+        if self.return_headings:
+            steering = np.concatenate([traj, headings[:, None]], axis=-1)  # (num_frames, 3)
+        else:
+            steering = traj  # (num_frames, 2)
+        return steering
+
+
+
+class MultiMP4DatasetMultiFrameIdxMappingNVIDIAPhysAI(OdometryHorizonMixin, MultiMP4DatasetMultiFrameIdxMapping):
+    """
+    Loads the odometry from the NVIDIAPhysAI dataset. The odometry is stored in a single HDF5 file. The HDF5 file contains a dataset for each video (named after the video id).
+    
+    """
+    def __init__(
+        self,
+        size,
+        mp4_paths_file,
+        odometry_h5_path,
+        num_frames,
+        num_frames_odo=None,
+        stored_data_frame_rate=5,
+        frame_rate=5,
+        odometry_frame_rate=None,
+        odometry_horizon="auto",
+        aug="resize_center",
+        backend=None,
+        odo_transform_config=None,
+        return_video_id=False,
+        subsample_interval=None,
+        intrinsics_h5_path=None,
+        spatial_transform_config=None,
+        validate_frame_rate_sample=True,
+    ):
+        self.odometry_frame_rate = frame_rate if odometry_frame_rate is None else odometry_frame_rate
+        self.odometry_horizon = odometry_horizon
+
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=frame_rate,
+            rate_name="frame_rate",
+        )
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=self.odometry_frame_rate,
+            rate_name="odometry_frame_rate",
+        )
+        self.odo_frame_interval = int(round(stored_data_frame_rate / self.odometry_frame_rate))
+        inferred_num_frames_odo = self._infer_num_frames_odo_for_matched_horizon(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            odometry_frame_rate=self.odometry_frame_rate,
+        )
+        self.num_frames_odo, self.odometry_horizon = self._resolve_odometry_horizon(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_frames_odo=num_frames_odo,
+            odometry_frame_rate=self.odometry_frame_rate,
+            odometry_horizon=odometry_horizon,
+            inferred_num_frames_odo=inferred_num_frames_odo,
+        )
+
+        super().__init__(
+            size=size,
+            mp4_paths_file=mp4_paths_file,
+            num_frames=num_frames,
+            stored_data_frame_rate=stored_data_frame_rate,
+            frame_rate=frame_rate,
+            aug=aug,
+            backend=backend,
+            subsample_interval=subsample_interval,
+            intrinsics_h5_path=intrinsics_h5_path,
+            spatial_transform_config=spatial_transform_config,
+            validate_frame_rate_sample=validate_frame_rate_sample,
+        )
+
+        self.odo_file = h5py.File(odometry_h5_path, "r")
+        self.return_video_id = return_video_id
+        
+        if odo_transform_config is None:
+            raise ValueError(
+                "MultiMP4DatasetMultiFrameIdxMappingNVIDIAPhysAI requires explicit "
+                "`odo_transform_config`. Use `data.steering_loaders.OdometryLoaderNVIDIAPhysAI` "
+                "for raw [speed, yaw_rate], or `data.steering_loaders."
+                "TrajectoryLoaderNVIDIAPhysAIFromSpeedCurvature` for trajectory targets."
+            )
+        self.odo_transform = instantiate_from_config(odo_transform_config)
+        self.steering_format = getattr(self.odo_transform, "steering_format", "unknown")
+        
+        # check that all videos have corresponding odometry data, and that the odometry data has the same number of frames as the video
+        for mp4_path in self.mp4_paths:
+            video_id = self._get_video_id(mp4_path)
+            if video_id not in self.odo_file:
+                raise KeyError(f"Video id {video_id} not found in odometry file {odometry_h5_path}")
+            odo_length = len(self.odo_file[video_id]['odometry'])
+            video_length = self.mp4_lengths_in_frames[mp4_path] if self.mp4_lengths_in_frames is not None else self.get_video_length(mp4_path)
+            # max difference of 1 frame is allowed to account for rounding issues when the frame rates are different
+            if abs(odo_length - video_length) > 1:
+                raise ValueError(f"Odometry length {odo_length} does not match video length {video_length} for video id {video_id}")
+
+    def scan_mp4_files(self):
+        self.index_to_starting_frame_map = []
+        required_span = max(
+            self.num_frames * self.frame_interval,
+            self.num_frames_odo * self.odo_frame_interval,
+        )
+        for path in self.mp4_paths:
+            if self.mp4_lengths_in_frames is not None and path in self.mp4_lengths_in_frames:
+                video_length = self.mp4_lengths_in_frames[path]
+            else:
+                video_length = self.get_video_length(path)
+
+            frame_interval = self.frame_interval if self.subsample_interval is None else self.subsample_interval * self.frame_interval
+            max_frame_index = video_length - required_span - 1
+            for i in range(0, max_frame_index + 1, frame_interval):
+                self.index_to_starting_frame_map.append((path, i))
+
+    def _get_odometry_indices(self, start_frame):
+        return list(range(
+            start_frame,
+            start_frame + self.num_frames_odo * self.odo_frame_interval,
+            self.odo_frame_interval,
+        ))
+
+    def _load_steering_for_video(self, video_id, start_frame):
+        indices = self._get_odometry_indices(start_frame)
+        odo_data = [self.odo_file[video_id]['odometry'][i] for i in indices]
+        return torch.from_numpy(self.odo_transform(odo_data)).float()
+
+    def __getitem__(self, idx):
+        images, (_, path, start_frame) = self.get_images_and_indices(idx)
+        images = self.apply_transforms(images, context=self.build_context((None, path, start_frame)))
+        video_id = self._get_video_id(path)
+        result = {
+            "images": images,
+            "frame_rate": self.frame_rate,
+            "steering": self._load_steering_for_video(video_id, start_frame),
+            "steering_format": self.steering_format,
+        }
+        if self.return_video_id:
+            result["video_id"] = video_id
+        return result
+
+
+class MultiHDF5DatasetMultiFrameIdxMappingOdometry(
+    OdometryHorizonMixin,
+    MultiHDF5DatasetMultiFrameIdxMapping,
+):
+    def __init__(
+        self,
+        size,
+        hdf5_paths_file,
+        num_frames,
+        num_frames_odo=None,
+        frames_file_suffix="frames.h5",
+        odo_files_suffix="odometry.h5",
+        stored_data_frame_rate=5,
+        frame_rate=5,
+        odometry_frame_rate=None,
+        odometry_horizon="auto",
+        aug="resize_center",
+        scale_min=0.15,
+        scale_max=0.5,
+        odo_transform_config=None,
+    ):
+        self.frames_file_suffix = frames_file_suffix
+        self.odo_files_suffix = odo_files_suffix if odo_files_suffix is None or odo_files_suffix.lower() != "none" else None
+        self.odometry_frame_rate = frame_rate if odometry_frame_rate is None else odometry_frame_rate
+        self.odometry_horizon = odometry_horizon
+
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=frame_rate,
+            rate_name="frame_rate",
+        )
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=self.odometry_frame_rate,
+            rate_name="odometry_frame_rate",
+        )
+        self.odo_frame_interval = int(round(stored_data_frame_rate / self.odometry_frame_rate))
+        inferred_num_frames_odo = self._infer_num_frames_odo_for_matched_horizon(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            odometry_frame_rate=self.odometry_frame_rate,
+        )
+        self.num_frames_odo, self.odometry_horizon = self._resolve_odometry_horizon(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_frames_odo=num_frames_odo,
+            odometry_frame_rate=self.odometry_frame_rate,
+            odometry_horizon=odometry_horizon,
+            inferred_num_frames_odo=inferred_num_frames_odo,
+        )
+
+        super().__init__(
+            size=size,
+            hdf5_paths_file=hdf5_paths_file,
+            num_frames=num_frames,
+            stored_data_frame_rate=stored_data_frame_rate,
+            frame_rate=frame_rate,
+            aug=aug,
+            scale_min=scale_min,
+            scale_max=scale_max,
+        )
+
+        if odo_transform_config is not None:
+            self.odo_transform = instantiate_from_config(odo_transform_config)
+        else:
+            self.odo_transform = OdometryLoaderNuPlan()
+        self.steering_format = getattr(self.odo_transform, "steering_format", "unknown")
+
+    @staticmethod
+    def frames_odo_matching_check(frames_h5, odo_h5):
+        for key in frames_h5.keys():
+            if "meta_data" in key:
+                continue
+            if key not in odo_h5:
+                raise KeyError(f"Odometry key {key} not found in frames data")
+            if len(odo_h5[key]) != len(frames_h5[key]):
+                raise ValueError(
+                    f"Odometry key {key} has different length than frames data: "
+                    f"{len(odo_h5[key])} != {len(frames_h5[key])}"
+                )
+
+    def scan_h5_files_odo(self):
+        self.files_odo = {}
+        if self.odo_files_suffix is None:
+            return
+
+        for h5_file_frames in self.hdf5_files:
+            odo_h5_path = h5_file_frames.filename.replace(self.frames_file_suffix, self.odo_files_suffix)
+            self.files_odo[odo_h5_path] = h5py.File(odo_h5_path, "r")
+            self.frames_odo_matching_check(h5_file_frames, self.files_odo[odo_h5_path])
+
+    def scan_h5_files(self):
+        self.index_to_starting_frame_map = []
+        required_span = max(
+            self.num_frames * self.frame_interval,
+            self.num_frames_odo * self.odo_frame_interval,
+        )
+        for file in self.hdf5_files:
+            for key in file.keys():
+                if "meta_data" in key:
+                    continue
+                video_length = len(file[key])
+                max_frame_index = video_length - required_span - 1
+                for i in range(0, max_frame_index + 1):
+                    self.index_to_starting_frame_map.append((file, key, i))
+        self.scan_h5_files_odo()
+
+    def _get_odometry_indices(self, start_frame):
+        return list(range(
+            start_frame,
+            start_frame + self.num_frames_odo * self.odo_frame_interval,
+            self.odo_frame_interval,
+        ))
+
+    def get_odometry(self, filename, key, start_idx):
+        odo_filename = filename.replace(self.frames_file_suffix, self.odo_files_suffix)
+        indices = self._get_odometry_indices(start_idx)
+        odo = [self.files_odo[odo_filename][key][i] for i in indices]
+        return torch.as_tensor(self.odo_transform(odo)).float()
+
+    def __getitem__(self, idx):
+        images, (filename, key, start_frame) = self.get_images_and_indices(idx)
+        images = self.apply_transforms(images)
+        if self.odo_files_suffix is not None:
+            odo = self.get_odometry(filename, key, start_frame)
+        else:
+            odo = torch.full((self.num_frames_odo, 3), float("nan"))
+        return {
+            "images": images,
+            "steering": odo,
+            "frame_rate": self.frame_rate,
+            "steering_format": self.steering_format,
+        }
+
+
+class MultiMP4DatasetMultiFrameIdxMappingNVIDIAPhysAIWithL2Context(
+    L2ContextMixin,
+    MultiMP4DatasetMultiFrameIdxMappingNVIDIAPhysAI,
+):
+    def __init__(self, *, num_l2_context, l2_frame_rate=1.0, l1_context_frames=1, **kwargs):
+        super().__init__(**kwargs)
+        self._init_l2_context(
+            num_l2_context=num_l2_context,
+            l2_frame_rate=l2_frame_rate,
+            l1_context_frames=l1_context_frames,
+            require_l2_context=True,
+        )
+        self.index_to_starting_frame_map = self.filter_index_map_with_l2_headroom(
+            self.index_to_starting_frame_map,
+        )
+
+    def _decode_indices(self, path, indices):
+        if self.backend == "decord":
+            from decord import VideoReader, cpu as decord_cpu
+
+            file = VideoReader(path, ctx=decord_cpu(0))
+            frames = file.get_batch(indices).asnumpy()
+        elif self.backend == "torchcodec":
+            from torchcodec.decoders import VideoDecoder
+
+            file = VideoDecoder(path)
+            frames = torch.stack([file[i] for i in indices])
+        else:
+            raise RuntimeError(f"Unknown backend {self.backend}")
+        return frames
+
+    def __getitem__(self, idx):
+        if idx >= len(self.index_to_starting_frame_map):
+            raise IndexError(f"Index {idx} out of range for dataset of length {len(self.index_to_starting_frame_map)}")
+
+        path, start_frame = self.index_to_starting_frame_map[idx]
+        l1_indices, l2_indices = self.get_l1_and_l2_indices(start_frame, self.num_frames)
+        all_frames = self._decode_indices(path, l1_indices + l2_indices)
+        all_transformed = self.augmenter(
+            all_frames,
+            context=self.build_context((None, path, start_frame)),
+        )
+
+        video_id = self._get_video_id(path)
+        return {
+            "images": all_transformed[:len(l1_indices)],
+            "l2_context": all_transformed[len(l1_indices):],
+            "frame_rate": self.frame_rate,
+            "steering": self._load_steering_for_video(video_id, start_frame),
+            "steering_format": self.steering_format,
+        }
+
+
+class VistaStyleNuScenesSteeringMixin(OdometryHorizonMixin):
+    def _init_vista_style_steering(
+        self,
+        *,
+        dbs_root,
+        annotation_key,
+        num_frames,
+        stored_data_frame_rate,
+        frame_rate,
+        num_frames_odo=None,
+        odometry_frame_rate=None,
+        odometry_horizon="auto",
+    ):
+        self.dbs_root = dbs_root
+        self.annotation_key = annotation_key
+        self.steering_format = annotation_key
+        self.odometry_frame_rate = frame_rate if odometry_frame_rate is None else odometry_frame_rate
+        self.odometry_horizon = odometry_horizon
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=frame_rate,
+            rate_name="frame_rate",
+        )
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=self.odometry_frame_rate,
+            rate_name="odometry_frame_rate",
+        )
+        self.odo_frame_interval = int(round(stored_data_frame_rate / self.odometry_frame_rate))
+        inferred_num_frames_odo = self._infer_num_frames_odo_for_matched_horizon(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            odometry_frame_rate=self.odometry_frame_rate,
+        )
+        self.num_frames_odo, self.odometry_horizon = self._resolve_odometry_horizon(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_frames_odo=num_frames_odo,
+            odometry_frame_rate=self.odometry_frame_rate,
+            odometry_horizon=odometry_horizon,
+            inferred_num_frames_odo=inferred_num_frames_odo,
+        )
+        assert self.annotation_key in ["speed_yawrate", "trajectory", "trajectory_with_heading"], (
+            f"Unsupported annotation key {self.annotation_key}"
+        )
+
+    @staticmethod
+    def _infer_num_frames_odo_for_matched_horizon(num_frames, frame_rate, odometry_frame_rate):
+        if num_frames is None:
+            raise ValueError("VistaStyleNuScenesLoaderSteering requires num_frames to resolve odometry horizon")
+        return OdometryHorizonMixin._infer_num_frames_odo_for_matched_horizon(
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            odometry_frame_rate=odometry_frame_rate,
+        )
+
+    def _get_vista_style_steering(self, sample):
+        pose_table = extract_pose_table(os.path.join(self.dbs_root, sample["db_name"]))
+        pose_tokens = sample["pose_tokens"]
+        poses = [get_pose(pose_table, pose_token) for pose_token in pose_tokens]
+        speeds = np.array([pose["vx"].values[0] for pose in poses])
+        yaw_rates = np.array([pose["angular_rate_z"].values[0] for pose in poses])
+
+        if self.annotation_key == "speed_yawrate":
+            speed_yawrate = np.stack([speeds, yaw_rates], axis=-1)
+            steering = speed_yawrate[::self.odo_frame_interval]
+            assert len(steering) >= self.num_frames_odo, (
+                f"Speed/YawRate length {len(steering)} is less than the required {self.num_frames_odo}"
+            )
+            steering = steering[:self.num_frames_odo]
+        elif self.annotation_key == "trajectory_with_heading":
+            traj, headings = get_trajectory_from_speeds_and_yaw_rates(
+                speeds,
+                yaw_rates,
+                dt=1 / self.stored_data_frame_rate,
+            )
+            traj, headings = traj[::self.odo_frame_interval], headings[::self.odo_frame_interval]
+            assert len(traj) >= self.num_frames_odo, (
+                f"Trajectory length {len(traj)} is less than the required {self.num_frames_odo}"
+            )
+            steering = np.concatenate([traj[:self.num_frames_odo], headings[:self.num_frames_odo, None]], axis=-1)
+        elif self.annotation_key == "trajectory":
+            traj, _ = get_trajectory_from_speeds_and_yaw_rates(
+                speeds,
+                yaw_rates,
+                dt=1 / self.stored_data_frame_rate,
+            )
+            traj = traj[::self.odo_frame_interval]
+            assert len(traj) >= self.num_frames_odo, (
+                f"Trajectory length {len(traj)} is less than the required {self.num_frames_odo}"
+            )
+            steering = traj[:self.num_frames_odo]
+        else:
+            raise ValueError(f"Unsupported annotation key {self.annotation_key}")
+
+        return torch.from_numpy(steering).float()
+
+
+class VistaStyleNuScenesLoaderSteering(VistaStyleNuScenesSteeringMixin, VistaStyleNuScenesLoader):
+    def __init__(
+        self,
+        *,
+        size,
+        json_path,
+        images_root,
+        dbs_root,
+        annotation_key,
+        num_frames=None,
+        num_frames_odo=None,
+        stored_data_frame_rate=10,
+        frame_rate=5,
+        odometry_frame_rate=None,
+        odometry_horizon="auto",
+        aug="resize_center",
+        sample_indices=None,
+    ):
+        self.frame_rate = frame_rate
+        self.stored_data_frame_rate = stored_data_frame_rate
+        self.odometry_frame_rate = frame_rate if odometry_frame_rate is None else odometry_frame_rate
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=frame_rate,
+            rate_name="frame_rate",
+        )
+        self._validate_frame_rate_ratio(
+            stored_rate=stored_data_frame_rate,
+            sampled_rate=self.odometry_frame_rate,
+            rate_name="odometry_frame_rate",
+        )
+        frame_rate_multiplier = frame_rate / stored_data_frame_rate
+        super().__init__(
+            size=size,
+            json_path=json_path,
+            images_root=images_root,
+            num_frames=num_frames,
+            frame_rate_multiplier=frame_rate_multiplier,
+            aug=aug,
+            sample_indices=sample_indices,
+        )
+        self._init_vista_style_steering(
+            dbs_root=dbs_root,
+            annotation_key=annotation_key,
+            num_frames=self.num_frames,
+            stored_data_frame_rate=stored_data_frame_rate,
+            frame_rate=frame_rate,
+            num_frames_odo=num_frames_odo,
+            odometry_frame_rate=odometry_frame_rate,
+            odometry_horizon=odometry_horizon,
+        )
+
+    def __getitem__(self, index):
+        images = super().__getitem__(index)
+        sample = self.data[index]
+        return {
+            "images": images,
+            "steering": self._get_vista_style_steering(sample),
+            "frame_rate": torch.tensor(self.frame_rate).float(),
+            "steering_format": self.steering_format,
+        }
+
+
+class VistaStyleNuScenesLoaderSteeringWithL2Context(
+    VistaStyleNuScenesSteeringMixin,
+    VistaStyleNuScenesLoaderWithL2Context,
+):
+    def __init__(
+        self,
+        *,
+        size,
+        json_path,
+        images_root,
+        dbs_root,
+        annotation_key,
+        num_frames=None,
+        num_frames_odo=None,
+        stored_data_frame_rate=10,
+        frame_rate=5,
+        odometry_frame_rate=None,
+        odometry_horizon="auto",
+        num_l2_context,
+        l2_frame_rate=1.0,
+        l1_context_frames=1,
+        aug="resize_center",
+        sample_indices=None,
+    ):
+        super().__init__(
+            size=size,
+            json_path=json_path,
+            images_root=images_root,
+            num_frames=num_frames,
+            stored_data_frame_rate=stored_data_frame_rate,
+            frame_rate=frame_rate,
+            num_l2_context=num_l2_context,
+            l2_frame_rate=l2_frame_rate,
+            l1_context_frames=l1_context_frames,
+            aug=aug,
+            sample_indices=sample_indices,
+        )
+        self._init_vista_style_steering(
+            dbs_root=dbs_root,
+            annotation_key=annotation_key,
+            num_frames=self.num_frames,
+            stored_data_frame_rate=stored_data_frame_rate,
+            frame_rate=frame_rate,
+            num_frames_odo=num_frames_odo,
+            odometry_frame_rate=odometry_frame_rate,
+            odometry_horizon=odometry_horizon,
+        )
+
+    def __getitem__(self, index):
+        batch = super().__getitem__(index)
+        sample = self.data[index]
+        batch["steering"] = self._get_vista_style_steering(sample)
+        batch["steering_format"] = self.steering_format
+        return batch
+
+
+class MultiMP4DatasetMultiFrameIdxMappingNoSteering(PlaceholderSteeringMixin, MultiMP4DatasetMultiFrameIdxMapping):
+    """MP4 dataset without real steering data.
+
+    Returns NaN steering placeholders so that rollout_steering_v2 can work with
+    --steering_file or --no_steering.  Using it without either option raises an
+    explicit error at rollout time.
+
+    Args:
+        num_frames_odo: Number of steering timesteps to expose per sample.  Set
+            this to at least the rollout odometry horizon, or rely on
+            reconfigure_params_for_required_odometry_horizon to expand it.
+        steering_dim: Feature dimension of the placeholder (default 2 for
+            speed/yaw-rate).
+        steering_format: The steering_format tag returned in the batch
+            (default "speed_yawrate").
+        All remaining args are forwarded to MultiMP4DatasetMultiFrameIdxMapping.
+    """
+
+    def __init__(
+        self,
+        size,
+        mp4_paths_file,
+        num_frames,
+        num_frames_odo,
+        steering_dim=2,
+        steering_format="speed_yawrate",
+        stored_data_frame_rate=5,
+        frame_rate=5,
+        aug="resize_center",
+        backend=None,
+        subsample_interval=None,
+        intrinsics_h5_path=None,
+        spatial_transform_config=None,
+    ):
+        self._init_placeholder_steering(num_frames_odo, steering_dim, steering_format)
+        MultiMP4DatasetMultiFrameIdxMapping.__init__(
+            self,
+            size=size,
+            mp4_paths_file=mp4_paths_file,
+            num_frames=num_frames,
+            stored_data_frame_rate=stored_data_frame_rate,
+            frame_rate=frame_rate,
+            aug=aug,
+            backend=backend,
+            subsample_interval=subsample_interval,
+            intrinsics_h5_path=intrinsics_h5_path,
+            spatial_transform_config=spatial_transform_config,
+        )
+
+    def __getitem__(self, idx):
+        item = MultiMP4DatasetMultiFrameIdxMapping.__getitem__(self, idx)
+        return self._add_placeholder_steering(item)
+
+
+class MultiHDF5DatasetMultiFrameIdxMappingNoSteering(PlaceholderSteeringMixin, MultiHDF5DatasetMultiFrameIdxMapping):
+    """HDF5 dataset without real steering data.
+
+    Returns NaN steering placeholders so that rollout_steering_v2 can work with
+    --steering_file or --no_steering.  Using it without either option raises an
+    explicit error at rollout time.
+
+    Args:
+        num_frames_odo: Number of steering timesteps to expose per sample.
+        steering_dim: Feature dimension of the placeholder (default 2).
+        steering_format: steering_format tag returned in the batch (default "speed_yawrate").
+        All remaining args are forwarded to MultiHDF5DatasetMultiFrameIdxMapping.
+    """
+
+    def __init__(
+        self,
+        size,
+        hdf5_paths_file,
+        num_frames,
+        num_frames_odo,
+        steering_dim=2,
+        steering_format="speed_yawrate",
+        stored_data_frame_rate=5,
+        frame_rate=5,
+        aug="resize_center",
+        scale_min=0.15,
+        scale_max=0.5,
+    ):
+        self._init_placeholder_steering(num_frames_odo, steering_dim, steering_format)
+        MultiHDF5DatasetMultiFrameIdxMapping.__init__(
+            self,
+            size=size,
+            hdf5_paths_file=hdf5_paths_file,
+            num_frames=num_frames,
+            stored_data_frame_rate=stored_data_frame_rate,
+            frame_rate=frame_rate,
+            aug=aug,
+            scale_min=scale_min,
+            scale_max=scale_max,
+        )
+
+    def __getitem__(self, idx):
+        item = MultiHDF5DatasetMultiFrameIdxMapping.__getitem__(self, idx)
+        return self._add_placeholder_steering(item)
+
+
+class VistaStyleNuScenesLoaderWithL2ContextNoSteering(
+    PlaceholderSteeringMixin,
+    VistaStyleNuScenesLoaderWithL2Context,
+):
+    """VistaStyle NuScenes + L2-context dataset without real steering data.
+
+    Returns NaN steering placeholders alongside the l2_context frames so that
+    rollout_steering_v2 can work with --steering_file or --no_steering.
+    Using it without either option raises an explicit error at rollout time.
+
+    Args:
+        num_frames_odo: Number of steering timesteps to expose per sample.
+        steering_dim: Feature dimension of the placeholder (default 2 for
+            speed/yaw-rate).
+        steering_format: steering_format tag returned in the batch
+            (default "speed_yawrate").
+        num_l2_context, l2_frame_rate, l1_context_frames: forwarded to
+            VistaStyleNuScenesLoaderWithL2Context / L2ContextMixin.
+        All remaining args forwarded to VistaStyleNuScenesLoaderWithL2Context.
+    """
+
+    def __init__(
+        self,
+        num_frames_odo,
+        steering_dim=2,
+        steering_format="speed_yawrate",
+        **kwargs,
+    ):
+        self._init_placeholder_steering(num_frames_odo, steering_dim, steering_format)
+        VistaStyleNuScenesLoaderWithL2Context.__init__(self, **kwargs)
+
+    def __getitem__(self, idx):
+        item = VistaStyleNuScenesLoaderWithL2Context.__getitem__(self, idx)
+        return self._add_placeholder_steering(item)
+
+
+class MultiMP4DatasetMultiFrameIdxMappingWithL2ContextNoSteering(
+    PlaceholderSteeringMixin,
+    MultiMP4DatasetMultiFrameIdxMappingWithL2Context,
+):
+    """MP4 + L2-context dataset without real steering data.
+
+    Returns NaN steering placeholders alongside the l2_context frames so that
+    rollout_steering_v2 can work with --steering_file or --no_steering.
+    Using it without either option raises an explicit error at rollout time.
+
+    Args:
+        num_frames_odo: Number of steering timesteps to expose per sample.
+        steering_dim: Feature dimension of the placeholder (default 2).
+        steering_format: steering_format tag returned in the batch (default "speed_yawrate").
+        num_l2_context, l2_frame_rate, l1_context_frames: forwarded to L2ContextMixin.
+        All remaining args forwarded to MultiMP4DatasetMultiFrameIdxMapping.
+    """
+
+    def __init__(
+        self,
+        num_frames_odo,
+        steering_dim=2,
+        steering_format="speed_yawrate",
+        **kwargs,
+    ):
+        self._init_placeholder_steering(num_frames_odo, steering_dim, steering_format)
+        MultiMP4DatasetMultiFrameIdxMappingWithL2Context.__init__(self, **kwargs)
+
+    def __getitem__(self, idx):
+        item = MultiMP4DatasetMultiFrameIdxMappingWithL2Context.__getitem__(self, idx)
+        return self._add_placeholder_steering(item)

orbis2/data/utils.py

@@ -0,0 +1,281 @@
+import collections
+import os
+import tarfile
+import urllib
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from data.helper_types import Annotation
+from torch.utils.data._utils.collate import np_str_obj_array_pattern, default_collate_err_msg_format
+from tqdm import tqdm
+
+
+def unpack(path):
+    if path.endswith("tar.gz"):
+        with tarfile.open(path, "r:gz") as tar:
+            tar.extractall(path=os.path.split(path)[0])
+    elif path.endswith("tar"):
+        with tarfile.open(path, "r:") as tar:
+            tar.extractall(path=os.path.split(path)[0])
+    elif path.endswith("zip"):
+        with zipfile.ZipFile(path, "r") as f:
+            f.extractall(path=os.path.split(path)[0])
+    else:
+        raise NotImplementedError(
+            "Unknown file extension: {}".format(os.path.splitext(path)[1])
+        )
+
+
+def reporthook(bar):
+    """tqdm progress bar for downloads."""
+
+    def hook(b=1, bsize=1, tsize=None):
+        if tsize is not None:
+            bar.total = tsize
+        bar.update(b * bsize - bar.n)
+
+    return hook
+
+
+def get_root(name):
+    base = "data/"
+    root = os.path.join(base, name)
+    os.makedirs(root, exist_ok=True)
+    return root
+
+
+def is_prepared(root):
+    return Path(root).joinpath(".ready").exists()
+
+
+def mark_prepared(root):
+    Path(root).joinpath(".ready").touch()
+
+
+def prompt_download(file_, source, target_dir, content_dir=None):
+    targetpath = os.path.join(target_dir, file_)
+    while not os.path.exists(targetpath):
+        if content_dir is not None and os.path.exists(
+            os.path.join(target_dir, content_dir)
+        ):
+            break
+        print(
+            "Please download '{}' from '{}' to '{}'.".format(file_, source, targetpath)
+        )
+        if content_dir is not None:
+            print(
+                "Or place its content into '{}'.".format(
+                    os.path.join(target_dir, content_dir)
+                )
+            )
+        input("Press Enter when done...")
+    return targetpath
+
+
+def download_url(file_, url, target_dir):
+    targetpath = os.path.join(target_dir, file_)
+    os.makedirs(target_dir, exist_ok=True)
+    with tqdm(
+        unit="B", unit_scale=True, unit_divisor=1024, miniters=1, desc=file_
+    ) as bar:
+        urllib.request.urlretrieve(url, targetpath, reporthook=reporthook(bar))
+    return targetpath
+
+
+def download_urls(urls, target_dir):
+    paths = dict()
+    for fname, url in urls.items():
+        outpath = download_url(fname, url, target_dir)
+        paths[fname] = outpath
+    return paths
+
+
+def quadratic_crop(x, bbox, alpha=1.0):
+    """bbox is xmin, ymin, xmax, ymax"""
+    im_h, im_w = x.shape[:2]
+    bbox = np.array(bbox, dtype=np.float32)
+    bbox = np.clip(bbox, 0, max(im_h, im_w))
+    center = 0.5 * (bbox[0] + bbox[2]), 0.5 * (bbox[1] + bbox[3])
+    w = bbox[2] - bbox[0]
+    h = bbox[3] - bbox[1]
+    l = int(alpha * max(w, h))
+    l = max(l, 2)
+
+    required_padding = -1 * min(
+        center[0] - l, center[1] - l, im_w - (center[0] + l), im_h - (center[1] + l)
+    )
+    required_padding = int(np.ceil(required_padding))
+    if required_padding > 0:
+        padding = [
+            [required_padding, required_padding],
+            [required_padding, required_padding],
+        ]
+        padding += [[0, 0]] * (len(x.shape) - 2)
+        x = np.pad(x, padding, "reflect")
+        center = center[0] + required_padding, center[1] + required_padding
+    xmin = int(center[0] - l / 2)
+    ymin = int(center[1] - l / 2)
+    return np.array(x[ymin : ymin + l, xmin : xmin + l, ...])
+
+
+def custom_collate(batch):
+    r"""source: pytorch 1.9.0, only one modification to original code """
+
+    elem = batch[0]
+    elem_type = type(elem)
+    if isinstance(elem, torch.Tensor):
+        out = None
+        if torch.utils.data.get_worker_info() is not None:
+            # If we're in a background process, concatenate directly into a
+            # shared memory tensor to avoid an extra copy
+            numel = sum([x.numel() for x in batch])
+            storage = elem.storage()._new_shared(numel)
+            out = elem.new(storage)
+        return torch.stack(batch, 0, out=out)
+    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
+            and elem_type.__name__ != 'string_':
+        if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap':
+            # array of string classes and object
+            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
+                raise TypeError(default_collate_err_msg_format.format(elem.dtype))
+
+            return custom_collate([torch.as_tensor(b) for b in batch])
+        elif elem.shape == ():  # scalars
+            return torch.as_tensor(batch)
+    elif isinstance(elem, float):
+        return torch.tensor(batch, dtype=torch.float64)
+    elif isinstance(elem, int):
+        return torch.tensor(batch)
+    elif isinstance(elem, str):
+        return batch
+    elif isinstance(elem, collections.abc.Mapping):
+        return {key: custom_collate([d[key] for d in batch]) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
+        return elem_type(*(custom_collate(samples) for samples in zip(*batch)))
+    if isinstance(elem, collections.abc.Sequence) and isinstance(elem[0], Annotation):  # added
+        return batch  # added
+    elif isinstance(elem, collections.abc.Sequence):
+        # check to make sure that the elements in batch have consistent size
+        it = iter(batch)
+        elem_size = len(next(it))
+        if not all(len(elem) == elem_size for elem in it):
+            raise RuntimeError('each element in list of batch should be of equal size')
+        transposed = zip(*batch)
+        return [custom_collate(samples) for samples in transposed]
+
+    raise TypeError(default_collate_err_msg_format.format(elem_type))
+
+
+def get_trajectory_from_speeds_and_yaw_rates(speeds, yaw_rates, dt):
+    heading_deltas = yaw_rates * dt
+    headings = np.cumsum(heading_deltas)
+    headings_for_translation = headings - heading_deltas
+    dx = speeds * np.cos(headings_for_translation) * dt
+    dy = speeds * np.sin(headings_for_translation) * dt
+    
+    x = np.cumsum(dx)
+    y = np.cumsum(dy)
+    traj = np.stack([x, y], axis=1)
+    
+    # Transform to local coordinates (first position is origin, first heading is along x-axis)
+    traj -= traj[0]  # translate to origin
+    initial_heading = headings_for_translation[0]
+    rotation_matrix = np.array([[np.cos(-initial_heading), -np.sin(-initial_heading)],
+                                    [np.sin(-initial_heading),  np.cos(-initial_heading)]])
+    local_traj = traj @ rotation_matrix.T  # rotate to align with initial heading
+    
+    return local_traj.astype(np.float32), headings.astype(np.float32)
+
+def get_trajectory_from_speeds_and_yaw_rates_batch(speeds, yaw_rates, dt):
+    """
+    Args:
+        speeds: Tensor of shape (B, N)
+        yaw_rates: Tensor of shape (B, N)
+        dt: Time step (scalar)
+    Returns:
+        local_traj: Tensor of shape (B, N, 2)
+        headings: Tensor of shape (B, N)
+    """
+    assert speeds.shape == yaw_rates.shape, f"Speeds shape {speeds.shape} and yaw rates shape {yaw_rates.shape} do not match"
+    B, N = speeds.shape
+    
+    # if dt is a scalar, ok, if dt is a tensor, make sure it has shape (B) and expand to (B, 1)
+    if isinstance(dt, torch.Tensor):
+        assert dt.shape == (B,), f"dt shape {dt.shape} does not match batch size {B}"
+        dt = dt.view(B, 1)  # Shape: (B, 1)
+    
+    heading_deltas = yaw_rates * dt  # Shape: (B, N)
+    headings = torch.cumsum(heading_deltas, dim=1)  # Shape: (B, N)
+    headings_for_translation = headings - heading_deltas
+
+    # Calculate dx and dy for each batch
+    dx = speeds * torch.cos(headings_for_translation) * dt  # Shape: (B, N)
+    dy = speeds * torch.sin(headings_for_translation) * dt  # Shape: (B, N)
+
+    # Calculate x and y for each batch
+    x = torch.cumsum(dx, dim=1)  # Shape: (B, N)
+    y = torch.cumsum(dy, dim=1)  # Shape: (B, N)
+
+    # Stack x and y to form the trajectory for each batch
+    traj = torch.stack([x, y], dim=2)  # Shape: (B, N, 2)
+
+    # Transform to local coordinates for each batch
+    traj = traj- traj[:, 0:1, :]  # Translate to origin for each batch
+    initial_heading = headings_for_translation[:, 0]  # Shape: (B,)
+
+    # Create rotation matrices for each batch
+    cos_theta = torch.cos(-initial_heading)  # Shape: (B,)
+    sin_theta = torch.sin(-initial_heading)  # Shape: (B,)
+
+    # Rotation matrix for each batch
+    rotation_matrix = torch.stack([
+        torch.stack([cos_theta, -sin_theta], dim=1),
+        torch.stack([sin_theta,  cos_theta], dim=1)
+    ], dim=1)  # Shape: (B, 2, 2)
+
+    # Rotate to align with initial heading for each batch
+    local_traj = torch.einsum('bni,bij->bnj', traj, rotation_matrix)  # Shape: (B, N, 2)
+
+    return torch.cat([local_traj, headings.unsqueeze(-1)], dim=-1).float()  # Return (B, N, 3)
+
+
+class RunningNorm(nn.Module):
+    def __init__(self, num_features, momentum=0.1, eps=1e-5):
+        super().__init__()
+        self.momentum = momentum
+        self.eps = eps
+
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_std', torch.ones(num_features))
+
+    def _reduce_dims(self, x):
+        # Dimensions to reduce: batch and spatial (leave channel/features alone)
+        return [0] + list(range(2, x.dim()))
+
+    def update_stats(self, x):
+        dims = self._reduce_dims(x)
+        batch_mean = x.mean(dim=dims)
+        batch_std = x.std(dim=dims, unbiased=False)
+
+        # Update running stats
+        with torch.no_grad():
+            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean
+            self.running_std = (1 - self.momentum) * self.running_std + self.momentum * batch_std
+
+    def normalize(self, x):
+        mean = self.running_mean.view(1, -1, *[1] * (x.dim() - 2))
+        std = self.running_std.view(1, -1, *[1] * (x.dim() - 2))
+        return (x - mean) / (std + self.eps)
+
+    def denormalize(self, x):
+        mean = self.running_mean.view(1, -1, *[1] * (x.dim() - 2))
+        std = self.running_std.view(1, -1, *[1] * (x.dim() - 2))
+        return x * (std + self.eps) + mean
+
+    def forward(self, x):
+        if self.training:
+            self.update_stats(x)
+        return self.normalize(x)

orbis2/data/video_loaders.py

@@ -0,0 +1,638 @@
+import json
+import os
+import random
+import signal
+import warnings
+
+from tqdm import tqdm
+import h5py
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, get_worker_info
+from torchvision import transforms
+import torch.distributed as dist
+
+from .l2_context import L2ContextMixin
+from util import instantiate_from_config
+
+try:
+    from torchcodec.decoders import VideoDecoder
+    TORCHCODEC_AVAILABLE = True
+except ImportError:
+    TORCHCODEC_AVAILABLE = False
+
+try:
+    from decord import VideoReader, cpu as decord_cpu
+    DECORD_AVAILABLE = True
+except ImportError:
+    DECORD_AVAILABLE = False
+
+
+
+def _normalize_size(size):
+    return (size, size) if isinstance(size, int) else tuple(size)
+
+
+class FrameAdapter:
+    def __call__(self, frames):
+        raise NotImplementedError
+
+
+class PILFrameAdapter(FrameAdapter):
+    def __call__(self, frames):
+        arrays = np.stack([np.asarray(frame) for frame in frames], axis=0)
+        return torch.from_numpy(arrays).permute(0, 3, 1, 2).float() / 255.0
+
+
+class DecordFrameAdapter(FrameAdapter):
+    def __call__(self, frames):
+        return torch.from_numpy(frames).permute(0, 3, 1, 2).float() / 255.0
+
+
+class TensorFrameAdapter(FrameAdapter):
+    def __call__(self, frames):
+        if not torch.is_tensor(frames):
+            raise TypeError(f"Expected torch.Tensor frames, got {type(frames).__name__}")
+        frames = frames.float()
+        if frames.max().item() > 1.0:
+            frames = frames / 255.0
+        return frames
+
+
+class SpatialTransformPolicy:
+    def __init__(self, size):
+        self.size = _normalize_size(size)
+
+    def sample_params(self, frames, context=None):
+        return None
+
+    def apply(self, frames, params, context=None):
+        raise NotImplementedError
+
+
+class ResizeCenterPolicy(SpatialTransformPolicy):
+    def apply(self, frames, params, context=None):
+        frames = transforms.functional.resize(frames, [min(self.size)], antialias=True)
+        return transforms.functional.center_crop(frames, list(self.size))
+
+
+class RandomShiftPolicy(SpatialTransformPolicy):
+    def __init__(self, size, max_shift_horizontal=60, max_shift_vertical=60):
+        super().__init__(size)
+        self.max_shift_horizontal = max_shift_horizontal
+        self.max_shift_vertical = max_shift_vertical
+
+    def sample_params(self, frames, context=None):
+        frames = transforms.functional.resize(frames, [min(self.size)], antialias=True)
+        crop_height, crop_width = self.size
+        width, height = frames.shape[-1], frames.shape[-2]
+        center_left = (width - crop_width) // 2
+        center_top = (height - crop_height) // 2
+
+        shift_horizontal = random.randint(-self.max_shift_horizontal, self.max_shift_horizontal)
+        shift_vertical = random.randint(-self.max_shift_vertical, self.max_shift_vertical)
+
+        left = max(0, min(center_left + shift_horizontal, width - crop_width))
+        top = max(0, min(center_top + shift_vertical, height - crop_height))
+        return left, top
+
+    def apply(self, frames, params, context=None):
+        frames = transforms.functional.resize(frames, [min(self.size)], antialias=True)
+        left, top = self.sample_params(frames) if params is None else params
+        crop_height, crop_width = self.size
+        return frames[..., top:top + crop_height, left:left + crop_width]
+
+
+class RandomResizedCenterPolicy(SpatialTransformPolicy):
+    def __init__(self, size, scale=(0.5, 1.0)):
+        super().__init__(size)
+        self.scale = scale
+
+    def sample_params(self, frames, context=None):
+        height, width = frames.shape[-2], frames.shape[-1]
+        area = height * width
+        aspect_ratio = width / height
+        target_area = random.uniform(*self.scale) * area
+
+        new_width = int(round((target_area * aspect_ratio) ** 0.5))
+        new_height = int(round((target_area / aspect_ratio) ** 0.5))
+        return new_height, new_width
+
+    def apply(self, frames, params, context=None):
+        new_height, new_width = self.sample_params(frames) if params is None else params
+        frames = transforms.functional.resize(frames, [new_height, new_width], antialias=True)
+        return transforms.functional.center_crop(frames, list(self.size))
+
+
+class CameraCalibrationPolicy(SpatialTransformPolicy):
+    def __init__(
+        self,
+        size,
+        transform=None,
+        transform_config=None,
+        calibration_key="intrinsics",
+        fallback_policy=None,
+    ):
+        super().__init__(size)
+        self.transform = transform if transform is not None else instantiate_from_config(transform_config)
+        self.calibration_key = calibration_key
+        self.fallback_policy = fallback_policy or ResizeCenterPolicy(size)
+
+    def apply(self, frames, params, context=None):
+        if context is None or self.calibration_key not in context:
+            # print(f"[CamCalib] FALLBACK — key={self.calibration_key!r} context_keys={list(context.keys()) if context else None}", flush=True)
+            return self.fallback_policy.apply(frames, params, context=context)
+        return self.transform(frames, context[self.calibration_key])
+
+
+class ClipAugmenter:
+    def __init__(self, adapter, spatial_policy):
+        self.adapter = adapter
+        self.spatial_policy = spatial_policy
+
+    def __call__(self, frames, context=None):
+        frames = self.adapter(frames)
+        params = self.spatial_policy.sample_params(frames, context=context)
+        frames = self.spatial_policy.apply(frames, params, context=context)
+        return frames * 2 - 1
+
+
+def _build_frame_adapter(input_format):
+    if input_format == "pil":
+        return PILFrameAdapter()
+    if input_format == "decord":
+        return DecordFrameAdapter()
+    if input_format in {"tensor", "torchcodec"}:
+        return TensorFrameAdapter()
+    raise ValueError(f"Unknown input format: {input_format}")
+
+
+def _build_spatial_policy(aug, size, scale_min=0.15, scale_max=0.5):
+    if aug == "resize_center":
+        return ResizeCenterPolicy(size)
+    if aug == "random_resize_center":
+        return RandomResizedCenterPolicy(size, scale=(scale_min, scale_max))
+    if aug == "random_shift":
+        return RandomShiftPolicy(size, max_shift_horizontal=60, max_shift_vertical=30)
+    raise ValueError(f"Unknown augmentation type: {aug}")
+
+
+def build_clip_augmenter(*, aug, size, input_format, scale_min=0.15, scale_max=0.5, spatial_transform_config=None):
+    spatial_policy = (
+        instantiate_from_config(spatial_transform_config)
+        if spatial_transform_config is not None
+        else _build_spatial_policy(aug, size, scale_min=scale_min, scale_max=scale_max)
+    )
+    return ClipAugmenter(
+        adapter=_build_frame_adapter(input_format),
+        spatial_policy=spatial_policy,
+    )
+
+
+def concat_frame_batches(*frame_batches):
+    first_batch = frame_batches[0]
+    if isinstance(first_batch, np.ndarray):
+        return np.concatenate(frame_batches, axis=0)
+    if torch.is_tensor(first_batch):
+        return torch.cat(frame_batches, dim=0)
+    raise TypeError(f"Unsupported frame batch type: {type(first_batch).__name__}")
+
+
+def _read_h5_value(node):
+    if isinstance(node, h5py.Dataset):
+        value = node[()]
+        if isinstance(value, bytes):
+            value = value.decode()
+        if isinstance(value, str) and value[:1] in "{[":
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError:
+                pass
+        if isinstance(value, np.ndarray):
+            return value.tolist()
+        return value.item() if hasattr(value, "item") else value
+    return {key: _read_h5_value(node[key]) for key in node.keys()}
+
+
+class DatasetMultiFrameIdxMapping(Dataset):
+    """
+    Template dataset that maps each index to a specific frame in a specific video.
+    Subclasses must implement get_images_and_indices() and apply_transforms().
+    """
+    def __init__(self):
+        super().__init__()
+        self.index_to_starting_frame_map = []
+
+    def __len__(self):
+        return len(self.index_to_starting_frame_map)
+
+    def get_images_and_indices(self, idx):
+        raise NotImplementedError
+
+    def apply_transforms(self, images, context=None):
+        raise NotImplementedError
+
+    def build_context(self, metadata):
+        return None
+
+    def __getitem__(self, idx):
+        images, metadata = self.get_images_and_indices(idx)
+        images = self.apply_transforms(images, context=self.build_context(metadata))
+        return {"images": images, "frame_rate": self.frame_rate}
+
+
+class MultiHDF5DatasetMultiFrameIdxMapping(DatasetMultiFrameIdxMapping):
+    """
+    This dataset maps each index to a specific frame in a specific video.
+    Useful for validation and selecting subsets of frames.
+    """
+    def __init__(self, size, hdf5_paths_file, num_frames, stored_data_frame_rate=5, frame_rate=5, aug="resize_center", scale_min=0.15, scale_max=0.5):
+        super().__init__()
+        self.frame_interval = int(stored_data_frame_rate / frame_rate)
+        self.frame_rate = frame_rate
+        self.stored_data_frame_rate = stored_data_frame_rate
+
+        self.size = (size, size) if isinstance(size, int) else size
+        self.num_frames = num_frames
+        self.hdf5_paths_file = hdf5_paths_file
+        with open(os.path.expandvars(hdf5_paths_file), "r") as f:
+            self.hdf5_paths = f.read().splitlines()
+
+        self.hdf5_files = [h5py.File(path, "r") for path in self.hdf5_paths]
+
+        self.scan_h5_files()
+
+        self.aug = aug
+        self.augmenter = build_clip_augmenter(
+            aug=self.aug,
+            size=self.size,
+            input_format="pil",
+            scale_min=scale_min,
+            scale_max=scale_max,
+        )
+
+    def scan_h5_files(self):
+        self.index_to_starting_frame_map = []
+        for file in self.hdf5_files:
+            keys = list(file.keys())
+            for key in keys:
+                video_length = len(file[key])
+                max_frame_index = video_length - self.num_frames * self.frame_interval - 1
+                for i in range(0, max_frame_index + 1):
+                    self.index_to_starting_frame_map.append((file, key, i))
+
+    def __str__(self):
+        return f"MultiHDF5DatasetMultiFrameIdxMapping({self.hdf5_paths_file}, num_samples={len(self)}, size={self.size}, num_frames={self.num_frames}, frame_interval={self.frame_interval})"
+
+    def get_images_and_indices(self, idx):
+        if idx >= len(self.index_to_starting_frame_map):
+            raise IndexError(f"Index {idx} out of range for dataset of length {len(self.index_to_starting_frame_map)}")
+        file, key, start_frame = self.index_to_starting_frame_map[idx]
+        images = [Image.fromarray(file[key][start_frame + i * self.frame_interval]) for i in range(self.num_frames)]
+        return images, (file.filename, key, start_frame)
+
+    def apply_transforms(self, images, context=None):
+        return self.augmenter(images, context=context)
+
+    def close(self):
+        for file in self.hdf5_files:
+            file.close()
+
+
+class MultiMP4DatasetMultiFrameIdxMapping(DatasetMultiFrameIdxMapping):
+    def __init__(self, size, mp4_paths_file, num_frames, stored_data_frame_rate=5, frame_rate=5, aug="resize_center", backend=None, subsample_interval=None, intrinsics_h5_path=None, spatial_transform_config=None, validate_frame_rate_sample=True):
+        super().__init__()
+        self.frame_interval = int(stored_data_frame_rate / frame_rate)
+        self.frame_rate = frame_rate
+        self.stored_data_frame_rate = stored_data_frame_rate
+        self.size = (size, size) if isinstance(size, int) else size
+        self.num_frames = num_frames
+        self.mp4_paths_file = mp4_paths_file
+        self.subsample_interval = subsample_interval
+        self.intrinsics_h5_path = os.path.expandvars(intrinsics_h5_path) if intrinsics_h5_path is not None else None
+        self.intrinsics_file = None
+        
+        with open(os.path.expandvars(mp4_paths_file), "r") as f:
+            # the list is either a list of mp4 file paths, or a csv with two columns: file_path,length_in_frames, with a header row.
+            lines = f.read().splitlines()
+            if len(lines) > 1 and len(lines[0].split(",")) == 2:
+                # CSV format
+                self.mp4_paths = [self._normalize_path(p.split(",")[0]) for p in lines[1:]]
+                mp4_lengths = [int(p.split(",")[1]) for p in lines[1:]]
+                # CSV format: create a {'path': length_in_frames} dict
+                self.mp4_lengths_in_frames = dict(zip(self.mp4_paths, mp4_lengths))
+            else:
+                # Plain list of file paths
+                self.mp4_paths = [self._normalize_path(p) for p in lines]
+                self.mp4_lengths_in_frames = None
+
+
+        if backend is not None:
+            self.backend = backend
+            assert self.backend in ["decord", "torchcodec"], f"Unknown backend {self.backend}"
+            assert (self.backend == "decord" and DECORD_AVAILABLE) or (self.backend == "torchcodec" and TORCHCODEC_AVAILABLE), f"Specified backend {self.backend} is not available"
+        else:
+            if DECORD_AVAILABLE:
+                self.backend = "decord"
+            elif TORCHCODEC_AVAILABLE:
+                self.backend = "torchcodec"
+            else:
+                raise ImportError("Either Decord or TorchCodec must be installed to use MultiMP4DatasetMultiFrameIdxMapping")
+
+        self.validate_frame_rate_sample = validate_frame_rate_sample
+        if self.validate_frame_rate_sample:
+            self._validate_frame_rate_sample()
+        self.scan_mp4_files_distributed()
+
+        if self.backend == "decord":
+            self.read_frames = self.read_frames_decord
+        elif self.backend == "torchcodec":
+            self.read_frames = self.read_frames_torchcodec
+        else:
+            raise ImportError("Either Decord or TorchCodec must be installed to use MultiMP4DatasetMultiFrameIdxMapping")
+
+        self.aug = aug
+        self.augmenter = build_clip_augmenter(
+            aug=self.aug,
+            size=self.size,
+            input_format=self.backend,
+            spatial_transform_config=spatial_transform_config,
+        )
+
+        # Decode robustness knobs (override via environment if needed).
+        self.decode_max_retries = max(int(os.environ.get("ORBIS_DECODE_MAX_RETRIES", "4")), 0)
+        self.decode_timeout_s = max(int(os.environ.get("ORBIS_DECODE_TIMEOUT_S", "20")), 0)
+        self.decode_warn_limit = max(int(os.environ.get("ORBIS_DECODE_WARN_LIMIT", "50")), 0)
+        self._decode_warn_count = 0
+
+    def _validate_frame_rate_sample(self, sample_size=100):
+        paths = random.sample(self.mp4_paths, min(sample_size, len(self.mp4_paths)))
+        for path in paths:
+            if self.backend == "decord":
+                fps = VideoReader(path, ctx=decord_cpu(0), num_threads=-1).get_avg_fps()
+            else:
+                fps = VideoDecoder(path).metadata.average_fps
+            assert self.stored_data_frame_rate == fps, (
+                f"Stored data frame rate {self.stored_data_frame_rate} does not match "
+                f"actual frame rate {fps} for file {path}"
+            )
+
+    def get_video_length(self, path):
+        if self.backend == "decord":
+            file = VideoReader(path, ctx=decord_cpu(0), num_threads=-1)
+            assert self.stored_data_frame_rate == file.get_avg_fps(), f"Stored data frame rate {self.stored_data_frame_rate} does not match actual frame rate {file.get_avg_fps()} for file {file}"
+            video_length = len(file)
+        elif self.backend == "torchcodec":
+            file = VideoDecoder(path)
+            assert self.stored_data_frame_rate == file.metadata.average_fps, f"Stored data frame rate {self.stored_data_frame_rate} does not match actual frame rate {file.metadata.average_fps} for file {file}"
+            video_length = file.metadata.num_frames
+        else:
+            raise ImportError("Either Decord or TorchCodec must be installed to use MultiMP4DatasetMultiFrameIdxMapping")
+        return video_length
+
+    def _normalize_path(self, path):
+        return os.path.normpath(os.path.expandvars(os.path.expanduser(path)))
+
+    def _get_video_id(self, path):
+        return os.path.basename(path).split(".")[0]
+
+    def _get_intrinsics_file(self):
+        if self.intrinsics_h5_path is None:
+            return None
+        if self.intrinsics_file is None:
+            self.intrinsics_file = h5py.File(self.intrinsics_h5_path, "r")
+        return self.intrinsics_file
+
+    def _load_intrinsics(self, path):
+        intrinsics_file = self._get_intrinsics_file()
+        if intrinsics_file is None:
+            return None
+        video_id = self._get_video_id(path)
+        if video_id not in intrinsics_file or "intrinsics" not in intrinsics_file[video_id]:
+            return None
+        return _read_h5_value(intrinsics_file[video_id]["intrinsics"])
+
+    def build_context(self, metadata):
+        _, path, start_frame = metadata
+        context = {"path": path, "video_id": self._get_video_id(path), "start_frame": start_frame}
+        intrinsics = self._load_intrinsics(path)
+        if intrinsics is not None:
+            context["intrinsics"] = intrinsics
+        return context
+
+    def scan_mp4_files(self):                
+        is_rank0 = not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0
+        self.index_to_starting_frame_map = []
+        for path in tqdm(
+            self.mp4_paths,
+            desc=f"Scanning MP4 files in {self.__class__.__name__}",
+            disable=not is_rank0,
+        ):
+            if self.mp4_lengths_in_frames is not None and path in self.mp4_lengths_in_frames:
+                video_length = self.mp4_lengths_in_frames[path]
+            else:
+                video_length = self.get_video_length(path)
+
+            frame_interval = self.frame_interval if self.subsample_interval is None else self.subsample_interval*self.frame_interval
+            
+            max_frame_index = video_length - self.num_frames * self.frame_interval - 1
+            # assert max_frame_index > 0
+            if not max_frame_index > 0: continue
+            for i in range(0, max_frame_index + 1, frame_interval):
+                self.index_to_starting_frame_map.append((path, i))
+
+    def scan_mp4_files_distributed(self):
+        if not (dist.is_available() and dist.is_initialized()):
+            self.scan_mp4_files()
+            return
+
+        # Scan only on rank 0 and broadcast index mapping to all other ranks.
+        if dist.get_rank() == 0:
+            self.scan_mp4_files()
+            index_map = self.index_to_starting_frame_map
+        else:
+            index_map = None
+
+        obj_list = [index_map]
+        dist.broadcast_object_list(obj_list, src=0)
+        self.index_to_starting_frame_map = obj_list[0]
+
+    def __str__(self):
+        return f"MultiMP4DatasetMultiFrameIdxMapping({self.mp4_paths_file}, num_samples={len(self)}, size={self.size}, num_frames={self.num_frames}, frame_interval={self.frame_interval})"
+
+    def apply_transforms(self, images, context=None):
+        return self.augmenter(images, context=context)
+
+    def read_frames_decord(self, path, start_frame):
+        indices = list(range(start_frame, start_frame + self.num_frames * self.frame_interval, self.frame_interval))
+        file = VideoReader(path, ctx=decord_cpu(0))
+        frames = file.get_batch(indices).asnumpy()
+        return frames
+
+    def read_frames_torchcodec(self, path, start_frame):
+        file = VideoDecoder(path)
+        frames = file[start_frame:start_frame + self.num_frames * self.frame_interval:self.frame_interval]
+        return frames
+
+    def _decode_with_timeout(self, path, start_frame):
+        if self.decode_timeout_s <= 0 or not hasattr(signal, "SIGALRM"):
+            return self.read_frames(path, start_frame)
+
+        def _timeout_handler(signum, frame):  # pragma: no cover - signal handler
+            raise TimeoutError(
+                f"Decode timed out after {self.decode_timeout_s}s for {path} at frame {start_frame}"
+            )
+
+        previous_handler = signal.signal(signal.SIGALRM, _timeout_handler)
+        signal.setitimer(signal.ITIMER_REAL, float(self.decode_timeout_s))
+        try:
+            return self.read_frames(path, start_frame)
+        finally:
+            signal.setitimer(signal.ITIMER_REAL, 0.0)
+            signal.signal(signal.SIGALRM, previous_handler)
+
+    def _warn_decode_error(self, msg):
+        if self._decode_warn_count < self.decode_warn_limit:
+            warnings.warn(msg)
+            self._decode_warn_count += 1
+        elif self._decode_warn_count == self.decode_warn_limit:
+            warnings.warn(
+                "Reached ORBIS_DECODE_WARN_LIMIT; suppressing further decode warnings."
+            )
+            self._decode_warn_count += 1
+
+    def get_images_and_indices(self, idx):
+        if idx >= len(self.index_to_starting_frame_map):
+            raise IndexError(f"Index {idx} out of range for dataset of length {len(self.index_to_starting_frame_map)}")
+        map_len = len(self.index_to_starting_frame_map)
+        worker_info = get_worker_info()
+        worker_id = worker_info.id if worker_info is not None else -1
+
+        last_exc = None
+        for attempt in range(self.decode_max_retries + 1):
+            attempt_idx = (idx + attempt) % map_len
+            path, start_frame = self.index_to_starting_frame_map[attempt_idx]
+            try:
+                frames = self._decode_with_timeout(path, start_frame)
+                if frames is None:
+                    raise RuntimeError("Decoder returned None")
+                num_decoded = frames.shape[0] if hasattr(frames, "shape") else len(frames)
+                if num_decoded != self.num_frames:
+                    raise RuntimeError(
+                        f"Decoder returned {num_decoded} frames, expected {self.num_frames}"
+                    )
+                return frames, (None, path, start_frame)
+            except Exception as exc:
+                last_exc = exc
+                self._warn_decode_error(
+                    f"[MultiMP4Dataset] decode failure worker={worker_id} attempt={attempt + 1}/"
+                    f"{self.decode_max_retries + 1} idx={attempt_idx} path={path} "
+                    f"start_frame={start_frame}: {type(exc).__name__}: {exc}"
+                )
+
+        raise RuntimeError(
+            f"Failed to decode sample idx={idx} after {self.decode_max_retries + 1} attempts"
+        ) from last_exc
+
+    def close(self):
+        if self.intrinsics_file is not None:
+            self.intrinsics_file.close()
+
+
+class MultiMP4DatasetMultiFrameIdxMappingWithL2Context(L2ContextMixin, MultiMP4DatasetMultiFrameIdxMapping):
+    """
+    Extends MultiMP4DatasetMultiFrameIdxMapping to also return a 'l2_context' key
+    containing num_l2_context frames at l2_frame_rate (typically 1 Hz) that end on
+    the first L1 frame.
+
+    The l2_context frames are sampled from the same video at the same spatial crop,
+    with l2_context[-1] matching images[0].
+
+    Args:
+        num_l2_context     : number of L2 context frames (default 3)
+        l2_frame_rate      : frame rate for L2 context frames in Hz (default 1.0)
+        All other args forwarded to MultiMP4DatasetMultiFrameIdxMapping.
+
+    Batch output adds:
+        'l2_context' : (num_l2_context, C, H, W) pixel tensor in [-1, 1],
+                       ordered oldest → newest (l2_context[-1] is the same
+                       frame as images[0]).
+    """
+
+    def __init__(self, num_l2_context=3, l2_frame_rate=1.0, l1_context_frames=1, **kwargs):
+        super().__init__(**kwargs)
+        self._init_l2_context(
+            num_l2_context=num_l2_context,
+            l2_frame_rate=l2_frame_rate,
+            l1_context_frames=l1_context_frames,
+        )
+        self.index_to_starting_frame_map = self.filter_index_map_with_l2_headroom(
+            self.index_to_starting_frame_map,
+        )
+
+    def get_images_and_indices(self, idx):
+        """Return (l1_frames, l2_frames, metadata)."""
+        if idx >= len(self.index_to_starting_frame_map):
+            raise IndexError(f"Index {idx} out of range.")
+        path, start_frame = self.index_to_starting_frame_map[idx]
+        l1_indices, l2_indices = self.get_l1_and_l2_indices(start_frame, self.num_frames)
+
+        all_indices = l1_indices + l2_indices
+        all_frames = self._decode_indices(path, all_indices)
+
+        l1_frames = all_frames[:len(l1_indices)]
+        l2_frames = all_frames[len(l1_indices):]
+        return l1_frames, l2_frames, (None, path, start_frame)
+
+    def _decode_indices(self, path, indices):
+        """Decode an arbitrary list of frame indices from an MP4."""
+        if self.backend == "decord":
+            file = VideoReader(path, ctx=decord_cpu(0))
+            frames = file.get_batch(indices).asnumpy()
+        elif self.backend == "torchcodec":
+            file = VideoDecoder(path)
+            frames = torch.stack([file[i] for i in indices])
+        else:
+            raise RuntimeError(f"Unknown backend {self.backend}")
+        return frames
+
+    def __getitem__(self, idx):
+        map_len = len(self.index_to_starting_frame_map)
+        worker_info = get_worker_info()
+        worker_id = worker_info.id if worker_info is not None else -1
+        last_exc = None
+
+        for attempt in range(self.decode_max_retries + 1):
+            attempt_idx = (idx + attempt) % map_len
+            path, start_frame = self.index_to_starting_frame_map[attempt_idx]
+            try:
+                l1_frames, l2_frames, _ = self.get_images_and_indices(attempt_idx)
+                all_frames = concat_frame_batches(l1_frames, l2_frames)
+                all_transformed = self.augmenter(
+                    all_frames,
+                    context=self.build_context((None, path, start_frame)),
+                )
+
+                n_l1 = self.num_frames
+                images_l1 = all_transformed[:n_l1]    # (F, C, H, W)
+                images_l2 = all_transformed[n_l1:]    # (num_l2_context, C, H, W)
+
+                return {
+                    "images": images_l1,
+                    "l2_context": images_l2,
+                    "frame_rate": self.frame_rate,
+                }
+            except Exception as exc:
+                last_exc = exc
+                self._warn_decode_error(
+                    f"[MultiMP4DatasetWithL2] decode failure worker={worker_id} "
+                    f"attempt={attempt + 1}/{self.decode_max_retries + 1} "
+                    f"idx={attempt_idx} path={path} start_frame={start_frame}: "
+                    f"{type(exc).__name__}: {exc}"
+                )
+
+        raise RuntimeError(
+            f"Failed to decode sample idx={idx} after {self.decode_max_retries + 1} attempts"
+        ) from last_exc
+

orbis2/data/vista_style.py

@@ -0,0 +1,183 @@
+from torch.utils.data import Dataset
+import json
+import os, json
+from collections import OrderedDict
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import cv2
+from PIL import Image
+import sqlite3
+import pandas as pd
+from .utils import get_trajectory_from_speeds_and_yaw_rates
+from .l2_context import L2ContextMixin
+from .video_loaders import build_clip_augmenter
+
+
+class VistaStyleNuScenesLoader(Dataset):
+    def __init__(self, *, size, json_path, images_root, num_frames=None, frame_rate_multiplier=1, aug="resize_center", sample_indices=None):
+        super().__init__()
+        self.size = (size, size) if isinstance(size, int) else size
+        self.json_path = json_path
+        self.num_frames = num_frames
+        self.images_root = images_root
+        self.aug = aug
+
+        assert frame_rate_multiplier <= 1, "Frame rate multiplier should be less than or equal to 1"
+        assert 1/frame_rate_multiplier == int(1/frame_rate_multiplier), 'reciprocal of frame_rate_multiplier must be an integer'
+        self.frame_interval = int(1/frame_rate_multiplier)        
+        
+        with open(json_path, 'r') as f:
+            self.data = json.load(f)
+        
+        self.sample_indices = sample_indices
+        if sample_indices is not None:
+            self.data = [self.data[i] for i in sample_indices]
+        self.augmenter = build_clip_augmenter(
+            aug=self.aug,
+            size=self.size,
+            input_format="pil",
+        )
+        
+    def __getitem__(self, index):
+        sample = self.data[index]
+        frame_paths = sample['frames'][::self.frame_interval]
+        if self.num_frames is not None:
+            if len(frame_paths) < self.num_frames:
+                print(f"Warning: Number of frames {len(frame_paths)} is less than the required {self.num_frames}")
+                raise ValueError(f"Number of frames {len(frame_paths)} is less than the required {self.num_frames}")
+            frame_paths = frame_paths[:self.num_frames]
+        images = [cv2.imread(os.path.join(self.images_root, frame_path)) for frame_path in frame_paths]
+        images = [Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) for image in images]
+        return self.augmenter(images)
+        
+    def __len__(self):
+        return len(self.data)
+    
+    
+    
+def extract_pose_table(db_path):
+    """
+    Connects to the nuPlan SQLite database and extracts the entire ego_pose table.
+    Returns a pandas DataFrame containing the pose metadata.
+    """
+    conn = sqlite3.connect(db_path)
+    query = "SELECT * FROM ego_pose;"
+    pose_df = pd.read_sql_query(query, conn)
+    conn.close()
+    # Ensure timestamp is numeric and sort by timestamp
+    pose_df['timestamp'] = pd.to_numeric(pose_df['timestamp'], errors='coerce')
+    pose_df.sort_values("timestamp", inplace=True)
+    return pose_df
+
+def get_pose(pose_df, pose_token):
+    tk = bytes.fromhex(pose_token)
+    pose = pose_df[pose_df['token'] == tk]
+    if pose.empty:
+        raise ValueError(f"Pose with token {pose_token} not found.")
+    return pose
+
+
+def __getattr__(name):
+    if name == "VistaStyleNuScenesLoaderSteering":
+        from data.steering_loaders import VistaStyleNuScenesLoaderSteering
+        return VistaStyleNuScenesLoaderSteering
+    if name == "VistaStyleNuScenesLoaderSteeringWithL2Context":
+        from data.steering_loaders import VistaStyleNuScenesLoaderSteeringWithL2Context
+        return VistaStyleNuScenesLoaderSteeringWithL2Context
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+
+class VistaStyleNuScenesLoaderWithL2Context(L2ContextMixin, Dataset):
+    def __init__(self, *, size, json_path, images_root, num_frames=None,
+                 frame_rate_multiplier=1, stored_data_frame_rate=None, frame_rate=None,
+                 num_l2_context=0, l2_frame_rate=1.0, l1_context_frames=1,
+                 aug=None, sample_indices=None):
+        super().__init__()
+        self.size = (size, size) if isinstance(size, int) else size
+        self.json_path = json_path
+        self.num_frames = num_frames
+        self.images_root = images_root
+        self.aug = aug or "resize_center"
+
+        # Compute frame_interval from stored_data_frame_rate/frame_rate when provided,
+        # otherwise fall back to frame_rate_multiplier.
+        if stored_data_frame_rate is not None and frame_rate is not None:
+            assert stored_data_frame_rate % frame_rate == 0, \
+                'stored_data_frame_rate must be divisible by frame_rate'
+            self.frame_interval = int(stored_data_frame_rate / frame_rate)
+            self.stored_data_frame_rate = stored_data_frame_rate
+            self.frame_rate = float(frame_rate)
+        else:
+            assert frame_rate_multiplier <= 1, "Frame rate multiplier should be less than or equal to 1"
+            assert 1/frame_rate_multiplier == int(1/frame_rate_multiplier), \
+                'reciprocal of frame_rate_multiplier must be an integer'
+            self.frame_interval = int(1/frame_rate_multiplier)
+            # Only set if not already provided by a subclass before calling super().__init__
+            if not hasattr(self, 'stored_data_frame_rate'):
+                self.stored_data_frame_rate = stored_data_frame_rate
+            if not hasattr(self, 'frame_rate'):
+                self.frame_rate = float(frame_rate) if frame_rate is not None else None
+
+        self._init_l2_context(
+            num_l2_context=num_l2_context,
+            l2_frame_rate=l2_frame_rate,
+            l1_context_frames=l1_context_frames,
+        )
+        self.l1_start_offset = self.get_required_l1_start_offset()
+
+        with open(json_path, 'r') as f:
+            self.data = json.load(f)
+
+        self.sample_indices = sample_indices
+        if sample_indices is not None:
+            self.data = [self.data[i] for i in sample_indices]
+        self.augmenter = build_clip_augmenter(
+            aug=self.aug,
+            size=self.size,
+            input_format="tensor",
+        )
+
+    def _load_raw_frames(self, frame_paths):
+        images = [cv2.imread(os.path.join(self.images_root, p)) for p in frame_paths]
+        images = [Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) for img in images]
+        return torch.stack(
+            [torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0 for img in images],
+            dim=0,
+        )
+
+    def _load_images(self, frame_paths):
+        frames = self._load_raw_frames(frame_paths)
+        return self.augmenter(frames)
+
+    def __getitem__(self, index):
+        sample = self.data[index]
+        all_frames = sample['frames']
+
+        l1_start = self.l1_start_offset
+        l1_indices = self.get_l1_indices(l1_start, self.num_frames)
+        l2_indices = self.get_l2_indices(l1_start)
+        all_indices = l1_indices + l2_indices
+
+        if max(all_indices) >= len(all_frames):
+            raise ValueError(
+                f"Number of available L1 frames {len(all_frames)} is insufficient for "
+                f"num_frames={self.num_frames} with l1_start_offset={l1_start}."
+            )
+
+        images_all = self._load_images([all_frames[i] for i in all_indices])
+        images = images_all[:len(l1_indices)]
+
+        if self.l2_context_enabled:
+            l2_context = images_all[len(l1_indices):]
+            return {
+                'images': images,
+                'l2_context': l2_context,
+                'frame_rate': torch.tensor(self.frame_rate).float(),
+            }
+
+        return images
+
+    def __len__(self):
+        return len(self.data)

orbis2/evaluate/rollout_demo_v2.py

@@ -0,0 +1,706 @@
+import argparse
+import os
+import sys
+import imageio
+import logging
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import cv2
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from omegaconf.errors import ConfigTypeError
+from PIL import Image
+from pytorch_lightning import seed_everything
+from torchvision.utils import save_image
+
+from data.l2_context import L2ContextMixin
+from data.video_loaders import ClipAugmenter, DecordFrameAdapter, ResizeCenterPolicy, TensorFrameAdapter
+from util import instantiate_from_config
+
+try:
+    from decord import VideoReader, cpu as decord_cpu
+    DECORD_AVAILABLE = True
+except ImportError:
+    DECORD_AVAILABLE = False
+
+try:
+    from torchcodec.decoders import VideoDecoder
+    TORCHCODEC_AVAILABLE = True
+except ImportError:
+    TORCHCODEC_AVAILABLE = False
+
+
+logger = logging.getLogger(__name__)
+
+STEERING_FORMAT = "speed_yawrate"
+
+
+class _L1L2FrameIndexer(L2ContextMixin):
+    """Computes L1/L2 frame indices into a single source video via `L2ContextMixin`."""
+
+    def __init__(self, frame_interval, stored_data_frame_rate, num_l2_context, l2_frame_rate, l1_context_frames):
+        self.frame_interval = frame_interval
+        self.stored_data_frame_rate = stored_data_frame_rate
+        self._init_l2_context(
+            num_l2_context=num_l2_context,
+            l2_frame_rate=l2_frame_rate,
+            l1_context_frames=l1_context_frames,
+            require_l2_context=True,
+        )
+
+
+def get_rollout_future_frame_count(model, num_gen_frames):
+    """Return the total number of future image frames produced by the rollout."""
+    return int(num_gen_frames) * int(model.num_pred_frames)
+
+
+def resolve_context_image_size(args, config):
+    """Return (height, width) to resize context images to, from CLI args or the training config."""
+    if args.height is not None and args.width is not None:
+        return int(args.height), int(args.width)
+
+    try:
+        size = OmegaConf.select(config, "data.params.validation.params.size")
+    except ConfigTypeError:
+        size = None
+    if size is None:
+        size = config.data.params.train[0].params.size
+    size = (size, size) if isinstance(size, int) else tuple(size)
+
+    height = int(args.height) if args.height is not None else int(size[0])
+    width = int(args.width) if args.width is not None else int(size[1])
+    return height, width
+
+
+def require_l1l2_model(model):
+    """Raise a clear error if the loaded model isn't an L1-L2 hierarchical model."""
+    condition_preprocessor = getattr(model, "condition_preprocessor", None)
+    if condition_preprocessor is None or not hasattr(condition_preprocessor, "num_context_frames") \
+            or not hasattr(condition_preprocessor, "l2_predictor_frame_rate"):
+        raise TypeError(
+            "This script requires an L1-L2 hierarchical model, i.e. a `condition_preprocessor` "
+            "with a frozen `l2_predictor` (num_context_frames/l2_predictor_frame_rate). "
+            f"Got {type(condition_preprocessor).__name__ if condition_preprocessor else None}."
+        )
+
+
+def resolve_video_backend():
+    """Pick decord as the primary video backend, falling back to the slower torchcodec if unavailable."""
+    if DECORD_AVAILABLE:
+        return "decord"
+    if TORCHCODEC_AVAILABLE:
+        return "torchcodec"
+    raise ImportError("Either `decord` or `torchcodec` must be installed to read video frames.")
+
+
+def get_video_fps_and_length(path, backend):
+    """Return (native_fps, num_frames) for a video file."""
+    if backend == "decord":
+        reader = VideoReader(path, ctx=decord_cpu(0))
+        return float(reader.get_avg_fps()), len(reader)
+    decoder = VideoDecoder(path)
+    return float(decoder.metadata.average_fps), int(decoder.metadata.num_frames)
+
+
+def decode_video_frames(path, indices, backend):
+    """Decode an arbitrary list of frame indices from a video file."""
+    if backend == "decord":
+        reader = VideoReader(path, ctx=decord_cpu(0))
+        return reader.get_batch(indices).asnumpy()
+    decoder = VideoDecoder(path)
+    return torch.stack([decoder[i] for i in indices])
+
+
+def compute_frame_interval(native_fps, target_frame_rate, label):
+    """Return the integer native-frame stride for `target_frame_rate`, or raise if not exact."""
+    ratio = native_fps / float(target_frame_rate)
+    rounded = round(ratio)
+    if abs(ratio - rounded) > 1e-6:
+        raise ValueError(
+            f"The video's native frame rate ({native_fps:g} Hz) must be an integer multiple of "
+            f"{label} ({target_frame_rate:g} Hz), got ratio={ratio:g}."
+        )
+    return int(rounded)
+
+
+def resolve_l2_frame_rate(args, model):
+    """Return the L2 sampling rate, defaulting to (and validating against) the frozen L2 predictor's own rate."""
+    model_rate = model.condition_preprocessor.l2_predictor_frame_rate
+    if args.l2_frame_rate is None:
+        return float(model_rate)
+    if abs(float(args.l2_frame_rate) - float(model_rate)) > 1e-6:
+        raise ValueError(
+            f"--l2_frame_rate={args.l2_frame_rate:g} does not match the frozen L2 predictor's "
+            f"trained frame rate ({model_rate:g}); the L2 image context must be sampled at the "
+            "rate the L2 predictor was trained on."
+        )
+    return float(args.l2_frame_rate)
+
+
+def load_l1_l2_context(video_path, start_frame, l1_frame_rate, l2_frame_rate, l1_context_frames, l2_context_frames, height, width, backend, device):
+    """Sample L1 (high-rate) and L2 (low-rate, further back) context windows from one video."""
+    native_fps, video_length = get_video_fps_and_length(video_path, backend)
+    frame_interval = compute_frame_interval(native_fps, l1_frame_rate, "--l1_frame_rate")
+    compute_frame_interval(native_fps, l2_frame_rate, "--l2_frame_rate")
+
+    indexer = _L1L2FrameIndexer(
+        frame_interval=frame_interval,
+        stored_data_frame_rate=native_fps,
+        num_l2_context=l2_context_frames,
+        l2_frame_rate=l2_frame_rate,
+        l1_context_frames=l1_context_frames,
+    )
+
+    l1_span = (l1_context_frames - 1) * frame_interval + 1
+    if start_frame is None:
+        start_frame = video_length - l1_span
+    if start_frame < 0 or start_frame + l1_span > video_length:
+        raise ValueError(
+            f"Video is too short for the requested L1 context: need frames [{start_frame}, "
+            f"{start_frame + l1_span - 1}], but the video only has {video_length} frames."
+        )
+
+    required_offset = indexer.get_required_l1_start_offset()
+    if start_frame < required_offset:
+        raise ValueError(
+            f"Video does not have enough lookback for L2 context: start_frame={start_frame} but "
+            f"at least {required_offset} frames of history are required before it. "
+            "Use a longer video, a later --start_frame, or a lower --l2_frame_rate."
+        )
+
+    l1_indices, l2_indices = indexer.get_l1_and_l2_indices(start_frame, l1_context_frames)
+    all_frames = decode_video_frames(video_path, l1_indices + l2_indices, backend)
+
+    adapter = DecordFrameAdapter() if backend == "decord" else TensorFrameAdapter()
+    augmenter = ClipAugmenter(adapter, ResizeCenterPolicy((height, width)))
+    all_tensor = augmenter(all_frames)  # [F, C, H, W] in [-1, 1]
+
+    l1_tensor = all_tensor[: len(l1_indices)].unsqueeze(0).to(device)
+    l2_tensor = all_tensor[len(l1_indices) :].unsqueeze(0).to(device)
+    return l1_tensor, l2_tensor
+
+
+def load_steering_trajectory(steering_file, min_odo_steps, dtype, device):
+    """Load a raw [speed, yaw_rate] steering trajectory from a .npy/.csv file as a [1, T, 2] tensor."""
+    if not os.path.isfile(steering_file):
+        raise FileNotFoundError(f"Steering file {steering_file} does not exist")
+
+    if steering_file.endswith(".npy"):
+        loaded = np.load(steering_file)
+    elif steering_file.endswith(".csv"):
+        loaded = np.loadtxt(steering_file, delimiter=",")
+    else:
+        raise ValueError("Steering file must end with .npy or .csv")
+
+    if loaded.ndim != 2 or loaded.shape[1] != 2:
+        raise ValueError(
+            f"Steering file must contain [T, 2] (speed, yaw_rate) rows, got shape {tuple(loaded.shape)}"
+        )
+    if min_odo_steps is not None and loaded.shape[0] < min_odo_steps:
+        raise ValueError(
+            f"Steering file has too few timesteps: got {loaded.shape[0]}, expected at least {min_odo_steps}"
+        )
+
+    return torch.as_tensor(loaded, dtype=dtype, device=device).unsqueeze(0)
+
+
+def make_unconditional_steering(min_odo_steps, dtype, device):
+    """Build an all-NaN [1, T, 2] steering placeholder, the framework's own 'no steering data' signal."""
+    if min_odo_steps is None:
+        raise ValueError(
+            "Cannot run unconditionally without --steering_file: the loaded condition_preprocessor "
+            "could not report a required odometry length (`get_required_rollout_odometry_steps` "
+            "returned None), so there is no safe length to build a NaN placeholder at. "
+            "Supply --steering_file explicitly."
+        )
+    return torch.full((1, int(min_odo_steps), 2), float("nan"), dtype=dtype, device=device)
+
+
+def maybe_apply_condition_preprocessor_scales(model, speed_scale, yaw_rate_scale):
+    condition_preprocessor = getattr(model, "condition_preprocessor", None)
+    if condition_preprocessor is None:
+        return
+
+    if hasattr(condition_preprocessor, "speed_scale"):
+        condition_preprocessor.speed_scale = float(speed_scale)
+    if hasattr(condition_preprocessor, "yaw_rate_scale"):
+        condition_preprocessor.yaw_rate_scale = float(yaw_rate_scale)
+
+
+def _rgb_to_cv2_color(color):
+    """Convert an RGB tuple in [0, 1] to OpenCV BGR channel order."""
+    return (color[2], color[1], color[0])
+
+
+def _resolve_rollout_cursor_index(cursor_index, shared_trajectory, ctx_size, t, num_points):
+    cursor_idx = 0
+    if cursor_index is not None:
+        cursor_idx = int(cursor_index.item())
+    elif shared_trajectory and t >= ctx_size:
+        cursor_idx = t - ctx_size
+    return min(max(cursor_idx, 0), num_points - 1)
+
+
+def _panel_coords_fit_trajectory(traj_xy, panel_w, panel_h, margin):
+    if traj_xy.shape[0] == 0:
+        return None
+
+    forward = traj_xy[:, 0]
+    lateral = traj_xy[:, 1]
+
+    forward_min = np.min(forward)
+    forward_max = np.max(forward)
+    lateral_min = np.min(lateral)
+    lateral_max = np.max(lateral)
+    forward_span = forward_max - forward_min
+    lateral_span = lateral_max - lateral_min
+    usable_w = max(1, panel_w - 2 * margin)
+    usable_h = max(1, panel_h - 2 * margin)
+
+    scales = []
+    if lateral_span > 1e-6:
+        scales.append(usable_w / lateral_span)
+    if forward_span > 1e-6:
+        scales.append(usable_h / forward_span)
+    scale = min(scales) if scales else 1.0
+
+    if lateral_span > 1e-6:
+        px = margin + (lateral_max - lateral) * scale
+    else:
+        px = np.full_like(lateral, panel_w / 2)
+    if forward_span > 1e-6:
+        py = margin + (forward_max - forward) * scale
+    else:
+        py = np.full_like(forward, panel_h / 2)
+    px = np.clip(px, 0, panel_w - 1)
+    py = np.clip(py, 0, panel_h - 1)
+    return np.stack([px, py], axis=1).astype(np.int32)
+
+
+def _transform_trajectory_to_ego_frame(traj_xy, traj_heading, cursor_idx):
+    """Translate and rotate the trajectory so the current ego pose is at the origin facing +forward."""
+    if traj_xy.shape[0] == 0:
+        return traj_xy
+
+    centered = traj_xy - traj_xy[cursor_idx : cursor_idx + 1]
+    heading = float(traj_heading[cursor_idx])
+    cos_heading = np.cos(heading)
+    sin_heading = np.sin(heading)
+    rotation = np.array(
+        [
+            [cos_heading, sin_heading],
+            [-sin_heading, cos_heading],
+        ],
+        dtype=np.float32,
+    )
+    return centered @ rotation.T
+
+
+def _panel_coords_ego_frame(traj_xy, panel_w, panel_h, margin):
+    if traj_xy.shape[0] == 0:
+        return None
+
+    forward = traj_xy[:, 0]
+    lateral = traj_xy[:, 1]
+    extent = max(
+        float(np.max(np.abs(forward))),
+        float(np.max(np.abs(lateral))),
+        1e-3,
+    )
+    usable_w = max(1, panel_w - 2 * margin)
+    usable_h = max(1, panel_h - 2 * margin)
+    scale = min(usable_w, usable_h) / (2.0 * extent)
+    center_x = panel_w / 2.0
+    center_y = panel_h / 2.0
+
+    px = center_x - lateral * scale
+    py = center_y - forward * scale
+    px = np.clip(px, 0, panel_w - 1)
+    py = np.clip(py, 0, panel_h - 1)
+    return np.stack([px, py], axis=1).astype(np.int32)
+
+
+def overlay_trajectory_on_images(images, visualization, mode="trajectory"):
+    """Draw a compact trajectory panel on top of each rollout frame."""
+    cursor_index = None
+    headings = None
+    trajectory = visualization
+    if isinstance(visualization, dict):
+        trajectory = visualization.get("trajectory")
+        cursor_index = visualization.get("cursor_index")
+        headings = visualization.get("heading")
+
+    if trajectory is None:
+        return images
+
+    if mode not in {"trajectory", "trajectory_ego"}:
+        raise ValueError(f"Unsupported trajectory visualization mode: {mode}")
+
+    shared_trajectory = trajectory.ndim == 3
+    if trajectory.ndim == 3:
+        trajectory = trajectory.unsqueeze(1).expand(-1, images.shape[1], -1, -1)
+    elif trajectory.ndim != 4:
+        raise ValueError(f"Expected trajectory with shape [B, T, N, 2] or [B, N, 2], got {tuple(trajectory.shape)}")
+    if trajectory.shape[0] != images.shape[0] or trajectory.shape[1] != images.shape[1]:
+        raise ValueError(
+            f"Trajectory/image batch mismatch: images={tuple(images.shape)}, trajectory={tuple(trajectory.shape)}"
+        )
+    if cursor_index is not None:
+        if not torch.is_tensor(cursor_index):
+            cursor_index = torch.as_tensor(cursor_index, dtype=torch.long)
+        if cursor_index.ndim == 1:
+            cursor_index = cursor_index.unsqueeze(0).expand(images.shape[0], -1)
+        if cursor_index.shape[0] != images.shape[0] or cursor_index.shape[1] != images.shape[1]:
+            raise ValueError(
+                "Cursor/image batch mismatch: "
+                f"images={tuple(images.shape)}, cursor_index={tuple(cursor_index.shape)}"
+            )
+        cursor_index = cursor_index.to(device=trajectory.device)
+    if mode == "trajectory_ego":
+        if headings is None:
+            raise ValueError("Ego-aligned trajectory visualization requires per-point heading data.")
+        if headings.ndim == 2:
+            headings = headings.unsqueeze(1).expand(-1, images.shape[1], -1)
+        elif headings.ndim != 3:
+            raise ValueError(f"Expected heading with shape [B, T, N] or [B, N], got {tuple(headings.shape)}")
+        if headings.shape[0] != images.shape[0] or headings.shape[1] != images.shape[1]:
+            raise ValueError(f"Heading/image batch mismatch: images={tuple(images.shape)}, heading={tuple(headings.shape)}")
+    ctx_size = images.shape[1] - trajectory.shape[2] if shared_trajectory else 0
+
+    height = images.shape[3]
+    width = images.shape[4]
+
+    panel_w = int(height * 0.35)
+    panel_h = int(height * 0.35)
+    margin = max(2, int(min(height, width) * 0.02))
+    panel_x0 = margin
+    panel_y0 = height - panel_h - margin
+
+    traj_color = tuple(c * 2 - 1 for c in _rgb_to_cv2_color((0.0, 1.0, 0.0)))
+    cursor_color = tuple(c * 2 - 1 for c in _rgb_to_cv2_color((1.0, 0.25, 0.0)))
+    border_color = tuple(c * 2 - 1 for c in _rgb_to_cv2_color((1.0, 1.0, 1.0)))
+
+    for b in range(images.shape[0]):
+        for t in range(images.shape[1]):
+            traj_bt = trajectory[b, t].detach().cpu().numpy()
+            valid = ~np.isnan(traj_bt).any(axis=1)
+            if not np.any(valid):
+                continue
+
+            traj_xy = traj_bt[valid, :2]
+            cursor_idx = _resolve_rollout_cursor_index(
+                None if cursor_index is None else cursor_index[b, t],
+                shared_trajectory=shared_trajectory,
+                ctx_size=ctx_size,
+                t=t,
+                num_points=traj_xy.shape[0],
+            )
+
+            if mode == "trajectory_ego":
+                heading_bt = headings[b, t].detach().cpu().numpy()
+                heading_valid = heading_bt[valid]
+                ego_traj = _transform_trajectory_to_ego_frame(traj_xy, heading_valid, cursor_idx)
+                traj_pts = _panel_coords_ego_frame(ego_traj, panel_w, panel_h, margin)
+            else:
+                traj_pts = _panel_coords_fit_trajectory(traj_xy, panel_w, panel_h, margin)
+
+            if traj_pts is None:
+                continue
+
+            panel_t = np.full((panel_h, panel_w, 3), -1.0, dtype=np.float32)
+            cv2.rectangle(panel_t, (0, 0), (panel_w - 1, panel_h - 1), border_color, 1)
+            if traj_pts.shape[0] > 1:
+                cv2.polylines(panel_t, [traj_pts.reshape(-1, 1, 2)], False, traj_color, 2)
+            else:
+                cv2.circle(panel_t, tuple(traj_pts[0]), 2, traj_color, -1)
+            if mode == "trajectory_ego":
+                ego_marker = np.array([panel_w / 2.0, panel_h / 2.0], dtype=np.float32).astype(np.int32)
+                cv2.circle(panel_t, tuple(ego_marker), 3, cursor_color, -1)
+            else:
+                cv2.circle(panel_t, tuple(traj_pts[cursor_idx]), 3, cursor_color, -1)
+
+            frame = images[b, t].permute(1, 2, 0).cpu().numpy().copy()
+            frame[panel_y0:panel_y0 + panel_h, panel_x0:panel_x0 + panel_w] = panel_t
+            images[b, t] = torch.from_numpy(frame).permute(2, 0, 1)
+
+    return images
+
+
+@torch.no_grad()
+def generate_images(args, unknown_args):
+    """Run an L1-L2 hierarchical rollout from a single input video and save the resulting frames."""
+    if args.seed > 0:
+        torch.backends.cudnn.enable = False
+        torch.backends.cudnn.deterministic = True
+        seed_everything(args.seed)
+
+    config = OmegaConf.load(args.config)
+    config = OmegaConf.merge(config, OmegaConf.from_dotlist(unknown_args))
+    model = instantiate_from_config(config.model)
+
+    _ckpt_result = model.load_state_dict(torch.load(args.ckpt)["state_dict"], strict=False)
+    _exempt_prefixes = tuple(getattr(model, "checkpoint_exempt_key_prefixes", ()))
+    _unexpected_missing_keys = [k for k in _ckpt_result.missing_keys if not k.startswith(_exempt_prefixes)]
+    assert _unexpected_missing_keys == [], _unexpected_missing_keys
+    model = model.to(args.device)
+    _ = model.eval()
+
+    require_l1l2_model(model)
+
+    if os.path.exists(args.output_dir):
+        print("Folder exists, new images will be saved to the same folder, delete it if you want to start from scratch")
+    else:
+        os.makedirs(args.output_dir)
+
+    if args.compile:
+        def _maybe_compile(module, attr):
+            net = getattr(module, attr, None)
+            if net is not None:
+                setattr(module, attr, torch.compile(net, mode=args.compile_mode))
+                logger.info(f"Compiled {type(module).__name__}.{attr} with mode={args.compile_mode!r}")
+
+        _maybe_compile(model, 'ema_vit' if args.evaluate_ema else 'vit')
+
+        _l2_predictor = getattr(getattr(model, 'condition_preprocessor', None), 'l2_predictor', None)
+        if _l2_predictor is not None:
+            _maybe_compile(_l2_predictor, 'ema_vit')
+
+        logger.info("First rollout step will be slow (compilation). Subsequent steps reuse the graph.")
+
+    if args.compile and args.compile_artifacts:
+        if os.path.exists(args.compile_artifacts):
+            with open(args.compile_artifacts, "rb") as _f:
+                torch.compiler.load_cache_artifacts(_f.read())
+            logger.info(f"Loaded compile artifacts from {args.compile_artifacts!r}")
+        else:
+            logger.info(
+                f"Compile artifacts not found at {args.compile_artifacts!r}; "
+                "will save after the first batch."
+            )
+
+    maybe_apply_condition_preprocessor_scales(model, args.speed_scale, args.yaw_rate_scale)
+
+    height, width = resolve_context_image_size(args, config)
+    backend = resolve_video_backend()
+    l2_frame_rate = resolve_l2_frame_rate(args, model)
+
+    l1_context_frames = int(model.vit.num_context_frames)
+    l2_context_frames = int(model.condition_preprocessor.num_context_frames)
+
+    l1_tensor, l2_tensor = load_l1_l2_context(
+        video_path=args.video,
+        start_frame=args.start_frame,
+        l1_frame_rate=args.l1_frame_rate,
+        l2_frame_rate=l2_frame_rate,
+        l1_context_frames=l1_context_frames,
+        l2_context_frames=l2_context_frames,
+        height=height,
+        width=width,
+        backend=backend,
+        device=args.device,
+    )
+
+    num_future_frames = get_rollout_future_frame_count(model, args.num_gen_frames)
+    frame_rate = torch.tensor(float(args.l1_frame_rate), device=args.device)
+    data_batch = {"images": l1_tensor, "l2_context": l2_tensor, "frame_rate": frame_rate}
+
+    get_required_steps = getattr(model.condition_preprocessor, "get_required_rollout_odometry_steps", None)
+    min_odo_steps = None
+    if callable(get_required_steps):
+        min_odo_steps = get_required_steps(
+            validation_params=None,
+            num_condition_frames=l1_context_frames,
+            num_gen_frames=num_future_frames,
+            rollout_steps=args.num_gen_frames,
+        )
+    if args.steering_file is not None:
+        data_batch["steering"] = load_steering_trajectory(
+            args.steering_file, min_odo_steps, dtype=l1_tensor.dtype, device=args.device
+        )
+    else:
+        data_batch["steering"] = make_unconditional_steering(min_odo_steps, dtype=l1_tensor.dtype, device=args.device)
+    data_batch["steering_format"] = STEERING_FORMAT
+
+    # Roll out `num_videos` futures in parallel from the same context, as a single
+    # minibatch. The context (and steering) is tiled along the batch dim; the sampler
+    # draws independent initial noise per element, so the futures diverge under one
+    # seed. Everything below is computed at the tiled batch size so condition_kwargs
+    # stays internally consistent.
+    num_videos = max(1, int(args.num_videos))
+    if num_videos > 1:
+        def _tile_batch(t):
+            return t.repeat(num_videos, *([1] * (t.dim() - 1)))
+
+        l1_tensor = _tile_batch(l1_tensor)
+        l2_tensor = _tile_batch(l2_tensor)
+        data_batch["images"] = l1_tensor
+        data_batch["l2_context"] = l2_tensor
+        data_batch["steering"] = _tile_batch(data_batch["steering"])
+
+    condition_kwargs = model.condition_preprocessor.get_condition_kwargs_from_batch(data_batch, split="rollout")
+
+    logger.info(f"Steering source: {'none' if args.steering_file is None else args.steering_file}")
+    logger.info(f"Steering scales: speed={args.speed_scale:g}, yaw_rate={args.yaw_rate_scale:g}")
+    logger.info(f"L1/L2 frame rates: {args.l1_frame_rate:g}/{l2_frame_rate:g} Hz")
+    logger.info(f"Saving generated images to {args.output_dir}")
+
+    autocast_enabled = args.device.startswith("cuda")
+    with torch.autocast(dtype=torch.float16, device_type="cuda", enabled=autocast_enabled):
+        _latents, gen_frames = model.roll_out(
+            x_0={"images": l1_tensor},
+            num_gen_frames=args.num_gen_frames,
+            latent_input=False,
+            NFE=args.num_steps,
+            eta=args.eta,
+            sample_with_ema=args.evaluate_ema,
+            num_samples=l1_tensor.size(0),
+            frame_rate=frame_rate.reshape(1).repeat(l1_tensor.size(0)),
+            condition_kwargs=condition_kwargs,
+            decode_device=args.decode_device,
+            num_condition_frames=l1_tensor.size(1),
+        )
+
+    if args.vis_mode in {"trajectory", "trajectory_ego"}:
+        overlay_trajectory = model.condition_preprocessor.get_rollout_visualization_trajectory(
+            condition_kwargs=model.condition_preprocessor.get_condition_kwargs_from_batch(data_batch, split="rollout"),
+            num_condition_frames=l1_context_frames,
+            num_gen_steps=args.num_gen_frames,
+            num_pred_frames=model.num_pred_frames,
+        )
+        if overlay_trajectory is not None:
+            gen_frames = overlay_trajectory_on_images(gen_frames, overlay_trajectory, mode=args.vis_mode)
+
+    # Release rollout-time latent state before CPU-side file I/O.
+    del _latents, condition_kwargs, l1_tensor, l2_tensor
+
+    # Save each rollout in the minibatch to its own sequence folder. The layout
+    # (fake_images/sequence_XXXX/frame_XXXX.jpg) matches what the demo app reads.
+    num_out = gen_frames.shape[0]
+    num_frames = gen_frames.shape[1]
+    for b in range(num_out):
+        seq_dir = os.path.join(args.output_dir, "fake_images", f"sequence_{b:04d}")
+        os.makedirs(seq_dir, exist_ok=True)
+        for f in range(num_frames):
+            save_image(
+                (gen_frames[b, f] + 1.0) / 2.0,
+                os.path.join(seq_dir, f"frame_{f:04d}.jpg"),
+            )
+
+        imageio.mimsave(
+            os.path.join(args.output_dir, f"rollout_{b:04d}.gif"),
+            [
+                np.array(Image.open(os.path.join(seq_dir, f"frame_{f:04d}.jpg")))
+                for f in range(num_frames)
+            ],
+            fps=args.l1_frame_rate,
+            loop=0,
+        )
+
+    if args.device.startswith("cuda"):
+        logger.info(f"Max memory: {torch.cuda.max_memory_allocated() / 1024**3:.02f} GB")
+
+    if args.compile and args.compile_artifacts and not os.path.exists(args.compile_artifacts):
+        _artifacts = torch.compiler.save_cache_artifacts()
+        if _artifacts is not None:
+            with open(args.compile_artifacts, "wb") as _f:
+                _f.write(_artifacts[0])
+            logger.info(f"Saved compile artifacts to {args.compile_artifacts!r}")
+
+
+def main(args, unknown_args):
+    """Entrypoint that launches rollout generation with resolved CLI arguments."""
+    generate_images(args, unknown_args)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+
+    def str2bool(v):
+        if isinstance(v, bool):
+            return v
+        if v.lower() in ("yes", "true", "t", "y", "1"):
+            return True
+        if v.lower() in ("no", "false", "f", "n", "0"):
+            return False
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--exp_dir", type=str, default=None, help="Path to the experiment directory, where the config and checkpoints are stored")
+    parser.add_argument("--ckpt", type=str, default="checkpoints/last.ckpt", help="Path to the checkpoint file, relative to exp_dir")
+    parser.add_argument("--config", type=str, default="config.yaml", help="Path to the config file, relative to exp_dir")
+    parser.add_argument("--video", type=str, required=True, help="Path to the input video file to sample L1/L2 context from.")
+    parser.add_argument("--l1_frame_rate", type=float, required=True, help="Frame rate (Hz) to sample L1 context/rollout frames at, and the generated GIF's fps.")
+    parser.add_argument("--l2_frame_rate", type=float, default=None, help="Frame rate (Hz) to sample L2 context frames at. Defaults to the frozen L2 predictor's own trained frame rate; must match it if given explicitly.")
+    parser.add_argument("--start_frame", type=int, default=None, help="Native-video frame index to start the L1 context window at. Defaults to the latest window that fits (the end of the video).")
+    parser.add_argument("--height", type=int, default=None, help="Height to resize context images to. Defaults to the training config's size.")
+    parser.add_argument("--width", type=int, default=None, help="Width to resize context images to. Defaults to the training config's size.")
+    parser.add_argument(
+        "--num_gen_frames",
+        type=int,
+        default=1,
+        help="Number of rollout steps to generate; each step predicts `model.num_pred_frames` future frames.",
+    )
+    parser.add_argument("--output_dir", type=str, required=True, help="Directory to save the generated frames and GIF to.")
+    parser.add_argument(
+        "--num_videos",
+        type=int,
+        default=1,
+        help="Number of futures to roll out in parallel from the same context, as one minibatch.",
+    )
+    parser.add_argument(
+        "--vis_mode",
+        type=str,
+        default="none",
+        choices=["none", "trajectory", "trajectory_ego"],
+        help="Visualization mode",
+    )
+    parser.add_argument("--steering_file", type=str, default=None, help="Optional .npy or .csv trajectory file (columns: speed, yaw_rate), already at the expected odometry rate, used as steering input.")
+    parser.add_argument("--speed_scale", type=float, default=1.0, help="Global multiplicative factor applied to raw speed conditioning")
+    parser.add_argument("--yaw_rate_scale", type=float, default=1.0, help="Global multiplicative factor applied to raw yaw-rate conditioning")
+
+    parser.add_argument("--seed", type=int, default=42, help="Seed for reproducibility")
+    parser.add_argument("--device", type=str, default="cuda", help="Device")
+    parser.add_argument(
+        "--decode_device",
+        type=str,
+        default="cpu",
+        help="Device used for decoded rollout frames. Use 'cpu' to reduce peak GPU memory during saving.",
+    )
+    parser.add_argument("--num_steps", type=int, default=30, help="Number of steps for sampling")
+    parser.add_argument("--eta", type=float, default=0.0, help="Stochasticity for sampling")
+    parser.add_argument("--evaluate_ema", "--use_ema", type=str2bool, default=True, help="If the evaluation happen with ema model")
+    parser.add_argument(
+        "--compile",
+        type=str2bool,
+        default=False,
+        help="Wrap the DiT network with torch.compile for faster inference (PyTorch 2.x).",
+    )
+    parser.add_argument(
+        "--compile_mode",
+        type=str,
+        default="reduce-overhead",
+        choices=["default", "reduce-overhead", "max-autotune"],
+        help="torch.compile mode. 'reduce-overhead' uses CUDA graphs; 'max-autotune' adds kernel autotuning.",
+    )
+    parser.add_argument(
+        "--compile_artifacts",
+        type=str,
+        default=None,
+        help=(
+            "Path to torch.compiler cache artifacts (.pkl), relative to --exp_dir. "
+            "If the file exists, artifacts are loaded before rollout (fast startup). "
+            "If it does not exist, artifacts are saved after rollout (for future runs). "
+            "Only effective when --compile is True."
+        ),
+    )
+
+    args, unknown = parser.parse_known_args()
+
+    args.ckpt = os.path.join(args.exp_dir, args.ckpt)
+    args.config = os.path.join(args.exp_dir, args.config)
+    if args.compile_artifacts:
+        args.compile_artifacts = os.path.join(args.exp_dir, args.compile_artifacts)
+
+    main(args, unknown)

orbis2/models/first_stage/vqgan.py

@@ -0,0 +1,1190 @@
+import math
+import random
+import pytorch_lightning as pl
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.layers.pos_embed import resample_abs_pos_embed
+from torch.optim.lr_scheduler import LambdaLR
+
+from util import instantiate_from_config
+
+class VQModel(pl.LightningModule):
+    def __init__(
+        self,
+        encoder_config,
+        decoder_config,
+        quantizer_config,
+        loss_config=None,
+        grad_acc_steps=1,
+        cont_ratio_trainig= 0.0,
+        ignore_keys=None,
+        monitor=None,
+        entropy_loss_weight_scheduler_config=None,
+        distill_model_type="VIT_DINOv2",  # Options: VIT_DINO, CNN, VIT_DINOv2, etc.
+        min_lr_multiplier=0.1,
+        only_decoder=False,
+        scale_equivariance=None,
+    ):
+        super().__init__()
+
+        ignore_keys = ignore_keys or []
+        self.automatic_optimization = False
+        self.grad_acc_steps = grad_acc_steps
+        self.monitor = monitor
+        self.distill_model_type = distill_model_type
+        self.cont_ratio_trainig = cont_ratio_trainig
+        self.only_decoder=only_decoder
+        self.min_lr_multiplier = min_lr_multiplier
+        
+        assert (not scale_equivariance) or len(scale_equivariance) == 2, "if defined, scale_equivariance should be a list of two lists"
+        self.scale_equivariance = scale_equivariance
+
+        params_cfg = getattr(loss_config, "params", None)
+        if params_cfg is not None and hasattr(params_cfg, "get"):
+            default_lr = params_cfg.get("learning_rate", None)
+        else:
+            default_lr = None
+        # Ensure optimizers have sensible defaults even if the caller overrides later.
+        self.learning_rate = default_lr if default_lr is not None else 1e-4
+        self.num_iters_per_epoch = 1
+
+
+        # Decoder uses encoder params if none provided
+        if not hasattr(decoder_config, "params"):
+            decoder_config.params = encoder_config.params
+
+        # Instantiate core components
+        self.encoder = instantiate_from_config(encoder_config)
+        self.decoder = instantiate_from_config(decoder_config)
+        self.quantize = instantiate_from_config(quantizer_config)
+        self.loss = instantiate_from_config(loss_config) if loss_config is not None else None
+        self.entropy_loss_weight_scheduler = (
+            instantiate_from_config(entropy_loss_weight_scheduler_config)
+            if entropy_loss_weight_scheduler_config is not None else None
+        )
+
+        # Convolutional layers for quantization
+        self.quant_conv = nn.Conv2d(encoder_config.params["z_channels"], quantizer_config.params["e_dim"], 1)
+        self.post_quant_conv = nn.Conv2d(quantizer_config.params["e_dim"], decoder_config.params["z_channels"], 1)
+
+        self.encoder_normalize_embedding = encoder_config.params.get("normalize_embedding", False)
+        self.quantizer_normalize_embedding = quantizer_config.params.get("normalize_embedding", False)
+
+        self.if_distill_loss = (
+            False if loss_config is None
+            else loss_config.params.get('distill_loss_weight', 0.0) != 0.0
+        )
+        
+        # Image and patch size
+        self.image_size = encoder_config.params["resolution"]
+        self.patch_size = encoder_config.params["patch_size"]
+
+        self._init_distill_model(distill_model_type, encoder_config, decoder_config, quantizer_config)
+
+    def _init_distill_model(self, distill_type, encoder_cfg, decoder_cfg, quantizer_cfg):
+        image_size = encoder_cfg.params["resolution"]
+        patch_size = encoder_cfg.params["patch_size"]
+        q_e_dim = quantizer_cfg.params["e_dim"]
+        z_channels = decoder_cfg.params["z_channels"]
+
+        def conv1x1(in_c, out_c): return nn.Conv2d(in_c, out_c, 1)
+
+        if distill_type == "VIT_DINO":
+            self.distill = timm.create_model("timm/vit_base_patch16_224.dino", img_size=image_size, pretrained=True).eval()
+            self.post_quant_conv_distill = conv1x1(q_e_dim, z_channels)
+        elif distill_type == "VIT_DINOv2":
+            img_size = self._compute_scaled_size(image_size, patch_size)
+            self.distill = timm.create_model("timm/vit_base_patch14_dinov2.lvd142m", img_size=img_size, pretrained=True).eval()
+            self.post_quant_conv_distill = conv1x1(q_e_dim, z_channels)
+        elif distill_type == "VIT_DINOv3":
+            img_size = self._compute_scaled_size(image_size, patch_size, ckpt_patch_size=16)
+            self.distill = torch.hub.load('../dinov3', 'dinov3_vitb16', source='local', weights='./pretrained_models/dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth').eval()
+            self.post_quant_conv_distill = conv1x1(q_e_dim, z_channels)
+        elif distill_type == "VIT_DINOv2g":
+            img_size = int(image_size * 14 / patch_size)
+            self.distill = timm.create_model("timm/vit_giant_patch14_dinov2.lvd142m", img_size=img_size, pretrained=True).eval()
+            self.post_quant_conv_distill = conv1x1(q_e_dim, 1536)
+        elif distill_type == "VIT_DINOv2_large":
+            img_size = int(image_size * 14 / patch_size)
+            self.distill = timm.create_model("timm/vit_large_patch14_dinov2.lvd142m", img_size=img_size, pretrained=True).eval()
+            self.post_quant_conv_distill = conv1x1(q_e_dim, z_channels)
+        elif distill_type == "VIT_DINOv2_large_reg4":
+            img_size = int(image_size * 14 / patch_size)
+            self.distill = timm.create_model("timm/vit_large_patch14_reg4_dinov2.lvd142m", img_size=img_size, pretrained=True).eval()
+            self.post_quant_conv_distill = conv1x1(q_e_dim, z_channels)
+        elif distill_type == "SAM_VIT":
+            self.distill = timm.create_model("samvit_large_patch16.sa1b", pretrained=True)
+            self.post_quant_conv_distill = nn.Identity()
+        elif distill_type == "SAM_VIT_w_conv":
+            self.distill = timm.create_model("samvit_large_patch16.sa1b", pretrained=True)
+            self.post_quant_conv_distill = conv1x1(q_e_dim, z_channels)
+
+        elif distill_type == "depth_anything_VIT_L14":
+            self.distill = timm.create_model("vit_large_patch14_dinov2.lvd142m", img_size=224, pretrained=False)
+            state_dict = torch.load("./pretrained_models/depth_anything_vitl14.pth")
+            state_dict = {k.replace("pretrained.", "", 1): v for k, v in state_dict.items()}
+            state_dict["pos_embed"] = resample_abs_pos_embed(state_dict["pos_embed"], new_size=(16, 16))
+            self.distill.load_state_dict(state_dict, strict=False)
+            self.post_quant_conv_distill = conv1x1(q_e_dim, z_channels)
+
+
+    @staticmethod
+    def _compute_scaled_size(image_size, patch_size, ckpt_patch_size=14):
+        if isinstance(image_size, int):
+            return [image_size * ckpt_patch_size // patch_size] * 2
+        return [image_size[0] * ckpt_patch_size // patch_size, image_size[1] * ckpt_patch_size // patch_size]
+    
+    def get_input(self, batch):
+        x = batch['images']
+        return x.float()
+
+    def entropy_loss_weight_scheduling(self):
+        self.loss.entropy_loss_weight = self.entropy_loss_weight_scheduler(self.global_step)
+
+    def entropy_loss_weight_scheduling(self):
+        self.loss.entropy_loss_weight = self.entropy_loss_weight_scheduler(self.global_step)
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        if self.encoder_normalize_embedding:
+            h = F.normalize(h, p=2, dim=1)
+        ret = self.quantize(h)
+        ret["continuous"] = h
+        return ret
+        
+    def decode(self, quant):
+        distill_conv_out = self.post_quant_conv_distill(quant)
+        quant2 = self.post_quant_conv(quant)
+        return self.decoder(quant2), distill_conv_out
+    
+    def forward(self, input):
+        encoded = self.encode(input)
+        if torch.rand(1) > self.cont_ratio_trainig:
+            dec, distill_conv_out = self.decode(encoded["quantized"])
+        else:
+            dec, distill_conv_out = self.decode(encoded["continuous"])
+        return dec, (encoded['quantization_loss'], encoded['entropy_loss']), distill_conv_out
+    
+    def forward_se(self, input):
+        random_scale = [random.choice(self.scale_equivariance[0]), random.choice(self.scale_equivariance[1])]
+        downscale_factor = [1/random_scale[0], 1/random_scale[1]]
+        encoded = self.encode(input)
+        if torch.rand(1) > self.cont_ratio_trainig:
+            dec, distill_conv_out = self.decode(encoded["quantized"])
+            quant_se = F.interpolate(encoded["quantized"], scale_factor=downscale_factor, mode='bilinear', align_corners=False)
+            dec_se = self.decode(quant_se)[0]
+        else:
+            dec, distill_conv_out = self.decode(encoded["continuous"])
+            latents_se = F.interpolate(encoded["continuous"], scale_factor=downscale_factor, mode='bilinear', align_corners=False)
+            dec_se = self.decode(latents_se)[0]
+
+        input_se = F.interpolate(input, scale_factor=downscale_factor, mode='bilinear', align_corners=False)
+        decs = [dec, dec_se]
+        inputs = [input, input_se]
+        return inputs, decs, (encoded['quantization_loss'], encoded['entropy_loss']), distill_conv_out
+
+    def distill_loss(self, distill_output, decoder_distill_output):
+        #print(f'DINO loss calculation')
+        if 'VIT' in self.distill_model_type:
+            if 'reg4' in self.distill_model_type:
+                distill_output = distill_output[:, 5:, :] # [CLS, Register*4, Embeddings]
+            elif 'reg4' not in self.distill_model_type and 'DINOv2' in self.distill_model_type:
+                distill_output = distill_output[:, 1:, :] # uncomment for DINOv1 
+            elif 'reg4' not in self.distill_model_type and 'DINOv3' in self.distill_model_type:
+                distill_output = distill_output['x_norm_patchtokens'] #distill_output[:, 1:, :] # uncomment for DINOv1 
+            elif 'depth_anything' in self.distill_model_type:
+                distill_output = distill_output[:, 1:, :]
+            elif self.distill_model_type == 'SAM_VIT':
+                distill_output = distill_output.permute(0, 2, 3, 1).contiguous().view(distill_output.shape[0], -1, distill_output.shape[1])
+                distill_output = F.normalize(distill_output, p=2, dim=2) # without post_conv layer
+            elif self.distill_model_type == 'SAM_VIT_w_conv':
+                distill_output = distill_output.permute(0, 2, 3, 1).contiguous().view(distill_output.shape[0], -1, distill_output.shape[1])
+                # without L2 normalization
+            distill_output = distill_output.permute(0, 2, 1).contiguous()
+        
+        elif self.distill_model_type == 'CNN':
+            distill_output = distill_output.view(distill_output.shape[0], distill_output.shape[1], -1)
+        decoder_distill_output = decoder_distill_output.view(decoder_distill_output.shape[0], decoder_distill_output.shape[1], -1)
+        cos_similarity = F.cosine_similarity(decoder_distill_output, distill_output, dim=1)
+        cosine_loss = 1 - cos_similarity
+        distill_loss = cosine_loss.mean()
+        return distill_loss
+
+    def get_warmup_scheduler(self, optimizer, warmup_steps, min_lr_multiplier):
+        min_lr = self.learning_rate * min_lr_multiplier
+        total_steps = self.trainer.max_epochs * self.num_iters_per_epoch
+        def lr_lambda(step):
+            if step < warmup_steps:
+                # Linear warmup
+                return step/warmup_steps
+            # After warmup_steps, we just return 1. This could be modified to implement your own schedule
+            else:
+                progress = (step - warmup_steps) / (total_steps - warmup_steps)
+                cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
+                decayed = (1 - min_lr) * cosine_decay + min_lr
+                return decayed
+            #return 1.0        
+        
+        return LambdaLR(optimizer, lr_lambda)
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quantize.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.post_quant_conv.parameters())+
+                                  list(self.post_quant_conv_distill.parameters()),
+                                  lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        
+        scheduler_ae_warmup = self.get_warmup_scheduler(opt_ae, self.loss.warmup_steps, self.min_lr_multiplier)
+        scheduler_disc_warmup = self.get_warmup_scheduler(opt_disc, self.loss.warmup_steps, self.min_lr_multiplier)
+        
+
+        return [opt_ae, opt_disc], [scheduler_ae_warmup, scheduler_disc_warmup]
+    
+    def validation_step(self, batch, batch_idx):
+        x = self.get_input(batch)
+        xrec, qloss, decoder_distill_output = self(x)
+
+        distill_loss = self.distill_loss(self.get_distill_gt(x), decoder_distill_output) if self.if_distill_loss else torch.tensor(0.0, device=x.device)
+        aeloss, log_dict_ae = self.loss(qloss, distill_loss, x, xrec, 0, self.global_step,
+                                        last_layer=self.get_last_layer(), split="val")
+        self.log("val/distill_loss", distill_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("val/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        _, log_dict_disc = self.loss(qloss, distill_loss, x, xrec, 1, self.global_step,
+                                            last_layer=self.get_last_layer(), split="val")
+        rec_loss = log_dict_ae["val/rec_loss"]
+        self.log("val/rec_loss", rec_loss,
+                   prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        return self.log_dict
+
+    def training_step(self, batch, batch_idx):
+        self.entropy_loss_weight_scheduling()
+        self.log("train/enropy_loss_weight", self.loss.entropy_loss_weight, 
+                 prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        opt_ae, opt_disc = self.optimizers()
+        [scheduler_ae_warmup, scheduler_disc_warmup] = self.lr_schedulers()
+        
+        if self.only_decoder:
+            for param in self.encoder.parameters():
+                param.requires_grad = False
+            for param in self.quant_conv.parameters():
+                param.requires_grad = False
+            for param in self.quantize.parameters():
+                param.requires_grad = False
+            for param in self.post_quant_conv_distill.parameters():
+                param.requires_grad = False
+
+        x = self.get_input(batch)
+        
+        if self.scale_equivariance:
+            xs, xrec, qloss, decoder_distill_output = self.forward_se(x)
+        else:
+            xrec, qloss, decoder_distill_output = self(x)
+            xs = x
+
+        distill_loss = self.distill_loss(self.get_distill_gt(x), decoder_distill_output) if self.if_distill_loss else torch.tensor(0.0, device=x.device)
+
+        optimizer_idx = 1
+        discloss, log_dict_disc = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        discloss = discloss / self.grad_acc_steps
+        self.manual_backward(discloss)
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_disc.step()
+            opt_disc.zero_grad()
+            scheduler_disc_warmup.step()
+
+        optimizer_idx = 0
+        aeloss, log_dict_ae = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/distill_loss", distill_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+
+        aeloss = aeloss / self.grad_acc_steps
+        self.manual_backward(aeloss) 
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_ae.step()
+            opt_ae.zero_grad()
+            scheduler_ae_warmup.step()
+
+    
+    def get_distill_gt(self, x):
+        with torch.no_grad():
+            if 'VIT' in self.distill_model_type:
+                # resize image x to 224x224
+                if 'VIT_DINOv2' in self.distill_model_type or 'depth_anything' in self.distill_model_type:
+                    image_size = (self.image_size*14//self.patch_size, self.image_size*14//self.patch_size) if isinstance(self.image_size, int) else (self.image_size[0]*14//self.patch_size, self.image_size[1]*14//self.patch_size)
+                    x_224 = F.interpolate(x, size=image_size, mode='bilinear', align_corners=False)
+                    distill_output = self.distill.forward_features(x_224)
+                elif 'VIT_DINOv3' in self.distill_model_type:
+                    image_size = self._compute_scaled_size(self.image_size, self.patch_size, ckpt_patch_size=16)
+                    x_224 = F.interpolate(x, size=image_size, mode='bilinear', align_corners=False)
+                    distill_output = self.distill.forward_features(x_224)
+                else: # for VIT-DINOv1, VIT-SAM models
+                    distill_output = self.distill.forward_features(x)
+
+            elif self.distill_model_type == 'CNN':
+                distill_output = self.distill(x)
+        return distill_output
+    
+    
+    def get_last_layer(self):
+        try:
+            return self.decoder.conv_out.weight
+        except:
+            return None
+        
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = self.get_input(batch)
+        x = x.to(self.device)
+        xrec, _, _ = self(x)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        return log
+
+
+class VQModelIF(VQModel):
+    def __init__(self, 
+                 encoder_config,
+                 decoder_config,
+                 quantizer_config,
+                 loss_config=None,
+                 grad_acc_steps=1,
+                 cont_ratio_trainig= 0.0,
+                 ignore_keys=[],
+                 monitor=None,
+                 entropy_loss_weight_scheduler_config=None,
+                 distill_model_type='VIT_DINOv2', # 'VIT_DINO' or 'CNN' or VIT_DINOv2, VIT_DINOv2_large_reg4, SAM_VIT
+                 min_lr_multiplier=0.1,
+                 only_decoder=False,
+                 scale_equivariance=[]
+                 ):
+        super().__init__(encoder_config, decoder_config, quantizer_config, loss_config, 
+                         grad_acc_steps, cont_ratio_trainig, ignore_keys, 
+                         monitor, 
+                         entropy_loss_weight_scheduler_config, 
+                         distill_model_type, min_lr_multiplier, only_decoder, scale_equivariance)
+    
+        self.encoder2 = instantiate_from_config(encoder_config)
+        self.post_quant_conv = torch.nn.Conv2d(quantizer_config.params['e_dim']*2, decoder_config.params["z_channels"], 1)
+        self.quant_conv2 = torch.nn.Conv2d(encoder_config.params["z_channels"], quantizer_config.params['e_dim'], 1)
+        self.quantize2 = instantiate_from_config(quantizer_config)
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.encoder2.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quantize.parameters())+
+                                  list(self.quantize2.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.quant_conv2.parameters())+
+                                  list(self.post_quant_conv.parameters())+
+                                  list(self.post_quant_conv_distill.parameters()),
+                                  lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        
+        scheduler_ae_warmup = self.get_warmup_scheduler(opt_ae, self.loss.warmup_steps, self.min_lr_multiplier)
+        scheduler_disc_warmup = self.get_warmup_scheduler(opt_disc, self.loss.warmup_steps, self.min_lr_multiplier)
+        
+
+        return [opt_ae, opt_disc], [scheduler_ae_warmup, scheduler_disc_warmup]
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        if self.encoder_normalize_embedding:
+            h = F.normalize(h, p=2, dim=1)
+
+        h2 = self.encoder2(x)
+        h2 = self.quant_conv2(h2)
+        if self.encoder_normalize_embedding:
+            h2 = F.normalize(h2, p=2, dim=1)
+
+        quant = self.quantize(h)
+        quant2 = self.quantize2(h2)
+        
+        quant_loss = quant['quantization_loss'] + quant2['quantization_loss']
+        entropy_loss = quant['entropy_loss'] + quant2['entropy_loss'] if quant['entropy_loss'] is not None and quant2['entropy_loss'] is not None else None
+        
+        ret = {
+            "quantized": (quant["quantized"], quant2["quantized"]),
+            "quantization_loss": quant_loss,
+            "entropy_loss": entropy_loss,
+            "indices": (quant["indices"], quant2["indices"]),
+            "continuous": (h, h2)
+        }
+        return ret
+    
+
+    def decode(self, quant):
+        if isinstance(quant, tuple):
+            quant_rec = quant[0]
+            quant_sem = quant[1]
+        else:
+            print('Error: quant should be a tuple')
+        distill_conv_out = self.post_quant_conv_distill(quant_sem).view(quant_sem.shape[0], -1, quant_sem.shape[2]*quant_sem.shape[3])
+        quant_cat = torch.cat((quant_rec, quant_sem), dim=1)
+        quant = self.post_quant_conv(quant_cat)
+        return self.decoder(quant), distill_conv_out
+    
+    def decode_code(self, code_b):
+        code_b_rec, code_b_sem = code_b
+        quant_b_rec = self.quantize.get_codebook_entry(code_b_rec, (-1, code_b_rec.size(1), code_b_rec.size(2), self.quantize.e_dim))
+        quant_b_sem = self.quantize2.get_codebook_entry(code_b_sem, (-1, code_b_sem.size(1), code_b_sem.size(2), self.quantize.e_dim))
+        quant_b = (quant_b_rec, quant_b_sem)
+        dec = self.decode(quant_b)
+        return dec
+    
+    def forward_se(self, input):
+        random_scale = [random.choice(self.scale_equivariance[0]), random.choice(self.scale_equivariance[1])]
+        downscale_factor = [1/random_scale[0], 1/random_scale[1]]
+        encoded = self.encode(input)
+        quantized = encoded["quantized"]
+        continuous = encoded["continuous"]
+        if torch.rand(1) > self.cont_ratio_trainig:
+            dec, distill_conv_out = self.decode(quantized)
+            quant_se = F.interpolate(quantized[0], scale_factor=downscale_factor, mode='bilinear', align_corners=False), \
+                       F.interpolate(quantized[1], scale_factor=downscale_factor, mode='bilinear', align_corners=False)
+            dec_se = self.decode(quant_se)[0]
+        else:
+            dec, distill_conv_out = self.decode(continuous)
+            latents_se =  F.interpolate(continuous[0], scale_factor=downscale_factor, mode='bilinear', align_corners=False), \
+                          F.interpolate(continuous[1], scale_factor=downscale_factor, mode='bilinear', align_corners=False)
+            dec_se = self.decode(latents_se)[0]
+
+        input_se = F.interpolate(input, scale_factor=downscale_factor, mode='bilinear', align_corners=False)
+        decs = [dec, dec_se]
+        inputs = [input, input_se]
+        return inputs, decs, (encoded["quantization_loss"], encoded["entropy_loss"]), distill_conv_out
+
+    def training_step(self, batch, batch_idx):
+        self.entropy_loss_weight_scheduling()
+        self.log("train/enropy_loss_weight", self.loss.entropy_loss_weight, 
+                 prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        opt_ae, opt_disc = self.optimizers()
+        [scheduler_ae_warmup, scheduler_disc_warmup] = self.lr_schedulers()
+        
+        if self.only_decoder:
+            for param in self.encoder.parameters():
+                param.requires_grad = False
+            for param in self.encoder2.parameters():
+                param.requires_grad = False
+            for param in self.quant_conv.parameters():
+                param.requires_grad = False
+            for param in self.quant_conv2.parameters():
+                param.requires_grad = False
+            for param in self.quantize.parameters():
+                param.requires_grad = False
+            for param in self.quantize2.parameters():
+                param.requires_grad = False
+            for param in self.post_quant_conv_distill.parameters():
+                param.requires_grad = False
+
+        x = self.get_input(batch)
+        
+        if self.scale_equivariance:
+            xs, xrec, qloss, decoder_distill_output = self.forward_se(x)
+        else:
+            xrec, qloss, decoder_distill_output = self(x)
+            xs = x
+
+        distill_loss = self.distill_loss(self.get_distill_gt(x), decoder_distill_output) if self.if_distill_loss else torch.tensor(0.0, device=x.device)
+
+        optimizer_idx = 1
+        discloss, log_dict_disc = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        discloss = discloss / self.grad_acc_steps
+        self.manual_backward(discloss)
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_disc.step()
+            opt_disc.zero_grad()
+            scheduler_disc_warmup.step()
+
+        optimizer_idx = 0
+        aeloss, log_dict_ae = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/distill_loss", distill_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+
+        aeloss = aeloss / self.grad_acc_steps
+        self.manual_backward(aeloss) 
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_ae.step()
+            opt_ae.zero_grad()
+            scheduler_ae_warmup.step()
+
+class VQModelIFExtra(VQModelIF):
+    def __init__(self, 
+                 encoder_config,
+                 decoder_config,
+                 quantizer_config,
+                 quantizer_config2,
+                 loss_config=None,
+                 grad_acc_steps=1,
+                 cont_ratio_trainig= 0.0,
+                 ignore_keys=[],
+                 monitor=None,
+                 entropy_loss_weight_scheduler_config=None,
+                 distill_model_type='VIT_DINOv2', # 'VIT_DINO' or 'CNN' or VIT_DINOv2, VIT_DINOv2_large_reg4, SAM_VIT
+                 min_lr_multiplier=0.1,
+                 only_decoder=False,
+                 scale_equivariance=[]
+                 ):
+        super().__init__(encoder_config, decoder_config, quantizer_config, loss_config, 
+                         grad_acc_steps, cont_ratio_trainig, ignore_keys, 
+                         monitor, 
+                         entropy_loss_weight_scheduler_config, 
+                         distill_model_type, min_lr_multiplier, only_decoder, scale_equivariance)
+    
+        self.encoder2 = instantiate_from_config(encoder_config)
+        self.post_quant_conv = torch.nn.Conv2d(quantizer_config.params['e_dim']+quantizer_config2.params['e_dim'], decoder_config.params["z_channels"], 1)
+        self.quant_conv2 = nn.Conv2d(encoder_config.params["z_channels"], quantizer_config2.params['e_dim'], 1)
+        self.quantize2 = instantiate_from_config(quantizer_config2)
+
+        self._init_distill_model(distill_model_type, encoder_config, decoder_config, quantizer_config2)
+
+class VQModelIFSepEnc(VQModelIF):
+    def __init__(self, 
+                 encoder_config,
+                 encoder_config2,
+                 decoder_config,
+                 quantizer_config,
+                 loss_config=None,
+                 grad_acc_steps=1,
+                 cont_ratio_trainig= 0.0,
+                 ignore_keys=[],
+                 monitor=None,
+                 entropy_loss_weight_scheduler_config=None,
+                 distill_model_type='VIT_DINOv2', # 'VIT_DINO' or 'CNN' or VIT_DINOv2, VIT_DINOv2_large_reg4, SAM_VIT
+                 min_lr_multiplier=0.1,
+                 only_decoder=False,
+                 scale_equivariance=[]
+                 ):
+        super().__init__(encoder_config, decoder_config, quantizer_config, loss_config, 
+                         grad_acc_steps, cont_ratio_trainig, ignore_keys, 
+                         monitor, 
+                         entropy_loss_weight_scheduler_config, 
+                         distill_model_type, min_lr_multiplier, only_decoder, scale_equivariance)
+    
+        self.encoder2 = instantiate_from_config(encoder_config2)
+        self.post_quant_conv = torch.nn.Conv2d(quantizer_config.params['e_dim'], decoder_config.params["z_channels"], 1)
+        self.quant_conv2 = nn.Conv2d(encoder_config2.params["z_channels"], quantizer_config.params['e_dim'], 1)
+        self.quantize2 = instantiate_from_config(quantizer_config)
+
+        self._init_distill_model(distill_model_type, encoder_config, decoder_config, quantizer_config)
+
+    def training_step(self, batch, batch_idx):
+        self.entropy_loss_weight_scheduling()
+        self.log("train/enropy_loss_weight", self.loss.entropy_loss_weight, 
+                 prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        opt_ae, opt_disc = self.optimizers()
+        [scheduler_ae_warmup, scheduler_disc_warmup] = self.lr_schedulers()
+
+        for param in self.encoder2.parameters():
+            param.requires_grad = False
+
+        x = self.get_input(batch)
+        
+        if self.scale_equivariance:
+            xs, xrec, qloss, decoder_distill_output = self.forward_se(x)
+        else:
+            xrec, qloss, decoder_distill_output = self(x)
+            xs = x
+
+        distill_loss = self.distill_loss(self.get_distill_gt(x), decoder_distill_output) if self.if_distill_loss else torch.tensor(0.0, device=x.device)
+
+        optimizer_idx = 1
+        discloss, log_dict_disc = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        discloss = discloss / self.grad_acc_steps
+        self.manual_backward(discloss)
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_disc.step()
+            opt_disc.zero_grad()
+            scheduler_disc_warmup.step()
+
+        optimizer_idx = 0
+        aeloss, log_dict_ae = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/distill_loss", distill_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+
+        aeloss = aeloss / self.grad_acc_steps
+        self.manual_backward(aeloss) 
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_ae.step()
+            opt_ae.zero_grad()
+            scheduler_ae_warmup.step()
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quantize.parameters())+
+                                  list(self.quantize2.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.quant_conv2.parameters())+
+                                  list(self.post_quant_conv.parameters())+
+                                  list(self.post_quant_conv_distill.parameters()),
+                                  lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        
+        scheduler_ae_warmup = self.get_warmup_scheduler(opt_ae, self.loss.warmup_steps, self.min_lr_multiplier)
+        scheduler_disc_warmup = self.get_warmup_scheduler(opt_disc, self.loss.warmup_steps, self.min_lr_multiplier)
+        
+
+        return [opt_ae, opt_disc], [scheduler_ae_warmup, scheduler_disc_warmup]
+    
+    def decode(self, quant):
+        if isinstance(quant, tuple):
+            quant_rec = quant[0]
+            quant_sem = quant[1]
+        else:
+            print('Error: quant should be a tuple')
+        distill_conv_out = self.post_quant_conv_distill(quant_sem).view(quant_sem.shape[0], -1, quant_sem.shape[2]*quant_sem.shape[3])
+        #quant_cat = torch.cat((quant_rec, quant_sem), dim=1)
+        quant = self.post_quant_conv(quant_rec)
+        return self.decoder(quant), distill_conv_out
+
+
+class VQModelIFExtraSepEnc(VQModelIF):
+    def __init__(self, 
+                 encoder_config,
+                 encoder_config2,
+                 decoder_config,
+                 quantizer_config,
+                 quantizer_config2,
+                 loss_config=None,
+                 grad_acc_steps=1,
+                 cont_ratio_trainig= 0.0,
+                 ignore_keys=[],
+                 monitor=None,
+                 entropy_loss_weight_scheduler_config=None,
+                 distill_model_type='VIT_DINOv2', # 'VIT_DINO' or 'CNN' or VIT_DINOv2, VIT_DINOv2_large_reg4, SAM_VIT
+                 min_lr_multiplier=0.1,
+                 only_decoder=False,
+                 scale_equivariance=[]
+                 ):
+        super().__init__(encoder_config, decoder_config, quantizer_config, loss_config, 
+                         grad_acc_steps, cont_ratio_trainig, ignore_keys, 
+                         monitor, 
+                         entropy_loss_weight_scheduler_config, 
+                         distill_model_type, min_lr_multiplier, only_decoder, scale_equivariance)
+    
+        self.encoder2 = instantiate_from_config(encoder_config2)
+        self.post_quant_conv = torch.nn.Conv2d(quantizer_config.params['e_dim']+quantizer_config2.params['e_dim'], decoder_config.params["z_channels"], 1)
+        self.quant_conv2 = nn.Conv2d(encoder_config2.params["z_channels"], quantizer_config2.params['e_dim'], 1)
+        self.quantize2 = instantiate_from_config(quantizer_config2)
+
+        self._init_distill_model(distill_model_type, encoder_config2, decoder_config, quantizer_config2)
+
+    def training_step(self, batch, batch_idx):
+        self.entropy_loss_weight_scheduling()
+        self.log("train/enropy_loss_weight", self.loss.entropy_loss_weight, 
+                 prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        opt_ae, opt_disc = self.optimizers()
+        [scheduler_ae_warmup, scheduler_disc_warmup] = self.lr_schedulers()
+
+        for param in self.encoder2.parameters():
+            param.requires_grad = False
+
+        x = self.get_input(batch)
+        
+        if self.scale_equivariance:
+            xs, xrec, qloss, decoder_distill_output = self.forward_se(x)
+        else:
+            xrec, qloss, decoder_distill_output = self(x)
+            xs = x
+
+        distill_loss = self.distill_loss(self.get_distill_gt(x), decoder_distill_output) if self.if_distill_loss else torch.tensor(0.0, device=x.device)
+
+        optimizer_idx = 1
+        discloss, log_dict_disc = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        discloss = discloss / self.grad_acc_steps
+        self.manual_backward(discloss)
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_disc.step()
+            opt_disc.zero_grad()
+            scheduler_disc_warmup.step()
+
+        optimizer_idx = 0
+        aeloss, log_dict_ae = self.loss(qloss, distill_loss, xs, xrec, optimizer_idx, self.global_step,
+                                        last_layer=self.get_last_layer(), split="train")
+        self.log("train/distill_loss", distill_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+
+        aeloss = aeloss / self.grad_acc_steps
+        self.manual_backward(aeloss) 
+        if (batch_idx+1) % self.grad_acc_steps == 0:
+            opt_ae.step()
+            opt_ae.zero_grad()
+            scheduler_ae_warmup.step()
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quantize.parameters())+
+                                  list(self.quantize2.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.quant_conv2.parameters())+
+                                  list(self.post_quant_conv.parameters())+
+                                  list(self.post_quant_conv_distill.parameters()),
+                                  lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr, betas=(self.loss.beta_1, self.loss.beta_2))
+        
+        scheduler_ae_warmup = self.get_warmup_scheduler(opt_ae, self.loss.warmup_steps, self.min_lr_multiplier)
+        scheduler_disc_warmup = self.get_warmup_scheduler(opt_disc, self.loss.warmup_steps, self.min_lr_multiplier)
+        
+
+        return [opt_ae, opt_disc], [scheduler_ae_warmup, scheduler_disc_warmup]
+
+    def decode(self, quant):
+        if isinstance(quant, tuple):
+            quant_rec = quant[0]
+            quant_sem = quant[1]
+        else:
+            print('Error: quant should be a tuple')
+        distill_conv_out = self.post_quant_conv_distill(quant_sem).view(quant_sem.shape[0], -1, quant_sem.shape[2]*quant_sem.shape[3])
+        quant_cat = torch.cat((quant_rec, quant_sem), dim=1)
+        quant = self.post_quant_conv(quant_cat)
+        return self.decoder(quant), distill_conv_out
+
+
+class VQModelIFExtraSepEncSepDec(VQModelIFExtraSepEnc):
+    """
+    IFExtraSepEnc variant with a dedicated semantic decoder.
+
+    The main reconstruction decoder sees only the reconstruction branch. The
+    semantic branch is decoded by a second decoder with the same architecture,
+    while the semantic latents are detached on that auxiliary reconstruction
+    path so decoder gradients do not flow back into the semantic
+    encoder/quantizer. The semantic decoder is trained with reconstruction /
+    perceptual losses only and does not participate in the adversarial path.
+    """
+
+    def __init__(self,
+                 encoder_config,
+                 encoder_config2,
+                 decoder_config,
+                 quantizer_config,
+                 quantizer_config2,
+                 loss_config=None,
+                 grad_acc_steps=1,
+                 cont_ratio_trainig=0.0,
+                 ignore_keys=[],
+                 monitor=None,
+                 entropy_loss_weight_scheduler_config=None,
+                 distill_model_type='VIT_DINOv2',
+                 min_lr_multiplier=0.1,
+                 only_decoder=False,
+                 scale_equivariance=[]
+                 ):
+        super().__init__(
+            encoder_config=encoder_config,
+            encoder_config2=encoder_config2,
+            decoder_config=decoder_config,
+            quantizer_config=quantizer_config,
+            quantizer_config2=quantizer_config2,
+            loss_config=loss_config,
+            grad_acc_steps=grad_acc_steps,
+            cont_ratio_trainig=cont_ratio_trainig,
+            ignore_keys=ignore_keys,
+            monitor=monitor,
+            entropy_loss_weight_scheduler_config=entropy_loss_weight_scheduler_config,
+            distill_model_type=distill_model_type,
+            min_lr_multiplier=min_lr_multiplier,
+            only_decoder=only_decoder,
+            scale_equivariance=scale_equivariance,
+        )
+
+        self.post_quant_conv = torch.nn.Conv2d(
+            quantizer_config.params['e_dim'],
+            decoder_config.params["z_channels"],
+            1,
+        )
+        self.post_quant_conv_sem = torch.nn.Conv2d(
+            quantizer_config2.params['e_dim'],
+            decoder_config.params["z_channels"],
+            1,
+        )
+        self.decoder_sem = instantiate_from_config(decoder_config)
+
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        opt_ae = torch.optim.Adam(
+            list(self.encoder.parameters())
+            + list(self.encoder2.parameters())
+            + list(self.decoder.parameters())
+            + list(self.decoder_sem.parameters())
+            + list(self.quantize.parameters())
+            + list(self.quantize2.parameters())
+            + list(self.quant_conv.parameters())
+            + list(self.quant_conv2.parameters())
+            + list(self.post_quant_conv.parameters())
+            + list(self.post_quant_conv_sem.parameters())
+            + list(self.post_quant_conv_distill.parameters()),
+            lr=lr,
+            betas=(self.loss.beta_1, self.loss.beta_2),
+        )
+        opt_disc = torch.optim.Adam(
+            self.loss.discriminator.parameters(),
+            lr=lr,
+            betas=(self.loss.beta_1, self.loss.beta_2),
+        )
+
+        scheduler_ae_warmup = self.get_warmup_scheduler(opt_ae, self.loss.warmup_steps, self.min_lr_multiplier)
+        scheduler_disc_warmup = self.get_warmup_scheduler(opt_disc, self.loss.warmup_steps, self.min_lr_multiplier)
+
+        return [opt_ae, opt_disc], [scheduler_ae_warmup, scheduler_disc_warmup]
+
+    def _decode_components(self, quant):
+        if isinstance(quant, tuple):
+            quant_rec = quant[0]
+            quant_sem = quant[1]
+        else:
+            raise ValueError('VQModelIFExtraSepEncSepDec expects quant to be a tuple')
+
+        distill_conv_out = self.post_quant_conv_distill(quant_sem).view(
+            quant_sem.shape[0], -1, quant_sem.shape[2] * quant_sem.shape[3]
+        )
+
+        rec = self.decoder(self.post_quant_conv(quant_rec))
+
+        # Stop semantic-decoder reconstruction gradients from reaching encoder2/quantize2.
+        sem_latent = quant_sem.detach()
+        sem = self.decoder_sem(self.post_quant_conv_sem(sem_latent))
+
+        return rec, sem, distill_conv_out
+
+    def _select_decode_input(self, encoded):
+        if torch.rand(1, device=self.device).item() > self.cont_ratio_trainig:
+            return encoded["quantized"]
+        return encoded["continuous"]
+
+    def _forward_branch_outputs(self, x):
+        encoded = self.encode(x)
+        quant = self._select_decode_input(encoded)
+        rec, sem, distill_conv_out = self._decode_components(quant)
+        qloss = (encoded["quantization_loss"], encoded["entropy_loss"])
+        return rec, sem, qloss, distill_conv_out
+
+    def _semantic_decoder_loss(self, inputs, reconstructions, split):
+        if isinstance(reconstructions, list):
+            device = reconstructions[0].device
+        else:
+            device = reconstructions.device
+
+        rec_loss = torch.tensor(0.0, device=device)
+        p_loss = torch.tensor(0.0, device=device)
+
+        if isinstance(inputs, list):
+            for idx, (input_img, recon_img) in enumerate(zip(inputs, reconstructions)):
+                se_weight = 1 if idx == 0 else self.loss.se_weight
+                l1 = torch.abs(input_img - recon_img).mean()
+                l2 = F.mse_loss(recon_img, input_img)
+                rec_loss += (self.loss.l1_loss_weight * l1 + self.loss.l2_loss_weight * l2) * se_weight
+                if self.loss.perceptual_weight > 0:
+                    p_loss += self.loss.perceptual_loss(input_img, recon_img).mean() * se_weight
+        else:
+            l1 = torch.abs(inputs - reconstructions).mean()
+            l2 = F.mse_loss(reconstructions, inputs)
+            rec_loss = self.loss.l1_loss_weight * l1 + self.loss.l2_loss_weight * l2
+            if self.loss.perceptual_weight > 0:
+                p_loss = self.loss.perceptual_loss(inputs, reconstructions).mean()
+
+        loss = rec_loss + self.loss.perceptual_weight * p_loss
+        log = {
+            f"{split}/sem_decoder_loss": loss.detach().mean(),
+            f"{split}/sem_decoder_rec_loss": rec_loss.detach().mean(),
+            f"{split}/sem_decoder_p_loss": p_loss.detach().mean(),
+        }
+        return loss, log
+
+    def decode(self, quant):
+        rec, _, distill_conv_out = self._decode_components(quant)
+        return rec, distill_conv_out
+
+    def validation_step(self, batch, batch_idx):
+        x = self.get_input(batch)
+
+        if self.scale_equivariance:
+            xs, xrec_rec, xrec_sem, qloss, decoder_distill_output = self._forward_se_branch_outputs(x)
+        else:
+            xrec_rec, xrec_sem, qloss, decoder_distill_output = self._forward_branch_outputs(x)
+            xs = x
+
+        distill_loss = self.distill_loss(self.get_distill_gt(x), decoder_distill_output) if self.if_distill_loss else torch.tensor(0.0, device=x.device)
+
+        rec_aeloss, log_dict_ae = self.loss(
+            qloss,
+            distill_loss,
+            xs,
+            xrec_rec,
+            0,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val",
+        )
+        _, log_dict_disc = self.loss(
+            qloss,
+            distill_loss,
+            xs,
+            xrec_rec,
+            1,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="val",
+        )
+        sem_aeloss, log_dict_sem = self._semantic_decoder_loss(xs, xrec_sem, split="val")
+        aeloss = rec_aeloss + sem_aeloss
+
+        self.log("val/distill_loss", distill_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("val/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("val/aeloss_rec", rec_aeloss, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("val/aeloss_sem", sem_aeloss, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_sem, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        return self.log_dict
+
+    def training_step(self, batch, batch_idx):
+        self.entropy_loss_weight_scheduling()
+        self.log("train/enropy_loss_weight", self.loss.entropy_loss_weight,
+                 prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+        opt_ae, opt_disc = self.optimizers()
+        [scheduler_ae_warmup, scheduler_disc_warmup] = self.lr_schedulers()
+
+        if self.only_decoder:
+            for param in self.encoder.parameters():
+                param.requires_grad = False
+            for param in self.encoder2.parameters():
+                param.requires_grad = False
+            for param in self.quant_conv.parameters():
+                param.requires_grad = False
+            for param in self.quant_conv2.parameters():
+                param.requires_grad = False
+            for param in self.quantize.parameters():
+                param.requires_grad = False
+            for param in self.quantize2.parameters():
+                param.requires_grad = False
+            for param in self.post_quant_conv_distill.parameters():
+                param.requires_grad = False
+
+        x = self.get_input(batch)
+
+        if self.scale_equivariance:
+            xs, xrec_rec, xrec_sem, qloss, decoder_distill_output = self._forward_se_branch_outputs(x)
+        else:
+            xrec_rec, xrec_sem, qloss, decoder_distill_output = self._forward_branch_outputs(x)
+            xs = x
+
+        distill_loss = self.distill_loss(self.get_distill_gt(x), decoder_distill_output) if self.if_distill_loss else torch.tensor(0.0, device=x.device)
+
+        discloss, log_dict_disc = self.loss(
+            qloss,
+            distill_loss,
+            xs,
+            xrec_rec,
+            1,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="train",
+        )
+        self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        discloss = discloss / self.grad_acc_steps
+        self.manual_backward(discloss)
+        if (batch_idx + 1) % self.grad_acc_steps == 0:
+            opt_disc.step()
+            opt_disc.zero_grad()
+            scheduler_disc_warmup.step()
+
+        rec_aeloss, log_dict_ae = self.loss(
+            qloss,
+            distill_loss,
+            xs,
+            xrec_rec,
+            0,
+            self.global_step,
+            last_layer=self.get_last_layer(),
+            split="train",
+        )
+        sem_aeloss, log_dict_sem = self._semantic_decoder_loss(xs, xrec_sem, split="train")
+        aeloss = rec_aeloss + sem_aeloss
+
+        self.log("train/distill_loss", distill_loss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("train/aeloss_rec", rec_aeloss, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("train/aeloss_sem", sem_aeloss, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_sem, prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+        aeloss = aeloss / self.grad_acc_steps
+        self.manual_backward(aeloss)
+        if (batch_idx + 1) % self.grad_acc_steps == 0:
+            opt_ae.step()
+            opt_ae.zero_grad()
+            scheduler_ae_warmup.step()
+
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = self.get_input(batch)
+        x = x.to(self.device)
+        rec, sem, _, _ = self._forward_branch_outputs(x)
+        log["inputs"] = x
+        log["reconstructions"] = rec
+        log["reconstructions_rec"] = rec
+        log["reconstructions_sem"] = sem
+        return log
+
+
+class VQModelIFExtraSepEncSepDecJointSem(VQModelIFExtraSepEncSepDec):
+    """
+    Variant of VQModelIFExtraSepEncSepDec where the semantic decoder is trained
+    jointly with the semantic encoder (encoder2 / quantize2).
+
+    Unlike the base class, quant_sem is NOT detached before being passed to
+    decoder_sem, so reconstruction gradients flow back into the semantic encoder
+    and quantizer.  The semantic decoder loss is weighted by
+    ``sem_dec_loss_weight`` (default 0.1) to avoid dominating the primary
+    reconstruction objective.
+    """
+
+    def __init__(self, sem_dec_loss_weight=0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.sem_dec_loss_weight = sem_dec_loss_weight
+
+    def _decode_components(self, quant):
+        if isinstance(quant, tuple):
+            quant_rec = quant[0]
+            quant_sem = quant[1]
+        else:
+            raise ValueError('VQModelIFExtraSepEncSepDecJointSem expects quant to be a tuple')
+
+        distill_conv_out = self.post_quant_conv_distill(quant_sem).view(
+            quant_sem.shape[0], -1, quant_sem.shape[2] * quant_sem.shape[3]
+        )
+
+        rec = self.decoder(self.post_quant_conv(quant_rec))
+        # No detach: gradients from decoder_sem flow back into encoder2/quantize2.
+        sem = self.decoder_sem(self.post_quant_conv_sem(quant_sem))
+
+        return rec, sem, distill_conv_out
+
+    def _combine_ae_losses(self, rec_aeloss, sem_aeloss):
+        return rec_aeloss + self.sem_dec_loss_weight * sem_aeloss
+
+
+class VQModelIFExtraSepEncSepDecJointSem(VQModelIFExtraSepEncSepDec):
+    """
+    Variant of VQModelIFExtraSepEncSepDec where the semantic decoder is trained
+    jointly with the semantic encoder (encoder2 / quantize2).
+
+    Unlike the base class, quant_sem is NOT detached before being passed to
+    decoder_sem, so reconstruction gradients flow back into the semantic encoder
+    and quantizer.  The semantic decoder loss is weighted by
+    ``sem_dec_loss_weight`` (default 0.1) to avoid dominating the primary
+    reconstruction objective.
+    """
+
+    def __init__(self, sem_dec_loss_weight=0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.sem_dec_loss_weight = sem_dec_loss_weight
+
+    def _decode_components(self, quant):
+        if isinstance(quant, tuple):
+            quant_rec = quant[0]
+            quant_sem = quant[1]
+        else:
+            raise ValueError('VQModelIFExtraSepEncSepDecJointSem expects quant to be a tuple')
+
+        distill_conv_out = self.post_quant_conv_distill(quant_sem).view(
+            quant_sem.shape[0], -1, quant_sem.shape[2] * quant_sem.shape[3]
+        )
+
+        rec = self.decoder(self.post_quant_conv(quant_rec))
+        # No detach: gradients from decoder_sem flow back into encoder2/quantize2.
+        sem = self.decoder_sem(self.post_quant_conv_sem(quant_sem))
+
+        return rec, sem, distill_conv_out
+
+    def _combine_ae_losses(self, rec_aeloss, sem_aeloss):
+        return rec_aeloss + self.sem_dec_loss_weight * sem_aeloss
+
+
+class VQModelIFSep(VQModelIF):
+    def __init__(self, 
+                 encoder_config,
+                 decoder_config,
+                 quantizer_config,
+                 loss_config=None,
+                 grad_acc_steps=1,
+                 cont_ratio_trainig= 0.0,
+                 ignore_keys=[],
+                 monitor=None,
+                 entropy_loss_weight_scheduler_config=None,
+                 distill_model_type='VIT_DINOv2', # 'VIT_DINO' or 'CNN' or VIT_DINOv2, VIT_DINOv2_large_reg4, SAM_VIT
+                 min_lr_multiplier=0.1,
+                 only_decoder=False,
+                 scale_equivariance=[]
+                 ):
+        super().__init__(encoder_config, decoder_config, quantizer_config, loss_config, 
+                         grad_acc_steps, cont_ratio_trainig, ignore_keys, 
+                         monitor, 
+                         entropy_loss_weight_scheduler_config, 
+                         distill_model_type, min_lr_multiplier, only_decoder, scale_equivariance)
+    
+        self.encoder2 = instantiate_from_config(encoder_config)
+        self.post_quant_conv = torch.nn.Conv2d(quantizer_config.params['e_dim'], decoder_config.params["z_channels"], 1)
+        self.quant_conv2 = nn.Conv2d(encoder_config.params["z_channels"], quantizer_config.params['e_dim'], 1)
+        self.quantize2 = instantiate_from_config(quantizer_config)
+
+
+    def decode(self, quant):
+        if isinstance(quant, tuple):
+            quant_rec = quant[0]
+            quant_sem = quant[1]
+        else:
+            print('Error: quant should be a tuple')
+        distill_conv_out = self.post_quant_conv_distill(quant_sem).view(quant_sem.shape[0], -1, quant_sem.shape[2]*quant_sem.shape[3])
+        #quant_cat = torch.cat((quant_rec, quant_sem), dim=1)
+        quant = self.post_quant_conv(quant_rec)
+        return self.decoder(quant), distill_conv_out

orbis2/models/second_stage/cd_model_v2.py

@@ -0,0 +1,279 @@
+import os
+from copy import deepcopy
+
+import torch
+
+from .fm_model_v2 import (
+    FlowMatchingObjective,
+    FlowMatchingSamplerTeacherForcing,
+    PredictorModule,
+    requires_grad,
+    update_ema,
+)
+
+_DEFAULT_SAMPLE_GRIDS = {
+    1: [0.99],
+    2: [0.99, 0.8333],
+    4: [0.99, 0.9375, 0.8333, 0.625],
+}
+
+
+def _endpoint(x, t, v):
+    """G(x, t) = x + t * v -- endpoint (x0) prediction."""
+    shape = [t.shape[0]] + [1] * (x.dim() - 1)
+    return x + t.view(*shape) * v
+
+
+class ConsistencyDistillationObjective(FlowMatchingObjective):
+    """
+    Causal consistency distillation loss for a teacher-forced flow-matching
+    predictor: distills a frozen `teacher_vit` into a few-step student `vit`,
+    whose EMA (`ema_vit`) is the deliverable.
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        sigma_min=1e-5,
+        cd_num_timesteps=48,
+        cd_t_min=0.01,
+        cd_t_max=0.99,
+        cd_grid="uniform",
+        cd_loss="l2",
+        cd_pseudo_huber_c=0.00054,
+    ):
+        super().__init__(predictor_module, sigma_min=sigma_min)
+        self.cd_loss = cd_loss
+        self.cd_pseudo_huber_c = cd_pseudo_huber_c
+
+        if cd_grid == "uniform":
+            grid = torch.linspace(cd_t_min, cd_t_max, cd_num_timesteps)
+        else:
+            grid = torch.as_tensor(list(cd_grid), dtype=torch.float32)
+            if grid.numel() != cd_num_timesteps:
+                raise ValueError(
+                    f"cd_grid has {grid.numel()} entries, expected "
+                    f"cd_num_timesteps={cd_num_timesteps}."
+                )
+        predictor_module.register_buffer("cd_t_grid", grid, persistent=False)
+
+    def _split_sequence(self, x):
+        context = x[:, : -self.module.num_pred_frames]
+        target = x[:, -self.module.num_pred_frames :]
+        return context, target
+
+    def _forward_v(self, net, context, target_t, t, frame_rate, model_condition_kwargs):
+        sampler = self.module.sampler
+        model_input = sampler._build_model_inputs(context, target_t, t)
+        model_t = sampler._build_model_t(context, target_t, t)
+        pred = net(model_input, t=model_t * sampler.timescale, frame_rate=frame_rate, **model_condition_kwargs)
+        return sampler._extract_target_prediction(pred)
+
+    def compute_loss(self, pred, target):
+        diff = pred.float() - target.float()
+        if self.cd_loss == "pseudo_huber":
+            c = self.cd_pseudo_huber_c * (diff[0].numel() ** 0.5)
+            return (diff.flatten(1).pow(2).sum(dim=1) + c * c).sqrt() - c
+        return diff.pow(2)
+
+    def compute_step(self, x, frame_rate, condition_kwargs=None):
+        if getattr(self.module, "teacher_vit", None) is None:
+            raise RuntimeError(
+                "ConsistencyDistillationObjective requires the predictor module to "
+                "have a `teacher_vit` (set teacher_ckpt_path on "
+                "ConsistencyDistillPredictorModule)."
+            )
+
+        context, target = self._split_sequence(x)
+        batch_size = target.shape[0]
+        model_condition_kwargs = self.module.condition_preprocessor.get_model_condition_kwargs(condition_kwargs)
+
+        grid = self.module.cd_t_grid
+        n = torch.randint(1, grid.numel(), (batch_size,), device=x.device)
+        t, t_prev = grid[n], grid[n - 1]
+
+        noise = torch.randn_like(target)
+        x_t, _ = self.add_noise(target, t, noise=noise)
+
+        with torch.no_grad():
+            v_teacher = self._forward_v(self.module.teacher_vit, context, x_t, t, frame_rate, model_condition_kwargs)
+            step_shape = [t.shape[0]] + [1] * (x_t.dim() - 1)
+            x_prev = x_t + (t - t_prev).view(*step_shape) * v_teacher
+            v_target = self._forward_v(self.module.ema_vit, context, x_prev, t_prev, frame_rate, model_condition_kwargs)
+            tgt = _endpoint(x_prev, t_prev, v_target)
+
+        v_student = self._forward_v(self.module.vit, context, x_t, t, frame_rate, model_condition_kwargs)
+        pred = _endpoint(x_t, t, v_student)
+
+        return self.compute_loss(pred, tgt)
+
+
+class ConsistencyDistillationSampler(FlowMatchingSamplerTeacherForcing):
+    """
+    Few-step consistency sampler: repeatedly renoises the current endpoint
+    estimate and re-predicts, over a short descending grid of timesteps,
+    instead of integrating a many-step Euler ODE.
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        timescale=1.0,
+        integration_t_eps=0.0,
+        timestep_conditioning="global",
+        sample_fresh_noise=True,
+        sample_t_grids=None,
+    ):
+        super().__init__(
+            predictor_module,
+            timescale=timescale,
+            integration_t_eps=integration_t_eps,
+            timestep_conditioning=timestep_conditioning,
+        )
+        self.sample_fresh_noise = sample_fresh_noise
+        self.sample_t_grids = dict(_DEFAULT_SAMPLE_GRIDS)
+        if sample_t_grids is not None:
+            self.sample_t_grids.update({int(k): list(v) for k, v in dict(sample_t_grids).items()})
+
+    def _forward_v(self, net, context, target_t, t, frame_rate, model_condition_kwargs):
+        model_input = self._build_model_inputs(context, target_t, t)
+        model_t = self._build_model_t(context, target_t, t)
+        pred = net(model_input, t=model_t * self.timescale, frame_rate=frame_rate, **model_condition_kwargs)
+        return self._extract_target_prediction(pred)
+
+    def _fewstep_grid(self, NFE):
+        if NFE in self.sample_t_grids:
+            return [float(v) for v in self.sample_t_grids[NFE]]
+        t_max = float(self.module.cd_t_grid[-1])
+        t_min = float(self.module.cd_t_grid[0])
+        return torch.linspace(t_max, t_min, NFE + 1)[:-1].tolist()
+
+    @torch.no_grad()
+    def sample(
+        self,
+        images=None,
+        latent=False,
+        eta=0.0,
+        NFE=2,
+        sample_with_ema=True,
+        num_samples=8,
+        frame_rate=None,
+        condition_kwargs=None,
+        return_sample=False,
+    ):
+        del eta  # accepted only so PredictorModule.sample/roll_out call sites keep working
+        net = self._get_net(sample_with_ema)
+        device = next(net.parameters()).device
+        context = self._prepare_context(images, latent)
+
+        condition_kwargs = self.module.condition_preprocessor.prepare_condition_kwargs(
+            condition_kwargs,
+            batch_size=num_samples,
+            device=device,
+            split="sample",
+        )
+        model_condition_kwargs = self.module.condition_preprocessor.get_model_condition_kwargs(condition_kwargs)
+
+        if frame_rate is None:
+            frame_rate = self._default_frame_rate(num_samples, device)
+
+        input_h, input_w = self._get_input_hw()
+        grid = self._fewstep_grid(NFE)
+
+        eps0 = torch.randn(
+            num_samples,
+            self.module.num_pred_frames,
+            self.module.vit.in_channels,
+            input_h,
+            input_w,
+            device=device,
+        )
+        x = eps0
+        x0_hat = None
+        for k, t_k in enumerate(grid):
+            t = torch.full((num_samples,), t_k, device=device)
+            if k > 0:
+                eps_k = torch.randn_like(x) if self.sample_fresh_noise else eps0
+                x = (1.0 - t_k) * x0_hat + t_k * eps_k
+            v = self._forward_v(net, context, x, t, frame_rate, model_condition_kwargs)
+            x0_hat = _endpoint(x, t, v)
+
+        if return_sample:
+            return x0_hat, self.module.decode_frames(x0_hat.clone())
+        return x0_hat
+
+
+class ConsistencyDistillPredictorModule(PredictorModule):
+    """
+    PredictorModule specialization that distills a frozen teacher (loaded from
+    `teacher_ckpt_path`) into a few-step student via consistency distillation.
+    The deliverable is `ema_vit` (evaluate with sample_with_ema=True, NFE 1/2/4).
+    """
+
+    # `teacher_vit.*` is intentionally absent from saved checkpoints (see
+    # on_save_checkpoint below) since the teacher is always reloaded fresh from
+    # teacher_ckpt_path at __init__ time. Eval scripts that assert on
+    # load_state_dict(...).missing_keys should exempt these prefixes.
+    checkpoint_exempt_key_prefixes = ("teacher_vit.",)
+
+    def __init__(
+        self,
+        *,
+        teacher_ckpt_path,
+        cd_ema_mu=0.999,
+        cd_weight_decay=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cd_ema_mu = cd_ema_mu
+        self.cd_weight_decay = cd_weight_decay
+        self.strict_loading = False
+
+        checkpoint_path = os.path.expandvars(teacher_ckpt_path)
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(f"teacher_ckpt_path {checkpoint_path} does not exist.")
+
+        try:
+            state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True, mmap=True)["state_dict"]
+        except (TypeError, RuntimeError):
+            state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)["state_dict"]
+
+        ema_state_dict = {
+            key[len("ema_vit.") :]: value.clone()
+            for key, value in state_dict.items()
+            if key.startswith("ema_vit.")
+        }
+        del state_dict
+        if not ema_state_dict:
+            raise ValueError(f"No ema_vit.* keys found in teacher checkpoint {checkpoint_path}.")
+
+        self.vit.load_state_dict(ema_state_dict, strict=True)
+        self.ema_vit.load_state_dict(ema_state_dict, strict=True)
+        self.teacher_vit = deepcopy(self.vit)
+        requires_grad(self.teacher_vit, False)
+        self.teacher_vit.eval()
+
+    def setup(self, stage=None):
+        super().setup(stage)
+        if hasattr(self, "teacher_vit") and self.teacher_vit is not None:
+            self.teacher_vit.requires_grad_(False)
+
+    def configure_optimizers(self):
+        params = [p for p in self.vit.parameters() if p.requires_grad]
+        optimizer = torch.optim.AdamW(params, lr=self.learning_rate, weight_decay=self.cd_weight_decay)
+        scheduler = self.get_warmup_scheduler(optimizer, self.warmup_steps, self.min_lr_multiplier)
+        return [optimizer], [{"scheduler": scheduler, "interval": "step"}]
+
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        if not hasattr(self, "_ema_stream"):
+            self._ema_stream = torch.cuda.Stream() if torch.cuda.is_available() else None
+        if self._ema_stream is not None:
+            with torch.cuda.stream(self._ema_stream):
+                update_ema(self.ema_vit, self.vit, decay=self.cd_ema_mu)
+        else:
+            update_ema(self.ema_vit, self.vit, decay=self.cd_ema_mu)
+
+    def on_save_checkpoint(self, checkpoint):
+        super().on_save_checkpoint(checkpoint)
+        for key in [k for k in checkpoint["state_dict"] if k.startswith("teacher_vit.")]:
+            del checkpoint["state_dict"][key]

orbis2/models/second_stage/fm_conditions_v2.py

@@ -0,0 +1,1414 @@
+import logging
+import os
+
+import numpy as np
+import PIL.Image
+import PIL.ImageDraw
+import torch
+from einops import rearrange
+from omegaconf import OmegaConf
+
+from data.utils import get_trajectory_from_speeds_and_yaw_rates_batch
+from util import instantiate_from_config
+
+
+def _requires_grad(model, flag=True):
+    for param in model.parameters():
+        param.requires_grad = flag
+
+
+def _text_on_image(images, texts):
+    for i in range(images.shape[0]):
+        image = images[i]
+        text = texts[i]
+        image = PIL.Image.fromarray((((image + 1) / 2).permute(1, 2, 0).cpu().numpy() * 255).astype("uint8"))
+        draw = PIL.ImageDraw.Draw(image)
+        draw.text((10, 10), text, fill="red", font_size=15)
+        image = (torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0) * 2 - 1
+        images[i] = image
+    return images
+
+
+class ConditionPreprocessor:
+    """
+    Default no-op condition preprocessor.
+
+    Concrete implementations can extract batch conditions, normalize sampling inputs,
+    and update autoregressive rollout state without changing predictor/objective/sampler
+    control flow.
+    """
+
+    def __init__(self, predictor_module):
+        self.module = predictor_module
+
+    def get_condition_kwargs_from_batch(self, batch, split):
+        return {}
+
+    def prepare_condition_kwargs(self, condition_kwargs=None, batch_size=None, device=None, split=None):
+        del batch_size, split
+        if condition_kwargs is None:
+            return {}
+
+        prepared = {}
+        for key, value in condition_kwargs.items():
+            if torch.is_tensor(value) and device is not None:
+                prepared[key] = value.to(device)
+            else:
+                prepared[key] = value
+        return prepared
+
+    def get_model_condition_kwargs(self, condition_kwargs):
+        if not condition_kwargs:
+            return {}
+        return {
+            key: value
+            for key, value in condition_kwargs.items()
+            if not key.startswith("_")
+        }
+
+    def slice_condition_kwargs(self, condition_kwargs, item):
+        if not condition_kwargs:
+            return {}
+
+        sliced = {}
+        for key, value in condition_kwargs.items():
+            if torch.is_tensor(value):
+                sliced[key] = value[item]
+            else:
+                sliced[key] = value
+        return sliced
+
+    def update_rollout_condition_kwargs(self, condition_kwargs, prediction, context, step_idx):
+        del prediction, context, step_idx
+        return condition_kwargs
+
+    def get_rollout_visualization_state(self, condition_kwargs):
+        del condition_kwargs
+        return None
+
+    def get_visualization_trajectory(self, visualization_state):
+        del visualization_state
+        return None
+
+    def update_rollout_visualization_state(self, visualization_state, step_idx):
+        del step_idx
+        return visualization_state
+
+    def get_rollout_visualization_trajectory(
+        self,
+        condition_kwargs,
+        num_gen_steps,
+        num_pred_frames=1,
+        num_condition_frames=None,
+    ):
+        del condition_kwargs, num_gen_steps, num_pred_frames, num_condition_frames
+        return None
+
+    def annotate_logged_images(self, images, batch, num_images):
+        del batch, num_images
+        return images
+
+    def to_device(self, device):
+        pass
+
+    def get_max_condition_odometry_offset(self):
+        """
+        Return the largest future raw-odometry offset needed to construct one
+        conditioning tensor from the current anchor position.
+
+        Preprocessors that do not depend on extra raw odometry should return
+        ``None``.
+        """
+        return None
+
+    def get_required_rollout_odometry_steps(
+        self,
+        *,
+        validation_params,
+        num_condition_frames,
+        num_gen_frames,
+        rollout_steps,
+    ):
+        del validation_params, num_condition_frames, num_gen_frames, rollout_steps
+        return None
+
+
+def _validate_speed_yaw_scale(value, name):
+    scale = float(value)
+    if not np.isfinite(scale):
+        raise ValueError(f"`{name}` must be finite, got {value!r}.")
+    return scale
+
+
+def _apply_speed_yaw_scales(steering, speed_scale=1.0, yaw_rate_scale=1.0):
+    if not torch.is_tensor(steering):
+        raise TypeError(f"`steering` must be a tensor, got {type(steering).__name__}.")
+    scaled = steering.clone()
+    scaled[..., 0] = scaled[..., 0] * float(speed_scale)
+    scaled[..., 1] = scaled[..., 1] * float(yaw_rate_scale)
+    return scaled
+
+
+class SpeedYawConditionPreprocessorBase(ConditionPreprocessor):
+    """
+    Shared steering-conditioning utilities for preprocessors that consume raw
+    `[speed, yaw_rate]` sequences and advance through them during rollout.
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        odometry_steps_per_image_frame,
+        context_frame_anchor_index=-1,
+        steering_format="speed_yawrate",
+        speed_scale=1.0,
+        yaw_rate_scale=1.0,
+        steering_drop=0.0,
+    ):
+        super().__init__(predictor_module)
+        self.steering_format = str(steering_format)
+        if self.steering_format != "speed_yawrate":
+            raise ValueError(
+                f"{type(self).__name__} only supports `steering_format='speed_yawrate'`, "
+                f"got {self.steering_format!r}."
+            )
+
+        self.odometry_steps_per_image_frame = int(odometry_steps_per_image_frame)
+        if self.odometry_steps_per_image_frame <= 0:
+            raise ValueError(
+                "`odometry_steps_per_image_frame` must be a positive integer, "
+                f"got {self.odometry_steps_per_image_frame}."
+            )
+        self.context_frame_anchor_index = int(context_frame_anchor_index)
+        self.speed_scale = _validate_speed_yaw_scale(speed_scale, "speed_scale")
+        self.yaw_rate_scale = _validate_speed_yaw_scale(yaw_rate_scale, "yaw_rate_scale")
+        self.steering_drop = self._validate_steering_drop(steering_drop)
+
+    def _validate_steering_drop(self, steering_drop):
+        drop_prob = float(steering_drop)
+        if not np.isfinite(drop_prob):
+            raise ValueError(f"`steering_drop` must be finite, got {steering_drop!r}.")
+        if not 0.0 <= drop_prob <= 1.0:
+            raise ValueError(f"`steering_drop` must be in [0, 1], got {drop_prob}.")
+        return drop_prob
+
+    def _validate_offsets(self, offsets, name):
+        if not offsets:
+            raise ValueError(f"`{name}` must contain at least one positive odometry offset.")
+        normalized_offsets = [int(offset) for offset in offsets]
+        if any(offset <= 0 for offset in normalized_offsets):
+            raise ValueError(f"`{name}` must be positive integers, got {normalized_offsets}.")
+        if normalized_offsets != sorted(normalized_offsets):
+            raise ValueError(f"`{name}` must be sorted in ascending order, got {normalized_offsets}.")
+        return normalized_offsets
+
+    def get_max_condition_odometry_offset(self):
+        offsets = None
+        if hasattr(self, "goal_offsets"):
+            offsets = self.goal_offsets
+        elif hasattr(self, "steering_offsets"):
+            offsets = self.steering_offsets
+        if not offsets:
+            return None
+        return int(max(offsets))
+
+    def get_required_rollout_odometry_steps(
+        self,
+        *,
+        validation_params,
+        num_condition_frames,
+        num_gen_frames,
+        rollout_steps,
+    ):
+        del validation_params, num_gen_frames
+        max_condition_offset = self.get_max_condition_odometry_offset()
+        if max_condition_offset is None:
+            return None
+
+        rollout_steps = int(rollout_steps)
+        if rollout_steps <= 0:
+            raise ValueError(f"`rollout_steps` must be positive, got {rollout_steps}.")
+
+        anchor_frame_index = self._resolve_context_frame_anchor_index(int(num_condition_frames))
+        initial_anchor_odo_index = anchor_frame_index * int(self.odometry_steps_per_image_frame)
+        rollout_step_odo = int(self.odometry_steps_per_image_frame) * int(self.module.num_pred_frames)
+        last_anchor_odo_index = initial_anchor_odo_index + (rollout_steps - 1) * rollout_step_odo
+        return last_anchor_odo_index + int(max_condition_offset) + 1
+
+    def _normalize_batch_steering_format(self, steering_format):
+        if steering_format is None:
+            return None
+        if isinstance(steering_format, str):
+            return steering_format
+        if isinstance(steering_format, (list, tuple)):
+            unique_formats = {str(item) for item in steering_format}
+            if len(unique_formats) != 1:
+                raise ValueError(f"Batch contains mixed steering formats: {sorted(unique_formats)}.")
+            return next(iter(unique_formats))
+        return str(steering_format)
+
+    def _validate_batch_steering_format(self, batch):
+        batch_format = self._normalize_batch_steering_format(batch.get("steering_format"))
+        if batch_format is None:
+            return
+        if batch_format != self.steering_format:
+            raise ValueError(
+                f"{type(self).__name__} expected batch steering format "
+                f"{self.steering_format!r}, got {batch_format!r}. Check the dataset "
+                "`odo_transform_config` or `annotation_key`."
+            )
+
+    def _validate_steering(self, steering):
+        """Validate and return raw steering with expected shape `[B, T, 2]`."""
+        if steering is None:
+            raise ValueError(f"`steering` is required for {type(self).__name__}.")
+        if steering.ndim != 3:
+            raise ValueError(
+                f"`steering` must have shape [B, T, 2] for {type(self).__name__}, "
+                f"got {tuple(steering.shape)}."
+            )
+        if steering.shape[2] != 2:
+            raise ValueError(
+                f"{type(self).__name__} expects raw steering features "
+                f"[speed, yaw_rate], got shape {tuple(steering.shape)}."
+            )
+        return steering
+
+    def _resolve_odometry_dt(self, batch, batch_size, device, required):
+        frame_rate = batch.get("frame_rate")
+        if frame_rate is None:
+            if required:
+                raise ValueError(
+                    f"`frame_rate` is required to compute physical steering trajectories for "
+                    f"{type(self).__name__}."
+                )
+            return None
+
+        frame_rate = torch.as_tensor(frame_rate, device=device, dtype=torch.float32)
+        if frame_rate.ndim == 0:
+            frame_rate = frame_rate.expand(batch_size)
+        elif frame_rate.ndim == 1 and frame_rate.shape[0] == 1 and batch_size != 1:
+            frame_rate = frame_rate.expand(batch_size)
+        elif frame_rate.ndim != 1 or frame_rate.shape[0] != batch_size:
+            raise ValueError(
+                "`frame_rate` must be broadcastable to shape [B], got "
+                f"{tuple(frame_rate.shape)} for batch_size={batch_size}."
+            )
+
+        if torch.any(frame_rate <= 0):
+            raise ValueError(f"`frame_rate` must be positive, got {frame_rate}.")
+
+        return 1.0 / (frame_rate * float(self.odometry_steps_per_image_frame))
+
+    def _get_context_num_frames(self, batch, split):
+        """Return the number of conditioning image frames available for the current split."""
+        if not isinstance(batch, dict) or "images" not in batch:
+            raise ValueError(f"`images` are required in the batch to compute {type(self).__name__} conditions.")
+        images = batch["images"]
+        if images.ndim != 5:
+            raise ValueError(f"`images` must have shape [B, F, C, H, W], got {tuple(images.shape)}.")
+
+        if split == "rollout":
+            return images.shape[1]
+        if images.shape[1] < self.module.num_pred_frames:
+            raise ValueError(
+                f"Need at least {self.module.num_pred_frames} image frames, got {images.shape[1]}."
+            )
+        return images.shape[1] - self.module.num_pred_frames
+
+    def _resolve_context_frame_anchor_index(self, context_num_frames):
+        """Resolve the configured anchor index against the current context length."""
+        anchor_index = self.context_frame_anchor_index
+        if anchor_index < 0:
+            anchor_index = context_num_frames + anchor_index
+        if not 0 <= anchor_index < context_num_frames:
+            raise IndexError(
+                f"Resolved context_frame_anchor_index={anchor_index} is out of bounds for "
+                f"context_num_frames={context_num_frames}."
+            )
+        return anchor_index
+
+    def _build_rollout_state(self, steering, context_num_frames):
+        steering = self._validate_steering(steering)
+        batch_size = steering.shape[0]
+        device = steering.device
+        steps_per_image_frame = torch.full(
+            (batch_size,),
+            self.odometry_steps_per_image_frame,
+            device=device,
+            dtype=torch.long,
+        )
+        context_frame_anchor_index = self._resolve_context_frame_anchor_index(context_num_frames)
+        anchor_odo_index = context_frame_anchor_index * steps_per_image_frame
+        rollout_step_odo = steps_per_image_frame * self.module.num_pred_frames
+        return {
+            "_raw_speed_yaw": steering,
+            "_anchor_odo_index": anchor_odo_index,
+            "_rollout_step_odo": rollout_step_odo,
+        }
+
+    def _transform_raw_steering(self, steering):
+        steering = self._validate_steering(steering)
+        if self.speed_scale == 1.0 and self.yaw_rate_scale == 1.0:
+            return steering
+        return _apply_speed_yaw_scales(
+            steering,
+            speed_scale=self.speed_scale,
+            yaw_rate_scale=self.yaw_rate_scale,
+        )
+
+    def _build_condition_kwargs(self, steering, context_num_frames, odometry_dt=None):
+        condition_kwargs = self._build_rollout_state(steering, context_num_frames)
+        if odometry_dt is not None:
+            condition_kwargs["_odometry_dt"] = odometry_dt.to(
+                device=steering.device,
+                dtype=steering.dtype,
+        )
+        condition_kwargs["steering"] = self._compute_condition_steering(
+            steering,
+            condition_kwargs["_anchor_odo_index"],
+            odometry_dt=condition_kwargs.get("_odometry_dt"),
+        )
+        return condition_kwargs
+
+    def _maybe_apply_steering_dropout(self, steering, split):
+        if split != "train" or self.steering_drop == 0.0:
+            return steering
+
+        drop_mask = torch.rand(steering.shape[0], device=steering.device) < self.steering_drop
+        if not drop_mask.any():
+            return steering
+
+        dropped = steering.clone()
+        dropped[drop_mask] = torch.nan
+        return dropped
+
+    def _compute_condition_steering(self, steering, anchor_odo_index, odometry_dt=None):
+        raise NotImplementedError
+
+    def _requires_odometry_dt_for_conditioning(self):
+        return False
+
+    def _make_no_steering_condition(self, batch_size, device):
+        """Return a fully-NaN steering condition tensor for batches with no steering data.
+
+        The NaN values cause LinearSteeringEmbedder to fall back to its no_value_embeddings,
+        making missing-steering batches equivalent to fully-dropped steering (steering_drop=1.0).
+        Returns None if the subclass does not support unconditional batches.
+        """
+        return None
+
+    def get_condition_kwargs_from_batch(self, batch, split):
+        if not isinstance(batch, dict) or "steering" not in batch:
+            if isinstance(batch, dict) and "images" in batch:
+                images = batch["images"]
+                no_steering = self._make_no_steering_condition(images.shape[0], images.device)
+                if no_steering is not None:
+                    return {"steering": no_steering}
+            return {}
+        self._validate_batch_steering_format(batch)
+        source_steering = self._validate_steering(batch["steering"])
+        steering = self._transform_raw_steering(source_steering)
+        context_num_frames = self._get_context_num_frames(batch, split)
+        odometry_dt = self._resolve_odometry_dt(
+            batch,
+            batch_size=steering.shape[0],
+            device=steering.device,
+            required=self._requires_odometry_dt_for_conditioning(),
+        )
+        condition_kwargs = self._build_condition_kwargs(steering, context_num_frames, odometry_dt=odometry_dt)
+        condition_kwargs["steering"] = self._maybe_apply_steering_dropout(condition_kwargs["steering"], split)
+        condition_kwargs["_source_raw_speed_yaw"] = source_steering
+        return condition_kwargs
+
+    def update_rollout_condition_kwargs(self, condition_kwargs, prediction, context, step_idx):
+        del prediction, context, step_idx
+        if not condition_kwargs:
+            return {}
+        updated = dict(condition_kwargs)
+        updated["_anchor_odo_index"] = condition_kwargs["_anchor_odo_index"] + condition_kwargs["_rollout_step_odo"]
+        updated["steering"] = self._compute_condition_steering(
+            condition_kwargs["_raw_speed_yaw"],
+            updated["_anchor_odo_index"],
+            odometry_dt=updated.get("_odometry_dt"),
+        )
+        return updated
+
+    def get_rollout_visualization_state(self, condition_kwargs):
+        if not condition_kwargs:
+            return None
+        return dict(condition_kwargs)
+
+    def get_visualization_trajectory(self, visualization_state):
+        if not visualization_state:
+            return None
+        return visualization_state.get("steering")
+
+    def update_rollout_visualization_state(self, visualization_state, step_idx):
+        if not visualization_state:
+            return None
+        return self.update_rollout_condition_kwargs(
+            visualization_state,
+            prediction=None,
+            context=None,
+            step_idx=step_idx,
+        )
+
+    def get_rollout_visualization_trajectory(
+        self,
+        condition_kwargs,
+        num_gen_steps,
+        num_pred_frames=1,
+        num_condition_frames=None,
+    ):
+        """Return shared trajectory geometry plus explicit per-frame cursor indices."""
+        if not condition_kwargs:
+            return None
+
+        odometry_dt = condition_kwargs.get("_odometry_dt")
+        if odometry_dt is None:
+            return None
+
+        rollout_steps = max(1, int(num_gen_steps))
+        rendered_frames_per_step = max(1, int(num_pred_frames))
+        num_context_frames = max(0, int(num_condition_frames or 0))
+
+        steering = self._validate_steering(condition_kwargs["_raw_speed_yaw"])
+        anchor_odo_index = condition_kwargs["_anchor_odo_index"]
+        rollout_step_odo = condition_kwargs["_rollout_step_odo"]
+        trajectories = []
+
+        for batch_idx in range(steering.shape[0]):
+            anchor_idx = int(anchor_odo_index[batch_idx].item())
+            step_odo = int(rollout_step_odo[batch_idx].item())
+            end_idx = anchor_idx + rollout_steps * step_odo
+            if end_idx >= steering.shape[1]:
+                raise ValueError(
+                    f"Need steering horizon of at least {end_idx + 1} odometry steps for visualization, "
+                    f"got {steering.shape[1]}. Anchor={anchor_idx}, rollout_step_odo={step_odo}, "
+                    f"num_gen_steps={num_gen_steps}, num_pred_frames={num_pred_frames}."
+                )
+
+            trajectory = get_trajectory_from_speeds_and_yaw_rates_batch(
+                speeds=steering[batch_idx : batch_idx + 1, anchor_idx : end_idx + 1, 0],
+                yaw_rates=steering[batch_idx : batch_idx + 1, anchor_idx : end_idx + 1, 1],
+                dt=odometry_dt[batch_idx : batch_idx + 1],
+            )
+            trajectories.append(trajectory)
+
+        trajectory = torch.cat(trajectories, dim=0).to(device=steering.device)
+        step_indices = torch.arange(
+            1,
+            rollout_steps + 1,
+            device=steering.device,
+            dtype=torch.long,
+        )
+        cursor_index = (step_indices * rollout_step_odo[0]).repeat_interleave(rendered_frames_per_step)
+        if num_context_frames > 0:
+            context_cursor_index = torch.zeros(
+                num_context_frames,
+                device=steering.device,
+                dtype=torch.long,
+            )
+            cursor_index = torch.cat([context_cursor_index, cursor_index], dim=0)
+        cursor_index = cursor_index.unsqueeze(0).expand(steering.shape[0], -1)
+        return {
+            "trajectory": trajectory[:, :, :2],
+            "heading": trajectory[:, :, 2],
+            "cursor_index": cursor_index,
+        }
+
+    def annotate_logged_images(self, images, batch, num_images):
+        if not isinstance(batch, dict) or "steering" not in batch:
+            return images
+
+        condition_batch = {
+            "images": batch["images"][:num_images],
+            "steering": batch["steering"][:num_images],
+        }
+        if "frame_rate" in batch:
+            frame_rate = batch["frame_rate"]
+            condition_batch["frame_rate"] = frame_rate[:num_images] if torch.is_tensor(frame_rate) else frame_rate
+        if "steering_format" in batch:
+            condition_batch["steering_format"] = batch["steering_format"]
+
+        condition_kwargs = self.get_condition_kwargs_from_batch(
+            condition_batch,
+            split="log_images",
+        )
+        steering = condition_kwargs["steering"]
+        steering_strings = []
+        for sample in steering.cpu().numpy():
+            step_strings = ["[" + ", ".join(f"{value:.1f}" for value in step) + "]" for step in sample]
+            steering_strings.append(" ".join(step_strings))
+        images[:, -1] = _text_on_image(images[:, -1], steering_strings)
+        return images
+
+
+class SpeedYawMovingGoalConditionPreprocessor(SpeedYawConditionPreprocessorBase):
+    """
+    Compute future anchored goal displacements from raw speed/yaw-rate steering using an
+    explicit anchor definition.
+
+    Assumptions:
+    - `batch["steering"]` has shape [B, T, 2] with features [speed, yaw_rate]
+    - odometry/image frame-rate ratio is fixed and integer
+    - the anchor is a configured image-frame index within the context window
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        goal_offsets,
+        odometry_steps_per_image_frame,
+        context_frame_anchor_index=-1,
+        steering_format="speed_yawrate",
+        speed_scale=1.0,
+        yaw_rate_scale=1.0,
+        **kwargs,
+    ):
+        """Configure moving-goal conditioning from raw `[speed, yaw_rate]` sequences."""
+        super().__init__(
+            predictor_module=predictor_module,
+            odometry_steps_per_image_frame=odometry_steps_per_image_frame,
+            context_frame_anchor_index=context_frame_anchor_index,
+            steering_format=steering_format,
+            speed_scale=speed_scale,
+            yaw_rate_scale=yaw_rate_scale,
+            **kwargs
+        )
+        self.goal_offsets = self._validate_offsets(goal_offsets, "goal_offsets")
+
+    def _make_no_steering_condition(self, batch_size, device):
+        return torch.full((batch_size, len(self.goal_offsets), 2), float("nan"), device=device)
+
+    def _requires_odometry_dt_for_conditioning(self):
+        return True
+
+    def _compute_condition_steering(self, steering, anchor_odo_index, odometry_dt=None):
+        """Convert raw speed/yaw windows into future goal points anchored at odometry indices."""
+        steering = self._validate_steering(steering)
+        batch_size, steering_len, _dims = steering.shape
+        device = steering.device
+        if odometry_dt is None:
+            raise ValueError(
+                "SpeedYawMovingGoalConditionPreprocessor requires `odometry_dt` to compute goal-point "
+                "conditioning."
+            )
+        odometry_dt = torch.as_tensor(odometry_dt, device=device, dtype=steering.dtype)
+        if odometry_dt.ndim == 0:
+            odometry_dt = odometry_dt.expand(batch_size)
+        elif odometry_dt.ndim != 1 or odometry_dt.shape[0] != batch_size:
+            raise ValueError(
+                "`odometry_dt` must be broadcastable to shape [B], got "
+                f"{tuple(odometry_dt.shape)} for batch_size={batch_size}."
+            )
+        max_offset = self.goal_offsets[-1]
+        goals = []
+
+        for batch_idx in range(batch_size):
+            current_anchor = int(anchor_odo_index[batch_idx].item())
+            goal_end_idx = current_anchor + max_offset
+            if current_anchor < 0:
+                raise ValueError(f"Anchor index must be non-negative, got {current_anchor}.")
+            if goal_end_idx >= steering_len:
+                raise ValueError(
+                    f"Need steering horizon of at least {goal_end_idx + 1} odometry steps, "
+                    f"got {steering_len}. Anchor={current_anchor}, max_offset={max_offset}."
+                )
+
+            centered_traj = get_trajectory_from_speeds_and_yaw_rates_batch(
+                speeds=steering[batch_idx : batch_idx + 1, current_anchor : goal_end_idx + 1, 0],
+                yaw_rates=steering[batch_idx : batch_idx + 1, current_anchor : goal_end_idx + 1, 1],
+                dt=odometry_dt[batch_idx : batch_idx + 1],
+            )
+            goals.append(centered_traj[:, self.goal_offsets, :2])
+
+        return torch.cat(goals, dim=0).to(device=device)
+
+
+class SpeedYawDirectConditionPreprocessor(SpeedYawConditionPreprocessorBase):
+    """
+    Slice future raw speed/yaw-rate commands directly from the odometry sequence.
+
+    This shares the same rollout alignment logic as moving-goal conditioning, but
+    returns raw `[speed, yaw_rate]` values instead of integrating them into XY goals.
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        steering_offsets,
+        odometry_steps_per_image_frame,
+        context_frame_anchor_index=-1,
+        steering_format="speed_yawrate",
+        speed_scale=1.0,
+        yaw_rate_scale=1.0,
+    ):
+        super().__init__(
+            predictor_module=predictor_module,
+            odometry_steps_per_image_frame=odometry_steps_per_image_frame,
+            context_frame_anchor_index=context_frame_anchor_index,
+            steering_format=steering_format,
+            speed_scale=speed_scale,
+            yaw_rate_scale=yaw_rate_scale,
+        )
+        self.steering_offsets = self._validate_offsets(steering_offsets, "steering_offsets")
+
+    def _make_no_steering_condition(self, batch_size, device):
+        return torch.full((batch_size, len(self.steering_offsets), 2), float("nan"), device=device)
+
+    def _compute_condition_steering(self, steering, anchor_odo_index, odometry_dt=None):
+        del odometry_dt
+        steering = self._validate_steering(steering)
+        batch_size, steering_len, _dims = steering.shape
+        commands = []
+        max_offset = self.steering_offsets[-1]
+        steering_offsets = torch.as_tensor(
+            self.steering_offsets,
+            device=steering.device,
+            dtype=torch.long,
+        )
+
+        for batch_idx in range(batch_size):
+            current_anchor = int(anchor_odo_index[batch_idx].item())
+            end_idx = current_anchor + max_offset
+            if current_anchor < 0:
+                raise ValueError(f"Anchor index must be non-negative, got {current_anchor}.")
+            if end_idx >= steering_len:
+                raise ValueError(
+                    f"Need steering horizon of at least {end_idx + 1} odometry steps, "
+                    f"got {steering_len}. Anchor={current_anchor}, max_offset={max_offset}."
+                )
+            commands.append(steering[batch_idx : batch_idx + 1, current_anchor + steering_offsets])
+
+        return torch.cat(commands, dim=0).to(device=steering.device)
+
+
+class L2EndpointConditionPreprocessor(ConditionPreprocessor):
+    """
+    Inference-only L2 endpoint conditioning for the CleanCtx PredL2 v1 model.
+
+    The preprocessor owns the frozen L2 predictor and supplies only
+    `z_l2_start` / `z_l2_end` to the DiT. Private keys are retained as rollout
+    state and filtered before model calls by ConditionPreprocessor.
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        l2_predictor_config,
+        l2_predictor_frame_rate=1.0,
+        l2_context_latent=False,
+        l2_pred_NFE=10,
+        num_context_frames=1,
+        use_z_l2_start=True,
+    ):
+        super().__init__(predictor_module)
+        self.l2_predictor_frame_rate = float(l2_predictor_frame_rate)
+        self.l2_context_latent = bool(l2_context_latent)
+        self.l2_pred_NFE = int(l2_pred_NFE)
+        self.num_context_frames = int(num_context_frames)
+        self.use_z_l2_start = bool(use_z_l2_start)
+        self.l2_predictor_config = l2_predictor_config
+        self.l2_predictor = self._build_l2_predictor(l2_predictor_config)
+        self.l2_predictor_encoder_branch = self._get_l2_predictor_encoder_branch(
+            predictor=self.l2_predictor,
+            predictor_config=l2_predictor_config,
+        )
+
+    # --- BEGIN: rollout debug path info (safe to remove) ---
+    def get_l2_predictor_path_info(self):
+        """Return (exp_folder_name, ckpt_name) describing where the L2 predictor was loaded from."""
+        folder = os.path.expandvars(self.l2_predictor_config.folder)
+        ckpt_path = self.l2_predictor_config.ckpt_path if self.l2_predictor_config.ckpt_path else "checkpoints/last.ckpt"
+        return os.path.basename(os.path.normpath(folder)), os.path.basename(ckpt_path)
+    # --- END: rollout debug path info (safe to remove) ---
+
+    @property
+    def ae(self):
+        return self.module.ae
+
+    @property
+    def enc_scale(self):
+        return self.module.enc_scale
+
+    @property
+    def enc_scale_dino(self):
+        return self.module.enc_scale_dino
+
+    def _build_l2_predictor(self, cfg):
+        folder = os.path.expandvars(cfg.folder)
+        ckpt_path = cfg.ckpt_path if cfg.ckpt_path else "checkpoints/last.ckpt"
+        logging.info(f"Loading L2 checkpoint from {os.path.join(folder, ckpt_path)}")
+        model_cfg = OmegaConf.load(os.path.join(folder, "config.yaml"))
+        predictor = instantiate_from_config(model_cfg.model)
+        state_dict = torch.load(
+            os.path.join(folder, ckpt_path),
+            map_location="cpu",
+            weights_only=True,
+        )["state_dict"]
+        predictor.load_state_dict(state_dict, strict=False)
+        predictor.eval()
+        _requires_grad(predictor, False)
+        return predictor
+
+    def _normalize_l2_predictor_encoder_branch(self, encoder_branch):
+        if isinstance(encoder_branch, int):
+            if encoder_branch == 0:
+                return "rec"
+            if encoder_branch == 1:
+                return "sem"
+
+        branch = str(encoder_branch).strip().lower()
+        aliases = {
+            "rec": "rec",
+            "x0": "rec",
+            "h": "rec",
+            "sem": "sem",
+            "x1": "sem",
+            "h2": "sem",
+        }
+        if branch not in aliases:
+            raise ValueError(
+                f"Unsupported delegated L2 encoder_branch={encoder_branch!r}. "
+                "Expected one of: rec, x0, h, sem, x1, h2."
+            )
+        return aliases[branch]
+
+    def _get_l2_predictor_encoder_branch(self, predictor, predictor_config):
+        first_stage = getattr(predictor, "first_stage", None)
+        if first_stage is not None and hasattr(first_stage, "encoder_branch"):
+            return self._normalize_l2_predictor_encoder_branch(first_stage.encoder_branch)
+
+        if hasattr(predictor, "encoder_branch"):
+            return self._normalize_l2_predictor_encoder_branch(predictor.encoder_branch)
+
+        folder = os.path.expandvars(predictor_config.folder)
+        model_cfg = OmegaConf.load(os.path.join(folder, "config.yaml"))
+        first_stage_cfg = model_cfg.model.params.get("first_stage_handler_config")
+        if first_stage_cfg and "params" in first_stage_cfg and "encoder_branch" in first_stage_cfg.params:
+            return self._normalize_l2_predictor_encoder_branch(first_stage_cfg.params.encoder_branch)
+
+        if "encoder_branch" in model_cfg.model.params:
+            return self._normalize_l2_predictor_encoder_branch(model_cfg.model.params.encoder_branch)
+
+        raise ValueError(
+            "Could not determine the delegated L2 predictor encoder branch from the loaded model "
+            "or its config. Expected a rec/sem branch selection."
+        )
+
+    def to_device(self, device):
+        self.l2_predictor.to(device)
+
+    def get_l2_predictor_encoder_branch(self):
+        return self.l2_predictor_encoder_branch
+
+    def get_l2_predictor_latent_scale(self):
+        if self.l2_predictor_encoder_branch == "rec":
+            return self.enc_scale
+        return self.enc_scale_dino
+
+    def _reject_training_split(self, split):
+        if split in {"train", "val"}:
+            raise RuntimeError(
+                "L2EndpointConditionPreprocessor is inference-only. "
+                "It supports sample/log_images/rollout conditioning, not train/val."
+            )
+
+    def _validate_guidance_scale(self, condition_kwargs):
+        if not condition_kwargs or "l2_guidance_scale" not in condition_kwargs:
+            return
+        guidance_scale = condition_kwargs["l2_guidance_scale"]
+        if torch.is_tensor(guidance_scale):
+            guidance_scale = float(guidance_scale.detach().cpu().item())
+        else:
+            guidance_scale = float(guidance_scale)
+        if abs(guidance_scale - 1.0) > 1e-6:
+            raise RuntimeError(
+                "L2EndpointConditionPreprocessor only supports "
+                "l2_guidance_scale == 1.0."
+            )
+
+    @torch.no_grad()
+    def _encode_l2_predictor_frames(self, images):
+        if images.ndim == 5:
+            b, f, _c, _h, _w = images.size()
+            images = rearrange(images, "b f c h w -> (b f) c h w")
+        else:
+            b, _c, _h, _w = images.size()
+            f = 1
+
+        continuous = self.ae.encode(images)["continuous"]
+        if not isinstance(continuous, tuple) or len(continuous) != 2:
+            raise ValueError(
+                "L2EndpointConditionPreprocessor expects tokenizer.encode(...)[\"continuous\"] "
+                f"to return a 2-tuple, got {type(continuous)}."
+            )
+        branch_index = 0 if self.l2_predictor_encoder_branch == "rec" else 1
+        latent_scale = self.get_l2_predictor_latent_scale()
+        latents = continuous[branch_index] * latent_scale
+        return rearrange(latents, "(b f) c h w -> b f c h w", b=b, f=f)
+
+    @torch.no_grad()
+    def _encode_l2_context(self, l2_context):
+        if self.l2_context_latent:
+            return l2_context * self.get_l2_predictor_latent_scale()
+        return self._encode_l2_predictor_frames(l2_context)
+
+    @torch.no_grad()
+    def _predict_l2_end(self, l2_context_latent):
+        batch_size = l2_context_latent.shape[0]
+        device = l2_context_latent.device
+        if next(self.l2_predictor.parameters()).device != device:
+            self.l2_predictor.to(device)
+        predictor = self.l2_predictor
+        frame_rate = torch.full(
+            (batch_size,),
+            self.l2_predictor_frame_rate,
+            device=device,
+        )
+        pred = predictor.sample(
+            images=l2_context_latent,
+            latent=True,
+            eta=0.0,
+            NFE=self.l2_pred_NFE,
+            sample_with_ema=True,
+            num_samples=batch_size,
+            frame_rate=frame_rate,
+        )
+        return pred[:, -1]
+
+    @torch.no_grad()
+    def _sample_l2_next_latent(self, l2_context_latent):
+        return self._predict_l2_end(l2_context_latent)
+
+    def _get_z_l2_start_from_context(self, l2_context_latent, z_l2_end):
+        if self.use_z_l2_start:
+            if l2_context_latent is None:
+                raise ValueError(
+                    "L2EndpointConditionPreprocessor requires z_l2_start or "
+                    "l2_context/_l2_context_latent when use_z_l2_start=True."
+                )
+            return l2_context_latent[:, -1]
+        return torch.zeros_like(z_l2_end)
+
+    def get_condition_kwargs_from_batch(self, batch, split):
+        self._reject_training_split(split)
+        if not isinstance(batch, dict) or "l2_context" not in batch:
+            raise ValueError(
+                "L2EndpointConditionPreprocessor requires batch['l2_context'] "
+                f"for split={split!r}."
+            )
+        return {"l2_context": batch["l2_context"]}
+
+    def _move_public_condition_kwargs(self, condition_kwargs, device):
+        prepared = {}
+        for key, value in condition_kwargs.items():
+            if key == "l2_guidance_scale":
+                continue
+            if torch.is_tensor(value) and device is not None:
+                prepared[key] = value.to(device)
+            else:
+                prepared[key] = value
+        return prepared
+
+    @torch.no_grad()
+    def prepare_condition_kwargs(self, condition_kwargs=None, batch_size=None, device=None, split=None):
+        self._reject_training_split(split)
+        condition_kwargs = condition_kwargs or {}
+        self._validate_guidance_scale(condition_kwargs)
+        prepared = self._move_public_condition_kwargs(condition_kwargs, device)
+
+        has_start = "z_l2_start" in prepared
+        has_end = "z_l2_end" in prepared
+        prepared["_z_l2_start_fixed"] = bool(prepared.get("_z_l2_start_fixed", has_start))
+        prepared["_z_l2_end_fixed"] = bool(prepared.get("_z_l2_end_fixed", has_end))
+
+        l2_context_latent = prepared.get("_l2_context_latent")
+        if l2_context_latent is None and "l2_context" in prepared:
+            l2_context_latent = self._encode_l2_context(prepared.pop("l2_context"))
+            prepared["_l2_context_latent"] = l2_context_latent
+        elif l2_context_latent is not None:
+            prepared["_l2_context_latent"] = l2_context_latent
+
+        if "z_l2_end" not in prepared:
+            if l2_context_latent is None:
+                raise ValueError(
+                    "L2EndpointConditionPreprocessor requires z_l2_end or "
+                    "l2_context/_l2_context_latent to predict it."
+                )
+            prepared["_l2_next_latent"] = self._sample_l2_next_latent(l2_context_latent)
+            prepared["z_l2_end"] = prepared["_l2_next_latent"]
+        elif "_l2_next_latent" not in prepared and not prepared["_z_l2_end_fixed"]:
+            prepared["_l2_next_latent"] = prepared["z_l2_end"]
+
+        if "z_l2_start" not in prepared:
+            prepared["z_l2_start"] = self._get_z_l2_start_from_context(
+                l2_context_latent,
+                prepared["z_l2_end"],
+            )
+
+        if batch_size is not None:
+            for key in ("z_l2_start", "z_l2_end"):
+                if prepared[key].shape[0] != batch_size:
+                    raise ValueError(
+                        f"{key} batch size {prepared[key].shape[0]} does not match "
+                        f"expected batch_size={batch_size}."
+                    )
+
+        return prepared
+
+    @torch.no_grad()
+    def update_rollout_condition_kwargs(self, condition_kwargs, prediction, context, step_idx):
+        del prediction, context, step_idx
+        if not condition_kwargs:
+            return {}
+
+        updated = dict(condition_kwargs)
+        l2_context_latent = updated.get("_l2_context_latent")
+        start_fixed = bool(updated.get("_z_l2_start_fixed", False))
+        end_fixed = bool(updated.get("_z_l2_end_fixed", False))
+        if l2_context_latent is None or (start_fixed and end_fixed):
+            return updated
+
+        l2_next_latent = updated.get("_l2_next_latent")
+        if l2_next_latent is None:
+            if end_fixed:
+                l2_next_latent = updated["z_l2_end"]
+            else:
+                raise ValueError(
+                    "L2EndpointConditionPreprocessor requires cached `_l2_next_latent` "
+                    "or explicit `z_l2_end` during rollout updates."
+                )
+        l2_next_latent = l2_next_latent.to(l2_context_latent.device)
+        updated["_l2_context_latent"] = torch.cat(
+            [l2_context_latent[:, 1:], l2_next_latent.unsqueeze(1)],
+            dim=1,
+        )
+        if not start_fixed:
+            updated["z_l2_start"] = self._get_z_l2_start_from_context(
+                updated["_l2_context_latent"],
+                updated["z_l2_end"],
+            )
+        if not end_fixed:
+            updated["_l2_next_latent"] = self._sample_l2_next_latent(updated["_l2_context_latent"])
+            updated["z_l2_end"] = updated["_l2_next_latent"]
+        return updated
+
+
+class L2SteeringEndpointConditionPreprocessor(L2EndpointConditionPreprocessor):
+    """
+    Steering-aware variant of L2EndpointConditionPreprocessor.
+
+    Raw speed/yaw-rate steering is converted into the conditioning expected by
+    the frozen L2 predictor. That steering state remains private to the L2
+    branch and is updated in lockstep with the autoregressive L2-context slide.
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        l2_predictor_config,
+        l2_predictor_frame_rate=1.0,
+        l2_context_latent=False,
+        l2_pred_NFE=10,
+        num_context_frames=1,
+        use_z_l2_start=True,
+        speed_scale=1.0,
+        yaw_rate_scale=1.0,
+    ):
+        super().__init__(
+            predictor_module=predictor_module,
+            l2_predictor_config=l2_predictor_config,
+            l2_predictor_frame_rate=l2_predictor_frame_rate,
+            l2_context_latent=l2_context_latent,
+            l2_pred_NFE=l2_pred_NFE,
+            num_context_frames=num_context_frames,
+            use_z_l2_start=use_z_l2_start,
+        )
+        self.speed_scale = _validate_speed_yaw_scale(speed_scale, "speed_scale")
+        self.yaw_rate_scale = _validate_speed_yaw_scale(yaw_rate_scale, "yaw_rate_scale")
+        self.l2_condition_preprocessor = self._get_l2_condition_preprocessor()
+        if hasattr(self.l2_condition_preprocessor, "goal_offsets"):
+            self.goal_offsets = [int(offset) for offset in self.l2_condition_preprocessor.goal_offsets]
+        if hasattr(self.l2_condition_preprocessor, "steering_offsets"):
+            self.steering_offsets = [int(offset) for offset in self.l2_condition_preprocessor.steering_offsets]
+        self._validate_l2_condition_preprocessor()
+
+    def _get_l2_condition_preprocessor(self):
+        preprocessor = getattr(self.l2_predictor, "condition_preprocessor", None)
+        if preprocessor is None:
+            raise TypeError(
+                "L2SteeringEndpointConditionPreprocessor requires the frozen L2 predictor "
+                "to expose `condition_preprocessor`."
+            )
+        return preprocessor
+
+    def _validate_l2_condition_preprocessor(self):
+        required_methods = (
+            "get_condition_kwargs_from_batch",
+            "prepare_condition_kwargs",
+            "update_rollout_condition_kwargs",
+        )
+        missing = [
+            method_name
+            for method_name in required_methods
+            if not callable(getattr(self.l2_condition_preprocessor, method_name, None))
+        ]
+        if missing:
+            raise TypeError(
+                "Frozen L2 condition preprocessor is incompatible with steering delegation; "
+                f"missing methods: {missing}."
+            )
+
+    def get_max_condition_odometry_offset(self):
+        if hasattr(self.l2_condition_preprocessor, "get_max_condition_odometry_offset"):
+            return self.l2_condition_preprocessor.get_max_condition_odometry_offset()
+        return None
+
+    def get_required_rollout_odometry_steps(
+        self,
+        *,
+        validation_params,
+        num_condition_frames,
+        num_gen_frames,
+        rollout_steps,
+    ):
+        del num_condition_frames
+        delegated_required_steps = getattr(self.l2_condition_preprocessor, "get_required_rollout_odometry_steps", None)
+        if callable(delegated_required_steps):
+            delegated_num_pred_frames = int(getattr(self.l2_predictor, "num_pred_frames", 1))
+            return delegated_required_steps(
+                validation_params=validation_params,
+                num_condition_frames=int(self.num_context_frames),
+                num_gen_frames=int(rollout_steps) * delegated_num_pred_frames,
+                rollout_steps=rollout_steps,
+            )
+
+        del num_gen_frames
+        max_condition_offset = self.get_max_condition_odometry_offset()
+        if max_condition_offset is None:
+            return None
+
+        source_odometry_rate = float(
+            getattr(validation_params, "odometry_frame_rate", getattr(validation_params, "frame_rate"))
+        )
+        target_odometry_rate = float(self.l2_predictor_frame_rate)
+        ratio = source_odometry_rate / target_odometry_rate
+        rounded_ratio = round(ratio)
+        if abs(ratio - rounded_ratio) > 1e-8:
+            raise ValueError(
+                "Delegated L2 rollout requires the parent odometry rate to be an integer multiple of "
+                f"the frozen L2 rate, got odometry_frame_rate={source_odometry_rate} and "
+                f"l2_predictor_frame_rate={target_odometry_rate}."
+            )
+
+        required_l2_steps = self.num_context_frames + int(rollout_steps) + int(max_condition_offset) - 1
+        return int((required_l2_steps - 1) * rounded_ratio + 1)
+
+    def _maybe_resample_raw_steering_for_l2(self, steering, batch):
+        if not torch.is_tensor(steering):
+            raise TypeError(f"`steering` must be a tensor, got {type(steering).__name__}.")
+        if "frame_rate" not in batch:
+            return steering
+
+        source_frame_rate = torch.as_tensor(
+            batch["frame_rate"],
+            device=steering.device,
+            dtype=torch.float32,
+        )
+        if source_frame_rate.ndim == 0:
+            source_frame_rate = source_frame_rate.expand(steering.shape[0])
+        elif source_frame_rate.ndim == 1 and source_frame_rate.shape[0] == 1 and steering.shape[0] != 1:
+            source_frame_rate = source_frame_rate.expand(steering.shape[0])
+        elif source_frame_rate.ndim != 1 or source_frame_rate.shape[0] != steering.shape[0]:
+            raise ValueError(
+                "`frame_rate` must be broadcastable to shape [B] for delegated L2 steering, got "
+                f"{tuple(source_frame_rate.shape)} for batch_size={steering.shape[0]}."
+            )
+
+        target_frame_rate = torch.full_like(source_frame_rate, float(self.l2_predictor_frame_rate))
+        if torch.any(target_frame_rate <= 0):
+            raise ValueError(f"`l2_predictor_frame_rate` must be positive, got {self.l2_predictor_frame_rate}.")
+
+        if torch.allclose(source_frame_rate, target_frame_rate):
+            return steering
+
+        ratios = source_frame_rate / target_frame_rate
+        rounded_ratios = torch.round(ratios)
+        if not torch.allclose(ratios, rounded_ratios, atol=1e-8, rtol=0.0):
+            raise ValueError(
+                "Delegated L2 steering requires the parent frame rate to be an integer multiple of "
+                f"the frozen L2 frame rate. Got source_frame_rate={source_frame_rate.tolist()} and "
+                f"l2_predictor_frame_rate={self.l2_predictor_frame_rate}."
+            )
+
+        resampled = []
+        for batch_idx in range(steering.shape[0]):
+            step = int(rounded_ratios[batch_idx].item())
+            if step <= 0:
+                raise ValueError(
+                    f"Delegated L2 steering subsample step must be positive, got {step}."
+                )
+            resampled.append(steering[batch_idx : batch_idx + 1, ::step])
+        return torch.cat(resampled, dim=0)
+
+    def _build_l2_rollout_batch(self, batch):
+        if "l2_context" not in batch:
+            raise ValueError(
+                "L2SteeringEndpointConditionPreprocessor requires batch['l2_context'] "
+                "to build delegated L2 conditions."
+            )
+        if "steering" not in batch:
+            raise ValueError(
+                "L2SteeringEndpointConditionPreprocessor requires batch['steering'] "
+                "to build delegated L2 conditions."
+            )
+        source_steering = batch["steering"]
+        scaled_steering = _apply_speed_yaw_scales(
+            source_steering,
+            speed_scale=self.speed_scale,
+            yaw_rate_scale=self.yaw_rate_scale,
+        )
+        l2_batch = {
+            "images": batch["l2_context"],
+            # Keep steering at raw odometry resolution and let the frozen L2
+            # preprocessor interpret it via its own odometry alignment config.
+            "steering": scaled_steering,
+        }
+        if "steering_format" in batch:
+            l2_batch["steering_format"] = batch["steering_format"]
+        if "frame_rate" in batch:
+            frame_rate = torch.as_tensor(
+                batch["frame_rate"],
+                device=batch["steering"].device,
+                dtype=torch.float32,
+            )
+            l2_batch["frame_rate"] = torch.full_like(
+                frame_rate,
+                float(self.l2_predictor_frame_rate),
+                dtype=torch.float32,
+            )
+        else:
+            batch_size = batch["steering"].shape[0]
+            l2_batch["frame_rate"] = torch.full(
+                (batch_size,),
+                float(self.l2_predictor_frame_rate),
+                device=batch["steering"].device,
+                dtype=torch.float32,
+            )
+        return l2_batch
+
+    def _prepare_l2_predictor_condition_kwargs(self, l2_condition_kwargs, device):
+        return self.l2_condition_preprocessor.prepare_condition_kwargs(
+            l2_condition_kwargs,
+            batch_size=None,
+            device=device,
+            split="rollout",
+        )
+
+    def _update_l2_predictor_condition_kwargs(self, l2_condition_kwargs):
+        return self.l2_condition_preprocessor.update_rollout_condition_kwargs(
+            l2_condition_kwargs,
+            prediction=None,
+            context=None,
+            step_idx=0,
+        )
+
+    @torch.no_grad()
+    def _predict_l2_end(self, l2_context_latent, l2_condition_kwargs=None):
+        batch_size = l2_context_latent.shape[0]
+        device = l2_context_latent.device
+        if next(self.l2_predictor.parameters()).device != device:
+            self.l2_predictor.to(device)
+        predictor = self.l2_predictor
+        frame_rate = torch.full(
+            (batch_size,),
+            self.l2_predictor_frame_rate,
+            device=device,
+        )
+        pred = predictor.sample(
+            images=l2_context_latent,
+            latent=True,
+            eta=0.0,
+            NFE=self.l2_pred_NFE,
+            sample_with_ema=True,
+            num_samples=batch_size,
+            frame_rate=frame_rate,
+            condition_kwargs=l2_condition_kwargs,
+        )
+        return pred[:, -1]
+
+    @torch.no_grad()
+    def _sample_l2_next_latent(self, l2_context_latent, l2_condition_kwargs=None):
+        return self._predict_l2_end(l2_context_latent, l2_condition_kwargs=l2_condition_kwargs)
+
+    def get_condition_kwargs_from_batch(self, batch, split):
+        self._reject_training_split(split)
+        if not isinstance(batch, dict) or "l2_context" not in batch:
+            raise ValueError(
+                "L2SteeringEndpointConditionPreprocessor requires batch['l2_context'] "
+                f"for split={split!r}."
+            )
+        if "steering" not in batch:
+            raise ValueError(
+                "L2SteeringEndpointConditionPreprocessor requires batch['steering'] "
+                f"for split={split!r}."
+            )
+        l2_context = batch["l2_context"]
+        if l2_context.ndim != 5:
+            raise ValueError(f"`l2_context` must have shape [B, F, C, H, W], got {tuple(l2_context.shape)}.")
+        source_l2_steering = self._maybe_resample_raw_steering_for_l2(batch["steering"], batch)
+        l2_condition_kwargs = self.l2_condition_preprocessor.get_condition_kwargs_from_batch(
+            self._build_l2_rollout_batch(batch),
+            split="rollout",
+        )
+        l2_condition_kwargs["_source_raw_speed_yaw"] = source_l2_steering
+        return {
+            "l2_context": l2_context,
+            "_l2_condition_kwargs": l2_condition_kwargs,
+        }
+
+    @torch.no_grad()
+    def prepare_condition_kwargs(self, condition_kwargs=None, batch_size=None, device=None, split=None):
+        self._reject_training_split(split)
+        condition_kwargs = condition_kwargs or {}
+        self._validate_guidance_scale(condition_kwargs)
+        prepared = self._move_public_condition_kwargs(condition_kwargs, device)
+
+        has_start = "z_l2_start" in prepared
+        has_end = "z_l2_end" in prepared
+        prepared["_z_l2_start_fixed"] = bool(prepared.get("_z_l2_start_fixed", has_start))
+        prepared["_z_l2_end_fixed"] = bool(prepared.get("_z_l2_end_fixed", has_end))
+
+        l2_context_latent = prepared.get("_l2_context_latent")
+        if l2_context_latent is None and "l2_context" in prepared:
+            l2_context_latent = self._encode_l2_context(prepared.pop("l2_context"))
+            prepared["_l2_context_latent"] = l2_context_latent
+        elif l2_context_latent is not None:
+            prepared["_l2_context_latent"] = l2_context_latent
+
+        l2_condition_kwargs = self._prepare_l2_predictor_condition_kwargs(
+            prepared.get("_l2_condition_kwargs"),
+            device=device,
+        )
+        prepared["_l2_condition_kwargs"] = l2_condition_kwargs
+
+        if "z_l2_end" not in prepared:
+            if l2_context_latent is None:
+                raise ValueError(
+                    "L2SteeringEndpointConditionPreprocessor requires z_l2_end or "
+                    "l2_context/_l2_context_latent to predict it."
+                )
+            prepared["_l2_next_latent"] = self._sample_l2_next_latent(
+                l2_context_latent,
+                l2_condition_kwargs=l2_condition_kwargs,
+            )
+            prepared["z_l2_end"] = prepared["_l2_next_latent"]
+        elif "_l2_next_latent" not in prepared and not prepared["_z_l2_end_fixed"]:
+            prepared["_l2_next_latent"] = prepared["z_l2_end"]
+
+        if "z_l2_start" not in prepared:
+            prepared["z_l2_start"] = self._get_z_l2_start_from_context(
+                l2_context_latent,
+                prepared["z_l2_end"],
+            )
+
+        if batch_size is not None:
+            for key in ("z_l2_start", "z_l2_end"):
+                if prepared[key].shape[0] != batch_size:
+                    raise ValueError(
+                        f"{key} batch size {prepared[key].shape[0]} does not match "
+                        f"expected batch_size={batch_size}."
+                    )
+
+        return prepared
+
+    @torch.no_grad()
+    def update_rollout_condition_kwargs(self, condition_kwargs, prediction, context, step_idx):
+        del prediction, context, step_idx
+        if not condition_kwargs:
+            return {}
+
+        updated = dict(condition_kwargs)
+        l2_context_latent = updated.get("_l2_context_latent")
+        start_fixed = bool(updated.get("_z_l2_start_fixed", False))
+        end_fixed = bool(updated.get("_z_l2_end_fixed", False))
+        if l2_context_latent is None:
+            return updated
+
+        if not (start_fixed and end_fixed):
+            l2_next_latent = updated.get("_l2_next_latent")
+            if l2_next_latent is None:
+                if end_fixed:
+                    l2_next_latent = updated["z_l2_end"]
+                else:
+                    raise ValueError(
+                        "L2SteeringEndpointConditionPreprocessor requires cached `_l2_next_latent` "
+                        "or explicit `z_l2_end` during rollout updates."
+                    )
+            l2_next_latent = l2_next_latent.to(l2_context_latent.device)
+            updated["_l2_context_latent"] = torch.cat(
+                [l2_context_latent[:, 1:], l2_next_latent.unsqueeze(1)],
+                dim=1,
+            )
+
+        if "_l2_condition_kwargs" in updated:
+            updated["_l2_condition_kwargs"] = self._update_l2_predictor_condition_kwargs(
+                updated["_l2_condition_kwargs"],
+            )
+
+        if not start_fixed:
+            updated["z_l2_start"] = self._get_z_l2_start_from_context(
+                updated["_l2_context_latent"],
+                updated["z_l2_end"],
+            )
+
+        if not end_fixed:
+            updated["_l2_next_latent"] = self._sample_l2_next_latent(
+                updated["_l2_context_latent"],
+                l2_condition_kwargs=updated.get("_l2_condition_kwargs"),
+            )
+            updated["z_l2_end"] = updated["_l2_next_latent"]
+        return updated
+
+    def get_rollout_visualization_state(self, condition_kwargs):
+        if not condition_kwargs:
+            return None
+        l2_condition_kwargs = condition_kwargs.get("_l2_condition_kwargs")
+        if not l2_condition_kwargs:
+            return None
+        return self.l2_condition_preprocessor.get_rollout_visualization_state(l2_condition_kwargs)
+
+    def get_visualization_trajectory(self, visualization_state):
+        if visualization_state is None:
+            return None
+        return self.l2_condition_preprocessor.get_visualization_trajectory(visualization_state)
+
+    def update_rollout_visualization_state(self, visualization_state, step_idx):
+        if visualization_state is None:
+            return None
+        return self.l2_condition_preprocessor.update_rollout_visualization_state(
+            visualization_state,
+            step_idx=step_idx,
+        )
+
+    def get_rollout_visualization_trajectory(
+        self,
+        condition_kwargs,
+        num_gen_steps,
+        num_pred_frames=1,
+        num_condition_frames=None,
+    ):
+        if not condition_kwargs:
+            return None
+        l2_condition_kwargs = condition_kwargs.get("_l2_condition_kwargs")
+        if not l2_condition_kwargs:
+            return None
+        return self.l2_condition_preprocessor.get_rollout_visualization_trajectory(
+            l2_condition_kwargs,
+            num_gen_steps=num_gen_steps,
+            num_pred_frames=num_pred_frames,
+            num_condition_frames=num_condition_frames,
+        )

orbis2/models/second_stage/fm_model_v2.py

@@ -0,0 +1,1319 @@
+import math
+import os
+from collections import OrderedDict
+from copy import deepcopy
+
+import pytorch_lightning as pl
+import torch
+import torchvision.utils as vutils
+from einops import rearrange
+from omegaconf import ListConfig, OmegaConf
+from timm.layers.pos_embed import resample_abs_pos_embed
+from torch.optim.lr_scheduler import LambdaLR
+from tqdm import tqdm
+
+from util import get_obj_from_str, instantiate_from_config
+from .fm_conditions_v2 import (
+    ConditionPreprocessor,
+    L2EndpointConditionPreprocessor,
+    SpeedYawDirectConditionPreprocessor,
+    SpeedYawMovingGoalConditionPreprocessor,
+)
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+def requires_grad(model, flag=True):
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def init_ema_model(model):
+    ema_model = deepcopy(model)
+    requires_grad(ema_model, False)
+    update_ema(ema_model, model, decay=0)
+    ema_model.eval()
+    return ema_model
+
+
+def _extract_tokenizer_decode_samples(decoded):
+    if isinstance(decoded, tuple):
+        return decoded[0]
+    return decoded
+
+
+def _decode_selected_if_branch(tokenizer, branch, latents):
+    if branch == "rec":
+        decoder = getattr(tokenizer, "decoder", None)
+        post_quant_conv = getattr(tokenizer, "post_quant_conv", None)
+    else:
+        decoder = getattr(tokenizer, "decoder_sem", None)
+        post_quant_conv = getattr(tokenizer, "post_quant_conv_sem", None)
+
+    if (
+        decoder is not None
+        and post_quant_conv is not None
+        and getattr(post_quant_conv, "in_channels", None) == latents.size(1)
+    ):
+        return _extract_tokenizer_decode_samples(decoder(post_quant_conv(latents)))
+
+    rec_channels = tokenizer.quant_conv.out_channels
+    sem_channels = tokenizer.quant_conv2.out_channels
+    if branch == "rec":
+        quant = (
+            latents,
+            torch.zeros(
+                latents.size(0),
+                sem_channels,
+                latents.size(2),
+                latents.size(3),
+                device=latents.device,
+                dtype=latents.dtype,
+            ),
+        )
+    else:
+        quant = (
+            torch.zeros(
+                latents.size(0),
+                rec_channels,
+                latents.size(2),
+                latents.size(3),
+                device=latents.device,
+                dtype=latents.dtype,
+            ),
+            latents,
+        )
+    return _extract_tokenizer_decode_samples(tokenizer.decode(quant))
+
+
+class FirstStageHandler:
+    def __init__(self, predictor_module):
+        self.module = predictor_module
+
+    @property
+    def ae(self):
+        return self.module.ae
+
+    @property
+    def vit(self):
+        return self.module.vit
+
+    @property
+    def enc_scale(self):
+        return self.module.enc_scale
+
+    @property
+    def use_precomputed_training_inputs(self):
+        return self.module.use_precomputed_training_inputs
+
+    @torch.no_grad()
+    def encode_frames(self, images):
+        if self.use_precomputed_training_inputs and images.shape[-2:] == self.vit.input_size:
+            return images * self.enc_scale
+
+        if images.ndim == 5:
+            b, f, _e, _h, _w = images.size()
+            images = rearrange(images, "b f e h w -> (b f) e h w")
+        else:
+            b, _e, _h, _w = images.size()
+            f = 1
+
+        x = self.ae.encode(images)["continuous"]
+        x = x * self.enc_scale
+        return rearrange(x, "(b f) e h w -> b f e h w", b=b, f=f)
+
+    @torch.no_grad()
+    def decode_frames(self, x, output_device=None):
+        samples = []
+        for idx in range(x.shape[1]):
+            frame = x[:, idx] / self.enc_scale
+            frame = self.ae.post_quant_conv(frame)
+            frame = self.ae.decoder(frame)
+            if output_device is not None:
+                frame = frame.to(output_device)
+            samples.append(frame.unsqueeze(1))
+        return torch.cat(samples, dim=1)
+
+
+class IFSelectBranchFirstStageHandler(FirstStageHandler):
+    """
+    First-stage handler that selects a single IF tokenizer branch, mirroring
+    `ModelIFSelectBranch` from the v1 codepath.
+    """
+
+    def __init__(self, predictor_module, encoder_branch="rec"):
+        super().__init__(predictor_module)
+        self.encoder_branch = self._normalize_encoder_branch(encoder_branch)
+        expected_channels = self._selected_branch_channels()
+        actual_channels = getattr(self.vit, "in_channels", None)
+        if actual_channels is not None and actual_channels != expected_channels:
+            raise ValueError(
+                f"generator_config.params.in_channels={actual_channels} does not match "
+                f"the selected {self.encoder_branch} branch channels ({expected_channels})."
+            )
+
+    @property
+    def enc_scale_dino(self):
+        return self.module.enc_scale_dino
+
+    def _normalize_encoder_branch(self, encoder_branch):
+        if isinstance(encoder_branch, int):
+            if encoder_branch == 0:
+                return "rec"
+            if encoder_branch == 1:
+                return "sem"
+
+        branch = str(encoder_branch).strip().lower()
+        aliases = {
+            "rec": "rec",
+            "x0": "rec",
+            "h": "rec",
+            "sem": "sem",
+            "x1": "sem",
+            "h2": "sem",
+        }
+        if branch not in aliases:
+            raise ValueError(
+                f"Unsupported encoder_branch={encoder_branch!r}. "
+                "Expected one of: rec, x0, h, sem, x1, h2."
+            )
+        return aliases[branch]
+
+    def _selected_branch_index(self):
+        return 0 if self.encoder_branch == "rec" else 1
+
+    def _selected_branch_scale(self):
+        return self.enc_scale if self.encoder_branch == "rec" else self.enc_scale_dino
+
+    def _selected_branch_channels(self):
+        if self.encoder_branch == "rec":
+            return self.ae.quant_conv.out_channels
+        return self.ae.quant_conv2.out_channels
+
+    @torch.no_grad()
+    def encode_frames(self, images):
+        if self.use_precomputed_training_inputs and images.shape[-2:] == self.vit.input_size:
+            return images * self._selected_branch_scale()
+
+        if images.ndim == 5:
+            b, f, _e, _h, _w = images.size()
+            images = rearrange(images, "b f e h w -> (b f) e h w")
+        else:
+            b, _e, _h, _w = images.size()
+            f = 1
+
+        continuous = self.ae.encode(images)["continuous"]
+        if not isinstance(continuous, tuple) or len(continuous) != 2:
+            raise ValueError(
+                "IFSelectBranchFirstStageHandler expects tokenizer.encode(...)[\"continuous\"] "
+                f"to return a 2-tuple, got {type(continuous)}."
+            )
+
+        x = continuous[self._selected_branch_index()] * self._selected_branch_scale()
+        expected_channels = self._selected_branch_channels()
+        if x.size(1) != expected_channels:
+            raise ValueError(
+                f"Selected {self.encoder_branch} branch has {x.size(1)} channels, "
+                f"expected {expected_channels}."
+            )
+
+        return rearrange(x, "(b f) e h w -> b f e h w", b=b, f=f)
+
+    @torch.no_grad()
+    def decode_frames(self, x, output_device=None):
+        b, f, c, _h, _w = x.size()
+
+        expected_channels = self._selected_branch_channels()
+        if c != expected_channels:
+            raise ValueError(
+                f"decode_frames expected {expected_channels} channels for "
+                f"encoder_branch={self.encoder_branch}, got {c}."
+            )
+
+        if output_device is None:
+            x = rearrange(x, "b f c h w -> (b f) c h w")
+            selected = x / self._selected_branch_scale()
+            samples = _decode_selected_if_branch(self.ae, self.encoder_branch, selected)
+            return rearrange(samples, "(b f) c h w -> b f c h w", b=b, f=f)
+
+        samples = []
+        for idx in range(f):
+            selected = x[:, idx] / self._selected_branch_scale()
+            frame = _decode_selected_if_branch(self.ae, self.encoder_branch, selected)
+            samples.append(frame.to(output_device).unsqueeze(1))
+        return torch.cat(samples, dim=1)
+
+
+class FlowMatchingObjective:
+    def __init__(self, predictor_module, sigma_min=1e-5):
+        self.module = predictor_module
+        self.sigma_min = sigma_min
+
+    def alpha(self, t):
+        return 1.0 - t
+
+    def sigma(self, t):
+        return self.sigma_min + t * (1.0 - self.sigma_min)
+
+    def A(self, t):
+        return 1.0
+
+    def B(self, t):
+        return -(1.0 - self.sigma_min)
+
+    def sample_t(self, x):
+        return torch.rand((x.shape[0],), device=x.device)
+
+    def add_noise(self, x, t, noise=None):
+        noise = torch.randn_like(x) if noise is None else noise
+        if t.dim() == 2:
+            shape = [x.shape[0], x.shape[1]] + [1] * (x.dim() - 2)
+        else:
+            shape = [x.shape[0]] + [1] * (x.dim() - 1)
+        x_t = self.alpha(t).view(*shape) * x + self.sigma(t).view(*shape) * noise
+        return x_t, noise
+
+    def get_supervision_target(self, target, noise, t):
+        return self.A(t) * target + self.B(t) * noise
+
+    def compute_loss(self, pred, target):
+        return (pred.float() - target.float()) ** 2
+
+    def prepare_training_inputs(self, x):
+        raise NotImplementedError
+
+    def compute_step(self, x, frame_rate, condition_kwargs=None):
+        model_input, t, supervision_target, prediction_slice = self.prepare_training_inputs(x)
+        model_condition_kwargs = self.module.condition_preprocessor.get_model_condition_kwargs(
+            condition_kwargs
+        )
+        pred = self.module.vit(model_input, t, frame_rate=frame_rate, **model_condition_kwargs)
+        if prediction_slice is not None:
+            pred = pred[:, prediction_slice]
+        return self.compute_loss(pred, supervision_target)
+
+
+class FlowMatchingObjectiveTeacherForcing(FlowMatchingObjective):
+    def __init__(
+        self,
+        predictor_module,
+        sigma_min=1e-5,
+        ctx_noise_aug_ratio=0.1,
+        ctx_noise_aug_prob=0.5,
+        drop_ctx_rate=0.2,
+    ):
+        super().__init__(predictor_module, sigma_min=sigma_min)
+        self.ctx_noise_aug_ratio = ctx_noise_aug_ratio
+        self.ctx_noise_aug_prob = ctx_noise_aug_prob
+        self.drop_ctx_rate = drop_ctx_rate
+
+    def _split_sequence(self, x):
+        context = x[:, :-self.module.num_pred_frames]
+        target = x[:, -self.module.num_pred_frames:]
+        return context, target
+
+    def _maybe_drop_context(self, context, target):
+        if context.size(1) == 0:
+            return context
+        if torch.rand(1, device=target.device) < self.drop_ctx_rate:
+            return context[:, :0]
+        return context
+
+    def _maybe_noise_context(self, context, t):
+        if context.size(1) == 0:
+            return context
+        if torch.rand(1, device=context.device) >= self.ctx_noise_aug_prob:
+            return context
+
+        mask = t >= self.ctx_noise_aug_ratio
+        if not mask.any():
+            return context
+
+        context = context.clone()
+        aug_noise = torch.randn_like(context)
+        context[mask] = context[mask] + aug_noise[mask] * self.ctx_noise_aug_ratio
+        return context
+
+    def _prepare_context(self, context, target, t):
+        if not self.module.training:
+            return context
+
+        context = self._maybe_drop_context(context, target)
+        if context.size(1) == 0:
+            return context
+        return self._maybe_noise_context(context, t)
+
+    def prepare_training_inputs(self, x):
+        context, target = self._split_sequence(x)
+        t = self.sample_t(target)
+        context = self._prepare_context(context, target, t)
+        target_t, noise = self.add_noise(target, t)
+
+        model_input = torch.cat([context, target_t], dim=1) if context.size(1) > 0 else target_t
+
+        supervision_target = self.get_supervision_target(target, noise, t)
+        prediction_slice = slice(-self.module.num_pred_frames, None)
+        model_t = self.module.sampler._build_model_t(context, target_t, t)
+        return model_input, model_t, supervision_target, prediction_slice
+
+
+class FlowMatchingObjectiveDiffusionForcing(FlowMatchingObjective):
+    def sample_t(self, x):
+        return torch.rand(x.shape[0], x.shape[1], device=x.device)
+
+    def prepare_training_inputs(self, x):
+        t = self.sample_t(x)
+        x_t, noise = self.add_noise(x, t)
+        supervision_target = self.get_supervision_target(x, noise, t)
+        return x_t, t, supervision_target, None
+
+
+class FlowMatchingSampler:
+    def __init__(self, predictor_module, timescale=1.0, integration_t_eps=0.0):
+        self.module = predictor_module
+        self.timescale = timescale
+        self.integration_t_eps = float(integration_t_eps)
+
+    def _get_net(self, sample_with_ema):
+        return self.module.ema_vit if sample_with_ema else self.module.vit
+
+    def _prepare_context(self, images, latent):
+        if images is None:
+            return None
+        if latent:
+            return images.clone()
+        return self.module.encode_frames(images)
+
+    def _default_frame_rate(self, num_samples, device):
+        return torch.full_like(torch.ones((num_samples,)), 5, device=device)
+
+    def _get_input_hw(self):
+        if isinstance(self.module.vit.input_size, (list, tuple, ListConfig)):
+            return self.module.vit.input_size[0], self.module.vit.input_size[1]
+        return self.module.vit.input_size, self.module.vit.input_size
+
+    def _build_model_inputs(self, context, target_t, t):
+        if context is not None and context.size(1) > 0:
+            model_input = torch.cat([context, target_t], dim=1)
+        else:
+            model_input = target_t
+        return model_input
+
+    def _build_model_t(self, context, target_t, t_scalar):
+        del context, target_t, t_scalar
+        raise NotImplementedError(
+            f"{self.__class__.__name__} must implement `_build_model_t`."
+        )
+
+    def _extract_target_prediction(self, pred):
+        return pred[:, -self.module.num_pred_frames :]
+
+    def _prepare_sampling_state(
+        self, images, latent, sample_with_ema, num_samples, frame_rate, condition_kwargs
+    ):
+        net = self._get_net(sample_with_ema)
+        device = next(net.parameters()).device
+        context = self._prepare_context(images, latent)
+        condition_kwargs = self.module.condition_preprocessor.prepare_condition_kwargs(
+            condition_kwargs,
+            batch_size=num_samples,
+            device=device,
+            split="sample",
+        )
+        model_condition_kwargs = self.module.condition_preprocessor.get_model_condition_kwargs(
+            condition_kwargs
+        )
+
+        if frame_rate is None:
+            frame_rate = self._default_frame_rate(num_samples, device)
+
+        input_h, input_w = self._get_input_hw()
+        target_t = torch.randn(
+            num_samples,
+            self.module.num_pred_frames,
+            self.module.vit.in_channels,
+            input_h,
+            input_w,
+            device=device,
+        )
+        return net, device, context, model_condition_kwargs, frame_rate, target_t
+
+    def _build_t_steps(self, NFE, device):
+        if not 0.0 <= self.integration_t_eps < 0.5:
+            raise ValueError(
+                "integration_t_eps must be in [0, 0.5), "
+                f"got {self.integration_t_eps}."
+            )
+        return torch.linspace(
+            1.0 - self.integration_t_eps,
+            self.integration_t_eps,
+            NFE + 1,
+            device=device,
+        )
+
+    def _eval_velocity(self, net, context, target_t, t_scalar, frame_rate, model_condition_kwargs):
+        model_input = self._build_model_inputs(context, target_t, t_scalar)
+        model_t = self._build_model_t(context, target_t, t_scalar)
+        pred = net(model_input, t=model_t * self.timescale, frame_rate=frame_rate, **model_condition_kwargs)
+        return self._extract_target_prediction(pred)
+
+    def _euler_maruyama_step(self, target_t, neg_v, t_i, t_ip1, eta):
+        dt = t_i - t_ip1
+        dw = torch.randn(target_t.size(), device=target_t.device) * torch.sqrt(dt)
+        diffusion = dt
+        return target_t + neg_v * dt + eta * torch.sqrt(2 * diffusion) * dw
+
+    def _heun_step(self, net, context, target_t, t_i, t_ip1, frame_rate, model_condition_kwargs):
+        t_i_scalar = t_i.repeat(target_t.shape[0])
+        t_ip1_scalar = t_ip1.repeat(target_t.shape[0])
+        dt = t_i - t_ip1
+        v1 = self._eval_velocity(net, context, target_t, t_i_scalar, frame_rate, model_condition_kwargs)
+        x_pred = target_t + v1 * dt
+        v2 = self._eval_velocity(net, context, x_pred, t_ip1_scalar, frame_rate, model_condition_kwargs)
+        return target_t + 0.5 * (v1 + v2) * dt
+
+    def _snapshot_condition_kwargs(self, condition_kwargs):
+        if not condition_kwargs:
+            return {}
+        snapshot = {}
+        for key, value in condition_kwargs.items():
+            if torch.is_tensor(value):
+                snapshot[key] = value.detach().cpu().clone()
+            else:
+                snapshot[key] = value
+        return snapshot
+
+    @torch.no_grad()
+    def sample(
+        self,
+        images=None,
+        latent=False,
+        eta=0.0,
+        NFE=20,
+        sample_with_ema=True,
+        num_samples=8,
+        frame_rate=None,
+        condition_kwargs=None,
+        return_sample=False,
+    ):
+        net, device, context, model_condition_kwargs, frame_rate, target_t = self._prepare_sampling_state(
+            images, latent, sample_with_ema, num_samples, frame_rate, condition_kwargs
+        )
+        t_steps = self._build_t_steps(NFE, device)
+        for i in range(NFE):
+            t_scalar = t_steps[i].repeat(target_t.shape[0])
+            neg_v = self._eval_velocity(net, context, target_t, t_scalar, frame_rate, model_condition_kwargs)
+            target_t = self._euler_maruyama_step(target_t, neg_v, t_steps[i], t_steps[i + 1], eta)
+
+        if return_sample:
+            return target_t, self.module.decode_frames(target_t.clone())
+        return target_t
+
+    def _update_rollout_context(self, context, prediction):
+        latest = prediction[:, -self.module.num_pred_frames:]
+        if self.module.num_pred_frames > context.size(1):
+            return latest
+        return torch.cat([context[:, self.module.num_pred_frames :], latest], dim=1)
+
+    @torch.no_grad()
+    def roll_out(
+        self,
+        x_0,
+        num_gen_frames=25,
+        latent_input=True,
+        eta=0.0,
+        NFE=20,
+        sample_with_ema=True,
+        num_samples=8,
+        frame_rate=None,
+        condition_kwargs=None,
+        decode_device=None,
+        return_condition_history=False,
+    ):
+        context = x_0.clone() if latent_input else self.module.encode_frames(x_0)
+        all_latents = context.clone()
+        condition_kwargs = self.module.condition_preprocessor.prepare_condition_kwargs(
+            condition_kwargs,
+            batch_size=context.size(0),
+            device=context.device,
+            split="rollout",
+        )
+        condition_history = []
+
+        for _idx in tqdm(range(num_gen_frames), desc="Rolling out frames"):
+            if return_condition_history:
+                condition_history.append(self._snapshot_condition_kwargs(condition_kwargs))
+            prediction = self.sample(
+                images=context,
+                latent=True,
+                eta=eta,
+                NFE=NFE,
+                sample_with_ema=sample_with_ema,
+                num_samples=num_samples,
+                frame_rate=frame_rate,
+                condition_kwargs=condition_kwargs,
+            )
+            all_latents = torch.cat([all_latents, prediction[:, -self.module.num_pred_frames :]], dim=1)
+            if _idx < num_gen_frames - 1:
+                condition_kwargs = self.module.condition_preprocessor.update_rollout_condition_kwargs(
+                    condition_kwargs,
+                    prediction=prediction,
+                    context=context,
+                    step_idx=_idx,
+                )
+            context = self._update_rollout_context(context, prediction)
+
+        result = (all_latents, self.module.decode_frames(all_latents, output_device=decode_device))
+        if return_condition_history:
+            return result + (condition_history,)
+        return result
+
+
+class FlowMatchingSamplerTeacherForcing(FlowMatchingSampler):
+    """
+    Inference sampler used for all v2 models.
+
+    The context frames are kept clean in `model_input`; only the target frame
+    block is initialized with noise and integrated during sampling.
+    """
+
+    def __init__(
+        self,
+        predictor_module,
+        timescale=1.0,
+        integration_t_eps=0.0,
+        timestep_conditioning="global",
+    ):
+        super().__init__(
+            predictor_module,
+            timescale=timescale,
+            integration_t_eps=integration_t_eps,
+        )
+        self.timestep_conditioning = self._normalize_timestep_conditioning(timestep_conditioning)
+
+    def _normalize_timestep_conditioning(self, timestep_conditioning):
+        mode = str(timestep_conditioning).strip().lower()
+        if mode not in {"global", "per_frame"}:
+            raise ValueError(
+                "timestep_conditioning must be one of {'global', 'per_frame'}, "
+                f"got {timestep_conditioning!r}."
+            )
+        return mode
+
+    def _build_model_t(self, context, target_t, t_scalar):
+        if self.timestep_conditioning == "global":
+            return t_scalar
+
+        target_t_full = t_scalar.unsqueeze(1).expand(-1, target_t.size(1))
+        if context is None or context.size(1) == 0:
+            return target_t_full
+        context_t = torch.zeros(
+            t_scalar.size(0),
+            context.size(1),
+            device=t_scalar.device,
+            dtype=t_scalar.dtype,
+        )
+        return torch.cat([context_t, target_t_full], dim=1)
+
+    def _validate_num_heun_steps(self, num_heun_steps, NFE, eta):
+        if isinstance(num_heun_steps, bool) or not isinstance(num_heun_steps, int):
+            raise TypeError(
+                f"num_heun_steps must be an int, got {type(num_heun_steps).__name__}."
+            )
+        if not 0 <= num_heun_steps <= NFE:
+            raise ValueError(
+                f"num_heun_steps must be in [0, NFE] (NFE={NFE}), got {num_heun_steps}."
+            )
+        if num_heun_steps > 0 and eta != 0.0:
+            raise ValueError(
+                "Heun sampling (num_heun_steps > 0) is deterministic-only and cannot "
+                f"be combined with the stochastic SDE term; got eta={eta}. "
+                "Set eta=0.0 when num_heun_steps > 0."
+            )
+
+    @torch.no_grad()
+    def sample(
+        self,
+        images=None,
+        latent=False,
+        eta=0.0,
+        NFE=20,
+        sample_with_ema=True,
+        num_samples=8,
+        frame_rate=None,
+        condition_kwargs=None,
+        return_sample=False,
+        num_heun_steps=0,
+    ):
+        self._validate_num_heun_steps(num_heun_steps, NFE, eta)
+        net, device, context, model_condition_kwargs, frame_rate, target_t = self._prepare_sampling_state(
+            images, latent, sample_with_ema, num_samples, frame_rate, condition_kwargs
+        )
+        t_steps = self._build_t_steps(NFE, device)
+        for i in range(NFE):
+            if i < num_heun_steps:
+                target_t = self._heun_step(
+                    net, context, target_t, t_steps[i], t_steps[i + 1], frame_rate, model_condition_kwargs
+                )
+            else:
+                t_scalar = t_steps[i].repeat(target_t.shape[0])
+                neg_v = self._eval_velocity(net, context, target_t, t_scalar, frame_rate, model_condition_kwargs)
+                target_t = self._euler_maruyama_step(target_t, neg_v, t_steps[i], t_steps[i + 1], eta)
+
+        if return_sample:
+            return target_t, self.module.decode_frames(target_t.clone())
+        return target_t
+
+    @torch.no_grad()
+    def roll_out(
+        self,
+        x_0,
+        num_gen_frames=25,
+        latent_input=True,
+        eta=0.0,
+        NFE=20,
+        sample_with_ema=True,
+        num_samples=8,
+        frame_rate=None,
+        condition_kwargs=None,
+        decode_device=None,
+        return_condition_history=False,
+        num_heun_steps=0,
+    ):
+        # num_heun_steps validation happens inside self.sample() on the first
+        # generated block, before any expensive work in this loop.
+        context = x_0.clone() if latent_input else self.module.encode_frames(x_0)
+        all_latents = context.clone()
+        condition_kwargs = self.module.condition_preprocessor.prepare_condition_kwargs(
+            condition_kwargs,
+            batch_size=context.size(0),
+            device=context.device,
+            split="rollout",
+        )
+        condition_history = []
+
+        for _idx in tqdm(range(num_gen_frames), desc="Rolling out frames"):
+            if return_condition_history:
+                condition_history.append(self._snapshot_condition_kwargs(condition_kwargs))
+            prediction = self.sample(
+                images=context,
+                latent=True,
+                eta=eta,
+                NFE=NFE,
+                sample_with_ema=sample_with_ema,
+                num_samples=num_samples,
+                frame_rate=frame_rate,
+                condition_kwargs=condition_kwargs,
+                num_heun_steps=num_heun_steps,
+            )
+            all_latents = torch.cat([all_latents, prediction[:, -self.module.num_pred_frames :]], dim=1)
+            if _idx < num_gen_frames - 1:
+                condition_kwargs = self.module.condition_preprocessor.update_rollout_condition_kwargs(
+                    condition_kwargs,
+                    prediction=prediction,
+                    context=context,
+                    step_idx=_idx,
+                )
+            context = self._update_rollout_context(context, prediction)
+
+        result = (all_latents, self.module.decode_frames(all_latents, output_device=decode_device))
+        if return_condition_history:
+            return result + (condition_history,)
+        return result
+
+
+class FlowMatchingSamplerDiffusionForcing(FlowMatchingSampler):
+    """
+    Deprecated: diffusion forcing is a training objective, not a separate
+    inference scheme in v2. Use FlowMatchingSamplerTeacherForcing for clean
+    context / noisy target sampling.
+    """
+
+    def __init__(self, *args, **kwargs):
+        raise RuntimeError(
+            "FlowMatchingSamplerDiffusionForcing is deprecated. "
+            "Use FlowMatchingSamplerTeacherForcing for inference."
+        )
+
+    def _build_model_t(self, context, target_t, t_scalar):
+        del context
+        return t_scalar.unsqueeze(1).expand(-1, target_t.size(1))
+
+
+class PredictorModule(pl.LightningModule):
+    def __init__(
+        self,
+        *,
+        tokenizer_config,
+        generator_config,
+        objective_config=None,
+        sampler_config=None,
+        condition_preprocessor_config=None,
+        first_stage_handler_config=None,
+        adjust_lr_to_batch_size=False,
+        num_pred_frames=1,
+        warmup_steps=5000,
+        min_lr_multiplier=0.1,
+        enc_scale=4,
+        enc_scale_dino=3.45062,
+        use_precomputed_training_inputs=False,
+        init_weights_path=None,
+        allow_different_resolution_checkpoint=False,
+    ):
+        super().__init__()
+
+        self.num_pred_frames = num_pred_frames
+        self.enc_scale = enc_scale
+        self.enc_scale_dino = enc_scale_dino
+        self.allow_different_resolution_checkpoint = allow_different_resolution_checkpoint
+        self.adjust_lr_to_batch_size = adjust_lr_to_batch_size
+        self.warmup_steps = warmup_steps
+        self.min_lr_multiplier = min_lr_multiplier
+        self.use_precomputed_training_inputs = use_precomputed_training_inputs
+
+        self.vit = self.build_generator(generator_config)
+        self.ae = self.build_tokenizer(tokenizer_config)
+        self.ema_vit = init_ema_model(self.vit)
+
+        self.first_stage = self.build_first_stage(first_stage_handler_config)
+        self.condition_preprocessor = self.build_condition_preprocessor(condition_preprocessor_config)
+        self.objective = self.build_objective(objective_config)
+        self.sampler = self.build_sampler(sampler_config)
+
+        if init_weights_path is not None:
+            self.init_weights_from_checkpoint(init_weights_path)
+
+    def setup(self, stage=None):
+        """Exclude 'unused' parameters from DDP gradient reduction."""
+        super().setup(stage)
+        # EMA
+        if hasattr(self, "ema_vit") and self.ema_vit is not None:
+            self.ema_vit.requires_grad_(False)
+        self.ae.requires_grad_(False)
+
+    def init_weights_from_checkpoint(self, init_weights_path):
+        checkpoint_path = os.path.expandvars(init_weights_path)
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(f"Initial weights {init_weights_path} does not exist.")
+
+        state_dict = torch.load(checkpoint_path, map_location="cpu")["state_dict"]
+        state_dict = self._prepare_second_stage_checkpoint(state_dict, checkpoint_path)
+        outcome = self.load_state_dict(state_dict, strict=False)
+        assert outcome.missing_keys == [], outcome.missing_keys
+        print(f"Loaded model from {init_weights_path}")
+
+    def build_tokenizer(self, tokenizer_config):
+        tokenizer_folder = os.path.expandvars(tokenizer_config.folder)
+        ckpt_path = tokenizer_config.ckpt_path if tokenizer_config.ckpt_path else "checkpoints/last.ckpt"
+
+        tokenizer_config = OmegaConf.load(os.path.join(tokenizer_folder, "config.yaml"))
+        model_cfg = OmegaConf.to_container(tokenizer_config.model, resolve=False)
+        for key in ("loss_config", "entropy_loss_weight_scheduler_config"):
+            model_cfg.get("params", {}).pop(key, None)
+        model_cfg.get("params", {})["distill_model_type"] = None
+        encoder_params = model_cfg.get("params", {}).get("encoder_config", {}).get("params", {})
+        if isinstance(encoder_params, dict):
+            encoder_params["use_pretrained_weights"] = False
+        model = instantiate_from_config(OmegaConf.create(model_cfg))
+
+        checkpoint_path = os.path.join(tokenizer_folder, ckpt_path)
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)["state_dict"]
+        checkpoint = self._prepare_tokenizer_checkpoint(model, checkpoint, tokenizer_config.model, checkpoint_path)
+        model.load_state_dict(checkpoint, strict=False)
+        model.eval()
+        return model
+
+    def build_generator(self, generator_config):
+        return instantiate_from_config(generator_config)
+
+    def _build_helper_from_config(self, config, default_target):
+        config = config or {"target": default_target, "params": {}}
+        params = dict(config.get("params", {}))
+        params["predictor_module"] = self
+        helper_cls = get_obj_from_str(config["target"])
+        return helper_cls(**params)
+
+    def build_objective(self, objective_config):
+        return self._build_helper_from_config(
+            objective_config,
+            "models.second_stage.fm_model_v2.FlowMatchingObjectiveTeacherForcing",
+        )
+
+    def build_sampler(self, sampler_config):
+        return self._build_helper_from_config(
+            sampler_config,
+            "models.second_stage.fm_model_v2.FlowMatchingSamplerTeacherForcing",
+        )
+
+    def build_condition_preprocessor(self, condition_preprocessor_config):
+        return self._build_helper_from_config(
+            condition_preprocessor_config,
+            "models.second_stage.fm_conditions_v2.ConditionPreprocessor",
+        )
+
+    def build_first_stage(self, first_stage_handler_config):
+        return self._build_helper_from_config(
+            first_stage_handler_config,
+            "models.second_stage.fm_model_v2.FirstStageHandler",
+        )
+
+    def _require_or_allow_resolution_mismatch(self, *, name, checkpoint_shape, model_shape, source):
+        if checkpoint_shape == model_shape:
+            return
+        if not self.allow_different_resolution_checkpoint:
+            raise ValueError(
+                f"{name} shape mismatch between checkpoint and current model: "
+                f"checkpoint={checkpoint_shape}, model={model_shape}. "
+                f"Source: {source}. "
+                "Set allow_different_resolution_checkpoint=True to adapt or skip "
+                "resolution-dependent positional embeddings."
+            )
+
+    def _prepare_tokenizer_checkpoint(self, model, checkpoint, tokenizer_model_config, checkpoint_path):
+        state_dict = checkpoint.copy()
+        target = tokenizer_model_config.get("target", "")
+        encoder_config = tokenizer_model_config.params.get("encoder_config")
+        encoder_target = encoder_config.get("target", "") if encoder_config else ""
+
+        if target != "models.first_stage.vqgan.VQModel" or encoder_target != "networks.tokenizer.pretrained_models.Encoder":
+            return state_dict
+
+        model_pos_embed = getattr(getattr(model.encoder, "encoder", None), "pos_embed", None)
+        checkpoint_key = "encoder.encoder.pos_embed"
+        checkpoint_pos_embed = state_dict.get(checkpoint_key)
+
+        if model_pos_embed is None or checkpoint_pos_embed is None:
+            return state_dict
+
+        self._require_or_allow_resolution_mismatch(
+            name=f"Tokenizer positional embedding ({checkpoint_key})",
+            checkpoint_shape=tuple(checkpoint_pos_embed.shape),
+            model_shape=tuple(model_pos_embed.shape),
+            source=checkpoint_path,
+        )
+
+        if checkpoint_pos_embed.shape != model_pos_embed.shape:
+            num_prefix_tokens = getattr(model.encoder.encoder, "num_prefix_tokens", 1)
+            target_grid = model.encoder.encoder.patch_embed.grid_size
+            state_dict[checkpoint_key] = resample_abs_pos_embed(
+                checkpoint_pos_embed,
+                new_size=target_grid,
+                num_prefix_tokens=num_prefix_tokens,
+            )
+
+        return state_dict
+
+    def _prepare_second_stage_checkpoint(self, checkpoint, checkpoint_path):
+        state_dict = checkpoint.copy()
+        for key in ("vit.pos_embed", "ema_vit.pos_embed"):
+            checkpoint_pos_embed = state_dict.get(key)
+            model_pos_embed = self.state_dict().get(key)
+            if checkpoint_pos_embed is None or model_pos_embed is None:
+                continue
+
+            self._require_or_allow_resolution_mismatch(
+                name=f"Second-stage positional embedding ({key})",
+                checkpoint_shape=tuple(checkpoint_pos_embed.shape),
+                model_shape=tuple(model_pos_embed.shape),
+                source=checkpoint_path,
+            )
+
+            if checkpoint_pos_embed.shape != model_pos_embed.shape:
+                del state_dict[key]
+
+        return state_dict
+
+    def get_warmup_scheduler(self, optimizer, warmup_steps=1, min_lr_multiplier=1.0):
+        batches = len(self.trainer.datamodule.train_dataloader())
+        steps_per_epoch = batches // self.trainer.accumulate_grad_batches
+        total_steps = self.trainer.max_epochs * steps_per_epoch
+
+        def lr_lambda(step):
+            if step < warmup_steps:
+                return step / warmup_steps
+            progress = (min(step, total_steps) - warmup_steps) / (total_steps - warmup_steps)
+            cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
+            return (1 - min_lr_multiplier) * cosine_decay + min_lr_multiplier
+
+        return LambdaLR(optimizer, lr_lambda)
+
+    def configure_optimizers(self):
+        params = [p for p in self.vit.parameters() if p.requires_grad]
+        optimizer = torch.optim.AdamW(params, lr=self.learning_rate, weight_decay=0.01)
+        scheduler = self.get_warmup_scheduler(optimizer, self.warmup_steps, self.min_lr_multiplier)
+        return [optimizer], [{"scheduler": scheduler, "interval": "step"}]
+
+    def get_input(self, batch, k):
+        if isinstance(batch, dict):
+            x = batch[k]
+            frame_rate = batch["frame_rate"]
+        else:
+            x = batch
+            frame_rate = None
+        assert len(x.shape) == 5 or self.use_precomputed_training_inputs, "When using images, input must be 5D tensor"
+        return x, frame_rate
+
+    @torch.no_grad()
+    def encode_frames(self, images):
+        return self.first_stage.encode_frames(images)
+
+    @torch.no_grad()
+    def decode_frames(self, x, output_device=None):
+        return self.first_stage.decode_frames(x, output_device=output_device)
+
+    def compute_prediction_loss(self, pred, target):
+        return self.objective.compute_loss(pred, target)
+
+    def log_training_losses(self, loss):
+        self.log("train/loss", loss.mean(), prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+    def log_validation_losses(self, loss):
+        self.log("val/loss", loss.mean(), prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+
+    def log_adaln_mean_abs(self, split, condition_kwargs=None):
+        stats_getter = getattr(self.vit, "get_last_adaln_mean_abs", None)
+        if not callable(stats_getter):
+            return
+
+        stats = stats_getter()
+        if not stats:
+            return
+
+        aggregated = {}
+        for name, value in stats.items():
+            _, stat_name = name.split("/", 1)
+            aggregated.setdefault(stat_name, []).append(value)
+
+        group_masks = self._get_adaln_stat_group_masks(condition_kwargs)
+        for stat_name, values in aggregated.items():
+            values_tensor = torch.stack(values)
+            for group_name, mask in group_masks.items():
+                group_values = values_tensor if mask is None else values_tensor[:, mask]
+                self._log_adaln_stat_values(split, group_name, stat_name, group_values)
+
+    def _get_adaln_stat_group_masks(self, condition_kwargs):
+        if not condition_kwargs or "steering" not in condition_kwargs:
+            return {"all": None}
+
+        steering = condition_kwargs["steering"]
+        if not torch.is_tensor(steering):
+            return {"all": None}
+
+        present_mask = ~torch.isnan(steering).flatten(1).all(dim=1)
+        return {
+            "present": present_mask,
+            "missing": ~present_mask,
+        }
+
+    def _log_adaln_stat_values(self, split, group_name, stat_name, values):
+        # Always emit exactly 2 log calls (mean + std) regardless of group size so that
+        # sync_dist=True allreduces are consistent across ranks even when different ranks
+        # receive batches with different steering availability (e.g. mixed-dataset training).
+        values = values.reshape(-1)
+        nan = torch.tensor(float("nan"), device=values.device)
+        mean_val = values.mean() if values.numel() > 0 else nan
+        std_val = values.std(unbiased=False) if values.numel() >= 2 else nan
+
+        self.log(
+            f"{split}/adaln/{group_name}/{stat_name}/mean",
+            mean_val,
+            prog_bar=False,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+            sync_dist=True,
+        )
+        self.log(
+            f"{split}/adaln/{group_name}/{stat_name}/std",
+            std_val,
+            prog_bar=False,
+            logger=True,
+            on_step=True,
+            on_epoch=True,
+            sync_dist=True,
+        )
+
+    def log_condition_embeddings(self, split, condition_kwargs=None):
+        getter = getattr(self.vit, 'get_last_condition_embeddings', None)
+        if not callable(getter):
+            return
+        embeddings = getter()
+        if not embeddings:
+            return
+        group_masks = self._get_adaln_stat_group_masks(condition_kwargs)
+        for emb_name, embedding in embeddings.items():
+            per_sample_mean_abs = embedding.abs().mean(dim=-1)
+            per_sample_norm = embedding.norm(dim=-1)
+            for stat_name, values in (("mean_abs", per_sample_mean_abs), ("norm", per_sample_norm)):
+                for group_name, mask in group_masks.items():
+                    group_values = values if mask is None else values[mask]
+                    group_values = group_values.reshape(-1)
+                    # Always emit exactly 2 log calls per group so sync_dist=True allreduces
+                    # are consistent across ranks regardless of per-rank batch composition.
+                    nan = torch.tensor(float("nan"), device=group_values.device)
+                    mean_val = group_values.mean() if group_values.numel() > 0 else nan
+                    std_val = group_values.std(unbiased=False) if group_values.numel() >= 2 else nan
+                    self.log(
+                        f"{split}/cond_emb/{emb_name}/{group_name}/{stat_name}/mean",
+                        mean_val,
+                        prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True,
+                    )
+                    self.log(
+                        f"{split}/cond_emb/{emb_name}/{group_name}/{stat_name}/std",
+                        std_val,
+                        prog_bar=False, logger=True, on_step=True, on_epoch=True, sync_dist=True,
+                    )
+
+    def log_adaln_gradients(self):
+        getter = getattr(self.vit, 'get_adaln_grad_stats', None)
+        if not callable(getter):
+            return
+        stats = getter()
+        for name, g in stats.items():
+            g_flat = g.float().flatten()
+            self.log(
+                f"train/adaln_grad/{name}/norm",
+                g_flat.norm(),
+                prog_bar=False, logger=True, on_step=True, on_epoch=False, sync_dist=True,
+            )
+            self.log(
+                f"train/adaln_grad/{name}/mean_abs",
+                g_flat.abs().mean(),
+                prog_bar=False, logger=True, on_step=True, on_epoch=False, sync_dist=True,
+            )
+
+    def on_after_backward(self):
+        self.log_adaln_gradients()
+
+    def _shared_step(self, batch, split):
+        images, frame_rate = self.get_input(batch, "images")
+        x = self.encode_frames(images)
+        condition_kwargs = self.condition_preprocessor.get_condition_kwargs_from_batch(batch, split=split)
+        loss = self.objective.compute_step(x, frame_rate, condition_kwargs=condition_kwargs)
+
+        if split == "train":
+            self.log_training_losses(loss)
+            self.log_adaln_mean_abs("train", condition_kwargs)
+            self.log_condition_embeddings("train", condition_kwargs)
+            return loss.mean()
+
+        self.log_validation_losses(loss)
+        self.log_adaln_mean_abs("val", condition_kwargs)
+        self.log_condition_embeddings("val", condition_kwargs)
+        return loss.mean()
+
+    def training_step(self, batch, batch_idx):
+        return self._shared_step(batch, "train")
+
+    def validation_step(self, batch, batch_idx):
+        self._shared_step(batch, "val")
+
+    def _sync_ema_stream(self):
+        """Block until the async EMA update finishes before reading ema_vit."""
+        stream = getattr(self, "_ema_stream", None)
+        if stream is not None:
+            stream.synchronize()
+    
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        """Update EMA asynchronously on a separate CUDA stream."""
+        if not hasattr(self, "_ema_stream"):
+            self._ema_stream = torch.cuda.Stream() if torch.cuda.is_available() else None
+        if self._ema_stream is not None:
+            with torch.cuda.stream(self._ema_stream):
+                update_ema(self.ema_vit, self.vit)
+        else:
+            update_ema(self.ema_vit, self.vit)
+
+    # def on_train_batch_end(self, outputs, batch, batch_idx):
+    #     update_ema(self.ema_vit, self.vit)
+
+    @torch.no_grad()
+    def sample(
+        self,
+        images=None,
+        latent=False,
+        eta=0.0,
+        NFE=20,
+        sample_with_ema=True,
+        num_samples=8,
+        frame_rate=None,
+        condition_kwargs=None,
+        return_sample=False,
+    ):
+        if sample_with_ema:
+            self._sync_ema_stream()
+
+        return self.sampler.sample(
+            images=images,
+            latent=latent,
+            eta=eta,
+            NFE=NFE,
+            sample_with_ema=sample_with_ema,
+            num_samples=num_samples,
+            frame_rate=frame_rate,
+            condition_kwargs=condition_kwargs,
+            return_sample=return_sample,
+        )
+
+    def _validate_rollout_context(self, x_0):
+        if not isinstance(x_0, dict):
+            raise TypeError(
+                f"{self.__class__.__name__}.roll_out expects x_0 to be a dict, "
+                f"got {type(x_0).__name__}."
+            )
+        if "images" not in x_0:
+            raise KeyError(f"{self.__class__.__name__}.roll_out requires x_0['images'].")
+
+    def _move_rollout_context_to_device(self, x_0, device):
+        moved = {}
+        for key, value in x_0.items():
+            if torch.is_tensor(value):
+                moved[key] = value.to(device)
+            else:
+                moved[key] = value
+        return moved
+
+    def _get_rollout_condition_frames(self, images, num_condition_frames=None):
+        if num_condition_frames is None:
+            num_condition_frames = getattr(self.vit, "num_context_frames", None)
+        if num_condition_frames is None:
+            num_condition_frames = images.size(1) - self.num_pred_frames
+        num_condition_frames = int(num_condition_frames)
+        if num_condition_frames <= 0:
+            raise ValueError(
+                f"rollout batch has {images.size(1)} frames, but "
+                f"num_pred_frames={self.num_pred_frames}; expected at least one "
+                "conditioning frame."
+            )
+        if num_condition_frames > images.size(1):
+            raise ValueError(
+                f"num_condition_frames={num_condition_frames} exceeds rollout "
+                f"batch length {images.size(1)}."
+            )
+        return num_condition_frames
+
+    def roll_out(
+        self,
+        x_0,
+        num_gen_frames=25,
+        latent_input=True,
+        eta=0.0,
+        NFE=20,
+        sample_with_ema=True,
+        num_samples=8,
+        frame_rate=None,
+        condition_kwargs=None,
+        decode_device=None,
+        return_condition_history=False,
+        num_condition_frames=None,
+    ):
+        device = next(self.parameters()).device
+        self._validate_rollout_context(x_0)
+        x_0 = self._move_rollout_context_to_device(x_0, device)
+        images = x_0["images"]
+        num_condition_frames = self._get_rollout_condition_frames(
+            images,
+            num_condition_frames=num_condition_frames,
+        )
+        cond_x = images[:, :num_condition_frames]
+        rollout_context = dict(x_0)
+        rollout_context["images"] = cond_x
+
+        if torch.is_tensor(frame_rate):
+            frame_rate = frame_rate.to(device)
+        elif frame_rate is None and torch.is_tensor(rollout_context.get("frame_rate")):
+            frame_rate = rollout_context["frame_rate"]
+
+        if condition_kwargs is None:
+            condition_kwargs = self.condition_preprocessor.get_condition_kwargs_from_batch(
+                rollout_context,
+                split="rollout",
+            )
+
+        return self.sampler.roll_out(
+            cond_x,
+            num_gen_frames=num_gen_frames,
+            latent_input=latent_input,
+            eta=eta,
+            NFE=NFE,
+            sample_with_ema=sample_with_ema,
+            num_samples=num_samples,
+            frame_rate=frame_rate,
+            condition_kwargs=condition_kwargs,
+            decode_device=decode_device,
+            return_condition_history=return_condition_history,
+        )
+
+    @torch.no_grad()
+    def log_images(self, batch, **kwargs):
+        log = {}
+        images, frame_rate = self.get_input(batch, "images")
+        N = min(5, images.size(0))
+        images = images[:N]
+        condition_kwargs = self.condition_preprocessor.get_condition_kwargs_from_batch(batch, split="log_images")
+        condition_kwargs = self.condition_preprocessor.slice_condition_kwargs(condition_kwargs, slice(0, N))
+
+        if self.use_precomputed_training_inputs and images.shape[-2:] == self.vit.input_size:
+            images = self.decode_frames(images)
+
+        images = self.condition_preprocessor.annotate_logged_images(images, batch=batch, num_images=N)
+        frame_rate = frame_rate[:N] if frame_rate is not None else None
+        num_frames = images.size(1)
+
+        visual = [images[:, frame_idx] for frame_idx in range(num_frames)]
+        visual_ema = [images[:, frame_idx] for frame_idx in range(num_frames)]
+
+        context = images[:, :-self.num_pred_frames] if num_frames - self.num_pred_frames > 0 else None
+
+        samples = self.sample(
+            context,
+            eta=0.0,
+            NFE=30,
+            sample_with_ema=False,
+            num_samples=N,
+            frame_rate=frame_rate,
+            condition_kwargs=condition_kwargs,
+            return_sample=True,
+        )[1]
+        for frame_idx in range(samples.size(1)):
+            visual.append(samples[:, frame_idx])
+
+        samples_ema = self.sample(
+            context,
+            eta=0.0,
+            NFE=30,
+            sample_with_ema=True,
+            num_samples=N,
+            frame_rate=frame_rate,
+            condition_kwargs=condition_kwargs,
+            return_sample=True,
+        )[1]
+        for frame_idx in range(samples_ema.size(1)):
+            visual_ema.append(samples_ema[:, frame_idx])
+
+        sampled = vutils.make_grid(torch.cat(torch.chunk(torch.cat(visual, dim=0), 4, dim=0), dim=0), nrow=N, padding=2, normalize=False)
+        sampled_ema = vutils.make_grid(torch.cat(torch.chunk(torch.cat(visual_ema, dim=0), 4, dim=0), dim=0), nrow=N, padding=2, normalize=False)
+
+        log["sampled"] = sampled
+        log["ema_sampled"] = sampled_ema
+        self.vit.train()
+        return log
+
+
+class L1L2PredL2PredictorModule(PredictorModule):
+    """
+    PredictorModule specialization for L1 rollout conditioned by an L2 endpoint
+    predictor.
+
+    The base PredictorModule keeps a tensor-level rollout API with explicit
+    condition_kwargs. This subclass treats rollout as a batch-level L1/L2
+    operation: extract the L1 context, derive the L2 condition through the
+    configured condition preprocessor, then delegate to the generic sampler.
+    """
+
+    pass

orbis2/modules/dit.py

@@ -0,0 +1,483 @@
+"""
+Reusable DiT building blocks needed for inference.
+Extracted from networks/DiT/dit.py so the inference set avoids importing
+the full training-only model classes in that file.
+"""
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from timm.layers.mlp import SwiGLU
+from timm.models.vision_transformer import Attention, Mlp
+
+
+def get_norm_layer(norm_layer):
+    if isinstance(norm_layer, str):
+        if norm_layer == 'layer_norm':
+            return nn.LayerNorm
+        elif norm_layer == 'rms_norm':
+            return nn.RMSNorm
+        else:
+            raise ValueError(f"Unsupported norm layer: {norm_layer}")
+    return norm_layer
+
+
+def modulate(x, shift, scale):
+    n = x.ndim - shift.ndim
+    for _ in range(n):
+        shift = shift.unsqueeze(-2)
+        scale = scale.unsqueeze(-2)
+    return x * (1 + scale) + shift
+
+
+def _broadcast_gate(gate, x):
+    n = x.ndim - gate.ndim
+    for _ in range(n):
+        gate = gate.unsqueeze(-2)
+    return gate
+
+
+class FrequencyEncoder:
+   def __init__(self, embed_dim, freq_min=1, freq_max=5):
+       """
+       Deterministic frequency encoder with fixed normalization.
+
+
+       Args:
+           embed_dim (int): Dimensionality of the token embeddings.
+           freq_min (float): Minimum frequency value.
+           freq_max (float): Maximum frequency value.
+       """
+       self.embed_dim = embed_dim
+       self.freq_min = freq_min
+       self.freq_max = freq_max
+
+
+   def encode(self, frequencies):
+       """
+       Encodes frequencies into embeddings using sine-cosine features.
+
+
+       Args:
+           frequencies (torch.Tensor): Tensor of shape (batch_size,) containing frequencies.
+
+
+       Returns:
+           torch.Tensor: Encoded frequency embeddings of shape (batch_size, embed_dim).
+       """
+       batch_size = frequencies.size(0)
+
+
+       # Fixed normalization: Scale frequencies to [0, 1]
+       normalized_freq = (frequencies - self.freq_min) / (self.freq_max - self.freq_min)
+
+
+       # Generate positional features using sine and cosine
+       positions = torch.arange(0, self.embed_dim, dtype=torch.float32, device=frequencies.device)
+       scaling_factors = 1 / (10000 ** (2 * (positions // 2) / self.embed_dim))
+       frequency_features = normalized_freq.unsqueeze(1) * scaling_factors  # Shape: (batch_size, embed_dim)
+
+
+       # Apply sine to even indices and cosine to odd indices
+       encoded_freq = torch.zeros(batch_size, self.embed_dim, device=frequencies.device)
+       encoded_freq[:, 0::2] = torch.sin(frequency_features[:, 0::2])  # Sine for even indices
+       encoded_freq[:, 1::2] = torch.cos(frequency_features[:, 1::2])  # Cosine for odd indices
+
+
+       return encoded_freq
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+   """
+   grid_size: int of the grid height and width
+   return:
+   pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+   """
+   grid_h = np.arange(grid_size[0], dtype=np.float32)
+   grid_w = np.arange(grid_size[1], dtype=np.float32)
+   grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+   grid = np.stack(grid, axis=0)
+
+
+   grid = grid.reshape([2, 1, grid_size[0], grid_size[1]])
+   pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+   if cls_token and extra_tokens > 0:
+       pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+   return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+   assert embed_dim % 2 == 0
+
+
+   # use half of dimensions to encode grid_h
+   emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+   emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+
+   emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+   return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+   """
+   embed_dim: output dimension for each position
+   pos: a list of positions to be encoded: size (M,)
+   out: (M, D)
+   """
+   assert embed_dim % 2 == 0
+   omega = np.arange(embed_dim // 2, dtype=np.float64)
+   omega /= embed_dim / 2.
+   omega = 1. / 10000**omega  # (D/2,)
+
+
+   pos = pos.reshape(-1)  # (M,)
+   out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+
+   emb_sin = np.sin(out) # (M, D/2)
+   emb_cos = np.cos(out) # (M, D/2)
+
+
+   emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+   return emb
+
+
+class TimestepEmbedder(nn.Module):
+   """
+   Embeds scalar timesteps into vector representations.
+   """
+   def __init__(self, hidden_size, frequency_embedding_size=256):
+       super().__init__()
+       self.mlp = nn.Sequential(
+           nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+           nn.SiLU(),
+           nn.Linear(hidden_size, hidden_size, bias=True),
+       )
+       self.frequency_embedding_size = frequency_embedding_size
+
+
+   @staticmethod
+   def timestep_embedding(t, dim, max_period=10000):
+       """
+       Create sinusoidal timestep embeddings.
+       :param t: a 1-D Tensor of N indices, one per batch element.
+                         These may be fractional.
+       :param dim: the dimension of the output.
+       :param max_period: controls the minimum frequency of the embeddings.
+       :return: an (N, D) Tensor of positional embeddings.
+       """
+       # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+       half = dim // 2
+       freqs = torch.exp(
+           -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+       ).to(device=t.device)
+       args = t[:, None].float() * freqs[None]
+       embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+       if dim % 2:
+           embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+       return embedding
+
+
+   def forward(self, t):
+       t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+       t_emb = self.mlp(t_freq)
+       return t_emb
+
+
+class STBlock(nn.Module):
+   # Used for temporal compression in context
+   def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, dropout_rate=0.0,
+                norm_layer=nn.LayerNorm,
+                mlp_block='mlp',
+                act_layer=lambda: nn.GELU(approximate="tanh"),
+                **block_kwargs):
+        super().__init__()
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.norm1 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm2 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm3 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm4 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.space_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True,
+                                    attn_drop=dropout_rate, proj_drop=dropout_rate, norm_layer=norm_layer, **block_kwargs)
+        self.time_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True,
+                                    attn_drop=dropout_rate, proj_drop=dropout_rate, norm_layer=norm_layer, **block_kwargs)
+        if mlp_block == 'mlp':
+            self.space_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=act_layer, norm_layer=norm_layer, drop=0)
+            self.time_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=act_layer, norm_layer=norm_layer, drop=0)
+        elif mlp_block == 'swiglu':
+            self.space_mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim*2)//3)
+            self.time_mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim*2)//3)
+        else:
+            raise NotImplementedError(f"mlp_block {mlp_block} not implemented")
+
+   def forward(self, x):
+       B, F, N, D = x.shape
+       x = rearrange(x, 'b f n d -> (b f) n d')
+       x = x + self.space_attn(self.norm1(x))
+       x = x + self.space_mlp(self.norm2(x))
+       x = rearrange(x, '(b f) n d -> (b n) f d', b=B, f=F, n=N)
+       x = x + self.time_attn(self.norm3(x))
+       x = x + self.time_mlp(self.norm4(x))
+       x = rearrange(x, '(b n) f d -> b f n d', b=B, n=N, f=F)
+       return x
+
+
+class DiTBlock(nn.Module):
+   """
+   A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+   """
+   def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, dropout_rate=0.0, norm_layer=nn.LayerNorm, mlp_block='mlp', **block_kwargs):
+        super().__init__()
+        if isinstance(norm_layer, str):
+            norm_layer = get_norm_layer(norm_layer)
+        self.norm1 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True, norm_layer=norm_layer, attn_drop=dropout_rate, proj_drop=dropout_rate, **block_kwargs)
+        self.norm2 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        if mlp_block == 'mlp':
+            approx_gelu = lambda: nn.GELU(approximate="tanh")
+            self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=norm_layer, drop=0)
+        elif mlp_block == 'swiglu':
+            self.mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim*2)//3, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+   def initialize_adaln_weights(self, gate_init_std=0.0):
+       nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+       nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+       if gate_init_std != 0.0:
+           hidden_size = self.adaLN_modulation[-1].out_features // 6
+           for gate_idx in (2, 5):
+               start = gate_idx * hidden_size
+               end = start + hidden_size
+               nn.init.normal_(self.adaLN_modulation[-1].weight[start:end], std=gate_init_std)
+
+   def forward(self, x, c):
+       shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+       x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+       x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+       return x
+
+
+class CDiTBlock(nn.Module):
+   """
+   A DiT block with cross-attention conditioning.
+   """
+   def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, norm_layer=nn.LayerNorm, mlp_block='mlp', **block_kwargs):
+        super().__init__()
+        self.norm1 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, norm_layer=norm_layer, **block_kwargs)
+        self.norm2 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_cond = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.cttn = nn.MultiheadAttention(hidden_size, num_heads=num_heads, add_bias_kv=True, bias=True, batch_first=True, **block_kwargs)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 11 * hidden_size, bias=True)
+        )
+        self.norm3 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        if mlp_block == 'mlp':
+            approx_gelu = lambda: nn.GELU(approximate="tanh")
+            self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=norm_layer, drop=0)
+        elif mlp_block == 'swiglu':
+            self.mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim*2)//3, bias=True)
+
+   def initialize_adaln_weights(self, gate_init_std=0.0):
+       nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+       nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+       if gate_init_std != 0.0:
+           hidden_size = self.adaLN_modulation[-1].out_features // 11
+           for gate_idx in (2, 7, 10):
+               start = gate_idx * hidden_size
+               end = start + hidden_size
+               nn.init.normal_(self.adaLN_modulation[-1].weight[start:end], std=gate_init_std)
+
+   def forward(self, x, c, x_cond):
+       shift_msa, scale_msa, gate_msa, shift_ca_xcond, scale_ca_xcond, shift_ca_x, scale_ca_x, gate_ca_x, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(11, dim=1)
+       x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+       x_cond_norm = modulate(self.norm_cond(x_cond), shift_ca_xcond, scale_ca_xcond)
+       x = x + gate_ca_x.unsqueeze(1) * self.cttn(query=modulate(self.norm2(x), shift_ca_x, scale_ca_x), key=x_cond_norm, value=x_cond_norm, need_weights=False)[0]
+       x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm3(x), shift_mlp, scale_mlp))
+       return x
+
+
+class STDiTBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, dropout_rate=0.0, causal_time_attn=False, modulate_time_attn=False, norm_layer=nn.LayerNorm, mlp_block='mlp', **block_kwargs):
+        super().__init__()
+        self.norm1 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.space_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True, norm_layer=norm_layer, attn_drop=dropout_rate, proj_drop=dropout_rate, **block_kwargs)
+        self.time_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True, norm_layer=norm_layer, attn_drop=dropout_rate, proj_drop=dropout_rate, **block_kwargs)
+        self.norm2 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm3 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm4 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        if mlp_block == 'mlp':
+            approx_gelu = lambda: nn.GELU(approximate="tanh")
+            self.space_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=None, drop=0)
+            self.time_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=None, drop=0)
+        elif mlp_block == 'swiglu':
+            self.space_mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim*2)//3)
+            self.time_mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim*2)//3)
+        else:
+            raise NotImplementedError(f"mlp_block {mlp_block} not implemented")
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 9 * hidden_size, bias=True)
+        )
+        self.causal_time_attn = causal_time_attn
+        self.modulate_time_attn = modulate_time_attn
+        self.layer_idx = block_kwargs.get("layer_idx")
+        self.log_adaln_mean_abs = False
+        self._last_adaln_mean_abs = None
+
+        if modulate_time_attn:
+            self.adaLN_time_attn_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, 3 * hidden_size, bias=True)
+            )
+            self.norm_time_attn = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+            # initialize
+            nn.init.constant_(self.adaLN_time_attn_modulation[-1].weight, 0)
+            nn.init.constant_(self.adaLN_time_attn_modulation[-1].bias, 0)
+        else:
+            self.norm_time_attn = nn.Identity()
+
+    def initialize_adaln_weights(self, gate_init_std=0.0):
+        nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+        if gate_init_std != 0.0:
+            hidden_size = self.adaLN_modulation[-1].out_features // 9
+            for gate_idx in (2, 5, 8):
+                start = gate_idx * hidden_size
+                end = start + hidden_size
+                nn.init.normal_(self.adaLN_modulation[-1].weight[start:end], std=gate_init_std)
+
+        if hasattr(self, "adaLN_time_attn_modulation"):
+            nn.init.constant_(self.adaLN_time_attn_modulation[-1].weight, 0)
+            nn.init.constant_(self.adaLN_time_attn_modulation[-1].bias, 0)
+            if gate_init_std != 0.0:
+                hidden_size = self.adaLN_time_attn_modulation[-1].out_features // 3
+                start = 2 * hidden_size
+                end = start + hidden_size
+                nn.init.normal_(self.adaLN_time_attn_modulation[-1].weight[start:end], std=gate_init_std)
+
+    def _collect_adaln_mean_abs(self, **tensors):
+        self._last_adaln_mean_abs = {
+            name: tensor.detach().abs().flatten(1).mean(dim=1)
+            for name, tensor in tensors.items()
+        }
+
+    def forward(self, x, c):
+        B, F, N, D = x.shape
+
+        # chunk into 9 [B, C] vectors
+        (shift_msa, scale_msa, gate_msa,
+        shift_mlp_s, scale_mlp_s, gate_mlp_s,
+        shift_mlp_t, scale_mlp_t, gate_mlp_t) = self.adaLN_modulation(c).chunk(9, dim=-1)
+
+        x_modulated = modulate(self.norm1(x), shift_msa, scale_msa)
+        x_modulated = rearrange(x_modulated, 'b f n d -> (b f) n d', b=B, f=F)
+        x_ = self.space_attn(x_modulated)
+        x_ = rearrange(x_, '(b f) n d -> b f n d', b=B, f=F)
+        x = x + _broadcast_gate(gate_msa, x) * x_
+
+        x_modulated = modulate(self.norm2(x), shift_mlp_s, scale_mlp_s)
+        x = x + _broadcast_gate(gate_mlp_s, x) * self.space_mlp(x_modulated)
+
+        # — temporal attention path —
+        if self.modulate_time_attn:
+            shift_mta, scale_mta, gate_mta = self.adaLN_time_attn_modulation(c).chunk(3, dim=-1)
+        else:
+            shift_mta, scale_mta, gate_mta = torch.zeros_like(shift_mlp_t), torch.zeros_like(scale_mlp_t), torch.ones_like(gate_mlp_t)
+
+        if self.log_adaln_mean_abs:
+            self._collect_adaln_mean_abs(
+                msa_shift=shift_msa,
+                msa_scale=scale_msa,
+                msa_gate=gate_msa,
+                mlp_s_shift=shift_mlp_s,
+                mlp_s_scale=scale_mlp_s,
+                mlp_s_gate=gate_mlp_s,
+                mta_shift=shift_mta,
+                mta_scale=scale_mta,
+                mta_gate=gate_mta,
+                mlp_t_shift=shift_mlp_t,
+                mlp_t_scale=scale_mlp_t,
+                mlp_t_gate=gate_mlp_t,
+            )
+
+        x_modulated = modulate(self.norm_time_attn(x), shift_mta, scale_mta)
+        x_modulated = rearrange(x_modulated, 'b f n d -> (b n) f d', b=B, f=F, n=N)
+        time_attn_mask = torch.tril(torch.ones(F, F, device=x.device)) if self.causal_time_attn else None
+        x_ = self.time_attn(x_modulated, attn_mask=time_attn_mask)
+        x_ = rearrange(x_, '(b n) f d -> b f n d', b=B, n=N, f=F)
+        x = x + _broadcast_gate(gate_mta, x) * x_
+
+        x_modulated = modulate(self.norm3(x), shift_mlp_t, scale_mlp_t)
+        x = x + _broadcast_gate(gate_mlp_t, x) * self.time_mlp(x_modulated)
+
+        return x
+
+
+class FinalLayer(nn.Module):
+   """
+   The final layer of DiT.
+   """
+   def __init__(self, hidden_size, patch_size, out_channels, norm_layer=nn.LayerNorm, act_layer=nn.SiLU):
+       super().__init__()
+       self.norm_final = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+       self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+       self.adaLN_modulation = nn.Sequential(
+           act_layer(),
+           nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+       )
+
+
+   def forward(self, x, c):
+       shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+       x = modulate(self.norm_final(x), shift, scale)
+       x = self.linear(x)
+       return x
+
+
+class STDiTBlockWithSpatialL2(STDiTBlock):
+    """
+    STDiTBlock augmented with per-token gated spatial conditioning from L2 tokens.
+
+    z_spatial (B, F, N, D) → (shift_s, scale_s, gate_s) per token via a small MLP.
+    Applied as an additive gated adaLN before the standard block operations:
+        x = x + gate_s * modulate(spatial_norm(x), shift_s, scale_s)
+
+    Zero-init of spatial_to_mod guarantees a no-op at the start of training,
+    so pre-trained STDiTDF weights can be fine-tuned without disruption.
+    """
+
+    def __init__(self, hidden_size, num_heads, norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__(hidden_size, num_heads, norm_layer=norm_layer, **kwargs)
+        self.spatial_norm = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.spatial_to_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 3 * hidden_size, bias=True),
+        )
+        # Small weight init so shift/scale start as genuine L2-dependent perturbations.
+        # Gate bias initialised to 2.0 → sigmoid(2.0) ≈ 0.88, so the gate starts open
+        # and the model must learn to close it rather than having to learn to open it.
+        nn.init.normal_(self.spatial_to_mod[-1].weight, std=0.02)
+        nn.init.zeros_(self.spatial_to_mod[-1].bias)
+        self.spatial_to_mod[-1].bias.data[2 * hidden_size:].fill_(2.0)
+
+    def forward(self, x, c, z_spatial):
+        # z_spatial: (B, F, N, D) — pre-computed L2 spatial tokens (patchified + position)
+        shift_s, scale_s, gate_s = self.spatial_to_mod(z_spatial).chunk(3, dim=-1)
+        x = x + torch.sigmoid(gate_s) * modulate(self.spatial_norm(x), shift_s, scale_s)
+        return super().forward(x, c)

orbis2/modules/quantize.py

@@ -0,0 +1,174 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, reduce
+
+
+def log(t, eps=1e-5):
+    return t.clamp(min=eps).log()
+
+
+def entropy(prob):
+    return (-prob * log(prob)).sum(dim=-1)
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        n_e,
+        e_dim,
+        beta,
+        normalize_embedding,
+        remap=None,
+        unknown_index="random",
+        sane_index_shape=False,
+        legacy=True,
+        diversity_gamma=1.0,
+        frac_per_sample_entropy=1.0,
+        token_noise=0.0,
+    ):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.legacy = legacy
+        self.normalize_embedding = normalize_embedding
+        self.diversity_gamma = diversity_gamma
+        self.frac_per_sample_entropy = frac_per_sample_entropy
+        self.token_noise = token_noise
+        self.sane_index_shape = sane_index_shape
+
+        # Codebook
+        self.embedding = nn.Embedding(n_e, e_dim)
+        self.embedding.weight.data.uniform_(-1.0 / n_e, 1.0 / n_e)
+
+        if self.normalize_embedding:
+            self.embedding.weight.data = F.normalize(self.embedding.weight.data, dim=1)
+
+        # Optional remapping
+        self.remap = remap
+        if remap is not None:
+            self.register_buffer("used", torch.tensor(np.load(remap)))
+            self.re_embed = self.used.shape[0]
+            self.unknown_index = unknown_index
+            if unknown_index == "extra":
+                self.unknown_index = self.re_embed
+                self.re_embed += 1
+            print(
+                f"Remapping {n_e} indices to {self.re_embed} indices. "
+                f"Using {self.unknown_index} for unknown indices."
+            )
+        else:
+            self.re_embed = n_e
+
+    def remap_to_used(self, indices):
+        ishape = indices.shape
+        indices = indices.view(ishape[0], -1)
+        used = self.used.to(indices)
+        match = (indices[:, :, None] == used[None, None, :]).long()
+        new = match.argmax(-1)
+        unknown = match.sum(2) < 1
+        if self.unknown_index == "random":
+            new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(indices.device)
+        else:
+            new[unknown] = self.unknown_index
+        return new.view(ishape)
+
+    def unmap_to_all(self, indices):
+        ishape = indices.shape
+        indices = indices.view(ishape[0], -1)
+        used = self.used.to(indices)
+        if self.re_embed > used.shape[0]:
+            indices[indices >= used.shape[0]] = 0
+        gathered = torch.gather(used.expand(indices.shape[0], -1), 1, indices)
+        return gathered.view(ishape)
+
+    def entropy_loss(self, distances, inv_temperature=100.0):
+        prob = (-distances * inv_temperature).softmax(dim=-1)
+
+        if self.frac_per_sample_entropy < 1.0:
+            num_tokens = prob.shape[0]
+            sample_size = int(num_tokens * self.frac_per_sample_entropy)
+            mask = torch.randperm(num_tokens, device=prob.device)[:sample_size]
+            per_sample_probs = prob[mask]
+        else:
+            per_sample_probs = prob
+
+        per_sample_entropy = entropy(per_sample_probs).mean()
+        avg_prob = reduce(per_sample_probs, "... d -> d", "mean")
+        codebook_entropy = entropy(avg_prob).mean()
+
+        return per_sample_entropy - self.diversity_gamma * codebook_entropy
+
+    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
+        assert temp in (None, 1.0)
+        assert not rescale_logits and not return_logits
+
+        if self.normalize_embedding:
+            self.embedding.weight.data = F.normalize(self.embedding.weight.data, dim=1)
+
+        # Flatten input
+        z = rearrange(z, "b c h w -> b h w c").contiguous()
+        z_flat = z.view(-1, self.e_dim)
+
+        # Compute distances
+        e = self.embedding.weight
+        d = (
+            torch.sum(z_flat ** 2, dim=1, keepdim=True)
+            + torch.sum(e ** 2, dim=1)
+            - 2 * torch.einsum("bd,dn->bn", z_flat, e.T)
+        )
+
+        min_indices = torch.argmin(d, dim=1)
+
+        # Optional token noise
+        if self.token_noise > 0.0 and self.training:
+            noise_mask = torch.rand_like(min_indices.float()) < self.token_noise
+            rand_indices = torch.randint(0, self.n_e, min_indices.shape, device=z.device)
+            min_indices[noise_mask] = rand_indices[noise_mask]
+
+        z_q = self.embedding(min_indices).view_as(z)
+
+        # Compute VQ loss
+        if self.legacy:
+            loss = F.mse_loss(z_q.detach(), z) + self.beta * F.mse_loss(z_q, z.detach())
+        else:
+            loss = self.beta * F.mse_loss(z_q.detach(), z) + F.mse_loss(z_q, z.detach())
+
+        # Optional entropy loss
+        entropy_aux = self.entropy_loss(d) if self.training else None
+
+        # Straight-through estimator
+        z_q = z + (z_q - z).detach()
+
+        # Reshape to original
+        z_q = rearrange(z_q, "b h w c -> b c h w")
+        z = rearrange(z, "b h w c -> b c h w")
+
+        # Remap if needed
+        if self.remap is not None:
+            min_indices = min_indices.view(z.shape[0], -1)
+            min_indices = self.remap_to_used(min_indices).view(-1, 1)
+
+        if self.sane_index_shape:
+            min_indices = min_indices.view(z_q.shape[0], z_q.shape[2], z_q.shape[3])
+
+        return {
+            "quantized": z_q,
+            "quantization_loss": loss,
+            "entropy_loss": entropy_aux,
+            "indices": min_indices,
+        }
+
+    def get_codebook_entry(self, indices, shape):
+        if self.remap is not None:
+            indices = indices.view(shape[0], -1)
+            indices = self.unmap_to_all(indices).view(-1)
+
+        z_q = self.embedding(indices)
+
+        if shape is not None:
+            z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous()
+
+        return z_q

orbis2/modules/steering.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn as nn
+
+
+def _make_scalar_embedder(hidden_size):
+    embedder = nn.Sequential(
+        nn.Linear(1, hidden_size, bias=True),
+        nn.SiLU(),
+        nn.Linear(hidden_size, hidden_size, bias=True),
+    )
+    nn.init.normal_(embedder[0].weight, std=0.02)
+    nn.init.normal_(embedder[0].bias, std=0.02)
+    nn.init.normal_(embedder[2].weight, std=0.02)
+    nn.init.normal_(embedder[2].bias, std=0.02)
+    return embedder
+
+
+class LinearSteeringEmbedder(torch.nn.Module):
+    """
+    Generic linear embedding layer for steering signals. When steering is missing, outputs a fixed zero embedding.
+    """
+    def __init__(self, num_input_features, hidden_size, learnable_no_value_embedding=False):
+        super().__init__()
+        self.embedders = nn.ModuleList()
+        self.register_buffer("no_value_embeddings", torch.zeros(num_input_features, hidden_size))
+        if learnable_no_value_embedding:
+            self.no_value_embeddings = nn.Parameter(torch.zeros(num_input_features, hidden_size))
+        for _i in range(num_input_features):
+            self.embedders.append(_make_scalar_embedder(hidden_size))
+
+    def forward(self, steering):
+        squeeze_k = False
+        if steering.ndim == 2:
+            steering = steering.unsqueeze(1)
+            squeeze_k = True
+        elif steering.ndim != 3:
+            raise ValueError(f"Expected steering to have shape [B, D] or [B, K, D], got {tuple(steering.shape)}")
+
+        _b, _k, num_features = steering.shape
+        assert num_features == len(self.embedders), f"Expected {len(self.embedders)} features, but got {num_features}"
+        embeddings = []
+        for i in range(num_features):
+            feature = steering[:, :, [i]]
+            missing = torch.isnan(feature).squeeze(-1).unsqueeze(-1)
+            embedding = self.embedders[i](torch.nan_to_num(feature, nan=0.0))
+            if missing.any():
+                embedding = torch.where(missing, self.no_value_embeddings[i].to(embedding.dtype).view(1, 1, -1), embedding)
+            embeddings.append(embedding)
+        steering_embedding = torch.stack(embeddings, dim=1).sum(dim=1)
+        if squeeze_k:
+            return steering_embedding[:, 0]
+        return steering_embedding

orbis2/networks/DiT/dit.py

@@ -0,0 +1,1026 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# GLIDE: https://github.com/openai/glide-text2im
+# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
+# --------------------------------------------------------
+import math
+import numpy as np
+
+
+from timm.layers.mlp import SwiGLU
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from torch import nn
+
+
+from einops import rearrange
+from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
+from omegaconf import ListConfig
+
+from ..swin.swin_free_aspect_ratio import SwinTransformerBlock, SwinAttention
+
+
+
+from modules.dit import (
+    get_norm_layer, modulate, FrequencyEncoder,
+    get_2d_sincos_pos_embed, get_2d_sincos_pos_embed_from_grid,
+    get_1d_sincos_pos_embed_from_grid, TimestepEmbedder,
+    STBlock, DiTBlock, CDiTBlock,
+    STDiTBlock, FinalLayer, STDiTBlockWithSpatialL2,
+)
+
+
+class Lambda(nn.Module):
+   def __init__(self, func):
+       super().__init__()
+       self.func = func
+
+
+   def forward(self, x):
+       return self.func(x)
+ 
+#################################################################################
+#                                 Cross Attention                               #
+#################################################################################
+
+
+class CrossAttention(nn.Module):
+   """
+   Cross-attention mechanism that computes attention between target (x) and context (y).
+   Args:
+       dim (int): Input dimension.
+       num_heads (int): Number of attention heads.
+       qkv_bias (bool): Whether to include bias terms.
+       attn_drop (float): Dropout rate for attention weights.
+       proj_drop (float): Dropout rate after projection.
+   """
+   def __init__(self, dim, num_heads=8, qkv_bias=True, qk_norm=True, attn_drop=0.0, proj_drop=0.0, norm_layer=nn.LayerNorm):
+       super(CrossAttention, self).__init__()
+       self.num_heads = num_heads
+       self.head_dim = dim // num_heads
+       self.scale = self.head_dim ** -0.5
+
+
+       self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
+       self.k_proj = nn.Linear(dim, dim, bias=qkv_bias)
+       self.v_proj = nn.Linear(dim, dim, bias=qkv_bias)
+       self.proj = nn.Linear(dim, dim)
+       self.attn_drop = nn.Dropout(attn_drop)
+       self.proj_drop = nn.Dropout(proj_drop)
+       self.qk_norm = qk_norm
+       if qk_norm:
+           self.q_norm = norm_layer(self.head_dim, eps=1e-6, elementwise_affine=False, bias=False)
+           self.k_norm = norm_layer(self.head_dim, eps=1e-6, elementwise_affine=False, bias=False)
+
+
+   def forward(self, x, y):
+       B, N, C = x.shape
+       B, M, _ = y.shape
+
+
+       q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+       k = self.k_proj(y).reshape(B, M, self.num_heads, self.head_dim).transpose(1, 2)
+       v = self.v_proj(y).reshape(B, M, self.num_heads, self.head_dim).transpose(1, 2)
+       if self.qk_norm:
+           q, k = self.q_norm(q), self.k_norm(k)
+       attn_output = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0)
+       attn_output = attn_output.transpose(1, 2).reshape(B, N, C)
+       return self.proj_drop(self.proj(attn_output))
+
+
+
+
+
+
+
+
+#################################################################################
+#                                 Core DiT Model                                #
+#################################################################################
+
+
+
+class STDiTBlock_tmpadaLN(nn.Module):
+   """
+   A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+   """
+   def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, dropout_rate=0.0, causal_time_attn=False, **block_kwargs):
+       raise NotImplementedError("This is a deprecated version of STDiTBlock with temporary adaLN for time attention. Use STDiTBlock instead.")
+       super().__init__()
+       self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+       self.space_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True, norm_layer=nn.LayerNorm, attn_drop=dropout_rate, proj_drop=dropout_rate, **block_kwargs)
+       self.time_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True, norm_layer=nn.LayerNorm, attn_drop=dropout_rate, proj_drop=dropout_rate, **block_kwargs)
+       self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+       self.norm3 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+       self.norm4 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+       mlp_hidden_dim = int(hidden_size * mlp_ratio)
+       approx_gelu = lambda: nn.GELU(approximate="tanh")
+       self.space_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+       self.time_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+       self.adaLN_modulation = nn.Sequential(
+           nn.SiLU(),
+           nn.Linear(hidden_size, 12 * hidden_size, bias=True)
+       )
+       self.causal_time_attn = causal_time_attn
+
+
+   def forward(self, x, c):
+       B, F, N, D = x.shape
+
+
+       # chunk into 12 [B, C] vectors
+       (shift_msa_s, scale_msa_s, gate_msa_s,
+        shift_msa_t, scale_msa_t, gate_msa_t,
+        shift_mlp_s, scale_mlp_s, gate_mlp_s,
+        shift_mlp_t, scale_mlp_t, gate_mlp_t) = self.adaLN_modulation(c).chunk(12, dim=1)
+
+
+       x_modulated = modulate(self.norm1(x), shift_msa_s, scale_msa_s)
+       x_modulated = rearrange(x_modulated, 'b f n d -> (b f) n d', b=B, f=F)
+       x_ = self.space_attn(x_modulated)
+       x = x + gate_msa_s.unsqueeze(1).unsqueeze(1) * rearrange(x_, '(b f) n d -> b f n d', b=B, f=F)
+
+
+       x_modulated = modulate(self.norm2(x), shift_mlp_s, scale_mlp_s)
+       x = x + gate_mlp_s.unsqueeze(1).unsqueeze(1) * self.space_mlp(x_modulated)
+
+
+       # — temporal attention path —
+       x_modulated = modulate(self.norm3(x), shift_msa_t, scale_msa_t)
+       x_modulated = rearrange(x_modulated, 'b f n d -> (b n) f d', b=B, f=F, n=N)
+       time_attn_mask = torch.tril(torch.ones(F, F, device=x.device)) if self.causal_time_attn else None
+       x_ = self.time_attn(x_modulated, attn_mask=time_attn_mask)
+       x = x + gate_msa_t.unsqueeze(1).unsqueeze(1) * rearrange(x_, '(b n) f d -> b f n d', b=B, f=F)
+       x_modulated = modulate(self.norm4(x), shift_mlp_t, scale_mlp_t)
+       x = x + gate_mlp_t.unsqueeze(1).unsqueeze(1) * self.time_mlp(x_modulated) 
+       return x
+
+
+class SwinSTDiTBlock(STDiTBlock):
+    def __init__(self, hidden_size, num_heads, input_shape, layer_idx, mlp_ratio=4.0, window_size=[6, 4], dropout_rate=0.0, 
+                 causal_time_attn=False, modulate_time_attn=False, 
+                 norm_layer=nn.LayerNorm, mlp_block='mlp', **block_kwargs):
+        super().__init__(hidden_size=hidden_size, num_heads=num_heads, mlp_ratio=mlp_ratio, dropout_rate=dropout_rate, 
+                         causal_time_attn=causal_time_attn, modulate_time_attn=modulate_time_attn, mlp_block=mlp_block)
+        self.norm1 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.space_attn = SwinTransformerBlock(
+                hidden_size,
+                input_resolution=input_shape,
+                num_heads=num_heads, 
+                window_size=window_size,
+                shift_size=(0,0) if (layer_idx % 2 == 0) else [ws//2 for ws in window_size],
+                qk_norm=False,
+                mlp_ratio=mlp_ratio,
+                drop=dropout_rate,
+                attn_drop=dropout_rate,
+                norm_layer=norm_layer,
+                mlp_block=mlp_block,
+                **block_kwargs
+                )
+
+
+class SwinSTDiTBlockNoExtraMLP(STDiTBlock):
+    def __init__(self, hidden_size, num_heads, input_shape, layer_idx, mlp_ratio=4.0, window_size=[6, 4], dropout_rate=0.0, norm_layer=nn.LayerNorm, mlp_block='mlp', **block_kwargs):
+        super().__init__(hidden_size, num_heads, input_shape, layer_idx, mlp_ratio, dropout_rate, norm_layer=norm_layer, mlp_block=mlp_block, **block_kwargs)
+        self.space_attn = SwinAttention(
+            hidden_size,
+            input_resolution=input_shape,
+            num_heads=num_heads,
+            window_size=window_size,
+            shift_size=(0,0) if (layer_idx % 2 == 0) else [ws//2 for ws in window_size],
+            proj_drop=dropout_rate,
+            attn_drop=dropout_rate,
+            qk_norm=True,
+            norm_layer=norm_layer,
+            **block_kwargs,
+        )
+
+
+
+class STDiTBlockWithRegisters(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        mlp_ratio=4.0,
+        dropout_rate=0.0,
+        causal_time_attn=False,
+        modulate_time_attn=True,
+        norm_layer=nn.LayerNorm,
+        mlp_block='mlp',
+        num_reg_tokens=1,   # must be > 0
+        **block_kwargs
+    ):
+        super().__init__()
+        if num_reg_tokens <= 0:
+            raise ValueError("STDiTBlockWithRegisters now assumes num_reg_tokens > 0.")
+        assert modulate_time_attn, "STDiTBlockWithRegisters currently requires modulate_time_attn=True."
+        self.num_reg_tokens = num_reg_tokens
+
+        # norms...
+        self.norm1 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm2 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm3 = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_reg_attn_s = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_reg_s = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_reg_attn_t = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.norm_reg_t = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+
+        self.space_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True,
+                                    norm_layer=norm_layer, attn_drop=dropout_rate, proj_drop=dropout_rate, **block_kwargs)
+        self.time_attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True,
+                                   norm_layer=norm_layer, attn_drop=dropout_rate, proj_drop=dropout_rate, **block_kwargs)
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        if mlp_block == 'mlp':
+            approx_gelu = lambda: nn.GELU(approximate="tanh")
+            self.space_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=None, drop=0)
+            self.time_mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=None, drop=0)
+            self.space_mlp_reg = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=None, drop=0)
+            self.time_mlp_reg = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, norm_layer=None, drop=0)
+        elif mlp_block == 'swiglu':
+            self.space_mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim * 2) // 3)
+            self.time_mlp = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim * 2) // 3)
+            self.space_mlp_reg = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim * 2) // 3)
+            self.time_mlp_reg = SwiGLU(in_features=hidden_size, hidden_features=(mlp_hidden_dim * 2) // 3)
+        else:
+            raise NotImplementedError(f"mlp_block {mlp_block} not implemented")
+
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 9 * hidden_size, bias=True))
+        self.adaLN_registers_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 12 * hidden_size, bias=True),
+        )
+
+        self.causal_time_attn = causal_time_attn
+        self.modulate_time_attn = modulate_time_attn
+
+        # Always-on time-attn modulation
+        self.adaLN_time_attn_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 3 * hidden_size, bias=True)
+        )
+        self.norm_time_attn = norm_layer(hidden_size, elementwise_affine=False, eps=1e-6)
+        nn.init.constant_(self.adaLN_time_attn_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_time_attn_modulation[-1].bias, 0)
+
+        self.log_adaln_mean_abs = False
+        self._last_adaln_mean_abs = {}
+
+    def initialize_adaln_weights(self, gate_init_std=0.0):
+        nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.adaLN_registers_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_registers_modulation[-1].bias, 0)
+        nn.init.constant_(self.adaLN_time_attn_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_time_attn_modulation[-1].bias, 0)
+
+        if gate_init_std != 0.0:
+            hidden_size = self.adaLN_modulation[-1].out_features // 9
+            for gate_idx in (2, 5, 8):
+                start = gate_idx * hidden_size
+                end = start + hidden_size
+                nn.init.normal_(self.adaLN_modulation[-1].weight[start:end], std=gate_init_std)
+
+            hidden_size = self.adaLN_registers_modulation[-1].out_features // 12
+            for gate_idx in (2, 5, 8, 11):
+                start = gate_idx * hidden_size
+                end = start + hidden_size
+                nn.init.normal_(self.adaLN_registers_modulation[-1].weight[start:end], std=gate_init_std)
+
+            hidden_size = self.adaLN_time_attn_modulation[-1].out_features // 3
+            start = 2 * hidden_size
+            end = start + hidden_size
+            nn.init.normal_(self.adaLN_time_attn_modulation[-1].weight[start:end], std=gate_init_std)
+
+    def _collect_adaln_mean_abs(self, **tensors):
+        self._last_adaln_mean_abs = {
+            name: tensor.detach().abs().flatten(1).mean(dim=1)
+            for name, tensor in tensors.items()
+        }
+
+    def forward(self, x, c, c_registers):
+        B, F, N, D = x.shape
+        R = self.num_reg_tokens
+
+        # always has registers
+        reg_tokens = x[:, :, :R, :]
+        x_patch = x[:, :, R:, :]
+
+        (shift_msa, scale_msa, gate_msa,
+         shift_mlp_s, scale_mlp_s, gate_mlp_s,
+         shift_mlp_t, scale_mlp_t, gate_mlp_t) = self.adaLN_modulation(c).chunk(9, dim=-1)
+        shift_mta, scale_mta, gate_mta = self.adaLN_time_attn_modulation(c).chunk(3, dim=-1)
+        (shift_reg_attn_s, scale_reg_attn_s, gate_reg_attn_s,
+         shift_reg_mlp_s, scale_reg_mlp_s, gate_reg_mlp_s,
+         shift_reg_attn_t, scale_reg_attn_t, gate_reg_attn_t,
+         shift_reg_mlp_t, scale_reg_mlp_t, gate_reg_mlp_t) = self.adaLN_registers_modulation(c_registers).chunk(12, dim=-1)
+
+        if self.log_adaln_mean_abs:
+            self._collect_adaln_mean_abs(
+                msa_shift=shift_msa, msa_scale=scale_msa, msa_gate=gate_msa,
+                mlp_s_shift=shift_mlp_s, mlp_s_scale=scale_mlp_s, mlp_s_gate=gate_mlp_s,
+                mta_shift=shift_mta, mta_scale=scale_mta, mta_gate=gate_mta,
+                mlp_t_shift=shift_mlp_t, mlp_t_scale=scale_mlp_t, mlp_t_gate=gate_mlp_t,
+                reg_attn_s_shift=shift_reg_attn_s, reg_attn_s_scale=scale_reg_attn_s, reg_attn_s_gate=gate_reg_attn_s,
+                reg_mlp_s_shift=shift_reg_mlp_s, reg_mlp_s_scale=scale_reg_mlp_s, reg_mlp_s_gate=gate_reg_mlp_s,
+                reg_attn_t_shift=shift_reg_attn_t, reg_attn_t_scale=scale_reg_attn_t, reg_attn_t_gate=gate_reg_attn_t,
+                reg_mlp_t_shift=shift_reg_mlp_t, reg_mlp_t_scale=scale_reg_mlp_t, reg_mlp_t_gate=gate_reg_mlp_t,
+            )
+        
+        # spatial attn (always concat regs)
+        x_mod = rearrange(modulate(self.norm1(x_patch), shift_msa, scale_msa), 'b f n d -> (b f) n d')
+        reg_in = rearrange(
+            modulate(self.norm_reg_attn_s(reg_tokens), shift_reg_attn_s, scale_reg_attn_s),
+            'b f r d -> (b f) r d'
+        )
+        x_mod = torch.cat([reg_in, x_mod], dim=1)
+        x_out = self.space_attn(x_mod)
+
+        reg_attn_out, x_out = torch.split(x_out, [R, x_patch.size(2)], dim=1)
+        reg_attn_out = rearrange(reg_attn_out, '(b f) r d -> b f r d', b=B, f=F)
+        x_out = rearrange(x_out, '(b f) n d -> b f n d', b=B, f=F)
+        x_patch = x_patch + gate_msa.unsqueeze(1).unsqueeze(1) * x_out if gate_msa.ndim == 2 else x_patch + gate_msa.unsqueeze(2) * x_out
+        reg_tokens = reg_tokens + gate_reg_attn_s.unsqueeze(1).unsqueeze(1) * reg_attn_out
+
+        # spatial mlp
+        x_patch = x_patch + gate_mlp_s.unsqueeze(1).unsqueeze(1) * self.space_mlp(modulate(self.norm2(x_patch), shift_mlp_s, scale_mlp_s)) if gate_mlp_s.ndim == 2 else x_patch + gate_mlp_s.unsqueeze(2) * self.space_mlp(modulate(self.norm2(x_patch), shift_mlp_s, scale_mlp_s))
+        reg_mlp_out = self.space_mlp_reg(modulate(self.norm_reg_s(reg_tokens), shift_reg_mlp_s, scale_reg_mlp_s))
+        reg_tokens = reg_tokens + gate_reg_mlp_s.unsqueeze(1).unsqueeze(1) * reg_mlp_out
+
+        # temporal attn patches
+        x_mod = modulate(self.norm_time_attn(x_patch), shift_mta, scale_mta)
+        x_mod = rearrange(x_mod, 'b f n d -> (b n) f d')
+        time_attn_mask = torch.tril(torch.ones(F, F, device=x.device)) if self.causal_time_attn else None
+        x_out = self.time_attn(x_mod, attn_mask=time_attn_mask)
+        x_out = rearrange(x_out, '(b n) f d -> b f n d', b=B, n=x_patch.size(2), f=F)
+        x_patch = x_patch + gate_mta.unsqueeze(1).unsqueeze(1) * x_out if gate_mta.ndim == 2 else x_patch + gate_mta.unsqueeze(2) * x_out
+
+        # temporal attn regs (always)
+        reg_time_in = rearrange(
+            modulate(self.norm_reg_attn_t(reg_tokens), shift_reg_attn_t, scale_reg_attn_t),
+            'b f r d -> (b r) f d'
+        )
+        reg_time_out = self.time_attn(reg_time_in, attn_mask=time_attn_mask)
+        reg_time_out = rearrange(reg_time_out, '(b r) f d -> b f r d', b=B, r=R, f=F)
+        reg_tokens = reg_tokens + gate_reg_attn_t.unsqueeze(1).unsqueeze(1) * reg_time_out
+
+        # temporal mlp
+        x_patch = x_patch + gate_mlp_t.unsqueeze(1).unsqueeze(1) * self.time_mlp(modulate(self.norm3(x_patch), shift_mlp_t, scale_mlp_t)) if gate_mlp_t.ndim == 2 else x_patch + gate_mlp_t.unsqueeze(2) * self.time_mlp(modulate(self.norm3(x_patch), shift_mlp_t, scale_mlp_t))
+        reg_mlp_out = self.time_mlp_reg(modulate(self.norm_reg_t(reg_tokens), shift_reg_mlp_t, scale_reg_mlp_t))
+        reg_tokens = reg_tokens + gate_reg_mlp_t.unsqueeze(1).unsqueeze(1) * reg_mlp_out
+        
+        return torch.cat([reg_tokens, x_patch], dim=2)
+
+
+
+
+
+class DiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        input_size=16,
+        patch_size=2,
+        in_channels=32,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        max_num_frames=6,
+        dropout=0.0,
+        ctx_noise_aug_ratio=0.1,
+        ctx_noise_aug_prob = 0.5,
+        drop_ctx_rate=0.2,
+        frequency_range=(2, 15),
+        learn_sigma=False,
+        norm_layer=nn.LayerNorm,
+        mlp_block='mlp',
+    ):
+        super().__init__()
+        self.input_size= input_size if isinstance(input_size, (list, tuple, ListConfig)) else [input_size, input_size]
+        self.in_channels = in_channels
+        self.out_channels = in_channels * 2 if learn_sigma else in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.ctx_noise_aug_ratio = ctx_noise_aug_ratio
+        self.ctx_noise_aug_prob = ctx_noise_aug_prob
+        self.drop_ctx_rate = drop_ctx_rate
+
+
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
+        self.num_patches = self.x_embedder.num_patches
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, hidden_size), requires_grad=False)
+
+        self.blocks = nn.ModuleList([
+            DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, dropout_rate=dropout, norm_layer=norm_layer, mlp_block=mlp_block) for _ in range(depth)
+        ])       
+        
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels, norm_layer=norm_layer)
+        self.max_num_frames = max_num_frames
+        self.frame_emb = nn.init.trunc_normal_(nn.Parameter(torch.zeros(1, self.max_num_frames, 1, hidden_size)), 0., 0.02)
+        self.frame_rate_encoder = FrequencyEncoder(hidden_size, freq_min=frequency_range[0], freq_max=frequency_range[1])
+
+        self.initialize_weights()
+
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], [self.input_size[0] // self.patch_size, self.input_size[1] // self.patch_size], cls_token=False, extra_tokens=0)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+
+    def unpatchify(self, x):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = self.x_embedder.grid_size[0]
+        w = self.x_embedder.grid_size[1]
+
+
+        x = x.reshape(shape=(x.shape[0], x.shape[1], h, w, p, p, c))
+        x = torch.einsum('bfhwpqc->bfchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], x.shape[1], c, h * p, w * p))
+        return imgs
+
+
+    def get_condition_embeddings(self, t):
+        """
+        Get the condition embeddings for the given timesteps.
+        t: (N,) tensor of diffusion timesteps
+        returns: (N, D) tensor of condition embeddings
+        """
+        return self.t_embedder(t)
+    
+    def preprocess_inputs(self, target, context, t, frame_rate):
+        b, f_target = target.size()[:2]
+        f_context = context.size(1)
+        
+        if self.training:
+            # Drop the context frame
+            if torch.rand(1, device=target.device)<self.drop_ctx_rate:
+                context = None
+                f_context = 0
+            elif torch.rand(1, device=target.device) < self.ctx_noise_aug_prob:
+                # Add noise to context frames (if t is less than ctx_noise_aug_ratio, we do not add noise)
+                mask = (t >= self.ctx_noise_aug_ratio)
+                aug_noise = torch.randn_like(context)
+                context[mask] = context[mask] + aug_noise[mask] * self.ctx_noise_aug_ratio
+
+
+        frame_embeddings = self.frame_rate_encoder.encode(frame_rate)
+        frame_embeddings = frame_embeddings.unsqueeze(1).unsqueeze(1).to(target.device)
+        x = torch.cat((context, target), dim=1) if context is not None else target
+        x = rearrange(x, 'b f c h w -> (b f) c h w')
+        x = self.x_embedder(x) + self.pos_embed.to(x.device)
+        x = rearrange(x, '(b f) hw c -> b f hw c', b=b)
+        x = x + self.frame_emb[:, self.max_num_frames-(f_target+f_context):].to(x.device) + frame_embeddings
+        return x
+
+    def get_condition_embeddings(self, t):
+        """
+        Get the condition embeddings for the given timesteps.
+        t: (N,) tensor of diffusion timesteps
+        returns: (N, D) tensor of condition embeddings
+        """
+        return self.t_embedder(t)
+    
+    def preprocess_inputs(self, target, context, t, frame_rate):
+        b, f_target = target.size()[:2]
+        f_context = context.size(1) if context is not None else 0
+        
+        if self.training:
+            # Drop the context frame
+            if torch.rand(1, device=target.device)<self.drop_ctx_rate:
+                context = None
+                f_context = 0
+            elif torch.rand(1, device=target.device) < self.ctx_noise_aug_prob:
+                # Add noise to context frames (if t is less than ctx_noise_aug_ratio, we do not add noise)
+                mask = (t >= self.ctx_noise_aug_ratio)
+                aug_noise = torch.randn_like(context)
+                context[mask] = context[mask] + aug_noise[mask] * self.ctx_noise_aug_ratio
+
+
+        frame_embeddings = self.frame_rate_encoder.encode(frame_rate)
+        frame_embeddings = frame_embeddings.unsqueeze(1).unsqueeze(1).to(target.device)
+        x = torch.cat((context, target), dim=1) if context is not None else target
+        x = rearrange(x, 'b f c h w -> (b f) c h w')
+        x = self.x_embedder(x) + self.pos_embed.to(x.device)
+        x = rearrange(x, '(b f) hw c -> b f hw c', b=b)
+        x = x + self.frame_emb[:, self.max_num_frames-(f_target+f_context):].to(x.device) + frame_embeddings
+        return x
+
+    def postprocess_outputs(self, out):
+        return self.unpatchify(out)
+
+    def forward(self, target, context, t, frame_rate, return_features=False):
+            """
+            Forward pass of DiT.
+            x: (N, F, C, H, W) tensor of spatial inputs (images or latent representations of images)
+            t: (N,) tensor of diffusion timesteps
+            y: (N,) tensor of class labels
+            """
+            
+            num_frames_ctx = context.size(1)
+            num_frames_pred = target.size(1)
+            
+            c = self.get_condition_embeddings(t)
+            
+            x = self.preprocess_inputs(target, context, t, frame_rate)
+            
+            x = rearrange(x,  'b f hw c -> b (f hw) c')
+            features = []
+            for block in self.blocks:
+                x = block(x, c)
+                features.append(x) if return_features else None
+            x = rearrange(x,  'b (f hw) c -> b f hw c', f=(num_frames_ctx+num_frames_pred))[:,-num_frames_pred:]
+            out = self.final_layer(x, c)
+                        
+            out = self.postprocess_outputs(out)
+            if return_features:
+                return out, features
+            return out
+
+
+class CDiT(DiT):
+    def __init__(self, input_size=16, patch_size=2, in_channels=32, hidden_size=1152, depth=28, num_heads=16, mlp_ratio=4.0, max_num_frames=6, dropout=0.1, ctx_noise_aug_ratio=0.1,ctx_noise_aug_prob=0.5, norm_layer=nn.LayerNorm, mlp_block='mlp', **kwargs):
+        if isinstance(norm_layer, str):
+            norm_layer = get_norm_layer(norm_layer)
+        super().__init__(input_size=input_size, patch_size=patch_size, in_channels=in_channels, hidden_size=hidden_size, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio, max_num_frames=max_num_frames, dropout=dropout, 
+                         ctx_noise_aug_ratio=ctx_noise_aug_ratio, ctx_noise_aug_prob=ctx_noise_aug_prob, norm_layer=norm_layer, mlp_block=mlp_block, **kwargs)
+        self.blocks = nn.ModuleList([
+                CDiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, norm_layer=norm_layer, mlp_block=mlp_block) for _ in range(depth)
+            ])
+
+
+    def forward(self, target, context, t, frame_rate, return_features=False):
+        """
+        Forward pass of DiT.
+        x: (N, F, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        """
+        num_frames_ctx = context.size(1)
+        num_frames_pred = target.size(1)
+        
+        c = self.get_condition_embeddings(t)                   # (N, D)
+        
+        x = self.preprocess_inputs(target, context, t, frame_rate)  # (B, F, N, D)
+
+
+        target = rearrange(x[:,-num_frames_pred:],  'b f hw c -> b (f hw) c')
+        ctx = rearrange(x[:,:-num_frames_pred],  'b f hw c -> b (f hw) c') if num_frames_ctx>1 else None
+
+        features = []
+        for block in self.blocks:
+            target = block(target, c, ctx)                      # (N, T, D)
+            features.append(target)
+
+        target = rearrange(target,  'b (f hw) c -> b f hw c', f=(num_frames_pred))
+        out = self.final_layer(target, c)                # (N, T, patch_size * out_channels)
+        out = self.postprocess_outputs(out)  # (N, T, patch_size ** 2 * out_channels)
+        if return_features:
+            return out, features
+        return out
+
+
+
+class STDiT(DiT):
+    def __init__(self, input_size=16, patch_size=2, in_channels=32, hidden_size=1152, depth=28, num_heads=16, mlp_ratio=4.0, max_num_frames=6, 
+                 dropout=0.1, ctx_noise_aug_ratio=0.1, ctx_noise_aug_prob=0.5, drop_ctx_rate=0.2, frequency_range=(2, 15), 
+                 causal_time_attn=False, modulate_time_attn=False, norm_layer=nn.LayerNorm, mlp_block='mlp',
+                 **kwargs):
+        
+        if isinstance(norm_layer, str):
+            norm_layer = get_norm_layer(norm_layer)
+            
+        super().__init__(input_size=input_size, patch_size=patch_size, in_channels=in_channels, hidden_size=hidden_size, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio, 
+                         max_num_frames=max_num_frames, dropout=dropout, ctx_noise_aug_ratio=ctx_noise_aug_ratio, ctx_noise_aug_prob=ctx_noise_aug_prob, drop_ctx_rate=drop_ctx_rate, 
+                         norm_layer=norm_layer, mlp_block=mlp_block, **kwargs)
+        self.blocks = nn.ModuleList([
+                STDiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, dropout_rate=dropout, 
+                           causal_time_attn=causal_time_attn, modulate_time_attn=modulate_time_attn, 
+                           norm_layer=norm_layer, mlp_block=mlp_block) for _ in range(depth)
+            ])
+
+    def forward(self, target, context, t, frame_rate, return_features=False):
+        """
+        Forward pass of DiT.
+        x: (N, F, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        """
+        f_pred = target.size(1)
+        
+        c = self.get_condition_embeddings(t)                   # (N, D)
+        
+        x = self.preprocess_inputs(target, context, t, frame_rate)  # (B, F, N, D)
+
+        features = []
+        for block in self.blocks:
+            x = block(x, c)
+            features.append(x) if return_features else None
+
+
+        out = self.final_layer(x[:,-f_pred:], c)                # (N, T, patch_size * out_channels)
+        out = self.postprocess_outputs(out)  # (N, T, patch_size ** 2 * out_channels)
+        if return_features:
+            return out, features
+        return out
+
+class STDiTDF(STDiT):
+    def __init__(self, input_size=16, patch_size=2, in_channels=32, hidden_size=1152, depth=28, num_heads=16, mlp_ratio=4.0, max_num_frames=6, 
+                 dropout=0.1, ctx_noise_aug_ratio=0.1, ctx_noise_aug_prob=0.5, drop_ctx_rate=0.2, frequency_range=(2, 15), 
+                 causal_time_attn=False, modulate_time_attn=False, norm_layer=nn.LayerNorm, mlp_block='mlp',
+                 **kwargs):
+        
+        if isinstance(norm_layer, str):
+            norm_layer = get_norm_layer(norm_layer)
+            
+        super().__init__(input_size=input_size, patch_size=patch_size, in_channels=in_channels, hidden_size=hidden_size, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio, 
+                         max_num_frames=max_num_frames, dropout=dropout, ctx_noise_aug_ratio=ctx_noise_aug_ratio, ctx_noise_aug_prob=ctx_noise_aug_prob, drop_ctx_rate=drop_ctx_rate, 
+                         frequency_range=frequency_range, causal_time_attn=causal_time_attn, modulate_time_attn=modulate_time_attn, norm_layer=norm_layer, mlp_block=mlp_block, **kwargs)
+
+        self.blocks = nn.ModuleList([
+            STDiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, dropout_rate=dropout, 
+                        causal_time_attn=causal_time_attn, modulate_time_attn=modulate_time_attn, 
+                        norm_layer=norm_layer, mlp_block=mlp_block) for _ in range(depth)
+        ])
+        
+    def preprocess_inputs(self, x, frame_rate):
+        b, f = x.size()[:2]
+
+        frame_embeddings = self.frame_rate_encoder.encode(frame_rate)
+        frame_embeddings = frame_embeddings.unsqueeze(1).unsqueeze(1).to(x.device)
+        x = rearrange(x, 'b f c h w -> (b f) c h w')
+        x = self.x_embedder(x) + self.pos_embed.to(x.device)
+        x = rearrange(x, '(b f) hw c -> b f hw c', b=b)
+        x = x + self.frame_emb[:, :f].to(x.device) + frame_embeddings
+        return x
+
+    def forward(self, x, t, frame_rate, return_features=False):
+        """
+        Forward pass of DiT.
+        x: (N, F, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        """
+        f_pred = x.size(1)
+        
+        t = rearrange(t, 'b f -> (b f)')
+        c = self.get_condition_embeddings(t)                   # (N, D)
+        c = rearrange(c, '(b f) d -> b f d', b=x.size(0), f=x.size(1))
+        
+        x = self.preprocess_inputs(x, frame_rate)  # (B, F, N, D)
+        
+        features = []
+        for block in self.blocks:
+            x = block(x, c)
+            features.append(x) if return_features else None
+
+        out = self.final_layer(x, c)                # (N, T, patch_size * out_channels)
+        out = self.postprocess_outputs(out)  # (N, T, patch_size ** 2 * out_channels)
+        if return_features:
+            return out, features
+        return out
+    
+    
+class STDiTWithGlobalBlocks(STDiT):
+    """
+    A DiT with extra global attention blocks (DiT blocks).
+    """
+    def __init__(self, *args, global_block_indices, **kwargs):
+        
+        class DiTBlockWrapper(nn.Module):
+            def __init__(self, block):
+                super().__init__()
+                self.block = block
+
+            def forward(self, x, c):
+                B, F, N, D = x.shape
+                x_ = rearrange(x, 'b f n d -> b (f n) d', b=B, f=F)
+                x_ = self.block(x_, c)
+                x_ = rearrange(x_, 'b (f n) d -> b f n d', b=B, f=F)
+                return x_
+        
+        super().__init__(*args, **kwargs)
+        self.global_block_indices = global_block_indices
+        for idx in global_block_indices:
+            dit_block = DiTBlock(hidden_size=kwargs.get('hidden_size',1152), num_heads=kwargs.get('num_heads',16),
+                                 mlp_ratio=kwargs.get('mlp_ratio',4.0), dropout_rate=kwargs.get('dropout',0.0), 
+                                 norm_layer=kwargs.get('norm_layer',nn.LayerNorm), mlp_block=kwargs.get('mlp_block','mlp'))
+            self.blocks[idx] = DiTBlockWrapper(dit_block)
+            
+            
+class STDiT_tmpadaLN(STDiT):
+    def __init__(self, input_size=16, patch_size=2, in_channels=32, hidden_size=1152, depth=28, num_heads=16, mlp_ratio=4.0, max_num_frames=6, dropout=0.1, ctx_noise_aug_ratio=0.1, ctx_noise_aug_prob=0.5, drop_ctx_rate=0.2, frequency_range=(2, 15), causal_time_attn=False, norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__(input_size=input_size, patch_size=patch_size, in_channels=in_channels, hidden_size=hidden_size, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio, max_num_frames=max_num_frames, dropout=dropout, ctx_noise_aug_ratio=ctx_noise_aug_ratio, ctx_noise_aug_prob=ctx_noise_aug_prob, drop_ctx_rate=drop_ctx_rate, norm_layer=norm_layer, **kwargs)
+        self.blocks = nn.ModuleList([
+                STDiTBlock_tmpadaLN(hidden_size, num_heads, mlp_ratio=mlp_ratio, dropout_rate=dropout, causal_time_attn=causal_time_attn, norm_layer=norm_layer) for _ in range(depth)
+            ])
+
+
+class SwinSTDiT(STDiT):
+    def __init__(self, input_size=16, patch_size=2, in_channels=32, hidden_size=1152, depth=28, num_heads=16, mlp_ratio=4.0, max_num_frames=6, window_size=[6, 4], dropout=0.1, ctx_noise_aug_ratio=0.1,ctx_noise_aug_prob=0.5, drop_ctx_rate=0.2, frequency_range=(2, 15), norm_layer=nn.LayerNorm, mlp_block='mlp', **kwargs):
+        super().__init__(input_size=input_size, patch_size=patch_size, in_channels=in_channels, hidden_size=hidden_size, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio, max_num_frames=max_num_frames, dropout=dropout, ctx_noise_aug_ratio=ctx_noise_aug_ratio, ctx_noise_aug_prob=ctx_noise_aug_prob, drop_ctx_rate=drop_ctx_rate, norm_layer=norm_layer, mlp_block=mlp_block, **kwargs)
+        self.blocks = nn.ModuleList([
+                SwinSTDiTBlock(hidden_size=hidden_size, num_heads=num_heads, input_shape=input_size, layer_idx=layer_idx, mlp_ratio=mlp_ratio, window_size=window_size, dropout_rate=dropout, norm_layer=norm_layer) for layer_idx in range(depth)
+            ])
+
+
+class SwinSTDiTNoExtraMLP(STDiT):
+    def __init__(self, input_size=16, patch_size=2, in_channels=32, hidden_size=1152, depth=28, num_heads=16, mlp_ratio=4.0, max_num_frames=6, window_size=[6, 4], 
+                 dropout=0.1, ctx_noise_aug_ratio=0.1,ctx_noise_aug_prob=0.5, drop_ctx_rate=0.2, frequency_range=(2, 15), 
+                 norm_layer='layer_norm', mlp_block='mlp', qk_norm=True, **kwargs):
+        
+        if isinstance(norm_layer, str):
+            norm_layer = get_norm_layer(norm_layer)
+            
+        super().__init__(input_size=input_size, patch_size=patch_size, in_channels=in_channels, hidden_size=hidden_size, depth=depth, 
+                         num_heads=num_heads, mlp_ratio=mlp_ratio, max_num_frames=max_num_frames, dropout=dropout, 
+                         ctx_noise_aug_ratio=ctx_noise_aug_ratio, ctx_noise_aug_prob=ctx_noise_aug_prob, drop_ctx_rate=drop_ctx_rate, 
+                         norm_layer=norm_layer, mlp_block=mlp_block, **kwargs)
+        
+        for block_idx, block in enumerate(self.blocks):
+            block.space_attn = SwinAttention(
+                hidden_size,
+                input_resolution=input_size,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift_size=(0,0) if (block_idx % 2 == 0) else [ws//2 for ws in window_size],
+                proj_drop=dropout,
+                attn_drop=dropout,
+                qk_norm=qk_norm,
+                norm_layer=norm_layer,
+            )
+
+
+
+class STDiTDFSpatialL2(STDiTDF):
+    """
+    STDiTDF with per-token spatial conditioning from L2 (semantic) latents.
+
+    The L2 latent (same spatial H×W as the rec latent, sem_in_channels channels) is
+    patchified into tokens, combined with a learned frame-position embedding, and
+    produces per-token (shift, scale, gate) modulation injected at every block via
+    STDiTBlockWithSpatialL2.
+
+    Conditioning scheme (num_frames = 1 context + num_pred_frames generated):
+        Frame 0 (input):     z_l2_start, position = 0.0
+        Frame k (generated): z_l2_end,   position = k / num_pred_frames
+
+    All new parameters (l2_patchify, position_proj, per-block spatial_to_mod) are
+    zero-initialized so that a pre-trained STDiTDF backbone is unchanged at init.
+
+    Args:
+        sem_in_channels: Number of channels in the L2 (sem) latent.
+        All other args are passed through to STDiTDF.
+    """
+
+    def __init__(
+        self,
+        sem_in_channels,
+        patch_size=2,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        max_num_frames=6,
+        dropout=0.1,
+        ctx_noise_aug_ratio=0.1,
+        ctx_noise_aug_prob=0.5,
+        drop_ctx_rate=0.2,
+        frequency_range=(2, 15),
+        causal_time_attn=False,
+        modulate_time_attn=False,
+        norm_layer=nn.LayerNorm,
+        mlp_block='mlp',
+        **kwargs,
+    ):
+        if isinstance(norm_layer, str):
+            norm_layer = get_norm_layer(norm_layer)
+
+        super().__init__(
+            patch_size=patch_size,
+            hidden_size=hidden_size,
+            depth=depth,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            max_num_frames=max_num_frames,
+            dropout=dropout,
+            ctx_noise_aug_ratio=ctx_noise_aug_ratio,
+            ctx_noise_aug_prob=ctx_noise_aug_prob,
+            drop_ctx_rate=drop_ctx_rate,
+            frequency_range=frequency_range,
+            causal_time_attn=causal_time_attn,
+            modulate_time_attn=modulate_time_attn,
+            norm_layer=norm_layer,
+            mlp_block=mlp_block,
+            **kwargs,
+        )
+
+        # Replace blocks with spatially-conditioned variants
+        self.blocks = nn.ModuleList([
+            STDiTBlockWithSpatialL2(
+                hidden_size, num_heads,
+                mlp_ratio=mlp_ratio,
+                dropout_rate=dropout,
+                causal_time_attn=causal_time_attn,
+                modulate_time_attn=modulate_time_attn,
+                norm_layer=norm_layer,
+                mlp_block=mlp_block,
+            )
+            for _ in range(depth)
+        ])
+
+        # Patchify L2 latent: patch_size^2 * sem_in_channels → hidden_size
+        # Xavier init so the patchified tokens carry real signal from step 1.
+        self.l2_patchify = nn.Linear(patch_size ** 2 * sem_in_channels, hidden_size, bias=True)
+        nn.init.xavier_uniform_(self.l2_patchify.weight)
+        nn.init.zeros_(self.l2_patchify.bias)
+
+        # Frame-position embedding: scalar in [0,1] → hidden_size
+        # Xavier init so position differences produce non-trivial embeddings from step 1.
+        self.position_proj = nn.Sequential(
+            nn.Linear(1, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+        for m in self.position_proj.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def _patchify_l2(self, z_l2):
+        """
+        Patchify L2 latent frames.
+
+        Args:
+            z_l2: (B, F, C_sem, H, W)
+        Returns:
+            (B, F, N, hidden_size)  where N = (H/patch_size) * (W/patch_size)
+        """
+        b, f, c, h, w = z_l2.shape
+        p = self.patch_size
+        z = rearrange(z_l2, 'b f c (h p1) (w p2) -> (b f) (h w) (p1 p2 c)', p1=p, p2=p)
+        z = self.l2_patchify(z)                          # (B*F, N, hidden_size)
+        return rearrange(z, '(b f) n d -> b f n d', b=b, f=f)
+
+    def _build_spatial_cond(self, z_l2_start, z_l2_end, num_frames):
+        """
+        Build per-frame L2 spatial tokens with position embedding.
+
+        Args:
+            z_l2_start: (B, C_sem, H, W) — sem latent for frame 0 (context/input)
+            z_l2_end:   (B, C_sem, H, W) — sem latent for the last frame (target end)
+            num_frames: total frames in x  (= 1 context + num_pred_frames generated)
+        Returns:
+            (B, num_frames, N, hidden_size)
+        """
+        device = z_l2_start.device
+        dtype  = z_l2_start.dtype
+        num_pred = num_frames - 1  # generated frames indexed 1..num_frames-1
+
+        z_start_tok = self._patchify_l2(z_l2_start.unsqueeze(1))  # (B, 1, N, D)
+        z_end_tok   = self._patchify_l2(z_l2_end.unsqueeze(1))    # (B, 1, N, D)
+
+        frames_cond = []
+        for i in range(num_frames):
+            if i == 0:
+                pos_val = 0.0
+                z_tok = z_start_tok
+            else:
+                pos_val = i / num_pred if num_pred > 0 else 1.0
+                z_tok = z_end_tok
+
+            pos = torch.tensor([[pos_val]], device=device, dtype=dtype)
+            pos_emb = self.position_proj(pos)                    # (1, D)
+            # Broadcast over batch and spatial tokens
+            z_frame = z_tok + pos_emb.unsqueeze(0)              # (B, 1, N, D)
+            frames_cond.append(z_frame)
+
+        return torch.cat(frames_cond, dim=1)                     # (B, F, N, D)
+
+    def forward(self, x, t, frame_rate, z_l2_start, z_l2_end, return_features=False):
+        """
+        Args:
+            x:          (B, F, C_rec, H, W) — rec latent (DF style: context + targets)
+            t:          (B, F) — per-frame diffusion timesteps
+            frame_rate: (B,)
+            z_l2_start: (B, C_sem, H, W) — clean L2 latent for frame 0
+            z_l2_end:   (B, C_sem, H, W) — clean L2 latent for last frame
+        """
+        b, f = x.shape[:2]
+
+        t_flat = rearrange(t, 'b f -> (b f)')
+        c = self.get_condition_embeddings(t_flat)            # (B*F, D)
+        c = rearrange(c, '(b f) d -> b f d', b=b, f=f)     # (B, F, D)
+
+        x = self.preprocess_inputs(x, frame_rate)            # (B, F, N, D)
+
+        z_spatial = self._build_spatial_cond(z_l2_start, z_l2_end, f)  # (B, F, N, D)
+
+        features = []
+        for block in self.blocks:
+            x = block(x, c, z_spatial)
+            features.append(x) if return_features else None
+
+        out = self.final_layer(x, c)
+        out = self.postprocess_outputs(out)
+        if return_features:
+            return out, features
+        return out
+
+
+class STDiTDFSpatialL2_CtxAndLast(STDiTDFSpatialL2):
+    """
+    Two-point L2 conditioning: only the last context frame and the last
+    predicted frame receive a non-zero L2 conditioning token.
+
+    Conditioning scheme (K = num_context_frames, F = num_frames):
+        Frame K-1 (last context frame):   z_l2_start, position = (K-1)/(F-1)
+        Frame F-1 (last predicted frame): z_l2_end,   position = 1.0
+        All other frames:                 zero token,  position = i/(F-1)
+
+    The zero-token frames still receive a position embedding so the network can
+    distinguish their temporal location, but carry no semantic conditioning.
+
+    Args:
+        num_context_frames: number of context frames in the L1 sequence (K).
+            Must match the NUM_CONTEXT_FRAMES of the model class.
+        All other args forwarded to STDiTDFSpatialL2.
+    """
+
+    def __init__(self, num_context_frames, **kwargs):
+        super().__init__(**kwargs)
+        self.num_context_frames = int(num_context_frames)
+
+    def _build_spatial_cond(self, z_l2_start, z_l2_end, num_frames):
+        device = z_l2_start.device
+        dtype  = z_l2_start.dtype
+
+        z_start_tok = self._patchify_l2(z_l2_start.unsqueeze(1))  # (B, 1, N, D)
+        z_end_tok   = self._patchify_l2(z_l2_end.unsqueeze(1))    # (B, 1, N, D)
+        z_zero_tok  = torch.zeros_like(z_start_tok)                # (B, 1, N, D)
+
+        ctx_idx = self.num_context_frames - 1
+
+        frames_cond = []
+        for i in range(num_frames):
+            pos_val = i / (num_frames - 1) if num_frames > 1 else 0.0
+            pos     = torch.tensor([[pos_val]], device=device, dtype=dtype)
+            pos_emb = self.position_proj(pos)                      # (1, D)
+
+            if i == ctx_idx:
+                z_tok = z_start_tok
+            elif i == num_frames - 1:
+                z_tok = z_end_tok
+            else:
+                z_tok = z_zero_tok
+
+            frames_cond.append(z_tok + pos_emb.unsqueeze(0))      # (B, 1, N, D)
+
+        return torch.cat(frames_cond, dim=1)                       # (B, F, N, D)

orbis2/networks/DiT/dit_v2.py

@@ -0,0 +1,740 @@
+import inspect
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from omegaconf import ListConfig
+from timm.models.vision_transformer import PatchEmbed
+
+# Returns True when executing inside a torch.compile region (PyTorch >= 2.1).
+# Guarding side-effect dict writes with this prevents graph breaks during compiled inference.
+_is_compiling = getattr(torch.compiler, 'is_compiling', lambda: False)
+
+from modules.dit import (
+    FinalLayer,
+    FrequencyEncoder,
+    TimestepEmbedder,
+    get_2d_sincos_pos_embed,
+    get_norm_layer,
+)
+from util import get_obj_from_str, instantiate_from_config
+
+
+def _signature_accepts_kwargs(cls):
+    return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in inspect.signature(cls.__init__).parameters.values())
+
+
+def _accepted_init_param_names(cls):
+    names = set()
+    for base_cls in inspect.getmro(cls):
+        if base_cls is object:
+            continue
+        signature = inspect.signature(base_cls.__init__)
+        for name, param in signature.parameters.items():
+            if name == "self" or param.kind == inspect.Parameter.VAR_KEYWORD:
+                continue
+            names.add(name)
+    return names
+
+
+def _embed_timesteps_with(t, t_embedder):
+    if t.ndim == 1:
+        return t_embedder(t)
+    if t.ndim == 2:
+        b, f = t.shape
+        c = t_embedder(rearrange(t, "b f -> (b f)"))
+        return rearrange(c, "(b f) d -> b f d", b=b, f=f)
+    raise ValueError(f"Unsupported timestep shape: {t.shape}")
+
+
+class GlobalAdditiveConditioner(nn.Module):
+    """
+    Sum optional timestep embeddings and extra global condition embeddings.
+
+    The DiT module continues to own `t_embedder` for checkpoint compatibility. Extra
+    condition embedders are instantiated here and consume keyworded condition tensors.
+    """
+
+    def __init__(self, hidden_size, condition_configs=None):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.condition_embedders = nn.ModuleDict()
+
+        condition_configs = condition_configs or {}
+        for condition_name, embedder_config in condition_configs.items():
+            self.condition_embedders[condition_name] = instantiate_from_config(embedder_config)
+
+    def _reduce_global_condition_embedding(self, condition_name, embedding):
+        if embedding.ndim == 2:
+            return embedding
+        if embedding.ndim == 3 and embedding.shape[1] == 1:
+            return embedding[:, 0]
+        raise ValueError(
+            f"GlobalAdditiveConditioner only supports global condition embeddings with shape [B, H] "
+            f"or [B, 1, H]. Condition `{condition_name}` produced {tuple(embedding.shape)}."
+        )
+
+    def _combine_condition_embedding(self, conditioning, condition_embedding):
+        if torch.is_tensor(conditioning) and conditioning.ndim == 3 and condition_embedding.ndim == 2:
+            raise ValueError(f"{type(self).__name__} got per-frame conditioning {tuple(conditioning.shape)} but a global embedding {tuple(condition_embedding.shape)}; use PerFrameSequenceGlobalAdditiveConditioner for per-frame timestep conditioning.")
+        return conditioning + condition_embedding
+
+    def forward(self, *, t=None, t_embedder=None, **condition_kwargs):
+        assert (t is None) == (t_embedder is None)
+        conditioning = 0 if t is None else _embed_timesteps_with(t, t_embedder)
+        if not _is_compiling():
+            self._last_condition_embeddings = {}
+        for condition_name, embedder in self.condition_embedders.items():
+            condition_embedding = embedder(condition_kwargs[condition_name])
+            condition_embedding = self._reduce_global_condition_embedding(condition_name, condition_embedding)
+            if not _is_compiling():
+                self._last_condition_embeddings[condition_name] = condition_embedding.detach()
+            conditioning = self._combine_condition_embedding(conditioning, condition_embedding)
+        return conditioning
+
+
+class SequenceGlobalAdditiveConditioner(GlobalAdditiveConditioner):
+    """
+    Create a single global conditioning vector from sequence-structured conditions.
+
+    Each condition embedder may return embeddings with shape [B, K, H]. 
+    Optionally, a positional embedding can be added to each of the K embeddings before aggregation.
+    This class aggregates across K before adding the result to the timestep embedding.
+    """
+
+    def __init__(self, hidden_size, condition_configs=None, aggregation="sum", add_positional_embedding=False, sequence_length=None):
+        super().__init__(hidden_size=hidden_size, condition_configs=condition_configs)
+        if aggregation not in {"sum", "mean"}:
+            raise ValueError(f"Unsupported aggregation `{aggregation}`. Expected one of: sum, mean.")
+        self.aggregation = aggregation
+        self.add_positional_embedding = add_positional_embedding
+        self.sequence_length = sequence_length
+        if add_positional_embedding:
+            if sequence_length is None: raise ValueError("sequence_length must be specified if add_positional_embedding is True.")
+            # One learnable positional embedding table per condition_name.
+            # Assumes self.condition_configs exists and is iterable/dict-like from parent.
+            self.positional_embeddings = nn.ParameterDict()
+            for condition_name in self.condition_embedders.keys():
+                pe = nn.Parameter(torch.zeros(sequence_length, hidden_size))
+                nn.init.normal_(pe, mean=0.0, std=0.02)
+                self.positional_embeddings[condition_name] = pe
+        else:
+            self.positional_embeddings = None
+
+    def _reduce_global_condition_embedding(self, condition_name, embedding):        
+        if embedding.ndim == 2:
+            return embedding
+        if embedding.ndim != 3:
+            raise ValueError(
+                f"SequenceGlobalAdditiveConditioner expects condition `{condition_name}` to produce "
+                f"shape [B, H] or [B, K, H], got {tuple(embedding.shape)}."
+            )
+        
+        # Enforce positional embeddings for true sequences (K > 1)
+        k = embedding.shape[1]
+        if k > 1 and not self.add_positional_embedding:
+            raise ValueError(
+                f"Condition `{condition_name}` produced sequence embedding with K={k}, but "
+                f"`add_positional_embedding=False`. Positional embedding must be enabled when K > 1."
+            )
+
+        if self.add_positional_embedding:
+            embedding = self._add_positional_embeddings(condition_name, embedding)
+
+        if self.aggregation == "sum":
+            return embedding.sum(dim=1)
+        return embedding.mean(dim=1)
+    
+    def _add_positional_embeddings(self, condition_name, embedding):
+        """
+        embedding: [B, K, H]
+        returns:   [B, K, H] with learnable positional embeddings added
+        """
+        if embedding.ndim != 3:
+            raise ValueError(
+                f"_add_positional_embeddings expects [B, K, H] for `{condition_name}`, "
+                f"got {tuple(embedding.shape)}."
+            )
+
+        b, k, h = embedding.shape
+
+        if h != self.hidden_size:
+            raise ValueError(
+                f"Hidden size mismatch for condition `{condition_name}`: "
+                f"embedding has H={h}, expected hidden_size={self.hidden_size}."
+            )
+
+        if self.positional_embeddings is None:
+            raise RuntimeError(
+                "Positional embedding is not initialized. "
+                "Set add_positional_embedding=True in constructor."
+            )
+
+        if k > self.sequence_length:
+            raise ValueError(
+                f"Condition `{condition_name}` has sequence length K={k}, "
+                f"but configured sequence_length={self.sequence_length}. "
+                f"Increase sequence_length or truncate inputs."
+            )
+
+        # [K, H] -> [1, K, H], broadcast over batch
+        pos = self.positional_embeddings[condition_name][:k].unsqueeze(0)
+        pos = pos.to(device=embedding.device, dtype=embedding.dtype)
+        return embedding + pos
+
+
+class PerFrameSequenceGlobalAdditiveConditioner(SequenceGlobalAdditiveConditioner):
+    """
+    Same as SequenceGlobalAdditiveConditioner, but broadcasts global [B, H] condition
+    embeddings across frames when timestep conditioning is per-frame (conditioning is
+    [B, F, H] instead of [B, H]), instead of colliding F against B.
+    """
+
+    def _combine_condition_embedding(self, conditioning, condition_embedding):
+        if torch.is_tensor(conditioning) and conditioning.ndim == 3:
+            condition_embedding = condition_embedding.unsqueeze(1)  # [B, H] -> [B, 1, H]
+        return conditioning + condition_embedding
+
+
+class DiT(nn.Module):
+    """
+    Generic DiT backbone with configurable block construction and execution layout.
+
+    The model does not distinguish target and context inputs. Forward accepts only a
+    generic sequence of frames `x` with shape [B, F, C, H, W].
+    """
+
+    def __init__(
+        self,
+        input_size=16,
+        patch_size=2,
+        in_channels=32,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        max_num_frames=6,
+        dropout=0.0,
+        frequency_range=(2, 15),
+        learn_sigma=False,
+        norm_layer="layer_norm",
+        mlp_block="mlp",
+        block_config=None,
+        conditioner_config=None,
+        frame_embedding_alignment="suffix",
+        ctx_noise_aug_ratio=0.0,
+        ctx_noise_aug_prob=0.0,
+        drop_ctx_rate=0.0,
+        log_adaln_mean_abs=False,
+        log_adaln_grad=False,
+        adaln_gate_init_std=0.0,
+        train_steering_adaln_only=False,
+    ):
+        super().__init__()
+        self._validate_legacy_context_args(
+            ctx_noise_aug_ratio=ctx_noise_aug_ratio,
+            ctx_noise_aug_prob=ctx_noise_aug_prob,
+            drop_ctx_rate=drop_ctx_rate,
+        )
+
+        if isinstance(norm_layer, str):
+            norm_layer = get_norm_layer(norm_layer)
+
+        self.input_size = input_size if isinstance(input_size, (list, tuple, ListConfig)) else [input_size, input_size]
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+        self.depth = depth
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.dropout = dropout
+        self.max_num_frames = max_num_frames
+        self.norm_layer = norm_layer
+        self.mlp_block = mlp_block
+        self.log_adaln_mean_abs = bool(log_adaln_mean_abs)
+        self.log_adaln_grad = bool(log_adaln_grad)
+        self.adaln_gate_init_std = float(adaln_gate_init_std)
+        self._last_adaln_mean_abs = {}
+        self.frame_embedding_alignment = self._normalize_frame_embedding_alignment(
+            frame_embedding_alignment
+        )
+
+        self.out_channels = in_channels * 2 if learn_sigma else in_channels
+        self.patch_size = patch_size
+
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
+        self.num_patches = self.x_embedder.num_patches
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.conditioner = self._build_conditioner(conditioner_config)
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, hidden_size), requires_grad=False)
+        self.frame_emb = nn.init.trunc_normal_(nn.Parameter(torch.zeros(1, self.max_num_frames, 1, hidden_size)), 0.0, 0.02)
+        self.frame_rate_encoder = FrequencyEncoder(hidden_size, freq_min=frequency_range[0], freq_max=frequency_range[1])
+
+        if block_config is None:
+            raise ValueError("DiT v2 expects block_config in instantiate_from_config format.")
+        if "target" not in block_config:
+            raise KeyError("block_config must define `target`.")
+
+        self.block_config = block_config
+        self.block_cls = get_obj_from_str(block_config["target"])
+        self.blocks = self._build_blocks()
+        self._configure_block_logging()
+
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels, norm_layer=norm_layer)
+        self.initialize_weights()
+
+        if train_steering_adaln_only:
+            self._freeze_non_steering_adaln_parameters()
+
+    def _freeze_non_steering_adaln_parameters(self):
+        self.requires_grad_(False)
+        self.t_embedder.requires_grad_(True)
+        cond_embedders = getattr(self.conditioner, "condition_embedders", {})
+        if "steering" in cond_embedders:
+            cond_embedders["steering"].requires_grad_(True)
+        for block in self.blocks:
+            block.adaLN_modulation.requires_grad_(True)
+            if hasattr(block, "adaLN_time_attn_modulation"):
+                block.adaLN_time_attn_modulation.requires_grad_(True)
+        self.final_layer.adaLN_modulation.requires_grad_(True)
+
+    def _validate_legacy_context_args(self, *, ctx_noise_aug_ratio, ctx_noise_aug_prob, drop_ctx_rate):
+        unsupported = {
+            "ctx_noise_aug_ratio": ctx_noise_aug_ratio,
+            "ctx_noise_aug_prob": ctx_noise_aug_prob,
+            "drop_ctx_rate": drop_ctx_rate,
+        }
+        non_zero = {name: value for name, value in unsupported.items() if value not in (None, 0, 0.0)}
+        if non_zero:
+            raise ValueError(
+                "DiT v2 does not implement context dropping or context noise augmentation. "
+                f"Received: {non_zero}"
+            )
+
+    def _normalize_frame_embedding_alignment(self, frame_embedding_alignment):
+        alignment = str(frame_embedding_alignment).strip().lower()
+        if alignment not in {"suffix", "prefix"}:
+            raise ValueError(
+                "frame_embedding_alignment must be one of {'suffix', 'prefix'}, "
+                f"got {frame_embedding_alignment!r}."
+            )
+        return alignment
+
+    def _get_frame_embeddings(self, num_frames, device):
+        if self.frame_embedding_alignment == "prefix":
+            return self.frame_emb[:, :num_frames].to(device)
+        return self.frame_emb[:, self.max_num_frames - num_frames :].to(device)
+
+    def _get_common_block_kwargs(self, layer_idx):
+        return {
+            "hidden_size": self.hidden_size,
+            "num_heads": self.num_heads,
+            "mlp_ratio": self.mlp_ratio,
+            "dropout_rate": self.dropout,
+            "norm_layer": self.norm_layer,
+            "mlp_block": self.mlp_block,
+            "input_shape": self.input_size,
+            "layer_idx": layer_idx,
+        }
+
+    def _build_single_block(self, layer_idx):
+        accepts_kwargs = _signature_accepts_kwargs(self.block_cls)
+        common_kwargs = self._get_common_block_kwargs(layer_idx)
+        user_kwargs = dict(self.block_config.get("params", {}))
+
+        valid_names = _accepted_init_param_names(self.block_cls)
+        if not accepts_kwargs:
+            unknown_user_kwargs = sorted(set(user_kwargs) - valid_names)
+            if unknown_user_kwargs:
+                raise ValueError(
+                    f"Unsupported block parameters for {self.block_cls.__name__}: {unknown_user_kwargs}"
+                )
+
+        init_kwargs = {}
+        for name, value in common_kwargs.items():
+            if name in valid_names:
+                init_kwargs[name] = value
+        init_kwargs.update(user_kwargs)
+        return instantiate_from_config(
+            {
+                "target": self.block_config["target"],
+                "params": init_kwargs,
+            }
+        )
+
+    def _build_blocks(self):
+        return nn.ModuleList([self._build_single_block(layer_idx) for layer_idx in range(self.depth)])
+
+    def _configure_block_logging(self):
+        for block in self.blocks:
+            if hasattr(block, "log_adaln_mean_abs"):
+                block.log_adaln_mean_abs = self.log_adaln_mean_abs
+
+    def _build_conditioner(self, conditioner_config):
+        if conditioner_config is None:
+            return None
+
+        params = dict(conditioner_config.get("params", {}))
+        params.setdefault("hidden_size", self.hidden_size)
+        return instantiate_from_config(
+            {
+                "target": conditioner_config["target"],
+                "params": params,
+            }
+        )
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        pos_embed = get_2d_sincos_pos_embed(
+            self.pos_embed.shape[-1],
+            [self.input_size[0] // self.patch_size, self.input_size[1] // self.patch_size],
+            cls_token=False,
+            extra_tokens=0,
+        )
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        for block in self.blocks:
+            if hasattr(block, "initialize_adaln_weights"):
+                block.initialize_adaln_weights(self.adaln_gate_init_std)
+
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x):
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = self.x_embedder.grid_size[0]
+        w = self.x_embedder.grid_size[1]
+
+        x = x.reshape(shape=(x.shape[0], x.shape[1], h, w, p, p, c))
+        x = torch.einsum("bfhwpqc->bfchpwq", x)
+        imgs = x.reshape(shape=(x.shape[0], x.shape[1], c, h * p, w * p))
+        return imgs
+
+    def postprocess_outputs(self, out):
+        return self.unpatchify(out)
+
+    def get_condition_embeddings(self, t, **_kwargs):
+        if self.conditioner is None:
+            return _embed_timesteps_with(t, self.t_embedder)
+        return self.conditioner(t=t, t_embedder=self.t_embedder, **_kwargs)
+
+    def _get_frame_rate_embeddings(self, frame_rate, batch_size, num_frames, device):
+        if frame_rate is None:
+            return torch.zeros(batch_size, 1, 1, self.hidden_size, device=device)
+
+        frame_rate_embeddings = self.frame_rate_encoder.encode(frame_rate).to(device)
+        if frame_rate_embeddings.ndim != 2 or frame_rate_embeddings.shape[0] != batch_size:
+            raise ValueError(
+                "frame_rate must have shape [B]. "
+                f"Received embedding source with shape {tuple(frame_rate.shape)}"
+            )
+        return frame_rate_embeddings.unsqueeze(1).unsqueeze(1)
+
+    def embed_inputs(self, x, frame_rate=None):
+        if x.ndim != 5:
+            raise ValueError(f"Expected x to have shape [B, F, C, H, W], got {tuple(x.shape)}")
+
+        b, f = x.shape[:2]
+        if f > self.max_num_frames:
+            raise ValueError(f"Received {f} frames, but max_num_frames={self.max_num_frames}")
+
+        frame_embeddings = self._get_frame_rate_embeddings(frame_rate, batch_size=b, num_frames=f, device=x.device)
+        x = rearrange(x, "b f c h w -> (b f) c h w")
+        x = self.x_embedder(x) + self.pos_embed.to(x.device)
+        x = rearrange(x, "(b f) hw c -> b f hw c", b=b, f=f)
+        x = x + self._get_frame_embeddings(f, x.device) + frame_embeddings
+        return x
+
+    def run_blocks(self, x, c, return_features=False):
+        features = []
+        adaln_stats = {}
+        for block_idx, block in enumerate(self.blocks):
+            x = block(x, c)
+            if self.log_adaln_mean_abs:
+                block_stats = getattr(block, "_last_adaln_mean_abs", None)
+                if block_stats:
+                    adaln_stats.update(
+                        {
+                            f"block_{block_idx:02d}/{name}": value
+                            for name, value in block_stats.items()
+                        }
+                    )
+            if return_features:
+                features.append(x)
+        if self.log_adaln_mean_abs:
+            self._last_adaln_mean_abs = adaln_stats
+        return x, features
+
+    def get_last_adaln_mean_abs(self):
+        return self._last_adaln_mean_abs
+
+    def get_last_condition_embeddings(self):
+        if self.conditioner is None:
+            return {}
+        return dict(getattr(self.conditioner, '_last_condition_embeddings', {}))
+
+    def _iter_conditioners(self):
+        if self.conditioner is not None:
+            yield ("", self.conditioner)
+
+    def get_adaln_grad_stats(self):
+        if not self.log_adaln_grad:
+            return {}
+        stats = {}
+        for block_idx, block in enumerate(self.blocks):
+            for mod_name in ("adaLN_modulation", "adaLN_time_attn_modulation", "adaLN_registers_modulation"):
+                mod = getattr(block, mod_name, None)
+                if mod is None:
+                    continue
+                linear = mod[-1]
+                for param_name in ("weight", "bias"):
+                    p = getattr(linear, param_name, None)
+                    if p is None or p.grad is None:
+                        continue
+                    stats[f"block_{block_idx:02d}/{mod_name}/{param_name}"] = p.grad.detach()
+        for cond_prefix, conditioner in self._iter_conditioners():
+            for cond_name, embedder in conditioner.condition_embedders.items():
+                for pname, p in embedder.named_parameters():
+                    if p.grad is not None:
+                        stats[f"{cond_prefix}cond_emb/{cond_name}/{pname}"] = p.grad.detach()
+        return stats
+
+    def forward(
+        self,
+        x,
+        t,
+        frame_rate=None,
+        return_features=False,
+        **condition_kwargs,
+    ):
+        c = self.get_condition_embeddings(t, **condition_kwargs)
+        x = self.embed_inputs(x, frame_rate=frame_rate)
+        x, features = self.run_blocks(x, c, return_features=return_features)
+        out = self.final_layer(x, c)
+        out = self.postprocess_outputs(out)
+        if return_features:
+            return out, features
+        return out
+
+
+class DiTWithRegisters(DiT):
+    def __init__(self, registers_conditioner_config=None, train_steering_adaln_only=False, **kwargs):
+        super().__init__(train_steering_adaln_only=False, **kwargs)
+        self.registers_conditioner = self._build_conditioner(registers_conditioner_config)
+        if self.registers_conditioner is not None:
+            self.registers_conditioner.apply(self._initialize_linear_module)
+        if train_steering_adaln_only:
+            self._freeze_non_steering_adaln_parameters()
+
+    @staticmethod
+    def _initialize_linear_module(module):
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+
+    def _freeze_non_steering_adaln_parameters(self):
+        self.requires_grad_(False)
+        if self.registers_conditioner is not None:
+            self.registers_conditioner.requires_grad_(True)
+        for block in self.blocks:
+            if hasattr(block, "adaLN_registers_modulation"):
+                block.adaLN_registers_modulation.requires_grad_(True)
+
+    def get_registers_condition_embeddings(self, **condition_kwargs):
+        if self.registers_conditioner is None:
+            return None
+        return self.registers_conditioner(**condition_kwargs)
+
+    def get_last_condition_embeddings(self):
+        result = super().get_last_condition_embeddings()
+        if self.registers_conditioner is not None:
+            reg_embs = getattr(self.registers_conditioner, '_last_condition_embeddings', {})
+            result.update({f"registers/{k}": v for k, v in reg_embs.items()})
+        return result
+
+    def _iter_conditioners(self):
+        yield from super()._iter_conditioners()
+        if self.registers_conditioner is not None:
+            yield ("registers/", self.registers_conditioner)
+
+    def run_blocks(self, x, c, c_registers=None, return_features=False):
+        features = []
+        adaln_stats = {}
+        for block_idx, block in enumerate(self.blocks):
+            if c_registers is None:
+                x = block(x, c=c)
+            else:
+                x = block(x, c=c, c_registers=c_registers)
+            if self.log_adaln_mean_abs:
+                block_stats = getattr(block, "_last_adaln_mean_abs", None)
+                if block_stats:
+                    adaln_stats.update(
+                        {
+                            f"block_{block_idx:02d}/{name}": value
+                            for name, value in block_stats.items()
+                        }
+                    )
+            if return_features:
+                features.append(x)
+        if self.log_adaln_mean_abs:
+            self._last_adaln_mean_abs = adaln_stats
+        return x, features
+
+    def forward(
+        self,
+        x,
+        t,
+        frame_rate=None,
+        return_features=False,
+        **condition_kwargs,
+    ):
+        c = self.get_condition_embeddings(t, **condition_kwargs)
+        c_registers = self.get_registers_condition_embeddings(**condition_kwargs)
+        x = self.embed_inputs(x, frame_rate=frame_rate)
+        x, features = self.run_blocks(x, c, c_registers=c_registers, return_features=return_features)
+        out = self.final_layer(x, c)
+        out = self.postprocess_outputs(out)
+        if return_features:
+            return out, features
+        return out
+
+
+class SpatialL2CtxAndLastDiT(DiT):
+    """
+    v2 DiT equivalent of the v1 STDiTDFSpatialL2_CtxAndLast network.
+
+    Spatial L2 endpoint latents are patchified and injected into blocks that
+    accept `block(x, c, z_spatial)`. Only the last context frame and last frame
+    receive non-zero L2 tokens; all frames receive a scalar position embedding.
+    """
+
+    def __init__(self, sem_in_channels, num_context_frames, frame_embedding_alignment, **kwargs):
+        super().__init__(frame_embedding_alignment=frame_embedding_alignment, **kwargs)
+        self.sem_in_channels = int(sem_in_channels)
+        self.num_context_frames = int(num_context_frames)
+
+        self.l2_patchify = nn.Linear(
+            self.patch_size ** 2 * self.sem_in_channels,
+            self.hidden_size,
+            bias=True,
+        )
+        nn.init.xavier_uniform_(self.l2_patchify.weight)
+        nn.init.zeros_(self.l2_patchify.bias)
+
+        self.position_proj = nn.Sequential(
+            nn.Linear(1, self.hidden_size),
+            nn.SiLU(),
+            nn.Linear(self.hidden_size, self.hidden_size),
+        )
+        for module in self.position_proj.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def _patchify_l2(self, z_l2):
+        if z_l2.ndim != 5:
+            raise ValueError(f"Expected z_l2 to have shape [B, F, C, H, W], got {tuple(z_l2.shape)}")
+        b, f, c, h, w = z_l2.shape
+        if c != self.sem_in_channels:
+            raise ValueError(
+                f"Expected L2 latents with {self.sem_in_channels} channels, got {c}."
+            )
+        if h % self.patch_size != 0 or w % self.patch_size != 0:
+            raise ValueError(
+                f"L2 latent spatial size {(h, w)} must be divisible by patch_size={self.patch_size}."
+            )
+
+        z = rearrange(
+            z_l2,
+            "b f c (h p1) (w p2) -> (b f) (h w) (p1 p2 c)",
+            p1=self.patch_size,
+            p2=self.patch_size,
+        )
+        z = self.l2_patchify(z)
+        return rearrange(z, "(b f) n d -> b f n d", b=b, f=f)
+
+    def _build_spatial_cond(self, z_l2_start, z_l2_end, num_frames):
+        if z_l2_start is None or z_l2_end is None:
+            raise ValueError("SpatialL2CtxAndLastDiT requires z_l2_start and z_l2_end.")
+        if z_l2_start.shape != z_l2_end.shape:
+            raise ValueError(
+                "z_l2_start and z_l2_end must have matching shapes, got "
+                f"{tuple(z_l2_start.shape)} and {tuple(z_l2_end.shape)}."
+            )
+        if self.num_context_frames < 1:
+            raise ValueError("num_context_frames must be at least 1.")
+        if self.num_context_frames > num_frames:
+            raise ValueError(
+                f"num_context_frames={self.num_context_frames} exceeds num_frames={num_frames}."
+            )
+
+        device = z_l2_start.device
+        dtype = z_l2_start.dtype
+        z_start_tok = self._patchify_l2(z_l2_start.unsqueeze(1))
+        z_end_tok = self._patchify_l2(z_l2_end.unsqueeze(1))
+        z_zero_tok = torch.zeros_like(z_start_tok)
+        ctx_idx = self.num_context_frames - 1
+
+        frames_cond = []
+        for frame_idx in range(num_frames):
+            pos_val = frame_idx / (num_frames - 1) if num_frames > 1 else 0.0
+            pos = torch.tensor([[pos_val]], device=device, dtype=dtype)
+            pos_emb = self.position_proj(pos)
+
+            if frame_idx == ctx_idx:
+                z_tok = z_start_tok
+            elif frame_idx == num_frames - 1:
+                z_tok = z_end_tok
+            else:
+                z_tok = z_zero_tok
+            frames_cond.append(z_tok + pos_emb.unsqueeze(0))
+
+        return torch.cat(frames_cond, dim=1)
+
+    def run_spatial_l2_blocks(self, x, c, z_spatial, return_features=False):
+        features = []
+        for block in self.blocks:
+            x = block(x, c, z_spatial)
+            if return_features:
+                features.append(x)
+        return x, features
+
+    def forward(
+        self,
+        x,
+        t,
+        frame_rate=None,
+        z_l2_start=None,
+        z_l2_end=None,
+        return_features=False,
+        **condition_kwargs,
+    ):
+        if z_l2_start is None or z_l2_end is None:
+            raise ValueError("SpatialL2CtxAndLastDiT requires z_l2_start and z_l2_end.")
+
+        c = self.get_condition_embeddings(t, **condition_kwargs)
+        x = self.embed_inputs(x, frame_rate=frame_rate)
+        z_spatial = self._build_spatial_cond(z_l2_start, z_l2_end, x.size(1))
+        x, features = self.run_spatial_l2_blocks(x, c, z_spatial, return_features=return_features)
+        out = self.final_layer(x, c)
+        out = self.postprocess_outputs(out)
+        if return_features:
+            return out, features
+        return out

orbis2/networks/swin/swin_free_aspect_ratio.py

@@ -0,0 +1,740 @@
+# --------------------------------------------------------
+# Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.layers.mlp import SwiGLU
+
+try:
+    import os, sys
+
+    kernel_path = os.path.abspath(os.path.join('..'))
+    sys.path.append(kernel_path)
+    from kernels.window_process.window_process import WindowProcess, WindowProcessReverse
+
+except:
+    WindowProcess = None
+    WindowProcessReverse = None
+    print("[Warning] Fused window process have not been installed. Please refer to get_started.md for installation.")
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    window_size = to_2tuple(window_size)
+    x = x.view(B, H // window_size[0], window_size[0], W // window_size[1], window_size[1], C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, *window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    window_size = to_2tuple(window_size)
+    B = int(windows.shape[0] / (H * W / window_size[0] / window_size[1]))
+    x = windows.view(B, H // window_size[0], W // window_size[1], window_size[0], window_size[1], -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, qk_norm=True, attn_drop=0., proj_drop=0., fused_attn=True, norm_layer=nn.LayerNorm):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        self.use_fused_attn = fused_attn
+        self.qk_norm = qk_norm
+        
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        
+        self.q_norm = norm_layer(head_dim, eps=1e-6) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(head_dim, eps=1e-6) if qk_norm else nn.Identity()
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def old_attn(self, q, k, v, B_, N, relative_position_bias, mask=None):
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+        attn = attn + relative_position_bias.unsqueeze(0)        
+    
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+        attn = self.attn_drop(attn)
+        x = (attn @ v)
+        return x
+    
+    def fused_attn(self, q, k, v, B_, N, relative_position_bias, mask=None):
+        attn_mask = relative_position_bias
+        if mask is not None:
+            num_win = mask.shape[0]
+            mask = mask.view(1, num_win, 1, N, N).expand(B_ // num_win, -1, self.num_heads, -1, -1)
+            attn_mask = attn_mask + mask.reshape(-1, self.num_heads, N, N)
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_drop.p if self.training else 0.,
+        )
+        return x
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+        q, k = self.q_norm(q), self.k_norm(k)
+        
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+        self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        
+        if self.use_fused_attn:  # test using scaled dot product attention
+            x = self.fused_attn(q, k, v, B_, N, relative_position_bias, mask=mask)
+        else:
+            x = self.old_attn(q, k, v, B_, N, relative_position_bias, mask=mask)
+            
+        x = x.transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+
+class SwinAttention(nn.Module):
+    """ Swin attention only, no MLP. """
+    def __init__(self, dim, input_resolution, num_heads, window_size=(7,7), shift_size=(0,0),
+                 qkv_bias=True, qk_scale=None, qk_norm=False, proj_drop=0., attn_drop=0., norm_layer=nn.LayerNorm,
+                 fused_window_process=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        
+        if min(self.input_resolution) <= min(self.window_size) or max(self.input_resolution) <= max(self.window_size):
+            raise ValueError("window size cannot be larger than input resolution")
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size[0] < self.window_size[0] and 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in [0, window_size)"
+
+        self.attn = WindowAttention(
+            dim, window_size=self.window_size, num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm,
+            attn_drop=attn_drop, proj_drop=proj_drop, fused_attn=True, norm_layer=norm_layer)
+
+        if self.shift_size[0] > 0 or self.shift_size[1] > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size[0]),
+                        slice(-self.window_size[0], -self.shift_size[0]),
+                        slice(-self.shift_size[0], None))
+            w_slices = (slice(0, -self.window_size[1]),
+                        slice(-self.window_size[1], -self.shift_size[1]),
+                        slice(-self.shift_size[1], None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = cnt
+                    cnt += 1
+
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size[0] * self.window_size[1])
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+        self.fused_window_process = fused_window_process
+
+    def forward(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size[0] > 0 or self.shift_size[1] > 0:
+            if not self.fused_window_process:
+                shifted_x = torch.roll(x, shifts=(-self.shift_size[0], -self.shift_size[1]), dims=(1, 2))
+                # partition windows
+                x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+            else:
+                x_windows = WindowProcess.apply(x, B, H, W, C, -self.shift_size[0], self.window_size[1])
+        else:
+            shifted_x = x
+            # partition windows
+            x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+
+        x_windows = x_windows.view(-1, self.window_size[0] * self.window_size[1], C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size[0], self.window_size[1], C)
+
+        # reverse cyclic shift
+        if self.shift_size[0] > 0 or self.shift_size[1] > 0:
+            if not self.fused_window_process:
+                shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+                x = torch.roll(shifted_x, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2))
+            else:
+                x = WindowProcessReverse.apply(attn_windows, B, H, W, C, self.shift_size, self.window_size)
+        else:
+            shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+            x = shifted_x
+        x = x.view(B, H * W, C)
+
+        return x
+
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+        fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, window_size=(7,7), shift_size=(0,0),
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, qk_norm=True, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 mlp_block='mlp',
+                 fused_window_process=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= min(self.window_size) or max(self.input_resolution) <= max(self.window_size):
+            raise ValueError("window size cannot be larger than input resolution")
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size[0] < self.window_size[0] and 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in [0, window_size)"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=self.window_size, num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        if mlp_block == 'mlp':
+            self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        elif mlp_block == 'swiglu':
+            self.mlp = SwiGLU(in_features=dim, hidden_features=(mlp_hidden_dim*2)//3)
+            
+        if self.shift_size[0] > 0 or self.shift_size[1] > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size[0]),
+                        slice(-self.window_size[0], -self.shift_size[0]),
+                        slice(-self.shift_size[0], None))
+            w_slices = (slice(0, -self.window_size[1]),
+                        slice(-self.window_size[1], -self.shift_size[1]),
+                        slice(-self.shift_size[1], None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = cnt
+                    cnt += 1
+
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size[0] * self.window_size[1])
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+        self.fused_window_process = fused_window_process
+
+    def forward(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size[0] > 0 or self.shift_size[1] > 0:
+            if not self.fused_window_process:
+                shifted_x = torch.roll(x, shifts=(-self.shift_size[0], -self.shift_size[1]), dims=(1, 2))
+                # partition windows
+                x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+            else:
+                x_windows = WindowProcess.apply(x, B, H, W, C, -self.shift_size[0], self.window_size[1])
+        else:
+            shifted_x = x
+            # partition windows
+            x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+
+        x_windows = x_windows.view(-1, self.window_size[0] * self.window_size[1], C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size[0], self.window_size[1], C)
+
+        # reverse cyclic shift
+        if self.shift_size[0] > 0 or self.shift_size[1] > 0:
+            if not self.fused_window_process:
+                shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+                x = torch.roll(shifted_x, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2))
+            else:
+                x = WindowProcessReverse.apply(attn_windows, B, H, W, C, self.shift_size, self.window_size)
+        else:
+            shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+            x = shifted_x
+        x = x.view(B, H * W, C)
+        x = shortcut + self.drop_path(x)
+
+        # FFN
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        nW = H * W / self.window_size / self.window_size
+        flops += nW * self.attn.flops(self.window_size * self.window_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        return flops
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+        fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 fused_window_process=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, window_size=window_size,
+                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer,
+                                 fused_window_process=fused_window_process)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+class PatchEmbed(nn.Module):
+    r""" Image to Patch Embedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+    def flops(self):
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class SwinTransformer(nn.Module):
+    r""" Swin Transformer
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+        fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
+                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                 use_checkpoint=False, fused_window_process=False, **kwargs):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=depths[i_layer],
+                               num_heads=num_heads[i_layer],
+                               window_size=window_size,
+                               mlp_ratio=self.mlp_ratio,
+                               qkv_bias=qkv_bias, qk_scale=qk_scale,
+                               drop=drop_rate, attn_drop=attn_drop_rate,
+                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                               norm_layer=norm_layer,
+                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                               use_checkpoint=use_checkpoint,
+                               fused_window_process=fused_window_process)
+            self.layers.append(layer)
+
+        self.norm = norm_layer(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = self.norm(x)  # B L C
+        x = self.avgpool(x.transpose(1, 2))  # B C 1
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+    def flops(self):
+        flops = 0
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            flops += layer.flops()
+        flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
+        flops += self.num_features * self.num_classes
+        return flops

orbis2/networks/tokenizer/ae.py

@@ -0,0 +1,852 @@
+# pytorch_diffusion + derived encoder decoder
+import math
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+from omegaconf import ListConfig
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models:
+    From Fairseq.
+    Build sinusoidal embeddings.
+    This matches the implementation in tensor2tensor, but differs slightly
+    from the description in Section 3.5 of "Attention Is All You Need".
+    """
+    assert len(timesteps.shape) == 1
+
+    half_dim = embedding_dim // 2
+    emb = math.log(10000) / (half_dim - 1)
+    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+    emb = emb.to(device=timesteps.device)
+    emb = timesteps.float()[:, None] * emb[None, :]
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+    if embedding_dim % 2 == 1:  # zero pad
+        emb = torch.nn.functional.pad(emb, (0,1,0,0))
+    return emb
+
+
+def nonlinearity(x):
+    # swish
+    return x*torch.sigmoid(x)
+
+
+class SpatialNorm(nn.Module):
+    def __init__(self, f_channels, zq_channels, norm_layer=nn.GroupNorm, freeze_norm_layer=False, add_conv=False, **norm_layer_params):
+        super().__init__()
+        self.norm_layer = norm_layer(num_channels=f_channels, **norm_layer_params)
+        if freeze_norm_layer:
+            for p in self.norm_layer.parameters: # for p in self.norm_layer.parameters():
+                p.requires_grad = False
+        self.add_conv = add_conv
+        if self.add_conv:
+            self.conv = nn.Conv2d(zq_channels, zq_channels, kernel_size=3, stride=1, padding=1)
+        self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)
+        self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)
+    def forward(self, f, zq):
+        #print(f'f shape: {f.shape}', f'zq shape: {zq.shape}')
+        f_size = f.shape[-2:]
+        zq = torch.nn.functional.interpolate(zq, size=f_size, mode="nearest")
+        if self.add_conv:
+            zq = self.conv(zq)
+        norm_f = self.norm_layer(f)
+        new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)
+        return new_f
+
+def SpatialNormalize(in_channels, zq_ch, add_conv):
+    return SpatialNorm(in_channels, zq_ch, norm_layer=nn.GroupNorm, freeze_norm_layer=False, add_conv=add_conv, num_groups=32, eps=1e-6, affine=True)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0,1,0,1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
+                 dropout, temb_channels=512):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x+h
+
+class SpatialResnetBlock(nn.Module):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
+                 dropout, temb_channels=512, zq_ch=None, add_conv=False):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = SpatialNormalize(in_channels, zq_ch, add_conv)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+        self.norm2 = SpatialNormalize(out_channels, zq_ch, add_conv)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x, temb, zq):
+        h = x
+        h = self.norm1(h, zq)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
+
+        h = self.norm2(h, zq)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x+h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = q.reshape(b,c,h*w)
+        q = q.permute(0,2,1)   # b,hw,c
+        k = k.reshape(b,c,h*w) # b,c,hw
+        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b,c,h*w)
+        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b,c,h,w)
+
+        h_ = self.proj_out(h_)
+
+        return x+h_
+    
+class SpatialAttnBlock(nn.Module):
+    def __init__(self, in_channels, zq_ch=None, add_conv=False):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = SpatialNormalize(in_channels, zq_ch, add_conv)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+
+    def forward(self, x, zq):
+        h_ = x
+        h_ = self.norm(h_, zq)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = q.reshape(b,c,h*w)
+        q = q.permute(0,2,1)   # b,hw,c
+        k = k.reshape(b,c,h*w) # b,c,hw
+        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b,c,h*w)
+        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b,c,h,w)
+
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+
+class Model(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, use_timestep=True):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = self.ch*4
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.use_timestep = use_timestep
+        if self.use_timestep:
+            # timestep embedding
+            self.temb = nn.Module()
+            self.temb.dense = nn.ModuleList([
+                torch.nn.Linear(self.ch,
+                                self.temb_ch),
+                torch.nn.Linear(self.temb_ch,
+                                self.temb_ch),
+            ])
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            skip_in = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                if i_block == self.num_res_blocks:
+                    skip_in = ch*in_ch_mult[i_level]
+                block.append(ResnetBlock(in_channels=block_in+skip_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+
+    def forward(self, x, t=None):
+        #assert x.shape[2] == x.shape[3] == self.resolution
+
+        if self.use_timestep:
+            # timestep embedding
+            assert t is not None
+            temb = get_timestep_embedding(t, self.ch)
+            temb = self.temb.dense[0](temb)
+            temb = nonlinearity(temb)
+            temb = self.temb.dense[1](temb)
+        else:
+            temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](
+                    torch.cat([h, hs.pop()], dim=1), temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size):
+    """
+    grid_size: int or (int, int) of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_size = (grid_size, grid_size) if type(grid_size) != tuple else grid_size
+    grid_h = np.arange(grid_size[0], dtype=np.float32)
+    grid_w = np.arange(grid_size[1], dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size[0], grid_size[1]])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+
+    return pos_embed
+
+def init_weights(m):
+    if isinstance(m, nn.Linear):
+        # we use xavier_uniform following official JAX ViT:
+        torch.nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.LayerNorm):
+        nn.init.constant_(m.bias, 0)
+        nn.init.constant_(m.weight, 1.0)
+    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+        with torch.no_grad():
+            w = m.weight
+            torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+class PreNorm(nn.Module):
+    def __init__(self, dim: int, fn: nn.Module) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        return self.fn(self.norm(x), **kwargs)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.Tanh(),
+            nn.Linear(hidden_dim, dim)
+        )
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim: int, heads: int = 8, dim_head: int = 64) -> None:
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim = -1)
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Linear(inner_dim, dim) if project_out else nn.Identity()
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        attn = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+        attn = self.attend(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(self, dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for idx in range(depth):
+            layer = nn.ModuleList([PreNorm(dim, Attention(dim, heads=heads, dim_head=dim_head)),
+                                   PreNorm(dim, FeedForward(dim, mlp_dim))])
+            self.layers.append(layer)
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+class EncoderVIT(nn.Module):
+    def __init__(self, image_size: Union[Tuple[int, int], int], patch_size: Union[Tuple[int, int], int],
+                 dim: int, depth: int, heads: int, mlp_dim: int, channels: int = 3, dim_head: int = 64) -> None:
+        super().__init__()
+        
+        image_height, image_width = image_size if isinstance(image_size, tuple) \
+                                    else (image_size, image_size)
+        patch_height, patch_width = patch_size if isinstance(patch_size, tuple) \
+                                    else (patch_size, patch_size)
+        self.image_height = image_height
+        self.image_width = image_width
+        self.patch_height = patch_height
+        self.patch_width = patch_width
+        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
+        en_pos_embedding = get_2d_sincos_pos_embed(dim, (image_height // patch_height, image_width // patch_width))
+
+        self.num_patches = (image_height // patch_height) * (image_width // patch_width)
+        self.patch_dim = channels * patch_height * patch_width
+
+        self.to_patch_embedding = nn.Sequential(
+            nn.Conv2d(channels, dim, kernel_size=patch_size, stride=patch_size),
+            Rearrange('b c h w -> b (h w) c'),
+        )
+        self.en_pos_embedding = nn.Parameter(torch.from_numpy(en_pos_embedding).float().unsqueeze(0), requires_grad=False)
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim)
+
+        self.apply(init_weights)
+
+    def forward(self, img: torch.FloatTensor) -> torch.FloatTensor:
+        x = self.to_patch_embedding(img)
+        x = x + self.en_pos_embedding
+        x = self.transformer(x) # (B, N, D)
+        #import pdb; pdb.set_trace()
+        x = x.reshape(x.shape[0], x.shape[1]//(self.image_height // self.patch_height), x.shape[1]//(self.image_width // self.patch_width), -1)
+        x = x.permute(0, 3, 1, 2).contiguous()
+        #import pdb; pdb.set_trace()
+        return x
+
+class Encoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, double_z=True, **ignore_kwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                if isinstance(curr_res, int):
+                    curr_res = curr_res // 2
+                elif isinstance(curr_res, (list, ListConfig)):
+                    curr_res = [r // 2 for r in curr_res]
+                else:
+                    raise TypeError(f"Unsupported type for curr_res: {type(curr_res)}")
+
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        2*z_channels if double_z else z_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+
+    def forward(self, x):
+        # timestep embedding
+        temb = None
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, give_pre_end=False, **ignorekwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1,)+tuple(ch_mult)
+        block_in = ch*ch_mult[self.num_resolutions-1]
+        
+        curr_res_h, curr_res_w = (resolution[0] // 2**(self.num_resolutions-1), resolution[1] // 2**(self.num_resolutions-1)) if isinstance(resolution, (list, tuple, ListConfig)) else (resolution// 2**(self.num_resolutions-1) ,resolution// 2**(self.num_resolutions-1))        
+        self.z_shape = (1,z_channels,curr_res_h,curr_res_w)
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(z_channels,
+                                       block_in,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        
+        #self.decoder_dino_conv = nn.Conv2d(block_in, 768, kernel_size=1, stride=1, padding=0) # remove hard-coded emb_dim=768
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res_h in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res_h = curr_res_h * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, z):
+        #assert z.shape[1:] == self.z_shape[1:]
+        self.last_z_shape = z.shape
+
+        # timestep embedding
+        temb = None
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+
+        return h
+    

orbis2/networks/tokenizer/pretrained_models.py

@@ -0,0 +1,158 @@
+import torch
+import timm
+from torch import nn
+from omegaconf import ListConfig
+from einops import rearrange
+
+
+from typing import Tuple, Union
+
+class Encoder(nn.Module):
+    def __init__(
+        self, 
+        resolution: Union[Tuple[int, int], int], 
+        channels: int = 3, 
+        pretrained_encoder = 'MAE',
+        patch_size: int = 16,
+        z_channels: int = 768,
+        e_dim: int = 8,
+        normalize_embedding: bool = True,
+        use_pretrained_weights = True,
+        # **ignore_kwargs
+    ) -> None:
+        # Initialize parent class with the first patch size
+        super().__init__()
+        self.image_size = resolution
+        self.patch_size = patch_size
+        self.channels = channels
+        self.normalize_embedding = normalize_embedding
+        self.z_channels = z_channels
+        self.e_dim = e_dim
+        self.pretrained_encoder = pretrained_encoder
+        
+        self.init_transformer(pretrained_encoder, use_pretrained_weights)
+
+    def init_transformer(self, pretrained_encoder, use_pretrained_weights):
+        if pretrained_encoder == 'VIT_DINO':
+            pretrained_encoder_model = 'timm/vit_base_patch16_224.dino'
+        elif pretrained_encoder == 'VIT_DINOv2':
+            pretrained_encoder_model = 'timm/vit_base_patch14_dinov2.lvd142m'
+        elif pretrained_encoder == 'MAE':
+            pretrained_encoder_model = 'timm/vit_base_patch16_224.mae'
+        elif pretrained_encoder == 'MAE_VIT_L':
+            pretrained_encoder_model = 'timm/vit_large_patch16_224.mae'
+        elif pretrained_encoder == 'VIT':
+            pretrained_encoder_model = 'timm/vit_large_patch32_224.orig_in21k'
+        elif pretrained_encoder == 'CLIP32':
+            pretrained_encoder_model = 'timm/vit_base_patch32_clip_224.openai'
+        elif pretrained_encoder == 'CLIP':
+            pretrained_encoder_model = 'timm/vit_base_patch16_clip_224.openai'
+        elif pretrained_encoder == 'base':
+            pretrained_encoder_model = 'timm/vit_base_patch16_224'
+        elif pretrained_encoder == 'large':
+            pretrained_encoder_model = 'timm/vit_large_patch16_224'
+
+        if pretrained_encoder == "VIT_DINOv3":
+            self.encoder = torch.hub.load('../dinov3', 'dinov3_vitb16', source='local', weights='./pretrained_models/dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth').train()
+        else:
+            self.encoder = timm.create_model(pretrained_encoder_model, img_size=self.image_size, patch_size=self.patch_size, pretrained=False, dynamic_img_size=True).train()
+            if use_pretrained_weights:
+                pretrained_model = timm.create_model(pretrained_encoder_model, img_size=self.image_size, patch_size=self.patch_size, pretrained=True)
+                """Initialize weights of target_model with weights from source_model."""
+                with torch.no_grad():
+                    for target_param, source_param in zip(self.encoder.parameters(), pretrained_model.parameters()):
+                        target_param.data.copy_(source_param.data)
+
+                # Clean up
+                del pretrained_model
+    
+    def forward(self, img: torch.FloatTensor) -> torch.FloatTensor:
+        if self.pretrained_encoder == "VIT_DINOv3":
+            h = self.encoder.forward_features(img)['x_norm_patchtokens']
+        else:
+            h = self.encoder.forward_features(img)[:,1:]
+        h = h.permute(0, 2, 1).contiguous()
+        h = h.reshape(h.shape[0], -1, img.size(2)//self.patch_size, img.size(3)//self.patch_size)
+        return h
+
+
+class MaskedEncoder(Encoder):
+    def __init__(
+        self,
+        *,
+        mask_ratio_min: float = 0.0,
+        mask_ratio_max: float = 0.5,
+        mask_during_eval: bool = False,
+        mask_token_init_std: float = 0.02,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        if not 0.0 <= mask_ratio_min <= mask_ratio_max <= 1.0:
+            raise ValueError(
+                f'Expected 0 <= mask_ratio_min <= mask_ratio_max <= 1, got '
+                f'{mask_ratio_min}, {mask_ratio_max}.'
+            )
+
+        self.mask_ratio_min = float(mask_ratio_min)
+        self.mask_ratio_max = float(mask_ratio_max)
+        self.mask_during_eval = bool(mask_during_eval)
+        embed_dim = getattr(self.encoder, 'embed_dim', self.z_channels)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        nn.init.normal_(self.mask_token, std=mask_token_init_std)
+
+    def _sample_patch_mask(self, batch_size: int, num_patches: int, device: torch.device):
+        if self.mask_ratio_max <= 0.0:
+            return None
+        if not self.training and not self.mask_during_eval:
+            return None
+
+        ratios = torch.empty(batch_size, device=device).uniform_(self.mask_ratio_min, self.mask_ratio_max)
+        return torch.rand(batch_size, num_patches, device=device) < ratios.unsqueeze(1)
+
+    def _apply_mask_after_patch_embed(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() != 3:
+            raise ValueError(f'Expected patch embeddings with shape (B, N, C), got {tuple(x.shape)}.')
+        mask = self._sample_patch_mask(x.shape[0], x.shape[1], x.device)
+        if mask is None:
+            return x
+        mask_token = self.mask_token.to(dtype=x.dtype).expand(x.shape[0], x.shape[1], -1)
+        return torch.where(mask.unsqueeze(-1), mask_token, x)
+
+    def _forward_timm_features(self, img: torch.FloatTensor) -> torch.Tensor:
+        x = self.encoder.patch_embed(img)
+        if x.dim() == 4:
+            embed_dim = getattr(self.encoder, 'embed_dim', x.shape[-1])
+            if x.shape[-1] == embed_dim:
+                b, h, w, c = x.shape
+                x = x.reshape(b, h * w, c)
+                x = self._apply_mask_after_patch_embed(x)
+                x = x.reshape(b, h, w, c)
+            else:
+                x = x.flatten(2).transpose(1, 2)
+                x = self._apply_mask_after_patch_embed(x)
+        else:
+            x = self._apply_mask_after_patch_embed(x)
+        x = self.encoder._pos_embed(x)
+        x = self.encoder.patch_drop(x)
+        x = self.encoder.norm_pre(x)
+
+        blocks = self.encoder.blocks
+        if isinstance(blocks, nn.ModuleList):
+            for blk in blocks:
+                x = blk(x)
+        else:
+            x = blocks(x)
+        x = self.encoder.norm(x)
+        return x
+
+    def forward(self, img: torch.FloatTensor) -> torch.FloatTensor:
+        if self.pretrained_encoder == 'VIT_DINOv3':
+            raise NotImplementedError('MaskedEncoder currently supports timm-based encoders only.')
+
+        h = self._forward_timm_features(img)
+        num_prefix_tokens = getattr(self.encoder, 'num_prefix_tokens', 1)
+        if num_prefix_tokens > 0:
+            h = h[:, num_prefix_tokens:]
+        h = h.permute(0, 2, 1).contiguous()
+        h = h.reshape(h.shape[0], -1, img.size(2) // self.patch_size, img.size(3) // self.patch_size)
+        return h

orbis2/util.py

@@ -0,0 +1,214 @@
+import argparse
+import hashlib
+import importlib
+import logging
+import os
+import requests
+
+from tqdm import tqdm
+from pytorch_lightning.trainer import Trainer
+
+
+def maybe_raise_nofile_limit(target_soft=32768, logger=None):
+    """
+    Raise RLIMIT_NOFILE soft limit for current process when possible.
+
+    Returns True if the limit was raised, False otherwise.
+    """
+    if logger is None:
+        logger = logging.getLogger(__name__)
+    try:
+        import resource
+    except ImportError:
+        logger.warning("resource module unavailable; cannot set RLIMIT_NOFILE")
+        return False
+    try:
+        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+        desired = target_soft
+        if hard not in (resource.RLIM_INFINITY, -1):
+            desired = min(desired, hard)
+        if soft >= desired:
+            return False
+        resource.setrlimit(resource.RLIMIT_NOFILE, (desired, hard))
+        logger.info("Raised RLIMIT_NOFILE from %s to %s", soft, desired)
+        return True
+    except (ValueError, PermissionError, OSError) as exc:
+        logger.warning("Unable to raise RLIMIT_NOFILE: %s", exc)
+        return False
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class KeyNotFoundError(Exception):
+    def __init__(self, cause, keys=None, visited=None):
+        self.cause = cause
+        self.keys = keys
+        self.visited = visited
+        messages = list()
+        if keys is not None:
+            messages.append("Key not found: {}".format(keys))
+        if visited is not None:
+            messages.append("Visited: {}".format(visited))
+        messages.append("Cause:\n{}".format(cause))
+        message = "\n".join(messages)
+        super().__init__(message)
+
+
+def retrieve(
+    list_or_dict, key, splitval="/", default=None, expand=True, pass_success=False
+):
+    """Given a nested list or dict return the desired value at key expanding
+    callable nodes if necessary and :attr:`expand` is ``True``. The expansion
+    is done in-place.
+
+    Parameters
+    ----------
+        list_or_dict : list or dict
+            Possibly nested list or dictionary.
+        key : str
+            key/to/value, path like string describing all keys necessary to
+            consider to get to the desired value. List indices can also be
+            passed here.
+        splitval : str
+            String that defines the delimiter between keys of the
+            different depth levels in `key`.
+        default : obj
+            Value returned if :attr:`key` is not found.
+        expand : bool
+            Whether to expand callable nodes on the path or not.
+
+    Returns
+    -------
+        The desired value or if :attr:`default` is not ``None`` and the
+        :attr:`key` is not found returns ``default``.
+
+    Raises
+    ------
+        Exception if ``key`` not in ``list_or_dict`` and :attr:`default` is
+        ``None``.
+    """
+
+    keys = key.split(splitval)
+
+    success = True
+    try:
+        visited = []
+        parent = None
+        last_key = None
+        for key in keys:
+            if callable(list_or_dict):
+                if not expand:
+                    raise KeyNotFoundError(
+                        ValueError(
+                            "Trying to get past callable node with expand=False."
+                        ),
+                        keys=keys,
+                        visited=visited,
+                    )
+                list_or_dict = list_or_dict()
+                parent[last_key] = list_or_dict
+
+            last_key = key
+            parent = list_or_dict
+
+            try:
+                if isinstance(list_or_dict, dict):
+                    list_or_dict = list_or_dict[key]
+                else:
+                    list_or_dict = list_or_dict[int(key)]
+            except (KeyError, IndexError, ValueError) as e:
+                raise KeyNotFoundError(e, keys=keys, visited=visited)
+
+            visited += [key]
+        # final expansion of retrieved value
+        if expand and callable(list_or_dict):
+            list_or_dict = list_or_dict()
+            parent[last_key] = list_or_dict
+    except KeyNotFoundError as e:
+        if default is None:
+            raise e
+        else:
+            list_or_dict = default
+            success = False
+
+    if not pass_success:
+        return list_or_dict
+    else:
+        return list_or_dict, success
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+
+def get_jobid():
+    try:
+        jobid = os.environ["PBS_JOBID"].replace(".lmbtorque.informatik.uni-freiburg.de", "")
+    except KeyError:
+        try:
+            jobid = "DLC" + os.environ["SLURM_JOB_ID"]
+        except KeyError:
+            try:
+                jobid = os.uname()[1]
+            except KeyError:
+                jobid = "local"
+    return jobid
+
+
+def nondefault_trainer_args(opt):
+    parser = argparse.ArgumentParser()
+    parser = Trainer.add_argparse_args(parser)
+    args = parser.parse_args([])
+    return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k))
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+

requirements.txt

@@ -12,6 +12,11 @@ imageio
 timm
 pillow
 
+# Orbis 2 (app2.py) extra deps — video-in rollout reads mp4s directly
+decord
+h5py
+requests
+
 # Space infrastructure
 gradio
 spaces