From 83be01cebcdbcf26a7d08d368ae4503c1f017d58 Mon Sep 17 00:00:00 2001 From: traugdor Date: Wed, 30 Oct 2024 09:07:10 -0500 Subject: [PATCH] Update nodes.py patched VAEDecode if fails due to lack of VRAM fallback to tiled decode method --- nodes.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/nodes.py b/nodes.py index c81a0af1e..d56d26383 100644 --- a/nodes.py +++ b/nodes.py @@ -280,11 +280,14 @@ class VAEDecode: CATEGORY = "latent" DESCRIPTION = "Decodes latent images back into pixel space images." - def decode(self, vae, samples): - images = vae.decode(samples["samples"]) - if len(images.shape) == 5: #Combine batches - images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) - return (images, ) + def decode(self, vae, samples, tile_size=512): + try: + images = vae.decode(samples["samples"]) + if len(images.shape) == 5: #Combine batches + images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) + return (images, ) + except (MemoryError, RuntimeError) as e: + return (vae.decode_tiled(samples["samples"], tile_x=tile_size // 8, tile_y=tile_size // 8, ), ) class VAEDecodeTiled: @classmethod