Spaces:
Running on Zero
Running on Zero
| """ | |
| MCP & Gradio Integration Module | |
| Provides: | |
| 1. register_high_level_mcp_apis: Expose only 7 high-level abstract API/MCP endpoints (using gr.api without polluting the visual UI structure) | |
| 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 | |
| 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", | |
| } | |
| 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/optional parameters and paste-and-run example_json_params. Recommended flow: get_task_list -> get_model_architecture_list -> get_model_list -> 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 = "") -> list | dict: | |
| """Get supported advanced features (LoRA, ControlNet, IPAdapter, etc.). Each feature includes ready-to-run example_chain_item and example_json_params. If feature_name is empty, returns summary of ALL features. Pass specific feature_name to retrieve complete parameters_schema.""" | |
| return sanitize_keys(handle_get_feature_list(feature_name.strip() if isinstance(feature_name, str) else feature_name)) | |
| def get_model_features(model: str = "") -> dict: | |
| """Query metadata for the specified model, including supported task types, extended features, default parameters, and paste-and-run example_json_params.""" | |
| return sanitize_keys(handle_get_model_features(model.strip())) | |
| def run(json_params: str = "{}") -> dict: | |
| """[Recommended Discovery Flow Step 4] Unified image generation task execution interface. Accepts a JSON string or dict of parameters. | |
| OPTIONAL CONTROL PARAMETERS: | |
| - seed (int): Random seed (-1 for random, >=0 for deterministic reproducibility). Default: -1. | |
| - batch_size (int): Number of images generated per batch (1 to 16). Default: 1. | |
| - zero_gpu_duration (int): GPU quota allocation in seconds on HuggingFace ZeroGPU spaces (default: 60, max: 120). | |
| - negative_prompt (str): Text prompt specifying elements to avoid. | |
| IMPORTANT FORMAT RULES FOR 'chain': | |
| - For LoRA/Civitai: 'lora_value' refers to the Civitai Version ID (modelVersionId), not the main Model ID. | |
| Paste-and-Run Example (Basic): | |
| {"task_type": "txt2img", "model": "stabilityai/SDXL-Base-1.0", "prompt": "A majestic lion jumping from a big stone at night", "width": 1024, "height": 1024} | |
| Paste-and-Run Example (With chain): | |
| {"task_type": "txt2img", "model": "stabilityai/SDXL-Base-1.0", "prompt": "A majestic lion jumping from a big stone at night", "width": 1024, "height": 1024, "chain": [{"injector_type": "lora", "source": "Civitai", "lora_value": "12345", "scale": 1.0}]} | |
| """ | |
| 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())) | |
| funcs = [ | |
| get_task_list, | |
| get_model_architecture_list, | |
| get_model_list, | |
| get_feature_list, | |
| get_model_features, | |
| run, | |
| get_task_status, | |
| ] | |
| 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] Successfully registered 7 High-Level Abstract MCP APIs via gr.api().") | |