From a803611ec1dd7227431059e85e278719490baa13 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 23 Jun 2026 18:38:52 -0600 Subject: [PATCH] Enable tiling on vae when decoding with low_vram flag on krea2 --- extensions_built_in/diffusion_models/krea2/krea2.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index e1638205..8df1f98c 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -467,7 +467,16 @@ class Krea2Model(BaseModel): ) latents = latents * latents_std + latents_mean - images = self.vae.decode(latents).sample + # Full-resolution decode spikes VRAM; tile it when low on VRAM (decode + # only -- encode stays untiled). + tiled = self.model_config.low_vram + if tiled: + self.vae.enable_tiling() + try: + images = self.vae.decode(latents).sample + finally: + if tiled: + self.vae.disable_tiling() images = images.squeeze(2) # drop frame dim return images.to(device, dtype=dtype)