WanVideo Set Attention Mode Override

This is a power-user surgical tool. Attention is the heaviest part of a Wan generation, and there are several implementations of it - the plain PyTorch one, FlashAttention, SageAttention and friends - trading speed against quality and against how fussy they are to install. Normally you pick one for the whole run. This node lets you override the attention backend for a specific range of steps and/or specific blocks, so you can run fast attention everywhere and swap to a safer one exactly where it causes problems.

The node description says it plainly: "Override the attention mode for the model for specific step and/or block range." You wire your Wan model through it and out comes the same model with the override patched in.

How it works

SageAttention is the community's go-to speedup - the KB flags it as a real accelerator, with the caveat that it's a pain to install on Windows (Triton) and can nibble at quality in places. Sometimes a particular block or the first couple of steps come out worse under an aggressive attention mode. Instead of giving up the speed for the whole run, you keep Sage globally and use this node to force, say, plain sdpa on the steps or blocks where Sage misbehaves. The node patches the model so that during the chosen window, attention runs in the mode you picked; everywhere else it uses whatever your main setting is.

This is not a node most people need. If you're not already chasing a specific artifact or squeezing performance, skip it - the default attention mode is fine.

The inputs and outputs that matter

• model (WANVIDEOMODEL) in, model out - insert it in your model chain before the sampler.
• attention_mode (default sdpa) - the backend to force in the window. Options run from sdpa (plain, always works) through flash_attn_2 / flash_attn_3, the sageattn family (sageattn, sageattn_3, radial_sage_attention, sageattn_compiled, sageattn_ultravico), and comfy (ComfyUI's own). The ones past sdpa/flash require those libraries actually installed.
• start_step (default 0) / end_step (default 10000) - the step window for the override. The high default end_step just means "to the end."
• blocks (optional INT) - restrict the override to specific transformer block(s) instead of all of them.
• verbose (default false) - log what it's doing, useful while dialing this in.

How to install it

ComfyUI Manager → search ComfyUI-WanVideoWrapper, install, restart. Manual:

cd ComfyUI/custom_nodes
git clone https://github.com/kijai/ComfyUI-WanVideoWrapper
pip install -r ComfyUI-WanVideoWrapper/requirements.txt

then restart. The node ships with the pack, but the attention backends don't - sdpa and comfy are always available, while the FlashAttention and SageAttention modes require their libraries installed in your environment. On Windows, SageAttention means wrangling Triton, which the KB (rightly) calls notoriously painful.

Common issues & troubleshooting

You pick a Sage/Flash mode and it errors. That backend isn't installed. sdpa always works; the fancy modes need their libraries present. If you haven't set up SageAttention/FlashAttention, those options will fail no matter what this node does.

No visible effect. Check your start_step/end_step and blocks - a window that doesn't overlap where the problem actually occurs does nothing. Turn on verbose to confirm the override is firing on the steps you think it is.

You're using this to fix quality but it's the wrong tool. If your whole output is soft under Sage, just set your global attention to sdpa and move on - this node is for localized overrides, not a substitute for choosing a sane default. Reach for it only when a specific block or step range is the culprit.