Skip to content

Commit

Permalink
Fix: do not break progressbar on non-job actions (add wrap_gradio_cal…
Browse files Browse the repository at this point in the history
  • Loading branch information
Panchovix committed Jul 13, 2024
1 parent f902bc6 commit 12435dd
Show file tree
Hide file tree
Showing 5 changed files with 29 additions and 21 deletions.
24 changes: 16 additions & 8 deletions modules/call_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,22 @@ def f(*args, **kwargs):


def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
@wraps(func)
def f(*args, **kwargs):
try:
res = func(*args, **kwargs)
finally:
shared.state.skipped = False
shared.state.interrupted = False
shared.state.stopping_generation = False
shared.state.job_count = 0
shared.state.job = ""
return res

return wrap_gradio_call_no_job(f, extra_outputs, add_stats)


def wrap_gradio_call_no_job(func, extra_outputs=None, add_stats=False):
@wraps(func)
def f(*args, extra_outputs_array=extra_outputs, **kwargs):
run_memmon = shared.opts.memmon_poll_rate > 0 and not shared.mem_mon.disabled and add_stats
Expand All @@ -66,9 +82,6 @@ def f(*args, extra_outputs_array=extra_outputs, **kwargs):
arg_str += f" (Argument list truncated at {max_debug_str_len}/{len(arg_str)} characters)"
errors.report(f"{message}\n{arg_str}", exc_info=True)

shared.state.job = ""
shared.state.job_count = 0

if extra_outputs_array is None:
extra_outputs_array = [None, '']

Expand All @@ -77,11 +90,6 @@ def f(*args, extra_outputs_array=extra_outputs, **kwargs):

devices.torch_gc()

shared.state.skipped = False
shared.state.interrupted = False
shared.state.stopping_generation = False
shared.state.job_count = 0

if not add_stats:
return tuple(res)

Expand Down
4 changes: 2 additions & 2 deletions modules/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
import gradio.utils
import numpy as np
from PIL import Image, PngImagePlugin # noqa: F401
from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call
from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call, wrap_gradio_call_no_job # noqa: F401

from modules import gradio_extensons, sd_schedulers # noqa: F401
from modules import sd_hijack, sd_models, script_callbacks, ui_extensions, deepbooru, extra_networks, ui_common, ui_postprocessing, progress, ui_loadsave, shared_items, ui_settings, timer, sysinfo, ui_checkpoint_merger, scripts, sd_samplers, processing, ui_extra_networks, ui_toprow, launch_utils
Expand Down Expand Up @@ -885,7 +885,7 @@ def select_img2img_tab(tab):
))

image.change(
fn=wrap_gradio_call(modules.extras.run_pnginfo),
fn=wrap_gradio_call_no_job(modules.extras.run_pnginfo),fn=wrap_gradio_call(modules.extras.run_pnginfo),
inputs=[image],
outputs=[html, generation_info, html2],
)
Expand Down
4 changes: 2 additions & 2 deletions modules/ui_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,7 @@ def open_folder(f, images=None, index=None):
)

save.click(
fn=call_queue.wrap_gradio_call(save_files),
fn=call_queue.wrap_gradio_call_no_job(save_files),
_js="(x, y, z, w) => [x, y, false, selected_gallery_index()]",
inputs=[
res.generation_info,
Expand All @@ -244,7 +244,7 @@ def open_folder(f, images=None, index=None):
)

save_zip.click(
fn=call_queue.wrap_gradio_call(save_files),
fn=call_queue.wrap_gradio_call_no_job(save_files),
_js="(x, y, z, w) => [x, y, true, selected_gallery_index()]",
inputs=[
res.generation_info,
Expand Down
14 changes: 7 additions & 7 deletions modules/ui_extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -624,37 +624,37 @@ def create_ui():
)

install_extension_button.click(
fn=modules.ui.wrap_gradio_call(install_extension_from_index, extra_outputs=[gr.update(), gr.update()]),
fn=modules.ui.wrap_gradio_call_no_job(install_extension_from_index, extra_outputs=[gr.update(), gr.update()]),
inputs=[extension_to_install, selected_tags, showing_type, filtering_type, sort_column, search_extensions_text],
outputs=[available_extensions_table, extensions_table, install_result],
)

search_extensions_text.change(
fn=modules.ui.wrap_gradio_call(search_extensions, extra_outputs=[gr.update()]),
fn=modules.ui.wrap_gradio_call_no_job(search_extensions, extra_outputs=[gr.update()]),
inputs=[search_extensions_text, selected_tags, showing_type, filtering_type, sort_column],
outputs=[available_extensions_table, install_result],
)

selected_tags.change(
fn=modules.ui.wrap_gradio_call(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
fn=modules.ui.wrap_gradio_call_no_job(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
inputs=[selected_tags, showing_type, filtering_type, sort_column, search_extensions_text],
outputs=[available_extensions_table, install_result]
)

showing_type.change(
fn=modules.ui.wrap_gradio_call(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
fn=modules.ui.wrap_gradio_call_no_job(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
inputs=[selected_tags, showing_type, filtering_type, sort_column, search_extensions_text],
outputs=[available_extensions_table, install_result]
)

filtering_type.change(
fn=modules.ui.wrap_gradio_call(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
fn=modules.ui.wrap_gradio_call_no_job(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
inputs=[selected_tags, showing_type, filtering_type, sort_column, search_extensions_text],
outputs=[available_extensions_table, install_result]
)

sort_column.change(
fn=modules.ui.wrap_gradio_call(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
fn=modules.ui.wrap_gradio_call_no_job(refresh_available_extensions_for_tags, extra_outputs=[gr.update()]),
inputs=[selected_tags, showing_type, filtering_type, sort_column, search_extensions_text],
outputs=[available_extensions_table, install_result]
)
Expand All @@ -667,7 +667,7 @@ def create_ui():
install_result = gr.HTML(elem_id="extension_install_result")

install_button.click(
fn=modules.ui.wrap_gradio_call(lambda *args: [gr.update(), *install_extension_from_url(*args)], extra_outputs=[gr.update(), gr.update()]),
fn=modules.ui.wrap_gradio_call_no_job(lambda *args: [gr.update(), *install_extension_from_url(*args)], extra_outputs=[gr.update(), gr.update()]),
inputs=[install_dirname, install_url, install_branch],
outputs=[install_url, extensions_table, install_result],
)
Expand Down
4 changes: 2 additions & 2 deletions modules/ui_settings.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import gradio as gr

from modules import ui_common, shared, script_callbacks, scripts, sd_models, sysinfo, timer
from modules.call_queue import wrap_gradio_call
from modules.call_queue import wrap_gradio_call_no_job
from modules.shared import opts
from modules.ui_components import FormRow
from modules.ui_gradio_extensions import reload_javascript
Expand Down Expand Up @@ -289,7 +289,7 @@ def add_quicksettings(self):

def add_functionality(self, demo):
self.submit.click(
fn=wrap_gradio_call(lambda *args: self.run_settings(*args), extra_outputs=[gr.update()]),
fn=wrap_gradio_call_no_job(lambda *args: self.run_settings(*args), extra_outputs=[gr.update()]),
inputs=self.components,
outputs=[self.text_settings, self.result],
)
Expand Down

0 comments on commit 12435dd

Please sign in to comment.