add page_size & auto requeue failed tasks settings

pull/127/head^2
Tung Nguyen 2023-11-07 15:13:21 +07:00
parent 11f04dd105
commit c32395e42a
6 changed files with 155 additions and 66 deletions

View File

@ -316,6 +316,19 @@ def regsiter_apis(app: App, task_runner: TaskRunner):
return {"success": True, "message": "Task requeued"}
@app.post("/agent-scheduler/v1/task/requeue-failed", dependencies=deps)
def requeue_failed_tasks():
failed_tasks = task_manager.get_tasks(status=TaskStatus.FAILED)
if (len(failed_tasks)) == 0:
return {"success": False, "message": "No failed tasks"}
for task in failed_tasks:
task.status = TaskStatus.PENDING
task.result = None
task_manager.update_task(task)
return {"success": True, "message": f"Requeued {len(failed_tasks)} failed tasks"}
@app.post("/agent-scheduler/v1/delete/{id}", dependencies=deps, deprecated=True)
@app.delete("/agent-scheduler/v1/task/{id}", dependencies=deps)
def delete_task(id: str):

View File

@ -68,7 +68,7 @@ class TaskRunner:
self.__current_thread: threading.Thread = None
self.__api = Api(FastAPI(), queue_lock)
self.__saved_images_path: List[Tuple[str, str]] = []
self.__saved_images_path: List[str] = []
script_callbacks.on_image_saved(self.__on_image_saved)
self.script_callbacks = {
@ -356,10 +356,15 @@ class TaskRunner:
log.error(f"[AgentScheduler] Task {task_id} failed: {res}")
log.debug(traceback.format_exc())
task.status = TaskStatus.FAILED
task.result = str(res) if res else None
task_manager.update_task(task)
self.__run_callbacks("task_finished", task_id, status=TaskStatus.FAILED, **task_meta)
if getattr(shared.opts, "queue_automatic_requeue_failed_task", False):
log.info(f"[AgentScheduler] Requeue task {task_id}")
task.status = TaskStatus.PENDING
task_manager.update_task(task)
else:
task.status = TaskStatus.FAILED
task.result = str(res) if res else None
task_manager.update_task(task)
self.__run_callbacks("task_finished", task_id, status=TaskStatus.FAILED, **task_meta)
else:
is_interrupted = self.interrupted == task_id
if is_interrupted:
@ -373,13 +378,11 @@ class TaskRunner:
**task_meta,
)
else:
geninfo = json.loads(res)
result = {
"images": [],
"infotexts": [],
"images": self.__saved_images_path.copy(),
"geninfo": geninfo,
}
for filename, pnginfo in self.__saved_images_path:
result["images"].append(filename)
result["infotexts"].append(pnginfo)
task.status = TaskStatus.DONE
task.result = json.dumps(result)
@ -516,7 +519,10 @@ class TaskRunner:
self.__run_callbacks("task_cleared")
def __on_image_saved(self, data: script_callbacks.ImageSaveParams):
self.__saved_images_path.append((data.filename, data.pnginfo.get("parameters", "")))
if data.filename.startswith(data.p.outpath_grids):
self.__saved_images_path.insert(0, data.filename)
else:
self.__saved_images_path.append(data.filename)
def on_task_registered(self, callback: Callable):
"""Callback when a task is registered

File diff suppressed because one or more lines are too long

View File

@ -281,20 +281,26 @@ def infotexts_to_geninfo(infotexts: List[str]):
geninfo = {"infotexts": infotexts, "all_prompts": all_promts, "all_seeds": all_seeds, "index_of_first_image": 0}
for infotext in infotexts:
# Dynamic prompt breaks layout of infotext
if "Template: " in infotext:
lines = infotext.split("\n")
lines = [l for l in lines if not (l.startswith("Template: ") or l.startswith("Negative Template: "))]
infotext = "\n".join(lines)
params = parse_generation_parameters(infotext)
if "prompt" not in params:
geninfo["prompt"] = params["Prompt"]
geninfo["negative_prompt"] = params["Negative prompt"]
geninfo["seed"] = params["Seed"]
geninfo["sampler_name"] = params["Sampler"]
geninfo["cfg_scale"] = params["CFG scale"]
geninfo["steps"] = params["Steps"]
geninfo["width"] = params["Size-1"]
geninfo["height"] = params["Size-2"]
if "prompt" not in geninfo:
geninfo["prompt"] = params.get("Prompt", "")
geninfo["negative_prompt"] = params.get("Negative prompt", "")
geninfo["seed"] = params.get("Seed", "-1")
geninfo["sampler_name"] = params.get("Sampler", "")
geninfo["cfg_scale"] = params.get("CFG scale", "")
geninfo["steps"] = params.get("Steps", "0")
geninfo["width"] = params.get("Size-1", "512")
geninfo["height"] = params.get("Size-2", "512")
all_promts.append(params["Prompt"])
all_seeds.append(params["Seed"])
all_promts.append(params.get("Prompt", ""))
all_seeds.append(params.get("Seed", "-1"))
return geninfo
@ -365,6 +371,8 @@ def remove_old_tasks():
def on_ui_tab(**_kwargs):
grid_page_size = getattr(shared.opts, "queue_grid_page_size", 0)
with gr.Blocks(analytics_enabled=False) as scheduler_tab:
with gr.Tabs(elem_id="agent_scheduler_tabs"):
with gr.Tab("Task Queue", id=0, elem_id="agent_scheduler_pending_tasks_tab"):
@ -405,7 +413,9 @@ def on_ui_tab(**_kwargs):
min_width=0,
elem_id="agent_scheduler_action_search",
)
gr.HTML('<div id="agent_scheduler_pending_tasks_grid" class="ag-theme-gradio"></div>')
gr.HTML(
f'<div id="agent_scheduler_pending_tasks_grid" class="ag-theme-gradio" data-page-size="{grid_page_size}"></div>'
)
with gr.Column(scale=1):
gr.Gallery(
elem_id="agent_scheduler_current_task_images",
@ -418,6 +428,11 @@ def on_ui_tab(**_kwargs):
with gr.Row(elem_id="agent_scheduler_history_wrapper"):
with gr.Column(scale=1):
with gr.Row(elem_id="agent_scheduler_history_actions", elem_classes="flex-row"):
gr.Button(
"Requeue Failed",
elem_id="agent_scheduler_action_requeue",
variant="primary",
)
gr.Button(
"Refresh",
elem_id="agent_scheduler_action_refresh_history",
@ -446,7 +461,9 @@ def on_ui_tab(**_kwargs):
min_width=0,
elem_id="agent_scheduler_action_search_history",
)
gr.HTML('<div id="agent_scheduler_history_tasks_grid" class="ag-theme-gradio"></div>')
gr.HTML(
f'<div id="agent_scheduler_history_tasks_grid" class="ag-theme-gradio" data-page-size="{grid_page_size}"></div>'
)
with gr.Column(scale=1, elem_id="agent_scheduler_history_results"):
galerry = gr.Gallery(
elem_id="agent_scheduler_history_gallery",
@ -567,6 +584,16 @@ def on_ui_settings():
section=section,
),
)
shared.opts.add_option(
"queue_button_hide_checkpoint",
shared.OptionInfo(
True,
"Hide the custom checkpoint dropdown",
gr.Checkbox,
{},
section=section,
),
)
shared.opts.add_option(
"queue_button_placement",
shared.OptionInfo(
@ -583,12 +610,17 @@ def on_ui_settings():
),
)
shared.opts.add_option(
"queue_button_hide_checkpoint",
"queue_ui_placement",
shared.OptionInfo(
True,
"Hide the checkpoint dropdown",
gr.Checkbox,
{},
ui_placement_as_tab,
"Task queue UI placement",
gr.Radio,
lambda: {
"choices": [
ui_placement_as_tab,
ui_placement_append_to_main,
]
},
section=section,
),
)
@ -604,6 +636,16 @@ def on_ui_settings():
section=section,
),
)
shared.opts.add_option(
"queue_automatic_requeue_failed_task",
shared.OptionInfo(
False,
"Auto requeue failed tasks",
gr.Checkbox,
{},
section=section,
),
)
def enqueue_keyboard_shortcut(disabled: bool, modifiers, key_code: str):
if disabled:
@ -679,17 +721,12 @@ def on_ui_settings():
)
shared.opts.add_option(
"queue_ui_placement",
"queue_grid_page_size",
shared.OptionInfo(
ui_placement_as_tab,
"Task queue UI placement",
gr.Radio,
lambda: {
"choices": [
ui_placement_as_tab,
ui_placement_append_to_main,
]
},
0,
"Task list page size (0 for auto)",
gr.Number,
{"minimum": 0, "maximum": 200},
section=section,
),
)

View File

@ -901,6 +901,16 @@ function initPendingTab() {
const eGridDiv = gradioApp().querySelector<HTMLDivElement>(
'#agent_scheduler_pending_tasks_grid'
)!;
if (typeof eGridDiv.dataset.pageSize === 'string') {
const pageSize = parseInt(eGridDiv.dataset.pageSize, 10);
if (pageSize > 0) {
gridOptions.paginationAutoPageSize = false;
gridOptions.paginationPageSize = pageSize;
}
}
new Grid(eGridDiv, gridOptions);
}
@ -919,6 +929,12 @@ function initHistoryTab() {
if (!confirm('Are you sure you want to clear the history?')) return;
store.clearHistory().then(notify);
});
const requeueButton = gradioApp().querySelector<HTMLButtonElement>(
'#agent_scheduler_action_requeue'
)!;
requeueButton.addEventListener('click', () => {
store.requeueFailedTasks().then(notify);
});
const resultTaskId = gradioApp().querySelector<HTMLTextAreaElement>(
'#agent_scheduler_history_selected_task textarea'
@ -1113,6 +1129,16 @@ function initHistoryTab() {
const eGridDiv = gradioApp().querySelector<HTMLDivElement>(
'#agent_scheduler_history_tasks_grid'
)!;
if (typeof eGridDiv.dataset.pageSize === 'string') {
const pageSize = parseInt(eGridDiv.dataset.pageSize, 10);
if (pageSize > 0) {
gridOptions.paginationAutoPageSize = false;
gridOptions.paginationPageSize = pageSize;
}
}
new Grid(eGridDiv, gridOptions);
}

View File

@ -14,6 +14,7 @@ type HistoryTasksActions = {
bookmarkTask: (id: string, bookmarked: boolean) => Promise<ResponseStatus>;
renameTask: (id: string, name: string) => Promise<ResponseStatus>;
requeueTask: (id: string) => Promise<ResponseStatus>;
requeueFailedTasks: () => Promise<ResponseStatus>;
clearHistory: () => Promise<ResponseStatus>;
};
@ -55,6 +56,12 @@ export const createHistoryTasksStore = (initialState: HistoryTasksState) => {
response.json()
);
},
requeueFailedTasks: async () => {
return fetch('/agent-scheduler/v1/task/requeue-failed', { method: 'POST' }).then(response => {
actions.refresh();
return response.json();
});
},
clearHistory: async () => {
return fetch('/agent-scheduler/v1/history/clear', { method: 'POST' }).then(response => {
actions.refresh();