diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index bf2e4500f..3ec87fb79 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -185,6 +185,7 @@ steps: - label: ":bar_chart: SSIM Tests" key: "ssim" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || build.env("TEST_SCOPE") == "scheduled" || @@ -211,6 +212,7 @@ steps: - label: ":test_tube: LoRA Inference Tests" key: "lora-inference" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -232,6 +234,7 @@ steps: - label: ":test_tube: LoRA Extraction Tests" key: "lora-extraction" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -253,6 +256,7 @@ steps: - label: ":test_tube: Training Tests" key: "training" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -276,6 +280,7 @@ steps: - label: ":test_tube: Distillation DMD Tests" key: "distillation" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -297,6 +302,7 @@ steps: - label: ":test_tube: Self-Forcing Tests" key: "self-forcing" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -318,6 +324,7 @@ steps: - label: ":test_tube: LoRA Training Tests" key: "lora-training" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -341,6 +348,7 @@ steps: - label: ":test_tube: Training Tests VSA" key: "training-vsa" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -364,6 +372,7 @@ steps: - label: ":test_tube: Inference Tests VMoBA" key: "inference-vmoba" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -385,6 +394,7 @@ steps: - label: ":test_tube: Performance Tests" key: "performance" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -406,6 +416,7 @@ steps: - label: ":test_tube: API Server Tests" key: "api-server" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -427,6 +438,7 @@ steps: - label: ":test_tube: Train Framework Tests" key: "train-framework" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && @@ -448,6 +460,7 @@ steps: - label: ":test_tube: Eval Metrics Tests" key: "eval" + depends_on: "golden-gate" if: | build.env("TEST_SCOPE") == "full" || (build.env("TEST_SCOPE") == "merge" && diff --git a/.buildkite/scripts/lanes/golden_gate.sh b/.buildkite/scripts/lanes/golden_gate.sh index eb233469b..1875d4d08 100755 --- a/.buildkite/scripts/lanes/golden_gate.sh +++ b/.buildkite/scripts/lanes/golden_gate.sh @@ -13,7 +13,7 @@ if [ -z "$selected" ]; then selected=all fi if [ "$selected" = all ]; then - exec pytest "$golden_root" -vs + exec pytest "$golden_root" -xvs fi [[ $selected =~ ^test_[a-z0-9_]+\.py(,test_[a-z0-9_]+\.py)*$ ]] || { @@ -32,4 +32,4 @@ for golden_file in "${golden_files[@]}"; do golden_paths+=("$golden_path") done -exec pytest "${golden_paths[@]}" -vs +exec pytest "${golden_paths[@]}" -xvs diff --git a/.buildkite/scripts/lanes/transformer.sh b/.buildkite/scripts/lanes/transformer.sh index 9041f10df..56b4d694b 100755 --- a/.buildkite/scripts/lanes/transformer.sh +++ b/.buildkite/scripts/lanes/transformer.sh @@ -2,4 +2,8 @@ # Canonical Slurm CI selection for the transformer lane. set -euo pipefail +# The existing block reference records an absent FASTVIDEO_FA4 (FA2). Keep +# that reference identity; the component lane also selects FA2 explicitly. +env -u FASTVIDEO_FA4 pytest ./fastvideo/tests/golden_gate/test_wan_t2v.py -xvs +pytest ./fastvideo/tests/golden_gate/test_wan_causal.py -xvs exec pytest ./fastvideo/tests/transformers -vs diff --git a/.buildkite/scripts/lanes/vae.sh b/.buildkite/scripts/lanes/vae.sh index 7cc60c451..abd6ac4d3 100755 --- a/.buildkite/scripts/lanes/vae.sh +++ b/.buildkite/scripts/lanes/vae.sh @@ -2,4 +2,5 @@ # Canonical Slurm CI selection for the VAE lane. set -euo pipefail +pytest ./fastvideo/tests/golden_gate/test_wan_vae.py -xvs exec pytest ./fastvideo/tests/vaes -vs diff --git a/.github/scripts/plan_merge_ci.py b/.github/scripts/plan_merge_ci.py index 074a3e652..db79723ce 100644 --- a/.github/scripts/plan_merge_ci.py +++ b/.github/scripts/plan_merge_ci.py @@ -211,8 +211,8 @@ FAMILY_COVERAGE = ( ("test_turbodiffusion_similarity.py", ), ), FamilyCoverage( - re.compile(r"(^|[/_.-])wan(video)?([/_.-]|$)"), - ("test_wan_t2v.py", ), + re.compile(r"(^|[/_.-])wan(video|vae)?([/_.-]|$)"), + ("test_wan_t2v.py", "test_wan_vae.py", "test_wan_causal.py", "test_wan_denoising.py"), ( "test_causal_similarity.py", "test_wan_i2v_similarity.py", @@ -291,6 +291,18 @@ def _family_coverage(path: str) -> tuple[set[str], set[str]]: if family.pattern.search(normalized): golden.update(family.golden_tests) ssim.update(family.ssim_tests) + # Select the component actually touched, including compatibility paths. + # Family configs/pipeline wiring can affect all four Wan gates. + if re.search(r"(^|[/_.-])wan(video|vae)?([/_.-]|$)", normalized): + if (normalized.endswith(("/wan/vae.py", "/wan/vae_config.py", "/vaes/wanvae.py")) + or normalized.endswith("/wan/stages/conditioning.py")): + golden = {"test_wan_vae.py"} + elif normalized.endswith(("/wan/causal_transformer.py", "/dits/causal_wanvideo.py", + "/wan/stages/causal_denoising.py")): + golden = {"test_wan_causal.py"} + elif (normalized == "fastvideo/models/dits/wanvideo.py" + or normalized.endswith(("/wan/transformer.py", "/wan/stages/denoising.py", "/wan/stages/dmd.py"))): + golden = {"test_wan_t2v.py", "test_wan_denoising.py"} return golden, ssim diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e21adcfe3..e10ea0b7b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ exclude: | tests/.*| scripts/.*| fastvideo/dataset/.*| - fastvideo/models/(?!wan/(config|vae_config|__init__)\.py$).*| + fastvideo/models/(?!wan/(config|vae_config|pipeline_config|definition|__init__)\.py$).*| ^apps/dreamverse/web/.*| examples/.*| \.agents/.*| diff --git a/docs/contributing/ci_architecture.md b/docs/contributing/ci_architecture.md index 1a630dd9f..d4868e957 100644 --- a/docs/contributing/ci_architecture.md +++ b/docs/contributing/ci_architecture.md @@ -141,6 +141,24 @@ it), while a GitHub outage or a >25 min wait lets it run anyway (fail open). The complete static graph remains available through `/test full`; path selection never deletes or dynamically invents a Buildkite step. +Integration steps depend on `golden-gate`. A failed selected golden prevents +the expensive downstream jobs from starting; `/test full` still selects all +twenty lanes. A condition-skipped golden satisfies the dependency, so direct +lane reruns, scheduled SSIM, and training-only merge plans keep their existing +meaning. This follows Buildkite's +[conditional dependency rules](https://buildkite.com/docs/pipelines/configure/depends-on). +No dependency allows failures. The trusted uploader accepts either the old +complete graph or the complete golden-first graph during rollout, and rejects +partial or arbitrary dependency changes. + +The six automatic Fastcheck lanes are unchanged. Within VAE and transformer +lanes, small Wan goldens run before independent component parity. Wan paths +select matching VAE, dense/trajectory, or causal-cache goldens; shared Wan +config and pipeline wiring select all four. Shared runtime changes continue +to select broader coverage. The existing block references retain their exact +environment identity; new tensor gates distinguish the effective FA2/FA4 +switch and VAE gates do not depend on an unused attention backend. + | Lane | Public `TEST_TYPE` | GPUs | Typical merge trigger | |---|---|---:|---| | Encoder | `encoder` | 1 | Universal Fastcheck | diff --git a/docs/contributing/coding_agents.md b/docs/contributing/coding_agents.md index 7475812d1..6c3dd335e 100644 --- a/docs/contributing/coding_agents.md +++ b/docs/contributing/coding_agents.md @@ -475,14 +475,17 @@ FastVideo integration. - Text encoder: `fastvideo/models/encoders/t5.py` 2. Parity test the core components. + - Start with `bash scripts/validate_wan.sh all`: contracts and tiny goldens. - Example tests: `fastvideo/tests/transformers/test_wanvideo.py`, `fastvideo/tests/vaes/test_wan_vae.py`, `fastvideo/tests/encoders/test_t5_encoder.py` 3. Pipeline wiring. - Pipeline: `fastvideo/pipelines/basic/wan/wan_pipeline.py` - - Pipeline config: `fastvideo/configs/pipelines/wan.py` - - Sampling defaults: `fastvideo/pipelines/basic/wan/profiles.py` + - Denoising and first-frame preparation: `fastvideo/pipelines/basic/wan/stages/` + - Variant definitions: `fastvideo/models/wan/definition.py` + - Pipeline config: `fastvideo/models/wan/pipeline_config.py` + - Sampling defaults: `fastvideo/pipelines/basic/wan/presets.py` 4. Minimal example. - Script: `examples/inference/basic/basic.py` diff --git a/docs/contributing/testing.md b/docs/contributing/testing.md index 1b65903fd..d9e7a4209 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -32,18 +32,56 @@ failure before starting heavier dependent checks: 4. Run full-quality renders or broad suites when required by the change, an explicit request, or CI policy, rather than on every edit. -A golden must cover the component being changed. Wan's existing golden checks -transformer block 0; it does not cover the VAE. Until a Wan VAE golden exists, -use its focused encode/decode parity test. Keep independent-reference parity -and end-to-end coverage for changes outside a golden's tested boundary. +A golden must cover the component being changed. Wan has four small gates: + +| Gate | Boundary | +|---|---| +| `test_wan_t2v.py` | Dense transformer block 0 | +| `test_wan_vae.py` | FP32 encode, BF16 decode, streaming, and cache reset | +| `test_wan_causal.py` | Real block weights, cache append/rewrite, and sink eviction | +| `test_wan_denoising.py` | Three real 1.3B DiT/UniPC steps with fixed prompt embeddings | + +These use an immutable Wan checkpoint revision and only download the required +component. The trajectory gate needs no tokenizer, text encoder, VAE, or video +reference. Weight-free stage tests also exercise every step of a 50-step UniPC +loop, CFG caching, expert switching, conditioning layouts, and DMD RNG order. +These checks do not replace independent Diffusers parity or end-to-end SSIM. Use the golden's matching GPU, dtype, backend, and runtime. A missing reference or environment mismatch is not a pass. Do not replace a reference with the candidate output just to clear a failure. For a relocation, the unchanged parent is the baseline; two imports of the same class are not numerical proof. -This is the developer iteration order, not a change to required CI checks or -their scheduling. See [CI/CD Architecture](ci_architecture.md) for current policy. +The VAE and transformer CI lanes run their Wan component goldens first. +Selected integration lanes wait for the golden lane in merge/full builds. +See [CI/CD Architecture](ci_architecture.md) for direct-rerun and skip semantics. + +For Wan, one command enforces the local ordering and stops on failure: + +The initial contracts include `tests/api/test_wan_definitions.py`: all registered +Wan aliases, local-manifest detector precedence, sampling/precision defaults, +config isolation, and legacy serialized-config compatibility. These require no +weights or Hub access; the current package import still needs its prepared +runtime environment. + +```bash +bash scripts/validate_wan.sh vae # contracts, then the VAE golden +bash scripts/validate_wan.sh dense parity # contracts, goldens, Diffusers parity +bash scripts/validate_wan.sh all default # then focused T2V/I2V/causal SSIM +``` + +The second argument is an upper validation tier, not a reference override. +Supply the GPU/backend/runtime and SSIM model/tier settings that match the +reference. The script never updates references. Causal component coverage is +a block-cache fingerprint, not independent full causal-pipeline parity. + +New named-tensor gates write a missing output to `*.candidate.pt` and fail; +running candidate code again cannot turn it into an approved reference. Seed +from unchanged, pushed source, verify two independent processes bit-for-bit, +then review and publish only the new reference files. Preserve the baseline +source SHA, test recipe SHA, checkpoint revision, runtime, and comparison +receipt with the run artifacts. Never overwrite an existing video reference +as a side effect of adding a tensor gate. Examples of focused checks: diff --git a/docs/design/overview.md b/docs/design/overview.md index c5e0f9986..f76039686 100644 --- a/docs/design/overview.md +++ b/docs/design/overview.md @@ -11,7 +11,7 @@ FastVideo maps a Diffusers-style repo into a pipeline like this: - `fastvideo/models/*`: model implementations (DiT, VAE, encoders, upsamplers). - `fastvideo/configs/models/*`: arch configs and `param_names_mapping` for weight name translation. -- `fastvideo/models/wan/`: Wan's dense transformer, VAE, and component configs +- `fastvideo/models/wan/`: Wan's transformers, VAE, configs, and variant definitions together; the old Wan component/config modules remain compatibility imports. - `fastvideo/configs/pipelines/*`: pipeline wiring (component classes + names). - `fastvideo/api/sampling_param.py`: runtime sampling parameters. @@ -54,13 +54,25 @@ runtime parameters consistent: co-located with `transformer.py`. - `fastvideo/models/wan/vae_config.py`: Wan's VAE architecture and runtime settings, co-located with `vae.py`. +- `fastvideo/models/wan/pipeline_config.py`: Wan component and pipeline defaults; + `configs/pipelines/wan.py` remains a compatibility import. +- `fastvideo/models/wan/definition.py`: data-only Wan variants linking HF aliases, + config classes, presets, workload metadata, and default sampling algorithms. - `fastvideo/configs/pipelines/`: pipeline wiring and required components. - `fastvideo/api/sampling_param.py`: sampling parameters (steps, frames, - guidance scale, resolution, fps). Defaults come from profiles in - `fastvideo/pipelines/basic//profiles.py`. + guidance scale, resolution, fps). Defaults come from presets in + `fastvideo/pipelines/basic//presets.py`. - `fastvideo/registry.py`: unified registry for pipeline config + sampling - defaults and model metadata resolution, defined via explicit - `register_configs(...)` blocks (no separate dict registries). + defaults and model metadata resolution. Wan registrations consume its + family-local definitions; other families use `register_configs(...)` blocks. + +Wan definitions reference existing defaults rather than copying them. Dense +UniPC, dense DMD, and causal DMD remain separate sampling algorithms. Dense +DMD uses a full training-noise scheduler with shift 8.0, separate from the +configurable scheduler mutated during timestep preparation. The catalog does +not override checkpoint manifests, user pipeline overrides, or component +precision settings. HF IDs, local checkpoints, and old config imports retain +their existing resolution behavior, including first-match detector ordering. `FastVideoArgs` (in `fastvideo/fastvideo_args.py`) provides runtime settings and is passed into pipeline construction and stages. @@ -151,8 +163,9 @@ How this maps to FastVideo: - `UMT5EncoderModel` -> `fastvideo/models/encoders/t5.py` - `T5TokenizerFast` -> loaded via HF in `fastvideo/models/loader/` - `UniPCMultistepScheduler` -> loaded via Diffusers scheduler utilities -- Pipeline defaults -> `fastvideo/configs/pipelines/wan.py` -- Sampling defaults -> `fastvideo/pipelines/basic/wan/profiles.py` +- Variant definitions -> `fastvideo/models/wan/definition.py` +- Pipeline defaults -> `fastvideo/models/wan/pipeline_config.py` +- Sampling defaults -> `fastvideo/pipelines/basic/wan/presets.py` ## Pipeline system diff --git a/docs/inference/architecture.md b/docs/inference/architecture.md index ed58c9a16..c86038750 100644 --- a/docs/inference/architecture.md +++ b/docs/inference/architecture.md @@ -86,7 +86,7 @@ from `model_index.json`. ``` PipelineConfig (fastvideo/configs/pipelines/base.py) -├── WanT2V480PConfig (fastvideo/configs/pipelines/wan.py) +├── WanT2V480PConfig (fastvideo/models/wan/pipeline_config.py) │ ├── WanT2V720PConfig │ └── WanI2V480PConfig ├── HunyuanConfig (fastvideo/configs/pipelines/hunyuan.py) @@ -107,6 +107,11 @@ Model-specific subclasses override defaults. For example, `WanT2V480PConfig` sets `flow_shift=3.0` and uses `WanVideoConfig` as its DiT config. +Wan's `models/wan/definition.py` links each registered variant to its pipeline +config and sampling preset. The shared registry consumes these definitions +without changing detector precedence or checkpoint/override-based pipeline +selection. `configs/pipelines/wan.py` remains a compatibility import. + ### ModelConfig / ArchConfig (`fastvideo/configs/models/base.py`) `ModelConfig` wraps an `ArchConfig` using `__getattr__` proxy — attribute @@ -256,6 +261,16 @@ Specialized variants: `CausalDenoisingStage`, `LTX2DenoisingStage`, `SRDenoisingStage`, `LTX2AudioDecodingStage`, `SD35ConditioningStage`, `LTX2TextEncodingStage`, `LTX2LatentPreparationStage`. +Wan owns its sampling recipes under `basic/wan/stages/`. `WanDenoisingStage` +specializes input packing, expert selection, timesteps, and first-frame +restoration around the shared dense loop. `WanFirstFrameEncodingStage` +produces normalized `ForwardBatch.first_frame_latent` before sampling; the +sampler no longer executes a VAE. Dense DMD and the two causal samplers have +family-local implementations and explicit scheduler ownership. Standard and +DMD causal sampling share cache allocation, not their sampling algorithm. +Legacy imports from `stages/` remain compatibility aliases. Sampling invariants +also live beside the code in `fastvideo/pipelines/basic/wan/AGENTS.md`. + ### Verification System (`fastvideo/pipelines/stages/validators.py`) `StageValidators` (aliased as `V`) provides static validators: diff --git a/fastvideo/AGENTS.md b/fastvideo/AGENTS.md index 07839f1e7..ed20c552f 100644 --- a/fastvideo/AGENTS.md +++ b/fastvideo/AGENTS.md @@ -48,6 +48,7 @@ fastvideo/ | Add a new model component | `models//.py` + `configs/models//.py` | | Edit Wan's dense transformer or arch config | `models/wan/transformer.py` + `models/wan/config.py`; old paths are compatibility shims | | Edit Wan's VAE or its config | `models/wan/vae.py` + `models/wan/vae_config.py`; old paths are compatibility shims | +| Edit a Wan variant or pipeline defaults | `models/wan/definition.py` + `models/wan/pipeline_config.py`; sampling presets remain pipeline-local | | Wire an existing model into a new pipeline | `pipelines/basic//presets.py` + reuse stages from `pipelines/stages/` | | Add a converter | `scripts/checkpoint_conversion/_to_*.py` (separate dir, separate AGENTS.md) | | Add an attention backend | `attention/backends/.py` + register in selector | @@ -68,6 +69,7 @@ These dirs are listed in `.pre-commit-config.yaml` `exclude`: Editing files there will NOT trigger yapf/ruff/mypy/codespell. Format manually if a sibling file shows clear style; do not introduce new violations. -Exception: `models/wan/config.py`, `models/wan/vae_config.py`, and -`models/wan/__init__.py` are checked, so the family-local configs retain their -prior lint coverage. +Exception: `models/wan/config.py`, `models/wan/vae_config.py`, +`models/wan/pipeline_config.py`, `models/wan/definition.py`, and +`models/wan/__init__.py` are checked, so family-local configs retain lint +coverage. diff --git a/fastvideo/configs/AGENTS.md b/fastvideo/configs/AGENTS.md index 0ebd8d772..7432796c4 100644 --- a/fastvideo/configs/AGENTS.md +++ b/fastvideo/configs/AGENTS.md @@ -7,8 +7,10 @@ Two layers of dataclass configs feed every pipeline: **arch configs** (what the Wan's transformer config is co-located in `fastvideo/models/wan/config.py`, and its VAE config in `fastvideo/models/wan/vae_config.py`. `models/dits/wanvideo.py` and `models/vaes/wanvae.py` under this directory are compatibility re-exports; -pipeline configs remain in `pipelines/wan.py`. Other model families retain the -layout below. +pipeline configs now live in `fastvideo/models/wan/pipeline_config.py`, with +`pipelines/wan.py` retaining compatibility exports. Wan variant registration +comes from `fastvideo/models/wan/definition.py`; other model families retain +the layout below. ## Layout diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 0322e78d5..821cf944a 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -13,7 +13,7 @@ from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config from fastvideo.configs.models.dits.mmaudio import MMAudioArchConfig, MMAudioTransformerConfig from fastvideo.configs.models.dits.stable_audio import StableAudioConfig -from fastvideo.configs.models.dits.wanvideo import WanVideoConfig +from fastvideo.models.wan.config import WanVideoConfig from fastvideo.configs.models.dits.zimage import ZImageDiTConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig diff --git a/fastvideo/configs/models/dits/dreamx_world.py b/fastvideo/configs/models/dits/dreamx_world.py index 45be46fe7..bbd5177bb 100644 --- a/fastvideo/configs/models/dits/dreamx_world.py +++ b/fastvideo/configs/models/dits/dreamx_world.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig -from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig +from fastvideo.models.wan.config import WanVideoArchConfig @dataclass diff --git a/fastvideo/configs/models/dits/matrixgame2.py b/fastvideo/configs/models/dits/matrixgame2.py index 33764ff03..c3091d489 100644 --- a/fastvideo/configs/models/dits/matrixgame2.py +++ b/fastvideo/configs/models/dits/matrixgame2.py @@ -1,7 +1,7 @@ from dataclasses import dataclass, field import torch -from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig +from fastvideo.models.wan.config import WanVideoArchConfig, WanVideoConfig @dataclass diff --git a/fastvideo/configs/models/dits/matrixgame3.py b/fastvideo/configs/models/dits/matrixgame3.py index bfe63654f..afe34f0cb 100644 --- a/fastvideo/configs/models/dits/matrixgame3.py +++ b/fastvideo/configs/models/dits/matrixgame3.py @@ -1,7 +1,7 @@ from dataclasses import dataclass, field import torch -from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig +from fastvideo.models.wan.config import WanVideoArchConfig, WanVideoConfig def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool: diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 7897e297c..fb97bbeb1 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -14,8 +14,8 @@ from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig from fastvideo.registry import get_pipeline_config_cls_from_name -from fastvideo.configs.pipelines.wan import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig, - WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig) +from fastvideo.models.wan.pipeline_config import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig, + WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig) __all__ = [ "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig", diff --git a/fastvideo/configs/pipelines/dreamx_world.py b/fastvideo/configs/pipelines/dreamx_world.py index 6c3c1e8ef..91158255f 100644 --- a/fastvideo/configs/pipelines/dreamx_world.py +++ b/fastvideo/configs/pipelines/dreamx_world.py @@ -13,7 +13,7 @@ from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config from fastvideo.configs.models.encoders.t5 import T5ArchConfig from fastvideo.configs.models.vaes import WanVAEConfig from fastvideo.configs.pipelines.base import PipelineConfig -from fastvideo.configs.pipelines.wan import LucyEditDevConfig, t5_postprocess_text +from fastvideo.models.wan.pipeline_config import LucyEditDevConfig, t5_postprocess_text def make_dreamx_world_5b_cam_dit_config() -> DreamXWorldConfig: diff --git a/fastvideo/configs/pipelines/lingbotworld.py b/fastvideo/configs/pipelines/lingbotworld.py index ce4af1cd6..56e50e45f 100644 --- a/fastvideo/configs/pipelines/lingbotworld.py +++ b/fastvideo/configs/pipelines/lingbotworld.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field from fastvideo.configs.models import DiTConfig -from fastvideo.configs.pipelines.wan import Wan2_2_I2V_A14B_Config +from fastvideo.models.wan.pipeline_config import Wan2_2_I2V_A14B_Config from fastvideo.configs.models.dits.lingbotworld import LingBotWorldVideoConfig diff --git a/fastvideo/configs/pipelines/lingbotworld2.py b/fastvideo/configs/pipelines/lingbotworld2.py index aa9d2cbd1..9d459d6b9 100644 --- a/fastvideo/configs/pipelines/lingbotworld2.py +++ b/fastvideo/configs/pipelines/lingbotworld2.py @@ -9,7 +9,7 @@ from fastvideo.configs.models import DiTConfig from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig from fastvideo.configs.models.encoders import BaseEncoderOutput, LingBotWorld2UMT5Config from fastvideo.configs.models.vaes import WanVAEConfig -from fastvideo.configs.pipelines.wan import Wan2_2_I2V_A14B_Config +from fastvideo.models.wan.pipeline_config import Wan2_2_I2V_A14B_Config def lingbotworld2_whitespace_preprocess(prompt: str) -> str: diff --git a/fastvideo/configs/pipelines/matrixgame2.py b/fastvideo/configs/pipelines/matrixgame2.py index e8c283d3b..854c3d20c 100644 --- a/fastvideo/configs/pipelines/matrixgame2.py +++ b/fastvideo/configs/pipelines/matrixgame2.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from fastvideo.configs.models import DiTConfig, EncoderConfig from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoConfig from fastvideo.configs.models.encoders import WAN2_1ControlCLIPVisionConfig -from fastvideo.configs.pipelines.wan import WanI2V480PConfig +from fastvideo.models.wan.pipeline_config import WanI2V480PConfig @dataclass diff --git a/fastvideo/configs/pipelines/matrixgame3.py b/fastvideo/configs/pipelines/matrixgame3.py index e226710b6..f083601b0 100644 --- a/fastvideo/configs/pipelines/matrixgame3.py +++ b/fastvideo/configs/pipelines/matrixgame3.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from fastvideo.configs.models import DiTConfig from fastvideo.configs.models.dits.matrixgame3 import MatrixGame3WanVideoConfig -from fastvideo.configs.pipelines.wan import WanT2V480PConfig +from fastvideo.models.wan.pipeline_config import WanT2V480PConfig @dataclass diff --git a/fastvideo/configs/pipelines/turbodiffusion.py b/fastvideo/configs/pipelines/turbodiffusion.py index 5a8f2541b..ae073f41b 100644 --- a/fastvideo/configs/pipelines/turbodiffusion.py +++ b/fastvideo/configs/pipelines/turbodiffusion.py @@ -12,7 +12,7 @@ from fastvideo.configs.models.dits import WanVideoConfig from fastvideo.configs.models.encoders import CLIPVisionConfig from fastvideo.configs.models.vaes import WanVAEConfig from fastvideo.configs.pipelines.base import PipelineConfig -from fastvideo.configs.pipelines.wan import t5_postprocess_text, T5Config, BaseEncoderOutput +from fastvideo.models.wan.pipeline_config import t5_postprocess_text, T5Config, BaseEncoderOutput import torch from collections.abc import Callable diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 17df6c231..6cd6c1d6b 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -1,324 +1,60 @@ # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Callable -from dataclasses import dataclass, field +"""Compatibility exports for the family-local Wan pipeline configuration.""" -import torch +from fastvideo.models.wan.pipeline_config import ( + BaseEncoderOutput, + CLIPVisionConfig, + DiTConfig, + EncoderConfig, + FastWan2_1_T2V_480P_Config, + FastWan2_2_TI2V_5B_Config, + LucyEditDevConfig, + PipelineConfig, + SelfForcingWan2_2_T2V480PConfig, + SelfForcingWanT2V480PConfig, + T5Config, + VAEConfig, + WAN2_1ControlCLIPVisionConfig, + WANV2VConfig, + Wan2_2_I2V_A14B_Config, + Wan2_2_T2V_A14B_Config, + Wan2_2_TI2V_5B_Config, + WanI2V480PConfig, + WanI2V720PConfig, + WanT2V480PConfig, + WanT2V720PConfig, + WanVAEArchConfig, + WanVAEConfig, + WanVideoArchConfig, + WanVideoConfig, + t5_postprocess_text, +) -from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig -from fastvideo.configs.models.dits import WanVideoConfig -from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig -from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config, - WAN2_1ControlCLIPVisionConfig) -from fastvideo.configs.pipelines.base import PipelineConfig -from fastvideo.models.wan.vae_config import WanVAEArchConfig, WanVAEConfig - - -def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: - mask: torch.Tensor = outputs.attention_mask - hidden_state: torch.Tensor = outputs.last_hidden_state - seq_lens = mask.gt(0).sum(dim=1).long() - assert torch.isnan(hidden_state).sum() == 0 - prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] - prompt_embeds_tensor: torch.Tensor = torch.stack( - [torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0) - return prompt_embeds_tensor - - -@dataclass -class WanT2V480PConfig(PipelineConfig): - """Base configuration for Wan T2V 1.3B pipeline architecture.""" - - # WanConfig-specific parameters with defaults - # DiT - dit_config: DiTConfig = field(default_factory=WanVideoConfig) - # VAE - vae_config: VAEConfig = field(default_factory=WanVAEConfig) - vae_tiling: bool = False - vae_sp: bool = False - - # Denoising stage - flow_shift: float | None = 3.0 - - # Text encoding stage - text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), )) - postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], - ...] = field(default_factory=lambda: (t5_postprocess_text, )) - - # Precision for each component - precision: str = "bf16" - # Keep the VAE *encode* in fp32: for I2V/causal Wan variants the image/video - # VAE encode seeds the denoising trajectory, and bf16 there shifts the output - # (e.g. Matrix-Game-2.0 SSIM 0.98 -> 0.76). Decode is output-only, so we lower - # just the decode below. - vae_precision: str = "fp32" - # bf16 VAE *decode* is effectively lossless (MS-SSIM 0.9999 vs fp32 on an - # identical latent) and faster; the same AutoencoderKLWan already runs bf16 - # in Cosmos-Predict2.5 and fp16 in Cosmos. Inherited by all Wan variants - # (incl. I2V/causal) — decode-only, so encode trajectories are untouched. - vae_decode_precision: str = "bf16" - text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", )) - - # self-forcing params - warp_denoising_step: bool = True - - # WanConfig-specific added parameters - - def __post_init__(self): - self.vae_config.load_encoder = False - self.vae_config.load_decoder = True - - -@dataclass -class WanT2V720PConfig(WanT2V480PConfig): - """Base configuration for Wan T2V 14B 720P pipeline architecture.""" - - # WanConfig-specific parameters with defaults - - # Denoising stage - flow_shift: float | None = 5.0 - - -@dataclass -class WanI2V480PConfig(WanT2V480PConfig): - """Base configuration for Wan I2V 14B 480P pipeline architecture.""" - - # WanConfig-specific parameters with defaults - - # Precision for each component - image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig) - image_encoder_precision: str = "fp32" - - def __post_init__(self) -> None: - self.vae_config.load_encoder = True - self.vae_config.load_decoder = True - - -@dataclass -class WanI2V720PConfig(WanI2V480PConfig): - """Base configuration for Wan I2V 14B 720P pipeline architecture.""" - - # WanConfig-specific parameters with defaults - - # Denoising stage - flow_shift: float | None = 5.0 - - -@dataclass -class WANV2VConfig(WanI2V480PConfig): - """Configuration for WAN2.1 1.3B Control pipeline.""" - - image_encoder_config: EncoderConfig = field(default_factory=WAN2_1ControlCLIPVisionConfig) - # CLIP encoder precision - image_encoder_precision: str = 'bf16' - - -@dataclass -class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): - """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD""" - - # WanConfig-specific parameters with defaults - - # Denoising stage - flow_shift: float | None = 8.0 - dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) - - -@dataclass -class Wan2_2_TI2V_5B_Config(WanT2V480PConfig): - flow_shift: float | None = 5.0 - ti2v_task: bool = True - expand_timesteps: bool = True - - def __post_init__(self) -> None: - assert not (self.ti2v_task and self.lucy_edit_task) - self.vae_config.load_encoder = True - self.vae_config.load_decoder = True - self.dit_config.expand_timesteps = self.expand_timesteps - - -@dataclass -class LucyEditDevConfig(Wan2_2_TI2V_5B_Config): - """Configuration for Decart Lucy Edit Dev video editing.""" - - dit_config: DiTConfig = field(default_factory=lambda: WanVideoConfig(arch_config=WanVideoArchConfig( - num_attention_heads=24, - in_channels=96, - out_channels=48, - ffn_dim=14336, - num_layers=30, - ))) - vae_config: VAEConfig = field(default_factory=lambda: WanVAEConfig(arch_config=WanVAEArchConfig( - base_dim=160, - decoder_base_dim=256, - z_dim=48, - in_channels=12, - out_channels=12, - scale_factor_spatial=16, - patch_size=2, - is_residual=True, - clip_output=False, - latents_mean=( - -0.2289, - -0.0052, - -0.1323, - -0.2339, - -0.2799, - 0.0174, - 0.1838, - 0.1557, - -0.1382, - 0.0542, - 0.2813, - 0.0891, - 0.1570, - -0.0098, - 0.0375, - -0.1825, - -0.2246, - -0.1207, - -0.0698, - 0.5109, - 0.2665, - -0.2108, - -0.2158, - 0.2502, - -0.2055, - -0.0322, - 0.1109, - 0.1567, - -0.0729, - 0.0899, - -0.2799, - -0.1230, - -0.0313, - -0.1649, - 0.0117, - 0.0723, - -0.2839, - -0.2083, - -0.0520, - 0.3748, - 0.0152, - 0.1957, - 0.1433, - -0.2944, - 0.3573, - -0.0548, - -0.1681, - -0.0667, - ), - latents_std=( - 0.4765, - 1.0364, - 0.4514, - 1.1677, - 0.5313, - 0.4990, - 0.4818, - 0.5013, - 0.8158, - 1.0344, - 0.5894, - 1.0901, - 0.6885, - 0.6165, - 0.8454, - 0.4978, - 0.5759, - 0.3523, - 0.7135, - 0.6804, - 0.5833, - 1.4146, - 0.8986, - 0.5659, - 0.7069, - 0.5338, - 0.4889, - 0.4917, - 0.4069, - 0.4999, - 0.6866, - 0.4093, - 0.5709, - 0.6065, - 0.6415, - 0.4944, - 0.5726, - 1.2042, - 0.5458, - 1.6887, - 0.3971, - 1.0600, - 0.3943, - 0.5537, - 0.5444, - 0.4089, - 0.7468, - 0.7744, - ), - ))) - ti2v_task: bool = False - lucy_edit_task: bool = True - - def __post_init__(self) -> None: - assert not (self.ti2v_task and self.lucy_edit_task) - # Lucy uses Wan2.2's enhanced 48-channel VAE latents. Denoising - # concatenates noise + video latents, matching the 96-channel - # transformer input declared above. - self.vae_config.load_encoder = True - self.vae_config.load_decoder = True - self.dit_config.expand_timesteps = self.expand_timesteps - - -@dataclass -class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): - flow_shift: float | None = 5.0 - dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) - - -@dataclass -class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): - flow_shift: float | None = 12.0 - boundary_ratio: float | None = 0.875 - - # self-forcing params - dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250]) - warp_denoising_step: bool = True - - def __post_init__(self) -> None: - self.dit_config.boundary_ratio = self.boundary_ratio - - -@dataclass -class Wan2_2_I2V_A14B_Config(WanI2V480PConfig): - flow_shift: float | None = 5.0 - boundary_ratio: float | None = 0.900 - - def __post_init__(self) -> None: - super().__post_init__() - self.dit_config.boundary_ratio = self.boundary_ratio - - -# ============================================= -# ============= Causal Self-Forcing ============= -# ============================================= -@dataclass -class SelfForcingWanT2V480PConfig(WanT2V480PConfig): - is_causal: bool = True - flow_shift: float | None = 5.0 - dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250]) - warp_denoising_step: bool = True - - -@dataclass -class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config): - is_causal: bool = True - flow_shift: float | None = 12.0 - boundary_ratio: float | None = 0.875 - dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 850, 700, 550, 350, 275, 200, 125]) - warp_denoising_step: bool = True - - def __post_init__(self) -> None: - self.vae_config.load_encoder = True - self.vae_config.load_decoder = True +__all__ = [ + "BaseEncoderOutput", + "CLIPVisionConfig", + "DiTConfig", + "EncoderConfig", + "FastWan2_1_T2V_480P_Config", + "FastWan2_2_TI2V_5B_Config", + "LucyEditDevConfig", + "PipelineConfig", + "SelfForcingWan2_2_T2V480PConfig", + "SelfForcingWanT2V480PConfig", + "T5Config", + "VAEConfig", + "WAN2_1ControlCLIPVisionConfig", + "WANV2VConfig", + "Wan2_2_I2V_A14B_Config", + "Wan2_2_T2V_A14B_Config", + "Wan2_2_TI2V_5B_Config", + "WanI2V480PConfig", + "WanI2V720PConfig", + "WanT2V480PConfig", + "WanT2V720PConfig", + "WanVAEArchConfig", + "WanVAEConfig", + "WanVideoArchConfig", + "WanVideoConfig", + "t5_postprocess_text", +] diff --git a/fastvideo/models/AGENTS.md b/fastvideo/models/AGENTS.md index e195cb25b..9a5aa009c 100644 --- a/fastvideo/models/AGENTS.md +++ b/fastvideo/models/AGENTS.md @@ -2,13 +2,13 @@ **Generated:** 2026-05-02 -DiT / VAE / encoder / scheduler / upsampler / audio model classes. **Pre-commit excludes this directory**, except `wan/config.py`, `wan/vae_config.py`, and `wan/__init__.py`. Match neighboring file style manually for excluded files. +DiT / VAE / encoder / scheduler / upsampler / audio model classes. **Pre-commit excludes this directory**, except Wan's configs, definition, and `__init__.py`. Match neighboring file style manually for excluded files. ## Layout ``` models/ -├── wan/ # Family-local dense Wan transformer, VAE, and configs +├── wan/ # Family-local dense/causal Wan transformers, VAE, and configs ├── dits/ │ ├── .py # Single-file DiT (ltx2, hunyuanvideo, cosmos, ...); wanvideo is a shim │ ├── hyworld/ # Multi-file DiT family @@ -28,10 +28,14 @@ to instantiate model components from a HF directory. New components plug in through `register_*` calls or by extending the `ComponentLoader` mappings. Wan is the first family-local package: edit `wan/transformer.py` with -`wan/config.py`, or `wan/vae.py` with `wan/vae_config.py`. The old -`dits/wanvideo.py`, `vaes/wanvae.py`, and matching `configs/models/` paths remain +`wan/config.py`, `wan/causal_transformer.py`, or `wan/vae.py` with +`wan/vae_config.py`. The old `dits/wanvideo.py`, `dits/causal_wanvideo.py`, +`vaes/wanvae.py`, and matching `configs/models/` paths remain compatibility re-exports. Shared encoders and VAE utilities stay in their current -locations. See `wan/AGENTS.md`; do not migrate other components as part of a Wan edit. +locations. Wan pipeline defaults and variant metadata live in +`wan/pipeline_config.py` and `wan/definition.py`; old pipeline config imports +remain aliases. See `wan/AGENTS.md`; do not migrate other components as part of +a Wan edit. ## Adding a Model Component (DiT / VAE / Encoder) diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 2875edec1..c981fe69e 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -1,817 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 - -import math -from typing import Any - -import numpy as np -import torch -import torch.nn as nn - -from torch.nn.attention.flex_attention import create_block_mask, flex_attention -from torch.nn.attention.flex_attention import BlockMask -# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention -# see https://github.com/pytorch/pytorch/issues/133254 -# change to default for other models -flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") -import torch.distributed as dist - -import fastvideo.envs as envs -from fastvideo.attention import (DistributedAttention, LocalAttention) -from fastvideo.configs.models.dits import WanVideoConfig -from fastvideo.distributed.parallel_state import get_sp_world_size -from fastvideo.forward_context import get_forward_context -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, - ScaleResidualLayerNormScaleShift) -from fastvideo.layers.linear import ReplicatedLinear -from fastvideo.layers.mlp import MLP -from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, get_rotary_pos_embed) -from fastvideo.layers.visual_embedding import (PatchEmbed) -from fastvideo.logger import init_logger -from fastvideo.models.dits._relative_rope import relativistic_window_offsets -from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding -from fastvideo.platforms import AttentionBackendEnum, current_platform - -logger = init_logger(__name__) - -GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES = 21 - - -class CausalWanSelfAttention(nn.Module): - - def __init__(self, - dim: int, - num_heads: int, - local_attn_size: int = -1, - sink_size: int = 0, - qk_norm=True, - eps=1e-6, - parallel_attention=False, - rope_cache_policy: str = "absolute") -> None: - assert dim % num_heads == 0 - super().__init__() - self.dim = dim - self.num_heads = num_heads - self.head_dim = dim // num_heads - self.local_attn_size = local_attn_size - self.sink_size = sink_size - self.qk_norm = qk_norm - self.eps = eps - self.parallel_attention = parallel_attention - self.rope_cache_policy = rope_cache_policy - - # Scaled dot product attention - self.attn = LocalAttention(num_heads=num_heads, - head_size=self.head_dim, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA)) - - def forward(self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], - block_mask: BlockMask, - kv_cache: dict | None = None, - current_start: int = 0, - cache_start: int | None = None, - frame_seqlen: int = 1560): - r""" - Args: - x(Tensor): Shape [B, L, num_heads, C / num_heads] - seq_lens(Tensor): Shape [B] - grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) - freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] - frame_seqlen (int): Number of tokens per latent frame, - e.g. 1560 for 480x832 resolution. - """ - if cache_start is None: - cache_start = current_start - - cos, sin = freqs_cis - # relativistic defers roping until the cache window is known (and caches raw k) - relativistic = self.rope_cache_policy == "relativistic" and kv_cache is not None - if not relativistic: - roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) - roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) - - if kv_cache is None: - # Padding for flex attention - padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] - padded_roped_query = torch.cat([ - roped_query, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], device=q.device, dtype=v.dtype) - ], - dim=1) - - padded_roped_key = torch.cat([ - roped_key, - torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], device=k.device, dtype=v.dtype) - ], - dim=1) - - padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype)], - dim=1) - - x = flex_attention(query=padded_roped_query.transpose(2, 1), - key=padded_roped_key.transpose(2, 1), - value=padded_v.transpose(2, 1), - block_mask=block_mask)[:, :, :-padded_length].transpose(2, 1) - else: - current_end = current_start + q.shape[1] - sink_tokens = self.sink_size * frame_seqlen - if self.local_attn_size == -1: - max_attention_size = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES * frame_seqlen) - else: - max_attention_size = self.local_attn_size * frame_seqlen - if self.local_attn_size == -1 and current_end > max_attention_size: - raise ValueError("Causal Wan local_attn_size=-1 keeps the previous " - f"{GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES}-latent-frame KV " - "window for compatibility. Set local_attn_size for " - f"longer rollouts; got current_end={current_end} tokens " - f"with frame_seqlen={frame_seqlen}.") - # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache - kv_cache_size = kv_cache["k"].shape[1] - num_new_tokens = q.shape[1] - stored_key = k if relativistic else roped_key # raw vs roped in cache - global_end_index = (int(kv_cache["global_end_index"].item()) if isinstance( - kv_cache["global_end_index"], torch.Tensor) else int(kv_cache["global_end_index"])) - local_end_index_prev = (int(kv_cache["local_end_index"].item()) if isinstance( - kv_cache["local_end_index"], torch.Tensor) else int(kv_cache["local_end_index"])) - if self.local_attn_size != -1 and (current_end - > global_end_index) and (num_new_tokens + local_end_index_prev - > kv_cache_size): - # Calculate the number of new tokens added in this step - # Shift existing cache content left to discard oldest tokens - # Clone the source slice to avoid overlapping memory error - num_evicted_tokens = num_new_tokens + local_end_index_prev - kv_cache_size - num_rolled_tokens = local_end_index_prev - num_evicted_tokens - sink_tokens - kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - # Insert the new keys/values at the end - local_end_index = local_end_index_prev + current_end - \ - global_end_index - num_evicted_tokens - local_start_index = local_end_index - num_new_tokens - kv_cache["k"][:, local_start_index:local_end_index] = stored_key - kv_cache["v"][:, local_start_index:local_end_index] = v - else: - # Assign new keys/values directly up to current_end - local_end_index = local_end_index_prev + current_end - global_end_index - local_start_index = local_end_index - num_new_tokens - kv_cache["k"] = kv_cache["k"].detach() - kv_cache["v"] = kv_cache["v"].detach() - # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) - kv_cache["k"][:, local_start_index:local_end_index] = stored_key - kv_cache["v"][:, local_start_index:local_end_index] = v - key_window = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index] - value_window = kv_cache["v"][:, max(0, local_end_index - max_attention_size):local_end_index] - if relativistic: - window_len, query_lo, query_hi = relativistic_window_offsets(local_end_index, num_new_tokens, - max_attention_size) - roped_query = _apply_rotary_emb(q, cos[query_lo:query_hi], sin[query_lo:query_hi], - is_neox_style=False).type_as(v) - key_window = _apply_rotary_emb(key_window, cos[:window_len], sin[:window_len], - is_neox_style=False).type_as(v) - x = self.attn(roped_query, key_window, value_window) - if isinstance(kv_cache["global_end_index"], torch.Tensor): - kv_cache["global_end_index"].fill_(current_end) - else: - kv_cache["global_end_index"] = current_end - if isinstance(kv_cache["local_end_index"], torch.Tensor): - kv_cache["local_end_index"].fill_(local_end_index) - else: - kv_cache["local_end_index"] = local_end_index - - return x - - -class CausalWanTransformerBlock(nn.Module): - - def __init__(self, - dim: int, - ffn_dim: int, - num_heads: int, - local_attn_size: int = -1, - sink_size: int = 0, - qk_norm: str = "rms_norm_across_heads", - cross_attn_norm: bool = False, - eps: float = 1e-6, - added_kv_proj_dim: int | None = None, - supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, - prefix: str = "", - rope_cache_policy: str = "absolute"): - super().__init__() - - # 1. Self-attention - self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False) - self.to_q = ReplicatedLinear(dim, dim, bias=True) - self.to_k = ReplicatedLinear(dim, dim, bias=True) - self.to_v = ReplicatedLinear(dim, dim, bias=True) - - self.to_out = ReplicatedLinear(dim, dim, bias=True) - self.attn1 = CausalWanSelfAttention(dim, - num_heads, - local_attn_size=local_attn_size, - sink_size=sink_size, - qk_norm=qk_norm, - eps=eps, - rope_cache_policy=rope_cache_policy) - self.hidden_dim = dim - self.num_attention_heads = num_heads - self.local_attn_size = local_attn_size - dim_head = dim // num_heads - if qk_norm == "rms_norm": - self.norm_q = RMSNorm(dim_head, eps=eps) - self.norm_k = RMSNorm(dim_head, eps=eps) - elif qk_norm == "rms_norm_across_heads": - # LTX applies qk norm across all heads - self.norm_q = RMSNorm(dim, eps=eps) - self.norm_k = RMSNorm(dim, eps=eps) - else: - print("QK Norm type not supported") - raise Exception - assert cross_attn_norm is True - self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32) - - # 2. Cross-attention - # Only T2V for now - self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) - self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32) - - # 3. Feed-forward - self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") - self.mlp_residual = ScaleResidual() - - self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - temb: torch.Tensor, - freqs_cis: tuple[torch.Tensor, torch.Tensor], - block_mask: BlockMask, - kv_cache: dict | None = None, - crossattn_cache: dict | None = None, - current_start: int = 0, - cache_start: int | None = None, - frame_seqlen: int | None = None, - ) -> torch.Tensor: - # hidden_states.shape: [batch_size, seq_length, inner_dim] - # temb.shape: [batch_size, temb_seq_len, 6, inner_dim] - if hidden_states.dim() == 4: - hidden_states = hidden_states.squeeze(1) - temb_seq_len = temb.shape[1] - tokens_per_temb = hidden_states.shape[1] // temb_seq_len - if frame_seqlen is None: - frame_seqlen = tokens_per_temb - else: - frame_seqlen = int(frame_seqlen) - bs, seq_length, _ = hidden_states.shape - orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb - # e.shape: [batch_size, temb_seq_len, 6, inner_dim] - assert e.shape == (bs, temb_seq_len, 6, self.hidden_dim) - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) - # *_msa.shape: [batch_size, temb_seq_len, 1, inner_dim] - # assert shift_msa.dtype == torch.float32 - - # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(temb_seq_len, tokens_per_temb)) * - (1 + scale_msa) + shift_msa).flatten(1, 2) - query, _ = self.to_q(norm_hidden_states) - key, _ = self.to_k(norm_hidden_states) - value, _ = self.to_v(norm_hidden_states) - - if self.norm_q is not None: - query = self.norm_q(query) - if self.norm_k is not None: - key = self.norm_k(key) - - query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - - attn_output = self.attn1( - query, - key, - value, - freqs_cis, - block_mask, - kv_cache, - current_start, - cache_start, - frame_seqlen=frame_seqlen, - ) - attn_output = attn_output.flatten(2) - attn_output, _ = self.to_out(attn_output) - attn_output = attn_output.squeeze(1) - - null_shift = null_scale = torch.tensor([0], device=hidden_states.device) - norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, - null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) - - # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None, - crossattn_cache=crossattn_cache) - norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, - c_scale_msa) - - # 3. Feed-forward - ff_output = self.ffn(norm_hidden_states) - hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) - - return hidden_states - - -class CausalWanTransformer3DModel(BaseDiT): - _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions - _compile_conditions = WanVideoConfig()._compile_conditions - _supported_attention_backends = WanVideoConfig()._supported_attention_backends - param_names_mapping = WanVideoConfig().param_names_mapping - reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping - lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping - - def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: - super().__init__(config=config, hf_config=hf_config) - - inner_dim = config.num_attention_heads * config.attention_head_dim - self.hidden_size = config.hidden_size - self.num_attention_heads = config.num_attention_heads - self.attention_head_dim = config.attention_head_dim - self.in_channels = config.in_channels - self.out_channels = config.out_channels - self.num_channels_latents = config.num_channels_latents - self.patch_size = config.patch_size - self.text_len = config.text_len - self.local_attn_size = config.local_attn_size - self.rope_cache_policy = config.arch_config.rope_cache_policy - - # 1. Patch & position embedding - self.patch_embedding = PatchEmbed(in_chans=config.in_channels, - embed_dim=inner_dim, - patch_size=config.patch_size, - flatten=False) - - # 2. Condition embeddings - self.condition_embedder = WanTimeTextImageEmbedding( - dim=inner_dim, - time_freq_dim=config.freq_dim, - text_embed_dim=config.text_dim, - image_embed_dim=config.image_dim, - ) - - # 3. Transformer blocks - self.blocks = nn.ModuleList([ - CausalWanTransformerBlock(inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.local_attn_size, - config.sink_size, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - prefix=f"{config.prefix}.blocks.{i}", - rope_cache_policy=config.arch_config.rope_cache_policy) - for i in range(config.num_layers) - ]) - - # 4. Output norm & projection - self.norm_out = LayerNormScaleShift(inner_dim, - norm_type="layer", - eps=config.eps, - elementwise_affine=False, - dtype=torch.float32) - self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) - self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) - - self.gradient_checkpointing = False - - # Causal-specific - self.block_mask = None - self.teacher_forcing_block_mask = None - self.num_frame_per_block = config.arch_config.num_frames_per_block - assert self.num_frame_per_block <= 3 - self.independent_first_frame = False - - self.__post_init__() - - @staticmethod - def _prepare_blockwise_causal_attn_mask(device: torch.device | str, - num_frames: int = 21, - frame_seqlen: int = 1560, - num_frame_per_block=1, - local_attn_size=-1) -> BlockMask: - """ - we will divide the token sequence into the following format - [1 latent frame] [1 latent frame] ... [1 latent frame] - We use flexattention to construct the attention mask - """ - total_length = num_frames * frame_seqlen - - # we do right padding to get to a multiple of 128 - padded_length = math.ceil(total_length / 128) * 128 - total_length - - ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) - - # Block-wise causal mask will attend to all elements that are before the end of the current chunk - frame_indices = torch.arange(start=0, end=total_length, step=frame_seqlen * num_frame_per_block, device=device) - - for tmp in frame_indices: - ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ - frame_seqlen * num_frame_per_block - - def attention_mask(b, h, q_idx, kv_idx): - if local_attn_size == -1: - return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) - else: - return ((kv_idx < ends[q_idx]) & - (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) - # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask - - block_mask = create_block_mask(attention_mask, - B=None, - H=None, - Q_LEN=total_length + padded_length, - KV_LEN=total_length + padded_length, - _compile=False, - device=device) - - if not dist.is_initialized() or dist.get_rank() == 0: - print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") - print(block_mask) - - # import imageio - # import numpy as np - # from torch.nn.attention.flex_attention import create_mask - - # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + - # padded_length, KV_LEN=total_length + padded_length, device=device) - # import cv2 - # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) - # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) - - return block_mask - - @staticmethod - def _prepare_teacher_forcing_mask(device: torch.device | str, - num_frames: int = 21, - frame_seqlen: int = 1560, - num_frame_per_block=1, - local_attn_size=-1) -> BlockMask: - """Attention mask for the teacher-forcing ``[clean | noisy]`` sequence. - - A noisy token attends to its own block plus the clean context of all - strictly previous blocks; clean tokens are block-wise causal. - """ - if local_attn_size != -1: - raise NotImplementedError(f"Teacher forcing ignores local_attn_size={local_attn_size}: " - "unlike the block-wise causal mask, this mask always attends " - "to the full clean context. Windowed teacher forcing is not " - "implemented; use local_attn_size=-1 for teacher-forcing " - "training.") - total_length = num_frames * frame_seqlen * 2 - padded_length = math.ceil(total_length / 128) * 128 - total_length - - clean_ends = num_frames * frame_seqlen - context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) - noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) - noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) - noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) - noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) - - attention_block_size = frame_seqlen * num_frame_per_block - frame_indices = torch.arange(start=0, - end=num_frames * frame_seqlen, - step=attention_block_size, - device=device, - dtype=torch.long) - for start in frame_indices: - context_ends[start:start + attention_block_size] = start + attention_block_size - - noisy_image_start_list = torch.arange(num_frames * frame_seqlen, - total_length, - step=attention_block_size, - device=device, - dtype=torch.long) - noisy_image_end_list = noisy_image_start_list + attention_block_size - for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): - noise_noise_starts[start:end] = start - noise_noise_ends[start:end] = end - noise_context_ends[start:end] = block_index * attention_block_size - - def attention_mask(b, h, q_idx, kv_idx): - clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) - c1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) - c2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) - noise_mask = (q_idx >= clean_ends) & (c1 | c2) - eye_mask = q_idx == kv_idx - return eye_mask | clean_mask | noise_mask - - block_mask = create_block_mask(attention_mask, - B=None, - H=None, - Q_LEN=total_length + padded_length, - KV_LEN=total_length + padded_length, - _compile=False, - device=device) - - if not dist.is_initialized() or dist.get_rank() == 0: - print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames") - print(block_mask) - - return block_mask - - def _forward_inference(self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] - | None = None, - kv_cache: dict = None, - crossattn_cache: dict = None, - current_start: int = 0, - cache_start: int = 0, - start_frame: int = 0, - **kwargs) -> torch.Tensor: - r""" - Run the diffusion model with kv caching. - See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. - This function will be run for num_frame times. - Process the latent frames one by one (1560 tokens each) - """ - - orig_dtype = hidden_states.dtype - if not isinstance(encoder_hidden_states, torch.Tensor): - encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, list): - encoder_hidden_states_image = (encoder_hidden_states_image[0] - if len(encoder_hidden_states_image) > 0 else None) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w - - # Get rotary embeddings - d = self.hidden_size // self.num_attention_heads - rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - if self.rope_cache_policy == "relativistic": - # fixed table over [0, max_attention_frames); attention slices it per step - max_attention_frames = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES - if self.local_attn_size == -1 else self.local_attn_size) - rope_num_frames = max_attention_frames * get_sp_world_size() - rope_start_frame = 0 - else: - rope_num_frames = post_patch_num_frames * get_sp_world_size() - rope_start_frame = start_frame # 0 when kv_cache is None - freqs_cos, freqs_sin = get_rotary_pos_embed((rope_num_frames, post_patch_height, post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000, - start_frame=rope_start_frame) - freqs_cos = freqs_cos.to(hidden_states.device) - freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None - - hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.stack([torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - - encoder_hidden_states = torch.cat([ - encoder_hidden_states, - encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), - encoder_hidden_states.size(2)) - ], - dim=1) - - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) - - if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - - encoder_hidden_states = encoder_hidden_states.to( - orig_dtype) if current_platform.is_mps() else encoder_hidden_states # cast to orig_dtype for MPS - - assert encoder_hidden_states.dtype == orig_dtype - - # 4. Transformer blocks - for block_index, block in enumerate(self.blocks): - if torch.is_grad_enabled() and self.gradient_checkpointing: - causal_kwargs = { - "kv_cache": kv_cache[block_index], - "current_start": current_start, - "cache_start": cache_start, - "block_mask": self.block_mask, - "frame_seqlen": post_patch_height * post_patch_width, - } - hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, - timestep_proj, freqs_cis, **causal_kwargs) - else: - causal_kwargs = { - "kv_cache": kv_cache[block_index], - "crossattn_cache": (crossattn_cache[block_index] if crossattn_cache is not None else None), - "current_start": current_start, - "cache_start": cache_start, - "block_mask": self.block_mask, - "frame_seqlen": post_patch_height * post_patch_width, - } - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, **causal_kwargs) - - # 5. Output norm, projection & unpatchify - temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) - hidden_states = self.norm_out(hidden_states, shift, scale) - hidden_states = self.proj_out(hidden_states) - - output = self.unpatchify(hidden_states, grid_sizes) - - return torch.stack(output) - - def _forward_train(self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] - | None = None, - start_frame: int = 0, - clean_x: torch.Tensor | None = None, - aug_t: torch.Tensor | None = None, - **kwargs) -> torch.Tensor: - - orig_dtype = hidden_states.dtype - teacher_forcing = clean_x is not None - if not isinstance(encoder_hidden_states, torch.Tensor): - encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, list): - encoder_hidden_states_image = (encoder_hidden_states_image[0] - if len(encoder_hidden_states_image) > 0 else None) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape - p_t, p_h, p_w = self.patch_size - post_patch_num_frames = num_frames // p_t - post_patch_height = height // p_h - post_patch_width = width // p_w - - # Get rotary embeddings - d = self.hidden_size // self.num_attention_heads - rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] - freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), - self.hidden_size, - self.num_attention_heads, - rope_dim_list, - dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000, - start_frame=start_frame) - freqs_cos = freqs_cos.to(hidden_states.device) - freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None - - if teacher_forcing: - if self.teacher_forcing_block_mask is None: - self.teacher_forcing_block_mask = self._prepare_teacher_forcing_mask( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size, - ) - block_mask = self.teacher_forcing_block_mask - else: - if self.block_mask is None: - self.block_mask = self._prepare_blockwise_causal_attn_mask(device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * - post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size) - block_mask = self.block_mask - - hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.stack([torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) - hidden_states = hidden_states.flatten(2).transpose(1, 2) - - encoder_hidden_states = torch.cat([ - encoder_hidden_states, - encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), - encoder_hidden_states.size(2)) - ], - dim=1) - encoder_hidden_states_text = encoder_hidden_states - - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) - - if encoder_hidden_states_image is not None: - encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) - - encoder_hidden_states = encoder_hidden_states.to( - orig_dtype) if current_platform.is_mps() else encoder_hidden_states # cast to orig_dtype for MPS - - assert encoder_hidden_states.dtype == orig_dtype - - if teacher_forcing: - # Tile RoPE/modulation so clean frame i and noisy frame i share a position. - clean_tokens = self.patch_embedding(clean_x).flatten(2).transpose(1, 2) - hidden_states = torch.cat([clean_tokens, hidden_states], dim=1) - if aug_t is None: - aug_t = torch.zeros_like(timestep) - _, timestep_proj_clean, _, _ = self.condition_embedder(aug_t.flatten(), encoder_hidden_states_text, None) - timestep_proj_clean = timestep_proj_clean.unflatten(1, - (6, self.hidden_size)).unflatten(dim=0, - sizes=timestep.shape) - timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1) - freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0), torch.cat([freqs_sin, freqs_sin], dim=0)) - - # 4. Transformer blocks - if torch.is_grad_enabled() and self.gradient_checkpointing: - for block in self.blocks: - hidden_states = self._gradient_checkpointing_func(block, - hidden_states, - encoder_hidden_states, - timestep_proj, - freqs_cis, - block_mask=block_mask) - else: - for block in self.blocks: - hidden_states = block(hidden_states, - encoder_hidden_states, - timestep_proj, - freqs_cis, - block_mask=block_mask) - - if teacher_forcing: - hidden_states = hidden_states[:, hidden_states.shape[1] // 2:] - - # 5. Output norm, projection & unpatchify - temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) - hidden_states = self.norm_out(hidden_states, shift, scale) - hidden_states = self.proj_out(hidden_states) - - output = self.unpatchify(hidden_states, grid_sizes) - - return torch.stack(output) - - def forward(self, *args, **kwargs): - if kwargs.get('kv_cache', None) is not None: - return self._forward_inference(*args, **kwargs) - else: - return self._forward_train(*args, **kwargs) - - def unpatchify(self, x, grid_sizes): - r""" - - - Args: - x (List[Tensor]): - List of patchified features, each with shape [L, C_out * prod(patch_size)] - grid_sizes (Tensor): - Original spatial-temporal grid dimensions before patching, - - - Returns: - Tensor: - Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] - """ - - c = self.out_channels - out = [] - for u, v in zip(x, grid_sizes.tolist()): - u = u[:math.prod(v)].view(*v, *self.patch_size, c) - u = u.permute(6, 0, 3, 1, 4, 2, 5) - # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) - u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) - out.append(u) - return out - - -# Entry point for model registry -EntryClass = CausalWanTransformer3DModel +"""Compatibility imports for the family-local causal Wan transformer.""" + +from fastvideo.models.wan.causal_transformer import ( + GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES, + CausalWanSelfAttention, + CausalWanTransformer3DModel, + CausalWanTransformerBlock, + EntryClass, +) + +__all__ = [ + "GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES", + "CausalWanSelfAttention", + "CausalWanTransformer3DModel", + "CausalWanTransformerBlock", + "EntryClass", +] diff --git a/fastvideo/models/dits/dreamx_world.py b/fastvideo/models/dits/dreamx_world.py index a53f50ec3..ba1d3ca8d 100644 --- a/fastvideo/models/dits/dreamx_world.py +++ b/fastvideo/models/dits/dreamx_world.py @@ -13,7 +13,7 @@ from fastvideo.distributed.parallel_state import get_sp_world_size from fastvideo.layers.layernorm import RMSNorm from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.rotary_embedding import get_rotary_pos_embed -from fastvideo.models.dits.wanvideo import (LayerNormScaleShift, PatchEmbed, WanTimeTextImageEmbedding, +from fastvideo.models.wan.transformer import (LayerNormScaleShift, PatchEmbed, WanTimeTextImageEmbedding, WanTransformer3DModel, WanTransformerBlock) from fastvideo.platforms import AttentionBackendEnum, current_platform from fastvideo.attention import LocalAttention diff --git a/fastvideo/models/dits/lingbotworld/model.py b/fastvideo/models/dits/lingbotworld/model.py index 26009396a..e88f7afc7 100644 --- a/fastvideo/models/dits/lingbotworld/model.py +++ b/fastvideo/models/dits/lingbotworld/model.py @@ -20,7 +20,7 @@ from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (PatchEmbed, WanCamControlPatchEmbedding) from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import ( +from fastvideo.models.wan.transformer import ( WanI2VCrossAttention, WanT2VCrossAttention, WanTimeTextImageEmbedding, diff --git a/fastvideo/models/dits/matrixgame2/causal_model.py b/fastvideo/models/dits/matrixgame2/causal_model.py index ba193dc18..54d6b0883 100644 --- a/fastvideo/models/dits/matrixgame2/causal_model.py +++ b/fastvideo/models/dits/matrixgame2/causal_model.py @@ -36,7 +36,7 @@ from fastvideo.layers.visual_embedding import ( from fastvideo.logger import init_logger from fastvideo.models.dits._relative_rope import relativistic_window_offsets from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import ( +from fastvideo.models.wan.transformer import ( WanT2VCrossAttention, WanImageEmbedding, ) diff --git a/fastvideo/models/dits/matrixgame2/model.py b/fastvideo/models/dits/matrixgame2/model.py index f3c680a9d..e53bbbb73 100644 --- a/fastvideo/models/dits/matrixgame2/model.py +++ b/fastvideo/models/dits/matrixgame2/model.py @@ -29,7 +29,7 @@ from fastvideo.layers.visual_embedding import ( ) from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import ( +from fastvideo.models.wan.transformer import ( WanSelfAttention, WanI2VCrossAttention, WanT2VCrossAttention, diff --git a/fastvideo/models/dits/matrixgame3/model.py b/fastvideo/models/dits/matrixgame3/model.py index 2d1e64749..fa29fac5f 100644 --- a/fastvideo/models/dits/matrixgame3/model.py +++ b/fastvideo/models/dits/matrixgame3/model.py @@ -25,7 +25,7 @@ from fastvideo.layers.visual_embedding import ( ) from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import WanSelfAttention +from fastvideo.models.wan.transformer import WanSelfAttention from fastvideo.platforms import AttentionBackendEnum # Import ActionModule diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 1f3ebc2bd..135730a83 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -31,7 +31,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "WanTransformer3DModel": ("wan", "transformer", "WanTransformer3DModel"), "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"), "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"), - "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), + "CausalWanTransformer3DModel": ("wan", "causal_transformer", "CausalWanTransformer3DModel"), "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"), "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"), "LongCatVideoTransformer3DModel": @@ -55,7 +55,7 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { "WanTransformer3DModel": ("wan", "transformer", "WanTransformer3DModel"), "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"), "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"), - "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), + "CausalWanTransformer3DModel": ("wan", "causal_transformer", "CausalWanTransformer3DModel"), "LingBotWorld2CausalFastTransformer3DModel": ( "dits", "lingbotworld2", diff --git a/fastvideo/models/wan/AGENTS.md b/fastvideo/models/wan/AGENTS.md index d6cf61a44..13f6ab795 100644 --- a/fastvideo/models/wan/AGENTS.md +++ b/fastvideo/models/wan/AGENTS.md @@ -1,15 +1,27 @@ # Wan model family -This package owns the dense Wan transformer (`transformer.py`), its architecture, +This package owns the dense and causal Wan transformers (`transformer.py`, +`causal_transformer.py`), their architecture, FSDP predicates, and checkpoint/LoRA mappings (`config.py`), and the Wan VAE (`vae.py`) with its config (`vae_config.py`). The VAE includes both the video -encoder and decoder. Shared text/image encoders, VAE utilities, pipelines, -pipeline configs, and training adapters remain in their existing directories. -Causal Wan and other families reuse these classes through compatibility imports. +encoder and decoder. `pipeline_config.py` owns component and pipeline defaults; +`definition.py` links supported variants to those configs and existing sampling +presets. Shared text/image encoders, VAE utilities, pipelines, sampling presets, +and training adapters remain in their existing directories. +Other families reuse these classes through canonical imports. Old paths remain +explicit compatibility exports. ## Invariants - Keep `__init__.py` lightweight: no eager model or pipeline imports. +- Keep `definition.py` data-only. Config and preset names reference their + existing owners; do not copy defaults into the catalog. Its registration + groups preserve first-match detector ordering around DreamX. Checkpoint + manifests and explicit pipeline overrides still choose the executable + pipeline; a definition must not silently pin that choice. +- Keep `configs.pipelines.wan` as explicit compatibility aliases, including + `t5_postprocess_text` for old serialized configs. Shared encoder classes + remain in `configs.models.encoders` and `models.encoders`. - Import configs directly from `fastvideo.models.wan.config` or `fastvideo.models.wan.vae_config` in the matching component. - Preserve the old `models.dits.wanvideo` and `configs.models.dits.wanvideo` @@ -19,11 +31,15 @@ Causal Wan and other families reuse these classes through compatibility imports. - Keep `EntryClass = WanTransformer3DModel` in `transformer.py` only. Registry architecture names, state-dict keys, layer names, and mappings are compatibility contracts. +- Keep `EntryClass = CausalWanTransformer3DModel` in `causal_transformer.py` + only; `dits/causal_wanvideo.py` is an alias. Preserve cache layout, sink + eviction, absolute/relativistic RoPE, and the global-attention window limit. - Keep `EntryClass = AutoencoderKLWan` in `vae.py` only. Preserve latent normalization, first-frame handling, cache reset, streaming, tiling, and encoder/decoder compile conditions. Do not merge the separate Cosmos25, Gen3C, or LingBotWorld2 VAE adapters into this implementation. -- `config.py`, `vae_config.py`, and `__init__.py` are pre-commit checked; +- `config.py`, `vae_config.py`, `pipeline_config.py`, `definition.py`, and + `__init__.py` are pre-commit checked; `transformer.py` and `vae.py` retain the existing model-code exclusion. Avoid unrelated reformatting. @@ -35,15 +51,17 @@ component, before default or full-quality renders. Imports may require GPU dependencies even when a check needs no weights. ```bash +pytest fastvideo/tests/api/test_wan_definitions.py -q pytest fastvideo/tests/loader/test_wan_family_imports.py -q pytest fastvideo/tests/contract/test_merge_ci_plan.py -q pytest fastvideo/tests/vaes/test_wan_vae_compile.py -q ``` -`fastvideo/tests/golden_gate/test_wan_t2v.py` checks dense transformer block 0, -not the VAE. Until a VAE golden exists, use `fastvideo/tests/vaes/test_wan_vae.py` -for numerical VAE changes; it requires one CUDA GPU and checkpoint access and -compares encode/decode against Diffusers plus streaming against full decode. -Use focused T2V/I2V SSIM for pipeline behavior not covered by component checks. -For a pure relocation, compare against the unchanged parent with identical -weights, settings, and runtime; aliases alone are not numerical parity evidence. +Use `bash scripts/validate_wan.sh [golden|parity|default]` +from the repo root; it stops at the first failed boundary. The default is the +small golden tier. The four Wan gates cover dense block 0, VAE encode/decode +and streaming, causal block cache updates, and a three-step dense trajectory. +The independent Diffusers component checks and focused SSIM remain available +at higher tiers. A relocation needs unchanged-parent numerical evidence; +aliases alone are not proof. See `basic/wan/AGENTS.md` under pipelines for +sampling ownership; shared T5/UMT5 and CLIP stay shared. diff --git a/fastvideo/models/wan/causal_transformer.py b/fastvideo/models/wan/causal_transformer.py new file mode 100644 index 000000000..08d992a6c --- /dev/null +++ b/fastvideo/models/wan/causal_transformer.py @@ -0,0 +1,817 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from torch.nn.attention.flex_attention import create_block_mask, flex_attention +from torch.nn.attention.flex_attention import BlockMask +# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention +# see https://github.com/pytorch/pytorch/issues/133254 +# change to default for other models +flex_attention = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") +import torch.distributed as dist + +import fastvideo.envs as envs +from fastvideo.attention import (DistributedAttention, LocalAttention) +from fastvideo.models.wan.config import WanVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.forward_context import get_forward_context +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import (PatchEmbed) +from fastvideo.logger import init_logger +from fastvideo.models.dits._relative_rope import relativistic_window_offsets +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.wan.transformer import WanT2VCrossAttention, WanTimeTextImageEmbedding +from fastvideo.platforms import AttentionBackendEnum, current_platform + +logger = init_logger(__name__) + +GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES = 21 + + +class CausalWanSelfAttention(nn.Module): + + def __init__(self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm=True, + eps=1e-6, + parallel_attention=False, + rope_cache_policy: str = "absolute") -> None: + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.eps = eps + self.parallel_attention = parallel_attention + self.rope_cache_policy = rope_cache_policy + + # Scaled dot product attention + self.attn = LocalAttention(num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) + + def forward(self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + kv_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + frame_seqlen: int = 1560): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + frame_seqlen (int): Number of tokens per latent frame, + e.g. 1560 for 480x832 resolution. + """ + if cache_start is None: + cache_start = current_start + + cos, sin = freqs_cis + # relativistic defers roping until the cache window is known (and caches raw k) + relativistic = self.rope_cache_policy == "relativistic" and kv_cache is not None + if not relativistic: + roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) + roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) + + if kv_cache is None: + # Padding for flex attention + padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] + padded_roped_query = torch.cat([ + roped_query, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], device=q.device, dtype=v.dtype) + ], + dim=1) + + padded_roped_key = torch.cat([ + roped_key, + torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], device=k.device, dtype=v.dtype) + ], + dim=1) + + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype)], + dim=1) + + x = flex_attention(query=padded_roped_query.transpose(2, 1), + key=padded_roped_key.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask)[:, :, :-padded_length].transpose(2, 1) + else: + current_end = current_start + q.shape[1] + sink_tokens = self.sink_size * frame_seqlen + if self.local_attn_size == -1: + max_attention_size = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES * frame_seqlen) + else: + max_attention_size = self.local_attn_size * frame_seqlen + if self.local_attn_size == -1 and current_end > max_attention_size: + raise ValueError("Causal Wan local_attn_size=-1 keeps the previous " + f"{GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES}-latent-frame KV " + "window for compatibility. Set local_attn_size for " + f"longer rollouts; got current_end={current_end} tokens " + f"with frame_seqlen={frame_seqlen}.") + # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache + kv_cache_size = kv_cache["k"].shape[1] + num_new_tokens = q.shape[1] + stored_key = k if relativistic else roped_key # raw vs roped in cache + global_end_index = (int(kv_cache["global_end_index"].item()) if isinstance( + kv_cache["global_end_index"], torch.Tensor) else int(kv_cache["global_end_index"])) + local_end_index_prev = (int(kv_cache["local_end_index"].item()) if isinstance( + kv_cache["local_end_index"], torch.Tensor) else int(kv_cache["local_end_index"])) + if self.local_attn_size != -1 and (current_end + > global_end_index) and (num_new_tokens + local_end_index_prev + > kv_cache_size): + # Calculate the number of new tokens added in this step + # Shift existing cache content left to discard oldest tokens + # Clone the source slice to avoid overlapping memory error + num_evicted_tokens = num_new_tokens + local_end_index_prev - kv_cache_size + num_rolled_tokens = local_end_index_prev - num_evicted_tokens - sink_tokens + kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = local_end_index_prev + current_end - \ + global_end_index - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = stored_key + kv_cache["v"][:, local_start_index:local_end_index] = v + else: + # Assign new keys/values directly up to current_end + local_end_index = local_end_index_prev + current_end - global_end_index + local_start_index = local_end_index - num_new_tokens + kv_cache["k"] = kv_cache["k"].detach() + kv_cache["v"] = kv_cache["v"].detach() + # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) + kv_cache["k"][:, local_start_index:local_end_index] = stored_key + kv_cache["v"][:, local_start_index:local_end_index] = v + key_window = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index] + value_window = kv_cache["v"][:, max(0, local_end_index - max_attention_size):local_end_index] + if relativistic: + window_len, query_lo, query_hi = relativistic_window_offsets(local_end_index, num_new_tokens, + max_attention_size) + roped_query = _apply_rotary_emb(q, cos[query_lo:query_hi], sin[query_lo:query_hi], + is_neox_style=False).type_as(v) + key_window = _apply_rotary_emb(key_window, cos[:window_len], sin[:window_len], + is_neox_style=False).type_as(v) + x = self.attn(roped_query, key_window, value_window) + if isinstance(kv_cache["global_end_index"], torch.Tensor): + kv_cache["global_end_index"].fill_(current_end) + else: + kv_cache["global_end_index"] = current_end + if isinstance(kv_cache["local_end_index"], torch.Tensor): + kv_cache["local_end_index"].fill_(local_end_index) + else: + kv_cache["local_end_index"] = local_end_index + + return x + + +class CausalWanTransformerBlock(nn.Module): + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = "", + rope_cache_policy: str = "absolute"): + super().__init__() + + # 1. Self-attention + self.norm1 = nn.LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + + self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.attn1 = CausalWanSelfAttention(dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps, + rope_cache_policy=rope_cache_policy) + self.hidden_dim = dim + self.num_attention_heads = num_heads + self.local_attn_size = local_attn_size + dim_head = dim // num_heads + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + # LTX applies qk norm across all heads + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + print("QK Norm type not supported") + raise Exception + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32) + + # 2. Cross-attention + # Only T2V for now + self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + kv_cache: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + frame_seqlen: int | None = None, + ) -> torch.Tensor: + # hidden_states.shape: [batch_size, seq_length, inner_dim] + # temb.shape: [batch_size, temb_seq_len, 6, inner_dim] + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + temb_seq_len = temb.shape[1] + tokens_per_temb = hidden_states.shape[1] // temb_seq_len + if frame_seqlen is None: + frame_seqlen = tokens_per_temb + else: + frame_seqlen = int(frame_seqlen) + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + # assert orig_dtype != torch.float32 + e = self.scale_shift_table + temb + # e.shape: [batch_size, temb_seq_len, 6, inner_dim] + assert e.shape == (bs, temb_seq_len, 6, self.hidden_dim) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) + # *_msa.shape: [batch_size, temb_seq_len, 1, inner_dim] + # assert shift_msa.dtype == torch.float32 + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(temb_seq_len, tokens_per_temb)) * + (1 + scale_msa) + shift_msa).flatten(1, 2) + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + attn_output = self.attn1( + query, + key, + value, + freqs_cis, + block_mask, + kv_cache, + current_start, + cache_start, + frame_seqlen=frame_seqlen, + ) + attn_output = attn_output.flatten(2) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + null_shift = null_scale = torch.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, + null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states, + context=encoder_hidden_states, + context_lens=None, + crossattn_cache=crossattn_cache) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa, + c_scale_msa) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + + return hidden_states + + +class CausalWanTransformer3DModel(BaseDiT): + _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions + _compile_conditions = WanVideoConfig()._compile_conditions + _supported_attention_backends = WanVideoConfig()._supported_attention_backends + param_names_mapping = WanVideoConfig().param_names_mapping + reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping + lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping + + def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.attention_head_dim = config.attention_head_dim + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.text_len = config.text_len + self.local_attn_size = config.local_attn_size + self.rope_cache_policy = config.arch_config.rope_cache_policy + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False) + + # 2. Condition embeddings + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + text_embed_dim=config.text_dim, + image_embed_dim=config.image_dim, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + CausalWanTransformerBlock(inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.local_attn_size, + config.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{config.prefix}.blocks.{i}", + rope_cache_policy=config.arch_config.rope_cache_policy) + for i in range(config.num_layers) + ]) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32) + self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + + # Causal-specific + self.block_mask = None + self.teacher_forcing_block_mask = None + self.num_frame_per_block = config.arch_config.num_frames_per_block + assert self.num_frame_per_block <= 3 + self.independent_first_frame = False + + self.__post_init__() + + @staticmethod + def _prepare_blockwise_causal_attn_mask(device: torch.device | str, + num_frames: int = 21, + frame_seqlen: int = 1560, + num_frame_per_block=1, + local_attn_size=-1) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange(start=0, end=total_length, step=frame_seqlen * num_frame_per_block, device=device) + + for tmp in frame_indices: + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & + (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) + # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask + + block_mask = create_block_mask(attention_mask, + B=None, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=False, + device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(block_mask) + + # import imageio + # import numpy as np + # from torch.nn.attention.flex_attention import create_mask + + # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + # padded_length, KV_LEN=total_length + padded_length, device=device) + # import cv2 + # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + + @staticmethod + def _prepare_teacher_forcing_mask(device: torch.device | str, + num_frames: int = 21, + frame_seqlen: int = 1560, + num_frame_per_block=1, + local_attn_size=-1) -> BlockMask: + """Attention mask for the teacher-forcing ``[clean | noisy]`` sequence. + + A noisy token attends to its own block plus the clean context of all + strictly previous blocks; clean tokens are block-wise causal. + """ + if local_attn_size != -1: + raise NotImplementedError(f"Teacher forcing ignores local_attn_size={local_attn_size}: " + "unlike the block-wise causal mask, this mask always attends " + "to the full clean context. Windowed teacher forcing is not " + "implemented; use local_attn_size=-1 for teacher-forcing " + "training.") + total_length = num_frames * frame_seqlen * 2 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + clean_ends = num_frames * frame_seqlen + context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + + attention_block_size = frame_seqlen * num_frame_per_block + frame_indices = torch.arange(start=0, + end=num_frames * frame_seqlen, + step=attention_block_size, + device=device, + dtype=torch.long) + for start in frame_indices: + context_ends[start:start + attention_block_size] = start + attention_block_size + + noisy_image_start_list = torch.arange(num_frames * frame_seqlen, + total_length, + step=attention_block_size, + device=device, + dtype=torch.long) + noisy_image_end_list = noisy_image_start_list + attention_block_size + for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): + noise_noise_starts[start:end] = start + noise_noise_ends[start:end] = end + noise_context_ends[start:end] = block_index * attention_block_size + + def attention_mask(b, h, q_idx, kv_idx): + clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) + c1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) + c2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) + noise_mask = (q_idx >= clean_ends) & (c1 | c2) + eye_mask = q_idx == kv_idx + return eye_mask | clean_mask | noise_mask + + block_mask = create_block_mask(attention_mask, + B=None, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=False, + device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames") + print(block_mask) + + return block_mask + + def _forward_inference(self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + kv_cache: dict = None, + crossattn_cache: dict = None, + current_start: int = 0, + cache_start: int = 0, + start_frame: int = 0, + **kwargs) -> torch.Tensor: + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (1560 tokens each) + """ + + orig_dtype = hidden_states.dtype + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = (encoder_hidden_states_image[0] + if len(encoder_hidden_states_image) > 0 else None) + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + if self.rope_cache_policy == "relativistic": + # fixed table over [0, max_attention_frames); attention slices it per step + max_attention_frames = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES + if self.local_attn_size == -1 else self.local_attn_size) + rope_num_frames = max_attention_frames * get_sp_world_size() + rope_start_frame = 0 + else: + rope_num_frames = post_patch_num_frames * get_sp_world_size() + rope_start_frame = start_frame # 0 when kv_cache is None + freqs_cos, freqs_sin = get_rotary_pos_embed((rope_num_frames, post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=rope_start_frame) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack([torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + encoder_hidden_states = torch.cat([ + encoder_hidden_states, + encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), + encoder_hidden_states.size(2)) + ], + dim=1) + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) + + encoder_hidden_states = encoder_hidden_states.to( + orig_dtype) if current_platform.is_mps() else encoder_hidden_states # cast to orig_dtype for MPS + + assert encoder_hidden_states.dtype == orig_dtype + + # 4. Transformer blocks + for block_index, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + causal_kwargs = { + "kv_cache": kv_cache[block_index], + "current_start": current_start, + "cache_start": cache_start, + "block_mask": self.block_mask, + "frame_seqlen": post_patch_height * post_patch_width, + } + hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states, + timestep_proj, freqs_cis, **causal_kwargs) + else: + causal_kwargs = { + "kv_cache": kv_cache[block_index], + "crossattn_cache": (crossattn_cache[block_index] if crossattn_cache is not None else None), + "current_start": current_start, + "cache_start": cache_start, + "block_mask": self.block_mask, + "frame_seqlen": post_patch_height * post_patch_width, + } + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, **causal_kwargs) + + # 5. Output norm, projection & unpatchify + temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + output = self.unpatchify(hidden_states, grid_sizes) + + return torch.stack(output) + + def _forward_train(self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] + | None = None, + start_frame: int = 0, + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, + **kwargs) -> torch.Tensor: + + orig_dtype = hidden_states.dtype + teacher_forcing = clean_x is not None + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, list): + encoder_hidden_states_image = (encoder_hidden_states_image[0] + if len(encoder_hidden_states_image) > 0 else None) + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=start_frame) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + if teacher_forcing: + if self.teacher_forcing_block_mask is None: + self.teacher_forcing_block_mask = self._prepare_teacher_forcing_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) + block_mask = self.teacher_forcing_block_mask + else: + if self.block_mask is None: + self.block_mask = self._prepare_blockwise_causal_attn_mask(device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * + post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size) + block_mask = self.block_mask + + hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack([torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + encoder_hidden_states = torch.cat([ + encoder_hidden_states, + encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), + encoder_hidden_states.size(2)) + ], + dim=1) + encoder_hidden_states_text = encoder_hidden_states + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1) + + encoder_hidden_states = encoder_hidden_states.to( + orig_dtype) if current_platform.is_mps() else encoder_hidden_states # cast to orig_dtype for MPS + + assert encoder_hidden_states.dtype == orig_dtype + + if teacher_forcing: + # Tile RoPE/modulation so clean frame i and noisy frame i share a position. + clean_tokens = self.patch_embedding(clean_x).flatten(2).transpose(1, 2) + hidden_states = torch.cat([clean_tokens, hidden_states], dim=1) + if aug_t is None: + aug_t = torch.zeros_like(timestep) + _, timestep_proj_clean, _, _ = self.condition_embedder(aug_t.flatten(), encoder_hidden_states_text, None) + timestep_proj_clean = timestep_proj_clean.unflatten(1, + (6, self.hidden_size)).unflatten(dim=0, + sizes=timestep.shape) + timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1) + freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0), torch.cat([freqs_sin, freqs_sin], dim=0)) + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.blocks: + hidden_states = self._gradient_checkpointing_func(block, + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + block_mask=block_mask) + else: + for block in self.blocks: + hidden_states = block(hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + block_mask=block_mask) + + if teacher_forcing: + hidden_states = hidden_states[:, hidden_states.shape[1] // 2:] + + # 5. Output norm, projection & unpatchify + temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + output = self.unpatchify(hidden_states, grid_sizes) + + return torch.stack(output) + + def forward(self, *args, **kwargs): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) + + def unpatchify(self, x, grid_sizes): + r""" + + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + + c = self.out_channels + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = u.permute(6, 0, 3, 1, 4, 2, 5) + # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + +# Entry point for model registry +EntryClass = CausalWanTransformer3DModel diff --git a/fastvideo/models/wan/definition.py b/fastvideo/models/wan/definition.py new file mode 100644 index 000000000..b0b673b65 --- /dev/null +++ b/fastvideo/models/wan/definition.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Data-only Wan variant definitions consumed by the shared registry. + +Config names reference ``pipeline_config.py``; preset names reference +``fastvideo/pipelines/basic/wan/presets.py``. Neither values nor factories are +duplicated here. Checkpoint manifests and explicit pipeline overrides still +select the executable pipeline; these definitions do not pin or replace them. +""" + +from dataclasses import dataclass +from typing import Literal + +# Dense DMD indexes the complete training-noise table. This is independent of +# the configurable inference flow_shift and must not become a shared scheduler +# instance: timestep preparation mutates its scheduler. +DMD_TRAINING_NOISE_SHIFT = 8.0 + + +@dataclass(frozen=True) +class WanModelDefinition: + pipeline_config: str + preset: str + sampling: Literal["unipc", "dmd", "causal_dmd"] + hf_model_paths: tuple[str, ...] + workload_types: tuple[Literal["t2v", "i2v"], ...] + match_any: tuple[str, ...] = () + match_all: tuple[str, ...] = () + exclude: tuple[str, ...] = () + + def matches(self, path_or_class: str) -> bool: + """Preserve legacy substring detectors, including their broad matches.""" + value = path_or_class.lower() + return (any(token in value for token in self.match_any) and all(token in value for token in self.match_all) + and not any(token in value for token in self.exclude)) + + +# The two groups preserve the registry's existing first-match ordering around +# DreamX. Keeping this boundary avoids changing local/custom-checkpoint +# resolution when both a Wan and DreamX detector match. Within each group, +# declaration order is registration order. +WAN_MODEL_DEFINITION_GROUPS = ( + ( + WanModelDefinition( + pipeline_config="WanT2V480PConfig", + preset="wan_t2v_1_3b", + sampling="unipc", + hf_model_paths=("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", ), + workload_types=("t2v", ), + match_any=("wanpipeline", ), + ), + WanModelDefinition( + pipeline_config="WanT2V720PConfig", + preset="wan_t2v_14b", + sampling="unipc", + hf_model_paths=("Wan-AI/Wan2.1-T2V-14B-Diffusers", "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers"), + workload_types=("t2v", ), + ), + WanModelDefinition( + pipeline_config="WanI2V480PConfig", + preset="wan_i2v_14b_480p", + sampling="unipc", + hf_model_paths=("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", ), + workload_types=("i2v", ), + match_any=("wanimagetovideo", ), + ), + WanModelDefinition( + pipeline_config="WanI2V720PConfig", + preset="wan_i2v_14b_720p", + sampling="unipc", + hf_model_paths=("Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", ), + workload_types=("i2v", ), + ), + WanModelDefinition( + pipeline_config="WanI2V480PConfig", + preset="wan_fun_1_3b_inp", + sampling="unipc", + hf_model_paths=("weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", ), + workload_types=("i2v", ), + ), + WanModelDefinition( + pipeline_config="WANV2VConfig", + preset="wan_fun_1_3b_control", + sampling="unipc", + hf_model_paths=("IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers", ), + workload_types=(), + ), + WanModelDefinition( + pipeline_config="FastWan2_1_T2V_480P_Config", + preset="fast_wan_t2v_480p", + sampling="dmd", + hf_model_paths=("FastVideo/FastWan2.1-T2V-1.3B-Diffusers", "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers"), + workload_types=("t2v", ), + match_any=("wandmdpipeline", ), + ), + WanModelDefinition( + pipeline_config="Wan2_2_TI2V_5B_Config", + preset="wan_2_2_ti2v_5b", + sampling="unipc", + hf_model_paths=("Wan-AI/Wan2.2-TI2V-5B-Diffusers", ), + workload_types=("t2v", "i2v"), + ), + ), + ( + WanModelDefinition( + pipeline_config="FastWan2_2_TI2V_5B_Config", + preset="fast_wan_2_2_ti2v_5b", + sampling="dmd", + hf_model_paths=("FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + "FastVideo/FastWan2.2-TI2V-5B-Diffusers"), + workload_types=("t2v", "i2v"), + ), + WanModelDefinition( + pipeline_config="LucyEditDevConfig", + preset="lucy_edit_dev", + sampling="unipc", + hf_model_paths=("decart-ai/Lucy-Edit-Dev", "decart-ai/Lucy-Edit-1.1-Dev"), + workload_types=(), + match_any=("lucy-edit", ), + ), + WanModelDefinition( + pipeline_config="Wan2_2_T2V_A14B_Config", + preset="wan_2_2_t2v_a14b", + sampling="unipc", + hf_model_paths=("Wan-AI/Wan2.2-T2V-A14B-Diffusers", ), + workload_types=("t2v", ), + ), + WanModelDefinition( + pipeline_config="Wan2_2_I2V_A14B_Config", + preset="wan_2_2_i2v_a14b", + sampling="unipc", + hf_model_paths=("Wan-AI/Wan2.2-I2V-A14B-Diffusers", ), + workload_types=("i2v", ), + ), + WanModelDefinition( + pipeline_config="SelfForcingWanT2V480PConfig", + preset="sf_wan_t2v_1_3b", + sampling="causal_dmd", + hf_model_paths=("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", ), + workload_types=("t2v", ), + match_any=("wancausaldmdpipeline", ), + ), + WanModelDefinition( + pipeline_config="SelfForcingWan2_2_T2V480PConfig", + preset="sf_wan_2_2_t2v_a14b", + sampling="causal_dmd", + hf_model_paths=("rand0nmr/SFWan2.2-T2V-A14B-Diffusers", ), + workload_types=("t2v", ), + match_any=("sfwan2.2", "sfwan2_2"), + exclude=("i2v", ), + ), + WanModelDefinition( + pipeline_config="SelfForcingWan2_2_T2V480PConfig", + preset="sf_wan_2_2_i2v_a14b", + sampling="causal_dmd", + hf_model_paths=("FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers", ), + workload_types=("i2v", ), + match_any=("sfwan2.2", "sfwan2_2"), + match_all=("i2v", ), + ), + ), +) + +WAN_MODEL_DEFINITIONS = tuple(definition for group in WAN_MODEL_DEFINITION_GROUPS for definition in group) diff --git a/fastvideo/models/wan/pipeline_config.py b/fastvideo/models/wan/pipeline_config.py new file mode 100644 index 000000000..66599c815 --- /dev/null +++ b/fastvideo/models/wan/pipeline_config.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan pipeline and component defaults; sampling presets stay pipeline-local.""" + +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.models.wan.config import WanVideoArchConfig, WanVideoConfig +from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config, + WAN2_1ControlCLIPVisionConfig) +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.models.wan.vae_config import WanVAEArchConfig, WanVAEConfig + + +def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: + mask: torch.Tensor = outputs.attention_mask + hidden_state: torch.Tensor = outputs.last_hidden_state + seq_lens = mask.gt(0).sum(dim=1).long() + assert torch.isnan(hidden_state).sum() == 0 + prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] + prompt_embeds_tensor: torch.Tensor = torch.stack( + [torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0) + return prompt_embeds_tensor + + +@dataclass +class WanT2V480PConfig(PipelineConfig): + """Base configuration for Wan T2V 1.3B pipeline architecture.""" + + # WanConfig-specific parameters with defaults + # DiT + dit_config: DiTConfig = field(default_factory=WanVideoConfig) + # VAE + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Denoising stage + flow_shift: float | None = 3.0 + + # Text encoding stage + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], + ...] = field(default_factory=lambda: (t5_postprocess_text, )) + + # Precision for each component + precision: str = "bf16" + # Keep the VAE *encode* in fp32: for I2V/causal Wan variants the image/video + # VAE encode seeds the denoising trajectory, and bf16 there shifts the output + # (e.g. Matrix-Game-2.0 SSIM 0.98 -> 0.76). Decode is output-only, so we lower + # just the decode below. + vae_precision: str = "fp32" + # bf16 VAE *decode* is effectively lossless (MS-SSIM 0.9999 vs fp32 on an + # identical latent) and faster; the same AutoencoderKLWan already runs bf16 + # in Cosmos-Predict2.5 and fp16 in Cosmos. Inherited by all Wan variants + # (incl. I2V/causal) — decode-only, so encode trajectories are untouched. + vae_decode_precision: str = "bf16" + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32", )) + + # self-forcing params + warp_denoising_step: bool = True + + # WanConfig-specific added parameters + + def __post_init__(self): + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + +@dataclass +class WanT2V720PConfig(WanT2V480PConfig): + """Base configuration for Wan T2V 14B 720P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: float | None = 5.0 + + +@dataclass +class WanI2V480PConfig(WanT2V480PConfig): + """Base configuration for Wan I2V 14B 480P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Precision for each component + image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig) + image_encoder_precision: str = "fp32" + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + +@dataclass +class WanI2V720PConfig(WanI2V480PConfig): + """Base configuration for Wan I2V 14B 720P pipeline architecture.""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: float | None = 5.0 + + +@dataclass +class WANV2VConfig(WanI2V480PConfig): + """Configuration for WAN2.1 1.3B Control pipeline.""" + + image_encoder_config: EncoderConfig = field(default_factory=WAN2_1ControlCLIPVisionConfig) + # CLIP encoder precision + image_encoder_precision: str = 'bf16' + + +@dataclass +class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): + """Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD""" + + # WanConfig-specific parameters with defaults + + # Denoising stage + flow_shift: float | None = 8.0 + dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) + + +@dataclass +class Wan2_2_TI2V_5B_Config(WanT2V480PConfig): + flow_shift: float | None = 5.0 + ti2v_task: bool = True + expand_timesteps: bool = True + + def __post_init__(self) -> None: + assert not (self.ti2v_task and self.lucy_edit_task) + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + self.dit_config.expand_timesteps = self.expand_timesteps + + +@dataclass +class LucyEditDevConfig(Wan2_2_TI2V_5B_Config): + """Configuration for Decart Lucy Edit Dev video editing.""" + + dit_config: DiTConfig = field(default_factory=lambda: WanVideoConfig(arch_config=WanVideoArchConfig( + num_attention_heads=24, + in_channels=96, + out_channels=48, + ffn_dim=14336, + num_layers=30, + ))) + vae_config: VAEConfig = field(default_factory=lambda: WanVAEConfig(arch_config=WanVAEArchConfig( + base_dim=160, + decoder_base_dim=256, + z_dim=48, + in_channels=12, + out_channels=12, + scale_factor_spatial=16, + patch_size=2, + is_residual=True, + clip_output=False, + latents_mean=( + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ), + latents_std=( + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ), + ))) + ti2v_task: bool = False + lucy_edit_task: bool = True + + def __post_init__(self) -> None: + assert not (self.ti2v_task and self.lucy_edit_task) + # Lucy uses Wan2.2's enhanced 48-channel VAE latents. Denoising + # concatenates noise + video latents, matching the 96-channel + # transformer input declared above. + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + self.dit_config.expand_timesteps = self.expand_timesteps + + +@dataclass +class FastWan2_2_TI2V_5B_Config(Wan2_2_TI2V_5B_Config): + flow_shift: float | None = 5.0 + dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 757, 522]) + + +@dataclass +class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): + flow_shift: float | None = 12.0 + boundary_ratio: float | None = 0.875 + + # self-forcing params + dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250]) + warp_denoising_step: bool = True + + def __post_init__(self) -> None: + self.dit_config.boundary_ratio = self.boundary_ratio + + +@dataclass +class Wan2_2_I2V_A14B_Config(WanI2V480PConfig): + flow_shift: float | None = 5.0 + boundary_ratio: float | None = 0.900 + + def __post_init__(self) -> None: + super().__post_init__() + self.dit_config.boundary_ratio = self.boundary_ratio + + +# ============================================= +# ============= Causal Self-Forcing ============= +# ============================================= +@dataclass +class SelfForcingWanT2V480PConfig(WanT2V480PConfig): + is_causal: bool = True + flow_shift: float | None = 5.0 + dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 750, 500, 250]) + warp_denoising_step: bool = True + + +@dataclass +class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config): + is_causal: bool = True + flow_shift: float | None = 12.0 + boundary_ratio: float | None = 0.875 + dmd_denoising_steps: list[int] | None = field(default_factory=lambda: [1000, 850, 700, 550, 350, 275, 200, 125]) + warp_denoising_step: bool = True + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True diff --git a/fastvideo/pipelines/AGENTS.md b/fastvideo/pipelines/AGENTS.md index 56803ee31..0ec613925 100644 --- a/fastvideo/pipelines/AGENTS.md +++ b/fastvideo/pipelines/AGENTS.md @@ -64,6 +64,12 @@ Reuse `stages/text_encoding.py` if your model takes text → embeddings via a st When forking, keep the file name model-prefixed (`longcat_*`, `gamecraft_*`) so the registry stays grep-able. +Wan's family-specific stages live in `basic/wan/stages/`; see +`basic/wan/AGENTS.md`. Shared dense scheduling/CFG stays in `stages/denoising.py`. +Do not put Wan expert selection, first-frame VAE execution, or DMD/causal +sampling recipes back into that shared loop. Legacy sampler imports remain +lazy aliases to avoid family/shared import cycles. + ## Per-Package `AGENTS.md` and `JOURNAL.md` (optional) Pipelines with non-trivial parity invariants, lazy-loaded shared components, diff --git a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py index 72656bb8b..e4bf62e84 100644 --- a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py +++ b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py @@ -2,6 +2,8 @@ """DreamX-World video pipeline entrypoint.""" from fastvideo.api.sampling_param import SamplingParam +from fastvideo.pipelines.basic.wan.stages.conditioning import WanFirstFrameEncodingStage +from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler @@ -11,7 +13,6 @@ from fastvideo.pipelines.basic.dreamx_world.stages import DreamXWorldCameraCondi from fastvideo.pipelines.stages import ( ConditioningStage, DecodingStage, - DenoisingStage, InputValidationStage, LatentPreparationStage, TextEncodingStage, @@ -59,13 +60,15 @@ class DreamXWorldPipeline(LoRAPipeline, ComposedPipelineBase): self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage()) + self.add_stage(stage_name="first_frame_encoding_stage", + stage=WanFirstFrameEncodingStage(vae=self.get_module("vae"))) + self.add_stage( stage_name="denoising_stage", - stage=DenoisingStage( + stage=WanDenoisingStage( transformer=self.get_module("transformer"), transformer_2=self.get_module("transformer_2", None), scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"), pipeline=self, ), ) diff --git a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py index f6a8dccce..0c591c715 100644 --- a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py +++ b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_i2v_pipeline.py @@ -12,13 +12,13 @@ Key differences from T2V: - Mask conditioning with encoded first frame """ +from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_rcm import RCMScheduler from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, DenoisingStage, ImageVAEEncodingStage, - InputValidationStage, LatentPreparationStage, TextEncodingStage, - TimestepPreparationStage) +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, ImageVAEEncodingStage, InputValidationStage, + LatentPreparationStage, TextEncodingStage, TimestepPreparationStage) logger = init_logger(__name__) @@ -63,11 +63,10 @@ class TurboDiffusionI2VPipeline(LoRAPipeline, ComposedPipelineBase): stage=ImageVAEEncodingStage(vae=self.get_module("vae"))) self.add_stage(stage_name="denoising_stage", - stage=DenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2", None), - scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"), - pipeline=self)) + stage=WanDenoisingStage(transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + pipeline=self)) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) diff --git a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py index f49f0c91c..f8a84a9ed 100644 --- a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py +++ b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py @@ -7,12 +7,13 @@ for 1-4 step video generation using rCM (recurrent Consistency Model) sampling with SLA (Sparse-Linear Attention). """ +from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_rcm import RCMScheduler from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage, - LatentPreparationStage, TextEncodingStage, TimestepPreparationStage) +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, InputValidationStage, LatentPreparationStage, + TextEncodingStage, TimestepPreparationStage) logger = init_logger(__name__) @@ -52,11 +53,10 @@ class TurboDiffusionPipeline(LoRAPipeline, ComposedPipelineBase): transformer=self.get_module("transformer", None))) self.add_stage(stage_name="denoising_stage", - stage=DenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2", None), - scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"), - pipeline=self)) + stage=WanDenoisingStage(transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + pipeline=self)) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) diff --git a/fastvideo/pipelines/basic/wan/AGENTS.md b/fastvideo/pipelines/basic/wan/AGENTS.md new file mode 100644 index 000000000..dd90e4d94 --- /dev/null +++ b/fastvideo/pipelines/basic/wan/AGENTS.md @@ -0,0 +1,42 @@ +# Wan pipelines + +Pipeline files assemble existing stages. Sampling defaults stay in `presets.py`. +Pipeline configs and variant definitions live in `fastvideo/models/wan/` as +`pipeline_config.py` and `definition.py`. The shared registry consumes these +definitions; old pipeline config imports remain compatibility aliases. Shared text/image +encoders stay shared; the Wan VAE video encoder and decoder live together in +`fastvideo/models/wan/vae.py`. + +## Sampling boundaries + +- `stages/conditioning.py`: encode and normalize `first_frame_latent` before + dense TI2V or causal DMD sampling. Clear it for each request. Preserve VAE + precision, normalization order, and device restoration. +- `stages/denoising.py`: Wan input packing, expert boundaries, guidance choice, + Lucy timesteps, and TI2V first-frame preservation around the shared loop. + Request state belongs in `WanDenoisingState`, never on a reusable stage. +- `stages/dmd.py`: dense DMD sampling. The caller passes a full training-noise + `FlowMatchEulerDiscreteScheduler` with `DMD_TRAINING_NOISE_SHIFT` from the + family definition (8.0). Do not share the scheduler + mutated by `TimestepPreparationStage`: DMD indexes the full timestep/sigma + table when converting noise to video and adding noise. +- `stages/causal_denoising.py`: standard and DMD causal samplers share cache + allocation, not a sampling-loop inheritance chain. Reset UniPC per block and + caches per request; preserve clean-context writes and RNG ordering. + +Do not change scheduler arithmetic, autocast placement, step counts, offload, +or model math as incidental cleanup. Keep legacy sampler exports as aliases; +`CausalDMDDenosingStage` retains its historical spelling for compatibility. + +The first-frame stage has its own timing. Its encode time is no longer inside +`dit_time_s`; total request latency still includes it. Do not reseed performance +baselines automatically because this boundary moved. + +## Checks + +See [the testing guide](../../../../docs/contributing/testing.md). Start with +`bash scripts/validate_wan.sh all` from the repo root in a GPU environment. +It runs weight-free contracts before small device/runtime-matched goldens. +Pass `parity` or `default` as the second argument for independent component +parity or focused SSIM after those gates. Cache goldens do not establish +long-rollout quality, distilled-model quality, or complete pipeline parity. diff --git a/fastvideo/pipelines/basic/wan/lucy_edit_pipeline.py b/fastvideo/pipelines/basic/wan/lucy_edit_pipeline.py index 3843ad8c4..28a32fb59 100644 --- a/fastvideo/pipelines/basic/wan/lucy_edit_pipeline.py +++ b/fastvideo/pipelines/basic/wan/lucy_edit_pipeline.py @@ -6,13 +6,13 @@ the noisy latent channels. The stage topology is therefore closest to Wan V2V, but the model repo does not include CLIP image-encoder components. """ +from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.basic.wan.wan_v2v_pipeline import WanVideoToVideoPipeline from fastvideo.pipelines.stages import ( ConditioningStage, DecodingStage, - DenoisingStage, InputValidationStage, LatentPreparationStage, TextEncodingStage, @@ -67,7 +67,7 @@ class LucyEditPipeline(WanVideoToVideoPipeline): self.add_stage( stage_name="denoising_stage", - stage=DenoisingStage( + stage=WanDenoisingStage( transformer=self.get_module("transformer"), transformer_2=self.get_module("transformer_2"), scheduler=self.get_module("scheduler"), diff --git a/fastvideo/pipelines/basic/wan/stages/__init__.py b/fastvideo/pipelines/basic/wan/stages/__init__.py new file mode 100644 index 000000000..882d692a2 --- /dev/null +++ b/fastvideo/pipelines/basic/wan/stages/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan sampling and conditioning. Import each stage from its defining module.""" diff --git a/fastvideo/pipelines/basic/wan/stages/causal_denoising.py b/fastvideo/pipelines/basic/wan/stages/causal_denoising.py new file mode 100644 index 000000000..f89161e2e --- /dev/null +++ b/fastvideo/pipelines/basic/wan/stages/causal_denoising.py @@ -0,0 +1,589 @@ +# SPDX-License-Identifier: Apache-2.0 +import torch # type: ignore + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +try: + from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionBackend) + vsa_available = True +except ImportError: + vsa_available = False + VideoSparseAttentionBackend = None # type: ignore + +logger = init_logger(__name__) + + +def _get_transformer_attr(transformer, name: str, default): + # Currently used only for causal KV-cache runtime settings: + # `local_attn_size` and `sink_size`. + # + # Causal Wan/MatrixGame expose these directly on the transformer. Some + # wrapper-style DiTs keep the executable module under transformer.model. Do + # not fall back to transformer.config for these two values: config is only + # the construction source and does not control them at runtime. + value = getattr(transformer, name, None) + if value is not None: + return value + + inner_model = getattr(transformer, "model", None) + if inner_model is not None: + value = getattr(inner_model, name, None) + if value is not None: + return value + + return default + + +class WanCausalDenoisingBase(DenoisingStage): + """Shared Wan cache layout; standard and DMD sampling remain separate algorithms.""" + + def __init__(self, transformer, scheduler, transformer_2=None, vae=None) -> None: + super().__init__(transformer, scheduler, transformer_2=transformer_2) + # KV and cross-attention cache state (initialized on first forward) + self.transformer = transformer + self.transformer_2 = transformer_2 + # Model-dependent constants (aligned with causal_inference.py assumptions) + self.num_transformer_blocks = len(self.transformer.blocks) + self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block + self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames + + self.local_attn_size = _get_transformer_attr(self.transformer, "local_attn_size", -1) + self.sink_size = _get_transformer_attr(self.transformer, "sink_size", 0) + + def _initialize_kv_cache(self, batch_size, dtype, device) -> list[dict]: + """ + Initialize a Per-GPU KV cache aligned with the Wan model assumptions. + """ + kv_cache1 = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = self.transformer.attention_head_dim + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size * self.frame_seq_length + else: + kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames + + for _ in range(self.num_transformer_blocks): + kv_cache1.append({ + "k": + torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + + return kv_cache1 + + def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype, device) -> list[dict]: + """ + Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. + """ + crossattn_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = self.transformer.attention_head_dim + for _ in range(self.num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "is_init": + False, + }) + return crossattn_cache + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + """Verify denoising stage inputs.""" + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("image_embeds", batch.image_embeds, V.is_list) + result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)) + result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) + result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) + result.add_check("eta", batch.eta, V.non_negative_float) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value) + result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds, + lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)) + return result + + +class CausalDMDDenosingStage(WanCausalDenoisingBase): + """ + Denoising stage for causal diffusion. + """ + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast + + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + patch_ratio = self.transformer.config.arch_config.patch_size[ + -1] * self.transformer.config.arch_config.patch_size[-2] + self.frame_seq_length = latent_seq_length // patch_ratio + # TODO(will): make this a parameter once we add i2v support + independent_first_frame = self.transformer.independent_first_frame if hasattr( + self.transformer, 'independent_first_frame') else False + # Timesteps for DMD + timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu() + if fastvideo_args.pipeline_config.warp_denoising_step: + scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) + timesteps = scheduler_timesteps[1000 - timesteps] + timesteps = timesteps.to(get_local_torch_device()) + + if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None: + boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps + high_noise_timesteps = timesteps[timesteps >= boundary_timestep] + else: + boundary_timestep = None + high_noise_timesteps = None + + # Image kwargs (kept empty unless caller provides compatible args) + image_kwargs: dict = {} + + pos_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + # "encoder_hidden_states_2": batch.clip_embedding_pos, + "encoder_attention_mask": batch.prompt_attention_mask, + }, + ) + + # Latents and prompts + assert batch.latents is not None, "latents must be provided" + latents = batch.latents # [B, C, T, H, W] + b, c, t, h, w = latents.shape + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + # Initialize or reset caches + kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) + kv_cache2 = None + if boundary_timestep is not None: + # Initialize the low noise kv cache + kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + def _get_kv_cache(timestep: float) -> list[dict]: + if boundary_timestep is not None: + if timestep >= boundary_timestep: + return kv_cache1 + else: + assert kv_cache2 is not None, "kv_cache2 is not initialized" + return kv_cache2 + return kv_cache1 + + crossattn_cache = self._initialize_crossattn_cache( + batch_size=latents.shape[0], + max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0].arch_config.text_len, + dtype=target_dtype, + device=latents.device) + + pos_start_base = 0 + + # Determine block sizes + if t % self.num_frames_per_block != 0: + raise ValueError("num_frames must be divisible by num_frames_per_block for causal DMD denoising") + num_blocks = t // self.num_frames_per_block + block_sizes = [self.num_frames_per_block] * num_blocks + start_index = 0 + + # For now hardcode the first block to be 1 frame assuming the model is Wan2.2-MoE + if boundary_timestep is not None: + block_sizes[0] = 1 + + first_frame_latent = batch.first_frame_latent + if batch.pil_image is not None: + assert first_frame_latent is not None, "Run WanFirstFrameEncodingStage before causal DMD denoising" + + # Fill the low noise and high noise kv cache with first_frame_latent and timestep 0 + t_zero = torch.zeros([latents.shape[0], 1], device=latents.device, dtype=torch.long) + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), \ + set_forward_context(current_timestep=0, + attn_metadata=None, + forward_batch=batch): + self.transformer( + first_frame_latent.to(target_dtype), + prompt_embeds, + t_zero, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ) + if boundary_timestep is not None: + self.transformer_2( + first_frame_latent.to(target_dtype), + prompt_embeds, + t_zero, + kv_cache=kv_cache2, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ) + + start_index += 1 + block_sizes.pop(0) + latents[:, :, :1, :, :] = first_frame_latent + + # DMD loop in causal blocks + with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: + for current_num_frames in block_sizes: + current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] + # use BTCHW for DMD conversion routines + noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) + video_raw_latent_shape = noise_latents_btchw.shape + + for i, t_cur in enumerate(timesteps): + if boundary_timestep is not None and t_cur < boundary_timestep: + current_model = self.transformer_2 + else: + current_model = self.transformer + # Copy for pred conversion + noise_latents = noise_latents_btchw.clone() + latent_model_input = current_latents.to(target_dtype) + + if batch.image_latent is not None and independent_first_frame and start_index == 0: + latent_model_input = torch.cat([latent_model_input, batch.image_latent.to(target_dtype)], dim=2) + + # Prepare inputs + t_expand = t_cur.repeat(latent_model_input.shape[0]) + + # Attention metadata if needed + if (vsa_available and self.attn_backend == VideoSparseAttentionBackend): + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls() + attn_metadata = self.attn_metadata_builder.build( # type: ignore + current_timestep=i, # type: ignore + raw_latent_shape=(current_num_frames, h, w), # type: ignore + patch_size=fastvideo_args.pipeline_config.dit_config.patch_size, # type: ignore + VSA_sparsity=fastvideo_args.VSA_sparsity, # type: ignore + device=get_local_torch_device(), # type: ignore + ) # type: ignore + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), \ + set_forward_context(current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch): + # Run transformer; follow DMD stage pattern + t_expanded_noise = t_cur * torch.ones( + (latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long) + pred_noise_btchw = current_model( + latent_model_input, + prompt_embeds, + t_expanded_noise, + kv_cache=_get_kv_cache(t_cur), + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ).permute(0, 2, 1, 3, 4) + + # Convert pred noise to pred video with FM Euler scheduler utilities + if boundary_timestep is not None and t_cur >= boundary_timestep: + pred_video_btchw = pred_noise_to_x_bound( + pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + boundary_timestep=torch.ones_like(t_expand) * boundary_timestep, + scheduler=self.scheduler).unflatten(0, pred_noise_btchw.shape[:2]) + else: + pred_video_btchw = pred_noise_to_pred_video(pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) + + if i < len(timesteps) - 1: + next_timestep = timesteps[i + 1] * torch.ones( + [1], dtype=torch.long, device=pred_video_btchw.device) + noise = torch.randn(video_raw_latent_shape, + dtype=pred_video_btchw.dtype, + generator=(batch.generator[0] if isinstance(batch.generator, list) else + batch.generator)).to(self.device) + noise_btchw = noise + if boundary_timestep is not None and i < len(high_noise_timesteps) - 1: + noise_latents_btchw = self.scheduler.add_noise_high( + pred_video_btchw.flatten(0, 1), noise_btchw.flatten(0, 1), next_timestep, + torch.ones_like(next_timestep) * boundary_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + elif boundary_timestep is not None and i == len(high_noise_timesteps) - 1: + noise_latents_btchw = pred_video_btchw + else: + noise_latents_btchw = self.scheduler.add_noise(pred_video_btchw.flatten(0, 1), + noise_btchw.flatten(0, 1), + next_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4) + else: + current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4) + + if progress_bar is not None: + progress_bar.update() + + # Write back and advance + latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents + + # Re-run with context timestep to update KV cache using clean context + context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0) + t_context = torch.ones([latents.shape[0]], device=latents.device, dtype=torch.long) * int(context_noise) + context_bcthw = current_latents.to(target_dtype) + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), \ + set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=batch): + t_expanded_context = t_context.unsqueeze(1) + + if boundary_timestep is not None: + self.transformer_2( + context_bcthw, + prompt_embeds, + t_expanded_context, + kv_cache=kv_cache2, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ) + + self.transformer( + context_bcthw, + prompt_embeds, + t_expanded_context, + kv_cache=kv_cache1, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ) + + start_index += current_num_frames + + if boundary_timestep is not None: + num_frames_to_remove = self.num_frames_per_block - 1 + latents = latents[:, :, :-num_frames_to_remove, :, :] + + batch.latents = latents + return batch + + +class CausalDenoisingStage(WanCausalDenoisingBase): + """Causal block-by-block denoising with standard multi-step + flow matching (scheduler.step), not DMD few-step. + + Each block is fully denoised through all scheduler timesteps + before moving to the next block. After each block is denoised, + the KV cache is updated with clean context so subsequent blocks + can attend to prior clean frames. + """ + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32 and not fastvideo_args.disable_autocast) + + latent_seq_length = (batch.latents.shape[-1] * batch.latents.shape[-2]) + patch_ratio = (self.transformer.config.arch_config.patch_size[-1] * + self.transformer.config.arch_config.patch_size[-2]) + self.frame_seq_length = latent_seq_length // patch_ratio + + pos_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_attention_mask": batch.prompt_attention_mask, + }, + ) + + assert batch.latents is not None, ("latents must be provided") + latents = batch.latents # [B, C, T, H, W] + b, c, t, h, w = latents.shape + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + num_inference_steps = batch.num_inference_steps + + # Initialize caches + kv_cache = self._initialize_kv_cache( + batch_size=b, + dtype=target_dtype, + device=latents.device, + ) + crossattn_cache = self._initialize_crossattn_cache( + batch_size=b, + max_text_len=(fastvideo_args.pipeline_config.text_encoder_configs[0].arch_config.text_len), + dtype=target_dtype, + device=latents.device, + ) + + # Determine block sizes + if t % self.num_frames_per_block != 0: + raise ValueError("num_frames must be divisible by " + "num_frames_per_block") + num_blocks = t // self.num_frames_per_block + block_sizes = [self.num_frames_per_block] * num_blocks + start_index = 0 + pos_start_base = 0 + + context_noise = getattr( + fastvideo_args.pipeline_config, + "context_noise", + 0, + ) + + with self.progress_bar(total=len(block_sizes) * num_inference_steps) as progress_bar: + for current_num_frames in block_sizes: + current_latents = latents[ + :, + :, + start_index:start_index + current_num_frames, + :, + :, + ] + + # Reset scheduler per block so its multi-step + # history starts fresh. + self.scheduler.set_timesteps( + num_inference_steps, + device=latents.device, + ) + timesteps = self.scheduler.timesteps + + for i, t_cur in enumerate(timesteps): + latent_model_input = current_latents.to(target_dtype) + t_expanded = t_cur * torch.ones( + (b, 1), + device=latents.device, + dtype=torch.long, + ) + + with ( + torch.autocast( + device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled, + ), + set_forward_context( + current_timestep=i, + attn_metadata=None, + forward_batch=batch, + ), + ): + # Transformer returns [B, C, T, H, W], + # permute to [B, T, C, H, W] then flatten + # B*T for the scheduler. + noise_pred = self.transformer( + latent_model_input, + prompt_embeds, + t_expanded, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=((pos_start_base + start_index) * self.frame_seq_length), + start_frame=start_index, + **pos_cond_kwargs, + ) + + # Flatten [B,C,T,H,W] -> [B*T,C,H,W] for + # scheduler, then unflatten back. + nf = current_num_frames + noise_pred_flat = (noise_pred.permute(0, 2, 1, 3, 4).flatten(0, 1)) + latents_flat = (current_latents.permute(0, 2, 1, 3, 4).flatten(0, 1)) + updated_flat = self.scheduler.step( + noise_pred_flat, + t_cur, + latents_flat, + return_dict=False, + )[0] + current_latents = (updated_flat.unflatten(0, (b, nf)).permute(0, 2, 1, 3, 4)) + + if progress_bar is not None: + progress_bar.update() + + # Write denoised block back + latents[ + :, + :, + start_index:start_index + current_num_frames, + :, + :, + ] = current_latents + + # Update KV cache with clean context + t_context = torch.ones( + [b], + device=latents.device, + dtype=torch.long, + ) * int(context_noise) + context_bcthw = current_latents.to(target_dtype) + with ( + torch.autocast( + device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled, + ), + set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=batch, + ), + ): + self.transformer( + context_bcthw, + prompt_embeds, + t_context.unsqueeze(1), + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=((pos_start_base + start_index) * self.frame_seq_length), + start_frame=start_index, + **pos_cond_kwargs, + ) + + start_index += current_num_frames + + batch.latents = latents + return batch diff --git a/fastvideo/pipelines/basic/wan/stages/conditioning.py b/fastvideo/pipelines/basic/wan/stages/conditioning.py new file mode 100644 index 000000000..8f16db2fd --- /dev/null +++ b/fastvideo/pipelines/basic/wan/stages/conditioning.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prepare normalized first-frame latents before Wan's sampling stages.""" + +import torch + +from fastvideo.distributed import get_local_torch_device +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + + +class WanFirstFrameEncodingStage(PipelineStage): + + def __init__(self, vae, *, causal: bool = False): + super().__init__() + self.vae = vae + self.causal = causal + + def forward(self, batch, fastvideo_args): + batch.first_frame_latent = None + if batch.pil_image is None or not (self.causal or fastvideo_args.pipeline_config.ti2v_task): + return batch + assert self.vae is not None, "VAE is required for first-frame conditioning" + original_device = next(self.vae.parameters()).device + self.vae = self.vae.to(get_local_torch_device()) + latent = self.vae.encode(batch.pil_image).mean.float() + # Preserve dense TI2V's restore-before-normalization order and causal + # DMD's explicit CPU offload after normalization. + offload = getattr(fastvideo_args, "vae_cpu_offload", False) + if offload and not self.causal: + self.vae = self.vae.to(original_device) + shift = getattr(self.vae, "shift_factor", None) + if shift is not None: + latent -= shift.to(latent.device, latent.dtype) if isinstance(shift, torch.Tensor) else shift + scale = self.vae.scaling_factor + latent = latent * (scale.to(latent.device, latent.dtype) if isinstance(scale, torch.Tensor) else scale) + if offload and self.causal: + self.vae = self.vae.to("cpu") + batch.first_frame_latent = latent + return batch + + def verify_input(self, batch, fastvideo_args): + result = VerificationResult() + if batch.pil_image is not None and (self.causal or fastvideo_args.pipeline_config.ti2v_task): + result.add_check("pil_image", batch.pil_image, V.is_tensor) + return result + + def verify_output(self, batch, fastvideo_args): + result = VerificationResult() + if batch.pil_image is not None and (self.causal or fastvideo_args.pipeline_config.ti2v_task): + result.add_check("first_frame_latent", batch.first_frame_latent, [V.is_tensor, V.with_dims(5)]) + return result diff --git a/fastvideo/pipelines/basic/wan/stages/denoising.py b/fastvideo/pipelines/basic/wan/stages/denoising.py new file mode 100644 index 000000000..52c976056 --- /dev/null +++ b/fastvideo/pipelines/basic/wan/stages/denoising.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan conditioning and expert selection around the shared dense sampling loop.""" + +from dataclasses import dataclass + +import torch + +from fastvideo.distributed import get_local_torch_device +from fastvideo.logger import init_logger +from fastvideo.pipelines.stages.denoising import DenoisingStage, DenoisingState +from fastvideo.utils import masks_like + +logger = init_logger(__name__) + + +@dataclass +class WanDenoisingState(DenoisingState): + boundary_timestep: float | None = None + first_frame: torch.Tensor | None = None + first_frame_mask: torch.Tensor | None = None + timestep_sequence_length: int | None = None + ti2v_task: bool = False + lucy_edit_task: bool = False + + +class WanDenoisingStage(DenoisingStage): + """Dense Wan sampling, including Wan2.2 experts, TI2V, and Lucy Edit. + + All conditioning state is per request. The stage receives normalized + first-frame latents; it does not load or execute a VAE. + """ + + def prepare_denoising(self, batch, fastvideo_args, target_dtype) -> WanDenoisingState: + config = fastvideo_args.pipeline_config + latents = batch.latents + assert latents is not None + assert latents.shape[0] == 1, "only support batch size 1" + boundary_ratio = config.dit_config.boundary_ratio + if batch.boundary_ratio is not None: + logger.info("Overriding boundary ratio from %s to %s", boundary_ratio, batch.boundary_ratio) + boundary_ratio = batch.boundary_ratio + state = WanDenoisingState( + latents=latents, + boundary_timestep=boundary_ratio * + self.scheduler.num_train_timesteps if boundary_ratio is not None else None, + ti2v_task=config.ti2v_task, + lucy_edit_task=config.lucy_edit_task, + ) + if config.ti2v_task and batch.pil_image is not None: + assert batch.image_latent is None, "TI2V task should not have image latents" + assert batch.first_frame_latent is not None, "Run WanFirstFrameEncodingStage before TI2V denoising" + state.first_frame = batch.first_frame_latent + latent_model_input = latents.to(target_dtype).squeeze(0) + _, masks = masks_like([latent_model_input], zero=True) + state.first_frame_mask = masks[0] + state.latents = ((1. - masks[0]) * state.first_frame + masks[0] * latent_model_input).to( + get_local_torch_device()) + temporal_scale = config.vae_config.arch_config.scale_factor_temporal + spatial_scale = config.vae_config.arch_config.scale_factor_spatial + patch_size = config.dit_config.arch_config.patch_size + state.timestep_sequence_length = (((batch.num_frames - 1) // temporal_scale + 1) * + (batch.height // spatial_scale) * (batch.width // spatial_scale) // + (patch_size[1] * patch_size[2])) + if state.lucy_edit_task: + patch_size = config.dit_config.arch_config.patch_size + assert patch_size[0] == 1, "Lucy Edit timestep expansion assumes temporal patch size 1" + state.timestep_sequence_length = (state.latents.shape[2] * (state.latents.shape[3] // patch_size[1]) * + (state.latents.shape[4] // patch_size[2])) + if batch.video_latent is not None: + state.video_padding = torch.zeros_like(state.latents) + return state + + def select_model(self, timestep, batch, fastvideo_args, state): + if state.boundary_timestep is None or timestep >= state.boundary_timestep: + model, inactive, guidance = self.transformer, self.transformer_2, batch.guidance_scale + else: + model, inactive, guidance = self.transformer_2, self.transformer, batch.guidance_scale_2 + self.activate_transformer(model, inactive, fastvideo_args) + return model, guidance + + def prepare_model_input(self, latents, batch, target_dtype, state): + if batch.video_latent is not None and state.lucy_edit_task: + return torch.cat([latents.to(target_dtype), batch.video_latent], dim=1).to(target_dtype) + if batch.image_latent is not None: + assert not state.ti2v_task, "image latents should not be provided for TI2V task" + return super().prepare_model_input(latents, batch, target_dtype, state) + + def prepare_timestep(self, timestep, latent_model_input, batch, state): + if state.lucy_edit_task: + return timestep.repeat(latent_model_input.shape[0], state.timestep_sequence_length) + if state.first_frame_mask is not None: + step = torch.stack([timestep]).to(get_local_torch_device()) + expanded = (state.first_frame_mask[0][:, ::2, ::2] * step).flatten() + expanded = torch.cat([ + expanded, + expanded.new_ones(state.timestep_sequence_length - expanded.size(0)) * step, + ]) + return expanded.unsqueeze(0).repeat(latent_model_input.shape[0], 1) + return super().prepare_timestep(timestep, latent_model_input, batch, state) + + def finish_step(self, latents, state): + if state.first_frame_mask is not None: + latents = latents.squeeze(0) + latents = (1. - state.first_frame_mask) * state.first_frame + state.first_frame_mask * latents + return latents diff --git a/fastvideo/pipelines/basic/wan/stages/dmd.py b/fastvideo/pipelines/basic/wan/stages/dmd.py new file mode 100644 index 000000000..ee19ac4e1 --- /dev/null +++ b/fastvideo/pipelines/basic/wan/stages/dmd.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan's non-causal DMD sampler. The pipeline supplies its training-noise scheduler.""" + +import torch + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.utils import dict_to_3d_list + +try: + from fastvideo.attention.backends.video_sparse_attn import VideoSparseAttentionBackend + vsa_available = True +except ImportError: + vsa_available = False + + +class DmdDenoisingStage(DenoisingStage): + """ + Denoising stage for DMD. + """ + + def __init__(self, transformer, scheduler) -> None: + super().__init__(transformer, scheduler) + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """ + Run the denoising loop. + + Args: + batch: The current batch information. + fastvideo_args: The inference arguments. + + Returns: + The batch with denoised latents. + """ + # Setup precision and autocast settings + # TODO(will): make the precision configurable for inference + # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast + + # Get timesteps and calculate warmup steps + timesteps = batch.timesteps + + # TODO(will): remove this once we add input/output validation for stages + if timesteps is None: + raise ValueError("Timesteps must be provided") + num_inference_steps = batch.num_inference_steps + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + + # Prepare image latents and embeddings for I2V generation + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert torch.isnan(image_embeds[0]).sum() == 0 + image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds] + + image_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_image": image_embeds, + "mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24) + }, + ) + + pos_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_2": batch.clip_embedding_pos, + "encoder_attention_mask": batch.prompt_attention_mask, + }, + ) + + # Get latents and embeddings + assert batch.latents is not None, "latents must be provided" + latents = batch.latents + + video_raw_latent_shape = latents.shape + prompt_embeds = batch.prompt_embeds + assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan" + timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, + dtype=torch.long, + device=get_local_torch_device()) + + # Run denoising loop + with self.progress_bar(total=len(timesteps)) as progress_bar: + for i, t in enumerate(timesteps): + # Skip if interrupted + if hasattr(self, 'interrupt') and self.interrupt: + continue + # Expand latents for I2V + noise_latents = latents.clone() + latent_model_input = latents.to(target_dtype) + + if batch.image_latent is not None: + latent_model_input = torch.cat( + [latent_model_input, batch.image_latent.permute(0, 2, 1, 3, 4)], dim=2).to(target_dtype) + assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan" + + # Prepare inputs for transformer + t_expand = t.repeat(latent_model_input.shape[0]) + guidance_expand = (torch.tensor( + [fastvideo_args.pipeline_config.embedded_cfg_scale] * latent_model_input.shape[0], + dtype=torch.float32, + device=get_local_torch_device(), + ).to(target_dtype) * 1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale is not None else None) + + # Predict noise residual + with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled): + if (vsa_available and self.attn_backend == VideoSparseAttentionBackend): + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() + + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls() + # TODO(will): clean this up + attn_metadata = self.attn_metadata_builder.build( # type: ignore + current_timestep=i, # type: ignore + raw_latent_shape=batch.raw_latent_shape[2:5], # type: ignore + patch_size=fastvideo_args.pipeline_config. # type: ignore + dit_config.patch_size, # type: ignore + VSA_sparsity=fastvideo_args.VSA_sparsity, # type: ignore + device=get_local_torch_device(), # type: ignore + ) # type: ignore + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + + batch.is_cfg_negative = False + with set_forward_context( + current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch, + # fastvideo_args=fastvideo_args + ): + # Run transformer + pred_noise = self.transformer( + latent_model_input.permute(0, 2, 1, 3, 4), + prompt_embeds, + t_expand, + guidance=guidance_expand, + **image_kwargs, + **pos_cond_kwargs, + ).permute(0, 2, 1, 3, 4) + + pred_video = pred_noise_to_pred_video(pred_noise=pred_noise.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + scheduler=self.scheduler).unflatten(0, pred_noise.shape[:2]) + + if i < len(timesteps) - 1: + next_timestep = timesteps[i + 1] * torch.ones([1], dtype=torch.long, device=pred_video.device) + noise = torch.randn(video_raw_latent_shape, + dtype=pred_video.dtype, + generator=batch.generator[0]).to(self.device) + latents = self.scheduler.add_noise(pred_video.flatten(0, 1), noise.flatten(0, 1), + next_timestep).unflatten(0, pred_video.shape[:2]) + else: + latents = pred_video + + # Update progress bar + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and + (i + 1) % self.scheduler.order == 0 and progress_bar is not None): + progress_bar.update() + + # Gather results if using sequence parallelism + latents = latents.permute(0, 2, 1, 3, 4) + # Update batch with final latents + batch.latents = latents + + return batch diff --git a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py index f05212127..a9f7df3bc 100644 --- a/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py @@ -5,13 +5,15 @@ Wan causal DMD pipeline implementation. This module wires the causal DMD denoising stage into the modular pipeline. """ +from fastvideo.pipelines.basic.wan.stages.conditioning import WanFirstFrameEncodingStage +from fastvideo.pipelines.basic.wan.stages.causal_denoising import CausalDMDDenosingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline # isort: off -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, CausalDMDDenosingStage, InputValidationStage, - LatentPreparationStage, TextEncodingStage) +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, InputValidationStage, LatentPreparationStage, + TextEncodingStage) # isort: on logger = init_logger(__name__) @@ -38,11 +40,13 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): stage=LatentPreparationStage(scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer", None))) + self.add_stage(stage_name="first_frame_encoding_stage", + stage=WanFirstFrameEncodingStage(vae=self.get_module("vae"), causal=True)) + self.add_stage(stage_name="denoising_stage", stage=CausalDMDDenosingStage(transformer=self.get_module("transformer"), transformer_2=self.get_module("transformer_2", None), - scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"))) + scheduler=self.get_module("scheduler"))) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) diff --git a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py index 2c6c4746a..080d8429c 100644 --- a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py @@ -5,6 +5,7 @@ Block-by-block causal inference with KV caching, using the full scheduler timestep schedule (40-50 steps) rather than DMD few-step. """ +from fastvideo.pipelines.basic.wan.stages.causal_denoising import CausalDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (FlowUniPCMultistepScheduler) @@ -12,7 +13,6 @@ from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline # isort: off from fastvideo.pipelines.stages import ( - CausalDenoisingStage, ConditioningStage, DecodingStage, InputValidationStage, diff --git a/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py index c7985c1fb..807a6d146 100644 --- a/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py @@ -6,14 +6,16 @@ This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ +from fastvideo.pipelines.basic.wan.stages.dmd import DmdDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.models.wan.definition import DMD_TRAINING_NOISE_SHIFT from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (FlowMatchEulerDiscreteScheduler) from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline # isort: off -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, DmdDenoisingStage, InputValidationStage, - LatentPreparationStage, TextEncodingStage, TimestepPreparationStage) +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, InputValidationStage, LatentPreparationStage, + TextEncodingStage, TimestepPreparationStage) # isort: on logger = init_logger(__name__) @@ -51,9 +53,12 @@ class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase): transformer=self.get_module("transformer", None), use_btchw_layout=True)) + # DMD needs the complete training-noise table, separate from the + # inference scheduler mutated by TimestepPreparationStage. self.add_stage(stage_name="denoising_stage", - stage=DmdDenoisingStage(transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"))) + stage=DmdDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=FlowMatchEulerDiscreteScheduler(shift=DMD_TRAINING_NOISE_SHIFT))) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) diff --git a/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py index 0a4fec4e3..b916dee80 100644 --- a/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py @@ -6,15 +6,17 @@ This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ +from fastvideo.pipelines.basic.wan.stages.dmd import DmdDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.models.wan.definition import DMD_TRAINING_NOISE_SHIFT from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline # isort: off -from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage, DecodingStage, DmdDenoisingStage, - ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage, - TextEncodingStage, TimestepPreparationStage) +from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage, DecodingStage, ImageVAEEncodingStage, + InputValidationStage, LatentPreparationStage, TextEncodingStage, + TimestepPreparationStage) # isort: on from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (FlowMatchEulerDiscreteScheduler) @@ -61,9 +63,12 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase): self.add_stage(stage_name="image_latent_preparation_stage", stage=ImageVAEEncodingStage(vae=self.get_module("vae"))) + # DMD needs the complete training-noise table, separate from the + # inference scheduler mutated by TimestepPreparationStage. self.add_stage(stage_name="denoising_stage", - stage=DmdDenoisingStage(transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"))) + stage=DmdDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=FlowMatchEulerDiscreteScheduler(shift=DMD_TRAINING_NOISE_SHIFT))) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) diff --git a/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py b/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py index d7728b3fc..a27cb7e74 100644 --- a/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py @@ -6,15 +6,16 @@ This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ +from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline # isort: off -from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage, - ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage, - TextEncodingStage, TimestepPreparationStage) +from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage, DecodingStage, ImageVAEEncodingStage, + InputValidationStage, LatentPreparationStage, TextEncodingStage, + TimestepPreparationStage) # isort: on from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (FlowUniPCMultistepScheduler) @@ -62,9 +63,9 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase): stage=ImageVAEEncodingStage(vae=self.get_module("vae"))) self.add_stage(stage_name="denoising_stage", - stage=DenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2"), - scheduler=self.get_module("scheduler"))) + stage=WanDenoisingStage(transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2"), + scheduler=self.get_module("scheduler"))) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) diff --git a/fastvideo/pipelines/basic/wan/wan_pipeline.py b/fastvideo/pipelines/basic/wan/wan_pipeline.py index 78aedba94..67a97fed5 100644 --- a/fastvideo/pipelines/basic/wan/wan_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_pipeline.py @@ -6,12 +6,14 @@ This module contains an implementation of the Wan video diffusion pipeline using the modular pipeline architecture. """ +from fastvideo.pipelines.basic.wan.stages.conditioning import WanFirstFrameEncodingStage +from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (FlowUniPCMultistepScheduler) from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage, - LatentPreparationStage, TextEncodingStage, TimestepPreparationStage) +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, InputValidationStage, LatentPreparationStage, + TextEncodingStage, TimestepPreparationStage) logger = init_logger(__name__) @@ -47,12 +49,14 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase): stage=LatentPreparationStage(scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer", None))) + self.add_stage(stage_name="first_frame_encoding_stage", + stage=WanFirstFrameEncodingStage(vae=self.get_module("vae"))) + self.add_stage(stage_name="denoising_stage", - stage=DenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2", None), - scheduler=self.get_module("scheduler"), - vae=self.get_module("vae"), - pipeline=self)) + stage=WanDenoisingStage(transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + pipeline=self)) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) diff --git a/fastvideo/pipelines/basic/wan/wan_v2v_pipeline.py b/fastvideo/pipelines/basic/wan/wan_v2v_pipeline.py index e8c0b1af3..95757cf70 100644 --- a/fastvideo/pipelines/basic/wan/wan_v2v_pipeline.py +++ b/fastvideo/pipelines/basic/wan/wan_v2v_pipeline.py @@ -6,15 +6,16 @@ This module contains an implementation of the Wan video-to-video diffusion pipel using the modular pipeline architecture. """ +from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline # isort: off -from fastvideo.pipelines.stages import (RefImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage, - VideoVAEEncodingStage, InputValidationStage, LatentPreparationStage, - TextEncodingStage, TimestepPreparationStage) +from fastvideo.pipelines.stages import (RefImageEncodingStage, ConditioningStage, DecodingStage, VideoVAEEncodingStage, + InputValidationStage, LatentPreparationStage, TextEncodingStage, + TimestepPreparationStage) # isort: on from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (FlowUniPCMultistepScheduler) @@ -62,9 +63,9 @@ class WanVideoToVideoPipeline(LoRAPipeline, ComposedPipelineBase): stage=VideoVAEEncodingStage(vae=self.get_module("vae"))) self.add_stage(stage_name="denoising_stage", - stage=DenoisingStage(transformer=self.get_module("transformer"), - transformer_2=self.get_module("transformer_2"), - scheduler=self.get_module("scheduler"))) + stage=WanDenoisingStage(transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2"), + scheduler=self.get_module("scheduler"))) self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 3a9ca1467..2a7d2533e 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -132,6 +132,8 @@ class ForwardBatch: raw_latent_shape: tuple[int, ...] | None = None noise_pred: torch.Tensor | None = None image_latent: torch.Tensor | None = None + # Normalized clean first frame for Wan TI2V and causal-DMD conditioning. + first_frame_latent: torch.Tensor | None = None # Action control inputs (Matrix-Game) mouse_cond: torch.Tensor | None = None # Shape: (B, T, 2) diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index ae4b701f6..639b1556e 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -7,12 +7,11 @@ complete diffusion pipelines. """ from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.causal_denoising import (CausalDMDDenosingStage, CausalDenoisingStage) from fastvideo.pipelines.stages.conditioning import ConditioningStage from fastvideo.pipelines.stages.decoding import DecodingStage from fastvideo.pipelines.stages.denoising import (Cosmos25AutoDenoisingStage, Cosmos25DenoisingStage, Cosmos25V2WDenoisingStage, Cosmos25T2WDenoisingStage, - CosmosDenoisingStage, DenoisingStage, DmdDenoisingStage) + CosmosDenoisingStage, DenoisingStage) from fastvideo.pipelines.stages.sr_denoising import SRDenoisingStage from fastvideo.pipelines.stages.encoding import EncodingStage from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, MatrixGame2ImageEncodingStage, @@ -102,3 +101,15 @@ __all__ = [ "LongCatKVCacheInitStage", "LongCatVCDenoisingStage", ] + + +def __getattr__(name): + # Family stages depend on shared stages. Resolve compatibility exports only + # when requested, so importing the canonical Wan modules cannot form a cycle. + if name == "DmdDenoisingStage": + from fastvideo.pipelines.basic.wan.stages.dmd import DmdDenoisingStage + return DmdDenoisingStage + if name in {"CausalDMDDenosingStage", "CausalDenoisingStage"}: + from fastvideo.pipelines.basic.wan.stages import causal_denoising + return getattr(causal_denoising, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index f6951f1d8..caabf2542 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -1,603 +1,10 @@ -import torch # type: ignore +# SPDX-License-Identifier: Apache-2.0 +"""Compatibility imports for Wan's family-owned causal sampling stages.""" -from fastvideo.distributed import get_local_torch_device -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.forward_context import set_forward_context -from fastvideo.logger import init_logger -from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.pipelines.stages.denoising import DenoisingStage -from fastvideo.pipelines.stages.validators import StageValidators as V -from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.pipelines.basic.wan.stages.causal_denoising import ( + CausalDMDDenosingStage, + CausalDenoisingStage, + _get_transformer_attr, +) -try: - from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionBackend) - vsa_available = True -except ImportError: - vsa_available = False - VideoSparseAttentionBackend = None # type: ignore - -logger = init_logger(__name__) - - -def _get_transformer_attr(transformer, name: str, default): - # Currently used only for causal KV-cache runtime settings: - # `local_attn_size` and `sink_size`. - # - # Causal Wan/MatrixGame expose these directly on the transformer. Some - # wrapper-style DiTs keep the executable module under transformer.model. Do - # not fall back to transformer.config for these two values: config is only - # the construction source and does not control them at runtime. - value = getattr(transformer, name, None) - if value is not None: - return value - - inner_model = getattr(transformer, "model", None) - if inner_model is not None: - value = getattr(inner_model, name, None) - if value is not None: - return value - - return default - - -class CausalDMDDenosingStage(DenoisingStage): - """ - Denoising stage for causal diffusion. - """ - - def __init__(self, transformer, scheduler, transformer_2=None, vae=None) -> None: - super().__init__(transformer, scheduler, transformer_2) - # KV and cross-attention cache state (initialized on first forward) - self.transformer = transformer - self.transformer_2 = transformer_2 - self.vae = vae - # Model-dependent constants (aligned with causal_inference.py assumptions) - self.num_transformer_blocks = len(self.transformer.blocks) - self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block - self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames - - self.local_attn_size = _get_transformer_attr(self.transformer, "local_attn_size", -1) - self.sink_size = _get_transformer_attr(self.transformer, "sink_size", 0) - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - - latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] - patch_ratio = self.transformer.config.arch_config.patch_size[ - -1] * self.transformer.config.arch_config.patch_size[-2] - self.frame_seq_length = latent_seq_length // patch_ratio - # TODO(will): make this a parameter once we add i2v support - independent_first_frame = self.transformer.independent_first_frame if hasattr( - self.transformer, 'independent_first_frame') else False - # Timesteps for DMD - timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu() - if fastvideo_args.pipeline_config.warp_denoising_step: - scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) - timesteps = scheduler_timesteps[1000 - timesteps] - timesteps = timesteps.to(get_local_torch_device()) - - if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None: - boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio * self.scheduler.num_train_timesteps - high_noise_timesteps = timesteps[timesteps >= boundary_timestep] - else: - boundary_timestep = None - high_noise_timesteps = None - - # Image kwargs (kept empty unless caller provides compatible args) - image_kwargs: dict = {} - - pos_cond_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - # "encoder_hidden_states_2": batch.clip_embedding_pos, - "encoder_attention_mask": batch.prompt_attention_mask, - }, - ) - - # Latents and prompts - assert batch.latents is not None, "latents must be provided" - latents = batch.latents # [B, C, T, H, W] - b, c, t, h, w = latents.shape - prompt_embeds = batch.prompt_embeds - assert torch.isnan(prompt_embeds[0]).sum() == 0 - - # Initialize or reset caches - kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) - kv_cache2 = None - if boundary_timestep is not None: - # Initialize the low noise kv cache - kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], - dtype=target_dtype, - device=latents.device) - - def _get_kv_cache(timestep: float) -> list[dict]: - if boundary_timestep is not None: - if timestep >= boundary_timestep: - return kv_cache1 - else: - assert kv_cache2 is not None, "kv_cache2 is not initialized" - return kv_cache2 - return kv_cache1 - - crossattn_cache = self._initialize_crossattn_cache( - batch_size=latents.shape[0], - max_text_len=fastvideo_args.pipeline_config.text_encoder_configs[0].arch_config.text_len, - dtype=target_dtype, - device=latents.device) - - pos_start_base = 0 - - # Determine block sizes - if t % self.num_frames_per_block != 0: - raise ValueError("num_frames must be divisible by num_frames_per_block for causal DMD denoising") - num_blocks = t // self.num_frames_per_block - block_sizes = [self.num_frames_per_block] * num_blocks - start_index = 0 - - # For now hardcode the first block to be 1 frame assuming the model is Wan2.2-MoE - if boundary_timestep is not None: - block_sizes[0] = 1 - - first_frame_latent = None - if batch.pil_image is not None: - # Causal video gen directly replaces the first frame of the latent with - # the image latent instead of appending along the channel dim - assert self.vae is not None, "VAE is not provided for causal video gen task" - self.vae = self.vae.to(get_local_torch_device()) - first_frame_latent = self.vae.encode(batch.pil_image).mean.float() - if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): - if isinstance(self.vae.shift_factor, torch.Tensor): - first_frame_latent -= self.vae.shift_factor.to(first_frame_latent.device, first_frame_latent.dtype) - else: - first_frame_latent -= self.vae.shift_factor - - if isinstance(self.vae.scaling_factor, torch.Tensor): - first_frame_latent = first_frame_latent * self.vae.scaling_factor.to( - first_frame_latent.device, first_frame_latent.dtype) - else: - first_frame_latent = first_frame_latent * self.vae.scaling_factor - - if fastvideo_args.vae_cpu_offload: - self.vae = self.vae.to("cpu") - - # Fill the low noise and high noise kv cache with first_frame_latent and timestep 0 - t_zero = torch.zeros([latents.shape[0], 1], device=latents.device, dtype=torch.long) - with torch.autocast(device_type="cuda", - dtype=target_dtype, - enabled=autocast_enabled), \ - set_forward_context(current_timestep=0, - attn_metadata=None, - forward_batch=batch): - self.transformer( - first_frame_latent.to(target_dtype), - prompt_embeds, - t_zero, - kv_cache=kv_cache1, - crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ) - if boundary_timestep is not None: - self.transformer_2( - first_frame_latent.to(target_dtype), - prompt_embeds, - t_zero, - kv_cache=kv_cache2, - crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ) - - start_index += 1 - block_sizes.pop(0) - latents[:, :, :1, :, :] = first_frame_latent - - # DMD loop in causal blocks - with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: - for current_num_frames in block_sizes: - current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] - # use BTCHW for DMD conversion routines - noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) - video_raw_latent_shape = noise_latents_btchw.shape - - for i, t_cur in enumerate(timesteps): - if boundary_timestep is not None and t_cur < boundary_timestep: - current_model = self.transformer_2 - else: - current_model = self.transformer - # Copy for pred conversion - noise_latents = noise_latents_btchw.clone() - latent_model_input = current_latents.to(target_dtype) - - if batch.image_latent is not None and independent_first_frame and start_index == 0: - latent_model_input = torch.cat([latent_model_input, batch.image_latent.to(target_dtype)], dim=2) - - # Prepare inputs - t_expand = t_cur.repeat(latent_model_input.shape[0]) - - # Attention metadata if needed - if (vsa_available and self.attn_backend == VideoSparseAttentionBackend): - self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() - if self.attn_metadata_builder_cls is not None: - self.attn_metadata_builder = self.attn_metadata_builder_cls() - attn_metadata = self.attn_metadata_builder.build( # type: ignore - current_timestep=i, # type: ignore - raw_latent_shape=(current_num_frames, h, w), # type: ignore - patch_size=fastvideo_args.pipeline_config.dit_config.patch_size, # type: ignore - VSA_sparsity=fastvideo_args.VSA_sparsity, # type: ignore - device=get_local_torch_device(), # type: ignore - ) # type: ignore - assert attn_metadata is not None, "attn_metadata cannot be None" - else: - attn_metadata = None - else: - attn_metadata = None - - with torch.autocast(device_type="cuda", - dtype=target_dtype, - enabled=autocast_enabled), \ - set_forward_context(current_timestep=i, - attn_metadata=attn_metadata, - forward_batch=batch): - # Run transformer; follow DMD stage pattern - t_expanded_noise = t_cur * torch.ones( - (latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long) - pred_noise_btchw = current_model( - latent_model_input, - prompt_embeds, - t_expanded_noise, - kv_cache=_get_kv_cache(t_cur), - crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ).permute(0, 2, 1, 3, 4) - - # Convert pred noise to pred video with FM Euler scheduler utilities - if boundary_timestep is not None and t_cur >= boundary_timestep: - pred_video_btchw = pred_noise_to_x_bound( - pred_noise=pred_noise_btchw.flatten(0, 1), - noise_input_latent=noise_latents.flatten(0, 1), - timestep=t_expand, - boundary_timestep=torch.ones_like(t_expand) * boundary_timestep, - scheduler=self.scheduler).unflatten(0, pred_noise_btchw.shape[:2]) - else: - pred_video_btchw = pred_noise_to_pred_video(pred_noise=pred_noise_btchw.flatten(0, 1), - noise_input_latent=noise_latents.flatten(0, 1), - timestep=t_expand, - scheduler=self.scheduler).unflatten( - 0, pred_noise_btchw.shape[:2]) - - if i < len(timesteps) - 1: - next_timestep = timesteps[i + 1] * torch.ones( - [1], dtype=torch.long, device=pred_video_btchw.device) - noise = torch.randn(video_raw_latent_shape, - dtype=pred_video_btchw.dtype, - generator=(batch.generator[0] if isinstance(batch.generator, list) else - batch.generator)).to(self.device) - noise_btchw = noise - if boundary_timestep is not None and i < len(high_noise_timesteps) - 1: - noise_latents_btchw = self.scheduler.add_noise_high( - pred_video_btchw.flatten(0, 1), noise_btchw.flatten(0, 1), next_timestep, - torch.ones_like(next_timestep) * boundary_timestep).unflatten( - 0, pred_video_btchw.shape[:2]) - elif boundary_timestep is not None and i == len(high_noise_timesteps) - 1: - noise_latents_btchw = pred_video_btchw - else: - noise_latents_btchw = self.scheduler.add_noise(pred_video_btchw.flatten(0, 1), - noise_btchw.flatten(0, 1), - next_timestep).unflatten( - 0, pred_video_btchw.shape[:2]) - current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4) - else: - current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4) - - if progress_bar is not None: - progress_bar.update() - - # Write back and advance - latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents - - # Re-run with context timestep to update KV cache using clean context - context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0) - t_context = torch.ones([latents.shape[0]], device=latents.device, dtype=torch.long) * int(context_noise) - context_bcthw = current_latents.to(target_dtype) - with torch.autocast(device_type="cuda", - dtype=target_dtype, - enabled=autocast_enabled), \ - set_forward_context(current_timestep=0, - attn_metadata=attn_metadata, - forward_batch=batch): - t_expanded_context = t_context.unsqueeze(1) - - if boundary_timestep is not None: - self.transformer_2( - context_bcthw, - prompt_embeds, - t_expanded_context, - kv_cache=kv_cache2, - crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ) - - self.transformer( - context_bcthw, - prompt_embeds, - t_expanded_context, - kv_cache=kv_cache1, - crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ) - - start_index += current_num_frames - - if boundary_timestep is not None: - num_frames_to_remove = self.num_frames_per_block - 1 - latents = latents[:, :, :-num_frames_to_remove, :, :] - - batch.latents = latents - return batch - - def _initialize_kv_cache(self, batch_size, dtype, device) -> list[dict]: - """ - Initialize a Per-GPU KV cache aligned with the Wan model assumptions. - """ - kv_cache1 = [] - num_attention_heads = self.transformer.num_attention_heads - attention_head_dim = self.transformer.attention_head_dim - if self.local_attn_size != -1: - kv_cache_size = self.local_attn_size * self.frame_seq_length - else: - kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames - - for _ in range(self.num_transformer_blocks): - kv_cache1.append({ - "k": - torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], - dtype=dtype, - device=device), - "v": - torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], - dtype=dtype, - device=device), - "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), - "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), - }) - - return kv_cache1 - - def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype, device) -> list[dict]: - """ - Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions. - """ - crossattn_cache = [] - num_attention_heads = self.transformer.num_attention_heads - attention_head_dim = self.transformer.attention_head_dim - for _ in range(self.num_transformer_blocks): - crossattn_cache.append({ - "k": - torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], - dtype=dtype, - device=device), - "v": - torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], - dtype=dtype, - device=device), - "is_init": - False, - }) - return crossattn_cache - - def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: - """Verify denoising stage inputs.""" - result = VerificationResult() - result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) - result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) - result.add_check("image_embeds", batch.image_embeds, V.is_list) - result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)) - result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) - result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) - result.add_check("eta", batch.eta, V.non_negative_float) - result.add_check("generator", batch.generator, V.generator_or_list_generators) - result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value) - result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds, - lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)) - return result - - -class CausalDenoisingStage(CausalDMDDenosingStage): - """Causal block-by-block denoising with standard multi-step - flow matching (scheduler.step), not DMD few-step. - - Each block is fully denoised through all scheduler timesteps - before moving to the next block. After each block is denoised, - the KV cache is updated with clean context so subsequent blocks - can attend to prior clean frames. - """ - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32 and not fastvideo_args.disable_autocast) - - latent_seq_length = (batch.latents.shape[-1] * batch.latents.shape[-2]) - patch_ratio = (self.transformer.config.arch_config.patch_size[-1] * - self.transformer.config.arch_config.patch_size[-2]) - self.frame_seq_length = latent_seq_length // patch_ratio - - pos_cond_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "encoder_attention_mask": batch.prompt_attention_mask, - }, - ) - - assert batch.latents is not None, ("latents must be provided") - latents = batch.latents # [B, C, T, H, W] - b, c, t, h, w = latents.shape - prompt_embeds = batch.prompt_embeds - assert torch.isnan(prompt_embeds[0]).sum() == 0 - - num_inference_steps = batch.num_inference_steps - - # Initialize caches - kv_cache = self._initialize_kv_cache( - batch_size=b, - dtype=target_dtype, - device=latents.device, - ) - crossattn_cache = self._initialize_crossattn_cache( - batch_size=b, - max_text_len=(fastvideo_args.pipeline_config.text_encoder_configs[0].arch_config.text_len), - dtype=target_dtype, - device=latents.device, - ) - - # Determine block sizes - if t % self.num_frames_per_block != 0: - raise ValueError("num_frames must be divisible by " - "num_frames_per_block") - num_blocks = t // self.num_frames_per_block - block_sizes = [self.num_frames_per_block] * num_blocks - start_index = 0 - pos_start_base = 0 - - context_noise = getattr( - fastvideo_args.pipeline_config, - "context_noise", - 0, - ) - - with self.progress_bar(total=len(block_sizes) * num_inference_steps) as progress_bar: - for current_num_frames in block_sizes: - current_latents = latents[ - :, - :, - start_index:start_index + current_num_frames, - :, - :, - ] - - # Reset scheduler per block so its multi-step - # history starts fresh. - self.scheduler.set_timesteps( - num_inference_steps, - device=latents.device, - ) - timesteps = self.scheduler.timesteps - - for i, t_cur in enumerate(timesteps): - latent_model_input = current_latents.to(target_dtype) - t_expanded = t_cur * torch.ones( - (b, 1), - device=latents.device, - dtype=torch.long, - ) - - with ( - torch.autocast( - device_type="cuda", - dtype=target_dtype, - enabled=autocast_enabled, - ), - set_forward_context( - current_timestep=i, - attn_metadata=None, - forward_batch=batch, - ), - ): - # Transformer returns [B, C, T, H, W], - # permute to [B, T, C, H, W] then flatten - # B*T for the scheduler. - noise_pred = self.transformer( - latent_model_input, - prompt_embeds, - t_expanded, - kv_cache=kv_cache, - crossattn_cache=crossattn_cache, - current_start=((pos_start_base + start_index) * self.frame_seq_length), - start_frame=start_index, - **pos_cond_kwargs, - ) - - # Flatten [B,C,T,H,W] -> [B*T,C,H,W] for - # scheduler, then unflatten back. - nf = current_num_frames - noise_pred_flat = (noise_pred.permute(0, 2, 1, 3, 4).flatten(0, 1)) - latents_flat = (current_latents.permute(0, 2, 1, 3, 4).flatten(0, 1)) - updated_flat = self.scheduler.step( - noise_pred_flat, - t_cur, - latents_flat, - return_dict=False, - )[0] - current_latents = (updated_flat.unflatten(0, (b, nf)).permute(0, 2, 1, 3, 4)) - - if progress_bar is not None: - progress_bar.update() - - # Write denoised block back - latents[ - :, - :, - start_index:start_index + current_num_frames, - :, - :, - ] = current_latents - - # Update KV cache with clean context - t_context = torch.ones( - [b], - device=latents.device, - dtype=torch.long, - ) * int(context_noise) - context_bcthw = current_latents.to(target_dtype) - with ( - torch.autocast( - device_type="cuda", - dtype=target_dtype, - enabled=autocast_enabled, - ), - set_forward_context( - current_timestep=0, - attn_metadata=None, - forward_batch=batch, - ), - ): - self.transformer( - context_bcthw, - prompt_embeds, - t_context.unsqueeze(1), - kv_cache=kv_cache, - crossattn_cache=crossattn_cache, - current_start=((pos_start_base + start_index) * self.frame_seq_length), - start_frame=start_index, - **pos_cond_kwargs, - ) - - start_index += current_num_frames - - batch.latents = latents - return batch +__all__ = ["CausalDMDDenosingStage", "CausalDenoisingStage", "_get_transformer_attr"] diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 771bc3c32..79dbca32c 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -7,6 +7,7 @@ import inspect import os import weakref from collections.abc import Iterable +from dataclasses import dataclass from typing import Any import torch @@ -20,14 +21,12 @@ from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.models.loader.component_loader import TransformerLoader -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (FlowMatchEulerDiscreteScheduler) -from fastvideo.models.utils import pred_noise_to_pred_video from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.platforms import AttentionBackendEnum -from fastvideo.utils import dict_to_3d_list, masks_like +from fastvideo.utils import dict_to_3d_list try: from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend @@ -45,6 +44,13 @@ except ImportError: logger = init_logger(__name__) +@dataclass +class DenoisingState: + """Request-local inputs reused by the dense sampling loop.""" + latents: torch.Tensor + video_padding: torch.Tensor | None = None + + class DenoisingStage(PipelineStage): """ Stage for running the denoising loop in diffusion pipelines. @@ -247,55 +253,12 @@ class DenoisingStage(PipelineStage): neg_prompt_embeds = batch.negative_prompt_embeds assert not torch.isnan(neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan" - # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert - boundary_ratio = fastvideo_args.pipeline_config.dit_config.boundary_ratio - if batch.boundary_ratio is not None: - logger.info("Overriding boundary ratio from %s to %s", boundary_ratio, batch.boundary_ratio) - boundary_ratio = batch.boundary_ratio - - boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps if boundary_ratio is not None else None - latent_model_input = latents.to(target_dtype) - assert latent_model_input.shape[0] == 1, "only support batch size 1" - - if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: - # TI2V directly replaces the first frame of the latent with - # the image latent instead of appending along the channel dim - assert batch.image_latent is None, "TI2V task should not have image latents" - assert self.vae is not None, "VAE is not provided for TI2V task" - vae_device = next(self.vae.parameters()).device - self.vae = self.vae.to(local_device) - z = self.vae.encode(batch.pil_image).mean.float() - if getattr(fastvideo_args, "vae_cpu_offload", False): - self.vae = self.vae.to(vae_device) - if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): - if isinstance(self.vae.shift_factor, torch.Tensor): - z -= self.vae.shift_factor.to(z.device, z.dtype) - else: - z -= self.vae.shift_factor - - if isinstance(self.vae.scaling_factor, torch.Tensor): - z = z * self.vae.scaling_factor.to(z.device, z.dtype) - else: - z = z * self.vae.scaling_factor - - latent_model_input = latent_model_input.squeeze(0) - _, mask2 = masks_like([latent_model_input], zero=True) - - latent_model_input = (1. - mask2[0]) * z + mask2[0] * latent_model_input - # latent_model_input = latent_model_input.unsqueeze(0) - latent_model_input = latent_model_input.to(get_local_torch_device()) - latents = latent_model_input - F = batch.num_frames - temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal - spatial_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial - patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size - seq_len = ((F - 1) // temporal_scale + 1) * (batch.height // spatial_scale) * ( - batch.width // spatial_scale) // (patch_size[1] * patch_size[2]) + state = self.prepare_denoising(batch, fastvideo_args, target_dtype) + latents = state.latents # Initialize lists for ODE trajectory trajectory_timesteps: list[torch.Tensor] = [] trajectory_latents: list[torch.Tensor] = [] - is_lucy_edit = fastvideo_args.pipeline_config.lucy_edit_task # Hoisted out of the per-step loop: depends only on inputs that # are constant across denoising steps. @@ -313,17 +276,6 @@ class DenoisingStage(PipelineStage): ).to(target_dtype) * (1.0 if _is_flux else 1000.0)) else: guidance_expand = None - # V2V padding: zero-filled tensor concatenated with each step's - # latent_model_input. Shape is fixed by latents and is never - # written to, so we allocate once. - v2v_zero_pad = torch.zeros_like(latents) if batch.video_latent is not None else None - lucy_timestep_seq_len = None - if is_lucy_edit: - patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size - assert patch_size[0] == 1, "Lucy Edit timestep expansion assumes temporal patch size 1" - lucy_timestep_seq_len = (latents.shape[2] * (latents.shape[3] // patch_size[1]) * - (latents.shape[4] // patch_size[2])) - # CFG gating / stale-uncond reuse setup (Adaptive Guidance LinearAG # variant, Castillo et al. 2023). When envs.FASTVIDEO_CFG_GATE_STEP # < 1.0, the uncond forward is skipped after the gating step and the @@ -374,61 +326,10 @@ class DenoisingStage(PipelineStage): if hasattr(self, 'interrupt') and self.interrupt: continue - if boundary_timestep is None or t >= boundary_timestep: - if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload - and self.transformer_2 is not None - and next(self.transformer_2.parameters()).device.type == 'cuda'): - self.transformer_2.to('cpu') - current_model = self.transformer - if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload - and not fastvideo_args.use_fsdp_inference and current_model is not None): - transformer_device = next(current_model.parameters()).device.type - if transformer_device == 'cpu': - current_model.to(get_local_torch_device()) - current_guidance_scale = batch.guidance_scale - else: - # low-noise stage in wan2.2 - if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload - and next(self.transformer.parameters()).device.type == 'cuda'): - self.transformer.to('cpu') - current_model = self.transformer_2 - if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload - and not fastvideo_args.use_fsdp_inference and current_model is not None): - transformer_2_device = next(current_model.parameters()).device.type - if transformer_2_device == 'cpu': - current_model.to(get_local_torch_device()) - current_guidance_scale = batch.guidance_scale_2 - assert current_model is not None, "current_model is None" - - # Expand latents for V2V/I2V - latent_model_input = latents.to(target_dtype) - if batch.video_latent is not None: - if is_lucy_edit: - latent_model_input = torch.cat( - [latent_model_input, batch.video_latent], - dim=1, - ).to(target_dtype) - else: - latent_model_input = torch.cat( - [latent_model_input, batch.video_latent, v2v_zero_pad], - dim=1, - ).to(target_dtype) - elif batch.image_latent is not None: - assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task" - latent_model_input = torch.cat([latent_model_input, batch.image_latent], dim=1).to(target_dtype) - + current_model, current_guidance_scale = self.select_model(t, batch, fastvideo_args, state) + latent_model_input = self.prepare_model_input(latents, batch, target_dtype, state) assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan" - if is_lucy_edit: - assert lucy_timestep_seq_len is not None - t_expand = t.repeat(latent_model_input.shape[0], lucy_timestep_seq_len) - elif fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: - timestep = torch.stack([t]).to(get_local_torch_device()) - temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten() - temp_ts = torch.cat([temp_ts, temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep]) - timestep = temp_ts.unsqueeze(0) - t_expand = timestep.repeat(latent_model_input.shape[0], 1) - else: - t_expand = t.repeat(latent_model_input.shape[0]) + t_expand = self.prepare_timestep(t, latent_model_input, batch, state) # Gate 4: Flux2 transformer multiplies timestep by 1000 internally, so # the pipeline must pass timestep/1000 (matching Diffusers). # Diffusers casts to the latent dtype before the division; doing @@ -539,6 +440,7 @@ class DenoisingStage(PipelineStage): noise_pred_text = noise_pred if _use_cached_delta: + assert delta_cached is not None # Reuse frozen delta = cond - uncond from the last # fresh compute. Algebra: # pred = uncond + s * (cond - uncond) @@ -594,10 +496,7 @@ class DenoisingStage(PipelineStage): else: with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] - if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None: - latents = latents.squeeze(0) - latents = (1. - mask2[0]) * z + mask2[0] * latents - # latents = latents.unsqueeze(0) + latents = self.finish_step(latents, state) # save trajectory latents if needed if batch.return_trajectory_latents: @@ -660,6 +559,43 @@ class DenoisingStage(PipelineStage): return batch + def prepare_denoising(self, batch, fastvideo_args, target_dtype) -> DenoisingState: + """Build request-local state. Family stages specialize latent conditioning.""" + latents = batch.latents + assert latents is not None + assert latents.shape[0] == 1, "only support batch size 1" + return DenoisingState( + latents=latents, + video_padding=torch.zeros_like(latents) if batch.video_latent is not None else None, + ) + + def activate_transformer(self, model, inactive_model, fastvideo_args) -> None: + """Keep CPU/layerwise/FSDP offload decisions independent of the sampling recipe.""" + assert model is not None, "current_model is None" + if fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload: + if inactive_model is not None and next(inactive_model.parameters()).device.type == "cuda": + inactive_model.to("cpu") + if not fastvideo_args.use_fsdp_inference and next(model.parameters()).device.type == "cpu": + model.to(get_local_torch_device()) + + def select_model(self, timestep, batch, fastvideo_args, state) -> tuple[Any, float]: + self.activate_transformer(self.transformer, self.transformer_2, fastvideo_args) + return self.transformer, batch.guidance_scale + + def prepare_model_input(self, latents, batch, target_dtype, state) -> torch.Tensor: + latent_model_input = latents.to(target_dtype) + if batch.video_latent is not None: + latent_model_input = torch.cat([latent_model_input, batch.video_latent, state.video_padding], dim=1) + elif batch.image_latent is not None: + latent_model_input = torch.cat([latent_model_input, batch.image_latent], dim=1) + return latent_model_input.to(target_dtype) + + def prepare_timestep(self, timestep, latent_model_input, batch, state) -> torch.Tensor: + return timestep.repeat(latent_model_input.shape[0]) + + def finish_step(self, latents, state) -> torch.Tensor: + return latents + def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]: """ Prepare extra kwargs for the scheduler step / denoise step. @@ -672,8 +608,9 @@ class DenoisingStage(PipelineStage): The prepared kwargs. """ extra_step_kwargs = {} + parameters = inspect.signature(func).parameters for k, v in kwargs.items(): - accepts = k in set(inspect.signature(func).parameters.keys()) + accepts = k in parameters if accepts: extra_step_kwargs[k] = v return extra_step_kwargs @@ -1247,163 +1184,9 @@ class Cosmos25AutoDenoisingStage(PipelineStage): return self._t2w.verify_output(batch, fastvideo_args) -class DmdDenoisingStage(DenoisingStage): - """ - Denoising stage for DMD. - """ - - def __init__(self, transformer, scheduler) -> None: - super().__init__(transformer, scheduler) - self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) - - def forward( - self, - batch: ForwardBatch, - fastvideo_args: FastVideoArgs, - ) -> ForwardBatch: - """ - Run the denoising loop. - - Args: - batch: The current batch information. - fastvideo_args: The inference arguments. - - Returns: - The batch with denoised latents. - """ - # Setup precision and autocast settings - # TODO(will): make the precision configurable for inference - # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] - target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - - # Get timesteps and calculate warmup steps - timesteps = batch.timesteps - - # TODO(will): remove this once we add input/output validation for stages - if timesteps is None: - raise ValueError("Timesteps must be provided") - num_inference_steps = batch.num_inference_steps - num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - - # Prepare image latents and embeddings for I2V generation - image_embeds = batch.image_embeds - if len(image_embeds) > 0: - assert torch.isnan(image_embeds[0]).sum() == 0 - image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds] - - image_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "encoder_hidden_states_image": image_embeds, - "mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24) - }, - ) - - pos_cond_kwargs = self.prepare_extra_func_kwargs( - self.transformer.forward, - { - "encoder_hidden_states_2": batch.clip_embedding_pos, - "encoder_attention_mask": batch.prompt_attention_mask, - }, - ) - - # Get latents and embeddings - assert batch.latents is not None, "latents must be provided" - latents = batch.latents - - video_raw_latent_shape = latents.shape - prompt_embeds = batch.prompt_embeds - assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan" - timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, - dtype=torch.long, - device=get_local_torch_device()) - - # Run denoising loop - with self.progress_bar(total=len(timesteps)) as progress_bar: - for i, t in enumerate(timesteps): - # Skip if interrupted - if hasattr(self, 'interrupt') and self.interrupt: - continue - # Expand latents for I2V - noise_latents = latents.clone() - latent_model_input = latents.to(target_dtype) - - if batch.image_latent is not None: - latent_model_input = torch.cat( - [latent_model_input, batch.image_latent.permute(0, 2, 1, 3, 4)], dim=2).to(target_dtype) - assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan" - - # Prepare inputs for transformer - t_expand = t.repeat(latent_model_input.shape[0]) - guidance_expand = (torch.tensor( - [fastvideo_args.pipeline_config.embedded_cfg_scale] * latent_model_input.shape[0], - dtype=torch.float32, - device=get_local_torch_device(), - ).to(target_dtype) * 1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale is not None else None) - - # Predict noise residual - with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled): - if (vsa_available and self.attn_backend == VideoSparseAttentionBackend): - self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() - - if self.attn_metadata_builder_cls is not None: - self.attn_metadata_builder = self.attn_metadata_builder_cls() - # TODO(will): clean this up - attn_metadata = self.attn_metadata_builder.build( # type: ignore - current_timestep=i, # type: ignore - raw_latent_shape=batch.raw_latent_shape[2:5], # type: ignore - patch_size=fastvideo_args.pipeline_config. # type: ignore - dit_config.patch_size, # type: ignore - VSA_sparsity=fastvideo_args.VSA_sparsity, # type: ignore - device=get_local_torch_device(), # type: ignore - ) # type: ignore - assert attn_metadata is not None, "attn_metadata cannot be None" - else: - attn_metadata = None - else: - attn_metadata = None - - batch.is_cfg_negative = False - with set_forward_context( - current_timestep=i, - attn_metadata=attn_metadata, - forward_batch=batch, - # fastvideo_args=fastvideo_args - ): - # Run transformer - pred_noise = self.transformer( - latent_model_input.permute(0, 2, 1, 3, 4), - prompt_embeds, - t_expand, - guidance=guidance_expand, - **image_kwargs, - **pos_cond_kwargs, - ).permute(0, 2, 1, 3, 4) - - pred_video = pred_noise_to_pred_video(pred_noise=pred_noise.flatten(0, 1), - noise_input_latent=noise_latents.flatten(0, 1), - timestep=t_expand, - scheduler=self.scheduler).unflatten(0, pred_noise.shape[:2]) - - if i < len(timesteps) - 1: - next_timestep = timesteps[i + 1] * torch.ones([1], dtype=torch.long, device=pred_video.device) - noise = torch.randn(video_raw_latent_shape, - dtype=pred_video.dtype, - generator=batch.generator[0]).to(self.device) - latents = self.scheduler.add_noise(pred_video.flatten(0, 1), noise.flatten(0, 1), - next_timestep).unflatten(0, pred_video.shape[:2]) - else: - latents = pred_video - - # Update progress bar - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and - (i + 1) % self.scheduler.order == 0 and progress_bar is not None): - progress_bar.update() - - # Gather results if using sequence parallelism - latents = latents.permute(0, 2, 1, 3, 4) - # Update batch with final latents - batch.latents = latents - - return batch +def __getattr__(name: str) -> Any: + # Keep legacy imports without an eager family -> shared -> family cycle. + if name == "DmdDenoisingStage": + from fastvideo.pipelines.basic.wan.stages.dmd import DmdDenoisingStage + return DmdDenoisingStage + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 083732901..755f6cef7 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -47,21 +47,8 @@ from fastvideo.configs.pipelines.turbodiffusion import ( TurboDiffusionT2V_14B_Config, TurboDiffusionT2V_1_3B_Config, ) -from fastvideo.configs.pipelines.wan import ( - FastWan2_1_T2V_480P_Config, - FastWan2_2_TI2V_5B_Config, - LucyEditDevConfig, - SelfForcingWan2_2_T2V480PConfig, - SelfForcingWanT2V480PConfig, - WANV2VConfig, - Wan2_2_I2V_A14B_Config, - Wan2_2_T2V_A14B_Config, - Wan2_2_TI2V_5B_Config, - WanI2V480PConfig, - WanI2V720PConfig, - WanT2V480PConfig, - WanT2V720PConfig, -) +from fastvideo.models.wan import pipeline_config as wan_pipeline_config +from fastvideo.models.wan.definition import WAN_MODEL_DEFINITION_GROUPS, WanModelDefinition from fastvideo.configs.pipelines.glm_image import GlmImageConfig from fastvideo.configs.pipelines.flux import FluxPipelineConfig from fastvideo.configs.pipelines.sd35 import SD35Config @@ -239,6 +226,19 @@ def _get_config_info( return None +def _register_wan_configs(definitions: tuple[WanModelDefinition, ...]) -> None: + for definition in definitions: + register_configs( + sampling_param_cls=None, + pipeline_config_cls=getattr(wan_pipeline_config, definition.pipeline_config), + workload_types=tuple(WorkloadType(value) for value in definition.workload_types), + hf_model_paths=list(definition.hf_model_paths), + model_detectors=[definition.matches] if definition.match_any else None, + model_family="wan", + default_preset=definition.preset, + ) + + def _register_configs() -> None: # MMAudio large-44k-v2 (video/text-to-audio). The checkpoint is converted # into standard per-component FastVideo/Diffusers-style directories by @@ -923,92 +923,8 @@ def _register_configs() -> None: default_preset="turbo_i2v_a14b", ) - # Wan — defaults provided by presets (no sampling_param_cls needed) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=WanT2V480PConfig, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=[ - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", - ], - model_detectors=[lambda path: "wanpipeline" in path.lower()], - model_family="wan", - default_preset="wan_t2v_1_3b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=WanT2V720PConfig, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=[ - "Wan-AI/Wan2.1-T2V-14B-Diffusers", - "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", - ], - model_family="wan", - default_preset="wan_t2v_14b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=WanI2V480PConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", - ], - model_detectors=[lambda path: "wanimagetovideo" in path.lower()], - model_family="wan", - default_preset="wan_i2v_14b_480p", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=WanI2V720PConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", - ], - model_family="wan", - default_preset="wan_i2v_14b_720p", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=WanI2V480PConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", - ], - model_family="wan", - default_preset="wan_fun_1_3b_inp", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=WANV2VConfig, - workload_types=(), - hf_model_paths=[ - "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers", - ], - model_family="wan", - default_preset="wan_fun_1_3b_control", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=FastWan2_1_T2V_480P_Config, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=[ - "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", - "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers", - ], - model_detectors=[lambda path: "wandmdpipeline" in path.lower()], - model_family="wan", - default_preset="fast_wan_t2v_480p", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=Wan2_2_TI2V_5B_Config, - workload_types=(WorkloadType.T2V, WorkloadType.I2V), - hf_model_paths=[ - "Wan-AI/Wan2.2-TI2V-5B-Diffusers", - ], - model_family="wan", - default_preset="wan_2_2_ti2v_5b", - ) + # Preserve first-match ordering around the DreamX registrations below. + _register_wan_configs(WAN_MODEL_DEFINITION_GROUPS[0]) register_configs( sampling_param_cls=None, @@ -1043,83 +959,7 @@ def _register_configs() -> None: model_family="dreamx_world", default_preset="dreamx_world_5b_ar", ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=FastWan2_2_TI2V_5B_Config, - workload_types=(WorkloadType.T2V, WorkloadType.I2V), - hf_model_paths=[ - "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", - "FastVideo/FastWan2.2-TI2V-5B-Diffusers", - ], - model_family="wan", - default_preset="fast_wan_2_2_ti2v_5b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=LucyEditDevConfig, - workload_types=(), - hf_model_paths=[ - "decart-ai/Lucy-Edit-Dev", - "decart-ai/Lucy-Edit-1.1-Dev", - ], - model_detectors=[lambda path: "lucy-edit" in path.lower()], - model_family="wan", - default_preset="lucy_edit_dev", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=Wan2_2_T2V_A14B_Config, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=[ - "Wan-AI/Wan2.2-T2V-A14B-Diffusers", - ], - model_family="wan", - default_preset="wan_2_2_t2v_a14b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=Wan2_2_I2V_A14B_Config, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=[ - "Wan-AI/Wan2.2-I2V-A14B-Diffusers", - ], - model_family="wan", - default_preset="wan_2_2_i2v_a14b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=SelfForcingWanT2V480PConfig, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=[ - "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", - ], - model_detectors=[lambda path: "wancausaldmdpipeline" in path.lower()], - model_family="wan", - default_preset="sf_wan_t2v_1_3b", - ) - # SFWan2.2: T2V and I2V variants by path - register_configs( - sampling_param_cls=None, - pipeline_config_cls=SelfForcingWan2_2_T2V480PConfig, - workload_types=(WorkloadType.T2V, ), - hf_model_paths=["rand0nmr/SFWan2.2-T2V-A14B-Diffusers"], - model_detectors=[ - lambda path: ("sfwan2.2" in path.lower() or "sfwan2_2" in path.lower()) and "i2v" not in path.lower(), - ], - model_family="wan", - default_preset="sf_wan_2_2_t2v_a14b", - ) - register_configs( - sampling_param_cls=None, - pipeline_config_cls=SelfForcingWan2_2_T2V480PConfig, - workload_types=(WorkloadType.I2V, ), - hf_model_paths=["FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers"], - model_detectors=[ - lambda path: ("sfwan2.2" in path.lower() or "sfwan2_2" in path.lower()) and "i2v" in path.lower(), - ], - model_family="wan", - default_preset="sf_wan_2_2_i2v_a14b", - ) + _register_wan_configs(WAN_MODEL_DEFINITION_GROUPS[1]) # MiniMax H3 register_configs( diff --git a/fastvideo/tests/api/test_lucy_edit.py b/fastvideo/tests/api/test_lucy_edit.py index 3dd28a21c..03db61890 100644 --- a/fastvideo/tests/api/test_lucy_edit.py +++ b/fastvideo/tests/api/test_lucy_edit.py @@ -2,7 +2,7 @@ from fastvideo.api.presets import get_preset from fastvideo.api.sampling_param import SamplingParam -from fastvideo.configs.pipelines.wan import LucyEditDevConfig +from fastvideo.models.wan.pipeline_config import LucyEditDevConfig from fastvideo.fastvideo_args import WorkloadType from fastvideo.pipelines.basic.wan.lucy_edit_pipeline import LucyEditPipeline from fastvideo.pipelines.pipeline_registry import PipelineType, get_pipeline_registry diff --git a/fastvideo/tests/api/test_wan_definitions.py b/fastvideo/tests/api/test_wan_definitions.py new file mode 100644 index 000000000..eaa295fac --- /dev/null +++ b/fastvideo/tests/api/test_wan_definitions.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Weight-free compatibility contracts for the Wan family catalog. + +Expected resolutions and numerical defaults below come from the registry and +configs at 275cf421, before the definition extraction. No Hub calls are needed. +""" + +import dataclasses +import importlib +import json +import pickle +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo import registry +from fastvideo.api.presets import get_preset +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.models.wan import pipeline_config +from fastvideo.models.wan.definition import DMD_TRAINING_NOISE_SHIFT, WAN_MODEL_DEFINITIONS +from fastvideo.pipelines.basic.wan.presets import ALL_PRESETS + + +# Keep this independent of the candidate catalog: dropping or misassigning an +# alias must fail even if the remaining definitions are internally consistent. +LEGACY_VARIANTS = ( + ("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "WanT2V480PConfig", "wan_t2v_1_3b", ("t2v",)), + ("Wan-AI/Wan2.1-T2V-14B-Diffusers", "WanT2V720PConfig", "wan_t2v_14b", ("t2v",)), + ("FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", "WanT2V720PConfig", "wan_t2v_14b", ("t2v",)), + ("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", "WanI2V480PConfig", "wan_i2v_14b_480p", ("i2v",)), + ("Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", "WanI2V720PConfig", "wan_i2v_14b_720p", ("i2v",)), + ("weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", "WanI2V480PConfig", "wan_fun_1_3b_inp", ("i2v",)), + ("IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers", "WANV2VConfig", "wan_fun_1_3b_control", ()), + ("FastVideo/FastWan2.1-T2V-1.3B-Diffusers", "FastWan2_1_T2V_480P_Config", "fast_wan_t2v_480p", ("t2v",)), + ("FastVideo/FastWan2.1-T2V-14B-480P-Diffusers", "FastWan2_1_T2V_480P_Config", "fast_wan_t2v_480p", ("t2v",)), + ("Wan-AI/Wan2.2-TI2V-5B-Diffusers", "Wan2_2_TI2V_5B_Config", "wan_2_2_ti2v_5b", ("t2v", "i2v")), + ("FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", "FastWan2_2_TI2V_5B_Config", "fast_wan_2_2_ti2v_5b", ("t2v", "i2v")), + ("FastVideo/FastWan2.2-TI2V-5B-Diffusers", "FastWan2_2_TI2V_5B_Config", "fast_wan_2_2_ti2v_5b", ("t2v", "i2v")), + ("decart-ai/Lucy-Edit-Dev", "LucyEditDevConfig", "lucy_edit_dev", ()), + ("decart-ai/Lucy-Edit-1.1-Dev", "LucyEditDevConfig", "lucy_edit_dev", ()), + ("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "Wan2_2_T2V_A14B_Config", "wan_2_2_t2v_a14b", ("t2v",)), + ("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "Wan2_2_I2V_A14B_Config", "wan_2_2_i2v_a14b", ("i2v",)), + ("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", "SelfForcingWanT2V480PConfig", "sf_wan_t2v_1_3b", ("t2v",)), + ("rand0nmr/SFWan2.2-T2V-A14B-Diffusers", "SelfForcingWan2_2_T2V480PConfig", "sf_wan_2_2_t2v_a14b", ("t2v",)), + ("FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers", "SelfForcingWan2_2_T2V480PConfig", "sf_wan_2_2_i2v_a14b", ("i2v",)), +) + + +@pytest.fixture(autouse=True) +def no_hub_access(monkeypatch): + def fail(*args, **kwargs): + pytest.fail("Wan config contracts must not access the Hub") + + monkeypatch.setattr(registry, "maybe_download_model_index", fail) + registry.get_model_info.cache_clear() + yield + registry.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model_path,config_name,preset_name,workloads", LEGACY_VARIANTS) +def test_legacy_aliases_and_sampling_defaults(model_path, config_name, preset_name, workloads): + expected_config = getattr(pipeline_config, config_name) + preset = get_preset(preset_name, "wan") + # Exact HF ID, mirror ID, and renamed parent directory of a local checkpoint + # all use the existing exact/short-name resolution before any manifest IO. + short_name = model_path.rsplit("/", 1)[-1] + for path in (model_path, f"mirror/{short_name}", f"/checkpoints/{short_name}"): + info = registry._get_config_info(path) + assert info.pipeline_config_cls is expected_config + assert info.model_family == "wan" + assert info.default_preset == preset_name + assert tuple(value.value for value in info.workload_types) == workloads + assert info.sampling_param_cls is None + assert info.pipeline_cls_name is None + sampling = SamplingParam.from_pretrained(path) + for name, value in preset.defaults.items(): + assert getattr(sampling, name) == value + + +@pytest.mark.parametrize("directory,manifest_class,expected_family,expected_preset", [ + ("custom", "WanPipeline", "wan", "wan_t2v_1_3b"), + ("custom", "WanImageToVideoPipeline", "wan", "wan_i2v_14b_480p"), + ("custom", "WanDMDPipeline", "wan", "fast_wan_t2v_480p"), + ("custom", "WanCausalDMDPipeline", "wan", "sf_wan_t2v_1_3b"), + ("Lucy-Edit-custom", "UnknownPipeline", "wan", "lucy_edit_dev"), + ("SFWan2.2-custom", "UnknownPipeline", "wan", "sf_wan_2_2_t2v_a14b"), + ("SFWan2_2-I2V-custom", "UnknownPipeline", "wan", "sf_wan_2_2_i2v_a14b"), + # Counterintuitive but supported legacy precedence: do not silently fix it + # during a relocation, including when path and class detectors disagree. + ("lucy-edit-custom", "WanPipeline", "wan", "wan_t2v_1_3b"), + ("sfwan2.2-i2v-custom", "WanCausalDMDPipeline", "wan", "sf_wan_t2v_1_3b"), + ("dreamx-world-5b-lucy-edit-custom", "UnknownPipeline", "dreamx_world", "dreamx_world_5b_ar"), + ("dreamx-world-5b-custom", "WanPipeline", "wan", "wan_t2v_1_3b"), +]) +def test_local_manifest_detectors_keep_first_match(tmp_path, directory, manifest_class, expected_family, expected_preset): + checkpoint = tmp_path / directory + checkpoint.mkdir() + (checkpoint / "model_index.json").write_text(json.dumps({ + "_class_name": manifest_class, "_diffusers_version": "0.33.0", + })) + info = registry._get_config_info(str(checkpoint)) + assert (info.model_family, info.default_preset) == (expected_family, expected_preset) + + +def test_manifest_selection_and_explicit_override_are_not_pinned(monkeypatch, tmp_path): + from fastvideo.pipelines import pipeline_registry + + selected = [] + def resolve(name, *args): + selected.append(name) + return type(name, (), {}) + + monkeypatch.setattr(pipeline_registry, "get_pipeline_registry", + lambda *args: SimpleNamespace(resolve_pipeline_cls=resolve)) + checkpoint = tmp_path / "Wan2.1-T2V-1.3B-Diffusers" + checkpoint.mkdir() + manifest = checkpoint / "model_index.json" + manifest.write_text(json.dumps({"_class_name": "CustomWanPipeline", "_diffusers_version": "0.33.0"})) + info = registry.get_model_info(str(checkpoint)) + assert info.pipeline_cls.__name__ == "CustomWanPipeline" + assert info.pipeline_config_cls is pipeline_config.WanT2V480PConfig + # An explicit override works even without a manifest, as before. + manifest.unlink() + info = registry.get_model_info(str(checkpoint), override_pipeline_cls_name="WanDMDPipeline") + assert info.pipeline_cls.__name__ == "WanDMDPipeline" + assert selected == ["CustomWanPipeline", "WanDMDPipeline"] + + +@pytest.mark.parametrize("config_name,flow_shift,dmd_steps,load_encoder", [ + ("WanT2V480PConfig", 3.0, None, False), + ("WanT2V720PConfig", 5.0, None, False), + ("WanI2V480PConfig", 3.0, None, True), + ("WanI2V720PConfig", 5.0, None, True), + ("WANV2VConfig", 3.0, None, True), + ("FastWan2_1_T2V_480P_Config", 8.0, [1000, 757, 522], False), + ("Wan2_2_TI2V_5B_Config", 5.0, None, True), + ("FastWan2_2_TI2V_5B_Config", 5.0, [1000, 757, 522], True), + ("LucyEditDevConfig", 5.0, None, True), + ("Wan2_2_T2V_A14B_Config", 12.0, [1000, 750, 500, 250], True), + ("Wan2_2_I2V_A14B_Config", 5.0, None, True), + ("SelfForcingWanT2V480PConfig", 5.0, [1000, 750, 500, 250], False), + ("SelfForcingWan2_2_T2V480PConfig", 12.0, [1000, 850, 700, 550, 350, 275, 200, 125], True), +]) +def test_component_precision_and_sampling_defaults(config_name, flow_shift, dmd_steps, load_encoder): + config_cls = getattr(pipeline_config, config_name) + first, second = config_cls(), config_cls() + assert first.flow_shift == flow_shift + assert first.dmd_denoising_steps == dmd_steps + assert first.dit_precision == "bf16" + assert first.vae_precision == "fp32" + assert first.vae_decode_precision == "bf16" + assert first.text_encoder_precisions == ("fp32",) + assert first.vae_config.load_encoder is load_encoder + assert first.vae_config.load_decoder is True + # Each request/instance still gets independent mutable component configs. + first.vae_config.load_encoder = not load_encoder + assert second.vae_config.load_encoder is load_encoder + first.dit_config.boundary_ratio = 0.123 + assert second.dit_config.boundary_ratio != 0.123 + if first.dmd_denoising_steps is not None: + first.dmd_denoising_steps.append(1) + assert second.dmd_denoising_steps == dmd_steps + + +def test_definitions_are_complete_data_only_and_reference_existing_configs(): + paths = [path for definition in WAN_MODEL_DEFINITIONS for path in definition.hf_model_paths] + assert len(paths) == len(set(paths)) + assert set(paths) == {case[0] for case in LEGACY_VARIANTS} + assert {definition.preset for definition in WAN_MODEL_DEFINITIONS} == {preset.name for preset in ALL_PRESETS} + assert len(WAN_MODEL_DEFINITIONS) == len(ALL_PRESETS) + assert DMD_TRAINING_NOISE_SHIFT == 8.0 + for definition in WAN_MODEL_DEFINITIONS: + assert json.loads(json.dumps(dataclasses.asdict(definition)))["preset"] == definition.preset + config = getattr(pipeline_config, definition.pipeline_config)() + if definition.sampling == "causal_dmd": + assert config.is_causal + assert config.warp_denoising_step + assert config.dmd_denoising_steps + elif definition.sampling == "dmd": + assert not config.is_causal + assert config.dmd_denoising_steps == [1000, 757, 522] + else: + assert definition.sampling == "unipc" + assert not config.is_causal + with pytest.raises(dataclasses.FrozenInstanceError): + definition.preset = "changed" + + +def test_legacy_pipeline_config_aliases_and_pickle_transport(): + def assert_same_state(actual, expected): + assert type(actual) is type(expected) + if dataclasses.is_dataclass(expected): + # vars also covers derived attributes omitted by dataclasses.asdict, + # including the VAE normalization shift_factor. + assert_same_state(vars(actual), vars(expected)) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys() + for key in expected: + assert_same_state(actual[key], expected[key]) + elif isinstance(expected, (tuple, list)): + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected, strict=True): + assert_same_state(actual_item, expected_item) + elif isinstance(expected, torch.Tensor): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + assert actual == expected + + legacy = importlib.import_module("fastvideo.configs.pipelines.wan") + for name in legacy.__all__: + assert getattr(legacy, name) is getattr(pipeline_config, name) + for name in {definition.pipeline_config for definition in WAN_MODEL_DEFINITIONS}: + original = getattr(pipeline_config, name)() + payload = pickle.dumps(original, protocol=0) + legacy_payload = payload.replace(b"fastvideo.models.wan.pipeline_config\n", b"fastvideo.configs.pipelines.wan\n") + assert legacy_payload != payload + restored = pickle.loads(legacy_payload) + assert type(restored) is type(original) + assert_same_state(restored, original) diff --git a/fastvideo/tests/contract/test_ci_test_collection.py b/fastvideo/tests/contract/test_ci_test_collection.py index 4a0e0b179..60c1fc4dc 100644 --- a/fastvideo/tests/contract/test_ci_test_collection.py +++ b/fastvideo/tests/contract/test_ci_test_collection.py @@ -186,6 +186,20 @@ def test_active_pipeline_has_no_modal_or_untrusted_compute_path(): assert all(str(step.get("command", "")).startswith("/opt/fastvideo-ci-runner/") for step in steps) +def test_expensive_integration_lanes_wait_for_selected_goldens(): + # A skipped conditional golden step satisfies the dependency (direct + # reruns/training-only plans remain available); a failed selected gate does + # not. Never permit allow_failure or a soft-fail dependency here. + for step in _pipeline_steps(): + key = step["key"] + fastcheck = key in {"encoder", "vae", "transformer", "kernel-tests", "unit", "dreamverse"} + if fastcheck or key == "golden-gate": + assert "depends_on" not in step + else: + assert step["depends_on"] == "golden-gate" + assert "allow_dependency_failure" not in step + + def test_gpu_tests_preserve_a_launcher_assigned_rendezvous_port(): hard_assignment = re.compile(r'''os\.environ\[\s*["']MASTER_PORT["']\s*\]\s*=''') overwrites = [] diff --git a/fastvideo/tests/contract/test_merge_ci_plan.py b/fastvideo/tests/contract/test_merge_ci_plan.py index 684a0c840..8cfce78ee 100644 --- a/fastvideo/tests/contract/test_merge_ci_plan.py +++ b/fastvideo/tests/contract/test_merge_ci_plan.py @@ -30,24 +30,40 @@ def test_docs_only_merge_adds_no_gpu_lanes(): assert plan.encoded_ssim_tests() == "none" -@pytest.mark.parametrize("path", [ - "fastvideo/models/dits/wanvideo.py", - "fastvideo/configs/models/dits/wanvideo.py", - "fastvideo/models/wan/__init__.py", - "fastvideo/models/wan/config.py", - "fastvideo/models/wan/transformer.py", - "fastvideo/models/wan/vae.py", - "fastvideo/models/wan/vae_config.py", +@pytest.mark.parametrize("path, goldens", [ + ("fastvideo/models/dits/wanvideo.py", "test_wan_denoising.py,test_wan_t2v.py"), + ("fastvideo/models/wan/transformer.py", "test_wan_denoising.py,test_wan_t2v.py"), + ("fastvideo/models/wan/vae.py", "test_wan_vae.py"), + ("fastvideo/models/wan/vae_config.py", "test_wan_vae.py"), + ("fastvideo/models/vaes/wanvae.py", "test_wan_vae.py"), + ("fastvideo/configs/models/vaes/wanvae.py", "test_wan_vae.py"), + ("fastvideo/models/wan/causal_transformer.py", "test_wan_causal.py"), + ("fastvideo/models/dits/causal_wanvideo.py", "test_wan_causal.py"), + ("fastvideo/pipelines/basic/wan/stages/conditioning.py", "test_wan_vae.py"), + ("fastvideo/pipelines/basic/wan/stages/causal_denoising.py", "test_wan_causal.py"), + ("fastvideo/pipelines/basic/wan/stages/denoising.py", "test_wan_denoising.py,test_wan_t2v.py"), + ("fastvideo/pipelines/basic/wan/stages/dmd.py", "test_wan_denoising.py,test_wan_t2v.py"), ]) -def test_model_family_change_selects_focused_golden_and_ssim(path): +def test_model_family_change_selects_focused_golden_and_ssim(path, goldens): plan = PLAN_MERGE_CI.classify_paths([path]) assert plan.encoded_lanes() == ",golden-gate,ssim," - assert plan.encoded_golden_tests() == "test_wan_t2v.py" + assert plan.encoded_golden_tests() == goldens assert plan.encoded_ssim_tests() == ( "test_causal_similarity.py,test_wan_i2v_similarity.py,test_wan_t2v_similarity.py") +@pytest.mark.parametrize("path", [ + "fastvideo/configs/models/dits/wanvideo.py", + "fastvideo/models/wan/config.py", + "fastvideo/models/wan/__init__.py", + "fastvideo/pipelines/basic/wan/wan_pipeline.py", +]) +def test_wan_shared_config_and_wiring_select_all_wan_gates(path): + plan = PLAN_MERGE_CI.classify_paths([path]) + assert plan.encoded_golden_tests() == "test_wan_causal.py,test_wan_denoising.py,test_wan_t2v.py,test_wan_vae.py" + + def test_wan_family_relocation_keeps_shared_registry_coverage(): plan = PLAN_MERGE_CI.classify_paths([ "fastvideo/models/dits/wanvideo.py", @@ -138,8 +154,8 @@ def test_performance_implementation_selects_only_performance_lane(): @pytest.mark.parametrize("path", [ "fastvideo/registry.py", - "fastvideo/models/vaes/wanvae.py", - "fastvideo/configs/models/vaes/wanvae.py", + "fastvideo/pipelines/stages/denoising.py", + "fastvideo/pipelines/stages/causal_denoising.py", ]) def test_shared_runtime_change_gets_focused_quality_smoke_not_every_lane(path): plan = PLAN_MERGE_CI.classify_paths([path]) diff --git a/fastvideo/tests/contract/test_tensor_golden_isolation.py b/fastvideo/tests/contract/test_tensor_golden_isolation.py new file mode 100644 index 000000000..34665c6d6 --- /dev/null +++ b/fastvideo/tests/contract/test_tensor_golden_isolation.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden recipes must not inherit or leak the model loader's compute policy.""" + +import os + +import pytest +import torch + +from fastvideo import utils +from fastvideo.tests.golden_gate._tensor_golden import deterministic_forward + + +@pytest.mark.parametrize("existing_policy", [False, True]) +@pytest.mark.parametrize("fails", [False, True]) +def test_tensor_golden_restores_runtime_state(monkeypatch, existing_policy, fails): + # No CUDA work is needed to check context ownership. + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + previous_policy = utils.MixedPrecisionState(param_dtype=torch.float16) + if existing_policy: + monkeypatch.setattr(utils._mixed_precision_state, "state", previous_policy, raising=False) + else: + monkeypatch.delattr(utils._mixed_precision_state, "state", raising=False) + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + benchmark = torch.backends.cudnn.benchmark + + class RecipeError(Exception): + pass + + try: + with deterministic_forward("FLASH_ATTN"): + assert utils.get_compute_dtype() == torch.get_default_dtype() + assert os.environ["FASTVIDEO_ATTENTION_BACKEND"] == "FLASH_ATTN" + assert torch.are_deterministic_algorithms_enabled() + assert not torch.backends.cudnn.benchmark + # Same process-state mutation performed by TransformerLoader. + utils.set_mixed_precision_policy(torch.bfloat16, torch.float32) + assert utils.get_compute_dtype() == torch.bfloat16 + if fails: + raise RecipeError + except RecipeError: + assert fails + + assert hasattr(utils._mixed_precision_state, "state") == existing_policy + if existing_policy: + assert utils.get_mixed_precision_state() is previous_policy + assert os.environ["FASTVIDEO_ATTENTION_BACKEND"] == "TORCH_SDPA" + assert torch.are_deterministic_algorithms_enabled() == deterministic + assert torch.is_deterministic_algorithms_warn_only_enabled() == warn_only + assert torch.backends.cudnn.benchmark == benchmark diff --git a/fastvideo/tests/contract/test_wan_validation_order.py b/fastvideo/tests/contract/test_wan_validation_order.py new file mode 100644 index 000000000..5db6a8503 --- /dev/null +++ b/fastvideo/tests/contract/test_wan_validation_order.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shell selection/order contracts; pytest is mocked, no GPU payload executes.""" + +import os +import subprocess +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[3] + + +@pytest.mark.parametrize("lane, expected", [ + ("vae", ["golden_gate/test_wan_vae.py", "tests/vaes"]), + ("transformer", ["golden_gate/test_wan_t2v.py", "golden_gate/test_wan_causal.py", "tests/transformers"]), +]) +@pytest.mark.parametrize("fail_first", [False, True]) +def test_component_lane_stops_after_failed_golden(tmp_path, lane, expected, fail_first): + log = tmp_path / "calls" + stub = tmp_path / "pytest" + stub.write_text('#!/bin/bash\nprintf "%s\\n" "$*" >> "$WAN_TEST_LOG"\nexit "${WAN_TEST_EXIT:-0}"\n') + stub.chmod(0o755) + env = {**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}", "WAN_TEST_LOG": str(log), + "WAN_TEST_EXIT": "1" if fail_first else "0"} + result = subprocess.run(["bash", f".buildkite/scripts/lanes/{lane}.sh"], cwd=ROOT, env=env, check=False) + calls = log.read_text().splitlines() + assert len(calls) == (1 if fail_first else len(expected)) + assert all(path in call for path, call in zip(expected, calls)) + assert result.returncode == int(fail_first) diff --git a/fastvideo/tests/distributed/test_sp_wan.py b/fastvideo/tests/distributed/test_sp_wan.py index 83369131c..aa3c19f6e 100644 --- a/fastvideo/tests/distributed/test_sp_wan.py +++ b/fastvideo/tests/distributed/test_sp_wan.py @@ -12,7 +12,7 @@ import pytest import torch import torch.distributed as dist -from fastvideo.configs.models.dits.wanvideo import ( +from fastvideo.models.wan.config import ( WanVideoArchConfig, WanVideoConfig, ) @@ -21,7 +21,7 @@ from fastvideo.distributed import ( maybe_init_distributed_environment_and_model_parallel, ) from fastvideo.forward_context import set_forward_context -from fastvideo.models.dits.wanvideo import WanTransformer3DModel +from fastvideo.models.wan.transformer import WanTransformer3DModel from fastvideo.pipelines.pipeline_batch_info import ForwardBatch SP_WORLD_SIZE = 2 diff --git a/fastvideo/tests/golden_gate/_harness.py b/fastvideo/tests/golden_gate/_harness.py index 0be65b4e0..3bc1969b8 100644 --- a/fastvideo/tests/golden_gate/_harness.py +++ b/fastvideo/tests/golden_gate/_harness.py @@ -57,25 +57,37 @@ class GateSpec: postprocess: Callable[[Any], torch.Tensor] = field(default=lambda out: out) model_root_env: str | None = None # env var naming a local checkpoint root weight_file: str = "diffusion_pytorch_model.safetensors" # single-file fallback + revision: str | None = None # Immutable checkpoint revision for pinned gates. -def env_fingerprint(spec: GateSpec) -> dict[str, str]: +def env_fingerprint(spec: GateSpec | str | None) -> dict[str, str]: try: import flash_attn fa_version = str(getattr(flash_attn, "__version__", "?")) except ImportError: fa_version = "" - return { + backend = spec.attention_backend if isinstance(spec, GateSpec) else spec + result = { "torch": str(torch.__version__), "cuda": str(torch.version.cuda), "cudnn": str(torch.backends.cudnn.version()), "flash_attn": fa_version, "device_name": torch.cuda.get_device_name(0), - "attention_backend": spec.attention_backend, + "attention_backend": backend or "NONE", "fastvideo_fa4": os.environ.get("FASTVIDEO_FA4", ""), "tf32_matmul": str(torch.backends.cuda.matmul.allow_tf32), "tf32_cudnn": str(torch.backends.cudnn.allow_tf32), } + if backend is None: + # Codec gates do not execute FlashAttention; its version/selection is not + # part of the codec compute path. Keep torch/CUDA/cuDNN/device/TF32 pins. + result.pop("flash_attn") + result.pop("fastvideo_fa4") + elif not isinstance(spec, GateSpec): + # New tensor gates record the effective switch. Unset and "0" both + # select FA2; retain the old block-gate metadata contract unchanged. + result["fastvideo_fa4"] = str(int(os.environ.get("FASTVIDEO_FA4", "0") != "0")) + return result def _component_dir(spec: GateSpec) -> Path: @@ -92,16 +104,16 @@ def _component_dir(spec: GateSpec) -> Path: index_name = f"{spec.subfolder}/{spec.weight_file}.index.json" prefix = spec.prefix_template.format(N=spec.layer) try: - index_path = Path(hf_hub_download(spec.repo_id, index_name)) + index_path = Path(hf_hub_download(spec.repo_id, index_name, revision=spec.revision)) except EntryNotFoundError: - single = Path(hf_hub_download(spec.repo_id, f"{spec.subfolder}/{spec.weight_file}")) + single = Path(hf_hub_download(spec.repo_id, f"{spec.subfolder}/{spec.weight_file}", revision=spec.revision)) return single.parent weight_map = json.loads(index_path.read_text())["weight_map"] shards = sorted({shard for name, shard in weight_map.items() if name.startswith(prefix)}) if not shards: pytest.fail(f"no shards found for prefix {prefix} in {spec.repo_id}", pytrace=False) for shard in shards: - hf_hub_download(spec.repo_id, f"{spec.subfolder}/{shard}") + hf_hub_download(spec.repo_id, f"{spec.subfolder}/{shard}", revision=spec.revision) return index_path.parent @@ -133,19 +145,26 @@ def load_layer_state(spec: GateSpec, component_dir: Path) -> dict[str, torch.Ten return state +def device_folder() -> str: + return re.sub(r"[^A-Za-z0-9]+", "_", torch.cuda.get_device_name(0)).strip("_") + + def _golden_relpath(spec: GateSpec) -> str: - device_slug = re.sub(r"[^A-Za-z0-9]+", "_", torch.cuda.get_device_name(0)).strip("_") - return f"{device_slug}/{spec.name}_layer{spec.layer}_{spec.attention_backend}_seed{spec.seed}.pt" + return f"{device_folder()}/{spec.name}_layer{spec.layer}_{spec.attention_backend}_seed{spec.seed}.pt" def _resolve_golden(spec: GateSpec) -> tuple[Path, bool]: - local = GOLDEN_ROOT / _golden_relpath(spec) + return resolve_golden_path(_golden_relpath(spec)) + + +def resolve_golden_path(relative_path: str) -> tuple[Path, bool]: + local = GOLDEN_ROOT / relative_path if local.exists(): return local, True from huggingface_hub import hf_hub_download from huggingface_hub.errors import EntryNotFoundError try: - remote = hf_hub_download(GOLDEN_REPO_ID, f"golden_gates/{_golden_relpath(spec)}", repo_type="dataset") + remote = hf_hub_download(GOLDEN_REPO_ID, f"golden_gates/{relative_path}", repo_type="dataset") return Path(remote), True except EntryNotFoundError: return local, False diff --git a/fastvideo/tests/golden_gate/_tensor_golden.py b/fastvideo/tests/golden_gate/_tensor_golden.py new file mode 100644 index 000000000..4c9d04470 --- /dev/null +++ b/fastvideo/tests/golden_gate/_tensor_golden.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Named tensor references for codecs and short sampling trajectories. + +Keep weight loading and forward calls in each test. This helper only controls +determinism and compares device/runtime-matched tensors; it is not a model runner. +""" + +import os +from contextlib import contextmanager + +import pytest +import torch + +from fastvideo.tests.golden_gate._harness import DEFAULT_SEED, device_folder, env_fingerprint, resolve_golden_path + + +@contextmanager +def deterministic_forward(attention_backend=None): + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("golden gate requires a bf16-capable CUDA GPU") + from fastvideo.utils import _mixed_precision_state + + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + benchmark = torch.backends.cudnn.benchmark + previous_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") + had_precision_state = hasattr(_mixed_precision_state, "state") + previous_precision_state = getattr(_mixed_precision_state, "state", None) + try: + # TransformerLoader sets a thread-local compute dtype even without + # FSDP. Let each recipe choose its own construction dtype, then restore + # the caller's policy: leaking bf16 changes later block gates' attention + # backend despite their unchanged torch default dtype and environment. + if had_precision_state: + del _mixed_precision_state.state + if attention_backend is not None: + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = attention_backend + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.benchmark = False + yield torch.device("cuda:0") + finally: + if had_precision_state: + _mixed_precision_state.state = previous_precision_state + elif hasattr(_mixed_precision_state, "state"): + del _mixed_precision_state.state + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + torch.backends.cudnn.benchmark = benchmark + if previous_backend is None: + os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None) + else: + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous_backend + + +def assert_tensor_golden(name, outputs, *, identity, attention_backend=None, seed=DEFAULT_SEED): + assert outputs, "A golden must contain at least one named tensor" + outputs = {key: value.detach().contiguous().cpu() for key, value in sorted(outputs.items())} + for key, value in outputs.items(): + assert value.numel() and torch.isfinite(value).all(), f"Invalid golden output: {key}" + metadata = { + "schema": 1, + "name": name, + "seed": seed, + "identity": identity, + "env": env_fingerprint(attention_backend), + "tensors": {key: {"shape": list(value.shape), "dtype": str(value.dtype)} for key, value in outputs.items()}, + } + relative = f"{device_folder()}/{name}_seed{seed}.pt" + path, exists = resolve_golden_path(relative) + if not exists: + # Never make a missing reference pass on the next run of candidate code. + candidate = path.with_suffix(".candidate.pt") + candidate.parent.mkdir(parents=True, exist_ok=True) + torch.save({"outputs": outputs, "metadata": metadata}, candidate) + pytest.fail( + f"Missing golden {relative}; candidate saved to {candidate}. " + "Generate from the unchanged main revision, verify in two separate processes, " + "then review and publish the reference. Candidate output is not an approved golden.", + pytrace=False, + ) + golden = torch.load(path, map_location="cpu", weights_only=True) + assert golden["metadata"] == metadata, "Golden input/checkpoint/runtime identity changed; review the reference" + assert golden["outputs"].keys() == outputs.keys() + for key, value in outputs.items(): + torch.testing.assert_close(value, golden["outputs"][key], atol=0, rtol=0, msg=f"{name}: {key}") diff --git a/fastvideo/tests/golden_gate/_wan_checkpoint.py b/fastvideo/tests/golden_gate/_wan_checkpoint.py new file mode 100644 index 000000000..a034a580f --- /dev/null +++ b/fastvideo/tests/golden_gate/_wan_checkpoint.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Immutable Wan 1.3B inputs shared by its small numerical gates.""" + +import hashlib +from pathlib import Path + +from huggingface_hub import snapshot_download + +WAN_REPO = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +WAN_REVISION = "0fad780a534b6463e45facd96134c9f345acfa5b" + + +def component_path(component): + root = snapshot_download(WAN_REPO, revision=WAN_REVISION, allow_patterns=[f"{component}/*"]) + return Path(root) / component + + +def checkpoint_identity(path): + return { + "repo": WAN_REPO, + "revision": WAN_REVISION, + "component": path.name, + "config_sha256": hashlib.sha256((path / "config.json").read_bytes()).hexdigest(), + } diff --git a/fastvideo/tests/golden_gate/test_wan_causal.py b/fastvideo/tests/golden_gate/test_wan_causal.py new file mode 100644 index 000000000..9170bfb86 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_wan_causal.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Real Wan block weights through causal append, rewrite, and sink eviction. + +This is a cache/attention fingerprint, not a quality test of distilled weights. +Only block 0 is loaded. Five tiny forwards replace a full causal video render. +""" + +from dataclasses import replace + +import torch + +from fastvideo.tests.golden_gate._harness import DEFAULT_SEED, _component_dir, distributed_runtime, load_layer_state +from fastvideo.tests.golden_gate._tensor_golden import assert_tensor_golden, deterministic_forward +from fastvideo.tests.golden_gate._wan_checkpoint import WAN_REVISION +from fastvideo.tests.golden_gate.test_wan_t2v import INNER_DIM, SPEC + +__all__ = ["distributed_runtime"] + + +def causal_outputs(device): + from fastvideo.forward_context import set_forward_context + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + from fastvideo.models.loader.utils import set_default_torch_dtype + from fastvideo.models.wan.causal_transformer import CausalWanTransformerBlock + from fastvideo.platforms import AttentionBackendEnum + + spec = replace(SPEC, revision=WAN_REVISION) + # Attention selects its backend during construction, not .to(bfloat16). + with set_default_torch_dtype(torch.bfloat16): + block = CausalWanTransformerBlock( + INNER_DIM, 8960, 12, local_attn_size=3, sink_size=1, cross_attn_norm=True, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,), + ) + assert block.attn1.attn.backend == AttentionBackendEnum.FLASH_ATTN + assert block.attn2.attn.backend == AttentionBackendEnum.FLASH_ATTN + block.load_state_dict(load_layer_state(spec, _component_dir(spec)), strict=True) + block = block.to(device=device, dtype=torch.bfloat16).eval() + generator = torch.Generator(device="cpu").manual_seed(DEFAULT_SEED) + + def randn(*shape): + return torch.randn(*shape, generator=generator).to(device=device, dtype=torch.bfloat16) + + encoder = randn(1, 16, INNER_DIM) + kv = { + "k": torch.zeros(1, 12, 12, 128, device=device, dtype=torch.bfloat16), + "v": torch.zeros(1, 12, 12, 128, device=device, dtype=torch.bfloat16), + "global_end_index": torch.zeros(1, device=device, dtype=torch.long), + "local_end_index": torch.zeros(1, device=device, dtype=torch.long), + } + cross = {"k": torch.zeros_like(encoder).reshape(1, 16, 12, 128), + "v": torch.zeros_like(encoder).reshape(1, 16, 12, 128), "is_init": False} + cos, sin = get_rotary_pos_embed((4, 2, 2), INNER_DIM, 12, [44, 42, 42], + rope_theta=10000, dtype=torch.float64) + outputs = {} + starts = [0, 0, 4, 8, 12] # Rewrite denoised context, append, then evict while preserving the sink. + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + for index, start in enumerate(starts): + outputs[f"block_{index}"] = block( + hidden_states=randn(1, 4, INNER_DIM), encoder_hidden_states=encoder, + temb=randn(1, 1, 6, INNER_DIM), + freqs_cis=(cos[start:start + 4].to(device).float(), sin[start:start + 4].to(device).float()), + block_mask=None, kv_cache=kv, crossattn_cache=cross, current_start=start, frame_seqlen=4, + ) + outputs[f"cache_k_{index}"] = kv["k"].clone() + outputs[f"cache_v_{index}"] = kv["v"].clone() + outputs[f"indices_{index}"] = torch.cat([kv["global_end_index"], kv["local_end_index"]]).clone() + assert cross["is_init"] + return outputs, {"repo": spec.repo_id, "revision": spec.revision, "layer": 0, + "starts": starts, "local_attn_size": 3, "sink_size": 1, + "frame_seqlen": 4, "text_length": 16, "rope_cache_policy": "absolute"} + + +def test_wan_causal_golden_gate(distributed_runtime): + with deterministic_forward("FLASH_ATTN") as device: + outputs, identity = causal_outputs(device) + assert_tensor_golden("wan_causal", outputs, identity=identity, attention_backend="FLASH_ATTN") diff --git a/fastvideo/tests/golden_gate/test_wan_denoising.py b/fastvideo/tests/golden_gate/test_wan_denoising.py new file mode 100644 index 000000000..45d8ee710 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_wan_denoising.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Three UniPC steps with the real 1.3B DiT, fixed embeddings, and tiny latents. + +No tokenizer, text encoder, VAE, video encode, or reference-video download. The +per-step tensor reference catches stage/scheduler drift after the block gate. +""" + +import torch + +from fastvideo.tests.golden_gate._harness import DEFAULT_SEED, distributed_runtime +from fastvideo.tests.golden_gate._tensor_golden import assert_tensor_golden, deterministic_forward +from fastvideo.tests.golden_gate._wan_checkpoint import checkpoint_identity, component_path + +__all__ = ["distributed_runtime"] + + +def denoising_outputs(device): + from fastvideo.models.wan.pipeline_config import WanT2V480PConfig + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.loader.component_loader import TransformerLoader + from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler + from fastvideo.pipelines.basic.wan.stages.denoising import WanDenoisingStage + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + path = component_path("transformer") + args = FastVideoArgs( + model_path=str(path), pipeline_config=WanT2V480PConfig(), + num_gpus=1, hsdp_shard_dim=1, use_fsdp_inference=False, dit_cpu_offload=False, + dit_layerwise_offload=False, enable_torch_compile=False, inference_torch_compile=False, + ) + model = TransformerLoader().load(str(path), args) + args.model_loaded["transformer"] = True + scheduler = FlowUniPCMultistepScheduler(shift=3.0) + scheduler.set_timesteps(3, device=device) + generator = torch.Generator(device="cpu").manual_seed(DEFAULT_SEED) + latents = torch.randn(1, 16, 3, 4, 4, generator=generator).to(device) + prompt = torch.randn(1, 16, 4096, generator=generator).to(device=device, dtype=torch.bfloat16) + negative = torch.randn(1, 16, 4096, generator=generator).to(device=device, dtype=torch.bfloat16) + batch = ForwardBatch( + data_type="video", latents=latents, prompt_embeds=[prompt], negative_prompt_embeds=[negative], + timesteps=scheduler.timesteps.clone(), num_inference_steps=3, guidance_scale=5.0, + do_classifier_free_guidance=True, height=32, width=32, num_frames=9, + raw_latent_shape=tuple(latents.shape), save_video=False, return_trajectory_latents=True, + ) + with torch.inference_mode(): + result = WanDenoisingStage(model, scheduler).forward(batch, args) + return {"latents": result.latents, "trajectory": result.trajectory_latents, + "timesteps": result.trajectory_timesteps}, { + **checkpoint_identity(path), "latent_shape": list(latents.shape), + "prompt_shape": list(prompt.shape), "steps": 3, "guidance_scale": 5.0, + "scheduler": "FlowUniPCMultistepScheduler", "shift": 3.0, "precision": "bf16", + "cfg_gate_step": 1.0, + } + + +def test_wan_denoising_golden_gate(distributed_runtime, monkeypatch): + monkeypatch.setenv("FASTVIDEO_CFG_GATE_STEP", "1.0") + with deterministic_forward("FLASH_ATTN") as device: + outputs, identity = denoising_outputs(device) + assert_tensor_golden("wan_denoising", outputs, identity=identity, attention_backend="FLASH_ATTN") diff --git a/fastvideo/tests/golden_gate/test_wan_t2v.py b/fastvideo/tests/golden_gate/test_wan_t2v.py index 8e8b99bd0..4bea1ee25 100644 --- a/fastvideo/tests/golden_gate/test_wan_t2v.py +++ b/fastvideo/tests/golden_gate/test_wan_t2v.py @@ -14,6 +14,7 @@ from __future__ import annotations import torch from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate +from fastvideo.tests.golden_gate._wan_checkpoint import WAN_REVISION __all__ = ["distributed_runtime"] @@ -22,8 +23,8 @@ TEXT_LEN = 512 def _build_block(layer: int) -> torch.nn.Module: - from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig - from fastvideo.models.dits.wanvideo import WanTransformerBlock + from fastvideo.models.wan.config import WanVideoArchConfig + from fastvideo.models.wan.transformer import WanTransformerBlock arch = WanVideoArchConfig( num_attention_heads=12, @@ -94,6 +95,7 @@ SPEC = GateSpec( (r"\.norm2\.", ".self_attn_residual_norm.norm."), ), attention_backend="FLASH_ATTN", + revision=WAN_REVISION, ) diff --git a/fastvideo/tests/golden_gate/test_wan_vae.py b/fastvideo/tests/golden_gate/test_wan_vae.py new file mode 100644 index 000000000..2965c6e90 --- /dev/null +++ b/fastvideo/tests/golden_gate/test_wan_vae.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan codec golden: FP32 encode, BF16 decode, streaming, and request reset. + +One CUDA GPU, only the pinned VAE weights. Nine tiny video frames exercise the +special first frame and two temporal-cache chunks; no DiT or encoder download. +""" + +import torch + +from fastvideo.tests.golden_gate._harness import DEFAULT_SEED, distributed_runtime +from fastvideo.tests.golden_gate._tensor_golden import assert_tensor_golden, deterministic_forward +from fastvideo.tests.golden_gate._wan_checkpoint import checkpoint_identity, component_path + +__all__ = ["distributed_runtime"] + + +def vae_outputs(device): + from fastvideo.configs.pipelines import PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.loader.component_loader import VAELoader + from fastvideo.models.wan.vae_config import WanVAEConfig + from fastvideo.pipelines.stages.decoding import DecodingStage + + path = component_path("vae") + args = FastVideoArgs( + model_path=str(path), + pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision="fp32", vae_decode_precision="bf16"), + vae_cpu_offload=False, + ) + vae = VAELoader().load(str(path), args) + assert vae.use_feature_cache and not vae.use_tiling + generator = torch.Generator(device="cpu").manual_seed(DEFAULT_SEED) + video = torch.randn(1, 3, 9, 16, 24, generator=generator).clamp(-1, 1).to(device) + latent = torch.randn(1, 16, 3, 2, 3, generator=generator).to(device) + with torch.inference_mode(): + posterior = vae.encode(video) + repeated = vae.encode(video) + torch.testing.assert_close(repeated.mean, posterior.mean, atol=0, rtol=0) + torch.testing.assert_close(repeated.logvar, posterior.logvar, atol=0, rtol=0) + decoder = DecodingStage(vae) + decoded = decoder.decode(latent, args) + cache, chunks = None, [] + for index, chunk in enumerate(latent.split(1, dim=2)): + frames, cache = decoder.streaming_decode(chunk, args, cache=cache, is_first_chunk=index == 0) + chunks.append(frames) + streamed = torch.cat(chunks, dim=2) + torch.testing.assert_close(streamed, decoded, atol=0, rtol=0) + torch.testing.assert_close(decoder.decode(latent, args), decoded, atol=0, rtol=0) + return {"mean": posterior.mean, "logvar": posterior.logvar, "decoded": decoded, "streamed": streamed}, { + **checkpoint_identity(path), + "video_shape": list(video.shape), + "latent_shape": list(latent.shape), + "encode_precision": "fp32", + "decode_precision": "bf16", + "normalization": "DecodingStage", + } + + +def test_wan_vae_golden_gate(distributed_runtime): + with deterministic_forward() as device: + outputs, identity = vae_outputs(device) + assert_tensor_golden("wan_vae", outputs, identity=identity) diff --git a/fastvideo/tests/loader/test_wan_family_imports.py b/fastvideo/tests/loader/test_wan_family_imports.py index dda2c36ac..8ec1d747b 100644 --- a/fastvideo/tests/loader/test_wan_family_imports.py +++ b/fastvideo/tests/loader/test_wan_family_imports.py @@ -12,7 +12,7 @@ import sys import pytest -from fastvideo.models.wan import config, transformer, vae, vae_config +from fastvideo.models.wan import causal_transformer, config, transformer, vae, vae_config TRANSFORMER_EXPORTS = ( "EntryClass", @@ -48,6 +48,13 @@ def test_legacy_transformer_exports_are_the_canonical_objects(): assert legacy.EntryClass is transformer.WanTransformer3DModel +def test_legacy_causal_transformer_exports_are_the_canonical_objects(): + legacy = importlib.import_module("fastvideo.models.dits.causal_wanvideo") + for name in legacy.__all__: + assert getattr(legacy, name) is getattr(causal_transformer, name) + assert legacy.EntryClass is causal_transformer.CausalWanTransformer3DModel + + def test_legacy_and_aggregate_config_exports_are_the_canonical_objects(): legacy = importlib.import_module("fastvideo.configs.models.dits.wanvideo") aggregate = importlib.import_module("fastvideo.configs.models.dits") @@ -131,6 +138,7 @@ def test_downstream_wan_consumers_import(module_name): @pytest.mark.parametrize("module_name, architecture", [ ("transformer", "WanTransformer3DModel"), + ("causal_transformer", "CausalWanTransformer3DModel"), ("vae", "AutoencoderKLWan"), ]) def test_registry_discovers_and_loads_the_canonical_wan_class(module_name, architecture, caplog): @@ -151,6 +159,8 @@ def test_registry_discovers_and_loads_the_canonical_wan_class(module_name, archi @pytest.mark.parametrize("first_module", [ "fastvideo.models.wan.config", "fastvideo.models.wan.transformer", + "fastvideo.models.wan.causal_transformer", + "fastvideo.models.dits.causal_wanvideo", "fastvideo.configs.models.dits.wanvideo", "fastvideo.models.dits.wanvideo", "fastvideo.models.wan.vae_config", diff --git a/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py b/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py index e53715900..027baf204 100644 --- a/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py +++ b/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py @@ -95,8 +95,8 @@ def _load_torch_wan22_from_diffusers(checkpoint: Path, config_path: Path, *, dty import torch from safetensors.torch import load_file - from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig - from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.wan.config import WanVideoArchConfig, WanVideoConfig + from fastvideo.models.wan.transformer import WanTransformer3DModel from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict hf_config = json.loads(config_path.read_text()) diff --git a/fastvideo/tests/mlx/tiny_wan.py b/fastvideo/tests/mlx/tiny_wan.py index 797d7fe15..2770ca31f 100644 --- a/fastvideo/tests/mlx/tiny_wan.py +++ b/fastvideo/tests/mlx/tiny_wan.py @@ -17,7 +17,7 @@ os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29513") -from fastvideo.configs.models.dits.wanvideo import ( # noqa: E402 +from fastvideo.models.wan.config import ( # noqa: E402 WanVideoArchConfig, WanVideoConfig, ) @@ -30,7 +30,7 @@ from fastvideo.mlx_runtime.fastwan import ( # noqa: E402 mlx_block_weights_from_torch, quantize_matrix, ) -from fastvideo.models.dits.wanvideo import WanTransformer3DModel # noqa: E402 +from fastvideo.models.wan.transformer import WanTransformer3DModel # noqa: E402 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch # noqa: E402 SEED = 2026 diff --git a/fastvideo/tests/modal/test_pr_test.py b/fastvideo/tests/modal/test_pr_test.py index f2360bfcc..76100a77d 100644 --- a/fastvideo/tests/modal/test_pr_test.py +++ b/fastvideo/tests/modal/test_pr_test.py @@ -367,7 +367,7 @@ def test_wave1_lane_functions_use_shared_scripts(monkeypatch): for script, payload in { "kernel_tests.sh": "pytest fastvideo-kernel/tests/ -vs", "inference_vmoba.sh": "python fastvideo/tests/inference/vmoba/test_vmoba_inference.py", - "golden_gate.sh": 'exec pytest "$golden_root" -vs', + "golden_gate.sh": 'exec pytest "$golden_root" -xvs', "encoder.sh": "pytest ./fastvideo/tests/encoders -vs", "vae.sh": "pytest ./fastvideo/tests/vaes -vs", "transformer.sh": "pytest ./fastvideo/tests/transformers -vs", diff --git a/fastvideo/tests/inference/test_cfg_gating.py b/fastvideo/tests/stages/_denoising_fixtures.py similarity index 64% rename from fastvideo/tests/inference/test_cfg_gating.py rename to fastvideo/tests/stages/_denoising_fixtures.py index 1f9f5e495..2cc589c99 100644 --- a/fastvideo/tests/inference/test_cfg_gating.py +++ b/fastvideo/tests/stages/_denoising_fixtures.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 +"""Weight-free fixtures for shared and Wan denoising contract tests.""" from contextlib import nullcontext from types import SimpleNamespace -import pytest import torch @@ -77,6 +77,7 @@ def _tiny_args(): pipeline_config=SimpleNamespace( embedded_cfg_scale=None, ti2v_task=False, + lucy_edit_task=False, dit_config=SimpleNamespace(boundary_ratio=None, patch_size=(1, 1, 1)), ), ) @@ -110,6 +111,7 @@ def _patch_denoising_module(monkeypatch, cfg_gate_step): expected_gate_step = float(cfg_gate_step) import fastvideo.pipelines.stages.denoising as denoising + from fastvideo.pipelines.stages.base import PipelineStage # envs.py evaluates FASTVIDEO_CFG_GATE_STEP lazily via __getattr__, so the # stage sees monkeypatched values without reloading the module. @@ -121,54 +123,56 @@ def _patch_denoising_module(monkeypatch, cfg_gate_step): monkeypatch.setattr(denoising, "get_world_group", lambda: SimpleNamespace(local_rank=0)) monkeypatch.setattr(denoising, "get_attn_backend", lambda **kwargs: object()) monkeypatch.setattr(denoising, "set_forward_context", lambda **kwargs: nullcontext()) + monkeypatch.setattr(PipelineStage, "device", property(lambda self: torch.device("cpu"))) return denoising, logger -def _run_stage(monkeypatch, cfg_gate_step): - denoising, logger = _patch_denoising_module(monkeypatch, cfg_gate_step) - model = TinyDenoiser() - stage = denoising.DenoisingStage(model, TinyScheduler()) - stage.progress_bar = lambda iterable=None, total=None: NullProgressBar() +class RecordingDenoiser(TinyDenoiser): - result = stage.forward(_tiny_batch(), _tiny_args()) - return result.latents, model, logger + def __init__(self, offset=0.0): + super().__init__() + self.offset = offset + self.inputs = [] + + def forward(self, hidden_states, prompt_embeds, timestep, guidance=None, encoder_hidden_states_image=None): + self.inputs.append((hidden_states.clone(), timestep.clone())) + prompt = prompt_embeds[0].float().mean() + self.calls.append("uncond" if prompt < 0 else "cond") + time = timestep.float().reshape(hidden_states.shape[0], -1).mean(dim=1).reshape(-1, 1, 1, 1, 1) + return hidden_states[:, :2].float() * 0.125 + prompt * 0.25 + time * 0.0001 + self.offset -def _run_legacy_two_pass(): +class TinyVAE(torch.nn.Module): + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(())) + self.shift_factor = torch.tensor([0.25, -0.25]).reshape(1, 2, 1, 1, 1) + self.scaling_factor = torch.tensor([2.0, 0.5]).reshape(1, 2, 1, 1, 1) + self.encode_calls = 0 + + def encode(self, image): + self.encode_calls += 1 + return SimpleNamespace(mean=torch.full((1, 2, 1, 2, 4), 0.5, dtype=torch.float32)) + + +def _batch(steps=4, cfg=True): batch = _tiny_batch() - model = TinyDenoiser() - scheduler = TinyScheduler() - assert batch.latents is not None - assert batch.timesteps is not None - latents = batch.latents.clone() - - for timestep in batch.timesteps: - latent_model_input = scheduler.scale_model_input(latents.to(torch.bfloat16), timestep) - timestep_expand = timestep.repeat(latent_model_input.shape[0]) - noise_pred_text = model(latent_model_input, batch.prompt_embeds, timestep_expand) - noise_pred_uncond = model(latent_model_input, batch.negative_prompt_embeds, timestep_expand) - noise_pred = noise_pred_uncond + batch.guidance_scale * (noise_pred_text - noise_pred_uncond) - latents = scheduler.step(noise_pred, timestep, latents, return_dict=False)[0] - - return latents + batch.latents = torch.linspace(-0.5, 0.5, 48).reshape(1, 2, 3, 2, 4) + batch.num_inference_steps = steps + batch.timesteps = torch.linspace(1000, 1, steps) + batch.do_classifier_free_guidance = cfg + batch.num_frames, batch.height, batch.width = 9, 16, 32 + batch.raw_latent_shape = tuple(batch.latents.shape) + batch.return_trajectory_latents = True + return batch -@pytest.mark.parametrize("cfg_gate_step", [None, "1.0"]) -def test_cfg_gating_default_off_matches_legacy_two_pass(monkeypatch, cfg_gate_step): - out, model, logger = _run_stage(monkeypatch, cfg_gate_step) - legacy_out = _run_legacy_two_pass() - - assert torch.equal(out, legacy_out) - assert model.calls == ["cond", "uncond"] * 4 - assert not any("CFG gating enabled" in msg for msg in logger.infos) - assert any("gate_step=-1/4" in msg and "reused=0" in msg for msg in logger.infos) - - -def test_cfg_gating_reuses_cached_delta_after_gate(monkeypatch): - out, model, logger = _run_stage(monkeypatch, "0.5") - legacy_out = _run_legacy_two_pass() - - assert model.calls == ["cond", "uncond", "cond", "uncond", "cond", "cond"] - assert any("CFG gating enabled: fraction=0.500, gate_step=2/4" in msg for msg in logger.infos) - assert any("fresh_uncond=2 reused=2 invalidations=0" in msg for msg in logger.infos) - assert torch.allclose(out, legacy_out, atol=1e-3, rtol=0.0) +def _args(): + args = _tiny_args() + args.vae_cpu_offload = True + arch = SimpleNamespace(patch_size=(1, 2, 2)) + args.pipeline_config.dit_config.arch_config = arch + args.pipeline_config.vae_config = SimpleNamespace( + arch_config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)) + return args diff --git a/fastvideo/tests/stages/test_cfg_gating.py b/fastvideo/tests/stages/test_cfg_gating.py new file mode 100644 index 000000000..9e20d0a30 --- /dev/null +++ b/fastvideo/tests/stages/test_cfg_gating.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared CFG regression tests; no weights or CUDA execution required.""" + +import pytest +import torch + +from fastvideo.tests.stages._denoising_fixtures import ( + NullProgressBar, TinyDenoiser, TinyScheduler, _patch_denoising_module, _tiny_args, _tiny_batch, +) + + +def _run_stage(monkeypatch, cfg_gate_step): + denoising, logger = _patch_denoising_module(monkeypatch, cfg_gate_step) + model = TinyDenoiser() + stage = denoising.DenoisingStage(model, TinyScheduler()) + stage.progress_bar = lambda iterable=None, total=None: NullProgressBar() + + result = stage.forward(_tiny_batch(), _tiny_args()) + return result.latents, model, logger + + +def _run_legacy_two_pass(): + batch = _tiny_batch() + model = TinyDenoiser() + scheduler = TinyScheduler() + assert batch.latents is not None + assert batch.timesteps is not None + latents = batch.latents.clone() + + for timestep in batch.timesteps: + latent_model_input = scheduler.scale_model_input(latents.to(torch.bfloat16), timestep) + timestep_expand = timestep.repeat(latent_model_input.shape[0]) + noise_pred_text = model(latent_model_input, batch.prompt_embeds, timestep_expand) + noise_pred_uncond = model(latent_model_input, batch.negative_prompt_embeds, timestep_expand) + noise_pred = noise_pred_uncond + batch.guidance_scale * (noise_pred_text - noise_pred_uncond) + latents = scheduler.step(noise_pred, timestep, latents, return_dict=False)[0] + + return latents + + +@pytest.mark.parametrize("cfg_gate_step", [None, "1.0"]) +def test_cfg_gating_default_off_matches_legacy_two_pass(monkeypatch, cfg_gate_step): + out, model, logger = _run_stage(monkeypatch, cfg_gate_step) + legacy_out = _run_legacy_two_pass() + + assert torch.equal(out, legacy_out) + assert model.calls == ["cond", "uncond"] * 4 + assert not any("CFG gating enabled" in msg for msg in logger.infos) + assert any("gate_step=-1/4" in msg and "reused=0" in msg for msg in logger.infos) + + +def test_cfg_gating_reuses_cached_delta_after_gate(monkeypatch): + out, model, logger = _run_stage(monkeypatch, "0.5") + legacy_out = _run_legacy_two_pass() + + assert model.calls == ["cond", "uncond", "cond", "uncond", "cond", "cond"] + assert any("CFG gating enabled: fraction=0.500, gate_step=2/4" in msg for msg in logger.infos) + assert any("fresh_uncond=2 reused=2 invalidations=0" in msg for msg in logger.infos) + assert torch.allclose(out, legacy_out, atol=1e-3, rtol=0.0) diff --git a/fastvideo/tests/stages/test_wan_causal_denoising.py b/fastvideo/tests/stages/test_wan_causal_denoising.py new file mode 100644 index 000000000..4e00fec30 --- /dev/null +++ b/fastvideo/tests/stages/test_wan_causal_denoising.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Weight-free contracts for causal Wan scheduling and cache lifetime.""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import torch + +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler +from fastvideo.tests.stages._denoising_fixtures import NullProgressBar, _patch_denoising_module +from fastvideo.tests.stages._denoising_fixtures import _args, _batch + + +class TinyCausalDenoiser(torch.nn.Module): + hidden_size = 2 + num_attention_heads = 1 + attention_head_dim = 2 + local_attn_size = 3 + sink_size = 1 + + def __init__(self): + super().__init__() + self.blocks = [None, None] + self.config = SimpleNamespace(arch_config=SimpleNamespace( + num_frames_per_block=3, sliding_window_num_frames=6, patch_size=(1, 1, 1))) + self.calls = [] + + def forward(self, latents, prompts, timestep, *, kv_cache, crossattn_cache, current_start, start_frame): + self.calls.append((timestep.clone(), start_frame, kv_cache, crossattn_cache)) + return latents.float() * 0.125 + + +class RecordingUniPC(FlowUniPCMultistepScheduler): + + def __init__(self): + super().__init__(shift=3.0) + self.resets = 0 + + def set_timesteps(self, *args, **kwargs): + self.resets += 1 + return super().set_timesteps(*args, **kwargs) + + +def test_causal_standard_resets_scheduler_per_block_and_caches_per_request(monkeypatch): + _patch_denoising_module(monkeypatch, "1.0") + from fastvideo.pipelines.basic.wan.stages import causal_denoising + monkeypatch.setattr(causal_denoising, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(causal_denoising, "set_forward_context", lambda **kwargs: nullcontext()) + model, scheduler = TinyCausalDenoiser(), RecordingUniPC() + stage = causal_denoising.CausalDenoisingStage(model, scheduler) + stage.progress_bar = lambda **kwargs: NullProgressBar() + args = _args() + args.pipeline_config.text_encoder_configs = [SimpleNamespace(arch_config=SimpleNamespace(text_len=8))] + args.pipeline_config.context_noise = 0 + + outputs = [] + for _ in range(2): + batch = _batch(steps=50, cfg=False) + batch.latents = batch.latents.repeat(1, 1, 2, 1, 1) + outputs.append(stage.forward(batch, args).latents.clone()) + + torch.testing.assert_close(outputs[0], outputs[1], atol=0, rtol=0) + assert scheduler.resets == 4 # Two blocks per request, including a fresh multi-step history. + assert len(model.calls) == 2 * 2 * (50 + 1) # Each block also writes its clean context to the cache. + first_cache, first_cross = model.calls[0][2:] + assert all(call[2] is first_cache and call[3] is first_cross for call in model.calls[:102]) + assert model.calls[102][2] is not first_cache + assert model.calls[102][3] is not first_cross + assert first_cache[0]["k"].shape == (1, 3 * 8, 1, 2) + assert first_cross[0]["k"].shape == (1, 8, 1, 2) + assert [model.calls[i][1] for i in (0, 50, 51, 101)] == [0, 0, 3, 3] + assert model.calls[50][0].eq(0).all() and model.calls[101][0].eq(0).all() + + +def test_causal_samplers_share_cache_layout_not_sampling_inheritance(): + from fastvideo.pipelines.basic.wan.stages.causal_denoising import ( + CausalDMDDenosingStage, CausalDenoisingStage, WanCausalDenoisingBase, + ) + assert issubclass(CausalDenoisingStage, WanCausalDenoisingBase) + assert issubclass(CausalDMDDenosingStage, WanCausalDenoisingBase) + assert not issubclass(CausalDenoisingStage, CausalDMDDenosingStage) diff --git a/fastvideo/tests/stages/test_wan_denoising.py b/fastvideo/tests/stages/test_wan_denoising.py new file mode 100644 index 000000000..083e09e0d --- /dev/null +++ b/fastvideo/tests/stages/test_wan_denoising.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan sampling contracts with tiny CPU tensors and no model weights. + +The real UniPC scheduler exercises the complete 50-step loop. These tests cover +orchestration, not transformer numerics; component goldens cover real weights. +""" + +import pytest +import torch + +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler +from fastvideo.tests.stages._denoising_fixtures import ( + NullProgressBar, RecordingDenoiser, TinyScheduler, TinyVAE, _args, _batch, _patch_denoising_module, +) + + +def _stage(monkeypatch, model, scheduler, *, second=None, gate="1.0"): + _, logger = _patch_denoising_module(monkeypatch, gate) + from fastvideo.pipelines.basic.wan.stages import denoising + monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu")) + stage = denoising.WanDenoisingStage(model, scheduler, transformer_2=second) + stage.progress_bar = lambda **kwargs: NullProgressBar() + return stage, logger + + +@pytest.mark.parametrize("steps", [3, 50]) +@pytest.mark.parametrize("cfg", [False, True]) +def test_wan_all_scheduler_steps_match_explicit_loop(monkeypatch, steps, cfg): + batch = _batch(steps, cfg) + scheduler = FlowUniPCMultistepScheduler(shift=3.0) + scheduler.set_timesteps(steps, device="cpu") + batch.timesteps = scheduler.timesteps.clone() + reference_scheduler = FlowUniPCMultistepScheduler(shift=3.0) + reference_scheduler.set_timesteps(steps, device="cpu") + expected = batch.latents.clone() + reference = RecordingDenoiser() + trajectory = [] + for t in reference_scheduler.timesteps: + x = reference_scheduler.scale_model_input(expected.to(torch.bfloat16), t) + cond = reference(x, batch.prompt_embeds, t.repeat(x.shape[0])) + if cfg: + uncond = reference(x, batch.negative_prompt_embeds, t.repeat(x.shape[0])) + cond = uncond + batch.guidance_scale * (cond - uncond) + expected = reference_scheduler.step(cond, t, expected, return_dict=False)[0] + trajectory.append(expected) + + model = RecordingDenoiser() + stage, _ = _stage(monkeypatch, model, scheduler) + result = stage.forward(batch, _args()) + + torch.testing.assert_close(result.latents, expected, atol=0, rtol=0) + torch.testing.assert_close(result.trajectory_latents, torch.stack(trajectory, dim=1), atol=0, rtol=0) + torch.testing.assert_close(result.trajectory_timesteps, reference_scheduler.timesteps, atol=0, rtol=0) + assert model.calls == (["cond", "uncond"] if cfg else ["cond"]) * steps + assert scheduler.step_index == steps + + +def test_wan_expert_boundary_invalidates_cfg_cache(monkeypatch): + primary, secondary = RecordingDenoiser(), RecordingDenoiser(offset=0.125) + stage, logger = _stage(monkeypatch, primary, TinyScheduler(), second=secondary, gate="0.0") + args, batch = _args(), _batch() + args.pipeline_config.dit_config.boundary_ratio = 0.9 + batch.boundary_ratio = 0.5 # Per-request override wins; equality still uses the high-noise expert. + batch.timesteps = torch.tensor([750.0, 500.0, 250.0, 1.0]) + batch.guidance_scale_2 = 3.0 + stage.forward(batch, args) + assert primary.calls == ["cond", "uncond", "cond"] + assert secondary.calls == ["cond", "uncond", "cond"] + assert any("invalidations=1" in line and "fresh_uncond=2" in line for line in logger.infos) + + +@pytest.mark.parametrize("kind,channels", [("t2v", 2), ("i2v", 4), ("v2v", 6), ("lucy", 4), ("ti2v", 2)]) +def test_wan_conditioning_layout_and_first_frame(monkeypatch, kind, channels): + args, batch = _args(), _batch(cfg=False) + model, vae = RecordingDenoiser(), TinyVAE() + if kind == "i2v": + batch.image_latent = torch.full_like(batch.latents, 0.25) + if kind in {"v2v", "lucy"}: + batch.video_latent = torch.full_like(batch.latents, 0.375) + args.pipeline_config.lucy_edit_task = kind == "lucy" + args.pipeline_config.ti2v_task = kind == "ti2v" + if kind == "ti2v": + batch.pil_image = torch.zeros(1, 3, 1, 16, 32) + stage, _ = _stage(monkeypatch, model, TinyScheduler()) + from fastvideo.pipelines.basic.wan.stages import conditioning + monkeypatch.setattr(conditioning, "get_local_torch_device", lambda: torch.device("cpu")) + conditioning.WanFirstFrameEncodingStage(vae).forward(batch, args) + result = stage.forward(batch, args) + assert all(x.shape[1] == channels for x, _ in model.inputs) + assert result.latents.shape == (1, 2, 3, 2, 4) + if kind in {"lucy", "ti2v"}: + assert all(t.shape == (1, 6) for _, t in model.inputs) + else: + assert all(t.shape == (1,) for _, t in model.inputs) + if kind == "v2v": + assert all(torch.count_nonzero(x[:, 4:]) == 0 for x, _ in model.inputs) + if kind == "ti2v": + expected = (torch.full((1, 2, 1, 2, 4), 0.5) - vae.shift_factor) * vae.scaling_factor + torch.testing.assert_close(result.latents[:, :, :1], expected, atol=0, rtol=0) + assert vae.encode_calls == 1 + assert model.inputs[0][1][0, :2].eq(0).all() + else: + assert vae.encode_calls == 0 + + +def test_wan_request_does_not_reuse_previous_cfg_delta(monkeypatch): + model = RecordingDenoiser() + stage, _ = _stage(monkeypatch, model, TinyScheduler(), gate="0.0") + first = stage.forward(_batch(), _args()).latents.clone() + second = stage.forward(_batch(), _args()).latents.clone() + torch.testing.assert_close(first, second, atol=0, rtol=0) + assert model.calls == ["cond", "uncond", "cond", "cond", "cond"] * 2 diff --git a/fastvideo/tests/stages/test_wan_dmd_denoising.py b/fastvideo/tests/stages/test_wan_dmd_denoising.py new file mode 100644 index 000000000..7d92834d0 --- /dev/null +++ b/fastvideo/tests/stages/test_wan_dmd_denoising.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DMD scheduler ownership, layout, and RNG ordering without checkpoints.""" + +from contextlib import nullcontext + +import torch + +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.tests.stages._denoising_fixtures import NullProgressBar, _patch_denoising_module +from fastvideo.tests.stages._denoising_fixtures import RecordingDenoiser, _args, _batch + + +def test_wan_dmd_uses_full_training_table_and_preserves_rng_order(monkeypatch): + _patch_denoising_module(monkeypatch, "1.0") + from fastvideo.pipelines.basic.wan.stages import dmd + monkeypatch.setattr(dmd, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(dmd, "set_forward_context", lambda **kwargs: nullcontext()) + args, batch = _args(), _batch(steps=3, cfg=False) + args.pipeline_config.dmd_denoising_steps = [1000, 750, 500] + batch.latents = batch.latents.permute(0, 2, 1, 3, 4) + batch.generator = [torch.Generator().manual_seed(123)] + expected_generator = torch.Generator().manual_seed(123) + expected = batch.latents.clone() + model, reference = RecordingDenoiser(), RecordingDenoiser() + scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) + expected_scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0) + table = scheduler.timesteps.clone() + sigmas = scheduler.sigmas.clone() + timesteps = args.pipeline_config.dmd_denoising_steps + for index, timestep in enumerate(timesteps): + t = torch.tensor([timestep], dtype=torch.long) + noise = reference(expected.to(torch.bfloat16).permute(0, 2, 1, 3, 4), batch.prompt_embeds, t) + noise = noise.permute(0, 2, 1, 3, 4) + clean = pred_noise_to_pred_video(noise.flatten(0, 1), expected.flatten(0, 1), t, expected_scheduler) + clean = clean.unflatten(0, noise.shape[:2]) + if index + 1 < len(timesteps): + random = torch.randn(expected.shape, generator=expected_generator, dtype=clean.dtype) + expected = expected_scheduler.add_noise( + clean.flatten(0, 1), random.flatten(0, 1), torch.tensor([timesteps[index + 1]], dtype=torch.long), + ).unflatten(0, noise.shape[:2]) + else: + expected = clean + + stage = dmd.DmdDenoisingStage(model, scheduler) + assert stage.scheduler is scheduler + stage.progress_bar = lambda **kwargs: NullProgressBar() + result = stage.forward(batch, args) + torch.testing.assert_close(result.latents, expected.permute(0, 2, 1, 3, 4), atol=0, rtol=0) + assert torch.equal(batch.generator[0].get_state(), expected_generator.get_state()) + assert torch.equal(scheduler.timesteps, table) and torch.equal(scheduler.sigmas, sigmas) + assert len(table) == 1000 + assert model.calls == ["cond"] * 3 + + +def test_legacy_sampling_exports_are_canonical(): + from fastvideo.pipelines import stages + from fastvideo.pipelines.basic.wan.stages import causal_denoising, dmd + from fastvideo.pipelines.stages import causal_denoising as legacy_causal + from fastvideo.pipelines.stages import denoising as legacy_dense + assert stages.DmdDenoisingStage is legacy_dense.DmdDenoisingStage is dmd.DmdDenoisingStage + for name in ("CausalDMDDenosingStage", "CausalDenoisingStage"): + assert getattr(stages, name) is getattr(legacy_causal, name) is getattr(causal_denoising, name) diff --git a/fastvideo/tests/stages/test_wan_pipeline_wiring.py b/fastvideo/tests/stages/test_wan_pipeline_wiring.py new file mode 100644 index 000000000..62d2040ab --- /dev/null +++ b/fastvideo/tests/stages/test_wan_pipeline_wiring.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run real pipeline assembly with weight-free modules, not source-text assertions.""" + +import importlib +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.tests.stages._denoising_fixtures import RecordingDenoiser, TinyVAE, _args, _batch, _patch_denoising_module + + +@pytest.mark.parametrize("module_name, sampler, first_frame", [ + ("wan.wan_pipeline", "WanDenoisingStage", True), + ("wan.wan_i2v_pipeline", "WanDenoisingStage", False), + ("wan.wan_v2v_pipeline", "WanDenoisingStage", False), + ("wan.lucy_edit_pipeline", "WanDenoisingStage", False), + ("wan.wan_dmd_pipeline", "DmdDenoisingStage", False), + ("wan.wan_i2v_dmd_pipeline", "DmdDenoisingStage", False), + ("wan.wan_causal_pipeline", "CausalDenoisingStage", False), + ("wan.wan_causal_dmd_pipeline", "CausalDMDDenosingStage", True), + ("dreamx_world.dreamx_world_pipeline", "WanDenoisingStage", True), + ("turbodiffusion.turbodiffusion_pipeline", "WanDenoisingStage", False), + ("turbodiffusion.turbodiffusion_i2v_pipeline", "WanDenoisingStage", False), +]) +def test_pipeline_wires_family_sampler_and_owns_scheduler(monkeypatch, module_name, sampler, first_frame): + _patch_denoising_module(monkeypatch, "1.0") + from fastvideo.pipelines.stages.base import PipelineStage + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + module = importlib.import_module(f"fastvideo.pipelines.basic.{module_name}") + + class PreparedStage: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + for name, value in list(vars(module).items()): + if isinstance(value, type) and issubclass(value, PipelineStage) and name not in {sampler, "WanFirstFrameEncodingStage"}: + monkeypatch.setattr(module, name, PreparedStage) + model = RecordingDenoiser() + model.blocks = [None] + model.config.arch_config = SimpleNamespace(num_frames_per_block=3, sliding_window_num_frames=21) + pipeline = module.EntryClass.__new__(module.EntryClass) + pipeline.modules = {"transformer": model, "transformer_2": None, "vae": TinyVAE(), + "scheduler": FlowMatchEulerDiscreteScheduler(shift=8.0)} + pipeline.get_module = lambda name, *default: pipeline.modules.get(name) + stages = {} + pipeline.add_stage = lambda *, stage_name, stage: stages.setdefault(stage_name, stage) + args = _args() + args.pipeline_config.flow_shift = 3.0 + pipeline.initialize_pipeline(args) + pipeline.create_pipeline_stages(args) + stage = stages["denoising_stage"] + assert type(stage).__name__ == sampler + assert type(stage).__module__.startswith("fastvideo.pipelines.basic.wan.stages.") + assert stage.vae is None + assert ("first_frame_encoding_stage" in stages) is first_frame + if first_frame: + assert list(stages).index("first_frame_encoding_stage") < list(stages).index("denoising_stage") + assert stages["first_frame_encoding_stage"].vae is pipeline.modules["vae"] + preparation = stages["latent_preparation_stage"].kwargs["scheduler"] + if sampler == "CausalDMDDenosingStage": + assert "timestep_preparation_stage" not in stages + else: + assert stages["timestep_preparation_stage"].kwargs["scheduler"] is preparation + if sampler == "DmdDenoisingStage": + assert stage.scheduler is not preparation + assert stage.scheduler.config.shift == 8.0 + preparation.set_timesteps(3) + assert len(stage.scheduler.timesteps) == 1000 + else: + assert stage.scheduler is preparation + + +@pytest.mark.parametrize("causal", [False, True]) +def test_first_frame_preparation_validates_and_clears_request_state(monkeypatch, causal): + from fastvideo.pipelines.basic.wan.stages import conditioning + monkeypatch.setattr(conditioning, "get_local_torch_device", lambda: torch.device("cpu")) + args, batch = _args(), _batch() + args.pipeline_config.ti2v_task = not causal + batch.pil_image = torch.zeros(1, 3, 1, 16, 32) + stage = conditioning.WanFirstFrameEncodingStage(TinyVAE(), causal=causal) + assert stage.verify_input(batch, args).is_valid() + stage.forward(batch, args) + assert stage.verify_output(batch, args).is_valid() + expected = (torch.full((1, 2, 1, 2, 4), 0.5) - stage.vae.shift_factor) * stage.vae.scaling_factor + torch.testing.assert_close(batch.first_frame_latent, expected, atol=0, rtol=0) + batch.pil_image = None + stage.forward(batch, args) + assert batch.first_frame_latent is None + assert stage.vae.encode_calls == 1 diff --git a/fastvideo/tests/training/distill/test_anyflow_pretrain.py b/fastvideo/tests/training/distill/test_anyflow_pretrain.py index b7239be28..ac6d4ecd7 100644 --- a/fastvideo/tests/training/distill/test_anyflow_pretrain.py +++ b/fastvideo/tests/training/distill/test_anyflow_pretrain.py @@ -67,7 +67,7 @@ def _make_embedder( deltatime_type: str = "r", init_seed: int = 0, ): - from fastvideo.models.dits.wanvideo import WanTimeTextImageEmbedding + from fastvideo.models.wan.transformer import WanTimeTextImageEmbedding emb = WanTimeTextImageEmbedding( dim=32, @@ -184,7 +184,7 @@ def test_wan_transformer_forward_signature_has_r_timestep() -> None: swallowed by **kwargs) so callers and type checkers can see it.""" import inspect - from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.wan.transformer import WanTransformer3DModel sig = inspect.signature(WanTransformer3DModel.forward) assert "r_timestep" in sig.parameters @@ -202,7 +202,7 @@ def test_wan_transformer_init_propagates_r_embedder_config() -> None: """ import inspect - from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.wan.transformer import WanTransformer3DModel src = inspect.getsource(WanTransformer3DModel.__init__) # All four arch config fields must be passed to WanTimeTextImageEmbedding. @@ -217,7 +217,7 @@ def test_wan_transformer_forward_threads_r_timestep_to_embedder() -> None: source inspection to avoid heavyweight distributed bring-up).""" import inspect - from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.wan.transformer import WanTransformer3DModel src = inspect.getsource(WanTransformer3DModel.forward) assert "r_timestep=r_timestep" in src, ("WanTransformer3DModel.forward must forward r_timestep into " diff --git a/fastvideo/tests/transformers/test_wanvideo.py b/fastvideo/tests/transformers/test_wanvideo.py index d3e068737..c8790c5f1 100644 --- a/fastvideo/tests/transformers/test_wanvideo.py +++ b/fastvideo/tests/transformers/test_wanvideo.py @@ -12,21 +12,18 @@ from fastvideo.configs.pipelines import PipelineConfig from fastvideo.forward_context import set_forward_context from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.loader.component_loader import TransformerLoader -from fastvideo.utils import maybe_download_model +from fastvideo.tests.golden_gate._wan_checkpoint import component_path from fastvideo.models.wan.config import WanVideoConfig from fastvideo.pipelines.pipeline_batch_info import ForwardBatch os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29503") -BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Wan transformer parity requires one CUDA GPU") @pytest.mark.usefixtures("distributed_setup") def test_wan_transformer(): - model_path = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) - transformer_path = os.path.join(model_path, "transformer") + transformer_path = str(component_path("transformer")) device = torch.device("cuda:0") precision = torch.bfloat16 args = FastVideoArgs( diff --git a/fastvideo/tests/vaes/test_wan_vae.py b/fastvideo/tests/vaes/test_wan_vae.py index a9fca444e..532060a99 100644 --- a/fastvideo/tests/vaes/test_wan_vae.py +++ b/fastvideo/tests/vaes/test_wan_vae.py @@ -12,19 +12,16 @@ from fastvideo.configs.pipelines import PipelineConfig from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.loader.component_loader import VAELoader from fastvideo.models.wan.vae_config import WanVAEConfig -from fastvideo.utils import maybe_download_model +from fastvideo.tests.golden_gate._wan_checkpoint import component_path os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29503") -BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Wan VAE parity requires one CUDA GPU") @pytest.mark.usefixtures("distributed_setup") def test_wan_vae(): - model_path = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) - vae_path = os.path.join(model_path, "vae") + vae_path = str(component_path("vae")) device = torch.device("cuda:0") precision = torch.float32 args = FastVideoArgs(model_path=vae_path, diff --git a/scripts/demo_anyflow_14b.py b/scripts/demo_anyflow_14b.py index a97f806aa..9671b5f4b 100644 --- a/scripts/demo_anyflow_14b.py +++ b/scripts/demo_anyflow_14b.py @@ -148,7 +148,7 @@ def encode_prompts(): def load_transformer(): banner("(2) Build FastVideo Wan-14B + load AnyFlow weights") from fastvideo.configs.models.dits import WanVideoConfig - from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.wan.transformer import WanTransformer3DModel cfg = WanVideoConfig() arch = cfg.arch_config diff --git a/scripts/validate_wan.sh b/scripts/validate_wan.sh new file mode 100644 index 000000000..e8ea8898f --- /dev/null +++ b/scripts/validate_wan.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Stop at the cheapest failing boundary. Run inside a prepared GPU environment. +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")/.." + +component=${1:-all} +level=${2:-golden} +case "$component" in all|vae|dense|causal) ;; *) echo 'component: all | vae | dense | causal' >&2; exit 2;; esac +case "$level" in golden|parity|default) ;; *) echo 'level: golden | parity | default' >&2; exit 2;; esac + +pytest fastvideo/tests/api/test_wan_definitions.py \ + fastvideo/tests/loader/test_wan_family_imports.py \ + fastvideo/tests/stages/test_cfg_gating.py fastvideo/tests/stages/test_wan_denoising.py \ + fastvideo/tests/stages/test_wan_dmd_denoising.py fastvideo/tests/stages/test_wan_causal_denoising.py \ + fastvideo/tests/stages/test_wan_pipeline_wiring.py -xq + +if [[ $component == all || $component == vae ]]; then + pytest fastvideo/tests/golden_gate/test_wan_vae.py -xq +fi +if [[ $component == all || $component == dense ]]; then + pytest fastvideo/tests/golden_gate/test_wan_t2v.py -xq + pytest fastvideo/tests/golden_gate/test_wan_denoising.py -xq +fi +if [[ $component == all || $component == causal ]]; then + pytest fastvideo/tests/golden_gate/test_wan_causal.py -xq +fi +[[ $level == golden ]] && exit 0 + +if [[ $component == all || $component == vae ]]; then + pytest fastvideo/tests/vaes/test_wan_vae.py -xq +fi +if [[ $component == all || $component == dense ]]; then + pytest fastvideo/tests/transformers/test_wanvideo.py -xq +fi +[[ $level == parity ]] && exit 0 + +# Preserve the selected model, prompt, backend, device and reference tier. The +# caller supplies matching SSIM environment settings; this command does not +# rewrite references or silently downgrade full-quality requests. +case "$component" in + vae|all) pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py fastvideo/tests/ssim/test_wan_i2v_similarity.py fastvideo/tests/ssim/test_causal_similarity.py -xq;; + dense) pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py fastvideo/tests/ssim/test_wan_i2v_similarity.py -xq;; + causal) pytest fastvideo/tests/ssim/test_causal_similarity.py -xq;; +esac diff --git a/scripts/verify_anyflow_fastvideo_parity.py b/scripts/verify_anyflow_fastvideo_parity.py index b000686c6..85a734bd8 100644 --- a/scripts/verify_anyflow_fastvideo_parity.py +++ b/scripts/verify_anyflow_fastvideo_parity.py @@ -107,7 +107,7 @@ def translate_keys( def build_fastvideo_transformer(): banner("(1) Build FastVideo WanTransformer3DModel + load AnyFlow weights") from fastvideo.configs.models.dits import WanVideoConfig - from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.wan.transformer import WanTransformer3DModel cfg = WanVideoConfig() arch = cfg.arch_config diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py index 5ab0f66b2..da8f8bf53 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py @@ -5,6 +5,7 @@ from types import SimpleNamespace import json import numpy as np +import pytest import torch from fastvideo.api.presets import get_preset @@ -29,6 +30,21 @@ from fastvideo.pipelines.stages.denoising import DenoisingStage from fastvideo.registry import get_default_preset, get_model_info, get_pipeline_config_cls_from_name +@pytest.fixture +def dreamx_model_manifests(monkeypatch): + """Registry/config checks do not need network access or model weights.""" + names = { + "GD-ML/DreamX-World-5B-Cam": "DreamXWorldPipeline", + "GD-ML/DreamX-World-5B": "DreamXWorldARPipeline", + } + + def manifest(model_path, revision=None): + name = names[model_path] + return {"_class_name": name, "pipeline_name": name, "_diffusers_version": "0.31.0"} + + monkeypatch.setattr("fastvideo.registry.maybe_download_model_index", manifest) + + def test_dreamx_world_5b_cam_pipeline_config_wires_first_scope_components(): config = DreamXWorld5BCamPipelineConfig() @@ -82,11 +98,11 @@ def test_dreamx_world_local_model_index_resolves_model_info(tmp_path): assert info.pipeline_config_cls is DreamXWorld5BCamPipelineConfig -def test_dreamx_world_model_path_resolves_pipeline_config(): +def test_dreamx_world_model_path_resolves_pipeline_config(dreamx_model_manifests): assert get_pipeline_config_cls_from_name("GD-ML/DreamX-World-5B-Cam") is DreamXWorld5BCamPipelineConfig -def test_dreamx_world_default_preset_is_registered(): +def test_dreamx_world_default_preset_is_registered(dreamx_model_manifests): preset_name = get_default_preset("GD-ML/DreamX-World-5B-Cam") preset = get_preset(preset_name, "dreamx_world") @@ -168,7 +184,7 @@ def test_dreamx_world_ar_pipeline_config_wires_components(): assert config.dit_config.arch_config.attn_compress == 4 -def test_dreamx_world_ar_pipeline_registry_and_preset(): +def test_dreamx_world_ar_pipeline_registry_and_preset(dreamx_model_manifests): from fastvideo.api.presets import get_preset from fastvideo.registry import get_pipeline_config_cls_from_name, get_preset_selection