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

flashdreams/flashdreams/core/attention/kvcache.py

@@ -365,3 +365,35 @@ def reset(self) -> None:
         self._prev_chunk_idx = -1
         self._curr_chunk_idx = None
         self._n_cached = 0
+
+    def clone_kv(self) -> tuple[Tensor, Tensor]:
+        """Return clones of the full physical K/V buffers.
+
+        Contents only — bookkeeping is not captured. Pair with
+        :meth:`overwrite_kv_` to snapshot/restore alternate contents for a
+        cache whose buffer addresses must stay stable (e.g. under CUDA
+        graphs).
+        """
+        return self._k.clone(), self._v.clone()
+
+    def overwrite_kv_(self, k: Tensor, v: Tensor) -> None:
+        """Overwrite the full physical K/V buffers in place.
+
+        Writes through ``copy_`` so the buffers keep their storage
+        addresses — required under CUDA graphs, whose captured kernels bake
+        in the buffer pointers. Bookkeeping is untouched, so this is only
+        meaningful for caches whose logical content spans the whole buffer
+        (e.g. the static cross-attention text cache built by
+        ``from_tensor``).
+
+        Args:
+            k: Replacement keys; must match the buffer shape exactly.
+            v: Replacement values; must match the buffer shape exactly.
+        """
+        assert k.shape == self._k.shape and v.shape == self._v.shape, (
+            f"overwrite_kv_ shape mismatch: got k {tuple(k.shape)} / "
+            f"v {tuple(v.shape)}, cache holds k {tuple(self._k.shape)} / "
+            f"v {tuple(self._v.shape)}"
+        )
+        self._k.copy_(k)
+        self._v.copy_(v)

integrations/omnidreams/guidance_distill/PLAN.md

@@ -0,0 +1,67 @@
+# Guidance self-distillation (Tier-2a of the live-edit hack)
+
+**Goal:** bake the two-prompt text-edit guidance (`TextEditGuidance`, s≈3) into a LoRA so
+a *plain* mid-stream prompt swap responds like a *guided* one — recovering the ~2x edit
+strength at **zero inference cost** (guidance doubles the DiT forwards while active).
+
+**Why it should work:** the teacher and student are the same network; the target is the
+network's own guided output on RNG-matched on-policy states. This is standard
+CFG-distillation, except the "CFG" here is the old-prompt/new-prompt axis and it only
+matters for a few chunks after a swap. No external data or models needed.
+
+## Recipe (on-policy, mirrors `drift_correction/train_v2.py`)
+
+Per training step:
+
+1. **Sample** a clip (32 local HF samples, `drift_correction/build_pairs._sample_files`),
+   a swap chunk `k ~ U[4, 20]`, and an edit prompt from the bank.
+2. **Roll the student** (LoRA active, plain swap at `k`) with the KV cache to a random
+   chunk `j >= k` — self-forcing-style on-policy states. History replay machinery:
+   `drift_correction/_host.py` (`reset_history`, `replay_history`, bracket helpers).
+3. **At chunk `j`, per denoise step** (timesteps 1000, 450):
+   - Teacher flow = frozen base (LoRA scale 0) with the guidance combine
+     (`kv_old`/`kv_new` loads + `flow_old + s*(flow_new - flow_old)`) — i.e. exactly
+     `CosmosTransformer._predict_with_text_edit_guidance` on unwrapped weights.
+   - Student flow = LoRA'd network, single branch, new-prompt KV only.
+   - Loss = MSE(student, teacher) in v-space; optionally also the context forward
+     (t=128) so committed history matches.
+4. **Backprop** through the student's step only (history detached — the KV buffer write
+   severs grads anyway; use `_train_attn.py` functional dual-branch attention +
+   per-block `torch.utils.checkpoint`, both proven on this host).
+
+**LoRA config:** start from the drift-corrector recipe — r16 on
+`blocks.*.self_attn.{q,k,v,output}_proj` — and add `cross_attn.{q,k,v,output}_proj`
+(the edit signal enters through cross-attn; likely where the capacity is needed).
+`_lora.py:apply_lora` handles both via substring match.
+
+**Prompt bank (v1):** the weather/lighting set from `scripts/sweep_text_edit.py`
+(incl. scene-native snow/rain phrasings) + per-clip base prompts as "no-op edits"
+(swap to the same prompt → teacher == plain flow → regularizes against drift).
+Precompute all text embeddings once (`pipeline.precompute_embeddings` pattern) so the
+14 GB text encoder is not resident during training.
+
+## Deployment: gate the LoRA like the guidance countdown
+
+Enable the LoRA **only for the N chunks after a swap** — the exact window
+`TextEditGuidance.chunks_remaining` covers today — via the drift corrector's per-chunk
+gating + premerge pattern (`_drift_corrector.py`; premerged weight swaps cost ~0 ms).
+Outside the window the base weights run untouched, so non-edit behavior carries zero
+regression risk by construction.
+
+## Eval / kill gate
+
+- Reuse `scripts/sweep_text_edit.py`: (LoRA + plain swap) vs (base + guided) divergence
+  curves on held-out clips x prompts; eyeball grids.
+- Pass: LoRA plain-swap reaches >=80% of guided divergence at matched chunks, with
+  no MUSIQ drop on no-swap rollouts (drift eval harness `eval_rollouts.py`).
+- Budget: ~1k steps eager w/ checkpointing; hours on the shared GB300 (fits the
+  ~65 GB share; full card is comfortable).
+
+## Open choices
+
+- Distill a *fixed* s (3.0) vs conditioning on s (start fixed; the wrapper default
+  becomes "swap = guided-strength swap").
+- Whether to include ReCache in the teacher rollout (probably yes — it is on by
+  default in serving).
+- Later (Tier-2b): extend the same loop with object/appearance edit pairs from
+  JoyAI-Video-Edit to push beyond what guidance alone can reach.

integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py

@@ -98,6 +98,9 @@ def __init__(
         resolution_wh: tuple[int, int],
         seed_for_every_rollout: int | None = None,
         device: torch.device = torch.device("cuda:0"),
+        text_edit_guidance_scale: float = 1.0,
+        text_edit_guidance_chunks: int = 0,
+        text_edit_recache: bool = True,
     ) -> None:
         """Instantiate the pipeline from a registered Omnidreams config.
 
@@ -113,6 +116,16 @@ def __init__(
             seed_for_every_rollout: Optional per-rollout RNG seed override. When
                 ``None``, each rollout draws a fresh OS-entropy seed.
             device: CUDA device the pipeline is moved to.
+            text_edit_guidance_scale: Edit strength applied when a mid-stream
+                prompt swap arrives via ``continue_generation``. ``1.0``
+                disables guidance (plain hot-swap); ``> 1.0`` amplifies the
+                edit for ``text_edit_guidance_chunks`` chunks at the cost of
+                one extra network forward per denoising step while active.
+            text_edit_guidance_chunks: Number of chunks to guide after a swap.
+            text_edit_recache: Re-commit the previous chunk's KV history
+                under the new prompt on every swap (one extra context
+                forward), so the attended window is consistent with the new
+                text.
 
         Raises:
             KeyError: ``pipeline_config`` is omitted and ``pipeline_config_name``
@@ -143,6 +156,9 @@ def __init__(
         self.video_resolution_wh = resolution_wh
         self._rollout_seed = seed_for_every_rollout
         self.fps = 30
+        self._text_edit_guidance_scale = text_edit_guidance_scale
+        self._text_edit_guidance_chunks = text_edit_guidance_chunks
+        self._text_edit_recache = text_edit_recache
 
         # ``len_t`` latent frames per AR block decode into ``len_t * 4`` pixel
         # frames for every continuation step; the first block emits a single
@@ -461,6 +477,36 @@ def start_generation(
             finalization_state={"autoregressive_index": 0},
         )
 
+    def apply_text_prompts(
+        self,
+        state: OmnidreamsConditioningState,
+        text_prompts: list[TextPrompt],
+    ) -> None:
+        """Mid-stream prompt swap at a chunk boundary.
+
+        Rebuilds the text cross-attention KV in place; the KV history
+        carries the generated scene forward under the new prompt. Only call
+        between a finalized chunk and the next ``continue_generation`` (or
+        pass ``text_prompts`` to ``continue_generation`` directly), and only
+        when the prompt actually changes — every call re-runs the 7B text
+        encoder.
+        """
+        assert len(text_prompts) == 1, (
+            "Only one text prompt (batch size == 1) is supported for now"
+        )
+        if state.pipeline_cache is None:
+            raise ValueError(
+                "Cannot swap the prompt: pipeline_cache is None "
+                "(session was started with skip_video_generation=True)"
+            )
+        self.pipeline.replace_text(
+            state.pipeline_cache,
+            self._build_text_batch(text_prompts),
+            guidance_scale=self._text_edit_guidance_scale,
+            guidance_chunks=self._text_edit_guidance_chunks,
+            recache_last_chunk=self._text_edit_recache,
+        )
+
     def continue_generation(
         self,
         state: OmnidreamsConditioningState,
@@ -525,12 +571,17 @@ def continue_generation(
         prev_block_idx = state.pipeline_cache.autoregressive_index
         block_idx = 0 if prev_block_idx is None else prev_block_idx + 1
 
+        if text_prompts is not None:
+            with profiler.measure(
+                "pipeline.replace_text", session_id=session_id, chunk_idx=chunk_idx
+            ):
+                self.apply_text_prompts(state, text_prompts)
+
         with profiler.measure(
             "pipeline.continue_generation",
             session_id=session_id,
             chunk_idx=chunk_idx,
         ):
-            del text_prompts  # Pipeline currently keeps prompts from initialize_cache.
             rgb_frames = self.pipeline.generate(
                 autoregressive_index=block_idx,
                 hdmap=condition,

integrations/omnidreams/omnidreams/pipeline.py

@@ -363,6 +363,100 @@ def precompute_embeddings(
             torch_module=torch,
         )
 
+    @torch.no_grad()
+    def replace_text(
+        self,
+        cache: OmnidreamsPipelineCache,
+        text: list[list[str]],
+        *,
+        guidance_scale: float = 1.0,
+        guidance_chunks: int = 0,
+        recache_last_chunk: bool = False,
+    ) -> None:
+        """Hot-swap the rollout's prompt between two AR steps.
+
+        Encodes ``text`` with the resident text encoder and rebuilds the
+        cross-attention text K/V in place; the self-attention history keeps
+        the generated scene, so the video continues seamlessly under the new
+        prompt. Call after ``finalize`` of one AR step and before
+        ``generate`` of the next.
+
+        Args:
+            cache: Live per-rollout cache.
+            text: ``[B, V]`` nested list of prompts, as in
+                ``initialize_cache``.
+            guidance_scale: Optional edit strength (``> 1.0`` pushes the
+                flow along the new-minus-old text direction for the next
+                ``guidance_chunks`` chunks at the cost of one extra network
+                forward per denoising step).
+            guidance_chunks: Number of upcoming chunks to guide.
+            recache_last_chunk: Re-commit the previous chunk's KV history
+                under the new prompt (one extra context forward), so the
+                window the next chunk attends to is already "explained" by
+                the new text. Helps the scene react faster after a swap.
+        """
+        assert self.text_encoder is not None, (
+            "replace_text requires the text encoder to be loaded; use "
+            "replace_text_from_embeddings with precomputed embeddings "
+            "otherwise."
+        )
+        assert isinstance(text, list) and len(text) > 0 and isinstance(text[0], list), (
+            f"text must be a [B, V] nested list of prompts, got {type(text)}"
+        )
+        text_embeddings = torch.stack(
+            [self.text_encoder(t) for t in text], dim=0
+        )  # [B, V, L, D]
+        self.replace_text_from_embeddings(
+            cache,
+            text_embeddings,
+            guidance_scale=guidance_scale,
+            guidance_chunks=guidance_chunks,
+            recache_last_chunk=recache_last_chunk,
+        )
+
+    @torch.no_grad()
+    def replace_text_from_embeddings(
+        self,
+        cache: OmnidreamsPipelineCache,
+        text_embeddings: Tensor,
+        *,
+        guidance_scale: float = 1.0,
+        guidance_chunks: int = 0,
+        recache_last_chunk: bool = False,
+    ) -> None:
+        """``replace_text`` for precomputed ``[B, V, L, D]`` embeddings."""
+        transformer = self.diffusion_model.transformer
+        assert isinstance(transformer, CosmosTransformer)
+        text_embeddings = text_embeddings.to(device=self.device)
+        text_embeddings = split_inputs_cp(
+            text_embeddings, seq_dim=1, cp_group=self.V_group
+        )
+        transformer.replace_text_embeddings(
+            cache.transformer_cache,
+            text_embeddings,
+            guidance_scale=guidance_scale,
+            guidance_chunks=guidance_chunks,
+        )
+        if recache_last_chunk:
+            self.recache_last_chunk(cache)
+
+    @torch.no_grad()
+    def recache_last_chunk(self, cache: OmnidreamsPipelineCache) -> None:
+        """Re-commit the previous chunk's KV history under the current text.
+
+        Re-opens the just-finalized AR step (``BlockKVCache`` permits
+        same-index rewrites: the window does not roll and the same physical
+        slots are overwritten) and re-runs the context forward, so the
+        cached history becomes consistent with a freshly swapped prompt.
+        Requires the step's ``finalize`` to have completed; a no-op before
+        the first ``generate``.
+        """
+        final_state = cache.final_state
+        if final_state is None:
+            return
+        final_state.cache.start(final_state.autoregressive_index)
+        self.diffusion_model.finalize(final_state=final_state)
+
     def _validate_image_resolution(self, image: Tensor) -> None:
         transformer = self.diffusion_model.transformer
         assert isinstance(transformer, CosmosTransformer), (