mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2025-08-03 19:02:27 +00:00
add allow specify the task id and get the location of task in the queue of pending task
This commit is contained in:
@@ -8,10 +8,13 @@ from pydantic import BaseModel, Field
|
||||
from modules.shared import opts
|
||||
|
||||
import modules.shared as shared
|
||||
|
||||
from collections import OrderedDict
|
||||
import string
|
||||
import random
|
||||
from typing import List
|
||||
|
||||
current_task = None
|
||||
pending_tasks = {}
|
||||
pending_tasks = OrderedDict()
|
||||
finished_tasks = []
|
||||
recorded_results = []
|
||||
recorded_results_limit = 2
|
||||
@@ -34,6 +37,11 @@ def finish_task(id_task):
|
||||
if len(finished_tasks) > 16:
|
||||
finished_tasks.pop(0)
|
||||
|
||||
def create_task_id(task_type):
|
||||
N = 7
|
||||
res = ''.join(random.choices(string.ascii_uppercase +
|
||||
string.digits, k=N))
|
||||
return f"task({task_type}-{res})"
|
||||
|
||||
def record_results(id_task, res):
|
||||
recorded_results.append((id_task, res))
|
||||
@@ -44,6 +52,9 @@ def record_results(id_task, res):
|
||||
def add_task_to_queue(id_job):
|
||||
pending_tasks[id_job] = time.time()
|
||||
|
||||
class PendingTasksResponse(BaseModel):
|
||||
size: int = Field(title="Pending task size")
|
||||
tasks: List[str] = Field(title="Pending task ids")
|
||||
|
||||
class ProgressRequest(BaseModel):
|
||||
id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for")
|
||||
@@ -63,8 +74,14 @@ class ProgressResponse(BaseModel):
|
||||
|
||||
|
||||
def setup_progress_api(app):
|
||||
app.add_api_route("/internal/pendingTasks", get_pending_tasks, methods=["GET"])
|
||||
return app.add_api_route("/internal/progress", progressapi, methods=["POST"], response_model=ProgressResponse)
|
||||
|
||||
def get_pending_tasks():
|
||||
pending_tasks_ids = [x for x in pending_tasks]
|
||||
pending_len = len(pending_tasks_ids)
|
||||
return PendingTasksResponse(size=pending_len, tasks=pending_tasks_ids)
|
||||
|
||||
|
||||
def progressapi(req: ProgressRequest):
|
||||
active = req.id_task == current_task
|
||||
|
Reference in New Issue
Block a user