CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-07-21 21:01:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

oom

Browse Source
This commit is contained in:
Haoming 2026-06-01 15:21:19 +08:00
parent 5d14e0c901
commit bb3fb1a935
6 changed files with 66 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
backend/attention.py
View File

@ -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

11
backend/memory_management.py
View File

@ -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

12
backend/patcher/vae.py
View File

@ -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:

32
modules/call_queue.py
View File

@ -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"<div class='error'>{html.escape(error_message)}</div>"]
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"<p class='vram'>{text_a}, <wbr>{text_r}, <wbr>{text_sys}</p>"
else:
vram_html = ''
vram_html = ""
if shared.opts.profiling_enable and os.path.exists(shared.opts.profiling_filename):
profiling_html = f"<p class='profile'> [ <a href='{profiling.webpath()}' download>Profile</a> ] </p>"
else:
profiling_html = ''
profiling_html = ""
# last item is always HTML
res[-1] += f"<div class='performance'><p class='time'>Time taken: <wbr><span class='measurement'>{elapsed_text}</span></p>{vram_html}{profiling_html}</div>"
@ -130,4 +129,3 @@ def wrap_gradio_call_no_job(func, extra_outputs=None, add_stats=False):
return tuple(res)
return f

15
modules/ui.py
View File

@ -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)

17
modules_forge/main_thread.py
View File

@ -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():