Omnidreams edit-SFT style skins: offline pair generation, VLM filter, mid-stream LoRA trainer

[wenqingw-nv]

Live prompt edits (#431) cover in-distribution weather and lighting; full game-style restyles (cel-shaded cartoon, arcade, comic ink) are outside the AV training distribution and no guidance scale reaches them. This PR adds the supervised path: an offline instruction-driven editor authors (original → restyled) pairs from omnidreams rollouts, a VLM filter gates them for structural fidelity, and an edit-timestamped LoRA trainer teaches the realtime model to apply the style itself mid-stream — one multi-style LoRA, selected by the prompt, deployed at zero added latency through the live-edit LoRA hook from #431. Everything lives under integrations/omnidreams/edit_sft/; nothing touches the serving path.

Contents

• generate_sources.py — RNG-matched source rollouts (videos + latents) from the sample-data scenes.
• filter_pairs.py — Cosmos-Reason1 VLM judge. MODE=style scores edit strength/persistence plus a style-agnostic road-layout criterion at an early and a late frame (describe-then-score phrasing: a bare cross-style JSON score collapses to a uniform 0). Heavy styles are demoted to early-window-only by policy — the judge over-credits late-frame layout on strong stylization (verified false pass), and offline editors drift layout over long clips (streaming error accumulation: the road is progressively replaced).
• style_prompts.py, precompute_style.py — style prompt bank; one-shot text embeddings + streaming latent encode of sources and styled targets on the AR schedule.
• train_style_sft.py — teacher-forced trainer: history replayed from source latents through the KV-commit path, flow targets switch to the styled latents at a random swap chunk while the text KV switches to the style prompt; pre-swap chunks and 10% no-op windows are supervised to stay unchanged. LoRA r64 on the attention projections; checkpoints load through TextEditLoRA (Omnidreams: live mid-stream editing — prompt hot-swap, edit guidance, actor spawn/despawn #431) unchanged.
• _host.py, _lora.py, _train_attn.py — training helpers vendored from the Clean Forcing infra (omnidreams: Clean Forcing drift corrector (training recipe + gated deploy) #398); consolidate when that lands.

Data + filter (measured, GB300)

• 124 restyled clips from 10 source scenes × 12 style instructions (JoyAI-Video-Edit batch, ~13 s per 7.4 s clip — an offline data factory at ~18 fps non-causal; it cannot serve realtime, which is what the LoRA is for).
• Filter verdicts: 35/124 full-clip passes (arcade-racer, comic-ink, low-poly, toy-world lead); heavy styles (cartoon-cel, anime, pixel-art) pass early-window only. Layout-preserving instruction suffixes helped some styles (anime 0/10 → 7/10) and hurt others (arcade-racer 5/10 → 0/10) — per-clip verdicts, not per-style assumptions.
• 20-step trainer smoke passes end-to-end (loss 0.58 → 0.16, deploy-format checkpoint).

Results (trained + evaluated)

• v3 LoRA (1600 steps, 30% style-maintenance episodes): pre-swap bit-exact (0.000 divergence with the window closed); a mid-drive "arcade mode" swap restyles the world at full 30 fps (per-chunk divergence ~55 vs ~15 for the plain prompt swap) with road layout preserved. Long holds drift: structure survives to +7–12 chunks, washes out by +18 — the style push compounds through the KV history (re-pulsing windows does not help; the drift lives in committed history).
• Style-drift corrector (clean-forcing recipe on the styled manifold): rank-16 corrector trained against the model's own +1..+4-chunk styled outputs as the clean reference (val dag-R² +0.42). At gain 0.25 a 20-chunk hold keeps road, lane marks, and vehicles intact where the uncorrected hold collapses; residual: background/roadside texture still softens (near objects hold). Deploys through the drift-corrector hook (omnidreams: Clean Forcing drift corrector (training recipe + gated deploy) #398) unfused, attached after TextEditLoRA.
• Checkpoints (fork releases, per the guidance-LoRA precedent): style-skin-lora-v3 and style-drift-corrector-v1. Corrector data-gen + trainer scripts import the live-edit deploy hook (Omnidreams: live mid-stream editing — prompt hot-swap, edit guidance, actor spawn/despawn #431) and land here once that merges.

Follow-ups: style-specific gate calibration (current alpha profile is photoreal-measured); background-weighted corrector retrain; a second corrector slot so the photoreal drift corrector can co-deploy.

Not included (follow-ups)

Trained checkpoint (attached on eval pass); spawn/object-materialization SFT (separate track); object add/remove pairs (fail the persistence gate in the offline editor).

Video (original | restyled training data)

Staged sample pairs (data-factory outputs, not yet the LoRA): arcade-racer and cartoon-cel side-by-sides.

sbs_style_arcade_racer.mp4

---

Update (bbe4d12d)

• Fix (review finding): the pre-swap text cache was built with the LoRA already active, so "base" rollout history was subtly styled. The cache is now initialized at LoRA scale 0 and the adapter enabled only at the swap.
• v4 self-consistency episodes (V4_MAINT/V4_EARLY): maintenance targets drawn from the run's own pre-drift styled window. Evaluated head-to-head vs v3: a wash — kept behind the knob for the record.
• gen_teacher_styled_targets.py (v5 recipe): re-rolls JoyAI-restyled clips through the 35-step bidirectional teacher (SDEdit at moderate σ) to produce sharp styled targets; precompute_style.py reads either corpus. A v5 LoRA trained on these targets is being evaluated — results will be posted here.

v5 result: the teacher-regen recipe works. On the 20-chunk held-style eval, v5 shows a flat post-swap divergence profile (54–61) where v3 climbs monotonically (57→69), and frame inspection at +10/+14/+19 chunks shows v5 retaining road, curb and solid vehicles at depths where v3 has fully melted — roughly 2× usable style-hold depth. Remaining weakness is slow texture smoothing deep in the hold; a corrector re-paired against v5 is being evaluated as the combined stack.

Final v5 stack: corrector re-paired against v5 (val dag-R² +0.368) and gain-swept. Final config: v5 LoRA + v5 corrector at gain 0.15 with the measured v5 gate profile + unsharp post — retains a populated styled world (houses, sidewalks, parked cars) at +19 chunks where the v3 stack keeps only the road corridor, with deep-window HF above the v3 stack. Known residual: faint corrector-induced sky banding (gain-proportional; a training-side fix would need sky-region loss masking). Checkpoints: https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/style-skin-v5-stack

Before vs after — final v5 serving stack

Side-by-side comparisons: left = base model (no edit), right = arcade skin via the final stack (v5 LoRA + re-paired corrector at gain 0.15, measured gate profile, unsharp post). Prompt-swap at chunk 8, 28-chunk rollouts, same seed both arms.

Highway scene (cleanest hold; lane markings and vehicles stay sharp end-to-end):

sbs_before_vs_v5cas_highway.mp4

Residential scene (style held to the end; houses/vehicles stay populated deep into the hold):

sbs_before_vs_v5cas.mp4

v6 — four skins, one adapter: retrained on teacher-regenerated corpora for arcade_racer + comic_ink + cyberpunk_neon + pixel_art (layout-preserving re-audition rescued the two styles the original filter rejected; the VLM judge's day→night miscalibration is documented in the filter report with eyes-on + edge-structure evidence). Cyberpunk is the headline skin (full 28-chunk neon-night hold); pixel_art holds full-range but is honest-partial on style identity. No arcade regression vs v5. Checkpoint: https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/style-skin-v6-multiskin — running live with mid-run key-switching inside the Crazy Robotaxi game in #494.

robotaxi_liveedit_demo3_cyberpunk.mp4

robotaxi_liveedit_demo3_annotated.mp4

Serving/deploy hook (new commits): this PR now also carries the canonical drift-corrector deploy hook. DRIFT_CORRECTOR_MODE=fused keeps the per-alpha pre-merged weight sets and CPU-side call-index gate but swaps weights by batched in-place copy_ into the original parameter storages, so the exact α*(t) profile survives under compile_network + use_cuda_graph — the unfused path had to disable both. DriftCorrectorDispatch adds per-state corrector switching (checkpoint + gate profile + gain per state), composing style-LoRA deltas via the new TextEditLoRA.release_targets() seam so one graph-safe writer carries both. Benchmarked at 38.7 fps vs 20.7 unfused; running live in the Robotaxi integration (#494). CPU equivalence/pointer-stability tests plus a GPU validation script (parity, graph capture, latency bench) included.

Corrector training stack + fp32 fix (new commits): the data-gen (gen_style_drift_pairs.py), trainer (train_style_corrector.py), and gate calibration (gate_style.py) now live under edit_sft/ — the full Clean-Forcing loop that produced the deployed corrector checkpoints. Also fixes a real fp32 aliasing bug caught during porting: base.to(torch.float32) is a no-copy alias on fp32 networks, so the in-place pre-merge corrupted the cached base set (bf16 deploys unaffected); now copy=True with a regression test.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

The PR adds an offline style-edit dataset pipeline, filtering and training utilities, and graph-safe LoRA/drift-corrector deployment support.

• Generates and filters source-to-restyled training pairs.
• Trains edit-timestamped style LoRAs and drift correctors.
• Adds fused per-state corrector dispatch with deployment validation and tests.

Confidence Score: 4/5

The PR is not yet safe to merge because pre-swap supervision still runs with the LoRA active instead of matching the inactive deployment path.

With PRE_CHUNKS enabled, the trainer restores LoRA scale one before processing j < k; only the subsequent KV commit is temporarily switched back to base weights, leaving the prediction and context losses on the wrong execution path.

Files Needing Attention: integrations/omnidreams/edit_sft/train_style_sft.py

Important Files Changed

Filename	Overview
integrations/omnidreams/edit_sft/train_style_sft.py	Adds edit-timestamped LoRA training and maintenance episodes; the previously reported pre-swap prediction-path mismatch remains.
integrations/omnidreams/omnidreams/_drift_corrector.py	Adds fused, graph-safe corrector weight dispatch and per-state corrector selection.
integrations/omnidreams/omnidreams/_edit_lora.py	Adds a target-release seam so the fused corrector can compose style-LoRA deltas.
integrations/omnidreams/tests/test_drift_corrector.py	Covers fused weight composition, dispatch behavior, pointer stability, and gate selection.
integrations/omnidreams/tests/test_edit_lora.py	Tests edit-LoRA weight activation and released-target behavior.

Sequence Diagram

sequenceDiagram
    participant Data as Offline pair pipeline
    participant Trainer as Style SFT trainer
    participant LoRA as Text-edit LoRA
    participant Corrector as Drift corrector
    participant Model as Realtime model
    Data->>Trainer: Filtered source and styled latents
    Trainer->>LoRA: Deploy-format checkpoint
    LoRA->>Model: Activate at prompt swap
    Corrector->>Model: Apply gated drift correction
    Model-->>Model: Continue styled autoregressive rollout

Loading

_{Reviews (6): Last reviewed commit: "Add style-drift corrector data-gen, trai..." | Re-trigger Greptile}

[wenqingw-nv]

Handoff note: my NVIDIA access ends 2026-08-28. Maintainer edits are enabled on this PR; all trained checkpoints are published (fork releases on wenqingw-nv/flashdreams-wq) and the full code+weights+demo-videos snapshot is mirrored internally at gitlab-master.nvidia.com/worldsim/flashdreams (branches wenqingw/*, artifacts in wenqingw/handoff-artifacts via LFS). Context owners going forward: @aidanfnv (game integration, #463/#494-merged) and Yong He (research review). Happy to answer questions async post-internship.

🤖 Generated with Claude Code