Skip to content

Commit 8e71a8a

Browse files
Pfannkuchensacklsteinclaude
authored
perf(qwen-image): add a tiling option to the Qwen-Image VAE nodes (#9427)
* feat(qwen-image): add a tiling option to the image-to-latents node The Qwen-Image i2l node hardcoded vae.disable_tiling(), so a full-frame encode was the only option. At 2560x1440 that peaks at 9.26 GiB — on top of a resident multi-GB transformer, which is what makes an upscale round-trip run out of headroom exactly at this node while every other node fits. Adds `tiled` / `tile_size` input fields following the SD/SDXL i2l node, OR'd with the global force_tiled_decode setting. Off by default, so behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter. Without it the change would be inert: the cache would keep reserving the full-frame figure (10.99 GiB at 2560x1440) and evict models to honour it, no matter what the VAE actually does. Tiled, it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Measured through the node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending; real images blend far better), which is why this stays opt-in. * Add test * feat(qwen-image): make VAE tiling usable on both Qwen-Image VAE nodes Both nodes reserve working memory for a full-frame operation, which at high resolutions exceeds a 24 GB card, so the model cache evicts everything else to honour it. On CUDA at 2560x1440: 19.91 GiB for the decode and 10.99 GiB for the encode. Tiling is the intended escape hatch, but it did not work on either node: - qwen_image_i2l hardcoded vae.disable_tiling(), so it could not be enabled. - qwen_image_l2i honoured the global force_tiled_decode, but computed its working-memory estimate before and independently of that flag. Tiling bounded the VAE while the cache still reserved the full-frame figure, so the memory was never freed for anything else — effectively inert. Adds `tiled` / `tile_size` input fields to both nodes following the SD/SDXL i2l/l2i nodes, OR'd with force_tiled_decode. Off by default; behaviour is unchanged unless enabled. estimate_vae_working_memory_qwen_image gains a matching tile_size parameter, and both nodes resolve tile_size=0 to the VAE default (256px) before estimating. Tiled it budgets one tile plus 25% overlap plus the resident RGB image, mirroring estimate_vae_working_memory_wan. Without this the change would be cosmetic on i2l and remain inert on l2i. Measured through the i2l node at 2560x1440: 10.99 -> 0.26 GiB reserved, 9.26 -> 0.17 GiB actual peak, identical latent shape. Verified across eight resolutions that tiled and untiled encodes produce the same latent dimensions. Tiled latents differ by ~1.4% relative L2 on noise input (worst case for tile blending), which is why this stays opt-in. Also fixes a crash in qwen_image_i2l: `width`/`height` are `int | None`, but the workflow UI sends 0 for an unset number input, and `0 is not None` reached `image.resize((0, 0))` -> "height and width must be > 0". Non-positive values are now treated as unset, matching how tile_size uses 0. * fix(qwen-image): pass a matched tile stride and scope the tiling state enable_tiling() was called with tile_sample_min_* only, leaving the stride at the module's 192px default. The tile loops step by stride but slice each accumulated tile to min, so any tile_size below 192 silently dropped whole bands of the image -- a 128px tile turned a 512x512 decode into 384x384 with no error -- while sizes above 256 grew every tile without removing any, making compute scale with tile_size^2 (8x a full frame at 512px). Pass all four parameters with the stock 4:3 ratio, rounding the stride down to a multiple of the 8x spatial compression so the pixel and latent steps agree. Tile sizes below 64px are clamped; the field carries the 0 "use default" sentinel and so cannot take a pydantic lower bound. enable_tiling() also writes straight onto the module, and disable_tiling() only clears use_tiling. That module is the model cache's own instance, so a tile size set once persisted for the lifetime of the cache entry and leaked across invocations and into anima_latents_to_image, which shares the instance. Apply the geometry through a context manager that restores it, and resolve the 0 sentinel against a constant instead of the module's current value. Also budget the pixel-space buffers tiled_decode holds simultaneously (~5 frames, not 1) -- the term that grows with output area, so it degraded in exactly the regime tiling exists for. * test(qwen-image): pin the multiple-of-8 tile-stride rounding The rounding in `_tile_stride_for` was load-bearing but uncovered: dropping it left the whole suite green, yet a raw 3/4 stride silently truncates the encode for any tile size whose 3/4 is not a multiple of 8. `tile_size` is `multiple_of=8`, so 72 and 80 are both reachable from the workflow UI and both land there (54 and 60). Cover it at the argument level (72 -> 48, 80 -> 56) and end-to-end against a real tiny VAE, plus a pin on the failure mode itself: min=72 with the un-rounded stride 54 turns a 512x512 encode into a 57x57 latent instead of 64x64, with no exception. Also correct the rationale on QWEN_IMAGE_VAE_MIN_TILE_SIZE. `_tile_stride_for` already floors the stride at 8, so the derived latent step never collapses to 0, and 8/24/32 all round to clean multiples of 8 and produce correctly sized output -- 64 is not the smallest valid tile. It is a cost floor: the tile count grows with the inverse square of the stride (1620 tiles at 64px versus 57,600 at 8px on a 2560x1440 frame). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
1 parent 77f00f1 commit 8e71a8a

7 files changed

Lines changed: 683 additions & 35 deletions

File tree

‎invokeai/app/invocations/qwen_image_image_to_latents.py‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,12 @@
1414
from invokeai.app.invocations.model import VAEField
1515
from invokeai.app.invocations.primitives import LatentsOutput
1616
from invokeai.app.services.shared.invocation_context import InvocationContext
17-
from invokeai.backend.krea2.vae_compat import as_qwen_image_vae
17+
from invokeai.backend.krea2.vae_compat import (
18+
QWEN_IMAGE_VAE_MIN_TILE_SIZE,
19+
as_qwen_image_vae,
20+
patch_qwen_image_vae_tiling,
21+
resolve_qwen_image_vae_tile_size,
22+
)
1823
from invokeai.backend.model_manager.load.load_base import LoadedModel
1924
from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
2025
from invokeai.backend.util.devices import TorchDevice
@@ -26,14 +31,23 @@
2631
title="Image to Latents - Qwen Image",
2732
tags=["image", "latents", "vae", "i2l", "qwen_image"],
2833
category="image",
29-
version="1.0.0",
34+
version="1.1.0",
3035
classification=Classification.Prototype,
3136
)
3237
class QwenImageImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard):
3338
"""Generates latents from an image using the Qwen Image VAE."""
3439

