From dbca512154341bb13e1b15d207176f2d403aff30 Mon Sep 17 00:00:00 2001 From: siutin Date: Fri, 3 Feb 2023 03:13:03 +0800 Subject: [PATCH 1/8] add an internal API for obtaining current task id --- modules/progress.py | 8 ++++++++ webui.py | 1 + 2 files changed, 9 insertions(+) diff --git a/modules/progress.py b/modules/progress.py index c69ecf3d..05032ac5 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -4,6 +4,7 @@ import time import gradio as gr from pydantic import BaseModel, Field +from typing import List from modules.shared import opts @@ -37,6 +38,9 @@ def add_task_to_queue(id_job): pending_tasks[id_job] = time.time() +class CurrentTaskResponse(BaseModel): + current_task: str = Field(default=None, title="Task ID", description="id of the current progress task") + 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") @@ -56,6 +60,8 @@ class ProgressResponse(BaseModel): def setup_progress_api(app): return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse) +def setup_current_task_api(app): + return app.add_api_route("/internal/current_task", current_task_api, methods=["GET"], response_model=CurrentTaskResponse) def progressapi(req: ProgressRequest): active = req.id_task == current_task @@ -97,3 +103,5 @@ def progressapi(req: ProgressRequest): return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo) +def current_task_api(): + return CurrentTaskResponse(current_task=current_task) \ No newline at end of file diff --git a/webui.py b/webui.py index b570895f..cca9c77f 100644 --- a/webui.py +++ b/webui.py @@ -279,6 +279,7 @@ def webui(): setup_middleware(app) modules.progress.setup_progress_api(app) + modules.progress.setup_current_task_api(app) if launch_api: create_api(app) From 9407f1731aa8c112ffc0efaa611a76f7fead3d0c Mon Sep 17 00:00:00 2001 From: siutin Date: Mon, 6 Feb 2023 03:53:05 +0800 Subject: [PATCH 2/8] store the last generated result --- modules/call_queue.py | 1 + modules/progress.py | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/modules/call_queue.py b/modules/call_queue.py index 92097c15..30ac26bc 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -37,6 +37,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None): res = func(*args, **kwargs) finally: progress.finish_task(id_task) + progress.set_last_task_result(id_task, res) shared.state.end() diff --git a/modules/progress.py b/modules/progress.py index 05032ac5..27a336ad 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -37,6 +37,16 @@ def finish_task(id_task): def add_task_to_queue(id_job): pending_tasks[id_job] = time.time() +last_task_id = None +last_task_result = None + +def set_last_task_result(id_job, result): + global last_task_id + global last_task_result + + last_task_id = id_job + last_task_result = result + class CurrentTaskResponse(BaseModel): current_task: str = Field(default=None, title="Task ID", description="id of the current progress task") From 4242e194e417ec5008d09ec6d756594ac65f77bd Mon Sep 17 00:00:00 2001 From: siutin Date: Mon, 6 Feb 2023 03:55:31 +0800 Subject: [PATCH 3/8] add a button to restore the current progress --- javascript/progressbar.js | 4 ++-- javascript/ui.js | 36 ++++++++++++++++++++++++++++++++++-- modules/progress.py | 14 ++++++++++++++ modules/ui.py | 34 ++++++++++++++++++++++++++++++---- 4 files changed, 80 insertions(+), 8 deletions(-) diff --git a/javascript/progressbar.js b/javascript/progressbar.js index 4ac9b8db..7ba14192 100644 --- a/javascript/progressbar.js +++ b/javascript/progressbar.js @@ -59,8 +59,8 @@ function setTitle(progress){ } -function randomId(){ - return "task(" + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7)+")" +function randomId(prefix=null){ + return "task(" + (prefix == null ? "" : prefix + "_") + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7)+")" } // starts sending progress requests to "/internal/progress" uri, creating progressbar above progressbarContainer element and diff --git a/javascript/ui.js b/javascript/ui.js index 4a440193..9fe884c0 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -163,7 +163,7 @@ function submit(){ rememberGallerySelection('txt2img_gallery') showSubmitButtons('txt2img', false) - var id = randomId() + var id = randomId("txt2img") requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function(){ showSubmitButtons('txt2img', true) @@ -180,7 +180,7 @@ function submit_img2img(){ rememberGallerySelection('img2img_gallery') showSubmitButtons('img2img', false) - var id = randomId() + var id = randomId("img2img") requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function(){ showSubmitButtons('img2img', true) }) @@ -361,3 +361,35 @@ function selectCheckpoint(name){ desiredCheckpointName = name; gradioApp().getElementById('change_checkpoint').click() } + +function restoreProgress (task_tag) { + + if (task_tag) { + let successHandler = ({ current_task }) => { + if (current_task) { + let _task_tag = ["txt2img", "img2img"].find(t => current_task.startsWith(`task(${t}_`) && current_task.endsWith(")")) + if (!_task_tag) { + console.warn(`task tag ${current_task} not implemented yet`) + return + } + if (task_tag != _task_tag) return + showSubmitButtons(task_tag, false) + requestProgress(current_task, gradioApp().getElementById(`${task_tag}_gallery_container`), gradioApp().getElementById(`${task_tag}_gallery`), function(){ + showSubmitButtons(task_tag, true) + }) + } + } + + let errorHandler = e => window.alert(`invalid internal api respsonse. message: ${e}`) + + fetch("./internal/current_task") + .then(res => res.json()) + .then(successHandler) + .catch(errorHandler) + } + + var res = create_submit_args(arguments) + res[0] = 0 + return res + +} \ No newline at end of file diff --git a/modules/progress.py b/modules/progress.py index 27a336ad..36963c92 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -48,6 +48,20 @@ def set_last_task_result(id_job, result): last_task_result = result +def restore_progress_call(task_tag): + if current_task is None or not current_task[5:-1].startswith(task_tag): + + # image, generation_info, html_info, html_log + return tuple(list([None, None, None, None])) + + else: + + t_task = current_task + while t_task != last_task_id: + time.sleep(2.5) + return last_task_result + + class CurrentTaskResponse(BaseModel): current_task: str = Field(default=None, title="Task ID", description="id of the current progress task") diff --git a/modules/ui.py b/modules/ui.py index 627fbe0b..0133ee12 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -41,6 +41,7 @@ from modules.textual_inversion import textual_inversion import modules.hypernetworks.ui from modules.generation_parameters_copypaste import image_from_url_text import modules.extras +from modules.progress import restore_progress_call warnings.filterwarnings("default" if opts.show_warnings else "ignore", category=UserWarning) @@ -293,6 +294,7 @@ def create_toprow(is_img2img): interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt", elem_classes="generate-box-interrupt") skip = gr.Button('Skip', elem_id=f"{id_part}_skip", elem_classes="generate-box-skip") submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary') + restore_progress = gr.Button('Restore Progress', elem_id=f"{id_part}_restore_progress") skip.click( fn=lambda: shared.state.skip(), @@ -329,7 +331,7 @@ def create_toprow(is_img2img): prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[k for k, v in shared.prompt_styles.styles.items()], value=[], multiselect=True) create_refresh_button(prompt_styles, shared.prompt_styles.reload, lambda: {"choices": [k for k, v in shared.prompt_styles.styles.items()]}, f"refresh_{id_part}_styles") - return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button + return prompt, prompt_styles, negative_prompt, submit, restore_progress, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button def setup_progressbar(*args, **kwargs): @@ -446,7 +448,7 @@ def create_ui(): modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False) with gr.Blocks(analytics_enabled=False) as txt2img_interface: - txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=False) + txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, restore_progress, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=False) dummy_component = gr.Label(visible=False) txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="binary", visible=False) @@ -578,6 +580,18 @@ def create_ui(): res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False) + restore_progress.click( + fn=lambda: restore_progress_call('txt2img'), + _js="() => restoreProgress('txt2img')", + inputs=[], + outputs=[ + txt2img_gallery, + generation_info, + html_info, + html_log, + ] + ) + txt_prompt_img.change( fn=modules.images.image_data, inputs=[ @@ -646,7 +660,7 @@ def create_ui(): modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True) with gr.Blocks(analytics_enabled=False) as img2img_interface: - img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True) + img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, restore_progress, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True) img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="binary", visible=False) @@ -898,6 +912,18 @@ def create_ui(): submit.click(**img2img_args) res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False) + restore_progress.click( + fn=lambda: restore_progress_call('img2img'), + _js="() => restoreProgress('img2img')", + inputs=[], + outputs=[ + img2img_gallery, + generation_info, + html_info, + html_log, + ] + ) + img2img_interrogate.click( fn=lambda *args: process_interrogate(interrogate, *args), **interrogate_args, @@ -1491,7 +1517,7 @@ def create_ui(): gr.HTML(shared.html("licenses.html"), elem_id="licenses") gr.Button(value="Show all pages", elem_id="settings_show_all_pages") - + def unload_sd_weights(): modules.sd_models.unload_model_weights() From e0b58527ff040f9c547ea45b5fcf1bfb7ab23cdd Mon Sep 17 00:00:00 2001 From: siutin Date: Mon, 6 Feb 2023 15:57:26 +0800 Subject: [PATCH 4/8] use condition to wait for result --- modules/call_queue.py | 2 +- modules/progress.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/modules/call_queue.py b/modules/call_queue.py index 30ac26bc..9888109e 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -7,7 +7,7 @@ import time from modules import shared, progress queue_lock = threading.Lock() - +queue_lock_condition = threading.Condition(lock=queue_lock) def wrap_queued_call(func): def f(*args, **kwargs): diff --git a/modules/progress.py b/modules/progress.py index 36963c92..1947c0fd 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -6,6 +6,7 @@ import gradio as gr from pydantic import BaseModel, Field from typing import List +from modules import call_queue from modules.shared import opts import modules.shared as shared @@ -57,8 +58,9 @@ def restore_progress_call(task_tag): else: t_task = current_task - while t_task != last_task_id: - time.sleep(2.5) + with call_queue.queue_lock_condition: + call_queue.queue_lock_condition.wait_for(lambda: t_task == last_task_id) + return last_task_result From 90366b8d8564c6fcbf5899fb31e426b68b04eb7b Mon Sep 17 00:00:00 2001 From: siutin Date: Wed, 29 Mar 2023 00:13:15 +0800 Subject: [PATCH 5/8] tool button --- javascript/hints.js | 2 +- modules/ui.py | 17 +++++++++-------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/javascript/hints.js b/javascript/hints.js index f48a0eb6..7c608311 100644 --- a/javascript/hints.js +++ b/javascript/hints.js @@ -22,7 +22,7 @@ titles = { "\u{1f4cb}": "Apply selected styles to current prompt", "\u{1f4d2}": "Paste available values into the field", "\u{1f3b4}": "Show/hide extra networks", - + "\u{1F300}": "Restore progress", "Inpaint a part of image": "Draw a mask over an image, and the script will regenerate the masked area with content according to prompt", "SD upscale": "Upscale image normally, split result into tiles, improve each tile using img2img, merge whole image back", diff --git a/modules/ui.py b/modules/ui.py index 0133ee12..8fc17ce7 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -82,6 +82,7 @@ apply_style_symbol = '\U0001f4cb' # 📋 clear_prompt_symbol = '\U0001f5d1\ufe0f' # 🗑️ extra_networks_symbol = '\U0001F3B4' # 🎴 switch_values_symbol = '\U000021C5' # ⇅ +restore_progress_symbol = '\U0001F300' # 🌀 def plaintext_to_html(text): @@ -294,7 +295,6 @@ def create_toprow(is_img2img): interrupt = gr.Button('Interrupt', elem_id=f"{id_part}_interrupt", elem_classes="generate-box-interrupt") skip = gr.Button('Skip', elem_id=f"{id_part}_skip", elem_classes="generate-box-skip") submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary') - restore_progress = gr.Button('Restore Progress', elem_id=f"{id_part}_restore_progress") skip.click( fn=lambda: shared.state.skip(), @@ -314,6 +314,7 @@ def create_toprow(is_img2img): extra_networks_button = ToolButton(value=extra_networks_symbol, elem_id=f"{id_part}_extra_networks") prompt_style_apply = ToolButton(value=apply_style_symbol, elem_id=f"{id_part}_style_apply") save_style = ToolButton(value=save_style_symbol, elem_id=f"{id_part}_style_create") + restore_progress_button = ToolButton(value=restore_progress_symbol, elem_id=f"{id_part}_restore_progress") token_counter = gr.HTML(value="0/75", elem_id=f"{id_part}_token_counter", elem_classes=["token-counter"]) token_button = gr.Button(visible=False, elem_id=f"{id_part}_token_button") @@ -331,7 +332,7 @@ def create_toprow(is_img2img): prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[k for k, v in shared.prompt_styles.styles.items()], value=[], multiselect=True) create_refresh_button(prompt_styles, shared.prompt_styles.reload, lambda: {"choices": [k for k, v in shared.prompt_styles.styles.items()]}, f"refresh_{id_part}_styles") - return prompt, prompt_styles, negative_prompt, submit, restore_progress, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button + return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button, restore_progress_button def setup_progressbar(*args, **kwargs): @@ -448,7 +449,7 @@ def create_ui(): modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False) with gr.Blocks(analytics_enabled=False) as txt2img_interface: - txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, restore_progress, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=False) + txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button, restore_progress_button = create_toprow(is_img2img=False) dummy_component = gr.Label(visible=False) txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="binary", visible=False) @@ -580,8 +581,8 @@ def create_ui(): res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False) - restore_progress.click( - fn=lambda: restore_progress_call('txt2img'), + restore_progress_button.click( + fn=lambda: restore_progress_call(), _js="() => restoreProgress('txt2img')", inputs=[], outputs=[ @@ -660,7 +661,7 @@ def create_ui(): modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True) with gr.Blocks(analytics_enabled=False) as img2img_interface: - img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, restore_progress, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True) + img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, extra_networks_button, token_counter, token_button, negative_token_counter, negative_token_button, restore_progress_button = create_toprow(is_img2img=True) img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="binary", visible=False) @@ -912,8 +913,8 @@ def create_ui(): submit.click(**img2img_args) res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False) - restore_progress.click( - fn=lambda: restore_progress_call('img2img'), + restore_progress_button.click( + fn=lambda: restore_progress_call(), _js="() => restoreProgress('img2img')", inputs=[], outputs=[ From 70ab21e67d128b953fbf4a360e02ac783f40dd55 Mon Sep 17 00:00:00 2001 From: siutin Date: Wed, 29 Mar 2023 00:17:19 +0800 Subject: [PATCH 6/8] keep randomId simpler --- javascript/progressbar.js | 4 ++-- javascript/ui.js | 10 ++-------- modules/progress.py | 4 ++-- 3 files changed, 6 insertions(+), 12 deletions(-) diff --git a/javascript/progressbar.js b/javascript/progressbar.js index 7ba14192..4ac9b8db 100644 --- a/javascript/progressbar.js +++ b/javascript/progressbar.js @@ -59,8 +59,8 @@ function setTitle(progress){ } -function randomId(prefix=null){ - return "task(" + (prefix == null ? "" : prefix + "_") + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7)+")" +function randomId(){ + return "task(" + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7)+")" } // starts sending progress requests to "/internal/progress" uri, creating progressbar above progressbarContainer element and diff --git a/javascript/ui.js b/javascript/ui.js index 9fe884c0..c9df066d 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -163,7 +163,7 @@ function submit(){ rememberGallerySelection('txt2img_gallery') showSubmitButtons('txt2img', false) - var id = randomId("txt2img") + var id = randomId() requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function(){ showSubmitButtons('txt2img', true) @@ -180,7 +180,7 @@ function submit_img2img(){ rememberGallerySelection('img2img_gallery') showSubmitButtons('img2img', false) - var id = randomId("img2img") + var id = randomId() requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function(){ showSubmitButtons('img2img', true) }) @@ -367,12 +367,6 @@ function restoreProgress (task_tag) { if (task_tag) { let successHandler = ({ current_task }) => { if (current_task) { - let _task_tag = ["txt2img", "img2img"].find(t => current_task.startsWith(`task(${t}_`) && current_task.endsWith(")")) - if (!_task_tag) { - console.warn(`task tag ${current_task} not implemented yet`) - return - } - if (task_tag != _task_tag) return showSubmitButtons(task_tag, false) requestProgress(current_task, gradioApp().getElementById(`${task_tag}_gallery_container`), gradioApp().getElementById(`${task_tag}_gallery`), function(){ showSubmitButtons(task_tag, true) diff --git a/modules/progress.py b/modules/progress.py index 1947c0fd..e99267f5 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -49,8 +49,8 @@ def set_last_task_result(id_job, result): last_task_result = result -def restore_progress_call(task_tag): - if current_task is None or not current_task[5:-1].startswith(task_tag): +def restore_progress_call(): + if current_task is None: # image, generation_info, html_info, html_log return tuple(list([None, None, None, None])) From 984970068c2bdc14cff266129ca25a26fbccbf2e Mon Sep 17 00:00:00 2001 From: siutin Date: Mon, 17 Apr 2023 01:06:28 +0800 Subject: [PATCH 7/8] multi users support --- modules/call_queue.py | 23 ++++++++++------- modules/progress.py | 60 +++++++++++++++++++++++++++++++------------ modules/ui.py | 4 +-- 3 files changed, 59 insertions(+), 28 deletions(-) diff --git a/modules/call_queue.py b/modules/call_queue.py index 9888109e..632afcdd 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -4,6 +4,7 @@ import threading import traceback import time +import gradio as gr from modules import shared, progress queue_lock = threading.Lock() @@ -20,41 +21,45 @@ def wrap_queued_call(func): def wrap_gradio_gpu_call(func, extra_outputs=None): - def f(*args, **kwargs): + def f(request: gr.Request, *args, **kwargs): + user = request.username # 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) + progress.add_task_to_queue(user, id_task) else: id_task = None with queue_lock: shared.state.begin() - progress.start_task(id_task) + progress.start_task(user, id_task) try: res = func(*args, **kwargs) finally: - progress.finish_task(id_task) - progress.set_last_task_result(id_task, res) + progress.finish_task(user, id_task) + progress.set_last_task_result(user, id_task, res) shared.state.end() return res - return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True) + return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True, add_request=True) -def wrap_gradio_call(func, extra_outputs=None, add_stats=False): - def f(*args, extra_outputs_array=extra_outputs, **kwargs): +def wrap_gradio_call(func, extra_outputs=None, add_stats=False, add_request=False): + def f(request: gr.Request, *args, extra_outputs_array=extra_outputs, **kwargs): run_memmon = shared.opts.memmon_poll_rate > 0 and not shared.mem_mon.disabled and add_stats if run_memmon: shared.mem_mon.monitor() t = time.perf_counter() try: - res = list(func(*args, **kwargs)) + if add_request: + res = list(func(request, *args, **kwargs)) + else: + res = list(func(*args, **kwargs)) except Exception as e: # When printing out our debug argument list, do not print out more than a MB of text max_debug_str_len = 131072 # (1024*1024)/8 diff --git a/modules/progress.py b/modules/progress.py index e99267f5..13568701 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -4,7 +4,9 @@ import time import gradio as gr from pydantic import BaseModel, Field -from typing import List +from typing import Optional +from fastapi import Depends, Security +from fastapi.security import APIKeyCookie from modules import call_queue from modules.shared import opts @@ -12,57 +14,71 @@ from modules.shared import opts import modules.shared as shared +current_task_user = None current_task = None pending_tasks = {} finished_tasks = [] -def start_task(id_task): +def start_task(user, id_task): global current_task + global current_task_user + current_task_user = user current_task = id_task - pending_tasks.pop(id_task, None) + pending_tasks.pop((user, id_task), None) -def finish_task(id_task): +def finish_task(user, id_task): global current_task + global current_task_user if current_task == id_task: current_task = None - finished_tasks.append(id_task) + if current_task_user == user: + current_task_user = None + + finished_tasks.append((user, id_task)) if len(finished_tasks) > 16: finished_tasks.pop(0) -def add_task_to_queue(id_job): - pending_tasks[id_job] = time.time() +def add_task_to_queue(user, id_job): + pending_tasks[(user, id_job)] = time.time() last_task_id = None last_task_result = None +last_task_user = None + +def set_last_task_result(user, id_job, result): -def set_last_task_result(id_job, result): global last_task_id global last_task_result + global last_task_user last_task_id = id_job last_task_result = result + last_task_user = user -def restore_progress_call(): +def restore_progress_call(request: gr.Request): if current_task is None: # image, generation_info, html_info, html_log return tuple(list([None, None, None, None])) else: + user = request.username - t_task = current_task - with call_queue.queue_lock_condition: - call_queue.queue_lock_condition.wait_for(lambda: t_task == last_task_id) + if current_task_user == user: + t_task = current_task + with call_queue.queue_lock_condition: + call_queue.queue_lock_condition.wait_for(lambda: t_task == last_task_id) - return last_task_result + return last_task_result + return tuple(list([None, None, None, None])) class CurrentTaskResponse(BaseModel): current_task: str = Field(default=None, title="Task ID", description="id of the current progress task") @@ -87,6 +103,19 @@ def setup_progress_api(app): return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse) def setup_current_task_api(app): + + def get_current_user(token: Optional[str] = Security(APIKeyCookie(name="access-token", auto_error=False))): + return None if token is None else app.tokens.get(token) + + def current_task_api(current_user: str = Depends(get_current_user)): + + if app.auth is None or current_task_user == current_user: + current_user_task = current_task + else: + current_user_task = None + + return CurrentTaskResponse(current_task=current_user_task) + return app.add_api_route("/internal/current_task", current_task_api, methods=["GET"], response_model=CurrentTaskResponse) def progressapi(req: ProgressRequest): @@ -127,7 +156,4 @@ def progressapi(req: ProgressRequest): else: live_preview = None - return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo) - -def current_task_api(): - return CurrentTaskResponse(current_task=current_task) \ No newline at end of file + return ProgressResponse(active=active, queued=queued, completed=completed, progress=progress, eta=eta, live_preview=live_preview, id_live_preview=id_live_preview, textinfo=shared.state.textinfo) \ No newline at end of file diff --git a/modules/ui.py b/modules/ui.py index 8fc17ce7..a7b3cccb 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -582,7 +582,7 @@ def create_ui(): res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False) restore_progress_button.click( - fn=lambda: restore_progress_call(), + fn=restore_progress_call, _js="() => restoreProgress('txt2img')", inputs=[], outputs=[ @@ -914,7 +914,7 @@ def create_ui(): res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False) restore_progress_button.click( - fn=lambda: restore_progress_call(), + fn=restore_progress_call, _js="() => restoreProgress('img2img')", inputs=[], outputs=[ From 3e5b3c79e49ec3a10174c250815dabce6efdddca Mon Sep 17 00:00:00 2001 From: siutin Date: Mon, 17 Apr 2023 11:50:08 +0800 Subject: [PATCH 8/8] replace with #wrap_session_call --- modules/call_queue.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/modules/call_queue.py b/modules/call_queue.py index 632afcdd..43f6ebe0 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -10,6 +10,11 @@ from modules import shared, progress queue_lock = threading.Lock() queue_lock_condition = threading.Condition(lock=queue_lock) +def wrap_session_call(func): + def f(request: gr.Request, *args, **kwargs): + return func(request, *args, **kwargs) + return f + def wrap_queued_call(func): def f(*args, **kwargs): with queue_lock: @@ -45,21 +50,18 @@ def wrap_gradio_gpu_call(func, extra_outputs=None): return res - return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True, add_request=True) + return wrap_session_call(wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True)) -def wrap_gradio_call(func, extra_outputs=None, add_stats=False, add_request=False): - def f(request: gr.Request, *args, extra_outputs_array=extra_outputs, **kwargs): +def wrap_gradio_call(func, extra_outputs=None, add_stats=False): + 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 if run_memmon: shared.mem_mon.monitor() t = time.perf_counter() try: - if add_request: - res = list(func(request, *args, **kwargs)) - else: - res = list(func(*args, **kwargs)) + res = list(func(*args, **kwargs)) except Exception as e: # When printing out our debug argument list, do not print out more than a MB of text max_debug_str_len = 131072 # (1024*1024)/8