fix(flux2): estimate working memory for denoise and both VAE directions

invokeai/app/invocations/flux2_denoise.py

@@ -49,8 +49,16 @@
 from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 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
 
+# FLUX.2 attention geometry. The head dim is 128 across every variant; the head count follows the
+# hidden size (Klein 4B: 24, Klein 9B: 32, FLUX.2 dev: 48). Only the head dim decides which SDPA
+# kernel is eligible; the head count scales the `math` fallback's score matrix, and since the
+# working-memory estimate is computed before the transformer is loaded, we use the largest.
+FLUX2_ATTENTION_HEAD_DIM = 128
+FLUX2_MAX_ATTENTION_HEADS = 48
+
 
 @invocation(
     "flux2_denoise",
@@ -458,10 +466,38 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                 bn_std=bn_std,
             )
 
+        # Estimate the peak activation memory the transformer forward will need and ask the model cache
+        # to keep that much VRAM free. Without this hint the cache reserves only the small default
+        # working memory and fills the rest of the card with the model, so anything beyond a plain
+        # low-resolution generation OOMs. Reference images are the dominant term: their latents are
+        # concatenated onto the image stream, so three 1024x1024 references quadruple the sequence
+        # (and with it the activation footprint) of a 1024x1024 generation.
+        ref_image_seq_len = ref_image_extension.ref_image_latents.shape[1] if ref_image_extension is not None else 0
+        # The additive bias is skipped entirely when reference images are present (see below), so the
+        # mask only costs anything -- storage, and possibly a materialized score matrix -- without them.
+        regional_attn_mask = regional_extension.restricted_attn_mask if ref_image_seq_len == 0 else None
+        estimated_working_memory = self._estimate_working_memory(
+            image_seq_len=packed_h * packed_w,
+            ref_image_seq_len=ref_image_seq_len,
+            text_seq_len=max(txt.shape[1], neg_txt.shape[1] if neg_txt is not None else 0),
+            num_loras=len(self.transformer.loras),
+            # The mask itself is already allocated; only the additive bias built per forward is new.
+            regional_attention_bias_bytes=(
+                regional_attn_mask.numel() * torch.empty((), dtype=inference_dtype).element_size()
+                if regional_attn_mask is not None
+                else 0
+            ),
+            has_regional_attention_mask=regional_attn_mask is not None,
+            device=device,
+            dtype=inference_dtype,
+        )
+
         with ExitStack() as exit_stack:
             # Load the transformer model
             (cached_weights, transformer) = exit_stack.enter_context(
-                context.models.load(self.transformer.transformer).model_on_device()
+                context.models.load(self.transformer.transformer).model_on_device(
+                    working_mem_bytes=estimated_working_memory
+                )
             )
             config = transformer_config
 
@@ -578,6 +614,67 @@ def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor)
         mask = mask.to(device=latents.device, dtype=latents.dtype)
         return mask.expand_as(latents)
 
