Omnidreams: live mid-stream editing — prompt hot-swap, edit guidance, actor spawn/despawn

[wenqingw-nv]

Omnidreams fixes the prompt at rollout start — the serving path encodes it once and discards later prompts — so the world can't be edited while driving. This PR adds live mid-stream editing: swap the prompt at any chunk boundary with full scene continuity, optionally amplify the edit with two-prompt guidance, and spawn/despawn grounded actors from user commands. Everything is opt-in at the event layer; a session that never edits runs the untouched hot path.

Contents

• core/attention/kvcache.py — clone_kv() / overwrite_kv_(): full-buffer KV overwrite through copy_, so storage addresses survive and captured CUDA graphs stay valid across a swap.
• omnidreams/transformer + pipeline.py — replace_text_embeddings / replace_text: rebuild the per-block text cross-attn KV in place at a chunk boundary; the self-attn history is untouched, so the scene carries forward under the new prompt. TextEditGuidance: for N post-swap chunks, each denoise step runs the cond branch under old and new text against the same history and extrapolates flow_old + s(flow_new − flow_old) (one captured graph, outputs cloned per replay; the KV commit always runs single-branch under the new prompt). recache_last_chunk: LongLive/GameCraft-2-style ReCache — a same-index cache-bracket re-open re-commits the previous chunk's KV under the new text.
• omnidreams/webrtc — datachannel event messages carry free-text prompts (a clear state restores the scene prompt); /spawn <preset> [dist_m] [speed_mps] [lateral_m] and /clear-actors place constant-velocity world-frame actor boxes into the HDMap conditioning through the same converter as the gRPC dynamic-actor path; /spawnt <fwd_m|auto> [lateral_m] [template_idx] clones a real parked-track template to a user target — the first spawn path that materializes anything (see the stress-tested placement rules below). Client gains a Scene Prompt panel + spawn buttons; typing in text fields no longer drives WASD.
• tests/ — 23 ci_cpu tests: swap reproduces a fresh cache init at stable addresses; guidance combine / countdown / finalize suppression; same-index rewrite semantics; spawn geometry and pool timing; template extraction filters, gap search, and clone rigidity.
• scripts/ — three GPU probes (RNG-matched per-chunk divergence vs control): smoke_text_edit.py, sweep_text_edit.py, smoke_spawn_actor.py (drives the WebRTC runtime headless).
• guidance_distill/PLAN.md — follow-up recipe to distill the guidance into a post-swap-gated LoRA (guided strength at plain-swap cost).

Measured (704×1280, 2-step distilled student, swap at chunk 8, RNG-matched vs control, GB300)

Swaps are bit-clean: 0.000 pre-swap pixel divergence. Post-swap, per-chunk mean |Δ| in uint8 units:

edit	s	gap @ +6 chunks	visual outcome
rain-night (training-style phrasing)	3	35	glossy wet road, streetlight reflections, windshield droplets
snow (training-style phrasing)	3	23	progressive full snow cover, artifact-free
night / sunset / fog	3	37 / 20 / 19	all land
snow (freeform phrasing)	5	52	dramatic but transition artifacts — s=3 + caption-style phrasing is the envelope

