diff --git a/invokeai/app/invocations/cogview4_image_to_latents.py b/invokeai/app/invocations/cogview4_image_to_latents.py index 23f1c13e262..db44c6d220a 100644 --- a/invokeai/app/invocations/cogview4_image_to_latents.py +++ b/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) diff --git a/invokeai/app/invocations/flux_denoise.py b/invokeai/app/invocations/flux_denoise.py index db73326706d..35d095e2799 100644 --- a/invokeai/app/invocations/flux_denoise.py +++ b/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 diff --git a/invokeai/app/invocations/flux_vae_encode.py b/invokeai/app/invocations/flux_vae_encode.py index daf039b80d2..a99e39bc05f 100644 --- a/invokeai/app/invocations/flux_vae_encode.py +++ b/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) diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index 7508c0716d8..98116e2d8d4 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/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") diff --git a/invokeai/app/invocations/sd3_image_to_latents.py b/invokeai/app/invocations/sd3_image_to_latents.py index fc88e85aa56..abe37d195fc 100644 --- a/invokeai/app/invocations/sd3_image_to_latents.py +++ b/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) diff --git a/invokeai/backend/flux/extensions/kontext_extension.py b/invokeai/backend/flux/extensions/kontext_extension.py index 6aabcb6cdad..b58c670115b 100644 --- a/invokeai/backend/flux/extensions/kontext_extension.py +++ b/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) diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts index fbbde7b97a3..b47244e5fc2 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts @@ -156,17 +156,24 @@ export const buildFLUXGraph = async (arg: GraphBuilderArg): Promise 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'); }