[review & refactor] Wan-Animate-2 - #14413
Open
yiyixuxu wants to merge 3 commits into
Open
Conversation
Model (`transformer_wan_animate_2.py`): - Replace the `forward(*args, method=...)` dispatch and the split `forward_ref`/`forward_gen` with a single documented `forward(..., kv_cache_mode="extract"|"cached")` returning `Transformer2DModelOutput`, following the Flux2 KV-cache precedent. The `SelfAttention`/`CrossAttention` pre/post split becomes a regular `WanAnimate2Attention` (`AttentionModuleMixin`) with processors that run through `dispatch_attention_fn` - native SDPA by default, any backend via `set_attention_backend`; only the in-context generation path is pinned to `flex`, since its attention pattern is expressed as a `BlockMask`. The hard `flash_attn` requirement is gone. - `IncontextAttentionBlock` was a pure pass-through around `AttentionBlock`; merged into one `WanAnimate2TransformerBlock` (checkpoint keys lose the `.block.` segment, handled in the single-file mapping). - KV cache is a `WanAnimate2KVCache` object instead of bare dicts passed through `forward`. Accelerate hooks copy dict arguments, so the dict version breaks under `enable_model_cpu_offload` (the reference pass fills a copy and the generation pass KeyErrors); the object passes through by reference, and `_skip_keys = ["kv_cache"]` covers group offloading. - Remove all autocast in favour of the `transformer_wan.py` dtype discipline (fp32 modulation with `.type_as` casts at block boundaries), so the model runs natively in bf16 and is no longer CUDA-only in principle. Replace the local float64 `sinusoidal_embedding_1d` with the existing `Timesteps` class. - Remove dead code: the unreachable padding mask in the reference path (the pipeline always fills `seq_len` exactly), `init_weights`, `load_from_official_state_dict`, and the unused `window_size`/`qk_norm`/ `sparse_type`/`log_scale` config flags (`log_scale` ships as 0.0, making the flex `score_mod` a no-op). Pipeline: call sites updated to the merged forward, autocast wrappers replaced with explicit casts at the call boundary. Also: - Fill in `convert_wan_animate_2_transformer_to_diffusers` with the actual key mapping (block unwrap + attention renames); it was a prefix-strip no-op. - Revert the `pipeline_utils.py` try/except import shims - the stub exception classes silently break real `except OfflineModeIsEnabled` handling; upgrade `huggingface_hub` to the `setup.py` pin instead. - Drop the modular pipeline for now; it needs its own pass and is not part of the initial release surface. Numerics: the refactored model matches the reference implementation at 2.47e-05 max relative difference in fp32 over all 40 layers on real weights, and end-to-end outputs match the reference pipeline to a max pixel difference of 7e-5 (PSNR 119 dB) when kernels and environment are held fixed. The checkpoint key rename is a pure rename - all 1303 tensors bitwise identical. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
size/L
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
part of #14412