3540
image: ImageField = InputField(description="The image to encode.")
3641
vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
42+
tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
43+
# NOTE: tile_size = 0 is a special value meaning "use the model's default", matching the
44+
# SD/SDXL i2l node. `int | None` is avoided because the workflow UI does not handle it well.
45+
tile_size: int = InputField(
46+
default=0,
47+
multiple_of=8,
48+
description=f"{FieldDescriptions.vae_tile_size} Values between 1 and "
49+
f"{QWEN_IMAGE_VAE_MIN_TILE_SIZE} are raised to {QWEN_IMAGE_VAE_MIN_TILE_SIZE}.",
50+
)
3751
width: int | None = InputField(
3852
default=None,
3953
description="Resize the image to this width before encoding. If not set, encodes at the image's original size.",
@@ -44,24 +58,36 @@ class QwenImageImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard)
4458
)
4559

4660
@staticmethod
47-
def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor:
61+
def vae_encode(
62+
vae_info: LoadedModel, image_tensor: torch.Tensor, tiled: bool = False, tile_size: int = 0
63+
) -> torch.Tensor:
4864
# NOTE: vae_info.model may be an AutoencoderKLWan (a native-layout qwen_image_vae single file is
4965
# classified with the Anima base); it is reinterpreted as AutoencoderKLQwenImage inside the
5066
# model_on_device context below. The working-memory estimate only reads tensor shape + element
5167
# size, so it is safe to run on either class here.
68+
# Resolve tile_size=0 ("model default") before estimating, so the reserved working memory
69+
# matches the tiles the VAE will actually use. Resolved against a constant rather than the
70+
# module's current tile_sample_min_height, which a previous invocation may have overwritten.
71+
effective_tile_size = resolve_qwen_image_vae_tile_size(tile_size) if tiled else None
72+
5273
estimated_working_memory = estimate_vae_working_memory_qwen_image(
5374
operation="encode",
5475
image_tensor=image_tensor,
5576
vae=vae_info.model,
77+
tile_size=effective_tile_size,
5678
)
5779
with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
5880
# Reinterpret an Anima-classified Wan VAE as AutoencoderKLQwenImage (identical weights).
5981
vae = as_qwen_image_vae(vae)
6082

