""" MCP Tool: run Unified image generation task submission and execution interface. """ import time import uuid from core.task_scheduler import QueueFullError, submit_background from core.runtime_config import CONFIG from .common import ( _load_yaml, _MODEL_LIST_PATH, _CHAIN_FEATURES_PATH, _TASK_DEFINITIONS, _TASKS_DB, _TASKS_LOCK, _get_task_snapshot, _execute_imagegen_pipeline, _get_model_feature_context, _canonicalize_chain_type, _resolve_feature_name, ) from .error_schema import make_error, make_validation_error, make_not_found_error from .image_requests import normalize_image_request, validate_image_plan from core.reference_inputs import ImageInputError def handle_run(params: dict) -> dict: """Unified image generation task execution interface.""" if not isinstance(params, dict): return make_validation_error("Request params must be an object.") params = params.copy() if isinstance(params.get("chain"), list): params["chain"] = [ item.copy() if isinstance(item, dict) else item for item in params["chain"] ] 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"] if ( not isinstance(task_type, str) or not isinstance(params["model"], str) or not isinstance(params["prompt"], str) or not task_type.strip() or not params["model"].strip() or not params["prompt"].strip() ): return make_validation_error( "task_type, model, and prompt must be strings.", invalid_fields={ key: "Expected a non-empty string" for key in ("task_type", "model", "prompt") if not isinstance(params.get(key), str) or not params.get(key, "").strip() }, ) task_type = task_type.strip() params["task_type"] = task_type params["model"] = params["model"].strip() 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"]) try: normalize_image_request(params) except ImageInputError as exc: return make_validation_error(str(exc), invalid_fields={"images": str(exc)}) task_requirements = { "txt2img": ("width", "height"), "reference": ("width", "height", "images"), "img2img": ("image",), "inpaint": ("image",), "outpaint": ("image", "pad_left", "pad_right", "pad_top", "pad_bottom"), "hires_fix": ("image", "upscale_by"), } missing_task_fields = [ field for field in task_requirements.get(task_type, ()) if not (field == "image" and params.get("images")) if field not in params or params[field] is None or params[field] == "" ] if missing_task_fields: return make_validation_error( f"Missing required parameter(s) for {task_type}: {', '.join(missing_task_fields)}", missing_fields=missing_task_fields, ) invalid_fields = {} chain_features = _load_yaml(_CHAIN_FEATURES_PATH) def _number(name, *, integer=False, minimum=None, maximum=None): if name not in params: return raw_value = params[name] if raw_value is None or (isinstance(raw_value, str) and not raw_value.strip()): if name in task_requirements.get(task_type, ()): invalid_fields[name] = "Expected a non-empty number" else: # 可选数值的空值等同于未提供，统一沿用执行层的任务 / 模型默认值。 params.pop(name) return if isinstance(raw_value, bool): invalid_fields[name] = "Expected a number" if not integer else "Expected an integer" return try: numeric_value = float(raw_value) except (TypeError, ValueError): invalid_fields[name] = "Expected a number" if not integer else "Expected an integer" return if numeric_value != numeric_value or numeric_value in (float("inf"), float("-inf")): invalid_fields[name] = "Must be finite" return if integer and not numeric_value.is_integer(): invalid_fields[name] = "Expected an integer" return value = int(numeric_value) if integer else numeric_value if minimum is not None and value < minimum: invalid_fields[name] = f"Must be >= {minimum}" elif maximum is not None and value > maximum: invalid_fields[name] = f"Must be <= {maximum}" else: params[name] = value for field in ("width", "height"): if task_type == "txt2img" or field in params: _number(field, integer=True, minimum=1, maximum=8192) for field in ("pad_left", "pad_right", "pad_top", "pad_bottom"): _number(field, integer=True, minimum=0, maximum=4096) _number("upscale_by", minimum=0.01, maximum=4.0) _number("denoise", minimum=0.0, maximum=1.0) _number("steps", integer=True, minimum=1, maximum=100) _number("cfg", minimum=0.0, maximum=100.0) _number("batch_size", integer=True, minimum=1, maximum=CONFIG.max_batch_size) _number("seed", integer=True, minimum=-1) for field in ("sampler", "scheduler", "negative_prompt"): if field in params and params[field] is not None and not isinstance(params[field], str): invalid_fields[field] = "Expected a string" if "async_execution" in params: async_execution = params["async_execution"] if isinstance(async_execution, str): normalized = async_execution.strip().casefold() if normalized in {"true", "1"}: params["async_execution"] = True elif normalized in {"false", "0"}: params["async_execution"] = False else: invalid_fields["async_execution"] = "Expected true/false or 1/0" elif not isinstance(async_execution, bool): invalid_fields["async_execution"] = "Expected a boolean" if "chain" in params: chain = params["chain"] if not isinstance(chain, list): invalid_fields["chain"] = "Expected a list of objects" else: context = _get_model_feature_context(params["model"]) for index, item in enumerate(chain): field_name = f"chain[{index}]" if not isinstance(item, dict): invalid_fields[field_name] = "Expected an object" continue injector_type = item.get("injector_type") if not isinstance(injector_type, str) or not injector_type.strip(): invalid_fields[field_name] = "injector_type must be a non-empty string" continue injector_type = injector_type.strip() injector_type = _canonicalize_chain_type(injector_type) item["injector_type"] = injector_type feature_name = _resolve_feature_name(injector_type) if feature_name is None: return make_validation_error( f"Unknown chain/injector type '{injector_type}'.", invalid_fields={field_name: "Unknown injector_type"}, ) feature_data = chain_features.get(feature_name, {}) aliases = feature_data.get("chains", feature_name) aliases = aliases if isinstance(aliases, list) else [aliases] enabled = context and context["enabled_chains"] supports = context and context["supports_chain"] if not enabled or not any(alias in enabled for alias in [feature_name, *aliases]): return make_error( "FEATURE_NOT_SUPPORTED", f"Model '{params['model']}' does not support chain '{injector_type}'.", {"model": params["model"], "chain": injector_type}, ) if supports and not any(supports(alias) for alias in [injector_type, feature_name, *aliases]): return make_error( "FEATURE_NOT_SUPPORTED", f"Model '{params['model']}' does not support chain '{injector_type}'.", {"model": params["model"], "chain": injector_type}, ) if invalid_fields: return make_validation_error("Request contains invalid parameters.", invalid_fields=invalid_fields) try: validate_image_plan(params) except ImageInputError as exc: return make_validation_error(str(exc), invalid_fields={"image_plan": str(exc)}) batch_size = params.get("batch_size", 1) if not 1 <= batch_size <= CONFIG.max_batch_size: return make_validation_error( f"batch_size must be between 1 and {CONFIG.max_batch_size}.", invalid_fields={"batch_size": f"Expected 1..{CONFIG.max_batch_size}"}, ) with _TASKS_LOCK: if len(_TASKS_DB) >= CONFIG.mcp_task_retention: finished = sorted( ( (old_task_id, task) for old_task_id, task in _TASKS_DB.items() if task.get("status") in {"completed", "failed"} ), key=lambda pair: pair[1].get("completed_at", pair[1].get("failed_at", 0)), ) remove_count = max(1, len(_TASKS_DB) - CONFIG.mcp_task_retention + 1) for old_task_id, _ in finished[:remove_count]: _TASKS_DB.pop(old_task_id, None) if len(_TASKS_DB) >= CONFIG.mcp_task_retention: return make_error( "QUEUE_FULL", "任务记录已满且当前任务均未结束，请稍后再试。", ) 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: try: submit_background(_execute_imagegen_pipeline, task_id, params) except QueueFullError as exc: with _TASKS_LOCK: _TASKS_DB.pop(task_id, None) return make_error("QUEUE_FULL", str(exc)) 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 _get_task_snapshot(task_id)