281 lines
9.9 KiB
Python
281 lines
9.9 KiB
Python
import io
|
|
import json
|
|
import threading
|
|
from uuid import uuid4
|
|
from zipfile import ZipFile
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
from gradio.routes import App
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from modules import shared, progress
|
|
|
|
from .db import TaskStatus, task_manager
|
|
from .models import (
|
|
Txt2ImgApiTaskArgs,
|
|
Img2ImgApiTaskArgs,
|
|
QueueTaskResponse,
|
|
QueueStatusResponse,
|
|
HistoryResponse,
|
|
TaskModel,
|
|
)
|
|
from .task_runner import TaskRunner
|
|
from .helpers import log
|
|
from .task_helpers import encode_image_to_base64
|
|
|
|
|
|
def regsiter_apis(app: App, task_runner: TaskRunner):
|
|
log.info("[AgentScheduler] Registering APIs")
|
|
|
|
@app.post("/agent-scheduler/v1/queue/txt2img", response_model=QueueTaskResponse)
|
|
def queue_txt2img(body: Txt2ImgApiTaskArgs):
|
|
task_id = str(uuid4())
|
|
args = body.dict()
|
|
checkpoint = args.pop("model_hash", None)
|
|
task_runner.register_api_task(
|
|
task_id,
|
|
api_task_id=None,
|
|
is_img2img=False,
|
|
args=args,
|
|
checkpoint=checkpoint,
|
|
)
|
|
task_runner.execute_pending_tasks_threading()
|
|
|
|
return QueueTaskResponse(task_id=task_id)
|
|
|
|
@app.post("/agent-scheduler/v1/queue/img2img", response_model=QueueTaskResponse)
|
|
def queue_img2img(body: Img2ImgApiTaskArgs):
|
|
task_id = str(uuid4())
|
|
args = body.dict()
|
|
checkpoint = args.pop("model_hash", None)
|
|
task_runner.register_api_task(
|
|
task_id,
|
|
api_task_id=None,
|
|
is_img2img=True,
|
|
args=args,
|
|
checkpoint=checkpoint,
|
|
)
|
|
task_runner.execute_pending_tasks_threading()
|
|
|
|
return QueueTaskResponse(task_id=task_id)
|
|
|
|
def format_task_args(task):
|
|
task_args = TaskRunner.instance.parse_task_args(task, deserialization=False)
|
|
named_args = task_args.named_args
|
|
named_args["checkpoint"] = task_args.checkpoint
|
|
# remove unused args to reduce payload size
|
|
named_args.pop("alwayson_scripts", None)
|
|
named_args.pop("script_args", None)
|
|
named_args.pop("init_images", None)
|
|
return named_args
|
|
|
|
@app.get("/agent-scheduler/v1/queue", response_model=QueueStatusResponse)
|
|
def queue_status_api(limit: int = 20, offset: int = 0):
|
|
current_task_id = progress.current_task
|
|
total_pending_tasks = task_manager.count_tasks(status="pending")
|
|
pending_tasks = task_manager.get_tasks(
|
|
status=TaskStatus.PENDING, limit=limit, offset=offset
|
|
)
|
|
parsed_tasks = []
|
|
for task in pending_tasks:
|
|
params = format_task_args(task)
|
|
task_data = task.dict()
|
|
task_data["params"] = params
|
|
if task.id == current_task_id:
|
|
task_data["status"] = "running"
|
|
|
|
parsed_tasks.append(TaskModel(**task_data))
|
|
|
|
return QueueStatusResponse(
|
|
current_task_id=current_task_id,
|
|
pending_tasks=parsed_tasks,
|
|
total_pending_tasks=total_pending_tasks,
|
|
paused=TaskRunner.instance.paused,
|
|
)
|
|
|
|
@app.get("/agent-scheduler/v1/history", response_model=HistoryResponse)
|
|
def history_api(status: str = None, limit: int = 20, offset: int = 0):
|
|
bookmarked = True if status == "bookmarked" else None
|
|
if not status or status == "all" or bookmarked:
|
|
status = [
|
|
TaskStatus.DONE,
|
|
TaskStatus.FAILED,
|
|
TaskStatus.INTERRUPTED,
|
|
]
|
|
|
|
total = task_manager.count_tasks(status=status)
|
|
tasks = task_manager.get_tasks(
|
|
status=status,
|
|
bookmarked=bookmarked,
|
|
limit=limit,
|
|
offset=offset,
|
|
order="desc",
|
|
)
|
|
parsed_tasks = []
|
|
for task in tasks:
|
|
params = format_task_args(task)
|
|
task_data = task.dict()
|
|
task_data["params"] = params
|
|
parsed_tasks.append(TaskModel(**task_data))
|
|
|
|
return HistoryResponse(
|
|
total=total,
|
|
tasks=parsed_tasks,
|
|
)
|
|
|
|
@app.post("/agent-scheduler/v1/run/{id}")
|
|
def run_task(id: str):
|
|
if progress.current_task is not None:
|
|
if progress.current_task == id:
|
|
return {"success": False, "message": f"Task {id} is already running"}
|
|
else:
|
|
# move task up in queue
|
|
task_manager.prioritize_task(id, 0)
|
|
return {
|
|
"success": True,
|
|
"message": f"Task {id} is scheduled to run next",
|
|
}
|
|
else:
|
|
# run task
|
|
task = task_manager.get_task(id)
|
|
current_thread = threading.Thread(
|
|
target=TaskRunner.instance.execute_task,
|
|
args=(
|
|
task,
|
|
lambda: None,
|
|
),
|
|
)
|
|
current_thread.daemon = True
|
|
current_thread.start()
|
|
|
|
return {"success": True, "message": f"Task {id} is executing"}
|
|
|
|
@app.post("/agent-scheduler/v1/requeue/{id}")
|
|
def requeue_task(id: str):
|
|
task = task_manager.get_task(id)
|
|
if task is None:
|
|
return {"success": False, "message": f"Task {id} not found"}
|
|
|
|
task.id = str(uuid4())
|
|
task.result = None
|
|
task.status = TaskStatus.PENDING
|
|
task.bookmarked = False
|
|
task.name = f"Copy of {task.name}" if task.name else None
|
|
task_manager.add_task(task)
|
|
task_runner.execute_pending_tasks_threading()
|
|
|
|
return {"success": True, "message": f"Task {id} is requeued"}
|
|
|
|
@app.post("/agent-scheduler/v1/delete/{id}")
|
|
def delete_task(id: str):
|
|
if progress.current_task == id:
|
|
shared.state.interrupt()
|
|
task_runner.interrupted = id
|
|
return {"success": True, "message": f"Task {id} is interrupted"}
|
|
|
|
task_manager.delete_task(id)
|
|
return {"success": True, "message": f"Task {id} is deleted"}
|
|
|
|
@app.post("/agent-scheduler/v1/move/{id}/{over_id}")
|
|
def move_task(id: str, over_id: str):
|
|
task = task_manager.get_task(id)
|
|
if task is None:
|
|
return {"success": False, "message": f"Task {id} not found"}
|
|
|
|
if over_id == "top":
|
|
task_manager.prioritize_task(id, 0)
|
|
return {"success": True, "message": f"Task {id} is moved to top"}
|
|
elif over_id == "bottom":
|
|
task_manager.prioritize_task(id, -1)
|
|
return {"success": True, "message": f"Task {id} is moved to bottom"}
|
|
else:
|
|
over_task = task_manager.get_task(over_id)
|
|
if over_task is None:
|
|
return {"success": False, "message": f"Task {over_id} not found"}
|
|
|
|
task_manager.prioritize_task(id, over_task.priority)
|
|
return {"success": True, "message": f"Task {id} is moved"}
|
|
|
|
@app.post("/agent-scheduler/v1/bookmark/{id}")
|
|
def pin_task(id: str):
|
|
task = task_manager.get_task(id)
|
|
if task is None:
|
|
return {"success": False, "message": f"Task {id} not found"}
|
|
|
|
task.bookmarked = True
|
|
task_manager.update_task(id, bookmarked=True)
|
|
return {"success": True, "message": f"Task {id} is bookmarked"}
|
|
|
|
@app.post("/agent-scheduler/v1/unbookmark/{id}")
|
|
def unpin_task(id: str):
|
|
task = task_manager.get_task(id)
|
|
if task is None:
|
|
return {"success": False, "message": f"Task {id} not found"}
|
|
|
|
task_manager.update_task(id, bookmarked=False)
|
|
return {"success": True, "message": f"Task {id} is unbookmarked"}
|
|
|
|
@app.post("/agent-scheduler/v1/rename/{id}")
|
|
def rename_task(id: str, name: str):
|
|
task = task_manager.get_task(id)
|
|
if task is None:
|
|
return {"success": False, "message": f"Task {id} not found"}
|
|
|
|
task_manager.update_task(id, name=name)
|
|
return {"success": True, "message": f"Task {id} is renamed"}
|
|
|
|
@app.get("/agent-scheduler/v1/results/{id}")
|
|
def get_task_results(id: str, zip: Optional[bool] = False):
|
|
task = task_manager.get_task(id)
|
|
if task is None:
|
|
return {"success": False, "message": f"Task not found"}
|
|
|
|
if task.status != TaskStatus.DONE:
|
|
return {"success": False, "message": f"Task is {task.status.value}"}
|
|
|
|
if task.result is None:
|
|
return {"success": False, "message": f"Task result is not available"}
|
|
|
|
result: dict = json.loads(task.result)
|
|
infotexts = result["infotexts"]
|
|
|
|
if zip:
|
|
zip_buffer = io.BytesIO()
|
|
|
|
# Create a new zip file in the in-memory buffer
|
|
with ZipFile(zip_buffer, "w") as zip_file:
|
|
# Loop through the files in the directory and add them to the zip file
|
|
for image in result["images"]:
|
|
if Path(image).is_file():
|
|
zip_file.write(Path(image), Path(image).name)
|
|
|
|
# Reset the buffer position to the beginning to avoid truncation issues
|
|
zip_buffer.seek(0)
|
|
|
|
# Return the in-memory buffer as a streaming response with the appropriate headers
|
|
return StreamingResponse(
|
|
zip_buffer,
|
|
media_type="application/zip",
|
|
headers={
|
|
"Content-Disposition": f"attachment; filename=results-{id}.zip"
|
|
},
|
|
)
|
|
else:
|
|
data = [
|
|
{"image": encode_image_to_base64(image), "infotext": infotexts[i]}
|
|
for i, image in enumerate(result["images"])
|
|
]
|
|
|
|
return {"success": True, "data": data}
|
|
|
|
@app.post("/agent-scheduler/v1/pause")
|
|
def pause_queue():
|
|
shared.opts.queue_paused = True
|
|
return {"success": True, "message": f"Queue is paused"}
|
|
|
|
@app.post("/agent-scheduler/v1/resume")
|
|
def resume_queue():
|
|
shared.opts.queue_paused = False
|
|
TaskRunner.instance.execute_pending_tasks_threading()
|
|
return {"success": True, "message": f"Queue is resumed"}
|