Spaces:
Running on Zero
Running on Zero
File size: 8,164 Bytes
97ac1ad b2cc4bf 97ac1ad b2cc4bf 97ac1ad b2cc4bf 97ac1ad b2cc4bf 97ac1ad b2cc4bf 97ac1ad b2cc4bf 97ac1ad | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """
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."""
if isinstance(model_architecture, dict):
model_architecture = model_architecture.get("model_architecture") or model_architecture.get("arch") or ""
arch = model_architecture.strip() if isinstance(model_architecture, str) and model_architecture.strip() 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."""
if isinstance(feature_name, dict):
feature_name = feature_name.get("feature_name") or feature_name.get("name") or ""
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."""
if isinstance(model, dict):
model = model.get("model") or model.get("name") or ""
model_str = model.strip() if isinstance(model, str) else ""
return sanitize_keys(handle_get_model_features(model_str))
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."""
if isinstance(task_id, dict):
task_id = task_id.get("task_id") or task_id.get("id") or ""
task_id_str = task_id.strip() if isinstance(task_id, str) else ""
return sanitize_keys(handle_get_task_status(task_id_str))
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().")
|