Spaces:
Running on Zero
Running on Zero
| """ | |
| MCP & Gradio Integration Module | |
| Provides: | |
| 1. register_high_level_mcp_apis: Expose the canonical API plus Fluxus-compatible aliases | |
| 2. cleanup_dependencies_api_names: Force cleanup of show_api attribute for non-high-level APIs in dependencies | |
| 3. patch_gradio_api_suppression: No-op implementation retained for backward compatibility | |
| """ | |
| import json | |
| import gradio as gr | |
| from .get_task_list import handle_get_task_list | |
| from .get_model_architecture_list import handle_get_model_architecture_list | |
| from .get_model_list import handle_get_model_list | |
| from .get_feature_list import handle_get_feature_list | |
| from .get_model_features import handle_get_model_features | |
| from .run import handle_run | |
| from .get_task_status import handle_get_task_status | |
| from .get_chain_schema import handle_get_chain_schema | |
| HIGH_LEVEL_MCP_API_NAMES = { | |
| "get_task_list", | |
| "get_model_architecture_list", | |
| "get_model_list", | |
| "get_feature_list", | |
| "get_model_features", | |
| "run", | |
| "get_task_status", | |
| "run_imagegen", | |
| "get_chain_schema", | |
| } | |
| def sanitize_keys(obj): | |
| """Recursively ensure all dictionary keys are converted to str type to avoid Gradio 5 orjson TypeError: Dict key must be str.""" | |
| if isinstance(obj, dict): | |
| return {str(k): sanitize_keys(v) for k, v in obj.items()} | |
| elif isinstance(obj, list): | |
| return [sanitize_keys(x) for x in obj] | |
| elif isinstance(obj, tuple): | |
| return tuple(sanitize_keys(x) for x in obj) | |
| return obj | |
| def patch_gradio_api_suppression(): | |
| """Retained for backward compatibility (no-op).""" | |
| pass | |
| def cleanup_dependencies_api_names(demo): | |
| """ | |
| Clean up residual auto-generated API names in demo.fns and demo.dependencies. | |
| Force only the 7 high-level abstract MCP APIs to be exposed as public endpoints. | |
| """ | |
| for fn in demo.fns.values(): | |
| api_name = getattr(fn, "api_name", None) | |
| if api_name not in HIGH_LEVEL_MCP_API_NAMES: | |
| fn.show_api = False | |
| deps = getattr(demo, "dependencies", None) | |
| if deps is None and hasattr(demo, "config") and isinstance(demo.config, dict): | |
| deps = demo.config.get("dependencies", []) | |
| if deps: | |
| for dep in deps: | |
| if isinstance(dep, dict): | |
| api_name = dep.get("api_name") | |
| if api_name not in HIGH_LEVEL_MCP_API_NAMES: | |
| dep["show_api"] = False | |
| print("[MCP Protection] Cleaned up demo dependencies. Suppressed atomic API endpoints.") | |
| def register_high_level_mcp_apis(demo): | |
| """ | |
| Explicitly register 7 high-level abstract MCP API endpoints on the Gradio demo using gr.api. | |
| Using gr.api() never adds any visual UI components (such as Row, Textbox, Button, etc.), avoiding duplicate interface rendering. | |
| """ | |
| def get_task_list() -> list: | |
| """[Recommended Discovery Flow Step 1] Get a list of all supported image generation task types (txt2img, img2img, inpaint, outpaint, hires_fix) along with their required and optional parameter lists. Recommended flow: get_task_list -> get_model_architecture_list -> get_model_list -> [Path 1: Call run directly (pass only required params) | Path 2: Call get_model_features to get official default hyperparams -> run].""" | |
| return sanitize_keys(handle_get_task_list()) | |
| def get_model_architecture_list() -> list: | |
| """[Recommended Discovery Flow Step 2] Get a list of all supported model architectures (e.g., SD1.5, SDXL, FLUX, etc.) along with their default resolutions. It is recommended to call this tool before get_model_list to obtain valid model_architecture parameters for precise model filtering.""" | |
| return sanitize_keys(handle_get_model_architecture_list()) | |
| def get_model_list(model_architecture: str = "") -> list | dict: | |
| """[Recommended Discovery Flow Step 3] Query the list of available image generation models. After obtaining models, choose one of two paths: 1. [Path 1 (Recommended - Minimal Mode)] Call run directly with only required parameters. Do NOT guess steps/cfg/sampler/scheduler from experience; the server will automatically apply the model's optimal default hyperparameters. 2. [Path 2 (Explicit Alignment Mode)] First call get_model_features to query the model's officially recommended hyperparameters, then pass them to run.""" | |
| arch = model_architecture.strip() if model_architecture else None | |
| return sanitize_keys(handle_get_model_list(arch)) | |
| def get_feature_list(feature_name: str = "", compact: bool = False) -> list | dict: | |
| """Get supported advanced features. Empty input returns Fluxus-compatible full schemas; set compact=true for a token-saving summary. Pass feature_name for one or more complete schemas.""" | |
| return sanitize_keys( | |
| handle_get_feature_list( | |
| feature_name.strip() if isinstance(feature_name, str) else feature_name, | |
| include_schema_on_empty=not compact, | |
| ) | |
| ) | |
| def get_model_features(model: str = "") -> dict: | |
| """Query the concrete model's task types, features, recommended sampling defaults and image_workflows. For editing/multi-image use image_workflows: supported run modes, reference roles, image count limits, parameter schema, aliases and usage rules. Do not infer editing support from architecture alone.""" | |
| return sanitize_keys(handle_get_model_features(model.strip())) | |
| def run(json_params: str = "{}") -> dict: | |
| """Submit a JSON object encoded as json_params. First discover get_task_list and get_model_features(model). Text generation: task_type=txt2img, model, prompt, width, height. Instruction editing/fusion: task_type=reference, model, prompt, width, height, images=[URL_or_DataURI,...]; reference_role defaults to auto. images order starts at 1: 图1 / img1 / image1 in prompt refer to the first image. Alternatively use top-level img1,img2 (or 图1,图2) fields; do not mix input forms. Ordinary repaint: task_type=img2img and image; batch repaint: images plus run_mode=multi_independent. See image_workflows for model-specific limits and other modes. Omit steps/cfg/sampler/scheduler to use recommended defaults. Prefer async_execution=true; poll the returned task_id via get_task_status (do not resubmit). Legacy chain and run_imagegen calls remain available.""" | |
| try: | |
| if isinstance(json_params, dict): | |
| params = json_params | |
| else: | |
| params = json.loads(json_params or "{}") | |
| except Exception as e: | |
| return {"error": {"code": "INVALID_JSON", "message": f"Failed to parse JSON params: {e}"}} | |
| return sanitize_keys(handle_run(params)) | |
| def get_task_status(task_id: str = "") -> dict: | |
| """Query the progress, status, and final generated results of an async image generation task.""" | |
| return sanitize_keys(handle_get_task_status(task_id.strip())) | |
| def run_imagegen(json_params: str = "{}") -> dict: | |
| """Backward-compatible alias for run, retained for Fluxus clients.""" | |
| return run(json_params) | |
| def get_chain_schema(chain_type: str = "") -> dict: | |
| """Backward-compatible Fluxus endpoint for one chain/injector schema.""" | |
| return sanitize_keys(handle_get_chain_schema(chain_type.strip())) | |
| funcs = [ | |
| get_task_list, | |
| get_model_architecture_list, | |
| get_model_list, | |
| get_feature_list, | |
| get_model_features, | |
| run, | |
| get_task_status, | |
| run_imagegen, | |
| get_chain_schema, | |
| ] | |
| for func in funcs: | |
| gr.api(func) | |
| for fn in demo.fns.values(): | |
| if getattr(fn, "api_name", None) in HIGH_LEVEL_MCP_API_NAMES: | |
| fn.show_api = True | |
| print("[MCP Integration] Registered 7 canonical MCP APIs and 2 Fluxus aliases.") | |