From 519e7760e71b135de4e78c14490481173db94539 Mon Sep 17 00:00:00 2001
From: traugdor <traugdor@users.noreply.github.com>
Date: Fri, 1 Nov 2024 16:18:32 -0500
Subject: [PATCH] Update sd.py

Move catching of RuntimeError and MemoryError to sd.py
---
 comfy/sd.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/comfy/sd.py b/comfy/sd.py
index e9a3579c9..f68858a8e 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -341,7 +341,7 @@ class VAE:
                 if pixel_samples is None:
                     pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device)
                 pixel_samples[x:x+batch_number] = out
-        except model_management.OOM_EXCEPTION as e:
+        except (model_management.OOM_EXCEPTION, RuntimeError, MemoryError) as e:
             logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
             dims = samples_in.ndim - 2
             if dims == 1: