feat(wan): run Wan 2.2 attention on SageAttention when enabled

docs/src/content/docs/configuration/Optimization/sage-attention.mdx

@@ -36,6 +36,8 @@ compared with PyTorch's default attention there:
 | Video model (as measured) | Resolution | Generation | Notes |
 | --- | --- | --- | --- |
 | LTX-2.5 distilled (int8) | 1248×704, 121 frames | 72.0 s → 63.4 s (−12 %) | Video self-attention only; text cross-attention and the audio streams stay on PyTorch. The audio reads the video through cross-attention, so it changes too, but measured no more than with an exact kernel swap: no shift against the picture, same loudness. |
+| Wan 2.2 TI2V-5B (GGUF Q8) | 1280×704, 121 frames | 380.9 s → 247.9 s (−35 %) | |
+| Wan 2.2 T2V A14B (GGUF Q4_K_M, 4-step Lightning LoRA) | 832×480, 81 frames | Denoising 121.5 s → 60.9 s (−50 %) | Measured with `device_working_mem_gb: 10`. At 1280×720 a 24 GB GPU is already full with PyTorch's attention, and SageAttention made denoising slower: 293 s → 662 s. |
 
 At 1024×1024 most of a step is spent outside attention, so the gain is small; at high resolutions and for video,
 attention dominates. GGUF models spend much of each step dequantizing weights, which SageAttention does not touch.

invokeai/app/invocations/wan/wan_denoise.py

@@ -53,6 +53,7 @@
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import WanConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.sage_attention import sage_attention_scope
 from invokeai.backend.wan.memory_optimization import wan_memory_optimization
 from invokeai.backend.wan.sampling_utils import get_spatial_scale_factor, make_noise
 
@@ -658,6 +659,7 @@ def low_lora_factory() -> Iterable[PatchSpec]:
                 ),
             )
             exit_stack.callback(swapper.close)
+            exit_stack.enter_context(sage_attention_scope())
 
             for step_idx, t in enumerate(tqdm(timesteps, desc="Denoising (Wan 2.2)", total=total_steps)):
                 timestep = t.expand(latents.shape[0])

invokeai/app/invocations/wan/wan_video_denoise.py

@@ -42,6 +42,7 @@
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import WanConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.sage_attention import sage_attention_scope
 from invokeai.backend.wan.memory_optimization import wan_memory_optimization
 from invokeai.backend.wan.sampling_utils import (
     get_default_latent_channels,
@@ -311,6 +312,7 @@ def low_lora_factory() -> Iterable[PatchSpec]:
                 ),
             )
             exit_stack.callback(swapper.close)
+            exit_stack.enter_context(sage_attention_scope())
 
             for step_idx, t in enumerate(
                 tqdm(timesteps, desc=f"Denoising Wan 2.2 video ({self.num_frames} frames)", total=total_steps)