61-
vae.disable_tiling()
62-
6383
image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype)
64-
with torch.inference_mode():
84+
85+
# Tiling bounds the encode's peak memory to a single tile, which is what makes large
86+
# inputs (e.g. a 2560x1440 upscale round-trip) encodable while a multi-GB transformer
87+
# is still resident. Off by default: full-frame is faster and avoids tile blending.
88+
# The tiling state is scoped to this block: the VAE module belongs to the model cache
89+
# and is shared with later invocations (and with the Anima decode node).
90+
with torch.inference_mode(), patch_qwen_image_vae_tiling(vae, effective_tile_size):
6591
# The Qwen Image VAE expects 5D input: (B, C, num_frames, H, W)
6692
if image_tensor.dim() == 4:
6793
image_tensor = image_tensor.unsqueeze(2)
@@ -91,7 +117,13 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
91117

92118
# If target dimensions are specified, resize the image BEFORE encoding
93119
# (matching the diffusers pipeline which resizes in pixel space, not latent space).
94-
if self.width is not None and self.height is not None:
120+
#
121+
# `width`/`height` are `int | None`, but the workflow UI cannot represent None in a number
122+
# input and sends 0 for "unset" — which `is not None`, so a naive check reached
123+
# `resize((0, 0))` and raised "height and width must be > 0". Treat any non-positive value
124+
# as unset, which is also how `tile_size` uses 0. Note this means a half-filled pair (e.g.
125+
# width=1024, height=0) encodes at the original size rather than raising.
126+
if (self.width or 0) > 0 and (self.height or 0) > 0:
95127
image = image.convert("RGB").resize((self.width, self.height), resample=PILImage.LANCZOS)
96128

97129
# multiple_of=16 ensures the post-VAE latents (vae_scale_factor=8) have even
@@ -102,7 +134,12 @@ def invoke(self, context: InvocationContext) -> LatentsOutput:
102134

103135
vae_info = context.models.load(self.vae.vae)
104136

105-
latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
137+
latents = self.vae_encode(
138+
vae_info=vae_info,
139+
image_tensor=image_tensor,
140+
tiled=self.tiled or context.config.get().force_tiled_decode,
141+
tile_size=self.tile_size,
142+
)
106143

107144
latents = latents.to("cpu")
108145
name = context.tensors.save(tensor=latents)

‎invokeai/app/invocations/qwen_image_latents_to_image.py‎

Lines changed: 33 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
from contextlib import nullcontext
2-
31
import torch
42
from einops import rearrange
53
from PIL import Image
@@ -16,7 +14,12 @@
1614
from invokeai.app.invocations.model import VAEField
1715
from invokeai.app.invocations.primitives import ImageOutput
1816
from invokeai.app.services.shared.invocation_context import InvocationContext
19-
from invokeai.backend.krea2.vae_compat import as_qwen_image_vae
17+
from invokeai.backend.krea2.vae_compat import (
18+
QWEN_IMAGE_VAE_MIN_TILE_SIZE,
19+
as_qwen_image_vae,
20+
patch_qwen_image_vae_tiling,
21+
resolve_qwen_image_vae_tile_size,
22+
)
2023
from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt
2124
from invokeai.backend.util.devices import TorchDevice
2225
from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_qwen_image
@@ -27,20 +30,38 @@
2730
title="Latents to Image - Qwen Image",
2831
tags=["latents", "image", "vae", "l2i", "qwen_image"],
2932
category="latents",
30-
version="1.0.0",
33+
version="1.1.0",
3134
classification=Classification.Prototype,
3235
)
3336
class QwenImageLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard):
3437
"""Generates an image from latents using the Qwen Image VAE."""
3538

