diff --git a/flashdreams/flashdreams/core/attention/kvcache.py b/flashdreams/flashdreams/core/attention/kvcache.py index 5673a32a9..79d88441d 100644 --- a/flashdreams/flashdreams/core/attention/kvcache.py +++ b/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) diff --git a/integrations/omnidreams/guidance_distill/PLAN.md b/integrations/omnidreams/guidance_distill/PLAN.md new file mode 100644 index 000000000..b6eb42696 --- /dev/null +++ b/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. diff --git a/integrations/omnidreams/omnidreams/_edit_lora.py b/integrations/omnidreams/omnidreams/_edit_lora.py new file mode 100644 index 000000000..13e56ddec --- /dev/null +++ b/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)" + ) diff --git a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py b/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py index 27500f7be..f28389c45 100644 --- a/integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py +++ b/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, diff --git a/integrations/omnidreams/omnidreams/pipeline.py b/integrations/omnidreams/omnidreams/pipeline.py index 5814f9926..fabd00c76 100644 --- a/integrations/omnidreams/omnidreams/pipeline.py +++ b/integrations/omnidreams/omnidreams/pipeline.py @@ -363,6 +363,125 @@ 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) + + _RECACHE_NOISE_SEED = 118_000 + """Base seed for the ReCache context-noise draw (offset by AR index).""" + + @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``. + + The context-noise draw comes from a dedicated generator seeded by + the AR index, not the model RNG — every noise rendition of the same + clean latent is in-distribution for the context forward (each + chunk's original commit already uses an independent draw), and + keeping the model RNG untouched means the rollout's subsequent + noise stream is identical with or without ReCache. Seedless + configurations (``DiffusionModelConfig.seed is None``) fall back to + the global RNG, matching their existing no-reproducibility + contract. + """ + final_state = cache.final_state + if final_state is None: + return + diffusion_model = self.diffusion_model + # Materialize the lazy model generator BEFORE snapshotting, so the + # restore never resets the rollout's noise stream to its seed. + seeded = diffusion_model.rng is not None + saved_rng = diffusion_model._rng + if seeded: + diffusion_model._rng = torch.Generator(device=self.device).manual_seed( + self._RECACHE_NOISE_SEED + final_state.autoregressive_index + ) + try: + final_state.cache.start(final_state.autoregressive_index) + diffusion_model.finalize(final_state=final_state) + finally: + diffusion_model._rng = saved_rng + def _validate_image_resolution(self, image: Tensor) -> None: transformer = self.diffusion_model.transformer assert isinstance(transformer, CosmosTransformer), ( diff --git a/integrations/omnidreams/omnidreams/transformer/__init__.py b/integrations/omnidreams/omnidreams/transformer/__init__.py index 886d40739..fe18af7fa 100644 --- a/integrations/omnidreams/omnidreams/transformer/__init__.py +++ b/integrations/omnidreams/omnidreams/transformer/__init__.py @@ -77,6 +77,48 @@ ## Per-rollout cache +@dataclass(kw_only=True) +class TextEditGuidance: + """Transient two-prompt guidance for a mid-rollout text edit. + + Built by :meth:`CosmosTransformer.replace_text_embeddings` when an edit + strength is requested. While active, ``predict_flow`` runs the cond + branch twice — once with the pre-edit ("old") text K/V and once with the + post-edit ("new") K/V — and combines them CFG-style: + + ``flow = flow_old + scale * (flow_new - flow_old)`` + + Both branches share the same self-attention history, so the guidance + direction is purely the text difference; the old-prompt branch anchors + scene identity while ``scale > 1`` amplifies the edit. KV contents are + loaded into the existing cross-attention buffers via ``overwrite_kv_``, + which preserves storage addresses and therefore composes with CUDA-graph + replay. The per-chunk KV commit (``finalize_kv_cache``) always runs + single-branch under the new prompt. + """ + + scale: float + """Edit strength: 1.0 reproduces the new prompt exactly (but wastes a + forward — callers should just not build this state); > 1.0 amplifies.""" + + chunks_remaining: int + """Number of upcoming AR chunks to apply guidance to. Decremented by + :meth:`CosmosTransformerCache.start`; the state clears itself after.""" + + kv_old: list[tuple[Tensor, Tensor]] = field(default_factory=list) + """Per-block (K, V) cross-attention contents of the pre-edit prompt + (unused for ``use_lora`` windows).""" + + kv_new: list[tuple[Tensor, Tensor]] = field(default_factory=list) + """Per-block (K, V) cross-attention contents of the post-edit prompt + (unused for ``use_lora`` windows).""" + + use_lora: bool = False + """Realize the window with the pre-merged edit LoRA weights instead of + the two-branch guidance combine: single forward per denoise step at + guided strength (the LoRA distilled the combine; see ``_edit_lora``).""" + + @dataclass(kw_only=True) class CosmosTransformerCache(TransformerAutoregressiveCache): """Long-lived AR cache for the Cosmos transformer.""" @@ -114,7 +156,20 @@ class CosmosTransformerCache(TransformerAutoregressiveCache): autoregressive_index: int = -1 """AR step index for the chunk currently being processed; ``-1`` before the first ``start``.""" + text_edit_guidance: TextEditGuidance | None = None + """Two-prompt guidance for an in-flight text edit; ``None`` when idle.""" + def start(self, autoregressive_index: int) -> None: + # Advance the text-edit guidance countdown on real chunk advances + # only (a same-index re-open, e.g. a post-swap KV re-commit of the + # previous chunk, must not consume a guidance chunk). + guidance = self.text_edit_guidance + if guidance is not None and autoregressive_index > self.autoregressive_index: + if guidance.chunks_remaining <= 0: + self.text_edit_guidance = None + else: + guidance.chunks_remaining -= 1 + # Hoist KV pre-update and RoPE shift out of the graph-captured forward # (predict_flow runs eager_mode=False; cond/uncond share rope_freqs). self.rope_freqs = self.rope_adapter.shift_t(autoregressive_index) @@ -346,6 +401,26 @@ def __init__(self, config: CosmosTransformerConfig) -> None: # directly. Multi-view: keep 5D [B, V, T, HW, D] for hierarchical CP. self.flatten_thw = config.num_views == 1 + # True while finalize_kv_cache runs its context forward; text-edit + # guidance is suppressed there so the KV commit is single-branch + # under the (new) post-edit prompt. + self._finalizing_kv_cache = False + + # Optional pre-merged edit LoRA (omnidreams._edit_lora.TextEditLoRA): + # when attached, replace_text_embeddings builds use_lora windows and + # predict_flow toggles the merged weights instead of double-branching. + self._text_edit_lora: Any | None = None + + def set_text_edit_lora(self, edit_lora: Any | None) -> None: + """Attach (or detach with ``None``) a pre-merged edit-LoRA hook. + + The hook must expose ``set_active(bool)`` and ``active`` + (:class:`omnidreams._edit_lora.TextEditLoRA`). While attached, edit + windows requested via :meth:`replace_text_embeddings` run at guided + strength through the merged weights — one forward per denoise step. + """ + self._text_edit_lora = edit_lora + def _configure_optimized_dit_from_config(self) -> None: from omnidreams.native import omnidreams_singleview @@ -600,6 +675,11 @@ def initialize_autoregressive_cache( mask_first_patched = self.patchify_and_maybe_split_cp(mask_first_block) mask_other_patched = self.patchify_and_maybe_split_cp(mask_other_blocks) + # A fresh rollout always starts on the base weights; a mid-window + # session teardown must not leak edit weights into the next session. + if self._text_edit_lora is not None: + self._text_edit_lora.set_active(False) + if self._use_cuda_graph: self._cuda_graph_dispatch.reset() @@ -616,6 +696,91 @@ def initialize_autoregressive_cache( self._optimized_dit_executor.after_initialize_autoregressive_cache(cache) return cache + @torch.no_grad() + def replace_text_embeddings( + self, + cache: CosmosTransformerCache, + text_embeddings: Tensor, + *, + guidance_scale: float = 1.0, + guidance_chunks: int = 0, + ) -> None: + """Hot-swap the rollout's text conditioning at a chunk boundary. + + Rebuilds the per-block cross-attention text K/V in place (storage + addresses survive, so captured CUDA graphs stay valid) while the + self-attention history, RoPE state, and image/mask conditioning are + untouched — the rollout continues under the new prompt with full + visual continuity. Call between ``finalize`` of one AR step and + ``generate`` of the next. + + Args: + cache: Live per-rollout cache. + text_embeddings: ``[B, V, L, D]`` replacement text embeddings + (same fixed ``L`` as the original prompt). + guidance_scale: Optional edit strength. Values ``> 1.0`` enable + two-prompt guidance for the next ``guidance_chunks`` chunks: + the old prompt anchors the scene and the flow is pushed + along the new-minus-old text direction (costs one extra + network forward per denoising step while active). ``1.0`` + disables guidance (plain hot-swap). + guidance_chunks: Number of upcoming chunks to guide; ``0`` + disables guidance. + """ + if self._optimized_dit_executor is not None: + raise NotImplementedError( + "replace_text_embeddings is not wired for the native " + "optimized-DiT path yet; run with " + "native_dit_acceleration='disabled'." + ) + cfg = self.config + text_embeddings = text_embeddings.to(device=self.device, dtype=cfg.dtype) + if self.cp_groups.V_group is not None: + text_embeddings = split_inputs_cp( + text_embeddings, seq_dim=1, cp_group=self.cp_groups.V_group + ) + + use_guidance = guidance_scale != 1.0 and guidance_chunks > 0 + assert not (use_guidance and cache.network_cache_uncond is not None), ( + "Text-edit guidance shares the cond branch's self-attention " + "history and is mutually exclusive with negative-prompt CFG " + "(guidance_scale > 1.0 configs)." + ) + + if use_guidance and self._text_edit_lora is not None: + # Distilled path: the pre-merged LoRA realizes the window at + # guided strength with a single branch — no KV snapshots needed. + self.network.replace_text_embeddings(cache.network_cache, text_embeddings) + self._text_edit_lora.set_active(True) + cache.text_edit_guidance = TextEditGuidance( + scale=guidance_scale, + chunks_remaining=guidance_chunks, + use_lora=True, + ) + return + + block_caches = cache.network_cache.block_caches + kv_old: list[tuple[Tensor, Tensor]] | None = None + if use_guidance: + kv_old = [bc.cross_attn.clone_kv() for bc in block_caches] + + self.network.replace_text_embeddings(cache.network_cache, text_embeddings) + + if use_guidance: + assert kv_old is not None + cache.text_edit_guidance = TextEditGuidance( + scale=guidance_scale, + chunks_remaining=guidance_chunks, + kv_old=kv_old, + kv_new=[bc.cross_attn.clone_kv() for bc in block_caches], + ) + else: + # A plain swap supersedes any in-flight guidance (whose old/new + # snapshots no longer match the buffers). + cache.text_edit_guidance = None + if self._text_edit_lora is not None: + self._text_edit_lora.set_active(False) + ## Mask-injection helpers def _maybe_inject_image( @@ -675,6 +840,45 @@ def _predict_branch( eager_mode=False, ) + def _predict_with_text_edit_guidance( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: CosmosTransformerCache, + input: Tensor | None, + guidance: TextEditGuidance, + ) -> Tensor: + """Two-prompt CFG for an in-flight text edit. + + Runs the cond branch under the old and the new text K/V against the + SAME self-attention history and combines CFG-style. The KV loads + write in place, so under CUDA graphs both calls are plain replays of + the already-captured cond graph (whose outputs are cloned per + replay). Buffers are left holding the new-prompt K/V. + """ + block_caches = cache.network_cache.block_caches + for bc, (k, v) in zip(block_caches, guidance.kv_old): + bc.cross_attn.overwrite_kv_(k, v) + flow_old = self._predict_branch( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + network_cache=cache.network_cache, + input=input, + uncond=False, + ) + for bc, (k, v) in zip(block_caches, guidance.kv_new): + bc.cross_attn.overwrite_kv_(k, v) + flow_new = self._predict_branch( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + network_cache=cache.network_cache, + input=input, + uncond=False, + ) + return flow_old + guidance.scale * (flow_new - flow_old) + def predict_flow( self, noisy_latent: Tensor, @@ -689,6 +893,30 @@ def predict_flow( cache=cache, input=input, ) + guidance = cache.text_edit_guidance + if guidance is not None and guidance.use_lora: + # Distilled edit window: merged weights, single branch. The + # KV-commit forwards inside the window also run merged (the + # LoRA was trained to match the guided context forward). + assert self._text_edit_lora is not None + self._text_edit_lora.set_active(True) + elif self._text_edit_lora is not None and self._text_edit_lora.active: + # Window expired (cache.start cleared the countdown): the first + # forward of the next chunk restores the base weights. + self._text_edit_lora.set_active(False) + if ( + guidance is not None + and not guidance.use_lora + and not self._finalizing_kv_cache + and cache.network_cache_uncond is None + ): + return self._predict_with_text_edit_guidance( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + input=input, + guidance=guidance, + ) flow_cond = self._predict_branch( noisy_latent=noisy_latent, timestep=timestep, @@ -724,7 +952,14 @@ def finalize_kv_cache( ) -> None: try: if not self.config.skip_finalize_kv_cache: - super().finalize_kv_cache(*args, **kwargs) + # The context forward commits KV history single-branch under + # the current (post-edit) prompt; text-edit guidance only + # shapes the denoising flow, never the committed history. + self._finalizing_kv_cache = True + try: + super().finalize_kv_cache(*args, **kwargs) + finally: + self._finalizing_kv_cache = False finally: if self._optimized_dit_executor is not None: self._optimized_dit_executor.after_finalize_kv_cache() diff --git a/integrations/omnidreams/omnidreams/transformer/impl/network.py b/integrations/omnidreams/omnidreams/transformer/impl/network.py index 307e88220..9ca92617a 100644 --- a/integrations/omnidreams/omnidreams/transformer/impl/network.py +++ b/integrations/omnidreams/omnidreams/transformer/impl/network.py @@ -408,6 +408,35 @@ def initialize_cache( ) return CosmosDiTNetworkCache(block_caches=block_caches) + @torch.no_grad() + def replace_text_embeddings( + self, + cache: CosmosDiTNetworkCache, + text_embeddings: Tensor, + ) -> None: + """Replace the cached cross-attention text K/V for all blocks in place. + + Mirrors the cross-attention half of :meth:`initialize_cache`, but + writes through ``copy_`` into the existing cache buffers so their + storage addresses survive — required under CUDA graphs, whose + captured kernels bake in the buffer pointers. Self-attention + history is untouched, so the rollout continues seamlessly under the + new prompt. + + Args: + cache: Live per-rollout network cache. + text_embeddings: ``[B, V, L, D]`` replacement text embeddings; + ``L`` must match the original prompt's token length (the + text encoder pads to a fixed ``max_length``). + """ + context = text_embeddings + if self.config.use_crossattn_projection: + context = self.crossattn_proj(context) + for block, block_cache in zip(self.blocks, cache.block_caches): + assert isinstance(block, Block) + fresh = block.cross_attn.compute_kv(context) + block_cache.cross_attn.overwrite_kv_(*fresh.clone_kv()) + def forward( self, x: Tensor, diff --git a/integrations/omnidreams/omnidreams/webrtc/actors.py b/integrations/omnidreams/omnidreams/webrtc/actors.py new file mode 100644 index 000000000..382c5a142 --- /dev/null +++ b/integrations/omnidreams/omnidreams/webrtc/actors.py @@ -0,0 +1,380 @@ +# 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. + +"""User-spawned dynamic actors for the Omnidreams WebRTC drive. + +The model's control branch was trained to materialize objects at rendered +HDMap bboxes, so "add an object mid-drive" is expressed as a wireframe cube +in the Ludus conditioning stream: spawn a box, the model paints an object +there (the prompt names its appearance). Actors follow a constant-velocity +world-frame motion model — enough for parked obstacles and lead vehicles. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch +from ludus_renderer import CubePool +from omnidreams.grpc.utils import dynamic_state_to_ludus_cube_pool +from scipy.spatial.transform import Rotation + +## Spawn presets + +RIG_HEIGHT_M = 1.5 +"""Ego rig-origin height above the road plane; spawn z-correction.""" + +ACTOR_PRESETS: dict[str, tuple[str, tuple[float, float, float]]] = { + # preset -> (actor class, FLU bbox size (length, width, height) in meters) + "car": ("CAR", (4.6, 2.0, 1.6)), + "truck": ("TRUCK", (8.0, 2.6, 3.2)), + "pedestrian": ("PEDESTRIAN", (0.6, 0.6, 1.8)), + "cyclist": ("CYCLIST", (1.8, 0.7, 1.7)), + "cone": ("OTHER", (0.4, 0.4, 0.8)), + # True-scale cones render too faintly (below the model's salience + # threshold for "Other" boxes) — oversized variants for obstacles. + "cone_big": ("OTHER", (1.0, 1.0, 1.2)), + "barrier": ("OTHER", (2.4, 0.6, 1.0)), +} + + +@dataclass +class SpawnedActor: + """One user-spawned actor with a constant-velocity world trajectory.""" + + class_id: str + """Actor class (drives the obstacle color), e.g. ``"CAR"``.""" + + size_xyz: tuple[float, float, float] + """FLU bbox dimensions in meters.""" + + spawn_timestamp_us: int + """First frame timestamp at which the actor exists.""" + + translation: np.ndarray + """``[3]`` world-frame bbox-center position at spawn time.""" + + quat_xyzw: np.ndarray + """``[4]`` world-frame orientation quaternion.""" + + velocity: np.ndarray + """``[3]`` world-frame velocity in m/s (zeros = parked).""" + + def translation_at(self, timestamp_us: int) -> np.ndarray: + dt_s = (timestamp_us - self.spawn_timestamp_us) * 1e-6 + return self.translation + self.velocity * dt_s + + +def spawn_actor_ahead( + *, + preset: str, + ego_pose: np.ndarray, + spawn_timestamp_us: int, + distance_m: float = 12.0, + speed_mps: float = 0.0, + lateral_m: float = 0.0, + yaw_offset_deg: float = 0.0, +) -> SpawnedActor: + """Place a preset actor relative to the ego vehicle. + + Args: + preset: Key into :data:`ACTOR_PRESETS`. + ego_pose: ``[4, 4]`` world-from-ego FLU pose (x forward, y left, + z up) to spawn relative to. + spawn_timestamp_us: Timestamp of the first frame the actor exists. + distance_m: Meters ahead of the ego along its heading. + speed_mps: Actor speed along the ego heading (0 = parked). + lateral_m: Meters to the left (+) / right (-) of the ego heading. + yaw_offset_deg: Box heading relative to the ego heading (0 = same + direction, 180 = oncoming). The rendered box's front/back face + colors encode heading, which the model reads as travel + direction. + + Raises: + KeyError: Unknown preset. + """ + class_id, size_xyz = ACTOR_PRESETS[preset] + + ego_pose = np.asarray(ego_pose, dtype=np.float64) + rotation = ego_pose[:3, :3] + # Ground-plane heading: project the ego forward axis onto XY so tilted + # camera poses don't pitch the spawned box into the road. + forward = rotation @ np.array([1.0, 0.0, 0.0]) + forward_xy = np.array([forward[0], forward[1], 0.0]) + norm = float(np.linalg.norm(forward_xy)) + if norm < 1e-6: + forward_xy = np.array([1.0, 0.0, 0.0]) + norm = 1.0 + forward_xy /= norm + left_xy = np.array([-forward_xy[1], forward_xy[0], 0.0]) + + center = ( + ego_pose[:3, 3] + + distance_m * forward_xy + + lateral_m * left_xy + # Bbox center sits half a height above the road. The ego pose is the + # RIG origin (~camera height above ground, empirically ~1.5 m on the + # HDMap scenes — verified against the scene's own actor boxes); + # without the correction spawned boxes float at eye level and the + # model under-renders them. + + np.array([0.0, 0.0, size_xyz[2] / 2.0 - RIG_HEIGHT_M]) + ) + yaw = float(np.arctan2(forward_xy[1], forward_xy[0])) + float( + np.deg2rad(yaw_offset_deg) + ) + quat_xyzw = Rotation.from_euler("z", yaw).as_quat().astype(np.float32) + + return SpawnedActor( + class_id=class_id, + size_xyz=size_xyz, + spawn_timestamp_us=int(spawn_timestamp_us), + translation=center.astype(np.float32), + quat_xyzw=quat_xyzw, + velocity=(speed_mps * forward_xy).astype(np.float32), + ) + + +def actors_to_cube_pool( + actors: list[SpawnedActor], + frame_timestamps_us: list[int], + device: torch.device | str, +) -> CubePool | None: + """Sample the actors at the chunk's frame timestamps as a Ludus pool. + + Reuses the gRPC ``DynamicWorldState`` conversion path (colors, + interpolation, category mapping) by building the equivalent actor dicts + with one exact pose per frame timestamp. Actors spawned mid-chunk simply + have no poses for the earlier frames. + """ + actor_dicts: list[dict] = [] + for actor in actors: + poses = [] + for ts in frame_timestamps_us: + ts = int(ts) + if ts < actor.spawn_timestamp_us: + continue + x, y, z = (float(v) for v in actor.translation_at(ts)) + qx, qy, qz, qw = (float(v) for v in actor.quat_xyzw) + poses.append( + { + "timestamp_us": ts, + "pose": { + "vec": {"x": x, "y": y, "z": z}, + "quat": {"x": qx, "y": qy, "z": qz, "w": qw}, + }, + } + ) + if not poses: + continue + size_x, size_y, size_z = actor.size_xyz + actor_dicts.append( + { + "class_id": actor.class_id, + "bbox_dims": {"size_x": size_x, "size_y": size_y, "size_z": size_z}, + "trajectory": {"poses": poses}, + } + ) + if not actor_dicts: + return None + return dynamic_state_to_ludus_cube_pool( + {"actors": actor_dicts}, frame_timestamps_us, device + ) + + +## Template-based spawning (cloned real perception tracks) + + +@dataclass +class TrackTemplate: + """A real scene-actor track extracted for cloning. + + Synthesized preset boxes are ignored by the model (both the distilled + student and the 35-step teacher — mask-verified 2026-08-11), while a + bit-for-bit clone of a real perception track materializes. Templates + carry everything the model may key on: per-frame jitter, real + dimensions, orientation, z, and the source pool's colors and render + flags. + """ + + timestamps_us: torch.Tensor + """``[n]`` original per-sample timestamps.""" + + translations: torch.Tensor + """``[n, 3]`` world positions (with the source's per-frame jitter).""" + + quaternions: torch.Tensor + """``[n, 4]`` world orientations.""" + + scale: torch.Tensor + """``[1, 3]`` bbox dimensions.""" + + colors: torch.Tensor + """``[1, 6]`` front/back face colors.""" + + prim_type_id: int + render_flags: int + source_fwd_m: float + source_lateral_m: float + + +def _pool_track_slices(pool: CubePool) -> list[tuple[int, int]]: + """Per-track (start, end) ranges into a pool's concatenated arrays.""" + prefix = pool.cube_ts_prefix_sum.cpu().numpy() + starts = np.concatenate([[0], prefix[:-1]]) + return [(int(a), int(b)) for a, b in zip(starts, prefix)] + + +def _ego_frame(ego_pose: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """(origin_xy, forward_xy, left_xy) of the ego ground frame.""" + ego_pose = np.asarray(ego_pose, dtype=np.float64) + forward = ego_pose[:3, 0].copy() + forward[2] = 0.0 + forward /= np.linalg.norm(forward) + left = np.array([-forward[1], forward[0]]) + return ego_pose[:2, 3], forward[:2], left + + +def extract_parked_templates( + pools: list[CubePool], + *, + ego_pose: np.ndarray, + t0_us: int, + min_coverage_s: float = 5.5, + length_range: tuple[float, float] = (3.4, 5.6), + max_drift_m: float = 1.5, +) -> list[TrackTemplate]: + """Extract parked car-sized tracks usable as spawn templates. + + Tracks start at their first perception frame, so coverage is measured + from up to 1 s after ``t0_us``. Sorted nearest-to-ego first. + """ + origin, forward, left = _ego_frame(ego_pose) + templates: list[tuple[float, TrackTemplate]] = [] + for pool in pools: + scales = pool.scales.cpu().numpy() + for track_index, (a, b) in enumerate(_pool_track_slices(pool)): + ts = pool.track_timestamps_us[a:b].cpu().numpy() + if ( + len(ts) < 8 + or ts[0] > t0_us + 1_000_000 + or ts[-1] < t0_us + int(min_coverage_s * 1e6) + ): + continue + length = float(scales[track_index].max()) + if not length_range[0] <= length <= length_range[1]: + continue + tr = pool.translations[a:b].cpu().numpy() + if float(np.linalg.norm(tr[-1, :2] - tr[0, :2])) > max_drift_m: + continue + rel = tr[0, :2] - origin + template = TrackTemplate( + timestamps_us=pool.track_timestamps_us[a:b].clone(), + translations=pool.translations[a:b].clone(), + quaternions=pool.quaternions[a:b].clone(), + scale=pool.scales[track_index : track_index + 1].clone(), + colors=pool.colors[track_index : track_index + 1].clone(), + prim_type_id=pool.prim_type_id, + render_flags=pool.render_flags, + source_fwd_m=float(rel @ forward), + source_lateral_m=float(rel @ left), + ) + templates.append((float(np.linalg.norm(rel)), template)) + templates.sort(key=lambda item: item[0]) + return [template for _, template in templates] + + +def find_empty_gap( + pools: list[CubePool], + *, + ego_pose: np.ndarray, + lateral_m: float, + fwd_range: tuple[float, float] = (20.0, 65.0), + lane_halfwidth_m: float = 2.0, + clearance_m: float = 1.5, +) -> tuple[float, float]: + """Center and width of the largest actor-free forward gap on a lateral line.""" + origin, forward, left = _ego_frame(ego_pose) + occupied: list[tuple[float, float]] = [] + for pool in pools: + scales = pool.scales.cpu().numpy() + for track_index, (a, b) in enumerate(_pool_track_slices(pool)): + rel = pool.translations[a].cpu().numpy()[:2] - origin + if abs(float(rel @ left) - lateral_m) > lane_halfwidth_m: + continue + half = float(scales[track_index].max()) / 2 + clearance_m + fwd = float(rel @ forward) + occupied.append((fwd - half, fwd + half)) + occupied.sort() + lo_bound, hi_bound = fwd_range + best_center, best_width = (lo_bound + hi_bound) / 2, 0.0 + cursor = lo_bound + spans = [s for s in occupied if s[1] > lo_bound and s[0] < hi_bound] + for lo, hi in spans + [(hi_bound, hi_bound)]: + width = min(lo, hi_bound) - cursor + if width > best_width: + best_width, best_center = width, cursor + width / 2 + cursor = max(cursor, hi) + return best_center, best_width + + +def clone_template_pool( + placements: list[tuple[TrackTemplate, float, float]], + *, + ego_pose: np.ndarray, +) -> CubePool: + """Merged CubePool of templates rigidly moved to (fwd, lateral) targets. + + Each template keeps its per-frame jitter, orientation, z, dimensions, + colors, and render flags; only its ground-plane position changes. + """ + assert placements, "clone_template_pool needs at least one placement" + origin, forward, left = _ego_frame(ego_pose) + device = placements[0][0].translations.device + track_ts, translations, quaternions, scales, colors, lengths = ( + [], + [], + [], + [], + [], + [], + ) + for template, fwd_m, lateral_m in placements: + target = origin + forward * fwd_m + left * lateral_m + src0 = template.translations[0].cpu().numpy()[:2] + shift = target - src0 + moved = template.translations.clone() + moved[:, 0] += float(shift[0]) + moved[:, 1] += float(shift[1]) + track_ts.append(template.timestamps_us) + translations.append(moved) + quaternions.append(template.quaternions) + scales.append(template.scale) + colors.append(template.colors) + lengths.append(template.timestamps_us.shape[0]) + all_ts = torch.cat(track_ts) + return CubePool( + timestamps_us=torch.unique(all_ts).sort()[0], + cube_ts_prefix_sum=torch.cumsum( + torch.tensor(lengths, dtype=torch.int32, device=device), dim=0 + ).to(torch.int32), + track_timestamps_us=all_ts, + translations=torch.cat(translations), + quaternions=torch.cat(quaternions), + scales=torch.cat(scales), + colors=torch.cat(colors), + prim_type_id=placements[0][0].prim_type_id, + render_flags=placements[0][0].render_flags, + ) diff --git a/integrations/omnidreams/omnidreams/webrtc/session.py b/integrations/omnidreams/omnidreams/webrtc/session.py index 86e4d020b..0ee1d8446 100644 --- a/integrations/omnidreams/omnidreams/webrtc/session.py +++ b/integrations/omnidreams/omnidreams/webrtc/session.py @@ -44,6 +44,16 @@ scenes_cache_root, ) from omnidreams.transformer import CosmosTransformerConfig +from omnidreams.webrtc.actors import ( + ACTOR_PRESETS, + SpawnedActor, + TrackTemplate, + actors_to_cube_pool, + clone_template_pool, + extract_parked_templates, + find_empty_gap, + spawn_actor_ahead, +) from flashdreams.core.distributed.rank_orchestration import ( RankCoordinator, @@ -462,6 +472,17 @@ class OmnidreamsRuntimeConfig: encoder_backend: EncoderBackend = "auto" encoder_bitrate_bps: int = 6_000_000 encoder_gop: int = 30 + # Mid-stream prompt-swap knobs (datachannel ``event`` messages); see + # OmnidreamsConditioningWrapper for semantics. Defaults from the + # 2026-08-08 calibration sweep: s=3 for 6 chunks is the sweet spot + # (edits land convincingly; s=5 causes transition artifacts). + text_edit_guidance_scale: float = 3.0 + text_edit_guidance_chunks: int = 6 + text_edit_recache: bool = True + # Optional guidance-distillation LoRA: edit windows run at guided + # strength through pre-merged weights (single forward per step) instead + # of the two-branch combine. + text_edit_lora_path: Path | None = None @dataclass(frozen=True, slots=True) @@ -512,6 +533,12 @@ def __init__(self, config: OmnidreamsRuntimeConfig | None = None) -> None: self._scene_data: Any | None = None self._initial_rgb_frames: torch.Tensor | None = None self._text_prompts: list[TextPrompt] | None = None + self._initial_prompt: str | None = None + self._active_prompt: str | None = None + self._spawned_actors: list[SpawnedActor] = [] + self._template_placements: list[tuple[TrackTemplate, float, float]] = [] + self._templates_cache: list[TrackTemplate] | None = None + self._last_ego_pose: np.ndarray | None = None self._camera_to_rig: torch.Tensor | None = None self._initial_ego_pose: np.ndarray | None = None self._next_timestamp_us: int = 0 @@ -590,6 +617,31 @@ async def close(self) -> None: finally: self._executor.shutdown(wait=False, cancel_futures=True) + async def trigger_event( + self, *, event_id: str, state: str = "trigger" + ) -> dict[str, str | None]: + """Mid-stream prompt swap driven by datachannel ``event`` messages. + + ``event_id`` carries the free-text prompt verbatim (there is no + fixed event vocabulary — the model takes arbitrary prompts). A + clearing ``state`` (``clear``/``release``/``off``/``none``) or an + empty prompt restores the scene's original prompt. + """ + if self._closed: + raise OmnidreamsRuntimeError("Runtime is closed.") + if self._wrapper is None: + raise OmnidreamsRuntimeError("Runtime is not initialized.") + async with self._step_lock: + if self._closed: + raise OmnidreamsRuntimeError("Runtime is closed.") + if self._wrapper is None: + raise OmnidreamsRuntimeError("Runtime is not initialized.") + return await self._run_on_runtime_thread( + self._trigger_event_sync_all_ranks, + event_id, + state, + ) + async def generate_chunk( self, *, @@ -665,10 +717,202 @@ def _generate_chunk_sync_all_ranks( ) -> WebRTCStepResult: return self._generate_one_chunk_sync(segments=segments, frame_times=frame_times) + @distributed_op(WebRTCControlSignal.EVENT) + def _trigger_event_sync_all_ranks( + self, + event_id: str, + state: str = "trigger", + ) -> dict[str, str | None]: + return self._trigger_event_sync(event_id=event_id, state=state) + @distributed_op(WebRTCControlSignal.CLOSE) def _close_sync_all_ranks(self) -> None: self._close_sync() + _EVENT_CLEAR_STATES = frozenset({"clear", "release", "off", "none"}) + + def _trigger_event_sync( + self, *, event_id: str, state: str + ) -> dict[str, str | None]: + if self._wrapper is None: + raise OmnidreamsRuntimeError("Runtime is not initialized.") + + if event_id.strip().startswith("/"): + return self._handle_actor_command_sync(event_id.strip()) + + prompt = event_id.strip() + if state.strip().lower() in self._EVENT_CLEAR_STATES or not prompt: + if self._initial_prompt is None: + raise OmnidreamsRuntimeError("No scene prompt available to restore.") + prompt = self._initial_prompt + + if prompt == self._active_prompt: + return {"prompt": prompt, "applied": "unchanged"} + + text_prompts = [TextPrompt(positive=prompt)] + if self._state is None or self._state.pipeline_cache is None: + # Rollout has not produced a chunk yet (or HDMap-only debug + # mode): stage the prompt for start_generation instead. + self._text_prompts = text_prompts + self._active_prompt = prompt + return {"prompt": prompt, "applied": "at_start"} + + swap_t0 = time.perf_counter() + self._wrapper.apply_text_prompts(self._state, text_prompts) + self._active_prompt = prompt + logger.info( + "Swapped Omnidreams prompt in {:.0f} ms (chunk={}): {}", + (time.perf_counter() - swap_t0) * 1000.0, + self.autoregressive_index, + prompt, + ) + return {"prompt": prompt, "applied": "immediate"} + + def _spawn_template_sync( + self, args: list[str], command: str + ) -> dict[str, str | None]: + """``/spawnt [lateral_m] [template_idx]``. + + Clones a real parked-vehicle track from the scene (dimensions, + orientation, per-frame jitter, colors) and places it at the target. + ``auto`` picks the largest actor-free forward gap on the lateral + line. Placements are anchored to the session's initial ego pose. + """ + if ( + self._renderer is None + or self._initial_ego_pose is None + or self._scene_data is None + ): + raise OmnidreamsRuntimeError("Scene state is not initialized.") + pools = list(self._renderer._base_timestamped_scene.cube_pools or []) + if self._templates_cache is None: + self._templates_cache = extract_parked_templates( + pools, + ego_pose=self._initial_ego_pose, + t0_us=int(self._scene_data.ego_poses[0].timestamp), + ) + templates = self._templates_cache + if not templates: + raise OmnidreamsRuntimeError("No parked-vehicle templates in this scene.") + + try: + template_idx = int(args[2]) if len(args) > 2 else 0 + template = templates[template_idx % len(templates)] + lateral_m = float(args[1]) if len(args) > 1 else template.source_lateral_m + if not args or args[0].lower() == "auto": + fwd_m, gap = find_empty_gap( + pools, ego_pose=self._initial_ego_pose, lateral_m=lateral_m + ) + if gap < float(template.scale.max()) + 2.0: + raise OmnidreamsRuntimeError( + f"No free gap on lateral {lateral_m:.1f} m " + f"(largest {gap:.1f} m)." + ) + else: + fwd_m = float(args[0]) + except (ValueError, IndexError) as exc: + raise OmnidreamsRuntimeError( + f"Bad arguments in {command!r}: {exc}. Use " + "/spawnt [lateral_m] [template_idx]" + ) from exc + + self._template_placements.append((template, fwd_m, lateral_m)) + logger.info( + "Template-spawned clone (template {} of {}) at fwd {:.1f} m, " + "lateral {:.1f} m; {} active.", + template_idx % len(templates), + len(templates), + fwd_m, + lateral_m, + len(self._template_placements), + ) + return { + "prompt": None, + "applied": ( + f"cloned template at {fwd_m:.1f}m fwd, {lateral_m:.1f}m lateral " + f"({len(self._template_placements)} active)" + ), + } + + def _handle_actor_command_sync(self, command: str) -> dict[str, str | None]: + """``/spawn [dist] [speed] [lateral]`` and ``/clear-actors``. + + Commands share the datachannel ``event`` path with prompt swaps + (anything starting with ``/`` is a command). Spawned actors become + wireframe bboxes in the HDMap conditioning from the next chunk on — + the model materializes an object there; the prompt names its look. + """ + parts = command.removeprefix("/").split() + name = parts[0].lower() if parts else "" + + if name in {"clear-actors", "clear_actors", "despawn", "clear"}: + cleared = len(self._spawned_actors) + len(self._template_placements) + self._spawned_actors.clear() + self._template_placements.clear() + return {"prompt": None, "applied": f"cleared {cleared} actors"} + + if name == "spawnt": + return self._spawn_template_sync(parts[1:], command) + + if name != "spawn": + raise OmnidreamsRuntimeError( + f"Unknown command {command!r}. Use " + "/spawn [dist_m] [speed_mps] [lateral_m] " + f"(presets: {', '.join(sorted(ACTOR_PRESETS))}) or /clear-actors." + ) + + preset = parts[1].lower() if len(parts) > 1 else "car" + if preset not in ACTOR_PRESETS: + raise OmnidreamsRuntimeError( + f"Unknown actor preset {preset!r}; " + f"available: {', '.join(sorted(ACTOR_PRESETS))}." + ) + try: + distance_m = float(parts[2]) if len(parts) > 2 else 12.0 + speed_mps = float(parts[3]) if len(parts) > 3 else 0.0 + lateral_m = float(parts[4]) if len(parts) > 4 else 0.0 + yaw_offset_deg = float(parts[5]) if len(parts) > 5 else 0.0 + except ValueError as exc: + raise OmnidreamsRuntimeError( + f"Non-numeric spawn argument in {command!r}: {exc}" + ) from exc + + ego_pose = ( + self._last_ego_pose + if self._last_ego_pose is not None + else self._initial_ego_pose + ) + if ego_pose is None: + raise OmnidreamsRuntimeError("Scene state is not initialized.") + + actor = spawn_actor_ahead( + preset=preset, + ego_pose=ego_pose, + spawn_timestamp_us=self._next_timestamp_us, + distance_m=distance_m, + speed_mps=speed_mps, + lateral_m=lateral_m, + yaw_offset_deg=yaw_offset_deg, + ) + self._spawned_actors.append(actor) + logger.info( + "Spawned actor {} at {:.1f} m ahead (speed {:.1f} m/s, lateral " + "{:.1f} m); {} active (chunk={}).", + preset, + distance_m, + speed_mps, + lateral_m, + len(self._spawned_actors), + self.autoregressive_index, + ) + return { + "prompt": None, + "applied": ( + f"spawned {preset} {distance_m:g}m ahead" + f" ({len(self._spawned_actors)} active)" + ), + } + def _initialize_sync(self) -> None: if self._wrapper is not None: return @@ -747,7 +991,9 @@ def _initialize_sync(self) -> None: ) prompt = prompt_path.read_text(encoding="utf-8").strip() or AV_POSITIVE_PROMPT + self._initial_prompt = prompt self._text_prompts = [TextPrompt(positive=prompt)] + self._active_prompt = prompt loadable_clipgt_dir = self._prepare_clipgt_dir(clipgt_dir) logger.info("Loading Omnidreams scene data from {}", loadable_clipgt_dir) @@ -797,6 +1043,10 @@ def _initialize_sync(self) -> None: resolution_wh=(cfg.video_width, cfg.video_height), seed_for_every_rollout=cfg.seed, device=self._device, + text_edit_guidance_scale=cfg.text_edit_guidance_scale, + text_edit_guidance_chunks=cfg.text_edit_guidance_chunks, + text_edit_recache=cfg.text_edit_recache, + text_edit_lora_path=cfg.text_edit_lora_path, ) logger.info( "Omnidreams pipeline setup complete in {:.1f}s.", @@ -918,6 +1168,14 @@ def _reset_rollout_sync( self.autoregressive_index = 0 self._next_timestamp_us = int(self._scene_data.ego_poses[0].timestamp) self._wrapper.set_rollout_seed(self.config.seed) + # A new session always starts from the scene's own prompt; mid-stream + # swaps and spawned actors from the previous session must not leak in. + if self._initial_prompt is not None: + self._text_prompts = [TextPrompt(positive=self._initial_prompt)] + self._active_prompt = self._initial_prompt + self._spawned_actors = [] + self._template_placements = [] + self._last_ego_pose = None def _close_sync(self) -> None: state = self._state @@ -928,6 +1186,8 @@ def _close_sync(self) -> None: self._scene_data = None self._initial_rgb_frames = None self._text_prompts = None + self._initial_prompt = None + self._active_prompt = None self._camera_to_rig = None self._initial_ego_pose = None self._close_postprocess_stream() @@ -1025,12 +1285,27 @@ def _generate_one_chunk_sync( ego_poses = self.pose_integrator.integrate_chunk( segments=segments, frame_times=frame_times ) + self._last_ego_pose = ego_poses[-1].copy() ego_poses_t = torch.from_numpy(ego_poses).to( device=self._device, dtype=torch.float32 ) camera_poses = torch.einsum("nij,jk->nik", ego_poses_t, self._camera_to_rig) frame_timestamps_us = self._consume_timestamps(num_frames) + dynamic_actor_pool = None + if self._spawned_actors: + dynamic_actor_pool = actors_to_cube_pool( + self._spawned_actors, frame_timestamps_us, self._device + ) + if self._template_placements: + # Cloned real tracks materialize where synthetic presets do not + # (2026-08-11 finding); template placements take precedence when + # both kinds are active. + assert self._initial_ego_pose is not None + dynamic_actor_pool = clone_template_pool( + self._template_placements, ego_pose=self._initial_ego_pose + ) + camera_names = [self.config.camera_name] camera_poses_per_view = {self.config.camera_name: camera_poses} serve_hdmaps = self.config.debug_serve_hdmaps @@ -1043,6 +1318,7 @@ def _generate_one_chunk_sync( camera_poses_per_view=camera_poses_per_view, frame_timestamps_us=frame_timestamps_us, skip_video_generation=serve_hdmaps, + dynamic_actor_pool=dynamic_actor_pool, ) self._state = output.state else: @@ -1052,6 +1328,7 @@ def _generate_one_chunk_sync( camera_poses_per_view=camera_poses_per_view, frame_timestamps_us=frame_timestamps_us, skip_video_generation=serve_hdmaps, + dynamic_actor_pool=dynamic_actor_pool, ) self._state = output.state diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css b/integrations/omnidreams/omnidreams/webrtc/web/request_session.css index 890c36147..b3047ffd4 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.css +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.css @@ -571,3 +571,79 @@ body[data-status="generating"] .statusLine strong { border-bottom: 0; } } + +.promptCard { + position: absolute; + left: clamp(18px, 3vw, 48px); + bottom: clamp(238px, 30vh, 300px); + width: min(380px, calc(100vw - 36px)); + padding: 18px 20px 20px; +} + +.promptCard h2 { + display: flex; + align-items: center; + gap: 10px; + margin: 0 0 12px; + font-size: 1.08rem; + font-weight: 740; + letter-spacing: 0; +} + +.promptCard h2 span { + width: 3px; + height: 22px; + border-radius: 999px; + background: var(--accent); + box-shadow: 0 0 14px rgba(142, 240, 28, 0.42); +} + +.promptInput { + width: 100%; + box-sizing: border-box; + resize: vertical; + min-height: 58px; + padding: 8px 10px; + border: 1px solid rgba(142, 240, 28, 0.30); + border-radius: 6px; + background: rgba(10, 14, 8, 0.55); + color: var(--text); + font: inherit; + font-size: 0.92rem; +} + +.promptInput:focus { + outline: none; + border-color: rgba(142, 240, 28, 0.6); +} + +.promptButtons { + display: flex; + gap: 8px; + margin-top: 10px; +} + +.promptButton { + flex: 1; + min-height: 32px; + border: 1px solid rgba(142, 240, 28, 0.45); + border-radius: 6px; + background: rgba(142, 240, 28, 0.12); + color: var(--text); + cursor: pointer; + font-weight: 700; +} + +.promptButton:hover { + background: rgba(142, 240, 28, 0.20); +} + +.promptButtonSecondary { + border-color: rgba(255, 255, 255, 0.28); + background: rgba(255, 255, 255, 0.06); + font-weight: 600; +} + +.promptButtonSecondary:hover { + background: rgba(255, 255, 255, 0.12); +} diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html b/integrations/omnidreams/omnidreams/webrtc/web/request_session.html index 263a1b299..8c2f4ffeb 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.html +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.html @@ -58,6 +58,48 @@

