diff --git a/integrations/omnidreams/drift_correction/.gitignore b/integrations/omnidreams/drift_correction/.gitignore
new file mode 100644
index 00000000..17aa483a
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/.gitignore
@@ -0,0 +1 @@
+outputs/
diff --git a/integrations/omnidreams/drift_correction/README.md b/integrations/omnidreams/drift_correction/README.md
new file mode 100644
index 00000000..e11ecbe5
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/README.md
@@ -0,0 +1,89 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# Clean Forcing drift corrector for Omnidreams
+
+Long autoregressive Omnidreams rollouts drift: each chunk conditions on
+self-generated history and small errors compound (color/saturation runaway,
+texture mush, content repetition). This module ships a trained corrector — a
+frozen-base LoRA (r16 on the self-attention q/k/v/o projections, ~7.3M
+params) trained from the model's own closed-loop rollouts with a
+counterfactual clean-history teacher — plus the full recipe that produced it.
+
+## Deploy
+
+Point the runner at a corrector checkpoint; everything else is automatic:
+
+```python
+config = OmnidreamsRunnerConfig(
+    ...,
+    drift_corrector=Path("lora_v2_v3_valpeak.pt"),  # None (default) = exact base behavior
+    # drift_corrector_gain defaults to the shipped 0.25
+)
+```
+
+The hook (`omnidreams/_drift_corrector.py`) wraps the transformer's
+self-attention projections with the LoRA and rescales it to
+``alpha*(t) x gain`` before every denoise step, where ``alpha*(t)`` is the
+measured per-timestep systematic share of the drift error (0.96 at t=1000,
+0.667 at t=803 on this host). The LoRA stays unfused (the per-timestep gate
+cannot merge into one weight set); ``drift_corrector=None`` skips module
+surgery entirely.
+
+The trained checkpoint (~88 MB ``.pt``,
+md5 ``84beb8c014346833ce182310373b5d32``) is distributed separately; see the
+PR / release notes for the download link.
+
+Measured on held-out 21.5 s rollouts (704x1280, 81 chunks): drift
+Delta-MUSIQ +0.99 vs base +2.44 (60% cut), late-window MUSIQ +5% over base,
+overall MUSIQ above base, RAFT dynamics −17%.
+
+## Reproduce the corrector
+
+All scripts run from the repo root on one GPU (~1 GPU-day end to end; every
+stage checkpoints and resumes). The pair rollouts need the
+``nvidia/omni-dreams-samples`` HF dataset (real dashcam first frames + 80 s
+HDMaps; gated — export ``HF_TOKEN``):
+
+```bash
+python drift_correction/build_pairs.py      # tiled-HDMap loop rollouts -> pair clips
+                                            # (LAP_MIX=4,5,6 = shipped mixed-lap layout)
+python drift_correction/gate_faithful.py    # step-0 alpha*(t) go/no-go diagnostic
+python drift_correction/train_v1.py         # counterfactual-teacher LoRA (velocity space)
+python drift_correction/build_pairs.py      # DAgger round: CORRECTOR_LORA=<v1 ckpt>
+python drift_correction/train_v2.py         # DAgger pool + drift-contraction round
+                                            # (SNAP_EVERY=200 keeps the val-peak snapshot
+                                            #  — the shipped checkpoint is a val-peak)
+python drift_correction/eval_rollouts.py    # dial-grid sweep on held-out scenes
+```
+
+Train only if the gate passes (the drift gap is systematic: ``alpha*``
+high); the measured ``alpha*(t)`` profile doubles as the deploy gate in
+``omnidreams/_drift_corrector.py``. Correctors encode a checkpoint-specific
+error map — re-instantiation on a new base checkpoint is ~1 GPU-day with
+this recipe.
+
+Ablation flags kept in the scripts: ``PAIR_SCHEME=fork`` +
+``build_pairs_v4.py`` (loop-free re-anchored fork counterparts; the pair
+design gate-validates at ``alpha* 0.89``), ``UW=1`` (per-token
+``alpha*``-weighted loss), ``TBIN_EVAL`` (per-timestep val R², feeds
+``gain_predict.py``).
+
+## Known issues (shipped config)
+
+- Some repeated rows of buildings/shops and repeated trucks can appear in
+  long rollouts (partial content-repetition prior from the loop-structured
+  training pairs; the mixed-lap layout removes the fixed-period component).
+- Mild softness vs base on some scenes (owner-reviewed and accepted; the
+  variance-of-Laplacian sharpness ratio is NOT a reliable proxy on this
+  host — judge with frame stacks).
+- RAFT dynamic degree −17% vs base.
+
+## Design notes and background
+
+- Method, gate math, result tables, and the paper-host reproduction:
+  https://gitlab-master.nvidia.com/wenqingw/clean_forcing
+- The HY-WorldPlay port of the same recipe (content-keyed deploy):
+  ``integrations/hy_worldplay/drift_correction/``.
diff --git a/integrations/omnidreams/drift_correction/_host.py b/integrations/omnidreams/drift_correction/_host.py
new file mode 100644
index 00000000..b08ece9c
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/_host.py
@@ -0,0 +1,320 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Omnidreams host adapter for the Clean Forcing drift corrector.
+
+Implements the three per-model functions from the library design
+(``LIBRARY_DESIGN.md`` in the HY port): rollout capture, and the
+counterfactual x0 probe (history substitution + matched-state predict).
+
+Host specifics vs HY-WorldPlay: history lives ONLY in the rolling
+self-attention KV window (``window_size_t`` latent frames; no explicit
+latent-history tensor and no long-range memory reads), so substituting a
+history means resetting the per-block KV caches and replaying the prefix
+chunks through ``finalize_kv_cache`` forwards at the model's own
+``context_noise`` timestep. RoPE is absolute (``shift_t(ar_idx)``), so a
+replay at the original chunk indices reproduces positions exactly.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import gc
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+# Must land before the first CUDA allocation: long captures fragment the
+# allocator; expandable segments keep the probe phase inside the VRAM share
+# left over by co-tenant jobs.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import torch
+from omnidreams.config import SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE
+from omnidreams.pipeline import OmnidreamsPipeline, OmnidreamsPipelineCache
+from omnidreams.transformer import CosmosTransformer, CosmosTransformerCache
+from torch import Tensor
+
+from flashdreams.infra.config import derive_config
+
+## Pipeline construction
+
+CONTEXT_NOISE_SEED = 77_000
+"""Base seed for the per-chunk context-noise draws used by history replays.
+Both branches of a counterfactual probe re-noise their history chunks with
+the SAME eps (seeded per chunk index), so the prediction gap isolates the
+history *content* difference."""
+
+
+def build_pipeline(
+    *,
+    with_oneshot_encoders: bool,
+    seed: int | None = 42,
+) -> OmnidreamsPipeline:
+    """Build the distilled chunk2 Omnidreams pipeline for capture / probes.
+
+    Compile and CUDA graphs are disabled everywhere: probes reset and replay
+    the KV caches out of the normal AR order, which a captured graph (whose
+    slot pointers bake in one cache object) cannot express — and eager keeps
+    capture and probe numerics identical.
+
+    Args:
+        with_oneshot_encoders: Load the Cosmos-Reason1 text encoder and the
+            first-frame VAE encoder. Pair-building needs them; the gate
+            rebuilds caches from embeddings stored in the clips and skips
+            the ~14 GB text encoder.
+        seed: Diffusion-model RNG seed (callers typically re-seed per
+            rollout via ``diffusion_model._rng``).
+
+    Returns:
+        Pipeline with the checkpoint loaded, on CUDA.
+    """
+    cfg = derive_config(
+        SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE,
+        enable_sync_and_profile=False,
+        diffusion_model=dict(
+            seed=seed,
+            transformer=dict(compile_network=False, use_cuda_graph=False),
+        ),
+    )
+    if not with_oneshot_encoders:
+        cfg = dataclasses.replace(cfg, text_encoder=None, image_encoder=None)
+    pipe = cfg.setup()
+    assert isinstance(pipe, OmnidreamsPipeline)
+    return pipe.to("cuda")
+
+
+## Rollout capture
+
+
+@dataclass
+class ChunkSnapshot:
+    """Per-chunk state captured from a rollout, sufficient to replay probes."""
+
+    clean_latent: Tensor
+    """Patchified x0 the sampler produced for this chunk (image latent
+    injected at chunk 0, as the AR cache consumed it)."""
+
+    hdmap: Tensor
+    """Patchified per-AR-step HDMap conditioning exactly as
+    ``predict_flow`` consumed it (post streaming-encoder, post patchify)."""
+
+    def to(self, device: torch.device | str) -> "ChunkSnapshot":
+        """Return a copy with both tensors on ``device``."""
+        return ChunkSnapshot(
+            clean_latent=self.clean_latent.to(device), hdmap=self.hdmap.to(device)
+        )
+
+
+def capture_rollout(
+    pipe: OmnidreamsPipeline,
+    cache: OmnidreamsPipelineCache,
+    *,
+    hdmap_video: Tensor,
+    num_chunk: int,
+    noise_seed: int,
+) -> tuple[list[ChunkSnapshot], Tensor]:
+    """Roll out ``num_chunk`` chunks and snapshot the per-chunk probe inputs.
+
+    Args:
+        pipe: Pipeline from :func:`build_pipeline`.
+        cache: Fresh per-rollout cache (``initialize_cache`` /
+            ``initialize_cache_from_embeddings``); consumed by the rollout.
+        hdmap_video: ``[B=1, V=1, T, 3, H, W]`` HDMap pixels in ``[-1, 1]``;
+            must cover ``5 + (num_chunk - 1) * 8`` frames on the chunk2 host.
+        noise_seed: Seed for the diffusion model's noise generator (initial
+            noise + context-noise draws), making captures reproducible.
+
+    Returns:
+        ``(snaps, video)`` — one :class:`ChunkSnapshot` per chunk (tensors
+        stay on the compute device) and the decoded rollout
+        ``[B, V, T, 3, H, W]`` on CPU for eyeballing.
+    """
+    device = pipe.device
+    pipe.diffusion_model._rng = torch.Generator(device=device).manual_seed(noise_seed)
+
+    snaps: list[ChunkSnapshot] = []
+    chunks: list[Tensor] = []
+    start = 0
+    for ar_idx in range(num_chunk):
+        num_frames = pipe.get_num_frames(ar_idx)
+        end = start + num_frames
+        assert end <= hdmap_video.shape[2], (
+            f"hdmap video too short: chunk {ar_idx} needs frames "
+            f"[{start}, {end}) but only {hdmap_video.shape[2]} available."
+        )
+        chunk = pipe.generate(ar_idx, cache, hdmap=hdmap_video[:, :, start:end])
+        fs = cache.final_state
+        assert fs is not None and isinstance(fs.input, Tensor)
+        snaps.append(
+            ChunkSnapshot(
+                clean_latent=fs.clean_latent.detach().clone(),
+                hdmap=fs.input.detach().clone(),
+            )
+        )
+        pipe.finalize(ar_idx, cache)
+        chunks.append(chunk.cpu())
+        start = end
+
+    video = torch.cat(chunks, dim=2)
+    del cache, chunks
+    gc.collect()
+    torch.cuda.empty_cache()
+    return snaps, video
+
+
+## Counterfactual x0 probes
+
+
+@torch.no_grad()
+def swap_text_kv(network, tc: CosmosTransformerCache, text_embeddings: Tensor) -> None:
+    """Swap the per-block cross-attention (text) KV for another clip's prompt.
+
+    Mirrors the cross-attn half of ``CosmosDiTNetwork.initialize_cache`` so
+    one long-lived transformer cache (rope, masks, rolling self-attn
+    buffers) can serve clips with different prompts.
+    """
+    w = network.blocks[0].cross_attn.k_proj.weight
+    ctx = text_embeddings.to(device=w.device, dtype=w.dtype)
+    if network.config.use_crossattn_projection:
+        ctx = network.crossattn_proj(ctx)
+    for block, bc in zip(network.blocks, tc.network_cache.block_caches):
+        bc.cross_attn = block.cross_attn.initialize_cache(ctx)
+
+
+def reset_history(tc: CosmosTransformerCache) -> None:
+    """Reset the rolling self-attention KV caches to the empty state.
+
+    The cross-attention (text) KV and the image/mask fields are per-rollout
+    constants and stay untouched. After a reset the next replay must start
+    at chunk 0 (``BlockKVCache`` enforces contiguous chunk indices).
+    """
+    for bc in tc.network_cache.block_caches:
+        bc.self_attn.reset()
+    if tc.network_cache_uncond is not None:
+        for bc in tc.network_cache_uncond.block_caches:
+            bc.self_attn.reset()
+
+
+def replay_history(
+    transformer: CosmosTransformer,
+    tc: CosmosTransformerCache,
+    *,
+    latents: list[Tensor],
+    hdmaps: list[Tensor],
+    context_timestep: Tensor,
+    add_noise,
+) -> None:
+    """Rebuild the KV history by replaying chunks ``0 .. len(latents) - 1``.
+
+    Each chunk runs one ``finalize_kv_cache``-style forward on its latent
+    re-noised at the host's ``context_noise`` timestep — the exact cache
+    write path of a real rollout. The eps draw is seeded per chunk index
+    (:data:`CONTEXT_NOISE_SEED`), so gen and clean replays of the same
+    positions consume identical noise and differ only in content.
+
+    Args:
+        transformer: The pipeline's Cosmos transformer.
+        tc: AR cache, freshly :func:`reset_history`-ed.
+        latents: Patchified per-chunk x0 history (branch-specific content).
+        hdmaps: Patchified per-chunk HDMap conditioning (shared).
+        context_timestep: 0-d context-noise timestep tensor (host config).
+        add_noise: ``scheduler.add_noise`` bound from the pipeline.
+    """
+    assert len(latents) == len(hdmaps)
+    device = latents[0].device
+    for j, (lat, hd) in enumerate(zip(latents, hdmaps)):
+        rng = torch.Generator(device=device).manual_seed(CONTEXT_NOISE_SEED + j)
+        noisy = add_noise(lat, context_timestep, rng=rng)
+        tc.start(j)
+        transformer.finalize_kv_cache(
+            noisy_latent=noisy, timestep=context_timestep, cache=tc, input=hd
+        )
+        tc.finalize(j)
+
+
+def start_probe_chunk(tc: CosmosTransformerCache, *, ar_idx: int) -> None:
+    """Open the probed chunk after a history replay up to ``ar_idx - 1``."""
+    tc.start(ar_idx)
+
+
+def finish_probe_chunk(tc: CosmosTransformerCache, *, ar_idx: int) -> None:
+    """Close the ``start`` / ``finalize`` bracket after a probe sweep."""
+    tc.finalize(ar_idx)
+
+
+def predict_v(
+    transformer: CosmosTransformer,
+    tc: CosmosTransformerCache,
+    *,
+    hdmap: Tensor,
+    z_t: Tensor,
+    timestep: Tensor,
+) -> Tensor:
+    """Predict the velocity at a matched noisy state under the current history.
+
+    This host is velocity/rectified-flow (``v = eps - x0``); the corrector's
+    target space follows the solver, so probes and training targets live in
+    v-space. ``x0_hat = z_t - sigma * v`` recovers the x0 view when needed
+    (per fixed ``(z_t, t)`` the two spaces differ by the factor ``-sigma``).
+
+    Args:
+        transformer: The pipeline's Cosmos transformer.
+        tc: AR cache positioned via :func:`start_probe_chunk`.
+        hdmap: The probed chunk's captured (patchified) HDMap conditioning.
+        z_t: Patchified noisy latent ``(1 - sigma) * x0 + sigma * eps``.
+        timestep: 0-d timestep tensor in the network dtype.
+
+    Returns:
+        fp32 predicted flow, same shape as ``z_t``.
+    """
+    flow = transformer.predict_flow(
+        noisy_latent=z_t, timestep=timestep, cache=tc, input=hdmap
+    )
+    return flow.float()
+
+
+## Clip I/O (shared by build_pairs / gate)
+
+
+def save_clip(
+    path: Path,
+    *,
+    snaps: list[ChunkSnapshot],
+    embeddings: dict,
+    meta: dict,
+) -> None:
+    """Save a captured pair clip (CPU tensors) for the gate / training."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "latents": [s.clean_latent.cpu() for s in snaps],
+            "hdmaps": [s.hdmap.cpu() for s in snaps],
+            "embeddings": {
+                k: (v.cpu() if isinstance(v, Tensor) else v)
+                for k, v in embeddings.items()
+            },
+            **meta,
+        },
+        path,
+    )
+
+
+def load_clip(path: Path, device: torch.device | str, dtype: torch.dtype) -> dict:
+    """Load a ``build_pairs.py`` clip; latents/hdmaps land on ``device``."""
+    d = torch.load(path, map_location="cpu", weights_only=False)
+    d["latents"] = [x.to(device, dtype) for x in d["latents"]]
+    d["hdmaps"] = [x.to(device, dtype) for x in d["hdmaps"]]
+    return d
diff --git a/integrations/omnidreams/drift_correction/_lora.py b/integrations/omnidreams/drift_correction/_lora.py
new file mode 100644
index 00000000..52a5eee0
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/_lora.py
@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""LoRA corrector r_phi as low-rank deltas on the frozen Cosmos DiT.
+
+Mirror of the HY port's ``_lora.py`` with the Cosmos module names: the
+blocks' self-attention projections are
+``blocks.{i}.self_attn.{q,k,v,output}_proj``.
+"""
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+DEFAULT_TARGETS = (
+    "self_attn.q_proj",
+    "self_attn.k_proj",
+    "self_attn.v_proj",
+    "self_attn.output_proj",
+)
+"""Module-name suffixes wrapped by :func:`apply_lora`."""
+
+
+class LoRALinear(nn.Module):
+    """Frozen base linear plus a runtime-gated low-rank delta.
+
+    ``B`` is zero-initialized so the wrapped module starts as an exact
+    identity over the base (any ``scale`` gives the base output until
+    training moves ``B``). ``scale`` is the runtime gain: ``0`` recovers the
+    frozen base (the clean-history teacher pass), ``1`` the corrected pass.
+    The A/B path runs in fp32 regardless of the base dtype.
+    """
+
+    def __init__(self, base: nn.Linear, rank: int = 16):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad_(False)
+        self.A = nn.Linear(base.in_features, rank, bias=False)
+        self.B = nn.Linear(rank, base.out_features, bias=False)
+        nn.init.normal_(self.A.weight, std=1.0 / rank)
+        nn.init.zeros_(self.B.weight)
+        self.scale = 1.0
+
+    def forward(self, x: Tensor) -> Tensor:
+        out = self.base(x)
+        if self.scale != 0:
+            delta = self.B(self.A(x.to(self.A.weight.dtype)))
+            out = out + self.scale * delta.to(out.dtype)
+        return out
+
+
+def unwrap_compiled(network: object) -> nn.Module:
+    """Return the eager module behind a ``torch.compile`` wrapper, if any.
+
+    Accepts ``object``: the DiT lives on the transformer as a plain
+    attribute, which the type checker resolves through ``nn.Module``'s
+    ``__getattr__`` union.
+    """
+    inner = getattr(network, "_orig_mod", network)
+    assert isinstance(inner, nn.Module), type(inner)
+    return inner
+
+
+def apply_lora(
+    model: nn.Module,
+    rank: int = 16,
+    targets: tuple[str, ...] = DEFAULT_TARGETS,
+) -> list[str]:
+    """Wrap every target linear in ``model`` with a :class:`LoRALinear`.
+
+    Args:
+        model: The (unwrapped) DiT network.
+        rank: LoRA rank.
+        targets: Module-name suffixes to wrap.
+
+    Returns:
+        Fully qualified names of the wrapped modules.
+    """
+    wrapped = []
+    for mname, module in list(model.named_modules()):
+        for cname, child in list(module.named_children()):
+            full = f"{mname}.{cname}" if mname else cname
+            if isinstance(child, nn.Linear) and any(t in full for t in targets):
+                setattr(module, cname, LoRALinear(child, rank).to(child.weight.device))
+                wrapped.append(full)
+    return wrapped
+
+
+def set_lora_scale(model: nn.Module, scale: float) -> None:
+    """Set the runtime gain on every :class:`LoRALinear` in ``model``."""
+    for m in model.modules():
+        if isinstance(m, LoRALinear):
+            m.scale = scale
+
+
+def lora_parameters(model: nn.Module) -> list[nn.Parameter]:
+    """Collect the trainable A/B parameters in deterministic module order."""
+    ps: list[nn.Parameter] = []
+    for m in model.modules():
+        if isinstance(m, LoRALinear):
+            ps += list(m.A.parameters()) + list(m.B.parameters())
+    return ps
+
+
+def save_lora(model: nn.Module, path) -> None:
+    """Save the LoRA parameters (index-keyed, CPU) to ``path``."""
+    torch.save(
+        {"lora": {i: p.detach().cpu() for i, p in enumerate(lora_parameters(model))}},
+        path,
+    )
+
+
+def load_lora(model: nn.Module, path) -> None:
+    """Load :func:`save_lora` output into an already-wrapped ``model``."""
+    sd = torch.load(path, map_location="cpu", weights_only=False)["lora"]
+    params = lora_parameters(model)
+    assert len(sd) == len(params), (
+        f"checkpoint has {len(sd)} LoRA tensors but the model exposes "
+        f"{len(params)}; rank or target mismatch."
+    )
+    for i, p in enumerate(params):
+        p.data.copy_(sd[i].to(p.device, p.dtype))
diff --git a/integrations/omnidreams/drift_correction/_train_attn.py b/integrations/omnidreams/drift_correction/_train_attn.py
new file mode 100644
index 00000000..2771fda4
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/_train_attn.py
@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Grad-friendly functional self-attention for corrector training (Omnidreams).
+
+The production ``SelfAttention.forward`` routes the current chunk's K / V
+through the rolling ``BlockKVCache`` (in-place write into a no-grad buffer,
+then a buffer read), which severs gradients into the ``k_proj`` / ``v_proj``
+projections; the fused RoPE Triton kernel is likewise inference-only (raw
+stores, no backward).
+
+The functional variant below computes Q / K / V directly (differentiable
+non-interleaved RoPE), reads the *history prefix* from the cache buffer
+(no grad -- replayed history is frozen in v1), and runs attention over
+``[prefix, current]`` without writing the cache. For a single probe forward
+per ``start`` / ``finalize`` bracket this is mathematically identical to the
+stock path. Deployment keeps the stock path; the weights are shared.
+
+The patch is a *toggle*, not a wholesale swap: history-replay forwards
+(``finalize_kv_cache``) must keep the stock cache-writing path, so probes
+run inside ``functional_attention()`` and everything else stays stock.
+
+Mirrors :class:`omnidreams.transformer.impl.modules.SelfAttention` and the
+``rope_kernel`` semantics (cos/sin in fp32, cast to ``x.dtype`` before the
+multiply); revisit on upstream changes.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import math
+
+import torch
+from omnidreams.transformer.impl.modules import SelfAttention
+from torch import Tensor
+
+from flashdreams.core.attention.kvcache import BlockKVCache
+
+_FUNCTIONAL = False
+"""Process-wide toggle read by the patched forward."""
+
+_RECORD: list | None = None
+"""When a list, each functional self-attn forward appends its (k, v)
+projection pair (block order) — the grad-carrying KV of a "committed"
+chunk for the contraction term."""
+
+_INJECT: list | None = None
+"""When a list (block order), each functional forward swaps the LAST
+``chunk_size`` prefix tokens' K/V for the injected grad-carrying pair
+(numerically identical to the buffer twin written by the stock finalize)."""
+
+_STOCK_FORWARD = SelfAttention.forward
+"""Original cache-writing forward, used whenever the toggle is off."""
+
+
+def _rope_half(x: Tensor, freqs: Tensor) -> Tensor:
+    """Differentiable equivalent of the fused non-interleaved RoPE kernel.
+
+    Rotates pair ``(d, d + D/2)`` by angle ``freqs[..., d]`` (the full-width
+    ``shift_t`` layout duplicates the angles across halves, so the first
+    ``D/2`` lanes carry them).
+
+    Args:
+        x: Activations ``[B, S, H, D]``.
+        freqs: Angles ``[S, 1, 1, D]`` from
+            :meth:`RotaryPositionEmbedding3D.shift_t`.
+
+    Returns:
+        Rotated tensor, same shape/dtype, fresh storage.
+    """
+    half = x.shape[-1] // 2
+    ang = freqs[:, 0, 0, :half].float()
+    cos = ang.cos().to(x.dtype)[None, :, None, :]
+    sin = ang.sin().to(x.dtype)[None, :, None, :]
+    a, b = x[..., :half], x[..., half:]
+    return torch.cat((a * cos - b * sin, b * cos + a * sin), dim=-1)
+
+
+def _functional_forward(
+    self: SelfAttention,
+    x: Tensor,
+    kv_cache: BlockKVCache,
+    rope_freqs: Tensor | None = None,
+) -> Tensor:
+    """Cache-free re-implementation of ``SelfAttention.forward``.
+
+    Reads the replayed-history prefix ``[0, write_start)`` from the cache
+    buffers (frozen) and uses the freshly projected current-chunk K / V
+    directly, so the visible attention span matches the stock read-back
+    exactly -- in the filling phase and (post-roll) in steady state alike.
+    """
+    if not _FUNCTIONAL:
+        return _STOCK_FORWARD(self, x, kv_cache, rope_freqs)  # ty: ignore[invalid-argument-type]
+    batch_shape = x.shape[:-2]
+    batch_size = math.prod(batch_shape)
+    L, D = x.shape[-2:]
+    n, d = self.n_heads, self.head_dim
+    assert n * d == D, "n * d must be equal to D"
+
+    q = self.q_norm(self.q_proj(x).reshape(batch_size, L, n, d))
+    k = self.k_norm(self.k_proj(x).reshape(batch_size, L, n, d))
+    v = self.v_proj(x).reshape(batch_size, L, n, d)
+    if rope_freqs is not None:
+        q = _rope_half(q, rope_freqs)
+        k = _rope_half(k, rope_freqs)
+
+    if _RECORD is not None:
+        _RECORD.append((k, v))
+
+    write_start, _ = kv_cache._current_write_bounds()
+    prefix = kv_cache._seq_slice(0, write_start)
+    pk, pv = kv_cache._k[prefix], kv_cache._v[prefix]
+    if _INJECT is not None:
+        ik, iv = _INJECT.pop(0)
+        span = ik.shape[-3]
+        pk = torch.cat([pk[..., :-span, :, :], ik], dim=-3)
+        pv = torch.cat([pv[..., :-span, :, :], iv], dim=-3)
+    keys = torch.cat([pk, k], dim=-3)
+    vals = torch.cat([pv, v], dim=-3)
+
+    out = self.attn_op(q, keys, vals)
+    return self.output_proj(out.reshape(batch_shape + (L, n * d)))
+
+
+def patch_functional_attention() -> None:
+    """Install the toggled functional forward, process-wide (idempotent)."""
+    SelfAttention.forward = _functional_forward  # ty: ignore[invalid-assignment]
+
+
+@contextlib.contextmanager
+def functional_attention():
+    """Run probe forwards through the grad-friendly, non-writing path."""
+    global _FUNCTIONAL
+    prev = _FUNCTIONAL
+    _FUNCTIONAL = True
+    try:
+        yield
+    finally:
+        _FUNCTIONAL = prev
+
+
+@contextlib.contextmanager
+def record_kv(store: list):
+    """Record each block's grad-carrying (k, v) during one forward."""
+    global _RECORD
+    prev = _RECORD
+    _RECORD = store
+    try:
+        yield
+    finally:
+        _RECORD = prev
+
+
+@contextlib.contextmanager
+def inject_kv(kvs: list):
+    """Swap the committed chunk's buffered KV for grad-carrying twins.
+
+    ``kvs`` must be the block-ordered list from :func:`record_kv`, and the
+    buffer's corresponding slots must hold the numerically identical
+    no-grad twins (same inputs through the stock finalize write).
+    """
+    global _INJECT
+    prev = _INJECT
+    _INJECT = list(kvs)
+    try:
+        yield
+    finally:
+        _INJECT = prev
diff --git a/integrations/omnidreams/drift_correction/build_pairs.py b/integrations/omnidreams/drift_correction/build_pairs.py
new file mode 100644
index 00000000..21eae4f2
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/build_pairs.py
@@ -0,0 +1,287 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tiled-HDMap loop rollouts: drifted-vs-clean pair clips for Omnidreams (v2).
+
+The Omnidreams action grammar is the HDMap conditioning video, so the
+strafe-loop trick ports as *HDMap tiling*: the rollout conditioning is
+``clip[:5] + tile(clip[5 : 5 + 8 * LAP_CHUNKS], LAPS)`` — every lap's chunk
+``i`` consumes pixel-identical HDMap frames, so a late (drifted) chunk's
+KV-window content can be swapped for its lap-aligned clean counterparts at
+the same conditioning and the same RoPE positions. The lap boundary is a
+conditioning teleport; probes/training only use chunks whose surviving
+window lies inside one lap (gate) or map counterfactuals by lap position
+(training). Rollouts run ~21.5 s (81 chunks) because this host's residual
+drift shows late.
+
+v2 (régime fix 2026-07-22): pairs v1 seeded rollouts with the local
+benchmarking corpus, whose "first frames" are HDMap renders — the rollouts
+stayed in a render-adjacent régime and collapsed to one scene. v2 sources
+the gated HF sample set (``nvidia/omni-dreams-samples``): 32 clips with
+REAL dashcam first frames, per-clip prompts, and 80 s authentic HDMaps.
+Requires an authenticated HF token.
+
+v3 (v3 layout, 2026-07-23): run with ``LAP_MIX=4,5,6`` — mixed lap
+lengths rotated per clip break the fixed 40-frame revisit period that let
+the v2-trained corrector learn a repeat prior (see README known issues). Default
+``LAP_MIX=5`` reproduces the v2 layout exactly.
+
+Run from the flashdreams repo root::
+
+    HF_TOKEN=$(cat ~/.cache/huggingface/token) \
+        .venv/bin/python integrations/omnidreams/drift_correction/build_pairs.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import torch
+from _host import build_pipeline, capture_rollout, save_clip
+from omnidreams.runner import (
+    DEFAULT_VIDEO_HEIGHT,
+    DEFAULT_VIDEO_WIDTH,
+    _ensure_hf_single_view_example_data_synced,
+    _load_video,
+)
+
+from flashdreams.infra.runner_io import (
+    load_first_frame_tensor,
+    write_video_tensor,
+)
+
+## Pair-build configuration
+
+OUT_DIR = Path(
+    os.environ.get(
+        "PAIRS_DIR", "integrations/omnidreams/drift_correction/outputs/pairs_v2"
+    )
+)
+"""Clip files consumed by ``gate_faithful.py`` / ``train_v1.py``."""
+
+N_CLIPS = int(os.environ.get("N_CLIPS", "6"))
+"""Sample clips to roll out (first ``N_CLIPS`` of the dataset, sorted)."""
+
+LAP_MIX = tuple(int(x) for x in os.environ.get("LAP_MIX", "5").split(","))
+"""AR chunks per lap, rotated per clip (clip ``c`` gets
+``LAP_MIX[c % len(LAP_MIX)]``). Pairs v1/v2 used a single ``5`` (40 decoded
+frames = 1.33 s). Pairs v3 uses ``4,5,6`` (32/40/48 frames): a FIXED lap
+period lets the corrector learn "content recurs at lag 40" as a shortcut —
+the HY v2 repeat-prior failure, reproduced here on the scen6 deploy sweep
+(hallucinated repeated trailers + suppressed conditioned crosswalk;
+see README known issues) — and mixing the lengths breaks the fixed
+revisit period."""
+
+TARGET_LAP_CHUNKS = 80
+"""Tiled chunks per rollout: laps per clip = ``round(80 / lap_chunks)``,
+so every clip runs ~81 chunks = ~21.5 s regardless of its lap length."""
+
+
+def clip_layout(c: int) -> tuple[int, int, int]:
+    """Per-clip ``(lap_chunks, laps, num_chunk)`` under the LAP_MIX rotation.
+
+    ``num_chunk`` = chunk 0 (image-anchored, 5 frames) + the tiled laps.
+    """
+    lap_chunks = LAP_MIX[c % len(LAP_MIX)]
+    laps = round(TARGET_LAP_CHUNKS / lap_chunks)
+    return lap_chunks, laps, 1 + lap_chunks * laps
+
+
+NOISE_SEED = int(os.environ.get("NOISE_SEED", "5042"))
+"""Diffusion RNG seed per rollout (offset by clip index)."""
+
+CORRECTOR_LORA = os.environ.get("CORRECTOR_LORA", "")
+"""Optional v1 checkpoint; when set, rollouts run with the corrector at the
+deploy dial (per-step ``alpha*(t) x 0.5``) so the DAgger pool reflects the
+states the deployed corrector actually visits."""
+
+
+def _list_sample_uuids(n: int) -> list[str]:
+    """Return the first ``n`` single-view sample UUIDs, alphabetically.
+
+    ``SAMPLE_UUIDS`` (comma list) bypasses the HF listing API — the shared
+    IP rate limit killed three pipeline stages on 2026-07-24; every consumer
+    already has the files in the local HF cache, only this listing needed
+    the network.
+    """
+    if os.environ.get("SAMPLE_UUIDS"):
+        uuids = os.environ["SAMPLE_UUIDS"].split(",")
+        assert len(uuids) >= n, f"SAMPLE_UUIDS lists only {len(uuids)} clips"
+        return uuids[:n]
+    from huggingface_hub import HfApi
+    from huggingface_hub.hf_api import RepoFolder
+
+    entries = HfApi().list_repo_tree(
+        repo_id="nvidia/omni-dreams-samples",
+        repo_type="dataset",
+        path_in_repo="data/single_view",
+        recursive=False,
+    )
+    uuids = sorted(
+        e.path.rsplit("/", 1)[-1] for e in entries if isinstance(e, RepoFolder)
+    )
+    assert len(uuids) >= n, f"dataset lists only {len(uuids)} single-view clips"
+    return uuids[:n]
+
+
+def _sample_files(uuid: str) -> tuple[tuple[Path, ...], tuple[Path, ...]]:
+    """Cache-first ``((hdmap,), (first_frame,))`` resolution for a sample.
+
+    The runner's ``_ensure_hf_single_view_example_data_synced`` lists the HF
+    repo per clip even when every file is already local; the shared IP rate
+    limit killed three pipeline stages on 2026-07-24. Fall back to the
+    network path only on a cache miss.
+    """
+    root = (
+        Path.home()
+        / ".cache/huggingface/hub/datasets--nvidia--omni-dreams-samples/snapshots"
+    )
+    hits = sorted(root.glob(f"*/data/single_view/{uuid}/*_hdmap.mp4"))
+    for h in hits:
+        frame = h.parent / "first_frame.png"
+        if frame.exists():
+            return (h,), (frame,)
+    return _ensure_hf_single_view_example_data_synced(uuid)
+
+
+def _clip_prompt(uuid: str) -> str:
+    """Fetch the clip's own prompt from the dataset (cache-first)."""
+    root = (
+        Path.home()
+        / ".cache/huggingface/hub/datasets--nvidia--omni-dreams-samples/snapshots"
+    )
+    hits = sorted(root.glob(f"*/data/single_view/{uuid}/prompt.txt"))
+    if hits:
+        return hits[0].read_text().strip()
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(
+        repo_id="nvidia/omni-dreams-samples",
+        repo_type="dataset",
+        filename=f"data/single_view/{uuid}/prompt.txt",
+    )
+    return Path(path).read_text().strip()
+
+
+def tile_hdmap(hdmap: torch.Tensor, lap_chunks: int, laps: int) -> torch.Tensor:
+    """Tile a ``[T, C, H, W]`` HDMap clip into the loop conditioning.
+
+    Layout: the first 5 frames (chunk 0) then ``laps`` copies of the next
+    ``8 * lap_chunks`` frames, so chunks ``1 + j * lap_chunks + i`` consume
+    identical pixels for every lap ``j``.
+    """
+    lap_frames = 8 * lap_chunks
+    assert hdmap.shape[0] >= 5 + lap_frames, (
+        f"clip has {hdmap.shape[0]} frames; need {5 + lap_frames}."
+    )
+    lap = hdmap[5 : 5 + lap_frames]
+    return torch.cat([hdmap[:5], lap.repeat(laps, 1, 1, 1)], dim=0)
+
+
+def main() -> None:
+    torch.set_grad_enabled(False)
+    dtype = torch.bfloat16
+
+    # Load every clip's inputs BEFORE any model/rollout work: video decode
+    # forks ffmpeg, which fails silently (empty read) once this process has
+    # grown to rollout size -- observed 3x on this box, always on the
+    # second decode of a job.
+    inputs: list[tuple[str, str, torch.Tensor, torch.Tensor] | None] = []
+    for c, uuid in enumerate(_list_sample_uuids(N_CLIPS)):
+        if (OUT_DIR / f"clip_{c:02d}.pt").exists():
+            inputs.append(None)
+            continue
+        (hdmap_path,), (frame_path,) = _sample_files(uuid)
+        hdmap = _load_video(
+            hdmap_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device="cpu",  # ty: ignore[invalid-argument-type]
+            dtype=dtype,
+        )
+        lap_chunks, laps, _ = clip_layout(c)
+        hdmap = tile_hdmap(hdmap, lap_chunks, laps)[None, None]  # [1, 1, T, C, H, W]
+        first = load_first_frame_tensor(
+            frame_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device="cpu",  # ty: ignore[invalid-argument-type]
+            dtype=dtype,
+        )[None, :, None]  # [1, V=1, 1, C, H, W]
+        inputs.append((uuid, _clip_prompt(uuid), hdmap, first))
+        print(f"loaded inputs for clip {c} ({uuid})", flush=True)
+
+    pipe = build_pipeline(with_oneshot_encoders=True)
+    device = pipe.device
+
+    if CORRECTOR_LORA:
+        from _lora import apply_lora, load_lora, unwrap_compiled
+        from eval_rollouts import install_alpha_gate
+
+        transformer = pipe.diffusion_model.transformer
+        network = unwrap_compiled(transformer.network)
+        apply_lora(network)
+        load_lora(network, CORRECTOR_LORA)
+        install_alpha_gate(transformer, network, {"gain": ("gate", 0.5)})
+        print(f"corrector active at gate x 0.5: {CORRECTOR_LORA}", flush=True)
+
+    for c, item in enumerate(inputs):
+        out = OUT_DIR / f"clip_{c:02d}.pt"
+        if item is None:
+            print(f"SKIP clip {c}: {out} exists", flush=True)
+            continue
+        uuid, prompt, hdmap, first = item
+        lap_chunks, laps, num_chunk = clip_layout(c)
+        hdmap = hdmap.to(device)
+        first = first.to(device)
+
+        embeddings = pipe.precompute_embeddings(text=[[prompt]], image=first)
+        cache = pipe.initialize_cache_from_embeddings(
+            text_embeddings=embeddings["text_embeddings"],  # ty: ignore[invalid-argument-type]
+            image_embeddings=embeddings["image_embeddings"],  # ty: ignore[invalid-argument-type]
+        )
+        snaps, video = capture_rollout(
+            pipe,
+            cache,
+            hdmap_video=hdmap,
+            num_chunk=num_chunk,
+            noise_seed=NOISE_SEED + c,
+        )
+        save_clip(
+            out,
+            snaps=snaps,
+            embeddings=embeddings,
+            meta={
+                "uuid": uuid,
+                "prompt": prompt,
+                "num_chunk": num_chunk,
+                "lap_chunks": lap_chunks,
+                "laps": laps,
+                "noise_seed": NOISE_SEED + c,
+            },
+        )
+        mp4 = OUT_DIR / f"clip_{c:02d}_loop.mp4"
+        write_video_tensor(video[0, 0].permute(0, 2, 3, 1), mp4, fps=30, layout="thwc")
+        print(f"clip {c} ({uuid}): saved {out} + {mp4}", flush=True)
+
+    print("PAIRS-DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/drift_correction/build_pairs_v4.py b/integrations/omnidreams/drift_correction/build_pairs_v4.py
new file mode 100644
index 00000000..5ef42b0d
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/build_pairs_v4.py
@@ -0,0 +1,216 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Non-looping drift pairs with re-anchored fork counterparts (pairs v4).
+
+Every loop-pair recipe (v1-v3) revisits its own content — the corrector
+learns a *loop prior* (self-copying / novel-content suppression; project
+history). v4 removes revisits entirely:
+
+- **Conditioning**: each rollout consumes its clip's OWN continuous 645-frame
+  HDMap (the HF samples ship ~80 s) — no tiling, no stitching, no teleports;
+  content never recurs by construction.
+- **Clean counterpart**: every :data:`SEG_CHUNKS` chunks the build forks a
+  fresh IMAGE-ANCHORED rollout B_s seeded by re-encoding the drifted rollout
+  A's own decoded frame at the fork, rolled over the same upcoming HDMap
+  frames. B is content-matched at the fork and early-horizon-clean (drift
+  shows late; anchored chunks regenerate crisp structure), so the pair
+  isolates the *structure/sharpness* share of drift. Known limitation
+  (pre-stated): the anchor inherits A's low-frequency color drift, so that
+  share is NOT in the target — the alpha*(t) gate on these pairs
+  (``gate_faithful.py`` with ``PAIR_SCHEME=fork``) must clear its bar before
+  any training (see README).
+
+Frame alignment: fork s anchors at decoded frame ``f0 - 5`` and consumes
+HDMap ``[f0-5, f0+160)`` — its 5-frame chunk 0 covers the tail of A's chunk
+``c_s - 1``, so B chunk ``1 + m`` consumes exactly A chunk ``c_s + m``'s
+8-frame window (no conditioning lag).
+
+Per-clip output adds ``fork_starts`` / ``fork_latents`` (B chunks 1..20 per
+fork; chunk 0 is the anchor and is never a counterpart) to the pairs-v2/v3
+format, so v3-format consumers keep working (they ignore the extra keys).
+
+Run from the flashdreams repo root::
+
+    HF_TOKEN=... .venv/bin/python integrations/omnidreams/drift_correction/build_pairs_v4.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import torch
+from _host import build_pipeline, capture_rollout, save_clip
+from build_pairs import _clip_prompt, _list_sample_uuids, _sample_files
+from omnidreams.runner import (
+    DEFAULT_VIDEO_HEIGHT,
+    DEFAULT_VIDEO_WIDTH,
+    _load_video,
+)
+
+from flashdreams.infra.runner_io import (
+    load_first_frame_tensor,
+    write_video_tensor,
+)
+
+## Pair-build configuration
+
+OUT_DIR = Path(
+    os.environ.get(
+        "PAIRS_DIR", "integrations/omnidreams/drift_correction/outputs/pairs_v4"
+    )
+)
+
+N_CLIPS = int(os.environ.get("N_CLIPS", "6"))
+"""Sample clips (first N of the dataset, sorted) — same roster as v2/v3."""
+
+NUM_CHUNK = int(os.environ.get("NUM_CHUNK", "81"))
+"""Drifted rollout A: chunk 0 + 80 chunks = 645 frames = 21.5 s."""
+
+SEG_CHUNKS = int(os.environ.get("SEG_CHUNKS", "20"))
+"""Fork cadence: B_s counterparts cover chunks ``c_s .. c_s + 19`` for
+``c_s in {1, 21, 41, 61}`` (defaults; env overrides are for smoke runs)."""
+
+FORK_STARTS = tuple(range(1, NUM_CHUNK, SEG_CHUNKS))
+
+NOISE_SEED = int(os.environ.get("NOISE_SEED", "5042"))
+
+CORRECTOR_LORA = os.environ.get("CORRECTOR_LORA", "")
+"""Optional checkpoint for DAgger rounds: rollouts (A and forks) run with
+the corrector at the deploy dial (``alpha*(t) x 0.5``)."""
+
+
+def main() -> None:
+    torch.set_grad_enabled(False)
+    dtype = torch.bfloat16
+
+    # Load all inputs before any model work (ffmpeg fails silently once the
+    # process reaches rollout size on this box).
+    n_frames = 5 + (NUM_CHUNK - 1) * 8
+    inputs: list[tuple[str, str, torch.Tensor, torch.Tensor] | None] = []
+    for c, uuid in enumerate(_list_sample_uuids(N_CLIPS)):
+        if (OUT_DIR / f"clip_{c:02d}.pt").exists():
+            inputs.append(None)
+            continue
+        (hdmap_path,), (frame_path,) = _sample_files(uuid)
+        hdmap = _load_video(
+            hdmap_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device="cpu",  # ty: ignore[invalid-argument-type]
+            dtype=dtype,
+        )
+        assert hdmap.shape[0] >= n_frames, (uuid, hdmap.shape)
+        first = load_first_frame_tensor(
+            frame_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device="cpu",  # ty: ignore[invalid-argument-type]
+            dtype=dtype,
+        )[None, :, None]
+        inputs.append((uuid, _clip_prompt(uuid), hdmap[:n_frames][None, None], first))
+        print(f"loaded inputs for clip {c} ({uuid})", flush=True)
+
+    pipe = build_pipeline(with_oneshot_encoders=True)
+    device = pipe.device
+
+    if CORRECTOR_LORA:
+        from _lora import apply_lora, load_lora, unwrap_compiled
+        from eval_rollouts import install_alpha_gate
+
+        transformer = pipe.diffusion_model.transformer
+        network = unwrap_compiled(transformer.network)
+        apply_lora(network)
+        load_lora(network, CORRECTOR_LORA)
+        install_alpha_gate(transformer, network, {"gain": ("gate", 0.5)})
+        print(f"corrector active at gate x 0.5: {CORRECTOR_LORA}", flush=True)
+
+    for c, item in enumerate(inputs):
+        out = OUT_DIR / f"clip_{c:02d}.pt"
+        if item is None:
+            print(f"SKIP clip {c}: {out} exists", flush=True)
+            continue
+        uuid, prompt, hdmap, first = item
+        hdmap = hdmap.to(device)
+        first = first.to(device)
+
+        # Drifted rollout A over the clip's own continuous HDMap.
+        embeddings = pipe.precompute_embeddings(text=[[prompt]], image=first)
+        cache = pipe.initialize_cache_from_embeddings(
+            text_embeddings=embeddings["text_embeddings"],  # ty: ignore[invalid-argument-type]
+            image_embeddings=embeddings["image_embeddings"],  # ty: ignore[invalid-argument-type]
+        )
+        snaps, video = capture_rollout(
+            pipe,
+            cache,
+            hdmap_video=hdmap,
+            num_chunk=NUM_CHUNK,
+            noise_seed=NOISE_SEED + c,
+        )
+
+        # Re-anchored fork counterparts B_s.
+        fork_latents: list[list[torch.Tensor]] = []
+        for s, c_s in enumerate(FORK_STARTS):
+            f0 = 5 + 8 * (c_s - 1)
+            anchor = video[0, 0, f0 - 5].to(device, dtype)[None, None, None]
+            emb_s = pipe.precompute_embeddings(text=[[prompt]], image=anchor)
+            cache_s = pipe.initialize_cache_from_embeddings(
+                text_embeddings=emb_s["text_embeddings"],  # ty: ignore[invalid-argument-type]
+                image_embeddings=emb_s["image_embeddings"],  # ty: ignore[invalid-argument-type]
+            )
+            snaps_s, video_s = capture_rollout(
+                pipe,
+                cache_s,
+                hdmap_video=hdmap[:, :, f0 - 5 : f0 - 5 + 5 + 8 * SEG_CHUNKS],
+                num_chunk=1 + SEG_CHUNKS,
+                noise_seed=NOISE_SEED + 100 * (s + 1) + c,
+            )
+            fork_latents.append([sn.clean_latent.detach().cpu() for sn in snaps_s[1:]])
+            if c == 0:
+                mp4 = OUT_DIR / f"clip_{c:02d}_fork{s}.mp4"
+                OUT_DIR.mkdir(parents=True, exist_ok=True)
+                write_video_tensor(
+                    video_s[0, 0].permute(0, 2, 3, 1), mp4, fps=30, layout="thwc"
+                )
+            del snaps_s, video_s, cache_s
+
+        save_clip(
+            out,
+            snaps=snaps,
+            embeddings=embeddings,
+            meta={
+                "uuid": uuid,
+                "prompt": prompt,
+                "num_chunk": NUM_CHUNK,
+                # lap fields kept for shape-compat; fork scheme ignores them.
+                "lap_chunks": SEG_CHUNKS,
+                "laps": (NUM_CHUNK - 1) // SEG_CHUNKS,
+                "noise_seed": NOISE_SEED + c,
+                "pair_scheme": "fork",
+                "fork_starts": list(FORK_STARTS),
+                "fork_latents": fork_latents,
+            },
+        )
+        mp4 = OUT_DIR / f"clip_{c:02d}_drifted.mp4"
+        write_video_tensor(video[0, 0].permute(0, 2, 3, 1), mp4, fps=30, layout="thwc")
+        print(f"clip {c} ({uuid}): saved {out} + {mp4}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/drift_correction/eval_rollouts.py b/integrations/omnidreams/drift_correction/eval_rollouts.py
new file mode 100644
index 00000000..0cb7a672
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/eval_rollouts.py
@@ -0,0 +1,241 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Closed-loop gain-sweep rollouts: frozen base vs +v1 corrector (Omnidreams).
+
+Deployment-realistic 21.5 s rollouts (non-tiled HDMap conditioning, real
+photo seeds, matched prompts) over the standard dial grid
+``{base, 0.5, alpha*(t) gate, gate x 0.5, 1.0}``. Scenarios: the
+collapse-prone clip-0 scene (stress case), the training-val scene, and two
+fresh sample clips never seen by training. MP4s land under
+``EVAL_OUT/{config}/`` for the HY host-agnostic scoring + sbs pass.
+
+Run from the flashdreams repo root::
+
+    HF_TOKEN=... .venv/bin/python integrations/omnidreams/drift_correction/eval_rollouts.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import torch
+from _host import build_pipeline, capture_rollout
+from _lora import apply_lora, load_lora, set_lora_scale, unwrap_compiled
+from build_pairs import _clip_prompt, _list_sample_uuids, _sample_files
+from omnidreams.runner import (
+    DEFAULT_VIDEO_HEIGHT,
+    DEFAULT_VIDEO_WIDTH,
+    _load_video,
+)
+
+from flashdreams.infra.runner_io import (
+    load_first_frame_tensor,
+    write_video_tensor,
+)
+
+## Sweep configuration
+
+_BASE = Path("integrations/omnidreams/drift_correction")
+
+LORA = os.environ.get("LORA", str(_BASE / "outputs/lora_v1.pt"))
+"""Corrector checkpoint (``train_v1.py`` format: dict with ``lora``)."""
+
+OUT_DIR = Path(os.environ.get("EVAL_OUT", str(_BASE / "outputs/eval_sweep")))
+
+NUM_CHUNK = int(os.environ.get("NUM_CHUNK", "81"))
+"""81 chunks = 645 frames = 21.5 s (matches the gate horizon)."""
+
+SCENARIO_IDS = tuple(
+    int(i) for i in os.environ.get("SCENARIO_IDS", "0,5,6,7").split(",")
+)
+"""Indices into the sorted sample-UUID list: 0 = collapse-prone stress
+scene, 5 = training val scene, 6-7 = never-seen scenarios."""
+
+SEED = int(os.environ.get("SEED", "5042"))
+
+ALPHA_STAR = {
+    float(kv.split(":")[0]): float(kv.split(":")[1])
+    for kv in os.environ.get("ALPHA_STAR", "1000:0.96,803:0.667").split(",")
+}
+"""Unbiased alpha* per warped solver timestep (default: pairs-v2 photoreal
+gate; override as ``ALPHA_STAR=t:a,t:a``); nearest-t lookup, matching the
+HY deploy convention."""
+
+LP_SIGMA = float(os.environ.get("LP_SIGMA", "0"))
+"""When > 0: low-pass correction test (analysis arm 2026-07-24, targets the
+class-a blur without a retrain). At each denoising step the gate configs
+run TWO forwards and blur only the correction delta in latent-space::
+
+    v_rect = v_base + alpha*(t) * gain * GaussianBlur_sigma(v_corr - v_base)
+
+Applied only at solver timesteps (nearest ALPHA_STAR key within 1); the
+``finalize_kv_cache`` context forward (t=128) keeps the plain single-pass
+LoRA scaling. 2x inference cost — test dial only."""
+
+CONFIGS = [c for c in os.environ.get("CONFIGS", "").split(",") if c]
+"""Optional subset of the dial grid (e.g. ``CONFIGS=corrgate050``)."""
+
+LAT_H, LAT_W = DEFAULT_VIDEO_HEIGHT // 16, DEFAULT_VIDEO_WIDTH // 16
+"""Patch-token grid (VAE /8 x patchify /2): 44 x 80 at 704x1280."""
+
+
+def lowpass_tokens(x: torch.Tensor, sigma: float) -> torch.Tensor:
+    """Separable gaussian blur over the token grid's spatial dims.
+
+    ``x`` is patchified ``[..., T*LAT_H*LAT_W, C]``; frames are blurred
+    independently (no temporal mixing).
+    """
+    import math
+
+    import torch.nn.functional as F
+
+    r = max(1, int(math.ceil(3 * sigma)))
+    k = torch.arange(-r, r + 1, device=x.device, dtype=torch.float32)
+    k = torch.exp(-0.5 * (k / sigma) ** 2)
+    k = (k / k.sum()).to(x.dtype)
+    orig = x.shape
+    n_frames = orig[-2] // (LAT_H * LAT_W)
+    assert n_frames * LAT_H * LAT_W == orig[-2], orig
+    y = x.reshape(-1, n_frames, LAT_H, LAT_W, orig[-1])
+    y = y.permute(0, 1, 4, 2, 3).reshape(-1, 1, LAT_H, LAT_W)
+    y = F.pad(y, (r, r, r, r), mode="reflect")
+    y = F.conv2d(y, k.view(1, 1, -1, 1))
+    y = F.conv2d(y, k.view(1, 1, 1, -1))
+    y = y.reshape(-1, n_frames, orig[-1], LAT_H, LAT_W).permute(0, 1, 3, 4, 2)
+    return y.reshape(orig)
+
+
+def install_alpha_gate(transformer, network, mode: dict) -> None:
+    """Wrap ``predict_flow`` so gate configs rescale the LoRA every step."""
+    orig_pf = transformer.predict_flow
+
+    def gated_pf(*args, **kwargs):
+        gain = mode["gain"]
+        if isinstance(gain, tuple):
+            # finalize_kv_cache calls positionally: (noisy_latent, timestep, ...)
+            ts = kwargs.get("timestep", args[1] if len(args) > 1 else None)
+            t = float(ts.reshape(-1).max())
+            t_near, alpha = min(ALPHA_STAR.items(), key=lambda kv: abs(kv[0] - t))
+            if LP_SIGMA > 0 and abs(t - t_near) < 1:
+                # Low-pass test at solver steps: blur the delta only. The
+                # denoise-step forward is cache-idempotent (chunk KV commits
+                # only at finalize), so the double forward is safe.
+                set_lora_scale(network, 0.0)
+                v_base = orig_pf(*args, **kwargs)
+                set_lora_scale(network, 1.0)
+                v_corr = orig_pf(*args, **kwargs)
+                delta = lowpass_tokens(v_corr - v_base, LP_SIGMA)
+                return v_base + alpha * gain[1] * delta
+            set_lora_scale(network, alpha * gain[1])
+        return orig_pf(*args, **kwargs)
+
+    transformer.predict_flow = gated_pf
+
+
+def main() -> None:
+    torch.set_grad_enabled(False)
+    dtype = torch.bfloat16
+    uuids = _list_sample_uuids(max(SCENARIO_IDS) + 1)
+
+    # Load every scenario's inputs BEFORE any model work (ffmpeg decode
+    # fails silently once the process reaches rollout size on this box).
+    inputs = []
+    for sid in SCENARIO_IDS:
+        uuid = uuids[sid]
+        (hdmap_path,), (frame_path,) = _sample_files(uuid)
+        hdmap = _load_video(
+            hdmap_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device="cpu",  # ty: ignore[invalid-argument-type]
+            dtype=dtype,
+        )
+        n_frames = 5 + (NUM_CHUNK - 1) * 8
+        assert hdmap.shape[0] >= n_frames, (sid, hdmap.shape)
+        first = load_first_frame_tensor(
+            frame_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device="cpu",  # ty: ignore[invalid-argument-type]
+            dtype=dtype,
+        )[None, :, None]
+        inputs.append(
+            (sid, uuid, _clip_prompt(uuid), hdmap[:n_frames][None, None], first)
+        )
+        print(f"loaded scenario {sid} ({uuid})", flush=True)
+
+    pipe = build_pipeline(with_oneshot_encoders=True)
+    device = pipe.device
+    transformer = pipe.diffusion_model.transformer
+    network = unwrap_compiled(transformer.network)
+    apply_lora(network)
+    load_lora(network, LORA)
+    mode: dict = {"gain": 0.0}
+    install_alpha_gate(transformer, network, mode)
+
+    configs: dict[str, float | tuple[str, float]] = {
+        "base": 0.0,
+        "corr050": 0.5,
+        "corrgate": ("gate", 1.0),
+        "corrgate050": ("gate", 0.5),
+        "corrgate025": ("gate", 0.25),
+        "corr": 1.0,
+    }
+    if CONFIGS:
+        configs = {k: v for k, v in configs.items() if k in CONFIGS}
+
+    # Cache per-scenario embeddings once (encoders stay loaded).
+    embeddings = {}
+    for sid, uuid, prompt, hdmap, first in inputs:
+        embeddings[sid] = pipe.precompute_embeddings(
+            text=[[prompt]], image=first.to(device)
+        )
+
+    for config, gain in configs.items():
+        for sid, uuid, prompt, hdmap, first in inputs:
+            mp4 = OUT_DIR / config / f"scen{sid}_s{SEED}.mp4"
+            if mp4.exists():
+                print(f"{config}/scen{sid}: exists, skipping", flush=True)
+                continue
+            mode["gain"] = gain
+            if not isinstance(gain, tuple):
+                set_lora_scale(network, gain)
+            emb = embeddings[sid]
+            cache = pipe.initialize_cache_from_embeddings(
+                text_embeddings=emb["text_embeddings"],  # ty: ignore[invalid-argument-type]
+                image_embeddings=emb["image_embeddings"],  # ty: ignore[invalid-argument-type]
+            )
+            print(f"{config}/scen{sid}: rolling {NUM_CHUNK} chunks ...", flush=True)
+            _, video = capture_rollout(
+                pipe,
+                cache,
+                hdmap_video=hdmap.to(device),
+                num_chunk=NUM_CHUNK,
+                noise_seed=SEED,
+            )
+            mp4.parent.mkdir(parents=True, exist_ok=True)
+            write_video_tensor(
+                video[0, 0].permute(0, 2, 3, 1), mp4, fps=30, layout="thwc"
+            )
+    print("SWEEP-DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/drift_correction/gain_predict.py b/integrations/omnidreams/drift_correction/gain_predict.py
new file mode 100644
index 00000000..aa1ce57b
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/gain_predict.py
@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Reliability-shrunk gain prediction: gain*(t) = alpha*(t) x rho(t).
+
+Combines a gate JSON (per-timestep unbiased alpha*, the systematic share of
+the drift error) with TBIN lines from a ``TBIN_EVAL`` trainer run (rho(t) =
+per-timestep val R^2 of the trained corrector) into the predicted deploy
+schedule, and ranks the standard dial arms by RMS distance to it
+(gain-prediction analysis, analysis arm 2026-07-24).
+
+Usage::
+
+    python gain_predict.py <gate.json> "t=1000 R2=+0.36; t=803 R2=+0.39"
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import sys
+
+
+def main() -> None:
+    gate = {
+        float(t): v["alpha_star_unbiased"]
+        for t, v in json.load(open(sys.argv[1]))["per_timestep"].items()
+    }
+    rho = {
+        float(m[0]): float(m[1])
+        for m in re.findall(r"t=(\d+)\s+R2=([+-][\d.]+)", sys.argv[2])
+    }
+    ts = sorted(gate, reverse=True)
+    pred = {t: gate[t] * rho[t] for t in ts}
+    print("predicted gain*(t):", {int(t): round(g, 3) for t, g in pred.items()})
+
+    arms = {
+        "corr (flat 1.0)": {t: 1.0 for t in ts},
+        "corr050 (flat 0.5)": {t: 0.5 for t in ts},
+        "corrgate (a*x1)": dict(gate),
+        "corrgate050 (a*x0.5)": {t: gate[t] * 0.5 for t in ts},
+        "corrgate025 (a*x0.25)": {t: gate[t] * 0.25 for t in ts},
+    }
+
+    def rms(s):
+        return (sum((s[t] - pred[t]) ** 2 for t in ts) / len(ts)) ** 0.5
+
+    for name, s in sorted(arms.items(), key=lambda kv: rms(kv[1])):
+        print(f"  {name:24s} rms-dist {rms(s):.3f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/drift_correction/gate_faithful.py b/integrations/omnidreams/drift_correction/gate_faithful.py
new file mode 100644
index 00000000..b4ffd53e
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/gate_faithful.py
@@ -0,0 +1,312 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Faithful step-0 gate: alpha* on real drifted-vs-clean loop pairs (Omnidreams).
+
+Measures the corrector's premise on ``build_pairs.py`` output: at a late
+chunk of a tiled-HDMap loop rollout, the KV-window history (the host's
+entire history mechanism) is swapped for its lap-1 clean counterparts —
+same HDMap conditioning, same RoPE positions, same context-noise eps,
+clean content — and the *velocity* prediction gap is decomposed over noise
+seeds into systematic bias vs variance. Probes span rollout depths up to
+~21.5 s because this host fights drift with Self-Forcing distillation and
+its residual accumulation shows late; the report resolves ``rel`` by depth.
+
+Kill bars (pre-stated): mean rel gap < 0.01 -> nothing to correct, STOP;
+alpha* below the reference bar -> report before any training.
+
+Run after ``build_pairs.py`` from the flashdreams repo root::
+
+    .venv/bin/python integrations/omnidreams/drift_correction/gate_faithful.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import torch
+from _host import (
+    build_pipeline,
+    finish_probe_chunk,
+    load_clip,
+    predict_v,
+    replay_history,
+    reset_history,
+    start_probe_chunk,
+)
+from torch import Tensor
+
+## Gate configuration
+
+PAIRS_DIR = Path(
+    os.environ.get(
+        "PAIRS_DIR", "integrations/omnidreams/drift_correction/outputs/pairs_v2"
+    )
+)
+"""Clip files from ``build_pairs.py``."""
+
+OUT_PATH = Path(
+    os.environ.get(
+        "GATE_OUT",
+        "integrations/omnidreams/drift_correction/outputs/gate/gate_faithful_v2.json",
+    )
+)
+"""Aggregated per-timestep and per-depth results."""
+
+PROBE_CHUNKS = tuple(
+    int(x) for x in os.environ.get("PROBE_CHUNKS", "24,39,54,69,79").split(",")
+)
+"""Probe chunks spanning ~6.6 s to ~21.4 s of rollout. Lap default: all
+satisfy ``(k - 1) % LAP_CHUNKS >= 3`` so the 3-chunk KV window lies inside
+one lap (no conditioning-teleport contamination), and all sit in laps >= 4
+(convention: drifted side well past the lap-2 clean
+reference). Fork pairs (v4) pass their own set with the window inside one
+fork segment."""
+
+PAIR_SCHEME = os.environ.get("PAIR_SCHEME", "lap")
+"""``lap`` (v1-v3 loop pairs) or ``fork`` (v4 re-anchored fork pairs:
+clean counterpart of chunk ``j`` = fork ``s`` with the largest
+``fork_starts[s] <= j``, chunk ``j - fork_starts[s]`` of its latents)."""
+
+WINDOW_CHUNKS = 3
+"""KV-window span in chunks (``window_size_t=6`` / ``len_t=2``)."""
+
+M_NOISE = 8
+"""Noise seeds per (clip, chunk, timestep) cell."""
+
+CLEAN_LAP = 2
+"""Lap supplying the clean counterpart content (convention 2026-07-22:
+lap 2, keeping the synthetic-seed transition tail of laps 0-1 out of the
+clean reference; the 2026-07-22 GO gate ran with lap 1)."""
+
+
+def lap_aligned(k: int, lap_chunks: int) -> int:
+    """Map chunk ``k`` (>= lap 1) to its lap-:data:`CLEAN_LAP` counterpart."""
+    assert k >= 1, "chunk 0 is image-anchored and never remapped"
+    return 1 + CLEAN_LAP * lap_chunks + (k - 1) % lap_chunks
+
+
+def window_indices(k: int) -> list[int]:
+    """Chunk indices whose KV survives in the window when probing chunk ``k``."""
+    return list(range(max(0, k - WINDOW_CHUNKS), k))
+
+
+def histories(d: dict, k: int) -> tuple[list[Tensor], list[Tensor], list[Tensor]]:
+    """Build the (gen latents, clean latents, hdmaps) prefix lists for chunk ``k``.
+
+    Both branches replay the full prefix ``0 .. k-1`` at original indices
+    (RoPE positions preserved; pre-window chunks are evicted and only feed
+    the roll). The clean branch swaps the surviving window chunks' content
+    for their lap-aligned lap-:data:`CLEAN_LAP` counterparts.
+    """
+    lap_chunks = d["lap_chunks"]
+    gen = [d["latents"][j] for j in range(k)]
+    clean = list(gen)
+    for j in window_indices(k):
+        if j == 0:
+            continue  # image-anchored chunk, never remapped
+        if PAIR_SCHEME == "fork":
+            starts = d["fork_starts"]
+            s = max(i for i, cs in enumerate(starts) if cs <= j)
+            clean[j] = d["fork_latents"][s][j - starts[s]].to(
+                gen[0].device, gen[0].dtype
+            )
+        else:
+            src = lap_aligned(j, lap_chunks)
+            if src != j:
+                clean[j] = d["latents"][src]
+    hdmaps = [d["hdmaps"][j] for j in range(k)]
+    return gen, clean, hdmaps
+
+
+def main() -> None:
+    torch.set_grad_enabled(False)
+    clips = sorted(PAIRS_DIR.glob("clip_*.pt"))
+    assert clips, f"no clips under {PAIRS_DIR}; run build_pairs.py first"
+
+    pipe = build_pipeline(with_oneshot_encoders=False)
+    device = pipe.device
+    dtype = torch.bfloat16
+    transformer = pipe.diffusion_model.transformer
+    scheduler = pipe.diffusion_model.scheduler
+    timesteps = scheduler.denoising_step_list
+    sigmas = scheduler.denoising_sigmas
+    n_steps = timesteps.shape[0]  # ty: ignore[not-subscriptable]
+    ctx_t = torch.tensor(
+        float(pipe.diffusion_model.config.context_noise), device=device, dtype=dtype
+    )
+
+    agg = {
+        i: {"alphas": [], "alphas_ub": [], "rels": [], "rels_x0": []}
+        for i in range(n_steps)
+    }
+    by_depth: dict[int, list[float]] = {}
+    for clip_path in clips:
+        d = load_clip(clip_path, device, dtype)
+        emb = d["embeddings"]
+        cache = pipe.initialize_cache_from_embeddings(
+            text_embeddings=emb["text_embeddings"],
+            image_embeddings=emb["image_embeddings"],
+        )
+        tc = cache.transformer_cache
+
+        for k in PROBE_CHUNKS:
+            if k >= d["num_chunk"]:
+                continue
+            gen_hist, clean_hist, hdmaps = histories(d, k)
+            hdmap_k = d["hdmaps"][k]
+            x0 = d["latents"][k]
+
+            z_ts: dict[int, list[Tensor]] = {}
+            for t_idx in range(n_steps):
+                sig = sigmas[t_idx].to(dtype)  # ty: ignore[not-subscriptable]
+                z_ts[t_idx] = []
+                for m in range(M_NOISE):
+                    g = torch.Generator(device=device).manual_seed(
+                        900_000 + 10_000 * k + 100 * t_idx + m
+                    )
+                    eps = torch.randn(x0.shape, device=device, dtype=dtype, generator=g)
+                    z_ts[t_idx].append((1 - sig) * x0 + sig * eps)
+
+            preds: dict[str, dict[int, list[Tensor]]] = {}
+            for name, hist in (("gen", gen_hist), ("clean", clean_hist)):
+                reset_history(tc)
+                replay_history(
+                    transformer,  # ty: ignore[invalid-argument-type]
+                    tc,
+                    latents=hist,
+                    hdmaps=hdmaps,
+                    context_timestep=ctx_t,
+                    add_noise=scheduler.add_noise,
+                )
+                start_probe_chunk(tc, ar_idx=k)
+                preds[name] = {
+                    t_idx: [
+                        predict_v(
+                            transformer,  # ty: ignore[invalid-argument-type]
+                            tc,
+                            hdmap=hdmap_k,
+                            z_t=z,
+                            timestep=timesteps[t_idx].to(device=device, dtype=dtype),  # ty: ignore[not-subscriptable]
+                        )
+                        for z in z_list
+                    ]
+                    for t_idx, z_list in z_ts.items()
+                }
+                finish_probe_chunk(tc, ar_idx=k)
+
+            line = [f"{clip_path.stem} k={k:2d} ({(5 + 8 * k) / 30:5.1f}s)"]
+            for t_idx in range(n_steps):
+                deltas = torch.stack(
+                    [a - b for a, b in zip(preds["gen"][t_idx], preds["clean"][t_idx])]
+                )
+                m = deltas.shape[0]
+                bias = deltas.mean(dim=0)
+                bias_sq = bias.square().sum().item()
+                sq_dev = (deltas - bias).square().sum(dim=tuple(range(1, deltas.ndim)))
+                var = sq_dev.mean().item()
+                var_ub = sq_dev.sum().item() / (m - 1)
+                bias_sq_ub = max(0.0, bias_sq - var_ub / m)
+                gen_norm = (
+                    torch.stack(preds["gen"][t_idx]).flatten(1).norm(dim=1).mean()
+                )
+                rel = (deltas.flatten(1).norm(dim=1).mean() / (gen_norm + 1e-12)).item()
+                # x0-space view: x0 = z_t - sigma * v with shared (z_t, sigma),
+                # so delta_x0 = -sigma * delta_v and only the denominator changes.
+                sig = float(sigmas[t_idx])  # ty: ignore[not-subscriptable]
+                x0_norm = (
+                    torch.stack(
+                        [
+                            z.float() - sig * v
+                            for z, v in zip(z_ts[t_idx], preds["gen"][t_idx])
+                        ]
+                    )
+                    .flatten(1)
+                    .norm(dim=1)
+                    .mean()
+                )
+                rel_x0 = (
+                    sig * deltas.flatten(1).norm(dim=1).mean() / (x0_norm + 1e-12)
+                ).item()
+                a = bias_sq / (bias_sq + var + 1e-12)
+                a_ub = bias_sq_ub / (bias_sq_ub + var_ub + 1e-12)
+                agg[t_idx]["alphas"].append(a)
+                agg[t_idx]["alphas_ub"].append(a_ub)
+                agg[t_idx]["rels"].append(rel)
+                agg[t_idx]["rels_x0"].append(rel_x0)
+                by_depth.setdefault(k, []).append(rel)
+                line.append(
+                    f"t={int(timesteps[t_idx]):4d} a*={a:.3f}/{a_ub:.3f}"  # ty: ignore[not-subscriptable]
+                    f" rel_v={rel:.4f} rel_x0={rel_x0:.4f}"
+                )
+            print(" | ".join(line), flush=True)
+        del cache, tc
+
+    summary = {
+        "per_timestep": {
+            str(int(timesteps[i])): {  # ty: ignore[not-subscriptable]
+                "alpha_star": sum(v["alphas"]) / len(v["alphas"]),
+                "alpha_star_unbiased": sum(v["alphas_ub"]) / len(v["alphas_ub"]),
+                "rel_v": sum(v["rels"]) / len(v["rels"]),
+                "rel_x0": sum(v["rels_x0"]) / len(v["rels_x0"]),
+                "cells": len(v["alphas"]),
+            }
+            for i, v in agg.items()
+            if v["alphas"]
+        },
+        "rel_v_by_depth": {
+            str(k): {"seconds": (5 + 8 * k) / 30, "rel_v": sum(r) / len(r)}
+            for k, r in sorted(by_depth.items())
+        },
+    }
+    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
+    OUT_PATH.write_text(json.dumps(summary, indent=2))
+
+    print(
+        "\n========== Omnidreams FAITHFUL gate (real drift pairs, v-space) =========="
+    )
+    for t, v in summary["per_timestep"].items():
+        print(
+            f"t={t:>4s}: alpha* {v['alpha_star']:.3f} (unbiased {v['alpha_star_unbiased']:.3f})"
+            f" | rel_v {v['rel_v']:.4f} rel_x0 {v['rel_x0']:.4f} | {v['cells']} cells"
+        )
+    for k, v in summary["rel_v_by_depth"].items():
+        print(f"depth k={k:>3s} ({v['seconds']:5.1f}s): rel_v {v['rel_v']:.4f}")
+    all_ub = [a for v in agg.values() for a in v["alphas_ub"]]
+    mean_rel = sum(r for v in agg.values() for r in v["rels"]) / max(
+        1, sum(len(v["rels"]) for v in agg.values())
+    )
+    frac = sum(a >= 0.7 for a in all_ub) / len(all_ub)
+    print(
+        f"cells with unbiased alpha* >= 0.7: {frac:.0%} | mean rel_v gap {mean_rel:.4f}"
+    )
+    if mean_rel < 0.01:
+        print("RESULT: drift gap ~zero -> nothing to correct. STOP.")
+    elif frac >= 0.7:
+        print("RESULT: real drift gap is systematic. Report headroom before training.")
+    else:
+        print(
+            "RESULT: below the reference bar -- report before committing to training."
+        )
+    print(f"saved {OUT_PATH}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/drift_correction/train_v1.py b/integrations/omnidreams/drift_correction/train_v1.py
new file mode 100644
index 00000000..8ae73486
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/train_v1.py
@@ -0,0 +1,480 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""v1 corrector trainer: counterfactual clean-history teacher in velocity space.
+
+Trains the LoRA corrector r_phi on the frozen distilled Omnidreams model so
+its velocity prediction under drifted history matches the frozen model's
+prediction under the lap-aligned clean history at the same ``z_t``::
+
+    L = ||v_{theta+LoRA}(z_t, h_gen, t) - v_theta(z_t, h_clean, t)||^2
+        / ||v_theta(z_t, h_clean, t) - v_theta(z_t, h_gen, t)||^2
+
+The drift-gap denominator is required (raw MSE diverges; reference
+finding); ``R^2 = 1 - L``. Both passes share the exact ``z_t`` re-noised
+from the rollout's own x0 (native anchoring). Timesteps are sampled from
+the 2 distillation steps with probability proportional to the measured
+``alpha*(t)`` (pairs-v2 gate: 0.96 @ t=1000, 0.667 @ t=803); fully-collapsed
+cells (``rel_v > 0.8``) are excluded at draw time per convention.
+
+Host adaptations vs the HY reference (``hy_worldplay`` ``train_v1.py``):
+velocity-space targets; history rebuilt per sample by a truncated
+forged-index replay (:data:`REPLAY_CHUNKS` warmup chunks -- deep-layer KV
+entangles history, so a bare window replay is NOT equivalent; checked
+numerically at startup); functional self-attention toggle
+(:mod:`_train_attn`) for the grad-carrying probe. Conventions (2026-07-22): clean reference = lap 2, training cells in laps >= 4,
+checkpoints every <= 200 steps with RESUME.
+
+Run from the repo root (resumable: re-run the same command)::
+
+    STEPS=1500 .venv/bin/python integrations/omnidreams/drift_correction/train_v1.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import numpy as np
+import torch
+from _host import build_pipeline, load_clip, reset_history, swap_text_kv
+from _lora import apply_lora, lora_parameters, set_lora_scale
+from _train_attn import functional_attention, patch_functional_attention
+from torch import Tensor
+
+## Training configuration
+
+BASE = Path("integrations/omnidreams/drift_correction")
+
+PAIRS_DIR = Path(os.environ.get("PAIRS_DIR", str(BASE / "outputs/pairs_v2")))
+"""Clip files from ``build_pairs.py`` (v2: real-photo-seeded rollouts)."""
+
+CKPT = Path(os.environ.get("CKPT", str(BASE / "outputs/lora_v1.pt")))
+"""Output checkpoint (LoRA + optimizer + step); saved every ``SAVE_EVERY``."""
+
+STEPS = int(os.environ.get("STEPS", "1500"))
+LR = float(os.environ.get("LR", "5e-4"))
+WARMUP = 60
+"""Linear LR warmup steps; required for stability (reference finding)."""
+
+GRAD_CLIP = 1.0
+RANK = 16
+EVAL_EVERY = 100
+SAVE_EVERY = 200
+"""Partial checkpoints every <= 200 steps (box wedges)."""
+
+CLEAN_LAP = 2
+"""Lap 2 supplies the clean reference (keeps the
+synthetic-seed transition tail of laps 0-1 out of it)."""
+
+MIN_LAP = 4
+"""Training cells live in laps >= 4 (drifted side well past the clean lap)."""
+
+REPLAY_CHUNKS = 15
+"""History chunks replayed before a probe. The KV cache holds only 3
+chunks, but deep-layer K/V entangle earlier history (each replayed chunk's
+projections depend on what that forward attended to), so a bare window
+replay diverges from the rollout state (rel 0.034 measured). The error
+decays with warmup depth; on the photoreal pairs-v2 rollouts the floor is
+higher than on the render-regime pairs (drift couples longer-range), and
+15 chunks reaches rel ~0.009-0.014 = 3-5% of the drift-gap signal
+(measured on this host 2026-07-22)."""
+
+ALPHA_STAR = tuple(
+    float(x) for x in os.environ.get("ALPHA_STAR", "0.96,0.667").split(",")
+)
+"""Measured unbiased alpha* per solver timestep (t=1000, t=803); used as
+the timestep sampling weights. Default = the pairs-v2 (photoreal) gate;
+override via env when a pair set gets its own gate run."""
+
+REL_V_EXCLUDE = 0.8
+"""Drop fully-collapsed cells. A sample whose
+velocity-space drift gap ``||v_clean - v_base|| / ||v_base||`` exceeds this
+is a degenerate state (e.g. the clip-0 late-horizon collapse) and is
+rejected at draw time (the probes needed for the check are computed anyway;
+exclusions are logged per clip)."""
+
+PAIR_SCHEME = os.environ.get("PAIR_SCHEME", "lap")
+"""``lap`` (v1-v3 loop pairs) or ``fork`` (pairs-v4, ``build_pairs_v4.py``):
+the clean counterpart of chunk ``j`` comes from the re-anchored fork
+covering ``j`` instead of a lap-aligned revisit. Training cells then live
+past the first fork segment (drifted side deep, counterpart anchored)."""
+
+UW = os.environ.get("UW", "0") == "1"
+"""Uncertainty-weighted loss (analysis arm 2026-07-24): estimate the
+per-token systematic share of the drift residual from ``UW_DRAWS`` noise
+draws at the same cell — ``w = bias^2 / (bias^2 + var)``, the per-token
+alpha* — and weight both loss numerator and denominator by ``w`` so the
+LoRA never learns to chase unpredictable content (foliage). Config flag:
+``UW=0`` keeps the unweighted objective as the ablation arm."""
+
+UW_DRAWS = 2
+"""Noise draws for the weight estimate (adds 2 no-grad teacher forwards)."""
+
+N_VAL_CLIPS = 1
+SEED = int(os.environ.get("SEED", "0"))
+
+TBIN_EVAL = int(os.environ.get("TBIN_EVAL", "0"))
+"""When > 0: no training — load CKPT, evaluate val R^2 binned per solver
+timestep (TBIN_EVAL cells per bin) and print ``TBIN t=<t> R2=<r>`` lines
+(the reliability factor rho(t) for the gain-prediction analysis, analysis arm 2026-07-24), then exit."""
+
+
+def lap_aligned(c: int, lap_chunks: int) -> int:
+    """Map chunk ``c`` (>= lap 1) to its lap-:data:`CLEAN_LAP` counterpart."""
+    assert c >= 1, "chunk 0 is image-anchored and never remapped"
+    return 1 + CLEAN_LAP * lap_chunks + (c - 1) % lap_chunks
+
+
+def training_ks(num_chunk: int, lap_chunks: int, laps: int) -> list[int]:
+    """All probe chunks in laps >= MIN_LAP.
+
+    Every lap position is usable: the clean swap maps each replayed chunk
+    to its lap-2 counterpart *by position*, so both branches share the lap
+    cycle's boundary structure (including the conditioning teleport) and
+    differ only in content cleanliness.
+    """
+    return [k for k in range(1 + MIN_LAP * lap_chunks, num_chunk)]
+
+
+def training_ks_fork(num_chunk: int, fork_starts: list[int]) -> list[int]:
+    """Fork-scheme cells: past the first segment (drifted side is deep
+    while the counterpart is at most SEG_CHUNKS from its anchor)."""
+    return [k for k in range(int(fork_starts[1]), num_chunk)]
+
+
+def clean_chunk(d: dict, gen: list[Tensor], j: int, lap_chunks: int) -> Tensor:
+    """Clean counterpart content for replayed chunk ``j`` (>= 1)."""
+    if PAIR_SCHEME == "fork":
+        starts = d["fork_starts"]
+        s = max(i for i, cs in enumerate(starts) if cs <= j)
+        return d["fork_latents"][s][j - int(starts[s])]
+    return gen[lap_aligned(j, lap_chunks)]
+
+
+def main() -> None:
+    clips = sorted(PAIRS_DIR.glob("clip_*.pt"))
+    assert clips, f"no clips under {PAIRS_DIR}; run build_pairs.py first"
+    rng = np.random.default_rng(SEED)
+
+    pipe = build_pipeline(with_oneshot_encoders=False)
+    device = pipe.device
+    dtype = torch.bfloat16
+    transformer = pipe.diffusion_model.transformer
+    scheduler = pipe.diffusion_model.scheduler
+    timesteps = scheduler.denoising_step_list
+    sigmas = scheduler.denoising_sigmas
+    n_steps = timesteps.shape[0]  # ty: ignore[not-subscriptable]
+    t_probs = np.array(ALPHA_STAR[:n_steps]) / sum(ALPHA_STAR[:n_steps])
+    ctx_t = torch.tensor(
+        float(pipe.diffusion_model.config.context_noise), device=device, dtype=dtype
+    )
+
+    datas = [load_clip(p, "cpu", dtype) for p in clips]
+    # Per-clip lap geometry: pairs v3 mixes lap lengths across clips (the
+    # repeat-prior fix), so nothing may assume a shared lap_chunks.
+    if PAIR_SCHEME == "fork":
+        ks_by_clip = [
+            training_ks_fork(int(d["num_chunk"]), d["fork_starts"]) for d in datas
+        ]
+    else:
+        ks_by_clip = [
+            training_ks(int(d["num_chunk"]), int(d["lap_chunks"]), int(d["laps"]))
+            for d in datas
+        ]
+    train_ids = list(range(len(datas) - N_VAL_CLIPS))
+    val_ids = list(range(len(datas) - N_VAL_CLIPS, len(datas)))
+    print(
+        f"{len(datas)} clips ({len(val_ids)} val) | lap_chunks "
+        f"{[int(d['lap_chunks']) for d in datas]} | cells/clip "
+        f"{[len(ks) for ks in ks_by_clip]} (laps >= {MIN_LAP}, "
+        f"clean lap {CLEAN_LAP})",
+        flush=True,
+    )
+
+    # One long-lived cache (rope, masks, rolling self-attn buffers) serves
+    # every clip; probes never touch chunk 0 (image), and the per-clip
+    # prompt is handled by swapping the cross-attn text KV on clip switch.
+    emb = datas[0]["embeddings"]
+    cache = pipe.initialize_cache_from_embeddings(
+        text_embeddings=emb["text_embeddings"],
+        image_embeddings=emb["image_embeddings"],
+    )
+    tc = cache.transformer_cache
+    _text_clip = [0]
+
+    def use_clip_text(c: int) -> None:
+        """Swap the cross-attn text KV when the sampled clip changes."""
+        if _text_clip[0] != c:
+            swap_text_kv(network, tc, datas[c]["embeddings"]["text_embeddings"])
+            _text_clip[0] = c
+
+    network = transformer.network
+    wrapped = apply_lora(network, rank=RANK)  # ty: ignore[invalid-argument-type]
+    params = lora_parameters(network)  # ty: ignore[invalid-argument-type]
+    print(
+        f"LoRA on {len(wrapped)} projections | "
+        f"{sum(p.numel() for p in params) / 1e6:.2f}M params",
+        flush=True,
+    )
+    patch_functional_attention()
+    opt = torch.optim.AdamW(params, lr=LR)
+
+    start_step = 0
+    if CKPT.exists():
+        state = torch.load(CKPT, map_location="cpu", weights_only=False)
+        for i, p in enumerate(params):
+            p.data.copy_(state["lora"][i].to(p.device, p.dtype))
+        opt.load_state_dict(state["opt"])
+        start_step = state["step"]
+        rng = np.random.default_rng(state["rng_seed"])
+        print(f"RESUMED from {CKPT} at step {start_step}", flush=True)
+
+    def save(step: int) -> None:
+        tmp = CKPT.with_suffix(".tmp")
+        torch.save(
+            {
+                "lora": {i: p.detach().cpu() for i, p in enumerate(params)},
+                "opt": opt.state_dict(),
+                "step": step,
+                "rng_seed": SEED + step,  # fresh stream on resume
+            },
+            tmp,
+        )
+        tmp.replace(CKPT)
+
+    _device_clips: dict[int, dict] = {}
+
+    def to_device(c: int) -> dict:
+        """Device-resident view of clip ``c`` (memoized; ~0.3 GB per clip)."""
+        if c not in _device_clips:
+
+            def move(v):
+                if isinstance(v, list):
+                    return [move(x) for x in v]
+                return v.to(device, dtype) if isinstance(v, Tensor) else v
+
+            _device_clips[c] = {
+                key: move(v)
+                for key, v in datas[c].items()
+                if key
+                in (
+                    "latents",
+                    "hdmaps",
+                    "lap_chunks",
+                    "num_chunk",
+                    "fork_starts",
+                    "fork_latents",
+                )
+            }
+        return _device_clips[c]
+
+    def replay_window(latents: list[Tensor], hdmaps: list[Tensor], k: int) -> None:
+        """Rebuild the KV state from chunks ``k - REPLAY_CHUNKS .. k - 1``.
+
+        Forged-index truncated replay at original absolute indices (RoPE
+        preserved) on a reset cache whose ``_prev_chunk_idx`` is set so the
+        ``BlockKVCache`` contiguity assert passes. Context-noise eps is
+        seeded per absolute index, shared across branches.
+        """
+        start = max(0, k - REPLAY_CHUNKS)
+        reset_history(tc)
+        for bc in tc.network_cache.block_caches:
+            bc.self_attn._prev_chunk_idx = start - 1
+        for j in range(start, k):
+            g = torch.Generator(device=device).manual_seed(77_000 + j)
+            noisy = scheduler.add_noise(latents[j], ctx_t, rng=g)
+            tc.start(j)
+            transformer.finalize_kv_cache(
+                noisy_latent=noisy, timestep=ctx_t, cache=tc, input=hdmaps[j]
+            )
+            tc.finalize(j)
+
+    def predict_v(z_t: Tensor, t_idx: int, hdmap: Tensor) -> Tensor:
+        t = timesteps[t_idx].to(device=device, dtype=dtype)  # ty: ignore[not-subscriptable]
+        with functional_attention():
+            flow = transformer.predict_flow(
+                noisy_latent=z_t, timestep=t, cache=tc, input=hdmap
+            )
+        return flow.float()
+
+    excl = {c: {"tested": 0, "excluded": 0} for c in range(len(datas))}
+
+    def sample_losses(
+        c: int, grad: bool, rng_: np.random.Generator, t_forced: int | None = None
+    ) -> tuple[Tensor, Tensor] | None:
+        """One drift-pair sample -> (normalized v-space loss, r_target sq-norm).
+
+        Returns ``None`` when the cell is degenerate (fully-collapsed state,
+        ``rel_v > REL_V_EXCLUDE``); callers redraw.
+        """
+        use_clip_text(c)
+        d = to_device(c)
+        lap_chunks = int(d["lap_chunks"])
+        k = int(rng_.choice(ks_by_clip[c]))
+        t_idx = int(rng_.choice(n_steps, p=t_probs)) if t_forced is None else t_forced
+        gen = d["latents"]
+        clean = list(gen)
+        for j in range(max(1, k - REPLAY_CHUNKS), k):
+            clean[j] = clean_chunk(d, gen, j, lap_chunks)
+        x0 = gen[k]
+        sig = sigmas[t_idx].to(dtype)  # ty: ignore[not-subscriptable]
+        n_draws = UW_DRAWS if UW else 1
+        z_ts = []
+        for _ in range(n_draws):
+            g = torch.Generator(device=device).manual_seed(int(rng_.integers(2**31)))
+            z_ts.append(
+                (1 - sig) * x0
+                + sig * torch.randn(x0.shape, device=device, dtype=dtype, generator=g)
+            )
+
+        with torch.no_grad():
+            set_lora_scale(network, 0.0)  # ty: ignore[invalid-argument-type]
+            replay_window(clean, d["hdmaps"], k)
+            tc.start(k)
+            v_cleans = [predict_v(z, t_idx, d["hdmaps"][k]) for z in z_ts]
+            tc.finalize(k)
+            replay_window(gen, d["hdmaps"], k)
+            tc.start(k)
+            v_bases = [predict_v(z, t_idx, d["hdmaps"][k]) for z in z_ts]
+            tc.finalize(k)
+        v_clean, v_base = v_cleans[0], v_bases[0]
+        r_target_sq = (v_clean - v_base).square().sum()
+
+        w = None
+        if UW:
+            rs = torch.stack([vc - vb for vc, vb in zip(v_cleans, v_bases)])
+            mean_r = rs.mean(0)
+            var = rs.var(0, unbiased=True)
+            bias2 = (mean_r.square() - var / n_draws).clamp_min(0.0)
+            w = (bias2 / (bias2 + var + 1e-12)).detach()
+
+        excl[c]["tested"] += 1
+        rel_v = (r_target_sq.sqrt() / (v_base.norm() + 1e-9)).item()
+        if rel_v > REL_V_EXCLUDE:
+            excl[c]["excluded"] += 1
+            return None
+
+        set_lora_scale(network, 1.0)  # ty: ignore[invalid-argument-type]
+        with torch.no_grad():
+            replay_window(gen, d["hdmaps"], k)  # LoRA-scaled replay, no grad
+        tc.start(k)
+        with torch.enable_grad() if grad else torch.no_grad():
+            v_corr = predict_v(z_ts[0], t_idx, d["hdmaps"][k])
+            err = (v_corr - v_clean).square()
+            gap = (v_clean - v_base).square()
+            if w is not None:
+                err, gap = err * w, gap * w
+            loss = err.sum() / (gap.sum() + 1e-8)
+        tc.finalize(k)
+        return loss, r_target_sq
+
+    def draw_losses(
+        ids: list[int],
+        grad: bool,
+        rng_: np.random.Generator,
+        t_forced: int | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        """Redraw until a non-degenerate cell is sampled."""
+        while True:
+            out = sample_losses(int(rng_.choice(ids)), grad, rng_, t_forced)
+            if out is not None:
+                return out
+
+    @torch.no_grad()
+    def val_r2(n: int = 12) -> float:
+        vrng = np.random.default_rng(1234)  # fixed cells/noise across evals
+        s = 0.0
+        for _ in range(n):
+            loss, _ = draw_losses(val_ids, False, vrng)
+            s += loss.item()
+        return 1 - s / n
+
+    def replay_equivalence_check() -> None:
+        """Assert the forged-index window replay matches a full-prefix replay."""
+        d = to_device(0)
+        k = ks_by_clip[0][0]
+        g = torch.Generator(device=device).manual_seed(1)
+        z_t = (1 - sigmas[0].to(dtype)) * d["latents"][k] + sigmas[0].to(  # ty: ignore[not-subscriptable]
+            dtype
+        ) * torch.randn(d["latents"][k].shape, device=device, dtype=dtype, generator=g)
+        set_lora_scale(network, 0.0)  # ty: ignore[invalid-argument-type]
+        with torch.no_grad():
+            replay_window(d["latents"], d["hdmaps"], k)
+            tc.start(k)
+            v_win = predict_v(z_t, 0, d["hdmaps"][k])
+            tc.finalize(k)
+            reset_history(tc)
+            for j in range(k):
+                gg = torch.Generator(device=device).manual_seed(77_000 + j)
+                noisy = scheduler.add_noise(d["latents"][j], ctx_t, rng=gg)
+                tc.start(j)
+                transformer.finalize_kv_cache(
+                    noisy_latent=noisy, timestep=ctx_t, cache=tc, input=d["hdmaps"][j]
+                )
+                tc.finalize(j)
+            tc.start(k)
+            v_full = predict_v(z_t, 0, d["hdmaps"][k])
+            tc.finalize(k)
+        rel = ((v_win - v_full).norm() / (v_full.norm() + 1e-9)).item()
+        print(f"replay equivalence: rel diff {rel:.2e}", flush=True)
+        assert rel < 2e-2, "truncated replay too far from full-prefix replay"
+
+    replay_equivalence_check()
+
+    if TBIN_EVAL:
+        # rho(t) for the gain-prediction analysis: val R^2 per solver step.
+        for t_idx in range(n_steps):
+            vrng = np.random.default_rng(1234)
+            s = 0.0
+            for _ in range(TBIN_EVAL):
+                loss, _ = draw_losses(val_ids, False, vrng, t_idx)
+                s += loss.item()
+            r2 = 1 - s / TBIN_EVAL
+            print(f"TBIN t={int(timesteps[t_idx])} R2={r2:+.4f}", flush=True)  # ty: ignore[not-subscriptable]
+        print("TBIN-DONE", flush=True)
+        return
+
+    torch.set_grad_enabled(True)
+    for step in range(start_step + 1, STEPS + 1):
+        for pg in opt.param_groups:
+            pg["lr"] = LR * min(1.0, step / WARMUP)
+        opt.zero_grad()
+        loss, rt_sq = draw_losses(train_ids, True, rng)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(params, GRAD_CLIP)
+        opt.step()
+        if step % EVAL_EVERY == 0 or step == 1:
+            excl_str = " ".join(
+                f"c{c}:{v['excluded']}/{v['tested']}" for c, v in excl.items()
+            )
+            print(
+                f"step {step:5d} | loss {loss.item():.4f}"
+                f" (train R^2 {1 - loss.item():+.3f})"
+                f" | val R^2 {val_r2():+.3f} | |r_t|^2 {rt_sq.item():.1f}"
+                f" | excluded {excl_str}",
+                flush=True,
+            )
+        if step % SAVE_EVERY == 0 or step == STEPS:
+            save(step)
+
+    print(f"TRAIN-V1-DONE | final val R^2 {val_r2(24):+.3f} | saved {CKPT}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/drift_correction/train_v2.py b/integrations/omnidreams/drift_correction/train_v2.py
new file mode 100644
index 00000000..33ad194d
--- /dev/null
+++ b/integrations/omnidreams/drift_correction/train_v2.py
@@ -0,0 +1,470 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""v2 corrector trainer: DAgger pool + drift-contraction (Omnidreams).
+
+The paper-final recipe on this host: the v1 counterfactual loss over the
+aggregated pair pool (round-0 pairs-v2 + corrected-rollout DAgger round-1),
+plus a drift-contraction term — the corrected chunk-``k`` prediction is
+*committed* into chunk ``k+1``'s KV history WITH gradient and chunk
+``k+1``'s gap to its clean teacher is penalized (weight ``CW_LOSS``)::
+
+    x0_corr(k)   = z_t - sigma * v_corr(k)                     (grad)
+    commit       : chunk-k context forward on noisy(x0_corr)   (grad KV,
+                   recorded functionally; buffer gets a no-grad twin)
+    L_con        = ||v_inject(k+1) - v_clean(k+1)||^2
+                   / ||v_clean(k+1) - v_base(k+1)||^2
+    L            = L_dag + CW_LOSS * L_con
+
+Host mechanics: the commit's grad-carrying per-block K/V are captured by
+``record_kv`` and swapped over their numerically identical buffered twins
+by ``inject_kv`` during the chunk-``k+1`` probe (``_train_attn``).
+
+Run from the repo root (resumable: re-run the same command)::
+
+    STEPS=1000 .venv/bin/python integrations/omnidreams/drift_correction/train_v2.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import numpy as np
+import torch
+from _host import build_pipeline, load_clip, reset_history, swap_text_kv
+from _lora import apply_lora, lora_parameters, set_lora_scale
+from _train_attn import (
+    functional_attention,
+    inject_kv,
+    patch_functional_attention,
+    record_kv,
+)
+from torch import Tensor
+
+## Training configuration
+
+BASE = Path("integrations/omnidreams/drift_correction")
+
+POOLS = [
+    BASE / "outputs" / p
+    for p in os.environ.get("POOLS", "pairs_v2,pairs_dagger1").split(",")
+]
+"""Pair pools, aggregated (round-0 + corrected-rollout DAgger rounds)."""
+
+INIT = os.environ.get("INIT", str(BASE / "outputs/lora_v1.pt"))
+"""Warm-start checkpoint (v1)."""
+
+CKPT = Path(os.environ.get("CKPT", str(BASE / "outputs/lora_v2.pt")))
+"""Output checkpoint (LoRA + optimizer + step); saved every ``SAVE_EVERY``."""
+
+STEPS = int(os.environ.get("STEPS", "1000"))
+LR = float(os.environ.get("LR", "2e-4"))
+WARMUP = 40
+GRAD_CLIP = 1.0
+RANK = 16
+EVAL_EVERY = 100
+SAVE_EVERY = 200
+CW_LOSS = float(os.environ.get("CW_LOSS", "0.5"))
+"""Drift-contraction weight (paper-final value)."""
+
+SNAP_EVERY = int(os.environ.get("SNAP_EVERY", "0"))
+"""When > 0: also keep step-tagged snapshot copies (``<ckpt>_stepN.pt``)
+every SNAP_EVERY steps plus a running val-peak snapshot
+(``<ckpt>_valpeak.pt``, refreshed whenever val dag-R^2 improves at an eval
+point), so the sweep can compare checkpoints without retraining (the
+checkpoint-selection comparison)."""
+
+CLEAN_LAP = 2
+MIN_LAP = 4
+REPLAY_CHUNKS = 15
+ALPHA_STAR = tuple(
+    float(x) for x in os.environ.get("ALPHA_STAR", "0.96,0.667").split(",")
+)
+REL_V_EXCLUDE = 0.8
+N_VAL_CLIPS = 1
+"""Same data conventions as ``train_v1`` (same conventions as train_v1); the
+val clip is the last clip of the FIRST pool (round-0 held-out scene)."""
+
+SEED = int(os.environ.get("SEED", "0"))
+
+PAIR_SCHEME = os.environ.get("PAIR_SCHEME", "lap")
+"""``lap`` (v1-v3 loop pairs) or ``fork`` (pairs-v4 re-anchored forks);
+see ``train_v1.py`` / ``build_pairs_v4.py``."""
+
+UW = os.environ.get("UW", "0") == "1"
+"""Uncertainty-weighted DAgger term (per-token alpha* weights from
+``UW_DRAWS`` noise draws; see ``train_v1.py``). The contraction term stays
+unweighted: it penalizes accumulation of whatever residual remains, while
+the weights' job is to keep unpredictable content out of the *target* —
+recorded design choice 2026-07-24. ``UW=0`` = ablation arm."""
+
+UW_DRAWS = 2
+
+
+def lap_aligned(c: int, lap_chunks: int) -> int:
+    """Map chunk ``c`` (>= lap 1) to its lap-:data:`CLEAN_LAP` counterpart."""
+    assert c >= 1, "chunk 0 is image-anchored and never remapped"
+    return 1 + CLEAN_LAP * lap_chunks + (c - 1) % lap_chunks
+
+
+def training_ks(num_chunk: int, lap_chunks: int, laps: int) -> list[int]:
+    """Probe chunks in laps >= MIN_LAP with a successor chunk available."""
+    return [k for k in range(1 + MIN_LAP * lap_chunks, num_chunk - 1)]
+
+
+def training_ks_fork(num_chunk: int, fork_starts: list[int]) -> list[int]:
+    """Fork-scheme cells past the first segment, successor available."""
+    return [k for k in range(int(fork_starts[1]), num_chunk - 1)]
+
+
+def clean_chunk(d: dict, gen: list[Tensor], j: int, lap_chunks: int) -> Tensor:
+    """Clean counterpart content for replayed chunk ``j`` (>= 1)."""
+    if PAIR_SCHEME == "fork":
+        starts = d["fork_starts"]
+        s = max(i for i, cs in enumerate(starts) if cs <= j)
+        return d["fork_latents"][s][j - int(starts[s])]
+    return gen[lap_aligned(j, lap_chunks)]
+
+
+def main() -> None:
+    rng = np.random.default_rng(SEED)
+    clip_paths: list[Path] = []
+    val_paths: list[Path] = []
+    for i, pool in enumerate(POOLS):
+        clips = sorted(pool.glob("clip_*.pt"))
+        assert clips, f"no clips under {pool}"
+        if i == 0:
+            val_paths = clips[-N_VAL_CLIPS:]
+            clips = clips[:-N_VAL_CLIPS]
+        clip_paths += clips
+    datas = [load_clip(p, "cpu", torch.bfloat16) for p in clip_paths + val_paths]
+    train_ids = list(range(len(clip_paths)))
+    val_ids = list(range(len(clip_paths), len(datas)))
+
+    pipe = build_pipeline(with_oneshot_encoders=False)
+    device = pipe.device
+    dtype = torch.bfloat16
+    transformer = pipe.diffusion_model.transformer
+    scheduler = pipe.diffusion_model.scheduler
+    timesteps = scheduler.denoising_step_list
+    sigmas = scheduler.denoising_sigmas
+    n_steps = timesteps.shape[0]  # ty: ignore[not-subscriptable]
+    t_probs = np.array(ALPHA_STAR[:n_steps]) / sum(ALPHA_STAR[:n_steps])
+    ctx_t = torch.tensor(
+        float(pipe.diffusion_model.config.context_noise), device=device, dtype=dtype
+    )
+
+    # Per-clip lap geometry: pairs v3 mixes lap lengths across clips (the
+    # repeat-prior fix), so nothing may assume a shared lap_chunks.
+    if PAIR_SCHEME == "fork":
+        ks_by_clip = [
+            training_ks_fork(int(d["num_chunk"]), d["fork_starts"]) for d in datas
+        ]
+    else:
+        ks_by_clip = [
+            training_ks(int(d["num_chunk"]), int(d["lap_chunks"]), int(d["laps"]))
+            for d in datas
+        ]
+    print(
+        f"{len(datas)} clips ({len(val_ids)} val) from {len(POOLS)} pools | "
+        f"lap_chunks {[int(d['lap_chunks']) for d in datas]} | cells/clip "
+        f"{[len(ks) for ks in ks_by_clip]} | contraction w={CW_LOSS}",
+        flush=True,
+    )
+
+    emb = datas[0]["embeddings"]
+    cache = pipe.initialize_cache_from_embeddings(
+        text_embeddings=emb["text_embeddings"],
+        image_embeddings=emb["image_embeddings"],
+    )
+    tc = cache.transformer_cache
+    _text_clip = [0]
+
+    network = transformer.network
+    wrapped = apply_lora(network, rank=RANK)  # ty: ignore[invalid-argument-type]
+    params = lora_parameters(network)  # ty: ignore[invalid-argument-type]
+    print(
+        f"LoRA on {len(wrapped)} projections | "
+        f"{sum(p.numel() for p in params) / 1e6:.2f}M params",
+        flush=True,
+    )
+    patch_functional_attention()
+    opt = torch.optim.AdamW(params, lr=LR)
+
+    def use_clip_text(c: int) -> None:
+        """Swap the cross-attn text KV when the sampled clip changes."""
+        if _text_clip[0] != c:
+            swap_text_kv(network, tc, datas[c]["embeddings"]["text_embeddings"])
+            _text_clip[0] = c
+
+    start_step = 0
+    if CKPT.exists():
+        state = torch.load(CKPT, map_location="cpu", weights_only=False)
+        for i, p in enumerate(params):
+            p.data.copy_(state["lora"][i].to(p.device, p.dtype))
+        opt.load_state_dict(state["opt"])
+        start_step = state["step"]
+        rng = np.random.default_rng(state["rng_seed"])
+        print(f"RESUMED from {CKPT} at step {start_step}", flush=True)
+    elif INIT and INIT != "scratch":
+        state = torch.load(INIT, map_location="cpu", weights_only=False)
+        for i, p in enumerate(params):
+            p.data.copy_(state["lora"][i].to(p.device, p.dtype))
+        print(f"warm start from {INIT}", flush=True)
+
+    def save(step: int, path: Path = CKPT) -> None:
+        tmp = path.with_suffix(".tmp")
+        torch.save(
+            {
+                "lora": {i: p.detach().cpu() for i, p in enumerate(params)},
+                "opt": opt.state_dict(),
+                "step": step,
+                "rng_seed": SEED + step,
+            },
+            tmp,
+        )
+        tmp.replace(path)
+
+    _device_clips: dict[int, dict] = {}
+
+    def to_device(c: int) -> dict:
+        """Device-resident view of clip ``c`` (memoized)."""
+        if c not in _device_clips:
+
+            def move(v):
+                if isinstance(v, list):
+                    return [move(x) for x in v]
+                return v.to(device, dtype) if isinstance(v, Tensor) else v
+
+            _device_clips[c] = {
+                key: move(v)
+                for key, v in datas[c].items()
+                if key
+                in (
+                    "latents",
+                    "hdmaps",
+                    "lap_chunks",
+                    "num_chunk",
+                    "fork_starts",
+                    "fork_latents",
+                )
+            }
+        return _device_clips[c]
+
+    def replay_window(latents: list[Tensor], hdmaps: list[Tensor], k: int) -> None:
+        """Rebuild the KV state from chunks ``k - REPLAY_CHUNKS .. k - 1``."""
+        start = max(0, k - REPLAY_CHUNKS)
+        reset_history(tc)
+        for bc in tc.network_cache.block_caches:
+            bc.self_attn._prev_chunk_idx = start - 1
+        for j in range(start, k):
+            g = torch.Generator(device=device).manual_seed(77_000 + j)
+            noisy = scheduler.add_noise(latents[j], ctx_t, rng=g)
+            tc.start(j)
+            transformer.finalize_kv_cache(
+                noisy_latent=noisy, timestep=ctx_t, cache=tc, input=hdmaps[j]
+            )
+            tc.finalize(j)
+
+    def predict_v(z_t: Tensor, t_idx: int, hdmap: Tensor) -> Tensor:
+        t = timesteps[t_idx].to(device=device, dtype=dtype)  # ty: ignore[not-subscriptable]
+        with functional_attention():
+            flow = transformer.predict_flow(
+                noisy_latent=z_t, timestep=t, cache=tc, input=hdmap
+            )
+        return flow.float()
+
+    def make_zt(x0: Tensor, t_idx: int, rng_: np.random.Generator) -> Tensor:
+        sig = sigmas[t_idx].to(dtype)  # ty: ignore[not-subscriptable]
+        g = torch.Generator(device=device).manual_seed(int(rng_.integers(2**31)))
+        return (1 - sig) * x0 + sig * torch.randn(
+            x0.shape, device=device, dtype=dtype, generator=g
+        )
+
+    excl = {c: {"tested": 0, "excluded": 0} for c in range(len(datas))}
+
+    def sample_losses(
+        c: int, grad: bool, rng_: np.random.Generator
+    ) -> tuple[Tensor, Tensor, Tensor] | None:
+        """One pooled sample -> (dag loss, contraction loss, r_target sq).
+
+        Returns ``None`` for degenerate (collapsed) cells; callers redraw.
+        """
+        use_clip_text(c)
+        d = to_device(c)
+        lap_chunks = int(d["lap_chunks"])
+        k = int(rng_.choice(ks_by_clip[c]))
+        t_idx = int(rng_.choice(n_steps, p=t_probs))
+        t2_idx = int(rng_.choice(n_steps, p=t_probs))
+        gen = d["latents"]
+        clean = list(gen)
+        for j in range(max(1, k - REPLAY_CHUNKS), k + 1):
+            clean[j] = clean_chunk(d, gen, j, lap_chunks)
+        n_draws = UW_DRAWS if UW else 1
+        z_ts = [make_zt(gen[k], t_idx, rng_) for _ in range(n_draws)]
+        z_t = z_ts[0]
+        z_t2 = make_zt(gen[k + 1], t2_idx, rng_)
+
+        with torch.no_grad():
+            set_lora_scale(network, 0.0)  # ty: ignore[invalid-argument-type]
+            # Teacher/base at k (clean swap over the replay span).
+            replay_window(clean, d["hdmaps"], k)
+            tc.start(k)
+            v_cleans = [predict_v(z, t_idx, d["hdmaps"][k]) for z in z_ts]
+            tc.finalize(k)
+            # Teacher/base at k+1 for the contraction target (the clean
+            # history now includes the counterpart chunk k).
+            replay_window(clean, d["hdmaps"], k + 1)
+            tc.start(k + 1)
+            v_clean2 = predict_v(z_t2, t2_idx, d["hdmaps"][k + 1])
+            tc.finalize(k + 1)
+            replay_window(gen, d["hdmaps"], k + 1)
+            tc.start(k + 1)
+            v_base2 = predict_v(z_t2, t2_idx, d["hdmaps"][k + 1])
+            tc.finalize(k + 1)
+            replay_window(gen, d["hdmaps"], k)
+            tc.start(k)
+            v_bases = [predict_v(z, t_idx, d["hdmaps"][k]) for z in z_ts]
+            tc.finalize(k)
+        v_clean, v_base = v_cleans[0], v_bases[0]
+        r_sq = (v_clean - v_base).square().sum()
+        r2_sq = (v_clean2 - v_base2).square().sum()
+
+        w = None
+        if UW:
+            rs = torch.stack([vc - vb for vc, vb in zip(v_cleans, v_bases)])
+            mean_r = rs.mean(0)
+            var = rs.var(0, unbiased=True)
+            bias2 = (mean_r.square() - var / n_draws).clamp_min(0.0)
+            w = (bias2 / (bias2 + var + 1e-12)).detach()
+
+        excl[c]["tested"] += 1
+        rel_v = (r_sq.sqrt() / (v_base.norm() + 1e-9)).item()
+        if rel_v > REL_V_EXCLUDE:
+            excl[c]["excluded"] += 1
+            return None
+
+        set_lora_scale(network, 1.0)  # ty: ignore[invalid-argument-type]
+        with torch.no_grad():
+            replay_window(gen, d["hdmaps"], k)  # LoRA-scaled replay, no grad
+        ctx = torch.enable_grad() if grad else torch.no_grad()
+        tc.start(k)
+        with ctx:
+            v_corr = predict_v(z_t, t_idx, d["hdmaps"][k])
+            err = (v_corr - v_clean).square()
+            gap = (v_clean - v_base).square()
+            if w is not None:
+                err, gap = err * w, gap * w
+            dag = err.sum() / (gap.sum() + 1e-8)
+            # Commit chunk k's corrected prediction with grad: recorded
+            # functional KV + a numerically identical no-grad buffer twin.
+            sig = sigmas[t_idx].to(dtype)  # ty: ignore[not-subscriptable]
+            x0_corr = (z_t.float() - float(sig) * v_corr).to(dtype)
+            g = torch.Generator(device=device).manual_seed(77_000 + k)
+            eps = torch.randn(x0_corr.shape, device=device, dtype=dtype, generator=g)
+            idx = torch.argmin(
+                (scheduler._full_timesteps - ctx_t.float()).abs()  # ty: ignore[unsupported-operator]
+            ).reshape(1)
+            sig_ctx = scheduler._full_sigmas.index_select(0, idx).reshape(()).to(dtype)  # ty: ignore[call-non-callable]
+            noisy_corr = (1 - sig_ctx) * x0_corr + sig_ctx * eps
+            recorded: list = []
+            with record_kv(recorded), functional_attention():
+                transformer.predict_flow(
+                    noisy_latent=noisy_corr,
+                    timestep=ctx_t,
+                    cache=tc,
+                    input=d["hdmaps"][k],
+                )
+        with torch.no_grad():  # buffer twin write + index advance
+            transformer.finalize_kv_cache(
+                noisy_latent=noisy_corr.detach(),
+                timestep=ctx_t,
+                cache=tc,
+                input=d["hdmaps"][k],
+            )
+        tc.finalize(k)
+        tc.start(k + 1)
+        with ctx:
+            with inject_kv(recorded):
+                v_con = predict_v(z_t2, t2_idx, d["hdmaps"][k + 1])
+            con = (v_con - v_clean2).square().sum() / (r2_sq + 1e-8)
+        tc.finalize(k + 1)
+        return dag, con, r_sq
+
+    def draw_losses(
+        ids: list[int], grad: bool, rng_: np.random.Generator
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Redraw until a non-degenerate cell is sampled."""
+        while True:
+            out = sample_losses(int(rng_.choice(ids)), grad, rng_)
+            if out is not None:
+                return out
+
+    @torch.no_grad()
+    def val_r2(n: int = 12) -> tuple[float, float]:
+        vrng = np.random.default_rng(1234)
+        s_dag = s_con = 0.0
+        for _ in range(n):
+            dag, con, _ = draw_losses(val_ids, False, vrng)
+            s_dag += dag.item()
+            s_con += con.item()
+        return 1 - s_dag / n, 1 - s_con / n
+
+    torch.set_grad_enabled(True)
+    best_vd = float("-inf")
+    for step in range(start_step + 1, STEPS + 1):
+        for pg in opt.param_groups:
+            pg["lr"] = LR * min(1.0, step / WARMUP)
+        opt.zero_grad()
+        dag, con, r_sq = draw_losses(train_ids, True, rng)
+        loss = dag + CW_LOSS * con
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(params, GRAD_CLIP)
+        opt.step()
+        if step % EVAL_EVERY == 0 or step == 1:
+            vd, vc = val_r2()
+            excl_str = " ".join(
+                f"c{c}:{v['excluded']}/{v['tested']}" for c, v in excl.items()
+            )
+            print(
+                f"step {step:5d} | dag {dag.item():.4f} con {con.item():.4f}"
+                f" | val dag-R^2 {vd:+.3f} con-R^2 {vc:+.3f}"
+                f" | |r|^2 {r_sq.item():.1f} | excluded {excl_str}",
+                flush=True,
+            )
+            if SNAP_EVERY and vd > best_vd:
+                best_vd = vd
+                save(step, CKPT.with_name(f"{CKPT.stem}_valpeak.pt"))
+                print(
+                    f"val-peak snapshot at step {step} (dag-R^2 {vd:+.3f})", flush=True
+                )
+        if step % SAVE_EVERY == 0 or step == STEPS:
+            save(step)
+        if SNAP_EVERY and step % SNAP_EVERY == 0:
+            save(step, CKPT.with_name(f"{CKPT.stem}_step{step}.pt"))
+
+    vd, vc = val_r2(24)
+    print(
+        f"TRAIN-V2-DONE | final val dag-R^2 {vd:+.3f} con-R^2 {vc:+.3f} | saved {CKPT}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/omnidreams/_drift_corrector.py b/integrations/omnidreams/omnidreams/_drift_corrector.py
new file mode 100644
index 00000000..b6b7b418
--- /dev/null
+++ b/integrations/omnidreams/omnidreams/_drift_corrector.py
@@ -0,0 +1,153 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Clean Forcing drift corrector for the Omnidreams runner.
+
+Deploys the trained corrector LoRA (``drift_correction/train_v2.py``
+checkpoints) on a built :class:`~omnidreams.runner.OmnidreamsRunner`'s
+pipeline at ``alpha*(t) * gain`` per denoise step. Shipped config (2026-07-24): ``lora_v2_v3_valpeak.pt`` at gain 0.25
+(``corrgate025`` — best trees/foliage detail and consistency, drift
+Delta +0.99 vs base +2.44). Mirrors the HY-WorldPlay deploy module
+(``hy_worldplay/_drift_corrector.py``); self-contained so the production
+runner does not import the research directory.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+## Deploy policy
+
+GATE_ALPHA = {1000.0: 0.96, 803.0: 0.667}
+"""Unbiased alpha*(t) from the step-0 systematicity gate on photoreal
+drift pairs (drift_correction's ``outputs/gate/gate_faithful_v2.json``):
+the systematic fraction of the drift-induced error at each of the two
+distilled solver timesteps. The corrector LoRA is rescaled to
+``alpha*(t) * gain`` before every denoise step (nearest-t lookup); the
+``finalize_kv_cache`` context forward (t=128) resolves to the t=803 entry,
+matching the evaluated deploy configs."""
+
+_LORA_TARGETS = (
+    "self_attn.q_proj",
+    "self_attn.k_proj",
+    "self_attn.v_proj",
+    "self_attn.output_proj",
+)
+"""Self-attention projections the corrector checkpoints were trained on."""
+
+_LORA_RANK = 16
+"""Rank of the shipped corrector checkpoints."""
+
+
+class _LoRALinear(nn.Module):
+    """Frozen base linear plus a runtime-gated low-rank delta.
+
+    Mirrors the training-side module in
+    ``integrations/omnidreams/drift_correction/_lora.py``: ``scale`` is the
+    runtime gain (``0`` = exact base output), and the A/B path runs in fp32
+    regardless of the base dtype.
+    """
+
+    def __init__(self, base: nn.Linear, rank: int):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad_(False)
+        self.A = nn.Linear(base.in_features, rank, bias=False)
+        self.B = nn.Linear(rank, base.out_features, bias=False)
+        nn.init.zeros_(self.B.weight)
+        self.scale = 0.0
+
+    def forward(self, x: Tensor) -> Tensor:
+        out = self.base(x)
+        if self.scale != 0:
+            delta = self.B(self.A(x.to(self.A.weight.dtype)))
+            out = out + self.scale * delta.to(out.dtype)
+        return out
+
+
+def _apply_lora(network: nn.Module) -> list[nn.Parameter]:
+    """Wrap the target linears and return the LoRA parameters in load order."""
+    for mname, module in list(network.named_modules()):
+        for cname, child in list(module.named_children()):
+            full = f"{mname}.{cname}" if mname else cname
+            # Substring match, exactly as the training-side apply_lora, so
+            # the wrap set and load order match the checkpoint indices.
+            if isinstance(child, nn.Linear) and any(t in full for t in _LORA_TARGETS):
+                setattr(
+                    module,
+                    cname,
+                    _LoRALinear(child, _LORA_RANK).to(child.weight.device),
+                )
+    params: list[nn.Parameter] = []
+    for m in network.modules():
+        if isinstance(m, _LoRALinear):
+            params += list(m.A.parameters()) + list(m.B.parameters())
+    return params
+
+
+def _set_scale(network: nn.Module, scale: float) -> None:
+    """Set the runtime gain on every wrapped linear."""
+    for m in network.modules():
+        if isinstance(m, _LoRALinear):
+            m.scale = scale
+
+
+def apply_drift_corrector(runner: Any, checkpoint: Path, gain: float) -> str:
+    """Deploy the corrector LoRA on ``runner`` with the alpha*(t) gate.
+
+    Args:
+        runner: A built ``OmnidreamsRunner``.
+        checkpoint: Corrector LoRA checkpoint (``train_v1``/``train_v2``
+            format: a dict whose ``"lora"`` entry maps load-order indices
+            to tensors).
+        gain: Global gain composed with the alpha*(t) profile; the
+            shipped configuration (``corrgate025``) is 0.25.
+
+    Returns:
+        A log-line string describing the deployed configuration.
+    """
+    network = runner.pipeline.diffusion_model.transformer.network
+    if hasattr(network, "_orig_mod"):  # unwrap torch.compile
+        network = network._orig_mod
+    params = _apply_lora(network)
+
+    sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"]
+    assert len(sd) == len(params), (
+        f"corrector checkpoint has {len(sd)} LoRA tensors but the network "
+        f"exposes {len(params)}; rank or target mismatch."
+    )
+    for i, p in enumerate(params):
+        p.data.copy_(sd[i].to(p.device, p.dtype))
+
+    # Per-step gate: rescale the LoRA to alpha*(t) x gain before every
+    # denoise step (nearest-t lookup; finalize_kv_cache calls positionally).
+    transformer = runner.pipeline.diffusion_model.transformer
+    orig_pf = transformer.predict_flow
+
+    def gated_pf(*args, **kwargs):
+        ts = kwargs.get("timestep", args[1] if len(args) > 1 else None)
+        t = float(ts.reshape(-1).max())
+        alpha = min(GATE_ALPHA.items(), key=lambda kv: abs(kv[0] - t))[1]
+        _set_scale(network, alpha * gain)
+        return orig_pf(*args, **kwargs)
+
+    transformer.predict_flow = gated_pf
+    return f"corrected (alpha*(t) x {gain})"
diff --git a/integrations/omnidreams/omnidreams/runner.py b/integrations/omnidreams/omnidreams/runner.py
index b0322e11..1fe49e0a 100644
--- a/integrations/omnidreams/omnidreams/runner.py
+++ b/integrations/omnidreams/omnidreams/runner.py
@@ -216,6 +216,17 @@ class OmnidreamsRunnerConfig(RunnerConfig):
     """Single-view example clip to pull from :data:`EXAMPLE_DATA_HF_REPO`.
     Ignored for multi-view or when paths are already populated."""
 
+    drift_corrector: Path | None = None
+    """Clean Forcing drift-corrector LoRA checkpoint. ``None`` (default)
+    disables correction and is byte-identical to current behavior. When
+    set, the corrector deploys at ``alpha*(t) * drift_corrector_gain`` per
+    denoise step (see ``omnidreams/_drift_corrector.py``). Shipped config:
+    the ``lora_v2_v3_valpeak.pt`` release checkpoint at the default gain."""
+
+    drift_corrector_gain: float = 0.25
+    """Global gain composed with the per-step ``alpha*(t)`` gate profile;
+    the shipped configuration (``corrgate025``) is 0.25."""
+
 
 class OmnidreamsRunner(Runner[OmnidreamsRunnerConfig, OmnidreamsPipeline]):
     """Streaming HDMap-conditioned I2V driver."""
@@ -231,6 +242,13 @@ def run(self) -> None:
         )
         if cfg.example_data:
             self._fill_example_data_defaults()
+        if cfg.drift_corrector is not None:
+            from omnidreams._drift_corrector import apply_drift_corrector
+
+            mode = apply_drift_corrector(
+                self, cfg.drift_corrector, cfg.drift_corrector_gain
+            )
+            logger.info(f"[{cfg.runner_name}] drift corrector: {mode}")
         if cfg.save_embeddings_path is not None:
             self._run_save_embeddings(cfg.save_embeddings_path)
             return
diff --git a/integrations/omnidreams/tests/test_drift_corrector.py b/integrations/omnidreams/tests/test_drift_corrector.py
new file mode 100644
index 00000000..916c8694
--- /dev/null
+++ b/integrations/omnidreams/tests/test_drift_corrector.py
@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU-only unit tests for the Clean Forcing drift-corrector deploy hook.
+
+Covers the behaviours a deployment depends on:
+
+* ``_LoRALinear`` is a strict identity at ``scale == 0`` and at the
+  zero-initialized ``B`` (so wrapping the network never changes base
+  outputs until a trained checkpoint is loaded and gated on).
+* ``_apply_lora`` wraps exactly the self-attention projections the
+  training-side module wraps (same match rule -> same checkpoint order),
+  and ``_set_scale`` reaches every wrapped linear.
+* The ``alpha*(t)`` gate profile resolves by nearest-t lookup, including
+  the context-noise forward.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+from omnidreams._drift_corrector import (
+    GATE_ALPHA,
+    _apply_lora,
+    _LoRALinear,
+    _set_scale,
+)
+
+pytestmark = pytest.mark.ci_cpu
+
+
+def test_lora_linear_is_identity_at_zero_scale():
+    torch.manual_seed(0)
+    base = nn.Linear(8, 6, bias=False)
+    lora = _LoRALinear(base, rank=2)
+    nn.init.normal_(lora.A.weight)
+    nn.init.normal_(lora.B.weight)  # non-zero delta path
+    x = torch.randn(3, 8)
+    lora.scale = 0.0
+    assert torch.equal(lora(x), base(x))
+
+
+def test_lora_linear_is_identity_at_zero_init_b():
+    torch.manual_seed(0)
+    base = nn.Linear(8, 6, bias=False)
+    lora = _LoRALinear(base, rank=2)  # B is zero-initialized
+    lora.scale = 1.0
+    x = torch.randn(3, 8)
+    assert torch.allclose(lora(x), base(x))
+
+
+def test_lora_linear_applies_scaled_delta():
+    torch.manual_seed(0)
+    base = nn.Linear(8, 6, bias=False)
+    lora = _LoRALinear(base, rank=2)
+    nn.init.normal_(lora.A.weight)
+    nn.init.normal_(lora.B.weight)
+    x = torch.randn(3, 8)
+    lora.scale = 0.5
+    delta = lora(x) - base(x)
+    lora.scale = 1.0
+    assert torch.allclose(2.0 * delta, lora(x) - base(x), atol=1e-5)
+
+
+def test_apply_lora_wraps_only_attention_targets_and_set_scale_reaches_all():
+    class Toy(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.self_attn = nn.ModuleDict(
+                {
+                    n: nn.Linear(4, 4, bias=False)
+                    for n in ("q_proj", "k_proj", "v_proj", "output_proj")
+                }
+            )
+            self.mlp = nn.Linear(4, 4, bias=False)
+
+    toy = Toy()
+    params = _apply_lora(toy)
+    wrapped = [m for m in toy.modules() if isinstance(m, _LoRALinear)]
+    assert len(wrapped) == 4  # q/k/v/output projections but not the mlp
+    assert len(params) == 8  # A + B per wrapped linear
+    assert not isinstance(toy.mlp, _LoRALinear)
+    _set_scale(toy, 0.25)
+    assert all(m.scale == 0.25 for m in wrapped)
+
+
+def test_gate_profile_nearest_t_lookup():
+    def alpha_at(t: float) -> float:
+        return min(GATE_ALPHA.items(), key=lambda kv: abs(kv[0] - t))[1]
+
+    assert alpha_at(1000.0) == GATE_ALPHA[1000.0]
+    assert alpha_at(803.0) == GATE_ALPHA[803.0]
+    # The context-noise forward (t=128) resolves to the low-t entry,
+    # matching the evaluated deploy configs.
+    assert alpha_at(128.0) == GATE_ALPHA[803.0]
+
+
+def test_gate_profile_is_a_strict_attenuation():
+    assert all(0.0 < a <= 1.0 for a in GATE_ALPHA.values())