feat(flux2): run FLUX.2 attention on SageAttention when enabled

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

@@ -27,6 +27,7 @@ compared with PyTorch's default attention there:
 | Model (as measured) | 1024×1024 | 2048×2048 | Notes |
 | --- | --- | --- | --- |
 | 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 %) | |
 
 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/flux2/flux2_denoise.py

@@ -59,6 +59,7 @@
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import FLUXConditioningInfo
 from invokeai.backend.util.attention import sdpa_score_matrix_bytes
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.sage_attention import sage_attention_scope
 
 # FLUX.2 attention geometry. The head dim is 128 across every variant and the head count follows
 # the hidden size (Klein 4B: 3072/24, Klein 9B: 4096/32, [dev] 6144/48), so the width is the single
@@ -623,25 +624,26 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                     "Regional masks will be ignored for this generation."
                 )
 
-            x = denoise(
-                model=transformer,
-                img=x,
-                img_ids=img_ids,
-                txt=txt,
-                txt_ids=txt_ids,
-                timesteps=timesteps,
-                step_callback=self._build_step_callback(context),
-                guidance=self.guidance,
-                cfg_scale=cfg_scale_list,
-                neg_txt=neg_txt,
-                neg_txt_ids=neg_txt_ids,
-                scheduler=scheduler,
-                mu=mu,
-                inpaint_extension=inpaint_extension,
-                img_cond_seq=img_cond_seq,
-                img_cond_seq_ids=img_cond_seq_ids,
-                pos_joint_attention_kwargs=pos_joint_attention_kwargs,
-            )
+            with sage_attention_scope():
+                x = denoise(
+                    model=transformer,
+                    img=x,
+                    img_ids=img_ids,
+                    txt=txt,
+                    txt_ids=txt_ids,
+                    timesteps=timesteps,
+                    step_callback=self._build_step_callback(context),
+                    guidance=self.guidance,
+                    cfg_scale=cfg_scale_list,
+                    neg_txt=neg_txt,
+                    neg_txt_ids=neg_txt_ids,
+                    scheduler=scheduler,
+                    mu=mu,
+                    inpaint_extension=inpaint_extension,
+                    img_cond_seq=img_cond_seq,
+                    img_cond_seq_ids=img_cond_seq_ids,
+                    pos_joint_attention_kwargs=pos_joint_attention_kwargs,
+                )
 
         # Apply BN denormalization if BN stats are available
         # The diffusers Flux2KleinPipeline applies: latents = latents * bn_std + bn_mean