File size: 2,396 Bytes
9dbb7e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """
MCP Tool: run_imagegen
Unified image generation task submission and execution interface.
"""
import time
import uuid
import threading
from .common import (
_load_yaml,
_MODEL_LIST_PATH,
_TASK_DEFINITIONS,
_TASKS_DB,
_execute_imagegen_pipeline,
)
from .error_schema import make_validation_error, make_not_found_error
def handle_run_imagegen(params: dict) -> dict:
"""Unified image generation task execution interface."""
if not isinstance(params, dict):
return make_validation_error("Request params must be an object.")
missing = []
for req_field in ["task_type", "model", "prompt"]:
if req_field not in params or not params[req_field]:
missing.append(req_field)
if missing:
return make_validation_error(
f"Missing required parameter(s): {', '.join(missing)}",
missing_fields=missing,
)
task_type = params["task_type"]
valid_tasks = [t["task_type"] for t in _TASK_DEFINITIONS]
if task_type not in valid_tasks:
return make_validation_error(
f"Invalid task_type '{task_type}'. Must be one of {valid_tasks}.",
invalid_fields={"task_type": f"Must be in {valid_tasks}"},
)
model_list = _load_yaml(_MODEL_LIST_PATH)
checkpoints = model_list.get("Checkpoint", {})
all_models = set()
for arch_name, arch_data in checkpoints.items():
if isinstance(arch_data, dict):
for m in arch_data.get("models", []):
all_models.add(m.get("display_name"))
if params["model"] not in all_models:
return make_not_found_error("model", params["model"])
task_id = f"img_task_{uuid.uuid4().hex[:10]}"
created_at = int(time.time())
_TASKS_DB[task_id] = {
"task_id": task_id,
"status": "queued",
"progress": 0,
"created_at": created_at,
}
async_exec = params.get("async_execution", False)
if async_exec:
t = threading.Thread(target=_execute_imagegen_pipeline, args=(task_id, params), daemon=True)
t.start()
return {
"status": "queued",
"task_id": task_id,
"poll_interval_ms": 2000,
"message": "Task queued successfully. Poll get_task_status for results.",
}
else:
_execute_imagegen_pipeline(task_id, params)
return _TASKS_DB[task_id]
|