Skip to content

Commit cf7fa5b

Browse files
perf(backend): clear torch cache after encoding each image in kontext extension
Slightly reduces VRAM allocations.
1 parent 39e4199 commit cf7fa5b

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

invokeai/backend/flux/extensions/kontext_extension.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
144144
image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
145145
# Use sample=False to get the distribution mean without noise
146146
kontext_latents_unpacked = vae.encode(image_tensor, sample=False)
147+
TorchDevice.empty_cache()
147148

148149
# Extract tensor dimensions
149150
batch_size, _, latent_height, latent_width = kontext_latents_unpacked.shape
@@ -169,11 +170,11 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
169170
# Option 1: Tile vertically (below existing content)
170171
potential_h_vertical = canvas_h + latent_height
171172
potential_w_vertical = max(canvas_w, latent_width)
172-
173+
173174
# Option 2: Tile horizontally (to the right of existing content)
174175
potential_h_horizontal = max(canvas_h, latent_height)
175176
potential_w_horizontal = canvas_w + latent_width
176-
177+
177178
# Choose arrangement that minimizes the maximum dimension
178179
# This keeps the canvas closer to square, optimizing attention computation
179180
if potential_h_vertical > potential_w_horizontal:

0 commit comments

Comments
 (0)