diff options
Diffstat (limited to 'modules/progress.py')
| -rw-r--r-- | modules/progress.py | 60 |
1 files changed, 43 insertions, 17 deletions
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 |