From bb3fb1a935b3d151e068aefdab805ee96e0406aa Mon Sep 17 00:00:00 2001 From: Haoming Date: Mon, 1 Jun 2026 15:21:19 +0800 Subject: [PATCH] oom --- backend/attention.py | 16 +++++++++++----- backend/memory_management.py | 11 +++++++++++ backend/patcher/vae.py | 12 ++++++++---- modules/call_queue.py | 32 +++++++++++++++----------------- modules/ui.py | 15 +++++++++++---- modules_forge/main_thread.py | 17 ++++++++++------- 6 files changed, 66 insertions(+), 37 deletions(-) diff --git a/backend/attention.py b/backend/attention.py index 7d652038..3fc79362 100644 --- a/backend/attention.py +++ b/backend/attention.py @@ -389,12 +389,15 @@ def slice_attention_vae(q, k, v): r1[:, :, i:end] = torch.bmm(v, s2) del s2 break - except memory_management.OOM_EXCEPTION as e: - memory_management.soft_empty_cache(True) - steps *= 2 + except Exception as e: + if not memory_management.is_oom(e): + raise e if steps > 128: raise e - logger.warning(f"Out of Memory Error; trying again... ({steps})") + + logger.warning("Out of Memory Error; retrying with higher steps...") + memory_management.soft_empty_cache() + steps *= 2 return r1 @@ -450,11 +453,14 @@ def pytorch_attention_vae(q, k, v): out = operations.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) _fallback = False - except memory_management.OOM_EXCEPTION: + except Exception as e: + if not memory_management.is_oom(e): + raise e logger.warning("Out of Memory Error; retrying with Slice Attention") _fallback = True if _fallback: + memory_management.soft_empty_cache() out = slice_attention_vae(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(orig_shape) return out diff --git a/backend/memory_management.py b/backend/memory_management.py index 900731e5..ae81a226 100644 --- a/backend/memory_management.py +++ b/backend/memory_management.py @@ -196,6 +196,17 @@ except Exception: pass OOM_EXCEPTION = getattr(torch, "OutOfMemoryError", Exception) +ACCELERATOR_ERROR = getattr(torch, "AcceleratorError", RuntimeError) + + +def is_oom(e: Exception) -> bool: + if isinstance(e, OOM_EXCEPTION): + return True + if isinstance(e, ACCELERATOR_ERROR) or "out of memory" in str(e).lower(): + discard_cuda_async_error() + return True + return False + if args.disable_xformers: XFORMERS_IS_AVAILABLE = False diff --git a/backend/patcher/vae.py b/backend/patcher/vae.py index adbb83a7..db4af748 100644 --- a/backend/patcher/vae.py +++ b/backend/patcher/vae.py @@ -219,8 +219,10 @@ 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 memory_management.OOM_EXCEPTION: - print("Warning: Encountered Out of Memory during VAE decoding; Retrying with Tiled VAE Decoding...") + except Exception as e: + if not memory_management.is_oom(e): + raise e + memory_management.logger.warning("Encountered Out of Memory during VAE decoding; Retrying with Tiled VAE Decoding...") _tile = True if _tile: @@ -270,8 +272,10 @@ class VAE: samples = torch.empty((_samples.shape[0],) + tuple(out.shape[1:]), device=self.output_device) samples[x : x + batch_number] = out _tile = False - except memory_management.OOM_EXCEPTION: - print("Warning: Encountered Out of Memory during VAE Encoding; Retrying with Tiled VAE Encoding...") + except Exception as e: + if not memory_management.is_oom(e): + raise e + memory_management.logger.warning("Encountered Out of Memory during VAE Encoding; Retrying with Tiled VAE Encoding...") _tile = True if _tile: diff --git a/modules/call_queue.py b/modules/call_queue.py index f2a7ed00..07c5b39a 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -1,11 +1,11 @@ -import os.path -from functools import wraps import html +import os.path import time import traceback +from functools import wraps +from modules import devices, fifo_lock, profiling, progress, shared from modules_forge import main_thread -from modules import shared, progress, errors, devices, fifo_lock, profiling queue_lock = fifo_lock.FIFOLock() @@ -76,15 +76,14 @@ def wrap_gradio_call_no_job(func, extra_outputs=None, add_stats=False): res = list(func(*args, **kwargs)) except Exception as e: if main_thread.last_exception is not None: - e = main_thread.last_exception + error_message = main_thread.last_exception else: + error_message = f"{type(e).__name__}: {e}" traceback.print_exc() - print(e) if extra_outputs_array is None: - extra_outputs_array = [None, ''] + extra_outputs_array = [None, ""] - error_message = f'{type(e).__name__}: {e}' res = extra_outputs_array + [f"
{html.escape(error_message)}
"] devices.torch_gc() @@ -97,15 +96,15 @@ def wrap_gradio_call_no_job(func, extra_outputs=None, add_stats=False): elapsed_s = elapsed % 60 elapsed_text = f"{elapsed_s:.1f} sec." if elapsed_m > 0: - elapsed_text = f"{elapsed_m} min. "+elapsed_text + elapsed_text = f"{elapsed_m} min. " + elapsed_text if run_memmon: - mem_stats = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.stop().items()} - active_peak = mem_stats['active_peak'] - reserved_peak = mem_stats['reserved_peak'] - sys_peak = mem_stats['system_peak'] - sys_total = mem_stats['total'] - sys_pct = sys_peak/max(sys_total, 1) * 100 + mem_stats = {k: -(v // -(1024 * 1024)) for k, v in shared.mem_mon.stop().items()} + active_peak = mem_stats["active_peak"] + reserved_peak = mem_stats["reserved_peak"] + sys_peak = mem_stats["system_peak"] + sys_total = mem_stats["total"] + sys_pct = sys_peak / max(sys_total, 1) * 100 toltip_a = "Active: peak amount of video memory used during generation (excluding cached data)" toltip_r = "Reserved: total amount of video memory allocated by the Torch library " @@ -117,12 +116,12 @@ def wrap_gradio_call_no_job(func, extra_outputs=None, add_stats=False): vram_html = f"

{text_a}, {text_r}, {text_sys}

" else: - vram_html = '' + vram_html = "" if shared.opts.profiling_enable and os.path.exists(shared.opts.profiling_filename): profiling_html = f"

[ Profile ]

" else: - profiling_html = '' + profiling_html = "" # last item is always HTML res[-1] += f"

Time taken: {elapsed_text}

{vram_html}{profiling_html}
" @@ -130,4 +129,3 @@ def wrap_gradio_call_no_job(func, extra_outputs=None, add_stats=False): return tuple(res) return f - diff --git a/modules/ui.py b/modules/ui.py index 6aa9cd28..cd8375ed 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -67,6 +67,13 @@ def no_config(*comps: gr.components.Component): setattr(comp, "_internal_preset_param", True) +def cleanup(): + from modules_forge.main_thread import last_exception + + if last_exception == "OOM": + sd_models.unload_model_weights() + + # Using constants for these since the variation selector isn't visible. # Important that they exactly match script.js for tooltip to work. random_symbol = "\U0001f3b2\ufe0f" # 🎲️ @@ -378,8 +385,8 @@ def create_ui(): show_progress=False, ) - toprow.prompt.submit(**txt2img_args) - toprow.submit.click(**txt2img_args) + toprow.prompt.submit(**txt2img_args).then(fn=cleanup) + toprow.submit.click(**txt2img_args).then(fn=cleanup) def select_gallery_image(index): index = int(index) @@ -749,8 +756,8 @@ def create_ui(): show_progress=False, ) - toprow.prompt.submit(**img2img_args) - toprow.submit.click(**img2img_args) + toprow.prompt.submit(**img2img_args).then(fn=cleanup) + toprow.submit.click(**img2img_args).then(fn=cleanup) res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False) diff --git a/modules_forge/main_thread.py b/modules_forge/main_thread.py index fac836bb..6c4dba6a 100644 --- a/modules_forge/main_thread.py +++ b/modules_forge/main_thread.py @@ -5,7 +5,7 @@ import threading import traceback from collections import deque -from typing import Callable +from typing import Callable, Optional lock = threading.Lock() condition = threading.Condition(lock) @@ -13,7 +13,7 @@ condition = threading.Condition(lock) last_id: int = 0 waiting_queue: deque["Task"] = deque() finished_tasks: dict[int, "Task"] = {} -last_exception: Exception = None +last_exception: Optional[str] = None class Task: @@ -23,18 +23,21 @@ class Task: self.args = args self.kwargs = kwargs self.result = None - self.exception = None def work(self): global last_exception try: self.result = self.func(*self.args, **self.kwargs) - self.exception = None last_exception = None except Exception as e: - traceback.print_exc() - self.exception = e - last_exception = e + from backend.memory_management import is_oom, logger + + if is_oom(e): + logger.error("Encountered Out of Memory during Sampling; Unloading all Models...") + last_exception = "OOM" + else: + traceback.print_exc() + last_exception = f"{type(e).__name__}: {e}" def loop():