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

invokeai/app/invocations/flux2_denoise.py

@@ -39,6 +39,7 @@
     unpack_flux2,
 )
 from invokeai.backend.flux2.text_conditioning import Flux2TextConditioning
+from invokeai.backend.model_manager.configs.flux2_variant import flux2_hidden_size
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType
 from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec
 from invokeai.backend.patches.lora_conversions.flux_bfl_peft_lora_conversion_utils import (
@@ -49,8 +50,25 @@
 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 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
+# number that describes both. Only the head dim decides which SDPA kernel is eligible; the head
+# count scales the `math` fallback's score matrix.
+FLUX2_ATTENTION_HEAD_DIM = 128
+# The width the per-token constant below was measured on. Estimates scale off this.
+FLUX2_REFERENCE_HIDDEN_SIZE = 4096
+# Peak reserved activation bytes per attended token at the reference width. Measured slope between
+# 4608 and 9216 tokens: 0.3859 MB/tok on CUDA (RTX 4090, torch 2.7.1) and 0.4067 on ROCm (RX 9070
+# XT, gfx1201, torch 2.10). 0.42 is an upper bound on both -- it was 0.4, which the ROCm point
+# exceeds. The margin is deliberately small because this is the dominant term at every resolution.
+FLUX2_BYTES_PER_TOKEN_AT_REFERENCE_WIDTH = int(0.42 * 1024**2)
+# The widest variant, used when the config does not tell us which one this is -- over-reserving on
+# an unknown model beats under-reserving on the largest one.
+FLUX2_MAX_HIDDEN_SIZE = 6144
+
 
 @invocation(
     "flux2_denoise",
@@ -458,10 +476,47 @@ 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),
+            # Taken from `x`, not from `b`. `b` is the *noise* tensor's batch, which this node builds
+            # at 1 from width/height/seed even when the init latents carry more; the img2img preblend
+            # above then broadcasts the two, so `x` is the only thing that knows how many samples
+            # actually go through the transformer. Reference latents are repeated to match it
+            # (`ensure_batch_size` below), so they scale with it too.
+            batch_size=x.shape[0],
+            # Activation cost per token scales with the transformer's width, and [dev] is 1.5x
+            # Klein 9B. Fall back to the widest variant when the config cannot tell us.
+            hidden_size=flux2_hidden_size(getattr(transformer_config, "variant", None)) or FLUX2_MAX_HIDDEN_SIZE,
+            # 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 +633,94 @@ 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,
+        batch_size: int = 1,
+        hidden_size: int = FLUX2_REFERENCE_HIDDEN_SIZE,
+        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 holds from 1.5k to 28k tokens and is independent of the block
+        count (a no-grad forward frees each block's intermediates). It is *not* independent of the
+        build: 0.3859 MB/token on CUDA, 0.4067 on ROCm/gfx1201, so the constant is 0.42.
+
+        It is *not* independent of the transformer's width, which is why ``hidden_size`` is a
+        parameter rather than a constant. Measured slope between 4608 and 9216 tokens, block count
+        and everything else held fixed:
+
+            CUDA          3072 (Klein 4B) 0.291 MB/tok   4096 (Klein 9B) 0.386   6144 ([dev]) 0.555
+            ROCm/gfx1201  3072            0.315           4096            0.407   6144         0.589
+
+        which is 0.755 / 1.000 / 1.438 on CUDA and 0.774 / 1.000 / 1.448 on ROCm, against width
+        ratios of 0.75 / 1.00 / 1.50 -- linear in width on both,
+        and slightly sub-linear at the top so scaling by width stays an upper bound. Calibrating on
+        Klein 9B alone would have under-reserved [dev] by a third: 1024x1024 with three 1024x1024
+        references is 16896 tokens, ~7.6GB reserved against ~10GB needed on both platforms. The head
+        count follows the same width, so the score-matrix term gets the real one instead of the
+        widest.
+
+        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.
+
+        Batch multiplies the token count and nothing else. A batch of B is B independent sequences,
+        so it enters the linear term exactly as extra sequence does -- measured on the Klein geometry
+        with a reduced block count: 4608 tokens at B=1 peaks at 2570MB, the same 4608 at B=2 (9216
+        tokens) at 5126MB, and 9728 tokens at B=1 at 5584MB. Batch and sequence are interchangeable
+        to within the noise. Reference latents are repeated per sample (`ensure_batch_size`), so they
+        scale with it too, and the score matrix is shaped (batch, heads, S, S). The fixed base does
+        not scale -- it is about weights, not activations -- and neither does the regional bias, which
+        is built as (1, 1, S, S) and broadcast across the batch.
+
+        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, leaving whatever else the build has -- and if that is
+        the ``math`` fallback, a materialized ``heads x S x S`` score matrix. Whether a mask forces
+        that is build-specific and not worth predicting: CUDA's memory-efficient kernel takes it, and
+        so does ROCm's on gfx1100. 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 term is zero (verified: peak stays linear with the
+        bias attached).
+        """
+        GB = 1024**3
+        per_token_bytes = int(FLUX2_BYTES_PER_TOKEN_AT_REFERENCE_WIDTH * hidden_size / FLUX2_REFERENCE_HIDDEN_SIZE)
+        total_seq_len = image_seq_len + ref_image_seq_len + text_seq_len
+        estimated = total_seq_len * batch_size * 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=(hidden_size // FLUX2_ATTENTION_HEAD_DIM) * batch_size,
+            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:
+            # A sidecar branch is an activation, so it scales with the batch like the rest of them.
+            estimated += int(0.5 * num_loras * batch_size * 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

@@ -18,7 +18,7 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 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,9 +46,22 @@ 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.
+        # Use the VAE's intended compute device for both the probe and the encode, matching the
+        # decode node. `choose_torch_device()` disagrees for a cpu_only VAE, which would both ask the
+        # dispatch probe about the wrong device and push the tensors onto a device the weights are
+        # not on (see #9373, and the same note in the decode node).
+        device = vae_info.compute_device
+        estimated_working_memory = estimate_vae_working_memory_flux2(
+            operation="encode",
+            image_tensor=image_tensor,
+            vae=vae_info.model,
+            device=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)
 
             # Encode using diffusers API

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]

invokeai/backend/model_manager/configs/flux2_variant.py

@@ -60,3 +60,23 @@ def flux2_variant_from_vec_dim(dim: int) -> Flux2VariantType | None:
 def flux2_variant_from_hidden_size(dim: int) -> Flux2VariantType | None:
     """Return the distilled FLUX.2 variant for a transformer hidden_size, or ``None`` if unrecognized."""
     return _HIDDEN_SIZE_TO_VARIANT.get(dim)
+
+
+# Base variants share their distilled twin's architecture (see the module docstring); only the
+# weights differ, so they resolve to the same geometry.
+_BASE_TO_DISTILLED: dict[Flux2VariantType, Flux2VariantType] = {
+    Flux2VariantType.Klein4BBase: Flux2VariantType.Klein4B,
+    Flux2VariantType.Klein9BBase: Flux2VariantType.Klein9B,
+}
+
+
+def flux2_hidden_size(variant: Flux2VariantType | None) -> int | None:
+    """Return the transformer hidden size for a variant, or ``None`` if unrecognized.
+
+    The forward direction of :func:`flux2_variant_from_hidden_size`. Callers that size activations
+    need this: the per-token activation cost scales with the transformer's width, so Klein 4B, Klein
+    9B and [dev] (3072 / 4096 / 6144) do not cost the same per token.
+    """
+    if variant is None:
+        return None
+    return _HIDDEN_SIZE.get(_BASE_TO_DISTILLED.get(variant, variant))