feat(mm): improved VAE encode VRAM usage

invokeai/app/invocations/cogview4_image_to_latents.py

@@ -36,9 +36,19 @@ class CogView4ImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard):
     image: ImageField = InputField(description="The image to encode.")
     vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
 
+    def _estimate_working_memory(self, image_tensor: torch.Tensor, vae: AutoencoderKL) -> int:
+        """Estimate the working memory required by the invocation in bytes."""
+        # Encode operations use approximately 50% of the memory required for decode operations
+        h = image_tensor.shape[-2]
+        w = image_tensor.shape[-1]
+        element_size = next(vae.parameters()).element_size()
+        scaling_constant = 1100  # 50% of decode scaling constant (2200)
+        working_memory = h * w * element_size * scaling_constant
+        return int(working_memory)
+
     @staticmethod
-    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
-        with vae_info as vae:
+    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor, estimated_working_memory: int) -> torch.Tensor:
+        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoencoderKL)
 
             vae.disable_tiling()
@@ -62,7 +72,12 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
             image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
 
         vae_info = context.models.load(self.vae.vae)
-        latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+        assert isinstance(vae_info.model, AutoencoderKL)
+
+        estimated_working_memory = self._estimate_working_memory(image_tensor, vae_info.model)
+        latents = self.vae_encode(
+            vae_info=vae_info, image_tensor=image_tensor, estimated_working_memory=estimated_working_memory
+        )
 
         latents = latents.to("cpu")
         name = context.tensors.save(tensor=latents)

invokeai/app/invocations/flux_denoise.py

@@ -328,6 +328,21 @@ def _run_diffusion(
             cfg_scale_end_step=self.cfg_scale_end_step,
         )
 
+        kontext_extension = None
+        if self.kontext_conditioning:
+            if not self.controlnet_vae:
+                raise ValueError("A VAE (e.g., controlnet_vae) must be provided to use Kontext conditioning.")
+
+            kontext_extension = KontextExtension(
+                context=context,
+                kontext_conditioning=self.kontext_conditioning
+                if isinstance(self.kontext_conditioning, list)
+                else [self.kontext_conditioning],
+                vae_field=self.controlnet_vae,
+                device=TorchDevice.choose_torch_device(),
+                dtype=inference_dtype,
+            )
+
         with ExitStack() as exit_stack:
             # Prepare ControlNet extensions.
             # Note: We do this before loading the transformer model to minimize peak memory (see implementation).
@@ -385,21 +400,6 @@ def _run_diffusion(
                 dtype=inference_dtype,
             )
 
-            kontext_extension = None
-            if self.kontext_conditioning:
-                if not self.controlnet_vae:
-                    raise ValueError("A VAE (e.g., controlnet_vae) must be provided to use Kontext conditioning.")
-
-                kontext_extension = KontextExtension(
-                    context=context,
-                    kontext_conditioning=self.kontext_conditioning
-                    if isinstance(self.kontext_conditioning, list)
-                    else [self.kontext_conditioning],
-                    vae_field=self.controlnet_vae,
-                    device=TorchDevice.choose_torch_device(),
-                    dtype=inference_dtype,
-                )
-
             # Prepare Kontext conditioning if provided
             img_cond_seq = None
             img_cond_seq_ids = None

invokeai/app/invocations/flux_vae_encode.py

@@ -35,14 +35,24 @@ class FluxVaeEncodeInvocation(BaseInvocation):
         input=Input.Connection,
     )
 
+    def _estimate_working_memory(self, image_tensor: torch.Tensor, vae: AutoEncoder) -> int:
+        """Estimate the working memory required by the invocation in bytes."""
+        # Encode operations use approximately 50% of the memory required for decode operations
+        h = image_tensor.shape[-2]
+        w = image_tensor.shape[-1]
+        element_size = next(vae.parameters()).element_size()
+        scaling_constant = 1100  # 50% of decode scaling constant (2200)
+        working_memory = h * w * element_size * scaling_constant
+        return int(working_memory)
+
     @staticmethod
-    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
+    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor, estimated_working_memory: int) -> torch.Tensor:
         # TODO(ryand): Expose seed parameter at the invocation level.
         # TODO(ryand): Write a util function for generating random tensors that is consistent across devices / dtypes.
         # There's a starting point in get_noise(...), but it needs to be extracted and generalized. This function
         # should be used for VAE encode sampling.
         generator = torch.Generator(device=TorchDevice.choose_torch_device()).manual_seed(0)
-        with vae_info as vae:
+        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoEncoder)
             vae_dtype = next(iter(vae.parameters())).dtype
             image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
@@ -60,7 +70,10 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
             image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
 
         context.util.signal_progress("Running VAE")
-        latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+        estimated_working_memory = self._estimate_working_memory(image_tensor, vae_info.model)
+        latents = self.vae_encode(
+            vae_info=vae_info, image_tensor=image_tensor, estimated_working_memory=estimated_working_memory
+        )
 
         latents = latents.to("cpu")
         name = context.tensors.save(tensor=latents)

