| """
|
| MCP Common Utilities & Data Structures
|
| Contains YAML loading utilities, config file paths, task definitions, and async task database.
|
| """
|
|
|
| import os
|
| import time
|
| import urllib.parse
|
| import urllib.request
|
| import base64
|
| import io
|
| import yaml
|
| from typing import Dict, Any
|
| from PIL import Image
|
|
|
| _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
| _YAML_DIR = os.path.join(_PROJECT_ROOT, "yaml")
|
|
|
| _MODEL_ARCHITECTURES_PATH = os.path.join(_YAML_DIR, "model_architectures.yaml")
|
| _MODEL_LIST_PATH = os.path.join(_YAML_DIR, "model_list.yaml")
|
| _MODEL_DEFAULTS_PATH = os.path.join(_YAML_DIR, "model_defaults.yaml")
|
| _IMAGE_GEN_FEATURES_PATH = os.path.join(_YAML_DIR, "image_gen_features.yaml")
|
| _CHAIN_FEATURES_PATH = os.path.join(_YAML_DIR, "chain_features.yaml")
|
| _CONSTANTS_PATH = os.path.join(_YAML_DIR, "constants.yaml")
|
|
|
|
|
| def _parse_image_param(image_param: Any) -> Any:
|
| """Parse a Base64 Data URI, local file path, or PIL.Image into a PIL Image object. HTTP URLs are not supported."""
|
| if isinstance(image_param, Image.Image):
|
| return image_param
|
|
|
| if not isinstance(image_param, str) or not image_param.strip():
|
| return None
|
|
|
| image_param = image_param.strip()
|
|
|
|
|
| if image_param.startswith("http://") or image_param.startswith("https://"):
|
| raise ValueError(
|
| "Image URLs are not supported. Please supply the image directly as a Base64 Data URI (e.g., 'data:image/png;base64,...')."
|
| )
|
|
|
|
|
| if image_param.startswith("data:image/"):
|
| _, encoded = image_param.split(",", 1) if "," in image_param else ("", image_param)
|
| data = base64.b64decode(encoded)
|
| return Image.open(io.BytesIO(data))
|
|
|
|
|
| if len(image_param) > 100 and not os.path.exists(image_param):
|
| try:
|
| data = base64.b64decode(image_param)
|
| return Image.open(io.BytesIO(data))
|
| except Exception:
|
| pass
|
|
|
|
|
| if os.path.exists(image_param):
|
| return Image.open(image_param)
|
|
|
| raise ValueError(
|
| "Invalid image parameter format. Expected a Base64 Data URI (e.g., 'data:image/png;base64,...') or local file path."
|
| )
|
|
|
|
|
| def _load_yaml(filepath: str) -> dict:
|
| """Safely load a YAML file, returning an empty dict if the file does not exist."""
|
| if not os.path.exists(filepath):
|
| print(f"Warning: YAML file not found: {filepath}")
|
| return {}
|
| with open(filepath, "r", encoding="utf-8") as f:
|
| return yaml.safe_load(f) or {}
|
|
|
|
|
| _COMMON_OPTIONAL_INPUTS = [
|
| "steps", "cfg", "sampler", "scheduler", "seed",
|
| "negative_prompt", "batch_size", "chain", "async_execution",
|
| ]
|
|
|
| _TASK_DEFINITIONS = [
|
| {
|
| "task_type": "txt2img",
|
| "display_name": "Text-to-Image",
|
| "description": "Generate images from text prompts. Canvas width and height must be specified.",
|
| "required_inputs": ["prompt", "width", "height"],
|
| "optional_inputs": _COMMON_OPTIONAL_INPUTS,
|
| },
|
| {
|
| "task_type": "img2img",
|
| "display_name": "Image-to-Image",
|
| "description": "Perform global repaint and style transfer based on a source image. Denoise strength must be specified.",
|
| "required_inputs": ["prompt", "image", "denoise"],
|
| "optional_inputs": _COMMON_OPTIONAL_INPUTS,
|
| },
|
| {
|
| "task_type": "inpaint",
|
| "display_name": "Inpaint",
|
| "description": "Repaint specified masked regions of the input image (with alpha mask/channel).",
|
| "required_inputs": ["prompt", "image"],
|
| "optional_inputs": ["denoise"] + _COMMON_OPTIONAL_INPUTS,
|
| },
|
| {
|
| "task_type": "outpaint",
|
| "display_name": "Outpaint",
|
| "description": "Extend the canvas outward from the source image. Padding pixel values for top, bottom, left, and right must be specified.",
|
| "required_inputs": ["prompt", "image", "pad_left", "pad_right", "pad_top", "pad_bottom"],
|
| "optional_inputs": _COMMON_OPTIONAL_INPUTS,
|
| },
|
| {
|
| "task_type": "hires_fix",
|
| "display_name": "Hi-Res Fix / Upscale",
|
| "description": "Enhance details and upscale an existing low-resolution image.",
|
| "required_inputs": ["prompt", "image", "upscale_by"],
|
| "optional_inputs": _COMMON_OPTIONAL_INPUTS,
|
| },
|
| ]
|
|
|
| _TASKS_DB: Dict[str, Dict[str, Any]] = {}
|
|
|
|
|
| class DummyProgress:
|
| def __call__(self, progress=0.0, desc=None):
|
| pass
|
|
|
|
|
| def _get_public_base_url() -> str:
|
| """Auto-resolve the publicly accessible base URL (including protocol and port)."""
|
|
|
| public_url = os.getenv("PUBLIC_URL") or os.getenv("BASE_URL")
|
| if public_url:
|
| return public_url.rstrip("/")
|
|
|
|
|
| space_host = os.getenv("SPACE_HOST")
|
| if space_host:
|
| if not space_host.startswith("http://") and not space_host.startswith("https://"):
|
| return f"https://{space_host}"
|
| return space_host.rstrip("/")
|
|
|
|
|
| try:
|
| from core.settings import GRADIO_SERVER_NAME, SERVER_PORT
|
| except ImportError:
|
| GRADIO_SERVER_NAME = "127.0.0.1"
|
| SERVER_PORT = 7860
|
|
|
| server_name = os.getenv("GRADIO_SERVER_NAME", GRADIO_SERVER_NAME)
|
| if server_name == "0.0.0.0":
|
| server_name = "127.0.0.1"
|
| port = os.getenv("GRADIO_SERVER_PORT", str(SERVER_PORT))
|
|
|
| return f"http://{server_name}:{port}"
|
|
|
|
|
| def _execute_imagegen_pipeline(task_id: str, params: dict):
|
| """Execute the image generation pipeline in the background and update _TASKS_DB."""
|
| start_time = time.time()
|
| try:
|
| _TASKS_DB[task_id]["status"] = "processing"
|
| _TASKS_DB[task_id]["progress"] = 10
|
| _TASKS_DB[task_id]["updated_at"] = int(start_time)
|
|
|
| from core.generation_logic import sd_image_pipeline
|
|
|
| task_type = params["task_type"]
|
| model = params["model"]
|
| prompt = params["prompt"]
|
|
|
| model_defaults = _load_yaml(_MODEL_DEFAULTS_PATH)
|
| model_list = _load_yaml(_MODEL_LIST_PATH)
|
| checkpoints = model_list.get("Checkpoint", {})
|
| found_arch = None
|
| for arch_name, arch_data in checkpoints.items():
|
| if isinstance(arch_data, dict):
|
| for m in arch_data.get("models", []):
|
| if m.get("display_name") == model:
|
| found_arch = arch_name
|
| break
|
| if found_arch:
|
| break
|
|
|
| arch_defaults_section = model_defaults.get(found_arch, {}) if found_arch else {}
|
| arch_level_defaults = arch_defaults_section.get("_defaults", {})
|
| model_specific_defaults = arch_defaults_section.get(model, {})
|
| global_defaults = model_defaults.get("Default", {})
|
| merged_defaults = {**global_defaults, **arch_level_defaults, **model_specific_defaults}
|
|
|
| steps = params.get("steps") if params.get("steps") is not None else merged_defaults.get("steps", 20)
|
| cfg = params.get("cfg") if params.get("cfg") is not None else merged_defaults.get("cfg", 1.0)
|
| sampler = params.get("sampler") or merged_defaults.get("sampler_name", "euler")
|
| scheduler = params.get("scheduler") or merged_defaults.get("scheduler", "simple")
|
|
|
| ui_inputs = {
|
| "task_type": task_type,
|
| "model_display_name": model,
|
| "base_model_" + task_type: model,
|
| "positive_prompt": prompt,
|
| "negative_prompt": params.get("negative_prompt", merged_defaults.get("negative_prompt", "")),
|
| "width": params.get("width", 1024),
|
| "height": params.get("height", 1024),
|
| "num_inference_steps": steps,
|
| "guidance_scale": cfg,
|
| "sampler": sampler,
|
| "scheduler": scheduler,
|
| "seed": params.get("seed", -1),
|
| "batch_size": params.get("batch_size", 1),
|
| "zero_gpu_duration": params.get("zero_gpu_duration"),
|
| "denoise": params.get("denoise", 1.0),
|
| }
|
|
|
| if "image" in params and params["image"]:
|
| pil_img = _parse_image_param(params["image"])
|
| if pil_img:
|
| if task_type == "img2img":
|
| ui_inputs["img2img_image"] = pil_img
|
| ui_inputs["img2img_denoise"] = params.get("denoise", 0.7)
|
| elif task_type == "inpaint":
|
| ui_inputs["inpaint_image"] = pil_img
|
| ui_inputs["inpaint_denoise"] = params.get("denoise", 1.0)
|
| elif task_type == "outpaint":
|
| ui_inputs["outpaint_image"] = pil_img
|
| ui_inputs["left"] = params.get("pad_left", 0)
|
| ui_inputs["right"] = params.get("pad_right", 0)
|
| ui_inputs["top"] = params.get("pad_top", 0)
|
| ui_inputs["bottom"] = params.get("pad_bottom", 0)
|
| ui_inputs["feathering"] = params.get("feathering", 10)
|
| elif task_type == "hires_fix":
|
| ui_inputs["hires_image"] = pil_img
|
| ui_inputs["hires_upscaler"] = params.get("upscaler", "latent")
|
| ui_inputs["hires_scale_by"] = params.get("upscale_by", 2.0)
|
| ui_inputs["hires_denoise"] = params.get("denoise", 0.55)
|
|
|
| chain = params.get("chain", [])
|
| if chain:
|
| lora_data = []
|
| controlnet_data = []
|
| ipadapter_data = []
|
| style_data = []
|
|
|
| for item in chain:
|
| itype = item.get("injector_type")
|
| if itype == "lora":
|
| lora_data.extend([
|
| item.get("lora_source", "Civitai"),
|
| item.get("lora_value", ""),
|
| item.get("scale", 1.0),
|
| None
|
| ])
|
| elif itype in ("controlnet", "krea2_controlnet", "anima_controlnet_lllite"):
|
| controlnet_data.extend([
|
| item.get("control_net_name", ""),
|
| _parse_image_param(item.get("image")),
|
| item.get("strength", 1.0)
|
| ])
|
| elif itype in ("ipadapter", "flux1_ipadapter", "sd3_ipadapter"):
|
| ipadapter_data.extend([
|
| item.get("preset", "STANDARD (medium strength)"),
|
| _parse_image_param(item.get("image")),
|
| item.get("weight", 1.0)
|
| ])
|
| elif itype == "style":
|
| style_data.extend([
|
| _parse_image_param(item.get("image")),
|
| item.get("strength", 1.0)
|
| ])
|
|
|
| if lora_data: ui_inputs["lora_data"] = lora_data
|
| if controlnet_data: ui_inputs["controlnet_data"] = controlnet_data
|
| if ipadapter_data: ui_inputs["ipadapter_data"] = ipadapter_data
|
| if style_data: ui_inputs["style_data"] = style_data
|
|
|
| _TASKS_DB[task_id]["progress"] = 50
|
|
|
|
|
| output = sd_image_pipeline.run(ui_inputs=ui_inputs, progress=DummyProgress())
|
|
|
| try:
|
| from core.settings import OUTPUT_DIR
|
| except ImportError:
|
| OUTPUT_DIR = os.path.join(_PROJECT_ROOT, "output")
|
|
|
| os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
| import tempfile
|
| import gradio.processing_utils as pu
|
|
|
| gradio_cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
|
| os.makedirs(gradio_cache_dir, exist_ok=True)
|
|
|
| base_url = _get_public_base_url()
|
| images = []
|
| raw_list = output if isinstance(output, list) else ([output] if output else [])
|
| for idx, item in enumerate(raw_list):
|
| target_path = None
|
| if hasattr(item, "save"):
|
| filename = f"mcp_{task_id}_{idx}.png"
|
| filepath = os.path.join(OUTPUT_DIR, filename)
|
| item.save(filepath)
|
| target_path = filepath
|
| elif isinstance(item, str) and os.path.exists(item):
|
| target_path = item
|
|
|
| if target_path:
|
| try:
|
| cached_path = pu.save_file_to_cache(target_path, cache_dir=gradio_cache_dir)
|
| abs_path = os.path.abspath(cached_path).replace("\\", "/")
|
| except Exception as e:
|
| print(f"Warning: Failed to cache image file to Gradio temp dir: {e}")
|
| abs_path = os.path.abspath(target_path).replace("\\", "/")
|
|
|
| url = f"{base_url}/gradio_api/file={urllib.parse.quote(abs_path)}"
|
| images.append(url)
|
| elif item:
|
| images.append(str(item))
|
|
|
| execution_time = round(time.time() - start_time, 2)
|
| _TASKS_DB[task_id]["status"] = "completed"
|
| _TASKS_DB[task_id]["progress"] = 100
|
| _TASKS_DB[task_id]["completed_at"] = int(time.time())
|
| _TASKS_DB[task_id]["result"] = {
|
| "images": images,
|
| "seed": params.get("seed", -1),
|
| "width": params.get("width", 1024),
|
| "height": params.get("height", 1024),
|
| "execution_time_seconds": execution_time,
|
| }
|
|
|
| except Exception as e:
|
| _TASKS_DB[task_id]["status"] = "failed"
|
| _TASKS_DB[task_id]["progress"] = 0
|
| _TASKS_DB[task_id]["failed_at"] = int(time.time())
|
| _TASKS_DB[task_id]["error"] = {
|
| "code": "EXECUTION_ERROR",
|
| "message": str(e),
|
| }
|
|
|
