415 lines
15 KiB
Python
415 lines
15 KiB
Python
import os
|
|
import json
|
|
import gradio as gr
|
|
from PIL import Image
|
|
|
|
from modules import shared, script_callbacks, scripts
|
|
from modules.shared import list_checkpoint_tiles, refresh_checkpoints
|
|
from modules.ui import create_refresh_button
|
|
from modules.generation_parameters_copypaste import (
|
|
create_buttons,
|
|
register_paste_params_button,
|
|
ParamBinding,
|
|
)
|
|
|
|
from scripts.task_runner import TaskRunner, get_instance, task_history_retenion_map
|
|
from scripts.helpers import log, compare_components_with_ids, get_components_by_ids
|
|
from scripts.db import init, task_manager, TaskStatus
|
|
|
|
task_runner: TaskRunner = None
|
|
initialized = False
|
|
|
|
checkpoint_current = "Current Checkpoint"
|
|
checkpoint_runtime = "Runtime Checkpoint"
|
|
|
|
placement_under_generate = "Under Generate button"
|
|
placement_between_prompt_and_generate = "Between Prompt and Generate button"
|
|
|
|
task_filter_choices = ["All", "Bookmarked", "Done", "Failed", "Interrupted"]
|
|
|
|
|
|
class Script(scripts.Script):
|
|
def __init__(self):
|
|
super().__init__()
|
|
script_callbacks.on_app_started(lambda block, _: self.on_app_started(block))
|
|
self.checkpoint_override = checkpoint_current
|
|
self.generate_button = None
|
|
|
|
def title(self):
|
|
return "Agent Scheduler"
|
|
|
|
def show(self, is_img2img):
|
|
return True
|
|
|
|
def on_checkpoint_changed(self, checkpoint):
|
|
self.checkpoint_override = checkpoint
|
|
|
|
def after_component(self, component, **_kwargs):
|
|
elem_id = "txt2img_generate" if self.is_txt2img else "img2img_generate"
|
|
|
|
if component.elem_id == elem_id:
|
|
self.generate_button = component
|
|
|
|
def on_app_started(self, block):
|
|
if self.generate_button is not None:
|
|
self.add_enqueue_button(block, self.generate_button)
|
|
|
|
def add_enqueue_button(self, root: gr.Blocks, generate: gr.Button):
|
|
is_img2img = self.is_img2img
|
|
dependencies: list[dict] = [
|
|
x
|
|
for x in root.dependencies
|
|
if x["trigger"] == "click" and generate._id in x["targets"]
|
|
]
|
|
|
|
dependency: dict = None
|
|
cnet_dependency: dict = None
|
|
UiControlNetUnit = None
|
|
for d in dependencies:
|
|
if len(d["outputs"]) == 1:
|
|
outputs = get_components_by_ids(root, d["outputs"])
|
|
output = outputs[0]
|
|
if (
|
|
isinstance(output, gr.State)
|
|
and type(output.value).__name__ == "UiControlNetUnit"
|
|
):
|
|
cnet_dependency = d
|
|
UiControlNetUnit = type(output.value)
|
|
|
|
elif len(d["outputs"]) == 4:
|
|
dependency = d
|
|
|
|
fn_block = next(
|
|
fn
|
|
for fn in root.fns
|
|
if compare_components_with_ids(fn.inputs, dependency["inputs"])
|
|
)
|
|
fn = self.wrap_register_ui_task()
|
|
args = dict(
|
|
fn=fn,
|
|
_js="submit_enqueue_img2img" if is_img2img else "submit_enqueue",
|
|
inputs=fn_block.inputs,
|
|
outputs=None,
|
|
show_progress=False,
|
|
)
|
|
|
|
id_part = "img2img" if is_img2img else "txt2img"
|
|
with root:
|
|
with gr.Row(elem_id=f"{id_part}_enqueue_wrapper") as row:
|
|
if not shared.opts.queue_button_hide_checkpoint:
|
|
checkpoint = gr.Dropdown(
|
|
choices=get_checkpoint_choices(),
|
|
value=checkpoint_current,
|
|
show_label=False,
|
|
interactive=True,
|
|
)
|
|
create_refresh_button(
|
|
checkpoint,
|
|
refresh_checkpoints,
|
|
lambda: {"choices": get_checkpoint_choices()},
|
|
f"refresh_{id_part}_checkpoint",
|
|
)
|
|
checkpoint.change(
|
|
fn=self.on_checkpoint_changed, inputs=[checkpoint]
|
|
)
|
|
|
|
submit = gr.Button(
|
|
"Enqueue", elem_id=f"{id_part}_enqueue", variant="primary"
|
|
)
|
|
submit.click(**args)
|
|
|
|
# relocation the enqueue button
|
|
root.children.pop()
|
|
if shared.opts.queue_button_placement == placement_between_prompt_and_generate:
|
|
if is_img2img:
|
|
# add to the iterrogate div
|
|
parent = generate.parent.parent.parent.children[1]
|
|
parent.add(row)
|
|
else:
|
|
# insert after the prompts
|
|
parent = generate.parent.parent.parent
|
|
row.parent = parent
|
|
parent.children.insert(1, row)
|
|
else:
|
|
# insert after the generate button
|
|
parent = generate.parent.parent
|
|
parent.children.insert(1, row)
|
|
|
|
if cnet_dependency is not None:
|
|
cnet_fn_block = next(
|
|
fn
|
|
for fn in root.fns
|
|
if compare_components_with_ids(fn.inputs, cnet_dependency["inputs"])
|
|
)
|
|
with root:
|
|
submit.click(
|
|
fn=UiControlNetUnit,
|
|
inputs=cnet_fn_block.inputs,
|
|
outputs=cnet_fn_block.outputs,
|
|
queue=False,
|
|
)
|
|
|
|
def wrap_register_ui_task(self):
|
|
def f(*args, **kwargs):
|
|
if len(args) == 0 and len(kwargs) == 0:
|
|
raise Exception("Invalid call")
|
|
|
|
if len(args) > 0 and type(args[0]) == str:
|
|
task_id = args[0]
|
|
else:
|
|
# not a task, exit
|
|
return (None, "", "<p>Invalid params</p>", "")
|
|
|
|
checkpoint = None
|
|
if self.checkpoint_override == checkpoint_current:
|
|
checkpoint = shared.sd_model.sd_checkpoint_info.title
|
|
elif self.checkpoint_override != checkpoint_runtime:
|
|
checkpoint = self.checkpoint_override
|
|
|
|
task_runner.register_ui_task(
|
|
task_id, self.is_img2img, *args, checkpoint=checkpoint
|
|
)
|
|
task_runner.execute_pending_tasks_threading()
|
|
|
|
return f
|
|
|
|
|
|
def get_checkpoint_choices():
|
|
choices = [checkpoint_current, checkpoint_runtime]
|
|
choices.extend(list_checkpoint_tiles())
|
|
return choices
|
|
|
|
|
|
def get_task_results(task_id: str, image_idx: int = None):
|
|
task = task_manager.get_task(task_id)
|
|
|
|
galerry = None
|
|
infotexts = None
|
|
if task is None:
|
|
pass
|
|
elif task.status != TaskStatus.DONE:
|
|
infotexts = f"Status: {task.status}"
|
|
if task.status == TaskStatus.FAILED and task.result:
|
|
infotexts += f"\nError: {task.result}"
|
|
elif task.status == TaskStatus.DONE:
|
|
try:
|
|
result: dict = json.loads(task.result)
|
|
images = result.get("images", [])
|
|
infos = result.get("infotexts", [])
|
|
galerry = (
|
|
[Image.open(i) for i in images if os.path.exists(i)]
|
|
if image_idx is None
|
|
else gr.update()
|
|
)
|
|
idx = image_idx if image_idx is not None else 0
|
|
if len(infos) == len(images):
|
|
infotexts = infos[idx]
|
|
else:
|
|
infotexts = "\n".join(infos).split("Prompt: ")[1:][idx]
|
|
|
|
except Exception as e:
|
|
log.error(f"[AgentScheduler] Failed to load task result")
|
|
log.error(e)
|
|
infotexts = f"Failed to load task result: {str(e)}"
|
|
|
|
res = (
|
|
gr.Textbox.update(infotexts, visible=infotexts is not None),
|
|
gr.Row.update(visible=galerry is not None),
|
|
)
|
|
return res if image_idx is not None else (galerry,) + res
|
|
|
|
|
|
def on_ui_tab(**_kwargs):
|
|
global initialized
|
|
if not initialized:
|
|
initialized = True
|
|
init()
|
|
|
|
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"
|
|
):
|
|
with gr.Row(elem_id="agent_scheduler_pending_tasks_wrapper"):
|
|
with gr.Column(scale=1):
|
|
with gr.Group(elem_id="agent_scheduler_pending_tasks_actions"):
|
|
paused = shared.opts.queue_paused
|
|
|
|
gr.Button(
|
|
"Pause",
|
|
elem_id="agent_scheduler_action_pause",
|
|
variant="stop",
|
|
visible=not paused,
|
|
)
|
|
gr.Button(
|
|
"Resume",
|
|
elem_id="agent_scheduler_action_resume",
|
|
variant="primary",
|
|
visible=paused,
|
|
)
|
|
gr.Button(
|
|
"Refresh",
|
|
elem_id="agent_scheduler_action_refresh",
|
|
elem_classes="agent_scheduler_action_refresh",
|
|
variant="secondary",
|
|
)
|
|
gr.HTML('<div id="agent_scheduler_action_search"></div>')
|
|
gr.HTML(
|
|
'<div id="agent_scheduler_pending_tasks_grid" class="ag-theme-alpine"></div>'
|
|
)
|
|
with gr.Column(scale=1):
|
|
gr.Gallery(
|
|
elem_id="agent_scheduler_current_task_images",
|
|
label="Output",
|
|
show_label=False,
|
|
).style(columns=2, object_fit="contain")
|
|
with gr.Tab("Task History", id=1, elem_id="agent_scheduler_history_tab"):
|
|
with gr.Row(elem_id="agent_scheduler_history_wrapper"):
|
|
with gr.Column(scale=1):
|
|
with gr.Group(elem_id="agent_scheduler_history_actions"):
|
|
gr.Button(
|
|
"Refresh",
|
|
elem_id="agent_scheduler_action_refresh_history",
|
|
elem_classes="agent_scheduler_action_refresh",
|
|
variant="secondary",
|
|
)
|
|
status = gr.Dropdown(
|
|
elem_id="agent_scheduler_status_filter",
|
|
choices=task_filter_choices,
|
|
value="All",
|
|
show_label=False,
|
|
)
|
|
gr.HTML(
|
|
'<div id="agent_scheduler_action_search_history"></div>'
|
|
)
|
|
gr.HTML(
|
|
'<div id="agent_scheduler_history_tasks_grid" class="ag-theme-alpine"></div>'
|
|
)
|
|
with gr.Column(scale=1, elem_id="agent_scheduler_history_results"):
|
|
galerry = gr.Gallery(
|
|
elem_id="agent_scheduler_history_gallery",
|
|
label="Output",
|
|
show_label=False,
|
|
).style(columns=2, object_fit="contain", preview=True)
|
|
gen_info = gr.Textbox(
|
|
label="Generation Info",
|
|
elem_id=f"agent_scheduler_history_gen_info",
|
|
interactive=False,
|
|
visible=True,
|
|
lines=3,
|
|
)
|
|
with gr.Row(
|
|
elem_id="agent_scheduler_history_result_actions",
|
|
visible=False,
|
|
) as result_actions:
|
|
try:
|
|
send_to_buttons = create_buttons(
|
|
["txt2img", "img2img", "inpaint", "extras"]
|
|
)
|
|
except:
|
|
pass
|
|
selected_task = gr.Textbox(
|
|
elem_id="agent_scheduler_history_selected_task",
|
|
visible=False,
|
|
show_label=False,
|
|
)
|
|
selected_task_id = gr.Textbox(
|
|
elem_id="agent_scheduler_history_selected_image",
|
|
visible=False,
|
|
show_label=False,
|
|
)
|
|
|
|
# register event handlers
|
|
status.change(
|
|
fn=lambda x: None,
|
|
_js="agent_scheduler_status_filter_changed",
|
|
inputs=[status],
|
|
)
|
|
selected_task.change(
|
|
fn=get_task_results,
|
|
inputs=[selected_task],
|
|
outputs=[galerry, gen_info, result_actions],
|
|
)
|
|
selected_task_id.change(
|
|
fn=lambda x, y: get_task_results(x, image_idx=int(y)),
|
|
inputs=[selected_task, selected_task_id],
|
|
outputs=[gen_info, result_actions],
|
|
)
|
|
try:
|
|
for paste_tabname, paste_button in send_to_buttons.items():
|
|
register_paste_params_button(
|
|
ParamBinding(
|
|
paste_button=paste_button,
|
|
tabname=paste_tabname,
|
|
source_text_component=gen_info,
|
|
source_image_component=galerry,
|
|
)
|
|
)
|
|
except:
|
|
pass
|
|
|
|
return [(scheduler_tab, "Agent Scheduler", "agent_scheduler")]
|
|
|
|
|
|
def on_ui_settings():
|
|
section = ("agent_scheduler", "Agent Scheduler")
|
|
shared.opts.add_option(
|
|
"queue_paused",
|
|
shared.OptionInfo(
|
|
False,
|
|
"Disable queue auto-processing",
|
|
gr.Checkbox,
|
|
{"interactive": True},
|
|
section=section,
|
|
),
|
|
)
|
|
shared.opts.add_option(
|
|
"queue_button_placement",
|
|
shared.OptionInfo(
|
|
placement_under_generate,
|
|
"Queue button placement",
|
|
gr.Radio,
|
|
lambda: {
|
|
"choices": [
|
|
placement_under_generate,
|
|
placement_between_prompt_and_generate,
|
|
]
|
|
},
|
|
section=section,
|
|
),
|
|
)
|
|
shared.opts.add_option(
|
|
"queue_button_hide_checkpoint",
|
|
shared.OptionInfo(
|
|
True,
|
|
"Hide the checkpoint dropdown",
|
|
gr.Checkbox,
|
|
{},
|
|
section=section,
|
|
),
|
|
)
|
|
shared.opts.add_option(
|
|
"queue_history_retention_days",
|
|
shared.OptionInfo(
|
|
"30 days",
|
|
"Auto delete queue history (bookmarked tasks excluded)",
|
|
gr.Radio,
|
|
lambda: {
|
|
"choices": list(task_history_retenion_map.keys()),
|
|
},
|
|
section=section,
|
|
),
|
|
)
|
|
|
|
|
|
def on_app_started(block, _):
|
|
if block is not None:
|
|
global task_runner
|
|
task_runner = get_instance(block)
|
|
task_runner.execute_pending_tasks_threading()
|
|
|
|
|
|
script_callbacks.on_ui_tabs(on_ui_tab)
|
|
script_callbacks.on_ui_settings(on_ui_settings)
|
|
script_callbacks.on_app_started(on_app_started)
|