From 6b58f6bd5ab0d876b912aaa9728ed29e4637dc04 Mon Sep 17 00:00:00 2001
From: Alex Butler <alexheretic@gmail.com>
Date: Wed, 23 Apr 2025 21:30:30 +0100
Subject: [PATCH] vae handle HIP OOM exceptions

---
 comfy/model_management.py | 8 ++++++++
 comfy/sd.py               | 5 +++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index 054291432..9935558b4 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -235,6 +235,14 @@ try:
 except:
     OOM_EXCEPTION = Exception
 
+
+def is_oom_exception(ex):
+    if isinstance(ex, OOM_EXCEPTION):
+        return True
+    # handle also other kinds of oom, e.g. "HIP error: out of memory"
+    msg = str(ex)
+    return "out of memory" in msg
+
 XFORMERS_VERSION = ""
 XFORMERS_ENABLED_VAE = True
 if args.disable_xformers:
diff --git a/comfy/sd.py b/comfy/sd.py
index cd13ab5f0..eb8379cc7 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -640,8 +640,9 @@ class VAE:
                 if samples is None:
                     samples = torch.empty((pixel_samples.shape[0],) + tuple(out.shape[1:]), device=self.output_device)
                 samples[x:x + batch_number] = out
-
-        except model_management.OOM_EXCEPTION:
+        except Exception as ex:
+            if not model_management.is_oom_exception(ex):
+                raise
             logging.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
             if self.latent_dim == 3:
                 tile = 256