diff --git a/.gitignore b/.gitignore index 08f109a0c..90c0342bd 100644 --- a/.gitignore +++ b/.gitignore @@ -124,6 +124,7 @@ compile_commands.json outputs/ data_local/ credentials/*.secret +slurm-logs/ CLAUDE.md AGENTS.md diff --git a/integrations/artifixer/README.md b/integrations/artifixer/README.md new file mode 100644 index 000000000..77f14b78f --- /dev/null +++ b/integrations/artifixer/README.md @@ -0,0 +1,70 @@ +# flashdreams-artifixer + +ArtiFixer reconstruction-enhanced T2V inference for Wan 2.1 1.3B, packaged as +a [`flashdreams`](../..) plugin. + +ArtiFixer extends Wan 2.1 1.3B with: + +- per-block opacity and Plucker-camera-ray MLPs; +- a third KV bank for neighbor cross-attention with PRoPE; +- opacity-weighted latent mixing of noise with the VAE-encoded + reconstruction-rendered frames; +- 4-step DMD distillation (`FlowMatchScheduler(shift=5)`). + +This plugin ports the ArtiFixer reference implementation to +flashdreams' faster Wan stack (RingAttention, cuDNN, `torch.compile`, +CUDA graphs). + +## Components + +| Component | Description | +| --- | --- | +| Recipe scaffold | AR rollout + 4-step DMD scheduler knobs match the ArtiFixer DMD stage-3 1.3B training config. | +| Per-block conditioning | Opacity + camera-ray MLPs, neighbor cross-attention, and PRoPE — parity-tested against the ArtiFixer reference. | +| Pipeline | Opacity-weighted latent mixing + self-forcing renoise loop inside `ArtifixerInferencePipeline.generate`. | +| External-driver surface | `initialize_cache` accepts pre-encoded UMT5 prompts + VAE-encoded condition / neighbor latents so an external driver can feed it directly. | +| Checkpoint loader | `state_dict_transform` for the merged ArtiFixer DMD safetensors (a consolidated single-file checkpoint built from the sharded FSDP training output). | + +Cross-backend parity (captured single-scene `final_video`): **51.34 dB** +PSNR vs the ArtiFixer reference's `ArtifixerKvCachePipeline` after the +fp32 AdaLN/norm/residual promotion and the no-op `finalize_kv_cache` +override on the flashdreams transformer. + +Set `ARTIFIXER_DMD_CHECKPOINT_PATH` to point at the merged ArtiFixer +DMD safetensors; alternatively set `ARTIFIXER_USE_BASE_WAN_WEIGHTS=1` +to fall back to vanilla Wan 2.1 1.3B HuggingFace weights (useful for +smoke-testing the recipe wiring before the merged safetensors are +available). + +## Shipped slugs + +| slug | description | +| --- | --- | +| `artifixer-dmd-wan2.1-t2v-1.3b` | ArtiFixer reconstruction-enhanced T2V (Wan 2.1 1.3B + opacity/camera/neighbor extensions, 4-step DMD). | + +## Install + +The plugin is registered as a `uv` workspace member in the repo-root +`pyproject.toml`, so a single `uv sync` from the repo root pulls it in: + +```bash +uv sync +``` + +## Run + +```bash +export HF_TOKEN= + +uv run flashdreams-run artifixer-dmd-wan2.1-t2v-1.3b --help + +uv run flashdreams-run artifixer-dmd-wan2.1-t2v-1.3b \ + --prompt "A photorealistic dolly-in shot of a modern living room." \ + --total-blocks 3 +``` + +## Tests + +```bash +uv run --extra dev pytest integrations/artifixer/tests +``` diff --git a/integrations/artifixer/artifixer/__init__.py b/integrations/artifixer/artifixer/__init__.py new file mode 100644 index 000000000..467079831 --- /dev/null +++ b/integrations/artifixer/artifixer/__init__.py @@ -0,0 +1,14 @@ +# 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. diff --git a/integrations/artifixer/artifixer/checkpoint.py b/integrations/artifixer/artifixer/checkpoint.py new file mode 100644 index 000000000..1853a3772 --- /dev/null +++ b/integrations/artifixer/artifixer/checkpoint.py @@ -0,0 +1,197 @@ +# 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. + +"""State-dict transforms for loading ArtiFixer checkpoints into flashdreams. + +Two source layouts are supported: + + * **Vanilla Wan 2.1 1.3B** from ``Wan-AI/Wan2.1-T2V-1.3B`` -- + :func:`zero_pad_artifixer_keys` zero-pads the ArtiFixer-only keys so + ``load_state_dict`` succeeds in strict mode. Zero-padding matches the + ArtiFixer reference's per-block initialization and produces a + behaviorally-identical-to-Wan model (the ArtiFixer extension paths + are zero-gated until trained weights are loaded). + + * **Merged ArtiFixer DMD safetensors** (consolidated from a sharded + FSDP training checkpoint). Built by + :func:`artifixer_dmd_state_dict_transform`, which applies the + HuggingFace diffusers -> WanDiTNetwork regex remap (same as + ``fastvideo_causal_wan22.config.state_dict_transform``) plus the + ArtiFixer ``attn2 -> cross_attn`` step that picks up + ``add_k_proj`` / ``add_v_proj`` / ``norm_added_k``. The 270 + ArtiFixer-only keys map cleanly onto the + ``ArtifixerBlock`` / ``ArtifixerCrossAttention`` submodules without + further per-key renames. +""" + +from __future__ import annotations + +import re +from typing import Callable + +import torch +from torch import Tensor + +from artifixer.network.dit import artifixer_embedding_dims + +# HF diffusers ``WanTransformer3DModel`` -> flashdreams ``WanDiTNetwork`` +# key remap. Verbatim copy of +# ``integrations/fastvideo_causal_wan22/.../config.py`` ``CHECKPOINT_KEY_MAPPING`` +# (Wan 2.1 / 2.2 share this layout). The ArtiFixer-only keys -- +# ``blocks.X.opacity_embedding.*``, ``blocks.X.camera_embedding.*``, +# ``blocks.X.attn2.add_k_proj.*``, ``blocks.X.attn2.add_v_proj.*``, +# ``blocks.X.attn2.norm_added_k.weight`` -- pass through the +# ``attn2 -> cross_attn`` substitution into the names +# :class:`ArtifixerBlock` / :class:`ArtifixerCrossAttention` register at +# ``__init__`` time. ``opacity_embedding`` / ``camera_embedding`` carry +# no ``attn2`` prefix so they fall through unchanged, which still +# matches the ArtifixerBlock attributes. +DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING: dict[str, str] = { + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"text_embedding.0.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"text_embedding.2.\1", + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"time_embedding.0.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"time_embedding.2.\1", + r"^condition_embedder\.time_proj\.(.*)$": r"time_projection.1.\1", + r"^scale_shift_table$": r"head.modulation", + r"^proj_out\.(.*)$": r"head.head.\1", + r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.self_attn.q.\2", + r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.self_attn.k.\2", + r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.self_attn.v.\2", + r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.self_attn.o.\2", + r"^blocks\.(\d+)\.attn2\.to_q\.(.*)$": r"blocks.\1.cross_attn.q.\2", + r"^blocks\.(\d+)\.attn2\.to_k\.(.*)$": r"blocks.\1.cross_attn.k.\2", + r"^blocks\.(\d+)\.attn2\.to_v\.(.*)$": r"blocks.\1.cross_attn.v.\2", + r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.cross_attn.o.\2", + # ArtiFixer-only attn2 keys: ``add_k_proj`` / ``add_v_proj`` / + # ``norm_added_k``. Same ``attn2 -> cross_attn`` substitution. + r"^blocks\.(\d+)\.attn2\.add_k_proj\.(.*)$": r"blocks.\1.cross_attn.add_k_proj.\2", + r"^blocks\.(\d+)\.attn2\.add_v_proj\.(.*)$": r"blocks.\1.cross_attn.add_v_proj.\2", + r"^blocks\.(\d+)\.attn2\.norm_added_k\.(.*)$": r"blocks.\1.cross_attn.norm_added_k.\2", + r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.self_attn.norm_q.\2", + r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.self_attn.norm_k.\2", + r"^blocks\.(\d+)\.attn2\.norm_q\.(.*)$": r"blocks.\1.cross_attn.norm_q.\2", + r"^blocks\.(\d+)\.attn2\.norm_k\.(.*)$": r"blocks.\1.cross_attn.norm_k.\2", + r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.norm3.\2", + r"^blocks\.(\d+)\.scale_shift_table$": r"blocks.\1.modulation", + r"^blocks\.(\d+)\.ffn\.fc_in\.(.*)$": r"blocks.\1.ffn.0.\2", + r"^blocks\.(\d+)\.ffn\.fc_out\.(.*)$": r"blocks.\1.ffn.2.\2", + r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.0.\2", + r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.2.\2", +} + + +def _remap_keys( + state_dict: dict[str, Tensor], mapping: dict[str, str] +) -> dict[str, Tensor]: + """Apply the first matching regex substitution to every key. + + Identical semantics to ``flashdreams.core.checkpoint.remap.remap_checkpoint_keys``; + inlined here so this module has no dependency on flashdreams' internal + helpers when used as a state_dict_transform callable. + """ + out: dict[str, Tensor] = {} + for k, v in state_dict.items(): + new_k = k + for old_pattern, new_pattern in mapping.items(): + if re.match(old_pattern, k): + new_k = re.sub(old_pattern, new_pattern, k) + break + out[new_k] = v + return out + + +def artifixer_dmd_state_dict_transform( + state_dict: dict[str, Tensor], +) -> dict[str, Tensor]: + """Remap a merged ArtiFixer DMD safetensors state_dict onto WanDiTNetwork. + + The merged safetensors carries HF diffusers ``WanTransformer3DModel`` + naming (e.g. ``blocks.X.attn1.to_q.weight``) plus 270 ArtiFixer-only + keys with the ``attn2`` cross-attention prefix + (``blocks.X.attn2.add_k_proj.weight``) and 60 keys without prefix + (``blocks.X.opacity_embedding.weight``, + ``blocks.X.camera_embedding.weight``). + + Apply :data:`DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING` to convert + diffusers names to the flashdreams ``WanDiTNetwork`` / + :class:`ArtifixerDiTNetwork` layout. The ArtiFixer-only keys land on + :class:`ArtifixerBlock` / :class:`ArtifixerCrossAttention` attributes + that are registered at ``__init__`` time, so ``load_state_dict`` + succeeds in strict mode with no remaining missing / unexpected keys. + """ + return _remap_keys(state_dict, DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING) + +def zero_pad_artifixer_keys( + *, + num_layers: int, + dim: int, + patch_size: tuple[int, int, int], + dtype: torch.dtype = torch.bfloat16, +) -> Callable[[dict[str, Tensor]], dict[str, Tensor]]: + """Build a state-dict transform that zero-fills ArtiFixer-only keys. + + Use when the source checkpoint is a vanilla Wan 2.1 base + (e.g. ``Wan-AI/Wan2.1-T2V-1.3B``): it has 825 base keys per block but + missing the ArtiFixer extensions, so :meth:`load_state_dict` in strict + mode would error. We pre-add zero tensors so loading succeeds and the + runtime behavior is unchanged versus vanilla Wan. + + Per block, the transform adds 9 keys (mirroring the 270 + ArtiFixer-only keys in the merged DMD safetensors): + + * Opacity + camera MLPs (4 keys per block): + ``opacity_embedding.{weight,bias}`` shape (dim, opacity_dim) / (dim,) + ``camera_embedding.{weight,bias}`` shape (dim, camera_dim) / (dim,) + * Neighbor cross-attention (5 keys per block): + ``cross_attn.add_k_proj.{weight,bias}`` shape (dim, dim) / (dim,) + ``cross_attn.add_v_proj.{weight,bias}`` shape (dim, dim) / (dim,) + ``cross_attn.norm_added_k.weight`` shape (dim,) + + Args: + num_layers: Number of transformer blocks (1.3B Wan: 30). + dim: Transformer hidden size (1.3B Wan: 1536). + patch_size: Network ``patch_size`` (default ``(1, 2, 2)``). + dtype: Dtype of the zero tensors (should match ``Wan21TransformerConfig.dtype``). + """ + opacity_dim, camera_dim = artifixer_embedding_dims(patch_size) + + shapes: dict[str, tuple[int, ...]] = {} + for layer in range(num_layers): + prefix = f"blocks.{layer}." + shapes[prefix + "opacity_embedding.weight"] = (dim, opacity_dim) + shapes[prefix + "opacity_embedding.bias"] = (dim,) + shapes[prefix + "camera_embedding.weight"] = (dim, camera_dim) + shapes[prefix + "camera_embedding.bias"] = (dim,) + shapes[prefix + "cross_attn.add_k_proj.weight"] = (dim, dim) + shapes[prefix + "cross_attn.add_k_proj.bias"] = (dim,) + shapes[prefix + "cross_attn.add_v_proj.weight"] = (dim, dim) + shapes[prefix + "cross_attn.add_v_proj.bias"] = (dim,) + shapes[prefix + "cross_attn.norm_added_k.weight"] = (dim,) + + def transform(state_dict: dict[str, Tensor]) -> dict[str, Tensor]: + out = dict(state_dict) + for key, shape in shapes.items(): + if key not in out: + out[key] = torch.zeros(shape, dtype=dtype) + return out + + return transform + + +__all__ = [ + "DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING", + "artifixer_dmd_state_dict_transform", + "zero_pad_artifixer_keys", +] diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py new file mode 100644 index 000000000..ea028af39 --- /dev/null +++ b/integrations/artifixer/artifixer/config.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. + +"""Configs for the ArtiFixer DMD-distilled inference recipe. + +ArtiFixer is a reconstruction-enhanced T2V model built on Wan 2.1 1.3B that +adds (a) per-block opacity and Plucker-camera-ray MLPs, (b) neighbor cross- +attention with PRoPE, and (c) opacity-weighted latent mixing. + +The network is :class:`ArtifixerDiTNetwork` whose blocks carry the +opacity + camera-ray MLPs (zero-initialized so behavior matches vanilla +Wan when ``opacity_extra`` / ``camera_extra`` are not provided). + +By default the recipe loads the merged ArtiFixer DMD safetensors via +:func:`artifixer_dmd_state_dict_transform`; set +``ARTIFIXER_USE_BASE_WAN_WEIGHTS=1`` to fall back to vanilla Wan 2.1 +1.3B base weights, in which case :func:`zero_pad_artifixer_keys` pads +the state dict (opacity / camera / neighbor cross-attn keys) so the +strict-mode load succeeds. +""" + +from __future__ import annotations + +import os + +import torch +from artifixer.checkpoint import ( + artifixer_dmd_state_dict_transform, + zero_pad_artifixer_keys, +) +from artifixer.network.dit import ArtifixerDiTNetwork1pt3BConfig +from artifixer.pipeline import ArtifixerInferencePipelineConfig +from artifixer.runner import ArtifixerDmdT2VRunnerConfig +from artifixer.transformer import ArtifixerWanTransformerConfig + +from flashdreams.infra.diffusion.model import DiffusionModelConfig +from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig +from flashdreams.infra.runner import RunnerConfig +from flashdreams.recipes.wan import WanVAEDecoderConfig + +# Mirrors the ArtiFixer DMD stage-3 1.3B training config: +# model_id = Wan-AI/Wan2.1-T2V-1.3B-Diffusers +# frames_per_block = 7 -> len_t +# local_attn_size = 21 -> window_size_t +# sink_size = 7 -> sink_size_t +# num_inference_steps = 4 +# timestep_shift = 5 -> FlowMatchSchedulerConfig.shift +ARTIFIXER_LEN_T = 7 +ARTIFIXER_WINDOW_SIZE_T = 21 +ARTIFIXER_SINK_SIZE_T = 7 +ARTIFIXER_NUM_INFERENCE_STEPS = 4 +ARTIFIXER_TIMESTEP_SHIFT = 5.0 + +# Default: load the merged ArtiFixer DMD safetensors (a consolidated +# single-file checkpoint built from the sharded FSDP training output). +# There is no committed default location, since the merged file is large +# and lives outside the repo; users MUST set +# ``ARTIFIXER_DMD_CHECKPOINT_PATH`` to the merged safetensors path, or +# set ``ARTIFIXER_USE_BASE_WAN_WEIGHTS=1`` to fall back to vanilla Wan +# 2.1 1.3B HuggingFace weights. +ARTIFIXER_DMD_CHECKPOINT_PATH: str | None = os.environ.get( + "ARTIFIXER_DMD_CHECKPOINT_PATH" +) + +# Fallback: vanilla Wan 2.1 1.3B base weights from HF, paired with +# ``zero_pad_artifixer_keys`` so ``load_state_dict`` succeeds in strict +# mode (the ArtiFixer extension paths sit at zero until trained weights +# are loaded). Useful for smoke-testing the recipe wiring when the +# merged safetensors are unavailable; flip ``ARTIFIXER_USE_BASE_WAN_WEIGHTS=1``. +BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH = ( + "https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B/blob/main/diffusion_pytorch_model.safetensors" +) + + +_BASE_NETWORK_CONFIG = ArtifixerDiTNetwork1pt3BConfig( + patch_embedding_type="conv3d", +) +_BASE_TRANSFORMER_DTYPE = torch.bfloat16 + +_USE_BASE_WAN_WEIGHTS = os.environ.get("ARTIFIXER_USE_BASE_WAN_WEIGHTS") == "1" +if _USE_BASE_WAN_WEIGHTS: + _CHECKPOINT_PATH: str = BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH + _STATE_DICT_TRANSFORM = zero_pad_artifixer_keys( + num_layers=_BASE_NETWORK_CONFIG.num_layers, + dim=_BASE_NETWORK_CONFIG.dim, + patch_size=_BASE_NETWORK_CONFIG.patch_size, + dtype=_BASE_TRANSFORMER_DTYPE, + ) +else: + if ARTIFIXER_DMD_CHECKPOINT_PATH is None: + raise RuntimeError( + "ArtiFixer recipe requires a checkpoint path. Set " + "``ARTIFIXER_DMD_CHECKPOINT_PATH`` to the merged DMD " + "safetensors file, or set " + "``ARTIFIXER_USE_BASE_WAN_WEIGHTS=1`` to fall back to " + "vanilla Wan 2.1 1.3B base weights." + ) + _CHECKPOINT_PATH = ARTIFIXER_DMD_CHECKPOINT_PATH + _STATE_DICT_TRANSFORM = artifixer_dmd_state_dict_transform + +PIPELINE_ARTIFIXER_DMD_T2V_1PT3B = ArtifixerInferencePipelineConfig( + recipe_name="artifixer-dmd-wan2.1-t2v-1.3b", + # Profiling is off for the initial bring-up. The base + # ``StreamInferencePipeline.generate`` is what installs + # ``cache.event_profiler`` when this is True, but our custom + # ``ArtifixerInferencePipeline.generate`` bypasses that path. Re-enable + # only alongside the corresponding ``EventProfiler`` setup in + # ``ArtifixerInferencePipeline.generate``. + enable_sync_and_profile=False, + encoder=None, + decoder=WanVAEDecoderConfig(), + diffusion_model=DiffusionModelConfig( + seed=42, + transformer=ArtifixerWanTransformerConfig( + network=_BASE_NETWORK_CONFIG, + dtype=_BASE_TRANSFORMER_DTYPE, + checkpoint_path=_CHECKPOINT_PATH, + state_dict_transform=_STATE_DICT_TRANSFORM, + batch_shape=(), + len_t=ARTIFIXER_LEN_T, + guidance_scale=1.0, + window_size_t=ARTIFIXER_WINDOW_SIZE_T, + sink_size_t=ARTIFIXER_SINK_SIZE_T, + stamp_image_latent=False, + compile_network=True, + ), + scheduler=FlowMatchSchedulerConfig( + num_inference_steps=ARTIFIXER_NUM_INFERENCE_STEPS, + denoising_timesteps=[1000, 750, 500, 250], + warp_denoising_step=True, + shift=ARTIFIXER_TIMESTEP_SHIFT, + sigma_min=0.0, + extra_one_step=True, + num_train_timesteps=1000, + ), + ), +) +RUNNER_ARTIFIXER_DMD_T2V_1PT3B = ArtifixerDmdT2VRunnerConfig( + runner_name=PIPELINE_ARTIFIXER_DMD_T2V_1PT3B.recipe_name, + description=( + "ArtiFixer reconstruction-enhanced T2V (Wan 2.1 1.3B + opacity/camera/neighbor " + "extensions, 4-step DMD)." + ), + pipeline=PIPELINE_ARTIFIXER_DMD_T2V_1PT3B, +) + + +RUNNER_CONFIGS: dict[str, RunnerConfig] = { + cfg.runner_name: cfg + for cfg in (RUNNER_ARTIFIXER_DMD_T2V_1PT3B,) +} diff --git a/integrations/artifixer/artifixer/latent_mix.py b/integrations/artifixer/artifixer/latent_mix.py new file mode 100644 index 000000000..eb19e0b11 --- /dev/null +++ b/integrations/artifixer/artifixer/latent_mix.py @@ -0,0 +1,106 @@ +# 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. + +"""Opacity-weighted latent mixing of noise + VAE-encoded reconstruction. + +Mirrors the ArtiFixer reference's ``ArtifixerPipelineBase.prepare_latents``. +The pipeline replaces the base ``DiffusionModel.generate`` initial-noise +sample with a per-AR-chunk mix:: + + latents = condition * opacity_lat + noise * (1 - opacity_lat) + +where ``opacity_lat`` is the per-pixel opacity max-pooled from the input +resolution down to the VAE latent grid. The mix is the *only* place +opacity influences the inference path before the network forward: the +per-block ``opacity_embedding`` MLPs feed the opacity into +attention through a separate channel. + +The helper is shape-checked and self-contained -- callers supply the noise +tensor so the pipeline can keep its existing RNG / CP-broadcast logic. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + + +def opacity_weighted_latent_mix( + *, + condition: Tensor, + opacity: Tensor, + noise: Tensor, + vae_scale_factor_temporal: int, + vae_scale_factor_spatial: int, + is_first_chunk: bool, +) -> Tensor: + """Compute the ArtiFixer opacity-weighted latent mix for one AR chunk. + + Args: + condition: VAE-encoded reconstruction-rendered RGB, shape + ``(B, in_dim, T_lat, Hl, Wl)``. Same layout the network consumes. + opacity: Per-pixel alpha at input resolution, shape + ``(B, T_input, H_input, W_input)``. The pipeline slices this + from the full-rollout opacity to the current chunk's frame range. + noise: Random noise tensor with the same shape as ``condition``. + vae_scale_factor_temporal: Wan VAE temporal stride (e.g. ``4``). + vae_scale_factor_spatial: Wan VAE spatial stride (e.g. ``8``). + is_first_chunk: When ``True``, the first input frame is left-padded + by 3 copies to absorb the Wan VAE's ``1 + 4`` temporal layout + (one latent frame covers one input frame on the boundary, every + other latent frame covers four). Mirrors the + ``is_first_chunk`` branch in the ArtiFixer reference. + + Returns: + Mixed latent, same shape as ``condition``. + """ + assert condition.shape == noise.shape, ( + f"condition.shape {tuple(condition.shape)} != noise.shape " + f"{tuple(noise.shape)}" + ) + + if is_first_chunk: + opacity = torch.cat( + [opacity[:, :1].repeat_interleave(3, dim=1), opacity], dim=1 + ) + + # max_pool3d expects (N, C, D, H, W); opacity is (B, T, H, W) -> add C=1. + opacity_5d = opacity.unsqueeze(1) + opacity_lat = torch.nn.functional.max_pool3d( + opacity_5d, + kernel_size=( + vae_scale_factor_temporal, + vae_scale_factor_spatial, + vae_scale_factor_spatial, + ), + ) + + expected_lat_shape = ( + opacity_lat.shape[0], + 1, + condition.shape[2], + condition.shape[3], + condition.shape[4], + ) + assert opacity_lat.shape == expected_lat_shape, ( + f"opacity_lat.shape {tuple(opacity_lat.shape)} != expected " + f"{expected_lat_shape}" + ) + + opacity_lat = opacity_lat.to(condition.dtype) + return condition * opacity_lat + noise * (1.0 - opacity_lat) + + +__all__ = ["opacity_weighted_latent_mix"] diff --git a/integrations/artifixer/artifixer/network/__init__.py b/integrations/artifixer/artifixer/network/__init__.py new file mode 100644 index 000000000..0bcd05531 --- /dev/null +++ b/integrations/artifixer/artifixer/network/__init__.py @@ -0,0 +1,24 @@ +# 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. + +# Intentionally empty: importing this package should not transitively pull +# in the full flashdreams dependency tree (RingAttention -> loguru -> boto3 +# etc.). Submodules are imported explicitly by the recipe and by tests:: +# +# from artifixer.network.block import ArtifixerBlock +# from artifixer.network.prope import PropeDotProductAttention +# +# Keeps the PRoPE module importable / testable with only ``torch`` +# installed (no flashdreams runtime required). diff --git a/integrations/artifixer/artifixer/network/block.py b/integrations/artifixer/artifixer/network/block.py new file mode 100644 index 000000000..fa74e7fd9 --- /dev/null +++ b/integrations/artifixer/artifixer/network/block.py @@ -0,0 +1,198 @@ +# 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. + +"""ArtiFixer transformer block with opacity + camera + neighbor extensions. + +Mirrors ``ArtifixerTransformerBlock`` in the ArtiFixer reference +implementation. Per-block extensions on top of :class:`Block`: + + * **Opacity + camera-ray MLPs.** Two ``nn.Linear`` heads project + per-token opacity and Plucker-camera-ray features into the + transformer hidden size; their outputs are added to the AdaLN-normed + hidden states *before* self-attention. Both heads are zero-initialized + so the wrapped block is a no-op extension of base Wan behavior at + load time. + + * **Neighbor cross-attention.** ``cross_attn`` is :class:`ArtifixerCrossAttention`, + which carries ``add_k_proj`` / ``add_v_proj`` / ``norm_added_k`` and a + separate ``attn_op_neighbor`` attention op for the neighbor-frame KV + branch. When the neighbor cache is unset / PRoPE modules are absent + the forward path is identical to ``CrossAttention.forward``. +""" + +from __future__ import annotations + +from typing import Any + +from artifixer.network.cross_attn import ArtifixerCrossAttention +from torch import Tensor, nn +from torch.nn import functional as F + +from flashdreams.recipes.wan.transformer.impl.modules import Block, BlockCache + + +def _layer_norm_fp32(x: Tensor, norm: nn.LayerNorm) -> Tensor: + """Run ``norm`` with fp32 input AND fp32 weight/bias, then cast back. + + Mirrors diffusers' ``FP32LayerNorm`` which the reference + ``WanTransformerBlock`` norms use under the hood: regardless of the + surrounding bf16 cast, the norm itself accumulates in fp32 to keep + the mean/std numerically stable. Plain + ``nn.LayerNorm(elementwise_affine=True)`` stores weight/bias in the + model dtype (bf16 after ``.to(bf16)``), so calling it with an fp32 + input raises ``expected scalar type Float but found BFloat16``. The + work-around is to call ``F.layer_norm`` directly with promoted + parameters. + """ + weight = norm.weight.float() if norm.weight is not None else None + bias = norm.bias.float() if norm.bias is not None else None + return F.layer_norm(x.float(), norm.normalized_shape, weight, bias, norm.eps) + + +class ArtifixerBlock(Block): + """Wan transformer block extended with opacity + camera-ray MLPs.""" + + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + opacity_embedding_dim: int, + camera_embedding_dim: int, + cross_attn_norm: bool = True, + eps: float = 1e-6, + i2v: bool = False, + ) -> None: + super().__init__( + dim=dim, + ffn_dim=ffn_dim, + num_heads=num_heads, + cross_attn_norm=cross_attn_norm, + eps=eps, + i2v=i2v, + ) + + # Replace the inherited cross_attn with the ArtiFixer variant that + # also carries add_k_proj / add_v_proj / norm_added_k for the + # neighbor-frame KV branch. + self.cross_attn = ArtifixerCrossAttention( + query_dim=dim, + n_heads=num_heads, + head_dim=dim // num_heads, + i2v=i2v, + eps=eps, + ) + + self.opacity_embedding = nn.Linear(opacity_embedding_dim, dim, bias=True) + self.camera_embedding = nn.Linear(camera_embedding_dim, dim, bias=True) + + # Zero-init so the wrapped block is a no-op extension of base Wan + # behavior at load time. + nn.init.zeros_(self.opacity_embedding.weight) + nn.init.zeros_(self.opacity_embedding.bias) + nn.init.zeros_(self.camera_embedding.weight) + nn.init.zeros_(self.camera_embedding.bias) + + def forward( # type: ignore[override] + self, + x: Tensor, + e: Tensor, + cache: BlockCache, + rope_freqs: Tensor, + opacity_extra: Tensor | None = None, + camera_extra: Tensor | None = None, + prope_src: Any = None, + prope_tgt: Any = None, + ignore_neighbors: bool = False, + ) -> Tensor: + """Run one transformer block update with ArtiFixer conditioning. + + Extras: + + - ``opacity_extra`` / ``camera_extra``: per-token opacity and + Plucker-camera-ray features added to the AdaLN-normed hidden + states before self-attention. + - ``prope_src`` / ``prope_tgt`` / ``ignore_neighbors``: forwarded + to :class:`ArtifixerCrossAttention.forward` to drive the PRoPE + neighbor branch. The actual neighbor K/V cache lives on the + cross_attn module itself (set via + ``cross_attn.initialize_neighbor_cache``). + + When every extra is ``None`` / ``False`` and the cross_attn's + neighbor cache is unset, this is identical to :meth:`Block.forward`. + """ + assert self._parameters_updated_after_loading_checkpoint, ( + "We expect to have called update_parameters_after_loading_checkpoint() " + "before running the forward pass" + ) + # Mirror the ArtiFixer reference's ``ArtifixerTransformerBlock.forward`` + # and promote the per-block AdaLN modulation, RMS-norms, and residual + # adds to fp32 before casting back to the input dtype. The base + # ``Block.forward`` keeps everything in ``x.dtype`` (typically bf16) + # which is fine in isolation, but with 30 blocks each contributing + # ~1 dB of bf16 noise the cross-backend PSNR drifts from 51 dB at + # block 0 down to 28 dB at block 29. The fp32 promotion below + # mirrors the reference: + # + # * modulation chunking: ``(scale_shift_table + temb).float().chunk(6)`` + # * self-attn pre-norm + AdaLN + # * self-attn residual (x + attn * gate) + # * cross-attn pre-norm + # * FFN pre-norm + AdaLN + # * FFN residual (x + ffn.float() * gate) -- note the reference + # also promotes ``ff_output`` to fp32 inside the residual. + # + # Cost: a handful of extra casts per block. Parity gain: closes + # the per-block drift so the cross-backend per-call PSNR stays + # >50 dB through all 30 layers. The flashdreams ``Block`` + # bf16-throughout path remains the default; this override only + # kicks in when ``ArtifixerBlock`` is used (the artifixer recipe). + e_chunks = (self.modulation.float() + e.float()).chunk(6, dim=-2) + x_dtype = x.dtype + + # ``norm1`` / ``norm2`` in the base Wan ``Block`` are + # ``elementwise_affine=False`` (no weight/bias) so calling them + # with an fp32 input "just works" -- they return fp32. ``norm3`` + # has ``elementwise_affine=True`` and stores its weight/bias in the + # model dtype, so we route through ``_layer_norm_fp32`` to also + # promote the weight/bias to fp32 (matches diffusers' + # FP32LayerNorm used by the reference ``WanTransformerBlock``). + y = (self.norm1(x.float()) * (1 + e_chunks[1]) + e_chunks[0]).to(x_dtype) + if opacity_extra is not None: + y = y + self.opacity_embedding(opacity_extra) + if camera_extra is not None: + y = y + self.camera_embedding(camera_extra) + y = self.self_attn( + y, + rope_freqs=rope_freqs, + kv_cache=cache.self_attn, + ) + x = (x.float() + y * e_chunks[2]).to(x_dtype) + + # Cross-attn pre-norm in fp32 (mirrors the reference). The + # post-cross-attn residual is NOT promoted in the reference so we + # keep that in the input dtype here too. + x = x + self.cross_attn( + _layer_norm_fp32(x, self.norm3).to(x_dtype), + kv_cache=cache.cross_attn, + prope_src=prope_src, + prope_tgt=prope_tgt, + ignore_neighbors=ignore_neighbors, + ) + y = (self.norm2(x.float()) * (1 + e_chunks[4]) + e_chunks[3]).to(x_dtype) + y = self.ffn(y) + # FFN residual with ff_output also promoted (matches the reference). + x = (x.float() + y.float() * e_chunks[5]).to(x_dtype) + return x diff --git a/integrations/artifixer/artifixer/network/cross_attn.py b/integrations/artifixer/artifixer/network/cross_attn.py new file mode 100644 index 000000000..33dd8599c --- /dev/null +++ b/integrations/artifixer/artifixer/network/cross_attn.py @@ -0,0 +1,209 @@ +# 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. + +"""ArtiFixer cross-attention with a neighbor-frame KV bank + PRoPE. + +Adds a third KV pathway (parallel to text and optional I2V image) for +attending over VAE-encoded neighbor frames. Module naming mirrors the +diffusers ``WanAttention`` extension in the ArtiFixer reference: + + - ``add_k_proj`` / ``add_v_proj`` — Linear projections from latent + neighbor context to K / V + - ``norm_added_k`` — RMSNorm on K (matches ``norm_k`` shape) + - ``attn_op_neighbor`` — separate RingAttention op for the neighbor + branch (lets q/k be PRoPE-transformed independently of the text + branch) + +``forward`` accepts optional ``neighbor_kv_cache`` + ``prope_src`` / +``prope_tgt`` modules. When all three are ``None`` (text-only path), +the call is identical to ``CrossAttention.forward``. Otherwise the +neighbor branch runs:: + + q_pr = prope_src._apply_to_q(q) + k_pr = prope_tgt._apply_to_kv(k_neighbor) + v_pr = prope_tgt._apply_to_kv(v_neighbor) + o_n = prope_src._apply_to_o(attn_op_neighbor(q_pr, k_pr, v_pr)) + out = out_text [+ out_img] + (0 if ignore_neighbors else 1) * o_n + +PRoPE math runs at fp32 with a ``.float() / .to(query.dtype)`` round +trip (matching the reference), since ``_rope_precompute_coeffs`` does +an int64 -> fp32 promotion internally. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +from torch import Tensor, nn + +from flashdreams.core.attention.kvcache import BlockKVCache +from flashdreams.core.attention.ring import RingAttention +from flashdreams.recipes.wan.transformer.impl.modules import ( + CrossAttention, + CrossAttnCache, +) + + +class ArtifixerCrossAttention(CrossAttention): + """Cross-attention extended with a neighbor-frame KV branch.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + # Parameter names match the ArtiFixer reference's WanAttention + # (``add_k_proj`` / ``add_v_proj`` / ``norm_added_k``) so loading + # a merged ArtiFixer DMD checkpoint only needs the diffusers + # ``attn2 -> cross_attn`` regex remap, not a per-key rename. + self.add_k_proj = nn.Linear(self.context_dim, self.inner_dim, bias=True) + self.add_v_proj = nn.Linear(self.context_dim, self.inner_dim, bias=True) + self.norm_added_k = nn.RMSNorm(self.inner_dim, eps=self.eps) + + # Zero-init add_v_proj only (matches the ArtiFixer reference). + # Because attention is softmax(Q K^T) @ V, V=0 forces the neighbor + # branch contribution to zero at load time even if add_k_proj is + # non-zero, so the wrapped block is a no-op extension of base Wan + # behavior. + nn.init.zeros_(self.add_v_proj.weight) + nn.init.zeros_(self.add_v_proj.bias) + + self.attn_op_neighbor = RingAttention(qkv_format="bshd", backend="cudnn") + + # Set once per rollout via :meth:`initialize_neighbor_cache`. Cleared + # to ``None`` when no neighbor context is provided. Non-parameter + # module state -- safe to mutate outside ``forward``. + self.neighbor_kv_cache: BlockKVCache | None = None + + def set_context_parallel_group(self, cp_group: Any) -> None: + """Set CP group on every attention op, including the neighbor branch.""" + super().set_context_parallel_group(cp_group) + self.attn_op_neighbor.set_context_parallel_group(cp_group=cp_group) + + def compute_kv_neighbor(self, context: Tensor) -> BlockKVCache: + """Project neighbor ``context`` into a static K/V cache. + + Args: + context: Neighbor context, shape ``[..., L_neighbor, context_dim]``. + + Returns: + ``BlockKVCache`` carrying ``(k, v)`` reshaped to + ``[batch_size, L, n_heads, head_dim]``. Static (no rolling). + """ + batch_shape = context.shape[:-2] + batch_size = math.prod(batch_shape) + L, _ = context.shape[-2:] + n, d = self.n_heads, self.head_dim + + k = self.norm_added_k(self.add_k_proj(context)).reshape(batch_size, L, n, d) + v = self.add_v_proj(context).reshape(batch_size, L, n, d) + return BlockKVCache.from_tensor(k, v, seq_dim=-3) + + def initialize_neighbor_cache(self, context: Tensor | None) -> None: + """Build and store the static neighbor KV cache for this module. + + Call once per rollout. ``context=None`` clears the cache so the + forward path skips the neighbor branch (text-only mode). + """ + self.neighbor_kv_cache = ( + None if context is None else self.compute_kv_neighbor(context) + ) + + def forward( # type: ignore[override] + self, + x: Tensor, + kv_cache: CrossAttnCache, + *, + prope_src: Any = None, + prope_tgt: Any = None, + ignore_neighbors: bool = False, + ) -> Tensor: + """Cross-attention with optional PRoPE neighbor branch. + + Uses the per-module ``self.neighbor_kv_cache`` populated by + :meth:`initialize_neighbor_cache`. When that is ``None`` (or PRoPE + modules are ``None``), behavior is identical to + :meth:`CrossAttention.forward`. Otherwise the PRoPE-modulated + neighbor branch is added to the output. + + Args: + x: Query tokens ``[..., L, query_dim]``. + kv_cache: Text (+ optional image) K/V cache from the base class. + prope_src: PRoPE module for the source / target cameras + (transforms q on read, o on write). + prope_tgt: PRoPE module for the neighbor cameras (transforms + k, v on read). + ignore_neighbors: If ``True``, zero out the neighbor contribution. + Matches the diffusion-forcing CFG dropout knob in the + ArtiFixer reference. + """ + neighbor_kv_cache = self.neighbor_kv_cache + batch_shape = x.shape[:-2] + batch_size = math.prod(batch_shape) + L, D = x.shape[-2:] + n, d = self.n_heads, self.head_dim + assert n * d == D, "n * d must be equal to D" + + q = self.norm_q(self.q(x)).reshape(batch_size, L, n, d) + out = self.attn_op(q, kv_cache.text.cached_k(), kv_cache.text.cached_v()) + + if self.i2v: + assert kv_cache.img is not None, ( + "kv_cache_img is expected to be provided for I2V cross-attention" + ) + out_img = self.attn_op_image( + q, kv_cache.img.cached_k(), kv_cache.img.cached_v() + ) + out = out + out_img + + if neighbor_kv_cache is not None: + assert prope_src is not None and prope_tgt is not None, ( + "prope_src and prope_tgt are required when neighbor_kv_cache is provided" + ) + q_dtype = q.dtype + + # PRoPE expects (B, H, L, D); attn ops use (B, L, H, D) (bshd). + # Run PRoPE math at fp32 with a dtype round-trip to mirror the + # ArtiFixer reference. + q_pr = ( + prope_src._apply_to_q(q.transpose(1, 2).float()) + .transpose(1, 2) + .to(q_dtype) + ) + k_n = neighbor_kv_cache.cached_k() + v_n = neighbor_kv_cache.cached_v() + k_pr = ( + prope_tgt._apply_to_kv(k_n.transpose(1, 2).float()) + .transpose(1, 2) + .to(q_dtype) + ) + v_pr = ( + prope_tgt._apply_to_kv(v_n.transpose(1, 2).float()) + .transpose(1, 2) + .to(q_dtype) + ) + + out_neighbor = self.attn_op_neighbor(q_pr, k_pr, v_pr) + out_neighbor = ( + prope_src._apply_to_o(out_neighbor.transpose(1, 2).float()) + .transpose(1, 2) + .to(q_dtype) + ) + + if not ignore_neighbors: + out = out + out_neighbor + + out = out.reshape(batch_shape + (L, n * d)) + return self.o(out) diff --git a/integrations/artifixer/artifixer/network/dit.py b/integrations/artifixer/artifixer/network/dit.py new file mode 100644 index 000000000..489b1f57d --- /dev/null +++ b/integrations/artifixer/artifixer/network/dit.py @@ -0,0 +1,116 @@ +# 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. + +"""ArtiFixer DiT network: ``WanDiTNetwork`` with opacity / camera blocks.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from artifixer.network.block import ArtifixerBlock +from torch import Tensor + +from flashdreams.recipes.wan.transformer.impl.network import ( + WanDiTNetwork, + WanDiTNetwork1pt3BConfig, + WanDiTNetworkConfig, +) + +# Wan2.1 VAE downsampling factors. Identical for T2V-1.3B and T2V-14B. +# Used to compute the per-block opacity / camera input dims. +WAN_VAE_SCALE_FACTOR_SPATIAL = 8 +WAN_VAE_SCALE_FACTOR_TEMPORAL = 4 + + +def artifixer_embedding_dims( + patch_size: tuple[int, int, int], + vae_scale_factor_spatial: int = WAN_VAE_SCALE_FACTOR_SPATIAL, + vae_scale_factor_temporal: int = WAN_VAE_SCALE_FACTOR_TEMPORAL, +) -> tuple[int, int]: + """Compute (opacity_embedding_dim, camera_embedding_dim) for ArtiFixer. + + opacity_embedding_dim = vae_t * vae_s * ph * vae_s * pw + camera_embedding_dim = vae_s * ph * vae_s * pw * 6 (6 = Plucker channels) + + For Wan 2.1 default ``patch_size = (1, 2, 2)`` with VAE strides + ``(4, 8, 8)`` this gives ``(1024, 1536)``. + """ + _, ph, pw = patch_size + opacity_dim = ( + vae_scale_factor_temporal + * vae_scale_factor_spatial + * ph + * vae_scale_factor_spatial + * pw + ) + camera_dim = vae_scale_factor_spatial * ph * vae_scale_factor_spatial * pw * 6 + return opacity_dim, camera_dim + + +@dataclass +class ArtifixerDiTNetworkConfig(WanDiTNetworkConfig): + """ArtiFixer DiT network config. + + Inherits every field from ``WanDiTNetworkConfig`` and overrides + ``_target`` to construct an :class:`ArtifixerDiTNetwork`. The + opacity / camera input dims are derived from the VAE strides + patch + size via :func:`artifixer_embedding_dims`, not stored on the config, + so they cannot drift from the model that produced the conditioning + tensors. + """ + + _target: type = field(default_factory=lambda: ArtifixerDiTNetwork) + + +@dataclass +class ArtifixerDiTNetwork1pt3BConfig(WanDiTNetwork1pt3BConfig): + """ArtiFixer 1.3B variant (matches the ArtiFixer reference model size).""" + + _target: type = field(default_factory=lambda: ArtifixerDiTNetwork) + + +class ArtifixerDiTNetwork(WanDiTNetwork): + """``WanDiTNetwork`` whose blocks are :class:`ArtifixerBlock`.""" + + def _build_block(self, layer_idx: int) -> ArtifixerBlock: + opacity_dim, camera_dim = artifixer_embedding_dims(self.patch_size) + return ArtifixerBlock( + dim=self.dim, + ffn_dim=self.ffn_dim, + num_heads=self.num_heads, + opacity_embedding_dim=opacity_dim, + camera_embedding_dim=camera_dim, + cross_attn_norm=self.cross_attn_norm, + eps=self.eps, + i2v=self.cross_attn_enable_img, + ) + + def initialize_neighbor_kv_caches(self, context: Tensor | None) -> None: + """Push neighbor context into every block's cross-attention. + + Call once per rollout, after :meth:`initialize_cache`. ``context=None`` + clears the per-module caches so subsequent forward passes skip the + PRoPE neighbor branch (vanilla T2V behavior). + + Args: + context: Neighbor latent context shape + ``[..., L_neighbor, dim]``. Same dim as the transformer + hidden size; provided by the pipeline after VAE-encoding + the neighbor frames and projecting them through + ``patch_embedding``. + """ + for block in self.blocks: + assert isinstance(block, ArtifixerBlock) + block.cross_attn.initialize_neighbor_cache(context) diff --git a/integrations/artifixer/artifixer/network/patches.py b/integrations/artifixer/artifixer/network/patches.py new file mode 100644 index 000000000..11d29cbae --- /dev/null +++ b/integrations/artifixer/artifixer/network/patches.py @@ -0,0 +1,131 @@ +# 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. + +"""Patchification of opacity / Plucker camera-ray conditioning tensors. + +Mirrors the ArtiFixer reference's per-chunk patchification: + + * opacity (B, T, H, W) -> opacity_extra_patches (B, L, opacity_embedding_dim) + * camera_rays (B, T, H, W, 6) -> camera_extra_patches (B, L, camera_embedding_dim) + +where ``L`` is the post-patch token count for the chunk. + +The Wan VAE has a ``1 + 4 * (n - 1)`` temporal layout: the first decoded +frame corresponds to a single latent frame, every subsequent latent frame +covers 4 input frames. When the first patchify call covers ``frame_offset == 0`` +(i.e. the first AR chunk), we left-pad the input frame stack by 3 copies of +the first frame so the rearrange treats every latent frame uniformly. Later +chunks already have the temporal downsampling baked in. +""" + +from __future__ import annotations + +import torch +from einops import rearrange +from torch import Tensor + + +def patchify_opacity( + opacity: Tensor, + *, + vae_scale_factor_temporal: int, + vae_scale_factor_spatial: int, + patch_size: tuple[int, int, int], + frame_offset: int, +) -> Tensor: + """Rearrange per-pixel opacity into per-token features. + + Args: + opacity: ``(B, T, H, W)`` -- raw per-frame alpha at input resolution. + vae_scale_factor_temporal: e.g. ``4`` for Wan 2.1. + vae_scale_factor_spatial: e.g. ``8`` for Wan 2.1. + patch_size: Network ``patch_size``, ``(kt, kh, kw)``. + frame_offset: Logical start of this chunk in latent-time tokens. + When ``0`` we pad the first frame by 3 copies to account for the +Wan VAE's ``1 + 4`` temporal layout (mirrors the reference). + + Returns: + ``(B, L, opacity_embedding_dim)`` where + ``L = (T / vae_t / kt) * (H / vae_s / kh) * (W / vae_s / kw)``. + """ + _, ph, pw = patch_size + + if frame_offset == 0: + opacity = torch.cat( + [opacity[:, :1].repeat_interleave(3, dim=1), opacity], dim=1 + ) + extra = rearrange( + opacity, + "b (t t4) (h h8) (w w8) -> b (h8 w8 t4) t h w", + h8=vae_scale_factor_spatial * ph, + w8=vae_scale_factor_spatial * pw, + t4=vae_scale_factor_temporal, + ) + return extra.flatten(2).transpose(1, 2) + + +def patchify_camera_rays( + camera_rays: Tensor, + *, + hidden_post_patch_t: int, + vae_scale_factor_temporal: int, + vae_scale_factor_spatial: int, + patch_size: tuple[int, int, int], + frame_offset: int, +) -> Tensor: + """Rearrange per-pixel Plucker rays into per-token features. + + Handles the reference's branch on whether ``camera_rays``' temporal + axis is already at the latent rate + (``camera_rays.shape[1] == hidden_post_patch_t``) or at the input rate + (needs the VAE ``1 + 4`` left-pad on the first AR chunk). + + Args: + camera_rays: ``(B, T, H, W, 6)`` Plucker rays. + hidden_post_patch_t: Number of post-patch latent frames in the chunk + (i.e. ``hidden_states.shape[2]`` after :meth:`patch_embedding`). + vae_scale_factor_temporal: e.g. ``4`` for Wan 2.1. + vae_scale_factor_spatial: e.g. ``8`` for Wan 2.1. + patch_size: Network ``patch_size``, ``(kt, kh, kw)``. + frame_offset: Logical start of this chunk in latent-time tokens. + + Returns: + ``(B, L, camera_embedding_dim)``. + """ + _, ph, pw = patch_size + + if camera_rays.shape[1] == hidden_post_patch_t: + extra = rearrange( + camera_rays, + "b t (h h8) (w w8) c -> b (c h8 w8) t h w", + h8=vae_scale_factor_spatial * ph, + w8=vae_scale_factor_spatial * pw, + ) + else: + if frame_offset == 0: + camera_rays = torch.cat( + [camera_rays[:, :1].repeat_interleave(3, dim=1), camera_rays], dim=1 + ) + extra = rearrange( + camera_rays, + "b (t t4) (h h8) (w w8) c -> b (c h8 w8 t4) t h w", + h8=vae_scale_factor_spatial * ph, + w8=vae_scale_factor_spatial * pw, + t4=vae_scale_factor_temporal, + ) + return extra.flatten(2).transpose(1, 2) + + +__all__ = ["patchify_opacity", "patchify_camera_rays"] diff --git a/integrations/artifixer/artifixer/network/prope.py b/integrations/artifixer/artifixer/network/prope.py new file mode 100644 index 000000000..e91fc82c7 --- /dev/null +++ b/integrations/artifixer/artifixer/network/prope.py @@ -0,0 +1,310 @@ +# MIT License +# +# Copyright (c) Authors of +# "Cameras as Relative Positional Encoding" https://arxiv.org/pdf/2507.10496 +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""PRoPE attention with precomputed RoPE coefficients. + +Port of the ArtiFixer reference's PRoPE attention module, matching the +paper "Cameras as Relative Positional Encoding" (arXiv 2507.10496). The +math is unchanged; the only difference vs the reference is that +``invert_SE3`` is inlined here so this module is self-contained. + +Usage for cross-attention:: + + attn_src = PropeDotProductAttention(...) + attn_tgt = PropeDotProductAttention(...) + attn_src._precompute_and_cache_apply_fns(viewmats_src, Ks_src) + attn_tgt._precompute_and_cache_apply_fns(viewmats_tgt, Ks_tgt) + q_src = attn_src._apply_to_q(q_src) + k_tgt = attn_tgt._apply_to_kv(k_tgt) + v_tgt = attn_tgt._apply_to_kv(v_tgt) + o_src = F.scaled_dot_product_attention(q_src, k_tgt, v_tgt, **kwargs) + o_src = attn_src._apply_to_o(o_src) + +A numerical-parity unit test against the ArtiFixer reference lives at +``integrations/artifixer/tests/test_prope_parity.py``. +""" + +from __future__ import annotations + +from functools import partial +from typing import Callable + +import torch +from torch import nn + + +class PropeDotProductAttention(nn.Module): + """PRoPE attention with precomputed RoPE coefficients.""" + + coeffs_x_0: torch.Tensor + coeffs_x_1: torch.Tensor + coeffs_y_0: torch.Tensor + coeffs_y_1: torch.Tensor + + def __init__( + self, + head_dim: int, + patches_x: int = 0, + patches_y: int = 0, + freq_base: float = 100.0, + freq_scale: float = 1.0, + ) -> None: + super().__init__() + self.head_dim = head_dim + self.patches_x = patches_x + self.patches_y = patches_y + self.freq_base = freq_base + self.freq_scale = freq_scale + + def update_coeffs( + self, patches_x: int, patches_y: int, device: str | torch.device + ) -> None: + self.patches_x = patches_x + self.patches_y = patches_y + coeffs_x = _rope_precompute_coeffs( + torch.tile(torch.arange(patches_x, device=device), (patches_y,)), + freq_base=self.freq_base, + freq_scale=self.freq_scale, + feat_dim=self.head_dim // 4, + ) + coeffs_y = _rope_precompute_coeffs( + torch.repeat_interleave(torch.arange(patches_y, device=device), patches_x), + freq_base=self.freq_base, + freq_scale=self.freq_scale, + feat_dim=self.head_dim // 4, + ) + # Non-persistent buffers: camera count may change between train/eval. + self.coeffs_x_0 = nn.Buffer(coeffs_x[0], persistent=False) + self.coeffs_x_1 = nn.Buffer(coeffs_x[1], persistent=False) + self.coeffs_y_0 = nn.Buffer(coeffs_y[0], persistent=False) + self.coeffs_y_1 = nn.Buffer(coeffs_y[1], persistent=False) + + def _precompute_and_cache_apply_fns( + self, viewmats: torch.Tensor, Ks_norm: torch.Tensor + ) -> None: + batch, cameras, _, _ = viewmats.shape + assert viewmats.shape == (batch, cameras, 4, 4) + assert Ks_norm.shape == (batch, cameras, 3, 3) + self.cameras = cameras + + self.apply_fn_q, self.apply_fn_kv, self.apply_fn_o = _prepare_apply_fns( + head_dim=self.head_dim, + viewmats=viewmats, + Ks_norm=Ks_norm, + coeffs_x=(self.coeffs_x_0, self.coeffs_x_1), + coeffs_y=(self.coeffs_y_0, self.coeffs_y_1), + ) + + def _apply_to_q(self, q: torch.Tensor) -> torch.Tensor: + batch, num_heads, seqlen, head_dim = q.shape + assert seqlen == self.cameras * self.patches_x * self.patches_y + assert head_dim == self.head_dim + return self.apply_fn_q(q) + + def _apply_to_kv(self, kv: torch.Tensor) -> torch.Tensor: + batch, num_heads, seqlen, head_dim = kv.shape + assert seqlen == self.cameras * self.patches_x * self.patches_y + assert head_dim == self.head_dim + return self.apply_fn_kv(kv) + + def _apply_to_o(self, o: torch.Tensor) -> torch.Tensor: + batch, num_heads, seqlen, head_dim = o.shape + assert seqlen == self.cameras * self.patches_x * self.patches_y + assert head_dim == self.head_dim + return self.apply_fn_o(o) + + +def _prepare_apply_fns( + head_dim: int, + viewmats: torch.Tensor, + Ks_norm: torch.Tensor, + coeffs_x: tuple[torch.Tensor, torch.Tensor], + coeffs_y: tuple[torch.Tensor, torch.Tensor], +) -> tuple[ + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], + Callable[[torch.Tensor], torch.Tensor], +]: + """Prepare transforms for PRoPE-style positional encoding. + + - ``K`` is an ``image<-camera`` transform. + - ``viewmats`` is a ``camera<-world`` transform. + - ``P = lift(K) @ viewmats`` is an ``image<-world`` transform. + """ + batch, cameras, _, _ = viewmats.shape + + P = torch.einsum("...ij,...jk->...ik", _lift_K(Ks_norm), viewmats) + P_T = P.transpose(-1, -2) + P_inv = torch.einsum( + "...ij,...jk->...ik", + _invert_SE3(viewmats), + _lift_K(_invert_K(Ks_norm)), + ) + assert P.shape == P_inv.shape == (batch, cameras, 4, 4) + assert head_dim % 4 == 0 + + transforms_q = [ + (partial(_apply_tiled_projmat, matrix=P_T), head_dim // 2), + (partial(_rope_apply_coeffs, coeffs=coeffs_x), head_dim // 4), + (partial(_rope_apply_coeffs, coeffs=coeffs_y), head_dim // 4), + ] + transforms_kv = [ + (partial(_apply_tiled_projmat, matrix=P_inv), head_dim // 2), + (partial(_rope_apply_coeffs, coeffs=coeffs_x), head_dim // 4), + (partial(_rope_apply_coeffs, coeffs=coeffs_y), head_dim // 4), + ] + transforms_o = [ + (partial(_apply_tiled_projmat, matrix=P), head_dim // 2), + (partial(_rope_apply_coeffs, coeffs=coeffs_x, inverse=True), head_dim // 4), + (partial(_rope_apply_coeffs, coeffs=coeffs_y, inverse=True), head_dim // 4), + ] + + apply_fn_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q) + apply_fn_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv) + apply_fn_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o) + return apply_fn_q, apply_fn_kv, apply_fn_o + + +def _apply_tiled_projmat( + feats: torch.Tensor, + matrix: torch.Tensor, +) -> torch.Tensor: + """Apply projection matrix to features. + + ``seqlen => (cameras, patches_x * patches_y)`` and + ``feat_dim => (feat_dim // 4, 4)``. + """ + batch, num_heads, seqlen, feat_dim = feats.shape + cameras = matrix.shape[1] + assert seqlen > cameras and seqlen % cameras == 0 + D = matrix.shape[-1] + assert matrix.shape == (batch, cameras, D, D) + assert feat_dim % D == 0 + return torch.einsum( + "bcij,bncpkj->bncpki", + matrix, + feats.reshape((batch, num_heads, cameras, -1, feat_dim // D, D)), + ).reshape(feats.shape) + + +def _rope_precompute_coeffs( + positions: torch.Tensor, + freq_base: float, + freq_scale: float, + feat_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Precompute RoPE (cos, sin) coefficients for given positions.""" + assert positions.ndim == 1 + assert feat_dim % 2 == 0 + num_freqs = feat_dim // 2 + freqs = freq_scale * ( + freq_base + ** ( + -torch.arange(num_freqs, device=positions.device)[None, None, None, :] + / num_freqs + ) + ) + angles = positions[None, None, :, None] * freqs + assert angles.shape == (1, 1, positions.shape[0], num_freqs) + return torch.cos(angles), torch.sin(angles) + + +def _rope_apply_coeffs( + feats: torch.Tensor, + coeffs: tuple[torch.Tensor, torch.Tensor], + inverse: bool = False, +) -> torch.Tensor: + """Apply RoPE coefficients to features with the *split* ordering convention.""" + cos, sin = coeffs + if cos.shape[2] != feats.shape[2]: + n_repeats = feats.shape[2] // cos.shape[2] + cos = cos.repeat(1, 1, n_repeats, 1) + sin = sin.repeat(1, 1, n_repeats, 1) + assert feats.ndim == cos.ndim == sin.ndim == 4 + assert cos.shape[-1] == sin.shape[-1] == feats.shape[-1] // 2 + x_in = feats[..., : feats.shape[-1] // 2] + y_in = feats[..., feats.shape[-1] // 2 :] + return torch.cat( + ( + [cos * x_in + sin * y_in, -sin * x_in + cos * y_in] + if not inverse + else [cos * x_in - sin * y_in, sin * x_in + cos * y_in] + ), + dim=-1, + ) + + +def _apply_block_diagonal( + feats: torch.Tensor, + func_size_pairs: list[tuple[Callable[[torch.Tensor], torch.Tensor], int]], +) -> torch.Tensor: + """Apply a block-diagonal function to an input array. + + Each function is specified as a tuple ``((Tensor) -> Tensor, int)`` where + the integer is the size of the input slice consumed by the function. + """ + funcs, block_sizes = zip(*func_size_pairs) + assert feats.shape[-1] == sum(block_sizes) + x_blocks = torch.split(feats, block_sizes, dim=-1) + out = torch.cat([f(x) for f, x in zip(funcs, x_blocks)], dim=-1) + assert out.shape == feats.shape + return out + + +def _lift_K(Ks: torch.Tensor) -> torch.Tensor: + """Lift 3x3 matrices to homogeneous 4x4 matrices.""" + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros(Ks.shape[:-2] + (4, 4), device=Ks.device, dtype=Ks.dtype) + out[..., :3, :3] = Ks + out[..., 3, 3] = 1.0 + return out + + +def _invert_K(Ks: torch.Tensor) -> torch.Tensor: + """Invert 3x3 intrinsics matrices. Assumes no skew.""" + assert Ks.shape[-2:] == (3, 3) + out = torch.zeros_like(Ks) + out[..., 0, 0] = 1.0 / Ks[..., 0, 0] + out[..., 1, 1] = 1.0 / Ks[..., 1, 1] + out[..., 0, 2] = -Ks[..., 0, 2] / Ks[..., 0, 0] + out[..., 1, 2] = -Ks[..., 1, 2] / Ks[..., 1, 1] + out[..., 2, 2] = 1.0 + return out + + +def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: + """Invert a 4x4 SE(3) matrix. + + Inlined from the reference's ``pose_utils.invert_SE3`` so this module + has no cross-repo dependencies. + """ + assert transforms.shape[-2:] == (4, 4) + Rinv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = Rinv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + return out + + +__all__ = ["PropeDotProductAttention"] diff --git a/integrations/artifixer/artifixer/pipeline.py b/integrations/artifixer/artifixer/pipeline.py new file mode 100644 index 000000000..ae21013df --- /dev/null +++ b/integrations/artifixer/artifixer/pipeline.py @@ -0,0 +1,535 @@ +# 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. + +"""ArtiFixer inference pipeline. + +Extends :class:`WanInferencePipeline` with the ArtiFixer-specific +conditioning surface: full-rollout opacity, Plucker camera rays, target +camera w2c/Ks, optional VAE-encoded neighbor frames + neighbor camera +matrices. + +``initialize_cache``: + + - VAE-encoded condition latent and neighbor latent arrive pre-computed + from the caller so this pipeline does not need its own VAE encoder. + - Runs the base ``WanInferencePipeline.initialize_cache`` (text + encoder + text K/V build) with ``image=None``. + - Projects the neighbor latent through ``patch_embedding`` to get the + per-token neighbor context, then pushes it into every block's + :class:`ArtifixerCrossAttention.neighbor_kv_cache` via + ``ArtifixerDiTNetwork.initialize_neighbor_kv_caches``. + - Updates the patches_x/patches_y RoPE coefficients on both PRoPE + modules and precomputes the neighbor-side ``apply_fns`` (the + neighbor cameras are static across AR steps). + - Stashes the full-rollout opacity / camera_rays / w2cs / Ks on the + cache for the per-AR ``generate`` slicing. + +``generate`` runs per-AR-chunk PRoPE-src precompute, opacity-weighted +latent mix, manual denoise loop with prepare_latents renoise, and decode. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +from artifixer.latent_mix import opacity_weighted_latent_mix +from artifixer.network.patches import patchify_camera_rays, patchify_opacity +from artifixer.transformer import ArtifixerCtrl, ArtifixerWanTransformer +from torch import Tensor + +from flashdreams.infra.diffusion.model import DiffusionModel +from flashdreams.infra.diffusion.scheduler.fm import FlowMatchScheduler +from flashdreams.recipes.wan.pipeline import ( + WanInferencePipeline, + WanInferencePipelineCache, + WanInferencePipelineConfig, +) + + +@dataclass(kw_only=True) +class ArtifixerInferencePipelineCache(WanInferencePipelineCache): + """Per-rollout cache for the ArtiFixer pipeline. + + Holds the full-rollout conditioning tensors so :meth:`generate` can + slice the chunk's frame range on demand. All optional fields default + to ``None`` so callers that do not have neighbors still construct a + valid cache. + """ + + condition_latent: Tensor | None = None + """VAE-encoded rendered RGB at latent resolution + ``[*batch_shape, in_dim, T_lat, Hl, Wl]``. Pre-encoded by the + caller.""" + + opacity: Tensor | None = None + """Per-pixel alpha at input resolution ``[*batch_shape, T_input, H, W]``, + sliced per AR chunk inside ``generate`` and fed to the + opacity-weighted latent mix.""" + + camera_rays: Tensor | None = None + """Plucker rays at input resolution + ``[*batch_shape, T_input, H, W, 6]`` (or already at the latent rate; + :func:`patchify_camera_rays` handles both).""" + + w2cs: Tensor | None = None + """Target camera world-to-camera matrices ``[*batch_shape, T_lat, 4, 4]`` + indexed by latent-frame.""" + + Ks: Tensor | None = None + """Target camera intrinsics ``[*batch_shape, T_lat, 3, 3]`` + indexed by latent-frame.""" + + neighbor_w2cs: Tensor | None = None + """Optional neighbor camera w2cs ``[*batch_shape, N_neighbors, 4, 4]``.""" + + neighbor_Ks: Tensor | None = None + """Optional neighbor camera intrinsics ``[*batch_shape, N_neighbors, 3, 3]``.""" + + +@dataclass(kw_only=True) +class ArtifixerInferencePipelineConfig(WanInferencePipelineConfig): + """Config for the ArtiFixer inference pipeline.""" + + _target: type["ArtifixerInferencePipeline"] = field( + default_factory=lambda: ArtifixerInferencePipeline + ) + + +class ArtifixerInferencePipeline(WanInferencePipeline): + """ArtiFixer inference pipeline. + + Subclasses :class:`WanInferencePipeline` to thread the per-rollout + conditioning (opacity, camera rays, target cameras, neighbor latent + + cameras) into the transformer and its per-block cross-attention + modules. + """ + + @torch.no_grad() + def initialize_cache( # type: ignore[override] + self, + text: list[str] | None = None, + *, + text_embeddings: Tensor | None = None, + condition_latent: Tensor, + opacity: Tensor, + camera_rays: Tensor, + w2cs: Tensor, + Ks: Tensor, + neighbor_latent: Tensor | None = None, + neighbor_w2cs: Tensor | None = None, + neighbor_Ks: Tensor | None = None, + height: int, + width: int, + ) -> ArtifixerInferencePipelineCache: + """Build a per-rollout cache and push static state into the network. + + Args: + text: One prompt per batch element. Mutually exclusive with + ``text_embeddings``: pass either the raw strings (we run + UMT5 internally) or the pre-encoded UMT5 embeddings (the + external driver may pass them directly). + text_embeddings: Pre-encoded UMT5 prompt embeddings + ``[B, L, D]``. Skips the in-pipeline UMT5 forward when set. + condition_latent: VAE-encoded reconstruction-rendered RGB, + ``[B, in_dim, T_lat, Hl, Wl]`` -- the caller handles VAE + encoding so this pipeline does not need its own encoder. + opacity: Per-pixel alpha at input resolution + ``[B, T_input, H, W]``. + camera_rays: Plucker rays + ``[B, T_lat or T_input, H, W, 6]``. + w2cs: Target camera w2c ``[B, T_lat, 4, 4]``. + Ks: Target camera intrinsics ``[B, T_lat, 3, 3]``. + neighbor_latent: Optional VAE-encoded neighbor frames, + ``[B, in_dim, T_n_lat, H_n_lat, W_n_lat]``. When ``None`` + the neighbor cross-attention branch is disabled. + neighbor_w2cs: Required iff ``neighbor_latent`` is set. + neighbor_Ks: Required iff ``neighbor_latent`` is set. + height: Pre-patchify latent height (post-VAE), matching + ``condition_latent.shape[-2]``. + width: Pre-patchify latent width (post-VAE). + """ + assert condition_latent.shape[-2] == height, ( + f"condition_latent height {condition_latent.shape[-2]} != " + f"argument height {height}" + ) + assert condition_latent.shape[-1] == width, ( + f"condition_latent width {condition_latent.shape[-1]} != " + f"argument width {width}" + ) + assert (text is None) ^ (text_embeddings is None), ( + "Pass exactly one of ``text`` (raw prompts) or " + "``text_embeddings`` (pre-encoded UMT5 output)." + ) + + # 1. Base text / image cross-attn cache (image=None for ArtiFixer). + if text is not None: + base_cache = super().initialize_cache( + text=text, height=height, width=width + ) + else: + assert text_embeddings is not None + base_cache = self._initialize_cache_from_text_embeddings( + text_embeddings=text_embeddings, height=height, width=width + ) + + # 2. Neighbor wiring. We patchify the neighbor latent through the + # network's ``patch_embedding`` so the cross-attention's + # ``add_k_proj`` sees per-token features in the transformer + # hidden dim, then push the result into every block's cross-attn + # neighbor cache. + transformer = self.diffusion_model.transformer + assert isinstance(transformer, ArtifixerWanTransformer), ( + f"ArtifixerInferencePipeline requires an ArtifixerWanTransformer " + f"diffusion_model.transformer, got {type(transformer).__name__}" + ) + + if neighbor_latent is not None: + assert neighbor_w2cs is not None and neighbor_Ks is not None, ( + "neighbor_w2cs and neighbor_Ks are required when " + "neighbor_latent is provided" + ) + neighbor_context = self._encode_neighbor_context(neighbor_latent) + transformer.network.initialize_neighbor_kv_caches(neighbor_context) + else: + assert neighbor_w2cs is None and neighbor_Ks is None, ( + "neighbor_w2cs / neighbor_Ks must be None when " + "neighbor_latent is None" + ) + transformer.network.initialize_neighbor_kv_caches(None) + + # 3. PRoPE coefficient setup. ``update_coeffs`` builds the 2D + # cos/sin tables for the patches_x x patches_y latent grid; + # ``_precompute_and_cache_apply_fns`` on the target side bakes + # the static neighbor cameras into the apply fn callables. + kt, kh, kw = transformer.config.network.patch_size + patches_x = width // kw + patches_y = height // kh + device = transformer.device + transformer.prope_cross_attn_src.update_coeffs( + patches_x=patches_x, patches_y=patches_y, device=device + ) + transformer.prope_cross_attn_tgt.update_coeffs( + patches_x=patches_x, patches_y=patches_y, device=device + ) + if neighbor_w2cs is not None: + transformer.prope_cross_attn_tgt._precompute_and_cache_apply_fns( + neighbor_w2cs, neighbor_Ks + ) + + return ArtifixerInferencePipelineCache( + transformer_cache=base_cache.transformer_cache, + encoder_cache=base_cache.encoder_cache, + decoder_cache=base_cache.decoder_cache, + image=base_cache.image, + condition_latent=condition_latent, + opacity=opacity, + camera_rays=camera_rays, + w2cs=w2cs, + Ks=Ks, + neighbor_w2cs=neighbor_w2cs, + neighbor_Ks=neighbor_Ks, + ) + + def _initialize_cache_from_text_embeddings( + self, + *, + text_embeddings: Tensor, + height: int, + width: int, + ) -> WanInferencePipelineCache: + """Build the base cache when prompts are pre-encoded. + + Mirrors :meth:`WanInferencePipeline.initialize_cache` minus the UMT5 + forward / negative-prompt encoding / I2V-image plumbing. CFG is + disabled per the ArtiFixer reference's kv-cache pipeline contract + (``negative_prompt`` is ignored), so we don't build + ``negative_text_embeddings`` even if ``guidance_scale > 1``. + """ + parent_cache = super( + WanInferencePipeline, self + ).initialize_cache( + transformer_context={ + "height": height, + "width": width, + "text_embeddings": text_embeddings, + "negative_text_embeddings": None, + "image_embeddings": None, + }, + ) + return WanInferencePipelineCache( + transformer_cache=parent_cache.transformer_cache, + encoder_cache=parent_cache.encoder_cache, + decoder_cache=parent_cache.decoder_cache, + image=None, + ) + + def _encode_neighbor_context(self, neighbor_latent: Tensor) -> Tensor: + """Run ``patch_embedding`` over neighbor latents and flatten to tokens. + + Mirrors the ArtiFixer reference's neighbor projection:: + + neighbor_hidden_states = self.patch_embedding(neighbor_hidden_states) + neighbor_hidden_states = neighbor_hidden_states.flatten(2).transpose(1, 2) + + Returns ``[*batch_shape, L_neighbor, dim]``. + """ + transformer = self.diffusion_model.transformer + assert isinstance(transformer, ArtifixerWanTransformer) + network = transformer.network + patched = network.patch_embedding(neighbor_latent) + return patched.flatten(2).transpose(1, 2) + + def _chunk_frame_ranges( + self, autoregressive_index: int, len_t: int, vae_t: int + ) -> tuple[int, int, int, int]: + """Return ``(lat_start, lat_end, input_start, input_end)`` for this AR chunk. + + Mirrors the per-chunk bookkeeping of the ArtiFixer reference kv-cache + pipeline: + + - first chunk covers ``1 + vae_t * (len_t - 1)`` input frames + (the Wan VAE 1+4 layout has one latent frame at the boundary + covering one input frame); + - subsequent chunks cover ``vae_t * len_t`` input frames each. + """ + lat_start = autoregressive_index * len_t + lat_end = lat_start + len_t + first_chunk_inputs = 1 + vae_t * (len_t - 1) + if autoregressive_index == 0: + input_start = 0 + input_end = first_chunk_inputs + else: + input_start = first_chunk_inputs + (autoregressive_index - 1) * vae_t * len_t + input_end = input_start + vae_t * len_t + return lat_start, lat_end, input_start, input_end + + def _build_ctrl( + self, + *, + chunk_opacity: Tensor, + chunk_camera_rays: Tensor, + chunk_post_patch_t: int, + frame_offset: int, + vae_t: int, + vae_s: int, + patch_size: tuple[int, int, int], + ) -> ArtifixerCtrl: + """Patchify the per-chunk opacity / camera_rays for ``ArtifixerCtrl``.""" + return ArtifixerCtrl( + opacity_extra=patchify_opacity( + chunk_opacity, + vae_scale_factor_temporal=vae_t, + vae_scale_factor_spatial=vae_s, + patch_size=patch_size, + frame_offset=frame_offset, + ), + camera_extra=patchify_camera_rays( + chunk_camera_rays, + hidden_post_patch_t=chunk_post_patch_t, + vae_scale_factor_temporal=vae_t, + vae_scale_factor_spatial=vae_s, + patch_size=patch_size, + frame_offset=frame_offset, + ), + ignore_neighbors=False, + ) + + @torch.no_grad() + def generate( # type: ignore[override] + self, + autoregressive_index: int, + cache: ArtifixerInferencePipelineCache, + ) -> Tensor: + """Generate one decoded video chunk for AR step ``autoregressive_index``. + + Custom denoise loop (bypasses ``DiffusionModel.generate``) so we can + renoise each step toward a fresh ``opacity_weighted_latent_mix`` of + the condition + new noise (matches the ArtiFixer reference's + kv-cache rollout). + + Steps: + 1. Slice the chunk's frame range out of the full-rollout cache. + 2. Update the source-side PRoPE ``apply_fns`` with the chunk's + target cameras. + 3. Build the ``ArtifixerCtrl`` payload (patchified opacity + + camera-ray features). + 4. Run the 4-step DMD denoise loop with prepare_latents renoise. + 5. Stash a ``FinalState`` on the cache so the standard + ``WanInferencePipeline.finalize`` path closes the AR cache. + 6. Decode and return the chunk. + """ + assert cache.condition_latent is not None, "initialize_cache must be called first" + assert cache.opacity is not None + assert cache.camera_rays is not None + assert cache.w2cs is not None + assert cache.Ks is not None + + transformer = self.diffusion_model.transformer + assert isinstance(transformer, ArtifixerWanTransformer) + scheduler = self.diffusion_model.scheduler + assert isinstance(scheduler, FlowMatchScheduler), ( + f"ArtiFixer pipeline requires a FlowMatchScheduler, got " + f"{type(scheduler).__name__}" + ) + + tcfg = transformer.config + len_t = tcfg.len_t + kt, kh, kw = tcfg.network.patch_size + # Wan VAE scale factors. The decoder owns the public name. + vae_t = self.decoder.temporal_compression_ratio if self.decoder is not None else 4 + vae_s = self.decoder.spatial_compression_ratio if self.decoder is not None else 8 + + lat_start, lat_end, input_start, input_end = self._chunk_frame_ranges( + autoregressive_index, len_t, vae_t + ) + chunk_condition = cache.condition_latent[..., lat_start:lat_end, :, :] + chunk_opacity = cache.opacity[..., input_start:input_end, :, :] + chunk_w2cs = cache.w2cs[..., lat_start:lat_end, :, :] + chunk_Ks = cache.Ks[..., lat_start:lat_end, :, :] + # camera_rays may be at the latent rate or input rate; slice the + # axis we're indexing on with the latent rate first, then fall back. + if cache.camera_rays.shape[1] == cache.condition_latent.shape[2]: + chunk_camera_rays = cache.camera_rays[:, lat_start:lat_end] + else: + chunk_camera_rays = cache.camera_rays[:, input_start:input_end] + + # Update the per-AR-chunk PRoPE source cameras (neighbor side was + # set once at initialize_cache; it is static across AR steps). + transformer.prope_cross_attn_src._precompute_and_cache_apply_fns( + chunk_w2cs, chunk_Ks + ) + + # Build the ArtifixerCtrl payload (patchified opacity / camera). + chunk_post_patch_t = (lat_end - lat_start) // kt + ctrl = self._build_ctrl( + chunk_opacity=chunk_opacity, + chunk_camera_rays=chunk_camera_rays, + chunk_post_patch_t=chunk_post_patch_t, + frame_offset=lat_start, + vae_t=vae_t, + vae_s=vae_s, + patch_size=tcfg.network.patch_size, + ) + + # Start the AR cache for this chunk (mirrors DiffusionModel.generate). + cache.transformer_cache.start(autoregressive_index) + + # Initial mix: condition * opacity_lat + noise * (1 - opacity_lat). + # We work in unpatchified latent space, then patchify before + # feeding into ``predict_flow``. + initial_noise = torch.randn( + chunk_condition.shape, + device=self.device, + dtype=self.diffusion_model.dtype, + generator=self.diffusion_model.rng, + ) + is_first = autoregressive_index == 0 + latent_unpatched = opacity_weighted_latent_mix( + condition=chunk_condition, + opacity=chunk_opacity, + noise=initial_noise, + vae_scale_factor_temporal=vae_t, + vae_scale_factor_spatial=vae_s, + is_first_chunk=is_first, + ) + # ``patchify_and_maybe_split_cp`` expects ``(B, T, C, H, W)`` (einops + # pattern ``... (t kt) c (h kh) (w kw)``), but a diffusers-convention + # VAE encoder emits ``(B, C, T, H, W)`` and we keep that convention + # for ``cache.condition_latent`` / ``latent_unpatched``. Permute + # before patchify. + latent = transformer.patchify_and_maybe_split_cp( + latent_unpatched.permute(0, 2, 1, 3, 4) + ) + + # 4-step DMD denoise with prepare_latents renoise. Matches the + # ArtiFixer reference: at every non-exit step we replace the + # next-iteration ``noisy`` with a fresh + # ``opacity_weighted_latent_mix(...)`` rather than the base + # FlowMatchScheduler.sample's ``clean``. + sigmas = scheduler.denoising_sigmas + timesteps = scheduler.denoising_step_list + n_steps = timesteps.shape[0] + input_dtype = latent.dtype + clean: Tensor | None = None + + for i in range(n_steps): + sigma = sigmas[i] + timestep = timesteps[i].to(dtype=input_dtype) + flow = transformer.predict_flow( + noisy_latent=latent, + timestep=timestep, + cache=cache.transformer_cache, + input=ctrl, + ) + clean = latent - sigma * flow + if i + 1 < n_steps: + sigma_next = sigmas[i + 1] + fresh_noise = torch.randn( + chunk_condition.shape, + device=self.device, + dtype=self.diffusion_model.dtype, + generator=self.diffusion_model.rng, + ) + fresh_mix_unpatched = opacity_weighted_latent_mix( + condition=chunk_condition, + opacity=chunk_opacity, + noise=fresh_noise, + vae_scale_factor_temporal=vae_t, + vae_scale_factor_spatial=vae_s, + is_first_chunk=is_first, + ) + # Same permute as the initial mix (see above): (B, C, T, H, W) + # -> (B, T, C, H, W) before patchify. + fresh_mix = transformer.patchify_and_maybe_split_cp( + fresh_mix_unpatched.permute(0, 2, 1, 3, 4) + ) + latent = ((1.0 - sigma_next) * clean + sigma_next * fresh_mix).to( + input_dtype + ) + + assert clean is not None, "denoise loop produced no clean estimate" + clean = transformer.postprocess_clean_latent( + clean_latent=clean, + cache=cache.transformer_cache, + input=None, # I2V stamp not used by ArtiFixer + ) + + # Stash a FinalState on the cache so the inherited ``finalize`` path + # (WanInferencePipeline.finalize -> DiffusionModel.finalize) closes + # the AR cache. + final_state = DiffusionModel.FinalState( + clean_latent=clean, + autoregressive_index=autoregressive_index, + cache=cache.transformer_cache, + input=ctrl, + ) + cache.final_state = final_state + cache.autoregressive_index = autoregressive_index + + clean_unpatched = transformer.unpatchify_and_maybe_gather_cp(clean) + + if self.decoder is None: + return clean_unpatched + return self.decoder( + clean_unpatched, autoregressive_index, cache.decoder_cache + ) + + +__all__ = [ + "ArtifixerInferencePipeline", + "ArtifixerInferencePipelineCache", + "ArtifixerInferencePipelineConfig", +] diff --git a/integrations/artifixer/artifixer/runner.py b/integrations/artifixer/artifixer/runner.py new file mode 100644 index 000000000..90dcf29e5 --- /dev/null +++ b/integrations/artifixer/artifixer/runner.py @@ -0,0 +1,144 @@ +# 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. + +"""ArtiFixer DMD-distilled streaming T2V runner. + +Text-only runner: a stripped clone of ``SelfForcingT2VRunner`` that +accepts a single text prompt. Used for smoke-testing the recipe wiring +and the ``flashdreams-run`` CLI entry point. + +The full ArtiFixer-specific conditioning surface (``rgb_rendered``, +``opacity``, neighbor frames, camera matrices) is wired in through the +:class:`ArtifixerInferencePipeline.initialize_cache` / ``generate`` +path; production drivers call those directly with conditioning tensors. +This runner stays text-only as a lightweight harness. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from pathlib import Path + +import mediapy as media +import torch +from einops import rearrange +from loguru import logger + +from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.recipes.wan import ( + WanInferencePipeline, + WanInferencePipelineCache, +) + +__all__ = [ + "ArtifixerDmdT2VRunnerConfig", + "ArtifixerDmdT2VRunner", +] + + +@dataclass(kw_only=True) +class ArtifixerDmdT2VRunnerConfig(RunnerConfig): + """Runner config for the ArtiFixer DMD T2V variants.""" + + _target: type = field(default_factory=lambda: ArtifixerDmdT2VRunner) + + prompt: str | Path = Path(__file__).resolve().parents[1] / "assets" / "prompt.txt" + """Either an inline prompt (``--prompt "..."``) or a path to a ``.txt`` + whose first non-empty line is read as the prompt.""" + + total_blocks: int = 3 + """Number of autoregressive chunks to generate. Matches the ArtiFixer + reference (``num_frames=81`` / ``frames_per_block=7``, i.e. 21 latent + frames / 7 = 3 chunks at the Wan VAE temporal stride of 4).""" + + pixel_height: int = 480 + pixel_width: int = 832 + fps: int = 16 + + +class ArtifixerDmdT2VRunner(Runner[ArtifixerDmdT2VRunnerConfig, WanInferencePipeline]): + """ArtiFixer DMD streaming T2V driver (text-only smoke runner).""" + + config: ArtifixerDmdT2VRunnerConfig + + def _resolve_prompt(self) -> str: + value = self.config.prompt + 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 _initialize_cache(self) -> WanInferencePipelineCache: + config = self.config + prompt = self._resolve_prompt() + + assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) + spatial_compression_ratio = self.pipeline.decoder.spatial_compression_ratio + assert config.pixel_height % spatial_compression_ratio == 0, ( + f"pixel_height={self.config.pixel_height} must divide " + f"{spatial_compression_ratio}." + ) + assert config.pixel_width % spatial_compression_ratio == 0, ( + f"pixel_width={self.config.pixel_width} must divide " + f"{spatial_compression_ratio}." + ) + latent_h = config.pixel_height // spatial_compression_ratio + latent_w = config.pixel_width // spatial_compression_ratio + + return self.pipeline.initialize_cache( + text=[prompt], image=None, height=latent_h, width=latent_w + ) + + def run(self) -> None: + config = self.config + + cache = self._initialize_cache() + + chunks: list[torch.Tensor] = [] + stats_history: list[dict[str, float]] = [] + for i in range(config.total_blocks): + video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache) + stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) + if stats is not None: + stats_history.append({"autoregressive_index": i, **stats}) + chunks.append(video_chunk.cpu()) + + generated = torch.cat(chunks, dim=0) + if not self.is_rank_zero: + return + + config.output_dir.mkdir(parents=True, exist_ok=True) + video_path = config.output_dir / f"{config.runner_name}.mp4" + canvas = rearrange(generated, "t c h w -> t h w c") + + arr = (canvas.float().numpy() + 1.0) / 2.0 + arr = (arr * 255).clip(0, 255).astype("uint8") + media.write_video(str(video_path), arr, fps=config.fps) + + logger.info( + f"[{config.runner_name}] wrote video {tuple(generated.shape)} " + f"-> {video_path.resolve()}" + ) + + if stats_history: + stats_path = config.output_dir / f"stats_{config.runner_name}.json" + stats_path.write_text(json.dumps(stats_history, indent=2)) + logger.info( + f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" + ) diff --git a/integrations/artifixer/artifixer/transformer.py b/integrations/artifixer/artifixer/transformer.py new file mode 100644 index 000000000..b8f87572c --- /dev/null +++ b/integrations/artifixer/artifixer/transformer.py @@ -0,0 +1,201 @@ +# 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. + +"""ArtiFixer transformer wrapper. + +Subclasses :class:`Wan21Transformer` so the recipe can: + + - own the two ``PropeDotProductAttention`` modules (target/source and + neighbor cameras) used by every transformer block's cross-attention + branch; + - hold per-rollout conditioning state (full-length opacity, Plucker rays, + camera matrices, optional neighbor cameras) populated once by the + pipeline at ``initialize_autoregressive_cache`` time; + - thread per-AR-chunk extras (patchified opacity / camera-ray features, + target-camera PRoPE precompute) into ``predict_flow`` -> + ``WanDiTNetwork.forward`` via an :class:`ArtifixerCtrl` payload. + +The pipeline owns the per-chunk slicing + patchification + PRoPE +precompute updates and packages them into an ``ArtifixerCtrl`` that +flows through ``DiffusionModel.generate(input=ctrl)``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from artifixer.network.prope import PropeDotProductAttention +from torch import Tensor + +from flashdreams.recipes.wan import Wan21Transformer, Wan21TransformerConfig +from flashdreams.recipes.wan.autoencoder.i2v import I2VCtrl +from flashdreams.recipes.wan.transformer.wan21 import Wan21TransformerCache + + +@dataclass(kw_only=True) +class ArtifixerCtrl: + """Per-AR-chunk conditioning payload for the ArtiFixer transformer. + + All tensors are already sliced to the current chunk's frame range; the + patchifiable ones are also already passed through + :func:`artifixer.network.patches.patchify_opacity` / + :func:`patchify_camera_rays`. The pipeline owns this preparation; the + transformer's ``predict_flow`` reads from here and forwards into + ``network_extra_kwargs``. + + The cross-attention "neighbor" KV bank is *not* carried here -- it is + set once per rollout on every block's cross-attention module via + ``ArtifixerCrossAttention.initialize_neighbor_cache`` and stays there + across AR chunks. + """ + + opacity_extra: Tensor + """Patchified per-token opacity features + ``[*batch_shape, L, opacity_embedding_dim]``.""" + + camera_extra: Tensor + """Patchified per-token Plucker camera-ray features + ``[*batch_shape, L, camera_embedding_dim]``.""" + + ignore_neighbors: bool = False + """Diffusion-forcing CFG-dropout knob. Inference leaves this ``False``; + training-time DMD distillation may toggle it.""" + + _is_patchified: bool = True + """Sentinel for ``Wan21Transformer.patchify_and_maybe_split_cp``: + ``opacity_extra`` and ``camera_extra`` arrive already patchified, so the + pipeline-level ``patchify_and_maybe_split_cp(input)`` dispatch in + ``DiffusionModel.generate`` is a no-op for us.""" + + +@dataclass(kw_only=True) +class ArtifixerWanTransformerConfig(Wan21TransformerConfig): + """Wan2.1 transformer config that constructs an :class:`ArtifixerWanTransformer`.""" + + _target: type = field(default_factory=lambda: ArtifixerWanTransformer) + + prope_freq_base: float = 100.0 + """RoPE frequency base for the PRoPE cross-attention transform.""" + + prope_freq_scale: float = 1.0 + """RoPE frequency scale for the PRoPE cross-attention transform.""" + + +class ArtifixerWanTransformer(Wan21Transformer): + """Wan 2.1 transformer with ArtiFixer's PRoPE cross-attention modules.""" + + config: ArtifixerWanTransformerConfig + + def __init__(self, config: ArtifixerWanTransformerConfig) -> None: + super().__init__(config) + + head_dim = config.network.dim // config.network.num_heads + self.prope_cross_attn_src = PropeDotProductAttention( + head_dim=head_dim, + freq_base=config.prope_freq_base, + freq_scale=config.prope_freq_scale, + ) + self.prope_cross_attn_tgt = PropeDotProductAttention( + head_dim=head_dim, + freq_base=config.prope_freq_base, + freq_scale=config.prope_freq_scale, + ) + + # Match the configured transformer dtype. + self.prope_cross_attn_src = self.prope_cross_attn_src.to(dtype=config.dtype) + self.prope_cross_attn_tgt = self.prope_cross_attn_tgt.to(dtype=config.dtype) + + def predict_flow( # type: ignore[override] + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: Wan21TransformerCache, + input: ArtifixerCtrl | I2VCtrl | None = None, + network_extra_kwargs: dict[str, Any] | None = None, + ) -> Tensor: + """Add ArtiFixer block_extra_kwargs to the base predict_flow. + + Wraps the ArtiFixer per-block extras inside a ``block_extra_kwargs`` + dict so :meth:`Wan21Transformer._predict_flow`'s + ``**network_extra_kwargs`` unpack lands them as + ``WanDiTNetwork.forward(... block_extra_kwargs=...)`` -- a single + kwarg, *not* individual ones. ``WanDiTNetwork.forward`` then forwards + each entry to every block as keyword args via ``**block_extra_kwargs``, + where :class:`ArtifixerBlock.forward` accepts ``opacity_extra`` / + ``camera_extra`` / ``prope_src`` / ``prope_tgt`` / + ``ignore_neighbors`` directly. + + The PRoPE source/target modules are bound at construction time, but + :class:`ArtifixerCrossAttention.forward` expects them as forward + kwargs so every per-block call gets them via the same path. + + When ``input`` is not an :class:`ArtifixerCtrl` (e.g. text-only smoke + without any conditioning), falls through to the base behavior; + :class:`ArtifixerBlock` treats every conditioning kwarg as optional. + """ + network_extra_kwargs = dict(network_extra_kwargs or {}) + if isinstance(input, ArtifixerCtrl): + block_extra_kwargs = dict(network_extra_kwargs.get("block_extra_kwargs", {})) + block_extra_kwargs.setdefault("opacity_extra", input.opacity_extra) + block_extra_kwargs.setdefault("camera_extra", input.camera_extra) + block_extra_kwargs.setdefault("ignore_neighbors", input.ignore_neighbors) + block_extra_kwargs.setdefault("prope_src", self.prope_cross_attn_src) + block_extra_kwargs.setdefault("prope_tgt", self.prope_cross_attn_tgt) + network_extra_kwargs["block_extra_kwargs"] = block_extra_kwargs + # The base ``predict_flow`` will dispatch ``_build_network_input`` + # on ``input`` -- ArtifixerCtrl isn't an I2VCtrl, so pass None + # for the I2V mask-inject leg. + i2v_input: I2VCtrl | None = None + else: + i2v_input = input + + return super().predict_flow( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + input=i2v_input, + network_extra_kwargs=network_extra_kwargs, + ) + + def finalize_kv_cache(self, *args: Any, **kwargs: Any) -> None: + """No-op: the ArtiFixer reference does not finalize the KV cache. + + The ArtiFixer reference kv-cache pipeline advances its KV cache + *in-place* during the regular denoise forwards and never runs an + extra "finalization" forward at AR-chunk boundaries. The + flashdreams default ``finalize_kv_cache`` runs one more + ``predict_flow`` at the context-noise timestep to advance the + cache; for ArtiFixer that extra forward writes a *different* KV + state than the reference's last in-loop forward, which empirically + opens a ~7-9 dB cross-backend PSNR gap at the start of every AR + chunk past chunk 0. + + Skipping the extra forward here aligns this pipeline's KV-cache + semantics with the reference: the cache going into AR chunk + ``N+1`` is the one written by the final denoise step of AR chunk + ``N``. ``cache.finalize`` is still called by + ``DiffusionModel.finalize`` after this (the bookkeeping that + increments ``autoregressive_index``), so the rollout state + advances correctly -- only the redundant predict is suppressed. + """ + del args, kwargs # intentionally no-op; see docstring + + +__all__ = [ + "ArtifixerCtrl", + "ArtifixerWanTransformer", + "ArtifixerWanTransformerConfig", +] diff --git a/integrations/artifixer/assets/prompt.txt b/integrations/artifixer/assets/prompt.txt new file mode 100644 index 000000000..198341c24 --- /dev/null +++ b/integrations/artifixer/assets/prompt.txt @@ -0,0 +1 @@ +A photorealistic dolly-in shot of a modern living room with warm afternoon sunlight streaming through floor-to-ceiling windows, a leather sofa in the foreground, hardwood floors, abstract artwork on the walls, and a bookshelf to the side. diff --git a/integrations/artifixer/pyproject.toml b/integrations/artifixer/pyproject.toml new file mode 100644 index 000000000..457bd1a42 --- /dev/null +++ b/integrations/artifixer/pyproject.toml @@ -0,0 +1,48 @@ +# 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. + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-artifixer" +version = "0.1.0" +description = "ArtiFixer reconstruction-enhanced T2V inference for Wan 2.1 1.3B, packaged as a flashdreams plugin." +readme = "README.md" +requires-python = ">=3.12" +dependencies = [ + "flashdreams", + "mediapy>=1.1", +] + +[tool.uv.sources] +flashdreams = { workspace = true } + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", +] + +# Each entry registers one ``runner_name`` slug with ``flashdreams-run``. +# The discovery layer (``flashdreams.plugins.registry.discover_runners``) +# scans this group at CLI startup; the entry-point name itself is purely +# informational, the registry key always comes from ``cfg.runner_name``. +[project.entry-points."flashdreams.runner_configs"] +"artifixer-dmd-wan2.1-t2v-1.3b" = "artifixer.config:RUNNER_ARTIFIXER_DMD_T2V_1PT3B" + +[tool.setuptools.packages.find] +include = ["artifixer*"] +exclude = ["tests"] diff --git a/integrations/artifixer/tests/test_latent_mix.py b/integrations/artifixer/tests/test_latent_mix.py new file mode 100644 index 000000000..1582fb10d --- /dev/null +++ b/integrations/artifixer/tests/test_latent_mix.py @@ -0,0 +1,129 @@ +# 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. + +"""Shape + invariance tests for the opacity-weighted latent mix. + +CPU-only; only ``torch`` is required. +""" + +from __future__ import annotations + +import torch +from artifixer.latent_mix import opacity_weighted_latent_mix + +VAE_T = 4 +VAE_S = 8 + + +def _condition_and_noise(batch=1, in_dim=16, t_lat=2, hl=4, wl=4): + """Build matching condition / noise tensors at the latent grid.""" + cond = torch.randn(batch, in_dim, t_lat, hl, wl) + noise = torch.randn(batch, in_dim, t_lat, hl, wl) + return cond, noise + + +def test_first_chunk_opacity_one_returns_condition_exactly() -> None: + """opacity = 1 everywhere -> latent_mix = condition (noise is dropped).""" + cond, noise = _condition_and_noise(t_lat=2, hl=4, wl=4) + # First chunk uses ``1 + 4 * (t_lat - 1)`` input frames after the 1+3 pad + # collapses to t_lat = 2 latent frames. + t_input_first = 1 + VAE_T * (cond.shape[2] - 1) + opacity = torch.ones(1, t_input_first, VAE_S * 4, VAE_S * 4) + + out = opacity_weighted_latent_mix( + condition=cond, + opacity=opacity, + noise=noise, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + is_first_chunk=True, + ) + torch.testing.assert_close(out, cond) + + +def test_first_chunk_opacity_zero_returns_noise_exactly() -> None: + """opacity = 0 everywhere -> latent_mix = noise (condition is dropped).""" + cond, noise = _condition_and_noise(t_lat=2, hl=4, wl=4) + t_input_first = 1 + VAE_T * (cond.shape[2] - 1) + opacity = torch.zeros(1, t_input_first, VAE_S * 4, VAE_S * 4) + + out = opacity_weighted_latent_mix( + condition=cond, + opacity=opacity, + noise=noise, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + is_first_chunk=True, + ) + torch.testing.assert_close(out, noise) + + +def test_later_chunk_shape_no_pad() -> None: + """Non-first chunk has ``vae_t * t_lat`` input frames already.""" + cond, noise = _condition_and_noise(batch=2, t_lat=2, hl=4, wl=4) + t_input = VAE_T * cond.shape[2] + opacity = torch.full((2, t_input, VAE_S * 4, VAE_S * 4), 0.5) + + out = opacity_weighted_latent_mix( + condition=cond, + opacity=opacity, + noise=noise, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + is_first_chunk=False, + ) + # opacity_lat = 0.5 everywhere -> 0.5 * cond + 0.5 * noise. + torch.testing.assert_close(out, 0.5 * cond + 0.5 * noise) + + +def test_first_chunk_handles_single_latent_frame() -> None: + """Edge case: t_lat = 1, which requires the 1+3 left-pad for vae_t=4.""" + cond, noise = _condition_and_noise(t_lat=1, hl=4, wl=4) + opacity = torch.full((1, 1, VAE_S * 4, VAE_S * 4), 0.7) + + out = opacity_weighted_latent_mix( + condition=cond, + opacity=opacity, + noise=noise, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + is_first_chunk=True, + ) + # After 1+3 pad: opacity is shape (1, 4, ...) all 0.7 (pad uses the + # repeated first frame which is 0.7). max_pool3d -> (1, 1, 1, hl, wl) + # all 0.7. Mix = 0.7 * cond + 0.3 * noise. + torch.testing.assert_close(out, 0.7 * cond + 0.3 * noise) + + +def test_max_pool_picks_largest_in_window() -> None: + """max_pool3d uses the MAX, not the average -- regression check.""" + cond, noise = _condition_and_noise(t_lat=1, hl=1, wl=1) + # 1 + 4 * 0 = 1 input frame. Pad to 4. Make first 3 pad-frames 0.0 and + # the last (original) frame 1.0 by overriding the source frame. + opacity = torch.zeros(1, 1, VAE_S, VAE_S) + opacity[..., 0, 0] = 1.0 # one pixel == 1 in the 8x8 input window + + out = opacity_weighted_latent_mix( + condition=cond, + opacity=opacity, + noise=noise, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + is_first_chunk=True, + ) + # The padded stack is 4 copies of the same 8x8 input. max_pool3d takes + # the max -- a single 1.0 pixel survives, so opacity_lat = 1.0 at the + # single latent pixel, mix = cond. + torch.testing.assert_close(out, cond) diff --git a/integrations/artifixer/tests/test_patches.py b/integrations/artifixer/tests/test_patches.py new file mode 100644 index 000000000..44dc86344 --- /dev/null +++ b/integrations/artifixer/tests/test_patches.py @@ -0,0 +1,129 @@ +# 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. + +"""Shape + invariance tests for the opacity / camera-ray patchification. + +CPU-only, torch + einops are the only deps. +""" + +from __future__ import annotations + +import torch +from artifixer.network.dit import artifixer_embedding_dims +from artifixer.network.patches import patchify_camera_rays, patchify_opacity + +VAE_T = 4 +VAE_S = 8 +PATCH = (1, 2, 2) + + +def _expected_L(t_latent: int, h: int, w: int, patch: tuple[int, int, int]) -> int: + kt, kh, kw = patch + assert t_latent % kt == 0 and h % (VAE_S * kh) == 0 and w % (VAE_S * kw) == 0 + return (t_latent // kt) * (h // (VAE_S * kh)) * (w // (VAE_S * kw)) + + +def test_patchify_opacity_first_chunk_shape() -> None: + """First AR chunk pads the first frame; output token count matches the + expected post-patch latent token count. + """ + # 1 latent frame for the first chunk: 1 + 3 padding -> 4 input frames + # covers 1 latent frame (VAE_T = 4). Spatial: 32x32 input -> 2x2 latent + # tokens after VAE (32 / 8 = 4) and patch (4 / 2 = 2). + opacity = torch.randn(2, 1, 32, 32) # (B, T, H, W) at INPUT resolution + out = patchify_opacity( + opacity, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + patch_size=PATCH, + frame_offset=0, + ) + opacity_dim, _ = artifixer_embedding_dims(PATCH) + assert out.shape == (2, _expected_L(1, 32, 32, PATCH), opacity_dim) + + +def test_patchify_opacity_later_chunk_no_pad() -> None: + """Non-zero frame_offset does not pad: input has already-paired frames.""" + # 4 input frames -> 1 latent frame (VAE_T = 4) + opacity = torch.randn(2, 4, 32, 32) + out = patchify_opacity( + opacity, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + patch_size=PATCH, + frame_offset=7, # any non-zero value + ) + opacity_dim, _ = artifixer_embedding_dims(PATCH) + assert out.shape == (2, _expected_L(1, 32, 32, PATCH), opacity_dim) + + +def test_patchify_camera_rays_latent_rate_shape() -> None: + """Plucker rays already at the latent temporal rate skip the pad/regroup.""" + # camera_rays.shape[1] == hidden_post_patch_t (1 latent frame) + rays = torch.randn(2, 1, 32, 32, 6) + out = patchify_camera_rays( + rays, + hidden_post_patch_t=1, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + patch_size=PATCH, + frame_offset=0, + ) + _, camera_dim = artifixer_embedding_dims(PATCH) + assert out.shape == (2, _expected_L(1, 32, 32, PATCH), camera_dim) + + +def test_patchify_camera_rays_input_rate_first_chunk_shape() -> None: + """Input-rate Plucker rays on the first AR chunk get the 1+3 left-pad. + + This branch in the ArtiFixer reference groups t4 input frames + per latent frame, so the per-token feature carries an extra ``vae_t`` + multiplier vs the latent-rate branch: + + per_token_dim_input_rate = vae_t * camera_dim_latent_rate + + At inference, camera_rays is always already at the latent rate (the + consuming driver pre-aligns it), so this branch is unused. We still + cover it for parity with the reference forward. + """ + rays = torch.randn(2, 1, 32, 32, 6) # 1 input frame + out = patchify_camera_rays( + rays, + hidden_post_patch_t=999, # mismatch -> not latent-rate + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + patch_size=PATCH, + frame_offset=0, + ) + _, camera_dim = artifixer_embedding_dims(PATCH) + expected_dim = VAE_T * camera_dim + assert out.shape == (2, _expected_L(1, 32, 32, PATCH), expected_dim) + + +def test_patchify_opacity_constant_input_constant_output() -> None: + """A constant opacity field produces a constant per-token feature. + + Invariant: any rearrangement of the same scalar must still be scalar + along the token axis. + """ + opacity = torch.full((1, 4, 16, 16), 0.42) + out = patchify_opacity( + opacity, + vae_scale_factor_temporal=VAE_T, + vae_scale_factor_spatial=VAE_S, + patch_size=PATCH, + frame_offset=7, + ) + assert torch.allclose(out, torch.full_like(out, 0.42)) diff --git a/integrations/artifixer/tests/test_prope_parity.py b/integrations/artifixer/tests/test_prope_parity.py new file mode 100644 index 000000000..705409e76 --- /dev/null +++ b/integrations/artifixer/tests/test_prope_parity.py @@ -0,0 +1,204 @@ +# 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. + +"""Numerical parity test for the ported PRoPE attention module. + +Compares :class:`artifixer.network.prope.PropeDotProductAttention` against +the ArtiFixer reference's ``model_training/net/prope.py``. When the +reference is not importable (CI images typically only ship flashdreams), +the test is split: + + * ``test_prope_internal_consistency`` runs against an inline copy of the + upstream math — i.e. tests that ``apply_to_o(apply_to_q(q) ...)`` round + trips correctly for the identity camera case. Always runs. + + * ``test_prope_matches_reference`` is skipped unless + ``model_training.net.prope`` from the ArtiFixer reference is on + ``sys.path``. Set ``ARTIFIXER_REFERENCE_REPO_ROOT`` to its checkout + root to enable. + +The intent is to catch any regression in the verbatim port the moment a +GPU + reference env is available, while still exercising the math today. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch + +from artifixer.network.prope import PropeDotProductAttention + + +@pytest.fixture(autouse=True) +def _deterministic_torch() -> None: + """Seed PyTorch for the file. + + Note: PRoPE's ``_rope_precompute_coeffs`` performs an implicit int64->fp32 + promotion (``torch.arange(...) / num_freqs``), so the RoPE coefficients + are always at fp32 precision regardless of the input dtype. The + reference has the same property. Tests therefore run at fp32 with + fp32-appropriate tolerances rather than fp64. The reference's + ``_lift_K`` also hard-codes float32 output (no ``dtype=`` in + ``torch.zeros``), so any cross-implementation comparison must use + fp32 inputs. + """ + torch.manual_seed(0) + + +def _build_inputs( + batch: int = 1, + cameras: int = 3, + patches_x: int = 4, + patches_y: int = 4, + num_heads: int = 4, + head_dim: int = 32, + dtype: torch.dtype = torch.float32, + device: torch.device | str = "cpu", +) -> dict[str, torch.Tensor]: + """Build a deterministic fixture for PRoPE parity tests.""" + seqlen = cameras * patches_x * patches_y + + # Random SE(3) viewmats: rotation = QR(randn), translation = randn. + R = torch.linalg.qr(torch.randn(batch, cameras, 3, 3, dtype=dtype, device=device)).Q + t = torch.randn(batch, cameras, 3, dtype=dtype, device=device) * 0.5 + viewmats = torch.zeros(batch, cameras, 4, 4, dtype=dtype, device=device) + viewmats[..., :3, :3] = R + viewmats[..., :3, 3] = t + viewmats[..., 3, 3] = 1.0 + + # Intrinsics in NDC-ish form: fx, fy ~ 1.0; cx, cy ~ 0.5. + Ks = torch.eye(3, dtype=dtype, device=device).expand(batch, cameras, 3, 3).clone() + Ks[..., 0, 0] = 1.0 + 0.1 * torch.randn(batch, cameras, dtype=dtype, device=device) + Ks[..., 1, 1] = 1.0 + 0.1 * torch.randn(batch, cameras, dtype=dtype, device=device) + Ks[..., 0, 2] = 0.5 + Ks[..., 1, 2] = 0.5 + + q = torch.randn(batch, num_heads, seqlen, head_dim, dtype=dtype, device=device) + k = torch.randn(batch, num_heads, seqlen, head_dim, dtype=dtype, device=device) + v = torch.randn(batch, num_heads, seqlen, head_dim, dtype=dtype, device=device) + + return { + "viewmats": viewmats, + "Ks": Ks, + "q": q, + "k": k, + "v": v, + "patches_x": patches_x, + "patches_y": patches_y, + "head_dim": head_dim, + } + + +def _instantiate(prope_cls, inp: dict) -> object: + """Build a PRoPE instance and precompute the apply-fns for ``inp``.""" + prope = prope_cls(head_dim=inp["head_dim"]) + prope.update_coeffs(inp["patches_x"], inp["patches_y"], device="cpu") + prope._precompute_and_cache_apply_fns(inp["viewmats"], inp["Ks"]) + return prope + + +def test_prope_internal_consistency() -> None: + """``apply_to_o`` undoes ``apply_to_q`` on identity cameras. + + With viewmats = identity and Ks = identity, ``P = P_T = P_inv = I`` + and the projection-matrix legs of the block-diagonal transform become + no-ops. The RoPE legs of ``apply_to_o`` are the *inverse* of those in + ``apply_to_q``, so ``apply_to_o(apply_to_q(x)) ~= x`` up to fp32 + round-off (RoPE coeffs are always at fp32 — see fixture docstring). + """ + inp = _build_inputs() + + inp["viewmats"] = ( + torch.eye(4, dtype=inp["viewmats"].dtype).expand_as(inp["viewmats"]).clone() + ) + inp["Ks"] = ( + torch.eye(3, dtype=inp["Ks"].dtype).expand_as(inp["Ks"]).clone() + ) + + prope = _instantiate(PropeDotProductAttention, inp) + q = inp["q"] + round_trip = prope._apply_to_o(prope._apply_to_q(q)) + torch.testing.assert_close(round_trip, q, atol=5e-6, rtol=5e-6) + + +def test_prope_apply_to_q_changes_input_for_nontrivial_cameras() -> None: + """Sanity: with real cameras, the PRoPE transform is not a no-op.""" + inp = _build_inputs() + prope = _instantiate(PropeDotProductAttention, inp) + transformed = prope._apply_to_q(inp["q"]) + diff = (transformed - inp["q"]).abs().max().item() + assert diff > 1e-3, f"expected non-trivial PRoPE transform, got max abs diff {diff}" + + +def _try_import_reference() -> object | None: + """Return the ArtiFixer reference's PropeDotProductAttention if importable.""" + repo_root = os.environ.get("ARTIFIXER_REFERENCE_REPO_ROOT") + if repo_root and Path(repo_root).is_dir() and repo_root not in sys.path: + sys.path.insert(0, repo_root) + try: + from model_training.net.prope import PropeDotProductAttention as RefPRoPE + except ImportError: + return None + return RefPRoPE + + +@pytest.mark.skipif( + _try_import_reference() is None, + reason=( + "ArtiFixer reference not importable. Set " + "ARTIFIXER_REFERENCE_REPO_ROOT to the reference checkout, or run " + "from an env that already has it on sys.path." + ), +) +def test_prope_matches_reference() -> None: + """Numerical parity vs the ArtiFixer reference at fp32. + + Asserts ``apply_to_q``, ``apply_to_kv``, ``apply_to_o`` produce + bit-identical outputs to the reference within fp32 round-off. + + Inputs are fp32 because the reference's ``_lift_K`` hard-codes float32 + output (no ``dtype=`` in ``torch.zeros``), so feeding fp64 viewmats + crashes its einsum. Our port passes ``dtype=Ks.dtype``, making it + slightly more dtype-robust; at fp32 the two are bit-identical. + """ + RefPRoPE = _try_import_reference() + assert RefPRoPE is not None + + inp = _build_inputs() + ours = _instantiate(PropeDotProductAttention, inp) + ref = _instantiate(RefPRoPE, inp) + + for name, x in (("q", inp["q"]), ("k", inp["k"]), ("v", inp["v"])): + ours_q = ours._apply_to_q(x) + ref_q = ref._apply_to_q(x) + torch.testing.assert_close( + ours_q, ref_q, atol=1e-6, rtol=1e-6, msg=f"apply_to_q on {name}" + ) + + ours_kv = ours._apply_to_kv(x) + ref_kv = ref._apply_to_kv(x) + torch.testing.assert_close( + ours_kv, ref_kv, atol=1e-6, rtol=1e-6, msg=f"apply_to_kv on {name}" + ) + + ours_o = ours._apply_to_o(x) + ref_o = ref._apply_to_o(x) + torch.testing.assert_close( + ours_o, ref_o, atol=1e-6, rtol=1e-6, msg=f"apply_to_o on {name}" + ) diff --git a/integrations/artifixer/tests/test_smoke.py b/integrations/artifixer/tests/test_smoke.py new file mode 100644 index 000000000..826a8305d --- /dev/null +++ b/integrations/artifixer/tests/test_smoke.py @@ -0,0 +1,272 @@ +# 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. + +"""Cheap import-time checks for the ``artifixer`` plugin. + +Heavy imports (``artifixer.config`` -> ``runner`` -> ``mediapy``; ``flashdreams``) +are kept *inside* the test functions that need them so that lighter unit +tests (e.g. ``test_compute_kv_neighbor_and_cache_init``) can run in +torch-only environments. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +ENTRY_POINT_GROUP = "flashdreams.runner_configs" + + +def _runner_configs() -> dict: + """Import RUNNER_CONFIGS lazily to avoid pulling in mediapy at collect time.""" + from artifixer.config import RUNNER_CONFIGS + + return RUNNER_CONFIGS + + +def test_runners_dict_is_non_empty() -> None: + assert _runner_configs(), "RUNNER_CONFIGS is empty" + + +def test_runner_name_mirrors_pipeline_recipe_name() -> None: + drifted = { + slug: (cfg.runner_name, cfg.pipeline.recipe_name) + for slug, cfg in _runner_configs().items() + if cfg.runner_name != cfg.pipeline.recipe_name + } + assert not drifted, f"runner_name != pipeline.recipe_name: {drifted}" + + +def test_runners_have_descriptions() -> None: + empty = [ + slug for slug, cfg in _runner_configs().items() if not cfg.description.strip() + ] + assert not empty, f"runners missing description: {empty}" + + +def test_entry_points_match_module_literals() -> None: + import tomllib + from typing import cast + + from artifixer import config as config_mod + + from flashdreams.infra.runner import RunnerConfig + + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + with pyproject.open("rb") as fh: + meta = tomllib.load(fh) + entries = meta["project"]["entry-points"][ENTRY_POINT_GROUP] + runner_configs = _runner_configs() + declared_slugs = set(entries) + module_slugs = set(runner_configs) + assert declared_slugs == module_slugs, ( + f"entry-point slugs ({sorted(declared_slugs)}) " + f"!= module runners ({sorted(module_slugs)})" + ) + + for slug, target in entries.items(): + module_name, attr = target.split(":", 1) + assert module_name == "artifixer.config", ( + f"unexpected module in entry point {slug!r}: {module_name}" + ) + cfg = cast(RunnerConfig, getattr(config_mod, attr)) + assert cfg.runner_name == slug, ( + f"entry point {slug!r} -> {attr} resolves to " + f"runner_name={cfg.runner_name!r}" + ) + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="entry-point discovery test relies on importlib.metadata 3.10+ shape", +) +def test_entry_points_discoverable_when_installed() -> None: + from importlib.metadata import entry_points + + eps = entry_points(group=ENTRY_POINT_GROUP) + discovered = {ep.name for ep in eps if ep.value.startswith("artifixer.")} + if not discovered: + pytest.skip("plugin not installed; run `uv sync` from the repo root first") + runner_configs = _runner_configs() + assert discovered == set(runner_configs), ( + f"discovered slugs ({sorted(discovered)}) != " + f"plugin runners ({sorted(runner_configs)})" + ) + + +def test_artifixer_hyperparams_match_reference_stage3() -> None: + """Sanity check: AR/scheduler knobs match the stage-3 DMD training config.""" + cfg = _runner_configs()["artifixer-dmd-wan2.1-t2v-1.3b"] + tcfg = cfg.pipeline.diffusion_model.transformer + scfg = cfg.pipeline.diffusion_model.scheduler + + assert tcfg.len_t == 7, "ArtiFixer frames_per_block is 7 latent frames" + assert tcfg.window_size_t == 21, "ArtiFixer local_attn_size is 21 latent frames" + assert tcfg.sink_size_t == 7, "ArtiFixer sink_size is 7 latent frames" + assert scfg.num_inference_steps == 4, "ArtiFixer DMD uses 4-step inference" + assert scfg.shift == 5.0, "ArtiFixer FlowMatchScheduler uses shift=5" + assert tcfg.guidance_scale == 1.0, "ArtiFixer KV-cache pipeline does not use CFG" + + +def test_artifixer_block_has_opacity_and_camera_mlps() -> None: + """ArtifixerBlock instances carry opacity + camera MLPs.""" + import torch + from artifixer.network.block import ArtifixerBlock + from artifixer.network.dit import artifixer_embedding_dims + + opacity_dim, camera_dim = artifixer_embedding_dims((1, 2, 2)) + block = ArtifixerBlock( + dim=1536, + ffn_dim=8960, + num_heads=12, + opacity_embedding_dim=opacity_dim, + camera_embedding_dim=camera_dim, + ) + + assert isinstance(block.opacity_embedding, torch.nn.Linear) + assert isinstance(block.camera_embedding, torch.nn.Linear) + assert block.opacity_embedding.weight.shape == (1536, opacity_dim) + assert block.camera_embedding.weight.shape == (1536, camera_dim) + + # Zero-init contract: the wrapped block is a no-op extension of base + # Wan behavior. Same invariant as the ArtiFixer reference. + assert torch.all(block.opacity_embedding.weight == 0) + assert torch.all(block.opacity_embedding.bias == 0) + assert torch.all(block.camera_embedding.weight == 0) + assert torch.all(block.camera_embedding.bias == 0) + + +def test_artifixer_recipe_uses_artifixer_network() -> None: + """The shipped recipe wires up ArtifixerDiTNetwork.""" + from artifixer.network.dit import ArtifixerDiTNetwork1pt3BConfig + + cfg = _runner_configs()["artifixer-dmd-wan2.1-t2v-1.3b"] + network_cfg = cfg.pipeline.diffusion_model.transformer.network + assert isinstance(network_cfg, ArtifixerDiTNetwork1pt3BConfig), ( + f"recipe network is {type(network_cfg).__name__}, expected ArtifixerDiTNetwork1pt3BConfig" + ) + + +def test_zero_pad_state_dict_transform_fills_missing_keys() -> None: + """The state_dict transform pads vanilla-Wan checkpoints. + + Covers all 9 ArtiFixer-only keys per block: + + * opacity_embedding.{weight,bias} + * camera_embedding.{weight,bias} + * cross_attn.{add_k_proj,add_v_proj}.{weight,bias} + * cross_attn.norm_added_k.weight + """ + import torch + from artifixer.checkpoint import zero_pad_artifixer_keys + from artifixer.network.dit import artifixer_embedding_dims + + transform = zero_pad_artifixer_keys( + num_layers=2, dim=1536, patch_size=(1, 2, 2), dtype=torch.bfloat16 + ) + padded = transform({"unrelated.weight": torch.zeros(8)}) + + opacity_dim, camera_dim = artifixer_embedding_dims((1, 2, 2)) + for layer in range(2): + prefix = f"blocks.{layer}." + assert padded[prefix + "opacity_embedding.weight"].shape == (1536, opacity_dim) + assert padded[prefix + "opacity_embedding.bias"].shape == (1536,) + assert padded[prefix + "camera_embedding.weight"].shape == (1536, camera_dim) + assert padded[prefix + "camera_embedding.bias"].shape == (1536,) + assert padded[prefix + "cross_attn.add_k_proj.weight"].shape == (1536, 1536) + assert padded[prefix + "cross_attn.add_k_proj.bias"].shape == (1536,) + assert padded[prefix + "cross_attn.add_v_proj.weight"].shape == (1536, 1536) + assert padded[prefix + "cross_attn.add_v_proj.bias"].shape == (1536,) + assert padded[prefix + "cross_attn.norm_added_k.weight"].shape == (1536,) + assert padded[prefix + "opacity_embedding.weight"].dtype == torch.bfloat16 + # Untouched keys pass through. + assert "unrelated.weight" in padded + # We pre-allocate exactly 9 keys per layer plus the untouched key. + assert len(padded) == 9 * 2 + 1 + + +def test_artifixer_block_has_neighbor_cross_attn_projections() -> None: + """ArtifixerBlock's cross_attn carries the neighbor branch.""" + import torch + from artifixer.network.block import ArtifixerBlock + from artifixer.network.cross_attn import ArtifixerCrossAttention + from artifixer.network.dit import artifixer_embedding_dims + + opacity_dim, camera_dim = artifixer_embedding_dims((1, 2, 2)) + block = ArtifixerBlock( + dim=1536, + ffn_dim=8960, + num_heads=12, + opacity_embedding_dim=opacity_dim, + camera_embedding_dim=camera_dim, + ) + + assert isinstance(block.cross_attn, ArtifixerCrossAttention), ( + f"cross_attn is {type(block.cross_attn).__name__}, expected ArtifixerCrossAttention" + ) + for name, expected_shape in ( + ("add_k_proj.weight", (1536, 1536)), + ("add_k_proj.bias", (1536,)), + ("add_v_proj.weight", (1536, 1536)), + ("add_v_proj.bias", (1536,)), + ("norm_added_k.weight", (1536,)), + ): + param = block.cross_attn.get_parameter(name) + assert param.shape == expected_shape, f"cross_attn.{name} shape {param.shape}" + + # Zero-init contract: add_v_proj is the gate that keeps the neighbor + # branch contribution zero at load time, matching the ArtiFixer + # reference. + assert torch.all(block.cross_attn.add_v_proj.weight == 0) + assert torch.all(block.cross_attn.add_v_proj.bias == 0) + + +def test_compute_kv_neighbor_and_cache_init() -> None: + """compute_kv_neighbor builds a static BlockKVCache, and + initialize_neighbor_cache toggles the per-module cache attribute. + """ + import torch + from artifixer.network.block import ArtifixerBlock + from artifixer.network.dit import artifixer_embedding_dims + + opacity_dim, camera_dim = artifixer_embedding_dims((1, 2, 2)) + block = ArtifixerBlock( + dim=128, # tiny for fast test + ffn_dim=256, + num_heads=4, + opacity_embedding_dim=opacity_dim, + camera_embedding_dim=camera_dim, + ) + cross_attn = block.cross_attn + + # Starts with no neighbor cache populated. + assert cross_attn.neighbor_kv_cache is None + + # Feed a fake neighbor context. + context = torch.randn(2, 16, 128) # (batch, L_neighbor, dim) + cache = cross_attn.compute_kv_neighbor(context) + assert cache.cached_k().shape == (2, 16, 4, 32) + assert cache.cached_v().shape == (2, 16, 4, 32) + + # initialize_neighbor_cache populates the per-module slot. + cross_attn.initialize_neighbor_cache(context) + assert cross_attn.neighbor_kv_cache is not None + assert cross_attn.neighbor_kv_cache.cached_k().shape == (2, 16, 4, 32) + + # Passing None clears it. + cross_attn.initialize_neighbor_cache(None) + assert cross_attn.neighbor_kv_cache is None diff --git a/integrations/artifixer/tests/test_state_dict_transform.py b/integrations/artifixer/tests/test_state_dict_transform.py new file mode 100644 index 000000000..20b20f590 --- /dev/null +++ b/integrations/artifixer/tests/test_state_dict_transform.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. + +"""Key-naming tests for the merged-DMD state_dict transform. + +CPU-only; uses a ``param_audit.json`` listing every key in the merged +ArtiFixer DMD safetensors to assert that +:func:`artifixer_dmd_state_dict_transform` rewrites the merged +diffusers-format keys into names that we expect to find on the +flashdreams ``ArtifixerDiTNetwork.state_dict()``. Provide the audit +JSON path via ``ARTIFIXER_PARAM_AUDIT_PATH``; the test is skipped if +unset. + +These tests do *not* load the actual safetensors. They synthesize the +expected key set from the audit JSON and run it through the transform. +A separate GPU-side test (later phase) will load the real safetensors +into an instantiated network with ``load_state_dict(strict=True)`` and +verify zero missing / unexpected keys. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import pytest +import torch +from artifixer.checkpoint import ( + DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING, + artifixer_dmd_state_dict_transform, +) + + +def _audit_path() -> Path | None: + raw = os.environ.get("ARTIFIXER_PARAM_AUDIT_PATH") + return Path(raw) if raw else None + + +def _load_audit() -> dict[str, dict]: + audit_path = _audit_path() + if audit_path is None: + pytest.skip( + "ARTIFIXER_PARAM_AUDIT_PATH is unset. Point it at the " + "param_audit.json listing every key in the merged " + "ArtiFixer DMD safetensors to exercise the full-merged-" + "checkpoint path." + ) + if not audit_path.exists(): + pytest.skip( + f"param audit JSON not found at {audit_path}; generate it " + f"from the merged ArtiFixer DMD safetensors first." + ) + return json.loads(audit_path.read_text()) + + +def _synthetic_state_dict(audit: dict[str, dict]) -> dict[str, torch.Tensor]: + """Build a dummy state_dict with one tensor per key in the audit.""" + state_dict: dict[str, torch.Tensor] = {} + for category in ("base_shared", "artifixer_only"): + for key, shape in audit[category].items(): + state_dict[key] = torch.zeros(shape) + return state_dict + + +def test_transform_renames_known_diffusers_keys() -> None: + """Spot-check a few diffusers -> WanDiTNetwork renames the mapping does.""" + sd = { + "blocks.0.attn1.to_q.weight": torch.zeros(1536, 1536), + "blocks.0.attn2.to_q.weight": torch.zeros(1536, 1536), + "blocks.0.attn2.add_k_proj.weight": torch.zeros(1536, 1536), + "blocks.0.attn2.add_v_proj.bias": torch.zeros(1536), + "blocks.0.attn2.norm_added_k.weight": torch.zeros(1536), + "blocks.0.ffn.net.0.proj.weight": torch.zeros(8960, 1536), + "blocks.0.ffn.net.2.weight": torch.zeros(1536, 8960), + "blocks.0.norm2.weight": torch.zeros(1536), + "scale_shift_table": torch.zeros(2, 1536), + "proj_out.weight": torch.zeros(64, 1536), + } + out = artifixer_dmd_state_dict_transform(sd) + + expected_renames = { + "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight", + "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight", + "blocks.0.attn2.add_k_proj.weight": "blocks.0.cross_attn.add_k_proj.weight", + "blocks.0.attn2.add_v_proj.bias": "blocks.0.cross_attn.add_v_proj.bias", + "blocks.0.attn2.norm_added_k.weight": "blocks.0.cross_attn.norm_added_k.weight", + "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight", + "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight", + "blocks.0.norm2.weight": "blocks.0.norm3.weight", + "scale_shift_table": "head.modulation", + "proj_out.weight": "head.head.weight", + } + for src, dst in expected_renames.items(): + assert src not in out, f"diffusers key {src!r} should have been renamed" + assert dst in out, f"expected remapped key {dst!r} missing from transform output" + + +def test_transform_preserves_opacity_and_camera_keys_unchanged() -> None: + """ArtiFixer-only keys without ``attn2`` prefix pass through unchanged. + + ``opacity_embedding`` and ``camera_embedding`` are registered directly + on :class:`ArtifixerBlock` (no nested wrapper), so their reference + names match the flashdreams names verbatim. + """ + sd = { + "blocks.0.opacity_embedding.weight": torch.zeros(1536, 1024), + "blocks.0.opacity_embedding.bias": torch.zeros(1536), + "blocks.5.camera_embedding.weight": torch.zeros(1536, 1536), + } + out = artifixer_dmd_state_dict_transform(sd) + assert set(out) == set(sd) + + +def test_transform_on_param_audit_yields_unique_keys() -> None: + """Apply the transform to the entire merged-checkpoint key set and + assert no two source keys collapse to the same target key. + """ + audit = _load_audit() + src_sd = _synthetic_state_dict(audit) + out = artifixer_dmd_state_dict_transform(src_sd) + assert len(out) == len(src_sd), ( + f"transform collapsed {len(src_sd) - len(out)} keys to duplicates" + ) + + +def test_transform_on_param_audit_matches_expected_block_layout() -> None: + """End-to-end shape check on the merged-DMD audit. + + Required post-transform invariants per block ``X`` (30 blocks total): + + - self_attn.{q,k,v,o,norm_q,norm_k} + - cross_attn.{q,k,v,o,norm_q,norm_k} + - cross_attn.{add_k_proj,add_v_proj,norm_added_k} (ArtiFixer-only) + - opacity_embedding, camera_embedding (ArtiFixer-only) + - norm3, ffn.0, ffn.2, modulation + """ + audit = _load_audit() + sd = _synthetic_state_dict(audit) + out = artifixer_dmd_state_dict_transform(sd) + + expected_per_block = [ + "self_attn.q.weight", + "self_attn.k.weight", + "self_attn.v.weight", + "self_attn.o.weight", + "self_attn.norm_q.weight", + "self_attn.norm_k.weight", + "cross_attn.q.weight", + "cross_attn.k.weight", + "cross_attn.v.weight", + "cross_attn.o.weight", + "cross_attn.norm_q.weight", + "cross_attn.norm_k.weight", + "cross_attn.add_k_proj.weight", + "cross_attn.add_v_proj.weight", + "cross_attn.norm_added_k.weight", + "opacity_embedding.weight", + "camera_embedding.weight", + "norm3.weight", + "ffn.0.weight", + "ffn.2.weight", + "modulation", + ] + num_blocks = max( + int(k.split(".", 2)[1]) for k in out if k.startswith("blocks.") + ) + 1 + assert num_blocks == 30, f"expected 30 transformer blocks, got {num_blocks}" + for block_idx in range(num_blocks): + for suffix in expected_per_block: + key = f"blocks.{block_idx}.{suffix}" + assert key in out, f"missing post-transform key {key!r}" + + +def test_transform_resolves_globals() -> None: + """Network-level keys land on the expected ``head.*`` / embedding names.""" + audit = _load_audit() + sd = _synthetic_state_dict(audit) + out = artifixer_dmd_state_dict_transform(sd) + expected_globals = [ + "patch_embedding.weight", + "patch_embedding.bias", + "text_embedding.0.weight", + "text_embedding.2.weight", + "time_embedding.0.weight", + "time_embedding.2.weight", + "time_projection.1.weight", + "head.modulation", + "head.head.weight", + "head.head.bias", + ] + for key in expected_globals: + assert key in out, f"missing post-transform global {key!r}" + + +def test_diffusers_mapping_has_no_overlap_in_capture_groups() -> None: + """Every regex in the mapping uses ``\\1`` / ``\\2`` consistently.""" + for old, new in DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING.items(): + # All groups referenced in ``new`` must exist in ``old``. + for ref in (1, 2, 3): + if f"\\{ref}" in new: + # Count ``(`` not preceded by ``?:``. + groups = old.count("(") - old.count("(?:") + assert ref <= groups, ( + f"pattern {old!r} -> {new!r} references group {ref} " + f"but pattern only has {groups} capture group(s)" + ) diff --git a/pyproject.toml b/pyproject.toml index 82c35b823..2a800cdb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ transformer-engine-torch = { git = "https://github.com/NVIDIA/TransformerEngine. extraPaths = [ "flashdreams", "integrations/alpadreams", + "integrations/artifixer", "integrations/causal_forcing", "integrations/fastvideo_causal_wan22", "integrations/lingbot", @@ -59,6 +60,7 @@ python-version = "3.12" extra-paths = [ "flashdreams", "integrations/alpadreams", + "integrations/artifixer", "integrations/causal_forcing", "integrations/fastvideo_causal_wan22", "integrations/lingbot", diff --git a/scripts/import_flashdreams_sqsh.sh b/scripts/import_flashdreams_sqsh.sh new file mode 100755 index 000000000..cbe330a65 --- /dev/null +++ b/scripts/import_flashdreams_sqsh.sh @@ -0,0 +1,47 @@ +#!/bin/bash +#SBATCH --job-name=fd-import +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=8 +#SBATCH --gpus-per-node=1 +#SBATCH --mem=64G +#SBATCH --time=01:00:00 +#SBATCH --partition=interactive_singlenode +#SBATCH --account=nvr_torontoai_videogen +#SBATCH --output=slurm-logs/fd-import-%j.out +#SBATCH --error=slurm-logs/fd-import-%j.err +# +# One-shot: pull the FlashDreams container from ghcr.io into a sqsh file +# under /lustre so later slurm jobs can reuse it via +# ``--container-image=/lustre/.../flashdreams-base-v0.3.sqsh`` instead of +# re-pulling from the registry on every worker. + +set -euo pipefail + +OUTPUT_SQSH="${1:-/lustre/fsw/portfolios/nvr/users/rdelutio/containers/flashdreams-base-v0.3.sqsh}" +IMAGE_URI="${2:-ghcr.io#nvidia/flashdreams:base-v0.3-20260424-55bd566}" + +mkdir -p "$(dirname "$OUTPUT_SQSH")" + +if [ -f "$OUTPUT_SQSH" ]; then + echo "[import] $OUTPUT_SQSH already exists; refusing to overwrite. Delete or pass a new path." >&2 + exit 0 +fi + +# Use a per-job ENROOT_TEMP / RUNTIME directory under the worker's /raid to +# avoid clashes with concurrent jobs that share /raid/containers. +JOB_RAID_DIR="/raid/containers/tmp/${USER}-${SLURM_JOB_ID:-import}" +mkdir -p "${JOB_RAID_DIR}" +export ENROOT_TEMP_PATH="${JOB_RAID_DIR}/tmp" +export ENROOT_RUNTIME_PATH="${JOB_RAID_DIR}/runtime" +export ENROOT_DATA_PATH="${JOB_RAID_DIR}/data" +mkdir -p "${ENROOT_TEMP_PATH}" "${ENROOT_RUNTIME_PATH}" "${ENROOT_DATA_PATH}" + +echo "[import] target sqsh: $OUTPUT_SQSH" +echo "[import] image uri: docker://${IMAGE_URI}" + +srun --ntasks=1 --cpus-per-task="${SLURM_CPUS_PER_TASK:-8}" \ + bash -c "enroot import --output '${OUTPUT_SQSH}' 'docker://${IMAGE_URI}'" + +echo "[import] done. Size:" +ls -lh "${OUTPUT_SQSH}" diff --git a/scripts/run_prope_parity.sh b/scripts/run_prope_parity.sh new file mode 100755 index 000000000..ba39ea604 --- /dev/null +++ b/scripts/run_prope_parity.sh @@ -0,0 +1,43 @@ +#!/bin/bash +#SBATCH --job-name=prope-parity +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-node=1 +#SBATCH --mem=16G +#SBATCH --time=00:10:00 +#SBATCH --partition=interactive_singlenode +#SBATCH --account=nvr_torontoai_videogen +#SBATCH --output=slurm-logs/prope-parity-%j.out +#SBATCH --error=slurm-logs/prope-parity-%j.err +# +# Run the PRoPE numerical-parity test that compares our port at +# integrations/artifixer/artifixer/network/prope.py against the ArtiFixer +# reference's model_training/net/prope.py. +# +# The artifixer container has torch installed at the system Python level +# and the reference repo is mounted via /lustre, so we don't need uv sync. + +set -euo pipefail + +REPO_DIR="${REPO_DIR:-$(git -C "${SLURM_SUBMIT_DIR:-$PWD}" rev-parse --show-toplevel 2>/dev/null || true)}" +ARTIFIXER_REFERENCE_REPO_ROOT="${ARTIFIXER_REFERENCE_REPO_ROOT:-}" +if [ -z "${ARTIFIXER_REFERENCE_REPO_ROOT}" ]; then + echo "Set ARTIFIXER_REFERENCE_REPO_ROOT to the ArtiFixer reference checkout." >&2 + exit 1 +fi +CONTAINER_IMAGE="${ARTIFIXER_CONTAINER_IMAGE:-/lustre/fsw/portfolios/nvr/users/hturki/containers/artifixer-cuda12.sqsh}" + +COMMAND="set -euo pipefail; \ +cd ${REPO_DIR} && \ +export PYTHONPATH=${REPO_DIR}/integrations/artifixer:${REPO_DIR}/flashdreams:${ARTIFIXER_REFERENCE_REPO_ROOT}:\${PYTHONPATH:-} && \ +export ARTIFIXER_REFERENCE_REPO_ROOT='${ARTIFIXER_REFERENCE_REPO_ROOT}' && \ +python3 -m pip install --user pytest loguru einops boto3 tyro 2>&1 | tail -3 && \ +python3 -m pytest -v integrations/artifixer/tests/test_prope_parity.py integrations/artifixer/tests/test_patches.py integrations/artifixer/tests/test_latent_mix.py integrations/artifixer/tests/test_state_dict_transform.py integrations/artifixer/tests/test_smoke.py::test_compute_kv_neighbor_and_cache_init" + +echo "[slurm] command: ${COMMAND}" + +srun --container-image="${CONTAINER_IMAGE}" \ + --container-mounts="$HOME:/home,/lustre:/lustre" \ + --container-workdir="${REPO_DIR}" \ + bash -c "${COMMAND}" diff --git a/scripts/smoke_test_artifixer.sh b/scripts/smoke_test_artifixer.sh new file mode 100755 index 000000000..cc9129b30 --- /dev/null +++ b/scripts/smoke_test_artifixer.sh @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH --job-name=fd-artifixer-smoke +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-node=1 +#SBATCH --mem=16G +#SBATCH --time=00:30:00 +#SBATCH --partition=interactive_singlenode +#SBATCH --account=nvr_torontoai_videogen +#SBATCH --output=slurm-logs/fd-artifixer-smoke-%j.out +#SBATCH --error=slurm-logs/fd-artifixer-smoke-%j.err +# +# Run the artifixer plugin smoke tests inside the FlashDreams container. +# +# Phase 1: smoke tests are dataclass / entry-point validation only; they +# need ``flashdreams`` importable + pytest. The flashdreams container ships +# with ``uv`` but not a pre-baked venv, so first invocation runs ``uv sync`` +# (slow). Subsequent runs in the same workspace reuse ``.venv/``. +# +# If ``uv sync`` is too slow for iteration, run the static checks in +# ``scripts/static_check_artifixer.sh`` instead (no env, no GPU, ~1s). +# +# Requires: +# - ``/lustre/fsw/portfolios/nvr/users/rdelutio/containers/flashdreams-base-v0.3.sqsh`` +# created via ``scripts/import_flashdreams_sqsh.sh`` (one-time) +# - GitHub PAT for ghcr.io (only for first sqsh import) + +set -euo pipefail + +REPO_DIR="${REPO_DIR:-$(git -C "${SLURM_SUBMIT_DIR:-$PWD}" rev-parse --show-toplevel 2>/dev/null || true)}" +if [ -z "$REPO_DIR" ]; then + echo "Could not determine REPO_DIR; set REPO_DIR or submit from inside the repo." >&2 + exit 1 +fi + +CONTAINER_IMAGE="${FLASHDREAMS_CONTAINER_IMAGE:-/lustre/fsw/portfolios/nvr/users/rdelutio/containers/flashdreams-base-v0.3.sqsh}" + +COMMAND="set -euo pipefail; \ +cd ${REPO_DIR} && \ +uv sync --extra dev --extra runners --group lint && \ +uv run --extra dev pytest -q integrations/artifixer/tests" + +echo "[slurm] command: ${COMMAND}" + +srun --container-image="${CONTAINER_IMAGE}" \ + --container-mounts="$HOME:/home,/lustre:/lustre" \ + --container-workdir="${REPO_DIR}" \ + bash -c "${COMMAND}" diff --git a/scripts/static_check_artifixer.sh b/scripts/static_check_artifixer.sh new file mode 100755 index 000000000..1e143e4a1 --- /dev/null +++ b/scripts/static_check_artifixer.sh @@ -0,0 +1,72 @@ +#!/bin/bash +# Fast static check for the artifixer plugin. Runs on any host with system +# python3 (no torch / flashdreams / uv required). Catches syntax errors, +# malformed pyproject.toml, and dangling imports against the flashdreams +# source tree. Designed as a sub-second iteration loop for Phase 1 scaffold +# changes before we go through the full slurm-based smoke test. +# +# Usage: +# bash scripts/static_check_artifixer.sh + +set -euo pipefail + +REPO_DIR="$(git -C "${PWD}" rev-parse --show-toplevel)" +cd "${REPO_DIR}" + +echo "[static-check] py_compile artifixer sources" +for f in $(find integrations/artifixer -name '*.py' -not -path '*/__pycache__/*'); do + python3 -m py_compile "$f" +done + +echo "[static-check] pyproject.toml entry-point matches a config.py attribute" +python3 - <<'PY' +import re + +with open("integrations/artifixer/pyproject.toml") as f: + pyproject = f.read() +m = re.search( + r'\[project\.entry-points\."flashdreams\.runner_configs"\]\s*\n((?:"[^"]+"\s*=\s*"[^"]+"\s*\n)+)', + pyproject, +) +assert m, "missing [project.entry-points.\"flashdreams.runner_configs\"] block" +entries = dict(re.findall(r'"([^"]+)"\s*=\s*"([^"]+)"', m.group(1))) + +with open("integrations/artifixer/artifixer/config.py") as f: + config_py = f.read() + +for slug, target in entries.items(): + module, attr = target.split(":", 1) + assert module == "artifixer.config", f"unexpected module in {slug!r}: {module}" + assert re.search(rf"^{re.escape(attr)}\s*=", config_py, flags=re.MULTILINE), ( + f"entry-point {slug!r} -> {attr} not defined in artifixer/config.py" + ) +print(f" ok: {len(entries)} entry-points all resolve to config.py attributes") +PY + +echo "[static-check] flashdreams imports resolve to classes in the source tree" +python3 - <<'PY' +import os +import re +from pathlib import Path + +import_re = re.compile( + r"from\s+(flashdreams\.[\w.]+)\s+import\s+(?:\(\s*([\s\S]*?)\s*\)|([^\n]+))", +) +flashdreams_files = list(Path("flashdreams").rglob("*.py")) +for source_path in sorted(Path("integrations/artifixer/artifixer").rglob("*.py")): + src = source_path.read_text() + for module, paren_blob, inline_blob in import_re.findall(src): + names_blob = paren_blob or inline_blob + for name in re.split(r"[,\s]+", names_blob): + name = name.strip() + if not name: + continue + pat = re.compile(rf"^class {re.escape(name)}\b|^{re.escape(name)}\s*=", re.MULTILINE) + hit = any(pat.search(p.read_text(errors="ignore")) for p in flashdreams_files) + assert hit, ( + f"{source_path}: from {module} import {name} — no class/def in flashdreams/" + ) + print(f" ok: from {module} import ({names_blob.strip()})") +PY + +echo "[static-check] all good"