From 8afdaf1dc9fb99a09786e5065e9b63027bb130bf Mon Sep 17 00:00:00 2001 From: Wenqing Wang Date: Wed, 22 Jul 2026 05:36:24 +0000 Subject: [PATCH] hy_worldplay: add Clean Forcing drift corrector (training recipe + content-keyed deploy) Adds a trained drift corrector for the HY-WorldPlay WAN-5B integration: a frozen-base LoRA (r16 q/k/v/o, ~0.3% params) trained from the model's own strafe-loop rollouts with a counterfactual clean-history teacher -- zero real videos end-to-end. Deploy: HyWorldPlayWanI2VRunnerConfig.drift_corrector (checkpoint path) + drift_corrector_gain. Selection is content-keyed per job: static (locked-off camera) poses run untouched base weights; commanded-motion poses apply the LoRA at alpha*(t) x gain per denoise step, where alpha*(t) is the measured systematic fraction of the drift-induced error at each solver timestep. Measured on 24-chunk (~19 s) rollouts: -48% delta-drift vs base with dynamics ~= base and no added seam pulse; the base's late-horizon saturation runaway is eliminated. Training/eval scripts and the step-0 diagnostic live under integrations/hy_worldplay/drift_correction/ with a README covering the method and the full reproduction recipe. Co-Authored-By: Claude Fable 5 Signed-off-by: Wenqing Wang --- .../hy_worldplay/drift_correction/.gitignore | 1 + .../hy_worldplay/drift_correction/README.md | 62 +++ .../hy_worldplay/drift_correction/_lora.py | 127 ++++++ .../hy_worldplay/drift_correction/_pairs.py | 102 +++++ .../hy_worldplay/drift_correction/_rollout.py | 325 +++++++++++++++ .../drift_correction/_train_attn.py | 157 ++++++++ .../drift_correction/build_pairs.py | 219 +++++++++++ .../drift_correction/demo_prompts.txt | 8 + .../drift_correction/demo_static.py | 166 ++++++++ .../drift_correction/eval_rollouts.py | 152 +++++++ .../drift_correction/gate_faithful.py | 212 ++++++++++ .../drift_correction/gate_systematicity.py | 323 +++++++++++++++ .../drift_correction/gen_first_frames.py | 139 +++++++ .../hy_worldplay/drift_correction/make_sbs.py | 87 ++++ .../drift_correction/score_drift.py | 362 +++++++++++++++++ .../hy_worldplay/drift_correction/train_v1.py | 274 +++++++++++++ .../hy_worldplay/drift_correction/train_v2.py | 370 ++++++++++++++++++ .../hy_worldplay/_drift_corrector.py | 163 ++++++++ .../hy_worldplay/hy_worldplay/runner.py | 74 ++-- .../tests/test_drift_corrector.py | 115 ++++++ 20 files changed, 3408 insertions(+), 30 deletions(-) create mode 100644 integrations/hy_worldplay/drift_correction/.gitignore create mode 100644 integrations/hy_worldplay/drift_correction/README.md create mode 100644 integrations/hy_worldplay/drift_correction/_lora.py create mode 100644 integrations/hy_worldplay/drift_correction/_pairs.py create mode 100644 integrations/hy_worldplay/drift_correction/_rollout.py create mode 100644 integrations/hy_worldplay/drift_correction/_train_attn.py create mode 100644 integrations/hy_worldplay/drift_correction/build_pairs.py create mode 100644 integrations/hy_worldplay/drift_correction/demo_prompts.txt create mode 100644 integrations/hy_worldplay/drift_correction/demo_static.py create mode 100644 integrations/hy_worldplay/drift_correction/eval_rollouts.py create mode 100644 integrations/hy_worldplay/drift_correction/gate_faithful.py create mode 100644 integrations/hy_worldplay/drift_correction/gate_systematicity.py create mode 100644 integrations/hy_worldplay/drift_correction/gen_first_frames.py create mode 100644 integrations/hy_worldplay/drift_correction/make_sbs.py create mode 100644 integrations/hy_worldplay/drift_correction/score_drift.py create mode 100644 integrations/hy_worldplay/drift_correction/train_v1.py create mode 100644 integrations/hy_worldplay/drift_correction/train_v2.py create mode 100644 integrations/hy_worldplay/hy_worldplay/_drift_corrector.py create mode 100644 integrations/hy_worldplay/tests/test_drift_corrector.py diff --git a/integrations/hy_worldplay/drift_correction/.gitignore b/integrations/hy_worldplay/drift_correction/.gitignore new file mode 100644 index 00000000..17aa483a --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/.gitignore @@ -0,0 +1 @@ +outputs/ diff --git a/integrations/hy_worldplay/drift_correction/README.md b/integrations/hy_worldplay/drift_correction/README.md new file mode 100644 index 00000000..2ec5c0d6 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/README.md @@ -0,0 +1,62 @@ + + +# Clean Forcing drift corrector for HY-WorldPlay + +Long autoregressive rollouts drift: each chunk conditions on self-generated +history and small errors compound (saturation runaway, texture mush). This +module ships a trained corrector — a frozen-base LoRA (r16 on the +self-attention q/k/v/o projections, ~0.3% params) trained from the model's +own closed-loop rollouts with a counterfactual clean-history teacher, zero +real videos — plus the full recipe that produced it. + +## Deploy + +Point the runner at a corrector checkpoint; everything else is automatic: + +```python +config = dataclasses.replace( + RUNNER_HY_WORLDPLAY_WAN_I2V_5B, + drift_corrector=Path("lora_v2.pt"), # None (default) = exact base behavior + drift_corrector_gain=0.5, # composed with the per-step alpha*(t) gate +) +``` + +Selection is content-keyed per job: static (locked-off camera) trajectories +run the untouched base weights — static scenes measure *negative* drift, so +correction there is pure artifact cost — while commanded-motion trajectories +apply the LoRA at ``alpha*(t) x gain`` per denoise step. The LoRA stays +unfused on motion jobs (a single-scale weight merge cannot express the +per-timestep gate); static jobs skip module surgery entirely. + +The trained v2 checkpoint (~30 MB ``.pt``) is distributed separately; see +the PR / release notes for the download link. + +## Reproduce the corrector + +All scripts run from the repo root on one GPU (~2.5 GPU-days end to end): + +```bash +python drift_correction/gen_first_frames.py # optional T2V seeds (PROMPTS_FILE=...) +python drift_correction/build_pairs.py # strafe-loop rollouts -> pair clips +python drift_correction/gate_faithful.py # step-0 alpha*(t) go/no-go diagnostic +python drift_correction/train_v1.py # counterfactual-teacher LoRA +python drift_correction/train_v2.py # DAgger pool + drift-contraction round +python drift_correction/eval_rollouts.py # gain-sweep rollouts +python drift_correction/score_drift.py # drift / dynamics / progression / seams +python drift_correction/demo_static.py # static-scene suite (+ make_sbs.py) +``` + +Train only if the gate passes (the drift gap is systematic: ``alpha*`` high); +the measured ``alpha*(t)`` profile doubles as the deploy gate in +``hy_worldplay/_drift_corrector.py``. + +## Design notes and background + +- Library shape (host adapter vs host-agnostic core), gate math, and the + cross-host playbook: ``LIBRARY_DESIGN.md`` on the working branch + (`wenqingw-nv/flashdreams-wq` @ ``hy-worldplay-counterfactual-forcing``). +- Method, result tables, and paper-host reproduction: + https://gitlab-master.nvidia.com/wenqingw/clean_forcing diff --git a/integrations/hy_worldplay/drift_correction/_lora.py b/integrations/hy_worldplay/drift_correction/_lora.py new file mode 100644 index 00000000..113a6f1a --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/_lora.py @@ -0,0 +1,127 @@ +# 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. + +"""LoRA corrector r_phi as low-rank deltas on the frozen DiT's self-attention.""" + +import torch +import torch.nn as nn +from torch import Tensor + +DEFAULT_TARGETS = ("self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o") +"""Module-name suffixes wrapped by :func:`apply_lora`; the HY DiT's +self-attention projections are ``blocks.{i}.self_attn.{q,k,v,o}``.""" + + +class LoRALinear(nn.Module): + """Frozen base linear plus a runtime-gated low-rank delta. + + ``B`` is zero-initialized so the wrapped module starts as an exact + identity over the base (any ``scale`` gives the base output until + training moves ``B``). ``scale`` is the runtime gain: ``0`` recovers the + frozen base (the clean-history teacher pass), ``1`` the corrected pass. + The A/B path runs in fp32 regardless of the base dtype. + """ + + def __init__(self, base: nn.Linear, rank: int = 16): + super().__init__() + self.base = base + for p in self.base.parameters(): + p.requires_grad_(False) + self.A = nn.Linear(base.in_features, rank, bias=False) + self.B = nn.Linear(rank, base.out_features, bias=False) + nn.init.normal_(self.A.weight, std=1.0 / rank) + nn.init.zeros_(self.B.weight) + self.scale = 1.0 + + def forward(self, x: Tensor) -> Tensor: + out = self.base(x) + if self.scale != 0: + delta = self.B(self.A(x.to(self.A.weight.dtype))) + out = out + self.scale * delta.to(out.dtype) + return out + + +def unwrap_compiled(network: object) -> nn.Module: + """Return the eager module behind a ``torch.compile`` wrapper, if any. + + Accepts ``object``: the DiT lives on the transformer as a plain + attribute, which the type checker resolves through ``nn.Module``'s + ``__getattr__`` union. + """ + inner = getattr(network, "_orig_mod", network) + assert isinstance(inner, nn.Module), type(inner) + return inner + + +def apply_lora( + model: nn.Module, + rank: int = 16, + targets: tuple[str, ...] = DEFAULT_TARGETS, +) -> list[str]: + """Wrap every target linear in ``model`` with a :class:`LoRALinear`. + + Args: + model: The (unwrapped) DiT network; pass ``network._orig_mod`` when + ``torch.compile`` is active so the wrapping survives recompiles. + rank: LoRA rank. + targets: Module-name suffixes to wrap. + + Returns: + Fully qualified names of the wrapped modules. + """ + wrapped = [] + for mname, module in list(model.named_modules()): + for cname, child in list(module.named_children()): + full = f"{mname}.{cname}" if mname else cname + if isinstance(child, nn.Linear) and any(t in full for t in targets): + setattr(module, cname, LoRALinear(child, rank).to(child.weight.device)) + wrapped.append(full) + return wrapped + + +def set_lora_scale(model: nn.Module, scale: float) -> None: + """Set the runtime gain on every :class:`LoRALinear` in ``model``.""" + for m in model.modules(): + if isinstance(m, LoRALinear): + m.scale = scale + + +def lora_parameters(model: nn.Module) -> list[nn.Parameter]: + """Collect the trainable A/B parameters in deterministic module order.""" + ps: list[nn.Parameter] = [] + for m in model.modules(): + if isinstance(m, LoRALinear): + ps += list(m.A.parameters()) + list(m.B.parameters()) + return ps + + +def save_lora(model: nn.Module, path) -> None: + """Save the LoRA parameters (index-keyed, CPU) to ``path``.""" + torch.save( + {"lora": {i: p.detach().cpu() for i, p in enumerate(lora_parameters(model))}}, + path, + ) + + +def load_lora(model: nn.Module, path) -> None: + """Load :func:`save_lora` output into an already-wrapped ``model``.""" + sd = torch.load(path, map_location="cpu", weights_only=False)["lora"] + params = lora_parameters(model) + assert len(sd) == len(params), ( + f"checkpoint has {len(sd)} LoRA tensors but the model exposes " + f"{len(params)}; rank or target mismatch." + ) + for i, p in enumerate(params): + p.data.copy_(sd[i].to(p.device, p.dtype)) diff --git a/integrations/hy_worldplay/drift_correction/_pairs.py b/integrations/hy_worldplay/drift_correction/_pairs.py new file mode 100644 index 00000000..6f2b4a4e --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/_pairs.py @@ -0,0 +1,102 @@ +# 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. + +"""Loading and counterfactual-history helpers for strafe-loop pair clips.""" + +from __future__ import annotations + +from pathlib import Path + +import torch +from hy_worldplay._action import HyWorldPlayCtrl +from torch import Tensor + +TOKENS_PER_FRAME = 880 +"""Post-patchify tokens per latent frame at 704x1280 (44/2 * 80/2).""" + +PATCH_DIM = 192 +"""Patchified feature width (48 channels * 1 * 2 * 2).""" + + +def load_clip(path: Path, device: torch.device | str, dtype: torch.dtype) -> dict: + """Load a ``build_pairs.py`` clip; latents land on ``device`` in ``dtype``.""" + d = torch.load(path, map_location="cpu", weights_only=False) + d["latents"] = d["latents"].to(device, dtype) + return d + + +def make_ctrl( + d: dict, k: int, *, device: torch.device | str, dtype: torch.dtype +) -> HyWorldPlayCtrl: + """Reconstruct the patchified per-AR-step ctrl payload for chunk ``k``. + + For chunks ``>= 1`` the captured ctrl carries an all-zero mask, so the + image-latent stamp is a no-op and zero latent/mask reconstruct the + rollout conditioning exactly. + """ + assert k >= 1, "chunk 0 carries the stamped image latent; not reconstructable" + sl = slice(k * 4, (k + 1) * 4) + zeros = torch.zeros(1, 4 * TOKENS_PER_FRAME, PATCH_DIM, device=device, dtype=dtype) + return HyWorldPlayCtrl( + latent=zeros, + mask=torch.zeros_like(zeros), + _is_patchified=True, + action=d["rollout_action"][..., sl].to(device), + viewmats=d["rollout_viewmats"][..., sl, :, :].to(device, dtype), + Ks=d["rollout_Ks"][..., sl, :, :].to(device, dtype), + memory_frame_indices=d["memory_frame_indices"][k], + rollout_viewmats=d["rollout_viewmats"].to(device, dtype), + rollout_Ks=d["rollout_Ks"].to(device, dtype), + rollout_action=d["rollout_action"].to(device), + ) + + +def history_of(d: dict, k: int) -> Tensor: + """Patchified drifted history for chunk ``k`` (frames ``0 .. 4k-1``).""" + return d["latents"][..., : k * 4 * TOKENS_PER_FRAME, :] + + +def chunk_x0(d: dict, k: int) -> Tensor: + """Patchified clean latent the rollout produced for chunk ``k``.""" + return d["latents"][ + ..., k * 4 * TOKENS_PER_FRAME : (k + 1) * 4 * TOKENS_PER_FRAME, : + ] + + +def clean_counterfactual( + history: Tensor, + *, + selected: list[int], + lap_latents: int, + clean_lap: int = 1, +) -> Tensor: + """Swap the memory-selected frames' content for lap-aligned clean frames. + + Frames already inside laps ``<= clean_lap`` map to themselves. The + prefill only reads the selected frames, so this substitution changes the + entire effective history while leaving positions and per-frame + conditioning untouched (the strafe loop makes actions and viewmats + identical across laps). + """ + out = history.clone() + horizon = (clean_lap + 1) * lap_latents + for idx in selected: + if idx < horizon: + continue + src = (idx % lap_latents) + clean_lap * lap_latents + s = slice(idx * TOKENS_PER_FRAME, (idx + 1) * TOKENS_PER_FRAME) + d = slice(src * TOKENS_PER_FRAME, (src + 1) * TOKENS_PER_FRAME) + out[..., s, :] = history[..., d, :] + return out diff --git a/integrations/hy_worldplay/drift_correction/_rollout.py b/integrations/hy_worldplay/drift_correction/_rollout.py new file mode 100644 index 00000000..1c6b4b0d --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/_rollout.py @@ -0,0 +1,325 @@ +# 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. + +"""Rollout capture and counterfactual x0 probes for the HY-WorldPlay drift corrector.""" + +from __future__ import annotations + +import dataclasses +import gc +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +# Must land before the first CUDA allocation: long captures fragment the +# allocator; expandable segments keep the probe phase inside the VRAM share +# left over by co-tenant jobs. +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import torch +from hy_worldplay._action import HyWorldPlayCtrl, HyWorldPlayWan21TransformerCache +from hy_worldplay.config import RUNNER_HY_WORLDPLAY_WAN_I2V_5B +from hy_worldplay.runner import ( + HyWorldPlayWanI2VRunner, + _resolve_prompt, + _write_mp4, + preprocess_first_frame, +) +from torch import Tensor + +## Runner construction + +_POINT_CLOUD_SEED = 0 +"""Global-RNG seed used right before ``_bind_memory_config``: the FOV +selector's Monte-Carlo point cloud draws from the global RNG, so pinning it +makes ``memory_frame_indices`` reproducible across processes.""" + + +def build_runner( + *, + num_chunk: int, + pose: str, + output_dir: Path, + image_path: Path | None = None, + compile_network: bool = True, +) -> HyWorldPlayWanI2VRunner: + """Build the HY-WorldPlay runner with the rollout geometry fixed up front. + + Args: + num_chunk: AR chunks per rollout; the pose source must cover + ``num_chunk * 4`` latents. + pose: Pose-string or trajectory-JSON path (upstream grammar). + output_dir: Directory the runner (and callers) write outputs into. + image_path: First-frame override; ``None`` lazy-downloads the + upstream sample image. + compile_network: Inductor-compile the DiT. Disable for training, + where LoRA module surgery + checkpointed autograd must run on + the raw module. + + Returns: + Runner with the pipeline built and the checkpoint loaded. + """ + from flashdreams.infra.config import derive_config + + cfg = dataclasses.replace( + RUNNER_HY_WORLDPLAY_WAN_I2V_5B, + example_data=image_path is None, + image_path=image_path, + pose=pose, + num_chunk=num_chunk, + output_dir=output_dir, + ) + # Eager VAE: the encoder/decoder CUDAGraphWrapper private pools cost + # ~35 GiB at 704x1280, which doesn't fit next to co-tenant jobs; these + # scripts trade decode speed for headroom. + cfg = derive_config( + cfg, + pipeline=dict( + encoder=dict(encoder=dict(use_cuda_graph=False)), + decoder=dict(use_cuda_graph=False), + diffusion_model=dict(transformer=dict(compile_network=compile_network)), + ), + ) + runner = cfg.setup() + assert isinstance(runner, HyWorldPlayWanI2VRunner) + if runner.config.example_data and runner.config.image_path is None: + runner.config.image_path = runner._fetch_example_image() + return runner + + +## Per-step alpha*(t) gating + +GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} +"""Unbiased alpha*(t) from the faithful step-0 gate +(``outputs/gate/gate_faithful.json``): the systematic fraction of the +drift-induced error per inference timestep. Gate configs deploy the LoRA at +``alpha*(t) * scale`` — correction strength follows how systematic the +error actually is at each step.""" + + +def parse_gain_token(token: str) -> float | tuple[str, float]: + """Parse a config gain token. + + ``"0.7"`` -> flat gain ``0.7``; ``"gate"`` -> ``("gate", 1.0)``; + ``"gate0.5"`` -> ``("gate", 0.5)`` (per-step ``alpha*(t) * 0.5``). + """ + token = token.strip() + if token.startswith("gate"): + return ("gate", float(token[4:]) if len(token) > 4 else 1.0) + return float(token) + + +def install_alpha_gate(runner: Any, network: Any, mode: dict) -> None: + """Wrap ``predict_flow`` so gate configs rescale the LoRA every step. + + When ``mode["gain"]`` is ``("gate", scale)``, each denoise step sets the + LoRA to ``alpha*(t) * scale`` via nearest-t lookup in + :data:`GATE_ALPHA`; flat-gain configs pass through untouched (the caller + sets the scale once per rollout). Per-token timesteps (AR0) include the + first-frame stabilization value; the max is always the scheduler step. + """ + from _lora import set_lora_scale + + transformer = runner.pipeline.diffusion_model.transformer + orig_pf = transformer.predict_flow + + def gated_pf(*args, **kwargs): + gain = mode["gain"] + if isinstance(gain, tuple): + t = float(kwargs["timestep"].reshape(-1).max()) + alpha = min(GATE_ALPHA.items(), key=lambda kv: abs(kv[0] - t))[1] + set_lora_scale(network, alpha * gain[1]) + return orig_pf(*args, **kwargs) + + transformer.predict_flow = gated_pf + + +## Rollout capture + + +@dataclass +class ChunkSnapshot: + """Per-chunk state captured from a rollout, sufficient to replay probes.""" + + history: Tensor | None + """Patchified clean-latent history the chunk was conditioned on + (``clean_latent_history`` right before this chunk's ``generate``); + ``None`` for chunk 0.""" + + clean_latent: Tensor + """Patchified x0 the sampler produced for this chunk.""" + + ctrl: HyWorldPlayCtrl + """Patchified per-AR-step control payload (action / camera / memory + indices) exactly as ``predict_flow`` consumed it.""" + + def to(self, device: torch.device | str) -> "ChunkSnapshot": + """Return a copy with every tensor (including ctrl fields) on ``device``.""" + moved = {} + for f in dataclasses.fields(self.ctrl): + v = getattr(self.ctrl, f.name) + moved[f.name] = v.to(device) if isinstance(v, Tensor) else v + return ChunkSnapshot( + history=None if self.history is None else self.history.to(device), + clean_latent=self.clean_latent.to(device), + ctrl=HyWorldPlayCtrl(**moved), + ) + + +def capture_rollout( + runner: HyWorldPlayWanI2VRunner, + *, + noise_seed: int, + save_path: Path | None = None, + mp4_path: Path | None = None, +) -> list[ChunkSnapshot]: + """Roll out ``num_chunk`` chunks and snapshot the per-chunk corrector inputs. + + Reproduces the runner's ``run()`` flow (bindings included) but records, + for every chunk, the conditioning history, the produced clean latent, + and the patchified ctrl payload. The diffusion RNG is re-seeded with + ``noise_seed`` so captures are reproducible per seed. + + Args: + runner: Runner from :func:`build_runner`. + noise_seed: Seed for the diffusion model's noise generator. + save_path: When set, snapshots are also saved (CPU tensors) here. + mp4_path: When set, the decoded rollout is written here for eyeballing. + + Returns: + One :class:`ChunkSnapshot` per chunk, tensors on the compute device. + """ + pipe = runner.pipeline + cfg = runner.config + device = next(pipe.parameters()).device + dtype = next(pipe.parameters()).dtype + + assert cfg.image_path is not None + image = preprocess_first_frame( + cfg.image_path, cfg.pixel_height, cfg.pixel_width + ).to(device=device, dtype=dtype) + cache = pipe.initialize_cache(text=[_resolve_prompt(cfg.prompt)], image=image) + + runner._bind_action_labels() + runner._bind_camera_data() + torch.manual_seed(_POINT_CLOUD_SEED) + runner._bind_memory_config(device=device) + + pipe.diffusion_model._rng = torch.Generator(device=device).manual_seed(noise_seed) + + tc = cache.transformer_cache + assert isinstance(tc, HyWorldPlayWan21TransformerCache) + snaps: list[ChunkSnapshot] = [] + chunks: list[Tensor] = [] + for ar_idx in range(cfg.num_chunk): + history = tc.clean_latent_history + history = None if history is None else history.detach().clone() + chunk = pipe.generate(ar_idx, cache) + chunks.append(chunk) + fs = cache.final_state + assert fs is not None and isinstance(fs.input, HyWorldPlayCtrl) + snaps.append( + ChunkSnapshot( + history=history, + clean_latent=fs.clean_latent.detach().clone(), + ctrl=fs.input, + ) + ) + pipe.finalize(ar_idx, cache) + + if mp4_path is not None: + mp4_path.parent.mkdir(parents=True, exist_ok=True) + _write_mp4(torch.cat(chunks, dim=-4), mp4_path, fps=cfg.fps) + if save_path is not None: + save_path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + {"snaps": [s.to("cpu") for s in snaps], "noise_seed": noise_seed}, + save_path, + ) + + # Drop this rollout's caches (rolling KV buffers, VAE stream state) + # before the caller starts the next one. + del cache, chunks + gc.collect() + torch.cuda.empty_cache() + return snaps + + +def load_rollout(path: Path, device: torch.device | str) -> list[ChunkSnapshot]: + """Load :func:`capture_rollout` snapshots and move them to ``device``.""" + data = torch.load(path, map_location="cpu", weights_only=False) + return [s.to(device) for s in data["snaps"]] + + +## Counterfactual x0 probes + + +def start_probe_chunk( + tc: HyWorldPlayWan21TransformerCache, + *, + ar_idx: int, + history: Tensor, +) -> None: + """Rewind the AR cache to chunk ``ar_idx`` conditioned on ``history``. + + Substitutes the patchified clean-latent history and resets the rolling + caches + prefill latch, so the next ``predict_flow`` call re-runs the + memory KV prefill against the substituted history. Subsequent calls at + the same chunk reuse the prefilled memory (matching in-chunk sampler + steps). + """ + assert ar_idx > 0, "probing requires a non-empty history, so ar_idx >= 1" + tc.clean_latent_history = history + tc.start(ar_idx) + + +def finish_probe_chunk(tc: HyWorldPlayWan21TransformerCache, *, ar_idx: int) -> None: + """Close the ``start`` / ``finalize`` bracket after a probe sweep. + + ``BlockKVCache`` enforces a strict ``before_update`` / ``after_update`` + alternation; :func:`start_probe_chunk` opened it, so every sweep must + call this before the next one (the pipeline's ``finalize`` plays this + role during generation). + """ + tc.finalize(ar_idx) + + +def predict_x0( + transformer: Any, + tc: HyWorldPlayWan21TransformerCache, + *, + ctrl: HyWorldPlayCtrl, + z_t: Tensor, + timestep: Tensor, + sigma: float, +) -> Tensor: + """Predict x0 at a matched noisy state ``z_t`` under the cache's current history. + + Args: + transformer: The pipeline's ``HyWorldPlayWan21Transformer``. + tc: AR cache positioned via :func:`start_probe_chunk`. + ctrl: The probed chunk's captured (patchified) ctrl payload. + z_t: Patchified noisy latent ``(1 - sigma) * x0 + sigma * eps``. + timestep: Scalar timestep tensor in the network dtype. + sigma: Noise level matching ``timestep`` on the inference schedule. + + Returns: + fp32 ``x0_hat = z_t - sigma * flow``. + """ + flow = transformer.predict_flow( + noisy_latent=z_t, timestep=timestep, cache=tc, input=ctrl + ) + return z_t.float() - sigma * flow.float() diff --git a/integrations/hy_worldplay/drift_correction/_train_attn.py b/integrations/hy_worldplay/drift_correction/_train_attn.py new file mode 100644 index 00000000..a4b6c07c --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/_train_attn.py @@ -0,0 +1,157 @@ +# 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. + +"""Grad-friendly functional dual-branch attention for corrector training. + +The production ``forward_dual_branch`` routes the current chunk's K / V +through the rolling ``BlockKVCache`` (in-place write into a no-grad buffer, +then a buffer read), which (a) severs gradients into the ``k`` / ``v`` +projections and (b) plants mutable storage inside the autograd tape, +breaking gradient-checkpoint recomputation. + +Training probes run exactly one forward per chunk on a freshly reset cache, +where ``cached_k()`` would return precisely the current chunk -- so the +functional equivalent below (use the projected K / V directly, prepend the +prefilled memory K / V) is mathematically identical while keeping the whole +path differentiable and side-effect free. Deployment keeps the stock path; +the corrector's weights are shared, so val metrics transfer. + +Must mirror :meth:`HyWorldPlayPRoPESelfAttention.forward_dual_branch` +(``hy_worldplay/_camera.py``); revisit on upstream changes. +""" + +from __future__ import annotations + +import math + +import torch +from hy_worldplay._camera import ( + HyWorldPlayMemoryKVCache, + HyWorldPlayPRoPESelfAttention, +) +from hy_worldplay._prope import prope_qkv +from torch import Tensor + + +def _rope_interleaved(x: Tensor, freqs: Tensor) -> Tensor: + """Differentiable equivalent of the fused interleaved-RoPE Triton kernel. + + The production kernel (``rope_kernel.apply_rotary_pos_emb``) is + inference-only: raw ``tl.store`` writes with no backward, so gradients + through it are silently dropped. Mirrors its arithmetic exactly: pair + ``(2k, 2k+1)`` rotated by angle ``freqs[..., k]``, cos/sin in fp32 then + cast to ``x.dtype`` before the multiply. + + Args: + x: Activations ``[B, S, H, D]``. + freqs: Angles in the full-width ``[S, 1, 1, D]`` ``shift_t`` layout, + where each pair's angle is duplicated; pair ``k`` reads lane + ``2k`` (the kernel's ``stride_fd * 2`` skip). + + Returns: + Rotated tensor, same shape/dtype, fresh storage. + """ + ang = freqs[:, 0, 0, 0::2].float() + cos = ang.cos().to(x.dtype)[None, :, None, :] + sin = ang.sin().to(x.dtype)[None, :, None, :] + pairs = x.unflatten(-1, (x.shape[-1] // 2, 2)) + a, b = pairs[..., 0], pairs[..., 1] + return torch.stack((a * cos - b * sin, b * cos + a * sin), dim=-1).flatten(-2) + + +def _functional_dual_branch( + self: HyWorldPlayPRoPESelfAttention, + x: Tensor, + kv_cache, + prope_kv_cache, + rope_freqs: Tensor, + viewmats: Tensor, + Ks: Tensor | None, + memory_kv_cache: HyWorldPlayMemoryKVCache | None = None, +) -> Tensor: + """Cache-free re-implementation of ``forward_dual_branch``. + + Valid only for the single-forward-per-chunk training regime: asserts the + rolling cache is empty so the functional K / V equal what the stock path + would have read back. + """ + assert kv_cache._n_cached == 0, ( + "functional attention expects a freshly reset rolling cache " + "(one forward per start/finalize bracket)." + ) + assert self.apply_rope_before_kvcache, ( + "cache-relative RoPE mode would need the cached-window freqs; the " + "TI2V-5B recipe applies RoPE before the cache." + ) + rope_freqs_q, rope_freqs_k = self._slice_rope_freqs(rope_freqs, kv_cache) + + batch_shape = x.shape[:-2] + batch_size = math.prod(batch_shape) + L, _ = x.shape[-2:] + n, d = self.n_heads, self.head_dim + + q_raw = self.norm_q(self.q(x)).reshape(batch_size, L, n, d) + k_raw = self.norm_k(self.k(x)).reshape(batch_size, L, n, d) + v_raw = self.v(x).reshape(batch_size, L, n, d) + + # The production kernel rotates in place, aliasing ``k_raw`` -- the + # stock PRoPE branch therefore consumes the *roped* K (and the un-roped + # Q, since the Q rotation happens after ``prope_qkv``). Replicate that + # dataflow explicitly, out of place and differentiably. + k_cur = k_raw + if rope_freqs_k is not None: + k_cur = _rope_interleaved(k_raw, rope_freqs_k) + + # PRoPE math runs in bhsd. + q_prope, k_prope_bhsd, v_prope_bhsd, apply_fn_o = prope_qkv( + q_raw.transpose(1, 2), + k_cur.transpose(1, 2), + v_raw.transpose(1, 2), + viewmats=viewmats, + Ks=Ks, + ) + + q_rope = q_raw + if rope_freqs_q is not None: + q_rope = _rope_interleaved(q_raw, rope_freqs_q) + + # Standard RoPE branch: [memory K/V, current K/V]. + keys, vals = k_cur, v_raw + if memory_kv_cache is not None and memory_kv_cache.has_rope_kv: + assert memory_kv_cache.k_rope is not None + assert memory_kv_cache.v_rope is not None + keys = torch.cat([memory_kv_cache.k_rope, keys], dim=-3) + vals = torch.cat([memory_kv_cache.v_rope, vals], dim=-3) + out_rope = self.attn_op(q_rope, keys, vals) + out_rope = self.o(out_rope.reshape(batch_shape + (L, n * d))) + + # PRoPE branch; same memory prepend on the camera side. + k_p = k_prope_bhsd.transpose(1, 2) + v_p = v_prope_bhsd.transpose(1, 2) + if memory_kv_cache is not None and memory_kv_cache.has_prope_kv: + assert memory_kv_cache.k_prope is not None + assert memory_kv_cache.v_prope is not None + k_p = torch.cat([memory_kv_cache.k_prope, k_p], dim=-3) + v_p = torch.cat([memory_kv_cache.v_prope, v_p], dim=-3) + out_prope = self.attn_op_prope(q_prope.transpose(1, 2), k_p, v_p) + out_prope = apply_fn_o(out_prope.transpose(1, 2)).transpose(1, 2) + out_prope = self.o_prope(out_prope.reshape(batch_shape + (L, n * d))) + + return out_rope + out_prope + + +def patch_functional_attention() -> None: + """Swap ``forward_dual_branch`` for the functional variant, process-wide.""" + HyWorldPlayPRoPESelfAttention.forward_dual_branch = _functional_dual_branch # type: ignore[assignment] # ty: ignore[invalid-assignment] diff --git a/integrations/hy_worldplay/drift_correction/build_pairs.py b/integrations/hy_worldplay/drift_correction/build_pairs.py new file mode 100644 index 00000000..4e7dd7b4 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/build_pairs.py @@ -0,0 +1,219 @@ +# 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. + +"""Paired drift data from strafe-loop rollouts (clean lap-1 teacher, zero real videos). + +A strafe loop (``w-N, d-N, s-N, a-N`` -- translations only, no yaw) returns +the camera to an identical world pose every ``4 * N`` latents, so every lap +revisits the same views with the same per-frame actions and viewmats. Frame +``j`` of a late (drifted) lap therefore has an exact condition-matched clean +counterpart at frame ``(j mod lap) + lap`` of lap 1 -- the handoff's +"ground-truth-seeded early-window reference", made position- and +camera-consistent by the loop. One long rollout per clip yields both sides +of the pair; no real videos are used anywhere. + +Per-clip output (``outputs/pairs/clip_XXXX.pt``): per-chunk patchified clean +latents (history is their concatenation), per-chunk ``memory_frame_indices``, +the rollout-scoped action / camera buffers, and the recipe metadata. +Resumable: existing clip files are skipped. Run from the repo root:: + + NUM_CLIPS=10 .venv/bin/python integrations/hy_worldplay/drift_correction/build_pairs.py +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import torch +from _rollout import build_runner, capture_rollout + +## Pair-set configuration + +NUM_CLIPS = int(os.environ.get("NUM_CLIPS", "10")) +"""Clips to capture in this invocation (resume-aware).""" + +NUM_LAPS = int(os.environ.get("NUM_LAPS", "6")) +"""Strafe-loop laps per rollout; laps >= 2 supply drifted windows against the +lap-1 teacher.""" + +LEGS = [int(x) for x in os.environ.get("LEGS", "4").split(",")] +"""Motion steps per loop leg, rotated per clip. A single ``4`` gives the +16-latent lap matching the memory budget; mixing (e.g. ``3,4,5`` -> laps of +12/16/20 latents) breaks the fixed revisit period so "content returns at lag +16" stops being a learnable shortcut (the v2 repeat-prior failure).""" + +LAP_PATTERNS = ( + "w-{n}, d-{n}, s-{n}, a-{n}", + "w-{n}, a-{n}, s-{n}, d-{n}", + "d-{n}, w-{n}, a-{n}, s-{n}", + "a-{n}, w-{n}, d-{n}, s-{n}", +) +"""Closed strafe loops (translation order permutations); rotated per clip for +trajectory variety.""" + +IMAGES_DIR = os.environ.get("IMAGES_DIR", "") +"""Directory of first-frame images; empty -> the upstream sample image for +every clip (seed/pattern variety only).""" + +PROMPTS_FILE = os.environ.get("PROMPTS_FILE", "") +"""Prompt source (one per line, e.g. Self-Forcing's +``MovieGenVideoBench_extended.txt``) for seed-image-matched text conditioning: +``frame_NNNN.png`` seeds (from ``gen_first_frames.py``) roll with line +``NNNN`` of this file instead of the integration's mismatched default +prompt. Other image names keep the default.""" + +CORRECTOR_LORA = os.environ.get("CORRECTOR_LORA", "") +"""Optional LoRA checkpoint; when set, rollouts run with the corrector +merged in at scale 1 (the DAgger round: pairs reflect the states the +deployed corrector actually visits).""" + +OUT_DIR = Path( + os.environ.get( + "OUT_DIR", "integrations/hy_worldplay/drift_correction/outputs/pairs" + ) +) +"""Per-clip pair files plus ``manifest.json``.""" + +_BASE_SEED = 31000 +"""Clip ``i`` rolls with diffusion seed ``_BASE_SEED + i``.""" + + +def loop_pose(pattern: str, leg: int, num_laps: int) -> str: + """Build a multi-lap pose string with ``4 * leg * num_laps - 1`` motion steps. + + The parser prepends an identity pose for the input frame, so the last + leg is shortened by one step to keep total latents = ``num_chunk * 4``. + """ + lap = pattern.format(n=leg) + laps = [lap] * num_laps + head, n = lap.rsplit("-", 1) + assert int(n) == leg + laps[-1] = f"{head}-{leg - 1}" if leg > 1 else ", ".join(lap.split(", ")[:-1]) + return ", ".join(laps) + + +def matched_prompt(image_path: Path | None) -> str | None: + """Return the prompt that generated ``frame_NNNN.png``, else ``None``.""" + if image_path is None or not image_path.stem.startswith("frame_"): + return None + try: + idx = int(image_path.stem.split("_")[1]) + lines = [ + ln.strip() + for ln in Path(PROMPTS_FILE).read_text().splitlines() + if ln.strip() + ] + return lines[idx] + except (ValueError, IndexError, OSError): + return None + + +def main() -> None: + torch.set_grad_enabled(False) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + if IMAGES_DIR: + images_dir = Path(IMAGES_DIR) + images: list[Path | None] = [ + *sorted(images_dir.glob("*.png")), + *sorted(images_dir.glob("*.jpg")), + ] + assert images, f"IMAGES_DIR {images_dir} contains no .png/.jpg files" + else: + images = [None] + + runner = None + manifest: dict[str, dict] = {} + manifest_path = OUT_DIR / "manifest.json" + if manifest_path.exists(): + manifest = json.loads(manifest_path.read_text()) + + for i in range(NUM_CLIPS): + clip_path = OUT_DIR / f"clip_{i:04d}.pt" + if clip_path.exists(): + print(f"clip {i:04d}: exists, skipping", flush=True) + continue + pattern = LAP_PATTERNS[i % len(LAP_PATTERNS)] + leg = LEGS[i % len(LEGS)] + lap_latents = 4 * leg + num_chunk = lap_latents * NUM_LAPS // 4 + pose = loop_pose(pattern, leg, NUM_LAPS) + image_path = images[i % len(images)] + if runner is None: + runner = build_runner( + num_chunk=num_chunk, + pose=pose, + output_dir=OUT_DIR, + image_path=image_path, + ) + if CORRECTOR_LORA: + from _lora import ( + apply_lora, + load_lora, + set_lora_scale, + unwrap_compiled, + ) + + network = unwrap_compiled( + runner.pipeline.diffusion_model.transformer.network + ) + apply_lora(network) + load_lora(network, CORRECTOR_LORA) + set_lora_scale(network, 1.0) + print(f"corrector active: {CORRECTOR_LORA}", flush=True) + else: + runner.config.pose = pose + runner.config.num_chunk = num_chunk + if image_path is not None: + runner.config.image_path = image_path + prompt = matched_prompt(image_path) + if prompt: + runner.config.prompt = prompt + seed = _BASE_SEED + i + print(f"clip {i:04d}: pose[{pattern}] seed={seed} ...", flush=True) + snaps = capture_rollout(runner, noise_seed=seed) + + ctrl0 = snaps[1].ctrl # chunk >= 1 carries the rollout-scoped buffers + assert ( + ctrl0.rollout_action is not None + and ctrl0.rollout_viewmats is not None + and ctrl0.rollout_Ks is not None + ), "chunk >= 1 ctrl must carry the rollout-scoped buffers" + torch.save( + { + "latents": torch.cat( + [s.clean_latent.to(torch.float16).cpu() for s in snaps], dim=-2 + ), + "memory_frame_indices": [s.ctrl.memory_frame_indices for s in snaps], + "rollout_action": ctrl0.rollout_action.cpu(), + "rollout_viewmats": ctrl0.rollout_viewmats.to(torch.float16).cpu(), + "rollout_Ks": ctrl0.rollout_Ks.to(torch.float16).cpu(), + "lap_latents": lap_latents, + "num_chunk": num_chunk, + "pose": pose, + "seed": seed, + "image": str(image_path) if image_path else "example", + }, + clip_path, + ) + manifest[f"{i:04d}"] = {"pose": pose, "seed": seed, "num_chunk": num_chunk} + manifest_path.write_text(json.dumps(manifest, indent=2)) + print(f"clip {i:04d}: saved {clip_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/demo_prompts.txt b/integrations/hy_worldplay/drift_correction/demo_prompts.txt new file mode 100644 index 00000000..16bfb3cb --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/demo_prompts.txt @@ -0,0 +1,8 @@ +A street performer juggles glowing pins on a stone plaza while pedestrians stroll past behind him, locked-off camera, late afternoon light. +A golden retriever sprints across a green lawn chasing a bouncing red ball, trees and a white fence in the background, static camera. +A large national flag ripples in strong wind on a pole in front of a weathered stone facade, fixed camera framing. +A tiered fountain in a sunlit courtyard, water jets pulsing and pigeons fluttering around the rim, camera locked off. +A cyclist in a yellow jacket rides across the frame on a riverside path, barges moving slowly on the river behind, static wide shot. +Two children fly a red kite in a park meadow, the kite looping against scattered clouds, camera fixed on a tripod. +A sailboat glides across a calm bay from right to left, a lighthouse and rocky shore in the background, static camera. +An empty cobblestone alley between old brick buildings, laundry lines overhead; a red delivery scooter enters from the left and drives through the frame, static camera. diff --git a/integrations/hy_worldplay/drift_correction/demo_static.py b/integrations/hy_worldplay/drift_correction/demo_static.py new file mode 100644 index 00000000..13f32440 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/demo_static.py @@ -0,0 +1,166 @@ +# 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. + +"""Static-background demo rollouts: locked-off camera, moving foreground. + +The flagship corrector use case: a static camera aligns *with* the anchoring +pull (background stability is the point), so the demo checks that foreground +motion and story flow survive it. Each scene from ``demo_prompts.txt`` rolls +under an all-identity pose (action class 0 throughout) for every gain in +``GAINS``, seeded by the matching ``first_frames/demo`` image. The last +scene is the new-element entrance probe (the anchoring pull resists novel +content — verify it doesn't). Scoring is a separate pass (``score_drift.py`` +with ``EVAL_OUT`` pointed here); labeled side-by-sides come from +``make_sbs.py``. + +Run from the repo root:: + + GAINS=0,0.7,1.0 NUM_CHUNK=24 \ + .venv/bin/python integrations/hy_worldplay/drift_correction/demo_static.py +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np +import torch +from _rollout import ( + build_runner, + capture_rollout, + install_alpha_gate, + parse_gain_token, +) + +## Demo configuration + +_BASE = Path("integrations/hy_worldplay/drift_correction") + +PROMPTS_FILE = Path(os.environ.get("PROMPTS_FILE", str(_BASE / "demo_prompts.txt"))) +"""One scene prompt per line; line ``i`` pairs with ``frame_{i:04d}.png``.""" + +FRAMES_DIR = Path( + os.environ.get("FRAMES_DIR", str(_BASE / "outputs/first_frames/demo")) +) +"""Seed frames from ``gen_first_frames.py`` run on ``PROMPTS_FILE``.""" + +LORA = os.environ.get("LORA", str(_BASE / "outputs/lora_v2.pt")) +"""Corrector checkpoint; gain 0 rows double as the base config.""" + +GAINS = tuple( + parse_gain_token(g) + for g in os.environ.get("GAINS", "0,0.7,1.0").split(",") + if g.strip() +) +"""Deployed LoRA gains; 0 = base, ``gate``/``gate0.5`` = per-step +``alpha*(t) x scale``. The entrance scene needs corrector-off/on/strong.""" + +NUM_CHUNK = int(os.environ.get("NUM_CHUNK", "24")) +"""Rollout horizon (24 chunks matches the T1 eval, ~19 s).""" + +SEED = int(os.environ.get("SEED", "5042")) +"""Diffusion seed, matched across configs.""" + +OUT_DIR = Path(os.environ.get("DEMO_OUT", str(_BASE / "outputs/demo_static"))) + +_DEFAULT_INTRINSIC = [ + [969.6969696969696, 0.0, 960.0], + [0.0, 969.6969696969696, 540.0], + [0.0, 0.0, 1.0], +] +"""Same 1920x1080 intrinsic the pose-string parser stamps on every frame.""" + + +def write_static_pose(n_latents: int, path: Path) -> Path: + """Write an all-identity pose JSON covering ``n_latents`` latents. + + Identity extrinsics give zero relative motion, so the action labels + resolve to class 0 (no translation, no rotation) — a locked-off camera + in the upstream grammar, which has no explicit "stay" token. + """ + eye = np.eye(4).tolist() + poses = { + str(i): {"extrinsic": eye, "K": _DEFAULT_INTRINSIC} + for i in range(n_latents + 1) + } + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(poses)) + return path + + +def main() -> None: + torch.set_grad_enabled(False) + prompts = [ + line.strip() for line in PROMPTS_FILE.read_text().splitlines() if line.strip() + ] + frames = [FRAMES_DIR / f"frame_{i:04d}.png" for i in range(len(prompts))] + missing = [f.name for f in frames if not f.exists()] + assert not missing, ( + f"seed frames missing under {FRAMES_DIR}: {missing}; run " + "gen_first_frames.py with PROMPTS_FILE/OUT_DIR pointed at the demo set." + ) + + pose_json = write_static_pose(NUM_CHUNK * 4, OUT_DIR / "static_pose.json") + + runner = None + network = None + mode: dict = {"gain": 0.0} + for gain in GAINS: + if isinstance(gain, tuple): + config = ( + "corrgate" + if gain[1] == 1.0 + else f"corrgate{gain[1]:.2f}".replace(".", "") + ) + else: + config = "base" if gain == 0 else f"corr{gain:.2f}".replace(".", "") + for i, (prompt, image_path) in enumerate(zip(prompts, frames)): + mp4 = OUT_DIR / config / f"scene{i}_s{SEED}.mp4" + if mp4.exists(): + print(f"{config}/{mp4.stem}: exists, skipping", flush=True) + continue + if runner is None: + runner = build_runner( + num_chunk=NUM_CHUNK, + pose=str(pose_json), + output_dir=OUT_DIR, + image_path=image_path, + ) + from _lora import apply_lora, load_lora, unwrap_compiled + + network = unwrap_compiled( + runner.pipeline.diffusion_model.transformer.network + ) + apply_lora(network) + load_lora(network, LORA) + install_alpha_gate(runner, network, mode) + else: + runner.config.image_path = image_path + from _lora import set_lora_scale + + mode["gain"] = gain + if not isinstance(gain, tuple): + assert network is not None # bound with the runner on first build + set_lora_scale(network, gain) + runner.config.prompt = prompt + print(f"{config}/{mp4.stem}: rolling {NUM_CHUNK} chunks ...", flush=True) + capture_rollout(runner, noise_seed=SEED, mp4_path=mp4) + print(f"done -> {OUT_DIR}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/eval_rollouts.py b/integrations/hy_worldplay/drift_correction/eval_rollouts.py new file mode 100644 index 00000000..68f71a92 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/eval_rollouts.py @@ -0,0 +1,152 @@ +# 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. + +"""Paired closed-loop eval rollouts: frozen base vs +corrector at scale 1. + +Generates long OOD rollouts (held-out first frames, non-loop trajectories, +matched seeds) for each config and writes MP4s under +``outputs/eval/{base,corr}/``. Scoring (Delta-MUSIQ, dynamic-degree guard, +frame strips) is a separate pass -- see ``score_drift.py``. + +Run from the repo root:: + + NUM_CHUNK=40 LORA=outputs/lora_v1_pilot.pt \ + .venv/bin/python integrations/hy_worldplay/drift_correction/eval_rollouts.py +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch +from _rollout import build_runner, capture_rollout, install_alpha_gate + +## Eval configuration + +_BASE = Path("integrations/hy_worldplay/drift_correction") + +LORA = os.environ.get("LORA", str(_BASE / "outputs/lora_v1_pilot.pt")) +"""Corrector checkpoint for the ``corr`` config.""" + +NUM_CHUNK = int(os.environ.get("NUM_CHUNK", "40")) +"""Rollout horizon (40 chunks = 160 latents = ~32s at 16 fps).""" + +POSES = ( + "w-40, right-3, w-40, left-3, w-73", + "w-30, d-10, w-30, a-10, w-79", + "w-60, right-2, w-50, right-2, w-45", +) +"""Non-loop eval trajectories (159 steps each = ``NUM_CHUNK * 4 - 1``); +distinct from the training strafe loops.""" + +IMAGES_DIR = os.environ.get("IMAGES_DIR", "") +"""Held-out first frames; empty -> the upstream sample image.""" + +SEEDS = tuple(int(s) for s in os.environ.get("SEEDS", "5042,5043").split(",")) +"""Diffusion seeds; every (pose, image, seed) cell runs in every config.""" + +N_IMAGES = int(os.environ.get("N_IMAGES", "0")) +"""Cap on held-out images (0 = all); finals trade breadth for horizon.""" + +PROMPT = os.environ.get("PROMPT", "") +"""Text-prompt override. The integration default ("ancient Athens") is +mismatched to most seed images and tugs content toward off-scene +structures/figures; pass a scene-matched or neutral prompt for evals.""" + +OUT_DIR = Path(os.environ.get("EVAL_OUT", str(_BASE / "outputs/eval"))) + + +def main() -> None: + torch.set_grad_enabled(False) + images = [None] + if IMAGES_DIR: + p = Path(IMAGES_DIR) + images = sorted(p.glob("*.png")) + sorted(p.glob("*.jpg")) + assert images, f"no images under {p}" + if N_IMAGES: + images = images[:N_IMAGES] + + # Config name -> deployed LoRA gain. ``GAINS=0.7,0.85`` adds partial-gain + # rows (named ``corr070`` etc.); ``TGATE=1`` adds the per-step + # alpha*(t)-gated row, ``TGATE=1,0.5`` also the gate x 0.5 composition + # (``corrgate050``). + configs: dict[str, float | tuple[str, float]] = {"base": 0.0, "corr": 1.0} + for g in os.environ.get("GAINS", "").split(","): + if g.strip(): + configs[f"corr{float(g):.2f}".replace(".", "")] = float(g) + for s in os.environ.get("TGATE", "").split(","): + if s.strip(): + scale = float(s) + name = ( + "corrgate" if scale == 1.0 else f"corrgate{scale:.2f}".replace(".", "") + ) + configs[name] = ("gate", scale) + + runner = None + network = None + mode: dict = {"gain": 0.0} + for config, gain in configs.items(): + for pi, pose in enumerate(POSES): + for ii, image_path in enumerate(images): + for seed in SEEDS: + name = f"p{pi}_i{ii}_s{seed}" + mp4 = OUT_DIR / config / f"{name}.mp4" + if mp4.exists(): + print(f"{config}/{name}: exists, skipping", flush=True) + continue + if runner is None: + runner = build_runner( + num_chunk=NUM_CHUNK, + pose=pose, + output_dir=OUT_DIR, + image_path=image_path, + ) + from _lora import ( + apply_lora, + load_lora, + set_lora_scale, + unwrap_compiled, + ) + + network = unwrap_compiled( + runner.pipeline.diffusion_model.transformer.network + ) + apply_lora(network) + load_lora(network, LORA) + if PROMPT: + runner.config.prompt = PROMPT + install_alpha_gate(runner, network, mode) + else: + runner.config.pose = pose + if image_path is not None: + runner.config.image_path = image_path + from _lora import set_lora_scale + + mode["gain"] = gain + if not isinstance(gain, tuple): + assert ( + network is not None + ) # bound with the runner on first build + set_lora_scale(network, gain) + print( + f"{config}/{name}: rolling {NUM_CHUNK} chunks ...", flush=True + ) + capture_rollout(runner, noise_seed=seed, mp4_path=mp4) + print(f"done -> {OUT_DIR}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/gate_faithful.py b/integrations/hy_worldplay/drift_correction/gate_faithful.py new file mode 100644 index 00000000..bbc282f5 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/gate_faithful.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. + +"""Faithful step-0 gate: alpha* on real drifted-vs-clean loop pairs. + +Unlike :mod:`gate_systematicity` (proxy perturbations), this measures the +method's actual premise on :mod:`build_pairs` output: at a late chunk of a +strafe-loop rollout, the drifted history's memory frames are swapped for +their lap-aligned lap-``CLEAN_LAP`` counterparts -- same actions, same +viewmats, same RoPE positions, clean content -- and the x0 prediction gap is +decomposed over noise seeds into systematic bias vs variance. Also reports +the drift-gap magnitude ``rel`` (the training loss denominator); near-zero +means no signal to train on. + +Run after ``build_pairs.py`` from the repo root:: + + .venv/bin/python integrations/hy_worldplay/drift_correction/gate_faithful.py +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import cast + +import torch +from _pairs import ( + chunk_x0, + clean_counterfactual, + history_of, + load_clip, + make_ctrl, +) +from _rollout import build_runner, finish_probe_chunk, predict_x0, start_probe_chunk +from hy_worldplay._action import HyWorldPlayWan21TransformerCache +from hy_worldplay.runner import _resolve_prompt, preprocess_first_frame +from torch import Tensor + +## Gate configuration + +PAIRS_DIR = Path("integrations/hy_worldplay/drift_correction/outputs/pairs") +"""Clip files from ``build_pairs.py``.""" + +OUT_PATH = Path( + "integrations/hy_worldplay/drift_correction/outputs/gate/gate_faithful.json" +) +"""Aggregated per-timestep results.""" + +PROBE_CHUNKS = (12, 16, 20, 23) +"""Late chunks probed per clip (laps 3+ of a 6-lap / 24-chunk rollout).""" + +M_NOISE = 8 +"""Noise seeds per (clip, chunk, timestep) cell.""" + +CLEAN_LAP = 1 +"""Lap supplying the clean teacher content. Lap 0 is even cleaner but its +first frame is the stamped real-image latent (different distribution).""" + + +def main() -> None: + torch.set_grad_enabled(False) + clips = sorted(PAIRS_DIR.glob("clip_*.pt")) + assert clips, f"no clips under {PAIRS_DIR}; run build_pairs.py first" + + meta = torch.load(clips[0], map_location="cpu", weights_only=False) + runner = build_runner( + num_chunk=meta["num_chunk"], pose=meta["pose"], output_dir=PAIRS_DIR + ) + pipe = runner.pipeline + device = next(pipe.parameters()).device + dtype = next(pipe.parameters()).dtype + cfg = runner.config + assert cfg.image_path is not None # build_runner resolves the sample image + image = preprocess_first_frame( + cfg.image_path, cfg.pixel_height, cfg.pixel_width + ).to(device=device, dtype=dtype) + cache = pipe.initialize_cache(text=[_resolve_prompt(cfg.prompt)], image=image) + tc = cache.transformer_cache + assert isinstance(tc, HyWorldPlayWan21TransformerCache) + transformer = pipe.diffusion_model.transformer + scheduler = pipe.diffusion_model.scheduler + timesteps = cast(Tensor, scheduler.timesteps) + sigmas = cast(Tensor, scheduler.sigmas) + n_steps = len(timesteps) - 1 + + # Accumulate squared-bias / variance / norms per timestep across all + # (clip, chunk) cells. + agg = {i: {"alphas": [], "alphas_ub": [], "rels": []} for i in range(n_steps)} + for clip_path in clips: + d = load_clip(clip_path, device, dtype) + lap = d["lap_latents"] + for k in PROBE_CHUNKS: + if k >= d["num_chunk"]: + continue + selected = d["memory_frame_indices"][k] + if not selected: + continue + ctrl = make_ctrl(d, k, device=device, dtype=dtype) + h_gen = history_of(d, k) + h_clean = clean_counterfactual( + h_gen, selected=selected, lap_latents=lap, clean_lap=CLEAN_LAP + ) + x0 = chunk_x0(d, k) + + z_ts: dict[int, list[Tensor]] = {} + for t_idx in range(n_steps): + sig = sigmas[t_idx].to(dtype) + z_ts[t_idx] = [] + for m in range(M_NOISE): + g = torch.Generator(device=device).manual_seed( + 900_000 + 10_000 * k + 100 * t_idx + m + ) + eps = torch.randn(x0.shape, device=device, dtype=dtype, generator=g) + z_ts[t_idx].append((1 - sig) * x0 + sig * eps) + + preds: dict[str, dict[int, list[Tensor]]] = {} + for name, h in (("gen", h_gen), ("clean", h_clean)): + start_probe_chunk(tc, ar_idx=k, history=h) + preds[name] = { + t_idx: [ + predict_x0( + transformer, + tc, + ctrl=ctrl, + z_t=z, + timestep=timesteps[t_idx].to(dtype), + sigma=float(sigmas[t_idx]), + ) + for z in z_list + ] + for t_idx, z_list in z_ts.items() + } + finish_probe_chunk(tc, ar_idx=k) + + line = [f"{clip_path.stem} k={k:2d}"] + for t_idx in range(n_steps): + deltas = torch.stack( + [a - b for a, b in zip(preds["gen"][t_idx], preds["clean"][t_idx])] + ) + m = deltas.shape[0] + bias = deltas.mean(dim=0) + bias_sq = bias.square().sum().item() + sq_dev = (deltas - bias).square().sum(dim=tuple(range(1, deltas.ndim))) + var = sq_dev.mean().item() + var_ub = sq_dev.sum().item() / (m - 1) + bias_sq_ub = max(0.0, bias_sq - var_ub / m) + gen_norm = ( + torch.stack(preds["gen"][t_idx]).flatten(1).norm(dim=1).mean() + ) + rel = (deltas.flatten(1).norm(dim=1).mean() / (gen_norm + 1e-12)).item() + a = bias_sq / (bias_sq + var + 1e-12) + a_ub = bias_sq_ub / (bias_sq_ub + var_ub + 1e-12) + agg[t_idx]["alphas"].append(a) + agg[t_idx]["alphas_ub"].append(a_ub) + agg[t_idx]["rels"].append(rel) + line.append( + f"t={int(timesteps[t_idx]):4d} a*={a:.3f}/{a_ub:.3f} rel={rel:.3f}" + ) + print(" | ".join(line), flush=True) + + summary = { + str(int(timesteps[i])): { + "alpha_star": sum(v["alphas"]) / len(v["alphas"]), + "alpha_star_unbiased": sum(v["alphas_ub"]) / len(v["alphas_ub"]), + "rel": sum(v["rels"]) / len(v["rels"]), + "cells": len(v["alphas"]), + } + for i, v in agg.items() + if v["alphas"] + } + OUT_PATH.parent.mkdir(parents=True, exist_ok=True) + OUT_PATH.write_text(json.dumps(summary, indent=2)) + + print("\n========== HY-WorldPlay FAITHFUL gate (real drift pairs) ==========") + for t, v in summary.items(): + print( + f"t={t:>4s}: alpha* {v['alpha_star']:.3f} (unbiased {v['alpha_star_unbiased']:.3f})" + f" | rel drift gap {v['rel']:.3f} | {v['cells']} cells" + ) + all_ub = [a for v in agg.values() for a in v["alphas_ub"]] + mean_rel = sum(r for v in agg.values() for r in v["rels"]) / max( + 1, sum(len(v["rels"]) for v in agg.values()) + ) + frac = sum(a >= 0.7 for a in all_ub) / len(all_ub) + print( + f"cells with unbiased alpha* >= 0.7: {frac:.0%} | mean rel gap {mean_rel:.3f}" + ) + if mean_rel < 0.01: + print("RESULT: drift gap ~zero -> nothing to correct. STOP.") + elif frac >= 0.7: + print("RESULT: real drift gap is systematic. GO for v1 training.") + else: + print( + "RESULT: below the reference bar -- report before committing to training." + ) + print(f"saved {OUT_PATH}") + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/gate_systematicity.py b/integrations/hy_worldplay/drift_correction/gate_systematicity.py new file mode 100644 index 00000000..5ca1f003 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/gate_systematicity.py @@ -0,0 +1,323 @@ +# 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. + +"""Step-0 go/no-go gate: history-sensitivity + systematicity (alpha*) on HY-WorldPlay. + +Counterfactual Forcing assumes the prediction error induced by a perturbed +history at a matched noisy state ``z_t`` is *systematic* -- reproducible +across the noise seeds of ``z_t`` -- so a deterministic corrector can remove +it. This gate measures, per distillation timestep ``t`` and chunk ``k``:: + + delta_m = x0(z_t^m, h_gen) - x0(z_t^m, h_alt) m = 1..M noise seeds + alpha*(t) = ||mean_m delta_m||^2 / (||mean_m delta_m||^2 + var_m) + rel(t) = mean_m ||delta_m|| / mean_m ||x0(z_t^m, h_gen)|| + +for three history perturbations ``h_alt``: a condition-matched history from an +independent same-pose rollout (``cross_seed``), and drift-like per-channel +scale/shift corruptions at two strengths. ``rel`` near zero means the +prediction ignores history (corrector cannot help -- STOP); ``alpha*`` >= 0.7 +at most cells means the history-induced gap is systematic (GO). Reference +host measured 0.91-0.99. + +Resumable: rollout captures are cached under ``outputs/gate/`` and reused. +Run from the repo root:: + + .venv/bin/python integrations/hy_worldplay/drift_correction/gate_systematicity.py +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import cast + +import torch +from _rollout import ( + ChunkSnapshot, + build_runner, + capture_rollout, + finish_probe_chunk, + load_rollout, + predict_x0, + start_probe_chunk, +) +from hy_worldplay._action import HyWorldPlayWan21TransformerCache +from hy_worldplay.runner import _resolve_prompt, preprocess_first_frame +from torch import Tensor + +## Gate configuration + +NUM_CHUNK = 12 +"""Rollout horizon in AR chunks (48 latent frames); long enough that late +chunks condition on FOV-selected memory with visible accumulated drift.""" + +POSE = "w-20, right-2, w-10, left-2, w-13" +"""47 motion steps -> 48 latents, exactly ``NUM_CHUNK * 4``. Mostly-forward +with two gentle turns so the FOV selector sees a production-like trajectory.""" + +ROLLOUT_SEEDS = (42, 1042) +"""Diffusion-noise seeds for the two independent rollouts; rollout A's history +is ``h_gen``, rollout B's supplies the condition-matched ``cross_seed`` +counterfactual.""" + +PROBE_CHUNKS = (4, 8, 11) +"""Chunks probed: first FOV-selected chunk (frame 16), mid, and last.""" + +M_NOISE = 8 +"""Noise seeds per (chunk, timestep) cell for the bias/variance decomposition. +The naive alpha* estimator inflates by ``(1 - alpha*) / M``; 8 seeds plus the +unbiased estimator below keep borderline cells out of threshold noise.""" + +CORRUPT_STRENGTHS = (0.15, 0.30) +"""Per-channel scale/shift corruption strengths (drift-like proxy, matches the +reference gate's ``drift_corrupt``).""" + +LATENT_CHANNELS = 48 +"""Wan 2.2 TI2V-5B latent channels; the patchified feature axis is laid out +``(c kt kh kw)``, so channel ``c`` owns the contiguous 4-slot span ``4c:4c+4``.""" + +OUT_DIR = Path("integrations/hy_worldplay/drift_correction/outputs/gate") +"""Gate artifacts: rollout captures, MP4s, and the results JSON.""" + + +## History perturbations + + +def corrupt_history(history: Tensor, strength: float, seed: int) -> Tensor: + """Apply drift-like per-latent-channel scale + shift to a patchified history. + + Args: + history: Patchified history ``[..., L, C * kt * kh * kw]``. + strength: Relative scale/shift magnitude. + seed: Generator seed so each (chunk, strength) cell perturbs identically + across noise seeds. + + Returns: + Corrupted history with the same shape/dtype. + """ + g = torch.Generator(device=history.device).manual_seed(seed) + spatial = history.shape[-1] // LATENT_CHANNELS + h = history.view(*history.shape[:-1], LATENT_CHANNELS, spatial) + scale = 1 + strength * torch.randn( + LATENT_CHANNELS, 1, device=history.device, dtype=history.dtype, generator=g + ) + shift = ( + strength + * history.float().std().to(history.dtype) + * torch.randn( + LATENT_CHANNELS, 1, device=history.device, dtype=history.dtype, generator=g + ) + ) + return (h * scale + shift).view_as(history) + + +## Probe sweep + + +@torch.no_grad() +def probe_cell( + transformer, + tc: HyWorldPlayWan21TransformerCache, + snap: ChunkSnapshot, + *, + ar_idx: int, + history: Tensor, + z_ts: dict[int, list[Tensor]], + timesteps: Tensor, + sigmas: Tensor, + dtype: torch.dtype, +) -> dict[int, list[Tensor]]: + """Predict x0 for every (timestep, noise-seed) probe under one history. + + One ``start_probe_chunk`` per history keeps a single memory-KV prefill + for the whole sweep, matching how the sampler reuses it across steps. + + Returns: + ``{t_idx: [x0 per noise seed]}`` fp32 tensors. + """ + start_probe_chunk(tc, ar_idx=ar_idx, history=history) + out: dict[int, list[Tensor]] = {} + for t_idx, z_list in z_ts.items(): + t = timesteps[t_idx].to(dtype) + sigma = float(sigmas[t_idx]) + out[t_idx] = [ + predict_x0(transformer, tc, ctrl=snap.ctrl, z_t=z, timestep=t, sigma=sigma) + for z in z_list + ] + finish_probe_chunk(tc, ar_idx=ar_idx) + return out + + +def main() -> None: + # The gate is inference-only; predict_flow is called outside the + # pipeline's no_grad-decorated entry points, and a taped 30-block + # forward retains ~50 GiB of activations. + torch.set_grad_enabled(False) + OUT_DIR.mkdir(parents=True, exist_ok=True) + runner = build_runner(num_chunk=NUM_CHUNK, pose=POSE, output_dir=OUT_DIR) + pipe = runner.pipeline + device = next(pipe.parameters()).device + dtype = next(pipe.parameters()).dtype + + # Capture (or reload) the two independent rollouts. + rollouts: dict[int, list[ChunkSnapshot]] = {} + for seed in ROLLOUT_SEEDS: + path = OUT_DIR / f"rollout_seed{seed}.pt" + if path.exists(): + print(f"reusing capture {path}", flush=True) + rollouts[seed] = load_rollout(path, device) + else: + print(f"capturing rollout seed={seed} ...", flush=True) + rollouts[seed] = capture_rollout( + runner, + noise_seed=seed, + save_path=path, + mp4_path=OUT_DIR / f"rollout_seed{seed}.mp4", + ) + snaps_a, snaps_b = rollouts[ROLLOUT_SEEDS[0]], rollouts[ROLLOUT_SEEDS[1]] + + # Fresh cache for probing (text embeddings + layout only; rolling caches + # are reset per probe chunk). + cfg = runner.config + assert cfg.image_path is not None + image = preprocess_first_frame( + cfg.image_path, cfg.pixel_height, cfg.pixel_width + ).to(device=device, dtype=dtype) + cache = pipe.initialize_cache(text=[_resolve_prompt(cfg.prompt)], image=image) + tc = cache.transformer_cache + assert isinstance(tc, HyWorldPlayWan21TransformerCache) + transformer = pipe.diffusion_model.transformer + scheduler = pipe.diffusion_model.scheduler + timesteps = cast(Tensor, scheduler.timesteps) + sigmas = cast(Tensor, scheduler.sigmas) + n_steps = len(timesteps) - 1 # trailing entry is the terminal t=0 + + results: dict[str, dict] = {} + for k in PROBE_CHUNKS: + h_gen = snaps_a[k].history + assert h_gen is not None + variants = { + "cross_seed": snaps_b[k].history, + **{ + f"corrupt_{s:g}": corrupt_history(h_gen, s, seed=1000 * k) + for s in CORRUPT_STRENGTHS + }, + } + + # Shared z_t per (t, m): the anchor state comes from rollout A's own + # x0 for this chunk, re-noised with the host's convention. + x0 = snaps_a[k].clean_latent + z_ts: dict[int, list[Tensor]] = {} + for t_idx in range(n_steps): + sig = sigmas[t_idx].to(dtype) + z_ts[t_idx] = [] + for m in range(M_NOISE): + g = torch.Generator(device=device).manual_seed( + 10_000 * k + 100 * t_idx + m + ) + eps = torch.randn(x0.shape, device=device, dtype=dtype, generator=g) + z_ts[t_idx].append((1 - sig) * x0 + sig * eps) + + x0_gen = probe_cell( + transformer, + tc, + snaps_a[k], + history=h_gen, + ar_idx=k, + z_ts=z_ts, + timesteps=timesteps, + sigmas=sigmas, + dtype=dtype, + ) + results[str(k)] = {} + for name, h_alt in variants.items(): + assert h_alt is not None + x0_alt = probe_cell( + transformer, + tc, + snaps_a[k], + history=h_alt, + ar_idx=k, + z_ts=z_ts, + timesteps=timesteps, + sigmas=sigmas, + dtype=dtype, + ) + alphas, alphas_ub, rels = [], [], [] + for t_idx in range(n_steps): + deltas = torch.stack( + [a - b for a, b in zip(x0_gen[t_idx], x0_alt[t_idx])] + ) + m = deltas.shape[0] + bias = deltas.mean(dim=0) + bias_sq = bias.square().sum().item() + sq_dev = (deltas - bias).square().sum(dim=tuple(range(1, deltas.ndim))) + var = sq_dev.mean().item() + # Unbiased decomposition: E||mean||^2 = ||bias||^2 + var/M and + # E[mean sq dev] = var (M-1)/M, so correct both before the ratio. + var_ub = sq_dev.sum().item() / (m - 1) + bias_sq_ub = max(0.0, bias_sq - var_ub / m) + gen_norm = torch.stack(x0_gen[t_idx]).flatten(1).norm(dim=1).mean() + alphas.append(bias_sq / (bias_sq + var + 1e-12)) + alphas_ub.append(bias_sq_ub / (bias_sq_ub + var_ub + 1e-12)) + rels.append( + (deltas.flatten(1).norm(dim=1).mean() / (gen_norm + 1e-12)).item() + ) + results[str(k)][name] = { + "t": [float(timesteps[i]) for i in range(n_steps)], + "alpha_star": alphas, + "alpha_star_unbiased": alphas_ub, + "rel": rels, + } + print( + f"k={k:2d} {name:14s} | " + + " | ".join( + f"t={int(timesteps[i]):4d} a*={alphas[i]:.3f}" + f"/{alphas_ub[i]:.3f} rel={rels[i]:.3f}" + for i in range(n_steps) + ), + flush=True, + ) + + (OUT_DIR / "gate_results.json").write_text(json.dumps(results, indent=2)) + + # GO / NO-GO summary over every (chunk, variant, timestep) cell. + cells = [ + (a, r) + for per_k in results.values() + for v in per_k.values() + for a, r in zip(v["alpha_star_unbiased"], v["rel"]) + ] + frac_go = sum(a >= 0.7 for a, _ in cells) / len(cells) + mean_alpha = sum(a for a, _ in cells) / len(cells) + mean_rel = sum(r for _, r in cells) / len(cells) + print("\n================ HY-WorldPlay systematicity gate ================") + print( + f"cells with unbiased alpha* >= 0.7: {frac_go:.0%} | mean alpha* " + f"{mean_alpha:.3f} | mean rel gap: {mean_rel:.3f}" + ) + if mean_rel < 0.01: + print("RESULT: x0 ~insensitive to history -> corrector cannot help. STOP.") + elif frac_go >= 0.7: + print("RESULT: history-induced gap is systematic at most steps. GO.") + else: + print( + "RESULT: gap is noise-dominated at many steps. Investigate before training." + ) + print(f"saved {OUT_DIR / 'gate_results.json'}") + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/gen_first_frames.py b/integrations/hy_worldplay/drift_correction/gen_first_frames.py new file mode 100644 index 00000000..0254bbde --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/gen_first_frames.py @@ -0,0 +1,139 @@ +# 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. + +"""Synthesize first-frame images from prompts with the Wan 2.2 TI2V-5B base. + +Keeps the whole corrector stack zero-real-data (the av2s regime): prompts -> +base-model T2V chunk 0 (no history, single-shot within the chunk) -> frame 0 +saved as PNG. These images seed :mod:`build_pairs` rollouts (train split) and +the held-out eval set. + +Prompt source: ``MovieGenVideoBench_extended.txt`` (same family the reference +av2s eval used); train frames draw from the file's head, eval frames from +index 700+ so neither overlaps the reference's 128-prompt eval indices. + +Run from the repo root:: + + N_FRAMES=40 SPLIT=train .venv/bin/python \ + integrations/hy_worldplay/drift_correction/gen_first_frames.py +""" + +from __future__ import annotations + +import os +from pathlib import Path + +# Must precede the first CUDA allocation (shared-GPU headroom). +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + +import torch + +## Generation configuration + +PROMPTS_FILE = Path(os.environ.get("PROMPTS_FILE", "prompts.txt")) +"""One prompt per line (e.g. Self-Forcing's +``MovieGenVideoBench_extended.txt``); set via the ``PROMPTS_FILE`` env var.""" + +N_FRAMES = int(os.environ.get("N_FRAMES", "40")) +SPLIT = os.environ.get("SPLIT", "train") +"""``train`` draws prompts from line 0 up; ``eval`` from line 700 up.""" + +OUT_DIR = Path( + os.environ.get( + "OUT_DIR", + f"integrations/hy_worldplay/drift_correction/outputs/first_frames/{SPLIT}", + ) +) + +HEIGHT, WIDTH = 480, 832 +"""Generation resolution. Below HY-WorldPlay's native 704x1280 on purpose: +the runner resizes/crops seed images anyway, and the full-res eager decode +does not fit next to co-tenant jobs (~17 GiB single permute).""" + +SEED = 7000 if SPLIT == "train" else 8000 +"""Base diffusion seed; frame ``i`` uses ``SEED + i``.""" + +_EVAL_PROMPT_OFFSET = 700 + + +def main() -> None: + torch.set_grad_enabled(False) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + prompts = [ + line.strip() for line in PROMPTS_FILE.read_text().splitlines() if line.strip() + ] + start = 0 if SPLIT == "train" else _EVAL_PROMPT_OFFSET + picked = prompts[start : start + N_FRAMES] + assert len(picked) == N_FRAMES, ( + f"prompt file has only {len(prompts)} lines; cannot draw {N_FRAMES} " + f"from offset {start}." + ) + + from wan22.config import PIPELINE_WAN22_TI2V_5B, WAN22_TI2V_5B_DIT_DIFFUSERS_PATH + + from flashdreams.infra.config import derive_config + + # T2V mode: the base pipeline asserts encoder-None when no image is + # given; VAE graphs off for co-tenant headroom. Upstream sharded the + # diffusers repo, so the single-file constant 404s -- point at the + # sharded index (already in the local HF cache), which the loader + # supports natively. + cfg = derive_config( + PIPELINE_WAN22_TI2V_5B, + encoder=None, + decoder=dict(use_cuda_graph=False), + diffusion_model=dict( + transformer=dict( + checkpoint_path=WAN22_TI2V_5B_DIT_DIFFUSERS_PATH + ".index.json", + # I2V-only mechanisms; both require an I2VCtrl input. + stamp_image_latent=False, + ti2v_first_frame_per_token_timestep=False, + ), + ), + ) + pipe = cfg.setup().to("cuda").eval() + dm = pipe.diffusion_model + device = next(pipe.parameters()).device + + from PIL import Image + + for i, prompt in enumerate(picked): + out_path = OUT_DIR / f"frame_{start + i:04d}.png" + if out_path.exists(): + print(f"{out_path.name}: exists, skipping", flush=True) + continue + dm._rng = torch.Generator(device=device).manual_seed(SEED + i) + # Release the ~11 GiB text encoder after each prompt: the 704x1280 + # eager VAE decode needs the headroom next to co-tenant jobs, and + # the per-prompt reload (~15 s) is cheap at this batch size. + cache = pipe.initialize_cache( + text=[prompt], + image=None, + height=HEIGHT // 16, + width=WIDTH // 16, + ) + video = pipe.generate(0, cache) # [*, T, C, H, W] in [-1, 1] + pipe.finalize(0, cache) + frame = video[0, 0] if video.ndim == 5 else video[0] + frame = ((frame.float().clamp(-1, 1) + 1) * 127.5).round().byte() + Image.fromarray(frame.permute(1, 2, 0).cpu().numpy()).save(out_path) + print(f"{out_path.name}: saved ({prompt[:60]}...)", flush=True) + del cache + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/make_sbs.py b/integrations/hy_worldplay/drift_correction/make_sbs.py new file mode 100644 index 00000000..96e593d6 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/make_sbs.py @@ -0,0 +1,87 @@ +# 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. + +"""Labeled side-by-side MP4s for owner eyeball verdicts (ffmpeg hstack). + +Pairs every non-``base`` config's videos against the matching ``base`` +video in an eval/demo directory and writes +``sbs_base_LEFT_vs_{config}_RIGHT_{name}.mp4`` next to them — which side +is which is in the filename (owner convention; the box ffmpeg has no +``drawtext``, so no burned-in labels). + +Run from the repo root:: + + EVAL_OUT=.../outputs/demo_static \ + python integrations/hy_worldplay/drift_correction/make_sbs.py +""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +EVAL_DIR = Path( + os.environ.get( + "EVAL_OUT", "integrations/hy_worldplay/drift_correction/outputs/eval" + ) +) + + +def hstack(left: Path, right: Path, out: Path) -> None: + """Write a horizontal side-by-side of two same-size videos.""" + subprocess.run( + [ + "ffmpeg", + "-y", + "-loglevel", + "error", + "-i", + str(left), + "-i", + str(right), + "-filter_complex", + "[0:v][1:v]hstack", + "-c:v", + "libx264", + "-crf", + "18", + str(out), + ], + check=True, + ) + + +def main() -> None: + base_dir = EVAL_DIR / "base" + assert base_dir.is_dir(), f"no base config under {EVAL_DIR}" + configs = [ + d.name for d in sorted(EVAL_DIR.iterdir()) if d.is_dir() and d.name != "base" + ] + for config in configs: + for mp4 in sorted((EVAL_DIR / config).glob("*.mp4")): + base_mp4 = base_dir / mp4.name + if not base_mp4.exists(): + print(f"{config}/{mp4.name}: no base counterpart, skipping") + continue + out = EVAL_DIR / f"sbs_base_LEFT_vs_{config}_RIGHT_{mp4.stem}.mp4" + if out.exists(): + continue + hstack(base_mp4, mp4, out) + print(f"wrote {out.name}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/score_drift.py b/integrations/hy_worldplay/drift_correction/score_drift.py new file mode 100644 index 00000000..4cef338a --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/score_drift.py @@ -0,0 +1,362 @@ +# 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. + +"""Drift + guard scoring for paired eval rollouts (Delta-MUSIQ, dynamics, strips). + +Per MP4: per-frame MUSIQ (pyiqa) at a fixed frame stride, Delta-drift = +MUSIQ(first 20%) - MUSIQ(all) (BAgger's metric), a RAFT mean-flow dynamic +degree (guard: a corrector that freezes motion "wins" every consistency +metric), the reference protocol metrics (DINO latesim anchoring, lag-2s +identity, full-rate cut count), and a 10-frame contact strip PNG. +Aggregates base-vs-corr into ``outputs/eval/scores.json``. + +pyiqa / timm are not part of the project deps; run with an ephemeral overlay:: + + uv run --with pyiqa --with timm python integrations/hy_worldplay/drift_correction/score_drift.py +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np +import torch + +EVAL_DIR = Path( + os.environ.get( + "EVAL_OUT", "integrations/hy_worldplay/drift_correction/outputs/eval" + ) +) +FRAME_STRIDE = 13 +"""Score every 13th frame (one per AR chunk at the decoded rate).""" + +DINO_FPS = 1.0 +"""DINO feature sampling rate — the reference host's ``posthoc_metrics.py`` +samples every 16th frame at 16 fps, so 1 fps keeps latesim / lag-2s on the +cross-host-comparable scale (SF 0.17 / av1s 0.20 / av2s 0.50).""" + + +def read_frames(path: Path, stride: int) -> np.ndarray: + """Decode every ``stride``-th frame of an MP4 to ``[N, H, W, 3]`` uint8.""" + import imageio.v3 as iio + + frames = iio.imread(path, plugin="pyav") + return frames[::stride] + + +def musiq_curve(frames: np.ndarray, model, device) -> list[float]: + """Per-frame MUSIQ scores.""" + scores = [] + for f in frames: + t = torch.from_numpy(f.copy()).permute(2, 0, 1)[None].float().to(device) / 255 + scores.append(float(model(t))) + return scores + + +def sharpness_curve(frames: np.ndarray) -> list[float]: + """Per-frame variance-of-Laplacian (grayscale) — catches blur/structure + loss that MUSIQ under-penalizes (observed on v1-pilot rollouts).""" + from scipy.ndimage import laplace + + out = [] + for f in frames: + gray = f.astype(np.float32).mean(axis=-1) + out.append(float(laplace(gray).var())) + return out + + +def sim_to_start(frames: np.ndarray) -> float: + """Mean last-20% frame-correlation to the first-2s mean frame. + + The reference repo's progression-bias probe (its + ``TODO_progression_bias.md`` convention, cross-host comparable scale: + SF 0.17 / av1s 0.20 / av2s 0.50). Lower = healthy progression, but + near-zero can be collapse — read alongside the quality guards. On a + forward trajectory the view should decorrelate from the start; a + corrector with a learned repeat prior plateaus high instead.""" + g = frames.astype(np.float32).mean(axis=-1)[:, ::4, ::4] + n2s = max(1, min(len(g), 3)) # first ~2s at the chunk-stride sampling + anchor = g[:n2s].mean(axis=0) + anchor = anchor - anchor.mean() + sims = [] + for f in g[-max(1, len(g) // 5) :]: + fc = f - f.mean() + sims.append( + float( + (anchor * fc).sum() + / (np.linalg.norm(anchor) * np.linalg.norm(fc) + 1e-8) + ) + ) + return float(np.mean(sims)) + + +def sat_drift(frames: np.ndarray) -> float: + """Mean |saturation - saturation(frame 0)| (the reference's in-domain + drift metric); catches the color-drift axis MUSIQ tolerates.""" + f = frames.astype(np.float32) / 255 + mx, mn = f.max(axis=-1), f.min(axis=-1) + sat = np.where(mx > 0, (mx - mn) / (mx + 1e-8), 0).mean(axis=(1, 2)) + return float(np.abs(sat - sat[0]).mean()) + + +def dino_features(frames: np.ndarray, model, mean, std, device) -> torch.Tensor: + """L2-normalized DINO features for ``[N, H, W, 3]`` uint8 frames.""" + import torch.nn.functional as F + + feats = [] + for f in frames: + x = torch.from_numpy(f.copy()).permute(2, 0, 1)[None].float().to(device) / 255 + x = F.interpolate(x, size=(224, 224), mode="bicubic", align_corners=False) + feats.append(F.normalize(model((x - mean) / std).float(), dim=-1)) + return torch.cat(feats) + + +def latesim(feats: torch.Tensor) -> float: + """Mean DINO similarity of the last fifth to the opening (first-3 mean). + + The reference protocol's anchoring metric (``posthoc_metrics.py``): a + corrector with an anchoring pull plateaus high while the base + decorrelates. Semantic-feature complement to the pixel-space + ``sim_to_start``; read alongside the dynamics guard. + """ + import torch.nn.functional as F + + ref = F.normalize(feats[:3].mean(0, keepdim=True), dim=-1) + sims = (feats @ ref.T).squeeze(-1).cpu().numpy() + return float(np.mean(sims[-max(1, len(sims) // 5) :])) + + +def lag2s_identity(feats: torch.Tensor, lag: int = 2) -> float: + """Mean DINO similarity between frames ``lag`` seconds apart. + + Identity-persistence probe (subjects morphing lowers it); ``lag`` is in + ``DINO_FPS`` samples. ``nan`` when the clip is shorter than the lag. + """ + if len(feats) <= lag: + return float("nan") + return float((feats[:-lag] * feats[lag:]).sum(-1).mean()) + + +CHUNK_FRAMES = 16 +"""Decoded frames per AR chunk after the first: Wan's 4x temporal VAE maps +the 4-latent seed chunk to 13 frames, then 16 per chunk (381 = 13 + 23*16 +at 24 chunks), so chunk boundaries sit at frame ``13 + 16k`` (per-frame +motion autocorrelation peaks at lag 16).""" + +FIRST_CHUNK_FRAMES = 13 +"""Decoded frames in the seed chunk (excluded from seam phase alignment).""" + + +def seam_motion_ratio(frames: np.ndarray) -> float: + """Motion at chunk-boundary transitions relative to chunk-interior motion. + + Phase-aligns full-rate adjacent-frame diffs to the decoded chunk cadence + (boundaries at frame ``13 + 16k``; seed chunk excluded) and returns + mean(boundary phases 0-1) / mean(interior phases 4-11). ~1.0 = motion + flows through the seams; above 1 = a boundary jump (anchoring kick / + statistics snap); below 1 = motion stalling at the seam. + """ + g = frames.astype(np.float32).mean(axis=-1)[:, ::2, ::2] + mot = np.abs(np.diff(g, axis=0)).mean(axis=(1, 2)) + seg = mot[FIRST_CHUNK_FRAMES - 1 :] # phase 0 = transition into chunk 1 + n = (len(seg) // CHUNK_FRAMES) * CHUNK_FRAMES + if n < 2 * CHUNK_FRAMES: + return float("nan") + phases = seg[:n].reshape(-1, CHUNK_FRAMES).mean(axis=0) + return float(phases[:2].mean() / (phases[4:12].mean() + 1e-9)) + + +def seam_sharpness_ratio(frames: np.ndarray) -> float: + """Sharpness of chunk-initial frames relative to chunk-interior frames. + + Variance-of-Laplacian per frame, phase-aligned to the decoded chunk + cadence (chunk-initial frames at ``13 + 16k``; seed chunk excluded): + mean(phases 0-1) / mean(phases 6-13). Well below 1 = post-boundary blur + (structural detail loss on each chunk's first frames). + """ + from scipy.ndimage import laplace + + g = frames.astype(np.float32).mean(axis=-1)[:, ::2, ::2] + sharp = np.array([float(laplace(f).var()) for f in g[FIRST_CHUNK_FRAMES:]]) + n = (len(sharp) // CHUNK_FRAMES) * CHUNK_FRAMES + if n < 2 * CHUNK_FRAMES: + return float("nan") + phases = sharp[:n].reshape(-1, CHUNK_FRAMES).mean(axis=0) + return float(phases[:2].mean() / (phases[6:14].mean() + 1e-9)) + + +def cut_count(frames: np.ndarray) -> int: + """Count hard cuts: adjacent full-rate mean-abs-diff above + ``max(0.10, mu + 5*sigma)`` (reference convention) — catches HY's + content-level jump-cut at the context/memory handoff.""" + diffs = np.empty(len(frames) - 1, dtype=np.float32) + for i in range(len(frames) - 1): + a = frames[i].astype(np.float32) + b = frames[i + 1].astype(np.float32) + diffs[i] = float(np.abs(b - a).mean() / 255.0) + thr = max(0.10, float(diffs.mean() + 5 * diffs.std())) + return int((diffs > thr).sum()) + + +def dynamic_degree(frames: np.ndarray, raft, device) -> float: + """Mean RAFT flow magnitude between scored frames (motion guard).""" + import torch.nn.functional as F + + mags = [] + for a, b in zip(frames[:-1], frames[1:]): + ta = torch.from_numpy(a.copy()).permute(2, 0, 1)[None].float().to(device) + tb = torch.from_numpy(b.copy()).permute(2, 0, 1)[None].float().to(device) + ta = F.interpolate(ta, size=(352, 640), mode="bilinear") / 127.5 - 1 + tb = F.interpolate(tb, size=(352, 640), mode="bilinear") / 127.5 - 1 + flow = raft(ta, tb)[-1] + mags.append(float(flow.square().sum(1).sqrt().mean())) + return float(np.mean(mags)) + + +def contact_strip(frames: np.ndarray, out_path: Path, n: int = 10) -> None: + """Write an n-frame horizontal strip PNG for eyeballing.""" + from PIL import Image + + idx = np.linspace(0, len(frames) - 1, n).astype(int) + strip = np.concatenate([frames[i] for i in idx], axis=1) + Image.fromarray(strip).save(out_path) + + +def main() -> None: + import imageio.v3 as iio + import pyiqa # ty: ignore[unresolved-import] + import timm # ty: ignore[unresolved-import] + from torchvision.models.optical_flow import Raft_Large_Weights, raft_large + + device = "cuda" if torch.cuda.is_available() else "cpu" + musiq = pyiqa.create_metric("musiq", device=device) + raft = raft_large(weights=Raft_Large_Weights.DEFAULT).to(device).eval() + dino = ( + timm.create_model("vit_base_patch16_224.dino", pretrained=True, num_classes=0) + .to(device) + .eval() + ) + dino_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device) + dino_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device) + + results: dict[str, dict] = {} + configs = [d.name for d in sorted(EVAL_DIR.iterdir()) if d.is_dir()] + for config in configs: + cfg_dir = EVAL_DIR / config + per_video = {} + for mp4 in sorted(cfg_dir.glob("*.mp4")): + if mp4.stem.startswith("sbs_"): + continue + # One full-rate decode per video; strided views feed the + # per-chunk metrics, ~1 fps feeds DINO, full rate feeds cuts. + frames_full = read_frames(mp4, 1) + fps = float(iio.immeta(mp4, plugin="pyav").get("fps", 16.0)) + frames = frames_full[::FRAME_STRIDE] + dino_stride = max(1, round(fps / DINO_FPS)) + with torch.no_grad(): + feats = dino_features( + frames_full[::dino_stride], dino, dino_mean, dino_std, device + ) + curve = musiq_curve(frames, musiq, device) + sharp = sharpness_curve(frames) + n20 = max(1, len(curve) // 5) + delta = float(np.mean(curve[:n20]) - np.mean(curve)) + # Late-vs-early sharpness ratio: ~1.0 = detail preserved; well + # below 1 = the compounding-blur failure mode. + sharp_ratio = float(np.mean(sharp[-n20:]) / (np.mean(sharp[:n20]) + 1e-9)) + with torch.no_grad(): + dyn = dynamic_degree(frames, raft, device) + contact_strip(frames, mp4.with_suffix(".strip.png")) + per_video[mp4.stem] = { + "musiq_overall": float(np.mean(curve)), + "musiq_late": float(np.mean(curve[-n20:])), + "delta_drift": delta, + "dynamic_degree": dyn, + "sharpness_ratio": sharp_ratio, + "sat_drift": sat_drift(frames), + "sim_to_start": sim_to_start(frames), + "latesim": latesim(feats), + "lag2s": lag2s_identity(feats), + "cuts": cut_count(frames_full), + "seam_motion_ratio": seam_motion_ratio(frames_full), + "seam_sharpness_ratio": seam_sharpness_ratio(frames_full), + "curve": curve, + "sharpness": sharp, + } + v = per_video[mp4.stem] + print( + f"{config}/{mp4.stem}: MUSIQ {np.mean(curve):.1f} " + f"(late {np.mean(curve[-n20:]):.1f}) | Delta {delta:+.2f} | " + f"dyn {dyn:.1f} | sharp-ratio {sharp_ratio:.2f} | " + f"latesim {v['latesim']:.3f} | lag2s {v['lag2s']:.3f} | " + f"cuts {v['cuts']}", + flush=True, + ) + agg = { + key: float(np.mean([v[key] for v in per_video.values()])) + for key in ( + "musiq_overall", + "musiq_late", + "delta_drift", + "dynamic_degree", + "sharpness_ratio", + "sat_drift", + "sim_to_start", + "latesim", + "lag2s", + "cuts", + "seam_motion_ratio", + "seam_sharpness_ratio", + ) + } + results[config] = {"videos": per_video, "aggregate": agg} + + (EVAL_DIR / "scores.json").write_text(json.dumps(results, indent=2)) + print("\n================ closed-loop drift eval ================") + print( + f"{'':10s} {'MUSIQ':>7s} {'late':>7s} {'Delta':>7s} {'dyn':>6s} " + f"{'sharp':>6s} {'sat':>7s} {'sim':>5s} {'lsim':>6s} {'lag2':>6s} " + f"{'cuts':>5s} {'seam':>6s} {'sseam':>6s}" + ) + for name in configs: + a = results[name]["aggregate"] + print( + f"{name:10s} {a['musiq_overall']:7.2f} {a['musiq_late']:7.2f} " + f"{a['delta_drift']:+7.2f} {a['dynamic_degree']:6.1f} " + f"{a['sharpness_ratio']:6.2f} {a['sat_drift']:7.4f} " + f"{a['sim_to_start']:5.2f} {a['latesim']:6.3f} {a['lag2s']:6.3f} " + f"{a['cuts']:5.1f} {a['seam_motion_ratio']:6.3f} " + f"{a['seam_sharpness_ratio']:6.3f}" + ) + b = results.get("base", {}).get("aggregate") + for name in configs: + if name == "base" or b is None or b["delta_drift"] <= 0: + continue + red = ( + 100 + * (b["delta_drift"] - results[name]["aggregate"]["delta_drift"]) + / b["delta_drift"] + ) + print(f"{name}: Delta-drift reduction vs base: {red:.0f}% (target >= 30%)") + print( + "guards: dynamic degree within ~20% of base (motion freeze) and " + "sharpness ratio not far below base (compounding blur)." + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/train_v1.py b/integrations/hy_worldplay/drift_correction/train_v1.py new file mode 100644 index 00000000..26741750 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/train_v1.py @@ -0,0 +1,274 @@ +# 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. + +"""v1 corrector trainer: counterfactual clean-history teacher in x0 space. + +Trains the LoRA corrector r_phi on the frozen distilled HY-WorldPlay model +so its x0 prediction under drifted history matches the frozen model's +prediction under the lap-aligned clean history at the same ``z_t``:: + + L = ||x0_{theta+LoRA}(z_t, h_gen, t) - x0_theta(z_t, h_clean, t)||^2 + / ||x0_theta(z_t, h_clean, t) - x0_theta(z_t, h_gen, t)||^2 + +The drift-gap denominator is required (raw MSE diverges); ``R^2 = 1 - L``. +Both passes share the exact ``z_t`` re-noised from the rollout's own x0 +(native anchoring). The student's memory-KV prefill runs LoRA-scaled but +under ``no_grad`` (gradients flow through the current-chunk forward only; +deploy-time behaviour is unaffected since the weights are shared). + +Host adaptations vs the Wan2.1 reference (``wan_train_synth.py``): +x0-space predictions, t sampled from the 4 distillation timesteps, eager +network (``compile_network=False``), functional dual-branch attention +(:mod:`_train_attn` -- restores k/v gradients, enables checkpointing), and +per-block gradient checkpointing (the fp32/dual-branch tape is ~40 GiB +without it). + +Run from the repo root (resumable via ``INIT``):: + + STEPS=1500 .venv/bin/python integrations/hy_worldplay/drift_correction/train_v1.py +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import cast + +import numpy as np +import torch +from _lora import ( + apply_lora, + load_lora, + lora_parameters, + save_lora, + set_lora_scale, + unwrap_compiled, +) +from _pairs import chunk_x0, clean_counterfactual, history_of, load_clip, make_ctrl +from _rollout import build_runner, finish_probe_chunk, start_probe_chunk +from _train_attn import patch_functional_attention +from hy_worldplay._action import ( + HyWorldPlayWan21Transformer, + HyWorldPlayWan21TransformerCache, +) +from hy_worldplay.runner import _resolve_prompt, preprocess_first_frame +from torch import Tensor + +## Training configuration + +PAIRS_DIR = Path( + os.environ.get( + "PAIRS_DIR", "integrations/hy_worldplay/drift_correction/outputs/pairs" + ) +) +"""Clip files from ``build_pairs.py``.""" + +CKPT = Path( + os.environ.get( + "CKPT", "integrations/hy_worldplay/drift_correction/outputs/lora_v1.pt" + ) +) +"""Output LoRA checkpoint; saved every ``SAVE_EVERY`` steps.""" + +INIT = os.environ.get("INIT", "") +"""Optional LoRA checkpoint to resume/init from.""" + +STEPS = int(os.environ.get("STEPS", "1500")) +LR = float(os.environ.get("LR", "5e-4")) +WARMUP = 60 +"""Linear LR warmup steps; required for stability (reference finding).""" + +GRAD_CLIP = 1.0 +RANK = 16 +EVAL_EVERY = 100 +SAVE_EVERY = 250 +CLEAN_LAP = 1 +SEED = int(os.environ.get("SEED", "0")) + + +def checkpoint_blocks(network) -> None: + """Route block ``forward`` and ``prefill_memory_kv`` through checkpointing. + + Per-instance overrides (not wrapper modules) so the network loop's + ``isinstance(block, Block)`` assertion keeps passing. Requires the + functional-attention patch: forward recomputation must be side-effect + free. The prefill's ``write_rope`` / ``write_prope`` side effects are + plain re-assignments of deterministically recomputed tensors, so its + recompute is value-stable. No-op under ``no_grad`` passes. + """ + from torch.utils.checkpoint import checkpoint + + def wrap(fn): + def ckpt_fn(*args, _inner=fn, **kwargs): + if not torch.is_grad_enabled(): + return _inner(*args, **kwargs) + return checkpoint(_inner, *args, use_reentrant=False, **kwargs) + + return ckpt_fn + + for block in network.blocks: + block.forward = wrap(block.forward) + block.prefill_memory_kv = wrap(block.prefill_memory_kv) + + +def main() -> None: + clips = sorted(PAIRS_DIR.glob("clip_*.pt")) + assert clips, f"no clips under {PAIRS_DIR}; run build_pairs.py first" + rng = np.random.default_rng(SEED) + + meta = torch.load(clips[0], map_location="cpu", weights_only=False) + runner = build_runner( + num_chunk=meta["num_chunk"], + pose=meta["pose"], + output_dir=CKPT.parent, + compile_network=False, + ) + pipe = runner.pipeline + device = next(pipe.parameters()).device + dtype = next(pipe.parameters()).dtype + transformer = cast(HyWorldPlayWan21Transformer, pipe.diffusion_model.transformer) + scheduler = pipe.diffusion_model.scheduler + timesteps = cast(Tensor, scheduler.timesteps) + sigmas = cast(Tensor, scheduler.sigmas) + n_steps = len(timesteps) - 1 + + cfg = runner.config + assert cfg.image_path is not None # build_runner resolves the sample image + image = preprocess_first_frame( + cfg.image_path, cfg.pixel_height, cfg.pixel_width + ).to(device=device, dtype=dtype) + cache = pipe.initialize_cache(text=[_resolve_prompt(cfg.prompt)], image=image) + tc = cache.transformer_cache + assert isinstance(tc, HyWorldPlayWan21TransformerCache) + + network = unwrap_compiled(transformer.network) + wrapped = apply_lora(network, rank=RANK) + if INIT: + load_lora(network, INIT) + print(f"LoRA init from {INIT}", flush=True) + # Functional attention: restores k/v gradients (the stock path severs + # them at the rolling-cache write) and makes per-block checkpointing + # sound (no mutable storage in the tape -- the fp32 residual/dual-branch + # tape is ~40 GiB unchunked). + patch_functional_attention() + checkpoint_blocks(network) + params = lora_parameters(network) + n_params = sum(p.numel() for p in params) + print( + f"LoRA on {len(wrapped)} projections | {n_params / 1e6:.2f}M params", flush=True + ) + opt = torch.optim.AdamW(params, lr=LR) + + # Clip latents stay on CPU; a draw moves one clip's tensors to device. + datas = [load_clip(p, "cpu", dtype) for p in clips] + nval = max(2, len(datas) // 8) + train_ids = list(range(len(datas) - nval)) + val_ids = list(range(len(datas) - nval, len(datas))) + lap_chunks = datas[0]["lap_latents"] // 4 + num_chunk = datas[0]["num_chunk"] + # Windows fully past the clean lap: selected recent-16 frames all drifted. + ks = list(range((CLEAN_LAP + 2) * lap_chunks, num_chunk)) + print( + f"{len(datas)} clips ({len(val_ids)} val) | k in {ks[0]}..{ks[-1]}", flush=True + ) + + def predict(ctrl, z_t: Tensor, t_idx: int) -> Tensor: + t = timesteps[t_idx].to(dtype) + flow = transformer.predict_flow( + noisy_latent=z_t, timestep=t, cache=tc, input=ctrl + ) + return z_t.float() - float(sigmas[t_idx]) * flow.float() + + def prefill_only(ctrl, history: Tensor, k: int) -> None: + """Open the chunk bracket and run the memory prefill without grad.""" + start_probe_chunk(tc, ar_idx=k, history=history) + with torch.no_grad(): + if ctrl.memory_frame_indices: + transformer.prefill_memory_kv_cache( + cache=tc, input=ctrl, timestep=timesteps[0].to(dtype) + ) + tc.prefill_completed_for_chunk = k + + def sample_losses(c: int, grad: bool) -> tuple[Tensor, Tensor]: + """One drift-pair sample -> (normalized loss, r_target sq-norm).""" + d = { + key: (v.to(device) if isinstance(v, Tensor) else v) + for key, v in datas[c].items() + } + k = int(rng.choice(ks)) + t_idx = int(rng.integers(n_steps)) + selected = d["memory_frame_indices"][k] + ctrl = make_ctrl(d, k, device=device, dtype=dtype) + h_gen = history_of(d, k) + h_clean = clean_counterfactual( + h_gen, selected=selected, lap_latents=d["lap_latents"], clean_lap=CLEAN_LAP + ) + x0 = chunk_x0(d, k) + sig = sigmas[t_idx].to(dtype) + z_t = (1 - sig) * x0 + sig * torch.randn(x0.shape, device=device, dtype=dtype) + + with torch.no_grad(): + set_lora_scale(network, 0.0) + prefill_only(ctrl, h_clean, k) + x0_clean = predict(ctrl, z_t, t_idx) + finish_probe_chunk(tc, ar_idx=k) + prefill_only(ctrl, h_gen, k) + x0_base = predict(ctrl, z_t, t_idx) + finish_probe_chunk(tc, ar_idx=k) + r_target_sq = (x0_clean - x0_base).square().sum() + + set_lora_scale(network, 1.0) + prefill_only(ctrl, h_gen, k) # LoRA-scaled prefill, no grad + with torch.enable_grad() if grad else torch.no_grad(): + x0_corr = predict(ctrl, z_t, t_idx) + loss = (x0_corr - x0_clean).square().sum() / (r_target_sq + 1e-8) + if not grad: + finish_probe_chunk(tc, ar_idx=k) + return loss, r_target_sq + + @torch.no_grad() + def val_r2(n: int = 8) -> float: + s = 0.0 + for _ in range(n): + loss, _ = sample_losses(int(rng.choice(val_ids)), grad=False) + s += loss.item() + return 1 - s / n + + torch.set_grad_enabled(True) + for step in range(1, STEPS + 1): + for pg in opt.param_groups: + pg["lr"] = LR * min(1.0, step / WARMUP) + opt.zero_grad() + loss, rt_sq = sample_losses(int(rng.choice(train_ids)), grad=True) + loss.backward() + # Backward precedes the bracket close: checkpoint recomputation + # needs the chunk's KV state alive. + finish_probe_chunk(tc, ar_idx=tc.autoregressive_index) + torch.nn.utils.clip_grad_norm_(params, GRAD_CLIP) + opt.step() + if step % EVAL_EVERY == 0 or step == 1: + print( + f"step {step:5d} | loss {loss.item():.4f} (train R^2 {1 - loss.item():+.3f})" + f" | val R^2 {val_r2():+.3f} | |r_t|^2 {rt_sq.item():.1f}", + flush=True, + ) + if step % SAVE_EVERY == 0 or step == STEPS: + save_lora(network, CKPT) + + print(f"final val R^2 {val_r2(16):+.3f} | saved {CKPT}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/drift_correction/train_v2.py b/integrations/hy_worldplay/drift_correction/train_v2.py new file mode 100644 index 00000000..2e512b21 --- /dev/null +++ b/integrations/hy_worldplay/drift_correction/train_v2.py @@ -0,0 +1,370 @@ +# 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. + +"""v2 corrector trainer: DAgger pool + drift-contraction (the av2s recipe). + +Initialized from the v1 LoRA, with the reference's two closed-loop upgrades +(``wan_train_v2.py``): + +- **DAgger** -- train on the aggregated pool of round-0 pairs (base-rollout + histories) and round-1 pairs regenerated with the v1 corrector active + (``CORRECTOR_LORA`` in ``build_pairs.py``), fixing the train/deploy + covariate shift. +- **Drift contraction** -- commit the corrected one-step + ``x0_hat = z_t - sigma * flow_corr`` *with gradient* as the newest history + frames for chunk ``k + 1`` and penalize that chunk's x0 gap against its + clean teacher (weight ``CW_LOSS``), optimizing the accumulation mechanism + directly. + +Run from the repo root:: + + POOLS=outputs/pairs,outputs/pairs_dagger1 INIT=outputs/lora_v1.pt \ + .venv/bin/python integrations/hy_worldplay/drift_correction/train_v2.py +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import cast + +import numpy as np +import torch +from _lora import ( + apply_lora, + load_lora, + lora_parameters, + save_lora, + set_lora_scale, + unwrap_compiled, +) +from _pairs import ( + TOKENS_PER_FRAME, + chunk_x0, + clean_counterfactual, + history_of, + load_clip, + make_ctrl, +) +from _rollout import build_runner, finish_probe_chunk, start_probe_chunk +from _train_attn import patch_functional_attention +from hy_worldplay._action import ( + HyWorldPlayWan21Transformer, + HyWorldPlayWan21TransformerCache, +) +from hy_worldplay.runner import _resolve_prompt, preprocess_first_frame +from torch import Tensor +from train_v1 import checkpoint_blocks + +## Training configuration + +_BASE = Path("integrations/hy_worldplay/drift_correction") + +POOLS = [ + _BASE / p + for p in os.environ.get("POOLS", "outputs/pairs,outputs/pairs_dagger1").split(",") +] +"""DAgger aggregation: round-0 (base rollouts) + round-1 (corrector-active +rollouts). Samples draw uniformly over pools.""" + +INIT = os.environ.get("INIT", str(_BASE / "outputs/lora_v1.pt")) +"""v1 LoRA checkpoint to initialize from (required for v2).""" + +CKPT = Path(os.environ.get("CKPT", str(_BASE / "outputs/lora_v2.pt"))) + +STEPS = int(os.environ.get("STEPS", "600")) +LR = float(os.environ.get("LR", "2e-4")) +"""Continue-training LR (below v1's 5e-4, per the reference).""" + +WARMUP = 40 +CW_LOSS = float(os.environ.get("CW_LOSS", "0.5")) +"""Drift-contraction term weight.""" + +T_WEIGHTS = [float(w) for w in os.environ.get("T_WEIGHTS", "1,1,1,1").split(",")] +"""Sampling weights over the 4 distillation timesteps (t=1000 first). +Oversampling t=1000 concentrates training where the faithful gate measured +the drift gap as most systematic (alpha* 0.81 vs ~0.55 elsewhere).""" + +FID_W = float(os.environ.get("FID_W", "0")) +"""Content-fidelity (trust-region) weight: penalizes +``||x0_corr - x0_base||^2`` (drift-gap normalized) so the corrector moves +*statistics* toward the clean teacher without rewriting *content* — the +anti-repeat / anti-hallucination seatbelt. Ramped linearly over +``FID_WARMUP`` steps.""" + +FID_WARMUP = 100 + +TARGETS = tuple( + t.strip() + for t in os.environ.get( + "TARGETS", "self_attn.q,self_attn.k,self_attn.v,self_attn.o" + ).split(",") +) +"""LoRA target projections. ``self_attn.q,self_attn.k`` gives the +attention-routing-only variant (no content-writing v/o adapters); note a +narrowed target set cannot init from a full-target checkpoint.""" + +GRAD_CLIP = 1.0 +RANK = 16 +EVAL_EVERY = 100 +SAVE_EVERY = 200 +CLEAN_LAP = 1 +SEED = int(os.environ.get("SEED", "0")) + + +def main() -> None: + pools = [sorted(p.glob("clip_*.pt")) for p in POOLS] + pools = [p for p in pools if p] + assert pools, f"no clips under any of {POOLS}" + rng = np.random.default_rng(SEED) + + t_probs = np.array(T_WEIGHTS, dtype=np.float64) + t_probs = t_probs / t_probs.sum() + + meta = torch.load(pools[0][0], map_location="cpu", weights_only=False) + runner = build_runner( + num_chunk=meta["num_chunk"], + pose=meta["pose"], + output_dir=CKPT.parent, + compile_network=False, + ) + pipe = runner.pipeline + device = next(pipe.parameters()).device + dtype = next(pipe.parameters()).dtype + transformer = cast(HyWorldPlayWan21Transformer, pipe.diffusion_model.transformer) + scheduler = pipe.diffusion_model.scheduler + timesteps = cast(Tensor, scheduler.timesteps) + sigmas = cast(Tensor, scheduler.sigmas) + n_steps = len(timesteps) - 1 + + cfg = runner.config + assert cfg.image_path is not None # build_runner resolves the sample image + image = preprocess_first_frame( + cfg.image_path, cfg.pixel_height, cfg.pixel_width + ).to(device=device, dtype=dtype) + cache = pipe.initialize_cache(text=[_resolve_prompt(cfg.prompt)], image=image) + tc = cache.transformer_cache + assert isinstance(tc, HyWorldPlayWan21TransformerCache) + + network = unwrap_compiled(transformer.network) + apply_lora(network, rank=RANK, targets=TARGETS) + if INIT and INIT != "scratch": + load_lora(network, INIT) + patch_functional_attention() + checkpoint_blocks(network) + params = lora_parameters(network) + print(f"v2 init from {INIT} | pools {[len(p) for p in pools]} clips", flush=True) + opt = torch.optim.AdamW(params, lr=LR) + + datas = [[load_clip(p, "cpu", dtype) for p in pool] for pool in pools] + n0 = len(datas[0]) + nval = max(2, n0 // 8) + train_ids = list(range(n0 - nval)) + val_ids = list(range(n0 - nval, n0)) + + def ks_for(d: dict) -> list[int]: + """Chunks whose memory window is fully past the clean lap (per-clip: + mixed-geometry pools carry different lap sizes / rollout lengths). + Leaves room for the k+1 contraction chunk.""" + lap_chunks = d["lap_latents"] // 4 + return list(range((CLEAN_LAP + 2) * lap_chunks, d["num_chunk"] - 1)) + + def predict(cachet, ctrl, z_t: Tensor, t_idx: int) -> Tensor: + flow = transformer.predict_flow( + noisy_latent=z_t, + timestep=timesteps[t_idx].to(dtype), + cache=cachet, + input=ctrl, + ) + return z_t.float() - float(sigmas[t_idx]) * flow.float() + + def prefill_only(cachet, ctrl, history: Tensor, k: int) -> None: + start_probe_chunk(cachet, ar_idx=k, history=history) + with torch.no_grad(): + if ctrl.memory_frame_indices: + transformer.prefill_memory_kv_cache( + cache=cachet, input=ctrl, timestep=timesteps[0].to(dtype) + ) + cachet.prefill_completed_for_chunk = k + + def teacher_and_base(cachet, ctrl, h_clean, h_gen, z_t, t_idx, k): + """Frozen-model x0 under clean and drifted history (no grad).""" + with torch.no_grad(): + set_lora_scale(network, 0.0) + prefill_only(cachet, ctrl, h_clean, k) + x0_clean = predict(cachet, ctrl, z_t, t_idx) + finish_probe_chunk(cachet, ar_idx=k) + prefill_only(cachet, ctrl, h_gen, k) + x0_base = predict(cachet, ctrl, z_t, t_idx) + finish_probe_chunk(cachet, ar_idx=k) + return x0_clean, x0_base + + def sample_losses( + pool_id: int, c: int, grad: bool, fid_w: float = 0.0 + ) -> tuple[Tensor, Tensor]: + """One sample -> (dagger loss, contraction loss).""" + # Pools differ in size (round-1 regenerates a train-split prefix); + # fold the pool-0 index into the smaller pool. Val ids exist only in + # pool 0, so this cannot leak validation clips into training. + c = c % len(datas[pool_id]) + d = { + key: (v.to(device) if isinstance(v, Tensor) else v) + for key, v in datas[pool_id][c].items() + } + lap = d["lap_latents"] + k = int(rng.choice(ks_for(d))) + t_idx = int(rng.choice(n_steps, p=t_probs)) + ctrl = make_ctrl(d, k, device=device, dtype=dtype) + h_gen = history_of(d, k) + h_clean = clean_counterfactual( + h_gen, + selected=d["memory_frame_indices"][k], + lap_latents=lap, + clean_lap=CLEAN_LAP, + ) + x0 = chunk_x0(d, k) + sig = sigmas[t_idx].to(dtype) + z_t = (1 - sig) * x0 + sig * torch.randn(x0.shape, device=device, dtype=dtype) + + x0_clean, x0_base = teacher_and_base(tc, ctrl, h_clean, h_gen, z_t, t_idx, k) + rt_sq = (x0_clean - x0_base).square().sum() + + if not grad: + set_lora_scale(network, 1.0) + prefill_only(tc, ctrl, h_gen, k) + with torch.no_grad(): + x0_corr = predict(tc, ctrl, z_t, t_idx) + finish_probe_chunk(tc, ar_idx=k) + l_dag = (x0_corr - x0_clean).square().sum() / (rt_sq + 1e-8) + return l_dag, torch.zeros((), device=device) + + # Grad path: two-phase backward so only one chunk's graph is ever + # live (a joint k / k+1 graph OOMs next to co-tenant jobs, and + # checkpoint recomputation forbids sharing one cache across both). + # + # Phase 1 -- contraction: run the student at k WITHOUT grad, make + # its x0 a leaf, splice it into chunk k+1's history, run the + # grad-capable committed prefill + k+1 forward, and backward + # ``CW_LOSS * l_con`` immediately. This accumulates the k+1-path + # LoRA grads and yields ``leaf.grad``. + set_lora_scale(network, 1.0) + prefill_only(tc, ctrl, h_gen, k) + with torch.no_grad(): + x0_ng = predict(tc, ctrl, z_t, t_idx) + finish_probe_chunk(tc, ar_idx=k) + leaf = x0_ng.detach().requires_grad_(True) + + k2 = k + 1 + sel2 = d["memory_frame_indices"][k2] or [] + touches_k = any(k * 4 <= idx < (k + 1) * 4 for idx in sel2) + l_con = torch.zeros((), device=device) + if touches_k: + ctrl2 = make_ctrl(d, k2, device=device, dtype=dtype) + t2_idx = int(rng.choice(n_steps, p=t_probs)) + x0_next = chunk_x0(d, k2) + sig2 = sigmas[t2_idx].to(dtype) + z2 = (1 - sig2) * x0_next + sig2 * torch.randn( + x0_next.shape, device=device, dtype=dtype + ) + h2_gen = history_of(d, k2) + h2_clean = clean_counterfactual( + h2_gen, + selected=sel2, + lap_latents=lap, + clean_lap=CLEAN_LAP, + ) + x0_clean2, x0_base2 = teacher_and_base( + tc, ctrl2, h2_clean, h2_gen, z2, t2_idx, k2 + ) + rt2_sq = (x0_clean2 - x0_base2).square().sum() + + h2_committed = h2_gen.clone() + s = slice(k * 4 * TOKENS_PER_FRAME, (k + 1) * 4 * TOKENS_PER_FRAME) + h2_committed[..., s, :] = leaf.to(dtype) + set_lora_scale(network, 1.0) + start_probe_chunk(tc, ar_idx=k2, history=h2_committed) + with torch.enable_grad(): + if ctrl2.memory_frame_indices: + transformer.prefill_memory_kv_cache( + cache=tc, input=ctrl2, timestep=timesteps[0].to(dtype) + ) + tc.prefill_completed_for_chunk = k2 + x0_corr2 = predict(tc, ctrl2, z2, t2_idx) + l_con = (x0_corr2 - x0_clean2).square().sum() / (rt2_sq + 1e-8) + (CW_LOSS * l_con).backward() + finish_probe_chunk(tc, ar_idx=k2) + l_con = l_con.detach() + + # Phase 2 -- dagger + chain rule: rerun the student at k WITH grad + # (deterministic, so it reproduces ``x0_ng`` exactly) and route the + # contraction gradient through it via the dot-product trick: + # d/dtheta [x0_corr . leaf.grad] == (dl_con/dx0) (dx0/dtheta). + set_lora_scale(network, 1.0) + prefill_only(tc, ctrl, h_gen, k) + with torch.enable_grad(): + x0_corr = predict(tc, ctrl, z_t, t_idx) + l_dag = (x0_corr - x0_clean).square().sum() / (rt_sq + 1e-8) + loss_k = l_dag + if fid_w > 0: + # Trust region: correct statistics without rewriting content. + loss_k = loss_k + fid_w * ( + (x0_corr - x0_base).square().sum() / (rt_sq + 1e-8) + ) + if leaf.grad is not None: + loss_k = loss_k + (x0_corr * leaf.grad.detach()).sum() + loss_k.backward() + finish_probe_chunk(tc, ar_idx=k) + return l_dag.detach(), l_con + + @torch.no_grad() + def val_loss(n: int = 8) -> float: + s = 0.0 + for _ in range(n): + l_dag, _ = sample_losses(0, int(rng.choice(val_ids)), grad=False) + s += l_dag.item() + return s / n + + torch.set_grad_enabled(True) + for step in range(1, STEPS + 1): + for pg in opt.param_groups: + pg["lr"] = LR * min(1.0, step / WARMUP) + opt.zero_grad() + pool_id = int(rng.integers(len(pools))) + c = int(rng.choice(train_ids)) + # Backwards happen inside sample_losses (two-phase; grads + # accumulate into .grad). + fid_w = FID_W * min(1.0, step / FID_WARMUP) + l_dag, l_con = sample_losses(pool_id, c, grad=True, fid_w=fid_w) + torch.nn.utils.clip_grad_norm_(params, GRAD_CLIP) + opt.step() + if step % EVAL_EVERY == 0 or step == 1: + vl = val_loss() + print( + f"step {step:4d} | dag {l_dag.item():.4f} | con {float(l_con):.4f} | " + f"val dag-loss {vl:.4f} (R^2 {1 - vl:+.3f})", + flush=True, + ) + if step % SAVE_EVERY == 0 or step == STEPS: + save_lora(network, CKPT) + + vl = val_loss(16) + print( + f"v2 done | final val dag-loss {vl:.4f} (R^2 {1 - vl:+.3f}) | saved {CKPT}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/hy_worldplay/hy_worldplay/_drift_corrector.py b/integrations/hy_worldplay/hy_worldplay/_drift_corrector.py new file mode 100644 index 00000000..cdca9197 --- /dev/null +++ b/integrations/hy_worldplay/hy_worldplay/_drift_corrector.py @@ -0,0 +1,163 @@ +# 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. + +"""Content-keyed Clean Forcing drift corrector for the HY-WorldPlay runner.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from torch import Tensor + +## Deploy policy + +GATE_ALPHA = {1000.0: 0.81, 960.0: 0.53, 888.8889: 0.53, 727.2728: 0.58} +"""Unbiased alpha*(t) from the step-0 systematicity gate (drift_correction's +``outputs/gate/gate_faithful.json``): the systematic fraction of the +drift-induced error at each of the distilled solver's timesteps. The +shipped config deploys the LoRA at ``alpha*(t) * gain`` per denoise step.""" + +STATIC_ACTION_CLASS = 0 +"""Action label for a no-translation / no-rotation step; a trajectory of +only this class is a locked-off camera.""" + +_LORA_TARGETS = ("self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o") +"""Self-attention projections the corrector checkpoint was trained on.""" + +_LORA_RANK = 16 +"""Rank of the shipped v2 corrector checkpoint.""" + + +class _LoRALinear(nn.Module): + """Frozen base linear plus a runtime-gated low-rank delta. + + Mirrors the training-side module in + ``integrations/hy_worldplay/drift_correction/_lora.py``: ``scale`` is + the runtime gain (``0`` = exact base output), and the A/B path runs in + fp32 regardless of the base dtype. + """ + + def __init__(self, base: nn.Linear, rank: int): + super().__init__() + self.base = base + for p in self.base.parameters(): + p.requires_grad_(False) + self.A = nn.Linear(base.in_features, rank, bias=False) + self.B = nn.Linear(rank, base.out_features, bias=False) + nn.init.zeros_(self.B.weight) + self.scale = 0.0 + + def forward(self, x: Tensor) -> Tensor: + out = self.base(x) + if self.scale != 0: + delta = self.B(self.A(x.to(self.A.weight.dtype))) + out = out + self.scale * delta.to(out.dtype) + return out + + +def _apply_lora(network: nn.Module) -> list[nn.Parameter]: + """Wrap the target linears and return the LoRA parameters in load order.""" + for mname, module in list(network.named_modules()): + for cname, child in list(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): + setattr( + module, + cname, + _LoRALinear(child, _LORA_RANK).to(child.weight.device), + ) + params: list[nn.Parameter] = [] + for m in network.modules(): + if isinstance(m, _LoRALinear): + params += list(m.A.parameters()) + list(m.B.parameters()) + return params + + +def _set_scale(network: nn.Module, scale: float) -> None: + """Set the runtime gain on every wrapped linear.""" + for m in network.modules(): + if isinstance(m, _LoRALinear): + m.scale = scale + + +def is_static_trajectory(pose: str | Path, n_latents: int) -> bool: + """Whether the job's trajectory is a locked-off camera. + + Args: + pose: Pose-string or trajectory-JSON path (upstream grammar). + n_latents: Rollout latent budget (``num_chunk * 4``). + + Returns: + ``True`` when every per-latent action label is + :data:`STATIC_ACTION_CLASS`. + """ + from hy_worldplay._pose import parse_pose_action_labels + + labels = parse_pose_action_labels(pose, n_latents) + return bool((labels == STATIC_ACTION_CLASS).all()) + + +def maybe_apply_drift_corrector(runner: Any, checkpoint: Path, gain: float) -> str: + """Deploy the corrector on ``runner`` keyed on the job's trajectory content. + + Ship rule (owner decision 2026-07-21): static scenes measure negative + drift on this host, so correction there is pure artifact cost — static + jobs run the untouched base weights. Commanded-motion jobs get the + corrector LoRA at ``alpha*(t) * gain`` per denoise step (the + ``corrgate050`` config at the default ``gain=0.5``). + + Args: + runner: A built ``HyWorldPlayWanI2VRunner``. + checkpoint: Corrector LoRA checkpoint (``save_lora`` format). + gain: Global gain multiplied into the alpha*(t) profile. + + Returns: + ``"base (static trajectory)"`` or ``"corrected (alpha*(t) x gain)"``, + for the runner's log line. + """ + cfg = runner.config + if is_static_trajectory(cfg.pose, cfg.num_chunk * 4): + return "base (static trajectory)" + + network = runner.pipeline.diffusion_model.transformer.network + if hasattr(network, "_orig_mod"): # unwrap torch.compile + network = network._orig_mod + params = _apply_lora(network) + + sd = torch.load(checkpoint, map_location="cpu", weights_only=False)["lora"] + assert len(sd) == len(params), ( + f"corrector checkpoint has {len(sd)} LoRA tensors but the network " + f"exposes {len(params)}; rank or target mismatch." + ) + for i, p in enumerate(params): + p.data.copy_(sd[i].to(p.device, p.dtype)) + + # Per-step gate: rescale the LoRA to alpha*(t) x gain before every + # denoise step (nearest-t lookup; per-token AR0 timesteps include the + # first-frame stabilization value, the max is always the scheduler step). + transformer = runner.pipeline.diffusion_model.transformer + orig_pf = transformer.predict_flow + + def gated_pf(*args, **kwargs): + t = float(kwargs["timestep"].reshape(-1).max()) + alpha = min(GATE_ALPHA.items(), key=lambda kv: abs(kv[0] - t))[1] + _set_scale(network, alpha * gain) + return orig_pf(*args, **kwargs) + + transformer.predict_flow = gated_pf + return "corrected (alpha*(t) x gain)" diff --git a/integrations/hy_worldplay/hy_worldplay/runner.py b/integrations/hy_worldplay/hy_worldplay/runner.py index 7999d437..3b24083a 100644 --- a/integrations/hy_worldplay/hy_worldplay/runner.py +++ b/integrations/hy_worldplay/hy_worldplay/runner.py @@ -18,6 +18,7 @@ from __future__ import annotations import contextlib +import json import time from dataclasses import dataclass, field from pathlib import Path @@ -28,14 +29,7 @@ from flashdreams.core.io.disk import default_flashdreams_cache_dir from flashdreams.core.io.download import download_to_cache -from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - ensure_output_dir, - resolve_prompt_value, - runner_artifact_path, - write_runner_stats, -) from flashdreams.recipes.wan.pipeline import WanInferencePipeline __all__ = [ @@ -126,6 +120,16 @@ def _pil_to_numpy(img: object) -> object: return np.asarray(img) +def _resolve_prompt(value: str | Path) -> str: + """Read an inline prompt or the first non-empty line of a prompt file.""" + if isinstance(value, Path): + lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] + assert lines, f"prompt file {value} has no non-empty lines" + return lines[0] + assert value, "--prompt must be a non-empty string or a path to a .txt file" + return value + + def _write_mp4(video: Tensor, out_path: Path, *, fps: int) -> None: """Persist a decoded video tensor as mp4. @@ -204,9 +208,6 @@ class HyWorldPlayWanI2VRunnerConfig(RunnerConfig): fps: int = 16 """Output video frame rate.""" - postprocess_output_layout: VideoTensorLayout | None = "btchw" - """Pipeline output layout for streaming post-processing.""" - context_window_length: int = 16 """Frame-count threshold below which the FOV-overlap memory selector is bypassed (AR steps with fewer accumulated frames emit @@ -248,6 +249,19 @@ class HyWorldPlayWanI2VRunnerConfig(RunnerConfig): memory_points_radius: float = 8.0 """Radius of the Monte-Carlo sphere.""" + drift_corrector: Path | None = None + """Clean Forcing drift-corrector LoRA checkpoint (v2). ``None`` + disables correction. When set, correction is keyed on the job's + trajectory content: static poses (all idle action labels) run the + untouched base weights — static scenes measure negative drift on this + host, so correction there is pure artifact cost — while + commanded-motion poses deploy the corrector at + ``alpha*(t) * drift_corrector_gain`` per denoise step.""" + + drift_corrector_gain: float = 0.5 + """Global gain composed with the per-step ``alpha*(t)`` gate profile; + the shipped configuration (``corrgate050``) is 0.5.""" + class HyWorldPlayWanI2VRunner( Runner[HyWorldPlayWanI2VRunnerConfig, WanInferencePipeline] @@ -339,6 +353,14 @@ def run(self) -> None: if not cfg.image_path.exists(): raise FileNotFoundError(f"image_path {cfg.image_path} does not exist") + if cfg.drift_corrector is not None: + from hy_worldplay._drift_corrector import maybe_apply_drift_corrector + + mode = maybe_apply_drift_corrector( + self, cfg.drift_corrector, cfg.drift_corrector_gain + ) + logger.info(f"[{cfg.runner_name}] drift corrector: {mode}") + first_param = next(self.pipeline.parameters()) device = first_param.device # Cast to the pipeline's parameter dtype so the VAE encoder's @@ -346,7 +368,7 @@ def run(self) -> None: image = preprocess_first_frame( cfg.image_path, cfg.pixel_height, cfg.pixel_width ).to(device=device, dtype=first_param.dtype) - prompt = resolve_prompt_value(cfg.prompt) + prompt = _resolve_prompt(cfg.prompt) cache = self.pipeline.initialize_cache( text=[prompt], @@ -371,38 +393,31 @@ def run(self) -> None: device=device, dtype=first_param.dtype ) - postprocess_stream = self.create_postprocess_stream( - fps=cfg.fps, move_to_cpu=False - ) + chunks: list[Tensor] = [] # Each ``finalize`` returns the per-stage ms dict for that AR # step; collect into ``stats_history`` and dump as JSON. - stats_history: list[dict[str, object]] = [] + stats_history: list[dict[str, float]] = [] if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() start_time = time.time() with vendor_noise_ctx: for ar_idx in range(cfg.num_chunk): chunk = self.pipeline.generate(ar_idx, cache) + chunks.append(chunk) # ``finalize`` records the chunk's CUDA events and # advances the KV cache; called on every chunk # (including the last) for consistent stats. stats = self.pipeline.finalize(ar_idx, cache) - postprocess_stream.process(chunk, autoregressive_index=ar_idx) - if postprocess_stream.collect_output and stats is not None: - stats_history.append( - { - "autoregressive_index": ar_idx, - **postprocess_stream.add_process_stats(stats), - } - ) - video = postprocess_stream.finish() + if stats is not None: + stats_history.append({"autoregressive_index": ar_idx, **stats}) elapsed = time.time() - start_time - if video is None: + if not self.is_rank_zero: return - ensure_output_dir(cfg.output_dir) - out_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") + video = torch.cat(chunks, dim=-4) # cat along T axis: [..., T, C, H, W] + cfg.output_dir.mkdir(parents=True, exist_ok=True) + out_path = cfg.output_dir / f"{cfg.runner_name}.mp4" _write_mp4(video, out_path, fps=cfg.fps) logger.info( f"[{cfg.runner_name}] wrote video " @@ -410,9 +425,8 @@ def run(self) -> None: ) if stats_history: - stats_path = write_runner_stats( - cfg.output_dir, cfg.runner_name, stats_history - ) + stats_path = cfg.output_dir / f"stats_{cfg.runner_name}.json" + stats_path.write_text(json.dumps(stats_history, indent=2)) logger.info( f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" ) diff --git a/integrations/hy_worldplay/tests/test_drift_corrector.py b/integrations/hy_worldplay/tests/test_drift_corrector.py new file mode 100644 index 00000000..17932a45 --- /dev/null +++ b/integrations/hy_worldplay/tests/test_drift_corrector.py @@ -0,0 +1,115 @@ +# 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 Clean Forcing drift-corrector deploy hook. + +Covers the two behaviours a deployment depends on: + +* ``_LoRALinear`` is a strict identity at ``scale == 0`` and at the + zero-initialized ``B`` (so wrapping the network never changes base + outputs until a trained checkpoint is loaded and gated on). +* ``is_static_trajectory`` keys the content-based selection: all-idle + pose strings bypass the corrector, any commanded motion enables it. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from hy_worldplay._drift_corrector import ( + _apply_lora, + _LoRALinear, + _set_scale, + is_static_trajectory, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_lora_linear_is_identity_at_zero_scale(): + torch.manual_seed(0) + base = nn.Linear(8, 6, bias=False) + lora = _LoRALinear(base, rank=2) + nn.init.normal_(lora.A.weight) + nn.init.normal_(lora.B.weight) # non-zero delta path + x = torch.randn(3, 8) + lora.scale = 0.0 + assert torch.equal(lora(x), base(x)) + + +def test_lora_linear_is_identity_at_zero_init_b(): + torch.manual_seed(0) + base = nn.Linear(8, 6, bias=False) + lora = _LoRALinear(base, rank=2) # B is zero-initialized + lora.scale = 1.0 + x = torch.randn(3, 8) + assert torch.allclose(lora(x), base(x)) + + +def test_lora_linear_applies_scaled_delta(): + torch.manual_seed(0) + base = nn.Linear(8, 6, bias=False) + lora = _LoRALinear(base, rank=2) + nn.init.normal_(lora.A.weight) + nn.init.normal_(lora.B.weight) + x = torch.randn(3, 8) + lora.scale = 0.5 + delta = lora(x) - base(x) + lora.scale = 1.0 + assert torch.allclose(2.0 * delta, lora(x) - base(x), atol=1e-5) + + +def test_apply_lora_wraps_only_attention_targets_and_set_scale_reaches_all(): + class Toy(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.ModuleDict( + {n: nn.Linear(4, 4, bias=False) for n in ("q", "k", "v", "o")} + ) + self.ffn = nn.Linear(4, 4, bias=False) + + toy = Toy() + params = _apply_lora(toy) + wrapped = [m for m in toy.modules() if isinstance(m, _LoRALinear)] + assert len(wrapped) == 4 # q/k/v/o but not ffn + assert len(params) == 8 # A + B per wrapped linear + assert not isinstance(toy.ffn, _LoRALinear) + _set_scale(toy, 0.25) + assert all(m.scale == 0.25 for m in wrapped) + + +def test_static_pose_json_bypasses_the_corrector(tmp_path): + # The upstream grammar has no explicit "stay" token; a locked-off + # camera is an all-identity trajectory JSON (see demo_static.py). + import json + + import numpy as np + + eye = np.eye(4).tolist() + intrinsic = [ + [1000.0, 0.0, 960.0], + [0.0, 1000.0, 540.0], + [0.0, 0.0, 1.0], + ] + n_latents = 16 + poses = {str(i): {"extrinsic": eye, "K": intrinsic} for i in range(n_latents + 1)} + pose_json = tmp_path / "static_pose.json" + pose_json.write_text(json.dumps(poses)) + assert is_static_trajectory(pose_json, n_latents=n_latents) is True + + +def test_commanded_motion_pose_enables_the_corrector(): + assert is_static_trajectory("w-8, s-8", n_latents=16) is False