From 28da6344ba33ab7bd44dc831199e7a9bd2f97ab6 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 18:56:37 -0700 Subject: [PATCH 01/22] integrations/artifixer: scaffold recipe for ArtiFixer DMD T2V Mirrors the integrations/self_forcing layout: pyproject.toml workspace member registering a single ``artifixer-dmd-wan2.1-t2v-1.3b`` entry point, artifixer/config.py with the AR / scheduler knobs that match the dreamfix stage-3 DMD training config (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), and a stripped clone of SelfForcingT2VRunner that takes a single text prompt. Phase 1 only: this loads vanilla Wan 2.1 1.3B base weights from HuggingFace so the recipe end-to-end is exercisable but the output is plain T2V. Later commits add the ArtiFixer architecture extensions (opacity/camera MLPs, neighbor cross-attention with PRoPE, opacity-weighted latent mixing) and the state_dict_transform that loads the merged DMD safetensors produced by dreamfix/scripts/merge_dcp_to_safetensors.py. Also adds three helper scripts under scripts/: - import_flashdreams_sqsh.sh: one-shot pull of the FlashDreams container from ghcr.io into a persistent sqsh under /lustre so later jobs avoid re-pulling and avoid the pyxis tmp-extraction failures. - smoke_test_artifixer.sh: slurm wrapper that runs the plugin smoke tests via ``uv sync`` + pytest inside the container. - static_check_artifixer.sh: sub-second static-only sanity check (py_compile + entry-point sanity + flashdreams import resolution) without needing torch/uv/GPU. Designed for fast Phase 1 iteration. Workspace pyproject.toml: add integrations/artifixer to the pyright / ty extra-paths lists. Verified: scripts/static_check_artifixer.sh passes (py_compile clean, 11 flashdreams imports all resolve to classes in flashdreams/, entry-point maps to RUNNER_ARTIFIXER_DMD_T2V_1PT3B in artifixer/config.py). --- .gitignore | 1 + integrations/artifixer/README.md | 65 +++++++++ integrations/artifixer/artifixer/__init__.py | 14 ++ integrations/artifixer/artifixer/config.py | 113 +++++++++++++++ integrations/artifixer/artifixer/runner.py | 141 +++++++++++++++++++ integrations/artifixer/assets/prompt.txt | 1 + integrations/artifixer/pyproject.toml | 48 +++++++ integrations/artifixer/tests/test_smoke.py | 106 ++++++++++++++ pyproject.toml | 2 + scripts/import_flashdreams_sqsh.sh | 47 +++++++ scripts/smoke_test_artifixer.sh | 49 +++++++ scripts/static_check_artifixer.sh | 79 +++++++++++ 12 files changed, 666 insertions(+) create mode 100644 integrations/artifixer/README.md create mode 100644 integrations/artifixer/artifixer/__init__.py create mode 100644 integrations/artifixer/artifixer/config.py create mode 100644 integrations/artifixer/artifixer/runner.py create mode 100644 integrations/artifixer/assets/prompt.txt create mode 100644 integrations/artifixer/pyproject.toml create mode 100644 integrations/artifixer/tests/test_smoke.py create mode 100755 scripts/import_flashdreams_sqsh.sh create mode 100755 scripts/smoke_test_artifixer.sh create mode 100755 scripts/static_check_artifixer.sh 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..cf6ae4d5d --- /dev/null +++ b/integrations/artifixer/README.md @@ -0,0 +1,65 @@ +# 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)`). + +The reference implementation lives in the +[dreamfix repo](https://gitlab-master.nvidia.com/hturki/dreamfix) under +`model_training/net/transformer.py` and `model_training/pipeline/`. This +plugin ports it to FlashDreams' faster Wan stack (RingAttention, cuDNN, +`torch.compile`, CUDA graphs). + +## Status + +| Phase | Scope | Status | +| --- | --- | --- | +| 1 | Recipe scaffold + AR/scheduler knobs match dreamfix stage-3 DMD | this commit | +| 2 | Per-block opacity + camera MLPs, neighbor cross-attn, PRoPE | upcoming | +| 3 | Opacity-weighted latent mixing + self-forcing renoise loop | upcoming | +| 4 | dreamfix-format conditioning surface (rgb_rendered, opacity, neighbors) | upcoming | +| 5 | `state_dict_transform` for the merged ArtiFixer DMD safetensors | upcoming | + +Phase 1 loads vanilla Wan 2.1 1.3B base weights from HuggingFace, so the +output is a plain T2V video. The recipe is wired up but the ArtiFixer +architectural extensions land in later commits. + +## 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/config.py b/integrations/artifixer/artifixer/config.py new file mode 100644 index 000000000..a3689c45e --- /dev/null +++ b/integrations/artifixer/artifixer/config.py @@ -0,0 +1,113 @@ +# 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. This file +ships the Phase 1 *recipe scaffold*: it wires up the runner / pipeline / +scheduler with the ArtiFixer hyperparameters (4-step flow-matching with +``shift=5``, chunked AR with ``len_t=7``, ``window_size_t=21``, +``sink_size_t=7``) but loads vanilla Wan 2.1 1.3B base weights from HF. + +Later commits replace the base ``Wan21TransformerConfig`` with an +``ArtifixerWanTransformerConfig`` carrying the architecture extensions and +the merged distilled-DMD safetensors. +""" + +from __future__ import annotations + +from artifixer.runner import ArtifixerDmdT2VRunnerConfig + +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 ( + Wan21TransformerConfig, + WanDiTNetwork1pt3BConfig, + WanInferencePipelineConfig, + WanVAEDecoderConfig, +) + +# Mirrors the dreamfix stage-3 run config (see +# wandb/.../run-*/files/config.yaml in +# artifixer-runs/artifixer-s3-dmd-1p3b-from-s2-10000-s1-15000-128g): +# 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 + +# Phase 1: load vanilla Wan 2.1 1.3B base weights from HF. Phase 5 replaces +# this with the merged ArtiFixer DMD safetensors plus a state_dict_transform +# that remaps diffusers naming and absorbs the ArtiFixer-only keys. +BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH = ( + "https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B/blob/main/diffusion_pytorch_model.safetensors" +) + + +PIPELINE_ARTIFIXER_DMD_T2V_1PT3B = WanInferencePipelineConfig( + recipe_name="artifixer-dmd-wan2.1-t2v-1.3b", + enable_sync_and_profile=True, + encoder=None, + decoder=WanVAEDecoderConfig(), + diffusion_model=DiffusionModelConfig( + seed=42, + transformer=Wan21TransformerConfig( + network=WanDiTNetwork1pt3BConfig( + patch_embedding_type="conv3d", + ), + checkpoint_path=BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH, + state_dict_transform=None, + 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). Phase 1 scaffold: loads vanilla Wan base weights." + ), + 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/runner.py b/integrations/artifixer/artifixer/runner.py new file mode 100644 index 000000000..c5320d55a --- /dev/null +++ b/integrations/artifixer/artifixer/runner.py @@ -0,0 +1,141 @@ +# 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. + +Phase 1 scaffold: this runner is a stripped clone of +``SelfForcingT2VRunner`` and accepts a single text prompt. Phase 3/4 will +introduce the ArtiFixer-specific conditioning surface (``rgb_rendered``, +``opacity``, neighbor frames, camera matrices) by accepting a path to a +dreamfix-format input bundle (e.g. an HDF5 reconstruction file plus a +caption file). +""" + +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 dreamfix's + ``num_frames=81`` / ``frames_per_block=7`` (21 latent frames / 7 = 3 + chunks for 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 (Phase 1: text-only scaffold).""" + + 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/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_smoke.py b/integrations/artifixer/tests/test_smoke.py new file mode 100644 index 000000000..330f5d382 --- /dev/null +++ b/integrations/artifixer/tests/test_smoke.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. + +"""Cheap import-time checks for the ``artifixer`` plugin.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import cast + +import pytest +import tomllib +from artifixer import config as config_mod +from artifixer.config import RUNNER_CONFIGS + +from flashdreams.infra.runner import RunnerConfig + +ENTRY_POINT_GROUP = "flashdreams.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: + 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] + 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") + assert discovered == set(RUNNER_CONFIGS), ( + f"discovered slugs ({sorted(discovered)}) != " + f"plugin runners ({sorted(RUNNER_CONFIGS)})" + ) + + +def test_artifixer_hyperparams_match_dreamfix_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" 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/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..ccb55c30e --- /dev/null +++ b/scripts/static_check_artifixer.sh @@ -0,0 +1,79 @@ +#!/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 \ + integrations/artifixer/artifixer/__init__.py \ + integrations/artifixer/artifixer/config.py \ + integrations/artifixer/artifixer/runner.py \ + integrations/artifixer/tests/test_smoke.py; 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 ( + "integrations/artifixer/artifixer/config.py", + "integrations/artifixer/artifixer/runner.py", +): + src = Path(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" From 8c3173b8aaccdefb5f507822eff755cc6bb54281 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 19:02:31 -0700 Subject: [PATCH 02/22] integrations/artifixer: add opacity + camera-ray MLPs (ArtifixerBlock) Phase 2.1 of the dreamfix port: each transformer block now carries the per-block opacity and Plucker-camera-ray projection heads from ``ArtifixerTransformerBlock`` (dreamfix/model_training/net/transformer.py L617-L767). The MLP outputs are added to the AdaLN-normed hidden states before self-attention. Both heads are zero-initialized so the block behaves identically to the underlying ``Block`` when ``opacity_extra`` / ``camera_extra`` are not passed (Phase 1 text-only path). New modules under integrations/artifixer/artifixer/: - network/block.py ArtifixerBlock(Block) with opacity/camera MLPs - network/dit.py ArtifixerDiTNetwork(WanDiTNetwork) + ArtifixerDiTNetwork1pt3BConfig, plus artifixer_embedding_dims() helper deriving opacity_dim=1024 and camera_dim=1536 from the Wan VAE strides + patch_size=(1,2,2) - checkpoint.py zero_pad_artifixer_keys() state_dict transform that zero-fills the 4 new keys per block when loading a vanilla Wan checkpoint so strict-mode load_state_dict still succeeds config.py wiring: - Replaces WanDiTNetwork1pt3BConfig with ArtifixerDiTNetwork1pt3BConfig - Sets state_dict_transform=zero_pad_artifixer_keys(...) for the BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH source Smoke tests extended: - ArtifixerBlock instantiates with correct linear shapes and zero-init invariants - Recipe wires up ArtifixerDiTNetwork1pt3BConfig (not vanilla) - zero_pad_artifixer_keys() fills missing keys with bf16 zeros Static-check (scripts/static_check_artifixer.sh) now scans every .py file under integrations/artifixer/artifixer/ and all 9 flashdreams imports across the recipe resolve cleanly. --- .../artifixer/artifixer/checkpoint.py | 97 ++++++++++++++++ integrations/artifixer/artifixer/config.py | 41 ++++--- .../artifixer/artifixer/network/__init__.py | 30 +++++ .../artifixer/artifixer/network/block.py | 109 ++++++++++++++++++ .../artifixer/artifixer/network/dit.py | 100 ++++++++++++++++ integrations/artifixer/tests/test_smoke.py | 61 ++++++++++ scripts/static_check_artifixer.sh | 13 +-- 7 files changed, 428 insertions(+), 23 deletions(-) create mode 100644 integrations/artifixer/artifixer/checkpoint.py create mode 100644 integrations/artifixer/artifixer/network/__init__.py create mode 100644 integrations/artifixer/artifixer/network/block.py create mode 100644 integrations/artifixer/artifixer/network/dit.py diff --git a/integrations/artifixer/artifixer/checkpoint.py b/integrations/artifixer/artifixer/checkpoint.py new file mode 100644 index 000000000..f44d47bd5 --- /dev/null +++ b/integrations/artifixer/artifixer/checkpoint.py @@ -0,0 +1,97 @@ +# 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`` (used in + Phase 1 / 2.1 while only the base architecture is implemented). The + transform zero-pads the ArtiFixer-only keys so ``load_state_dict`` + succeeds in strict mode. Zero-padding matches the dreamfix + initialization in ``ArtifixerTransformerBlock.__init__`` L637-651. + + * **Merged ArtiFixer DMD safetensors** produced by + ``dreamfix/scripts/merge_dcp_to_safetensors.py``. This still uses the + HuggingFace diffusers naming (e.g. ``blocks.X.attn1.to_q.weight``); + Phase 5 will add the diffusers -> ``WanDiTNetwork`` regex remap on + top of the zero-pad pass. +""" + +from __future__ import annotations + +from typing import Callable + +import torch +from torch import Tensor + +from artifixer.network.dit import artifixer_embedding_dims + +# Suffixes of the keys that the dreamfix ``ArtifixerTransformerBlock`` adds +# per transformer block. Phase 2.1 covers the opacity + camera MLPs; Phase +# 2.2 will add ``attn2.add_k_proj``, ``attn2.add_v_proj``, +# ``attn2.norm_added_k`` to this list. +_PHASE2_1_PER_BLOCK_SUFFIXES: tuple[str, ...] = ( + "opacity_embedding.weight", + "opacity_embedding.bias", + "camera_embedding.weight", + "camera_embedding.bias", +) + + +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. + + 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,) + + 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__ = [ + "zero_pad_artifixer_keys", +] diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index a3689c45e..009c83251 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -17,19 +17,26 @@ 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. This file -ships the Phase 1 *recipe scaffold*: it wires up the runner / pipeline / -scheduler with the ArtiFixer hyperparameters (4-step flow-matching with -``shift=5``, chunked AR with ``len_t=7``, ``window_size_t=21``, -``sink_size_t=7``) but loads vanilla Wan 2.1 1.3B base weights from HF. +attention with PRoPE, and (c) opacity-weighted latent mixing. -Later commits replace the base ``Wan21TransformerConfig`` with an -``ArtifixerWanTransformerConfig`` carrying the architecture extensions and -the merged distilled-DMD safetensors. +Phase 2.1: the network now uses :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). +The recipe still loads vanilla Wan 2.1 1.3B base weights from HuggingFace; +:func:`zero_pad_artifixer_keys` pads the state dict so strict-mode load +succeeds. + +Later commits add the neighbor cross-attention third KV bank (Phase 2.2), +PRoPE (Phase 2.3), the opacity-weighted latent mixing pipeline (Phase 3), +and the ``state_dict_transform`` for the merged ArtiFixer DMD safetensors +(Phase 5). """ from __future__ import annotations +import torch +from artifixer.checkpoint import zero_pad_artifixer_keys +from artifixer.network import ArtifixerDiTNetwork1pt3BConfig from artifixer.runner import ArtifixerDmdT2VRunnerConfig from flashdreams.infra.diffusion.model import DiffusionModelConfig @@ -37,7 +44,6 @@ from flashdreams.infra.runner import RunnerConfig from flashdreams.recipes.wan import ( Wan21TransformerConfig, - WanDiTNetwork1pt3BConfig, WanInferencePipelineConfig, WanVAEDecoderConfig, ) @@ -65,6 +71,11 @@ ) +_BASE_NETWORK_CONFIG = ArtifixerDiTNetwork1pt3BConfig( + patch_embedding_type="conv3d", +) +_BASE_TRANSFORMER_DTYPE = torch.bfloat16 + PIPELINE_ARTIFIXER_DMD_T2V_1PT3B = WanInferencePipelineConfig( recipe_name="artifixer-dmd-wan2.1-t2v-1.3b", enable_sync_and_profile=True, @@ -73,11 +84,15 @@ diffusion_model=DiffusionModelConfig( seed=42, transformer=Wan21TransformerConfig( - network=WanDiTNetwork1pt3BConfig( - patch_embedding_type="conv3d", - ), + network=_BASE_NETWORK_CONFIG, + dtype=_BASE_TRANSFORMER_DTYPE, checkpoint_path=BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH, - state_dict_transform=None, + 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, + ), batch_shape=(), len_t=ARTIFIXER_LEN_T, guidance_scale=1.0, diff --git a/integrations/artifixer/artifixer/network/__init__.py b/integrations/artifixer/artifixer/network/__init__.py new file mode 100644 index 000000000..c6d63cbd9 --- /dev/null +++ b/integrations/artifixer/artifixer/network/__init__.py @@ -0,0 +1,30 @@ +# 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. + +from artifixer.network.block import ArtifixerBlock +from artifixer.network.dit import ( + ArtifixerDiTNetwork, + ArtifixerDiTNetwork1pt3BConfig, + ArtifixerDiTNetworkConfig, + artifixer_embedding_dims, +) + +__all__ = [ + "ArtifixerBlock", + "ArtifixerDiTNetwork", + "ArtifixerDiTNetwork1pt3BConfig", + "ArtifixerDiTNetworkConfig", + "artifixer_embedding_dims", +] diff --git a/integrations/artifixer/artifixer/network/block.py b/integrations/artifixer/artifixer/network/block.py new file mode 100644 index 000000000..0f98e5c25 --- /dev/null +++ b/integrations/artifixer/artifixer/network/block.py @@ -0,0 +1,109 @@ +# 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-ray conditioning MLPs. + +Mirrors ``ArtifixerTransformerBlock`` in the dreamfix reference +(``model_training/net/transformer.py`` L617-L767). Two ``nn.Linear`` heads +project per-token opacity and Plucker-camera-ray features into the +transformer hidden size, and their outputs are added to the AdaLN-normed +hidden states **before** self-attention. Both heads are zero-initialized so +loading a vanilla Wan checkpoint leaves the block behavior unchanged. +""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn + +from flashdreams.recipes.wan.transformer.impl.modules import Block, BlockCache + + +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, + ) + + 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. Matches dreamfix transformer.py L637-651. + 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, + ) -> Tensor: + """Run one transformer block update with ArtiFixer conditioning. + + ``opacity_extra`` and ``camera_extra`` are added to the AdaLN-normed + hidden states before self-attention (matches dreamfix + ``ArtifixerTransformerBlock.forward`` L763-767). When both are + ``None`` (Phase 1 text-only path), the block is identical to the + underlying ``Block``. + """ + assert self._parameters_updated_after_loading_checkpoint, ( + "We expect to have called update_parameters_after_loading_checkpoint() " + "before running the forward pass" + ) + e_chunks = (self.modulation + e).chunk(6, dim=-2) + + y = self.norm1(x) * (1 + e_chunks[1]) + e_chunks[0] + 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 + (y * e_chunks[2]) + + x = x + self.cross_attn( + self.norm3(x), + kv_cache=cache.cross_attn, + ) + y = self.norm2(x) * (1 + e_chunks[4]) + e_chunks[3] + y = self.ffn(y) + x = x + (y * e_chunks[5]) + return x diff --git a/integrations/artifixer/artifixer/network/dit.py b/integrations/artifixer/artifixer/network/dit.py new file mode 100644 index 000000000..79f373250 --- /dev/null +++ b/integrations/artifixer/artifixer/network/dit.py @@ -0,0 +1,100 @@ +# 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 flashdreams.recipes.wan.transformer.impl.network import ( + WanDiTNetwork, + WanDiTNetwork1pt3BConfig, + WanDiTNetworkConfig, +) + +# Wan2.1 VAE downsampling factors. Identical for T2V-1.3B and T2V-14B and +# what the dreamfix reference (``ArtifixerTransformer.__init__`` L228-237) +# uses 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. + + Mirrors ``ArtifixerTransformer.__init__`` L228-237. + + 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 dreamfix stage-3 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, + ) diff --git a/integrations/artifixer/tests/test_smoke.py b/integrations/artifixer/tests/test_smoke.py index 330f5d382..ea38abee5 100644 --- a/integrations/artifixer/tests/test_smoke.py +++ b/integrations/artifixer/tests/test_smoke.py @@ -104,3 +104,64 @@ def test_artifixer_hyperparams_match_dreamfix_stage3() -> None: 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: + """Phase 2.1: ArtifixerBlock instances carry opacity + camera MLPs.""" + import torch + from artifixer.network import ArtifixerBlock, 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 dreamfix transformer.py L637-651. + 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: + """Phase 2.1: the shipped recipe wires up ArtifixerDiTNetwork.""" + from artifixer.network 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: + """Phase 2.1: the state_dict transform pads vanilla-Wan checkpoints.""" + import torch + from artifixer.checkpoint import zero_pad_artifixer_keys + from artifixer.network 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 + "opacity_embedding.weight"].dtype == torch.bfloat16 + # Untouched keys pass through. + assert "unrelated.weight" in padded diff --git a/scripts/static_check_artifixer.sh b/scripts/static_check_artifixer.sh index ccb55c30e..1e143e4a1 100755 --- a/scripts/static_check_artifixer.sh +++ b/scripts/static_check_artifixer.sh @@ -14,11 +14,7 @@ REPO_DIR="$(git -C "${PWD}" rev-parse --show-toplevel)" cd "${REPO_DIR}" echo "[static-check] py_compile artifixer sources" -for f in \ - integrations/artifixer/artifixer/__init__.py \ - integrations/artifixer/artifixer/config.py \ - integrations/artifixer/artifixer/runner.py \ - integrations/artifixer/tests/test_smoke.py; do +for f in $(find integrations/artifixer -name '*.py' -not -path '*/__pycache__/*'); do python3 -m py_compile "$f" done @@ -57,11 +53,8 @@ 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 ( - "integrations/artifixer/artifixer/config.py", - "integrations/artifixer/artifixer/runner.py", -): - src = Path(source_path).read_text() +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): From 5c54f5bdf79c773154495a463e9ed88e87a40cb7 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 19:06:15 -0700 Subject: [PATCH 03/22] integrations/artifixer: add neighbor cross-attention third KV bank Phase 2.2: ArtifixerCrossAttention(CrossAttention) adds the parameters that drive the neighbor-frame KV bank in dreamfix (model_training/net/transformer.py L670-L699). Names are kept identical to the dreamfix layout so the merged ArtiFixer DMD safetensors only need the diffusers ``attn2 -> cross_attn`` regex remap (added in Phase 5) at Phase 5, not a per-key rename: cross_attn.add_k_proj.{weight,bias} Linear(inner_dim, inner_dim) cross_attn.add_v_proj.{weight,bias} Linear(inner_dim, inner_dim) (zero-init) cross_attn.norm_added_k.weight RMSNorm(inner_dim) Plus a separate ``attn_op_neighbor`` RingAttention op so the neighbor branch can later carry PRoPE-transformed q/k without entangling the text branch. ArtifixerBlock now swaps the inherited ``self.cross_attn`` for the ArtifixerCrossAttention variant after super().__init__(). The forward path is unchanged in this commit (inherited CrossAttention.forward ignores the neighbor branch when no neighbor context is supplied), so behavior is identical to Phase 2.1 until the pipeline-level neighbor wiring lands in Phase 3 and PRoPE in Phase 2.3. zero_pad_artifixer_keys extended to fill all 9 ArtiFixer-only keys per block (4 from Phase 2.1 + 5 from Phase 2.2), matching the 270 extras identified by dreamfix/scripts/dump_artifixer_param_names.py. Smoke tests now also assert the cross_attn neighbor projections have the correct shape and add_v_proj is zero-initialized. --- .../artifixer/artifixer/checkpoint.py | 28 +++--- .../artifixer/artifixer/network/__init__.py | 2 + .../artifixer/artifixer/network/block.py | 33 +++++-- .../artifixer/artifixer/network/cross_attn.py | 96 +++++++++++++++++++ integrations/artifixer/tests/test_smoke.py | 55 ++++++++++- 5 files changed, 195 insertions(+), 19 deletions(-) create mode 100644 integrations/artifixer/artifixer/network/cross_attn.py diff --git a/integrations/artifixer/artifixer/checkpoint.py b/integrations/artifixer/artifixer/checkpoint.py index f44d47bd5..7c59d6497 100644 --- a/integrations/artifixer/artifixer/checkpoint.py +++ b/integrations/artifixer/artifixer/checkpoint.py @@ -39,18 +39,6 @@ from artifixer.network.dit import artifixer_embedding_dims -# Suffixes of the keys that the dreamfix ``ArtifixerTransformerBlock`` adds -# per transformer block. Phase 2.1 covers the opacity + camera MLPs; Phase -# 2.2 will add ``attn2.add_k_proj``, ``attn2.add_v_proj``, -# ``attn2.norm_added_k`` to this list. -_PHASE2_1_PER_BLOCK_SUFFIXES: tuple[str, ...] = ( - "opacity_embedding.weight", - "opacity_embedding.bias", - "camera_embedding.weight", - "camera_embedding.bias", -) - - def zero_pad_artifixer_keys( *, num_layers: int, @@ -66,6 +54,17 @@ def zero_pad_artifixer_keys( 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 identified by ``dreamfix/scripts/dump_artifixer_param_names.py``): + + * Phase 2.1 — 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,) + * Phase 2.2 — 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). @@ -81,6 +80,11 @@ def zero_pad_artifixer_keys( 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) diff --git a/integrations/artifixer/artifixer/network/__init__.py b/integrations/artifixer/artifixer/network/__init__.py index c6d63cbd9..3e0a0ec1a 100644 --- a/integrations/artifixer/artifixer/network/__init__.py +++ b/integrations/artifixer/artifixer/network/__init__.py @@ -14,6 +14,7 @@ # limitations under the License. from artifixer.network.block import ArtifixerBlock +from artifixer.network.cross_attn import ArtifixerCrossAttention from artifixer.network.dit import ( ArtifixerDiTNetwork, ArtifixerDiTNetwork1pt3BConfig, @@ -23,6 +24,7 @@ __all__ = [ "ArtifixerBlock", + "ArtifixerCrossAttention", "ArtifixerDiTNetwork", "ArtifixerDiTNetwork1pt3BConfig", "ArtifixerDiTNetworkConfig", diff --git a/integrations/artifixer/artifixer/network/block.py b/integrations/artifixer/artifixer/network/block.py index 0f98e5c25..9aa0db36d 100644 --- a/integrations/artifixer/artifixer/network/block.py +++ b/integrations/artifixer/artifixer/network/block.py @@ -13,19 +13,29 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ArtiFixer transformer block with opacity + camera-ray conditioning MLPs. +"""ArtiFixer transformer block with opacity + camera + neighbor extensions. Mirrors ``ArtifixerTransformerBlock`` in the dreamfix reference -(``model_training/net/transformer.py`` L617-L767). Two ``nn.Linear`` heads -project per-token opacity and Plucker-camera-ray features into the -transformer hidden size, and their outputs are added to the AdaLN-normed -hidden states **before** self-attention. Both heads are zero-initialized so -loading a vanilla Wan checkpoint leaves the block behavior unchanged. +(``model_training/net/transformer.py`` L617-L767). Per-block extensions on +top of :class:`Block`: + + * **Phase 2.1** — 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. + + * **Phase 2.2** — ``cross_attn`` is replaced with + :class:`ArtifixerCrossAttention`, which carries ``add_k_proj`` / + ``add_v_proj`` / ``norm_added_k`` and a separate ``attn_op_neighbor`` + RingAttention op for a future neighbor-frame KV branch. The forward + path is unchanged (parent ``CrossAttention.forward`` is called); the + pipeline-level neighbor wiring lands in Phase 3. """ from __future__ import annotations import torch +from artifixer.network.cross_attn import ArtifixerCrossAttention from torch import Tensor, nn from flashdreams.recipes.wan.transformer.impl.modules import Block, BlockCache @@ -54,6 +64,17 @@ def __init__( 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 + # future neighbor-frame KV branch (Phase 3 wiring). + 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) diff --git a/integrations/artifixer/artifixer/network/cross_attn.py b/integrations/artifixer/artifixer/network/cross_attn.py new file mode 100644 index 000000000..3d19750a4 --- /dev/null +++ b/integrations/artifixer/artifixer/network/cross_attn.py @@ -0,0 +1,96 @@ +# 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. + +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 used in dreamfix +(``model_training/net/transformer.py`` L670-L699): + + - ``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 so the neighbor + branch can later carry PRoPE-transformed q/k without entangling + the text branch + +Phase 2.2 only adds these parameters. Forward currently inherits from +``CrossAttention`` and does *not* yet wire the neighbor branch: callers +that do not pass a neighbor context observe identical behavior. The +wiring lands with the pipeline changes in Phase 3, and PRoPE q/k/v/o +transforms layer on in Phase 2.3. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +from torch import nn + +from flashdreams.core.attention.ring import RingAttention +from flashdreams.recipes.wan.transformer.impl.modules import CrossAttention + + +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 dreamfix (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 (added in Phase 5), 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 dreamfix transformer.py + # L687-L688). 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") + + 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: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Project neighbor ``context`` into K/V tensors. + + Args: + context: Neighbor context, shape ``[..., L_neighbor, context_dim]``. + + Returns: + ``(k, v)`` reshaped to ``[batch_size, L, n_heads, head_dim]``. + """ + 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 k, v diff --git a/integrations/artifixer/tests/test_smoke.py b/integrations/artifixer/tests/test_smoke.py index ea38abee5..9411773f0 100644 --- a/integrations/artifixer/tests/test_smoke.py +++ b/integrations/artifixer/tests/test_smoke.py @@ -145,7 +145,15 @@ def test_artifixer_recipe_uses_artifixer_network() -> None: def test_zero_pad_state_dict_transform_fills_missing_keys() -> None: - """Phase 2.1: the state_dict transform pads vanilla-Wan checkpoints.""" + """Phase 2.1/2.2: 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 import artifixer_embedding_dims @@ -162,6 +170,51 @@ def test_zero_pad_state_dict_transform_fills_missing_keys() -> None: 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: + """Phase 2.2: ArtifixerBlock's cross_attn carries the neighbor branch.""" + import torch + from artifixer.network import ( + ArtifixerBlock, + ArtifixerCrossAttention, + 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}" + + # Phase 2.2 zero-init contract: add_v_proj is the gate that keeps the + # neighbor branch contribution zero at load time, matching dreamfix + # transformer.py L687-688. + assert torch.all(block.cross_attn.add_v_proj.weight == 0) + assert torch.all(block.cross_attn.add_v_proj.bias == 0) From e6b7714f285bacd5b82b919f07d6c4d59c305a9e Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 19:09:12 -0700 Subject: [PATCH 04/22] integrations/artifixer: port PRoPE attention + parity test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 2.3: verbatim port of dreamfix ``model_training/net/prope.py`` (paper "Cameras as Relative Positional Encoding", arXiv 2507.10496). The module is self-contained — the only dependency, ``model_training/utils/pose_utils.invert_SE3``, is inlined as the private helper ``_invert_SE3`` so there are no cross-repo runtime imports. PRoPE applies block-diagonal SE(3) projection-matrix + 2D RoPE transforms to q/k/v/o on cross-attention. ``apply_to_o`` is the inverse of the RoPE legs in ``apply_to_q`` / ``apply_to_kv``. Used twice per ArtifixerBlock in dreamfix: - prope_cross_attn_src on target camera (w2cs / Ks) - prope_cross_attn_tgt on neighbor cameras (neighbor_w2cs / neighbor_Ks) This commit ports the module; the wiring inside ArtifixerCrossAttention.forward (calling ``_apply_to_q``, ``_apply_to_kv``, ``_apply_to_o`` around the RingAttention SDPA call) lands in Phase 2.4 with the network forward changes. Tests: - tests/test_prope_parity.py:test_prope_internal_consistency exercises the identity-camera invariant ``apply_to_o(apply_to_q(x)) == x`` exactly at fp64. - tests/test_prope_parity.py:test_prope_apply_to_q_changes_input_for_nontrivial_cameras sanity-checks that real cameras do produce a non-identity transform. - tests/test_prope_parity.py:test_prope_matches_dreamfix_reference is a numerical-parity test against dreamfix at fp64. Skipped when DREAMFIX_REPO_ROOT is not set / dreamfix is not on PYTHONPATH. The intent is to gate the architecture changes on bit-identical math the moment a dual env is available. PRoPE module is exported from artifixer.network for convenience. Static check passes (12 flashdreams imports resolve, py_compile clean). --- .../artifixer/artifixer/network/__init__.py | 2 + .../artifixer/artifixer/network/prope.py | 310 ++++++++++++++++++ .../artifixer/tests/test_prope_parity.py | 187 +++++++++++ 3 files changed, 499 insertions(+) create mode 100644 integrations/artifixer/artifixer/network/prope.py create mode 100644 integrations/artifixer/tests/test_prope_parity.py diff --git a/integrations/artifixer/artifixer/network/__init__.py b/integrations/artifixer/artifixer/network/__init__.py index 3e0a0ec1a..e7b1dad27 100644 --- a/integrations/artifixer/artifixer/network/__init__.py +++ b/integrations/artifixer/artifixer/network/__init__.py @@ -21,6 +21,7 @@ ArtifixerDiTNetworkConfig, artifixer_embedding_dims, ) +from artifixer.network.prope import PropeDotProductAttention __all__ = [ "ArtifixerBlock", @@ -28,5 +29,6 @@ "ArtifixerDiTNetwork", "ArtifixerDiTNetwork1pt3BConfig", "ArtifixerDiTNetworkConfig", + "PropeDotProductAttention", "artifixer_embedding_dims", ] diff --git a/integrations/artifixer/artifixer/network/prope.py b/integrations/artifixer/artifixer/network/prope.py new file mode 100644 index 000000000..1c8f9153c --- /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. + +Verbatim port of dreamfix ``model_training/net/prope.py`` (matching paper +"Cameras as Relative Positional Encoding", arXiv 2507.10496). The math is +unchanged; the only difference is that ``invert_SE3`` is inlined here so +this module is self-contained. + +Usage for cross-attention follows the upstream docstring:: + + 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 dreamfix 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. + + Mirrors ``dreamfix/model_training/utils/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/tests/test_prope_parity.py b/integrations/artifixer/tests/test_prope_parity.py new file mode 100644 index 000000000..7aae50392 --- /dev/null +++ b/integrations/artifixer/tests/test_prope_parity.py @@ -0,0 +1,187 @@ +# 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 dreamfix reference at ``model_training/net/prope.py``. When dreamfix is +not importable (the CI image only ships 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_dreamfix_reference`` is skipped unless + ``model_training.net.prope`` can be imported (set + ``DREAMFIX_REPO_ROOT`` to its repo root to enable). + +The intent is to catch any regression in the verbatim port the moment a +GPU + dreamfix 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: + """Use fp64 by default in this file: PRoPE math is bit-exact at fp64.""" + 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.float64, + 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`` exactly. + """ + inp = _build_inputs() + + # Replace viewmats / Ks with identity so the projmat legs cancel. + 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=1e-10, rtol=1e-10) + + +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_dreamfix_reference() -> object | None: + """Return dreamfix's PropeDotProductAttention if importable.""" + repo_root = os.environ.get("DREAMFIX_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_dreamfix_reference() is None, + reason=( + "dreamfix not importable. Set DREAMFIX_REPO_ROOT to the dreamfix " + "checkout, or run from an env that already has it on sys.path." + ), +) +def test_prope_matches_dreamfix_reference() -> None: + """Numerical parity vs dreamfix at fp64. + + Asserts ``apply_to_q``, ``apply_to_kv``, ``apply_to_o`` produce + bit-identical outputs to the reference within ``atol=1e-10`` + (fp64 round-off floor). + """ + RefPRoPE = _try_import_dreamfix_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-10, rtol=1e-10, 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-10, rtol=1e-10, 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-10, rtol=1e-10, msg=f"apply_to_o on {name}" + ) From 731a3efaf7397bc0b3d66af490f09c6c9b8c0623 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 19:23:09 -0700 Subject: [PATCH 05/22] integrations/artifixer: keep network __init__ lean; fix PRoPE parity tests Two fixes: 1. ``artifixer/network/__init__.py`` was eagerly importing every submodule, which transitively pulled in ``flashdreams.core.attention.ring`` and the rest of the flashdreams tree just to use PRoPE. The PRoPE unit test cannot run without the full flashdreams env in place because of this. Made the __init__ empty (mirroring dreamfix's ``model_training/net/__init__.py``) so ``from artifixer.network.prope import ...`` only needs torch. The recipe and tests now import submodules explicitly. 2. ``tests/test_prope_parity.py``: - Default fixture dtype switched from fp64 -> fp32. PRoPE's ``_rope_precompute_coeffs`` does an implicit int64 -> fp32 promotion (``torch.arange(...) / num_freqs``) so the RoPE coefficients are always at fp32 precision regardless of input dtype. dreamfix's ``_lift_K`` also hard-codes float32 output (no ``dtype=``), so any cross-implementation comparison must use fp32 inputs to avoid an einsum dtype crash on the reference. - Tolerances relaxed to fp32 (atol/rtol=1e-6 for parity, 5e-6 for the round-trip identity check). Our port already carries the ``dtype=Ks.dtype`` fix in ``_lift_K``, so it is slightly more dtype-robust than upstream at fp64, but bit-identical at fp32. Verified: 3/3 tests pass in the artifixer-cuda12 container with DREAMFIX_REPO_ROOT pointing at the dreamfix checkout: - test_prope_internal_consistency PASSED - test_prope_apply_to_q_changes_input_... PASSED - test_prope_matches_dreamfix_reference PASSED --- integrations/artifixer/artifixer/config.py | 2 +- .../artifixer/artifixer/network/__init__.py | 28 +++++---------- .../artifixer/tests/test_prope_parity.py | 36 +++++++++++++------ integrations/artifixer/tests/test_smoke.py | 15 ++++---- 4 files changed, 42 insertions(+), 39 deletions(-) diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index 009c83251..bdaedc582 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -36,7 +36,7 @@ import torch from artifixer.checkpoint import zero_pad_artifixer_keys -from artifixer.network import ArtifixerDiTNetwork1pt3BConfig +from artifixer.network.dit import ArtifixerDiTNetwork1pt3BConfig from artifixer.runner import ArtifixerDmdT2VRunnerConfig from flashdreams.infra.diffusion.model import DiffusionModelConfig diff --git a/integrations/artifixer/artifixer/network/__init__.py b/integrations/artifixer/artifixer/network/__init__.py index e7b1dad27..bf63d7456 100644 --- a/integrations/artifixer/artifixer/network/__init__.py +++ b/integrations/artifixer/artifixer/network/__init__.py @@ -13,22 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from artifixer.network.block import ArtifixerBlock -from artifixer.network.cross_attn import ArtifixerCrossAttention -from artifixer.network.dit import ( - ArtifixerDiTNetwork, - ArtifixerDiTNetwork1pt3BConfig, - ArtifixerDiTNetworkConfig, - artifixer_embedding_dims, -) -from artifixer.network.prope import PropeDotProductAttention - -__all__ = [ - "ArtifixerBlock", - "ArtifixerCrossAttention", - "ArtifixerDiTNetwork", - "ArtifixerDiTNetwork1pt3BConfig", - "ArtifixerDiTNetworkConfig", - "PropeDotProductAttention", - "artifixer_embedding_dims", -] +# 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 +# +# Mirrors dreamfix's ``model_training/net/__init__.py`` (also empty) so the +# PRoPE module is testable with only ``torch`` installed. diff --git a/integrations/artifixer/tests/test_prope_parity.py b/integrations/artifixer/tests/test_prope_parity.py index 7aae50392..81b5be27d 100644 --- a/integrations/artifixer/tests/test_prope_parity.py +++ b/integrations/artifixer/tests/test_prope_parity.py @@ -45,7 +45,16 @@ @pytest.fixture(autouse=True) def _deterministic_torch() -> None: - """Use fp64 by default in this file: PRoPE math is bit-exact at fp64.""" + """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 dreamfix + reference has the same property. Tests therefore run at fp32 with fp32- + appropriate tolerances rather than fp64. dreamfix'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) @@ -56,7 +65,7 @@ def _build_inputs( patches_y: int = 4, num_heads: int = 4, head_dim: int = 32, - dtype: torch.dtype = torch.float64, + dtype: torch.dtype = torch.float32, device: torch.device | str = "cpu", ) -> dict[str, torch.Tensor]: """Build a deterministic fixture for PRoPE parity tests.""" @@ -107,11 +116,11 @@ def test_prope_internal_consistency() -> None: 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`` exactly. + ``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() - # Replace viewmats / Ks with identity so the projmat legs cancel. inp["viewmats"] = ( torch.eye(4, dtype=inp["viewmats"].dtype).expand_as(inp["viewmats"]).clone() ) @@ -122,7 +131,7 @@ def test_prope_internal_consistency() -> None: 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=1e-10, rtol=1e-10) + 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: @@ -154,11 +163,16 @@ def _try_import_dreamfix_reference() -> object | None: ), ) def test_prope_matches_dreamfix_reference() -> None: - """Numerical parity vs dreamfix at fp64. + """Numerical parity vs dreamfix at fp32. Asserts ``apply_to_q``, ``apply_to_kv``, ``apply_to_o`` produce - bit-identical outputs to the reference within ``atol=1e-10`` - (fp64 round-off floor). + bit-identical outputs to the reference within fp32 round-off. + + Inputs are fp32 because dreamfix's ``_lift_K`` hard-codes float32 + output (no ``dtype=`` in ``torch.zeros``), so feeding fp64 viewmats + crashes its einsum. Our port carries the dreamfix dtype bug fix + (passes ``dtype=Ks.dtype``), making it slightly more dtype-robust; + at fp32 the two are bit-identical. """ RefPRoPE = _try_import_dreamfix_reference() assert RefPRoPE is not None @@ -171,17 +185,17 @@ def test_prope_matches_dreamfix_reference() -> None: ours_q = ours._apply_to_q(x) ref_q = ref._apply_to_q(x) torch.testing.assert_close( - ours_q, ref_q, atol=1e-10, rtol=1e-10, msg=f"apply_to_q on {name}" + 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-10, rtol=1e-10, msg=f"apply_to_kv on {name}" + 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-10, rtol=1e-10, msg=f"apply_to_o on {name}" + 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 index 9411773f0..9c3634722 100644 --- a/integrations/artifixer/tests/test_smoke.py +++ b/integrations/artifixer/tests/test_smoke.py @@ -109,7 +109,8 @@ def test_artifixer_hyperparams_match_dreamfix_stage3() -> None: def test_artifixer_block_has_opacity_and_camera_mlps() -> None: """Phase 2.1: ArtifixerBlock instances carry opacity + camera MLPs.""" import torch - from artifixer.network import ArtifixerBlock, artifixer_embedding_dims + 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( @@ -135,7 +136,7 @@ def test_artifixer_block_has_opacity_and_camera_mlps() -> None: def test_artifixer_recipe_uses_artifixer_network() -> None: """Phase 2.1: the shipped recipe wires up ArtifixerDiTNetwork.""" - from artifixer.network import ArtifixerDiTNetwork1pt3BConfig + from artifixer.network.dit import ArtifixerDiTNetwork1pt3BConfig cfg = RUNNER_CONFIGS["artifixer-dmd-wan2.1-t2v-1.3b"] network_cfg = cfg.pipeline.diffusion_model.transformer.network @@ -156,7 +157,7 @@ def test_zero_pad_state_dict_transform_fills_missing_keys() -> None: """ import torch from artifixer.checkpoint import zero_pad_artifixer_keys - from artifixer.network import artifixer_embedding_dims + 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 @@ -185,11 +186,9 @@ def test_zero_pad_state_dict_transform_fills_missing_keys() -> None: def test_artifixer_block_has_neighbor_cross_attn_projections() -> None: """Phase 2.2: ArtifixerBlock's cross_attn carries the neighbor branch.""" import torch - from artifixer.network import ( - ArtifixerBlock, - ArtifixerCrossAttention, - artifixer_embedding_dims, - ) + 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( From 8e86c61488224228272290d542f735ffbc1850a7 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 19:37:40 -0700 Subject: [PATCH 06/22] integrations/artifixer: wire PRoPE neighbor branch through cross-attn + block + network Phase 2.4: complete the architecture. ArtifixerCrossAttention, ArtifixerBlock, and ArtifixerDiTNetwork now plumb the per-block opacity + camera-ray MLPs, the neighbor-frame KV bank, and PRoPE q/k/v/o transforms end-to-end. With every extra ``None`` / ``False`` the path is a no-op extension of vanilla Wan, matching Phase 2.1 / 2.2 behavior. artifixer/network/cross_attn.py: - Added ``initialize_neighbor_cache(context)`` which calls ``compute_kv_neighbor`` and stores a per-module ``neighbor_kv_cache: BlockKVCache | None``. ``context=None`` clears it so subsequent forward passes skip the neighbor branch. - Added ``forward(x, kv_cache, *, prope_src=None, prope_tgt=None, ignore_neighbors=False)``. When the neighbor cache is set and PRoPE modules are supplied, adds the PRoPE-modulated neighbor branch on top of the text (+ optional I2V image) branch (matches dreamfix transformer.py L833-L940). PRoPE math runs at fp32 with a ``.float() / .to(query.dtype)`` round trip. - ``compute_kv_neighbor`` now returns a ``BlockKVCache`` (was a ``(k, v)`` tuple) -- consistent with the I2V ``compute_kv_image`` pattern and ready to be reused across denoise steps. artifixer/network/block.py: ``ArtifixerBlock.forward`` accepts ``prope_src`` / ``prope_tgt`` / ``ignore_neighbors`` kwargs and forwards them to ``cross_attn``. The neighbor KV cache itself lives on the cross_attn module rather than in ``block_extra_kwargs`` so that one ``block_extra_kwargs`` dict can be passed identically to every block by ``WanDiTNetwork.forward`` (while each block still has its own cache). artifixer/network/dit.py: ``ArtifixerDiTNetwork.initialize_neighbor_kv_caches(context)`` walks every block and propagates the (optional) neighbor latent context to each ``cross_attn``. Called once per rollout by the pipeline (Phase 3). tests/test_smoke.py: heavy imports (mediapy, full flashdreams) deferred into the test functions that need them; lighter tests (``test_compute_kv_neighbor_and_cache_init``, the PRoPE parity ones) now collect in torch-only environments without pulling boto3 / mediapy. Added ``test_compute_kv_neighbor_and_cache_init`` covering the cache toggle. Verified inside the artifixer container with dreamfix on PYTHONPATH: test_prope_internal_consistency PASSED test_prope_apply_to_q_changes_input_for_... PASSED test_prope_matches_dreamfix_reference PASSED (fp32 atol 1e-6) test_compute_kv_neighbor_and_cache_init PASSED scripts/run_prope_parity.sh also installs boto3 + tyro now and includes the new cache test in the run. --- .../artifixer/artifixer/network/block.py | 27 +++- .../artifixer/artifixer/network/cross_attn.py | 153 +++++++++++++++--- .../artifixer/artifixer/network/dit.py | 19 +++ integrations/artifixer/tests/test_smoke.py | 83 ++++++++-- scripts/run_prope_parity.sh | 39 +++++ 5 files changed, 282 insertions(+), 39 deletions(-) create mode 100755 scripts/run_prope_parity.sh diff --git a/integrations/artifixer/artifixer/network/block.py b/integrations/artifixer/artifixer/network/block.py index 9aa0db36d..07345c846 100644 --- a/integrations/artifixer/artifixer/network/block.py +++ b/integrations/artifixer/artifixer/network/block.py @@ -34,6 +34,8 @@ from __future__ import annotations +from typing import Any + import torch from artifixer.network.cross_attn import ArtifixerCrossAttention from torch import Tensor, nn @@ -93,14 +95,26 @@ def forward( # type: ignore[override] 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. - ``opacity_extra`` and ``camera_extra`` are added to the AdaLN-normed - hidden states before self-attention (matches dreamfix - ``ArtifixerTransformerBlock.forward`` L763-767). When both are - ``None`` (Phase 1 text-only path), the block is identical to the - underlying ``Block``. + Extras (Phase 2.1 + 2.4): + + - ``opacity_extra`` / ``camera_extra``: per-token opacity and + Plucker-camera-ray features added to the AdaLN-normed hidden + states before self-attention (matches dreamfix + ``ArtifixerTransformerBlock.forward`` L763-L767). + - ``prope_src`` / ``prope_tgt`` / ``ignore_neighbors``: forwarded + to :class:`ArtifixerCrossAttention.forward` to drive the PRoPE + neighbor branch (matches dreamfix L795-L807). 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() " @@ -123,6 +137,9 @@ def forward( # type: ignore[override] x = x + self.cross_attn( self.norm3(x), kv_cache=cache.cross_attn, + prope_src=prope_src, + prope_tgt=prope_tgt, + ignore_neighbors=ignore_neighbors, ) y = self.norm2(x) * (1 + e_chunks[4]) + e_chunks[3] y = self.ffn(y) diff --git a/integrations/artifixer/artifixer/network/cross_attn.py b/integrations/artifixer/artifixer/network/cross_attn.py index 3d19750a4..e2323cc56 100644 --- a/integrations/artifixer/artifixer/network/cross_attn.py +++ b/integrations/artifixer/artifixer/network/cross_attn.py @@ -13,25 +13,35 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ArtiFixer cross-attention with a neighbor-frame KV bank. +"""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 used in dreamfix -(``model_training/net/transformer.py`` L670-L699): +(``model_training/net/transformer.py`` L670-L699 + the +``ArtifixerCrossAttnProcessor`` at L833-L940 that drives the forward): - ``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 so the neighbor - branch can later carry PRoPE-transformed q/k without entangling - the text branch - -Phase 2.2 only adds these parameters. Forward currently inherits from -``CrossAttention`` and does *not* yet wire the neighbor branch: callers -that do not pass a neighbor context observe identical behavior. The -wiring lands with the pipeline changes in Phase 3, and PRoPE q/k/v/o -transforms layer on in Phase 2.3. + - ``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`` (Phase 1 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 dreamfix L893-L916), since +``_rope_precompute_coeffs`` does an int64 -> fp32 promotion internally. """ from __future__ import annotations @@ -40,10 +50,14 @@ from typing import Any import torch -from torch import nn +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 +from flashdreams.recipes.wan.transformer.impl.modules import ( + CrossAttention, + CrossAttnCache, +) class ArtifixerCrossAttention(CrossAttention): @@ -70,21 +84,25 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: 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: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """Project neighbor ``context`` into K/V tensors. + 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: - ``(k, v)`` reshaped to ``[batch_size, L, n_heads, head_dim]``. + ``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) @@ -93,4 +111,101 @@ def compute_kv_neighbor( 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 k, v + 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 (Phase 1 / 2 text-only). + """ + 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 dreamfix + (transformer.py L936). + """ + 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 + # dreamfix transformer.py L893-L916. + 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 index 79f373250..d228ac8fa 100644 --- a/integrations/artifixer/artifixer/network/dit.py +++ b/integrations/artifixer/artifixer/network/dit.py @@ -20,6 +20,7 @@ from dataclasses import dataclass, field from artifixer.network.block import ArtifixerBlock +from torch import Tensor from flashdreams.recipes.wan.transformer.impl.network import ( WanDiTNetwork, @@ -98,3 +99,21 @@ def _build_block(self, layer_idx: int) -> ArtifixerBlock: 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/tests/test_smoke.py b/integrations/artifixer/tests/test_smoke.py index 9c3634722..04a8ede67 100644 --- a/integrations/artifixer/tests/test_smoke.py +++ b/integrations/artifixer/tests/test_smoke.py @@ -13,32 +13,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Cheap import-time checks for the ``artifixer`` plugin.""" +"""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 -from typing import cast import pytest -import tomllib -from artifixer import config as config_mod -from artifixer.config import RUNNER_CONFIGS - -from flashdreams.infra.runner import RunnerConfig 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" + 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() + 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}" @@ -46,18 +53,26 @@ def test_runner_name_mirrors_pipeline_recipe_name() -> None: def test_runners_have_descriptions() -> None: empty = [ - slug for slug, cfg in RUNNER_CONFIGS.items() if not cfg.description.strip() + 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) + module_slugs = set(runner_configs) assert declared_slugs == module_slugs, ( f"entry-point slugs ({sorted(declared_slugs)}) " f"!= module runners ({sorted(module_slugs)})" @@ -86,15 +101,16 @@ def test_entry_points_discoverable_when_installed() -> None: 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") - assert discovered == set(RUNNER_CONFIGS), ( + runner_configs = _runner_configs() + assert discovered == set(runner_configs), ( f"discovered slugs ({sorted(discovered)}) != " - f"plugin runners ({sorted(RUNNER_CONFIGS)})" + f"plugin runners ({sorted(runner_configs)})" ) def test_artifixer_hyperparams_match_dreamfix_stage3() -> None: """Sanity check: AR/scheduler knobs match the stage-3 DMD training config.""" - cfg = RUNNER_CONFIGS["artifixer-dmd-wan2.1-t2v-1.3b"] + cfg = _runner_configs()["artifixer-dmd-wan2.1-t2v-1.3b"] tcfg = cfg.pipeline.diffusion_model.transformer scfg = cfg.pipeline.diffusion_model.scheduler @@ -138,7 +154,7 @@ def test_artifixer_recipe_uses_artifixer_network() -> None: """Phase 2.1: the shipped recipe wires up ArtifixerDiTNetwork.""" from artifixer.network.dit import ArtifixerDiTNetwork1pt3BConfig - cfg = RUNNER_CONFIGS["artifixer-dmd-wan2.1-t2v-1.3b"] + 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" @@ -217,3 +233,40 @@ def test_artifixer_block_has_neighbor_cross_attn_projections() -> None: # transformer.py L687-688. 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: + """Phase 2.4: 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/scripts/run_prope_parity.sh b/scripts/run_prope_parity.sh new file mode 100755 index 000000000..3f4247364 --- /dev/null +++ b/scripts/run_prope_parity.sh @@ -0,0 +1,39 @@ +#!/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 dreamfix +# reference at model_training/net/prope.py. +# +# The artifixer container has torch installed at the system Python level +# and the dreamfix 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)}" +DREAMFIX_REPO_ROOT="${DREAMFIX_REPO_ROOT:-/lustre/fsw/portfolios/nvr/users/rdelutio/code/dreamfix}" +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:${DREAMFIX_REPO_ROOT}:\${PYTHONPATH:-} && \ +export DREAMFIX_REPO_ROOT='${DREAMFIX_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_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}" From 2c66def1d0157e2829ab5feef1fd74446087fa76 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 19:40:39 -0700 Subject: [PATCH 07/22] integrations/artifixer: introduce ArtifixerWanTransformer with PRoPE modules Phase 3.1: add a Wan21Transformer subclass that owns the two PropeDotProductAttention modules (target/source camera + neighbor camera) used by ArtifixerCrossAttention. The modules are constructed once at transformer-init time at the correct ``head_dim`` and dtype; their per-rollout ``_precompute_and_cache_apply_fns(...)`` calls will be driven by the pipeline (Phase 3.2 / 3.3). Wiring: ``artifixer.config.PIPELINE_ARTIFIXER_DMD_T2V_1PT3B`` now uses ``ArtifixerWanTransformerConfig`` in place of vanilla ``Wan21TransformerConfig``. No behavioral change yet (the per-rollout state plumbing lands next); the PRoPE modules sit dormant. Static check + linter clean. --- integrations/artifixer/artifixer/config.py | 9 +- .../artifixer/artifixer/transformer.py | 84 +++++++++++++++++++ 2 files changed, 87 insertions(+), 6 deletions(-) create mode 100644 integrations/artifixer/artifixer/transformer.py diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index bdaedc582..3bac28495 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -38,15 +38,12 @@ from artifixer.checkpoint import zero_pad_artifixer_keys from artifixer.network.dit import ArtifixerDiTNetwork1pt3BConfig 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 ( - Wan21TransformerConfig, - WanInferencePipelineConfig, - WanVAEDecoderConfig, -) +from flashdreams.recipes.wan import WanInferencePipelineConfig, WanVAEDecoderConfig # Mirrors the dreamfix stage-3 run config (see # wandb/.../run-*/files/config.yaml in @@ -83,7 +80,7 @@ decoder=WanVAEDecoderConfig(), diffusion_model=DiffusionModelConfig( seed=42, - transformer=Wan21TransformerConfig( + transformer=ArtifixerWanTransformerConfig( network=_BASE_NETWORK_CONFIG, dtype=_BASE_TRANSFORMER_DTYPE, checkpoint_path=BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH, diff --git a/integrations/artifixer/artifixer/transformer.py b/integrations/artifixer/artifixer/transformer.py new file mode 100644 index 000000000..d40b5fea1 --- /dev/null +++ b/integrations/artifixer/artifixer/transformer.py @@ -0,0 +1,84 @@ +# 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 that state into ``predict_flow`` -> ``WanDiTNetwork.forward`` + via ``network_extra_kwargs`` (Phase 3.3 wires the slicing). + +This commit lands the static structure only -- the per-rollout state and +the ``predict_flow`` override that slices it per AR chunk are added in +later Phase 3 sub-commits. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from artifixer.network.prope import PropeDotProductAttention + +from flashdreams.recipes.wan import Wan21Transformer, Wan21TransformerConfig + + +@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 dreamfix dtype convention. + 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) + + +__all__ = [ + "ArtifixerWanTransformer", + "ArtifixerWanTransformerConfig", +] From 9142c9b5b9716baa5e7d7914ec495f02387a0da1 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 19:46:55 -0700 Subject: [PATCH 08/22] integrations/artifixer: opacity + camera-ray patchification helpers Phase 3.2: port the patchification logic from dreamfix ``model_training/net/transformer.py`` L309-L341 as two stand-alone functions: * ``patchify_opacity`` turns ``(B, T, H, W)`` per-pixel alpha into ``(B, L, opacity_embedding_dim)`` per-token features. On the first AR chunk (``frame_offset == 0``) the first input frame is left-padded by 3 copies so the rearrange treats every latent frame uniformly under Wan VAE's ``1 + 4`` temporal layout. * ``patchify_camera_rays`` turns ``(B, T, H, W, 6)`` Plucker rays into ``(B, L, camera_embedding_dim)``. Branches on whether the temporal axis is already at the post-patch latent rate (``hidden_post_patch_t`` == camera_rays.shape[1]) or at the input rate, matching dreamfix exactly. The input-rate branch carries an extra ``vae_t`` multiplier in the per-token feature and is unused at inference (``kv_cache_pipeline.py`` L213 slices at the latent rate); we port both for parity with the dreamfix forward. tests/test_patches.py covers shape correctness on both branches plus a constant-field invariance check. 9/9 tests pass (4 PRoPE + 5 patches) inside the artifixer container with dreamfix on PYTHONPATH. scripts/run_prope_parity.sh now also runs the patches tests. --- .../artifixer/artifixer/network/patches.py | 131 ++++++++++++++++++ integrations/artifixer/tests/test_patches.py | 129 +++++++++++++++++ scripts/run_prope_parity.sh | 2 +- 3 files changed, 261 insertions(+), 1 deletion(-) create mode 100644 integrations/artifixer/artifixer/network/patches.py create mode 100644 integrations/artifixer/tests/test_patches.py diff --git a/integrations/artifixer/artifixer/network/patches.py b/integrations/artifixer/artifixer/network/patches.py new file mode 100644 index 000000000..ac82679e5 --- /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 ``ArtifixerTransformer.forward`` L309-L341 in dreamfix: + + * 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 dreamfix L309-L311). + + 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 dreamfix'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/tests/test_patches.py b/integrations/artifixer/tests/test_patches.py new file mode 100644 index 000000000..7c5a41211 --- /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 dreamfix (transformer.py L330-L340) 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 (see + ``model_training/pipeline/kv_cache_pipeline.py`` L213), so this branch + is unused. We still cover it for parity with the dreamfix 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/scripts/run_prope_parity.sh b/scripts/run_prope_parity.sh index 3f4247364..80212dac7 100755 --- a/scripts/run_prope_parity.sh +++ b/scripts/run_prope_parity.sh @@ -29,7 +29,7 @@ cd ${REPO_DIR} && \ export PYTHONPATH=${REPO_DIR}/integrations/artifixer:${REPO_DIR}/flashdreams:${DREAMFIX_REPO_ROOT}:\${PYTHONPATH:-} && \ export DREAMFIX_REPO_ROOT='${DREAMFIX_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_smoke.py::test_compute_kv_neighbor_and_cache_init" +python3 -m pytest -v integrations/artifixer/tests/test_prope_parity.py integrations/artifixer/tests/test_patches.py integrations/artifixer/tests/test_smoke.py::test_compute_kv_neighbor_and_cache_init" echo "[slurm] command: ${COMMAND}" From 4347539d901466361b5dcaeaa1b340f4ecd7885a Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 21:57:49 -0700 Subject: [PATCH 09/22] integrations/artifixer: ArtifixerCtrl + predict_flow override Phase 3.3: define :class:`ArtifixerCtrl` as a per-AR-chunk conditioning payload carrying patchified opacity / camera features and the ``ignore_neighbors`` flag, and override :meth:`ArtifixerWanTransformer.predict_flow` to repackage those into ``network_extra_kwargs`` so they reach every transformer block via ``WanDiTNetwork.forward(... **block_extra_kwargs)``. The override also forwards ``self.prope_cross_attn_src`` and ``self.prope_cross_attn_tgt`` so ``ArtifixerCrossAttention.forward`` gets the PRoPE modules uniformly. The pipeline (Phase 3.5) owns the per-chunk slicing, patchification, and the per-chunk ``prope_cross_attn_src._precompute_and_cache_apply_fns(...)`` update; this commit just plumbs the kwargs. ``ArtifixerCtrl`` sets ``_is_patchified=True`` so the ``DiffusionModel.generate`` patchify-on-input dispatch (base.py L163-164) treats it as a no-op for our tensors. When the input is not an ``ArtifixerCtrl`` (e.g. pre-Phase-3 T2V smoke) the override falls through to the base ``predict_flow`` unchanged -- every conditioning kwarg in :class:`ArtifixerBlock` is optional. Static check + linter clean. No behavioral test yet -- exercising ``predict_flow`` end-to-end requires a GPU + the pipeline state from Phase 3.4 / 3.5. --- .../artifixer/artifixer/transformer.py | 95 ++++++++++++++++++- 1 file changed, 90 insertions(+), 5 deletions(-) diff --git a/integrations/artifixer/artifixer/transformer.py b/integrations/artifixer/artifixer/transformer.py index d40b5fea1..af174dfa6 100644 --- a/integrations/artifixer/artifixer/transformer.py +++ b/integrations/artifixer/artifixer/transformer.py @@ -23,21 +23,62 @@ - 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 that state into ``predict_flow`` -> ``WanDiTNetwork.forward`` - via ``network_extra_kwargs`` (Phase 3.3 wires the slicing). + - thread per-AR-chunk extras (patchified opacity / camera-ray features, + target-camera PRoPE precompute) into ``predict_flow`` -> + ``WanDiTNetwork.forward`` via an :class:`ArtifixerCtrl` payload. -This commit lands the static structure only -- the per-rollout state and -the ``predict_flow`` override that slices it per AR chunk are added in -later Phase 3 sub-commits. +The pipeline (Phase 3.5) 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 + (Phase 3.5); 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 (matches dreamfix L936). 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`` (base.py L163-164) is a no-op for us.""" @dataclass(kw_only=True) @@ -77,8 +118,52 @@ def __init__(self, config: ArtifixerWanTransformerConfig) -> None: 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. + + Reads patchified opacity/camera features and the ``ignore_neighbors`` + flag from :class:`ArtifixerCtrl` and forwards them to every block via + ``network_extra_kwargs``. Also passes the PRoPE source/target modules + themselves -- they are bound at construction time, but + :class:`ArtifixerCrossAttention.forward` expects them as forward args + so they show up uniformly under ``block_extra_kwargs``. + + When ``input`` is not an :class:`ArtifixerCtrl` (e.g. T2V 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): + network_extra_kwargs.setdefault("opacity_extra", input.opacity_extra) + network_extra_kwargs.setdefault("camera_extra", input.camera_extra) + network_extra_kwargs.setdefault("ignore_neighbors", input.ignore_neighbors) + network_extra_kwargs.setdefault("prope_src", self.prope_cross_attn_src) + network_extra_kwargs.setdefault("prope_tgt", self.prope_cross_attn_tgt) + # 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, + ) + __all__ = [ + "ArtifixerCtrl", "ArtifixerWanTransformer", "ArtifixerWanTransformerConfig", ] From 4fcd2756051df1b900eeee8bb1386ee992b1b48c Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:01:43 -0700 Subject: [PATCH 10/22] integrations/artifixer: opacity-weighted latent mix helper Phase 3.4: port the ``ArtifixerPipelineBase.prepare_latents`` mix from dreamfix (pipeline_base.py L98-L122) as a stand-alone, side-effect-free helper ``opacity_weighted_latent_mix``. The mix replaces the base ``DiffusionModel.generate`` initial-noise draw with a per-AR-chunk weighted blend:: latents = condition * opacity_lat + noise * (1 - opacity_lat) where ``opacity_lat`` is the per-pixel alpha max-pooled from the input resolution to the VAE latent grid. The pipeline (Phase 3.5) will pull the chunk's slice of the full-rollout opacity, sample noise via the scheduler's RNG (so CP-broadcast / determinism stay intact), and feed both to this helper. The first AR chunk left-pads the input opacity by 3 copies of the first frame to absorb the Wan VAE's ``1 + 4`` temporal layout. Later chunks already arrive at ``vae_t * t_lat`` input frames. tests/test_latent_mix.py: 5 CPU-only invariant tests (opacity=1 -> drop noise, opacity=0 -> drop condition, t_lat=1 edge case, max-pool semantics, non-first-chunk shape). All 14/14 tests now pass on the slurm node: PRoPE (3) + patches (5) + latent_mix (5) + cache init (1). --- .../artifixer/artifixer/latent_mix.py | 107 +++++++++++++++ .../artifixer/tests/test_latent_mix.py | 129 ++++++++++++++++++ scripts/run_prope_parity.sh | 2 +- 3 files changed, 237 insertions(+), 1 deletion(-) create mode 100644 integrations/artifixer/artifixer/latent_mix.py create mode 100644 integrations/artifixer/tests/test_latent_mix.py diff --git a/integrations/artifixer/artifixer/latent_mix.py b/integrations/artifixer/artifixer/latent_mix.py new file mode 100644 index 000000000..1588a1d22 --- /dev/null +++ b/integrations/artifixer/artifixer/latent_mix.py @@ -0,0 +1,107 @@ +# 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 ``ArtifixerPipelineBase.prepare_latents`` at +``dreamfix/model_training/pipeline/pipeline_base.py`` L98-L122. 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 (Phase 2.1) 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 dreamfix 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 dreamfix. + + 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/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/scripts/run_prope_parity.sh b/scripts/run_prope_parity.sh index 80212dac7..dac2aee3a 100755 --- a/scripts/run_prope_parity.sh +++ b/scripts/run_prope_parity.sh @@ -29,7 +29,7 @@ cd ${REPO_DIR} && \ export PYTHONPATH=${REPO_DIR}/integrations/artifixer:${REPO_DIR}/flashdreams:${DREAMFIX_REPO_ROOT}:\${PYTHONPATH:-} && \ export DREAMFIX_REPO_ROOT='${DREAMFIX_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_smoke.py::test_compute_kv_neighbor_and_cache_init" +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_smoke.py::test_compute_kv_neighbor_and_cache_init" echo "[slurm] command: ${COMMAND}" From ba5ec9ecf66bbfba62c28d3af2cb23077503f933 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:06:05 -0700 Subject: [PATCH 11/22] integrations/artifixer: pipeline cache + initialize_cache Phase 3.5a: ``ArtifixerInferencePipeline`` extends ``WanInferencePipeline`` with the ArtiFixer conditioning surface. The cache adds full-rollout fields for the per-AR-chunk slicing in Phase 3.5b (``condition_latent``, ``opacity``, ``camera_rays``, ``w2cs``, ``Ks``, optional ``neighbor_w2cs`` / ``neighbor_Ks``). ``initialize_cache`` takes pre-VAE-encoded latents from the caller (the dreamfix-side Phase 4 driver owns VAE encoding) and: 1. Runs ``WanInferencePipeline.initialize_cache(text, image=None, height, width)`` for text embeddings + base text K/V build. 2. Projects neighbor latents through ``patch_embedding`` to per-token ``neighbor_context`` (mirrors dreamfix transformer.py L361-362) and pushes it into every block's ``ArtifixerCrossAttention.neighbor_kv_cache`` via the ``ArtifixerDiTNetwork.initialize_neighbor_kv_caches`` hook. 3. Updates the per-grid ``patches_x`` / ``patches_y`` RoPE coefficients on both PRoPE modules. 4. Precomputes the neighbor-side PRoPE ``apply_fns`` once per rollout (neighbor cameras are static across AR steps; the source-side ``apply_fns`` get updated per AR chunk in Phase 3.5b). Wires ``ArtifixerInferencePipelineConfig`` into the shipped recipe in ``config.py``. Phase 3.5b adds ``generate``: per-AR-chunk PRoPE-src precompute, opacity-weighted latent mix, manual denoise loop with ``prepare_latents`` renoise, and decode. Static check + linter clean. --- integrations/artifixer/artifixer/config.py | 5 +- integrations/artifixer/artifixer/pipeline.py | 249 +++++++++++++++++++ 2 files changed, 252 insertions(+), 2 deletions(-) create mode 100644 integrations/artifixer/artifixer/pipeline.py diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index 3bac28495..5edc51602 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -37,13 +37,14 @@ import torch from artifixer.checkpoint import 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 WanInferencePipelineConfig, WanVAEDecoderConfig +from flashdreams.recipes.wan import WanVAEDecoderConfig # Mirrors the dreamfix stage-3 run config (see # wandb/.../run-*/files/config.yaml in @@ -73,7 +74,7 @@ ) _BASE_TRANSFORMER_DTYPE = torch.bfloat16 -PIPELINE_ARTIFIXER_DMD_T2V_1PT3B = WanInferencePipelineConfig( +PIPELINE_ARTIFIXER_DMD_T2V_1PT3B = ArtifixerInferencePipelineConfig( recipe_name="artifixer-dmd-wan2.1-t2v-1.3b", enable_sync_and_profile=True, encoder=None, diff --git a/integrations/artifixer/artifixer/pipeline.py b/integrations/artifixer/artifixer/pipeline.py new file mode 100644 index 000000000..85a357a18 --- /dev/null +++ b/integrations/artifixer/artifixer/pipeline.py @@ -0,0 +1,249 @@ +# 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. + +Phase 3.5a (this commit) lands the cache + ``initialize_cache``: + + - VAE-encoded condition latent and neighbor latent arrive pre-computed + from the caller (the dreamfix-side driver in Phase 4) so this + pipeline does not need its own VAE encoder. + - On ``initialize_cache`` we: + 1. Run the base ``WanInferencePipeline.initialize_cache`` (text + encoder + text K/V build) with ``image=None``. + 2. Project the neighbor latent through ``patch_embedding`` to get + the per-token neighbor context, then push it into every block's + :class:`ArtifixerCrossAttention.neighbor_kv_cache` via + ``ArtifixerDiTNetwork.initialize_neighbor_kv_caches``. + 3. Update the patches_x/patches_y RoPE coefficients on both PRoPE + modules and precompute the neighbor-side ``apply_fns`` (the + neighbor cameras are static across AR steps). + 4. Stash the full-rollout opacity / camera_rays / w2cs / Ks on the + cache for the per-AR ``generate`` slicing (Phase 3.5b). + +Phase 3.5b adds ``generate``: 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.transformer import ArtifixerWanTransformer +from torch import Tensor + +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], + *, + 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. + condition_latent: VAE-encoded reconstruction-rendered RGB, + ``[B, in_dim, T_lat, Hl, Wl]`` -- the caller (Phase 4 + driver) 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}" + ) + + # 1. Base text / image cross-attn cache (image=None for ArtiFixer). + base_cache = super().initialize_cache(text=text, 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 _encode_neighbor_context(self, neighbor_latent: Tensor) -> Tensor: + """Run ``patch_embedding`` over neighbor latents and flatten to tokens. + + Mirrors dreamfix ``ArtifixerTransformer.forward`` L361-L362:: + + 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) + + +__all__ = [ + "ArtifixerInferencePipeline", + "ArtifixerInferencePipelineCache", + "ArtifixerInferencePipelineConfig", +] From f2321ad09fb07e40408e501a8bccea0fc812ac16 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:08:50 -0700 Subject: [PATCH 12/22] integrations/artifixer: pipeline generate with self-forcing renoise loop Phase 3.5b: ``ArtifixerInferencePipeline.generate`` runs the per-AR-chunk denoise loop and decodes a video chunk. Bypasses ``DiffusionModel.generate`` so we can renoise each step toward a fresh ``opacity_weighted_latent_mix`` of the condition + new noise (matching ``ArtifixerKvCachePipeline.generate_samples_from_batch`` L211-L264 in dreamfix), instead of ``FlowMatchScheduler.sample``'s default renoise toward the predicted clean. Per-AR-step orchestration: 1. ``_chunk_frame_ranges`` translates the AR index into latent-rate and input-rate slice boundaries, accounting for the Wan VAE's ``1 + 4`` temporal layout on the first chunk. 2. Slice the chunk's ``condition_latent``, ``opacity``, ``w2cs``, ``Ks``, ``camera_rays`` (auto-detects latent vs input rate) out of the full-rollout cache. 3. ``transformer.prope_cross_attn_src._precompute_and_cache_apply_fns (chunk_w2cs, chunk_Ks)`` -- the neighbor-side cameras are static and were already precomputed in ``initialize_cache``. 4. ``_build_ctrl`` patchifies the chunk's opacity / camera_rays via :func:`patchify_opacity` / :func:`patchify_camera_rays` into an :class:`ArtifixerCtrl` payload threaded through ``predict_flow``. 5. ``transformer_cache.start(ar_idx)``. 6. Initial latent = ``opacity_weighted_latent_mix(condition, opacity, randn)`` (unpatchified) -> patchify. 7. For each of the ``num_inference_steps`` denoise steps: - ``flow = transformer.predict_flow(latent, t, cache, input=ctrl)`` - ``clean = latent - sigma * flow`` - If not the last step: sample fresh noise, re-mix, patchify, ``latent = (1 - sigma_next) * clean + sigma_next * fresh_mix``. 8. ``postprocess_clean_latent`` (no-op for non-I2V). 9. Stash a ``DiffusionModel.FinalState`` on the cache so the inherited ``finalize`` path closes the AR cache. 10. ``unpatchify_and_maybe_gather_cp`` + decode -> return. Static check + linter clean. End-to-end testability requires the merged DMD safetensors (Phase 5) and the dreamfix-side driver that prepares condition_latent + neighbor_latent (Phase 4). --- integrations/artifixer/artifixer/pipeline.py | 230 ++++++++++++++++++- 1 file changed, 229 insertions(+), 1 deletion(-) diff --git a/integrations/artifixer/artifixer/pipeline.py b/integrations/artifixer/artifixer/pipeline.py index 85a357a18..607d0bffd 100644 --- a/integrations/artifixer/artifixer/pipeline.py +++ b/integrations/artifixer/artifixer/pipeline.py @@ -48,9 +48,13 @@ from dataclasses import dataclass, field import torch -from artifixer.transformer import ArtifixerWanTransformer +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, @@ -241,6 +245,230 @@ def _encode_neighbor_context(self, neighbor_latent: Tensor) -> Tensor: 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 dreamfix bookkeeping in + ``kv_cache_pipeline.generate_samples_from_batch`` L201-L204: + + - 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 dreamfix + ``kv_cache_pipeline.generate_samples_from_batch`` L211-L264). + + 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 L165). + 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, + ) + latent = transformer.patchify_and_maybe_split_cp(latent_unpatched) + + # 4-step DMD denoise with prepare_latents renoise. Mirrors + # ``ArtifixerKvCachePipeline.generate_samples_from_batch`` L238-L264: + # at every non-exit step we replace the next-iteration ``noisy`` + # with a fresh ``opacity_weighted_latent_mix(...)``, not 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, + ) + fresh_mix = transformer.patchify_and_maybe_split_cp(fresh_mix_unpatched) + 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", From a520a8427551930dc18ddaf8d8f46b1f0f4d6261 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:14:37 -0700 Subject: [PATCH 13/22] integrations/artifixer: state_dict_transform for merged DMD safetensors Phase 5a: ``artifixer_dmd_state_dict_transform`` remaps the merged checkpoint produced by ``dreamfix/scripts/merge_dcp_to_safetensors.py`` from HF diffusers naming to the flashdreams ``WanDiTNetwork`` / ``ArtifixerDiTNetwork`` naming. The regex mapping is a superset of ``integrations/fastvideo_causal_wan22/.../config.CHECKPOINT_KEY_MAPPING`` (Wan 2.1 / 2.2 share that layout). Three ArtiFixer-specific extra substitutions land the 270 ArtiFixer-only keys on the right modules: blocks.X.attn2.add_k_proj.* -> blocks.X.cross_attn.add_k_proj.* blocks.X.attn2.add_v_proj.* -> blocks.X.cross_attn.add_v_proj.* blocks.X.attn2.norm_added_k.* -> blocks.X.cross_attn.norm_added_k.* The other 60 ArtiFixer-only keys (``blocks.X.opacity_embedding.*``, ``blocks.X.camera_embedding.*``) have no ``attn2`` prefix and pass through unchanged -- they already match the :class:`ArtifixerBlock` attribute names registered at __init__ time. tests/test_state_dict_transform.py: 6 CPU-only key-naming tests that load the ``param_audit.json`` produced by Phase 0 and assert: - Diffusers-named keys are renamed correctly (spot check). - ``opacity_embedding`` / ``camera_embedding`` pass through unchanged. - The full 1095-key audit transforms to 1095 unique keys (no collisions). - Every block (0..29) has the expected 21 attributes (self_attn / cross_attn / norms / ffn / modulation / add_k_proj / add_v_proj / norm_added_k / opacity_embedding / camera_embedding). - Network-level globals (``head.*``, embeddings) all present. - All regex backreferences are well-formed. 20/20 tests pass (3 PRoPE + 5 patches + 5 latent_mix + 6 state_dict + 1 cache_init). This commit only adds the transform helper; switching the recipe's checkpoint_path + state_dict_transform to use it (and dropping ``zero_pad_artifixer_keys``) is Phase 5b once the merged safetensors are at a portable path. --- .../artifixer/artifixer/checkpoint.py | 117 +++++++++- .../tests/test_state_dict_transform.py | 208 ++++++++++++++++++ scripts/run_prope_parity.sh | 2 +- 3 files changed, 317 insertions(+), 10 deletions(-) create mode 100644 integrations/artifixer/tests/test_state_dict_transform.py diff --git a/integrations/artifixer/artifixer/checkpoint.py b/integrations/artifixer/artifixer/checkpoint.py index 7c59d6497..41e649a19 100644 --- a/integrations/artifixer/artifixer/checkpoint.py +++ b/integrations/artifixer/artifixer/checkpoint.py @@ -17,21 +17,29 @@ Two source layouts are supported: - * **Vanilla Wan 2.1 1.3B** from ``Wan-AI/Wan2.1-T2V-1.3B`` (used in - Phase 1 / 2.1 while only the base architecture is implemented). The - transform zero-pads the ArtiFixer-only keys so ``load_state_dict`` - succeeds in strict mode. Zero-padding matches the dreamfix - initialization in ``ArtifixerTransformerBlock.__init__`` L637-651. + * **Vanilla Wan 2.1 1.3B** from ``Wan-AI/Wan2.1-T2V-1.3B`` -- the + transform :func:`zero_pad_artifixer_keys` zero-pads the ArtiFixer-only + keys so ``load_state_dict`` succeeds in strict mode. Zero-padding + matches dreamfix's initialization in + ``ArtifixerTransformerBlock.__init__`` L637-651 and produces a + behaviorally-identical-to-Wan model (the ArtiFixer extension paths + are zero-gated until trained weights are loaded). * **Merged ArtiFixer DMD safetensors** produced by - ``dreamfix/scripts/merge_dcp_to_safetensors.py``. This still uses the - HuggingFace diffusers naming (e.g. ``blocks.X.attn1.to_q.weight``); - Phase 5 will add the diffusers -> ``WanDiTNetwork`` regex remap on - top of the zero-pad pass. + ``dreamfix/scripts/merge_dcp_to_safetensors.py``. 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 our new + ``ArtifixerBlock`` / ``ArtifixerCrossAttention`` submodules without + further per-key renames. """ from __future__ import annotations +import re from typing import Callable import torch @@ -39,6 +47,95 @@ 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 produced by + ``dreamfix/scripts/merge_dcp_to_safetensors.py`` carries the 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, @@ -97,5 +194,7 @@ def transform(state_dict: dict[str, Tensor]) -> dict[str, Tensor]: __all__ = [ + "DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING", + "artifixer_dmd_state_dict_transform", "zero_pad_artifixer_keys", ] 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..7ad4d5a26 --- /dev/null +++ b/integrations/artifixer/tests/test_state_dict_transform.py @@ -0,0 +1,208 @@ +# 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 the ``param_audit.json`` produced by +``dreamfix/scripts/dump_artifixer_param_names.py`` 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()``. + +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 +from pathlib import Path + +import pytest +import torch +from artifixer.checkpoint import ( + DIFFUSERS_TO_WAN_DIT_NETWORK_KEY_MAPPING, + artifixer_dmd_state_dict_transform, +) + +AUDIT_PATH = Path( + "/lustre/fsw/portfolios/nvr/users/rdelutio/code/dreamfix/" + "merged_checkpoints/artifixer_dmd_1p3b_s3_2000/param_audit.json" +) + + +def _load_audit() -> dict[str, dict]: + if not AUDIT_PATH.exists(): + pytest.skip( + f"param audit JSON not found at {AUDIT_PATH}; run " + f"dreamfix/scripts/dump_artifixer_param_names.py 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 dreamfix 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/scripts/run_prope_parity.sh b/scripts/run_prope_parity.sh index dac2aee3a..0971d67a6 100755 --- a/scripts/run_prope_parity.sh +++ b/scripts/run_prope_parity.sh @@ -29,7 +29,7 @@ cd ${REPO_DIR} && \ export PYTHONPATH=${REPO_DIR}/integrations/artifixer:${REPO_DIR}/flashdreams:${DREAMFIX_REPO_ROOT}:\${PYTHONPATH:-} && \ export DREAMFIX_REPO_ROOT='${DREAMFIX_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_smoke.py::test_compute_kv_neighbor_and_cache_init" +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}" From eebc1d21f837d0d043efc71be4c545feb13d6724 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:17:32 -0700 Subject: [PATCH 14/22] integrations/artifixer: wire merged DMD safetensors as the default Phase 5b: ``PIPELINE_ARTIFIXER_DMD_T2V_1PT3B`` now defaults to the merged ArtiFixer DMD safetensors produced by ``dreamfix/scripts/merge_dcp_to_safetensors.py``, paired with ``artifixer_dmd_state_dict_transform`` (added in Phase 5a) for the diffusers -> ``WanDiTNetwork`` regex remap and the ArtiFixer-only ``attn2.add_k_proj`` / ``attn2.add_v_proj`` / ``attn2.norm_added_k`` -> ``cross_attn.*`` substitutions. Two env vars cover the deployment surface: * ``ARTIFIXER_DMD_CHECKPOINT_PATH`` -- path to the merged safetensors; defaults to the dreamfix repo's ``merged_checkpoints/`` layout (``/lustre/fsw/.../artifixer_dmd_1p3b_s3_2000/model.safetensors``). * ``ARTIFIXER_USE_BASE_WAN_WEIGHTS=1`` -- fall back to vanilla Wan 2.1 1.3B HuggingFace weights + ``zero_pad_artifixer_keys``; useful for wiring smoke-tests when the merged safetensors are unavailable. The 20 existing CPU-only tests still pass; live load of the merged safetensors into the network is exercised by Phase 4's end-to-end run (driver in dreamfix). --- integrations/artifixer/README.md | 19 +++++---- integrations/artifixer/artifixer/config.py | 47 +++++++++++++++++----- 2 files changed, 47 insertions(+), 19 deletions(-) diff --git a/integrations/artifixer/README.md b/integrations/artifixer/README.md index cf6ae4d5d..ee2699081 100644 --- a/integrations/artifixer/README.md +++ b/integrations/artifixer/README.md @@ -21,15 +21,18 @@ plugin ports it to FlashDreams' faster Wan stack (RingAttention, cuDNN, | Phase | Scope | Status | | --- | --- | --- | -| 1 | Recipe scaffold + AR/scheduler knobs match dreamfix stage-3 DMD | this commit | -| 2 | Per-block opacity + camera MLPs, neighbor cross-attn, PRoPE | upcoming | -| 3 | Opacity-weighted latent mixing + self-forcing renoise loop | upcoming | +| 1 | Recipe scaffold + AR/scheduler knobs match dreamfix stage-3 DMD | done | +| 2 | Per-block opacity + camera MLPs, neighbor cross-attn, PRoPE (parity-tested vs dreamfix) | done | +| 3 | Opacity-weighted latent mixing + self-forcing renoise loop | done | | 4 | dreamfix-format conditioning surface (rgb_rendered, opacity, neighbors) | upcoming | -| 5 | `state_dict_transform` for the merged ArtiFixer DMD safetensors | upcoming | - -Phase 1 loads vanilla Wan 2.1 1.3B base weights from HuggingFace, so the -output is a plain T2V video. The recipe is wired up but the ArtiFixer -architectural extensions land in later commits. +| 5 | `state_dict_transform` for the merged ArtiFixer DMD safetensors | done | + +By default the recipe loads the merged ArtiFixer DMD safetensors from +`ARTIFIXER_DMD_CHECKPOINT_PATH` (defaults to a `/lustre` path that +matches the dreamfix repo's `merged_checkpoints/`); 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 diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index 5edc51602..7b79ef630 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -34,8 +34,13 @@ from __future__ import annotations +import os + import torch -from artifixer.checkpoint import zero_pad_artifixer_keys +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 @@ -61,9 +66,21 @@ ARTIFIXER_NUM_INFERENCE_STEPS = 4 ARTIFIXER_TIMESTEP_SHIFT = 5.0 -# Phase 1: load vanilla Wan 2.1 1.3B base weights from HF. Phase 5 replaces -# this with the merged ArtiFixer DMD safetensors plus a state_dict_transform -# that remaps diffusers naming and absorbs the ArtiFixer-only keys. +# Phase 5b default: load the merged ArtiFixer DMD safetensors produced by +# ``dreamfix/scripts/merge_dcp_to_safetensors.py``. The path is overridable +# via ``ARTIFIXER_DMD_CHECKPOINT_PATH`` so deployments outside the original +# /lustre layout don't have to fork the recipe. +ARTIFIXER_DMD_CHECKPOINT_PATH = os.environ.get( + "ARTIFIXER_DMD_CHECKPOINT_PATH", + "/lustre/fsw/portfolios/nvr/users/rdelutio/code/dreamfix/" + "merged_checkpoints/artifixer_dmd_1p3b_s3_2000/model.safetensors", +) + +# Pre-Phase-5b 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" ) @@ -74,6 +91,19 @@ ) _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 = 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: + _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", enable_sync_and_profile=True, @@ -84,13 +114,8 @@ transformer=ArtifixerWanTransformerConfig( network=_BASE_NETWORK_CONFIG, dtype=_BASE_TRANSFORMER_DTYPE, - checkpoint_path=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, - ), + checkpoint_path=_CHECKPOINT_PATH, + state_dict_transform=_STATE_DICT_TRANSFORM, batch_shape=(), len_t=ARTIFIXER_LEN_T, guidance_scale=1.0, From 593c454f1c225003d49115a8bfe4b7005fe6b010 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:20:01 -0700 Subject: [PATCH 15/22] integrations/artifixer: pipeline initialize_cache accepts pre-encoded prompts Phase 4a (flashdreams side): ``ArtifixerInferencePipeline.initialize_cache`` now takes either ``text: list[str]`` (raw prompts, encoded internally by UMT5) or ``text_embeddings: Tensor`` (pre-encoded UMT5 output). The two are mutually exclusive. The dreamfix-side driver (Phase 4b) feeds the pre-encoded prompt tensor that dreamfix's data pipeline produces (each eval item carries an ``encoded_prompt`` field), avoiding a second UMT5 forward when we already have the embeddings. The ``text_embeddings`` path bypasses the UMT5-only call site in :meth:`WanInferencePipeline.initialize_cache` and slots the embeddings directly into the transformer cache via ``StreamInferencePipeline.initialize_cache(transformer_context=...)``. CFG is intentionally disabled in this branch (no ``negative_text_embeddings`` even when ``guidance_scale > 1``) to match the dreamfix ``ArtifixerKvCachePipeline`` contract -- the KV-cache pipeline ignores ``negative_prompt`` (see ``kv_cache_pipeline.py`` L104-L109). Static check + linter clean. --- integrations/artifixer/artifixer/pipeline.py | 59 +++++++++++++++++++- 1 file changed, 56 insertions(+), 3 deletions(-) diff --git a/integrations/artifixer/artifixer/pipeline.py b/integrations/artifixer/artifixer/pipeline.py index 607d0bffd..c8f8aa423 100644 --- a/integrations/artifixer/artifixer/pipeline.py +++ b/integrations/artifixer/artifixer/pipeline.py @@ -123,8 +123,9 @@ class ArtifixerInferencePipeline(WanInferencePipeline): @torch.no_grad() def initialize_cache( # type: ignore[override] self, - text: list[str], + text: list[str] | None = None, *, + text_embeddings: Tensor | None = None, condition_latent: Tensor, opacity: Tensor, camera_rays: Tensor, @@ -139,7 +140,13 @@ def initialize_cache( # type: ignore[override] """Build a per-rollout cache and push static state into the network. Args: - text: One prompt per batch element. + 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 + dreamfix-backed driver in :mod:`model_eval` passes 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 (Phase 4 driver) handles VAE encoding so this pipeline does not @@ -167,9 +174,21 @@ def initialize_cache( # type: ignore[override] 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). - base_cache = super().initialize_cache(text=text, height=height, width=width) + 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 @@ -229,6 +248,40 @@ def initialize_cache( # type: ignore[override] 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 -- the + dreamfix driver already has the encoded prompts and never uses + I2V images. CFG is disabled per the dreamfix ``ArtifixerKvCachePipeline`` + 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. From 64b5e92a00ec93bdb445be1783867c0778eb576e Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:48:51 -0700 Subject: [PATCH 16/22] integrations/artifixer: pass extras as block_extra_kwargs dict, not unpacked MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live-run fix: ``WanDiTNetwork.forward`` accepts a SINGLE ``block_extra_kwargs`` parameter (a dict), not arbitrary kwargs. The base ``Wan21Transformer._predict_flow`` does ``self._select_network(...) (**network_extra_kwargs)`` which unpacks ``network_extra_kwargs`` into the network forward call. Our override was flattening the extras into ``network_extra_kwargs`` directly: network_extra_kwargs = { "opacity_extra": ..., "camera_extra": ..., ... } …which after unpack became ``network(opacity_extra=..., camera_extra=...)`` and crashed with ``TypeError: WanDiTNetwork.forward() got an unexpected keyword argument 'opacity_extra'``. Wrap them in a single ``block_extra_kwargs`` dict so the unpack lands on ``network(block_extra_kwargs={"opacity_extra": ..., ...})``, which ``WanDiTNetwork.forward`` then forwards to every block as ``block(... **block_extra_kwargs)`` (network.py L438-449). The :class:`ArtifixerBlock.forward` signature accepts each extra as a keyword arg. Bug surfaced on the first end-to-end run with ``--inference_backend flashdreams`` on a DL3DV scene; before the crash, the FlashDreams pipeline ``.setup()`` loaded all 1095 keys of the merged DMD safetensors cleanly (validating Phase 5's state_dict_transform on the real network). --- .../artifixer/artifixer/transformer.py | 35 ++++++++++++------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/integrations/artifixer/artifixer/transformer.py b/integrations/artifixer/artifixer/transformer.py index af174dfa6..e59a659f6 100644 --- a/integrations/artifixer/artifixer/transformer.py +++ b/integrations/artifixer/artifixer/transformer.py @@ -128,24 +128,33 @@ def predict_flow( # type: ignore[override] ) -> Tensor: """Add ArtiFixer block_extra_kwargs to the base predict_flow. - Reads patchified opacity/camera features and the ``ignore_neighbors`` - flag from :class:`ArtifixerCtrl` and forwards them to every block via - ``network_extra_kwargs``. Also passes the PRoPE source/target modules - themselves -- they are bound at construction time, but - :class:`ArtifixerCrossAttention.forward` expects them as forward args - so they show up uniformly under ``block_extra_kwargs``. - - When ``input`` is not an :class:`ArtifixerCtrl` (e.g. T2V smoke + 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`` + (network.py L438-449), 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): - network_extra_kwargs.setdefault("opacity_extra", input.opacity_extra) - network_extra_kwargs.setdefault("camera_extra", input.camera_extra) - network_extra_kwargs.setdefault("ignore_neighbors", input.ignore_neighbors) - network_extra_kwargs.setdefault("prope_src", self.prope_cross_attn_src) - network_extra_kwargs.setdefault("prope_tgt", self.prope_cross_attn_tgt) + 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. From eec9eb328d60631fbffa9326d8545ee4f3fca0ee Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 22:52:28 -0700 Subject: [PATCH 17/22] integrations/artifixer: permute condition latent (B,C,T,H,W) -> (B,T,C,H,W) before patchify Second live-run bug fix. dreamfix's ``encode_video_frames`` returns latents in diffusers convention ``(B, in_dim, T_lat, Hl, Wl)`` (the same ordering ``cache. condition_latent`` and ``opacity_weighted_latent_mix`` use), but FlashDreams' ``WanDiTNetwork.patchify_and_maybe_split_cp`` uses the einops pattern: "... (t kt) c (h kh) (w kw) -> ... (t h w) (c kt kh kw)" i.e. it expects ``(B, T, C, H, W)`` with C *after* T. Feeding the diffusers-shaped latent in put C=16 into the t-axis position, so the per-token feature dim came out as ``T_chunk * kt * kh * kw = 7 * 1 * 2 * 2 = 28`` instead of ``C * kt * kh * kw = 16 * 1 * 2 * 2 = 64``, and the post-patchify linear (which fuses the Conv3d patch_embedding) failed with:: a and b must have same reduction dim, but got [.., 28] X [64, 1536] Permute ``(B, C, T, H, W) -> (B, T, C, H, W)`` for both the initial mix and the per-step ``fresh_mix`` before calling ``patchify_and_maybe_split_cp``. The unpatchify side already returns ``(B, T, C, H, W)`` (the FlashDreams convention) which is what the ``WanVAEDecoder`` expects, so no symmetric permute is needed at the decode call site. Static check + linter clean. Bug surfaced on the second ``--inference_backend flashdreams`` end-to-end run (the previous predict_flow kwarg fix unblocked us to reach the patchify). --- integrations/artifixer/artifixer/pipeline.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/integrations/artifixer/artifixer/pipeline.py b/integrations/artifixer/artifixer/pipeline.py index c8f8aa423..05ffb2c41 100644 --- a/integrations/artifixer/artifixer/pipeline.py +++ b/integrations/artifixer/artifixer/pipeline.py @@ -451,7 +451,13 @@ def generate( # type: ignore[override] vae_scale_factor_spatial=vae_s, is_first_chunk=is_first, ) - latent = transformer.patchify_and_maybe_split_cp(latent_unpatched) + # FlashDreams ``patchify_and_maybe_split_cp`` expects ``(B, T, C, H, W)`` + # (pattern ``... (t kt) c (h kh) (w kw)``), but dreamfix's VAE encoder + # (and therefore ``cache.condition_latent`` / ``latent_unpatched``) is + # in diffusers convention ``(B, C, T, H, W)``. 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. Mirrors # ``ArtifixerKvCachePipeline.generate_samples_from_batch`` L238-L264: @@ -490,7 +496,11 @@ def generate( # type: ignore[override] vae_scale_factor_spatial=vae_s, is_first_chunk=is_first, ) - fresh_mix = transformer.patchify_and_maybe_split_cp(fresh_mix_unpatched) + # 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 ) From 90ddfc61117dc31449793fa4d69f41c40f0b5771 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Mon, 11 May 2026 23:27:40 -0700 Subject: [PATCH 18/22] integrations/artifixer: disable enable_sync_and_profile on the DMD pipeline The base StreamInferencePipeline.generate installs cache.event_profiler when this flag is True, but our custom ArtifixerInferencePipeline.generate bypasses that path -- so finalize() asserts on a None profiler. Turn the flag off until our generate() grows the matching EventProfiler setup; the inline comment documents the contract so the next person knows what to wire up before flipping it back on. Unblocks the live DL3DV rollout end-to-end through the FlashDreams backend. --- integrations/artifixer/artifixer/config.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index 7b79ef630..72ab5350c 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -106,7 +106,13 @@ PIPELINE_ARTIFIXER_DMD_T2V_1PT3B = ArtifixerInferencePipelineConfig( recipe_name="artifixer-dmd-wan2.1-t2v-1.3b", - enable_sync_and_profile=True, + # 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( From 6202a95d925d79806c1121357b9404745ff8eacf Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Tue, 12 May 2026 14:04:07 -0700 Subject: [PATCH 19/22] artifixer: no-op finalize_kv_cache to match dreamfix KV semantics The dreamfix-native ``ArtifixerKvCachePipeline.generate_samples_from_batch`` advances its KV cache *in-place* during the regular denoise forwards and never runs an extra forward at AR-chunk boundaries. The FlashDreams default ``Transformer.finalize_kv_cache`` runs one more ``predict_flow`` at the context-noise timestep to advance the cache (the extra forward discarded by ``_ = ...``). For the artifixer recipe this extra forward writes a *different* KV state than dreamfix'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: ``scripts/parity_harness.py`` ``transformer_call_*`` diff, before this change: chunk 0 step 0: 43.55 dB ... chunk 1 step 0: 7.53 dB <- cliff chunk 1 step 1: 18.15 dB chunk 2 step 0: 7.57 dB <- cliff after this change: chunk 0 step 0: 43.55 dB (unchanged, no AR boundary yet) chunk 1 step 0: 37.37 dB chunk 1 step 1: 37.20 dB chunk 2 step 0: 37.44 dB End-to-end ``final_video`` PSNR jumps from 30.49 dB -> 47.72 dB on the parity-harness item. ``cache.finalize(autoregressive_index)`` (the bookkeeping that increments AR index) is still invoked by ``DiffusionModel.finalize`` after this no-op, so the rollout state advances correctly; only the redundant predict is suppressed. Mirrors the pattern in ``flashdreams/recipes/alpadreams/transformer/__init__.py::finalize_kv_cache`` which conditionally skips the same call. --- .../artifixer/artifixer/transformer.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/integrations/artifixer/artifixer/transformer.py b/integrations/artifixer/artifixer/transformer.py index e59a659f6..d55bb0e1d 100644 --- a/integrations/artifixer/artifixer/transformer.py +++ b/integrations/artifixer/artifixer/transformer.py @@ -170,6 +170,30 @@ def predict_flow( # type: ignore[override] network_extra_kwargs=network_extra_kwargs, ) + def finalize_kv_cache(self, *args: Any, **kwargs: Any) -> None: + """No-op: artifixer (dreamfix reference) does not finalize the KV cache. + + The dreamfix ``ArtifixerKvCachePipeline.generate_samples_from_batch`` + 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; that extra forward writes a *different* KV state than + dreamfix'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 + (calls 4, 8 in ``scripts/parity_harness.py``'s per-step diff). + + Skipping the extra forward here aligns FlashDreams' KV-cache + semantics with the dreamfix reference: the cache going into AR + chunk ``N+1`` is the one written by the final denoise step of AR + chunk ``N``, identical to what dreamfix uses. ``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", From 0a480e7382a01efb42fdbce710b331aef844085d Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Tue, 12 May 2026 14:49:52 -0700 Subject: [PATCH 20/22] artifixer block: promote AdaLN / norms / residuals to fp32 (mirror dreamfix) The dreamfix ``ArtifixerTransformerBlock.forward`` (``model_training/net/transformer.py`` L725-816) runs every per-block AdaLN modulation, RMSNorm, and residual add in fp32, then casts the result back to the input dtype. The base FlashDreams ``Block.forward`` keeps everything in ``x.dtype`` (typically bf16). 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 -- a layer-by-layer accumulation visible in ``scripts/parity_harness.py --capture_blocks``. This change mirrors the six dreamfix promotions: - modulation chunking: ``(modulation.float() + e.float()).chunk(6)`` - self-attn pre-norm + AdaLN - self-attn residual (x.float() + y * gate) - cross-attn pre-norm - FFN pre-norm + AdaLN - FFN residual (x.float() + ff.float() * gate) ``norm1`` / ``norm2`` are ``elementwise_affine=False`` so calling them with an fp32 input is a free upcast (no weight/bias dtype check). ``norm3`` is ``elementwise_affine=True`` with bf16 weight/bias after ``.to(bf16)``; passing fp32 input through ``nn.LayerNorm`` raises ``expected scalar type Float but found BFloat16``, so we route through a ``_layer_norm_fp32`` helper that promotes weight + bias to fp32 too. This matches diffusers' ``FP32LayerNorm`` which the dreamfix WanTransformerBlock norms use under the hood. End-to-end ``final_video`` cross-backend PSNR jumps 47.72 dB -> 51.34 dB on the parity-harness item. Block-by-block (call 0): block 0 went from 51.08 dB -> 55.35 dB; L1_max outliers in the middle of the stack (blocks 17, 21) went from 60+ -> 28-45. The remaining residual is attention-impl op-ordering noise inside the bf16 SDPA / FFN paths, which is irreducible without forcing the same attention backend on both implementations. Performance cost: a handful of extra casts per block. The fp32-promoted norms / residuals run on small tensors compared to the attention QKV / FFN GEMMs, so wall-clock impact is negligible. --- .../artifixer/artifixer/network/block.py | 68 ++++++++++++++++--- 1 file changed, 60 insertions(+), 8 deletions(-) diff --git a/integrations/artifixer/artifixer/network/block.py b/integrations/artifixer/artifixer/network/block.py index 07345c846..d320e1c3a 100644 --- a/integrations/artifixer/artifixer/network/block.py +++ b/integrations/artifixer/artifixer/network/block.py @@ -36,13 +36,30 @@ from typing import Any -import torch 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 dreamfix ``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)`` on the FD side 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.""" @@ -120,9 +137,40 @@ def forward( # type: ignore[override] "We expect to have called update_parameters_after_loading_checkpoint() " "before running the forward pass" ) - e_chunks = (self.modulation + e).chunk(6, dim=-2) - - y = self.norm1(x) * (1 + e_chunks[1]) + e_chunks[0] + # Mirror dreamfix ``ArtifixerTransformerBlock.forward`` (transformer.py + # L727, L763, L787, L790, L810, L814) 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 (see ``scripts/parity_harness.py --capture_blocks``). + # The fp32 promotion below mirrors the dreamfix reference exactly: + # + # * 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 dreamfix also + # promotes ``ff_output`` to fp32 inside the residual (L814). + # + # 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. Note FD's bf16-throughout path + # remains the default; this override only kicks in when the + # ``ArtifixerBlock`` subclass is used (i.e. 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 which dreamfix's WanTransformerBlock norms use). + 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: @@ -132,16 +180,20 @@ def forward( # type: ignore[override] rope_freqs=rope_freqs, kv_cache=cache.self_attn, ) - x = x + (y * e_chunks[2]) + x = (x.float() + y * e_chunks[2]).to(x_dtype) + # Cross-attn pre-norm in fp32 (mirrors dreamfix norm2 promotion at L790). + # The post-cross-attn residual is NOT promoted in dreamfix (L807) so + # we keep that in the input dtype here too. x = x + self.cross_attn( - self.norm3(x), + _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) * (1 + e_chunks[4]) + e_chunks[3] + y = (self.norm2(x.float()) * (1 + e_chunks[4]) + e_chunks[3]).to(x_dtype) y = self.ffn(y) - x = x + (y * e_chunks[5]) + # FFN residual with ff_output also promoted (matches dreamfix L814). + x = (x.float() + y.float() * e_chunks[5]).to(x_dtype) return x From ac4ab0a80874a70066b19b4419265718a0bf8ed0 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Tue, 12 May 2026 20:37:24 -0700 Subject: [PATCH 21/22] integrations/artifixer: polish for MR (docstrings, env-var checkpoint, README) Three small cleanups in preparation for opening the upstream MR: 1. ``ARTIFIXER_DMD_CHECKPOINT_PATH`` is now required (no committed ``/lustre/.../rdelutio/...`` default). Either set the env var or set ``ARTIFIXER_USE_BASE_WAN_WEIGHTS=1`` -- otherwise ``config.py`` raises a clear ``RuntimeError`` at import time pointing the user to ``dreamfix/scripts/merge_dcp_to_safetensors.py``. The same env-var pattern is applied to ``tests/test_state_dict_transform.py``'s ``ARTIFIXER_PARAM_AUDIT_PATH``; the existing ``pytest.skip`` path keeps CI green without the audit JSON. 2. Drop ``Phase X.Y (this commit)`` narrative wording from docstrings and comments across ``block.py`` / ``transformer.py`` / ``pipeline.py`` / ``config.py`` / ``runner.py`` / ``cross_attn.py`` / ``latent_mix.py`` / ``checkpoint.py`` / ``tests/test_smoke.py``. The Phase numbering was commit-narrative bleed-through; now reads as plain past-tense prose without dating the code to its development order. 3. Rework ``integrations/artifixer/README.md``: drop the ``| Phase | Scope | Status |`` table (Phase 4 was still ``upcoming`` despite the dreamfix-side adapter closing it), replace with a plain ``| Component | Description |`` table that describes the five recipe pieces. Add a cross-backend parity dB note so reviewers see the validation summary up front. No behavior change. Phase references in pre-existing ``.pyc`` files under ``__pycache__/`` will refresh on next collect. --- integrations/artifixer/README.md | 36 +++++++----- .../artifixer/artifixer/checkpoint.py | 4 +- integrations/artifixer/artifixer/config.py | 58 +++++++++++-------- .../artifixer/artifixer/latent_mix.py | 2 +- .../artifixer/artifixer/network/block.py | 25 ++++---- .../artifixer/artifixer/network/cross_attn.py | 10 ++-- integrations/artifixer/artifixer/pipeline.py | 42 +++++++------- integrations/artifixer/artifixer/runner.py | 19 +++--- .../artifixer/artifixer/transformer.py | 12 ++-- integrations/artifixer/tests/test_smoke.py | 14 ++--- .../tests/test_state_dict_transform.py | 23 +++++--- 11 files changed, 135 insertions(+), 110 deletions(-) diff --git a/integrations/artifixer/README.md b/integrations/artifixer/README.md index ee2699081..a2bda4ec2 100644 --- a/integrations/artifixer/README.md +++ b/integrations/artifixer/README.md @@ -17,22 +17,26 @@ The reference implementation lives in the plugin ports it to FlashDreams' faster Wan stack (RingAttention, cuDNN, `torch.compile`, CUDA graphs). -## Status - -| Phase | Scope | Status | -| --- | --- | --- | -| 1 | Recipe scaffold + AR/scheduler knobs match dreamfix stage-3 DMD | done | -| 2 | Per-block opacity + camera MLPs, neighbor cross-attn, PRoPE (parity-tested vs dreamfix) | done | -| 3 | Opacity-weighted latent mixing + self-forcing renoise loop | done | -| 4 | dreamfix-format conditioning surface (rgb_rendered, opacity, neighbors) | upcoming | -| 5 | `state_dict_transform` for the merged ArtiFixer DMD safetensors | done | - -By default the recipe loads the merged ArtiFixer DMD safetensors from -`ARTIFIXER_DMD_CHECKPOINT_PATH` (defaults to a `/lustre` path that -matches the dreamfix repo's `merged_checkpoints/`); 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). +## Components + +| Component | Description | +| --- | --- | +| Recipe scaffold | AR rollout + 4-step DMD scheduler knobs match the dreamfix stage-3 run config. | +| Per-block conditioning | Opacity + camera-ray MLPs, neighbor cross-attention, and PRoPE — parity-tested against dreamfix. | +| Pipeline | Opacity-weighted latent mixing + self-forcing renoise loop inside `ArtifixerInferencePipeline.generate`. | +| dreamfix-compatible surface | `initialize_cache` accepts pre-encoded UMT5 prompts + VAE-encoded condition / neighbor latents so the dreamfix-side driver (`model_eval/flashdreams_backend.py`) can feed it directly. | +| Checkpoint loader | `state_dict_transform` for the merged ArtiFixer DMD safetensors produced by `dreamfix/scripts/merge_dcp_to_safetensors.py`. | + +Cross-backend parity (captured single-scene `final_video`): **51.34 dB** +PSNR vs the dreamfix-native `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 diff --git a/integrations/artifixer/artifixer/checkpoint.py b/integrations/artifixer/artifixer/checkpoint.py index 41e649a19..c944e347e 100644 --- a/integrations/artifixer/artifixer/checkpoint.py +++ b/integrations/artifixer/artifixer/checkpoint.py @@ -154,10 +154,10 @@ def zero_pad_artifixer_keys( Per block, the transform adds 9 keys (mirroring the 270 ArtiFixer-only keys identified by ``dreamfix/scripts/dump_artifixer_param_names.py``): - * Phase 2.1 — opacity + camera MLPs (4 keys per block): + * 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,) - * Phase 2.2 — neighbor cross-attention (5 keys per block): + * 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,) diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index 72ab5350c..56a9fba43 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -19,17 +19,16 @@ adds (a) per-block opacity and Plucker-camera-ray MLPs, (b) neighbor cross- attention with PRoPE, and (c) opacity-weighted latent mixing. -Phase 2.1: the network now uses :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). -The recipe still loads vanilla Wan 2.1 1.3B base weights from HuggingFace; -:func:`zero_pad_artifixer_keys` pads the state dict so strict-mode load -succeeds. - -Later commits add the neighbor cross-attention third KV bank (Phase 2.2), -PRoPE (Phase 2.3), the opacity-weighted latent mixing pipeline (Phase 3), -and the ``state_dict_transform`` for the merged ArtiFixer DMD safetensors -(Phase 5). +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 @@ -66,21 +65,21 @@ ARTIFIXER_NUM_INFERENCE_STEPS = 4 ARTIFIXER_TIMESTEP_SHIFT = 5.0 -# Phase 5b default: load the merged ArtiFixer DMD safetensors produced by -# ``dreamfix/scripts/merge_dcp_to_safetensors.py``. The path is overridable -# via ``ARTIFIXER_DMD_CHECKPOINT_PATH`` so deployments outside the original -# /lustre layout don't have to fork the recipe. -ARTIFIXER_DMD_CHECKPOINT_PATH = os.environ.get( - "ARTIFIXER_DMD_CHECKPOINT_PATH", - "/lustre/fsw/portfolios/nvr/users/rdelutio/code/dreamfix/" - "merged_checkpoints/artifixer_dmd_1p3b_s3_2000/model.safetensors", +# Default: load the merged ArtiFixer DMD safetensors produced by +# ``dreamfix/scripts/merge_dcp_to_safetensors.py``. 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" ) -# Pre-Phase-5b fallback: vanilla Wan 2.1 1.3B base weights from HF, paired -# with ``zero_pad_artifixer_keys`` so ``load_state_dict`` succeeds in strict +# 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``. +# 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" ) @@ -93,7 +92,7 @@ _USE_BASE_WAN_WEIGHTS = os.environ.get("ARTIFIXER_USE_BASE_WAN_WEIGHTS") == "1" if _USE_BASE_WAN_WEIGHTS: - _CHECKPOINT_PATH = BASE_WAN_T2V_1PT3B_CHECKPOINT_PATH + _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, @@ -101,6 +100,15 @@ 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 produced by " + "``dreamfix/scripts/merge_dcp_to_safetensors.py``, 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 @@ -145,7 +153,7 @@ 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). Phase 1 scaffold: loads vanilla Wan base weights." + "extensions, 4-step DMD)." ), pipeline=PIPELINE_ARTIFIXER_DMD_T2V_1PT3B, ) diff --git a/integrations/artifixer/artifixer/latent_mix.py b/integrations/artifixer/artifixer/latent_mix.py index 1588a1d22..c007ea957 100644 --- a/integrations/artifixer/artifixer/latent_mix.py +++ b/integrations/artifixer/artifixer/latent_mix.py @@ -25,7 +25,7 @@ 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 (Phase 2.1) feed the opacity into +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 diff --git a/integrations/artifixer/artifixer/network/block.py b/integrations/artifixer/artifixer/network/block.py index d320e1c3a..93283dab6 100644 --- a/integrations/artifixer/artifixer/network/block.py +++ b/integrations/artifixer/artifixer/network/block.py @@ -19,17 +19,18 @@ (``model_training/net/transformer.py`` L617-L767). Per-block extensions on top of :class:`Block`: - * **Phase 2.1** — opacity + camera-ray MLPs. Two ``nn.Linear`` heads - project per-token opacity and Plucker-camera-ray features into the + * **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. - - * **Phase 2.2** — ``cross_attn`` is replaced with - :class:`ArtifixerCrossAttention`, which carries ``add_k_proj`` / - ``add_v_proj`` / ``norm_added_k`` and a separate ``attn_op_neighbor`` - RingAttention op for a future neighbor-frame KV branch. The forward - path is unchanged (parent ``CrossAttention.forward`` is called); the - pipeline-level neighbor wiring lands in Phase 3. + 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 @@ -85,7 +86,7 @@ def __init__( # Replace the inherited cross_attn with the ArtiFixer variant that # also carries add_k_proj / add_v_proj / norm_added_k for the - # future neighbor-frame KV branch (Phase 3 wiring). + # neighbor-frame KV branch. self.cross_attn = ArtifixerCrossAttention( query_dim=dim, n_heads=num_heads, @@ -118,7 +119,7 @@ def forward( # type: ignore[override] ) -> Tensor: """Run one transformer block update with ArtiFixer conditioning. - Extras (Phase 2.1 + 2.4): + Extras: - ``opacity_extra`` / ``camera_extra``: per-token opacity and Plucker-camera-ray features added to the AdaLN-normed hidden diff --git a/integrations/artifixer/artifixer/network/cross_attn.py b/integrations/artifixer/artifixer/network/cross_attn.py index e2323cc56..f1de5ffec 100644 --- a/integrations/artifixer/artifixer/network/cross_attn.py +++ b/integrations/artifixer/artifixer/network/cross_attn.py @@ -29,9 +29,9 @@ branch) ``forward`` accepts optional ``neighbor_kv_cache`` + ``prope_src`` / -``prope_tgt`` modules. When all three are ``None`` (Phase 1 text-only -path), the call is identical to ``CrossAttention.forward``. Otherwise -the neighbor branch runs:: +``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) @@ -69,7 +69,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # Parameter names match dreamfix (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 (added in Phase 5), not a per-key rename. + # 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) @@ -117,7 +117,7 @@ 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 (Phase 1 / 2 text-only). + forward path skips the neighbor branch (text-only mode). """ self.neighbor_kv_cache = ( None if context is None else self.compute_kv_neighbor(context) diff --git a/integrations/artifixer/artifixer/pipeline.py b/integrations/artifixer/artifixer/pipeline.py index 05ffb2c41..0dd21634d 100644 --- a/integrations/artifixer/artifixer/pipeline.py +++ b/integrations/artifixer/artifixer/pipeline.py @@ -20,27 +20,26 @@ camera w2c/Ks, optional VAE-encoded neighbor frames + neighbor camera matrices. -Phase 3.5a (this commit) lands the cache + ``initialize_cache``: +``initialize_cache``: - VAE-encoded condition latent and neighbor latent arrive pre-computed - from the caller (the dreamfix-side driver in Phase 4) so this - pipeline does not need its own VAE encoder. - - On ``initialize_cache`` we: - 1. Run the base ``WanInferencePipeline.initialize_cache`` (text - encoder + text K/V build) with ``image=None``. - 2. Project the neighbor latent through ``patch_embedding`` to get - the per-token neighbor context, then push it into every block's - :class:`ArtifixerCrossAttention.neighbor_kv_cache` via - ``ArtifixerDiTNetwork.initialize_neighbor_kv_caches``. - 3. Update the patches_x/patches_y RoPE coefficients on both PRoPE - modules and precompute the neighbor-side ``apply_fns`` (the - neighbor cameras are static across AR steps). - 4. Stash the full-rollout opacity / camera_rays / w2cs / Ks on the - cache for the per-AR ``generate`` slicing (Phase 3.5b). - -Phase 3.5b adds ``generate``: per-AR-chunk PRoPE-src precompute, -opacity-weighted latent mix, manual denoise loop with prepare_latents -renoise, and decode. + from the caller (the dreamfix-side driver in + ``model_eval/flashdreams_backend.py``) 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 @@ -148,9 +147,8 @@ def initialize_cache( # type: ignore[override] 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 (Phase 4 - driver) handles VAE encoding so this pipeline does not - need its own encoder. + ``[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 diff --git a/integrations/artifixer/artifixer/runner.py b/integrations/artifixer/artifixer/runner.py index c5320d55a..45c014b6e 100644 --- a/integrations/artifixer/artifixer/runner.py +++ b/integrations/artifixer/artifixer/runner.py @@ -15,12 +15,17 @@ """ArtiFixer DMD-distilled streaming T2V runner. -Phase 1 scaffold: this runner is a stripped clone of -``SelfForcingT2VRunner`` and accepts a single text prompt. Phase 3/4 will -introduce the ArtiFixer-specific conditioning surface (``rgb_rendered``, -``opacity``, neighbor frames, camera matrices) by accepting a path to a -dreamfix-format input bundle (e.g. an HDF5 reconstruction file plus a -caption file). +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; the dreamfix-side driver in +``dreamfix/model_eval/flashdreams_backend.py`` is the production entry +point that feeds those conditioning tensors. This runner stays +text-only as a lightweight harness. """ from __future__ import annotations @@ -68,7 +73,7 @@ class ArtifixerDmdT2VRunnerConfig(RunnerConfig): class ArtifixerDmdT2VRunner(Runner[ArtifixerDmdT2VRunnerConfig, WanInferencePipeline]): - """ArtiFixer DMD streaming T2V driver (Phase 1: text-only scaffold).""" + """ArtiFixer DMD streaming T2V driver (text-only smoke runner).""" config: ArtifixerDmdT2VRunnerConfig diff --git a/integrations/artifixer/artifixer/transformer.py b/integrations/artifixer/artifixer/transformer.py index d55bb0e1d..c5fc9eb6f 100644 --- a/integrations/artifixer/artifixer/transformer.py +++ b/integrations/artifixer/artifixer/transformer.py @@ -27,9 +27,9 @@ target-camera PRoPE precompute) into ``predict_flow`` -> ``WanDiTNetwork.forward`` via an :class:`ArtifixerCtrl` payload. -The pipeline (Phase 3.5) owns the per-chunk slicing + patchification + -PRoPE precompute updates and packages them into an ``ArtifixerCtrl`` -that flows through ``DiffusionModel.generate(input=ctrl)``. +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 @@ -52,9 +52,9 @@ class ArtifixerCtrl: 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 - (Phase 3.5); the transformer's ``predict_flow`` reads from here and - forwards into ``network_extra_kwargs``. + :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 diff --git a/integrations/artifixer/tests/test_smoke.py b/integrations/artifixer/tests/test_smoke.py index 04a8ede67..2455ea825 100644 --- a/integrations/artifixer/tests/test_smoke.py +++ b/integrations/artifixer/tests/test_smoke.py @@ -123,7 +123,7 @@ def test_artifixer_hyperparams_match_dreamfix_stage3() -> None: def test_artifixer_block_has_opacity_and_camera_mlps() -> None: - """Phase 2.1: ArtifixerBlock instances carry opacity + camera MLPs.""" + """ArtifixerBlock instances carry opacity + camera MLPs.""" import torch from artifixer.network.block import ArtifixerBlock from artifixer.network.dit import artifixer_embedding_dims @@ -151,7 +151,7 @@ def test_artifixer_block_has_opacity_and_camera_mlps() -> None: def test_artifixer_recipe_uses_artifixer_network() -> None: - """Phase 2.1: the shipped recipe wires up ArtifixerDiTNetwork.""" + """The shipped recipe wires up ArtifixerDiTNetwork.""" from artifixer.network.dit import ArtifixerDiTNetwork1pt3BConfig cfg = _runner_configs()["artifixer-dmd-wan2.1-t2v-1.3b"] @@ -162,7 +162,7 @@ def test_artifixer_recipe_uses_artifixer_network() -> None: def test_zero_pad_state_dict_transform_fills_missing_keys() -> None: - """Phase 2.1/2.2: the state_dict transform pads vanilla-Wan checkpoints. + """The state_dict transform pads vanilla-Wan checkpoints. Covers all 9 ArtiFixer-only keys per block: @@ -200,7 +200,7 @@ def test_zero_pad_state_dict_transform_fills_missing_keys() -> None: def test_artifixer_block_has_neighbor_cross_attn_projections() -> None: - """Phase 2.2: ArtifixerBlock's cross_attn carries the neighbor branch.""" + """ArtifixerBlock's cross_attn carries the neighbor branch.""" import torch from artifixer.network.block import ArtifixerBlock from artifixer.network.cross_attn import ArtifixerCrossAttention @@ -228,15 +228,15 @@ def test_artifixer_block_has_neighbor_cross_attn_projections() -> None: param = block.cross_attn.get_parameter(name) assert param.shape == expected_shape, f"cross_attn.{name} shape {param.shape}" - # Phase 2.2 zero-init contract: add_v_proj is the gate that keeps the - # neighbor branch contribution zero at load time, matching dreamfix + # Zero-init contract: add_v_proj is the gate that keeps the neighbor + # branch contribution zero at load time, matching dreamfix # transformer.py L687-688. 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: - """Phase 2.4: compute_kv_neighbor builds a static BlockKVCache, and + """compute_kv_neighbor builds a static BlockKVCache, and initialize_neighbor_cache toggles the per-module cache attribute. """ import torch diff --git a/integrations/artifixer/tests/test_state_dict_transform.py b/integrations/artifixer/tests/test_state_dict_transform.py index 7ad4d5a26..464e73a0a 100644 --- a/integrations/artifixer/tests/test_state_dict_transform.py +++ b/integrations/artifixer/tests/test_state_dict_transform.py @@ -31,6 +31,7 @@ from __future__ import annotations import json +import os from pathlib import Path import pytest @@ -40,19 +41,27 @@ artifixer_dmd_state_dict_transform, ) -AUDIT_PATH = Path( - "/lustre/fsw/portfolios/nvr/users/rdelutio/code/dreamfix/" - "merged_checkpoints/artifixer_dmd_1p3b_s3_2000/param_audit.json" -) + +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]: - if not AUDIT_PATH.exists(): + audit_path = _audit_path() + if audit_path is None: + pytest.skip( + "ARTIFIXER_PARAM_AUDIT_PATH is unset. Point it at the " + "param_audit.json produced by " + "``dreamfix/scripts/dump_artifixer_param_names.py`` to " + "exercise the full-merged-checkpoint path." + ) + if not audit_path.exists(): pytest.skip( - f"param audit JSON not found at {AUDIT_PATH}; run " + f"param audit JSON not found at {audit_path}; run " f"dreamfix/scripts/dump_artifixer_param_names.py first" ) - return json.loads(AUDIT_PATH.read_text()) + return json.loads(audit_path.read_text()) def _synthetic_state_dict(audit: dict[str, dict]) -> dict[str, torch.Tensor]: From 50bb9b0435ee8305c59c68fa4eebc5e0b21b87c6 Mon Sep 17 00:00:00 2001 From: Riccardo de Lutio Date: Tue, 12 May 2026 22:23:32 -0700 Subject: [PATCH 22/22] integrations/artifixer: rename dreamfix->artifixer + drop too-specific refs Public-friendly cleanup for the upstream MR: * "dreamfix" rewritten to "the ArtiFixer reference" / "the reference" throughout module docstrings, comments, and tests. The name "dreamfix" is an internal project handle that has no place in a public flashdreams plugin; the model name "ArtiFixer" is what readers actually need. * Dropped every "L###-L###" line-number citation pointing into the reference repo's source -- those rot the moment the reference moves a line, and they were diagnostic rather than substantive. * Removed two stale references to the deleted ``scripts/parity_harness.py`` (in block.py and transformer.py). * Renamed the test-side env var ``DREAMFIX_REPO_ROOT`` -> ``ARTIFIXER_REFERENCE_REPO_ROOT`` and the test function ``test_prope_matches_dreamfix_reference`` -> ``test_prope_matches_reference``. The companion ``scripts/run_prope_parity.sh`` is updated to match. * README.md drops the internal gitlab-master URL and reframes the "components" table without referring to the reference repo by name. No behaviour change. Static check (``scripts/static_check_artifixer.sh``) still passes. --- integrations/artifixer/README.md | 20 +++--- .../artifixer/artifixer/checkpoint.py | 29 ++++---- integrations/artifixer/artifixer/config.py | 20 +++--- .../artifixer/artifixer/latent_mix.py | 11 ++- .../artifixer/artifixer/network/__init__.py | 4 +- .../artifixer/artifixer/network/block.py | 70 +++++++++---------- .../artifixer/artifixer/network/cross_attn.py | 34 +++++---- .../artifixer/artifixer/network/dit.py | 9 +-- .../artifixer/artifixer/network/patches.py | 12 ++-- .../artifixer/artifixer/network/prope.py | 16 ++--- integrations/artifixer/artifixer/pipeline.py | 45 ++++++------ integrations/artifixer/artifixer/runner.py | 12 ++-- .../artifixer/artifixer/transformer.py | 53 +++++++------- integrations/artifixer/tests/test_patches.py | 8 +-- .../artifixer/tests/test_prope_parity.py | 51 +++++++------- integrations/artifixer/tests/test_smoke.py | 8 +-- .../tests/test_state_dict_transform.py | 22 +++--- scripts/run_prope_parity.sh | 16 +++-- 18 files changed, 214 insertions(+), 226 deletions(-) diff --git a/integrations/artifixer/README.md b/integrations/artifixer/README.md index a2bda4ec2..77f14b78f 100644 --- a/integrations/artifixer/README.md +++ b/integrations/artifixer/README.md @@ -11,26 +11,24 @@ ArtiFixer extends Wan 2.1 1.3B with: reconstruction-rendered frames; - 4-step DMD distillation (`FlowMatchScheduler(shift=5)`). -The reference implementation lives in the -[dreamfix repo](https://gitlab-master.nvidia.com/hturki/dreamfix) under -`model_training/net/transformer.py` and `model_training/pipeline/`. This -plugin ports it to FlashDreams' faster Wan stack (RingAttention, cuDNN, -`torch.compile`, CUDA graphs). +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 dreamfix stage-3 run config. | -| Per-block conditioning | Opacity + camera-ray MLPs, neighbor cross-attention, and PRoPE — parity-tested against dreamfix. | +| 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`. | -| dreamfix-compatible surface | `initialize_cache` accepts pre-encoded UMT5 prompts + VAE-encoded condition / neighbor latents so the dreamfix-side driver (`model_eval/flashdreams_backend.py`) can feed it directly. | -| Checkpoint loader | `state_dict_transform` for the merged ArtiFixer DMD safetensors produced by `dreamfix/scripts/merge_dcp_to_safetensors.py`. | +| 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 dreamfix-native `ArtifixerKvCachePipeline` after the +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. +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` diff --git a/integrations/artifixer/artifixer/checkpoint.py b/integrations/artifixer/artifixer/checkpoint.py index c944e347e..1853a3772 100644 --- a/integrations/artifixer/artifixer/checkpoint.py +++ b/integrations/artifixer/artifixer/checkpoint.py @@ -13,26 +13,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""State-dict transforms for loading ArtiFixer checkpoints into FlashDreams. +"""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`` -- the - transform :func:`zero_pad_artifixer_keys` zero-pads the ArtiFixer-only - keys so ``load_state_dict`` succeeds in strict mode. Zero-padding - matches dreamfix's initialization in - ``ArtifixerTransformerBlock.__init__`` L637-651 and produces a + * **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** produced by - ``dreamfix/scripts/merge_dcp_to_safetensors.py``. Built by + * **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 our new + ArtiFixer-only keys map cleanly onto the ``ArtifixerBlock`` / ``ArtifixerCrossAttention`` submodules without further per-key renames. """ @@ -118,11 +117,9 @@ def artifixer_dmd_state_dict_transform( ) -> dict[str, Tensor]: """Remap a merged ArtiFixer DMD safetensors state_dict onto WanDiTNetwork. - The merged safetensors produced by - ``dreamfix/scripts/merge_dcp_to_safetensors.py`` carries the HF - diffusers ``WanTransformer3DModel`` naming (e.g. - ``blocks.X.attn1.to_q.weight``) plus 270 ArtiFixer-only keys with - the ``attn2`` cross-attention prefix + 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``). @@ -151,8 +148,8 @@ def zero_pad_artifixer_keys( 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 identified by ``dreamfix/scripts/dump_artifixer_param_names.py``): + 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,) diff --git a/integrations/artifixer/artifixer/config.py b/integrations/artifixer/artifixer/config.py index 56a9fba43..ea028af39 100644 --- a/integrations/artifixer/artifixer/config.py +++ b/integrations/artifixer/artifixer/config.py @@ -50,9 +50,7 @@ from flashdreams.infra.runner import RunnerConfig from flashdreams.recipes.wan import WanVAEDecoderConfig -# Mirrors the dreamfix stage-3 run config (see -# wandb/.../run-*/files/config.yaml in -# artifixer-runs/artifixer-s3-dmd-1p3b-from-s2-10000-s1-15000-128g): +# 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 @@ -65,12 +63,13 @@ ARTIFIXER_NUM_INFERENCE_STEPS = 4 ARTIFIXER_TIMESTEP_SHIFT = 5.0 -# Default: load the merged ArtiFixer DMD safetensors produced by -# ``dreamfix/scripts/merge_dcp_to_safetensors.py``. 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. +# 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" ) @@ -104,8 +103,7 @@ raise RuntimeError( "ArtiFixer recipe requires a checkpoint path. Set " "``ARTIFIXER_DMD_CHECKPOINT_PATH`` to the merged DMD " - "safetensors produced by " - "``dreamfix/scripts/merge_dcp_to_safetensors.py``, or set " + "safetensors file, or set " "``ARTIFIXER_USE_BASE_WAN_WEIGHTS=1`` to fall back to " "vanilla Wan 2.1 1.3B base weights." ) diff --git a/integrations/artifixer/artifixer/latent_mix.py b/integrations/artifixer/artifixer/latent_mix.py index c007ea957..eb19e0b11 100644 --- a/integrations/artifixer/artifixer/latent_mix.py +++ b/integrations/artifixer/artifixer/latent_mix.py @@ -15,10 +15,9 @@ """Opacity-weighted latent mixing of noise + VAE-encoded reconstruction. -Mirrors ``ArtifixerPipelineBase.prepare_latents`` at -``dreamfix/model_training/pipeline/pipeline_base.py`` L98-L122. The pipeline -replaces the base ``DiffusionModel.generate`` initial-noise sample with a -per-AR-chunk mix:: +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) @@ -47,7 +46,7 @@ def opacity_weighted_latent_mix( vae_scale_factor_spatial: int, is_first_chunk: bool, ) -> Tensor: - """Compute the dreamfix opacity-weighted latent mix for one AR chunk. + """Compute the ArtiFixer opacity-weighted latent mix for one AR chunk. Args: condition: VAE-encoded reconstruction-rendered RGB, shape @@ -62,7 +61,7 @@ def opacity_weighted_latent_mix( 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 dreamfix. + ``is_first_chunk`` branch in the ArtiFixer reference. Returns: Mixed latent, same shape as ``condition``. diff --git a/integrations/artifixer/artifixer/network/__init__.py b/integrations/artifixer/artifixer/network/__init__.py index bf63d7456..0bcd05531 100644 --- a/integrations/artifixer/artifixer/network/__init__.py +++ b/integrations/artifixer/artifixer/network/__init__.py @@ -20,5 +20,5 @@ # from artifixer.network.block import ArtifixerBlock # from artifixer.network.prope import PropeDotProductAttention # -# Mirrors dreamfix's ``model_training/net/__init__.py`` (also empty) so the -# PRoPE module is testable with only ``torch`` installed. +# 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 index 93283dab6..fa74e7fd9 100644 --- a/integrations/artifixer/artifixer/network/block.py +++ b/integrations/artifixer/artifixer/network/block.py @@ -15,9 +15,8 @@ """ArtiFixer transformer block with opacity + camera + neighbor extensions. -Mirrors ``ArtifixerTransformerBlock`` in the dreamfix reference -(``model_training/net/transformer.py`` L617-L767). Per-block extensions on -top of :class:`Block`: +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 @@ -47,14 +46,15 @@ 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 dreamfix ``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)`` on the FD side 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. + 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 @@ -99,7 +99,7 @@ def __init__( 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. Matches dreamfix transformer.py L637-651. + # 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) @@ -123,13 +123,12 @@ def forward( # type: ignore[override] - ``opacity_extra`` / ``camera_extra``: per-token opacity and Plucker-camera-ray features added to the AdaLN-normed hidden - states before self-attention (matches dreamfix - ``ArtifixerTransformerBlock.forward`` L763-L767). + states before self-attention. - ``prope_src`` / ``prope_tgt`` / ``ignore_neighbors``: forwarded to :class:`ArtifixerCrossAttention.forward` to drive the PRoPE - neighbor branch (matches dreamfix L795-L807). The actual - neighbor K/V cache lives on the cross_attn module itself - (set via ``cross_attn.initialize_neighbor_cache``). + 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`. @@ -138,29 +137,28 @@ def forward( # type: ignore[override] "We expect to have called update_parameters_after_loading_checkpoint() " "before running the forward pass" ) - # Mirror dreamfix ``ArtifixerTransformerBlock.forward`` (transformer.py - # L727, L763, L787, L790, L810, L814) 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 (see ``scripts/parity_harness.py --capture_blocks``). - # The fp32 promotion below mirrors the dreamfix reference exactly: + # 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 dreamfix also - # promotes ``ff_output`` to fp32 inside the residual (L814). + # * 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. Note FD's bf16-throughout path - # remains the default; this override only kicks in when the - # ``ArtifixerBlock`` subclass is used (i.e. the artifixer recipe). + # >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 @@ -170,7 +168,7 @@ def forward( # type: ignore[override] # 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 which dreamfix's WanTransformerBlock norms use). + # 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) @@ -183,9 +181,9 @@ def forward( # type: ignore[override] ) x = (x.float() + y * e_chunks[2]).to(x_dtype) - # Cross-attn pre-norm in fp32 (mirrors dreamfix norm2 promotion at L790). - # The post-cross-attn residual is NOT promoted in dreamfix (L807) so - # we keep that in the input dtype here too. + # 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, @@ -195,6 +193,6 @@ def forward( # type: ignore[override] ) 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 dreamfix L814). + # 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 index f1de5ffec..33dd8599c 100644 --- a/integrations/artifixer/artifixer/network/cross_attn.py +++ b/integrations/artifixer/artifixer/network/cross_attn.py @@ -17,9 +17,7 @@ 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 used in dreamfix -(``model_training/net/transformer.py`` L670-L699 + the -``ArtifixerCrossAttnProcessor`` at L833-L940 that drives the forward): +diffusers ``WanAttention`` extension in the ArtiFixer reference: - ``add_k_proj`` / ``add_v_proj`` — Linear projections from latent neighbor context to K / V @@ -40,8 +38,8 @@ 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 dreamfix L893-L916), since -``_rope_precompute_coeffs`` does an int64 -> fp32 promotion internally. +trip (matching the reference), since ``_rope_precompute_coeffs`` does +an int64 -> fp32 promotion internally. """ from __future__ import annotations @@ -66,19 +64,19 @@ class ArtifixerCrossAttention(CrossAttention): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - # Parameter names match dreamfix (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. + # 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 dreamfix transformer.py - # L687-L688). 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. + # 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) @@ -148,8 +146,8 @@ def forward( # type: ignore[override] 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 dreamfix - (transformer.py L936). + Matches the diffusion-forcing CFG dropout knob in the + ArtiFixer reference. """ neighbor_kv_cache = self.neighbor_kv_cache batch_shape = x.shape[:-2] @@ -177,8 +175,8 @@ def forward( # type: ignore[override] 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 - # dreamfix transformer.py L893-L916. + # 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) diff --git a/integrations/artifixer/artifixer/network/dit.py b/integrations/artifixer/artifixer/network/dit.py index d228ac8fa..489b1f57d 100644 --- a/integrations/artifixer/artifixer/network/dit.py +++ b/integrations/artifixer/artifixer/network/dit.py @@ -28,9 +28,8 @@ WanDiTNetworkConfig, ) -# Wan2.1 VAE downsampling factors. Identical for T2V-1.3B and T2V-14B and -# what the dreamfix reference (``ArtifixerTransformer.__init__`` L228-237) -# uses to compute the per-block opacity / camera input dims. +# 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 @@ -42,8 +41,6 @@ def artifixer_embedding_dims( ) -> tuple[int, int]: """Compute (opacity_embedding_dim, camera_embedding_dim) for ArtiFixer. - Mirrors ``ArtifixerTransformer.__init__`` L228-237. - opacity_embedding_dim = vae_t * vae_s * ph * vae_s * pw camera_embedding_dim = vae_s * ph * vae_s * pw * 6 (6 = Plucker channels) @@ -79,7 +76,7 @@ class ArtifixerDiTNetworkConfig(WanDiTNetworkConfig): @dataclass class ArtifixerDiTNetwork1pt3BConfig(WanDiTNetwork1pt3BConfig): - """ArtiFixer 1.3B variant (matches dreamfix stage-3 model size).""" + """ArtiFixer 1.3B variant (matches the ArtiFixer reference model size).""" _target: type = field(default_factory=lambda: ArtifixerDiTNetwork) diff --git a/integrations/artifixer/artifixer/network/patches.py b/integrations/artifixer/artifixer/network/patches.py index ac82679e5..11d29cbae 100644 --- a/integrations/artifixer/artifixer/network/patches.py +++ b/integrations/artifixer/artifixer/network/patches.py @@ -15,7 +15,7 @@ """Patchification of opacity / Plucker camera-ray conditioning tensors. -Mirrors ``ArtifixerTransformer.forward`` L309-L341 in dreamfix: +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) @@ -54,7 +54,7 @@ def patchify_opacity( 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 dreamfix L309-L311). +Wan VAE's ``1 + 4`` temporal layout (mirrors the reference). Returns: ``(B, L, opacity_embedding_dim)`` where @@ -87,10 +87,10 @@ def patchify_camera_rays( ) -> Tensor: """Rearrange per-pixel Plucker rays into per-token features. - Handles dreamfix'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). + 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. diff --git a/integrations/artifixer/artifixer/network/prope.py b/integrations/artifixer/artifixer/network/prope.py index 1c8f9153c..e91fc82c7 100644 --- a/integrations/artifixer/artifixer/network/prope.py +++ b/integrations/artifixer/artifixer/network/prope.py @@ -23,12 +23,12 @@ """PRoPE attention with precomputed RoPE coefficients. -Verbatim port of dreamfix ``model_training/net/prope.py`` (matching paper -"Cameras as Relative Positional Encoding", arXiv 2507.10496). The math is -unchanged; the only difference is that ``invert_SE3`` is inlined here so -this module is self-contained. +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 follows the upstream docstring:: +Usage for cross-attention:: attn_src = PropeDotProductAttention(...) attn_tgt = PropeDotProductAttention(...) @@ -40,7 +40,7 @@ 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 dreamfix reference lives at +A numerical-parity unit test against the ArtiFixer reference lives at ``integrations/artifixer/tests/test_prope_parity.py``. """ @@ -295,8 +295,8 @@ def _invert_K(Ks: torch.Tensor) -> torch.Tensor: def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor: """Invert a 4x4 SE(3) matrix. - Mirrors ``dreamfix/model_training/utils/pose_utils.invert_SE3`` so this - module has no cross-repo dependencies. + 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) diff --git a/integrations/artifixer/artifixer/pipeline.py b/integrations/artifixer/artifixer/pipeline.py index 0dd21634d..ae21013df 100644 --- a/integrations/artifixer/artifixer/pipeline.py +++ b/integrations/artifixer/artifixer/pipeline.py @@ -23,9 +23,7 @@ ``initialize_cache``: - VAE-encoded condition latent and neighbor latent arrive pre-computed - from the caller (the dreamfix-side driver in - ``model_eval/flashdreams_backend.py``) so this pipeline does not - need its own VAE encoder. + 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 @@ -142,8 +140,7 @@ def initialize_cache( # type: ignore[override] 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 - dreamfix-backed driver in :mod:`model_eval` passes them - directly). + 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, @@ -256,11 +253,10 @@ def _initialize_cache_from_text_embeddings( """Build the base cache when prompts are pre-encoded. Mirrors :meth:`WanInferencePipeline.initialize_cache` minus the UMT5 - forward / negative-prompt encoding / I2V-image plumbing -- the - dreamfix driver already has the encoded prompts and never uses - I2V images. CFG is disabled per the dreamfix ``ArtifixerKvCachePipeline`` - contract (negative_prompt is ignored), so we don't build - ``negative_text_embeddings`` even if guidance_scale > 1. + 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 @@ -283,7 +279,7 @@ def _initialize_cache_from_text_embeddings( def _encode_neighbor_context(self, neighbor_latent: Tensor) -> Tensor: """Run ``patch_embedding`` over neighbor latents and flatten to tokens. - Mirrors dreamfix ``ArtifixerTransformer.forward`` L361-L362:: + 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) @@ -301,8 +297,8 @@ def _chunk_frame_ranges( ) -> tuple[int, int, int, int]: """Return ``(lat_start, lat_end, input_start, input_end)`` for this AR chunk. - Mirrors the dreamfix bookkeeping in - ``kv_cache_pipeline.generate_samples_from_batch`` L201-L204: + 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 @@ -361,8 +357,8 @@ def generate( # type: ignore[override] 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 dreamfix - ``kv_cache_pipeline.generate_samples_from_batch`` L211-L264). + 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. @@ -428,7 +424,7 @@ def generate( # type: ignore[override] patch_size=tcfg.network.patch_size, ) - # Start the AR cache for this chunk (mirrors DiffusionModel.generate L165). + # 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). @@ -449,18 +445,19 @@ def generate( # type: ignore[override] vae_scale_factor_spatial=vae_s, is_first_chunk=is_first, ) - # FlashDreams ``patchify_and_maybe_split_cp`` expects ``(B, T, C, H, W)`` - # (pattern ``... (t kt) c (h kh) (w kw)``), but dreamfix's VAE encoder - # (and therefore ``cache.condition_latent`` / ``latent_unpatched``) is - # in diffusers convention ``(B, C, T, H, W)``. Permute before patchify. + # ``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. Mirrors - # ``ArtifixerKvCachePipeline.generate_samples_from_batch`` L238-L264: - # at every non-exit step we replace the next-iteration ``noisy`` - # with a fresh ``opacity_weighted_latent_mix(...)``, not the base + # 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 diff --git a/integrations/artifixer/artifixer/runner.py b/integrations/artifixer/artifixer/runner.py index 45c014b6e..90dcf29e5 100644 --- a/integrations/artifixer/artifixer/runner.py +++ b/integrations/artifixer/artifixer/runner.py @@ -22,10 +22,8 @@ The full ArtiFixer-specific conditioning surface (``rgb_rendered``, ``opacity``, neighbor frames, camera matrices) is wired in through the :class:`ArtifixerInferencePipeline.initialize_cache` / ``generate`` -path; the dreamfix-side driver in -``dreamfix/model_eval/flashdreams_backend.py`` is the production entry -point that feeds those conditioning tensors. This runner stays -text-only as a lightweight harness. +path; production drivers call those directly with conditioning tensors. +This runner stays text-only as a lightweight harness. """ from __future__ import annotations @@ -63,9 +61,9 @@ class ArtifixerDmdT2VRunnerConfig(RunnerConfig): whose first non-empty line is read as the prompt.""" total_blocks: int = 3 - """Number of autoregressive chunks to generate. Matches dreamfix's - ``num_frames=81`` / ``frames_per_block=7`` (21 latent frames / 7 = 3 - chunks for the Wan VAE temporal stride of 4).""" + """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 diff --git a/integrations/artifixer/artifixer/transformer.py b/integrations/artifixer/artifixer/transformer.py index c5fc9eb6f..b8f87572c 100644 --- a/integrations/artifixer/artifixer/transformer.py +++ b/integrations/artifixer/artifixer/transformer.py @@ -71,14 +71,14 @@ class ArtifixerCtrl: ``[*batch_shape, L, camera_embedding_dim]``.""" ignore_neighbors: bool = False - """Diffusion-forcing CFG-dropout knob (matches dreamfix L936). Inference - leaves this ``False``; training-time DMD distillation may toggle it.""" + """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`` (base.py L163-164) is a no-op for us.""" + ``DiffusionModel.generate`` is a no-op for us.""" @dataclass(kw_only=True) @@ -114,7 +114,7 @@ def __init__(self, config: ArtifixerWanTransformerConfig) -> None: freq_scale=config.prope_freq_scale, ) - # Match the dreamfix dtype convention. + # 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) @@ -133,9 +133,9 @@ def predict_flow( # type: ignore[override] ``**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`` - (network.py L438-449), where :class:`ArtifixerBlock.forward` accepts - ``opacity_extra`` / ``camera_extra`` / ``prope_src`` / ``prope_tgt`` / + 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 @@ -171,26 +171,25 @@ def predict_flow( # type: ignore[override] ) def finalize_kv_cache(self, *args: Any, **kwargs: Any) -> None: - """No-op: artifixer (dreamfix reference) does not finalize the KV cache. - - The dreamfix ``ArtifixerKvCachePipeline.generate_samples_from_batch`` - 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; that extra forward writes a *different* KV state than - dreamfix'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 - (calls 4, 8 in ``scripts/parity_harness.py``'s per-step diff). - - Skipping the extra forward here aligns FlashDreams' KV-cache - semantics with the dreamfix reference: the cache going into AR - chunk ``N+1`` is the one written by the final denoise step of AR - chunk ``N``, identical to what dreamfix uses. ``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. + """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 diff --git a/integrations/artifixer/tests/test_patches.py b/integrations/artifixer/tests/test_patches.py index 7c5a41211..44dc86344 100644 --- a/integrations/artifixer/tests/test_patches.py +++ b/integrations/artifixer/tests/test_patches.py @@ -88,15 +88,15 @@ def test_patchify_camera_rays_latent_rate_shape() -> None: 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 dreamfix (transformer.py L330-L340) groups t4 input frames + 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 (see - ``model_training/pipeline/kv_cache_pipeline.py`` L213), so this branch - is unused. We still cover it for parity with the dreamfix forward. + 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( diff --git a/integrations/artifixer/tests/test_prope_parity.py b/integrations/artifixer/tests/test_prope_parity.py index 81b5be27d..705409e76 100644 --- a/integrations/artifixer/tests/test_prope_parity.py +++ b/integrations/artifixer/tests/test_prope_parity.py @@ -16,19 +16,21 @@ """Numerical parity test for the ported PRoPE attention module. Compares :class:`artifixer.network.prope.PropeDotProductAttention` against -the dreamfix reference at ``model_training/net/prope.py``. When dreamfix is -not importable (the CI image only ships flashdreams), the test is split: +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_dreamfix_reference`` is skipped unless - ``model_training.net.prope`` can be imported (set - ``DREAMFIX_REPO_ROOT`` to its repo root to enable). + * ``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 + dreamfix env is available, while still exercising the math today. +GPU + reference env is available, while still exercising the math today. """ from __future__ import annotations @@ -49,11 +51,12 @@ def _deterministic_torch() -> None: 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 dreamfix - reference has the same property. Tests therefore run at fp32 with fp32- - appropriate tolerances rather than fp64. dreamfix's ``_lift_K`` also - hard-codes float32 output (no ``dtype=`` in ``torch.zeros``), so any - cross-implementation comparison must use fp32 inputs. + 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) @@ -143,9 +146,9 @@ def test_prope_apply_to_q_changes_input_for_nontrivial_cameras() -> None: assert diff > 1e-3, f"expected non-trivial PRoPE transform, got max abs diff {diff}" -def _try_import_dreamfix_reference() -> object | None: - """Return dreamfix's PropeDotProductAttention if importable.""" - repo_root = os.environ.get("DREAMFIX_REPO_ROOT") +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: @@ -156,25 +159,25 @@ def _try_import_dreamfix_reference() -> object | None: @pytest.mark.skipif( - _try_import_dreamfix_reference() is None, + _try_import_reference() is None, reason=( - "dreamfix not importable. Set DREAMFIX_REPO_ROOT to the dreamfix " - "checkout, or run from an env that already has it on sys.path." + "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_dreamfix_reference() -> None: - """Numerical parity vs dreamfix at fp32. +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 dreamfix's ``_lift_K`` hard-codes float32 + 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 carries the dreamfix dtype bug fix - (passes ``dtype=Ks.dtype``), making it slightly more dtype-robust; - at fp32 the two are bit-identical. + 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_dreamfix_reference() + RefPRoPE = _try_import_reference() assert RefPRoPE is not None inp = _build_inputs() diff --git a/integrations/artifixer/tests/test_smoke.py b/integrations/artifixer/tests/test_smoke.py index 2455ea825..826a8305d 100644 --- a/integrations/artifixer/tests/test_smoke.py +++ b/integrations/artifixer/tests/test_smoke.py @@ -108,7 +108,7 @@ def test_entry_points_discoverable_when_installed() -> None: ) -def test_artifixer_hyperparams_match_dreamfix_stage3() -> None: +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 @@ -143,7 +143,7 @@ def test_artifixer_block_has_opacity_and_camera_mlps() -> None: 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 dreamfix transformer.py L637-651. + # 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) @@ -229,8 +229,8 @@ def test_artifixer_block_has_neighbor_cross_attn_projections() -> None: 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 dreamfix - # transformer.py L687-688. + # 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) diff --git a/integrations/artifixer/tests/test_state_dict_transform.py b/integrations/artifixer/tests/test_state_dict_transform.py index 464e73a0a..20b20f590 100644 --- a/integrations/artifixer/tests/test_state_dict_transform.py +++ b/integrations/artifixer/tests/test_state_dict_transform.py @@ -15,11 +15,13 @@ """Key-naming tests for the merged-DMD state_dict transform. -CPU-only; uses the ``param_audit.json`` produced by -``dreamfix/scripts/dump_artifixer_param_names.py`` to assert that +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()``. +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. @@ -52,14 +54,14 @@ def _load_audit() -> dict[str, dict]: if audit_path is None: pytest.skip( "ARTIFIXER_PARAM_AUDIT_PATH is unset. Point it at the " - "param_audit.json produced by " - "``dreamfix/scripts/dump_artifixer_param_names.py`` to " - "exercise the full-merged-checkpoint path." + "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}; run " - f"dreamfix/scripts/dump_artifixer_param_names.py first" + 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()) @@ -110,8 +112,8 @@ 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 dreamfix names - match the flashdreams names verbatim. + 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), diff --git a/scripts/run_prope_parity.sh b/scripts/run_prope_parity.sh index 0971d67a6..ba39ea604 100755 --- a/scripts/run_prope_parity.sh +++ b/scripts/run_prope_parity.sh @@ -12,22 +12,26 @@ #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 dreamfix -# reference at model_training/net/prope.py. +# 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 dreamfix repo is mounted via /lustre, so we don't need uv sync. +# 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)}" -DREAMFIX_REPO_ROOT="${DREAMFIX_REPO_ROOT:-/lustre/fsw/portfolios/nvr/users/rdelutio/code/dreamfix}" +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:${DREAMFIX_REPO_ROOT}:\${PYTHONPATH:-} && \ -export DREAMFIX_REPO_ROOT='${DREAMFIX_REPO_ROOT}' && \ +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"