diff --git a/modules/call_queue.py b/modules/call_queue.py index 114395661..466f44843 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -7,20 +7,41 @@ from modules import shared, progress, errors, timer from modules.logger import log -queue_lock = threading.Lock() -debug = os.environ.get('SD_QUEUE_DEBUG', None) is not None +_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}') + return _queue_lock + + +queue_lock = Queue() # public lock for external use def get_lock(): - if 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={queue_lock.locked()} fn={fn}') return queue_lock +def is_locked(): + return _queue_lock.locked() + + def wrap_queued_call(func): def f(*args, **kwargs): - with get_lock(): + with Queue(): res = func(*args, **kwargs) return res return f @@ -35,7 +56,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None, name=None): progress.add_task_to_queue(id_task) else: id_task = None - with get_lock(): + with Queue(): progress.start_task(id_task) try: res = func(*args, **kwargs) diff --git a/modules/control/run.py b/modules/control/run.py index 1157bb276..f5dae8c2b 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -1,5 +1,6 @@ import os import sys +from functools import partial import cv2 from PIL import Image from modules.logger import log @@ -336,6 +337,12 @@ def control_process(p: StableDiffusionProcessingControl, return output, info, script_run +def generate(*args, **kwargs): + generate_call = partial(control_run, *args, **kwargs) + # return [], '', '', 'Queued' + return generate_call() + + def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg units: list[unit.Unit] | None = None, inputs: list[Image.Image] | None = None, inits: list[Image.Image] | None = None, mask: Image.Image = None, unit_type: str | None = None, is_generator: bool = True, input_type: int = 0, diff --git a/modules/postprocess/aurasr_model.py b/modules/postprocess/aurasr_model.py index c5c72c7a8..d259ab982 100644 --- a/modules/postprocess/aurasr_model.py +++ b/modules/postprocess/aurasr_model.py @@ -17,7 +17,7 @@ class UpscalerAuraSR(Upscaler): def callback(self, _step: int, _timestep: int, _latents: torch.FloatTensor): pass - def do_upscale(self, img: Image.Image, selected_model): + def do_upscale(self, img: Image.Image, selected_model): # pylint: disable=arguments-differ from modules.postprocess.aurasr_arch import AuraSR if self.model is None: self.model = AuraSR.from_pretrained("vladmandic/aurasr", use_safetensors=False) diff --git a/modules/postprocess/esrgan_model.py b/modules/postprocess/esrgan_model.py index 7aaffc124..9ca7d7c69 100644 --- a/modules/postprocess/esrgan_model.py +++ b/modules/postprocess/esrgan_model.py @@ -125,7 +125,7 @@ class UpscalerESRGAN(Upscaler): self.scalers = self.find_scalers() self.models = {} - def do_upscale(self, img, selected_model): + def do_upscale(self, img, selected_model): # pylint: disable=arguments-differ model = self.load_model(selected_model) if model is None: return img diff --git a/modules/postprocess/realesrgan_model.py b/modules/postprocess/realesrgan_model.py index 5aa89fcb1..51a51bae9 100644 --- a/modules/postprocess/realesrgan_model.py +++ b/modules/postprocess/realesrgan_model.py @@ -39,7 +39,7 @@ class UpscalerRealESRGAN(Upscaler): def load_model(self, path): # pylint: disable=unused-argument pass - def do_upscale(self, img, selected_model): + def do_upscale(self, img, selected_model): # pylint: disable=arguments-differ if not self.enable: return img try: diff --git a/modules/postprocess/scunet_model.py b/modules/postprocess/scunet_model.py index d66dece39..9be3c3be8 100644 --- a/modules/postprocess/scunet_model.py +++ b/modules/postprocess/scunet_model.py @@ -68,7 +68,7 @@ class UpscalerSCUNet(Upscaler): output = E.div_(W) return output - def do_upscale(self, img: Image.Image, selected_file): + def do_upscale(self, img: Image.Image, selected_file): # pylint: disable=arguments-differ devices.torch_gc() model = self.load_model(selected_file) if model is None: diff --git a/modules/postprocess/sdupscaler_model.py b/modules/postprocess/sdupscaler_model.py index 01a6efc73..17917191d 100644 --- a/modules/postprocess/sdupscaler_model.py +++ b/modules/postprocess/sdupscaler_model.py @@ -41,7 +41,7 @@ class UpscalerDiffusion(Upscaler): def callback(self, _step: int, _timestep: int, _latents: torch.FloatTensor): pass - def do_upscale(self, img: Image.Image, selected_model): + def do_upscale(self, img: Image.Image, selected_model): # pylint: disable=arguments-differ devices.torch_gc() model = self.load_model(selected_model) if model is None: diff --git a/modules/postprocess/swinir_model.py b/modules/postprocess/swinir_model.py index c9326f4f6..6f1f10476 100644 --- a/modules/postprocess/swinir_model.py +++ b/modules/postprocess/swinir_model.py @@ -66,7 +66,7 @@ class UpscalerSwinIR(Upscaler): log.error(f'Upscaler invalid parameters: type={self.name} model={info.local_data_path} {e}') return model - def do_upscale(self, img, selected_model): + def do_upscale(self, img, selected_model): # pylint: disable=arguments-differ model = self.load_model(selected_model) if model is None: return img diff --git a/modules/processing.py b/modules/processing.py index de4d33a21..b72c8e151 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -6,6 +6,7 @@ from PIL import Image, ImageOps from modules import shared, devices, errors, images, scripts_manager, memstats, script_callbacks, extra_networks, sd_models, sd_checkpoint, sd_vae, processing_helpers, processing_grading, timer, masking from modules.logger import log from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet +from modules.processing_info import create_infotext from modules.processing_class import ( # pylint: disable=unused-import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, @@ -13,7 +14,6 @@ from modules.processing_class import ( # pylint: disable=unused-import StableDiffusionProcessingControl, StableDiffusionProcessingVideo, ) -from modules.processing_info import create_infotext opt_C = 4 @@ -484,6 +484,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if not p.prompts: break p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts) + + if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner): p.scripts.process_batch(p, batch_number=n, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds) diff --git a/modules/progress.py b/modules/progress.py index 292e28728..cd6a52d4f 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -45,6 +45,16 @@ def add_task_to_queue(id_job): pending_tasks[id_job] = time.time() +def get_tasks(): + return { + "current": current_task, + "pending": list(pending_tasks.keys()), + "finished": finished_tasks, + "results": recorded_results, + } + + + class ProgressRequest(BaseModel): id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for") id_live_preview: int = Field(default=-1, title="Live preview image ID", description="id of last received last preview image") diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 2a9e8c863..93cb6a965 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -38,7 +38,7 @@ def setup_img2img_steps(p, steps=None): def single_sample_to_image(sample, approximation=None, fast=False): - with queue_lock: + with queue_lock: # only one preview can run at a time t0 = time.time() approximation = approximation or shared.opts.show_progress_type if debug: diff --git a/modules/ui_control.py b/modules/ui_control.py index f31edb989..11bda73bf 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -1,5 +1,6 @@ import os import time +import asyncio import gradio as gr from modules.control import unit from modules import errors, shared, progress, generation_parameters_copypaste, call_queue, scripts_manager, masking, images, processing_vae, timer # pylint: disable=ungrouped-imports @@ -17,6 +18,7 @@ controls: list[gr.components.Component] = [] # list of gr controls debug = log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None debug('Trace: CONTROL') use_generator = os.environ.get('SD_USE_GENERATOR', None) is not None +use_async = os.environ.get('SD_USE_ASYNC', None) is not None def return_stats(t: float | None = None): @@ -96,7 +98,7 @@ def generate_click_generator(job_id: str, state: str, active_tab: str, *args): # while helpers.busy: debug(f'Control: tab="{active_tab}" job={job_id} busy') time.sleep(0.1) - from modules.control.run import control_run + from modules.control.run import generate debug(f'Control: tab="{active_tab}" job={job_id} args={args}') progress.add_task_to_queue(job_id) with call_queue.get_lock(): @@ -107,7 +109,7 @@ def generate_click_generator(job_id: str, state: str, active_tab: str, *args): # t = time.perf_counter() results = {} try: - for results in control_run(state, units, helpers.input_source, helpers.input_init, helpers.input_mask, active_tab, True, *args): + for results in generate(state, units, helpers.input_source, helpers.input_init, helpers.input_mask, active_tab, True, *args): progress.record_results(job_id, results) yield return_controls(results, t) except GeneratorExit: @@ -126,7 +128,7 @@ def generate_click(job_id: str, state: str, active_tab: str, *args): while helpers.busy: debug(f'Control: tab="{active_tab}" job={job_id} busy') time.sleep(0.1) - from modules.control.run import control_run + from modules.control.run import generate debug(f'Control: tab="{active_tab}" job={job_id} args={args}') progress.add_task_to_queue(job_id) with call_queue.get_lock(): @@ -136,7 +138,7 @@ def generate_click(job_id: str, state: str, active_tab: str, *args): progress.start_task(job_id) try: t = time.perf_counter() - for results in control_run(state, units, helpers.input_source, helpers.input_init, helpers.input_mask, active_tab, True, *args): + for results in generate(state, units, helpers.input_source, helpers.input_init, helpers.input_mask, active_tab, True, *args): progress.record_results(job_id, results) except GeneratorExit: log.error("Control: generator exit") @@ -150,6 +152,10 @@ def generate_click(job_id: str, state: str, active_tab: str, *args): return return_controls(results, t) +async def generate_click_async(job_id: str, state: str, active_tab: str, *args): + return await asyncio.to_thread(generate_click, job_id, state, active_tab, *args) + + def create_ui(_blocks: gr.Blocks=None): helpers.initialize() @@ -335,14 +341,19 @@ def create_ui(_blocks: gr.Blocks=None): output_html_log, ] - generate_fn = generate_click_generator if use_generator else generate_click + if use_generator: + generate_fn = generate_click_generator + elif use_async: + generate_fn = generate_click_async + else: + generate_fn = generate_click + control_dict = dict( fn=generate_fn, _js="submit_control", inputs=[tabs_state, state, tabs_state] + input_fields + input_script_args, outputs=output_fields, show_progress='hidden', - # queue=not shared.cmd_opts.listen, ) prompt.submit(**control_dict) negative.submit(**control_dict) diff --git a/modules/upscaler_algo.py b/modules/upscaler_algo.py index ea4c8ce6b..24e45cbdc 100644 --- a/modules/upscaler_algo.py +++ b/modules/upscaler_algo.py @@ -13,7 +13,7 @@ class UpscalerDCC(Upscaler): UpscalerData("DCC Interpolation", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=unused-argument,arguments-differ import math import numpy as np from modules.postprocess.dcc import DCC @@ -41,7 +41,7 @@ class UpscalerVIPS(Upscaler): UpscalerData("VIPS MagicKernelSharp 2021", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=unused-argument,arguments-differ if selected_model is None: return img from installer import install @@ -85,7 +85,7 @@ class UpscalerHQX(Upscaler): UpscalerData("HQX Interpolation", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=unused-argument,arguments-differ import numpy as np from modules.postprocess.hqx import hqx t0 = time.time() @@ -106,7 +106,7 @@ class UpscalerICBI(Upscaler): UpscalerData("ICB Interpolation", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=unused-argument,arguments-differ import numpy as np from modules.postprocess.icbi import icbi t0 = time.time() diff --git a/modules/upscaler_nvvfx.py b/modules/upscaler_nvvfx.py index fb9cd1552..592e5cabd 100644 --- a/modules/upscaler_nvvfx.py +++ b/modules/upscaler_nvvfx.py @@ -32,7 +32,8 @@ class UpscalerNVVFX(Upscaler): UpscalerData("nVidia VFX highbitrate ultra", None, self, scale=19), ] - def upscale(self, img: Image.Image | torch.Tensor, scale, selected_model: str | None = None): # nvvfx overrides upscale instead of do_upscale because it handles scale directly + # nvvfx overrides upscale instead of do_upscale because it handles scale directly + def upscale(self, img: Image.Image | torch.Tensor, scale, selected_model: str | None = None): if selected_model is None: return img from installer import install diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py index bffcec827..0aa267a15 100644 --- a/modules/upscaler_simple.py +++ b/modules/upscaler_simple.py @@ -12,7 +12,7 @@ class UpscalerNone(Upscaler): def load_model(self, path): pass - def do_upscale(self, img, selected_model=None): + def do_upscale(self, img, selected_model=None): # pylint: disable=unused-argument,arguments-differ return img @@ -31,7 +31,7 @@ class UpscalerResize(Upscaler): UpscalerData("Resize Sharpfin Lanczos3", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=arguments-differ if selected_model is None: return img elif selected_model == "Resize Nearest": @@ -74,7 +74,7 @@ class UpscalerLatent(Upscaler): UpscalerData("Latent Bicubic antialias", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=arguments-differ import torch import torch.nn.functional as F if isinstance(img, torch.Tensor) and (len(img.shape) == 4): diff --git a/modules/upscaler_vae.py b/modules/upscaler_vae.py index a0318302f..dca3ef9b1 100644 --- a/modules/upscaler_vae.py +++ b/modules/upscaler_vae.py @@ -15,7 +15,7 @@ class UpscalerAsymmetricVAE(Upscaler): UpscalerData("Asymmetric VAE v2", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=arguments-differ if selected_model is None: return img import diffusers @@ -55,7 +55,7 @@ class UpscalerWanUpscale(Upscaler): UpscalerData("WAN Asymmetric Upscale", None, self), ] - def do_upscale(self, img: Image.Image, selected_model=None): + def do_upscale(self, img: Image.Image, selected_model=None): # pylint: disable=arguments-differ if selected_model is None: return img import torch.nn.functional as FN diff --git a/webui.py b/webui.py index 1b7cfe97b..911d35ccf 100644 --- a/webui.py +++ b/webui.py @@ -19,7 +19,7 @@ import modules.paths import modules.devices import modules.migrate from modules import shared -from modules.call_queue import queue_lock, wrap_queued_call, wrap_gradio_gpu_call # pylint: disable=unused-import +from modules import call_queue import modules.gr_tempdir import modules.modeldata import modules.extensions @@ -196,22 +196,22 @@ def load_model(): thread_refiner.join() shared.state.end(jobid) timer.startup.record("checkpoint") - shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='model')), call=False) - shared.opts.onchange("sd_model_refiner", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='refiner')), call=False) - shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False) - shared.opts.onchange("sd_unet", wrap_queued_call(lambda: modules.sd_unet.load_unet(shared.sd_model)), call=False) - shared.opts.onchange("sd_unet_secondary", wrap_queued_call(lambda: modules.sd_unet.load_unet_secondary(shared.sd_model)), call=False) - shared.opts.onchange("sd_text_encoder", wrap_queued_call(lambda: modules.sd_models.reload_text_encoder()), call=False) + shared.opts.onchange("sd_model_checkpoint", call_queue.wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='model')), call=False) + shared.opts.onchange("sd_model_refiner", call_queue.wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='refiner')), call=False) + shared.opts.onchange("sd_vae", call_queue.wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False) + shared.opts.onchange("sd_unet", call_queue.wrap_queued_call(lambda: modules.sd_unet.load_unet(shared.sd_model)), call=False) + shared.opts.onchange("sd_unet_secondary", call_queue.wrap_queued_call(lambda: modules.sd_unet.load_unet_secondary(shared.sd_model)), call=False) + shared.opts.onchange("sd_text_encoder", call_queue.wrap_queued_call(lambda: modules.sd_models.reload_text_encoder()), call=False) shared.opts.onchange("temp_dir", modules.gr_tempdir.on_tmpdir_changed) for opt in modules.sd_offload_state.offload_reapply_options: - shared.opts.onchange(opt, wrap_queued_call(modules.sd_models.reapply_offload), call=False) + shared.opts.onchange(opt, call_queue.wrap_queued_call(modules.sd_models.reapply_offload), call=False) timer.startup.record("onchange") def create_api(app): log.debug('API initialize') from modules.api.api import Api - api = Api(app, queue_lock) + api = Api(app, call_queue.queue_lock) return api