mcp_tools/common.py

"""

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()

    # Reject HTTP / HTTPS URL
    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,...')."
        )

    # 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))

    # Base64 string without header
    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

    # Local file path
    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)."""
    # 1. Explicit environment variable override
    public_url = os.getenv("PUBLIC_URL") or os.getenv("BASE_URL")
    if public_url:
        return public_url.rstrip("/")

    # 2. Hugging Face Space environment variable
    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("/")

    # 3. Local Gradio config fallback
    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

        # Execute Pipeline
        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"):  # PIL Image
                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),
        }