mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 07:31:01 +02:00
CalamitousFelicitousness
78a49eab33
Queue.__exit__ returned _queue_lock, and a truthy __exit__ return suppresses the exception in flight, so every `with queue_lock:` block discarded exceptions and resumed with locals from the aborted block unassigned. The bare threading.Lock it replaced returned None.
120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
import os
|
|
import sys
|
|
import html
|
|
import threading
|
|
import time
|
|
from modules import shared, progress, errors, timer
|
|
from modules.logger import log
|
|
|
|
|
|
_queue_lock = threading.Lock() # internal
|
|
_queue_debug = os.environ.get('SD_QUEUE_DEBUG', None) is not None
|
|
|
|
|
|
class Queue:
|
|
def __enter__(self):
|
|
_queue_lock.acquire()
|
|
if _queue_debug:
|
|
fn = f'{sys._getframe(3).f_code.co_name}:{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
|
|
log.debug(f'Queue: lock state={_queue_lock.locked()} fn={fn}')
|
|
return _queue_lock
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb): # pylint: disable=unused-argument
|
|
if _queue_lock.locked():
|
|
_queue_lock.release()
|
|
if _queue_debug:
|
|
fn = f'{sys._getframe(3).f_code.co_name}:{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
|
|
log.debug(f'Queue: unlock state={_queue_lock.locked()} fn={fn}')
|
|
# no return: a truthy __exit__ suppresses the exception in flight
|
|
|
|
|
|
queue_lock = Queue() # public lock for external use
|
|
|
|
|
|
def get_lock():
|
|
return queue_lock
|
|
|
|
|
|
def is_locked():
|
|
return _queue_lock.locked()
|
|
|
|
|
|
def wrap_queued_call(func):
|
|
def f(*args, **kwargs):
|
|
with Queue():
|
|
res = func(*args, **kwargs)
|
|
return res
|
|
return f
|
|
|
|
|
|
def wrap_gradio_gpu_call(func, extra_outputs=None, name=None):
|
|
name = name or func.__name__
|
|
def f(*args, **kwargs):
|
|
# if the first argument is a string that says "task(...)", it is treated as a job id
|
|
if len(args) > 0 and type(args[0]) == str and args[0][0:5] == "task(" and args[0][-1] == ")":
|
|
id_task = args[0]
|
|
progress.add_task_to_queue(id_task)
|
|
else:
|
|
id_task = None
|
|
with Queue():
|
|
progress.start_task(id_task)
|
|
try:
|
|
res = func(*args, **kwargs)
|
|
progress.record_results(id_task, res)
|
|
except Exception as e:
|
|
log.error(f"Exception: {e}")
|
|
log.error(f"Arguments: args={str(args)[:10240]} kwargs={str(kwargs)[:10240]}")
|
|
errors.display(e, 'gradio call')
|
|
res = extra_outputs or []
|
|
res.append(f"<div class='error'>{html.escape(str(e))}</div>")
|
|
finally:
|
|
progress.finish_task(id_task)
|
|
return res
|
|
return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True, name=name)
|
|
|
|
|
|
def wrap_gradio_call(func, extra_outputs=None, add_stats=False, name=None):
|
|
job_name = name if name is not None else func.__name__
|
|
def f(*args, extra_outputs_array=extra_outputs, **kwargs):
|
|
t = time.perf_counter()
|
|
shared.mem_mon.reset()
|
|
if len(args) > 0 and type(args[0]) == str and args[0][0:5] == "task(" and args[0][-1] == ")":
|
|
task_id = args[0]
|
|
else:
|
|
task_id = 0
|
|
jobid = shared.state.begin(job_name, task_id=task_id)
|
|
try:
|
|
if shared.cmd_opts.profile:
|
|
errors.profile_stop()
|
|
errors.profile_print('BeforeWrapGradioCall')
|
|
errors.profile_start()
|
|
res = func(*args, **kwargs)
|
|
if res is None:
|
|
msg = "No result returned from function"
|
|
log.warning(msg)
|
|
res = extra_outputs_array or []
|
|
res.append(f"<div class='error'>{html.escape(msg)}</div>")
|
|
else:
|
|
res = list(res)
|
|
if shared.cmd_opts.profile:
|
|
errors.profile_stop()
|
|
errors.profile_print('AfterWrapGradioCall')
|
|
errors.profile_start()
|
|
except Exception as e:
|
|
errors.display(e, 'gradio call')
|
|
res = extra_outputs_array or []
|
|
res.append(f"<div class='error'>{html.escape(type(e).__name__+': '+str(e))}</div>")
|
|
shared.state.end(jobid)
|
|
if not add_stats:
|
|
return tuple(res)
|
|
elapsed = time.perf_counter() - t
|
|
elapsed_m = int(elapsed // 60)
|
|
elapsed_s = elapsed % 60
|
|
elapsed_text = f"{elapsed_m}m {elapsed_s:.2f}s" if elapsed_m > 0 else f"{elapsed_s:.2f}s"
|
|
summary = timer.process.summary(min_time=0.25, total=False).replace('=', ' ')
|
|
memory = shared.mem_mon.summary()
|
|
if isinstance(res, list) and isinstance(res[-1], str):
|
|
res[-1] += f"<div class='performance'><p>Time: {elapsed_text} | {summary} {memory}</p></div>"
|
|
return tuple(res)
|
|
return f
|