Controls

+
+

Scene Prompt

+ +
+ + +
+
+ + + +
+
+

Client Logs

diff --git a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js b/integrations/omnidreams/omnidreams/webrtc/web/request_session.js index 11b7302e8..4d22d4710 100644 --- a/integrations/omnidreams/omnidreams/webrtc/web/request_session.js +++ b/integrations/omnidreams/omnidreams/webrtc/web/request_session.js @@ -16,6 +16,12 @@ const modelValue = document.getElementById("modelValue") const postprocessField = document.getElementById("postprocessField") const postprocessSelect = document.getElementById("postprocessSelect") const controlButtons = Array.from(document.querySelectorAll("[data-control-key]")) +const promptInput = document.getElementById("promptInput") +const promptApplyButton = document.getElementById("promptApplyButton") +const promptResetButton = document.getElementById("promptResetButton") +const spawnCarButton = document.getElementById("spawnCarButton") +const spawnConeButton = document.getElementById("spawnConeButton") +const clearActorsButton = document.getElementById("clearActorsButton") const allowedKeys = new Set(["w", "a", "s", "d"]) const keyAliases = new Map([ @@ -437,6 +443,36 @@ function enqueueAction(action) { } } +function sendPromptEvent(prompt, state) { + if (!connected || !controlChannel || controlChannel.readyState !== "open") { + logEvent("prompt not sent: connect session first", { level: "error" }) + return false + } + controlChannel.send( + JSON.stringify({ + type: "event", + event_id: prompt, + state, + }) + ) + logEvent( + state === "trigger" ? `prompt sent: ${prompt}` : "prompt reset to scene default", + { source: "client" } + ) + return true +} + +function applyPromptFromInput() { + const prompt = (promptInput.value || "").trim() + if (!prompt) { + logEvent("prompt is empty; use Reset to restore the scene prompt", { + level: "error", + }) + return + } + sendPromptEvent(prompt, "trigger") +} + function enqueueHeldKeyRepeats() { const heldKeys = Array.from(activeKeys).sort((a, b) => { return (heldKeyOrder.get(a) || 0) - (heldKeyOrder.get(b) || 0) @@ -533,6 +569,13 @@ function handleControlMessage(rawMessage) { return } + if (payload.type === "event_ack") { + const applied = payload.applied || "ok" + const promptText = payload.prompt ? `: ${payload.prompt}` : "" + logEvent(`prompt ${applied}${promptText}`) + return + } + if (payload.type === "server_log") { logEvent(payload.message || "server log") return @@ -843,7 +886,19 @@ async function connectSession() { } } +function isTextEntryTarget(event) { + const target = event.target + if (!target) { + return false + } + const tag = String(target.tagName || "").toLowerCase() + return tag === "textarea" || tag === "input" || target.isContentEditable === true +} + function handleKeyDown(event) { + if (isTextEntryTarget(event)) { + return + } const key = normalizeKey(event.key) if (!allowedKeys.has(key)) { return @@ -857,6 +912,9 @@ function handleKeyDown(event) { } function handleKeyUp(event) { + if (isTextEntryTarget(event)) { + return + } const key = normalizeKey(event.key) if (!allowedKeys.has(key)) { return @@ -933,6 +991,26 @@ remoteVideo.addEventListener("playing", () => { remoteVideo.addEventListener("emptied", () => { setVideoVisible(false) }) +promptApplyButton.addEventListener("click", applyPromptFromInput) +promptResetButton.addEventListener("click", () => { + sendPromptEvent("", "clear") +}) +spawnCarButton.addEventListener("click", () => { + sendPromptEvent("/spawn car 12", "trigger") +}) +spawnConeButton.addEventListener("click", () => { + sendPromptEvent("/spawn cone 8", "trigger") +}) +clearActorsButton.addEventListener("click", () => { + sendPromptEvent("/clear-actors", "trigger") +}) +promptInput.addEventListener("keydown", (event) => { + if (event.key === "Enter" && (event.ctrlKey || event.metaKey)) { + event.preventDefault() + applyPromptFromInput() + } +}) + window.addEventListener("keydown", handleKeyDown) window.addEventListener("keyup", handleKeyUp) window.addEventListener("blur", releaseAllKeys) diff --git a/integrations/omnidreams/scripts/probe_moving_clone.py b/integrations/omnidreams/scripts/probe_moving_clone.py new file mode 100644 index 000000000..66bb850dd --- /dev/null +++ b/integrations/omnidreams/scripts/probe_moving_clone.py @@ -0,0 +1,244 @@ +# 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. + +"""Probe: does a cloned MOVING track materialize more solidly than a parked one? + +Parked-template clones materialize at only ~20% render strength (box +darkening 9-10 uint8 vs ~52 for a scene-native car). Hypothesis: a parked +clone is pinned against contradicting history at one spot (the first-frame +photo and every prior render show that curb empty), while recorded moving +traffic materializes solidly from boxes alone — the object is never at odds +with the same pixels for long. This probe clones a real MOVING car track +(drift >= ``MIN_DRIFT_M``) and rigidly shifts it along the ego heading so it +repeats its recorded motion offset in space. + +Env knobs: ``SHIFT_FWD_M`` (default 25), ``MTEMPLATE_IDX`` (default 0 = +nearest at window start), ``MIN_DRIFT_M`` (default 15), ``GUIDE_SCALE`` +(default 0 = off; >0 adds the box-axis guidance combine over the clone, +same recipe as ``probe_spawn_guidance``), ``HDMAP_ONLY``, ``N_CHUNKS``, +``OUT_DIR``. + +Run from the repo root (venv bin on PATH for the Ludus ninja build):: + + HDMAP_ONLY=1 OUT_DIR=.../mclone_hdmap python probe_moving_clone.py + OUT_DIR=.../mclone python probe_moving_clone.py +""" + +from __future__ import annotations + +import os +from pathlib import Path + +# Must land before the first CUDA allocation (co-tenant VRAM share). +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import numpy as np +import torch +from omnidreams.config import ( + OMNIDREAMS_CONFIGS, + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, +) +from omnidreams.webrtc.actors import ( + TrackTemplate, + _ego_frame, + _pool_track_slices, + clone_template_pool, +) + +from flashdreams.infra.config import derive_config + +_EAGER_NAME = "omnidreams-sv-2steps-chunk2-moving-clone-probe" +OMNIDREAMS_CONFIGS[_EAGER_NAME] = derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, + name=_EAGER_NAME, + enable_sync_and_profile=False, + diffusion_model=dict( + seed=42, + transformer=dict(compile_network=False, use_cuda_graph=False), + ), +) + +from flashdreams.infra.runner_io import write_video_tensor # noqa: E402 +from omnidreams.webrtc import session as webrtc_session # noqa: E402 +from omnidreams.webrtc.session import ( # noqa: E402 + OmnidreamsInferenceRuntime, + OmnidreamsRuntimeConfig, +) + + +FPS = 30 +N_CHUNKS = int(os.environ.get("N_CHUNKS", "26")) +SHIFT_FWD_M = float(os.environ.get("SHIFT_FWD_M", "25")) +MTEMPLATE_IDX = int(os.environ.get("MTEMPLATE_IDX", "0")) +MIN_DRIFT_M = float(os.environ.get("MIN_DRIFT_M", "15")) +GUIDE_SCALE = float(os.environ.get("GUIDE_SCALE", "0")) +HDMAP_ONLY = os.environ.get("HDMAP_ONLY", "0") == "1" +OUT_DIR = Path( + os.environ.get("OUT_DIR", "integrations/omnidreams/scripts/outputs/mclone") +) + + +def _extract_moving_templates( + pools, *, ego_pose: np.ndarray, t0_us: int +) -> list[TrackTemplate]: + """Car-sized tracks that MOVE >= MIN_DRIFT_M and cover the window.""" + origin, forward, left = _ego_frame(ego_pose) + out: list[tuple[float, TrackTemplate]] = [] + for pool in pools: + scales = pool.scales.cpu().numpy() + for track_index, (a, b) in enumerate(_pool_track_slices(pool)): + ts = pool.track_timestamps_us[a:b].cpu().numpy() + if len(ts) < 8 or ts[0] > t0_us + 1_000_000 or ts[-1] < t0_us + 5_500_000: + continue + length = float(scales[track_index].max()) + if not 3.4 <= length <= 5.6: + continue + tr = pool.translations[a:b].cpu().numpy() + drift = float(np.linalg.norm(tr[-1, :2] - tr[0, :2])) + if drift < MIN_DRIFT_M: + continue + rel = tr[0, :2] - origin + template = TrackTemplate( + timestamps_us=pool.track_timestamps_us[a:b].clone(), + translations=pool.translations[a:b].clone(), + quaternions=pool.quaternions[a:b].clone(), + scale=pool.scales[track_index : track_index + 1].clone(), + colors=pool.colors[track_index : track_index + 1].clone(), + prim_type_id=pool.prim_type_id, + render_flags=pool.render_flags, + source_fwd_m=float(rel @ forward), + source_lateral_m=float(rel @ left), + ) + out.append((float(np.linalg.norm(rel)), template)) + out.sort(key=lambda item: item[0]) + return [t for _, t in out] + + +def main() -> None: + config = OmnidreamsRuntimeConfig( + pipeline_config_name=_EAGER_NAME, debug_serve_hdmaps=HDMAP_ONLY + ) + runtime = OmnidreamsInferenceRuntime(config) + print("initializing runtime (scene + pipeline)...", flush=True) + runtime._initialize_sync() + + renderer = runtime._renderer + assert renderer is not None + pools = list(renderer._base_timestamped_scene.cube_pools or []) + ego0 = runtime._initial_ego_pose + assert ego0 is not None + assert runtime._scene_data is not None + t0_us = int(runtime._scene_data.ego_poses[0].timestamp) + origin, forward, left = _ego_frame(ego0) + + templates = _extract_moving_templates(pools, ego_pose=ego0, t0_us=t0_us) + assert templates, "no moving car-sized tracks cover the rollout window" + for i, tpl in enumerate(templates[:6]): + tr = tpl.translations.cpu().numpy() + rel_end = tr[-1, :2] - origin + print( + f"moving template {i}: start fwd {tpl.source_fwd_m:.1f} " + f"lat {tpl.source_lateral_m:.1f} -> end fwd " + f"{float(rel_end @ forward):.1f} lat {float(rel_end @ left):.1f} " + f"({tr.shape[0]} samples, len {float(tpl.scale.max()):.1f} m)", + flush=True, + ) + template = templates[MTEMPLATE_IDX % len(templates)] + + # Rigid shift along the ego heading: same recorded motion, offset start. + target_fwd = template.source_fwd_m + SHIFT_FWD_M + clone = clone_template_pool( + [(template, target_fwd, template.source_lateral_m)], ego_pose=ego0 + ) + print( + f"cloned moving template {MTEMPLATE_IDX % len(templates)} shifted " + f"+{SHIFT_FWD_M} m fwd: starts fwd {target_fwd:.1f} m, " + f"lat {template.source_lateral_m:.1f} m", + flush=True, + ) + + # Route through the standard overlay path: sentinel actor list so the + # session builds a pool, patched builder returns the clone. + runtime._spawned_actors = [object()] # ty: ignore[invalid-assignment] + webrtc_session.actors_to_cube_pool = ( # ty: ignore[invalid-assignment] + lambda actors, ts, device: clone + ) + + if GUIDE_SCALE > 0 and not HDMAP_ONLY: + wrapper = runtime._wrapper + assert wrapper is not None + pipe = wrapper.pipeline + transformer = pipe.diffusion_model.transformer + encoder = pipe.encoder + assert encoder is not None + shadow_encoder_cache = encoder.initialize_autoregressive_cache() + state: dict = {"alt_input": None, "ar_idx": 0} + orig_render = wrapper._render_condition_frames + + def dual_render(renderer, camera_names, poses, timestamps, pool=None): + frames_box = orig_render(renderer, camera_names, poses, timestamps, pool) + frames_nobox = orig_render(renderer, camera_names, poses, timestamps, None) + with torch.no_grad(): + model_in = wrapper._to_model_range( + wrapper._normalize_condition_input(frames_nobox) + ) + encoded = encoder( + input=model_in, + autoregressive_index=state["ar_idx"], + cache=shadow_encoder_cache, + ) + state["alt_input"] = transformer.patchify_and_maybe_split_cp(encoded) + state["ar_idx"] += 1 + return frames_box + + wrapper._render_condition_frames = dual_render # ty: ignore[invalid-assignment] + orig_pf = transformer.predict_flow + + def guided_pf(noisy_latent, timestep, cache, input=None): + if transformer._finalizing_kv_cache or state["alt_input"] is None: + return orig_pf(noisy_latent, timestep, cache, input=input) + flow_box = orig_pf(noisy_latent, timestep, cache, input=input) + flow_nobox = orig_pf( + noisy_latent, timestep, cache, input=state["alt_input"] + ) + return flow_nobox + GUIDE_SCALE * (flow_box - flow_nobox) + + transformer.predict_flow = guided_pf # ty: ignore[invalid-assignment] + print(f"box-axis guidance active at s={GUIDE_SCALE}", flush=True) + + chunks: list[torch.Tensor] = [] + t = 0.0 + for ar_idx in range(N_CHUNKS): + num_frames = runtime.peek_next_chunk_num_frames() + t_end = t + num_frames / FPS + segments = [(t, t_end, frozenset({"w"}))] + frame_times = [t + i / FPS for i in range(num_frames)] + result = runtime._generate_one_chunk_sync( + segments=segments, frame_times=frame_times + ) + chunks.append(result.video_chunk[0, 0]) + t = t_end + if ar_idx % 4 == 0: + print(f"chunk {ar_idx} done", flush=True) + + video = torch.cat(chunks, dim=0).float() / 127.5 - 1.0 + OUT_DIR.mkdir(parents=True, exist_ok=True) + name = "hdmap.mp4" if HDMAP_ONLY else "drive.mp4" + write_video_tensor(video, OUT_DIR / name, fps=FPS, layout="tchw") + print(f"{video.shape[0]} frames -> {OUT_DIR / name}") + + +if __name__ == "__main__": + main() diff --git a/integrations/omnidreams/scripts/probe_pedestrians.py b/integrations/omnidreams/scripts/probe_pedestrians.py new file mode 100644 index 000000000..4e168602e --- /dev/null +++ b/integrations/omnidreams/scripts/probe_pedestrians.py @@ -0,0 +1,212 @@ +# 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. + +"""Probe: pedestrian-clone density ladder (1 / 5 / 20 ahead of the ego). + +Car-sized template clones materialize (ghost-to-solid depending on motion +and guidance); this probe asks whether PEDESTRIAN tracks cloned from the +scene's own pools materialize at all, and where the density ceiling is — +a dense crowd on a residential road is far off the AV training manifold, +so the expectation is a few near materializations and mush beyond. + +Placements form a deterministic grid ahead of the ego (rows every 6 m +from ``FWD0``, lateral spread across the lane), cycling the available +pedestrian templates. Masks come from the HDMAP arm as usual. + +Env knobs: ``PED_COUNT`` (default 5), ``FWD0`` (default 22), ``HDMAP_ONLY``, +``N_CHUNKS``, ``OUT_DIR``, ``DEBUG_TRACKS``. + +Run from the repo root (venv bin on PATH for the Ludus ninja build):: + + HDMAP_ONLY=1 PED_COUNT=5 OUT_DIR=.../ped5_hdmap python probe_pedestrians.py + PED_COUNT=5 OUT_DIR=.../ped5 python probe_pedestrians.py +""" + +from __future__ import annotations + +import os +from pathlib import Path + +# Must land before the first CUDA allocation (co-tenant VRAM share). +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import numpy as np +import torch +from omnidreams.config import ( + OMNIDREAMS_CONFIGS, + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, +) +from omnidreams.webrtc.actors import ( + TrackTemplate, + _ego_frame, + _pool_track_slices, + clone_template_pool, +) + +from flashdreams.infra.config import derive_config + +_EAGER_NAME = "omnidreams-sv-2steps-chunk2-ped-probe" +OMNIDREAMS_CONFIGS[_EAGER_NAME] = derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, + name=_EAGER_NAME, + enable_sync_and_profile=False, + diffusion_model=dict( + seed=42, + transformer=dict(compile_network=False, use_cuda_graph=False), + ), +) + +from flashdreams.infra.runner_io import write_video_tensor # noqa: E402 +from omnidreams.webrtc import session as webrtc_session # noqa: E402 +from omnidreams.webrtc.session import ( # noqa: E402 + OmnidreamsInferenceRuntime, + OmnidreamsRuntimeConfig, +) + + +FPS = 30 +N_CHUNKS = int(os.environ.get("N_CHUNKS", "26")) +PED_COUNT = int(os.environ.get("PED_COUNT", "5")) +FWD0 = float(os.environ.get("FWD0", "22")) +HDMAP_ONLY = os.environ.get("HDMAP_ONLY", "0") == "1" +DEBUG_TRACKS = os.environ.get("DEBUG_TRACKS", "0") == "1" +OUT_DIR = Path( + os.environ.get("OUT_DIR", "integrations/omnidreams/scripts/outputs/ped_probe") +) + +_LATERALS = tuple( + float(x) for x in os.environ.get("PED_LATERALS", "-3.2,-1.1,1.1,3.2").split(",") +) +"""Lateral columns of the placement grid (default: across the ego lane; +override with PED_LATERALS for sidewalk-band placements).""" + + +def _extract_ped_templates( + pools, *, ego_pose: np.ndarray, t0_us: int +) -> list[TrackTemplate]: + """Pedestrian-sized tracks covering the early window (walking allowed).""" + origin, forward, left = _ego_frame(ego_pose) + out: list[tuple[float, TrackTemplate]] = [] + for pool_index, pool in enumerate(pools): + scales = pool.scales.cpu().numpy() + for track_index, (a, b) in enumerate(_pool_track_slices(pool)): + ts = pool.track_timestamps_us[a:b].cpu().numpy() + length = float(scales[track_index].max()) + if DEBUG_TRACKS and length <= 2.2: + tr0 = pool.translations[a].cpu().numpy()[:2] - origin + print( + f"pool{pool_index} track{track_index}: n={b - a} " + f"len={length:.2f} t=[{(ts[0] - t0_us) / 1e6:.1f}," + f"{(ts[-1] - t0_us) / 1e6:.1f}]s " + f"fwd={float(tr0 @ forward):.1f} lat={float(tr0 @ left):.1f}", + flush=True, + ) + if len(ts) < 6 or ts[0] > t0_us + 1_500_000 or ts[-1] < t0_us + 3_500_000: + continue + # Person-sized: tallest dim is the ~1.7 m height (scale.max() + # is NOT the footprint), the other dims are sub-metre. + dims = np.sort(scales[track_index]) + if not (1.2 <= dims[-1] <= 2.1 and dims[-2] <= 1.2): + continue + rel = pool.translations[a].cpu().numpy()[:2] - origin + template = TrackTemplate( + timestamps_us=pool.track_timestamps_us[a:b].clone(), + translations=pool.translations[a:b].clone(), + quaternions=pool.quaternions[a:b].clone(), + scale=pool.scales[track_index : track_index + 1].clone(), + colors=pool.colors[track_index : track_index + 1].clone(), + prim_type_id=pool.prim_type_id, + render_flags=pool.render_flags, + source_fwd_m=float(rel @ forward), + source_lateral_m=float(rel @ left), + ) + out.append((float(np.linalg.norm(rel)), template)) + out.sort(key=lambda item: item[0]) + return [t for _, t in out] + + +def main() -> None: + config = OmnidreamsRuntimeConfig( + pipeline_config_name=_EAGER_NAME, debug_serve_hdmaps=HDMAP_ONLY + ) + runtime = OmnidreamsInferenceRuntime(config) + print("initializing runtime (scene + pipeline)...", flush=True) + runtime._initialize_sync() + + renderer = runtime._renderer + assert renderer is not None and runtime._initial_ego_pose is not None + pools = list(renderer._base_timestamped_scene.cube_pools or []) + ego0 = runtime._initial_ego_pose + assert runtime._scene_data is not None + t0_us = int(runtime._scene_data.ego_poses[0].timestamp) + + templates = _extract_ped_templates(pools, ego_pose=ego0, t0_us=t0_us) + assert templates, "no pedestrian-sized tracks cover the window (try DEBUG_TRACKS=1)" + print( + f"{len(templates)} pedestrian templates; nearest at " + f"fwd {templates[0].source_fwd_m:.1f} lat {templates[0].source_lateral_m:.1f}", + flush=True, + ) + + # Deterministic grid: rows every 6 m, columns across the lane. + placements = [] + i = 0 + while len(placements) < PED_COUNT: + fwd = FWD0 + 6.0 * (i // len(_LATERALS)) + lateral = _LATERALS[i % len(_LATERALS)] + placements.append((templates[i % len(templates)], fwd, lateral)) + i += 1 + clone = clone_template_pool(placements, ego_pose=ego0) + print(f"placed {len(placements)} pedestrian clones from fwd {FWD0} m", flush=True) + + runtime._spawned_actors = [object()] # ty: ignore[invalid-assignment] + webrtc_session.actors_to_cube_pool = ( # ty: ignore[invalid-assignment] + lambda actors, ts, device: clone + ) + + if os.environ.get("EDIT_PROMPT"): + # Scene-class synergy: align the text channel with the box channel + # (e.g. a street-festival prompt to make mid-road crowds plausible). + print( + runtime._trigger_event_sync( + event_id=os.environ["EDIT_PROMPT"], state="trigger" + ), + flush=True, + ) + + chunks: list[torch.Tensor] = [] + t = 0.0 + for ar_idx in range(N_CHUNKS): + num_frames = runtime.peek_next_chunk_num_frames() + t_end = t + num_frames / FPS + segments = [(t, t_end, frozenset({"w"}))] + frame_times = [t + i / FPS for i in range(num_frames)] + result = runtime._generate_one_chunk_sync( + segments=segments, frame_times=frame_times + ) + chunks.append(result.video_chunk[0, 0]) + t = t_end + if ar_idx % 4 == 0: + print(f"chunk {ar_idx} done", flush=True) + + video = torch.cat(chunks, dim=0).float() / 127.5 - 1.0 + OUT_DIR.mkdir(parents=True, exist_ok=True) + name = "hdmap.mp4" if HDMAP_ONLY else "drive.mp4" + write_video_tensor(video, OUT_DIR / name, fps=FPS, layout="tchw") + print(f"{video.shape[0]} frames -> {OUT_DIR / name}") + + +if __name__ == "__main__": + main() diff --git a/integrations/omnidreams/scripts/smoke_spawn_actor.py b/integrations/omnidreams/scripts/smoke_spawn_actor.py new file mode 100644 index 000000000..5bcee4506 --- /dev/null +++ b/integrations/omnidreams/scripts/smoke_spawn_actor.py @@ -0,0 +1,133 @@ +# 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. + +"""Smoke test: user-spawned actors in a headless WebRTC-runtime drive. + +Drives the Omnidreams WebRTC runtime synchronously (no browser, no +networking): hold W, spawn a car ahead mid-drive via the same +``/spawn`` command the datachannel uses, spawn a cone later, and save the +rollout. Verifies the full chain scene -> Ludus bbox render -> HDMap +conditioning -> model materializes an object. + +Env knobs: ``N_CHUNKS``, ``SPAWN_AT``, ``SPAWN_CMD``, ``SPAWN2_AT``, +``SPAWN2_CMD``, ``EDIT_PROMPT`` (optional prompt swap alongside the first +spawn), ``HDMAP_ONLY=1`` (skip the model, save the rendered conditioning — +fast check that the bbox actually lands in the HDMap stream), ``OUT_DIR``. + +Run from the repo root:: + + .venv/bin/python integrations/omnidreams/scripts/smoke_spawn_actor.py +""" + +from __future__ import annotations + +import os +from pathlib import Path + +# Must land before the first CUDA allocation (co-tenant VRAM share). +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import torch +from omnidreams.config import ( + OMNIDREAMS_CONFIGS, + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, +) + +from flashdreams.infra.config import derive_config +from flashdreams.infra.runner_io import write_video_tensor + +# Register an eager variant before the runtime resolves the name: probing +# scripts skip compile / CUDA graphs to trade steady-state latency for +# startup time. +_EAGER_NAME = "omnidreams-sv-2steps-chunk2-smoke-eager" +OMNIDREAMS_CONFIGS[_EAGER_NAME] = derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, + name=_EAGER_NAME, + enable_sync_and_profile=False, + diffusion_model=dict( + seed=42, + transformer=dict(compile_network=False, use_cuda_graph=False), + ), +) + +from omnidreams.webrtc.session import ( # noqa: E402 (needs the config registered) + OmnidreamsInferenceRuntime, + OmnidreamsRuntimeConfig, +) + +FPS = 30 +N_CHUNKS = int(os.environ.get("N_CHUNKS", "24")) +SPAWN_AT = int(os.environ.get("SPAWN_AT", "6")) +SPAWN_CMD = os.environ.get("SPAWN_CMD", "/spawn car 16 0 0") +SPAWN2_AT = int(os.environ.get("SPAWN2_AT", "14")) +SPAWN2_CMD = os.environ.get("SPAWN2_CMD", "/spawn cone 10 0 -2") +EDIT_PROMPT = os.environ.get("EDIT_PROMPT", "") +HDMAP_ONLY = os.environ.get("HDMAP_ONLY", "0") == "1" +OUT_DIR = Path( + os.environ.get("OUT_DIR", "integrations/omnidreams/scripts/outputs/spawn_smoke") +) + + +def main() -> None: + config = OmnidreamsRuntimeConfig( + pipeline_config_name=_EAGER_NAME, + debug_serve_hdmaps=HDMAP_ONLY, + ) + runtime = OmnidreamsInferenceRuntime(config) + print("initializing runtime (scene + pipeline)...", flush=True) + runtime._initialize_sync() + + chunks: list[torch.Tensor] = [] + t = 0.0 + for ar_idx in range(N_CHUNKS): + if ar_idx == SPAWN_AT: + for command in SPAWN_CMD.split(";"): + print( + runtime._trigger_event_sync( + event_id=command.strip(), state="trigger" + ) + ) + if EDIT_PROMPT: + print( + runtime._trigger_event_sync(event_id=EDIT_PROMPT, state="trigger") + ) + if ar_idx == SPAWN2_AT: + print(runtime._trigger_event_sync(event_id=SPAWN2_CMD, state="trigger")) + + num_frames = runtime.peek_next_chunk_num_frames() + t_end = t + num_frames / FPS + segments = [(t, t_end, frozenset({"w"}))] # hold W: drive forward + frame_times = [t + i / FPS for i in range(num_frames)] + result = runtime._generate_one_chunk_sync( + segments=segments, frame_times=frame_times + ) + chunks.append(result.video_chunk[0, 0]) # [T, 3, H, W] uint8 + t = t_end + if ar_idx % 4 == 0: + print(f"chunk {ar_idx} done", flush=True) + + video = torch.cat(chunks, dim=0).float() / 127.5 - 1.0 + OUT_DIR.mkdir(parents=True, exist_ok=True) + name = "hdmap.mp4" if HDMAP_ONLY else "drive.mp4" + write_video_tensor(video, OUT_DIR / name, fps=FPS, layout="tchw") + print( + f"{video.shape[0]} frames -> {OUT_DIR / name} " + f"(spawn at chunk {SPAWN_AT}: {SPAWN_CMD!r}; " + f"chunk {SPAWN2_AT}: {SPAWN2_CMD!r})" + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/omnidreams/scripts/smoke_text_edit.py b/integrations/omnidreams/scripts/smoke_text_edit.py new file mode 100644 index 000000000..8a50c23bf --- /dev/null +++ b/integrations/omnidreams/scripts/smoke_text_edit.py @@ -0,0 +1,266 @@ +# 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. + +"""Smoke test: mid-stream prompt swap on the real distilled model. + +Rolls the same seed / HDMap / first frame several ways and reports how +strongly the video diverges after the swap chunk: + + A control original clip prompt throughout + B swap hot-swap to ``EDIT_PROMPT`` at chunk ``SWAP_AT`` + C swap+guide same swap with two-prompt edit guidance + D swap+recache same swap plus previous-chunk KV re-commit + +B and C consume the identical RNG stream as A (the swap itself draws no +noise), so the per-chunk ``|B - A|`` pixel gap is a pure measure of prompt +responsiveness: ~0 before the swap (sanity check), and the post-swap +magnitude/growth is the signal. D draws one extra context-noise sample at +the recache, so its pre-swap sanity still holds but its post-swap gap is +noise-shifted — judge D visually against B. + +Env knobs: ``UUID``, ``EDIT_PROMPT``, ``N_CHUNKS``, ``SWAP_AT``, +``GUIDE_SCALE``, ``GUIDE_CHUNKS``, ``SEED``, ``OUT_DIR``. + +Run from the repo root:: + + .venv/bin/python integrations/omnidreams/scripts/smoke_text_edit.py +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +# Must land before the first CUDA allocation (co-tenant VRAM share). +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import torch +from omnidreams.config import SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE +from omnidreams.pipeline import OmnidreamsPipeline +from omnidreams.runner import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH +from torch import Tensor + +from flashdreams.infra.config import derive_config +from flashdreams.infra.runner_io import ( + load_first_frame_tensor, + load_video_tensor, + write_video_tensor, +) + +SAMPLES_ROOT = ( + Path.home() + / ".cache/huggingface/hub/datasets--nvidia--omni-dreams-samples/snapshots" +) + +UUID = os.environ.get("UUID", "23599139-948f-4681-b7f4-74794113086d") +N_CHUNKS = int(os.environ.get("N_CHUNKS", "16")) +SWAP_AT = int(os.environ.get("SWAP_AT", "8")) +GUIDE_SCALE = float(os.environ.get("GUIDE_SCALE", "2.5")) +GUIDE_CHUNKS = int(os.environ.get("GUIDE_CHUNKS", "4")) +SEED = int(os.environ.get("SEED", "42")) +OUT_DIR = Path( + os.environ.get("OUT_DIR", "integrations/omnidreams/scripts/outputs/text_edit_smoke") +) +EDIT_PROMPT = os.environ.get( + "EDIT_PROMPT", + "Driving scene from a front-facing car camera at night in a heavy " + "snowstorm. Thick snow falling, snow-covered road and buildings, " + "headlights and streetlights glowing through the snow. Photorealistic " + "dashcam footage.", +) + + +def _sample_paths(uuid: str) -> tuple[Path, Path, str]: + hdmaps = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/*_hdmap.mp4")) + frames = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/first_frame.png")) + prompts = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/prompt.txt")) + assert hdmaps and frames and prompts, ( + f"sample {uuid} not in the local HF cache under {SAMPLES_ROOT}" + ) + return hdmaps[0], frames[0], prompts[0].read_text().strip() + + +def _build_pipeline() -> OmnidreamsPipeline: + cfg = derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, + enable_sync_and_profile=False, + diffusion_model=dict( + seed=SEED, + transformer=dict(compile_network=False, use_cuda_graph=False), + ), + ) + pipe = cfg.setup() + assert isinstance(pipe, OmnidreamsPipeline) + pipe = pipe.to("cuda") + # EDIT_LORA=: deploy the pre-merged guidance-distillation LoRA, so + # the guided variants exercise the production use_lora window instead of + # the two-branch combine. + if os.environ.get("EDIT_LORA"): + from omnidreams._edit_lora import TextEditLoRA + + transformer = pipe.diffusion_model.transformer + edit_lora = TextEditLoRA(transformer.network, os.environ["EDIT_LORA"]) + transformer.set_text_edit_lora(edit_lora) + print(f"deployed {edit_lora.describe()}", flush=True) + return pipe + + +@torch.no_grad() +def _rollout( + pipe: OmnidreamsPipeline, + *, + hdmap: Tensor, + first: Tensor, + base_prompt: str, + swap: dict | None = None, +) -> Tensor: + """Return the decoded rollout ``[T, 3, H, W]`` in ``[-1, 1]`` on CPU.""" + device = pipe.device + pipe.diffusion_model._rng = torch.Generator(device=device).manual_seed(SEED) + cache = pipe.initialize_cache(text=[[base_prompt]], image=first) + chunks: list[Tensor] = [] + start = 0 + for ar_idx in range(N_CHUNKS): + repulse = int(os.environ.get("REPULSE_EVERY", "0")) + if ( + swap is not None + and repulse > 0 + and ar_idx > swap["at"] + and (ar_idx - swap["at"]) % repulse == 0 + ): + # Re-open the edit window before the previous one's style fades: + # duty-cycled skin for LoRAs whose long-hold drifts. + pipe.replace_text( + cache, + [[swap["prompt"]]], + guidance_scale=swap.get("scale", 1.0), + guidance_chunks=swap.get("chunks", 0), + recache_last_chunk=False, + ) + if swap is not None and ar_idx == swap["at"]: + pipe.replace_text( + cache, + [[swap["prompt"]]], + guidance_scale=swap.get("scale", 1.0), + guidance_chunks=swap.get("chunks", 0), + recache_last_chunk=swap.get("recache", False), + ) + num_frames = pipe.get_num_frames(ar_idx) + end = start + num_frames + assert end <= hdmap.shape[2], f"hdmap too short at chunk {ar_idx}" + chunk = pipe.generate(ar_idx, cache, hdmap=hdmap[:, :, start:end]) + pipe.finalize(ar_idx, cache) + chunks.append(chunk[0, 0].float().cpu()) + start = end + del cache + torch.cuda.empty_cache() + return torch.cat(chunks, dim=0) + + +def _chunk_bounds() -> list[tuple[int, int]]: + bounds, start = [], 0 + for ar_idx in range(N_CHUNKS): + n = 5 if ar_idx == 0 else 8 + bounds.append((start, start + n)) + start += n + return bounds + + +def _per_chunk_gap(a: Tensor, b: Tensor) -> list[float]: + """Mean |a - b| per chunk in uint8 units (0..255).""" + return [float((a[s:e] - b[s:e]).abs().mean() * 127.5) for s, e in _chunk_bounds()] + + +def main() -> None: + hdmap_path, frame_path, clip_prompt = _sample_paths(UUID) + total_frames = 5 + (N_CHUNKS - 1) * 8 + print(f"clip {UUID}\n prompt: {clip_prompt}\n edit: {EDIT_PROMPT}") + print(f" chunks={N_CHUNKS} swap_at={SWAP_AT} frames={total_frames}") + + device = torch.device("cuda") + hdmap = load_video_tensor( + hdmap_path, + pixel_height=DEFAULT_VIDEO_HEIGHT, + pixel_width=DEFAULT_VIDEO_WIDTH, + device=device, + dtype=torch.bfloat16, + )[:total_frames][None, None] + first = load_first_frame_tensor( + frame_path, + pixel_height=DEFAULT_VIDEO_HEIGHT, + pixel_width=DEFAULT_VIDEO_WIDTH, + device=device, + dtype=torch.bfloat16, + )[None, None] # [B=1, V=1, 1, C, H, W] + + pipe = _build_pipeline() + + variants: dict[str, dict | None] = { + "control": None, + "swap": {"at": SWAP_AT, "prompt": EDIT_PROMPT}, + "swap_guided": { + "at": SWAP_AT, + "prompt": EDIT_PROMPT, + "scale": GUIDE_SCALE, + "chunks": GUIDE_CHUNKS, + }, + "swap_recache": {"at": SWAP_AT, "prompt": EDIT_PROMPT, "recache": True}, + } + + OUT_DIR.mkdir(parents=True, exist_ok=True) + videos: dict[str, Tensor] = {} + for name, swap in variants.items(): + print(f"rolling out {name} ...", flush=True) + videos[name] = _rollout( + pipe, hdmap=hdmap, first=first, base_prompt=clip_prompt, swap=swap + ) + write_video_tensor(videos[name], OUT_DIR / f"{name}.mp4", fps=30, layout="tchw") + + control = videos["control"] + report: dict[str, list[float]] = {} + for name in ("swap", "swap_guided", "swap_recache"): + gaps = _per_chunk_gap(videos[name], control) + report[name] = gaps + pre = max(gaps[:SWAP_AT]) + post = gaps[SWAP_AT:] + print( + f"{name:>13}: pre-swap max gap {pre:6.3f} " + f"post-swap per-chunk {' '.join(f'{g:6.2f}' for g in post)}" + ) + + # Side-by-side [control | swap | swap_guided] for eyeballing. + sbs = torch.cat( + [control, videos["swap"], videos["swap_guided"]], dim=3 + ) # widths concat + write_video_tensor(sbs, OUT_DIR / "sbs.mp4", fps=30, layout="tchw") + + meta = { + "uuid": UUID, + "clip_prompt": clip_prompt, + "edit_prompt": EDIT_PROMPT, + "n_chunks": N_CHUNKS, + "swap_at": SWAP_AT, + "guide_scale": GUIDE_SCALE, + "guide_chunks": GUIDE_CHUNKS, + "seed": SEED, + "per_chunk_gap_uint8": report, + } + (OUT_DIR / "report.json").write_text(json.dumps(meta, indent=2)) + print(f"videos + report under {OUT_DIR}/") + + +if __name__ == "__main__": + main() diff --git a/integrations/omnidreams/scripts/sweep_text_edit.py b/integrations/omnidreams/scripts/sweep_text_edit.py new file mode 100644 index 000000000..beb427bd6 --- /dev/null +++ b/integrations/omnidreams/scripts/sweep_text_edit.py @@ -0,0 +1,271 @@ +# 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. + +"""Calibration sweep: which mid-stream edits land, and at what guidance. + +One pipeline load, then RNG-matched rollouts for a bank of edit prompts x +guidance scales against a shared control. The snow prompts include the +scene bundle's own snowstorm phrasing (training-distribution wording) to +separate "snow is OOD" from "my prompt was OOD". Writes per-combo videos, +a per-chunk divergence report, and a comparison grid. + +Run from the repo root:: + + .venv/bin/python integrations/omnidreams/scripts/sweep_text_edit.py +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +# Must land before the first CUDA allocation (co-tenant VRAM share). +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import mediapy as media +import numpy as np +import torch +from omnidreams.config import SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE +from omnidreams.pipeline import OmnidreamsPipeline +from omnidreams.runner import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH +from torch import Tensor + +from flashdreams.infra.config import derive_config +from flashdreams.infra.runner_io import ( + load_first_frame_tensor, + load_video_tensor, + write_video_tensor, +) + +SAMPLES_ROOT = ( + Path.home() + / ".cache/huggingface/hub/datasets--nvidia--omni-dreams-samples/snapshots" +) +UUID = os.environ.get("UUID", "23599139-948f-4681-b7f4-74794113086d") +N_CHUNKS = int(os.environ.get("N_CHUNKS", "28")) +SWAP_AT = int(os.environ.get("SWAP_AT", "8")) +SEED = int(os.environ.get("SEED", "42")) +OUT_DIR = Path( + os.environ.get("OUT_DIR", "integrations/omnidreams/scripts/outputs/edit_sweep") +) + +# The scene bundle's own weather phrasings (training-distribution wording), +# lightly de-scene-specified (drop the named parked cars). +SNOW_NATIVE = ( + "A dashcam perspective from inside a vehicle driving down a wide suburban " + "residential street during a snowstorm. The road is heavily covered in " + "white snow with visible parallel tire tracks. Vehicles parked along the " + "curb are coated in a layer of snow. The surrounding houses, lawns, and " + "large trees are completely blanketed in winter snow. The sky is overcast " + "and gray with snowflakes visibly falling. In the foreground, the bottom " + "of the windshield and the car's hood are visible, with snow accumulating " + "around the windshield wipers." +) +SNOW_MINE = ( + "Driving scene from a front-facing car camera at night in a heavy " + "snowstorm. Thick snow falling, snow-covered road and buildings, " + "headlights and streetlights glowing through the snow. Photorealistic " + "dashcam footage." +) +RAIN_NIGHT_NATIVE = ( + "A deep night sky of dark blue and grey is heavy with persistent, visible " + "rain streaks. The overall atmosphere is dark and thoroughly wet. An " + "asphalt road, marked by double yellow center lines, extends into the " + "distance, its surface completely saturated with sheeting water, creating " + "a glossy mirror that breaks and complexifies the reflections of multiple " + "warm-toned overhead streetlights. In the immediate lower foreground, the " + "car's wet hood is covered with rain droplets and reflecting light." +) +FOG = ( + "A dashcam perspective of a suburban street in extremely dense fog. " + "Visibility is very low; buildings and trees fade into a uniform white-" + "gray haze within tens of meters. Faint silhouettes of parked cars line " + "the curb, headlights diffuse into soft glows. Muted, desaturated colors." +) +NIGHT = ( + "A dashcam perspective of a suburban street late at night. Dark sky, the " + "road lit by warm streetlights and the car's headlights, parked cars in " + "shadow along the curb, illuminated house windows, deep shadows under the " + "trees. Photorealistic night dashcam footage." +) +SUNSET = ( + "A dashcam perspective of a suburban street at golden-hour sunset. Warm " + "orange low sun ahead near the horizon, long shadows across the road, " + "golden light on the trees and house facades, glowing warm sky with a few " + "pink clouds. Photorealistic dashcam footage." +) + +# (name, prompt, guidance_scale, guidance_chunks); scale 1.0 = plain swap. +COMBOS: list[tuple[str, str, float, int]] = [ + ("snow_native_plain", SNOW_NATIVE, 1.0, 0), + ("snow_native_g3", SNOW_NATIVE, 3.0, 6), + ("snow_native_g5", SNOW_NATIVE, 5.0, 6), + ("snow_mine_g3", SNOW_MINE, 3.0, 6), + ("snow_mine_g5", SNOW_MINE, 5.0, 6), + ("rain_night_g3", RAIN_NIGHT_NATIVE, 3.0, 6), + ("fog_g3", FOG, 3.0, 6), + ("night_g3", NIGHT, 3.0, 6), + ("sunset_g3", SUNSET, 3.0, 6), +] + + +def _sample_paths(uuid: str) -> tuple[Path, Path, str]: + hdmaps = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/*_hdmap.mp4")) + frames = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/first_frame.png")) + prompts = sorted(SAMPLES_ROOT.glob(f"*/data/single_view/{uuid}/prompt.txt")) + assert hdmaps and frames and prompts, f"sample {uuid} missing from local HF cache" + return hdmaps[0], frames[0], prompts[0].read_text().strip() + + +@torch.no_grad() +def _rollout( + pipe: OmnidreamsPipeline, + *, + hdmap: Tensor, + first: Tensor, + base_prompt: str, + edit: tuple[str, float, int] | None, +) -> Tensor: + pipe.diffusion_model._rng = torch.Generator(device=pipe.device).manual_seed(SEED) + cache = pipe.initialize_cache(text=[[base_prompt]], image=first) + chunks: list[Tensor] = [] + start = 0 + for ar_idx in range(N_CHUNKS): + if edit is not None and ar_idx == SWAP_AT: + prompt, scale, guide_chunks = edit + pipe.replace_text( + cache, + [[prompt]], + guidance_scale=scale, + guidance_chunks=guide_chunks, + ) + num_frames = pipe.get_num_frames(ar_idx) + chunk = pipe.generate( + ar_idx, cache, hdmap=hdmap[:, :, start : start + num_frames] + ) + pipe.finalize(ar_idx, cache) + chunks.append(chunk[0, 0].float().cpu()) + start += num_frames + del cache + torch.cuda.empty_cache() + return torch.cat(chunks, dim=0) + + +def _per_chunk_gap(a: Tensor, b: Tensor) -> list[float]: + gaps, start = [], 0 + for ar_idx in range(N_CHUNKS): + n = 5 if ar_idx == 0 else 8 + gaps.append( + float((a[start : start + n] - b[start : start + n]).abs().mean() * 127.5) + ) + start += n + return gaps + + +def main() -> None: + hdmap_path, frame_path, clip_prompt = _sample_paths(UUID) + total_frames = 5 + (N_CHUNKS - 1) * 8 + device = torch.device("cuda") + hdmap = load_video_tensor( + hdmap_path, + pixel_height=DEFAULT_VIDEO_HEIGHT, + pixel_width=DEFAULT_VIDEO_WIDTH, + device=device, + dtype=torch.bfloat16, + )[:total_frames][None, None] + first = load_first_frame_tensor( + frame_path, + pixel_height=DEFAULT_VIDEO_HEIGHT, + pixel_width=DEFAULT_VIDEO_WIDTH, + device=device, + dtype=torch.bfloat16, + )[None, None] + + cfg = derive_config( + SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE, + enable_sync_and_profile=False, + diffusion_model=dict( + seed=SEED, + transformer=dict(compile_network=False, use_cuda_graph=False), + ), + ) + pipe = cfg.setup() + assert isinstance(pipe, OmnidreamsPipeline) + pipe = pipe.to("cuda") + + OUT_DIR.mkdir(parents=True, exist_ok=True) + print(f"clip {UUID}: {clip_prompt[:100]}...") + print("rolling out control ...", flush=True) + control = _rollout( + pipe, hdmap=hdmap, first=first, base_prompt=clip_prompt, edit=None + ) + write_video_tensor(control, OUT_DIR / "control.mp4", fps=30, layout="tchw") + + report: dict[str, dict] = {} + videos: dict[str, Tensor] = {"control": control} + for name, prompt, scale, guide_chunks in COMBOS: + print(f"rolling out {name} ...", flush=True) + video = _rollout( + pipe, + hdmap=hdmap, + first=first, + base_prompt=clip_prompt, + edit=(prompt, scale, guide_chunks), + ) + videos[name] = video + write_video_tensor(video, OUT_DIR / f"{name}.mp4", fps=30, layout="tchw") + gaps = _per_chunk_gap(video, control) + report[name] = { + "prompt": prompt, + "guidance_scale": scale, + "guidance_chunks": guide_chunks, + "pre_swap_max_gap": max(gaps[:SWAP_AT]), + "post_swap_gaps": gaps[SWAP_AT:], + } + post = gaps[SWAP_AT:] + print( + f"{name:>18}: pre {max(gaps[:SWAP_AT]):5.3f} " + f"post first/mid/last {post[0]:6.2f} {post[len(post) // 2]:6.2f} {post[-1]:6.2f}" + ) + + # Grid: rows = [control, *combos], cols = pre-swap / +6 / +12 / last. + frame_cols = [SWAP_AT * 8 - 8, SWAP_AT * 8 + 45, SWAP_AT * 8 + 93, total_frames - 1] + row_names = ["control", *(name for name, *_ in COMBOS)] + rows = [] + for name in row_names: + arr = ((videos[name].numpy() + 1.0) * 127.5).clip(0, 255).astype("uint8") + rows.append( + np.concatenate([arr[c].transpose(1, 2, 0) for c in frame_cols], axis=1) + ) + grid = np.concatenate(rows, axis=0)[::2, ::2] + media.write_image(OUT_DIR / "grid.png", grid) + + meta = { + "uuid": UUID, + "clip_prompt": clip_prompt, + "n_chunks": N_CHUNKS, + "swap_at": SWAP_AT, + "seed": SEED, + "grid_row_order": row_names, + "grid_frame_cols": frame_cols, + "combos": report, + } + (OUT_DIR / "report.json").write_text(json.dumps(meta, indent=2)) + print(f"done -> {OUT_DIR}/ (grid rows: {', '.join(row_names)})") + + +if __name__ == "__main__": + main() diff --git a/integrations/omnidreams/tests/test_edit_lora.py b/integrations/omnidreams/tests/test_edit_lora.py new file mode 100644 index 000000000..f53a4cb34 --- /dev/null +++ b/integrations/omnidreams/tests/test_edit_lora.py @@ -0,0 +1,160 @@ +# 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. + +"""CPU-only unit tests for the pre-merged text-edit LoRA deploy hook. + +Covers the deploy invariants: + +* ``TextEditLoRA`` merges ``W + B @ A`` correctly, toggles by in-place + ``copy_`` (stable storage addresses), restores the base bit-exactly, + and is idempotent. +* With the hook attached, ``replace_text_embeddings`` builds a + ``use_lora`` window (no KV snapshots), ``predict_flow`` runs a single + branch on merged weights, the window expiry restores base weights, and + a fresh rollout resets the hook. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import torch +from omnidreams._edit_lora import TextEditLoRA, _target_linears +from omnidreams.transformer import CosmosTransformer + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from test_text_edit import _init_cache, _tiny_transformer # noqa: E402 + +pytestmark = pytest.mark.ci_cpu + + +def _fake_checkpoint(network, *, rank: int = 4, path: Path): + torch.manual_seed(3) + linears = _target_linears(network) + sd = {} + for i, lin in enumerate(linears): + sd[2 * i] = torch.randn(rank, lin.in_features) * 0.02 # A + sd[2 * i + 1] = torch.randn(lin.out_features, rank) * 0.02 # B + torch.save({"lora": sd}, path) + return linears, sd + + +def _make_hooked_transformer(tmp_path) -> tuple[CosmosTransformer, TextEditLoRA]: + transformer = _tiny_transformer() + ckpt = tmp_path / "edit_lora.pt" + _fake_checkpoint(transformer.network, path=ckpt) + edit_lora = TextEditLoRA(transformer.network, ckpt) + transformer.set_text_edit_lora(edit_lora) + return transformer, edit_lora + + +def test_merge_toggle_and_bit_exact_restore(tmp_path): + transformer = _tiny_transformer() + ckpt = tmp_path / "edit_lora.pt" + linears, sd = _fake_checkpoint(transformer.network, path=ckpt) + base = [lin.weight.detach().clone() for lin in linears] + ptrs = [lin.weight.data_ptr() for lin in linears] + + edit_lora = TextEditLoRA(transformer.network, ckpt) + assert edit_lora.rank == 4 + assert len(linears) == 2 * 8 # 2 tiny blocks x 8 projections + + edit_lora.set_active(True) + for i, lin in enumerate(linears): + expected = ( + base[i].to(torch.float32) + sd[2 * i + 1].float() @ sd[2 * i].float() + ).to(base[i].dtype) + assert torch.equal(lin.weight, expected) + assert lin.weight.data_ptr() == ptrs[i] # in place: CUDA-graph safe + edit_lora.set_active(True) # idempotent + + edit_lora.set_active(False) + for i, lin in enumerate(linears): + assert torch.equal(lin.weight, base[i]) + assert lin.weight.data_ptr() == ptrs[i] + + +def test_checkpoint_shape_mismatch_rejected(tmp_path): + transformer = _tiny_transformer() + ckpt = tmp_path / "bad.pt" + torch.save({"lora": {0: torch.zeros(4, 8), 1: torch.zeros(8, 4)}}, ckpt) + with pytest.raises(AssertionError, match="target-list mismatch"): + TextEditLoRA(transformer.network, ckpt) + + +def test_replace_builds_lora_window_and_expiry_restores(tmp_path): + transformer, edit_lora = _make_hooked_transformer(tmp_path) + cache, _ = _init_cache(transformer) + + transformer.replace_text_embeddings( + cache, torch.randn(1, 1, 10, 32), guidance_scale=3.0, guidance_chunks=2 + ) + guidance = cache.text_edit_guidance + assert guidance is not None and guidance.use_lora + assert guidance.kv_old == [] and guidance.kv_new == [] # no snapshots + assert edit_lora.active + + # predict_flow runs a single branch (the stub counts calls). + calls = [] + + def fake_branch(**kwargs): + calls.append(kwargs["network_cache"]) + return torch.zeros(4) + + transformer._predict_branch = fake_branch # ty: ignore[invalid-assignment] + cache.start(0) + transformer.predict_flow( + noisy_latent=torch.zeros(4), timestep=torch.tensor(1000.0), cache=cache + ) + assert len(calls) == 1 # no double branch + assert edit_lora.active + cache.finalize(0) + + cache.start(1) # second (last) guided chunk + assert cache.text_edit_guidance is not None + cache.finalize(1) + + cache.start(2) # countdown expired -> cleared by the cache... + assert cache.text_edit_guidance is None + transformer.predict_flow( + noisy_latent=torch.zeros(4), timestep=torch.tensor(1000.0), cache=cache + ) + assert not edit_lora.active # ...and the first forward restores base + cache.finalize(2) + + +def test_plain_swap_and_new_rollout_deactivate(tmp_path): + transformer, edit_lora = _make_hooked_transformer(tmp_path) + cache, _ = _init_cache(transformer) + + transformer.replace_text_embeddings( + cache, torch.randn(1, 1, 10, 32), guidance_scale=3.0, guidance_chunks=4 + ) + assert edit_lora.active + + # A plain swap (no guidance) mid-window supersedes it and restores base. + transformer.replace_text_embeddings(cache, torch.randn(1, 1, 10, 32)) + assert cache.text_edit_guidance is None + assert not edit_lora.active + + # Mid-window session teardown: a fresh rollout resets the hook. + transformer.replace_text_embeddings( + cache, torch.randn(1, 1, 10, 32), guidance_scale=3.0, guidance_chunks=4 + ) + assert edit_lora.active + _init_cache(transformer, seed=2) + assert not edit_lora.active diff --git a/integrations/omnidreams/tests/test_text_edit.py b/integrations/omnidreams/tests/test_text_edit.py new file mode 100644 index 000000000..ddbef59db --- /dev/null +++ b/integrations/omnidreams/tests/test_text_edit.py @@ -0,0 +1,412 @@ +# 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. + +"""CPU-only unit tests for the mid-stream text-edit path. + +Covers the invariants a live prompt swap depends on: + +* ``BlockKVCache.overwrite_kv_`` replaces contents without moving storage + (CUDA-graph safety) and rejects shape drift. +* Same-index cache rewrites (the ReCache primitive) overwrite only the last + chunk's slots and leave bookkeeping untouched. +* ``CosmosDiTNetwork.replace_text_embeddings`` reproduces exactly the + cross-attn K/V a fresh ``initialize_cache`` would build for the new + prompt, in place, without touching self-attention history. +* ``CosmosTransformer.replace_text_embeddings`` snapshots old/new K/V for + text-edit guidance, and ``predict_flow`` combines the two branches + CFG-style, leaving the buffers on the new prompt. +* The guidance countdown clears after the requested number of chunks and + ignores same-index re-opens. +""" + +from __future__ import annotations + +import pytest +import torch +from omnidreams.transformer import ( + CosmosTransformer, + CosmosTransformerConfig, + TextEditGuidance, +) +from omnidreams.transformer.impl.network import ( + CosmosDiTNetwork, + CosmosDiTNetworkConfig, +) + +from flashdreams.core.attention.kvcache import BlockKVCache + +pytestmark = pytest.mark.ci_cpu + + +## BlockKVCache primitives + + +def _make_cross_attn_cache(L: int = 6, n: int = 2, d: int = 4) -> BlockKVCache: + k = torch.randn(1, L, n, d) + v = torch.randn(1, L, n, d) + return BlockKVCache.from_tensor(k, v, seq_dim=-3) + + +def test_overwrite_kv_preserves_addresses_and_content(): + torch.manual_seed(0) + cache = _make_cross_attn_cache() + k_ptr = cache._k.data_ptr() + v_ptr = cache._v.data_ptr() + + new_k = torch.randn_like(cache._k) + new_v = torch.randn_like(cache._v) + cache.overwrite_kv_(new_k, new_v) + + assert cache._k.data_ptr() == k_ptr + assert cache._v.data_ptr() == v_ptr + assert torch.equal(cache.cached_k(), new_k) + assert torch.equal(cache.cached_v(), new_v) + + +def test_overwrite_kv_rejects_shape_mismatch(): + cache = _make_cross_attn_cache(L=6) + bad_k = torch.randn(1, 5, 2, 4) + bad_v = torch.randn(1, 5, 2, 4) + with pytest.raises(AssertionError, match="shape mismatch"): + cache.overwrite_kv_(bad_k, bad_v) + + +def test_clone_kv_returns_detached_copies(): + torch.manual_seed(0) + cache = _make_cross_attn_cache() + k_clone, v_clone = cache.clone_kv() + assert k_clone.data_ptr() != cache._k.data_ptr() + k_before = cache.cached_k().clone() + k_clone.fill_(0.0) + v_clone.fill_(0.0) + assert torch.equal(cache.cached_k(), k_before) + + +def test_same_index_rewrite_overwrites_last_chunk_only(): + """ReCache primitive: re-opening the just-committed chunk index rewrites + the same physical slots without rolling the window or advancing + bookkeeping.""" + torch.manual_seed(0) + chunk, n_chunks = 4, 4 + cache = BlockKVCache( + k_shape=(1, chunk * n_chunks, 2, 4), + v_shape=(1, chunk * n_chunks, 2, 4), + seq_dim=-3, + chunk_size=chunk, + window_size=chunk * n_chunks, + sink_size=0, + device="cpu", + dtype=torch.float32, + ) + chunks = [torch.randn(1, chunk, 2, 4) for _ in range(3)] + for idx, c in enumerate(chunks): + cache.before_update(idx) + cache.update(c, c) + cache.after_update(idx) + n_cached, prev_idx = cache._n_cached, cache._prev_chunk_idx + + replacement = torch.randn(1, chunk, 2, 4) + cache.before_update(2) + cache.update(replacement, replacement) + cache.after_update(2) + + assert cache._n_cached == n_cached + assert cache._prev_chunk_idx == prev_idx + got_k = cache._k[:, : 3 * chunk] + assert torch.equal(got_k[:, :chunk], chunks[0]) + assert torch.equal(got_k[:, chunk : 2 * chunk], chunks[1]) + assert torch.equal(got_k[:, 2 * chunk :], replacement) + + # The rollout continues normally afterwards. + cache.before_update(3) + cache.update(chunks[0], chunks[0]) + cache.after_update(3) + assert cache._prev_chunk_idx == 3 + + +## Network-level replace + + +def _tiny_network(seed: int = 0) -> CosmosDiTNetwork: + torch.manual_seed(seed) + config = CosmosDiTNetworkConfig( + in_channels=16, + out_channels=16, + patch_spatial=2, + patch_temporal=1, + model_channels=64, + num_blocks=2, + num_heads=4, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=0, + enable_cross_view_attn=False, + ) + return CosmosDiTNetwork(config) + + +def test_network_replace_matches_fresh_init_and_keeps_self_attn(): + torch.manual_seed(0) + network = _tiny_network() + ctx1 = torch.randn(1, 1, 10, 32) + ctx2 = torch.randn(1, 1, 10, 32) + + cache = network.initialize_cache( + chunk_size=32, window_size=96, sink_size=0, context=ctx1 + ) + reference = network.initialize_cache( + chunk_size=32, window_size=96, sink_size=0, context=ctx2 + ) + + cross_ptrs = [bc.cross_attn._k.data_ptr() for bc in cache.block_caches] + self_ptrs = [bc.self_attn._k.data_ptr() for bc in cache.block_caches] + self_snapshot = [bc.self_attn.clone_kv() for bc in cache.block_caches] + + network.replace_text_embeddings(cache, ctx2) + + for bc, ref, cross_ptr, self_ptr, (self_k, self_v) in zip( + cache.block_caches, reference.block_caches, cross_ptrs, self_ptrs, self_snapshot + ): + assert torch.equal(bc.cross_attn._k, ref.cross_attn._k) + assert torch.equal(bc.cross_attn._v, ref.cross_attn._v) + assert bc.cross_attn._k.data_ptr() == cross_ptr + assert bc.self_attn._k.data_ptr() == self_ptr + assert torch.equal(bc.self_attn._k, self_k) + assert torch.equal(bc.self_attn._v, self_v) + + +## Transformer-level replace + guidance + + +def _tiny_transformer(seed: int = 0) -> CosmosTransformer: + torch.manual_seed(seed) + config = CosmosTransformerConfig( + network=CosmosDiTNetworkConfig( + in_channels=16, + out_channels=16, + patch_spatial=2, + patch_temporal=1, + model_channels=64, + num_blocks=2, + num_heads=4, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=0, + enable_cross_view_attn=False, + ), + checkpoint_path=None, + batch_shape=(1,), + num_views=1, + len_t=2, + window_size_t=6, + sink_size_t=0, + compile_network=False, + use_cuda_graph=False, + guidance_scale=1.0, + ) + return CosmosTransformer(config) + + +def _init_cache(transformer: CosmosTransformer, seed: int = 1): + torch.manual_seed(seed) + text = torch.randn(1, 1, 10, 32) + image = torch.randn(1, 1, 1, 16, 8, 8) + cache = transformer.initialize_autoregressive_cache( + height=8, width=8, text_embeddings=text, image_embeddings=image + ) + return cache, text + + +def test_transformer_replace_snapshots_old_and_new_kv(): + transformer = _tiny_transformer() + cache, _ = _init_cache(transformer) + old_kv = [bc.cross_attn.clone_kv() for bc in cache.network_cache.block_caches] + + new_text = torch.randn(1, 1, 10, 32) + transformer.replace_text_embeddings( + cache, new_text, guidance_scale=2.0, guidance_chunks=3 + ) + + guidance = cache.text_edit_guidance + assert guidance is not None + assert guidance.scale == 2.0 and guidance.chunks_remaining == 3 + for (k_old, v_old), (k_ref, v_ref) in zip(guidance.kv_old, old_kv): + assert torch.equal(k_old, k_ref) + assert torch.equal(v_old, v_ref) + # Buffers and the "new" snapshot both hold the new prompt's K/V. + for (k_new, v_new), bc in zip(guidance.kv_new, cache.network_cache.block_caches): + assert torch.equal(k_new, bc.cross_attn.cached_k()) + assert torch.equal(v_new, bc.cross_attn.cached_v()) + assert not torch.equal(k_new, guidance.kv_old[0][0]) + + # A follow-up plain swap (no guidance) clears the guidance state. + transformer.replace_text_embeddings(cache, torch.randn(1, 1, 10, 32)) + assert cache.text_edit_guidance is None + + +def test_predict_flow_guidance_combines_and_lands_on_new_kv(): + transformer = _tiny_transformer() + cache, _ = _init_cache(transformer) + block_caches = cache.network_cache.block_caches + + kv_old = [ + (torch.zeros_like(bc.cross_attn._k), torch.zeros_like(bc.cross_attn._v)) + for bc in block_caches + ] + kv_new = [ + (torch.ones_like(bc.cross_attn._k), torch.ones_like(bc.cross_attn._v)) + for bc in block_caches + ] + cache.text_edit_guidance = TextEditGuidance( + scale=3.0, chunks_remaining=1, kv_old=kv_old, kv_new=kv_new + ) + + # Stub the branch forward: report the current block-0 cross-K content so + # the test observes which prompt each branch ran under (old=0, new=1). + def fake_branch(**kwargs): + return block_caches[0].cross_attn.cached_k().mean() * torch.ones(4) + + transformer._predict_branch = fake_branch # ty: ignore[invalid-assignment] + + flow = transformer.predict_flow( + noisy_latent=torch.zeros(4), + timestep=torch.tensor(1000.0), + cache=cache, + ) + # flow_old + scale * (flow_new - flow_old) = 0 + 3 * (1 - 0) + assert torch.allclose(flow, torch.full((4,), 3.0)) + for bc, (k_new, v_new) in zip(block_caches, kv_new): + assert torch.equal(bc.cross_attn._k, k_new) + assert torch.equal(bc.cross_attn._v, v_new) + + # The KV-commit forward must run single-branch under the new prompt. + transformer._finalizing_kv_cache = True + flow = transformer.predict_flow( + noisy_latent=torch.zeros(4), + timestep=torch.tensor(128.0), + cache=cache, + ) + assert torch.allclose(flow, torch.ones(4)) + + +def test_guidance_countdown_clears_after_n_chunks(): + transformer = _tiny_transformer() + cache, _ = _init_cache(transformer) + transformer.replace_text_embeddings( + cache, + torch.randn(1, 1, 10, 32), + guidance_scale=2.0, + guidance_chunks=2, + ) + assert cache.text_edit_guidance is not None + + cache.start(0) + assert cache.text_edit_guidance is not None # guided chunk 1 of 2 + assert cache.text_edit_guidance.chunks_remaining == 1 + cache.finalize(0) + + # A same-index re-open (ReCache of chunk 0) must not consume a chunk. + cache.start(0) + assert cache.text_edit_guidance.chunks_remaining == 1 + cache.finalize(0) + + cache.start(1) + assert cache.text_edit_guidance is not None # guided chunk 2 of 2 + assert cache.text_edit_guidance.chunks_remaining == 0 + cache.finalize(1) + + cache.start(2) + assert cache.text_edit_guidance is None # guidance expired + cache.finalize(2) + + +def test_replace_rejects_native_dit_and_cfg_guidance_combination(): + transformer = _tiny_transformer() + cache, _ = _init_cache(transformer) + + transformer._optimized_dit_executor = object() + with pytest.raises(NotImplementedError): + transformer.replace_text_embeddings(cache, torch.randn(1, 1, 10, 32)) + transformer._optimized_dit_executor = None + + cache.network_cache_uncond = cache.network_cache # any non-None sentinel + with pytest.raises(AssertionError, match="mutually exclusive"): + transformer.replace_text_embeddings( + cache, + torch.randn(1, 1, 10, 32), + guidance_scale=2.0, + guidance_chunks=1, + ) + # A plain swap (no guidance) is still fine with CFG configs. + cache.network_cache_uncond = None + transformer.replace_text_embeddings(cache, torch.randn(1, 1, 10, 32)) + + +## ReCache RNG neutrality + + +def test_recache_uses_dedicated_rng_and_restores_model_stream(): + """ReCache draws its context noise from a per-index seeded generator and + leaves the model RNG stream exactly where it was.""" + from omnidreams.pipeline import OmnidreamsPipeline + + pipe = OmnidreamsPipeline.__new__(OmnidreamsPipeline) + + class FakeCache: + autoregressive_index = 7 + started = None + + def start(self, idx): + self.started = idx + + class FakeFinalState: + autoregressive_index = 7 + cache = FakeCache() + + class FakeDM: + device = torch.device("cpu") + + def __init__(self): + self._rng = torch.Generator().manual_seed(42) + self.seen_seed = None + + @property + def rng(self): + return self._rng + + def finalize(self, final_state): + self.seen_seed = self._rng.initial_seed() + + dm = FakeDM() + rollout_rng = dm._rng + state_before = rollout_rng.get_state().clone() + pipe.diffusion_model = dm + + class FakePipelineCache: + final_state = FakeFinalState() + + pipe.recache_last_chunk(FakePipelineCache()) + assert dm.seen_seed == OmnidreamsPipeline._RECACHE_NOISE_SEED + 7 + assert dm._rng is rollout_rng # restored, same object + assert torch.equal(rollout_rng.get_state(), state_before) # untouched + assert FakeFinalState.cache.started == 7 + + # No final state -> no-op. + class EmptyCache: + final_state = None + + pipe.recache_last_chunk(EmptyCache()) diff --git a/integrations/omnidreams/tests/test_webrtc_actors.py b/integrations/omnidreams/tests/test_webrtc_actors.py new file mode 100644 index 000000000..953762434 --- /dev/null +++ b/integrations/omnidreams/tests/test_webrtc_actors.py @@ -0,0 +1,243 @@ +# 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. + +"""CPU-only unit tests for user-spawned WebRTC actors.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from ludus_renderer import CubePool +from omnidreams.webrtc.actors import ( + ACTOR_PRESETS, + RIG_HEIGHT_M, + actors_to_cube_pool, + clone_template_pool, + extract_parked_templates, + find_empty_gap, + spawn_actor_ahead, +) +from scipy.spatial.transform import Rotation + +pytestmark = pytest.mark.ci_cpu + + +def _ego_pose(x: float = 0.0, y: float = 0.0, yaw_deg: float = 0.0) -> np.ndarray: + pose = np.eye(4, dtype=np.float64) + pose[:3, :3] = Rotation.from_euler("z", np.deg2rad(yaw_deg)).as_matrix() + pose[:3, 3] = [x, y, 0.0] + return pose + + +def test_spawn_ahead_places_actor_along_heading(): + actor = spawn_actor_ahead( + preset="car", + ego_pose=_ego_pose(x=5.0, y=2.0, yaw_deg=90.0), + spawn_timestamp_us=1_000_000, + distance_m=10.0, + lateral_m=1.0, + ) + # Heading +90deg: forward is +y, left is -x. + np.testing.assert_allclose(actor.translation[0], 4.0, atol=1e-5) + np.testing.assert_allclose(actor.translation[1], 12.0, atol=1e-5) + # Bbox center sits half its height above the road plane (the ego pose is + # the rig origin, RIG_HEIGHT_M above the road). + np.testing.assert_allclose( + actor.translation[2], + ACTOR_PRESETS["car"][1][2] / 2.0 - RIG_HEIGHT_M, + atol=1e-6, + ) + np.testing.assert_allclose(actor.velocity, np.zeros(3), atol=1e-6) + + +def test_spawn_with_speed_moves_along_heading(): + actor = spawn_actor_ahead( + preset="truck", + ego_pose=_ego_pose(), + spawn_timestamp_us=0, + distance_m=20.0, + speed_mps=5.0, + ) + later = actor.translation_at(2_000_000) # +2 s + np.testing.assert_allclose(later[0] - actor.translation[0], 10.0, atol=1e-4) + np.testing.assert_allclose(later[1], actor.translation[1], atol=1e-6) + + +def test_spawn_heading_ignores_camera_pitch(): + pose = _ego_pose() + pose[:3, :3] = Rotation.from_euler("y", np.deg2rad(-20.0)).as_matrix() + actor = spawn_actor_ahead( + preset="cone", ego_pose=pose, spawn_timestamp_us=0, distance_m=8.0 + ) + # Forward projected to the ground plane: full 8 m in x, none in z beyond + # the half-height-minus-rig offset. + np.testing.assert_allclose(actor.translation[0], 8.0, atol=1e-5) + np.testing.assert_allclose( + actor.translation[2], + ACTOR_PRESETS["cone"][1][2] / 2.0 - RIG_HEIGHT_M, + atol=1e-6, + ) + + +def test_unknown_preset_raises(): + with pytest.raises(KeyError): + spawn_actor_ahead(preset="dragon", ego_pose=_ego_pose(), spawn_timestamp_us=0) + + +def test_actors_to_cube_pool_respects_spawn_time(): + frame_ts = [0, 33_333, 66_666, 99_999] + early = spawn_actor_ahead( + preset="car", ego_pose=_ego_pose(), spawn_timestamp_us=0, distance_m=10.0 + ) + late = spawn_actor_ahead( + preset="cone", + ego_pose=_ego_pose(), + spawn_timestamp_us=66_666, + distance_m=5.0, + ) + pool = actors_to_cube_pool([early, late], frame_ts, device="cpu") + assert pool is not None + # Track lengths: early actor has all 4 frames, late actor only the last 2. + lengths = np.diff(np.concatenate([[0], pool.cube_ts_prefix_sum.cpu().numpy()])) + assert lengths.tolist() == [4, 2] + assert pool.scales.shape[0] == 2 + + # Not-yet-spawned actors produce no pool at all. + future = spawn_actor_ahead( + preset="car", ego_pose=_ego_pose(), spawn_timestamp_us=10_000_000 + ) + assert actors_to_cube_pool([future], frame_ts, device="cpu") is None + + +def test_pool_positions_track_constant_velocity(): + frame_ts = [0, 1_000_000] + actor = spawn_actor_ahead( + preset="car", + ego_pose=_ego_pose(), + spawn_timestamp_us=0, + distance_m=10.0, + speed_mps=3.0, + ) + pool = actors_to_cube_pool([actor], frame_ts, device="cpu") + assert pool is not None + translations = pool.translations.cpu().numpy() + np.testing.assert_allclose(translations[0][0], 10.0, atol=1e-4) + np.testing.assert_allclose(translations[1][0], 13.0, atol=1e-4) + + +def _scene_pool(tracks: list[dict]) -> CubePool: + """Synthetic CubePool from per-track specs (xy, n, t0_s, length, drift).""" + track_ts, translations, quaternions, scales, colors, lengths = ( + [], + [], + [], + [], + [], + [], + ) + for track in tracks: + n = track.get("n", 12) + t0 = track.get("t0_s", 0.0) + ts = torch.tensor( + [int((t0 + 0.5 * i) * 1e6) for i in range(n)], dtype=torch.int64 + ) + xy = np.asarray(track["xy"], dtype=np.float64) + drift = track.get("drift", 0.0) + pos = torch.tensor( + [[xy[0] + drift * i / max(n - 1, 1), xy[1], 0.7] for i in range(n)], + dtype=torch.float64, + ) + # Deterministic sub-centimeter jitter, like real perception tracks. + pos[:, 0] += 0.01 * torch.sin(torch.arange(n, dtype=torch.float64)) + length = track.get("length", 4.5) + track_ts.append(ts) + translations.append(pos) + quaternions.append( + torch.tensor([[0.0, 0.0, 0.2, 0.98]]).repeat(n, 1).to(torch.float64) + ) + scales.append(torch.tensor([[length, 1.9, 1.5]], dtype=torch.float64)) + colors.append(torch.rand(1, 6, dtype=torch.float64)) + lengths.append(n) + all_ts = torch.cat(track_ts) + return CubePool( + timestamps_us=torch.unique(all_ts).sort()[0], + cube_ts_prefix_sum=torch.cumsum( + torch.tensor(lengths, dtype=torch.int32), dim=0 + ).to(torch.int32), + track_timestamps_us=all_ts, + translations=torch.cat(translations), + quaternions=torch.cat(quaternions), + scales=torch.cat(scales), + colors=torch.cat(colors), + ) + + +def test_extract_templates_filters_and_sorts_by_distance(): + pool = _scene_pool( + [ + {"xy": (40.0, -7.0)}, # good, farther + {"xy": (15.0, -7.0)}, # good, nearest -> first + {"xy": (20.0, -7.0), "drift": 4.0}, # moving: rejected + {"xy": (25.0, -7.0), "n": 4}, # short coverage: rejected + {"xy": (30.0, -7.0), "length": 8.0}, # truck-sized: rejected + {"xy": (35.0, -7.0), "t0_s": 2.0}, # starts too late: rejected + ] + ) + templates = extract_parked_templates([pool], ego_pose=_ego_pose(), t0_us=0) + assert [t.source_fwd_m for t in templates] == pytest.approx([15.0, 40.0], abs=0.05) + assert templates[0].source_lateral_m == pytest.approx(-7.0, abs=0.05) + assert templates[0].translations.shape == (12, 3) + + +def test_find_empty_gap_targets_largest_free_span(): + pool = _scene_pool([{"xy": (25.0, -7.0)}, {"xy": (40.0, -7.2)}]) + center, width = find_empty_gap( + [pool], ego_pose=_ego_pose(), lateral_m=-7.0, fwd_range=(20.0, 65.0) + ) + # Occupied: 25 and 40, each +-(4.5/2 + 1.5) -> largest gap is (43.75, 65). + assert center == pytest.approx((43.75 + 65.0) / 2, abs=0.05) + assert width == pytest.approx(65.0 - 43.75, abs=0.05) + + # Actors on other lateral lines do not shrink the gap. + _, full_width = find_empty_gap( + [pool], ego_pose=_ego_pose(), lateral_m=7.0, fwd_range=(20.0, 65.0) + ) + assert full_width == pytest.approx(45.0, abs=1e-6) + + +def test_clone_template_pool_moves_rigidly_and_preserves_track(): + ego = _ego_pose(x=3.0, y=-2.0, yaw_deg=90.0) # forward +y, left -x + source = _scene_pool([{"xy": (10.0, 30.0)}]) + (template,) = extract_parked_templates([source], ego_pose=ego, t0_us=0) + + pool = clone_template_pool([(template, 30.0, -7.0)], ego_pose=ego) + first = pool.translations[0].cpu().numpy() + np.testing.assert_allclose(first[:2], [3.0 + 7.0, -2.0 + 30.0], atol=1e-6) + # Rigid shift: per-frame jitter, z, orientation, size, colors all survive. + np.testing.assert_allclose( + (pool.translations - pool.translations[0]).cpu().numpy(), + (template.translations - template.translations[0]).cpu().numpy(), + atol=1e-9, + ) + assert torch.equal(pool.quaternions, template.quaternions) + assert torch.equal(pool.scales, template.scale) + assert torch.equal(pool.colors, template.colors) + + two = clone_template_pool( + [(template, 30.0, -7.0), (template, 40.0, -7.0)], ego_pose=ego + ) + assert two.cube_ts_prefix_sum.cpu().tolist() == [12, 24] + assert two.scales.shape[0] == 2