Actor spawn — corrected after a reviewer-caught misattribution, now with a partially materializing path. The /spawn preset plumbing verifiably places grounded, correctly-classed boxes into the HDMap conditioning, but synthetic preset boxes do not materialize objects — on the distilled student or on the 35-step bidirectional teacher (mask-level checks against an RNG-matched baseline show empty road at the box locations; earlier "spawned car" observations were the scene's own recorded traffic re-rendered after divergence). The model keys on perception-track statistics that hand-built presets lack, so /spawnt clones a real parked-vehicle track out of the scene's own cube pools — per-frame jitter, true dimensions, orientation, colors — and rigidly moves it to the user target; auto picks the largest actor-free gap on the requested lateral line, so a materialized clone provably isn't recorded traffic. Honest quality bar (mask-verified across three configs — right curb auto, left curb auto, two stacked clones): 2 of 3 placements materialize a semi-transparent, blurry vehicle where the RNG-matched control shows empty road — box-interior darkening ~9–10 uint8 vs ~52 for a scene-native parked car, i.e. roughly 20% of full render strength — and the left-curb placement failed to materialize at all (near-zero darkening; in-box divergence 13–16 is texture shift, not an object). Cloning is the existence proof that user boxes can materialize training-free; solid rendering is the open problem. Leading explanation: the context frames and all KV history predate the box, so the 2-step student splits the difference between conditioning and history (scene-native parked cars are in the first-frame pixels; recorded moving traffic materializes solidly from boxes alone, suggesting a fixed-spot history conflict rather than a conditioning limitation). Preset /spawn remains as game world-state (collision logic, gRPC parity); solid materialization is the follow-up (history re-commit on spawn, mild box-axis guidance over the clone signal, moving-track clones, instruction-SFT). Spawned boxes are grounded on the road plane (the ego pose is the rig origin, ~1.5 m up).

Spawn stress tests (later commits: probe_moving_clone.py, probe_pedestrians.py)

Placement-manifold rule, mask-verified against RNG-matched no-spawn controls: clones materialize only where their class plausibly lives in the AV training distribution.

scenario	in-box divergence (uint8, vs ~52 real-car contrast)	outcome
parked-car clone, curb	9–10 darkening (~20%)	semi-transparent ghost
parked clone + box-axis guidance s=2	~20 (~40%)	solid by late frames
moving-track clone (oncoming)	18–21 from entry	solid but smeared
moving clone + s=2/s=3 guidance	41–53 (≈ benchmark)	fully opaque; smooth "toy-car" texture
pedestrians mid-road, 1/5/20	2–8 (noise)	never materialize, any density
pedestrians on sidewalks, 3/20/40	19–23	materialize and scale

Box-axis guidance (GUIDE_SCALE on the probes) is the two-prompt-guidance analog on the HDMap axis — it amplifies a clone's real signal (s=2–3 usable) but cannot create one for off-manifold placements (mid-road crowds, synthetic presets). Texture realism beyond this is SFT territory (separate track).

sbs_spawnt_ped40_annotated.mp4

Overhead

• Idle: zero — no new work on the hot path, and text_edit_guidance=None short-circuits to the existing single-branch forward.
• On swap: one text-encoder forward + in-place KV copies at a chunk boundary; ReCache adds one off-hot-path context forward (~42 ms).
• Guidance active (default s=3 for 6 chunks): +1 DiT forward per denoise step, 84 → ~168 ms/chunk — within the 266 ms/chunk playback budget.

Guidance distillation (Tier-2a, later commits)

guidance_distill/PLAN.md documents the recipe (the trainer scripts build on the Clean Forcing training infra and land with #398); the deploy hook here is self-contained. The recipe trains a LoRA to reproduce the guided flow from a plain swap — teacher is the frozen base running the guidance combine on the same on-policy states, so no external data or models. The r64 / 1600-step checkpoint passes the eval gate: the LoRA'd plain swap reaches 0.854 of guided divergence on held-out clips (bar 0.8; untrained base 0.376), visually clean. Deployed via omnidreams/_edit_lora.py: both weight sets pre-merged at load, edit windows toggle by in-place copy_ (CUDA-graph-safe, ~0.2 ms at window boundaries only) — guided-strength edits at plain-swap cost, replacing the +84 ms/chunk two-branch combine. Opt-in via text_edit_lora_path; without it, plain swaps are bit-identical to the hookless run (GPU-verified). Checkpoint (211 MB, md5 1fc7b33ea037f745eeb8d5a1bee59371): https://github.com/wenqingw-nv/flashdreams-wq/releases/tag/live-edit-guidance-lora-v1 (fork release, per the drift-corrector precedent; happy to move to HF on maintainer preference).

Not included (follow-ups)

gRPC field for mid-stream prompts; the native optimized-DiT (FP8) path raises NotImplementedError on swap (default configs unaffected); instruction-SFT for grounded props, object-by-prompt edits, and trajectory-locked moving actors.

---

Video comparison (original | edited)

All clips are RNG-matched: identical seed, controls, and conditioning until the edit; labels burned in. Left: base rollout. Right: same rollout with the live edit.

Prompt swap → heavy rain (guided s=3, swap at 2 s): wet glossy road, streetlight/taillight reflections, windshield droplets — same street, same trajectory.

sbs_rain_edit.mp4

Prompt swap → snowstorm (training-style phrasing, s=3): progressive snow cover on road, trees, and parked cars, artifact-free.

sbs_snow_edit.mp4

/spawn car 16 mid-drive (superseded — kept for the correction record): the car visible here is the scene's own traffic re-rendered after divergence, not the commanded box; see the corrected Actor-spawn paragraph above.

sbs_spawnt_moving_s3_annotated.mp4

Spawn + /clear-actors (superseded — kept for the correction record): same confound as above; the appearing/vanishing car tracks the scene's recorded traffic under divergence rather than the commanded box.

sbs_spawn_remove.mp4

Props via prompt assist: an oversized "Other"-class cone box alone under-renders, but combined with a construction-zone prompt swap the model paints real orange traffic cones along the road — visual dressing rather than grounded objects, which is the documented props limit.

sbs_cone_prompt_assist.mp4

/spawnt template clone (verified, ghost-strength): a real parked-track clone placed in an auto-selected empty gap materializes as a semi-transparent vehicle on the right curb (~20% of a scene-native car's contrast); the RNG-matched control (left) shows empty curb at the same location. Three reliability clips staged in integrations/omnidreams/scripts/outputs/pr_videos/sbs_spawnt_rel{1,2,3}.mp4 — rel1/rel3 show the ghost vehicle, rel2 (left curb) failed to materialize and is kept for the record.

/spawnt moving clone + s=3 guidance (annotated): the red box tracks the user-spawned clone — a fully opaque car passing in the left lane; scene traffic is unmarked.

Sidewalk pedestrians (annotated, 20 and 40): person-sized clones on both sidewalks materialize as standing figures; the same grids placed mid-road never materialize (the placement-manifold rule above).

---

Update: composite overlay stack (cf686c8a)

Spawn-SFT probes established that the student will not materialize arbitrary new actors from boxes alone (placement-manifold rule above). The practical route we landed on for game-style object/actor insertion is composite, then renoise-refine:

• composite_assets.py — pastes photoreal cutouts (or custom sprites) at box-projected screen positions. Placement masks come from diffing the boxed vs baseline conditioning renders, so positions are exact in screen space (fisheye included); IoU tracking + EMA smoothing keep sprites glued to their boxes, and contact shadows + luminance/chroma harmonization match the plate.
• composite_track_items.py — item courses without hdmap or GPU: a pinhole camera is fitted from the clip's lane geometry, then animated pickups and moving traffic are projected along the lane.
• probe_composite_refine.py — re-rolls the composited clip through the student from moderate noise (sweet spot σ≈0.47): pasted content inherits model texture and lighting while the scene stays pinned (reported as in-box divergence vs out-of-box cleanliness).
• enhance_frames.py — optional learned restoration pass (4 Real-ESRGAN variants with measured fps).
• probe_pedestrians.py gains the teacher/guidance/crowd-follow/ego-stop/walk-speed modes used for the crowd-persistence findings.

[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 live prompt replacement, transient edit guidance and LoRA deployment, plus WebRTC actor editing. The attempted ReCache correction preserves the rollout RNG stream but still re-finalizes history using a different context-noise realization.

• Adds pointer-stable text cross-attention KV replacement and edit-guidance windows.
• Adds prompt and actor commands to the WebRTC session and browser UI.
• Adds template-based actor placement, probes, and CPU tests.
• Adds an optional pre-merged edit LoRA deployment path.

Confidence Score: 4/5

The PR is not yet safe to merge because ReCache can replace finalized KV history with a representation derived from context noise different from the rendered chunk.

ReCache preserves the rollout generator state but FinalState retains only the clean latent, so repeated finalization under the dedicated generator necessarily samples a new noisy latent and overwrites same-index history with an inconsistent representation.

Files Needing Attention: integrations/omnidreams/omnidreams/pipeline.py

Important Files Changed

Filename	Overview
integrations/omnidreams/omnidreams/pipeline.py	Adds prompt replacement and ReCache, but ReCache still rewrites finalized history from newly sampled context noise.
integrations/omnidreams/omnidreams/transformer/init.py	Adds text-edit guidance state, prompt-KV replacement, and LoRA window handling.
flashdreams/flashdreams/core/attention/kvcache.py	Adds pointer-stable full-buffer cloning and overwrite operations for captured CUDA graphs.
integrations/omnidreams/omnidreams/webrtc/session.py	Integrates prompt-edit and actor commands into live WebRTC sessions.
integrations/omnidreams/omnidreams/webrtc/actors.py	Adds synthetic and template-derived actor placement and lifecycle logic.
integrations/omnidreams/omnidreams/_edit_lora.py	Adds in-place toggling between cached base and LoRA-merged projection weights.

Sequence Diagram

sequenceDiagram
  participant Client
  participant Pipeline
  participant Transformer
  participant FinalState
  participant KVCache
  Client->>Pipeline: replace_text(new prompt)
  Pipeline->>Transformer: replace_text_embeddings()
  Transformer->>KVCache: overwrite text cross-attention KV
  Pipeline->>FinalState: recache_last_chunk()
  FinalState->>Pipeline: retained clean latent
  Pipeline->>Pipeline: draw newly seeded context noise
  Pipeline->>Transformer: finalize_kv_cache(new noisy latent)
  Transformer->>KVCache: same-index history overwrite

Loading

_{Reviews (9): Last reviewed commit: "Add composite overlay stack; extend pede..." | Re-trigger Greptile}

[wenqingw-nv]

Re the ReCache noise concern: addressed in 5541985, with context on why the behavior was in-distribution either way.

The context forward is trained on randomly re-noised latents, and every chunk's original KV commit already uses an independent noise draw — so a ReCache re-commit under a fresh draw produces a valid representation of the same clean latent (our GPU probes showed improved post-swap compliance and no artifacts). The commit that mattered: ReCache previously consumed the model RNG, shifting every subsequent noise draw relative to a no-ReCache rollout. It now draws from a dedicated per-AR-index seeded generator, making the re-commit deterministic and leaving the rollout's noise stream identical with or without ReCache (unit-tested).

[wenqingw-nv]

/ok to test 5541985

[greptile-apps]

ReCache resamples finalized history

When a prompt swap runs after a chunk finalized with nonzero context noise, ReCache installs a newly seeded generator and finalizes the retained clean latent again. This overwrites the rendered chunk's KV with a representation derived from different context noise, causing continuity loss or transition artifacts in the next chunk.

Knowledge Base Used: Core Engine (flashdreams/flashdreams/core/)

[greptile-apps]

ReCache still resamples finalized history

When a mid-stream prompt swap enables ReCache after a chunk finalized with nonzero context_noise, this dedicated generator makes the repeated finalization draw a different noise realization because FinalState retains only the clean latent. The same-index rewrite therefore replaces the rendered chunk's KV entries with a representation derived from a different noisy latent, causing boundary discontinuity or transition artifacts in the next chunk.

Knowledge Base Used:

• Core Engine (flashdreams/flashdreams/core/)
• Infra Pipeline: Building and Running a FlashDreams Model

[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