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, though so far only at ghost strength (see 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).

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_spawn_car.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.

[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 or distilled-LoRA edit windows, and WebRTC actor controls. It also introduces same-index ReCache for recommitting the previous chunk, but that path still reconstructs history with newly sampled context noise.

• Adds in-place text cross-attention KV replacement while preserving self-attention history and CUDA-graph buffer addresses.
• Adds two-branch edit guidance and an optional pre-merged edit LoRA.
• Adds WebRTC prompt events and spawned-actor conditioning.
• Adds CPU tests and GPU smoke probes for text edits and actor conditioning.

Confidence Score: 4/5

The PR is not yet safe to merge because ReCache still replaces finalized history with KV derived from a different context-noise realization.

ReCache only retains the clean final latent, installs a newly seeded generator, and reruns finalization, so the rewritten KV does not represent the same noisy context used when the prior chunk was rendered.

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 recommits the previous latent using a different context-noise realization.
integrations/omnidreams/omnidreams/transformer/init.py	Adds text-KV replacement, guidance countdown, dual-prompt flow combination, and optional edit-LoRA activation.
flashdreams/flashdreams/core/attention/kvcache.py	Adds in-place full-buffer KV cloning and overwrite primitives while preserving storage addresses.
integrations/omnidreams/omnidreams/webrtc/actors.py	Adds spawned-actor state, geometry, trajectory sampling, and conversion into Ludus conditioning pools.
integrations/omnidreams/omnidreams/webrtc/session.py	Wires prompt-edit and actor commands into live WebRTC session state.

Sequence Diagram

sequenceDiagram
  participant Event as WebRTC edit event
  participant Pipeline
  participant Transformer
  participant Cache as KV cache
  Event->>Pipeline: replace_text(new prompt)
  Pipeline->>Transformer: replace_text_embeddings
  Transformer->>Cache: overwrite text cross-attention KV
  Pipeline->>Cache: start(previous AR index)
  Pipeline->>Pipeline: finalize(clean latent, newly seeded noise)
  Pipeline->>Cache: overwrite previous chunk KV
  Note over Cache: Rewritten history differs from rendered chunk

Loading

_{Reviews (7): Last reviewed commit: "Add /spawnt template spawner: clone real..." | 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/)