3639
latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection)
3740
vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection)
41+
tiled: bool = InputField(default=False, description=FieldDescriptions.tiled)
42+
# NOTE: tile_size = 0 is a special value meaning "use the model's default", matching the
43+
# SD/SDXL l2i node. `int | None` is avoided because the workflow UI does not handle it well.
44+
tile_size: int = InputField(
45+
default=0,
46+
multiple_of=8,
47+
description=f"{FieldDescriptions.vae_tile_size} Values between 1 and "
48+
f"{QWEN_IMAGE_VAE_MIN_TILE_SIZE} are raised to {QWEN_IMAGE_VAE_MIN_TILE_SIZE}.",
49+
)
3850

3951
@torch.no_grad()
4052
def invoke(self, context: InvocationContext) -> ImageOutput:
4153
latents = context.tensors.load(self.latents.latents_name)
4254

4355
vae_info = context.models.load(self.vae.vae)
56+
tiled = self.tiled or context.config.get().force_tiled_decode
57+
# Resolve tile_size=0 ("model default") before estimating, so the memory the cache reserves
58+
# matches the tiles the VAE will actually use. Without this the estimate stays at the
59+
# full-frame figure (~21 GB at 2560x1440 on CUDA) and tiling frees nothing: the VAE is
60+
# bounded, but the cache still evicts other models to honour the reservation.
61+
# Resolved against a constant rather than the module's current tile_sample_min_height, which a
62+
# previous invocation (including the Anima decode node, which shares this VAE instance) may
63+
# have overwritten.
64+
effective_tile_size = resolve_qwen_image_vae_tile_size(self.tile_size) if tiled else None
4465
# NOTE: vae_info.model may be an AutoencoderKLWan (a native-layout qwen_image_vae single file is
4566
# classified with the Anima base); it is reinterpreted as AutoencoderKLQwenImage inside the
4667
# model_on_device context below. The working-memory estimate only reads tensor shape + element
@@ -49,6 +70,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
4970
operation="decode",
5071
image_tensor=latents,
5172
vae=vae_info.model,
73+
tile_size=effective_tile_size,
5274
)
5375
with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
5476
context.util.signal_progress("Running VAE")
@@ -62,17 +84,13 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
6284
# which would wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
6385
latents = latents.to(device=vae_info.compute_device, dtype=vae.dtype)
6486

65-
# Honor the global force_tiled_decode setting, like the SD/SDXL l2i node. Tiling bounds the
66-
# VAE's per-tile memory, which is the scalable way to decode very large outputs that would
67-
# exceed VRAM even after offloading the transformer/text encoder. For normal sizes, leave
68-
# it off (faster, no tile blending) — the reserved working memory offloads other models so
69-
# the full-frame decode fits.
70-
if context.config.get().force_tiled_decode:
71-
vae.enable_tiling()
72-
else:
73-
vae.disable_tiling()
74-
75-
tiling_context = nullcontext()
87+
# Tiling bounds the VAE's per-tile memory, which is the scalable way to decode very
88+
# large outputs that would exceed VRAM even after offloading the transformer/text
89+
# encoder. For normal sizes, leave it off (faster, no tile blending) — the reserved
90+
# working memory offloads other models so the full-frame decode fits.
91+
# The tiling state is scoped to this block: the VAE module belongs to the model cache
92+
# and is shared with later invocations (and with the Anima decode node).
93+
tiling_context = patch_qwen_image_vae_tiling(vae, effective_tile_size)
7694

7795
TorchDevice.empty_cache()
7896