invokeai/app/invocations/image_to_latents.py

@@ -52,11 +52,48 @@ class ImageToLatentsInvocation(BaseInvocation):
     tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size)
     fp32: bool = InputField(default=False, description=FieldDescriptions.fp32)
 
+    def _estimate_working_memory(
+        self, image_tensor: torch.Tensor, use_tiling: bool, vae: AutoencoderKL | AutoencoderTiny
+    ) -> int:
+        """Estimate the working memory required by the invocation in bytes."""
+        # Encode operations use approximately 50% of the memory required for decode operations
+        element_size = 4 if self.fp32 else 2
+        scaling_constant = 1100  # 50% of decode scaling constant (2200)
+
+        if use_tiling:
+            tile_size = self.tile_size
+            if tile_size == 0:
+                tile_size = vae.tile_sample_min_size
+                assert isinstance(tile_size, int)
+            h = tile_size
+            w = tile_size
+            working_memory = h * w * element_size * scaling_constant
+
+            # We add 25% to the working memory estimate when tiling is enabled to account for factors like tile overlap
+            # and number of tiles. We could make this more precise in the future, but this should be good enough for
+            # most use cases.
+            working_memory = working_memory * 1.25
+        else:
+            h = image_tensor.shape[-2]
+            w = image_tensor.shape[-1]
+            working_memory = h * w * element_size * scaling_constant
+
+        if self.fp32:
+            # If we are running in FP32, then we should account for the likely increase in model size (~250MB).
+            working_memory += 250 * 2**20
+
+        return int(working_memory)
+
     @staticmethod
     def vae_encode(
-        vae_info: LoadedModel, upcast: bool, tiled: bool, image_tensor: torch.Tensor, tile_size: int = 0
+        vae_info: LoadedModel,
+        upcast: bool,
+        tiled: bool,
+        image_tensor: torch.Tensor,
+        tile_size: int = 0,
+        estimated_working_memory: int = 0,
     ) -> torch.Tensor:
-        with vae_info as vae:
+        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
             orig_dtype = vae.dtype
             if upcast:
@@ -113,14 +150,23 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
         image = context.images.get_pil(self.image.image_name)
 
         vae_info = context.models.load(self.vae.vae)
+        assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny))
 
         image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB"))
         if image_tensor.dim() == 3:
             image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
 
+        use_tiling = self.tiled or context.config.get().force_tiled_decode
+        estimated_working_memory = self._estimate_working_memory(image_tensor, use_tiling, vae_info.model)
+
         context.util.signal_progress("Running VAE encoder")
         latents = self.vae_encode(
-            vae_info=vae_info, upcast=self.fp32, tiled=self.tiled, image_tensor=image_tensor, tile_size=self.tile_size
+            vae_info=vae_info,
+            upcast=self.fp32,
+            tiled=self.tiled,
+            image_tensor=image_tensor,
+            tile_size=self.tile_size,
+            estimated_working_memory=estimated_working_memory,
         )
 
         latents = latents.to("cpu")

invokeai/app/invocations/sd3_image_to_latents.py

@@ -32,9 +32,19 @@ class SD3ImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard):
     image: ImageField = InputField(description="The image to encode")
     vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
 
+    def _estimate_working_memory(self, image_tensor: torch.Tensor, vae: AutoencoderKL) -> int:
+        """Estimate the working memory required by the invocation in bytes."""
+        # Encode operations use approximately 50% of the memory required for decode operations
+        h = image_tensor.shape[-2]
+        w = image_tensor.shape[-1]
+        element_size = next(vae.parameters()).element_size()
+        scaling_constant = 1100  # 50% of decode scaling constant (2200)
+        working_memory = h * w * element_size * scaling_constant
+        return int(working_memory)
+
     @staticmethod
-    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
-        with vae_info as vae:
+    def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor, estimated_working_memory: int) -> torch.Tensor:
+        with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoencoderKL)
 
             vae.disable_tiling()
@@ -58,7 +68,12 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
             image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
 
         vae_info = context.models.load(self.vae.vae)
-        latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+        assert isinstance(vae_info.model, AutoencoderKL)
+
+        estimated_working_memory = self._estimate_working_memory(image_tensor, vae_info.model)
+        latents = self.vae_encode(
+            vae_info=vae_info, image_tensor=image_tensor, estimated_working_memory=estimated_working_memory
+        )
 
         latents = latents.to("cpu")
         name = context.tensors.save(tensor=latents)

invokeai/backend/flux/extensions/kontext_extension.py

@@ -106,8 +106,8 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
 
         # Track cumulative dimensions for spatial tiling
         # These track the running extent of the virtual canvas in latent space
-        h = 0  # Running height extent
-        w = 0  # Running width extent
+        canvas_h = 0  # Running canvas height
+        canvas_w = 0  # Running canvas width
 
         vae_info = self._context.models.load(self._vae_field.vae)
 
