From b749e0db491cf87fa7a2a7936943e60de8d3adf7 Mon Sep 17 00:00:00 2001 From: yash-711 Date: Wed, 25 Feb 2026 13:04:24 +0530 Subject: [PATCH] Add cleanup only on the error exit (when tile < 128) before re-raising the OOM exception --- comfy_extras/nodes_upscale_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py index d8bf55e2b..484ff246a 100644 --- a/comfy_extras/nodes_upscale_model.py +++ b/comfy_extras/nodes_upscale_model.py @@ -89,6 +89,7 @@ class ImageUpscaleWithModel(io.ComfyNode): except model_management.OOM_EXCEPTION as e: tile //= 2 if tile < 128: + upscale_model.to("cpu") raise e # upscale_model.to("cpu") # Commented out to keep model on GPU because when processing batch images then model unnecessarily moves to CPU