feat(z-image): run Z-Image attention on SageAttention when enabled

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

@@ -28,6 +28,7 @@ compared with PyTorch's default attention there:
 | --- | --- | --- | --- |
 | FLUX.1 dev (FP8) | 16.1 s → 14.0 s (−13 %) | 84.2 s → 54.8 s (−35 %) | Also serves ControlNets inside the denoising loop. |
 | FLUX.2 Klein 9B (FP8) | 15.6 s → 15.1 s (−3 %) | 27.1 s → 23.1 s (−15 %) | |
+| Z-Image (FP8) | 10.4 s → 10.0 s (−5 %) | 33.7 s → 25.9 s (−23 %) | Z-Image Control masks its attention and keeps PyTorch'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/z_image/z_image_denoise.py

@@ -39,6 +39,7 @@
 from invokeai.backend.util.attention import sdpa_score_matrix_bytes
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.fp8 import get_model_compute_dtype
+from invokeai.backend.util.sage_attention import sage_attention_scope
 from invokeai.backend.z_image.extensions.regional_prompting_extension import ZImageRegionalPromptingExtension
 from invokeai.backend.z_image.text_conditioning import ZImageTextConditioning
 from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter
@@ -702,6 +703,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                 )
             )
 
+            exit_stack.enter_context(sage_attention_scope())
+
             # Denoising loop - supports both built-in Euler and diffusers schedulers
             # Track user-facing step for progress (accounts for Heun's double steps)
             user_step = 0