| """
|
| 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 urllib.error
|
| 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")
|
|
|
|
|
| _MAX_IMAGE_DOWNLOAD_BYTES = 50 * 1024 * 1024
|
| _IMAGE_DOWNLOAD_TIMEOUT = 30
|
| _ALLOWED_IMAGE_CONTENT_TYPES = frozenset([
|
| "image/png", "image/jpeg", "image/jpg", "image/gif",
|
| "image/webp", "image/bmp", "image/tiff",
|
| ])
|
|
|
|
|
| def _get_ipadapter_presets_by_arch() -> Dict[str, list]:
|
| """Load IPAdapter presets from yaml/ipadapter.yaml for SD1.5 and SDXL."""
|
| ipadapter_yaml_path = os.path.join(_YAML_DIR, "ipadapter.yaml")
|
| data = _load_yaml(ipadapter_yaml_path)
|
| res = {}
|
| for arch in ("SD1.5", "SDXL"):
|
| std = data.get("IPAdapter_presets", {}).get(arch, [])
|
| face = data.get("IPAdapter_FaceID_presets", {}).get(arch, [])
|
| res[arch] = list(std) + list(face)
|
| return res
|
|
|
|
|
| def _parse_image_param(image_param: Any) -> Any:
|
| """Parse a Base64 Data URI, HTTP/HTTPS URL, or PIL.Image into a PIL Image object."""
|
| 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://"):
|
| return _download_image_from_url(image_param)
|
|
|
|
|
| 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:
|
| try:
|
| data = base64.b64decode(image_param)
|
| return Image.open(io.BytesIO(data))
|
| except Exception:
|
| pass
|
|
|
| raise ValueError(
|
| "Invalid image parameter format. Expected a Base64 Data URI (e.g., 'data:image/png;base64,...') "
|
| "or an HTTP/HTTPS URL."
|
| )
|
|
|
|
|
| def _download_image_from_url(url: str) -> Image.Image:
|
| """Download an image from an HTTP/HTTPS URL and return it as a PIL Image.
|
|
|
| Security measures:
|
| - Timeout to prevent hanging on slow/malicious servers.
|
| - Response size cap to prevent memory exhaustion.
|
| - Content-Type validation to reject non-image responses.
|
| """
|
| req = urllib.request.Request(url, headers={"User-Agent": "ImageGen-MCP/1.0"})
|
| try:
|
| with urllib.request.urlopen(req, timeout=_IMAGE_DOWNLOAD_TIMEOUT) as resp:
|
|
|
| content_type = resp.headers.get("Content-Type", "").split(";")[0].strip().lower()
|
| if content_type and content_type not in _ALLOWED_IMAGE_CONTENT_TYPES:
|
| raise ValueError(
|
| f"URL returned non-image Content-Type '{content_type}'. "
|
| f"Expected one of: {', '.join(sorted(_ALLOWED_IMAGE_CONTENT_TYPES))}."
|
| )
|
|
|
|
|
| content_length = resp.headers.get("Content-Length")
|
| if content_length and int(content_length) > _MAX_IMAGE_DOWNLOAD_BYTES:
|
| raise ValueError(
|
| f"Image at URL is too large ({int(content_length)} bytes). "
|
| f"Maximum allowed size is {_MAX_IMAGE_DOWNLOAD_BYTES} bytes."
|
| )
|
|
|
|
|
| chunks = []
|
| total = 0
|
| while True:
|
| chunk = resp.read(8192)
|
| if not chunk:
|
| break
|
| total += len(chunk)
|
| if total > _MAX_IMAGE_DOWNLOAD_BYTES:
|
| raise ValueError(
|
| f"Image download exceeded maximum allowed size of "
|
| f"{_MAX_IMAGE_DOWNLOAD_BYTES} bytes."
|
| )
|
| chunks.append(chunk)
|
|
|
| data = b"".join(chunks)
|
|
|
| except urllib.error.URLError as e:
|
| raise ValueError(f"Failed to download image from URL: {e}") from e
|
| except urllib.error.HTTPError as e:
|
| raise ValueError(f"HTTP error {e.code} when downloading image from URL: {e.reason}") from e
|
|
|
| if not data:
|
| raise ValueError("Downloaded image data is empty.")
|
|
|
| return Image.open(io.BytesIO(data))
|
|
|
|
|
| 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", "zero_gpu_duration", "chain", "async_execution",
|
| ]
|
|
|
| _COMMON_OPTIONAL_INPUTS_SCHEMA = {
|
| "seed": {
|
| "type": "integer",
|
| "default": -1,
|
| "description": "Random seed for image generation. Default: -1 (random seed). Specify a non-negative integer for deterministic reproducible generation."
|
| },
|
| "batch_size": {
|
| "type": "integer",
|
| "default": 1,
|
| "minimum": 1,
|
| "maximum": 16,
|
| "description": "Number of images generated in a single inference batch (1 to 16, default: 1)."
|
| },
|
| "zero_gpu_duration": {
|
| "type": "integer",
|
| "default": 60,
|
| "minimum": 1,
|
| "maximum": 120,
|
| "description": "GPU execution time quota allocation in seconds on HuggingFace ZeroGPU spaces (default: 60, maximum: 120)."
|
| },
|
| "steps": {
|
| "type": "integer",
|
| "description": "Number of inference steps. If omitted, model optimal default is automatically applied."
|
| },
|
| "cfg": {
|
| "type": "number",
|
| "description": "Classifier-Free Guidance scale. If omitted, model optimal default is automatically applied."
|
| },
|
| "sampler": {
|
| "type": "string",
|
| "description": "Sampling algorithm (e.g., 'euler', 'dpmpp_2m'). If omitted, model optimal default is applied."
|
| },
|
| "scheduler": {
|
| "type": "string",
|
| "description": "Noise scheduler type (e.g., 'simple', 'karras'). If omitted, model optimal default is applied."
|
| },
|
| "negative_prompt": {
|
| "type": "string",
|
| "default": "",
|
| "description": "Negative prompt specifying undesired attributes or quality flaws to avoid."
|
| },
|
| "async_execution": {
|
| "type": "boolean",
|
| "default": False,
|
| "description": "If true, submits task asynchronously and returns task_id immediately for polling via get_task_status."
|
| },
|
| "chain": {
|
| "type": "array",
|
| "description": "List of injector objects for extended features (LoRA, ControlNet, IPAdapter, etc.). Must be a JSON array of dicts with 'injector_type'."
|
| }
|
| }
|
|
|
| _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,
|
| "optional_inputs_schema": _COMMON_OPTIONAL_INPUTS_SCHEMA,
|
| "example_json_params": {
|
| "task_type": "txt2img",
|
| "model": "stabilityai/SDXL-Base-1.0",
|
| "prompt": "A majestic lion jumping from a big stone at night",
|
| "width": 1024,
|
| "height": 1024
|
| },
|
| "example_json_params_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": 0.8
|
| }
|
| ]
|
| }
|
| },
|
| {
|
| "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,
|
| "optional_inputs_schema": _COMMON_OPTIONAL_INPUTS_SCHEMA,
|
| "example_json_params": {
|
| "task_type": "img2img",
|
| "model": "stabilityai/SDXL-Base-1.0",
|
| "prompt": "A majestic lion jumping from a big stone at night",
|
| "image": "https://example.com/source_image.png",
|
| "denoise": 0.7
|
| }
|
| },
|
| {
|
| "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,
|
| "optional_inputs_schema": _COMMON_OPTIONAL_INPUTS_SCHEMA,
|
| "example_json_params": {
|
| "task_type": "inpaint",
|
| "model": "stabilityai/SDXL-Base-1.0",
|
| "prompt": "red floral dress, detailed lace",
|
| "image": "https://example.com/image_with_alpha_mask.png",
|
| "denoise": 0.95
|
| }
|
| },
|
| {
|
| "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,
|
| "optional_inputs_schema": _COMMON_OPTIONAL_INPUTS_SCHEMA,
|
| "example_json_params": {
|
| "task_type": "outpaint",
|
| "model": "stabilityai/SDXL-Base-1.0",
|
| "prompt": "beautiful scenery background, high quality",
|
| "image": "https://example.com/source_image.png",
|
| "pad_left": 128,
|
| "pad_right": 128,
|
| "pad_top": 0,
|
| "pad_bottom": 0
|
| }
|
| },
|
| {
|
| "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": ["upscaler", "denoise"] + _COMMON_OPTIONAL_INPUTS,
|
| "optional_inputs_schema": _COMMON_OPTIONAL_INPUTS_SCHEMA,
|
| "example_json_params": {
|
| "task_type": "hires_fix",
|
| "model": "stabilityai/SDXL-Base-1.0",
|
| "prompt": "masterpiece, best quality, sharp focus",
|
| "image": "https://example.com/low_res_image.png",
|
| "upscale_by": 2.0,
|
| "upscaler": "nearest-exact",
|
| "denoise": 0.55
|
| }
|
| },
|
| ]
|
|
|
| _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
|
| upscaler = params.get("upscaler", "nearest-exact")
|
| if upscaler == "latent" or upscaler not in ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]:
|
| upscaler = "nearest-exact"
|
| ui_inputs["hires_upscaler"] = upscaler
|
| 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 = []
|
| embedding_data = []
|
| controlnet_data = []
|
| diffsynth_controlnet_data = []
|
| ipadapter_data = []
|
| ipadapter_images = []
|
| ipadapter_weights = []
|
| ipadapter_lora_strengths = []
|
| ipadapter_global_preset = params.get("ipadapter_preset") or params.get("preset")
|
| ipadapter_global_embeds_scaling = params.get("ipadapter_embeds_scaling") or params.get("embeds_scaling")
|
| ipadapter_global_combine_method = params.get("ipadapter_combine_method") or params.get("combine_method")
|
| ipadapter_global_final_weight = params.get("ipadapter_final_weight") or params.get("final_weight")
|
| flux1_ipadapter_images = []
|
| flux1_ipadapter_weights = []
|
| flux1_ipadapter_starts = []
|
| flux1_ipadapter_ends = []
|
| sd3_ipadapter_images = []
|
| sd3_ipadapter_weights = []
|
| sd3_ipadapter_starts = []
|
| sd3_ipadapter_ends = []
|
| style_images = []
|
| style_strengths = []
|
| krea2_identity_edit_data = []
|
| krea2_reference_edit_data = []
|
| krea2_controlnet_data = []
|
| anima_controlnet_lllite_data = []
|
| reference_latent_data = []
|
| reference_image_data = []
|
| joyai_reference_data = []
|
| boogu_edit_data = []
|
| qwen_image_edit_data = []
|
| hidream_o1_reference_data = []
|
| cond_prompts = []
|
| cond_widths = []
|
| cond_heights = []
|
| cond_xs = []
|
| cond_ys = []
|
| cond_strengths = []
|
|
|
| for item in chain:
|
| itype = item.get("injector_type")
|
| if itype == "lora":
|
| lora_data.extend([
|
| item.get("source", item.get("lora_source", "Civitai")),
|
| item.get("lora_value", ""),
|
| item.get("scale", 1.0),
|
| None
|
| ])
|
| elif itype == "embedding":
|
| e_source = item.get("source", item.get("embedding_source", "Civitai"))
|
| e_val = item.get("embedding_value", item.get("value", item.get("embedding_id", "")))
|
| if e_source and e_val:
|
| embedding_data.extend([
|
| e_source,
|
| str(e_val),
|
| None
|
| ])
|
| elif itype == "conditioning":
|
| p = item.get("prompt", "")
|
| if p:
|
| cond_prompts.append(p)
|
| cond_widths.append(int(item.get("width", 512)))
|
| cond_heights.append(int(item.get("height", 512)))
|
| cond_xs.append(int(item.get("x", 0)))
|
| cond_ys.append(int(item.get("y", 0)))
|
| cond_strengths.append(float(item.get("strength", 1.0)))
|
| elif itype == "controlnet":
|
| cn_type = item.get("type", item.get("Type", ""))
|
| cn_series = item.get("series", item.get("Series", ""))
|
| cn_strength = float(item.get("strength", 1.0))
|
| cn_img = _parse_image_param(item.get("image"))
|
|
|
| cn_filepath = item.get("control_net_name", "None")
|
| cn_raw = _load_yaml(os.path.join(_YAML_DIR, "controlnet_models.yaml")).get("ControlNet", {})
|
|
|
| cn_arch_key = None
|
| if found_arch:
|
| arch_cfg = _load_yaml(os.path.join(_YAML_DIR, "model_architectures.yaml")).get("architectures", {})
|
| cn_arch_key = arch_cfg.get(found_arch, {}).get("controlnet_key", found_arch)
|
|
|
| arch_entries = []
|
| if cn_arch_key and cn_arch_key in cn_raw:
|
| arch_entries = cn_raw[cn_arch_key]
|
| elif found_arch and found_arch in cn_raw:
|
| arch_entries = cn_raw[found_arch]
|
| else:
|
| for val in cn_raw.values():
|
| if isinstance(val, list):
|
| arch_entries.extend(val)
|
| elif isinstance(val, dict):
|
| arch_entries.append(val)
|
|
|
| if arch_entries:
|
| for entry in arch_entries:
|
| entry_types = entry.get("Type", [])
|
| if isinstance(entry_types, str):
|
| entry_types = [entry_types]
|
| if not cn_type or cn_type in entry_types:
|
| if not cn_series or entry.get("Series") == cn_series:
|
| cn_filepath = entry.get("Filepath", cn_filepath)
|
| if not cn_series:
|
| cn_series = entry.get("Series", "")
|
| if not cn_type and entry_types:
|
| cn_type = entry_types[0]
|
| break
|
|
|
| controlnet_data.extend([
|
| cn_img,
|
| cn_type,
|
| cn_series,
|
| cn_strength,
|
| cn_filepath
|
| ])
|
| elif itype == "anima_controlnet_lllite":
|
| cn_type = item.get("type", item.get("Type", ""))
|
| cn_series = item.get("series", item.get("Series", ""))
|
| cn_strength = float(item.get("strength", 1.0))
|
| cn_start = float(item.get("start_percent", 0.0))
|
| cn_end = float(item.get("end_percent", 1.0))
|
| cn_img = _parse_image_param(item.get("image"))
|
|
|
| cn_filepath = item.get("control_net_name", "None")
|
| anima_cfg = _load_yaml(os.path.join(_YAML_DIR, "anima_controlnet_lllite_models.yaml")).get("Anima_ControlNet_Lllite", [])
|
| if anima_cfg:
|
| for entry in anima_cfg:
|
| entry_types = entry.get("Type", [])
|
| if isinstance(entry_types, str):
|
| entry_types = [entry_types]
|
| if not cn_type or cn_type in entry_types:
|
| if not cn_series or entry.get("Series") == cn_series:
|
| cn_filepath = entry.get("Filepath", cn_filepath)
|
| if not cn_series:
|
| cn_series = entry.get("Series", "")
|
| if not cn_type and entry_types:
|
| cn_type = entry_types[0]
|
| break
|
|
|
| anima_controlnet_lllite_data.extend([
|
| cn_img,
|
| cn_type,
|
| cn_series,
|
| cn_strength,
|
| cn_filepath,
|
| cn_start,
|
| cn_end
|
| ])
|
| elif itype == "diffsynth_controlnet":
|
| cn_type = item.get("type", "")
|
| cn_series = item.get("series", "")
|
| cn_strength = float(item.get("strength", 1.0))
|
| cn_img = _parse_image_param(item.get("image"))
|
|
|
| cn_filepath = "None"
|
| diffsynth_raw = _load_yaml(os.path.join(_YAML_DIR, "diffsynth_controlnet_models.yaml")).get("DiffSynth_ControlNet", {})
|
| diffsynth_entries = []
|
| if isinstance(diffsynth_raw, dict):
|
| for val in diffsynth_raw.values():
|
| if isinstance(val, list):
|
| diffsynth_entries.extend(val)
|
| elif isinstance(val, dict):
|
| diffsynth_entries.append(val)
|
| elif isinstance(diffsynth_raw, list):
|
| diffsynth_entries = diffsynth_raw
|
|
|
| if diffsynth_entries:
|
| for entry in diffsynth_entries:
|
| if not cn_type or cn_type in entry.get("Type", []):
|
| if not cn_series or entry.get("Series") == cn_series:
|
| cn_filepath = entry.get("Filepath", cn_filepath)
|
| if not cn_series:
|
| cn_series = entry.get("Series", cn_series)
|
| if not cn_type and entry.get("Type"):
|
| cn_type = entry.get("Type")[0]
|
| break
|
|
|
| diffsynth_controlnet_data.extend([
|
| cn_img,
|
| cn_type,
|
| cn_series,
|
| cn_strength,
|
| cn_filepath
|
| ])
|
| elif itype == "krea2_controlnet":
|
| cn_type = item.get("type", "Depth")
|
| cn_series = item.get("series", "Patil")
|
| cn_strength = float(item.get("strength", 1.0))
|
| cn_img = _parse_image_param(item.get("image"))
|
|
|
| cn_filepath = "depth-control-lora.safetensors"
|
| krea2_cfg = _load_yaml(os.path.join(_YAML_DIR, "krea2_controlnet_models.yaml")).get("Krea2_ControlNet", [])
|
| if krea2_cfg:
|
| for entry in krea2_cfg:
|
| if cn_type in entry.get("Type", []):
|
| if not cn_series or entry.get("Series") == cn_series:
|
| cn_filepath = entry.get("Filepath", cn_filepath)
|
| cn_series = entry.get("Series", cn_series)
|
| break
|
|
|
| krea2_controlnet_data.extend([
|
| cn_img,
|
| cn_type,
|
| cn_series,
|
| cn_strength,
|
| cn_filepath
|
| ])
|
| elif itype == "flux1_ipadapter":
|
| if len(flux1_ipadapter_images) < 5:
|
| img = _parse_image_param(item.get("image"))
|
| weight = float(item.get("weight", 1.0))
|
| start_at = float(item.get("start_at", item.get("start_percent", item.get("start", 0.0))))
|
| end_at = float(item.get("end_at", item.get("end_percent", item.get("end", 1.0))))
|
| flux1_ipadapter_images.append(img)
|
| flux1_ipadapter_weights.append(weight)
|
| flux1_ipadapter_starts.append(start_at)
|
| flux1_ipadapter_ends.append(end_at)
|
| elif itype == "sd3_ipadapter":
|
| if len(sd3_ipadapter_images) < 5:
|
| img = _parse_image_param(item.get("image"))
|
| weight = float(item.get("weight", 1.0))
|
| start_at = float(item.get("start_at", item.get("start_percent", item.get("start", 0.0))))
|
| end_at = float(item.get("end_at", item.get("end_percent", item.get("end", 1.0))))
|
| sd3_ipadapter_images.append(img)
|
| sd3_ipadapter_weights.append(weight)
|
| sd3_ipadapter_starts.append(start_at)
|
| sd3_ipadapter_ends.append(end_at)
|
| elif itype == "ipadapter":
|
| if len(ipadapter_images) < 5:
|
| img = _parse_image_param(item.get("image"))
|
| weight = float(item.get("weight", 1.0))
|
| lora_str = float(item.get("lora_strength", 0.6))
|
| ipadapter_images.append(img)
|
| ipadapter_weights.append(weight)
|
| ipadapter_lora_strengths.append(lora_str)
|
|
|
| if "preset" in item and not ipadapter_global_preset:
|
| ipadapter_global_preset = item["preset"]
|
| if "embeds_scaling" in item and not ipadapter_global_embeds_scaling:
|
| ipadapter_global_embeds_scaling = item["embeds_scaling"]
|
| if "combine_method" in item and not ipadapter_global_combine_method:
|
| ipadapter_global_combine_method = item["combine_method"]
|
| if "final_weight" in item and ipadapter_global_final_weight is None:
|
| ipadapter_global_final_weight = float(item["final_weight"])
|
| elif itype in ("style", "flux1_style"):
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| style_images.append(img)
|
| style_strengths.append(float(item.get("strength", item.get("weight", 1.0))))
|
| elif itype == "pid":
|
| is_enabled = item.get("enabled", True)
|
| if isinstance(is_enabled, str):
|
| is_enabled = is_enabled.upper() in ("ON", "TRUE", "1")
|
| ui_inputs["pid_settings"] = "ON" if is_enabled else "OFF"
|
| elif itype == "krea2_identity_edit":
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| krea2_identity_edit_data.append(img)
|
| elif itype == "krea2_style_reference":
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| krea2_reference_edit_data.append(img)
|
| elif itype in ("reference_latent", "reference_edit"):
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| reference_latent_data.append(img)
|
| elif itype in ("reference_image", "mage_flow_reference_edit"):
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| reference_image_data.append(img)
|
| elif itype in ("joyai_image", "joyai_reference_edit"):
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| joyai_reference_data.append(img)
|
| elif itype in ("boogu_image_edit", "boogu_edit"):
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| boogu_edit_data.append(img)
|
| elif itype == "qwen_image_edit":
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| qwen_image_edit_data.append(img)
|
| elif itype == "hidream_o1_reference":
|
| img = _parse_image_param(item.get("image"))
|
| if img:
|
| hidream_o1_reference_data.append(img)
|
| elif itype == "vae":
|
| v_source = item.get("source", item.get("vae_source", "Civitai"))
|
| v_val = item.get("vae_value", item.get("value", item.get("vae_id", item.get("vae_name", ""))))
|
| if v_source and v_val:
|
| ui_inputs["vae_source"] = v_source
|
| ui_inputs["vae_id"] = str(v_val)
|
|
|
| if lora_data: ui_inputs["lora_data"] = lora_data
|
| if embedding_data: ui_inputs["embedding_data"] = embedding_data
|
| if controlnet_data: ui_inputs["controlnet_data"] = controlnet_data
|
| if anima_controlnet_lllite_data: ui_inputs["anima_controlnet_lllite_data"] = anima_controlnet_lllite_data
|
| if diffsynth_controlnet_data: ui_inputs["diffsynth_controlnet_data"] = diffsynth_controlnet_data
|
| if krea2_controlnet_data: ui_inputs["krea2_controlnet_data"] = krea2_controlnet_data
|
| if ipadapter_images:
|
| preset = ipadapter_global_preset or "STANDARD (medium strength)"
|
| embeds_scaling = ipadapter_global_embeds_scaling or "V only"
|
| combine_method = ipadapter_global_combine_method or "concat"
|
| final_weight = float(ipadapter_global_final_weight) if ipadapter_global_final_weight is not None else 1.0
|
| final_lora_strength = 0.6
|
|
|
| presets_by_arch = _get_ipadapter_presets_by_arch()
|
| target_arch = "SD1.5" if found_arch in ("sd15", "SD1.5") else "SDXL"
|
| allowed_presets = presets_by_arch.get(target_arch, [])
|
|
|
| if preset not in allowed_presets:
|
| raise ValueError(
|
| f"Invalid IPAdapter preset '{preset}' for model architecture '{target_arch}'. "
|
| f"Preset must match the target model architecture. Allowed presets for {target_arch}: {allowed_presets}"
|
| )
|
|
|
| ui_inputs["ipadapter_data"] = (
|
| ipadapter_images + ipadapter_weights + ipadapter_lora_strengths +
|
| [preset, final_weight, final_lora_strength, embeds_scaling, combine_method]
|
| )
|
| elif ipadapter_data:
|
| ui_inputs["ipadapter_data"] = ipadapter_data
|
| if flux1_ipadapter_images:
|
| ui_inputs["flux1_ipadapter_data"] = (
|
| flux1_ipadapter_images + flux1_ipadapter_weights + flux1_ipadapter_starts + flux1_ipadapter_ends
|
| )
|
| if sd3_ipadapter_images:
|
| ui_inputs["sd3_ipadapter_chain"] = (
|
| sd3_ipadapter_images + sd3_ipadapter_weights + sd3_ipadapter_starts + sd3_ipadapter_ends
|
| )
|
| if style_images: ui_inputs["style_data"] = style_images + style_strengths
|
| if krea2_identity_edit_data: ui_inputs["krea2_identity_edit_data"] = krea2_identity_edit_data
|
| if krea2_reference_edit_data: ui_inputs["krea2_reference_edit_data"] = krea2_reference_edit_data
|
| if reference_latent_data: ui_inputs["reference_latent_data"] = reference_latent_data
|
| if reference_image_data: ui_inputs["reference_image_data"] = reference_image_data
|
| if joyai_reference_data: ui_inputs["joyai_reference_data"] = joyai_reference_data
|
| if boogu_edit_data: ui_inputs["boogu_edit_data"] = boogu_edit_data
|
| if qwen_image_edit_data: ui_inputs["qwen_image_edit_data"] = qwen_image_edit_data
|
| if hidream_o1_reference_data: ui_inputs["hidream_o1_reference_data"] = hidream_o1_reference_data
|
| if cond_prompts:
|
| ui_inputs["conditioning_data"] = (
|
| cond_prompts + cond_widths + cond_heights + cond_xs + cond_ys + cond_strengths
|
| )
|
|
|
| if "vae_source" in params and "vae_id" in params:
|
| ui_inputs["vae_source"] = params["vae_source"]
|
| ui_inputs["vae_id"] = str(params["vae_id"])
|
|
|
| pid_val = params.get("pid") if params.get("pid") is not None else params.get("pid_settings")
|
| if pid_val is not None:
|
| if isinstance(pid_val, bool):
|
| ui_inputs["pid_settings"] = "ON" if pid_val else "OFF"
|
| elif str(pid_val).upper() in ("ON", "TRUE", "1"):
|
| ui_inputs["pid_settings"] = "ON"
|
| else:
|
| ui_inputs["pid_settings"] = "OFF"
|
|
|
| _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),
|
| }
|
|
|
