From d151d573138ac424b2f3d18b5bfed5d46d553609 Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Thu, 13 Aug 2026 22:13:36 +0000
Subject: [PATCH 1/8] Add edit-SFT style-skin pipeline: JoyAI pair generation,
 VLM filter, LoRA trainer

Offline JoyAI-Video-Edit restyles of the model's own rollouts become
(source, style prompt, styled) training pairs. A style-mode VLM filter
scores edit strength/persistence plus a style-agnostic road-layout
criterion; heavy styles whose layout drifts late train early-window-only
instead of being dropped. A one-shot precompute encodes prompts, first
frames, HDMaps, and styled targets into the pipeline latent space. The
trainer does edit-timestamped teacher-forced flow matching (source-latent
context replay, plain text-KV swap at chunk k, styled targets across the
window) into a multi-style LoRA, checkpoint-compatible with the live-edit
deploy hook (#431). Training helpers are vendored pending #398.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 integrations/omnidreams/edit_sft/PLAN.md      |  58 ++
 integrations/omnidreams/edit_sft/README.md    |  97 ++++
 integrations/omnidreams/edit_sft/_host.py     | 327 +++++++++++
 integrations/omnidreams/edit_sft/_lora.py     | 138 +++++
 .../omnidreams/edit_sft/_train_attn.py        | 186 ++++++
 .../omnidreams/edit_sft/filter_pairs.py       | 259 +++++++++
 .../omnidreams/edit_sft/generate_sources.py   | 162 ++++++
 .../omnidreams/edit_sft/precompute_style.py   | 264 +++++++++
 .../omnidreams/edit_sft/style_prompts.py      |  86 +++
 .../omnidreams/edit_sft/train_style_sft.py    | 546 ++++++++++++++++++
 10 files changed, 2123 insertions(+)
 create mode 100644 integrations/omnidreams/edit_sft/PLAN.md
 create mode 100644 integrations/omnidreams/edit_sft/README.md
 create mode 100644 integrations/omnidreams/edit_sft/_host.py
 create mode 100644 integrations/omnidreams/edit_sft/_lora.py
 create mode 100644 integrations/omnidreams/edit_sft/_train_attn.py
 create mode 100644 integrations/omnidreams/edit_sft/filter_pairs.py
 create mode 100644 integrations/omnidreams/edit_sft/generate_sources.py
 create mode 100644 integrations/omnidreams/edit_sft/precompute_style.py
 create mode 100644 integrations/omnidreams/edit_sft/style_prompts.py
 create mode 100644 integrations/omnidreams/edit_sft/train_style_sft.py

diff --git a/integrations/omnidreams/edit_sft/PLAN.md b/integrations/omnidreams/edit_sft/PLAN.md
new file mode 100644
index 000000000..5537d67ae
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/PLAN.md
@@ -0,0 +1,58 @@
+# Instruction-SFT for grounded edits (Tier-2b of the live-edit hack)
+
+**Goal:** push mid-stream editing beyond what the base model's priors give us
+training-free. Three targets, ranked by demo value / feasibility:
+
+1. **Object-by-prompt** — "a red sports car parked ahead", "boxes on the road":
+   prompt-driven object addition that actually materializes (today only weather/
+   lighting/atmosphere edits are reliable).
+2. **Grounded props** — "Other"-class bboxes (cones, barriers) render as visible
+   obstacles (today they under-render even oversized and grounded).
+3. **Trajectory-locked moving actors** — a commanded constant-gap lead vehicle stays a
+   lead vehicle (today the residential scene prior overrides commanded box motion and
+   paints an oncoming pass instead). Hardest; data-dependent; may slip.
+
+## Data recipes
+
+- **A. JoyAI edit pairs (targets 1, partially 2).** Source videos = plain rollouts of
+  our own model over the 32 local HF sample clips (`generate_sources.py`). Push each
+  through JoyAI-Video-Edit (Apache-2.0, 30 FPS streaming editor, released weights)
+  with a driving-edit instruction bank (add parked vehicles / obstacles / cones,
+  appearance changes) → temporally aligned (source, instruction, edited) triplets.
+  Filter: VLM edit-correctness + preservation + temporal consistency (JoyAI's own
+  criteria); recaption with the achieved edit.
+- **B. Cone/prop self-training (target 2).** Our model + construction-zone prompt
+  assist demonstrably paints cones. Generate rollouts with oversized "Other" boxes +
+  prompt assist; VLM-filter frames where a prop actually appears near the box
+  projection; the surviving clips are (HDMap-with-box → video-with-prop) pairs that
+  teach the box→prop mapping *without* the prompt crutch.
+- **C. Teacher-regenerated continuations (targets 1, 3; optional).** The verified
+  bidirectional teacher (HF `single_view/teacher/...`, same latent space) regenerates
+  post-edit continuations at 35 steps + CFG under edit prompts / modified HDMaps —
+  higher-compliance targets than the student can produce for itself. Use where JoyAI's
+  weak spot (local add) shows.
+
+## Training recipe
+
+Teacher-forced flow-matching SFT on **edit-timestamped clips** (prompt A for chunks
+[0, k), prompt/HDMap B after), LoRA on the same 8-projection target set as the
+guidance-distillation trainer (its infra — functional attention, per-block
+checkpointing, premerge deploy — is vendored into this directory). Two phases:
+
+1. **Context build:** replay chunks [0, k) from the SOURCE latents through
+   `finalize_kv_cache` (the vendored `_host.py:replay_history` machinery from the
+   Clean Forcing training infra, PR #398) — history is real source content.
+2. **Edit supervision:** for chunks [k, k+m), flow-matching MSE on the EDITED latents
+   under the new conditioning, per-chunk independent timesteps; commit edited latents
+   into the KV as teacher forcing proceeds.
+
+Composes with the Tier-2a LoRA (train on top of, or merge and re-gate — decide after
+first results). Eval: the InterBench-style VLM judge (Trigger/Align/Consistency) on a
+held-out instruction bank + the existing divergence harness.
+
+## Milestones
+
+- M1: source corpus generated; JoyAI running on one clip end-to-end.
+- M2: ≥500 filtered pairs across ~10 instruction types; recipe B pilot (~200 clips).
+- M3: SFT v1 trained; object-by-prompt trigger rate measured vs base.
+- M4: props from boxes without prompt assist; decision point on target 3.
diff --git a/integrations/omnidreams/edit_sft/README.md b/integrations/omnidreams/edit_sft/README.md
new file mode 100644
index 000000000..2515841fc
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/README.md
@@ -0,0 +1,97 @@
+# Edit-SFT (Tier-2b) — style restyles; see PLAN.md
+
+Teach the realtime distilled student to apply a **global visual style**
+("arcade racing game world", "comic ink", ...) mid-stream on a plain prompt
+swap, imitating offline JoyAI restyles. Training is teacher-forced flow
+matching on edit-timestamped clips: history = the model's own source rollout,
+targets after the swap chunk = the JoyAI-styled video encoded into the
+model's latent space. The checkpoint format is the guidance-distillation
+trainer's, so it deploys through the live-edit deploy hook (`TextEditLoRA`,
+PR #431) unchanged. The training helpers (`_host.py`, `_lora.py`,
+`_train_attn.py`) are vendored from the Clean Forcing training infra
+(PR #398) so this directory is self-contained on `main`; consolidate when
+that lands.
+
+Run from the repo root, in order:
+
+1. `.venv/bin/python integrations/omnidreams/edit_sft/generate_sources.py`
+   — roll the source corpus (videos + per-chunk latents + manifest,
+   `outputs/sources/`). *Done: 10 clips x 28 chunks.*
+2. JoyAI-Video-Edit batch (separate venv, see the project notes) over the
+   style instruction bank -> `outputs/style_pairs/<uuid>__<style>.mp4`.
+   *Done: 62 pairs across 11 slugs.*
+3. `PAIRS_DIR=integrations/omnidreams/edit_sft/outputs/style_pairs \
+   JOBS=style_jobs.jsonl,audition2_jobs.jsonl MODE=style \
+   .venv/bin/python integrations/omnidreams/edit_sft/filter_pairs.py`
+   — style-mode VLM gate (edit applied + persistent + road-layout at an
+   early and a late frame) -> `style_pairs/filter_report.json`. `passed`
+   pairs train anywhere; `early_window_ok` pairs train only in the first
+   ~4 s (before streaming layout drift). *Running.*
+4. `.venv/bin/python integrations/omnidreams/edit_sft/precompute_style.py`
+   — one shot, ~20 GB VRAM: style + clip prompt embeddings, first-frame
+   latents, per-chunk HDMap latents, and per-chunk styled-target latents
+   (pipeline Wan-VAE encoder, AR chunk schedule 5->2 / 8->2 latent frames,
+   patchified). Resumable; rerun if a decode assert fires.
+5. `STEPS=1000 .venv/bin/python integrations/omnidreams/edit_sft/train_style_sft.py`
+   — the trainer; **~30 GB VRAM eager** (2B DiT + one per-episode KV cache
+   + immediate per-term backwards). Per step: replay 3 source chunks (KV
+   window) at base weights, plain prompt swap at `k ~ U[4, 20]`, then
+   `PRE_CHUNKS=1` no-op chunk (source target, original prompt) +
+   `SPAN=4` styled chunks (styled targets, style prompt), committing the
+   teacher-forced latents into the KV as it proceeds. Checkpoints
+   (`{"lora": {i: tensor}}`, loadable by `TextEditLoRA`) land in
+   `outputs/lora_style_stepN.pt` every 100 steps.
+
+Knobs: `STEPS`, `LR` (2e-4), `RANK` (64), `SEED`, `HOLDOUT` (2),
+`NOOP_PROB` (0.1), `PRE_CHUNKS` (1), `SPAN` (4), `SWAP_MIN`/`SWAP_MAX`
+(4/20), `EARLY_CHUNKS` (15), `REPLAY_CHUNKS` (3), `SKIP_STYLES`
+(`cyberpunk_neon,watercolor`).
+
+## Style filtering (`MODE=style`) and the heavy-style early-window policy
+
+`filter_pairs.py MODE=style` adapts the VLM gate to global restyles, where
+the object-mode `scene_preserved` axis is meaningless (everything changes
+by design). It keeps `edit_applied` / `persistent` and adds a
+style-agnostic **road-layout** check — the judge compares a source frame
+against the styled frame and scores only WHERE the road, lane lines, and
+vehicles sit (describe-then-score prompting; a bare cross-style JSON score
+collapses to 0) — at an early (frame 60) and a late (frame 190) point:
+
+- **`passed`** (`edit_applied`, `persistent`, and `layout_late` all >= 3):
+  the restyle holds and the layout survives the whole clip — the pair
+  trains at any swap position.
+- **`early_window_ok`** (`edit_applied` and `layout_early` >= 3, but the
+  late check fails): the observed heavy-style failure mode — streaming
+  error accumulation progressively replaces the road layout — sets in
+  only after a few seconds. Instead of dropping these pairs,
+  `train_style_sft.py` *demotes* them: they train only within the first
+  `EARLY_CHUNKS` (15) chunks (~4 s), the pre-drift window the filter
+  certified. This keeps heavy styles (cartoon cel, anime, pixel art) in
+  the training set without teaching layout drift.
+- Everything else is dropped. `SKIP_STYLES` (default
+  `cyberpunk_neon,watercolor` — scene re-imagined / washed out at the
+  audition) are excluded regardless of their filter verdict.
+
+## Eval plan (gate before any deploy claim)
+
+Reuse the guidance-distillation eval protocol (its `eval_guidance.py`, not
+yet upstreamed) on the 2 held-out clips x the filter-passing styles,
+RNG-matched against a shared no-edit control rollout:
+
+- `lora_plain`: LoRA gated to `[SWAP_AT, SWAP_AT + SPAN)` + plain swap to
+  the style prompt (deployment semantics; `RANK` must match the ckpt).
+- `base_plain`: base weights + plain swap (floor).
+- Reference divergence: the JoyAI styled clip vs the source clip over the
+  same window (the offline target this SFT imitates — there is no "guided
+  teacher" arm for styles).
+
+Gate: styled-swap divergence ratio
+`sum(gap_lora) / sum(gap_styled_target) >= 0.8` over the window, averaged
+across combos, with `pre_swap_max_gap ~ 0` (the LoRA is inactive before the
+swap, so any pre-swap divergence is a bug). Then the two checks divergence
+cannot prove: (a) the style-mode VLM judge (`filter_pairs.py MODE=style`)
+on (control, lora rollout) — `edit_applied` and `layout` scores, catching
+"diverged but not styled" and layout damage; (b) a human visual pass on
+side-by-sides (project lesson: divergence alone proves nothing). For
+early-window styles, evaluate with `SWAP_AT <= 10` so the window ends
+inside the certified range.
diff --git a/integrations/omnidreams/edit_sft/_host.py b/integrations/omnidreams/edit_sft/_host.py
new file mode 100644
index 000000000..00536c091
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/_host.py
@@ -0,0 +1,327 @@
+# 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.
+
+Vendored from the Clean Forcing training infra (PR #398) so this pipeline
+is self-contained; consolidate when that lands.
+
+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. K/V are computed
+    through the live projection modules, so ``_lora``-wrapped projections
+    contribute their delta at the current scale.
+    """
+    k_proj = network.blocks[0].cross_attn.k_proj
+    # A LoRALinear wrap hides ``.weight`` behind ``.base``.
+    w = getattr(k_proj, "base", 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 the pair builder / gate in the PR #398 infra)
+
+
+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 :func:`save_clip` 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/edit_sft/_lora.py b/integrations/omnidreams/edit_sft/_lora.py
new file mode 100644
index 000000000..a31d97424
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/_lora.py
@@ -0,0 +1,138 @@
+# 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.
+
+Vendored from the Clean Forcing training infra (PR #398) so this pipeline
+is self-contained; consolidate when that lands.
+
+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/edit_sft/_train_attn.py b/integrations/omnidreams/edit_sft/_train_attn.py
new file mode 100644
index 000000000..fa63434fe
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/_train_attn.py
@@ -0,0 +1,186 @@
+# 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).
+
+Vendored from the Clean Forcing training infra (PR #398) so this pipeline
+is self-contained; consolidate when that lands.
+
+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
+from typing import cast
+
+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:
+        # Stock callers always pass rope_freqs (required by the original
+        # signature); the wider Optional here only serves the mirror.
+        return _STOCK_FORWARD(self, x, kv_cache, cast(Tensor, rope_freqs))
+    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/edit_sft/filter_pairs.py b/integrations/omnidreams/edit_sft/filter_pairs.py
new file mode 100644
index 000000000..367b67ad4
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/filter_pairs.py
@@ -0,0 +1,259 @@
+# 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.
+
+"""VLM filter for the edit-SFT pair corpus (Tier-2b, quality gate).
+
+Scores every (source, instruction, edited) triplet with Cosmos-Reason1
+(a Qwen2.5-VL — already cached as the pipeline's text encoder) on three
+axes, 0-5 each:
+
+* ``edit_applied`` — is the instructed edit visible in the edited video?
+* ``persistent`` — is it present across early / mid / late frames?
+* ``scene_preserved`` — is everything else (layout, road, buildings)
+  unchanged from the source?
+
+A pair passes at >= 3 on all three. Writes ``filter_report.json`` next to
+the pairs plus a per-instruction summary table, so weak instruction types
+are visible before training.
+
+Run from the repo root::
+
+    .venv/bin/python integrations/omnidreams/edit_sft/filter_pairs.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from pathlib import Path
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import mediapy as media
+import numpy as np
+import torch
+from PIL import Image
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+
+MODEL_NAME = "nvidia/Cosmos-Reason1-7B"
+REVISION = "3210bec0495fdc7a8d3dbb8d58da5711eab4b423"
+"""Same pin as the pipeline's text encoder — weights already cached."""
+
+PAIRS_DIR = Path(
+    os.environ.get("PAIRS_DIR", "integrations/omnidreams/edit_sft/outputs/pairs")
+)
+SOURCES_DIR = Path(
+    os.environ.get("SOURCES_DIR", "integrations/omnidreams/edit_sft/outputs/sources")
+)
+JOBS_FILES = os.environ.get("JOBS", "jobs.jsonl").split(",")
+FRAME_IDS = (30, 110, 190)
+"""Early / mid / late sample frames (221-frame clips)."""
+
+MODE = os.environ.get("MODE", "object")
+"""``object`` (default): original three-axis gate. ``style``: for global
+restyles — replaces ``scene_preserved`` with style-agnostic ROAD-LAYOUT
+checks at an early and a late frame, so streaming layout drift (the
+observed failure mode: the road is progressively replaced) is caught and
+heavy styles can still qualify for early-window-only training."""
+
+PASS_BAR = 3
+MAX_SIDE = 640
+"""Longest image side fed to the judge (speed/VRAM)."""
+
+
+def _frames(path: Path, ids=FRAME_IDS) -> list[Image.Image]:
+    video = media.read_video(path)
+    out = []
+    for i in ids:
+        frame = video[min(i, len(video) - 1)]
+        img = Image.fromarray(np.asarray(frame))
+        scale = MAX_SIDE / max(img.size)
+        out.append(img.resize((int(img.width * scale), int(img.height * scale))))
+    return out
+
+
+def _parse_scores(text: str, keys: tuple[str, ...]) -> dict[str, int]:
+    """Extract integer scores from the judge's reply (JSON or loose text)."""
+    scores: dict[str, int] = {}
+    try:
+        hit = re.search(r"\{.*\}", text, re.S)
+        blob = json.loads(hit.group(0) if hit else "")
+        for key in keys:
+            if key in blob:
+                scores[key] = max(0, min(5, int(blob[key])))
+    except (AttributeError, ValueError, TypeError, json.JSONDecodeError):
+        pass
+    for key in keys:
+        if key not in scores:
+            hit = re.search(rf"{key}\D{{0,12}}([0-5])", text)
+            scores[key] = int(hit.group(1)) if hit else 0
+    return scores
+
+
+class Judge:
+    """Thin greedy-decoding wrapper around Cosmos-Reason1 as a VLM judge."""
+
+    def __init__(self) -> None:
+        self.processor = AutoProcessor.from_pretrained(
+            MODEL_NAME, revision=REVISION, local_files_only=True
+        )
+        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            MODEL_NAME, revision=REVISION, local_files_only=True, dtype=torch.bfloat16
+        ).to("cuda")  # ty: ignore[invalid-argument-type]  # stub loses the bound self
+        self.model.eval().requires_grad_(False)
+
+    @torch.no_grad()
+    def ask(self, images: list[Image.Image], question: str) -> str:
+        content = [{"type": "image"} for _ in images] + [
+            {"type": "text", "text": question}
+        ]
+        prompt = self.processor.apply_chat_template(
+            [{"role": "user", "content": content}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        inputs = self.processor(text=[prompt], images=images, return_tensors="pt").to(
+            "cuda"
+        )
+        out = self.model.generate(
+            **inputs, max_new_tokens=96, do_sample=False, temperature=None, top_p=None
+        )
+        reply = self.processor.batch_decode(
+            out[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
+        )[0]
+        return reply
+
+
+_LAYOUT_QUESTION = (
+    "Image 1 is a real driving photo. Image 2 is a stylized version of the "
+    "same scene. Look only at WHERE the road, lane lines, and vehicles are "
+    "placed in the picture. First, in one sentence, say whether the road in "
+    "image 2 runs through the same part of the picture in the same direction "
+    "as in image 1. Then on the last line write SCORE: n where n is 0-5 "
+    "(5 = road and objects in the same places, 0 = road moved or missing)."
+)
+"""Describe-then-score framing: a bare cross-style JSON score collapses to a
+uniform 0 (the style gap dominates); letting the judge state the comparison
+first produces calibrated scores. A reply with no SCORE line counts as 0 —
+in testing that degenerate mode only appeared on true layout failures."""
+
+
+def _judge_layout(judge: "Judge", source: Path, edited: Path, frame: int) -> int:
+    reply = judge.ask(
+        _frames(source, ids=(frame,)) + _frames(edited, ids=(frame,)),
+        _LAYOUT_QUESTION,
+    )
+    hit = re.search(r"SCORE:\s*([0-5])", reply)
+    return int(hit.group(1)) if hit else 0
+
+
+def main() -> None:
+    jobs = [
+        json.loads(line)
+        for name in JOBS_FILES
+        for line in (PAIRS_DIR / name).read_text().splitlines()
+        if line.strip()
+    ]
+    judge = Judge()
+    report: list[dict] = []
+
+    for i, job in enumerate(jobs):
+        edited_path = Path(job["output"])
+        source_path = Path(job["input"])
+        if not edited_path.exists():
+            report.append({"output": edited_path.name, "error": "missing"})
+            continue
+        instruction = job["instruction"]
+        slug = edited_path.stem.split("__", 1)[1]
+
+        edited = _frames(edited_path)
+        reply1 = judge.ask(
+            edited,
+            "These are three frames (early, middle, late) from an edited "
+            f'driving video. The requested edit was: "{instruction}". '
+            'Score strictly and answer ONLY JSON like {"edit_applied": 0-5, '
+            '"persistent": 0-5}: edit_applied = how clearly the edit is '
+            "visible; persistent = whether it is present in ALL three frames.",
+        )
+        s1 = _parse_scores(reply1, ("edit_applied", "persistent"))
+
+        entry = {
+            "output": edited_path.name,
+            "uuid": source_path.stem,
+            "slug": slug,
+            "instruction": instruction,
+            **s1,
+        }
+        if MODE == "style":
+            entry["layout_early"] = _judge_layout(judge, source_path, edited_path, 60)
+            entry["layout_late"] = _judge_layout(judge, source_path, edited_path, 190)
+            entry["passed"] = (
+                min(entry["edit_applied"], entry["persistent"]) >= PASS_BAR
+                and entry["layout_late"] >= PASS_BAR
+            )
+            entry["early_window_ok"] = (
+                entry["edit_applied"] >= PASS_BAR and entry["layout_early"] >= PASS_BAR
+            )
+        else:
+            src_mid = _frames(source_path, ids=(110,))
+            edit_mid = _frames(edited_path, ids=(110,))
+            reply2 = judge.ask(
+                src_mid + edit_mid,
+                "Image 1 is a frame from an original driving video; image 2 is "
+                "the same frame after a video edit was applied. Ignoring the "
+                f'requested edit ("{instruction}"), score how well everything '
+                "else is preserved (same street layout, buildings, lighting, "
+                'other vehicles). Answer ONLY JSON like {"scene_preserved": 0-5}.',
+            )
+            entry.update(_parse_scores(reply2, ("scene_preserved",)))
+            entry["passed"] = all(
+                entry[k] >= PASS_BAR
+                for k in ("edit_applied", "persistent", "scene_preserved")
+            )
+        report.append(entry)
+        if (i + 1) % 15 == 0:
+            print(f"{i + 1}/{len(jobs)} judged", flush=True)
+
+    (PAIRS_DIR / "filter_report.json").write_text(json.dumps(report, indent=2))
+
+    # Per-instruction summary.
+    by_slug: dict[str, list[dict]] = {}
+    for entry in report:
+        if "error" not in entry:
+            by_slug.setdefault(entry["slug"], []).append(entry)
+    extra = (
+        ("layout_early", "layout_late", "early_window_ok")
+        if MODE == "style"
+        else ("scene_preserved",)
+    )
+    header = " ".join(f"{k[:9]:>9s}" for k in ("edit_applied", "persistent") + extra)
+    print(f"\n{'instruction':22s} {'pass':>5s} {header}")
+    for slug in sorted(by_slug):
+        rows = by_slug[slug]
+        n_pass = sum(r["passed"] for r in rows)
+        means = " ".join(
+            f"{float(np.mean([r[k] for r in rows])):9.1f}"
+            for k in ("edit_applied", "persistent") + extra
+        )
+        print(f"{slug:22s} {n_pass:2d}/{len(rows):<2d} {means}")
+    total_pass = sum(e.get("passed", False) for e in report)
+    print(
+        f"\nFILTER-DONE | {total_pass}/{len(report)} pairs pass -> filter_report.json"
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/edit_sft/generate_sources.py b/integrations/omnidreams/edit_sft/generate_sources.py
new file mode 100644
index 000000000..40fa62f08
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/generate_sources.py
@@ -0,0 +1,162 @@
+# 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.
+
+"""Source-video corpus for the edit-SFT data pipeline (Tier-2b, recipe A).
+
+Rolls the distilled model over every locally cached sample clip under its
+own prompt and saves the decoded videos plus a manifest — the *source*
+side of the (source, instruction, edited) triplets that JoyAI-Video-Edit
+produces downstream. Per-chunk latents are saved alongside so the SFT
+context-replay phase does not need to re-encode the videos.
+
+Env knobs: ``N_CLIPS``, ``N_CHUNKS``, ``SEED``, ``OUT_DIR``.
+
+Run from the repo root::
+
+    .venv/bin/python integrations/omnidreams/edit_sft/generate_sources.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+# Must land before the first CUDA allocation (co-tenant VRAM share).
+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
+from omnidreams.runner import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH
+
+from flashdreams.infra.config import derive_config
+from flashdreams.infra.runner_io import (
+    load_first_frame_tensor,
+    load_video_tensor,
+    write_video_tensor,
+)
+
+SAMPLES_ROOT = (
+    Path.home()
+    / ".cache/huggingface/hub/datasets--nvidia--omni-dreams-samples/snapshots"
+)
+N_CLIPS = int(os.environ.get("N_CLIPS", "20"))
+N_CHUNKS = int(os.environ.get("N_CHUNKS", "28"))
+SEED = int(os.environ.get("SEED", "42"))
+OUT_DIR = Path(
+    os.environ.get("OUT_DIR", "integrations/omnidreams/edit_sft/outputs/sources")
+)
+
+
+def _local_clips() -> list[tuple[str, Path, Path, str]]:
+    """(uuid, hdmap, first_frame, prompt) for every fully cached sample clip."""
+    clips = []
+    for prompt_path in sorted(SAMPLES_ROOT.glob("*/data/single_view/*/prompt.txt")):
+        clip_dir = prompt_path.parent
+        hdmaps = sorted(clip_dir.glob("*_hdmap.mp4"))
+        frames = sorted(clip_dir.glob("first_frame.png"))
+        if hdmaps and frames:
+            clips.append(
+                (clip_dir.name, hdmaps[0], frames[0], prompt_path.read_text().strip())
+            )
+    return clips
+
+
+@torch.no_grad()
+def main() -> None:
+    """Roll and save the source corpus."""
+    clips = _local_clips()[:N_CLIPS]
+    assert clips, f"no cached sample clips under {SAMPLES_ROOT}"
+    total_frames = 5 + (N_CHUNKS - 1) * 8
+
+    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),
+        ),
+    )
+    pipe = cfg.setup()
+    assert isinstance(pipe, OmnidreamsPipeline)
+    pipe = pipe.to("cuda")
+    device = pipe.device
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    manifest: list[dict] = []
+    for uuid, hdmap_path, frame_path, prompt in clips:
+        out_mp4 = OUT_DIR / f"{uuid}.mp4"
+        out_lat = OUT_DIR / f"{uuid}_latents.pt"
+        if out_mp4.exists() and out_lat.exists():
+            print(f"skip {uuid} (exists)", flush=True)
+        else:
+            hdmap = load_video_tensor(
+                hdmap_path,
+                pixel_height=DEFAULT_VIDEO_HEIGHT,
+                pixel_width=DEFAULT_VIDEO_WIDTH,
+                device=device,
+                dtype=torch.bfloat16,
+            )
+            if hdmap.shape[0] < total_frames:
+                print(f"skip {uuid}: only {hdmap.shape[0]} HDMap frames", flush=True)
+                continue
+            hdmap = hdmap[:total_frames][None, None]
+            first = load_first_frame_tensor(
+                frame_path,
+                pixel_height=DEFAULT_VIDEO_HEIGHT,
+                pixel_width=DEFAULT_VIDEO_WIDTH,
+                device=device,
+                dtype=torch.bfloat16,
+            )[None, None]
+
+            pipe.diffusion_model._rng = torch.Generator(device=device).manual_seed(SEED)
+            cache = pipe.initialize_cache(text=[[prompt]], image=first)
+            chunks, latents, start = [], [], 0
+            for ar_idx in range(N_CHUNKS):
+                num_frames = pipe.get_num_frames(ar_idx)
+                chunk = pipe.generate(
+                    ar_idx, cache, hdmap=hdmap[:, :, start : start + num_frames]
+                )
+                final_state = cache.final_state
+                assert final_state is not None
+                latents.append(final_state.clean_latent.detach().to("cpu"))
+                pipe.finalize(ar_idx, cache)
+                chunks.append(chunk[0, 0].float().cpu())
+                start += num_frames
+            video = torch.cat(chunks, dim=0)
+            write_video_tensor(video, out_mp4, fps=30, layout="tchw")
+            torch.save({"latents": latents, "prompt": prompt, "seed": SEED}, out_lat)
+            del cache
+            torch.cuda.empty_cache()
+            print(f"rolled {uuid}: {video.shape[0]} frames", flush=True)
+
+        manifest.append(
+            {
+                "uuid": uuid,
+                "video": out_mp4.name,
+                "latents": out_lat.name,
+                "prompt": prompt,
+                "n_chunks": N_CHUNKS,
+                "seed": SEED,
+            }
+        )
+
+    (OUT_DIR / "manifest.json").write_text(json.dumps(manifest, indent=2))
+    print(f"SOURCES-DONE | {len(manifest)} clips -> {OUT_DIR}/", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/edit_sft/precompute_style.py b/integrations/omnidreams/edit_sft/precompute_style.py
new file mode 100644
index 000000000..e4df5d539
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/precompute_style.py
@@ -0,0 +1,264 @@
+# 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.
+
+"""Precompute the style-SFT embeddings and latents (one shot).
+
+Three phases, so ``train_style_sft.py`` never loads the ~14 GB text
+encoder or touches ffmpeg (the precompute-embeddings pattern from the
+guidance-distillation trainer, extended with video encoding):
+
+A. Text + first-frame embeddings — every style-slug prompt
+   (:data:`~style_prompts.STYLE_PROMPTS`) and every source clip's own
+   prompt / first frame through the resident one-shot encoders.
+B. Per-clip HDMap latents — the clip's HDMap video stream-encoded with the
+   pipeline's per-AR-step Wan-VAE encoder in the AR chunk schedule (5
+   frames -> 2 latent frames, then 8 -> 2), patchified per chunk.
+C. Styled-target latents — every ``style_pairs/<uuid>__<slug>.mp4``
+   through the SAME encoder/schedule, so the flow-matching targets live in
+   the exact latent space the pipeline itself maps RGB into (the encoder
+   checkpoint is the one the first-frame I2V injection uses).
+
+All styled pairs found on disk are encoded; the VLM gate
+(``style_pairs/filter_report.json``, possibly still being written) is
+applied at training time, not here.
+
+Outputs (under ``edit_sft/outputs/``):
+
+- ``style_embeddings.pt``: ``{slug | "clip:<uuid>": [1, 1, L, D] bf16}``.
+- ``style_clip_assets.pt``: ``{"uuids", "prompts",
+  "image_embeddings": {uuid: [1, 1, 1, Cl, Hl, Wl] bf16}}``.
+- ``latents/<uuid>_hdmap_latents.pt``: ``{"hdmaps": [n_chunks x
+  [1, 1, L, D] bf16]}`` (patchified, ``transformer.latent_shape``).
+- ``latents/<uuid>__<slug>_latents.pt``: ``{"latents": [...]}``, same
+  layout.
+
+Latent files are skipped when present, so the script is resumable — the
+frame-count assert catches the silent-empty ffmpeg reads observed on this
+box (they hit once the process has grown to rollout size; the one-shot
+encoders are released before the first decode to shrink the process).
+
+Run from the flashdreams repo root (~20 GB VRAM during phase A)::
+
+    .venv/bin/python integrations/omnidreams/edit_sft/precompute_style.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+# Must land before the first CUDA allocation (co-tenant VRAM share).
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import torch
+from _host import build_pipeline
+from omnidreams.pipeline import OmnidreamsPipeline
+from omnidreams.runner import (
+    DEFAULT_VIDEO_HEIGHT,
+    DEFAULT_VIDEO_WIDTH,
+    _ensure_hf_single_view_example_data_synced,
+)
+from style_prompts import STYLE_PROMPTS, clip_key
+from torch import Tensor
+
+from flashdreams.infra.runner_io import load_first_frame_tensor, load_video_tensor
+
+## Configuration
+
+BASE = Path("integrations/omnidreams/edit_sft")
+OUT_DIR = Path(os.environ.get("OUT_DIR", str(BASE / "outputs")))
+SRC_DIR = OUT_DIR / "sources"
+STYLE_DIR = OUT_DIR / "style_pairs"
+LAT_DIR = OUT_DIR / "latents"
+
+
+def _sample_files(uuid: str) -> tuple[tuple[Path, ...], tuple[Path, ...]]:
+    """Cache-first ``((hdmap,), (first_frame,))`` resolution for a sample.
+
+    Vendored from the Clean Forcing pair builder (PR #398): the runner's
+    ``_ensure_hf_single_view_example_data_synced`` lists the HF repo per
+    clip even when every file is already local, and the shared IP rate
+    limit has killed pipeline stages before. 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 _styled_videos(known_uuids: set[str]) -> list[tuple[str, str, Path]]:
+    """``(uuid, slug, path)`` for every recognized styled pair on disk."""
+    out: list[tuple[str, str, Path]] = []
+    for path in sorted(STYLE_DIR.glob("*__*.mp4")):
+        uuid, slug = path.stem.split("__", 1)
+        if uuid not in known_uuids:
+            print(f"skip {path.name}: uuid not in the sources manifest", flush=True)
+            continue
+        if slug not in STYLE_PROMPTS:
+            print(f"skip {path.name}: slug {slug!r} not in STYLE_PROMPTS", flush=True)
+            continue
+        out.append((uuid, slug, path))
+    return out
+
+
+def encode_video_chunks(
+    pipe: OmnidreamsPipeline, video: Tensor, n_chunks: int
+) -> list[Tensor]:
+    """Stream-encode an RGB video in the AR chunk schedule.
+
+    Args:
+        pipe: Pipeline whose per-AR-step Wan-VAE encoder does the work
+            (fresh streaming cache per video).
+        video: ``[T, 3, H, W]`` pixels in ``[-1, 1]`` on CPU; ``T`` must
+            cover ``n_chunks`` chunks (5 + 8 * (n_chunks - 1) frames).
+        n_chunks: AR chunks to encode.
+
+    Returns:
+        Per-chunk patchified latents (``transformer.latent_shape`` layout,
+        the space ``generate_sources.py`` saved the source latents in),
+        bf16 on CPU.
+    """
+    transformer = pipe.diffusion_model.transformer
+    assert pipe.encoder is not None
+    enc_cache = pipe.encoder.initialize_autoregressive_cache()
+    chunks: list[Tensor] = []
+    start = 0
+    for ar_idx in range(n_chunks):
+        num_frames = pipe.get_num_frames(ar_idx)
+        pixels = video[start : start + num_frames][None, None].to(pipe.device)
+        z = pipe.encoder(input=pixels, autoregressive_index=ar_idx, cache=enc_cache)
+        chunks.append(transformer.patchify_and_maybe_split_cp(z).cpu())
+        start += num_frames
+    return chunks
+
+
+@torch.no_grad()
+def main() -> None:
+    """Encode embeddings (phase A), HDMap latents (B), styled latents (C)."""
+    dtype = torch.bfloat16
+    manifest = json.loads((SRC_DIR / "manifest.json").read_text())
+    uuids: list[str] = [e["uuid"] for e in manifest]
+    prompts: dict[str, str] = {e["uuid"]: e["prompt"] for e in manifest}
+    n_chunks = int(manifest[0]["n_chunks"])
+    total_frames = 5 + (n_chunks - 1) * 8
+    styled = _styled_videos(set(uuids))
+    assert styled, f"no recognized styled pairs under {STYLE_DIR}"
+
+    # First frames are PNGs (cv2, no ffmpeg fork) — safe to load lazily,
+    # but front-load them anyway so phase A never touches the filesystem
+    # mid-encoding.
+    firsts: dict[str, Tensor] = {}
+    for uuid in uuids:
+        _, (frame_path,) = _sample_files(uuid)
+        firsts[uuid] = load_first_frame_tensor(
+            frame_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device=torch.device("cpu"),
+            dtype=dtype,
+        )[None, :, None]  # [1, V=1, 1, C, H, W]
+
+    pipe = build_pipeline(with_oneshot_encoders=True)
+    device = pipe.device
+    assert pipe.text_encoder is not None  # with_oneshot_encoders=True
+
+    # Phase A: style + clip prompt embeddings and first-frame latents.
+    emb_by_text: dict[str, Tensor] = {}
+    prompt_embeddings: dict[str, Tensor] = {}
+    for slug in sorted({slug for _, slug, _ in styled}):
+        text = STYLE_PROMPTS[slug]
+        if text not in emb_by_text:  # _v2 slugs alias their base prompt
+            emb = torch.stack([pipe.text_encoder([text])], dim=0)  # [1, 1, L, D]
+            emb_by_text[text] = emb.to("cpu", dtype)
+        prompt_embeddings[slug] = emb_by_text[text]
+        print(
+            f"encoded style prompt {slug}: {tuple(prompt_embeddings[slug].shape)}",
+            flush=True,
+        )
+    image_embeddings: dict[str, Tensor] = {}
+    for uuid in uuids:
+        emb = pipe.precompute_embeddings(
+            text=[[prompts[uuid]]], image=firsts[uuid].to(device)
+        )
+        text_emb = emb["text_embeddings"]
+        image_emb = emb["image_embeddings"]
+        assert text_emb is not None and image_emb is not None
+        prompt_embeddings[clip_key(uuid)] = text_emb.to("cpu", dtype)
+        image_embeddings[uuid] = image_emb.to("cpu", dtype)
+        print(f"encoded clip {uuid}", flush=True)
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    torch.save(prompt_embeddings, OUT_DIR / "style_embeddings.pt")
+    torch.save(
+        {"uuids": uuids, "prompts": prompts, "image_embeddings": image_embeddings},
+        OUT_DIR / "style_clip_assets.pt",
+    )
+    # Shrink the process before the ffmpeg decodes below (fork hazard).
+    pipe.release_oneshot_encoders()
+
+    def encode_to(path: Path, video_path: Path, key: str) -> None:
+        video = load_video_tensor(
+            video_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device=torch.device("cpu"),
+            dtype=dtype,
+        )
+        assert video.shape[0] >= total_frames, (
+            f"{video_path.name}: {video.shape[0]} frames < {total_frames} "
+            "(short clip, or a silent-empty ffmpeg read — re-run to resume)."
+        )
+        chunks = encode_video_chunks(pipe, video[:total_frames], n_chunks)
+        torch.save({key: chunks, "n_chunks": n_chunks}, path)
+        print(f"encoded {path.name}: {n_chunks} chunks", flush=True)
+
+    LAT_DIR.mkdir(parents=True, exist_ok=True)
+    # Phase B: per-clip HDMap latents (conditioning for replay + training).
+    for uuid in uuids:
+        out = LAT_DIR / f"{uuid}_hdmap_latents.pt"
+        if out.exists():
+            print(f"skip {out.name} (exists)", flush=True)
+            continue
+        (hdmap_path,), _ = _sample_files(uuid)
+        encode_to(out, hdmap_path, "hdmaps")
+
+    # Phase C: styled-target latents.
+    for uuid, slug, path in styled:
+        out = LAT_DIR / f"{uuid}__{slug}_latents.pt"
+        if out.exists():
+            print(f"skip {out.name} (exists)", flush=True)
+            continue
+        encode_to(out, path, "latents")
+
+    print(
+        f"PRECOMPUTE-STYLE-DONE | {len(prompt_embeddings)} prompt embeddings, "
+        f"{len(uuids)} clips, {len(styled)} styled pairs -> {OUT_DIR}/",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/edit_sft/style_prompts.py b/integrations/omnidreams/edit_sft/style_prompts.py
new file mode 100644
index 000000000..7aafa4ea9
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/style_prompts.py
@@ -0,0 +1,86 @@
+# 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.
+
+"""Style-slug -> edit-prompt bank for the style-SFT run (Tier-2b, restyles).
+
+Each JoyAI style slug (``style_pairs/<uuid>__<slug>.mp4``) maps to the
+DECLARATIVE scene-description phrasing of its edit instruction
+(``style_pairs/instructions.json`` / the audition job files): the DiT's
+text conditioning was trained on scene descriptions, not imperative edit
+commands, so "Restyle the scene as a bright arcade racing game world ..."
+becomes "A bright arcade racing game world ...". The ``_v2`` audition
+slugs reran the same styles with a layout-preservation suffix; the suffix
+is a JoyAI-side control, not a scene property, so they share their base
+slug's prompt (the tensors alias, so ``torch.save`` stores them once).
+"""
+
+from __future__ import annotations
+
+STYLE_PROMPTS: dict[str, str] = {
+    "arcade_racer": (
+        "A bright arcade racing game world with exaggerated saturated "
+        "colors, clean stylized surfaces, and a cheerful sunny palette."
+    ),
+    "cartoon_cel": (
+        "A colorful cartoon video game world with cel-shaded surfaces, "
+        "bold outlines, and vivid saturated colors."
+    ),
+    "anime": (
+        "A Japanese anime style world with painted textures, soft "
+        "shading, and vibrant colors."
+    ),
+    "toy_world": (
+        "A miniature toy world where everything has glossy plastic "
+        "surfaces and bright playful colors."
+    ),
+    "lowpoly": (
+        "A low-poly stylized video game environment with flat-shaded "
+        "geometric surfaces and vivid colors."
+    ),
+    "comic_ink": (
+        "A comic book style world with bold black ink outlines, halftone "
+        "shading, and vivid flat colors."
+    ),
+    "pixel_art": (
+        "Retro 16-bit pixel art video game graphics with visible pixels "
+        "and a bright limited color palette."
+    ),
+    "cyberpunk_neon": (
+        "A neon-lit cyberpunk night city with glowing signs and rain-slicked streets."
+    ),
+    "watercolor": (
+        "A soft watercolor painting of the scene with visible brush "
+        "strokes and paper texture."
+    ),
+}
+STYLE_PROMPTS["anime_v2"] = STYLE_PROMPTS["anime"]
+STYLE_PROMPTS["cartoon_cel_v2"] = STYLE_PROMPTS["cartoon_cel"]
+
+DEFAULT_SKIP_STYLES: tuple[str, ...] = ("cyberpunk_neon", "watercolor")
+"""Styles rejected at the audition (scene re-imagined / washed out) —
+excluded from training by default even if the VLM filter passes them."""
+
+
+def clip_key(uuid: str) -> str:
+    """Return the embedding-dict key of a sample clip's own prompt.
+
+    Args:
+        uuid: ``nvidia/omni-dreams-samples`` single-view clip UUID.
+
+    Returns:
+        The key under which ``precompute_style.py`` stores the clip
+        prompt's text embeddings (style prompts live under the raw slug).
+    """
+    return f"clip:{uuid}"
diff --git a/integrations/omnidreams/edit_sft/train_style_sft.py b/integrations/omnidreams/edit_sft/train_style_sft.py
new file mode 100644
index 000000000..c3724f259
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/train_style_sft.py
@@ -0,0 +1,546 @@
+# 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.
+
+"""Edit-timestamped style-SFT trainer (Tier-2b of the live-edit hack).
+
+Teaches the realtime distilled student (2-step causal Cosmos DiT) to apply
+a global visual style mid-stream when the text prompt is swapped,
+imitating the offline JoyAI restyles (``PLAN.md`` recipe A, restricted to
+the global-style instruction bank). Supervision is teacher-forced
+flow matching in latent space on edit-timestamped clips:
+
+- **Context build**: chunks before the sampled swap chunk ``k`` are the
+  SOURCE rollout's own latents (``generate_sources.py``), replayed through
+  ``finalize_kv_cache`` at the host's ``context_noise`` (the vendored
+  ``_host.replay_history`` machinery) under base weights —
+  at deploy the LoRA window opens only at the swap. RoPE is absolute and
+  the KV window holds ``REPLAY_CHUNKS`` chunks, so replaying just the
+  window prefix reproduces a full rollout's visible state exactly.
+- **Edit supervision**: at ``k`` the text KV is rebuilt to the style
+  prompt (plain ``swap_text_kv`` rebuild — exactly the deploy swap of the
+  live-edit hook, PR #431); chunks ``[k, k + SPAN)`` are then supervised
+  toward the STYLED latents (the JoyAI output encoded by
+  ``precompute_style.py``). States
+  stay on-policy within each chunk (the student's own detached flow
+  advances the two-step trajectory, mirroring ``scheduler.sample``); the
+  v-target at every state is ``(z_t - x0_styled) / sigma``. A 0.5-weighted
+  context term (t=128) supervises the same forward whose K/V the commit
+  writes, and the STYLED latent is committed into the rolling KV — teacher
+  forcing proceeds on edited history, merged-weight commits matching the
+  deploy window semantics (the live-edit deploy hook, PR #431).
+- **No-op regularization**: ``PRE_CHUNKS`` supervised chunks right before
+  ``k`` target the SOURCE latents under the clip's own prompt (keeps
+  non-edit behavior identical), and with ``NOOP_PROB`` the whole window is
+  a no-op (swap to the clip's own prompt, source targets) so the style
+  must come from the text KV, not the weights.
+
+Data is gated by the style-mode VLM filter
+(``style_pairs/filter_report.json``) when present: ``passed`` pairs train
+anywhere, ``early_window_ok`` pairs only within the first
+:data:`EARLY_CHUNKS` chunks (~4 s — the window before streaming layout
+drift sets in on heavy styles). Without a report every pair on disk
+trains, with a warning.
+
+Host mechanics are the guidance-distillation trainer's (vendored into
+this directory): eager pipeline, LoRA r64 on the 8 attention projections
+(the checkpoint loads through the live-edit deploy hook's
+``TextEditLoRA``, PR #431, unchanged), functional self-attention on grad
+forwards, per-block checkpointing, and every term's backward
+immediately after its forward inside ``functional_attention()`` (the
+recompute must retake the non-writing path, and later stock forwards
+in-place mutate buffers the tape saved). No teacher forward exists here —
+the target is a fixed latent — so steps are lighter than Tier-2a's.
+
+VRAM (fits the ~65 GB co-tenant share): bf16 2B DiT ~4 GB + per-episode
+KV caches ~19 GB (28 blocks x 6-latent-frame window at 88x160 latents) +
+fp32 r64 LoRA/AdamW ~0.5 GB + one grad forward's checkpointed block
+inputs ~1.6 GB (immediate per-term backward keeps only one tape alive)
+— ~30 GB peak, eager.
+
+Run from the flashdreams repo root (after ``precompute_style.py``)::
+
+    STEPS=1000 .venv/bin/python \
+        integrations/omnidreams/edit_sft/train_style_sft.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+from typing import cast
+
+# Must land before the first CUDA allocation (co-tenant VRAM share).
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import numpy as np
+import torch
+from _host import CONTEXT_NOISE_SEED, build_pipeline, swap_text_kv
+from _lora import (
+    apply_lora,
+    lora_parameters,
+    save_lora,
+    set_lora_scale,
+    unwrap_compiled,
+)
+from _train_attn import functional_attention, patch_functional_attention
+from style_prompts import DEFAULT_SKIP_STYLES, STYLE_PROMPTS, clip_key
+from torch import Tensor
+
+from flashdreams.infra.diffusion.scheduler.fm import FlowMatchScheduler
+
+## Training configuration
+
+BASE = Path("integrations/omnidreams/edit_sft")
+OUT_DIR = BASE / "outputs"
+SRC_DIR = OUT_DIR / "sources"
+STYLE_DIR = OUT_DIR / "style_pairs"
+LAT_DIR = OUT_DIR / "latents"
+
+STEPS = int(os.environ.get("STEPS", "1000"))
+LR = float(os.environ.get("LR", "2e-4"))
+RANK = int(os.environ.get("RANK", "64"))
+"""r64 = the Tier-2a capacity that passed its gate; style is a stronger
+transform than weather, so start at the proven upper rank."""
+
+SEED = int(os.environ.get("SEED", "0"))
+HOLDOUT = int(os.environ.get("HOLDOUT", "2"))
+"""Clips (last of the sources manifest) reserved for the eval gate."""
+
+WARMUP = 40
+GRAD_CLIP = 1.0
+SAVE_EVERY = 100
+LOG_EVERY = 10
+EMA_DECAY = 0.98
+
+NOOP_PROB = float(os.environ.get("NOOP_PROB", "0.1"))
+"""Probability of a no-op window (swap to the clip's own prompt, source
+targets): anchors the restyle to the text KV instead of the weights."""
+
+PRE_CHUNKS = int(os.environ.get("PRE_CHUNKS", "1"))
+"""Supervised chunks immediately before ``k`` (source targets, original
+prompt) — no-op regularization at the swap boundary."""
+
+SPAN = int(os.environ.get("SPAN", "4"))
+"""Supervised chunks per window, ``[k, k + SPAN)`` (the deploy edit-window
+width; Tier-2a shipped 6-chunk windows, styled pairs support any width)."""
+
+SWAP_MIN, SWAP_MAX = (
+    int(os.environ.get("SWAP_MIN", "4")),
+    int(os.environ.get("SWAP_MAX", "20")),
+)
+"""Swap chunk ``k`` range — past the 3-chunk KV window fill; the upper end
+is clamped per pair so the window fits the pair's usable range."""
+
+EARLY_CHUNKS = int(os.environ.get("EARLY_CHUNKS", "15"))
+"""Usable chunks for ``early_window_ok``-only pairs: chunks 0..14 cover
+frames 0..116 (~3.9 s) — the pre-drift window the style filter certified."""
+
+REPLAY_CHUNKS = int(os.environ.get("REPLAY_CHUNKS", "3"))
+"""Source chunks replayed before the first supervised chunk. 3 chunks =
+the full ``window_size_t=6`` / ``len_t=2`` KV window, so the visible state
+equals a full-history replay (the drift corrector's mid-stream-start
+machinery, PR #398); raise only if the host's window grows."""
+
+CTX_WEIGHT = 0.5
+"""Weight of the finalize/context-forward (t=128) term."""
+
+SKIP_STYLES = frozenset(
+    s
+    for s in os.environ.get("SKIP_STYLES", ",".join(DEFAULT_SKIP_STYLES)).split(",")
+    if s
+)
+"""Slugs excluded from training regardless of the filter verdict."""
+
+LORA_TARGETS = (
+    "self_attn.q_proj",
+    "self_attn.k_proj",
+    "self_attn.v_proj",
+    "self_attn.output_proj",
+    "cross_attn.q_proj",
+    "cross_attn.k_proj",
+    "cross_attn.v_proj",
+    "cross_attn.output_proj",
+)
+"""Drift-corrector recipe + cross-attn (the edit signal enters there).
+Vendored from the guidance-distillation trainer (deferred from PR #431) so
+this pipeline is self-contained; consolidate when that trainer lands."""
+
+
+def checkpoint_blocks(network) -> None:
+    """Route every DiT block ``forward`` through gradient checkpointing.
+
+    Vendored from the guidance-distillation trainer (deferred from
+    PR #431); consolidate when that trainer lands. Per-instance overrides
+    (not wrapper modules) so the network loop's ``isinstance(block, Block)``
+    assertion keeps passing. Requires the functional-attention toggle on
+    grad passes: recomputation must be side-effect free and must retake the
+    same code path, so every backward runs inside
+    ``functional_attention()``. No-op under ``no_grad`` passes (rollout,
+    finalize).
+
+    Args:
+        network: The unwrapped ``CosmosDiTNetwork``.
+    """
+    from torch.utils.checkpoint import checkpoint
+
+    def wrap(fn):
+        def ckpt_fn(*args, _inner=fn, **kwargs):
+            if not torch.is_grad_enabled():
+                return _inner(*args, **kwargs)
+            return checkpoint(_inner, *args, use_reentrant=False, **kwargs)
+
+        return ckpt_fn
+
+    for block in network.blocks:
+        block.forward = wrap(block.forward)
+
+
+def load_pairs(train_uuids: set[str], n_chunks: int) -> list[tuple[str, str, int]]:
+    """Collect trainable ``(uuid, slug, max_chunk)`` styled pairs.
+
+    ``max_chunk`` is the exclusive bound on supervised chunk indices:
+    the clip length for filter-``passed`` pairs, :data:`EARLY_CHUNKS` for
+    ``early_window_ok``-only pairs. Pairs missing precomputed latents, on
+    held-out clips, in :data:`SKIP_STYLES`, or rejected by the filter are
+    dropped. Without a filter report every remaining pair trains at full
+    range (with a warning) — rerun once the filter job lands.
+    """
+    report_path = STYLE_DIR / "filter_report.json"
+    report: dict[str, dict] = {}
+    if report_path.exists():
+        for entry in json.loads(report_path.read_text()):
+            if "output" in entry:
+                report[entry["output"]] = entry
+    else:
+        print(
+            f"WARNING: {report_path} not found — training on ALL styled "
+            "pairs unfiltered.",
+            flush=True,
+        )
+
+    pairs: list[tuple[str, str, int]] = []
+    for path in sorted(STYLE_DIR.glob("*__*.mp4")):
+        uuid, slug = path.stem.split("__", 1)
+        if uuid not in train_uuids or slug in SKIP_STYLES:
+            continue
+        if slug not in STYLE_PROMPTS:
+            print(f"skip {path.name}: slug {slug!r} not in STYLE_PROMPTS", flush=True)
+            continue
+        if not (LAT_DIR / f"{path.stem}_latents.pt").exists():
+            print(f"skip {path.name}: no precomputed latents", flush=True)
+            continue
+        if report_path.exists():
+            entry = report.get(path.name)
+            if entry is None or "error" in entry:
+                continue
+            if entry.get("passed"):
+                max_chunk = n_chunks
+            elif entry.get("early_window_ok"):
+                max_chunk = EARLY_CHUNKS
+            else:
+                continue
+        else:
+            max_chunk = n_chunks
+        if min(SWAP_MAX, max_chunk - SPAN) < SWAP_MIN:
+            print(f"skip {path.name}: usable range < SWAP_MIN", flush=True)
+            continue
+        pairs.append((uuid, slug, max_chunk))
+    return pairs
+
+
+def main() -> None:
+    """Run the teacher-forced style-SFT loop."""
+    rng = np.random.default_rng(SEED)
+    manifest = json.loads((SRC_DIR / "manifest.json").read_text())
+    uuids: list[str] = [e["uuid"] for e in manifest]
+    n_chunks = int(manifest[0]["n_chunks"])
+    assert len(uuids) > HOLDOUT, f"{len(uuids)} clips can't spare {HOLDOUT} held out"
+    train_uuids = uuids[: len(uuids) - HOLDOUT]
+    pairs = load_pairs(set(train_uuids), n_chunks)
+    assert pairs, (
+        f"no trainable styled pairs under {STYLE_DIR} (run precompute_style.py "
+        "first; check filter_report.json / SKIP_STYLES)."
+    )
+
+    prompt_emb = torch.load(
+        OUT_DIR / "style_embeddings.pt", map_location="cpu", weights_only=False
+    )
+    assets = torch.load(
+        OUT_DIR / "style_clip_assets.pt", map_location="cpu", weights_only=False
+    )
+
+    # Latents are small (~0.9 MB per chunk) — front-load everything on CPU.
+    src_lat: dict[str, list[Tensor]] = {}
+    hd_lat: dict[str, list[Tensor]] = {}
+    for uuid in sorted({uuid for uuid, _, _ in pairs}):
+        src = torch.load(
+            SRC_DIR / f"{uuid}_latents.pt", map_location="cpu", weights_only=False
+        )["latents"]
+        hd = torch.load(
+            LAT_DIR / f"{uuid}_hdmap_latents.pt",
+            map_location="cpu",
+            weights_only=False,
+        )["hdmaps"]
+        assert len(src) == n_chunks and len(hd) == n_chunks, (
+            f"clip {uuid}: {len(src)} source / {len(hd)} hdmap chunks != {n_chunks}"
+        )
+        src_lat[uuid], hd_lat[uuid] = src, hd
+    tgt_lat: dict[tuple[str, str], list[Tensor]] = {}
+    for uuid, slug, _ in pairs:
+        tgt = torch.load(
+            LAT_DIR / f"{uuid}__{slug}_latents.pt",
+            map_location="cpu",
+            weights_only=False,
+        )["latents"]
+        assert len(tgt) == n_chunks, (
+            f"pair {uuid}__{slug}: {len(tgt)} styled chunks != {n_chunks}"
+        )
+        tgt_lat[(uuid, slug)] = tgt
+
+    pipe = build_pipeline(with_oneshot_encoders=False)
+    assert pipe.V_group is None, "single-GPU trainer; run without CP"
+    device = pipe.device
+    dtype = torch.bfloat16
+    dm = pipe.diffusion_model
+    transformer = dm.transformer
+    # The base ``Scheduler`` annotation hides the FM schedule buffers.
+    scheduler = cast(FlowMatchScheduler, dm.scheduler)
+    timesteps = scheduler.denoising_step_list  # [1000, 450] on the chunk2 host
+    sigmas = scheduler.denoising_sigmas
+    n_steps = int(timesteps.shape[0])
+    ctx_t = torch.tensor(float(dm.config.context_noise), device=device, dtype=dtype)
+    # Context-noise sigma from the warped training table (the add_noise
+    # snap, resolved once so the v-target can be built from a known eps).
+    ctx_idx = torch.argmin((scheduler._full_timesteps - ctx_t.float()).abs())
+    sigma_ctx = scheduler._full_sigmas[ctx_idx]
+
+    network = unwrap_compiled(transformer.network)
+    network.requires_grad_(False)  # frozen base; only the LoRA A/B path trains
+    wrapped = apply_lora(network, rank=RANK, targets=LORA_TARGETS)
+    params = lora_parameters(network)
+    n_pass = sum(mc == n_chunks for _, _, mc in pairs)
+    print(
+        f"LoRA r{RANK} on {len(wrapped)} projections | "
+        f"{sum(p.numel() for p in params) / 1e6:.2f}M params | "
+        f"{len(pairs)} pairs ({n_pass} full + {len(pairs) - n_pass} early-window) "
+        f"over {len(src_lat)} clips ({HOLDOUT} held out) | "
+        f"styles {sorted({slug for _, slug, _ in pairs})}",
+        flush=True,
+    )
+    patch_functional_attention()
+    checkpoint_blocks(network)
+    opt = torch.optim.AdamW(params, lr=LR)
+
+    def predict_v(tc, z_t: Tensor, timestep: Tensor, hd_p: Tensor) -> Tensor:
+        """One functional-attention (non-writing) forward -> fp32 flow."""
+        with functional_attention():
+            flow = transformer.predict_flow(
+                noisy_latent=z_t, timestep=timestep, cache=tc, input=hd_p
+            )
+        return flow.float()
+
+    def train_chunk(
+        tc,
+        j: int,
+        x0_tgt: Tensor,
+        hd_p: Tensor,
+        weight: float,
+        in_window: bool,
+        model_rng: torch.Generator,
+    ) -> dict[str, float]:
+        """Supervise chunk ``j`` toward ``x0_tgt`` and commit it to the KV.
+
+        Mirrors ``scheduler.sample``: the student's own (detached) flow
+        advances the two-step trajectory, and each state's v-target points
+        at the teacher-forced ``x0_tgt`` (``(z_t - x0) / sigma``; at
+        t=1000, sigma=1, that is exactly ``eps - x0``). Each term backwards
+        immediately inside ``functional_attention()``. The KV commit is
+        the stock finalize on ``x0_tgt`` re-noised at t=128 — the same
+        state the 0.5-weighted context term supervises — run at LoRA
+        scale 1 inside the edit window (deploy commits merged) and scale 0
+        before it (deploy commits base pre-swap).
+
+        Args:
+            tc: Transformer cache, positioned at chunk ``j - 1``.
+            j: Supervised AR chunk index.
+            x0_tgt: fp32 patchified target latent (styled or source).
+            hd_p: Patchified HDMap conditioning for chunk ``j``.
+            weight: Per-chunk loss weight (1 / supervised chunks).
+            in_window: Whether ``j`` is inside the edit window (>= k).
+            model_rng: The rollout's noise generator.
+
+        Returns:
+            Unweighted per-term losses (t1000 / t450 / ctx floats).
+        """
+        tc.start(j)
+        losses: dict[str, float] = {}
+        noisy = torch.randn(
+            transformer.latent_shape, device=device, dtype=dtype, generator=model_rng
+        )
+        clean: Tensor | None = None
+        for i in range(n_steps):
+            sigma = sigmas[i]
+            timestep = timesteps[i].to(dtype=dtype)
+            if i > 0:
+                assert clean is not None
+                noise = torch.empty_like(noisy).normal_(generator=model_rng)
+                noisy = ((1.0 - sigma) * clean + sigma * noise).to(dtype)
+            v_student = predict_v(tc, noisy, timestep, hd_p)
+            v_target = (noisy.float() - x0_tgt) / sigma
+            term = (v_student - v_target).square().mean()
+            with functional_attention():
+                (weight * term / n_steps).backward()
+            losses[f"t{int(timesteps[i].item())}"] = float(term)
+            clean = noisy - sigma * v_student.detach()  # fp32 via promotion
+
+        # Context term at the state the commit below writes: for a known
+        # eps, the v-target is exactly eps - x0.
+        eps = torch.randn(
+            transformer.latent_shape,
+            device=device,
+            dtype=torch.float32,
+            generator=model_rng,
+        )
+        z_ctx = ((1.0 - sigma_ctx) * x0_tgt + sigma_ctx * eps).to(dtype)
+        v_ctx = predict_v(tc, z_ctx, ctx_t, hd_p)
+        term = (v_ctx - (eps - x0_tgt)).square().mean()
+        with functional_attention():
+            (weight * CTX_WEIGHT * term).backward()
+        losses["ctx"] = float(term)
+
+        # All backwards done -> the stock finalize (buffer write + index
+        # advance) commits the teacher-forced target into the history.
+        if not in_window:
+            set_lora_scale(network, 0.0)
+        with torch.no_grad():
+            transformer.finalize_kv_cache(
+                noisy_latent=z_ctx.detach(), timestep=ctx_t, cache=tc, input=hd_p
+            )
+        if not in_window:
+            set_lora_scale(network, 1.0)
+        tc.finalize(j)
+        return losses
+
+    def train_step() -> tuple[dict[str, float], str]:
+        """One edit-timestamped episode -> (losses, episode).
+
+        Backwards happen inside (per term); the caller owns ``zero_grad``
+        / clip / ``opt.step``.
+        """
+        uuid, slug, max_chunk = pairs[int(rng.integers(len(pairs)))]
+        no_op = bool(rng.random() < NOOP_PROB)
+        k_hi = min(SWAP_MAX, max_chunk - SPAN)
+        k = int(rng.integers(SWAP_MIN, k_hi + 1))
+        n_pre = min(PRE_CHUNKS, k)
+        j0 = k - n_pre
+
+        dm._rng = torch.Generator(device=device).manual_seed(int(rng.integers(2**31)))
+        cache = pipe.initialize_cache_from_embeddings(
+            text_embeddings=prompt_emb[clip_key(uuid)],
+            image_embeddings=assets["image_embeddings"][uuid],
+        )
+        tc = cache.transformer_cache
+        end = k + SPAN
+        src = [x.to(device, dtype) for x in src_lat[uuid][:end]]
+        tgt = [x.to(device, dtype) for x in tgt_lat[(uuid, slug)][:end]]
+        hd = [x.to(device, dtype) for x in hd_lat[uuid][:end]]
+
+        # Context build: replay the KV-window prefix from SOURCE latents at
+        # base weights (deploy: pre-window chunks run the unmodified
+        # network). Per-chunk seeded context noise = the replay_history
+        # contract; the mid-stream start is the drift corrector's machinery.
+        set_lora_scale(network, 0.0)
+        start = max(0, j0 - REPLAY_CHUNKS)
+        with torch.no_grad():
+            for bc in tc.network_cache.block_caches:
+                bc.self_attn._prev_chunk_idx = start - 1
+            for j in range(start, j0):
+                g = torch.Generator(device=device).manual_seed(CONTEXT_NOISE_SEED + j)
+                noisy = scheduler.add_noise(src[j], ctx_t, rng=g)
+                tc.start(j)
+                transformer.finalize_kv_cache(
+                    noisy_latent=noisy, timestep=ctx_t, cache=tc, input=hd[j]
+                )
+                tc.finalize(j)
+        set_lora_scale(network, 1.0)
+
+        model_rng = dm.rng
+        assert model_rng is not None
+        weight = 1.0 / (n_pre + SPAN)
+        sums: dict[str, float] = {}
+        total = 0.0
+        for j in range(j0, end):
+            if j == k:
+                # The deploy swap: plain cross-attn text-KV rebuild, the
+                # same rebuild the live-edit deploy hook (PR #431) runs on
+                # a plain prompt swap (no-op windows swap to the clip's own
+                # prompt — same machinery, identical KV — so style must be
+                # read from the text).
+                emb = prompt_emb[clip_key(uuid) if no_op else slug]
+                swap_text_kv(network, tc, emb)
+            x0_tgt = (src[j] if (no_op or j < k) else tgt[j]).float()
+            losses = train_chunk(tc, j, x0_tgt, hd[j], weight, j >= k, model_rng)
+            for key, value in losses.items():
+                sums[key] = sums.get(key, 0.0) + weight * value
+            total += weight * (
+                sum(v for key, v in losses.items() if key != "ctx") / n_steps
+                + CTX_WEIGHT * losses["ctx"]
+            )
+
+        del cache
+        sums["total"] = total
+        episode = (
+            f"{uuid[:8]} {'no_op:' if no_op else ''}{slug[:16]} "
+            f"k={k} pre={n_pre} span={SPAN} max={max_chunk}"
+        )
+        return sums, episode
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    torch.set_grad_enabled(True)
+    ema: float | None = None
+    for step in range(1, STEPS + 1):
+        for pg in opt.param_groups:
+            pg["lr"] = LR * min(1.0, step / WARMUP)
+        opt.zero_grad()
+        losses, episode = train_step()
+        torch.nn.utils.clip_grad_norm_(params, GRAD_CLIP)
+        opt.step()
+        torch.cuda.empty_cache()  # the episode's KV caches died with its frame
+        ema = (
+            losses["total"]
+            if ema is None
+            else EMA_DECAY * ema + (1.0 - EMA_DECAY) * losses["total"]
+        )
+        if step % LOG_EVERY == 0 or step == 1:
+            terms = " ".join(f"{k} {v:.4f}" for k, v in losses.items() if k != "total")
+            print(
+                f"step {step:5d} | loss {losses['total']:.4f} (ema {ema:.4f})"
+                f" | {terms} | {episode}",
+                flush=True,
+            )
+        if step % SAVE_EVERY == 0 or step == STEPS:
+            path = OUT_DIR / f"lora_style_step{step}.pt"
+            save_lora(network, path)
+            print(f"saved {path}", flush=True)
+
+    print(f"TRAIN-STYLE-SFT-DONE | final loss ema {ema:.4f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

From 0ce92eb5866c313651ba98175d216d0991636f8f Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Thu, 13 Aug 2026 22:29:36 +0000
Subject: [PATCH 2/8] Fix pre-swap text cache built at LoRA scale 1

initialize_cache_from_embeddings ran before the scale reset, so from the
second optimizer step the original prompt's cross-attn KV carried LoRA
deltas that the closed-window deploy path never has (reviewer-caught).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 integrations/omnidreams/edit_sft/train_style_sft.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/integrations/omnidreams/edit_sft/train_style_sft.py b/integrations/omnidreams/edit_sft/train_style_sft.py
index c3724f259..dcae6100c 100644
--- a/integrations/omnidreams/edit_sft/train_style_sft.py
+++ b/integrations/omnidreams/edit_sft/train_style_sft.py
@@ -451,6 +451,11 @@ def train_step() -> tuple[dict[str, float], str]:
         j0 = k - n_pre
 
         dm._rng = torch.Generator(device=device).manual_seed(int(rng.integers(2**31)))
+        # Build the original prompt's text KV at base weights: at deploy the
+        # LoRA window is closed until the swap, and a scale left at 1 from
+        # the previous episode would bake LoRA deltas into the pre-swap
+        # cache (reviewer-caught).
+        set_lora_scale(network, 0.0)
         cache = pipe.initialize_cache_from_embeddings(
             text_embeddings=prompt_emb[clip_key(uuid)],
             image_embeddings=assets["image_embeddings"][uuid],

From d1b97d9497d55a099f30aca6d92634c94f103ac1 Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Fri, 14 Aug 2026 04:49:19 +0000
Subject: [PATCH 3/8] Add style-maintenance episodes to keep long windows
 stable

The step-1600 LoRA nails the style inside its trained 4-chunk span
(bit-exact pre-swap, strong clean restyle at +1..+4 chunks) but blurs
out at +7-10: trained only on unstyled history, it learns a constant
style push that compounds through the KV commits. Maintenance episodes
(MAINT_PROB=0.3) replay styled history at scale 1 with styled targets,
teaching the window's fixed point.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../omnidreams/edit_sft/train_style_sft.py    | 37 ++++++++++++++-----
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/integrations/omnidreams/edit_sft/train_style_sft.py b/integrations/omnidreams/edit_sft/train_style_sft.py
index dcae6100c..397a48cca 100644
--- a/integrations/omnidreams/edit_sft/train_style_sft.py
+++ b/integrations/omnidreams/edit_sft/train_style_sft.py
@@ -129,6 +129,7 @@
 EMA_DECAY = 0.98
 
 NOOP_PROB = float(os.environ.get("NOOP_PROB", "0.1"))
+MAINT_PROB = float(os.environ.get("MAINT_PROB", "0.3"))
 """Probability of a no-op window (swap to the clip's own prompt, source
 targets): anchors the restyle to the text KV instead of the weights."""
 
@@ -445,9 +446,16 @@ def train_step() -> tuple[dict[str, float], str]:
         """
         uuid, slug, max_chunk = pairs[int(rng.integers(len(pairs)))]
         no_op = bool(rng.random() < NOOP_PROB)
+        # Steady-state ("maintenance") episode: the swap happened long ago —
+        # styled history, style text, LoRA on, styled targets. Without this
+        # the LoRA only ever sees unstyled history and learns a constant
+        # style PUSH; deployed past the trained span the push compounds
+        # through the KV commits and the world blurs out (observed at
+        # +7-10 chunks on the step-1600 run).
+        maint = (not no_op) and bool(rng.random() < MAINT_PROB)
         k_hi = min(SWAP_MAX, max_chunk - SPAN)
         k = int(rng.integers(SWAP_MIN, k_hi + 1))
-        n_pre = min(PRE_CHUNKS, k)
+        n_pre = 0 if maint else min(PRE_CHUNKS, k)
         j0 = k - n_pre
 
         dm._rng = torch.Generator(device=device).manual_seed(int(rng.integers(2**31)))
@@ -466,18 +474,25 @@ def train_step() -> tuple[dict[str, float], str]:
         tgt = [x.to(device, dtype) for x in tgt_lat[(uuid, slug)][:end]]
         hd = [x.to(device, dtype) for x in hd_lat[uuid][:end]]
 
-        # Context build: replay the KV-window prefix from SOURCE latents at
-        # base weights (deploy: pre-window chunks run the unmodified
-        # network). Per-chunk seeded context noise = the replay_history
+        if maint:
+            # In-window steady state: text already swapped, history styled,
+            # commits at LoRA scale 1 (deploy's use_lora window semantics).
+            swap_text_kv(network, tc, prompt_emb[slug])
+
+        # Context build: replay the KV-window prefix (SOURCE latents at base
+        # weights for swap episodes — deploy runs pre-window chunks on the
+        # unmodified network; STYLED latents at scale 1 for maintenance
+        # episodes). Per-chunk seeded context noise = the replay_history
         # contract; the mid-stream start is the drift corrector's machinery.
-        set_lora_scale(network, 0.0)
+        replay = tgt if maint else src
+        set_lora_scale(network, 1.0 if maint else 0.0)
         start = max(0, j0 - REPLAY_CHUNKS)
         with torch.no_grad():
             for bc in tc.network_cache.block_caches:
                 bc.self_attn._prev_chunk_idx = start - 1
             for j in range(start, j0):
                 g = torch.Generator(device=device).manual_seed(CONTEXT_NOISE_SEED + j)
-                noisy = scheduler.add_noise(src[j], ctx_t, rng=g)
+                noisy = scheduler.add_noise(replay[j], ctx_t, rng=g)
                 tc.start(j)
                 transformer.finalize_kv_cache(
                     noisy_latent=noisy, timestep=ctx_t, cache=tc, input=hd[j]
@@ -491,7 +506,7 @@ def train_step() -> tuple[dict[str, float], str]:
         sums: dict[str, float] = {}
         total = 0.0
         for j in range(j0, end):
-            if j == k:
+            if j == k and not maint:
                 # The deploy swap: plain cross-attn text-KV rebuild, the
                 # same rebuild the live-edit deploy hook (PR #431) runs on
                 # a plain prompt swap (no-op windows swap to the clip's own
@@ -499,8 +514,10 @@ def train_step() -> tuple[dict[str, float], str]:
                 # read from the text).
                 emb = prompt_emb[clip_key(uuid) if no_op else slug]
                 swap_text_kv(network, tc, emb)
-            x0_tgt = (src[j] if (no_op or j < k) else tgt[j]).float()
-            losses = train_chunk(tc, j, x0_tgt, hd[j], weight, j >= k, model_rng)
+            x0_tgt = (src[j] if (no_op or (not maint and j < k)) else tgt[j]).float()
+            losses = train_chunk(
+                tc, j, x0_tgt, hd[j], weight, maint or j >= k, model_rng
+            )
             for key, value in losses.items():
                 sums[key] = sums.get(key, 0.0) + weight * value
             total += weight * (
@@ -511,7 +528,7 @@ def train_step() -> tuple[dict[str, float], str]:
         del cache
         sums["total"] = total
         episode = (
-            f"{uuid[:8]} {'no_op:' if no_op else ''}{slug[:16]} "
+            f"{uuid[:8]} {'no_op:' if no_op else 'maint:' if maint else ''}{slug[:16]} "
             f"k={k} pre={n_pre} span={SPAN} max={max_chunk}"
         )
         return sums, episode

From bbe4d12d384bf83013ba029fdf53be5ffe324249 Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Tue, 18 Aug 2026 20:04:06 +0000
Subject: [PATCH 4/8] Fix pre-swap cache LoRA-scale leak; add v4 episodes and
 teacher-regen v5 recipe

The pre-swap text cache was built with the LoRA already at scale 1, so
'base' rollout history was subtly styled (review finding): initialize
the cache at scale 0 and only enable the adapter at the swap. Add v4
self-consistency maintenance episodes (V4_MAINT/V4_EARLY; evaluated as
a wash vs v3 - kept behind the knob, default documented) and
gen_teacher_styled_targets.py: re-roll JoyAI-restyled clips through the
35-step bidirectional teacher (SDEdit at moderate sigma) to produce
sharp styled targets for the v5 corpus; precompute reads either corpus
directory.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../edit_sft/gen_teacher_styled_targets.py    | 446 ++++++++++++++++++
 .../omnidreams/edit_sft/precompute_style.py   | 201 ++++----
 .../omnidreams/edit_sft/train_style_sft.py    | 240 ++++++----
 3 files changed, 712 insertions(+), 175 deletions(-)
 create mode 100644 integrations/omnidreams/edit_sft/gen_teacher_styled_targets.py

diff --git a/integrations/omnidreams/edit_sft/gen_teacher_styled_targets.py b/integrations/omnidreams/edit_sft/gen_teacher_styled_targets.py
new file mode 100644
index 000000000..b873d1923
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/gen_teacher_styled_targets.py
@@ -0,0 +1,446 @@
+# 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.
+
+"""Regenerate the JoyAI styled targets through the 35-step bidirectional teacher.
+
+Issue 1's final quality lever: every style-LoRA version was capped by SOFT
+training targets — the JoyAI RGB restyles re-encoded through the VAE. This
+script produces SHARP, temporally consistent styled targets by SDEdit-style
+renoise-refine through the TEACHER (which is photoreal-trained and will not
+restyle from a prompt alone — the low-noise styled latent carries the style
+and layout, the teacher re-renders them at its native sharpness):
+
+1. Encode a filter-passed JoyAI styled video
+   (``style_pairs/<uuid>__<slug>.mp4``) with the teacher's Wan-VAE encoder
+   into ONE bidirectional chunk48 window. ``len_t = 48`` LATENT frames =
+   189 pixel frames (``1 + 47 * 4``) — exactly student chunks 0..23, which
+   covers the trainer's maximum supervised chunk
+   (``SWAP_MAX + SPAN - 1 = 23``), so a single window per clip suffices.
+2. Renoise the styled latent at a mid sigma of the teacher's 35-step
+   shift-5 UniPC schedule and run the REMAINDER of the schedule (partial
+   UniPC restart: order-1 warmup at the entry step, exactly the baked
+   coefficients afterwards), CFG 3.0 with the style prompt
+   (``style_prompts.py`` declarative phrasing) + the clip's HDMap
+   conditioning.
+3. Decode through the light TAE decoder (mandatory — the full Wan-VAE
+   decode OOMs on 48-latent-frame chunks) and write
+   ``<uuid>__<slug>.mp4`` + a styled-vs-teacher contact sheet.
+
+The teacher pipeline setup (chunk48 config, checkpoint key filter, light
+TAE decoder) mirrors ``spawn_distill/gen_teacher_pairs.py``; the
+renoise-refine mechanics mirror ``scripts/probe_composite_refine.py``.
+
+Env knobs: ``SIGMAS`` (comma list; >1 value = sigma sweep, suffixing
+outputs with ``__s<sigma>``), ``STYLES`` (default
+``arcade_racer,comic_ink``), ``LIMIT`` (max pairs, 0 = all), ``ONLY``
+(comma list of ``<uuid>__<slug>`` stems), ``OUT_DIR`` (default
+``outputs/style_pairs_teacher``), ``SEED``, ``MAX_USED_GB`` (VRAM
+co-tenant gate, default 100).
+
+Run from the repo root (sweep first, then batch)::
+
+    SIGMAS=0.35,0.5,0.65 LIMIT=1 OUT_DIR=.../style_pairs_teacher_sweep \\
+        .venv/bin/python integrations/omnidreams/edit_sft/gen_teacher_styled_targets.py
+    SIGMAS=<chosen> \\
+        .venv/bin/python integrations/omnidreams/edit_sft/gen_teacher_styled_targets.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import time
+import zlib
+from pathlib import Path
+
+# Must land before the first CUDA allocation (co-tenant VRAM share).
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import numpy as np
+import torch
+from omnidreams.config import (
+    SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE,
+    SV_35STEPS_CHUNK48_LOC48_COSMOS2_2B_RES720P_30FPS_HDMAP_VAE_MADS1M,
+)
+from omnidreams.pipeline import OmnidreamsPipeline
+from omnidreams.transformer.impl.network import (
+    CosmosDiTNetwork,
+    CosmosDiTNetworkConfig,
+)
+from style_prompts import STYLE_PROMPTS
+from torch import Tensor
+
+from flashdreams.infra.config import derive_config
+from flashdreams.infra.runner_io import read_video_rgb, write_video_tensor
+
+DEFAULT_VIDEO_HEIGHT = 704
+DEFAULT_VIDEO_WIDTH = 1280
+
+
+def _load_video(
+    path: Path,
+    *,
+    pixel_height: int,
+    pixel_width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> Tensor:
+    """Load + resize a video to ``[T, C, H, W]`` in ``[-1, 1]``."""
+    import cv2  # noqa: PLC0415
+
+    video_np = read_video_rgb(path)
+    if video_np.shape[1:3] != (pixel_height, pixel_width):
+        video_np = np.stack(
+            [cv2.resize(f, (pixel_width, pixel_height)) for f in video_np], axis=0
+        )
+    tensor = torch.from_numpy(video_np).to(dtype=dtype, device=device) / 127.5 - 1.0
+    return tensor.permute(0, 3, 1, 2).contiguous()
+
+
+## Configuration
+
+BASE = Path("integrations/omnidreams/edit_sft")
+STYLE_DIR = Path(os.environ.get("STYLE_DIR", str(BASE / "outputs/style_pairs")))
+OUT_DIR = Path(os.environ.get("OUT_DIR", str(BASE / "outputs/style_pairs_teacher")))
+
+SIGMAS = [float(s) for s in os.environ.get("SIGMAS", "0.35,0.5,0.65").split(",")]
+STYLES = frozenset(os.environ.get("STYLES", "arcade_racer,comic_ink").split(","))
+LIMIT = int(os.environ.get("LIMIT", "0"))
+ONLY = frozenset(s for s in os.environ.get("ONLY", "").split(",") if s)
+SEED = int(os.environ.get("SEED", "42"))
+MAX_USED_GB = float(os.environ.get("MAX_USED_GB", "100"))
+
+TEACHER_FRAMES = 189
+"""One bidirectional chunk48 window: 48 latent frames = 1 + 47 * 4 pixels."""
+
+CONTACT_FRAMES = (0, 40, 90, 140, 188)
+
+_HF_SNAPSHOT = (
+    Path.home()
+    / ".cache/huggingface/hub/models--nvidia--omni-dreams-models/snapshots"
+    / "253701787e2f99efec31aaab665d0d9e0cc1eb4a/single_view"
+)
+_TEACHER_CKPT = _HF_SNAPSHOT / "teacher/3b4c21d0-7b77-4694-9d9d-6ac9b6dbba51_model.pt"
+
+
+def _network_key_filter_transform():
+    """``net.``-prefix strip + drop training-stack extras (teacher probe)."""
+    with torch.device("meta"):
+        reference = CosmosDiTNetwork(CosmosDiTNetworkConfig(additional_concat_ch=16))
+    valid = set(reference.state_dict().keys())
+    del reference
+
+    def transform(state_dict):
+        out = {}
+        for key, value in state_dict.items():
+            key = key[len("net.") :] if key.startswith("net.") else key
+            if key in valid:
+                out[key] = value
+        assert len(out) == len(valid), (
+            f"teacher checkpoint covers {len(out)}/{len(valid)} network keys"
+        )
+        return out
+
+    return transform
+
+
+def _wait_for_vram(max_used_gb: float) -> None:
+    """Block until the co-tenant leaves enough VRAM for the teacher (~150 GB)."""
+    while True:
+        out = subprocess.run(
+            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
+            capture_output=True,
+            text=True,
+            check=True,
+        ).stdout
+        used_gb = max(float(line) for line in out.strip().splitlines()) / 1024.0
+        if used_gb <= max_used_gb:
+            return
+        print(f"GPU busy ({used_gb:.0f} GB used > {max_used_gb:.0f}); waiting 60 s")
+        time.sleep(60)
+
+
+def _build_pipeline() -> OmnidreamsPipeline:
+    cfg = derive_config(
+        SV_35STEPS_CHUNK48_LOC48_COSMOS2_2B_RES720P_30FPS_HDMAP_VAE_MADS1M,
+        name="omnidreams-sv-teacher-styled-targets-chunk48",
+        enable_sync_and_profile=False,
+        # Full Wan-VAE decode OOMs on 48-latent-frame chunks; the light TAE
+        # decoder shares the latent space and decode is output-only.
+        decoder=SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE.decoder,
+        diffusion_model=dict(
+            seed=SEED,
+            transformer=dict(
+                checkpoint_path=str(_TEACHER_CKPT),
+                state_dict_transform=_network_key_filter_transform(),
+                compile_network=False,
+                use_cuda_graph=False,
+            ),
+        ),
+    )
+    pipe = cfg.setup()
+    assert isinstance(pipe, OmnidreamsPipeline)
+    return pipe.to("cuda")
+
+
+def _hdmap_path(uuid: str) -> Path:
+    root = (
+        Path.home()
+        / ".cache/huggingface/hub/datasets--nvidia--omni-dreams-samples/snapshots"
+    )
+    hits = sorted(root.glob(f"*/data/single_view/{uuid}/*_hdmap.mp4"))
+    assert hits, f"sample {uuid} hdmap not in the local HF cache"
+    return hits[0]
+
+
+def _passed_pairs() -> list[tuple[str, str, Path]]:
+    """Filter-passed ``(uuid, slug, path)`` styled pairs in ``STYLES``."""
+    report = json.loads((STYLE_DIR / "filter_report.json").read_text())
+    passed = {e["output"] for e in report if e.get("passed")}
+    out: list[tuple[str, str, Path]] = []
+    for path in sorted(STYLE_DIR.glob("*__*.mp4")):
+        uuid, slug = path.stem.split("__", 1)
+        if path.name not in passed or slug not in STYLES:
+            continue
+        if ONLY and path.stem not in ONLY:
+            continue
+        out.append((uuid, slug, path))
+    return out
+
+
+@torch.no_grad()
+def _encode_window(pipe: OmnidreamsPipeline, video: Tensor) -> Tensor:
+    """Encode ``[189, 3, H, W]`` pixels into the patchified chunk-0 latent."""
+    transformer = pipe.diffusion_model.transformer
+    assert pipe.encoder is not None
+    enc_cache = pipe.encoder.initialize_autoregressive_cache()
+    pixels = video[None, None].to(pipe.device)
+    z = pipe.encoder(input=pixels, autoregressive_index=0, cache=enc_cache)
+    return transformer.patchify_and_maybe_split_cp(z)
+
+
+@torch.no_grad()
+def _sdedit_window(
+    pipe: OmnidreamsPipeline,
+    cache,
+    *,
+    styled_latent: Tensor,
+    hdmap_latent: Tensor,
+    sigma_target: float,
+    rng: torch.Generator,
+) -> tuple[Tensor, dict]:
+    """Renoise-refine one chunk48 window from ``sigma_target``.
+
+    Partial restart of the scheduler's own order-2 UniPC loop: enter the
+    schedule at the step whose sigma is nearest ``sigma_target``, renoise
+    the styled latent there, and run the remaining steps with the baked
+    per-step coefficients (order-1 warmup at the entry step — no
+    predictor/corrector history exists yet — matching ``sample()``'s own
+    step-0 behavior). Returns the decoded ``[T, 3, H, W]`` video on CPU.
+    """
+    transformer = pipe.diffusion_model.transformer
+    sched = pipe.diffusion_model.scheduler
+    tc = cache.transformer_cache
+    assert pipe.decoder is not None and cache.decoder_cache is not None
+    tc.start(0)
+
+    n_steps = int(sched.timesteps.shape[0])
+    i0 = int(torch.argmin((sched.sigmas - sigma_target).abs()).item())
+    sigma0 = float(sched.sigmas[i0].item())
+    dtype = styled_latent.dtype
+
+    noise = torch.randn(
+        styled_latent.shape, device=styled_latent.device, dtype=dtype, generator=rng
+    )
+    sample = ((1.0 - sigma0) * styled_latent.float() + sigma0 * noise.float()).to(dtype)
+
+    m_prev: Tensor | None = None
+    m_prev_prev: Tensor | None = None
+    last_sample: Tensor | None = None
+    t_start = time.perf_counter()
+    for i in range(i0, n_steps):
+        timestep = sched.timesteps[i].to(dtype=dtype)
+        flow = transformer.predict_flow(
+            noisy_latent=sample, timestep=timestep, cache=tc, input=hdmap_latent
+        )
+        m_curr = sample.to(torch.float32) - sched.sigmas[i] * flow.to(torch.float32)
+        if i > i0:
+            assert last_sample is not None and m_prev is not None
+            m_pp = m_prev_prev if m_prev_prev is not None else m_prev
+            sample = (
+                sched.a_corr[i] * last_sample.to(torch.float32)
+                + sched.b_corr_m0[i] * m_prev
+                + sched.b_corr_dprev[i] * (m_pp - m_prev)
+                + sched.b_corr_dt[i] * (m_curr - m_prev)
+            ).to(dtype)
+        last_sample = sample
+        m_p = m_prev if m_prev is not None else m_curr
+        sample = (
+            sched.a_pred[i] * sample.to(torch.float32)
+            + sched.b_pred_m0[i] * m_curr
+            + sched.b_pred_dprev[i] * (m_p - m_curr)
+        ).to(dtype)
+        m_prev_prev, m_prev = m_prev, m_curr
+
+    clean = transformer.postprocess_clean_latent(
+        clean_latent=sample, cache=tc, input=hdmap_latent
+    )
+    decoded = pipe.decoder(
+        input=transformer.unpatchify_and_maybe_gather_cp(clean),
+        autoregressive_index=0,
+        cache=cache.decoder_cache,
+    )
+    info = {
+        "sigma_target": sigma_target,
+        "sigma_snapped": sigma0,
+        "start_step": i0,
+        "steps_run": n_steps - i0,
+        "seconds": round(time.perf_counter() - t_start, 1),
+    }
+    return decoded[0, 0].float().cpu(), info
+
+
+def _to_uint8(frame: Tensor) -> np.ndarray:
+    return (
+        ((frame.permute(1, 2, 0).numpy() + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
+    )
+
+
+def _write_contact_sheet(path: Path, rows: dict[str, Tensor]) -> None:
+    """Grid PNG: rows = named videos, columns = ``CONTACT_FRAMES``."""
+    from PIL import Image, ImageDraw
+
+    stacked = []
+    for name, video in rows.items():
+        row = np.concatenate(
+            [_to_uint8(video[min(f, video.shape[0] - 1)]) for f in CONTACT_FRAMES],
+            axis=1,
+        )
+        img = Image.fromarray(row)
+        ImageDraw.Draw(img).text((8, 8), name, fill=(255, 255, 0))
+        stacked.append(np.asarray(img))
+    Image.fromarray(np.concatenate(stacked, axis=0)).save(path)
+
+
+@torch.no_grad()
+def main() -> None:
+    pairs = _passed_pairs()
+    if LIMIT:
+        pairs = pairs[:LIMIT]
+    assert pairs, f"no filter-passed pairs for styles {sorted(STYLES)} in {STYLE_DIR}"
+    print(
+        f"{len(pairs)} pairs x sigmas {SIGMAS} -> {OUT_DIR}",
+        flush=True,
+    )
+    _wait_for_vram(MAX_USED_GB)
+
+    pipe = _build_pipeline()
+    dtype = torch.bfloat16
+    device = pipe.device
+
+    # Phase A: one-shot embeddings for every pair (style prompt + negative +
+    # styled first frame), then drop the ~14 GB text encoder.
+    styled_videos: dict[str, Tensor] = {}
+    emb: dict[str, dict] = {}
+    emb_by_prompt: dict[str, Tensor] = {}
+    for uuid, slug, path in pairs:
+        video = _load_video(
+            path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device=torch.device("cpu"),
+            dtype=dtype,
+        )
+        assert video.shape[0] >= TEACHER_FRAMES, (
+            f"{path.name}: {video.shape[0]} frames < {TEACHER_FRAMES}"
+        )
+        styled_videos[path.stem] = video[:TEACHER_FRAMES].contiguous()
+        first = styled_videos[path.stem][:1][None, None].to(device)
+        e = pipe.precompute_embeddings(text=[[STYLE_PROMPTS[slug]]], image=first)
+        # Text/negative embeddings only depend on the slug; share them.
+        if STYLE_PROMPTS[slug] not in emb_by_prompt:
+            emb_by_prompt[STYLE_PROMPTS[slug]] = e["text_embeddings"]
+        emb[path.stem] = {
+            "text": emb_by_prompt[STYLE_PROMPTS[slug]],
+            "negative": e["negative_text_embeddings"],
+            "image": e["image_embeddings"],
+        }
+        print(f"embedded {path.stem}", flush=True)
+    pipe.release_oneshot_encoders()
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    report: dict[str, dict] = {}
+    report_path = OUT_DIR / "regen_report.json"
+    for uuid, slug, path in pairs:
+        stem = path.stem
+        styled = styled_videos[stem]
+        outputs = {
+            (f"{stem}__s{sigma:g}" if len(SIGMAS) > 1 else stem): sigma
+            for sigma in SIGMAS
+        }
+        if all((OUT_DIR / f"{name}.mp4").exists() for name in outputs):
+            print(f"{stem}: exists, skipped", flush=True)
+            continue
+
+        styled_latent = _encode_window(pipe, styled)
+        hdmap = _load_video(
+            _hdmap_path(uuid),
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device=torch.device("cpu"),
+            dtype=dtype,
+        )
+        assert hdmap.shape[0] >= TEACHER_FRAMES, f"{uuid}: hdmap too short"
+        hdmap_latent = _encode_window(pipe, hdmap[:TEACHER_FRAMES].contiguous())
+
+        sheet_rows: dict[str, Tensor] = {"styled (JoyAI)": styled.float()}
+        for name, sigma in outputs.items():
+            rng = torch.Generator(device=device).manual_seed(
+                SEED + zlib.crc32(name.encode()) % 100_000
+            )
+            cache = pipe.initialize_cache_from_embeddings(
+                text_embeddings=emb[stem]["text"],
+                image_embeddings=emb[stem]["image"],
+                negative_text_embeddings=emb[stem]["negative"],
+            )
+            refined, info = _sdedit_window(
+                pipe,
+                cache,
+                styled_latent=styled_latent,
+                hdmap_latent=hdmap_latent,
+                sigma_target=sigma,
+                rng=rng,
+            )
+            del cache
+            torch.cuda.empty_cache()
+            write_video_tensor(refined, OUT_DIR / f"{name}.mp4", fps=30, layout="tchw")
+            sheet_rows[f"teacher s={info['sigma_snapped']:.3f}"] = refined
+            report[name] = info
+            print(f"{name}: {info}", flush=True)
+
+        _write_contact_sheet(OUT_DIR / f"{stem}_contact.png", sheet_rows)
+        report_path.write_text(json.dumps(report, indent=2))
+        del styled_latent, hdmap_latent
+        torch.cuda.empty_cache()
+
+    print(f"TEACHER-STYLED-TARGETS-DONE -> {OUT_DIR}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/edit_sft/precompute_style.py b/integrations/omnidreams/edit_sft/precompute_style.py
index e4df5d539..cbd37e84c 100644
--- a/integrations/omnidreams/edit_sft/precompute_style.py
+++ b/integrations/omnidreams/edit_sft/precompute_style.py
@@ -16,8 +16,8 @@
 """Precompute the style-SFT embeddings and latents (one shot).
 
 Three phases, so ``train_style_sft.py`` never loads the ~14 GB text
-encoder or touches ffmpeg (the precompute-embeddings pattern from the
-guidance-distillation trainer, extended with video encoding):
+encoder or touches ffmpeg (the ``guidance_distill/precompute_embeddings``
+pattern, extended with video encoding):
 
 A. Text + first-frame embeddings — every style-slug prompt
    (:data:`~style_prompts.STYLE_PROMPTS`) and every source clip's own
@@ -46,8 +46,8 @@
 
 Latent files are skipped when present, so the script is resumable — the
 frame-count assert catches the silent-empty ffmpeg reads observed on this
-box (they hit once the process has grown to rollout size; the one-shot
-encoders are released before the first decode to shrink the process).
+box (``build_pairs.py`` note; the one-shot encoders are released before
+the first decode to shrink the process).
 
 Run from the flashdreams repo root (~20 GB VRAM during phase A)::
 
@@ -65,48 +65,31 @@
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "drift_correction"))
 
 import torch
 from _host import build_pipeline
+from build_pairs import _sample_files
 from omnidreams.pipeline import OmnidreamsPipeline
 from omnidreams.runner import (
     DEFAULT_VIDEO_HEIGHT,
     DEFAULT_VIDEO_WIDTH,
-    _ensure_hf_single_view_example_data_synced,
+    _load_first_frame,
+    _load_video,
 )
 from style_prompts import STYLE_PROMPTS, clip_key
 from torch import Tensor
 
-from flashdreams.infra.runner_io import load_first_frame_tensor, load_video_tensor
-
 ## Configuration
 
 BASE = Path("integrations/omnidreams/edit_sft")
 OUT_DIR = Path(os.environ.get("OUT_DIR", str(BASE / "outputs")))
 SRC_DIR = OUT_DIR / "sources"
-STYLE_DIR = OUT_DIR / "style_pairs"
-LAT_DIR = OUT_DIR / "latents"
-
-
-def _sample_files(uuid: str) -> tuple[tuple[Path, ...], tuple[Path, ...]]:
-    """Cache-first ``((hdmap,), (first_frame,))`` resolution for a sample.
-
-    Vendored from the Clean Forcing pair builder (PR #398): the runner's
-    ``_ensure_hf_single_view_example_data_synced`` lists the HF repo per
-    clip even when every file is already local, and the shared IP rate
-    limit has killed pipeline stages before. 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)
+STYLE_DIR = Path(os.environ.get("STYLE_DIR", str(OUT_DIR / "style_pairs")))
+"""Styled-target videos; point at ``style_pairs_teacher`` for the
+teacher-regenerated corpus (``gen_teacher_styled_targets.py``)."""
+LAT_DIR = Path(os.environ.get("LAT_DIR", str(OUT_DIR / "latents")))
+"""Latent output dir; use a separate dir per styled corpus (don't clobber)."""
 
 
 def _styled_videos(known_uuids: set[str]) -> list[tuple[str, str, Path]]:
@@ -167,73 +150,106 @@ def main() -> None:
     styled = _styled_videos(set(uuids))
     assert styled, f"no recognized styled pairs under {STYLE_DIR}"
 
-    # First frames are PNGs (cv2, no ffmpeg fork) — safe to load lazily,
-    # but front-load them anyway so phase A never touches the filesystem
-    # mid-encoding.
-    firsts: dict[str, Tensor] = {}
-    for uuid in uuids:
-        _, (frame_path,) = _sample_files(uuid)
-        firsts[uuid] = load_first_frame_tensor(
-            frame_path,
-            pixel_height=DEFAULT_VIDEO_HEIGHT,
-            pixel_width=DEFAULT_VIDEO_WIDTH,
-            device=torch.device("cpu"),
-            dtype=dtype,
-        )[None, :, None]  # [1, V=1, 1, C, H, W]
-
-    pipe = build_pipeline(with_oneshot_encoders=True)
-    device = pipe.device
-    assert pipe.text_encoder is not None  # with_oneshot_encoders=True
-
-    # Phase A: style + clip prompt embeddings and first-frame latents.
-    emb_by_text: dict[str, Tensor] = {}
-    prompt_embeddings: dict[str, Tensor] = {}
-    for slug in sorted({slug for _, slug, _ in styled}):
-        text = STYLE_PROMPTS[slug]
-        if text not in emb_by_text:  # _v2 slugs alias their base prompt
-            emb = torch.stack([pipe.text_encoder([text])], dim=0)  # [1, 1, L, D]
-            emb_by_text[text] = emb.to("cpu", dtype)
-        prompt_embeddings[slug] = emb_by_text[text]
-        print(
-            f"encoded style prompt {slug}: {tuple(prompt_embeddings[slug].shape)}",
-            flush=True,
+    # Phase A is skippable when the embedding files already cover every
+    # needed key (reruns against a new STYLE_DIR/LAT_DIR corpus): the
+    # ~14 GB text encoder never loads.
+    emb_path = OUT_DIR / "style_embeddings.pt"
+    assets_path = OUT_DIR / "style_clip_assets.pt"
+    skip_phase_a = False
+    if emb_path.exists() and assets_path.exists():
+        prev_emb = torch.load(emb_path, map_location="cpu", weights_only=False)
+        prev_assets = torch.load(assets_path, map_location="cpu", weights_only=False)
+        skip_phase_a = all(slug in prev_emb for _, slug, _ in styled) and all(
+            clip_key(u) in prev_emb and u in prev_assets["image_embeddings"]
+            for u in uuids
         )
-    image_embeddings: dict[str, Tensor] = {}
-    for uuid in uuids:
-        emb = pipe.precompute_embeddings(
-            text=[[prompts[uuid]]], image=firsts[uuid].to(device)
+        del prev_emb, prev_assets
+
+    if skip_phase_a:
+        print("phase A: existing embedding files cover all keys — skipped", flush=True)
+        pipe = build_pipeline(with_oneshot_encoders=False)
+    else:
+        # First frames are PNGs (cv2, no ffmpeg fork) — safe to load lazily,
+        # but front-load them anyway so phase A never touches the filesystem
+        # mid-encoding.
+        firsts: dict[str, Tensor] = {}
+        for uuid in uuids:
+            _, (frame_path,) = _sample_files(uuid)
+            firsts[uuid] = _load_first_frame(
+                frame_path,
+                pixel_height=DEFAULT_VIDEO_HEIGHT,
+                pixel_width=DEFAULT_VIDEO_WIDTH,
+                device="cpu",
+                dtype=dtype,
+            )[None, :, None]  # [1, V=1, 1, C, H, W]
+
+        pipe = build_pipeline(with_oneshot_encoders=True)
+        device = pipe.device
+        assert pipe.text_encoder is not None  # with_oneshot_encoders=True
+
+        # Phase A: style + clip prompt embeddings and first-frame latents.
+        emb_by_text: dict[str, Tensor] = {}
+        prompt_embeddings: dict[str, Tensor] = {}
+        for slug in sorted({slug for _, slug, _ in styled}):
+            text = STYLE_PROMPTS[slug]
+            if text not in emb_by_text:  # _v2 slugs alias their base prompt
+                emb = torch.stack([pipe.text_encoder([text])], dim=0)  # [1, 1, L, D]
+                emb_by_text[text] = emb.to("cpu", dtype)
+            prompt_embeddings[slug] = emb_by_text[text]
+            print(
+                f"encoded style prompt {slug}: {tuple(prompt_embeddings[slug].shape)}",
+                flush=True,
+            )
+        image_embeddings: dict[str, Tensor] = {}
+        for uuid in uuids:
+            emb = pipe.precompute_embeddings(
+                text=[[prompts[uuid]]], image=firsts[uuid].to(device)
+            )
+            text_emb = emb["text_embeddings"]
+            image_emb = emb["image_embeddings"]
+            assert text_emb is not None and image_emb is not None
+            prompt_embeddings[clip_key(uuid)] = text_emb.to("cpu", dtype)
+            image_embeddings[uuid] = image_emb.to("cpu", dtype)
+            print(f"encoded clip {uuid}", flush=True)
+
+        OUT_DIR.mkdir(parents=True, exist_ok=True)
+        torch.save(prompt_embeddings, emb_path)
+        torch.save(
+            {"uuids": uuids, "prompts": prompts, "image_embeddings": image_embeddings},
+            assets_path,
         )
-        text_emb = emb["text_embeddings"]
-        image_emb = emb["image_embeddings"]
-        assert text_emb is not None and image_emb is not None
-        prompt_embeddings[clip_key(uuid)] = text_emb.to("cpu", dtype)
-        image_embeddings[uuid] = image_emb.to("cpu", dtype)
-        print(f"encoded clip {uuid}", flush=True)
-
-    OUT_DIR.mkdir(parents=True, exist_ok=True)
-    torch.save(prompt_embeddings, OUT_DIR / "style_embeddings.pt")
-    torch.save(
-        {"uuids": uuids, "prompts": prompts, "image_embeddings": image_embeddings},
-        OUT_DIR / "style_clip_assets.pt",
-    )
-    # Shrink the process before the ffmpeg decodes below (fork hazard).
-    pipe.release_oneshot_encoders()
+        # Shrink the process before the ffmpeg decodes below (fork hazard).
+        pipe.release_oneshot_encoders()
 
-    def encode_to(path: Path, video_path: Path, key: str) -> None:
-        video = load_video_tensor(
+    def encode_to(
+        path: Path, video_path: Path, key: str, allow_short: bool = False
+    ) -> None:
+        video = _load_video(
             video_path,
             pixel_height=DEFAULT_VIDEO_HEIGHT,
             pixel_width=DEFAULT_VIDEO_WIDTH,
-            device=torch.device("cpu"),
+            device="cpu",
             dtype=dtype,
         )
-        assert video.shape[0] >= total_frames, (
-            f"{video_path.name}: {video.shape[0]} frames < {total_frames} "
-            "(short clip, or a silent-empty ffmpeg read — re-run to resume)."
-        )
-        chunks = encode_video_chunks(pipe, video[:total_frames], n_chunks)
-        torch.save({key: chunks, "n_chunks": n_chunks}, path)
-        print(f"encoded {path.name}: {n_chunks} chunks", flush=True)
+        if allow_short and video.shape[0] < total_frames:
+            # Teacher-regenerated targets cover one bidirectional chunk48
+            # window (189 frames = student chunks 0..23); encode the full
+            # chunks available. Still assert non-empty (silent ffmpeg reads).
+            avail = (int(video.shape[0]) - 5) // 8 + 1
+            assert avail >= 1, (
+                f"{video_path.name}: {video.shape[0]} frames < one chunk "
+                "(silent-empty ffmpeg read? re-run to resume)."
+            )
+        else:
+            assert video.shape[0] >= total_frames, (
+                f"{video_path.name}: {video.shape[0]} frames < {total_frames} "
+                "(short clip, or a silent-empty ffmpeg read — re-run to resume)."
+            )
+            avail = n_chunks
+        frames_used = 5 + (avail - 1) * 8
+        chunks = encode_video_chunks(pipe, video[:frames_used], avail)
+        torch.save({key: chunks, "n_chunks": avail}, path)
+        print(f"encoded {path.name}: {avail} chunks", flush=True)
 
     LAT_DIR.mkdir(parents=True, exist_ok=True)
     # Phase B: per-clip HDMap latents (conditioning for replay + training).
@@ -251,11 +267,12 @@ def encode_to(path: Path, video_path: Path, key: str) -> None:
         if out.exists():
             print(f"skip {out.name} (exists)", flush=True)
             continue
-        encode_to(out, path, "latents")
+        encode_to(out, path, "latents", allow_short=True)
 
+    phase_a = "reused" if skip_phase_a else "encoded"
     print(
-        f"PRECOMPUTE-STYLE-DONE | {len(prompt_embeddings)} prompt embeddings, "
-        f"{len(uuids)} clips, {len(styled)} styled pairs -> {OUT_DIR}/",
+        f"PRECOMPUTE-STYLE-DONE | embeddings {phase_a}, {len(uuids)} clips, "
+        f"{len(styled)} styled pairs -> {LAT_DIR}/",
         flush=True,
     )
 
diff --git a/integrations/omnidreams/edit_sft/train_style_sft.py b/integrations/omnidreams/edit_sft/train_style_sft.py
index 397a48cca..dafa71f23 100644
--- a/integrations/omnidreams/edit_sft/train_style_sft.py
+++ b/integrations/omnidreams/edit_sft/train_style_sft.py
@@ -23,28 +23,42 @@
 
 - **Context build**: chunks before the sampled swap chunk ``k`` are the
   SOURCE rollout's own latents (``generate_sources.py``), replayed through
-  ``finalize_kv_cache`` at the host's ``context_noise`` (the vendored
-  ``_host.replay_history`` machinery) under base weights —
+  ``finalize_kv_cache`` at the host's ``context_noise`` (the
+  ``drift_correction`` ``replay_history`` machinery) under base weights —
   at deploy the LoRA window opens only at the swap. RoPE is absolute and
   the KV window holds ``REPLAY_CHUNKS`` chunks, so replaying just the
   window prefix reproduces a full rollout's visible state exactly.
 - **Edit supervision**: at ``k`` the text KV is rebuilt to the style
-  prompt (plain ``swap_text_kv`` rebuild — exactly the deploy swap of the
-  live-edit hook, PR #431); chunks ``[k, k + SPAN)`` are then supervised
-  toward the STYLED latents (the JoyAI output encoded by
-  ``precompute_style.py``). States
+  prompt (plain ``replace_text_from_embeddings``, exactly the deploy
+  swap); chunks ``[k, k + SPAN)`` are then supervised toward the STYLED
+  latents (the JoyAI output encoded by ``precompute_style.py``). States
   stay on-policy within each chunk (the student's own detached flow
   advances the two-step trajectory, mirroring ``scheduler.sample``); the
   v-target at every state is ``(z_t - x0_styled) / sigma``. A 0.5-weighted
   context term (t=128) supervises the same forward whose K/V the commit
   writes, and the STYLED latent is committed into the rolling KV — teacher
   forcing proceeds on edited history, merged-weight commits matching the
-  deploy window semantics (the live-edit deploy hook, PR #431).
+  deploy window semantics (``_edit_lora``).
 - **No-op regularization**: ``PRE_CHUNKS`` supervised chunks right before
   ``k`` target the SOURCE latents under the clip's own prompt (keeps
   non-edit behavior identical), and with ``NOOP_PROB`` the whole window is
   a no-op (swap to the clip's own prompt, source targets) so the style
   must come from the text KV, not the weights.
+- **v4 self-consistency maintenance** (:data:`V4_MAINT`, default on): the
+  v3 maintenance episodes (``MAINT_PROB``: styled history at scale 1,
+  styled targets) taught a SOFT steady state — their JoyAI targets are RGB
+  re-encoded through the VAE, so the fixed point the LoRA converges to is
+  blurred and the drift corrector can only pull back to that soft point
+  (``issues_and_fixes.md`` Issue 1). v4 replaces the JoyAI latents in
+  maintenance episodes ONLY with the model's OWN early-window styled
+  rollouts — the ``gen_style_drift_pairs.py`` clean reference branches:
+  history and targets both come from a branch segment held at most
+  :data:`V4_EARLY` chunks past its swap (the pre-drift styled manifold),
+  so the steady state supervised is the model's own sharp styled
+  continuation. Swap-transition and no-op episodes keep their JoyAI /
+  source targets unchanged — they still define the style identity at the
+  transition. Pairs without a usable branch segment fall back to the v3
+  maintenance path (counted, reported in the DONE line).
 
 Data is gated by the style-mode VLM filter
 (``style_pairs/filter_report.json``) when present: ``passed`` pairs train
@@ -53,11 +67,10 @@
 drift sets in on heavy styles). Without a report every pair on disk
 trains, with a warning.
 
-Host mechanics are the guidance-distillation trainer's (vendored into
-this directory): eager pipeline, LoRA r64 on the 8 attention projections
-(the checkpoint loads through the live-edit deploy hook's
-``TextEditLoRA``, PR #431, unchanged), functional self-attention on grad
-forwards, per-block checkpointing, and every term's backward
+Host mechanics are the ``guidance_distill/train_guidance.py`` ones: eager
+pipeline, LoRA r64 on the 8 attention projections (checkpoint loads
+through ``omnidreams/_edit_lora.py`` unchanged), functional self-attention
+on grad forwards, per-block checkpointing, and every term's backward
 immediately after its forward inside ``functional_attention()`` (the
 recompute must retake the non-writing path, and later stock forwards
 in-place mutate buffers the tape saved). No teacher forward exists here —
@@ -81,16 +94,17 @@
 import os
 import sys
 from pathlib import Path
-from typing import cast
 
 # Must land before the first CUDA allocation (co-tenant VRAM share).
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "drift_correction"))
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "guidance_distill"))
 
 import numpy as np
 import torch
-from _host import CONTEXT_NOISE_SEED, build_pipeline, swap_text_kv
+from _host import CONTEXT_NOISE_SEED, build_pipeline
 from _lora import (
     apply_lora,
     lora_parameters,
@@ -101,16 +115,22 @@
 from _train_attn import functional_attention, patch_functional_attention
 from style_prompts import DEFAULT_SKIP_STYLES, STYLE_PROMPTS, clip_key
 from torch import Tensor
-
-from flashdreams.infra.diffusion.scheduler.fm import FlowMatchScheduler
+from train_guidance import LORA_TARGETS, checkpoint_blocks
 
 ## Training configuration
 
 BASE = Path("integrations/omnidreams/edit_sft")
 OUT_DIR = BASE / "outputs"
 SRC_DIR = OUT_DIR / "sources"
-STYLE_DIR = OUT_DIR / "style_pairs"
-LAT_DIR = OUT_DIR / "latents"
+STYLE_DIR = Path(os.environ.get("STYLE_DIR", str(OUT_DIR / "style_pairs")))
+"""Styled-target corpus; ``style_pairs_teacher`` = the teacher-regenerated
+targets (``gen_teacher_styled_targets.py``, v5 recipe)."""
+LAT_DIR = Path(os.environ.get("LAT_DIR", str(OUT_DIR / "latents")))
+"""Latents matching STYLE_DIR (``precompute_style.py`` with the same env)."""
+CKPT_PREFIX = os.environ.get("CKPT_PREFIX", "lora_style")
+"""Checkpoint filename prefix (``<prefix>_step<N>.pt``)."""
+PAIRS_DIR = Path(os.environ.get("PAIRS_DIR", str(OUT_DIR / "style_drift_pairs")))
+"""Clean-branch corpus (``gen_style_drift_pairs.py``) for v4 maintenance."""
 
 STEPS = int(os.environ.get("STEPS", "1000"))
 LR = float(os.environ.get("LR", "2e-4"))
@@ -133,6 +153,20 @@
 """Probability of a no-op window (swap to the clip's own prompt, source
 targets): anchors the restyle to the text KV instead of the weights."""
 
+V4_MAINT = os.environ.get("V4_MAINT", "1") == "1"
+"""Maintenance episodes replay/target the model's OWN early-window styled
+branches (self-consistency) instead of the JoyAI re-encodes; 0 = exact v3
+maintenance behavior. ``MAINT_PROB`` semantics are unchanged."""
+
+V4_EARLY = int(os.environ.get("V4_EARLY", "7"))
+"""Max style-hold depth (chunks since swap, 1-based) a v4 episode may
+touch: history ``[k - REPLAY_CHUNKS, k)`` and targets ``[k, k + SPAN)``
+must fit inside branch depths ``1..V4_EARLY``. 7 is the tightest feasible
+with ``REPLAY_CHUNKS=3`` / ``SPAN=4`` (one k per branch, ``k = s + 3``);
+Issue 1 puts "excellent" at +1..+4 and structure-holds at +7..+12, so 7
+keeps targets on the sharp side of the drift. Raising it trades target
+sharpness for more (s, k) cells."""
+
 PRE_CHUNKS = int(os.environ.get("PRE_CHUNKS", "1"))
 """Supervised chunks immediately before ``k`` (source targets, original
 prompt) — no-op regularization at the swap boundary."""
@@ -155,8 +189,8 @@
 REPLAY_CHUNKS = int(os.environ.get("REPLAY_CHUNKS", "3"))
 """Source chunks replayed before the first supervised chunk. 3 chunks =
 the full ``window_size_t=6`` / ``len_t=2`` KV window, so the visible state
-equals a full-history replay (the drift corrector's mid-stream-start
-machinery, PR #398); raise only if the host's window grows."""
+equals a full-history replay (``train_v2.py`` mid-stream-start machinery);
+raise only if the host's window grows."""
 
 CTX_WEIGHT = 0.5
 """Weight of the finalize/context-forward (t=128) term."""
@@ -168,49 +202,6 @@
 )
 """Slugs excluded from training regardless of the filter verdict."""
 
-LORA_TARGETS = (
-    "self_attn.q_proj",
-    "self_attn.k_proj",
-    "self_attn.v_proj",
-    "self_attn.output_proj",
-    "cross_attn.q_proj",
-    "cross_attn.k_proj",
-    "cross_attn.v_proj",
-    "cross_attn.output_proj",
-)
-"""Drift-corrector recipe + cross-attn (the edit signal enters there).
-Vendored from the guidance-distillation trainer (deferred from PR #431) so
-this pipeline is self-contained; consolidate when that trainer lands."""
-
-
-def checkpoint_blocks(network) -> None:
-    """Route every DiT block ``forward`` through gradient checkpointing.
-
-    Vendored from the guidance-distillation trainer (deferred from
-    PR #431); consolidate when that trainer lands. Per-instance overrides
-    (not wrapper modules) so the network loop's ``isinstance(block, Block)``
-    assertion keeps passing. Requires the functional-attention toggle on
-    grad passes: recomputation must be side-effect free and must retake the
-    same code path, so every backward runs inside
-    ``functional_attention()``. No-op under ``no_grad`` passes (rollout,
-    finalize).
-
-    Args:
-        network: The unwrapped ``CosmosDiTNetwork``.
-    """
-    from torch.utils.checkpoint import checkpoint
-
-    def wrap(fn):
-        def ckpt_fn(*args, _inner=fn, **kwargs):
-            if not torch.is_grad_enabled():
-                return _inner(*args, **kwargs)
-            return checkpoint(_inner, *args, use_reentrant=False, **kwargs)
-
-        return ckpt_fn
-
-    for block in network.blocks:
-        block.forward = wrap(block.forward)
-
 
 def load_pairs(train_uuids: set[str], n_chunks: int) -> list[tuple[str, str, int]]:
     """Collect trainable ``(uuid, slug, max_chunk)`` styled pairs.
@@ -265,6 +256,49 @@ def load_pairs(train_uuids: set[str], n_chunks: int) -> list[tuple[str, str, int
     return pairs
 
 
+def load_v4_branches(pair_keys: set[tuple[str, str]]) -> dict[tuple[str, str], dict]:
+    """Load the clean-branch segments usable for v4 maintenance episodes.
+
+    A ``gen_style_drift_pairs.py`` branch file stores per swap offset ``s``
+    the styled rollout's chunks ``s..n_chunks-1``. A maintenance swap chunk
+    ``k`` is usable when the whole episode — history ``[k - REPLAY_CHUNKS,
+    k)`` plus targets ``[k, k + SPAN)`` — lies at hold depth <=
+    :data:`V4_EARLY` inside the branch::
+
+        s + REPLAY_CHUNKS <= k <= min(s + V4_EARLY, n_chunks) - SPAN
+
+    Returns ``{(uuid, slug): {"cells": [(s, k), ...], "branches":
+    {s: latents for chunks s..}}}`` for the trainable ``pair_keys``,
+    branch latents trimmed to the deepest chunk any usable ``k`` touches.
+    Files or offsets too short for any ``k`` are skipped (reported).
+    """
+    corpus: dict[tuple[str, str], dict] = {}
+    for path in sorted(PAIRS_DIR.glob("*__*.pt")):
+        uuid, slug = path.stem.split("__", 1)
+        if (uuid, slug) not in pair_keys:
+            continue
+        d = torch.load(path, map_location="cpu", weights_only=False)
+        n = int(d["n_chunks"])
+        cells: list[tuple[int, int]] = []
+        branches: dict[int, list[Tensor]] = {}
+        for s, lats in sorted((int(s), lat) for s, lat in d["branches"].items()):
+            assert len(lats) == n - s, f"{path.name} s={s}: {len(lats)} chunks"
+            ks = list(range(s + REPLAY_CHUNKS, min(s + V4_EARLY, n) - SPAN + 1))
+            if not ks:
+                continue
+            cells += [(s, k) for k in ks]
+            branches[s] = lats[: max(ks) + SPAN - s]
+        if not cells:
+            print(f"v4: skip {path.name} — no usable k at V4_EARLY={V4_EARLY}")
+            continue
+        corpus[(uuid, slug)] = {"cells": cells, "branches": branches}
+        print(
+            f"v4: {uuid[:8]}__{slug} usable k {[k for _, k in cells]}",
+            flush=True,
+        )
+    return corpus
+
+
 def main() -> None:
     """Run the teacher-forced style-SFT loop."""
     rng = np.random.default_rng(SEED)
@@ -278,6 +312,16 @@ def main() -> None:
         f"no trainable styled pairs under {STYLE_DIR} (run precompute_style.py "
         "first; check filter_report.json / SKIP_STYLES)."
     )
+    v4_corpus = load_v4_branches({(u, s) for u, s, _ in pairs}) if V4_MAINT else {}
+    v4_fallbacks = [0]  # maintenance draws without a branch -> v3 path
+    if V4_MAINT:
+        n_cells = sum(len(e["cells"]) for e in v4_corpus.values())
+        print(
+            f"v4 maintenance ON: {n_cells} self-consistency cells over "
+            f"{len(v4_corpus)}/{len(pairs)} pairs (V4_EARLY={V4_EARLY}); "
+            "pairs without branches fall back to v3 maintenance.",
+            flush=True,
+        )
 
     prompt_emb = torch.load(
         OUT_DIR / "style_embeddings.pt", map_location="cpu", weights_only=False
@@ -309,10 +353,17 @@ def main() -> None:
             map_location="cpu",
             weights_only=False,
         )["latents"]
-        assert len(tgt) == n_chunks, (
-            f"pair {uuid}__{slug}: {len(tgt)} styled chunks != {n_chunks}"
+        # Teacher-regenerated targets cover one chunk48 window (24 student
+        # chunks) — shorter than the clip; clamp the pair's usable range.
+        assert SWAP_MIN + SPAN <= len(tgt) <= n_chunks, (
+            f"pair {uuid}__{slug}: {len(tgt)} styled chunks "
+            f"(need {SWAP_MIN + SPAN}..{n_chunks})"
         )
         tgt_lat[(uuid, slug)] = tgt
+    pairs = [
+        (uuid, slug, min(max_chunk, len(tgt_lat[(uuid, slug)])))
+        for uuid, slug, max_chunk in pairs
+    ]
 
     pipe = build_pipeline(with_oneshot_encoders=False)
     assert pipe.V_group is None, "single-GPU trainer; run without CP"
@@ -320,8 +371,7 @@ def main() -> None:
     dtype = torch.bfloat16
     dm = pipe.diffusion_model
     transformer = dm.transformer
-    # The base ``Scheduler`` annotation hides the FM schedule buffers.
-    scheduler = cast(FlowMatchScheduler, dm.scheduler)
+    scheduler = dm.scheduler
     timesteps = scheduler.denoising_step_list  # [1000, 450] on the chunk2 host
     sigmas = scheduler.denoising_sigmas
     n_steps = int(timesteps.shape[0])
@@ -453,8 +503,18 @@ def train_step() -> tuple[dict[str, float], str]:
         # through the KV commits and the world blurs out (observed at
         # +7-10 chunks on the step-1600 run).
         maint = (not no_op) and bool(rng.random() < MAINT_PROB)
-        k_hi = min(SWAP_MAX, max_chunk - SPAN)
-        k = int(rng.integers(SWAP_MIN, k_hi + 1))
+        # v4: self-consistency maintenance — replay/target a clean
+        # early-window branch of the model's OWN styled rollout instead of
+        # the (VAE-soft) JoyAI latents. Only maintenance changes; swap and
+        # no-op episodes keep JoyAI/source targets.
+        v4 = v4_corpus.get((uuid, slug)) if maint else None
+        if maint and V4_MAINT and v4 is None:
+            v4_fallbacks[0] += 1
+        if v4 is not None:
+            s, k = v4["cells"][int(rng.integers(len(v4["cells"])))]
+        else:
+            k_hi = min(SWAP_MAX, max_chunk - SPAN)
+            k = int(rng.integers(SWAP_MIN, k_hi + 1))
         n_pre = 0 if maint else min(PRE_CHUNKS, k)
         j0 = k - n_pre
 
@@ -470,20 +530,28 @@ def train_step() -> tuple[dict[str, float], str]:
         )
         tc = cache.transformer_cache
         end = k + SPAN
-        src = [x.to(device, dtype) for x in src_lat[uuid][:end]]
-        tgt = [x.to(device, dtype) for x in tgt_lat[(uuid, slug)][:end]]
         hd = [x.to(device, dtype) for x in hd_lat[uuid][:end]]
+        if v4 is not None:
+            # Absolute-indexed view of the branch (chunks < s are never
+            # touched: history starts at k - REPLAY_CHUNKS >= s and
+            # maintenance has no pre-window / no-op chunks).
+            blat = v4["branches"][s]
+            tgt = [None] * s + [x.to(device, dtype) for x in blat[: end - s]]
+            src = tgt
+        else:
+            src = [x.to(device, dtype) for x in src_lat[uuid][:end]]
+            tgt = [x.to(device, dtype) for x in tgt_lat[(uuid, slug)][:end]]
 
         if maint:
             # In-window steady state: text already swapped, history styled,
             # commits at LoRA scale 1 (deploy's use_lora window semantics).
-            swap_text_kv(network, tc, prompt_emb[slug])
+            pipe.replace_text_from_embeddings(cache, prompt_emb[slug])
 
         # Context build: replay the KV-window prefix (SOURCE latents at base
         # weights for swap episodes — deploy runs pre-window chunks on the
         # unmodified network; STYLED latents at scale 1 for maintenance
         # episodes). Per-chunk seeded context noise = the replay_history
-        # contract; the mid-stream start is the drift corrector's machinery.
+        # contract; the mid-stream start is the train_v2 machinery.
         replay = tgt if maint else src
         set_lora_scale(network, 1.0 if maint else 0.0)
         start = max(0, j0 - REPLAY_CHUNKS)
@@ -507,13 +575,11 @@ def train_step() -> tuple[dict[str, float], str]:
         total = 0.0
         for j in range(j0, end):
             if j == k and not maint:
-                # The deploy swap: plain cross-attn text-KV rebuild, the
-                # same rebuild the live-edit deploy hook (PR #431) runs on
-                # a plain prompt swap (no-op windows swap to the clip's own
-                # prompt — same machinery, identical KV — so style must be
-                # read from the text).
+                # The deploy swap: plain cross-attn KV rebuild (no-op
+                # windows swap to the clip's own prompt — same machinery,
+                # identical KV — so style must be read from the text).
                 emb = prompt_emb[clip_key(uuid) if no_op else slug]
-                swap_text_kv(network, tc, emb)
+                pipe.replace_text_from_embeddings(cache, emb)
             x0_tgt = (src[j] if (no_op or (not maint and j < k)) else tgt[j]).float()
             losses = train_chunk(
                 tc, j, x0_tgt, hd[j], weight, maint or j >= k, model_rng
@@ -527,9 +593,16 @@ def train_step() -> tuple[dict[str, float], str]:
 
         del cache
         sums["total"] = total
+        if no_op:
+            tag = "no_op:"
+        elif v4 is not None:
+            tag = f"maintv4[s={s}]:"
+        elif maint:
+            tag = "maintfb:" if V4_MAINT else "maint:"
+        else:
+            tag = ""
         episode = (
-            f"{uuid[:8]} {'no_op:' if no_op else 'maint:' if maint else ''}{slug[:16]} "
-            f"k={k} pre={n_pre} span={SPAN} max={max_chunk}"
+            f"{uuid[:8]} {tag}{slug[:16]} k={k} pre={n_pre} span={SPAN} max={max_chunk}"
         )
         return sums, episode
 
@@ -557,11 +630,12 @@ def train_step() -> tuple[dict[str, float], str]:
                 flush=True,
             )
         if step % SAVE_EVERY == 0 or step == STEPS:
-            path = OUT_DIR / f"lora_style_step{step}.pt"
+            path = OUT_DIR / f"{CKPT_PREFIX}_step{step}.pt"
             save_lora(network, path)
             print(f"saved {path}", flush=True)
 
-    print(f"TRAIN-STYLE-SFT-DONE | final loss ema {ema:.4f}", flush=True)
+    fb = f" | v4 fallbacks {v4_fallbacks[0]}" if V4_MAINT else ""
+    print(f"TRAIN-STYLE-SFT-DONE | final loss ema {ema:.4f}{fb}", flush=True)
 
 
 if __name__ == "__main__":

From 17135956dd0f1518aa1ce0548600c4556de3cadb Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Sat, 22 Aug 2026 08:19:39 +0000
Subject: [PATCH 5/8] Add pre-merged text-edit LoRA deploy hook with
 release_targets seam

TextEditLoRA caches base and base-plus-delta weight sets at load and
toggles edit windows by in-place copy_ into the live projection
weights, so storage addresses survive and captured CUDA graphs stay
valid. release_targets() hands selected projections over to another
writer (the fused drift-corrector dispatch) and returns their fp32
deltas for composition. Ported from omnidreams-live-edit; the
transformer-level edit-window tests stay on that branch since this
base lacks the replace_text_embeddings machinery, so the tests here
cover the merge/toggle/restore and release invariants with an inlined
tiny-transformer helper.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../omnidreams/omnidreams/_edit_lora.py       | 169 ++++++++++++++++++
 .../omnidreams/tests/test_edit_lora.py        | 150 ++++++++++++++++
 2 files changed, 319 insertions(+)
 create mode 100644 integrations/omnidreams/omnidreams/_edit_lora.py
 create mode 100644 integrations/omnidreams/tests/test_edit_lora.py

diff --git a/integrations/omnidreams/omnidreams/_edit_lora.py b/integrations/omnidreams/omnidreams/_edit_lora.py
new file mode 100644
index 000000000..f6193eaff
--- /dev/null
+++ b/integrations/omnidreams/omnidreams/_edit_lora.py
@@ -0,0 +1,169 @@
+# 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.
+
+"""Pre-merged text-edit LoRA deploy hook for mid-stream prompt swaps.
+
+Deploys a ``guidance_distill/train_guidance.py`` checkpoint — a LoRA
+distilled from the two-prompt edit guidance — so a plain prompt swap
+responds at guided strength without the guidance's extra forward per
+denoise step. Both weight sets (base and base-plus-delta) are cached at
+load; toggling an edit window ``copy_``s the right set into the live
+projection weights, so storage addresses survive and captured CUDA graphs
+stay valid (the drift corrector's pointer-rebinding swap is not
+graph-safe). Toggles happen only at edit-window boundaries — a few chunks
+apart — so the copy cost (~1.6 GiB, sub-millisecond) is off the hot path.
+
+Window semantics live in :class:`~omnidreams.transformer.TextEditGuidance`:
+``CosmosTransformer.replace_text_embeddings`` builds a ``use_lora`` window
+when a hook is attached, ``predict_flow`` activates the merged weights for
+the window's chunks (including the KV-commit context forwards — the
+checkpoint was trained to match the guided context forward too), and the
+first forward after the countdown expires restores the base weights.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+_LORA_TARGETS = (
+    "self_attn.q_proj",
+    "self_attn.k_proj",
+    "self_attn.v_proj",
+    "self_attn.output_proj",
+    "cross_attn.q_proj",
+    "cross_attn.k_proj",
+    "cross_attn.v_proj",
+    "cross_attn.output_proj",
+)
+"""Projections the guidance-distillation checkpoints were trained on.
+
+Must match ``guidance_distill/train_guidance.py``'s ``LORA_TARGETS`` (same
+substring rule, same ``named_modules`` walk) so the checkpoint's
+load-order indices line up. ``cross_attn.`` does not match the multi-view
+``cross_view_attn.`` modules.
+"""
+
+
+def _target_linears(network: nn.Module) -> list[nn.Linear]:
+    """Target linears in checkpoint load order (the training-side walk)."""
+    linears: list[nn.Linear] = []
+    for mname, module in network.named_modules():
+        for cname, child in module.named_children():
+            full = f"{mname}.{cname}" if mname else cname
+            if isinstance(child, nn.Linear) and any(t in full for t in _LORA_TARGETS):
+                linears.append(child)
+    return linears
+
+
+class TextEditLoRA:
+    """Two cached weight sets (base / edit) toggled per edit window.
+
+    Args:
+        network: The unwrapped ``CosmosDiTNetwork`` whose projection
+            weights are toggled in place.
+        checkpoint: ``train_guidance.py`` checkpoint (a dict whose
+            ``"lora"`` entry maps load-order indices to A/B tensors;
+            ``A_i`` at ``2i``, ``B_i`` at ``2i + 1``).
+        scale: Gain on the LoRA delta. The checkpoint distills a fixed
+            teacher strength, so ``1.0`` reproduces the evaluated deploy.
+    """
+
+    def __init__(
+        self,
+        network: nn.Module,
+        checkpoint: Path | str,
+        *,
+        scale: float = 1.0,
+    ) -> None:
+        if hasattr(network, "_orig_mod"):  # unwrap torch.compile
+            network = network._orig_mod
+        linears = _target_linears(network)
+        sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"]
+        assert len(sd) == 2 * len(linears), (
+            f"edit-LoRA checkpoint has {len(sd)} tensors but the network "
+            f"exposes {2 * len(linears)} ({len(linears)} target projections); "
+            "target-list mismatch with the training recipe."
+        )
+
+        self._linears = linears
+        self._base: list[Tensor] = []
+        self._edit: list[Tensor] = []
+        added_bytes = 0
+        for i, lin in enumerate(linears):
+            a = sd[2 * i].to(lin.weight.device, torch.float32)
+            b = sd[2 * i + 1].to(lin.weight.device, torch.float32)
+            base = lin.weight.detach().clone()
+            w32 = base.to(torch.float32)
+            edit = w32.addmm_(b, a, alpha=scale).to(base.dtype)
+            self._base.append(base)
+            self._edit.append(edit)
+            added_bytes += 2 * base.numel() * base.element_size()
+        self.rank = int(sd[0].shape[0])
+        self.added_bytes = added_bytes
+        self.active = False
+
+    def set_active(self, active: bool) -> None:
+        """Copy the requested weight set into the live buffers (idempotent).
+
+        In-place ``copy_`` so the weight storage addresses never change —
+        captured CUDA graphs keep reading the same buffers and only the
+        contents differ.
+        """
+        if active == self.active:
+            return
+        source = self._edit if active else self._base
+        for lin, w in zip(self._linears, source):
+            lin.weight.data.copy_(w)
+        self.active = active
+
+    def release_targets(self, linears: list[nn.Linear]) -> list[Tensor]:
+        """Stop toggling the given live linears; return their fp32 deltas.
+
+        Composition seam for the fused drift-corrector dispatch
+        (:class:`omnidreams._drift_corrector.DriftCorrectorDispatch`): the
+        dispatch becomes the sole writer of the released projections and
+        folds the returned ``edit - base`` deltas into its per-state
+        pre-merged weight sets, while this hook keeps toggling only the
+        remaining (cross-attention) projections. Deltas are returned in
+        the order of ``linears``.
+        """
+        assert not self.active, "release targets while the base weights are live"
+        index = {id(lin): i for i, lin in enumerate(self._linears)}
+        drop: set[int] = set()
+        deltas: list[Tensor] = []
+        for lin in linears:
+            assert id(lin) in index, "linear is not one of this LoRA's targets"
+            i = index[id(lin)]
+            deltas.append(
+                self._edit[i].to(torch.float32) - self._base[i].to(torch.float32)
+            )
+            drop.add(i)
+        keep = [i for i in range(len(self._linears)) if i not in drop]
+        self._linears = [self._linears[i] for i in keep]
+        self._base = [self._base[i] for i in keep]
+        self._edit = [self._edit[i] for i in keep]
+        return deltas
+
+    def describe(self) -> str:
+        """One-line deploy description for startup logs."""
+        return (
+            f"text-edit LoRA r{self.rank} pre-merged on "
+            f"{len(self._linears)} projections "
+            f"(+{self.added_bytes / 2**20:.0f} MiB weight sets)"
+        )
diff --git a/integrations/omnidreams/tests/test_edit_lora.py b/integrations/omnidreams/tests/test_edit_lora.py
new file mode 100644
index 000000000..e6ae7da90
--- /dev/null
+++ b/integrations/omnidreams/tests/test_edit_lora.py
@@ -0,0 +1,150 @@
+# 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 pre-merged text-edit LoRA deploy hook.
+
+Covers the deploy invariants: ``TextEditLoRA`` merges ``W + B @ A``
+correctly, toggles by in-place ``copy_`` (stable storage addresses),
+restores the base bit-exactly, is idempotent, and ``release_targets``
+hands projections over to the fused drift-corrector dispatch.
+
+The transformer-level window tests (``replace_text_embeddings`` building
+a ``use_lora`` window, ``predict_flow`` expiry) live on the branch that
+carries the text-edit transformer machinery.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from omnidreams._edit_lora import TextEditLoRA, _target_linears
+from omnidreams.transformer import CosmosTransformer, CosmosTransformerConfig
+from omnidreams.transformer.impl.network import CosmosDiTNetworkConfig
+
+pytestmark = pytest.mark.ci_cpu
+
+
+def _tiny_transformer(seed: int = 0) -> CosmosTransformer:
+    torch.manual_seed(seed)
+    config = CosmosTransformerConfig(
+        network=CosmosDiTNetworkConfig(
+            in_channels=16,
+            out_channels=16,
+            patch_spatial=2,
+            patch_temporal=1,
+            model_channels=64,
+            num_blocks=2,
+            num_heads=4,
+            adaln_lora_dim=8,
+            crossattn_proj_in_channels=32,
+            crossattn_emb_channels=16,
+            additional_concat_ch=0,
+            enable_cross_view_attn=False,
+        ),
+        checkpoint_path=None,
+        batch_shape=(1,),
+        num_views=1,
+        len_t=2,
+        window_size_t=6,
+        sink_size_t=0,
+        compile_network=False,
+        use_cuda_graph=False,
+        guidance_scale=1.0,
+    )
+    return CosmosTransformer(config)
+
+
+def _fake_checkpoint(network, rank: int = 4, path=None):
+    torch.manual_seed(3)
+    linears = _target_linears(network)
+    sd = {}
+    for i, lin in enumerate(linears):
+        sd[2 * i] = torch.randn(rank, lin.in_features) * 0.02  # A
+        sd[2 * i + 1] = torch.randn(lin.out_features, rank) * 0.02  # B
+    torch.save({"lora": sd}, path)
+    return linears, sd
+
+
+def test_merge_toggle_and_bit_exact_restore(tmp_path):
+    transformer = _tiny_transformer()
+    ckpt = tmp_path / "edit_lora.pt"
+    linears, sd = _fake_checkpoint(transformer.network, path=ckpt)
+    base = [lin.weight.detach().clone() for lin in linears]
+    ptrs = [lin.weight.data_ptr() for lin in linears]
+
+    edit_lora = TextEditLoRA(transformer.network, ckpt)
+    assert edit_lora.rank == 4
+    assert len(linears) == 2 * 8  # 2 tiny blocks x 8 projections
+
+    edit_lora.set_active(True)
+    for i, lin in enumerate(linears):
+        expected = (
+            base[i].to(torch.float32) + sd[2 * i + 1].float() @ sd[2 * i].float()
+        ).to(base[i].dtype)
+        assert torch.equal(lin.weight, expected)
+        assert lin.weight.data_ptr() == ptrs[i]  # in place: CUDA-graph safe
+    edit_lora.set_active(True)  # idempotent
+
+    edit_lora.set_active(False)
+    for i, lin in enumerate(linears):
+        assert torch.equal(lin.weight, base[i])
+        assert lin.weight.data_ptr() == ptrs[i]
+
+
+def test_checkpoint_shape_mismatch_rejected(tmp_path):
+    transformer = _tiny_transformer()
+    ckpt = tmp_path / "bad.pt"
+    torch.save({"lora": {0: torch.zeros(4, 8), 1: torch.zeros(8, 4)}}, ckpt)
+    with pytest.raises(AssertionError, match="target-list mismatch"):
+        TextEditLoRA(transformer.network, ckpt)
+
+
+def test_release_targets_returns_deltas_and_stops_toggling_them(tmp_path):
+    transformer = _tiny_transformer()
+    ckpt = tmp_path / "edit_lora.pt"
+    linears, sd = _fake_checkpoint(transformer.network, path=ckpt)
+    base = [lin.weight.detach().clone() for lin in linears]
+    edit_lora = TextEditLoRA(transformer.network, ckpt)
+
+    released = [lin for i, lin in enumerate(linears) if i % 2 == 0]
+    kept = [(i, lin) for i, lin in enumerate(linears) if i % 2 == 1]
+    deltas = edit_lora.release_targets(released)
+
+    # Deltas are the merged edit minus base (== B @ A here, up to the base
+    # dtype's rounding — the tiny network is bf16), in released order.
+    for delta, lin in zip(deltas, released):
+        i = linears.index(lin)
+        expected = sd[2 * i + 1].float() @ sd[2 * i].float()
+        assert torch.allclose(delta, expected, atol=1e-3)
+        assert delta.abs().max() > 0
+
+    # Toggling now touches only the kept projections.
+    edit_lora.set_active(True)
+    for lin in released:
+        assert torch.equal(lin.weight, base[linears.index(lin)])
+    for i, lin in kept:
+        assert not torch.equal(lin.weight, base[i])
+    edit_lora.set_active(False)
+    for i, lin in enumerate(linears):
+        assert torch.equal(lin.weight, base[i])
+
+    # Releasing a foreign linear or releasing while active is rejected.
+    import torch.nn as nn
+
+    with pytest.raises(AssertionError, match="not one of"):
+        edit_lora.release_targets([nn.Linear(4, 4)])
+    edit_lora.set_active(True)
+    with pytest.raises(AssertionError, match="base weights"):
+        edit_lora.release_targets([kept[0][1]])

From 4cd5e0cec8b245de74d8cfa8ee30166f97464cd7 Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Sat, 22 Aug 2026 08:19:52 +0000
Subject: [PATCH 6/8] Add CUDA-graph-safe fused drift corrector with per-state
 dispatch

Brings the canonical drift-corrector deploy hook from
omnidreams-live-edit. DRIFT_CORRECTOR_MODE=fused keeps the per-alpha
pre-merged weight sets and the CPU-side call-index gate but swaps via
batched in-place copy_ into the original parameter storages, so the
exact alpha*(t) profile survives under compile_network + use_cuda_graph
(the unfused path forced both off). DriftCorrectorDispatch registers
multiple corrector states (checkpoint + gate profile + gain), each
pre-merged from a pristine base-weight snapshot and optionally composed
with a style-LoRA delta via TextEditLoRA.release_targets(), with a VRAM
budget guard and chunk-boundary state switching. Includes CPU
equivalence/pointer-stability tests and a GPU validation script
(parity, graph-capture check, chunk-latency bench, eyes-on frames).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../omnidreams/omnidreams/_drift_corrector.py | 615 ++++++++++++++++++
 .../scripts/validate_fused_corrector.py       | 347 ++++++++++
 .../omnidreams/tests/test_drift_corrector.py  | 538 +++++++++++++++
 3 files changed, 1500 insertions(+)
 create mode 100644 integrations/omnidreams/omnidreams/_drift_corrector.py
 create mode 100644 integrations/omnidreams/scripts/validate_fused_corrector.py
 create mode 100644 integrations/omnidreams/tests/test_drift_corrector.py

diff --git a/integrations/omnidreams/omnidreams/_drift_corrector.py b/integrations/omnidreams/omnidreams/_drift_corrector.py
new file mode 100644
index 000000000..52fae2dbd
--- /dev/null
+++ b/integrations/omnidreams/omnidreams/_drift_corrector.py
@@ -0,0 +1,615 @@
+# 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 (owner
+decision 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.
+
+By default the LoRA is **pre-merged**: at load time each discrete
+``alpha*(t) * gain`` value gets its own cached copy of the target
+projection weights with the scaled delta folded in, and the per-step gate
+just swaps the cached set in — zero extra work in the hot path. The gate
+is driven CPU-side from the load-time solver schedule (one
+``predict_flow`` call per solver step), so the corrected forward issues
+the same kernels as base with no GPU timestep readback. Set
+``DRIFT_CORRECTOR_UNFUSED=1`` to fall back to the runtime A/B-matmul path
+(the pre-2026-07-25 behavior).
+
+CUDA-graph-safe ``fused`` mode (``DRIFT_CORRECTOR_MODE=fused``)
+---------------------------------------------------------------
+
+Neither default works under the accelerated serving stack
+(``compile_network=True`` + ``use_cuda_graph=True``): the unfused path
+adds live gated matmuls (new ops the captured graph never saw), and the
+pre-merged path rebinds ``lin.weight.data`` to a cached tensor — captured
+kernels reference the *original* storage address, so after capture the
+gate silently stops changing what the graph computes. Serving therefore
+had to disable acceleration to run the corrector (~6 fps vs 30 fps).
+
+The fused mode keeps the same pre-merged weight sets and the same
+CPU-side call-index gate, but swaps by ``copy_``-ing the cached set into
+the original parameter storages (batched ``torch._foreach_copy_``) — the
+:class:`~omnidreams._edit_lora.TextEditLoRA` mechanism, extended from a
+per-window toggle to a per-denoise-step one. This is graph-safe by
+construction: the transformer's ``CUDAGraphWrapper`` captures ONE network
+forward, and each solver step plus the ``finalize_kv_cache`` context
+forward is a separate replay of that graph (``timestep`` is a staged
+input), so the exact per-step ``alpha*(t) * gain`` profile survives — no
+gate collapse to a constant is needed, and the graph sees only fixed
+parameter addresses whose values change between replays.
+
+Tradeoff, stated honestly: each within-chunk alpha change is a
+device-to-device copy of the four attention projections' weights (one
+copy per distinct consecutive alpha; the default two-entry profile costs
+two swaps per chunk, a profile with its own context-noise entry costs
+three). That is HBM bandwidth the pointer-rebind mode does not spend —
+sub-millisecond per swap on datacenter parts, and orders of magnitude
+cheaper than the acceleration the corrector previously forfeited by
+forcing the eager stack (real-time ~30 fps down to ~6 fps). An in-place ``addmm_`` of the rank-16 delta
+difference would avoid the cached sets' VRAM, but repeated bf16
+accumulation drifts over long live-game rollouts; copying from pristine
+fp32-merged sets restores exact values on every swap. Fused mode does
+not compose with a concurrently attached ``TextEditLoRA`` (both ``copy_``
+into the same self-attention projections; the last writer wins) — same
+restriction as the pre-merged mode; use ``unfused`` for stacked deploys,
+or the per-state :class:`DriftCorrectorDispatch` below, whose composed
+weight sets carry the style-LoRA delta and the corrector delta in one
+``copy_`` source (resolving the last-writer-wins conflict for the
+self-attention projections; the LoRA hook must then stop toggling them).
+
+Mode selection: ``DRIFT_CORRECTOR_MODE`` = ``premerged`` (default) |
+``fused`` | ``unfused``; the legacy ``DRIFT_CORRECTOR_UNFUSED=1`` still
+forces ``unfused``.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+## Deploy policy
+
+
+def _gate_alpha() -> dict[float, float]:
+    """Resolve the gate profile: ``GATE_ALPHA_JSON`` override or the default.
+
+    The override file holds either a flat ``{timestep: alpha}`` mapping or
+    an object with a ``"gate_alpha"`` entry (the ``edit_sft/gate_style.py``
+    output format). Read once at import time, so set the variable before
+    importing this module.
+    """
+    path = os.environ.get("GATE_ALPHA_JSON", "")
+    if not path:
+        return {1000.0: 0.96, 803.0: 0.667}
+    return _load_gate_json(path)
+
+
+def _load_gate_json(path: Path | str) -> dict[float, float]:
+    """Load a ``{timestep: alpha}`` profile from a gate JSON file.
+
+    Accepts either a flat mapping or an object with a ``"gate_alpha"``
+    entry (the ``edit_sft/gate_style.py`` output format).
+    """
+    table = json.loads(Path(path).read_text())
+    table = table.get("gate_alpha", table)
+    profile = {float(t): float(a) for t, a in table.items()}
+    assert profile and all(0.0 < a <= 1.0 for a in profile.values()), (
+        f"gate profile {str(path)!r} must map timesteps to alphas in (0, 1]"
+    )
+    return profile
+
+
+GATE_ALPHA = _gate_alpha()
+"""Unbiased alpha*(t) from the step-0 systematicity gate. Default: the
+photoreal drift-pair profile (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 nearest entry (t=803 in the default profile), matching the
+evaluated deploy configs. ``GATE_ALPHA_JSON`` swaps in a measured profile
+(e.g. ``edit_sft/outputs/gate_style.json`` for styled worlds), which may
+add its own low-t entry for the context forward."""
+
+_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 _nearest_alpha(t: float, profile: dict[float, float] | None = None) -> float:
+    """Return the profile entry (default :data:`GATE_ALPHA`) nearest in t."""
+    profile = GATE_ALPHA if profile is None else profile
+    return min(profile.items(), key=lambda kv: abs(kv[0] - t))[1]
+
+
+def _target_linears(network: nn.Module) -> list[nn.Linear]:
+    """Target linears in checkpoint load order (same walk as ``_apply_lora``)."""
+    linears: list[nn.Linear] = []
+    for mname, module in network.named_modules():
+        for cname, child in module.named_children():
+            full = f"{mname}.{cname}" if mname else cname
+            if isinstance(child, nn.Linear) and any(t in full for t in _LORA_TARGETS):
+                linears.append(child)
+    return linears
+
+
+def _premerge_weight_sets(
+    linears: list[nn.Linear], sd: dict, gain: float
+) -> tuple[dict[float, list[Tensor]], int]:
+    """Cache ``W + gain*alpha*(B @ A)`` per distinct gate value.
+
+    ``sd`` holds the checkpoint tensors in load order (``A_i`` at ``2i``,
+    ``B_i`` at ``2i + 1``). The merge runs in fp32 (matching the unfused
+    path's fp32 delta) and is cast back to the base weight dtype.
+
+    Returns:
+        The per-alpha weight sets and the total cached bytes.
+    """
+    sets: dict[float, list[Tensor]] = {}
+    added_bytes = 0
+    for alpha in sorted(set(GATE_ALPHA.values())):
+        merged: list[Tensor] = []
+        for i, lin in enumerate(linears):
+            a = sd[2 * i].to(lin.weight.device, torch.float32)
+            b = sd[2 * i + 1].to(lin.weight.device, torch.float32)
+            w32 = lin.weight.detach().to(torch.float32, copy=True)
+            w = w32.addmm_(b, a, alpha=gain * alpha).to(lin.weight.dtype)
+            merged.append(w)
+            added_bytes += w.numel() * w.element_size()
+        sets[alpha] = merged
+    return sets, added_bytes
+
+
+_MODES = ("premerged", "fused", "unfused")
+
+
+def _resolve_mode(mode: str | None, unfused: bool | None) -> str:
+    """Resolve the deploy mode from explicit args, then the environment.
+
+    Precedence: explicit ``mode`` > explicit ``unfused`` bool >
+    ``DRIFT_CORRECTOR_MODE`` > legacy ``DRIFT_CORRECTOR_UNFUSED=1`` >
+    ``premerged``.
+    """
+    if mode is None and unfused is not None:
+        mode = "unfused" if unfused else "premerged"
+    if mode is None:
+        mode = os.environ.get("DRIFT_CORRECTOR_MODE", "")
+    if not mode:
+        legacy = os.environ.get("DRIFT_CORRECTOR_UNFUSED", "0") == "1"
+        mode = "unfused" if legacy else "premerged"
+    assert mode in _MODES, f"drift-corrector mode {mode!r} not in {_MODES}"
+    return mode
+
+
+def apply_drift_corrector(
+    runner: Any,
+    checkpoint: Path,
+    gain: float,
+    *,
+    unfused: bool | None = None,
+    mode: str | None = None,
+) -> 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.
+        unfused: Force the runtime A/B-matmul path instead of the default
+            per-step pre-merged weights. ``None`` reads the environment
+            (see :func:`_resolve_mode`).
+        mode: Explicit deploy mode (``premerged`` | ``fused`` |
+            ``unfused``); overrides ``unfused`` and the environment.
+            ``fused`` is the CUDA-graph-safe in-place-``copy_`` variant —
+            required whenever the serving stack runs with
+            ``compile_network`` / ``use_cuda_graph`` enabled.
+
+    Returns:
+        A log-line string describing the deployed configuration.
+    """
+    mode = _resolve_mode(mode, unfused)
+    unfused = mode == "unfused"
+    network = runner.pipeline.diffusion_model.transformer.network
+    if hasattr(network, "_orig_mod"):  # unwrap torch.compile
+        network = network._orig_mod
+    transformer = runner.pipeline.diffusion_model.transformer
+    sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"]
+
+    if unfused:
+        params = _apply_lora(network)
+        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))
+        orig_pf = transformer.predict_flow
+
+        # Per-step gate: rescale the LoRA to alpha*(t) x gain before every
+        # denoise step (nearest-t lookup; finalize_kv_cache calls positionally).
+        def gated_pf(*args, **kwargs):
+            ts = kwargs.get("timestep", args[1] if len(args) > 1 else None)
+            t = float(ts.reshape(-1).max())
+            _set_scale(network, _nearest_alpha(t) * gain)
+            return orig_pf(*args, **kwargs)
+
+        transformer.predict_flow = gated_pf
+        return f"corrected (alpha*(t) x {gain}, unfused)"
+
+    # Pre-merged paths: one cached weight set per distinct alpha*(t)
+    # value; the per-step gate installs the cached set — no LoRA matmuls
+    # in the hot path. "premerged" re-points ``lin.weight.data`` (zero
+    # copy, NOT CUDA-graph-safe); "fused" ``copy_``s into the original
+    # parameter storages (graph-safe: captured kernels keep reading the
+    # same addresses and only the values change between replays).
+    linears = _target_linears(network)
+    assert len(sd) == 2 * len(linears), (
+        f"corrector checkpoint has {len(sd)} LoRA tensors but the network "
+        f"exposes {2 * len(linears)}; rank or target mismatch."
+    )
+    if mode == "fused":
+        assert getattr(transformer, "_optimized_dit_executor", None) is None, (
+            "fused drift corrector merges into the PyTorch network's "
+            "weights, which the native optimized-DiT executor bypasses; "
+            "run with native_dit_acceleration='disabled'."
+        )
+    weight_sets, added_bytes = _premerge_weight_sets(linears, sd, gain)
+    current: list[float | None] = [None]
+
+    if mode == "fused":
+        live = [lin.weight.data for lin in linears]
+
+        def _swap(alpha: float) -> None:
+            if alpha != current[0]:
+                torch._foreach_copy_(live, weight_sets[alpha])
+                current[0] = alpha
+    else:
+
+        def _swap(alpha: float) -> None:
+            if alpha != current[0]:
+                for lin, w in zip(linears, weight_sets[alpha]):
+                    lin.weight.data = w
+                current[0] = alpha
+
+    # Drive the gate CPU-side. The scheduler makes exactly one
+    # ``predict_flow`` call per solver step in a Python loop, so each
+    # step's alpha resolves from the load-time schedule by call index —
+    # reading the timestep tensor back per step (the unfused path's
+    # ``float(timestep.max())``) would stall the CPU launch queue every
+    # solver step.
+    scheduler = runner.pipeline.diffusion_model.scheduler
+    step_alphas = [_nearest_alpha(t) for t in scheduler.denoising_step_list.tolist()]
+    ctx_alpha = _nearest_alpha(
+        float(runner.pipeline.diffusion_model.config.context_noise)
+    )
+    orig_sample = scheduler.sample
+
+    def gated_sample(initial_noise, predict_flow, rng=None):
+        calls = [0]
+
+        def pf(noisy, timestep):
+            assert calls[0] < len(step_alphas), "predict_flow calls > solver steps"
+            _swap(step_alphas[calls[0]])
+            calls[0] += 1
+            return predict_flow(noisy, timestep)
+
+        return orig_sample(initial_noise=initial_noise, predict_flow=pf, rng=rng)
+
+    scheduler.sample = gated_sample
+    orig_finalize = transformer.finalize_kv_cache
+
+    def gated_finalize(*args, **kwargs):
+        _swap(ctx_alpha)
+        return orig_finalize(*args, **kwargs)
+
+    transformer.finalize_kv_cache = gated_finalize
+    kind = "graph-safe fused" if mode == "fused" else "pre-merged"
+    return (
+        f"corrected (alpha*(t) x {gain}, {kind} {len(weight_sets)} weight "
+        f"sets, +{added_bytes / 2**20:.0f} MiB)"
+    )
+
+
+## Per-state dispatch (fused mode, multiple correctors)
+
+_VRAM_WARN_GIB = 8.0
+"""Warn when the dispatch's cached weight sets exceed this budget."""
+
+
+@dataclass
+class _CorrectorState:
+    """Pre-merged weight sets and the per-step gate schedule for one state."""
+
+    sets: dict[float, list[Tensor]]
+    step_alphas: list[float]
+    ctx_alpha: float
+    added_bytes: int
+
+
+class DriftCorrectorDispatch:
+    """Multiple pre-merged corrector states behind one graph-safe selector.
+
+    Extends the ``fused`` mode of :func:`apply_drift_corrector` from one
+    corrector to a registry of named states, each pre-merged from a
+    pristine base-weight snapshot taken at construction:
+
+        ``sets[alpha] = pristine (+ lora_delta) + alpha * gain * (B @ A)``
+
+    The optional ``lora_delta`` (e.g. a style ``TextEditLoRA``'s
+    self-attention deltas) rides the same ``copy_`` source as the
+    corrector delta, which resolves the fused mode's last-writer-wins
+    conflict: this dispatch becomes the SOLE writer of the self-attention
+    projection weights, so a concurrently attached edit LoRA must be
+    restricted to projections outside :data:`_LORA_TARGETS` (e.g.
+    cross-attention only) while the dispatch is installed.
+
+    A ``"base"`` state (pristine weights, corrector off) is registered at
+    construction and is the initial selection; re-register it to give the
+    base world its own corrector (e.g. the shipped photoreal checkpoint).
+
+    :meth:`set_active_corrector` is safe between rollouts and at chunk
+    boundaries (between ``scheduler.sample`` calls) — the recommended call
+    sites. A mid-rollout call is still graph-safe (weights only change
+    between graph replays) but takes effect at the next denoise step, so
+    one chunk mixes two states; keep swaps at chunk boundaries for clean
+    visuals.
+    """
+
+    def __init__(self, runner: Any) -> None:
+        diffusion_model = runner.pipeline.diffusion_model
+        transformer = diffusion_model.transformer
+        assert getattr(transformer, "_optimized_dit_executor", None) is None, (
+            "the fused drift-corrector dispatch merges into the PyTorch "
+            "network's weights, which the native optimized-DiT executor "
+            "bypasses; run with native_dit_acceleration='disabled'."
+        )
+        network = transformer.network
+        if hasattr(network, "_orig_mod"):  # unwrap torch.compile
+            network = network._orig_mod
+        self._linears = _target_linears(network)
+        self._live = [lin.weight.data for lin in self._linears]
+        self._pristine32 = [
+            lin.weight.detach().to(torch.float32, copy=True) for lin in self._linears
+        ]
+        self._step_ts = [
+            float(t) for t in diffusion_model.scheduler.denoising_step_list
+        ]
+        self._ctx_t = float(diffusion_model.config.context_noise)
+        self._states: dict[str, _CorrectorState] = {}
+        self._active = "base"
+        self._current: tuple[str, float] | None = None
+        self.register_state("base")
+        self._install_gate_driver(diffusion_model.scheduler, transformer)
+
+    def register_state(
+        self,
+        name: str,
+        *,
+        checkpoint: Path | str | None = None,
+        gain: float = 0.0,
+        gate_alpha: dict[float, float] | Path | str | None = None,
+        lora_delta: list[Tensor] | None = None,
+    ) -> str:
+        """Pre-merge and cache the weight sets for one named state.
+
+        Args:
+            name: State name for :meth:`set_active_corrector`.
+            checkpoint: Corrector LoRA checkpoint (``train_v2`` format);
+                ``None`` (or ``gain == 0``) makes this an off state.
+            gain: Global gain composed with the alpha*(t) profile.
+            gate_alpha: Per-state gate profile — a ``{timestep: alpha}``
+                dict, a gate-JSON path, or ``None`` for :data:`GATE_ALPHA`.
+            lora_delta: Optional fp32 weight deltas (one per target linear,
+                checkpoint load order) folded into EVERY set of this state,
+                including its within-state off set — the state-aware base.
+
+        Returns:
+            A log-line string describing the registered state.
+        """
+        if isinstance(gate_alpha, (str, Path)):
+            gate_alpha = _load_gate_json(gate_alpha)
+        profile = GATE_ALPHA if gate_alpha is None else gate_alpha
+        base32 = [w.clone() for w in self._pristine32]
+        if lora_delta is not None:
+            assert len(lora_delta) == len(self._linears), (
+                f"lora_delta has {len(lora_delta)} tensors for "
+                f"{len(self._linears)} target projections"
+            )
+            for w, d in zip(base32, lora_delta):
+                w.add_(d.to(w.device, w.dtype))
+
+        added_bytes = 0
+        if checkpoint is None or gain == 0.0:
+            merged = [w.to(lin.weight.dtype) for w, lin in zip(base32, self._linears)]
+            added_bytes = sum(w.numel() * w.element_size() for w in merged)
+            state = _CorrectorState(
+                sets={0.0: merged},
+                step_alphas=[0.0] * len(self._step_ts),
+                ctx_alpha=0.0,
+                added_bytes=added_bytes,
+            )
+        else:
+            sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"]
+            assert len(sd) == 2 * len(self._linears), (
+                f"corrector checkpoint has {len(sd)} LoRA tensors but the "
+                f"network exposes {2 * len(self._linears)}; rank or target "
+                "mismatch."
+            )
+            sets: dict[float, list[Tensor]] = {}
+            for alpha in sorted(set(profile.values())):
+                merged = []
+                for i, (lin, w32) in enumerate(zip(self._linears, base32)):
+                    a = sd[2 * i].to(w32.device, torch.float32)
+                    b = sd[2 * i + 1].to(w32.device, torch.float32)
+                    w = w32.clone().addmm_(b, a, alpha=gain * alpha)
+                    merged.append(w.to(lin.weight.dtype))
+                    added_bytes += merged[-1].numel() * merged[-1].element_size()
+                sets[alpha] = merged
+            state = _CorrectorState(
+                sets=sets,
+                step_alphas=[_nearest_alpha(t, profile) for t in self._step_ts],
+                ctx_alpha=_nearest_alpha(self._ctx_t, profile),
+                added_bytes=added_bytes,
+            )
+
+        self._states[name] = state
+        if name == self._active:
+            self._current = None  # re-registration invalidates live weights
+        total = sum(s.added_bytes for s in self._states.values())
+        if total > _VRAM_WARN_GIB * 2**30:
+            warnings.warn(
+                f"drift-corrector dispatch caches {total / 2**30:.1f} GiB of "
+                f"weight sets across {len(self._states)} states, over the "
+                f"~{_VRAM_WARN_GIB:.0f} GiB budget; drop states or gate "
+                "entries.",
+                ResourceWarning,
+                stacklevel=2,
+            )
+        return (
+            f"corrector state {name!r}: {len(state.sets)} weight sets "
+            f"(gain {gain}, +{state.added_bytes / 2**20:.0f} MiB, "
+            f"total {total / 2**20:.0f} MiB)"
+        )
+
+    def set_active_corrector(self, name: str) -> None:
+        """Select the state whose weight sets the gate driver installs.
+
+        Takes effect at the next gated forward (next denoise step /
+        context forward); call at chunk or rollout boundaries. Forces a
+        copy on that forward even if the alpha value matches, so a
+        re-selected state always restores exact pre-merged values.
+        """
+        assert name in self._states, (
+            f"unknown corrector state {name!r}; registered: {sorted(self._states)}"
+        )
+        self._active = name
+        self._current = None
+
+    @property
+    def active_state(self) -> str:
+        """Name of the currently selected state."""
+        return self._active
+
+    def _swap(self, alpha: float) -> None:
+        key = (self._active, alpha)
+        if key != self._current:
+            torch._foreach_copy_(self._live, self._states[self._active].sets[alpha])
+            self._current = key
+
+    def _install_gate_driver(self, scheduler: Any, transformer: Any) -> None:
+        """Drive the gate CPU-side by call index (see the fused-mode notes)."""
+        orig_sample = scheduler.sample
+
+        def gated_sample(initial_noise, predict_flow, rng=None):
+            calls = [0]
+
+            def pf(noisy, timestep):
+                # Resolve the state per call so a mid-rollout selector swap
+                # stays consistent between step_alphas and the weight sets.
+                step_alphas = self._states[self._active].step_alphas
+                assert calls[0] < len(step_alphas), "predict_flow calls > solver steps"
+                self._swap(step_alphas[calls[0]])
+                calls[0] += 1
+                return predict_flow(noisy, timestep)
+
+            return orig_sample(initial_noise=initial_noise, predict_flow=pf, rng=rng)
+
+        scheduler.sample = gated_sample
+        orig_finalize = transformer.finalize_kv_cache
+
+        def gated_finalize(*args, **kwargs):
+            self._swap(self._states[self._active].ctx_alpha)
+            return orig_finalize(*args, **kwargs)
+
+        transformer.finalize_kv_cache = gated_finalize
diff --git a/integrations/omnidreams/scripts/validate_fused_corrector.py b/integrations/omnidreams/scripts/validate_fused_corrector.py
new file mode 100644
index 000000000..4f1e0b5d1
--- /dev/null
+++ b/integrations/omnidreams/scripts/validate_fused_corrector.py
@@ -0,0 +1,347 @@
+# 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.
+
+"""GPU validation for the CUDA-graph-safe fused drift corrector.
+
+One phase per process (the deploy hooks mutate the network irreversibly);
+all phases share the same styled-rollout protocol: base clip prompt, hot
+prompt-swap to the arcade style at chunk ``SWAP_AT``, style-drift
+corrector attached at build with the styled gate profile, fixed seed.
+
+Phases (``PHASE`` env):
+
+* ``parity_unfused`` / ``parity_fused`` — eager pipeline (no compile, no
+  graphs), corrector in the named mode. Saves frames + per-chunk stats;
+  once both exist, writes ``parity.json`` (max/mean frame diff) and an
+  sbs frame stack for eyeballing.
+* ``bench_base`` — compile_network + use_cuda_graph ON, no corrector.
+* ``bench_fused`` — compile_network + use_cuda_graph ON, fused corrector.
+  Asserts the cond CUDA graph actually captured (no silent eager
+  fallback) after the rollout.
+* ``bench_unfused_eager`` — the pre-existing serving compromise: eager
+  pipeline with the unfused corrector.
+
+Bench phases run one untimed warmup rollout (compile/autotune) and
+``RUNS`` timed rollouts; per-chunk wall latency is synchronized, and
+steady-state stats skip the first ``WARM_CHUNKS`` chunks (graph capture
+for this config finishes around chunk 5: capture_ar_idx=3 plus two
+warmup forwards). fps = 8 frames per steady chunk / median latency.
+
+Run each phase from the flashdreams repo root, e.g.::
+
+    PHASE=parity_unfused .venv/bin/python \
+        integrations/omnidreams/scripts/validate_fused_corrector.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from pathlib import Path
+from types import SimpleNamespace
+
+# The gate profile is read at _drift_corrector import time.
+os.environ.setdefault(
+    "GATE_ALPHA_JSON",
+    str(Path(__file__).resolve().parents[1] / "edit_sft/outputs/gate_style.json"),
+)
+
+import numpy as np
+import torch
+from einops import rearrange
+from omnidreams._drift_corrector import GATE_ALPHA, apply_drift_corrector
+from omnidreams.config import SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE
+from omnidreams.pipeline import OmnidreamsPipeline
+from omnidreams.runner import (
+    DEFAULT_VIDEO_HEIGHT,
+    DEFAULT_VIDEO_WIDTH,
+    _load_first_frame,
+    _load_video,
+    _write_video,
+)
+from torch import Tensor
+
+from flashdreams.infra.config import derive_config
+from flashdreams.infra.cuda_graph import CUDAGraphWrapper
+
+_OMNI = Path(__file__).resolve().parents[1]
+
+PHASES = (
+    "parity_unfused",
+    "parity_fused",
+    "bench_base",
+    "bench_fused",
+    "bench_unfused_eager",
+)
+PHASE = os.environ.get("PHASE", "")
+SEED = int(os.environ.get("SEED", "42"))
+N_CHUNKS = int(os.environ.get("N_CHUNKS", "16"))
+SWAP_AT = int(os.environ.get("SWAP_AT", "2"))
+RUNS = int(os.environ.get("RUNS", "3"))
+WARM_CHUNKS = int(os.environ.get("WARM_CHUNKS", "8"))
+UUID = os.environ.get("UUID", "23599139-948f-4681-b7f4-74794113086d")
+CORRECTOR = Path(
+    os.environ.get(
+        "CORRECTOR", str(_OMNI / "edit_sft/outputs/lora_style_corrector_valpeak.pt")
+    )
+)
+GAIN = float(os.environ.get("CORRECTOR_GAIN", "0.25"))
+OUT_DIR = Path(
+    os.environ.get("OUT_DIR", str(_OMNI / "scripts/outputs/fused_corrector_validation"))
+)
+STYLE_PROMPT = (
+    "A bright arcade racing game world with exaggerated saturated "
+    "colors, clean stylized surfaces, and a cheerful sunny palette."
+)
+
+SAMPLES_ROOT = (
+    Path.home()
+    / ".cache/huggingface/hub/datasets--nvidia--omni-dreams-samples/snapshots"
+)
+
+
+def _sample_paths(uuid: str) -> tuple[Path, Path, str]:
+    hdmaps = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/*_hdmap.mp4"))
+    frames = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/first_frame.png"))
+    prompts = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/prompt.txt"))
+    assert hdmaps and frames and prompts, (
+        f"sample {uuid} not in the local HF cache under {SAMPLES_ROOT}"
+    )
+    return hdmaps[0], frames[0], prompts[0].read_text().strip()
+
+
+def _build_pipeline(*, accel: bool) -> OmnidreamsPipeline:
+    cfg = derive_config(
+        SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE,
+        enable_sync_and_profile=False,
+        diffusion_model=dict(
+            seed=SEED,
+            transformer=dict(compile_network=accel, use_cuda_graph=accel),
+        ),
+    )
+    pipe = cfg.setup()
+    assert isinstance(pipe, OmnidreamsPipeline)
+    return pipe.to("cuda")
+
+
+@torch.no_grad()
+def _rollout(
+    pipe: OmnidreamsPipeline,
+    *,
+    hdmap: Tensor,
+    first: Tensor,
+    base_prompt: str,
+    decode: bool = True,
+) -> tuple[Tensor | None, list[float]]:
+    """Styled rollout; returns (decoded video or None, per-chunk seconds)."""
+    device = pipe.device
+    pipe.diffusion_model._rng = torch.Generator(device=device).manual_seed(SEED)
+    cache = pipe.initialize_cache(text=[[base_prompt]], image=first)
+    chunks: list[Tensor] = []
+    latencies: list[float] = []
+    start = 0
+    for ar_idx in range(N_CHUNKS):
+        if ar_idx == SWAP_AT:
+            pipe.replace_text(cache, [[STYLE_PROMPT]])
+        num_frames = pipe.get_num_frames(ar_idx)
+        end = start + num_frames
+        assert end <= hdmap.shape[2], f"hdmap too short at chunk {ar_idx}"
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        chunk = pipe.generate(ar_idx, cache, hdmap=hdmap[:, :, start:end])
+        pipe.finalize(ar_idx, cache)
+        torch.cuda.synchronize()
+        latencies.append(time.perf_counter() - t0)
+        if decode:
+            chunks.append(chunk[0, 0].float().cpu())
+        start = end
+    del cache
+    torch.cuda.empty_cache()
+    video = torch.cat(chunks, dim=0) if decode else None
+    return video, latencies
+
+
+def _chunk_bounds() -> list[tuple[int, int]]:
+    bounds, start = [], 0
+    for ar_idx in range(N_CHUNKS):
+        n = 5 if ar_idx == 0 else 8
+        bounds.append((start, start + n))
+        start += n
+    return bounds
+
+
+def _to_uint8(video: Tensor) -> np.ndarray:
+    arr = rearrange(video, "t c h w -> t h w c").numpy()
+    return (((arr + 1) / 2 * 255).clip(0, 255)).astype(np.uint8)
+
+
+def _compare_parity() -> None:
+    frames = {m: np.load(OUT_DIR / f"frames_{m}.npy") for m in ("unfused", "fused")}
+    a = frames["fused"].astype(np.int16)
+    b = frames["unfused"].astype(np.int16)
+    diff = np.abs(a - b)
+    per_chunk = [
+        {
+            "chunk": i,
+            "mean_abs": float(diff[s:e].mean()),
+            "max_abs": int(diff[s:e].max()),
+        }
+        for i, (s, e) in enumerate(_chunk_bounds())
+    ]
+    report = {
+        "protocol": {
+            "uuid": UUID,
+            "seed": SEED,
+            "n_chunks": N_CHUNKS,
+            "swap_at": SWAP_AT,
+            "gain": GAIN,
+            "gate_alpha": {str(k): v for k, v in GATE_ALPHA.items()},
+        },
+        "max_abs_pixel_diff_uint8": int(diff.max()),
+        "mean_abs_pixel_diff_uint8": float(diff.mean()),
+        "per_chunk": per_chunk,
+    }
+    (OUT_DIR / "parity.json").write_text(json.dumps(report, indent=2))
+    # 5-frame sbs stack (unfused top, fused bottom) for the eyes-on check.
+    import mediapy as media
+
+    n = len(frames["unfused"])
+    t_idx = [0, n // 4, n // 2, 3 * n // 4, n - 1]
+    rows = [
+        np.concatenate([frames[m][t] for t in t_idx], axis=1)
+        for m in ("unfused", "fused")
+    ]
+    media.write_image(str(OUT_DIR / "parity_sbs.png"), np.concatenate(rows, axis=0))
+    print(json.dumps({k: v for k, v in report.items() if k != "per_chunk"}, indent=2))
+
+
+def main() -> None:
+    assert PHASE in PHASES, f"PHASE={PHASE!r} not in {PHASES}"
+    torch.set_grad_enabled(False)
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    hdmap_path, frame_path, clip_prompt = _sample_paths(UUID)
+    total_frames = 5 + (N_CHUNKS - 1) * 8
+    device = torch.device("cuda")
+    hdmap = _load_video(
+        hdmap_path,
+        pixel_height=DEFAULT_VIDEO_HEIGHT,
+        pixel_width=DEFAULT_VIDEO_WIDTH,
+        device=device,
+        dtype=torch.bfloat16,
+    )[:total_frames][None, None]
+    first = _load_first_frame(
+        frame_path,
+        pixel_height=DEFAULT_VIDEO_HEIGHT,
+        pixel_width=DEFAULT_VIDEO_WIDTH,
+        device=device,
+        dtype=torch.bfloat16,
+    )[None, None]
+
+    accel = PHASE in ("bench_base", "bench_fused")
+    pipe = _build_pipeline(accel=accel)
+
+    corrector_mode = {
+        "parity_unfused": "unfused",
+        "parity_fused": "fused",
+        "bench_fused": "fused",
+        "bench_unfused_eager": "unfused",
+        "bench_base": None,
+    }[PHASE]
+    if corrector_mode is not None:
+        mode_str = apply_drift_corrector(
+            SimpleNamespace(pipeline=pipe), CORRECTOR, GAIN, mode=corrector_mode
+        )
+        print(f"{PHASE}: {mode_str}", flush=True)
+
+    if PHASE.startswith("parity"):
+        video, _ = _rollout(pipe, hdmap=hdmap, first=first, base_prompt=clip_prompt)
+        assert video is not None
+        arr = _to_uint8(video)
+        np.save(OUT_DIR / f"frames_{corrector_mode}.npy", arr)
+        _write_video(
+            rearrange(video, "t c h w -> t h w c"),
+            OUT_DIR / f"video_{corrector_mode}.mp4",
+            fps=30,
+        )
+        print("ROLLOUT-DONE", flush=True)
+        if all((OUT_DIR / f"frames_{m}.npy").exists() for m in ("unfused", "fused")):
+            _compare_parity()
+        return
+
+    # Bench phases: one warmup rollout (compile/autotune/graph shakeout),
+    # then RUNS timed rollouts. The video of the last run is kept for the
+    # eyes-on style-hold check.
+    print(f"{PHASE}: warmup rollout ...", flush=True)
+    _rollout(pipe, hdmap=hdmap, first=first, base_prompt=clip_prompt, decode=False)
+    all_lat: list[list[float]] = []
+    video = None
+    for r in range(RUNS):
+        video, lat = _rollout(
+            pipe,
+            hdmap=hdmap,
+            first=first,
+            base_prompt=clip_prompt,
+            decode=r == RUNS - 1,
+        )
+        all_lat.append(lat)
+        print(f"run {r}: per-chunk ms = {[round(x * 1e3, 1) for x in lat]}", flush=True)
+
+    if accel:
+        transformer = pipe.diffusion_model.transformer
+        wrapper = transformer._network_call
+        assert isinstance(wrapper, CUDAGraphWrapper) and wrapper._graph is not None, (
+            "cond CUDA graph did not capture — graph break or eager fallback"
+        )
+        print("cond CUDA graph captured and replayed OK", flush=True)
+
+    steady = np.array([run[WARM_CHUNKS:] for run in all_lat])  # seconds
+    med = float(np.median(steady))
+    report = {
+        "phase": PHASE,
+        "accel": accel,
+        "corrector_mode": corrector_mode,
+        "protocol": {
+            "uuid": UUID,
+            "seed": SEED,
+            "n_chunks": N_CHUNKS,
+            "swap_at": SWAP_AT,
+            "runs": RUNS,
+            "warm_chunks": WARM_CHUNKS,
+            "gain": GAIN,
+        },
+        "steady_chunk_ms": {
+            "median": med * 1e3,
+            "mean": float(steady.mean()) * 1e3,
+            "p90": float(np.quantile(steady, 0.9)) * 1e3,
+            "min": float(steady.min()) * 1e3,
+            "max": float(steady.max()) * 1e3,
+        },
+        "fps_at_8_frames_per_chunk": 8.0 / med,
+        "peak_vram_gib": torch.cuda.max_memory_allocated() / 2**30,
+    }
+    (OUT_DIR / f"bench_{PHASE}.json").write_text(json.dumps(report, indent=2))
+    print(json.dumps(report["steady_chunk_ms"], indent=2))
+    print(f"fps ~= {report['fps_at_8_frames_per_chunk']:.1f}", flush=True)
+    if video is not None:
+        _write_video(
+            rearrange(video, "t c h w -> t h w c"),
+            OUT_DIR / f"video_{PHASE}.mp4",
+            fps=30,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/tests/test_drift_corrector.py b/integrations/omnidreams/tests/test_drift_corrector.py
new file mode 100644
index 000000000..697f91285
--- /dev/null
+++ b/integrations/omnidreams/tests/test_drift_corrector.py
@@ -0,0 +1,538 @@
+# 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 copy
+
+import pytest
+import torch
+import torch.nn as nn
+from omnidreams._drift_corrector import (
+    _LORA_RANK,
+    GATE_ALPHA,
+    DriftCorrectorDispatch,
+    _apply_lora,
+    _gate_alpha,
+    _LoRALinear,
+    _nearest_alpha,
+    _premerge_weight_sets,
+    _resolve_mode,
+    _set_scale,
+    _target_linears,
+    apply_drift_corrector,
+)
+
+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():
+    assert _nearest_alpha(1000.0) == GATE_ALPHA[1000.0]
+    assert _nearest_alpha(803.0) == GATE_ALPHA[803.0]
+    # The context-noise forward (t=128) resolves to the low-t entry,
+    # matching the evaluated deploy configs.
+    assert _nearest_alpha(128.0) == GATE_ALPHA[803.0]
+    # The deployed 2-step solver schedule (warped [1000, 350] -> ~[1000,
+    # 803]) resolves to the two gate entries in order.
+    assert [_nearest_alpha(t) for t in (1000.0, 802.9)] == [
+        GATE_ALPHA[1000.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())
+
+
+## GATE_ALPHA_JSON override
+
+
+def test_gate_alpha_defaults_to_photoreal_profile(monkeypatch):
+    monkeypatch.delenv("GATE_ALPHA_JSON", raising=False)
+    assert _gate_alpha() == {1000.0: 0.96, 803.0: 0.667}
+
+
+def test_gate_alpha_json_accepts_a_flat_mapping(tmp_path, monkeypatch):
+    path = tmp_path / "gate.json"
+    path.write_text('{"1000": 0.9, "803": 0.5, "128": 0.3}')
+    monkeypatch.setenv("GATE_ALPHA_JSON", str(path))
+    assert _gate_alpha() == {1000.0: 0.9, 803.0: 0.5, 128.0: 0.3}
+
+
+def test_gate_alpha_json_accepts_the_gate_style_output_format(tmp_path, monkeypatch):
+    path = tmp_path / "gate_style.json"
+    path.write_text(
+        '{"per_timestep": {"1000": {"alpha_star_unbiased": 0.9}},'
+        ' "gate_alpha": {"1000": 0.9, "803": 0.5}}'
+    )
+    monkeypatch.setenv("GATE_ALPHA_JSON", str(path))
+    assert _gate_alpha() == {1000.0: 0.9, 803.0: 0.5}
+
+
+def test_gate_alpha_json_rejects_out_of_range_alphas(tmp_path, monkeypatch):
+    path = tmp_path / "gate.json"
+    path.write_text('{"1000": 1.5}')
+    monkeypatch.setenv("GATE_ALPHA_JSON", str(path))
+    with pytest.raises(AssertionError):
+        _gate_alpha()
+
+
+def test_gate_alpha_json_missing_file_fails_loudly(tmp_path, monkeypatch):
+    monkeypatch.setenv("GATE_ALPHA_JSON", str(tmp_path / "absent.json"))
+    with pytest.raises(FileNotFoundError):
+        _gate_alpha()
+
+
+## Pre-merged weight path
+
+
+def _toy_network() -> nn.Module:
+    class Block(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)
+
+    return nn.Sequential(Block(), Block())
+
+
+def _random_checkpoint(linears) -> dict[int, torch.Tensor]:
+    sd = {}
+    for i, lin in enumerate(linears):
+        sd[2 * i] = 0.1 * torch.randn(_LORA_RANK, lin.in_features)
+        sd[2 * i + 1] = 0.1 * torch.randn(lin.out_features, _LORA_RANK)
+    return sd
+
+
+def test_target_linears_match_apply_lora_load_order():
+    torch.manual_seed(0)
+    net = _toy_network()
+    linears = _target_linears(net)
+    wrapped_net = copy.deepcopy(net)
+    params = _apply_lora(wrapped_net)
+    wrapped = [m for m in wrapped_net.modules() if isinstance(m, _LoRALinear)]
+    assert len(params) == 2 * len(linears) == 2 * len(wrapped)
+    for lin, w in zip(linears, wrapped):
+        assert torch.equal(lin.weight, w.base.weight)  # same walk order
+
+
+def test_premerged_weights_match_the_unfused_delta():
+    torch.manual_seed(0)
+    net = _toy_network()
+    linears = _target_linears(net)
+    sd = _random_checkpoint(linears)
+    gain = 0.25
+    sets, added_bytes = _premerge_weight_sets(linears, sd, gain)
+    assert set(sets) == set(GATE_ALPHA.values())
+    x = torch.randn(3, 4)
+    for alpha, merged in sets.items():
+        for i, lin in enumerate(linears):
+            unfused = lin(x) + gain * alpha * (x @ sd[2 * i].T @ sd[2 * i + 1].T)
+            premerged = torch.nn.functional.linear(x, merged[i])
+            assert torch.allclose(premerged, unfused, atol=1e-5)
+    n_weights = sum(lin.weight.numel() for lin in linears)
+    assert added_bytes == len(sets) * n_weights * 4  # fp32 toy weights
+
+
+def test_premerge_does_not_mutate_base_weights():
+    torch.manual_seed(0)
+    net = _toy_network()
+    linears = _target_linears(net)
+    before = [lin.weight.detach().clone() for lin in linears]
+    _premerge_weight_sets(linears, _random_checkpoint(linears), gain=1.0)
+    for lin, w in zip(linears, before):
+        assert torch.equal(lin.weight, w)
+
+
+## Mode resolution
+
+
+def test_resolve_mode_defaults_to_premerged(monkeypatch):
+    monkeypatch.delenv("DRIFT_CORRECTOR_MODE", raising=False)
+    monkeypatch.delenv("DRIFT_CORRECTOR_UNFUSED", raising=False)
+    assert _resolve_mode(None, None) == "premerged"
+
+
+def test_resolve_mode_env_and_legacy_env(monkeypatch):
+    monkeypatch.setenv("DRIFT_CORRECTOR_MODE", "fused")
+    assert _resolve_mode(None, None) == "fused"
+    monkeypatch.delenv("DRIFT_CORRECTOR_MODE", raising=False)
+    monkeypatch.setenv("DRIFT_CORRECTOR_UNFUSED", "1")
+    assert _resolve_mode(None, None) == "unfused"
+
+
+def test_resolve_mode_explicit_args_override_env(monkeypatch):
+    monkeypatch.setenv("DRIFT_CORRECTOR_MODE", "fused")
+    assert _resolve_mode("unfused", None) == "unfused"
+    assert _resolve_mode(None, unfused=True) == "unfused"
+    assert _resolve_mode(None, unfused=False) == "premerged"
+    assert _resolve_mode("fused", unfused=True) == "fused"  # mode wins
+
+
+def test_resolve_mode_rejects_unknown_modes():
+    with pytest.raises(AssertionError):
+        _resolve_mode("turbo", None)
+
+
+## End-to-end deploy-hook harness (toy runner)
+
+
+class _ToyTransformer(nn.Module):
+    """Minimal stand-in for CosmosTransformer's deploy surface."""
+
+    def __init__(self):
+        super().__init__()
+        torch.manual_seed(0)
+        self.network = _toy_network()
+
+    def predict_flow(self, noisy, timestep):
+        x = noisy
+        for block in self.network:
+            attn = block.self_attn
+            x = attn["output_proj"](
+                attn["q_proj"](x) + attn["k_proj"](x) + attn["v_proj"](x)
+            ) + block.mlp(x)
+        return x * (1.0 + timestep / 1000.0)
+
+    def finalize_kv_cache(self, noisy, timestep):
+        # Mirrors the base transformer: the context forward is one
+        # predict_flow call whose flow is discarded; return it here so
+        # tests can compare the corrected context forward across modes.
+        return self.predict_flow(noisy, timestep)
+
+
+class _ToyScheduler:
+    denoising_step_list = torch.tensor([1000.0, 803.0])
+
+    def sample(self, initial_noise, predict_flow, rng=None):
+        flows = []
+        for t in self.denoising_step_list:
+            flows.append(predict_flow(initial_noise, t))
+        return torch.stack(flows)
+
+
+def _toy_runner(tmp_path):
+    """Fake runner + checkpoint exposing what apply_drift_corrector uses."""
+    from types import SimpleNamespace
+
+    transformer = _ToyTransformer()
+    sd = _random_checkpoint(_target_linears(transformer.network))
+    ckpt = tmp_path / "corrector.pt"
+    torch.save({"lora": sd}, ckpt)
+    diffusion_model = SimpleNamespace(
+        transformer=transformer,
+        scheduler=_ToyScheduler(),
+        config=SimpleNamespace(context_noise=128),
+    )
+    runner = SimpleNamespace(pipeline=SimpleNamespace(diffusion_model=diffusion_model))
+    return runner, ckpt
+
+
+def test_fused_matches_unfused_per_step_and_context(tmp_path):
+    """The graph-safe fused mode reproduces the unfused gated outputs."""
+    runner_u, ckpt = _toy_runner(tmp_path)
+    runner_f, _ = _toy_runner(tmp_path)  # same seed -> identical weights
+
+    apply_drift_corrector(runner_u, ckpt, gain=0.25, mode="unfused")
+    tr_u = runner_u.pipeline.diffusion_model.transformer
+    apply_drift_corrector(runner_f, ckpt, gain=0.25, mode="fused")
+
+    torch.manual_seed(1)
+    noise = torch.randn(3, 4)
+    sched_f = runner_f.pipeline.diffusion_model.scheduler
+    flows_f = sched_f.sample(
+        initial_noise=noise,
+        predict_flow=runner_f.pipeline.diffusion_model.transformer.predict_flow,
+    )
+    flows_u = torch.stack(
+        [tr_u.predict_flow(noise, t) for t in _ToyScheduler.denoising_step_list]
+    )
+    assert torch.allclose(flows_f, flows_u, atol=1e-5)
+
+    ctx_t = torch.tensor(128.0)
+    ctx_f = runner_f.pipeline.diffusion_model.transformer.finalize_kv_cache(
+        noise, ctx_t
+    )
+    ctx_u = tr_u.finalize_kv_cache(noise, ctx_t)
+    assert torch.allclose(ctx_f, ctx_u, atol=1e-5)
+
+
+def test_fused_swaps_change_outputs_across_steps(tmp_path):
+    """The two solver steps see different alphas (gate actually swaps)."""
+    runner, ckpt = _toy_runner(tmp_path)
+    dm = runner.pipeline.diffusion_model
+    apply_drift_corrector(runner, ckpt, gain=1.0, mode="fused")
+    torch.manual_seed(1)
+    noise = torch.randn(3, 4)
+    flows = dm.scheduler.sample(
+        initial_noise=noise, predict_flow=dm.transformer.predict_flow
+    )
+    # Undo the deterministic timestep modulation, leaving only the
+    # weight-set difference between the two steps.
+    step0 = flows[0] / (1.0 + 1000.0 / 1000.0)
+    step1 = flows[1] / (1.0 + 803.0 / 1000.0)
+    assert not torch.allclose(step0, step1, atol=1e-6)
+
+
+def test_fused_preserves_weight_storage_premerged_does_not(tmp_path):
+    """Graph-safety property: fused copies in place, premerged rebinds.
+
+    Captured CUDA graphs reference parameter storage addresses; only the
+    fused mode keeps them stable across gate swaps.
+    """
+
+    def drive_and_collect_ptrs(mode):
+        runner, ckpt = _toy_runner(tmp_path)
+        dm = runner.pipeline.diffusion_model
+        linears = _target_linears(dm.transformer.network)
+        before = [lin.weight.data_ptr() for lin in linears]
+        apply_drift_corrector(runner, ckpt, gain=0.25, mode=mode)
+        torch.manual_seed(1)
+        noise = torch.randn(3, 4)
+        dm.scheduler.sample(
+            initial_noise=noise, predict_flow=dm.transformer.predict_flow
+        )
+        dm.transformer.finalize_kv_cache(noise, torch.tensor(128.0))
+        after = [lin.weight.data_ptr() for lin in linears]
+        return before, after
+
+    before, after = drive_and_collect_ptrs("fused")
+    assert before == after
+    before, after = drive_and_collect_ptrs("premerged")
+    assert before != after
+
+
+def test_fused_refuses_native_dit_executor(tmp_path):
+    runner, ckpt = _toy_runner(tmp_path)
+    runner.pipeline.diffusion_model.transformer._optimized_dit_executor = object()
+    with pytest.raises(AssertionError, match="optimized-DiT"):
+        apply_drift_corrector(runner, ckpt, gain=0.25, mode="fused")
+
+
+## Per-state dispatch
+
+
+def _sample_flows(runner, noise):
+    dm = runner.pipeline.diffusion_model
+    return dm.scheduler.sample(
+        initial_noise=noise, predict_flow=dm.transformer.predict_flow
+    )
+
+
+def _lora_delta(linears, scale=0.05):
+    gen = torch.Generator().manual_seed(7)
+    return [scale * torch.randn(lin.weight.shape, generator=gen) for lin in linears]
+
+
+def test_dispatch_registers_base_and_rejects_unknown_states(tmp_path):
+    runner, ckpt = _toy_runner(tmp_path)
+    dispatch = DriftCorrectorDispatch(runner)
+    assert dispatch.active_state == "base"
+    summary = dispatch.register_state("skin", checkpoint=ckpt, gain=0.25)
+    assert "skin" in summary and "weight sets" in summary
+    dispatch.set_active_corrector("skin")
+    assert dispatch.active_state == "skin"
+    with pytest.raises(AssertionError, match="unknown corrector state"):
+        dispatch.set_active_corrector("nope")
+
+
+def test_dispatch_base_state_reproduces_pristine_outputs(tmp_path):
+    runner, ckpt = _toy_runner(tmp_path)
+    pristine, _ = _toy_runner(tmp_path)  # same seed -> identical weights
+    dispatch = DriftCorrectorDispatch(runner)
+    dispatch.register_state("skin", checkpoint=ckpt, gain=1.0)
+    torch.manual_seed(1)
+    noise = torch.randn(3, 4)
+    assert torch.allclose(_sample_flows(runner, noise), _sample_flows(pristine, noise))
+
+
+def test_dispatch_composed_state_matches_sequential_lora_then_corrector(tmp_path):
+    """base + lora_delta + alpha*gain*corrector == applying them in sequence."""
+    runner, ckpt = _toy_runner(tmp_path)
+    reference, _ = _toy_runner(tmp_path)  # same seed -> identical weights
+    gain = 0.25
+
+    ref_linears = _target_linears(
+        reference.pipeline.diffusion_model.transformer.network
+    )
+    delta = _lora_delta(ref_linears)
+    with torch.no_grad():
+        for lin, d in zip(ref_linears, delta):  # sequential: merge LoRA first
+            lin.weight.add_(d)
+    apply_drift_corrector(reference, ckpt, gain=gain, mode="unfused")
+    tr_ref = reference.pipeline.diffusion_model.transformer
+
+    dispatch = DriftCorrectorDispatch(runner)
+    dispatch.register_state("skin", checkpoint=ckpt, gain=gain, lora_delta=delta)
+    dispatch.set_active_corrector("skin")
+
+    torch.manual_seed(1)
+    noise = torch.randn(3, 4)
+    flows = _sample_flows(runner, noise)
+    flows_ref = torch.stack(
+        [tr_ref.predict_flow(noise, t) for t in _ToyScheduler.denoising_step_list]
+    )
+    assert torch.allclose(flows, flows_ref, atol=1e-5)
+    ctx_t = torch.tensor(128.0)
+    ctx = runner.pipeline.diffusion_model.transformer.finalize_kv_cache(noise, ctx_t)
+    assert torch.allclose(ctx, tr_ref.finalize_kv_cache(noise, ctx_t), atol=1e-5)
+
+
+def test_dispatch_selector_switching_restores_exact_state_weights(tmp_path):
+    runner, ckpt = _toy_runner(tmp_path)
+    pristine, _ = _toy_runner(tmp_path)
+    dispatch = DriftCorrectorDispatch(runner)
+    dispatch.register_state("skin", checkpoint=ckpt, gain=1.0)
+    torch.manual_seed(1)
+    noise = torch.randn(3, 4)
+    base_flows = _sample_flows(pristine, noise)
+
+    dispatch.set_active_corrector("skin")
+    skin_flows = _sample_flows(runner, noise)
+    assert not torch.allclose(skin_flows, base_flows, atol=1e-6)
+    dispatch.set_active_corrector("base")
+    assert torch.allclose(_sample_flows(runner, noise), base_flows)
+    dispatch.set_active_corrector("skin")
+    assert torch.allclose(_sample_flows(runner, noise), skin_flows)
+
+
+def test_dispatch_keeps_weight_storage_addresses_across_states(tmp_path):
+    runner, ckpt = _toy_runner(tmp_path)
+    linears = _target_linears(runner.pipeline.diffusion_model.transformer.network)
+    before = [lin.weight.data_ptr() for lin in linears]
+    dispatch = DriftCorrectorDispatch(runner)
+    dispatch.register_state("skin", checkpoint=ckpt, gain=0.5)
+    torch.manual_seed(1)
+    noise = torch.randn(3, 4)
+    for state in ("skin", "base", "skin"):
+        dispatch.set_active_corrector(state)
+        _sample_flows(runner, noise)
+    assert [lin.weight.data_ptr() for lin in linears] == before
+
+
+def test_dispatch_gate_profiles_are_per_state(tmp_path):
+    runner, ckpt = _toy_runner(tmp_path)
+    dispatch = DriftCorrectorDispatch(runner)
+    dispatch.register_state(
+        "skin", checkpoint=ckpt, gain=1.0, gate_alpha={1000.0: 0.9, 803.0: 0.4}
+    )
+    dispatch.register_state(
+        "weather", checkpoint=ckpt, gain=1.0, gate_alpha={1000.0: 0.3}
+    )
+    assert dispatch._states["skin"].step_alphas == [0.9, 0.4]
+    assert dispatch._states["skin"].ctx_alpha == 0.4  # t=128 -> nearest 803
+    assert dispatch._states["weather"].step_alphas == [0.3, 0.3]
+    assert dispatch._states["weather"].ctx_alpha == 0.3
+    assert dispatch._states["base"].step_alphas == [0.0, 0.0]
+
+
+def test_dispatch_gate_profile_accepts_a_json_path(tmp_path):
+    runner, ckpt = _toy_runner(tmp_path)
+    gate = tmp_path / "gate.json"
+    gate.write_text('{"1000": 0.8, "803": 0.2}')
+    dispatch = DriftCorrectorDispatch(runner)
+    dispatch.register_state("skin", checkpoint=ckpt, gain=1.0, gate_alpha=gate)
+    assert dispatch._states["skin"].step_alphas == [0.8, 0.2]
+
+
+def test_dispatch_warns_over_the_vram_budget(tmp_path, monkeypatch):
+    import omnidreams._drift_corrector as dc
+
+    runner, ckpt = _toy_runner(tmp_path)
+    dispatch = DriftCorrectorDispatch(runner)
+    monkeypatch.setattr(dc, "_VRAM_WARN_GIB", 1e-9)
+    with pytest.warns(ResourceWarning, match="over the"):
+        dispatch.register_state("skin", checkpoint=ckpt, gain=0.25)
+
+
+def test_dispatch_refuses_native_dit_executor(tmp_path):
+    runner, _ = _toy_runner(tmp_path)
+    runner.pipeline.diffusion_model.transformer._optimized_dit_executor = object()
+    with pytest.raises(AssertionError, match="optimized-DiT"):
+        DriftCorrectorDispatch(runner)

From 630b66b417882c3857b7df92be1df2b30d60f01d Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Sat, 22 Aug 2026 08:27:33 +0000
Subject: [PATCH 7/8] Fix fp32 base-weight aliasing in TextEditLoRA's pre-merge

On an fp32 network base.to(torch.float32) returns base itself,
so the in-place addmm_ merge corrupted the cached base weight set
and deactivating the edit could not restore the original weights.
Deploys are bf16 (where .to copies), so serving was unaffected;
merge with an explicit copy and add an fp32 regression test.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../omnidreams/omnidreams/_edit_lora.py       |  4 ++-
 .../omnidreams/tests/test_edit_lora.py        | 28 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/integrations/omnidreams/omnidreams/_edit_lora.py b/integrations/omnidreams/omnidreams/_edit_lora.py
index f6193eaff..f19f39783 100644
--- a/integrations/omnidreams/omnidreams/_edit_lora.py
+++ b/integrations/omnidreams/omnidreams/_edit_lora.py
@@ -109,7 +109,9 @@ def __init__(
             a = sd[2 * i].to(lin.weight.device, torch.float32)
             b = sd[2 * i + 1].to(lin.weight.device, torch.float32)
             base = lin.weight.detach().clone()
-            w32 = base.to(torch.float32)
+            # copy=True: on an fp32 network ``.to`` would alias ``base`` and
+            # the in-place merge below would corrupt the cached base set.
+            w32 = base.to(torch.float32, copy=True)
             edit = w32.addmm_(b, a, alpha=scale).to(base.dtype)
             self._base.append(base)
             self._edit.append(edit)
diff --git a/integrations/omnidreams/tests/test_edit_lora.py b/integrations/omnidreams/tests/test_edit_lora.py
index e6ae7da90..b9c084a93 100644
--- a/integrations/omnidreams/tests/test_edit_lora.py
+++ b/integrations/omnidreams/tests/test_edit_lora.py
@@ -148,3 +148,31 @@ def test_release_targets_returns_deltas_and_stops_toggling_them(tmp_path):
     edit_lora.set_active(True)
     with pytest.raises(AssertionError, match="base weights"):
         edit_lora.release_targets([kept[0][1]])
+
+
+def test_fp32_network_restores_base_bit_exactly(tmp_path):
+    """Regression: ``base.to(float32)`` on an fp32 network aliased ``base``,
+    so the in-place merge corrupted the cached base set and deactivating
+    the edit could not restore the original weights."""
+    import torch.nn as nn
+
+    torch.manual_seed(0)
+    network = nn.ModuleDict(
+        {
+            "self_attn": nn.ModuleDict(
+                {n: nn.Linear(8, 8, bias=False) for n in ("q_proj", "k_proj")}
+            )
+        }
+    )
+    assert all(lin.weight.dtype == torch.float32 for lin in _target_linears(network))
+    ckpt = tmp_path / "edit_lora.pt"
+    linears, _ = _fake_checkpoint(network, path=ckpt)
+    base = [lin.weight.detach().clone() for lin in linears]
+
+    edit_lora = TextEditLoRA(network, ckpt)
+    edit_lora.set_active(True)
+    for lin, w in zip(linears, base):
+        assert not torch.equal(lin.weight, w)  # the merge actually applied
+    edit_lora.set_active(False)
+    for lin, w in zip(linears, base):
+        assert torch.equal(lin.weight, w)

From 6221f0c9eb7f5fc7689848a8311e5f2a78bed116 Mon Sep 17 00:00:00 2001
From: wenqingw <wenqingw@nvidia.com>
Date: Sat, 22 Aug 2026 08:27:52 +0000
Subject: [PATCH 8/8] Add style-drift corrector data-gen, trainer, and gate
 calibration

The training side of the deployed style-drift corrector, ported from
omnidreams-live-edit:

- gen_style_drift_pairs.py rolls the counterfactual branch corpus: per
  clip and style, one branch per swap offset under the same seed, so
  every late-window (drifted, +8..+24) chunk has a clean early-window
  (+1..+4) counterpart at the same absolute index.
- train_style_corrector.py trains the rank-16 corrector on those pairs
  with the Clean Forcing recipe (DAG ratio + drift-contraction term,
  no-op episodes pinning identity off the drifted manifold); checkpoints
  load through omnidreams/_drift_corrector.py unchanged.
- gate_style.py measures the styled alpha*(t) profile (systematic-bias
  fraction of the drift gap per timestep) and writes the gate JSON the
  deploy hook's GATE_ALPHA_JSON override consumes.

Port adaptations vs the working branch: _host/_lora/_train_attn resolve
from the edit_sft vendored copies (no drift_correction sys.path);
LORA_TARGETS comes from omnidreams._edit_lora (byte-identical to the
guidance_distill list); _sample_files is inlined on top of the runner's
HF sync; video io goes through flashdreams.infra.runner_io like the
sibling edit_sft scripts.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../omnidreams/edit_sft/gate_style.py         | 473 +++++++++++++
 .../edit_sft/gen_style_drift_pairs.py         | 336 +++++++++
 .../edit_sft/train_style_corrector.py         | 656 ++++++++++++++++++
 3 files changed, 1465 insertions(+)
 create mode 100644 integrations/omnidreams/edit_sft/gate_style.py
 create mode 100644 integrations/omnidreams/edit_sft/gen_style_drift_pairs.py
 create mode 100644 integrations/omnidreams/edit_sft/train_style_corrector.py

diff --git a/integrations/omnidreams/edit_sft/gate_style.py b/integrations/omnidreams/edit_sft/gate_style.py
new file mode 100644
index 000000000..d63509f8e
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/gate_style.py
@@ -0,0 +1,473 @@
+# 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.
+
+"""Style-specific alpha*(t) gate for the style-drift corrector.
+
+``drift_correction/gate_faithful.py``'s measurement applied to the
+``gen_style_drift_pairs.py`` branch corpus: the deploy hook
+(``omnidreams/_drift_corrector.py``) currently gates the style corrector
+with the PHOTOREAL profile (``GATE_ALPHA`` = {1000: 0.96, 803: 0.667},
+measured on tiled-HDMap loop drift); this script measures the styled
+counterpart so styled worlds can be gated on their own numbers.
+
+Cells mirror ``train_style_corrector.py`` exactly: probe absolute chunk
+``k`` with the DRIFTED branch swapped ``k - a + 1`` chunks earlier
+(:data:`DEPTHS`, the +8..+24 long-hold blur regime) against the CLEAN
+styled reference branch swapped ``1..REF_MAX`` chunks earlier — same
+absolute index, same HDMap, same seed, matched RoPE by construction. Both
+branches replay a 3-chunk KV window at deploy semantics (style LoRA scale
+1 in-window, style text KV rebuilt at base weights) and are probed at the
+SAME ``z_t`` (built from the drifted chunk's x0) over :data:`M_NOISE`
+seeds; the velocity gap decomposes into systematic bias vs seed variance,
+``alpha*(t) = bias^2 / (bias^2 + var)`` (gate_faithful's estimator, both
+plug-in and unbiased). Probed timesteps: the 2-step solver schedule
+(t=1000 and the 450-warped t=803) plus the ``finalize_kv_cache`` context
+forward's t=128, which the deploy gate resolves by nearest-t lookup.
+
+Outputs ``edit_sft/outputs/gate_style.json`` (per-timestep, per-hold-depth
+breakdown, and a flat ``gate_alpha`` table consumable by the deploy hook's
+``GATE_ALPHA_JSON`` override), prints the comparison against the deployed
+photoreal profile and a one-line recommendation.
+
+Env knobs: ``PAIRS_DIR``, ``GATE_OUT``, ``STYLE_LORA``, ``STYLE_RANK``,
+``STYLES``, ``DEPTHS``, ``CELLS_PER_DEPTH``, ``REF_MAX``, ``M_NOISE``,
+``SEED``.
+
+Run from the flashdreams repo root (forward-only, fits the co-tenant
+VRAM share)::
+
+    .venv/bin/python integrations/omnidreams/edit_sft/gate_style.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+# Must land before the first CUDA allocation (co-tenant VRAM share).
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import numpy as np
+import torch
+from _host import CONTEXT_NOISE_SEED, build_pipeline, reset_history
+from _lora import apply_lora, load_lora, set_lora_scale, unwrap_compiled
+from _train_attn import functional_attention, patch_functional_attention
+from omnidreams._drift_corrector import GATE_ALPHA as PHOTOREAL_GATE
+from omnidreams._drift_corrector import _nearest_alpha
+from omnidreams._edit_lora import _LORA_TARGETS as LORA_TARGETS
+from style_prompts import clip_key
+from torch import Tensor
+
+## Gate configuration
+
+BASE = Path("integrations/omnidreams/edit_sft")
+PAIRS_DIR = Path(os.environ.get("PAIRS_DIR", str(BASE / "outputs/style_drift_pairs")))
+OUT_PATH = Path(os.environ.get("GATE_OUT", str(BASE / "outputs/gate_style.json")))
+
+STYLE_LORA = Path(
+    os.environ.get("STYLE_LORA", str(BASE / "outputs/lora_style_step1600.pt"))
+)
+"""Style-skin checkpoint (must be the one the branches were rolled with)."""
+
+STYLE_RANK = int(os.environ.get("STYLE_RANK", "64"))
+"""Rank of the style checkpoint (the eval RANK-must-match gotcha)."""
+
+STYLES = frozenset(s for s in os.environ.get("STYLES", "").split(",") if s)
+"""Optional slug filter; empty = every style found in the pairs dir."""
+
+DEPTHS = tuple(int(x) for x in os.environ.get("DEPTHS", "8,12,16,20,24").split(","))
+"""Style-hold depths (chunks since swap, 1-based) of the drifted branch to
+probe — the +8..+24 regime the corrector deploys on (its training band is
++8..+20; +24 checks the tail)."""
+
+CELLS_PER_DEPTH = int(os.environ.get("CELLS_PER_DEPTH", "1"))
+"""Probe cells per (clip, style, depth); 1 keeps the sweep ~30 min."""
+
+REF_MAX = int(os.environ.get("REF_MAX", "4"))
+"""Max hold depth of the clean styled reference (+1..+4 = the manifold
+``train_style_corrector.py`` teaches toward)."""
+
+REPLAY_CHUNKS = 3
+"""History chunks replayed before a probe = the full ``window_size_t=6`` /
+``len_t=2`` KV window (the trainer's setting)."""
+
+M_NOISE = int(os.environ.get("M_NOISE", "8"))
+"""Noise seeds per (cell, timestep) — gate_faithful's setting."""
+
+DEV_BAR = 0.05
+"""|style alpha* - deployed photoreal alpha| that matters: at the shipped
+gain 0.25 this is a >=7% relative change of the smaller effective per-step
+LoRA scale."""
+
+SEED = int(os.environ.get("SEED", "0"))
+
+
+def load_pairs() -> list[dict]:
+    """Load the branch corpus; per-branch full latent lists as the trainer."""
+    bases: dict[str, dict] = {}
+    pairs: list[dict] = []
+    for path in sorted(PAIRS_DIR.glob("*__*.pt")):
+        uuid, slug = path.stem.split("__", 1)
+        if STYLES and slug not in STYLES:
+            continue
+        if uuid not in bases:
+            base_path = PAIRS_DIR / f"{uuid}_base.pt"
+            assert base_path.exists(), f"missing {base_path} for {path.name}"
+            bases[uuid] = torch.load(base_path, map_location="cpu", weights_only=False)
+        base = bases[uuid]
+        n = int(base["n_chunks"])
+        assert len(base["latents"]) == n and len(base["hdmaps"]) == n
+        d = torch.load(path, map_location="cpu", weights_only=False)
+        full = {
+            int(s): base["latents"][: int(s)] + lat for s, lat in d["branches"].items()
+        }
+        pairs.append(
+            {
+                "uuid": uuid,
+                "slug": slug,
+                "n_chunks": n,
+                "hdmaps": base["hdmaps"],
+                "full": full,
+                "offsets": sorted(full),
+            }
+        )
+    assert pairs, f"no styled branch files under {PAIRS_DIR}"
+    return pairs
+
+
+def select_cells(p: dict, rng: np.random.Generator) -> list[tuple[int, int, int, int]]:
+    """Stratified probe cells ``(k, a, b, depth)`` for one (clip, style) pair.
+
+    Per drift depth in :data:`DEPTHS`, up to :data:`CELLS_PER_DEPTH` chunks
+    ``k`` with the drifted branch swapped at ``a = k - depth + 1`` and a
+    clean reference branch ``b`` at hold depth ``1..REF_MAX`` — the
+    trainer's cell rule, stratified by depth instead of pooled. With the
+    corpus offsets spaced 4 apart, multiple-of-4 depths pin the reference
+    at hold depth +4: the reference KV window is then fully styled content
+    (no unstyled pre-swap chunks), the cleanest form of the premise.
+    """
+    offsets, n = p["offsets"], p["n_chunks"]
+    cells: list[tuple[int, int, int, int]] = []
+    for depth in DEPTHS:
+        ks = [
+            k
+            for k in range(1 + REPLAY_CHUNKS, n)
+            if (k - depth + 1) in offsets
+            and any(1 <= k - b + 1 <= REF_MAX for b in offsets)
+        ]
+        if not ks:
+            print(f"{p['uuid'][:8]}__{p['slug']}: no cells at depth {depth}")
+            continue
+        picks = rng.choice(len(ks), size=min(CELLS_PER_DEPTH, len(ks)), replace=False)
+        for i in sorted(int(x) for x in picks):
+            k = ks[i]
+            refs = [b for b in offsets if 1 <= k - b + 1 <= REF_MAX]
+            b = refs[int(rng.integers(len(refs)))]
+            cells.append((k, k - depth + 1, b, depth))
+    return cells
+
+
+def main() -> None:
+    torch.set_grad_enabled(False)
+    rng = np.random.default_rng(SEED)
+    pairs = load_pairs()
+
+    prompt_emb = torch.load(
+        BASE / "outputs/style_embeddings.pt", map_location="cpu", weights_only=False
+    )
+    assets = torch.load(
+        BASE / "outputs/style_clip_assets.pt", map_location="cpu", weights_only=False
+    )
+
+    pipe = build_pipeline(with_oneshot_encoders=False)
+    device = pipe.device
+    dtype = torch.bfloat16
+    dm = pipe.diffusion_model
+    transformer = dm.transformer
+    scheduler = dm.scheduler
+    ctx_t = torch.tensor(float(dm.config.context_noise), device=device, dtype=dtype)
+    ctx_idx = torch.argmin((scheduler._full_timesteps - ctx_t.float()).abs())
+    sigma_ctx = float(scheduler._full_sigmas[ctx_idx])
+    # Probe points: the solver schedule (1000 and the 450-warped 803) plus
+    # the context forward (t=128) the deploy gate also rescales.
+    probe_ts: list[tuple[int, Tensor, float]] = [
+        (
+            int(t),
+            scheduler.denoising_step_list[i].to(device=device, dtype=dtype),
+            float(scheduler.denoising_sigmas[i]),
+        )
+        for i, t in enumerate(scheduler.denoising_step_list.tolist())
+    ] + [(int(ctx_t), ctx_t, sigma_ctx)]
+
+    # Style LoRA at deploy semantics (train_style_corrector's setup): frozen
+    # weights, per-chunk scale toggling; NO corrector — the gate measures
+    # the uncorrected styled drift gap the corrector is premised on.
+    network = unwrap_compiled(transformer.network)
+    network.requires_grad_(False)
+    style_names = apply_lora(network, rank=STYLE_RANK, targets=LORA_TARGETS)
+    load_lora(network, STYLE_LORA)
+    patch_functional_attention()
+    print(
+        f"style r{STYLE_RANK} on {len(style_names)} projections | "
+        f"{len(pairs)} (clip, style) pairs | depths {DEPTHS} x "
+        f"{CELLS_PER_DEPTH}/depth, {M_NOISE} seeds, "
+        f"timesteps {[t for t, _, _ in probe_ts]}",
+        flush=True,
+    )
+
+    # One long-lived cache; probes never replay chunk 0 (start >= 1), so
+    # the image/mask fields can come from any clip.
+    set_lora_scale(network, 0.0)
+    first_uuid = pairs[0]["uuid"]
+    cache = pipe.initialize_cache_from_embeddings(
+        text_embeddings=prompt_emb[clip_key(first_uuid)],
+        image_embeddings=assets["image_embeddings"][first_uuid],
+    )
+    tc = cache.transformer_cache
+    _cur_text = [clip_key(first_uuid)]
+
+    def swap_text(key: str) -> None:
+        """Rebuild the cross-attn text KV at BASE weights (deploy truth)."""
+        if _cur_text[0] == key:
+            return
+        set_lora_scale(network, 0.0)
+        pipe.replace_text_from_embeddings(cache, prompt_emb[key])
+        _cur_text[0] = key
+
+    def replay(
+        latents: list[Tensor],
+        hdmaps: list[Tensor],
+        upto: int,
+        swap_at: int,
+        slug: str,
+        uuid: str,
+    ) -> None:
+        """Rebuild the KV window for chunks ``upto - 3 .. upto - 1``.
+
+        The trainer's semantics: pre-swap chunks commit at base weights
+        under the clip's own prompt; chunks at/after ``swap_at`` commit at
+        style scale 1 under the style prompt.
+        """
+        start = max(0, upto - REPLAY_CHUNKS)
+        reset_history(tc)
+        for bc in tc.network_cache.block_caches:
+            bc.self_attn._prev_chunk_idx = start - 1
+        swapped = swap_at <= start
+        swap_text(slug if swapped else clip_key(uuid))
+        for j in range(start, upto):
+            if j == swap_at and not swapped:
+                swap_text(slug)
+                swapped = True
+            set_lora_scale(network, 1.0 if j >= swap_at else 0.0)
+            g = torch.Generator(device=device).manual_seed(CONTEXT_NOISE_SEED + 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)
+
+    agg = {
+        t: {"alphas": [], "alphas_ub": [], "rels": [], "rels_x0": []}
+        for t, _, _ in probe_ts
+    }
+    by_depth: dict[int, dict[int, dict[str, list[float]]]] = {}
+    for p in pairs:
+        uuid, slug = p["uuid"], p["slug"]
+        for k, a, b, depth in select_cells(p, rng):
+            end = k + 1
+            lat_d = [x.to(device, dtype) for x in p["full"][a][:end]]
+            lat_r = [x.to(device, dtype) for x in p["full"][b][:end]]
+            hd = [x.to(device, dtype) for x in p["hdmaps"][:end]]
+            x0 = lat_d[k]
+
+            z_ts: dict[int, list[Tensor]] = {}
+            for t_idx, (_, _, sig) in enumerate(probe_ts):
+                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, (lat, swap) in (
+                ("drift", (lat_d, a)),
+                ("clean", (lat_r, b)),
+            ):
+                replay(lat, hd, k, swap, slug, uuid)
+                # The probed chunk is in-window on both branches (hold
+                # depth >= 1): style text + scale 1, even when the swap
+                # happened AT k (depth-1 references) and the replay loop
+                # never reached it.
+                swap_text(slug)
+                set_lora_scale(network, 1.0)
+                tc.start(k)
+                with functional_attention():
+                    preds[name] = {
+                        t_idx: [
+                            transformer.predict_flow(
+                                noisy_latent=z, timestep=tt, cache=tc, input=hd[k]
+                            ).float()
+                            for z in z_ts[t_idx]
+                        ]
+                        for t_idx, (_, tt, _) in enumerate(probe_ts)
+                    }
+                tc.finalize(k)
+
+            line = [f"{uuid[:8]}__{slug} k={k:2d} d=+{depth:2d} (ref +{k - b + 1})"]
+            for t_idx, (t_lab, _, sig) in enumerate(probe_ts):
+                deltas = torch.stack(
+                    [
+                        d_ - c_
+                        for d_, c_ in zip(preds["drift"][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)
+                drift_norm = (
+                    torch.stack(preds["drift"][t_idx]).flatten(1).norm(dim=1).mean()
+                )
+                rel = (
+                    deltas.flatten(1).norm(dim=1).mean() / (drift_norm + 1e-12)
+                ).item()
+                x0_norm = (
+                    torch.stack(
+                        [
+                            z.float() - sig * v
+                            for z, v in zip(z_ts[t_idx], preds["drift"][t_idx])
+                        ]
+                    )
+                    .flatten(1)
+                    .norm(dim=1)
+                    .mean()
+                )
+                rel_x0 = (
+                    sig * deltas.flatten(1).norm(dim=1).mean() / (x0_norm + 1e-12)
+                ).item()
+                a_hat = bias_sq / (bias_sq + var + 1e-12)
+                a_ub = bias_sq_ub / (bias_sq_ub + var_ub + 1e-12)
+                agg[t_lab]["alphas"].append(a_hat)
+                agg[t_lab]["alphas_ub"].append(a_ub)
+                agg[t_lab]["rels"].append(rel)
+                agg[t_lab]["rels_x0"].append(rel_x0)
+                slot = by_depth.setdefault(depth, {}).setdefault(
+                    t_lab, {"alphas_ub": [], "rels": []}
+                )
+                slot["alphas_ub"].append(a_ub)
+                slot["rels"].append(rel)
+                line.append(f"t={t_lab:4d} a*={a_hat:.3f}/{a_ub:.3f} rel_v={rel:.4f}")
+            print(" | ".join(line), flush=True)
+
+    summary = {
+        "style_lora": str(STYLE_LORA),
+        "config": {
+            "depths": list(DEPTHS),
+            "cells_per_depth": CELLS_PER_DEPTH,
+            "ref_max": REF_MAX,
+            "m_noise": M_NOISE,
+            "seed": SEED,
+            "styles": sorted({p["slug"] for p in pairs}),
+            "clips": len({p["uuid"] for p in pairs}),
+        },
+        "per_timestep": {
+            str(t): {
+                "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 t, v in agg.items()
+            if v["alphas"]
+        },
+        "per_depth": {
+            str(d): {
+                str(t): {
+                    "alpha_star_unbiased": sum(s["alphas_ub"]) / len(s["alphas_ub"]),
+                    "rel_v": sum(s["rels"]) / len(s["rels"]),
+                    "cells": len(s["alphas_ub"]),
+                }
+                for t, s in sorted(slots.items())
+            }
+            for d, slots in sorted(by_depth.items())
+        },
+        "photoreal_gate_alpha": {str(int(t)): a for t, a in PHOTOREAL_GATE.items()},
+    }
+    # The deploy-facing profile (GATE_ALPHA_JSON consumes this entry):
+    # unbiased alpha*, the same estimator the photoreal 0.96/0.667 came from.
+    summary["gate_alpha"] = {
+        t: round(v["alpha_star_unbiased"], 3)
+        for t, v in summary["per_timestep"].items()
+    }
+    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
+    OUT_PATH.write_text(json.dumps(summary, indent=2))
+
+    print("\n========== Omnidreams STYLE gate (styled drift pairs, v-space) ==========")
+    for t, v in summary["per_timestep"].items():
+        print(
+            f"t={t:>4s}: alpha* {v['alpha_star']:.3f} (unbiased"
+            f" {v['alpha_star_unbiased']:.3f}) | rel_v {v['rel_v']:.4f}"
+            f" rel_x0 {v['rel_x0']:.4f} | {v['cells']} cells"
+        )
+    for d, slots in summary["per_depth"].items():
+        parts = [
+            f"t={t} a*_ub={s['alpha_star_unbiased']:.3f} rel_v={s['rel_v']:.3f}"
+            for t, s in slots.items()
+        ]
+        print(f"depth +{d:>2s}: " + " | ".join(parts))
+
+    print("\n========== style profile vs deployed photoreal GATE_ALPHA ==========")
+    deviating: list[str] = []
+    for t, a_style in summary["gate_alpha"].items():
+        a_photo = _nearest_alpha(float(t))
+        delta = a_style - a_photo
+        if abs(delta) >= DEV_BAR:
+            deviating.append(f"t={t} {a_photo:.3f}->{a_style:.3f} ({delta:+.3f})")
+        print(
+            f"t={t:>4s}: style {a_style:.3f} | photoreal deploy {a_photo:.3f}"
+            f" | delta {delta:+.3f}"
+        )
+    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())
+    )
+    if mean_rel < 0.01:
+        rec = "styled drift gap ~zero -> nothing to gate differently. STOP."
+    elif deviating:
+        rec = (
+            f"override for styled worlds: GATE_ALPHA_JSON={OUT_PATH} "
+            f"({'; '.join(deviating)} exceed the {DEV_BAR} bar)"
+        )
+    else:
+        rec = (
+            f"keep the photoreal profile (every timestep within {DEV_BAR} "
+            "of the deployed values)"
+        )
+    print(f"RECOMMENDATION: {rec}")
+    print(f"saved {OUT_PATH}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/edit_sft/gen_style_drift_pairs.py b/integrations/omnidreams/edit_sft/gen_style_drift_pairs.py
new file mode 100644
index 000000000..117925f75
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/gen_style_drift_pairs.py
@@ -0,0 +1,336 @@
+# 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.
+
+"""Style-drift pair generation for the style-drift corrector (GPU, one shot).
+
+The style-skin LoRA (``train_style_sft.py``) compounds through the KV
+history on long holds: the first ~1-4 chunks after a swap are the clean
+styled manifold, but by +8..+20 chunks the re-styled history has blurred
+the world out (``issues_and_fixes.md`` Issue 1). This script rolls the
+branch corpus the corrector trains on:
+
+- The style LoRA is deployed exactly as in serving: pre-merged
+  :class:`~omnidreams._edit_lora.TextEditLoRA` weights plus a ``use_lora``
+  edit window held open to the end of the rollout (the
+  ``scripts/smoke_text_edit.py`` ``EDIT_LORA`` + swap machinery).
+- Per source clip (``generate_sources.py`` manifest) and per style slug,
+  one branch per swap offset in :data:`SWAP_OFFSETS`, all under the SAME
+  seed. The swap draws no noise on this host, so branches are bit-equal
+  before their swap (asserted) and diverge only in style-hold depth.
+- Per chunk, the patchified x0 latent is captured exactly as the AR cache
+  consumed it (the ``drift_correction/_host.py:capture_rollout`` fields),
+  plus the patchified HDMap conditioning once per clip.
+
+Offsets spaced 4 apart make every late-window chunk of one branch (swap
+8-20 chunks in the past — drifted) share its absolute index with an
+early-window chunk of another branch (swap 1-4 chunks in the past — the
+clean styled manifold). ``train_style_corrector.py`` builds counterfactual
+probes from those (drifted, clean) branch pairs.
+
+Outputs (under ``OUT_DIR``):
+
+- ``<uuid>_base.pt``: ``{"latents": [n_chunks x [1, 1, L, D]],
+  "hdmaps": [...], "n_chunks", "seed"}`` — the unswapped rollout.
+- ``<uuid>__<slug>.pt``: ``{"branches": {offset: [latents for chunks
+  offset..n_chunks-1]}, "swap_offsets", "n_chunks", "seed",
+  "style_lora"}`` — pre-swap chunks live in the base file.
+
+Env knobs: ``STYLE_LORA``, ``STYLES``, ``SWAP_OFFSETS``, ``N_CLIPS``,
+``N_CHUNKS``, ``SEED``, ``OUT_DIR``, ``SAVE_MP4``.
+
+Run from the flashdreams repo root (after ``precompute_style.py``; needs
+no one-shot encoders — prompts come from ``style_embeddings.pt``)::
+
+    .venv/bin/python integrations/omnidreams/edit_sft/gen_style_drift_pairs.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+# Must land before the first CUDA allocation (co-tenant VRAM share).
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import torch
+from _host import build_pipeline
+from omnidreams._edit_lora import TextEditLoRA
+from omnidreams.pipeline import OmnidreamsPipeline
+from omnidreams.runner import (
+    DEFAULT_VIDEO_HEIGHT,
+    DEFAULT_VIDEO_WIDTH,
+    _ensure_hf_single_view_example_data_synced,
+)
+from style_prompts import STYLE_PROMPTS, clip_key
+from torch import Tensor
+
+from flashdreams.infra.runner_io import load_video_tensor, write_video_tensor
+
+## Configuration
+
+BASE = Path("integrations/omnidreams/edit_sft")
+OUT_DIR = Path(os.environ.get("OUT_DIR", str(BASE / "outputs/style_drift_pairs")))
+SRC_DIR = BASE / "outputs" / "sources"
+
+STYLE_LORA = Path(
+    os.environ.get("STYLE_LORA", str(BASE / "outputs/lora_style_step1600.pt"))
+)
+"""Style-skin checkpoint to deploy (``train_style_sft.py`` format)."""
+
+STYLES = tuple(
+    s for s in os.environ.get("STYLES", "arcade_racer,comic_ink").split(",") if s
+)
+"""Style slugs to roll branches for (the full-strength training styles)."""
+
+SWAP_OFFSETS = tuple(
+    int(x) for x in os.environ.get("SWAP_OFFSETS", "4,8,12,16,20,24").split(",")
+)
+"""Swap chunks, spaced 4 apart so every probe chunk has both a drifted
+(+8..+20) and a clean (+1..+4) branch counterpart at the same index."""
+
+N_CLIPS = int(os.environ.get("N_CLIPS", "0"))
+"""Source clips to roll (manifest order); 0 = all."""
+
+SEED = int(os.environ.get("SEED", "42"))
+SAVE_MP4 = os.environ.get("SAVE_MP4", "0") == "1"
+"""Also write the decoded branch rollouts for eyeballing the drift."""
+
+
+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)
+
+
+@torch.no_grad()
+def rollout(
+    pipe: OmnidreamsPipeline,
+    *,
+    text_emb: Tensor,
+    image_emb: Tensor,
+    style_emb: Tensor | None,
+    hdmap: Tensor,
+    n_chunks: int,
+    swap_at: int | None,
+) -> tuple[list[Tensor], list[Tensor], Tensor]:
+    """Roll ``n_chunks`` with an optional held-open style window at ``swap_at``.
+
+    Args:
+        pipe: Pipeline with the pre-merged style ``TextEditLoRA`` attached.
+        text_emb: ``[1, 1, L, D]`` clip-prompt embeddings.
+        image_emb: ``[1, 1, 1, Cl, Hl, Wl]`` first-frame embeddings.
+        style_emb: ``[1, 1, L, D]`` style-prompt embeddings (``swap_at``
+            branches only).
+        hdmap: ``[1, 1, T, 3, H, W]`` HDMap pixels covering ``n_chunks``.
+        n_chunks: AR chunks to roll.
+        swap_at: Chunk index of the style swap; ``None`` = base branch.
+
+    Returns:
+        ``(latents, hdmaps, video)`` — per-chunk patchified x0 and HDMap
+        conditioning exactly as the AR cache consumed them (CPU), plus the
+        decoded rollout ``[T, 3, H, W]`` on CPU.
+    """
+    device = pipe.device
+    pipe.diffusion_model._rng = torch.Generator(device=device).manual_seed(SEED)
+    cache = pipe.initialize_cache_from_embeddings(
+        text_embeddings=text_emb, image_embeddings=image_emb
+    )
+    latents: list[Tensor] = []
+    hdmaps: list[Tensor] = []
+    chunks: list[Tensor] = []
+    start = 0
+    for ar_idx in range(n_chunks):
+        if swap_at is not None and ar_idx == swap_at:
+            # The deploy swap: guidance_scale is unused on the use_lora
+            # path but must be != 1.0 to open the window; the countdown
+            # covers every remaining chunk (held-open long-hold regime).
+            assert style_emb is not None
+            pipe.replace_text_from_embeddings(
+                cache,
+                style_emb,
+                guidance_scale=2.0,
+                guidance_chunks=n_chunks - swap_at,
+            )
+            g = cache.transformer_cache.text_edit_guidance
+            assert g is not None and g.use_lora, (
+                "style window must run the pre-merged LoRA path "
+                "(is the TextEditLoRA attached?)"
+            )
+        num_frames = pipe.get_num_frames(ar_idx)
+        end = start + num_frames
+        assert end <= hdmap.shape[2], f"hdmap too short at chunk {ar_idx}"
+        chunk = pipe.generate(ar_idx, cache, hdmap=hdmap[:, :, start:end])
+        fs = cache.final_state
+        assert fs is not None and isinstance(fs.input, Tensor)
+        latents.append(fs.clean_latent.detach().to("cpu"))
+        hdmaps.append(fs.input.detach().to("cpu"))
+        pipe.finalize(ar_idx, cache)
+        chunks.append(chunk[0, 0].float().cpu())
+        start = end
+    del cache
+    torch.cuda.empty_cache()
+    return latents, hdmaps, torch.cat(chunks, dim=0)
+
+
+def main() -> None:
+    """Roll the base + per-(style, offset) branch corpus and save it."""
+    manifest = json.loads((SRC_DIR / "manifest.json").read_text())
+    uuids = [e["uuid"] for e in manifest]
+    if N_CLIPS:
+        uuids = uuids[:N_CLIPS]
+    n_chunks = int(os.environ.get("N_CHUNKS", str(manifest[0]["n_chunks"])))
+    total_frames = 5 + (n_chunks - 1) * 8
+    assert max(SWAP_OFFSETS) < n_chunks, (
+        f"swap offset {max(SWAP_OFFSETS)} outside the {n_chunks}-chunk rollout"
+    )
+
+    prompt_emb = torch.load(
+        BASE / "outputs/style_embeddings.pt", map_location="cpu", weights_only=False
+    )
+    assets = torch.load(
+        BASE / "outputs/style_clip_assets.pt", map_location="cpu", weights_only=False
+    )
+    for slug in STYLES:
+        assert slug in STYLE_PROMPTS and slug in prompt_emb, (
+            f"style {slug!r} missing from STYLE_PROMPTS / style_embeddings.pt"
+        )
+
+    pipe = build_pipeline(with_oneshot_encoders=False)
+    transformer = pipe.diffusion_model.transformer
+    edit_lora = TextEditLoRA(transformer.network, STYLE_LORA)
+    transformer.set_text_edit_lora(edit_lora)
+    print(f"deployed {edit_lora.describe()} from {STYLE_LORA}", flush=True)
+    print(
+        f"{len(uuids)} clips x {len(STYLES)} styles x {len(SWAP_OFFSETS)} "
+        f"offsets, {n_chunks} chunks each",
+        flush=True,
+    )
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    for uuid in uuids:
+        base_path = OUT_DIR / f"{uuid}_base.pt"
+        todo = [s for s in STYLES if not (OUT_DIR / f"{uuid}__{s}.pt").exists()]
+        if base_path.exists() and not todo:
+            print(f"skip {uuid} (exists)", flush=True)
+            continue
+
+        (hdmap_path,), _ = _sample_files(uuid)
+        hdmap = load_video_tensor(
+            hdmap_path,
+            pixel_height=DEFAULT_VIDEO_HEIGHT,
+            pixel_width=DEFAULT_VIDEO_WIDTH,
+            device=pipe.device,
+            dtype=torch.bfloat16,
+        )
+        assert hdmap.shape[0] >= total_frames, (
+            f"{uuid}: {hdmap.shape[0]} HDMap frames < {total_frames}"
+        )
+        hdmap = hdmap[:total_frames][None, None]
+        text_emb = prompt_emb[clip_key(uuid)]
+        image_emb = assets["image_embeddings"][uuid]
+
+        if base_path.exists():
+            base = torch.load(base_path, map_location="cpu", weights_only=False)[
+                "latents"
+            ]
+        else:
+            base, hdmaps, _ = rollout(
+                pipe,
+                text_emb=text_emb,
+                image_emb=image_emb,
+                style_emb=None,
+                hdmap=hdmap,
+                n_chunks=n_chunks,
+                swap_at=None,
+            )
+            torch.save(
+                {
+                    "uuid": uuid,
+                    "latents": base,
+                    "hdmaps": hdmaps,
+                    "n_chunks": n_chunks,
+                    "seed": SEED,
+                },
+                base_path,
+            )
+            print(f"rolled {uuid} base ({n_chunks} chunks)", flush=True)
+
+        for slug in todo:
+            branches: dict[int, list[Tensor]] = {}
+            for s in SWAP_OFFSETS:
+                latents, _, video = rollout(
+                    pipe,
+                    text_emb=text_emb,
+                    image_emb=image_emb,
+                    style_emb=prompt_emb[slug],
+                    hdmap=hdmap,
+                    n_chunks=n_chunks,
+                    swap_at=s,
+                )
+                # The swap draws no noise -> branches are bit-equal to the
+                # base rollout before their swap (verified host property);
+                # only the post-swap chunks are stored.
+                assert torch.equal(latents[0], base[0]) and torch.equal(
+                    latents[s - 1], base[s - 1]
+                ), f"{uuid}__{slug} s={s}: pre-swap chunks diverged from base"
+                branches[s] = latents[s:]
+                if SAVE_MP4:
+                    write_video_tensor(
+                        video,
+                        OUT_DIR / f"{uuid[:8]}__{slug}_s{s:02d}.mp4",
+                        fps=30,
+                        layout="tchw",
+                    )
+                print(f"rolled {uuid}__{slug} s={s}", flush=True)
+            torch.save(
+                {
+                    "uuid": uuid,
+                    "slug": slug,
+                    "branches": branches,
+                    "swap_offsets": list(SWAP_OFFSETS),
+                    "n_chunks": n_chunks,
+                    "seed": SEED,
+                    "style_lora": str(STYLE_LORA),
+                },
+                OUT_DIR / f"{uuid}__{slug}.pt",
+            )
+
+    print(
+        f"STYLE-DRIFT-PAIRS-DONE | {len(uuids)} clips x {len(STYLES)} styles "
+        f"x {len(SWAP_OFFSETS)} offsets -> {OUT_DIR}/",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/integrations/omnidreams/edit_sft/train_style_corrector.py b/integrations/omnidreams/edit_sft/train_style_corrector.py
new file mode 100644
index 000000000..c880fcf6d
--- /dev/null
+++ b/integrations/omnidreams/edit_sft/train_style_corrector.py
@@ -0,0 +1,656 @@
+# 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.
+
+"""Style-drift corrector trainer (Issue 1 fix ladder item 1).
+
+The Clean Forcing recipe (``drift_correction/train_v2.py``) retargeted at
+the style-skin long-hold blur: a SEPARATE rank-16 LoRA on the self-attn
+q/k/v/output projections, trained WITH the style LoRA active, that maps
+late-window styled predictions back toward the early-window styled
+manifold. Like the shipped drift corrector, it corrects the *flow
+prediction under the live (drifted) KV history* — every denoise-step
+``predict_flow`` and the KV-commit context forward — not the latent
+post hoc, so the checkpoint deploys through the exact same hook
+(``omnidreams/_drift_corrector.py``: same 4 targets, same rank, same
+``{"lora": {2i: A_i, 2i+1: B_i}}`` load order).
+
+Counterfactual pairs come from ``gen_style_drift_pairs.py`` branches: for
+a probe chunk ``k``, the DRIFTED branch swapped ``DRIFT_MIN..DRIFT_MAX``
+chunks earlier (its window has compounded into blur) and the REFERENCE
+branch swapped ``1..REF_MAX`` chunks earlier (the model's own clean styled
+output at the SAME absolute index, same HDMap, same seed — a reference
+that never round-trips through JoyAI, sidestepping the VAE-blur bias).
+Per pooled sample, at a matched noisy state ``z_t`` built from the drifted
+chunk::
+
+    v_clean = style-LoRA'd model | reference history, corrector 0
+    v_base  = style-LoRA'd model | drifted history,   corrector 0
+    v_corr  = style-LoRA'd model | drifted history,   corrector 1  (grad)
+    L_dag   = ||v_corr - v_clean||^2 / ||v_clean - v_base||^2
+
+plus the ``train_v2`` drift-contraction term (weight :data:`CW_LOSS`): the
+corrected chunk-``k`` prediction is committed into chunk ``k+1``'s KV with
+grad (``record_kv`` / ``inject_kv``) and chunk ``k+1``'s gap to its clean
+styled teacher is penalized. No-op episodes (:data:`NOOP_PROB`) pin the
+corrector to identity on unstyled and on early-window styled history, so
+the always-on deploy hook is safe outside and at the start of style
+windows.
+
+Style-LoRA semantics mirror deploy (``_edit_lora`` / ``TextEditLoRA``):
+text KV is always rebuilt at base weights; in-window forwards and commits
+run at style scale 1; pre-swap replay chunks run at base. The corrector
+trains at scale 1 and deploys at ``alpha*(t) * gain`` (the hook's gate;
+sweep ``drift_corrector_gain`` at eval, as for the photoreal corrector).
+
+Deploy note: compose with the pre-merged style LoRA via the corrector's
+UNFUSED path (``DRIFT_CORRECTOR_UNFUSED=1``), attaching the
+``TextEditLoRA`` BEFORE ``apply_drift_corrector``. The default pre-merged
+corrector path and ``TextEditLoRA.set_active`` both assume exclusive
+ownership of the projection ``weight.data`` and corrupt each other.
+
+Run from the flashdreams repo root (resumable: re-run the same command)::
+
+    STEPS=1000 .venv/bin/python \
+        integrations/omnidreams/edit_sft/train_style_corrector.py
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+# Must land before the first CUDA allocation (co-tenant VRAM share).
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from _host import CONTEXT_NOISE_SEED, build_pipeline, reset_history
+from _lora import (
+    DEFAULT_TARGETS,
+    LoRALinear,
+    apply_lora,
+    load_lora,
+    lora_parameters,
+    unwrap_compiled,
+)
+from _train_attn import (
+    functional_attention,
+    inject_kv,
+    patch_functional_attention,
+    record_kv,
+)
+from omnidreams._edit_lora import _LORA_TARGETS as LORA_TARGETS
+from omnidreams.runner import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH
+from style_prompts import clip_key
+from torch import Tensor
+
+## Training configuration
+
+BASE = Path("integrations/omnidreams/edit_sft")
+OUT_DIR = BASE / "outputs"
+PAIRS_DIR = Path(os.environ.get("PAIRS_DIR", str(OUT_DIR / "style_drift_pairs")))
+
+STYLE_LORA = Path(os.environ.get("STYLE_LORA", str(OUT_DIR / "lora_style_step1600.pt")))
+"""Style-skin checkpoint the corrector trains under (must be the one whose
+branches ``gen_style_drift_pairs.py`` rolled)."""
+
+STYLE_RANK = int(os.environ.get("STYLE_RANK", "64"))
+"""Rank of the style checkpoint (the eval RANK-must-match gotcha)."""
+
+CKPT = Path(os.environ.get("CKPT", str(OUT_DIR / "lora_style_corrector.pt")))
+"""Output checkpoint (corrector LoRA + optimizer + step); loads through
+``omnidreams/_drift_corrector.py`` unchanged."""
+
+INIT = os.environ.get("INIT", "")
+"""Optional warm-start corrector checkpoint (same format; the photoreal
+``drift_correction/outputs/lora_v2_v3_valpeak.pt`` is shape-compatible).
+Empty = scratch."""
+
+STEPS = int(os.environ.get("STEPS", "1000"))
+LR = float(os.environ.get("LR", "2e-4"))
+WARMUP = 40
+GRAD_CLIP = 1.0
+EVAL_EVERY = 100
+SAVE_EVERY = 200
+
+RANK = int(os.environ.get("RANK", "16"))
+"""Corrector rank. The deploy hook hardcodes 16 — other values need a
+matching ``_drift_corrector._LORA_RANK`` change."""
+
+SNAP_EVERY = int(os.environ.get("SNAP_EVERY", "0"))
+"""When > 0: step-tagged snapshots every SNAP_EVERY steps plus a running
+val-peak snapshot (``<ckpt>_valpeak.pt``), the ``train_v2`` convention."""
+
+CW_LOSS = float(os.environ.get("CW_LOSS", "0.5"))
+"""Drift-contraction weight (the ``train_v2`` paper-final value)."""
+
+DRIFT_WEIGHT = os.environ.get("DRIFT_WEIGHT", "0") == "1"
+"""Per-token loss weighting by drift magnitude. The style drift
+concentrates in background tokens that the uniform DAG numerator
+under-weights (roadside texture stays soft after correction); with the
+knob on, each token's squared error in the DAG and contraction numerators
+is multiplied by that token's clean-vs-base gap norm (mean-1 per frame,
+clamped to [0.5, 4.0], detached). 0 = exact unweighted behavior."""
+
+TOKENS_PER_FRAME = (DEFAULT_VIDEO_HEIGHT // 16) * (DEFAULT_VIDEO_WIDTH // 16)
+"""Patch-token grid per latent frame (VAE /8 x patchify /2): 44*80."""
+
+NOOP_PROB = float(os.environ.get("NOOP_PROB", "0.15"))
+"""Probability of a no-op episode: corrector at scale 1 must match the
+corrector-free prediction on unstyled history (50%) or on early-window
+styled history (50%) — keeps the always-on deploy gate safe off the
+drifted manifold."""
+
+NOOP_WEIGHT = float(os.environ.get("NOOP_WEIGHT", "1.0"))
+
+DRIFT_MIN = int(os.environ.get("DRIFT_MIN", "8"))
+DRIFT_MAX = int(os.environ.get("DRIFT_MAX", "20"))
+"""Style-hold depth (chunks since swap, 1-based) of the drifted branch:
++8..+20 spans "structure still holds" through "washed out" on the v3
+skin."""
+
+REF_MAX = int(os.environ.get("REF_MAX", "4"))
+"""Max hold depth of the reference branch (+1..+4 = the clean manifold)."""
+
+REPLAY_CHUNKS = int(os.environ.get("REPLAY_CHUNKS", "3"))
+"""History chunks replayed before a probe. 3 chunks = the full
+``window_size_t=6`` / ``len_t=2`` KV window, so the visible state equals a
+full-history replay (the ``train_style_sft`` mid-stream-start setting)."""
+
+REL_V_EXCLUDE = 0.8
+"""Drop cells whose clean-vs-base gap exceeds this fraction of ``|v_base|``
+(unpredictable content difference, the ``train_v2`` convention)."""
+
+ALPHA_STAR = tuple(
+    float(x) for x in os.environ.get("ALPHA_STAR", "0.96,0.667").split(",")
+)
+"""Per-timestep sampling weights (the deploy gate's relative profile)."""
+
+N_VAL_CLIPS = int(os.environ.get("N_VAL_CLIPS", "1"))
+"""Clips (last of the sources manifest with pairs) held out for val."""
+
+STYLES = frozenset(s for s in os.environ.get("STYLES", "").split(",") if s)
+"""Optional slug filter; empty = every style found in the pairs dir."""
+
+SEED = int(os.environ.get("SEED", "0"))
+
+
+def drift_weights(gap: Tensor) -> Tensor:
+    """Per-token weights from a clean-vs-base flow gap ``[..., T*HW, C]``.
+
+    Weight = the token's channel-norm of ``gap``, normalized to mean 1
+    over each frame's tokens, clamped to [0.5, 4.0], detached — so tokens
+    the drift actually moved dominate the numerator without changing the
+    loss scale or blowing up on near-uniform frames.
+    """
+    w = gap.detach().norm(dim=-1, keepdim=True)
+    assert w.shape[-2] % TOKENS_PER_FRAME == 0, tuple(gap.shape)
+    lead = w.shape[:-2]
+    w = w.reshape(*lead, -1, TOKENS_PER_FRAME, 1)
+    w = w / (w.mean(dim=-2, keepdim=True) + 1e-8)
+    return w.reshape(*lead, -1, 1).clamp(0.5, 4.0)
+
+
+def load_pairs() -> list[dict]:
+    """Load the branch corpus and precompute the probe cells per pair.
+
+    A cell ``(k, a, b)`` probes absolute chunk ``k`` with the drifted
+    branch swapped at ``a`` (depth ``k - a + 1`` in
+    ``[DRIFT_MIN, DRIFT_MAX]``) against the reference branch swapped at
+    ``b`` (depth in ``[1, REF_MAX]``); ``k + 1`` must exist for the
+    contraction successor.
+    """
+    bases: dict[str, dict] = {}
+    pairs: list[dict] = []
+    for path in sorted(PAIRS_DIR.glob("*__*.pt")):
+        uuid, slug = path.stem.split("__", 1)
+        if STYLES and slug not in STYLES:
+            continue
+        if uuid not in bases:
+            base_path = PAIRS_DIR / f"{uuid}_base.pt"
+            assert base_path.exists(), f"missing {base_path} for {path.name}"
+            bases[uuid] = torch.load(base_path, map_location="cpu", weights_only=False)
+        base = bases[uuid]
+        n = int(base["n_chunks"])
+        assert len(base["latents"]) == n and len(base["hdmaps"]) == n
+        d = torch.load(path, map_location="cpu", weights_only=False)
+        full = {
+            int(s): base["latents"][: int(s)] + lat for s, lat in d["branches"].items()
+        }
+        offsets = sorted(full)
+        cells: list[tuple[int, int, int]] = []
+        for k in range(1 + REPLAY_CHUNKS, n - 1):
+            refs = [b for b in offsets if 1 <= k - b + 1 <= REF_MAX]
+            drifts = [a for a in offsets if DRIFT_MIN <= k - a + 1 <= DRIFT_MAX]
+            cells += [(k, a, b) for a in drifts for b in refs]
+        if not cells:
+            print(f"skip {path.name}: no probe cells in range", flush=True)
+            continue
+        pairs.append(
+            {
+                "uuid": uuid,
+                "slug": slug,
+                "n_chunks": n,
+                "base": base["latents"],
+                "hdmaps": base["hdmaps"],
+                "full": full,
+                "offsets": offsets,
+                "cells": cells,
+            }
+        )
+    assert pairs, f"no styled branch files under {PAIRS_DIR}"
+    return pairs
+
+
+def main() -> None:
+    """Run the counterfactual style-corrector training loop."""
+    rng = np.random.default_rng(SEED)
+    pairs = load_pairs()
+    manifest = json.loads((OUT_DIR / "sources" / "manifest.json").read_text())
+    uuids_in_order = [
+        e["uuid"] for e in manifest if any(p["uuid"] == e["uuid"] for p in pairs)
+    ]
+    assert len(uuids_in_order) > N_VAL_CLIPS, (
+        f"{len(uuids_in_order)} clips can't spare {N_VAL_CLIPS} for val"
+    )
+    val_uuids = set(uuids_in_order[-N_VAL_CLIPS:])
+    train_ids = [i for i, p in enumerate(pairs) if p["uuid"] not in val_uuids]
+    val_ids = [i for i, p in enumerate(pairs) if p["uuid"] in val_uuids]
+
+    prompt_emb = torch.load(
+        OUT_DIR / "style_embeddings.pt", map_location="cpu", weights_only=False
+    )
+    assets = torch.load(
+        OUT_DIR / "style_clip_assets.pt", map_location="cpu", weights_only=False
+    )
+
+    pipe = build_pipeline(with_oneshot_encoders=False)
+    assert pipe.V_group is None, "single-GPU trainer; run without CP"
+    device = pipe.device
+    dtype = torch.bfloat16
+    dm = pipe.diffusion_model
+    transformer = dm.transformer
+    scheduler = dm.scheduler
+    timesteps = scheduler.denoising_step_list  # [1000, 450] on the chunk2 host
+    sigmas = scheduler.denoising_sigmas
+    n_steps = int(timesteps.shape[0])
+    t_probs = np.array(ALPHA_STAR[:n_steps]) / sum(ALPHA_STAR[:n_steps])
+    ctx_t = torch.tensor(float(dm.config.context_noise), device=device, dtype=dtype)
+    ctx_idx = torch.argmin((scheduler._full_timesteps - ctx_t.float()).abs())
+    sigma_ctx = scheduler._full_sigmas[ctx_idx].to(dtype)
+
+    # Style LoRA (frozen, deploy-window semantics) + corrector on top. The
+    # corrector wraps the inner base linears of the 4 self-attn targets in
+    # the SAME named_modules walk order as the deploy hook's _apply_lora,
+    # so the saved index order loads through _drift_corrector unchanged.
+    network = unwrap_compiled(transformer.network)
+    network.requires_grad_(False)
+    style_names = apply_lora(network, rank=STYLE_RANK, targets=LORA_TARGETS)
+    load_lora(network, STYLE_LORA)
+    for p in lora_parameters(network):  # style-only at this point
+        p.requires_grad_(False)
+    style_mods = [network.get_submodule(n) for n in style_names]
+
+    corr_mods: list[LoRALinear] = []
+    for mname, module in list(network.named_modules()):
+        if not isinstance(module, LoRALinear):
+            continue
+        if not any(mname.endswith(t) for t in DEFAULT_TARGETS):
+            continue
+        inner = module.base
+        assert isinstance(inner, nn.Linear)
+        corr = LoRALinear(inner, rank=RANK).to(inner.weight.device)
+        corr.scale = 0.0
+        module.base = corr
+        corr_mods.append(corr)
+    assert 2 * len(corr_mods) == len(style_mods), (
+        f"{len(corr_mods)} corrector wraps vs {len(style_mods)} style wraps"
+    )
+    corr_params = [p for m in corr_mods for p in (m.A.weight, m.B.weight)]
+    if RANK != 16:
+        print(
+            f"WARNING: RANK={RANK} will not load through the deploy hook "
+            "(_drift_corrector._LORA_RANK is 16).",
+            flush=True,
+        )
+    print(
+        f"style r{STYLE_RANK} on {len(style_mods)} projections (frozen) | "
+        f"corrector r{RANK} on {len(corr_mods)} projections, "
+        f"{sum(p.numel() for p in corr_params) / 1e6:.2f}M params | "
+        f"drift-weight {'on' if DRIFT_WEIGHT else 'off'} | "
+        f"{len(pairs)} pairs ({len(val_ids)} val), cells/pair "
+        f"{[len(p['cells']) for p in pairs]}",
+        flush=True,
+    )
+    patch_functional_attention()
+    opt = torch.optim.AdamW(corr_params, lr=LR)
+
+    def set_scales(style: float, corr: float) -> None:
+        """Independent runtime gains for the two LoRA stacks."""
+        for m in style_mods:
+            m.scale = style
+        for m in corr_mods:
+            m.scale = corr
+
+    # One long-lived cache; probes never replay chunk 0 with content that
+    # matters (start >= 1), so the image/mask fields can come from any clip.
+    set_scales(0.0, 0.0)
+    first_uuid = pairs[0]["uuid"]
+    cache = pipe.initialize_cache_from_embeddings(
+        text_embeddings=prompt_emb[clip_key(first_uuid)],
+        image_embeddings=assets["image_embeddings"][first_uuid],
+    )
+    tc = cache.transformer_cache
+    _cur_text = [clip_key(first_uuid)]
+
+    def swap_text(key: str) -> None:
+        """Rebuild the cross-attn text KV at BASE weights (deploy truth)."""
+        if _cur_text[0] == key:
+            return
+        set_scales(0.0, 0.0)
+        pipe.replace_text_from_embeddings(cache, prompt_emb[key])
+        _cur_text[0] = key
+
+    start_step = 0
+    if CKPT.exists():
+        state = torch.load(CKPT, map_location="cpu", weights_only=False)
+        for i, p in enumerate(corr_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:
+        state = torch.load(INIT, map_location="cpu", weights_only=False)
+        for i, p in enumerate(corr_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(corr_params)},
+                "opt": opt.state_dict(),
+                "step": step,
+                "rng_seed": SEED + step,
+            },
+            tmp,
+        )
+        tmp.replace(path)
+
+    def replay(
+        latents: list[Tensor],
+        hdmaps: list[Tensor],
+        upto: int,
+        swap_at: int | None,
+        slug: str | None,
+        uuid: str,
+        corr: float,
+    ) -> None:
+        """Rebuild the KV window for chunks ``max(0, upto - 3) .. upto - 1``.
+
+        Pre-swap chunks commit at base weights under the clip's own prompt;
+        chunks at/after ``swap_at`` commit at style scale 1 under the style
+        prompt (the deploy window semantics), with the corrector at ``corr``
+        (1 for the corrected branch — deploy commits corrected — else 0).
+        """
+        start = max(0, upto - REPLAY_CHUNKS)
+        reset_history(tc)
+        for bc in tc.network_cache.block_caches:
+            bc.self_attn._prev_chunk_idx = start - 1
+        swapped = swap_at is not None and swap_at <= start
+        assert slug is not None or swap_at is None
+        swap_text(slug if (swapped and slug is not None) else clip_key(uuid))
+        with torch.no_grad():
+            for j in range(start, upto):
+                if swap_at is not None and j == swap_at and not swapped:
+                    assert slug is not None
+                    swap_text(slug)
+                    swapped = True
+                in_win = swap_at is not None and j >= swap_at
+                set_scales(1.0 if in_win else 0.0, corr if in_win else 0.0)
+                g = torch.Generator(device=device).manual_seed(CONTEXT_NOISE_SEED + 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:
+        """One functional-attention (non-writing) forward -> fp32 flow."""
+        t = timesteps[t_idx].to(device=device, dtype=dtype)
+        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)
+        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(pairs))}
+
+    def styled_losses(
+        c: int, grad: bool, rng_: np.random.Generator
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor] | None:
+        """One pooled cell -> (dag, contraction, r_target sq, unweighted dag).
+
+        Under :data:`DRIFT_WEIGHT` the dag/contraction numerators are
+        drift-weighted per token; the unweighted dag ratio is returned
+        alongside (equal to dag when off) so val stays comparable across
+        weighted and unweighted runs. Returns ``None`` for excluded
+        (unpredictable) cells; callers redraw.
+        """
+        p = pairs[c]
+        uuid, slug = p["uuid"], p["slug"]
+        k, a, b = p["cells"][int(rng_.integers(len(p["cells"])))]
+        t_idx = int(rng_.choice(n_steps, p=t_probs))
+        t2_idx = int(rng_.choice(n_steps, p=t_probs))
+        end = k + 2
+        lat_d = [x.to(device, dtype) for x in p["full"][a][:end]]
+        lat_r = [x.to(device, dtype) for x in p["full"][b][:end]]
+        hd = [x.to(device, dtype) for x in p["hdmaps"][:end]]
+        z_t = make_zt(lat_d[k], t_idx, rng_)
+        z_t2 = make_zt(lat_d[k + 1], t2_idx, rng_)
+
+        with torch.no_grad():
+            # Teacher: the clean styled manifold at k and (for the
+            # contraction target) at k+1 — reference-branch history.
+            replay(lat_r, hd, k, b, slug, uuid, corr=0.0)
+            set_scales(1.0, 0.0)
+            tc.start(k)
+            v_clean = predict_v(z_t, t_idx, hd[k])
+            tc.finalize(k)
+            replay(lat_r, hd, k + 1, b, slug, uuid, corr=0.0)
+            set_scales(1.0, 0.0)
+            tc.start(k + 1)
+            v_clean2 = predict_v(z_t2, t2_idx, hd[k + 1])
+            tc.finalize(k + 1)
+            # Base: the drifted branch without corrector.
+            replay(lat_d, hd, k + 1, a, slug, uuid, corr=0.0)
+            set_scales(1.0, 0.0)
+            tc.start(k + 1)
+            v_base2 = predict_v(z_t2, t2_idx, hd[k + 1])
+            tc.finalize(k + 1)
+            replay(lat_d, hd, k, a, slug, uuid, corr=0.0)
+            set_scales(1.0, 0.0)
+            tc.start(k)
+            v_base = predict_v(z_t, t_idx, hd[k])
+            tc.finalize(k)
+        r_sq = (v_clean - v_base).square().sum()
+        r2_sq = (v_clean2 - v_base2).square().sum()
+        w = drift_weights(v_clean - v_base) if DRIFT_WEIGHT else None
+        w2 = drift_weights(v_clean2 - v_base2) if DRIFT_WEIGHT else None
+
+        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
+
+        # Corrected branch: drifted history committed WITH the corrector
+        # (deploy commits corrected), probe + commit + successor probe.
+        replay(lat_d, hd, k, a, slug, uuid, corr=1.0)
+        set_scales(1.0, 1.0)
+        ctx = torch.enable_grad() if grad else torch.no_grad()
+        tc.start(k)
+        with ctx:
+            v_corr = predict_v(z_t, t_idx, hd[k])
+            err = (v_corr - v_clean).square()
+            dag_raw = err.sum() / (r_sq + 1e-8)
+            dag = (err * w).sum() / (r_sq + 1e-8) if w is not None else dag_raw
+            # Commit chunk k's corrected prediction with grad: recorded
+            # functional KV + a numerically identical no-grad buffer twin
+            # (the train_v2 contraction machinery).
+            sig = sigmas[t_idx].to(dtype)
+            x0_corr = (z_t.float() - float(sig) * v_corr).to(dtype)
+            g = torch.Generator(device=device).manual_seed(CONTEXT_NOISE_SEED + k)
+            eps = torch.randn(x0_corr.shape, device=device, dtype=dtype, generator=g)
+            noisy_corr = (1 - sigma_ctx) * x0_corr + sigma_ctx * eps
+            recorded: list = []
+            with record_kv(recorded), functional_attention():
+                transformer.predict_flow(
+                    noisy_latent=noisy_corr, timestep=ctx_t, cache=tc, input=hd[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=hd[k]
+            )
+        tc.finalize(k)
+        tc.start(k + 1)
+        with ctx:
+            with inject_kv(recorded):
+                v_con = predict_v(z_t2, t2_idx, hd[k + 1])
+            err2 = (v_con - v_clean2).square()
+            con = (err2 * w2 if w2 is not None else err2).sum() / (r2_sq + 1e-8)
+        tc.finalize(k + 1)
+        return dag, con, r_sq, dag_raw
+
+    def noop_loss(c: int, grad: bool, rng_: np.random.Generator) -> Tensor:
+        """Corrector-at-1 vs corrector-at-0 gap off the drifted manifold."""
+        p = pairs[c]
+        uuid, slug, n = p["uuid"], p["slug"], p["n_chunks"]
+        styled = bool(rng_.random() < 0.5)
+        if styled:
+            b = int(p["offsets"][int(rng_.integers(len(p["offsets"])))])
+            depth = int(rng_.integers(1, REF_MAX + 1))
+            k = min(b + depth - 1, n - 1)
+            latents, swap_at, key = p["full"][b], b, slug
+        else:
+            k = int(rng_.integers(1 + REPLAY_CHUNKS, n))
+            latents, swap_at, key = p["base"], None, None
+        t_idx = int(rng_.choice(n_steps, p=t_probs))
+        lat = [x.to(device, dtype) for x in latents[: k + 1]]
+        hd = [x.to(device, dtype) for x in p["hdmaps"][: k + 1]]
+        z_t = make_zt(lat[k], t_idx, rng_)
+        sty = 1.0 if styled else 0.0
+        replay(lat, hd, k, swap_at, key, uuid, corr=0.0)
+        tc.start(k)
+        with torch.no_grad():
+            set_scales(sty, 0.0)
+            v0 = predict_v(z_t, t_idx, hd[k])
+        ctx = torch.enable_grad() if grad else torch.no_grad()
+        with ctx:
+            set_scales(sty, 1.0)
+            v1 = predict_v(z_t, t_idx, hd[k])
+            loss = (v1 - v0).square().sum() / (v0.square().sum() + 1e-8)
+        tc.finalize(k)
+        return loss
+
+    def draw_styled(
+        ids: list[int], grad: bool, rng_: np.random.Generator
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        """Redraw until a non-excluded cell is sampled."""
+        while True:
+            out = styled_losses(int(rng_.choice(ids)), grad, rng_)
+            if out is not None:
+                return out
+
+    @torch.no_grad()
+    def val_metrics(n: int = 12) -> tuple[float, float, float, float]:
+        vrng = np.random.default_rng(1234)
+        s_dag = s_con = s_noop = s_raw = 0.0
+        for _ in range(n):
+            dag, con, _, raw = draw_styled(val_ids, False, vrng)
+            s_dag += dag.item()
+            s_con += con.item()
+            s_raw += raw.item()
+        n_noop = max(1, n // 2)
+        for _ in range(n_noop):
+            s_noop += noop_loss(int(vrng.choice(val_ids)), False, vrng).item()
+        return 1 - s_dag / n, 1 - s_con / n, s_noop / n_noop, 1 - s_raw / 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()
+        if rng.random() < NOOP_PROB:
+            noop = noop_loss(int(rng.choice(train_ids)), True, rng)
+            loss = NOOP_WEIGHT * noop
+            terms = f"noop {noop.item():.4f}"
+        else:
+            dag, con, r_sq, _ = draw_styled(train_ids, True, rng)
+            loss = dag + CW_LOSS * con
+            terms = f"dag {dag.item():.4f} con {con.item():.4f} |r|^2 {r_sq.item():.1f}"
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(corr_params, GRAD_CLIP)
+        opt.step()
+        if step % EVAL_EVERY == 0 or step == 1:
+            vd, vc, vn, vraw = val_metrics()
+            excl_str = " ".join(
+                f"c{c}:{v['excluded']}/{v['tested']}" for c, v in excl.items()
+            )
+            raw_str = f" (unw {vraw:+.3f})" if DRIFT_WEIGHT else ""
+            print(
+                f"step {step:5d} | {terms} | val dag-R^2 {vd:+.3f}{raw_str} "
+                f"con-R^2 {vc:+.3f} noop {vn:.4f} | 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, vn, vraw = val_metrics(24)
+    raw_str = f" (unw {vraw:+.3f})" if DRIFT_WEIGHT else ""
+    print(
+        f"TRAIN-STYLE-CORRECTOR-DONE | final val dag-R^2 {vd:+.3f}{raw_str} "
+        f"con-R^2 {vc:+.3f} noop {vn:.4f} | saved {CKPT}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()