‎invokeai/backend/krea2/vae_compat.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,13 @@
66
share the exact same diffusers state-dict (identical keys and shapes), so a Wan-loaded VAE can be
77
used through the same encode/decode path without rebuilding it. Both default configs carry the same
88
Qwen-Image ``latents_mean`` / ``latents_std`` / ``z_dim`` values read by the Qwen encode/decode nodes.
9+
10+
Also holds the tiling helpers those nodes share, since the tile geometry has to be applied identically
11+
on both classes and restored afterwards — see ``patch_qwen_image_vae_tiling``.
912
"""
1013

14+
from collections.abc import Iterator
15+
from contextlib import contextmanager
1116
from typing import Any
1217

1318
from diffusers.models.autoencoders import AutoencoderKLWan
@@ -55,3 +60,92 @@ def as_qwen_image_vae(model: Any) -> QwenImageCompatibleVAE:
5560
)
5661

5762
return model
63+
64+
65+
# The stock AutoencoderKLQwenImage tile geometry: 256px tiles advancing in 192px steps, i.e. a 3/4
66+
# stride ratio with a 64px blend band. Both nodes resolve tile_size=0 to QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE
67+
# rather than reading the module's current value, which another invocation may have overwritten.
68+
QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE = 256
69+
_QWEN_IMAGE_VAE_TILE_STRIDE_NUMERATOR = 3
70+
_QWEN_IMAGE_VAE_TILE_STRIDE_DENOMINATOR = 4
71+
72+
# A cost floor, not a correctness one: `_tile_stride_for` keeps the geometry valid all the way down
73+
# (smaller tiles decode and encode to the right size), but the tile *count* grows with the inverse
74+
# square of the stride. At 2560x1440 a 64px tile already emits 1620 tiles; an 8px tile would emit
75+
# 57,600, and the per-tile Python/kernel-launch overhead dominates long before that. Tiles this small
76+
# also blend badly, so the field's low end is clamped rather than honoured literally.
77+
QWEN_IMAGE_VAE_MIN_TILE_SIZE = 64
78+
79+
80+
def resolve_qwen_image_vae_tile_size(tile_size: int) -> int:
81+
"""Resolve a node's ``tile_size`` field to the tile size the VAE will actually use.
82+
83+
``tile_size <= 0`` is the nodes' "use the default" sentinel (the workflow UI cannot represent
84+
``None`` in a number input and sends 0). Values below ``QWEN_IMAGE_VAE_MIN_TILE_SIZE`` are clamped
85+
rather than rejected, because the field also has to accept the 0 sentinel and so cannot carry a
86+
pydantic lower bound. The clamp is about cost, not validity -- see the constant.
87+
"""
88+
if tile_size <= 0:
89+
return QWEN_IMAGE_VAE_DEFAULT_TILE_SIZE
90+
return max(tile_size, QWEN_IMAGE_VAE_MIN_TILE_SIZE)
91+
92+
93+
def _tile_stride_for(tile_size: int) -> int:
94+
"""Return the tile stride to pair with ``tile_size``, keeping the stock 3/4 ratio.
95+
96+
Rounded down to a multiple of the VAE's 8x spatial compression: ``tiled_encode``/``tiled_decode``
97+
step the tile loop in one space (pixels for encode, latents for decode) while slicing the
98+
accumulated tile in the other, so the pixel stride must be exactly 8x the latent stride or the
99+
two disagree and the output is misaligned.
100+
"""
101+
stride = tile_size * _QWEN_IMAGE_VAE_TILE_STRIDE_NUMERATOR // _QWEN_IMAGE_VAE_TILE_STRIDE_DENOMINATOR
102+
return max(_QWEN_IMAGE_VAE_SPATIAL_SCALE, stride // _QWEN_IMAGE_VAE_SPATIAL_SCALE * _QWEN_IMAGE_VAE_SPATIAL_SCALE)
103+
104+
105+
@contextmanager
106+
def patch_qwen_image_vae_tiling(vae: QwenImageCompatibleVAE, tile_size: int | None) -> Iterator[None]:
107+
"""Set the VAE's tiling state for the duration of the block, then restore it.
108+
109+
Two things make this a context manager rather than a bare ``enable_tiling()`` call:
110+
111+
- ``enable_tiling`` writes the tile geometry straight onto the module, and ``disable_tiling`` only
112+
clears ``use_tiling`` — it does not restore the sizes. The module here is the model cache's own
113+
instance (``as_qwen_image_vae`` deliberately returns it unchanged to keep partial-loading hooks
114+
intact), so without a restore a tile size set once would persist for the lifetime of the cache
115+
entry and leak into later invocations — including ``anima_latents_to_image``, which shares the
116+
same VAE instance when a native-layout ``qwen_image_vae`` single file is loaded.
117+
- All four parameters are always passed explicitly. ``enable_tiling`` falls back to the module's
118+
current value for any argument left out, and its ``min``/``stride`` pair must stay consistent:
119+
the tile loops advance by *stride* but slice each accumulated tile to *min*. A ``min`` below the
120+
inherited 192px stride silently drops whole bands of the image, and a ``min`` above it grows
121+
every tile without removing any, making compute scale with ``tile_size**2``.
122+
123+
``tile_size=None`` disables tiling for the block.
124+
"""
125+
original = (
126+
vae.use_tiling,
127+
vae.tile_sample_min_height,
128+
vae.tile_sample_min_width,
129+
vae.tile_sample_stride_height,
130+
vae.tile_sample_stride_width,
131+
)
132+
try:
133+
if tile_size is None:
134+
vae.disable_tiling()
135+
else:
136+
stride = _tile_stride_for(tile_size)
137+
vae.enable_tiling(
138+
tile_sample_min_height=tile_size,
139+
tile_sample_min_width=tile_size,
140+
tile_sample_stride_height=stride,
141+
tile_sample_stride_width=stride,
142+
)
143+
yield
144+
finally:
145+
(
146+
vae.use_tiling,
147+
vae.tile_sample_min_height,
148+
vae.tile_sample_min_width,
149+
vae.tile_sample_stride_height,
150+
vae.tile_sample_stride_width,
151+
) = original

‎invokeai/backend/util/vae_working_memory.py‎

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -189,14 +189,25 @@ def estimate_vae_working_memory_wan(
189189

190190

191191
def estimate_vae_working_memory_qwen_image(
192-
operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKLQwenImage
192+
operation: Literal["encode", "decode"],
193+
image_tensor: torch.Tensor,
194+
vae: AutoencoderKLQwenImage,
195+
tile_size: int | None = None,
193196
) -> int:
194197
"""Estimate the working memory required by the invocation in bytes.
195198
196199
The Qwen Image VAE is a video-style autoencoder that operates on 5D tensors of shape
197-
(B, C, num_frames, H, W). Tiling is not used, so peak working memory scales with the full
198-
spatial output. The two trailing dimensions are the spatial H/W in latent space (decode) or
199-
pixel space (encode), matching the convention used by the other estimators here.
200+
(B, C, num_frames, H, W). The two trailing dimensions are the spatial H/W in latent space
201+
(decode) or pixel space (encode), matching the convention used by the other estimators here.
202+
203+
Without tiling, peak working memory scales with the full spatial extent. With tiling it is
204+
bounded by a single tile instead, so the estimate must follow suit — otherwise the cache keeps
205+
reserving the full-frame figure (~11.8 GB for a 2560x1440 encode on CUDA) and tiling buys
206+
nothing. Mirrors ``estimate_vae_working_memory_wan``: one tile plus 25% for the tile overlap,
207+
plus the pixel-space buffers, which stay resident on the execution device either way.
208+
209+
``tile_size`` is the resolved tile size (the nodes' 0 sentinel already substituted), and assumes
210+
the 4:3 tile-to-stride ratio applied by ``patch_qwen_image_vae_tiling``.
200211
"""
201212
latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1
202213

@@ -240,7 +251,22 @@ def estimate_vae_working_memory_qwen_image(
240251
else: # encode
241252
scaling_constant = 6300 if is_rocm else 1600
242253

243-
working_memory = h * w * element_size * scaling_constant
254+
if tile_size is not None and tile_size > 0:
255+
# Bounded by one tile (plus overlap) rather than the full frame.
256+
working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25
257+
# The full RGB image is the encode input / decode output and stays resident regardless. Unlike
258+
# the per-tile term this scales with the output area, so it is the term that decides whether the
259+
# estimate still holds at the resolutions tiling exists for.
260+
#
261+
# `tiled_decode` holds several pixel-space copies at once: every decoded tile in `rows`
262+
# ((tile_min / tile_stride)^2 ~ 1.8 frames at the 4:3 ratio the nodes set), the blended and
263+
# cropped `result_rows` (~1 frame) and the final `torch.cat` output (~1 frame). Measured at
264+
# ~5 frames on a 2560x1440 fp16 decode. Encode consumes its input image without duplicating it,
265+
# and accumulates only latents (16 channels at 1/64 the area — negligible).
266+
image_copies = 5 if operation == "decode" else 1
267+
working_memory += image_copies * 3 * h * w * element_size
268+
else:
269+
working_memory = h * w * element_size * scaling_constant
244270

245271
return int(working_memory)
246272

0 commit comments

Comments
 (0)