Merge a19cb1a13b into 532e285079

comfy/ldm/modules/attention.py

@@ -365,7 +365,9 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
                 r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
                 del s2
             break
-        except model_management.OOM_EXCEPTION as e:
+        except Exception as e:
+            if not model_management.is_oom_exception(e):
+                raise
             if first_op_done == False:
                 model_management.soft_empty_cache(True)
                 if cleared_cache == False:

comfy/ldm/modules/diffusionmodules/model.py

@@ -267,7 +267,9 @@ def slice_attention(q, k, v):
                 r1[:, :, i:end] = torch.bmm(v, s2)
                 del s2
             break
-        except model_management.OOM_EXCEPTION as e:
+        except Exception as e:
+            if not model_management.is_oom_exception(e):
+                raise
             model_management.soft_empty_cache(True)
             steps *= 2
             if steps > 128:
@@ -323,7 +325,9 @@ def pytorch_attention(q, k, v):
     try:
         out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
         out = out.transpose(2, 3).reshape(orig_shape)
-    except model_management.OOM_EXCEPTION:
+    except Exception as ex:
+        if not model_management.is_oom_exception(ex):
+            raise
         logging.warning("scaled_dot_product_attention OOMed: switched to slice attention")
         oom_fallback = True
     if oom_fallback:

comfy/ldm/modules/sub_quadratic_attention.py

@@ -169,7 +169,9 @@ def _get_attention_scores_no_kv_chunking(
     try:
         attn_probs = attn_scores.softmax(dim=-1)
         del attn_scores
-    except model_management.OOM_EXCEPTION:
+    except Exception as ex:
+        if not model_management.is_oom_exception(ex):
+            raise
         logging.warning("ran out of memory while running softmax in  _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
         attn_scores -= attn_scores.max(dim=-1, keepdim=True).values # noqa: F821 attn_scores is not defined
         torch.exp(attn_scores, out=attn_scores)

comfy/model_management.py

@@ -253,6 +253,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:

comfy/sd.py

@@ -804,7 +804,9 @@ 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:
+        except Exception as ex:
+            if not model_management.is_oom_exception(ex):
+                raise
             logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
             #NOTE: We don't know what tensors were allocated to stack variables at the time of the
             #exception and the exception itself refs them all until we get out of this except block.
@@ -878,8 +880,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.")
             #NOTE: We don't know what tensors were allocated to stack variables at the time of the
             #exception and the exception itself refs them all until we get out of this except block.

comfy_extras/nodes_upscale_model.py

@@ -84,7 +84,9 @@ class ImageUpscaleWithModel(io.ComfyNode):
                 pbar = comfy.utils.ProgressBar(steps)
                 s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
                 oom = False
-            except model_management.OOM_EXCEPTION as e:
+            except Exception as e:
+                if not model_management.is_oom_exception(e):
+                    raise
                 tile //= 2
                 if tile < 128:
                     raise e

execution.py

@@ -597,7 +597,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
         logging.error(traceback.format_exc())
         tips = ""
 
-        if isinstance(ex, comfy.model_management.OOM_EXCEPTION):
+        if comfy.model_management.is_oom_exception(ex):
             tips = "This error means you ran out of memory on your GPU.\n\nTIPS: If the workflow worked before you might have accidentally set the batch_size to a large number."
             logging.error("Got an OOM, unloading all loaded models.")
             comfy.model_management.unload_all_models()