Spaces:
Running on Zero
Running on Zero
| """ | |
| 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 ipaddress | |
| import socket | |
| import base64 | |
| import io | |
| import copy | |
| import threading | |
| import yaml | |
| from typing import Dict, Any | |
| from PIL import Image | |
| from core.runtime_config import CONFIG | |
| _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 # 50 MB | |
| _IMAGE_DOWNLOAD_TIMEOUT = 30 # seconds | |
| _ALLOWED_IMAGE_CONTENT_TYPES = frozenset([ | |
| "image/png", "image/jpeg", "image/jpg", "image/gif", | |
| "image/webp", "image/bmp", "image/tiff", | |
| ]) | |
| def _validate_public_image_url(url: str) -> None: | |
| parsed = urllib.parse.urlsplit(url) | |
| if parsed.scheme not in {"http", "https"} or not parsed.hostname: | |
| raise ValueError("Image URL must use HTTP or HTTPS and include a host.") | |
| if parsed.username or parsed.password: | |
| raise ValueError("Image URLs containing credentials are not accepted.") | |
| try: | |
| addresses = { | |
| item[4][0] | |
| for item in socket.getaddrinfo( | |
| parsed.hostname, | |
| parsed.port or (443 if parsed.scheme == "https" else 80), | |
| ) | |
| } | |
| except socket.gaierror as exc: | |
| raise ValueError(f"Could not resolve image URL host: {exc}") from exc | |
| for address in addresses: | |
| ip = ipaddress.ip_address(address.split("%", 1)[0]) | |
| if not ip.is_global: | |
| raise ValueError("Image URL must resolve to a public network address.") | |
| class _PublicOnlyRedirectHandler(urllib.request.HTTPRedirectHandler): | |
| def redirect_request(self, req, fp, code, msg, headers, newurl): | |
| _validate_public_image_url(newurl) | |
| return super().redirect_request(req, fp, code, msg, headers, newurl) | |
| 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 _validate_decoded_image(image_param) | |
| if not isinstance(image_param, str) or not image_param.strip(): | |
| return None | |
| image_param = image_param.strip() | |
| # HTTP / HTTPS URL — download the image | |
| if image_param.startswith("http://") or image_param.startswith("https://"): | |
| return _validate_decoded_image(_download_image_from_url(image_param)) | |
| # 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) | |
| if len(encoded) > (_MAX_IMAGE_DOWNLOAD_BYTES * 4 // 3) + 8: | |
| raise ValueError("Base64 image exceeds the 50 MB decoded-size limit.") | |
| data = base64.b64decode(encoded) | |
| if len(data) > _MAX_IMAGE_DOWNLOAD_BYTES: | |
| raise ValueError("Decoded image exceeds the 50 MB size limit.") | |
| return _validate_decoded_image(Image.open(io.BytesIO(data))) | |
| # Base64 string without header | |
| if len(image_param) > 100: | |
| try: | |
| if len(image_param) > (_MAX_IMAGE_DOWNLOAD_BYTES * 4 // 3) + 8: | |
| raise ValueError("Base64 image exceeds the 50 MB decoded-size limit.") | |
| data = base64.b64decode(image_param) | |
| if len(data) > _MAX_IMAGE_DOWNLOAD_BYTES: | |
| raise ValueError("Decoded image exceeds the 50 MB size limit.") | |
| return _validate_decoded_image(Image.open(io.BytesIO(data))) | |
| except ValueError: | |
| raise | |
| 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 _validate_decoded_image(image: Image.Image) -> Image.Image: | |
| image.load() | |
| megapixels = (image.width * image.height) / 1_000_000 | |
| if megapixels > CONFIG.max_input_megapixels: | |
| raise ValueError( | |
| f"Decoded image is {megapixels:.1f} MP; maximum is " | |
| f"{CONFIG.max_input_megapixels:g} MP." | |
| ) | |
| return image.copy() | |
| 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. | |
| """ | |
| _validate_public_image_url(url) | |
| req = urllib.request.Request(url, headers={"User-Agent": "ImageGen-MCP/1.0"}) | |
| opener = urllib.request.build_opener(_PublicOnlyRedirectHandler()) | |
| try: | |
| with opener.open(req, timeout=_IMAGE_DOWNLOAD_TIMEOUT) as resp: | |
| # Validate content type | |
| 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))}." | |
| ) | |
| # Enforce size limit | |
| 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." | |
| ) | |
| # Stream-read with size cap | |
| 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.HTTPError as e: | |
| raise ValueError(f"HTTP error {e.code} when downloading image from URL: {e.reason}") from e | |
| except urllib.error.URLError as e: | |
| raise ValueError(f"Failed to download image from URL: {e}") 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", "chain", "async_execution", | |
| "images", "run_mode", "extra_models", "reference_role", "use_model_defaults", | |
| ] | |
| _TASK_DEFINITIONS = [ | |
| { | |
| "task_type": "reference", | |
| "display_name": "参考图编辑 / 融合", | |
| "description": "一张图按指令编辑，多张图整组参考。images 按数组顺序从 1 编号；图1 / img1 / image1 等价。先查询 get_model_features.image_workflows 的用途和张数上限，不需要填写 chain。", | |
| "required_inputs": ["prompt", "images", "width", "height"], | |
| "optional_inputs": [field for field in _COMMON_OPTIONAL_INPUTS if field != "images"], | |
| }, | |
| { | |
| "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 defaults to 0.7 when omitted.", | |
| "required_inputs": ["prompt", "image"], | |
| "optional_inputs": ["denoise"] + _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": ["upscaler", "denoise"] + _COMMON_OPTIONAL_INPUTS, | |
| }, | |
| ] | |
| def _find_model_record(model: str) -> tuple[str | None, dict | None]: | |
| """Return the registry architecture and record for a display name.""" | |
| model_list = _load_yaml(_MODEL_LIST_PATH) | |
| for architecture, architecture_data in (model_list.get("Checkpoint", {}) or {}).items(): | |
| if not isinstance(architecture_data, dict): | |
| continue | |
| for record in architecture_data.get("models", []) or []: | |
| if isinstance(record, dict) and record.get("display_name") == model: | |
| return architecture, record | |
| return None, None | |
| def _resolve_feature_name(chain_type: str) -> str | None: | |
| """Resolve a canonical chain feature from its name or configured alias.""" | |
| chain_features = _load_yaml(_CHAIN_FEATURES_PATH) | |
| if chain_type in chain_features: | |
| return chain_type | |
| for feature_name, feature_data in chain_features.items(): | |
| aliases = feature_data.get("chains") if isinstance(feature_data, dict) else None | |
| if aliases == chain_type or (isinstance(aliases, list) and chain_type in aliases): | |
| return feature_name | |
| return None | |
| _FEATURE_ADAPTER_ALIASES = { | |
| "boogu_edit": "boogu_image_edit", | |
| } | |
| def _canonicalize_chain_type(chain_type: str) -> str: | |
| """Normalize explicitly supported legacy request aliases.""" | |
| return _FEATURE_ADAPTER_ALIASES.get(chain_type, chain_type) | |
| def _get_model_feature_context(model: str) -> dict | None: | |
| """Return registry-backed architecture and enabled-chain information.""" | |
| architecture, _ = _find_model_record(model) | |
| if not architecture: | |
| return None | |
| architecture_config = _load_yaml(_MODEL_ARCHITECTURES_PATH).get("architectures", {}) or {} | |
| architecture_info = architecture_config.get(architecture, {}) or {} | |
| model_type = architecture_info.get("model_type", architecture.lower()) | |
| features_config = _load_yaml(_IMAGE_GEN_FEATURES_PATH) | |
| enabled = features_config.get(model_type, {}) or {} | |
| from core.model_capabilities import supports_chain_for_model | |
| return { | |
| "architecture": architecture, | |
| "model_type": model_type, | |
| "enabled_chains": set(enabled.get("enabled_chains", []) or []), | |
| "supports_chain": lambda chain: supports_chain_for_model(model, chain), | |
| } | |
| _TASKS_DB: Dict[str, Dict[str, Any]] = {} | |
| _TASKS_LOCK = threading.RLock() | |
| def _update_task(task_id: str, **changes) -> None: | |
| with _TASKS_LOCK: | |
| if task_id in _TASKS_DB: | |
| _TASKS_DB[task_id].update(changes) | |
| def _get_task_snapshot(task_id: str) -> Dict[str, Any] | None: | |
| with _TASKS_LOCK: | |
| task = _TASKS_DB.get(task_id) | |
| return copy.deepcopy(task) if task is not None else None | |
| 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: | |
| _update_task( | |
| task_id, | |
| status="processing", | |
| progress=10, | |
| 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 task_type in {"img2img", "inpaint", "outpaint", "hires_fix"}: | |
| pil_img = _parse_image_param(params.get("image")) | |
| if pil_img is not None or params.get("images"): | |
| 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" | |
| _update_task(task_id, progress=50) | |
| from core.execution_plan import build_execution_plan, execute_generation_plan | |
| plan = build_execution_plan( | |
| ui_inputs, mode=params.get("run_mode", "single"), | |
| images=[_parse_image_param(value) for value in params.get("images", [])], | |
| extra_models=params.get("extra_models"), | |
| reference_role=params.get("reference_role", "auto"), | |
| use_model_defaults=params.get("use_model_defaults", True), | |
| ) | |
| jobs = [] | |
| def generate(values, progress): | |
| job = {"model": values["model_display_name"], "task_type": values["task_type"], "status": "processing"} | |
| if "_source_image_index" in values: | |
| job["source_index"] = values["_source_image_index"] | |
| jobs.append(job) | |
| try: | |
| result = sd_image_pipeline.run(ui_inputs=values, progress=progress) | |
| if result is None or (isinstance(result, list) and not result): | |
| raise ValueError("生成器没有返回图片。") | |
| job.update(status="completed", output_count=len(result) if isinstance(result, list) else int(result is not None)) | |
| return result | |
| except Exception as exc: | |
| job.update(status="failed", error=str(exc)) | |
| raise | |
| finally: | |
| job.update( | |
| seed=values.get("seed"), width=values.get("width"), height=values.get("height"), | |
| effective_prompt=values.get("positive_prompt", ""), | |
| image_references=values.get("_image_references", []), | |
| ) | |
| output, summary = execute_generation_plan(plan, generate, progress=DummyProgress()) | |
| ui_inputs = plan[0].inputs | |
| 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 = [] | |
| captions = [] | |
| for idx, (item, caption) in enumerate(output): | |
| captions.append(caption) | |
| 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) | |
| _update_task( | |
| task_id, | |
| status="completed", | |
| progress=100, | |
| completed_at=int(time.time()), | |
| result={ | |
| "images": images, | |
| "seed": ui_inputs.get("seed", params.get("seed", -1)), | |
| "width": ui_inputs.get("width", params.get("width", 1024)), | |
| "height": ui_inputs.get("height", params.get("height", 1024)), | |
| "execution_time_seconds": execution_time, | |
| "effective_prompt": ui_inputs.get("positive_prompt", ""), | |
| "image_references": ui_inputs.get("_image_references", []), | |
| "captions": captions, | |
| "jobs": jobs, | |
| "partial_success": any(job["status"] == "failed" for job in jobs), | |
| "summary": summary, | |
| }, | |
| ) | |
| except Exception as e: | |
| _update_task( | |
| task_id, | |
| status="failed", | |
| progress=0, | |
| failed_at=int(time.time()), | |
| error={"code": "EXECUTION_ERROR", "message": str(e)}, | |
| ) | |