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=[