Spaces:
Running on Zero
Running on Zero
| """ | |
| 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) | |