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/_edit_lora.py

@@ -0,0 +1,142 @@
+# 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
+from typing import cast
+
+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 = cast(nn.Module, 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 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)"
+        )

integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py

@@ -24,10 +24,12 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any
 
 import numpy as np
 import torch
+from loguru import logger
 from ludus_renderer import CubePool
 from omnidreams.conditioning.renderer import LudusRenderer
 from omnidreams.conditioning.world_scenario.data_types import SceneData
@@ -98,6 +100,10 @@ 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,
+        text_edit_lora_path: "str | Path | None" = None,
     ) -> None:
         """Instantiate the pipeline from a registered Omnidreams config.
 
@@ -113,6 +119,21 @@ 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.
+            text_edit_lora_path: Optional ``guidance_distill`` LoRA
+                checkpoint. When set, edit windows run through the
+                pre-merged distilled weights (guided strength, single
+                forward per denoise step) instead of the two-branch
+                guidance combine.
 
         Raises:
             KeyError: ``pipeline_config`` is omitted and ``pipeline_config_name``
@@ -143,6 +164,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
@@ -155,6 +179,14 @@ def __init__(
         assert isinstance(pipeline, OmnidreamsPipeline)  # for type checking
         self.pipeline: OmnidreamsPipeline = pipeline
 
+        if text_edit_lora_path is not None:
+            from omnidreams._edit_lora import TextEditLoRA
+
+            transformer = pipeline.diffusion_model.transformer
+            edit_lora = TextEditLoRA(transformer.network, text_edit_lora_path)
+            transformer.set_text_edit_lora(edit_lora)
+            logger.info("Deployed {}", edit_lora.describe())
+
     @property
     def V_group(self) -> torch.distributed.ProcessGroup | None:
         # Pipeline backend handles CP internally, so server-side split/gather
@@ -461,6 +493,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 +587,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,