+    def _estimate_working_memory(
+        self,
+        image_seq_len: int,
+        ref_image_seq_len: int,
+        text_seq_len: int,
+        num_loras: int,
+        regional_attention_bias_bytes: int = 0,
+        has_regional_attention_mask: bool = False,
+        device: torch.device | None = None,
+        dtype: torch.dtype = torch.bfloat16,
+    ) -> int:
+        """Estimate peak transformer activation memory (bytes) so the model cache reserves enough headroom.
+
+        FLUX.2 attention runs through SDPA without materializing the O(seq^2) score matrix, so the
+        activation footprint scales *linearly* with the total attended sequence -- text tokens, image
+        tokens, and reference-image tokens alike. Measured on the Klein 9B geometry in bf16 as peak
+        reserved memory, that slope is ~0.39 MB per token and holds from 1.5k to 28k tokens; it is
+        also independent of the block count (a no-grad forward frees each block's intermediates), so
+        the constant applies to both the 4B and 9B variants.
+
+        The reference-image term is what makes this estimate necessary rather than merely nice to
+        have: a 1024x1024 generation is 4096 image tokens (~1.7GB), but attaching three 1024x1024
+        references adds 12288 more for ~6.5GB, and a 1328px tile with three 1328px references reaches
+        ~10.9GB -- against a default ``device_working_mem_gb`` of 3.
+
+        A fixed base covers resolution-independent overhead (transient fp8/GGUF -> bf16 weight casts
+        during the forward, and allocator slack across many steps). LoRA sidecar patches add an extra
+        activation branch per patched layer, so we add a per-LoRA margin.
+
+        The linear model holds only while attention runs on a fused kernel. Regional prompting is
+        where that stops being a given: it hands the transformer a dense additive ``S x S`` bias,
+        which flash attention never accepts and which ROCm's memory-efficient kernel rejects as
+        well, leaving the ``math`` fallback and its materialized ``heads x S x S`` score matrix. The
+        device decides too -- MPS has no fused SDPA kernel at all -- and so does the diffusers
+        attention backend this build dispatches through. ``sdpa_score_matrix_bytes`` asks all three
+        and adds the score matrix only where it is really built: on CUDA with the stock backend the
+        memory-efficient kernel takes the bias and the term is zero (verified: peak stays linear
+        with the bias attached).
+        """
+        GB = 1024**3
+        MB = 1024**2
+        per_token_bytes = int(0.4 * MB)
+        total_seq_len = image_seq_len + ref_image_seq_len + text_seq_len
+        estimated = total_seq_len * per_token_bytes
+        estimated += int(1.0 * GB)
+        estimated += regional_attention_bias_bytes
+        estimated += sdpa_score_matrix_bytes(
+            device=device if device is not None else TorchDevice.choose_torch_device(),
+            dtype=dtype,
+            num_heads=FLUX2_MAX_ATTENTION_HEADS,
+            head_dim=FLUX2_ATTENTION_HEAD_DIM,
+            seq_len=total_seq_len,
+            has_attn_mask=has_regional_attention_mask,
+            # The FLUX.2 transformer's attention goes through diffusers' `dispatch_attention_fn`,
+            # which can route around torch's SDPA entirely -- including to a forced `math` backend.
+            via_diffusers_dispatch=True,
+        )
+        if num_loras > 0:
+            estimated += int(0.5 * num_loras * GB)
+        return estimated
+
     def _load_text_conditioning(
         self,
         context: InvocationContext,

invokeai/app/invocations/flux2_vae_decode.py

@@ -21,6 +21,7 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2
 
 
 @invocation(
@@ -49,7 +50,14 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima
         Input latents should already be in the correct space after BN denormalization
         was applied in the denoiser. The VAE expects (B, 32, H, W) format.
         """
-        with vae_info.model_on_device() as (_, vae):
+        # Decoding at FLUX.2 resolutions costs multiple GB of activations (~4.3GB at 1024x1024),
+        # far above the default working memory the cache would otherwise reserve. Tell it up front so
+        # it offloads enough of the (possibly still resident) transformer to leave room.
+        estimated_working_memory = estimate_vae_working_memory_flux2(
+            operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device
+        )
+
+        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             vae_dtype = next(iter(vae.parameters())).dtype
             # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
             # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would

invokeai/app/invocations/flux2_vae_encode.py

@@ -19,6 +19,7 @@
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2
 
 
 @invocation(
@@ -46,7 +47,16 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc
         The VAE encodes to 32-channel latent space.
         Output latents shape: (B, 32, H/8, W/8).
         """
-        with vae_info.model_on_device() as (_, vae):
+        # See the decode node: FLUX.2 VAE activations are multi-GB, so the cache needs the estimate to
+        # free room rather than discovering the shortfall as an OOM.
+        estimated_working_memory = estimate_vae_working_memory_flux2(
+            operation="encode",
+            image_tensor=image_tensor,
+            vae=vae_info.model,
+            device=TorchDevice.choose_torch_device(),
+        )
+
+        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             vae_dtype = next(iter(vae.parameters())).dtype
             device = TorchDevice.choose_torch_device()
             image_tensor = image_tensor.to(device=device, dtype=vae_dtype)

invokeai/backend/flux2/ref_image_extension.py

@@ -21,12 +21,16 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux2.sampling_utils import pack_flux2
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux2
 
 # Maximum pixel counts for reference images (matches BFL FLUX.2 sampling.py)
 # Single reference image: 2024² pixels, Multiple: 1024² pixels
 MAX_PIXELS_SINGLE_REF = 2024**2  # ~4.1M pixels
 MAX_PIXELS_MULTI_REF = 1024**2  # ~1M pixels
 
+# Tile size (in pixels) forced on the VAE for reference-image encoding, see _prepare_ref_images().
+REF_ENCODE_TILE_SIZE = 512
+
 
 def resize_image_to_max_pixels(image: Image.Image, max_pixels: int) -> Image.Image:
     """Resize image to fit within max_pixels while preserving aspect ratio.
@@ -203,8 +207,18 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]:
             image_tensor = image_tensor * 2.0 - 1.0
             image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension
 
-            # Encode using FLUX.2 VAE
-            with vae_info.model_on_device() as (_, vae):
+            # Encode using FLUX.2 VAE. The encode below forces REF_ENCODE_TILE_SIZE tiling, so the
+            # peak is bounded by one tile; tell the cache that up front so it frees the room instead
+            # of hitting the shortfall as an OOM.
+            estimated_working_memory = estimate_vae_working_memory_flux2(
+                operation="encode",
+                image_tensor=image_tensor,
+                vae=vae_info.model,
+                tile_size=REF_ENCODE_TILE_SIZE,
+                device=TorchDevice.choose_torch_device(),
+            )
+
+            with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
                 vae_dtype = next(iter(vae.parameters())).dtype
                 image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
 
@@ -219,8 +233,8 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]:
                 downsample = 2 ** (len(vae.config.block_out_channels) - 1)
                 prev_tiling = (vae.use_tiling, vae.tile_sample_min_size, vae.tile_latent_min_size)
                 vae.use_tiling = True
-                vae.tile_sample_min_size = 512
-                vae.tile_latent_min_size = 512 // downsample
+                vae.tile_sample_min_size = REF_ENCODE_TILE_SIZE
+                vae.tile_latent_min_size = REF_ENCODE_TILE_SIZE // downsample
                 try:
                     # FLUX.2 VAE uses diffusers API
                     latent_dist = vae.encode(image_tensor, return_dict=False)[0]