@@ -131,12 +131,20 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
 
             # Continue with VAE encoding
             # Don't sample from the distribution for reference images - use the mean (matching ComfyUI)
-            with vae_info as vae:
+            # Estimate working memory for encode operation (50% of decode memory requirements)
+            img_h = image_tensor.shape[-2]
+            img_w = image_tensor.shape[-1]
+            element_size = next(vae_info.model.parameters()).element_size()
+            scaling_constant = 1100  # 50% of decode scaling constant (2200)
+            estimated_working_memory = int(img_h * img_w * element_size * scaling_constant)
+
+            with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
                 assert isinstance(vae, AutoEncoder)
                 vae_dtype = next(iter(vae.parameters())).dtype
                 image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
                 # Use sample=False to get the distribution mean without noise
                 kontext_latents_unpacked = vae.encode(image_tensor, sample=False)
+                TorchDevice.empty_cache()
 
             # Extract tensor dimensions
             batch_size, _, latent_height, latent_width = kontext_latents_unpacked.shape
@@ -154,21 +162,33 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
             kontext_latents_packed = pack(kontext_latents_unpacked).to(self._device, self._dtype)
 
             # Determine spatial offsets for this reference image
-            # - Compare the potential new canvas dimensions if we add the image vertically vs horizontally
-            # - Choose the placement that results in a more square-like canvas
             h_offset = 0
             w_offset = 0
 
             if idx > 0:  # First image starts at (0, 0)
-                # Check which placement would result in better canvas dimensions
-                # If adding to height would make the canvas taller than wide, tile horizontally
-                # Otherwise, tile vertically
-                if latent_height + h > latent_width + w:
+                # Calculate potential canvas dimensions for each tiling option
+                # Option 1: Tile vertically (below existing content)
+                potential_h_vertical = canvas_h + latent_height
+
+                # Option 2: Tile horizontally (to the right of existing content)
+                potential_w_horizontal = canvas_w + latent_width
+
+                # Choose arrangement that minimizes the maximum dimension
+                # This keeps the canvas closer to square, optimizing attention computation
+                if potential_h_vertical > potential_w_horizontal:
                     # Tile horizontally (to the right of existing images)
-                    w_offset = w
+                    w_offset = canvas_w
+                    canvas_w = canvas_w + latent_width
+                    canvas_h = max(canvas_h, latent_height)
                 else:
                     # Tile vertically (below existing images)
-                    h_offset = h
+                    h_offset = canvas_h
+                    canvas_h = canvas_h + latent_height
+                    canvas_w = max(canvas_w, latent_width)
+            else:
+                # First image - just set canvas dimensions
+                canvas_h = latent_height
+                canvas_w = latent_width
 
             # Generate IDs with both index offset and spatial offsets
             kontext_ids = generate_img_ids_with_offset(
@@ -182,11 +202,6 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
                 w_offset=w_offset,
             )
 
-            # Update cumulative dimensions
-            # Track the maximum extent of the virtual canvas after placing this image
-            h = max(h, latent_height + h_offset)
-            w = max(w, latent_width + w_offset)
-
             all_latents.append(kontext_latents_packed)
             all_ids.append(kontext_ids)
 

invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts

@@ -156,17 +156,24 @@ export const buildFLUXGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
       .filter((entity) => getGlobalReferenceImageWarnings(entity, model).length === 0);
 
     if (validFLUXKontextConfigs.length > 0) {
-      const kontextConcatenator = g.addNode({
-        id: getPrefixedId('flux_kontext_image_prep'),
-        type: 'flux_kontext_image_prep',
-        images: validFLUXKontextConfigs.map(({ config }) => zImageField.parse(config.image)),
+      const fluxKontextCollect = g.addNode({
+        type: 'collect',
+        id: getPrefixedId('flux_kontext_collect'),
       });
-      const kontextConditioning = g.addNode({
-        type: 'flux_kontext',
-        id: getPrefixedId('flux_kontext'),
-      });
-      g.addEdge(kontextConcatenator, 'image', kontextConditioning, 'image');
-      g.addEdge(kontextConditioning, 'kontext_cond', denoise, 'kontext_conditioning');
+      for (const { config } of validFLUXKontextConfigs) {
+        const kontextImagePrep = g.addNode({
+          id: getPrefixedId('flux_kontext_image_prep'),
+          type: 'flux_kontext_image_prep',
+          images: [zImageField.parse(config.image)],
+        });
+        const kontextConditioning = g.addNode({
+          type: 'flux_kontext',
+          id: getPrefixedId('flux_kontext'),
+        });
+        g.addEdge(kontextImagePrep, 'image', kontextConditioning, 'image');
+        g.addEdge(kontextConditioning, 'kontext_cond', fluxKontextCollect, 'item');
+      }
+      g.addEdge(fluxKontextCollect, 'collection', denoise, 'kontext_conditioning');
 
       g.upsertMetadata({ ref_images: [validFLUXKontextConfigs] }, 'merge');
     }