"""Panel-style o1key image generation node.""" from __future__ import annotations import asyncio import json import os from typing import Any import folder_paths from PIL import Image, ImageOps from comfy_api.latest import io, ui import numpy as np import torch from .nano_banana import NanoBanana from ..clients.gpt_image_client import GptImageClient, resolve_gpt_image_model from ..clients.seedream_image_client import ( SeedreamImageClient, resolve_seedream_model, validate_seedream_reference_images, ) from ..utils.config import get_api_key_or_raise, get_base_url_by_route from ..utils.http2_client import create_http_client from ..utils.image_utils import ( IMAGE_BATCH_MODE_GROUP_TO_MODELS, IMAGE_BATCH_MODE_SINGLE_REFERENCES, IMAGE_BATCH_MODES, expand_image_generation_tasks, pil_to_tensor, tensor_to_pil, ) from ..utils.http_error import format_o1key_image_error from ..utils.o1key_image_catalog import ( BANANA_ASPECT_RATIO_OPTIONS, BANANA_RESOLUTION_OPTIONS, GPT_IMAGE_ASPECT_RATIO_OPTIONS, GPT_IMAGE_BACKGROUND_OPTIONS, GPT_IMAGE_COUNTS, GPT_IMAGE_EXACT_SIZE_OPTIONS, GPT_IMAGE_OUTPUT_FORMAT_OPTIONS, GPT_IMAGE_25_QUALITY_OPTIONS, GPT_IMAGE_RESOLUTION_OPTIONS, MAX_UNIFIED_BATCH_IMAGES, MAX_UNIFIED_IMAGE_TASKS, MAX_UNIFIED_REFERENCE_IMAGES, UNIFIED_IMAGE_COUNTS, UNIFIED_IMAGE_MODEL_OPTIONS, UNIFIED_IMAGE_ROUTE_OPTIONS, UNIFIED_IMAGE_SMART_RESOLUTION, SEEDREAM_ASPECT_RATIO_OPTIONS, SEEDREAM_LAYER_RESOLUTION_OPTIONS, SEEDREAM_OUTPUT_FORMAT_OPTIONS, SEEDREAM_RESOLUTION_OPTIONS, is_gpt_image_model, is_seedream_model, resolve_gpt_image_size, resolve_gpt_image_quality, resolve_seedream_layer_size, resolve_seedream_size, ) from ..utils.o1key_image_save import ( DEFAULT_SAVE_NAMING_RULE, SAVE_FORMAT_OPTIONS, SAVE_NAMING_RULE_OPTIONS, normalize_naming_rule, normalize_save_format, normalize_save_location, save_tensor_images, ) _IMAGE_COUNTS = [str(value) for value in sorted(set(UNIFIED_IMAGE_COUNTS + GPT_IMAGE_COUNTS))] _ALL_RESOLUTIONS = [ UNIFIED_IMAGE_SMART_RESOLUTION, *BANANA_RESOLUTION_OPTIONS, *(value for value in GPT_IMAGE_RESOLUTION_OPTIONS if value not in BANANA_RESOLUTION_OPTIONS), *(value for value in SEEDREAM_LAYER_RESOLUTION_OPTIONS if value not in BANANA_RESOLUTION_OPTIONS), ] def _is_enabled(value: Any) -> bool: return value is True or str(value or "").strip().lower() in {"1", "true", "开启"} def _attach_main_reference_filename(output_tensor, references: dict[str, Any]) -> None: first_reference = next(iter(references.values()), None) metadata = getattr(first_reference, "_o1key_source_metadata", None) if not isinstance(metadata, list) or not metadata or not isinstance(metadata[0], dict): return filename = metadata[0].get("filename") if isinstance(filename, str) and filename.strip(): setattr(output_tensor, "_o1key_main_filename", filename) def _normalized_save_settings( model: str, filename_prefix: str, save_format: str, save_location: str, naming_rule: str, ) -> dict[str, str]: """Validate generator-owned save settings before any paid request.""" selected_rule = normalize_naming_rule(naming_rule) selected_prefix = str(filename_prefix or "o1key").strip() if selected_rule == "自定义前缀" and ( not selected_prefix or len(selected_prefix) > 512 ): raise ValueError("文件名前缀无效") return { "filename_prefix": selected_prefix or "o1key", "format": ( "原始" if is_gpt_image_model(model) or is_seedream_model(model) else normalize_save_format(save_format) ), "save_location": normalize_save_location(save_location), "naming_rule": selected_rule, } def _attach_save_settings( output_tensor, references: dict[str, Any], save_settings: dict[str, str], ) -> None: _attach_main_reference_filename(output_tensor, references) setattr(output_tensor, "_o1key_save_settings", dict(save_settings)) def _parse_upload_manifest( value: str | list[dict[str, Any]] | None, label: str = "参考图", max_images: int = MAX_UNIFIED_REFERENCE_IMAGES, ) -> list[dict[str, str]]: """Parse the serialized upload list and keep only the fields the backend uses.""" if value in (None, ""): return [] if isinstance(value, str): try: value = json.loads(value) except json.JSONDecodeError as exc: raise ValueError(f"{label}清单格式无效，请重新上传图片") from exc if not isinstance(value, list): raise ValueError(f"{label}清单必须是数组") if len(value) > max_images: raise ValueError(f"{label}最多支持 {max_images} 张") result: list[dict[str, str]] = [] for item in value: if not isinstance(item, dict): raise ValueError(f"{label}清单包含无效项目") name = str(item.get("name") or "").strip() subfolder = str(item.get("subfolder") or "").strip() folder_type = str(item.get("type") or "input").strip() if not name: raise ValueError(f"{label}文件名不能为空") if folder_type != "input": raise ValueError(f"{label}必须来自 ComfyUI input 目录") result.append({"name": name, "subfolder": subfolder, "type": "input"}) return result def _parse_mask_manifest(value: str | dict[str, Any] | None) -> dict[str, str] | None: if value in (None, "", "{}"): return None if isinstance(value, str): try: value = json.loads(value) except json.JSONDecodeError as exc: raise ValueError("蒙版清单格式无效，请重新上传蒙版") from exc if not isinstance(value, dict): raise ValueError("蒙版清单必须是对象") name = str(value.get("name") or "").strip() subfolder = str(value.get("subfolder") or "").strip() folder_type = str(value.get("type") or "input").strip() if not name: return None if folder_type != "input": raise ValueError("蒙版必须来自 ComfyUI input 目录") return {"name": name, "subfolder": subfolder, "type": "input"} def _resolve_input_image(item: dict[str, str], label: str = "参考图") -> str: """Resolve an upload descriptor without allowing it to escape the input folder.""" input_root = os.path.abspath(folder_paths.get_input_directory()) candidate = os.path.abspath( os.path.join(input_root, item.get("subfolder", ""), item["name"]) ) try: inside_input = os.path.commonpath([input_root, candidate]) == input_root except ValueError: inside_input = False if not inside_input: raise ValueError(f"{label}路径不安全：{item['name']}") if not os.path.isfile(candidate): raise ValueError(f"{label}不存在：{item['name']}") return candidate def _load_reference_tensors( manifest: str | list[dict[str, Any]] | None, *, label: str = "参考图", max_images: int = MAX_UNIFIED_REFERENCE_IMAGES, ) -> dict[str, Any]: """Load uploads as independent tensors so mixed image dimensions remain supported.""" tensors: dict[str, Any] = {} for index, item in enumerate( _parse_upload_manifest(manifest, label, max_images), start=1, ): path = _resolve_input_image(item, label) try: with Image.open(path) as opened: source_format = str(opened.format or "").upper() orientation = opened.getexif().get(274, 1) image = ImageOps.exif_transpose(opened).convert("RGB").copy() image.format = source_format or None setattr(image, "_o1key_original_format", source_format or None) setattr(image, "_o1key_original_filename", item["name"]) # Only reuse the original bytes when no EXIF rotation/flip is needed. # Otherwise upload the already transposed pixels in the same format. if orientation in (None, 1): setattr(image, "_o1key_original_path", path) except Exception as exc: raise ValueError(f"无法读取{label}：{item['name']}") from exc tensors[f"{label}{index}"] = pil_to_tensor([image]) return tensors def _load_mask_tensor(manifest: str | dict[str, Any] | None): item = _parse_mask_manifest(manifest) if item is None: return None path = _resolve_input_image(item, "蒙版") try: with Image.open(path) as opened: image = ImageOps.exif_transpose(opened).convert("RGB").copy() except Exception as exc: raise ValueError(f"无法读取蒙版：{item['name']}") from exc try: return pil_to_tensor([image])[..., :3].mean(dim=3) finally: image.close() async def _generate_gpt_images( *, tasks: list[dict[str, Any]], model: str, model_route: str, resolution: str, aspect_ratio: str, quality: str, seed: int, references: list[Any], model_references: list[Any], mask, resize_mode: str, output_format: str, background: str, ): if mask is not None and not references: raise ValueError("提供了蒙版但未提供参考图，请同时上传参考图") client = GptImageClient() client.base_url = get_base_url_by_route() client.response_log_enabled = False client.poll_log_enabled = False actual_model = resolve_gpt_image_model(model, model_route) quality_value = resolve_gpt_image_quality(model, quality) request_semaphore = asyncio.Semaphore(max(UNIFIED_IMAGE_COUNTS)) async def _single_request(request_index: int, task: dict[str, Any]): task_references = [ references[index] for index in task["reference_indices"] ] + [ model_references[index] for index in task["model_reference_indices"] ] async with request_semaphore: images = await client.generate_image_async( prompt=task["prompt"], model=actual_model, quality=quality_value, size=( None if resolution == UNIFIED_IMAGE_SMART_RESOLUTION else resolve_gpt_image_size(resolution, aspect_ratio) ), n=1, seed=seed, image_tensor=task_references or None, mask_tensor=mask, output_format=output_format, background=background, resize_mode=resize_mode, special_price_parallel=False, log_downloads=True, log_request_start=False, log_prefix=f"[o1key 图片生成 GPT #{request_index + 1}]", ) if not images: raise RuntimeError(f"第 {request_index + 1} 次请求没有返回图片") selected = images[0] for extra in images[1:]: extra.close() return selected, task_references results = await asyncio.gather( *( _single_request(index, task) for index, task in enumerate(tasks) ), return_exceptions=True, ) images_with_references = [] errors = [] for result in results: if isinstance(result, BaseException): errors.append(str(result)) continue if result: images_with_references.append(result) if not images_with_references: raise RuntimeError(errors[0] if errors else "生成完成但没有可用图片") return images_with_references async def _generate_seedream_images( *, tasks: list[dict[str, Any]], model: str, model_route: str, resolution: str, aspect_ratio: str, output_format: str, references: list[Any], model_references: list[Any], layer_decomposition: bool = False, ): request_semaphore = asyncio.Semaphore(max(UNIFIED_IMAGE_COUNTS)) upload_cache: dict[int, Any] = {} reference_images: list[Image.Image] = [] model_reference_images: list[Image.Image] = [] results: list[Any] = [] def _convert_references(values: list[Any]) -> list[Image.Image]: converted: list[Image.Image] = [] for value in values: images = tensor_to_pil(value) if not images: raise ValueError("Seedream 参考图为空") converted.append(images[0]) for extra in images[1:]: extra.close() return converted try: reference_images = _convert_references(references) model_reference_images = _convert_references(model_references) validate_seedream_reference_images( [*reference_images, *model_reference_images], layer_decomposition=layer_decomposition, ) api_key = get_api_key_or_raise("O1KEY_API_KEY") base_url = get_base_url_by_route() actual_model = resolve_seedream_model(model, model_route) client = SeedreamImageClient(base_url=base_url, api_key=api_key) async with create_http_client( http2=True, max_connections=32, max_keepalive_connections=16, ) as session: async def _single_request(request_index: int, task: dict[str, Any]): task_images = [ reference_images[index] for index in task["reference_indices"] ] + [ model_reference_images[index] for index in task["model_reference_indices"] ] task_references = [ references[index] for index in task["reference_indices"] ] + [ model_references[index] for index in task["model_reference_indices"] ] async with request_semaphore: images, _timing = await client.generate_async( session=session, prompt=task["prompt"], model=actual_model, size=( None if resolution == UNIFIED_IMAGE_SMART_RESOLUTION else ( resolve_seedream_layer_size(resolution) if layer_decomposition else resolve_seedream_size(resolution, aspect_ratio) ) ), output_format=output_format, images=task_images, layer_decomposition=layer_decomposition, upload_cache=upload_cache, log_downloads=True, log_task_success=False, ) if not images: raise RuntimeError(f"第 {request_index + 1} 次请求没有返回图片") selected = images if layer_decomposition else images[:1] for extra in images[len(selected):]: extra.close() return [(image, task_references) for image in selected] results = await asyncio.gather( *( _single_request(index, task) for index, task in enumerate(tasks) ), return_exceptions=True, ) finally: for image in [*reference_images, *model_reference_images]: image.close() images_with_references = [] errors = [] for result in results: if isinstance(result, BaseException): errors.append(str(result)) continue if result: images_with_references.extend(result) if not images_with_references: raise RuntimeError(errors[0] if errors else "生成完成但没有可用图片") return images_with_references def _task_reference_values( task: dict[str, Any], references: list[Any], model_references: list[Any], ) -> list[Any]: return [references[index] for index in task["reference_indices"]] + [ model_references[index] for index in task["model_reference_indices"] ] async def _generate_nano_batch( *, tasks: list[dict[str, Any]], model: str, model_route: str, thinking_level: str, resolution: str, aspect_ratio: str, seed: int, references: list[Any], model_references: list[Any], resize_mode: str, google_search: bool, ) -> torch.Tensor: """Run task-specific Nano requests while retaining partial successes.""" request_semaphore = asyncio.Semaphore(max(UNIFIED_IMAGE_COUNTS)) async def _single_request(task: dict[str, Any]) -> torch.Tensor: task_references = _task_reference_values(task, references, model_references) reference_group = { f"参考图{index}": value for index, value in enumerate(task_references, start=1) } async with request_semaphore: result = await asyncio.to_thread( NanoBanana.execute, prompt=task["prompt"], 模型=model, 模型线路=model_route, 思考等级=thinking_level, 分辨率=resolution, 宽高比=aspect_ratio, 生图数量=1, seed=seed, 参考图组=reference_group, 缩放图片=resize_mode, _o1key_google_search=google_search, _o1key_unlimited_downloads=True, ) output = result[0] return output results = await asyncio.gather( *(_single_request(task) for task in tasks), return_exceptions=True, ) tensors = [result for result in results if not isinstance(result, BaseException)] if not tensors: error = next( (str(result) for result in results if isinstance(result, BaseException)), "生成完成但没有可用图片", ) raise RuntimeError(error) return torch.cat(tensors, dim=0) class O1keyImageGenerator(io.ComfyNode): """A full image-generation panel backed by multiple image providers.""" @classmethod def define_schema(cls): return io.Schema( node_id="O1keyImageGenerator", display_name="o1key 图片生成", category="o1key/image", description="节点内上传参考图、设置参数并生成图片。", inputs=[ io.String.Input( "prompt", display_name="提示词", default="", multiline=True, tooltip="多个提示词请用独占一行的 --- 分隔。总任务数 = 提示词数量 × 生图数量。", socketless=True, ), io.Combo.Input( "模型", options=UNIFIED_IMAGE_MODEL_OPTIONS, default="Nano Banana 2", socketless=True, ), io.Combo.Input( "模型线路", options=UNIFIED_IMAGE_ROUTE_OPTIONS, default="畅速", socketless=True, ), io.Combo.Input( "思考等级", options=["低", "高"], default="低", socketless=True, ), io.Combo.Input( "分辨率", options=_ALL_RESOLUTIONS, default=UNIFIED_IMAGE_SMART_RESOLUTION, socketless=True, ), io.Combo.Input( "宽高比", options=BANANA_ASPECT_RATIO_OPTIONS, default="智能", socketless=True, ), io.Combo.Input( "生图数量", options=_IMAGE_COUNTS, default="1", tooltip="选择本次生成的图像数量。", socketless=True, ), io.Int.Input( "seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, step=1, display_mode=io.NumberDisplay.number, socketless=True, ), io.String.Input( "参考图清单", default="[]", multiline=True, socketless=True, ), io.Combo.Input( "质量", options=GPT_IMAGE_25_QUALITY_OPTIONS, default="自动", socketless=True, ), io.Combo.Input( "输出格式", options=GPT_IMAGE_OUTPUT_FORMAT_OPTIONS, default="jpeg", socketless=True, ), io.String.Input( "蒙版清单", default="{}", multiline=True, socketless=True, ), io.Combo.Input( "缩放图片", options=["不缩放", "智能缩放"], default="不缩放", tooltip="Nano Banana 与 GPT Image 生效；智能缩放可能发生像素偏移。", socketless=True, ), io.Combo.Input( "背景", options=GPT_IMAGE_BACKGROUND_OPTIONS, default="auto", tooltip="仅 GPT Image 生效；透明背景仅支持 PNG 或 WebP。", socketless=True, ), io.Boolean.Input( "批量出图", default=False, label_on="开启", label_off="关闭", tooltip="关闭时保持普通参考图生成；开启后按素材图和目标图展开组合。", socketless=True, ), io.Combo.Input( "批量模式", options=list(IMAGE_BATCH_MODES), default=IMAGE_BATCH_MODE_GROUP_TO_MODELS, tooltip="整组素材：全部素材图作为一组并依次应用到多个目标；全匹配：每张素材图分别与每张目标图组合。", socketless=True, ), io.String.Input( "模特图清单", default="[]", multiline=True, socketless=True, ), io.Combo.Input( "命名规则", options=list(SAVE_NAMING_RULE_OPTIONS), default=DEFAULT_SAVE_NAMING_RULE, tooltip="自定义：使用文件名前缀；也可按主图文件名或自然数字命名。", socketless=True, ), io.String.Input( "filename_prefix", display_name="文件名前缀", default="o1key", socketless=True, ), io.Combo.Input( "格式", options=list(SAVE_FORMAT_OPTIONS), default="原始", tooltip="仅 Nano Banana 生效；GPT Image 和 Seedream 使用各自的 API 输出格式。", socketless=True, ), io.String.Input( "保存位置", default="", placeholder="留空为 output；也可填写 D:/图片", tooltip="留空保存到 ComfyUI output 根目录；相对路径是 output 内的子文件夹；绝对路径可保存到任意磁盘目录。", socketless=True, ), io.Combo.Input( "在线搜索", options=["关闭", "打开"], default="关闭", tooltip="仅 Nano Banana 2 生效；打开时发送 google_search=true。", socketless=True, ), io.Boolean.Input( "图层拆分", default=False, label_on="开启", label_off="关闭", tooltip="仅 Seedream 生效；必须且只能上传1张参考图，返回底图和最多16个透明图层。", socketless=True, ), ], outputs=[ io.Image.Output("IMAGE", display_name="IMAGE"), io.Image.Output("LAYERS", display_name="图层", is_output_list=True), io.Mask.Output("LAYER_MASKS", display_name="图层遮罩", is_output_list=True), io.String.Output("LAYER_INFO", display_name="图层信息"), ], not_idempotent=True, ) @classmethod async def execute( cls, prompt: str, 模型: str = "Nano Banana 2", 模型线路: str = "畅速", 思考等级: str = "低", 分辨率: str = UNIFIED_IMAGE_SMART_RESOLUTION, 宽高比: str = "智能", 生图数量: str | int = "1", seed: int = 0, 参考图清单: str = "[]", 质量: str = "自动", 输出格式: str | None = None, 蒙版清单: str = "{}", 缩放图片: str | None = None, 背景: str = "auto", 批量出图: bool = False, 批量模式: str = IMAGE_BATCH_MODE_GROUP_TO_MODELS, 模特图清单: str = "[]", 命名规则: str = DEFAULT_SAVE_NAMING_RULE, filename_prefix: str = "o1key", 格式: str = "原始", 保存位置: str = "", 在线搜索: str = "关闭", 图层拆分: bool = False, ) -> io.NodeOutput: layer_decomposition = _is_enabled(图层拆分) if (not prompt or not prompt.strip()) and not ( is_seedream_model(模型) and layer_decomposition ): raise ValueError("请输入提示词") if 输出格式 is None: 输出格式 = "png" if is_gpt_image_model(模型) else "jpeg" if 缩放图片 is None: 缩放图片 = "智能缩放" if is_gpt_image_model(模型) else "不缩放" save_settings = _normalized_save_settings( 模型, filename_prefix, 格式, 保存位置, 命名规则, ) if 缩放图片 not in {"不缩放", "智能缩放"}: raise ValueError("缩放图片参数无效") if 在线搜索 not in {"关闭", "打开"}: raise ValueError("在线搜索参数无效") google_search = 模型 == "Nano Banana 2" and 在线搜索 == "打开" image_count = int(生图数量) if is_gpt_image_model(模型): if image_count not in GPT_IMAGE_COUNTS and image_count != 9: raise ValueError("GPT Image 生图数量仅支持 1–8；旧工作流的 9 张仍可执行") elif image_count not in UNIFIED_IMAGE_COUNTS: raise ValueError("生图数量仅支持：1、2、4、9") batch_enabled = _is_enabled(批量出图) if 批量模式 not in IMAGE_BATCH_MODES: raise ValueError("批量模式无效") references = _load_reference_tensors( 参考图清单, max_images=( MAX_UNIFIED_BATCH_IMAGES if batch_enabled else MAX_UNIFIED_REFERENCE_IMAGES ), ) model_references = ( _load_reference_tensors( 模特图清单, label="目标图", max_images=MAX_UNIFIED_BATCH_IMAGES, ) if batch_enabled and 批量模式 != IMAGE_BATCH_MODE_SINGLE_REFERENCES else {} ) tasks = expand_image_generation_tasks( prompt, image_count, batch_enabled=batch_enabled, batch_mode=批量模式, reference_count=len(references), model_reference_count=len(model_references), ) if len(tasks) > MAX_UNIFIED_IMAGE_TASKS: raise ValueError(f"单次生成任务最多支持 {MAX_UNIFIED_IMAGE_TASKS} 个") if ( batch_enabled and 批量模式 == IMAGE_BATCH_MODE_GROUP_TO_MODELS and len(references) + 1 > MAX_UNIFIED_REFERENCE_IMAGES ): raise ValueError( f"整组素材模式每次请求最多支持 {MAX_UNIFIED_REFERENCE_IMAGES - 1} 张素材图和1张目标图" ) if batch_enabled and _parse_mask_manifest(蒙版清单) is not None: raise ValueError("批量出图暂不支持蒙版，请先移除蒙版") if is_seedream_model(模型): if layer_decomposition: if 分辨率 != UNIFIED_IMAGE_SMART_RESOLUTION: resolve_seedream_layer_size(分辨率) if batch_enabled: raise ValueError("Seedream 图层拆分不支持批量出图") if image_count != 1: raise ValueError("Seedream 图层拆分的生图数量必须为1") if len(references) != 1 or model_references: raise ValueError("Seedream 图层拆分必须且只能上传1张参考图") 输出格式 = "png" else: if 分辨率 not in { UNIFIED_IMAGE_SMART_RESOLUTION, *SEEDREAM_RESOLUTION_OPTIONS, }: raise ValueError("Seedream 分辨率无效") if 宽高比 not in SEEDREAM_ASPECT_RATIO_OPTIONS: raise ValueError("Seedream 宽高比无效") if 输出格式 not in SEEDREAM_OUTPUT_FORMAT_OPTIONS: raise ValueError("Seedream 输出格式仅支持 png 或 jpeg") if _parse_mask_manifest(蒙版清单) is not None: raise ValueError("Seedream 不支持蒙版") images_with_references: list[tuple[Image.Image, list[Any]]] = [] try: images_with_references = await _generate_seedream_images( tasks=tasks, model=模型, model_route=模型线路, resolution=分辨率, aspect_ratio=宽高比, output_format=输出格式, references=list(references.values()), model_references=list(model_references.values()), layer_decomposition=layer_decomposition, ) except Exception as exc: message = format_o1key_image_error(exc) if message == str(exc): raise raise RuntimeError(message) from None try: output_parts = [] for image, _task_references in images_with_references: output_part = pil_to_tensor([image]) output_parts.append(output_part) if layer_decomposition: base_tensor = output_parts[0] layer_tensors = output_parts[1:] layer_masks = [] layer_info = [] for image, _task_references in images_with_references[1:]: if "A" in image.getbands(): alpha = torch.from_numpy( np.asarray(image.getchannel("A"), dtype=np.float32) / 255.0 ).unsqueeze(0) else: alpha = torch.ones((1, image.height, image.width), dtype=torch.float32) layer_masks.append(alpha) metadata = getattr(image, "_o1key_seedream_layer", None) layer_info.append(metadata if isinstance(metadata, dict) else {}) _attach_save_settings(base_tensor, references, save_settings) for layer_tensor in layer_tensors: _attach_save_settings(layer_tensor, references, save_settings) return io.NodeOutput( base_tensor, layer_tensors, layer_masks, json.dumps(layer_info, ensure_ascii=False), ) output_tensor = torch.cat(output_parts, dim=0) source_metadata = [] for output_part in output_parts: metadata = getattr(output_part, "_o1key_source_metadata", None) if isinstance(metadata, list): source_metadata.extend(metadata) if len(source_metadata) == output_tensor.shape[0]: setattr(output_tensor, "_o1key_source_metadata", source_metadata) _attach_save_settings(output_tensor, references, save_settings) return io.NodeOutput(output_tensor, [], [], "[]") finally: for image, _task_references in images_with_references: image.close() if is_gpt_image_model(模型): if 分辨率 not in { UNIFIED_IMAGE_SMART_RESOLUTION, *GPT_IMAGE_RESOLUTION_OPTIONS, *GPT_IMAGE_EXACT_SIZE_OPTIONS, }: raise ValueError("GPT Image 分辨率无效") if 宽高比 not in GPT_IMAGE_ASPECT_RATIO_OPTIONS: raise ValueError("GPT Image 宽高比无效") resolve_gpt_image_quality(模型, 质量) if 输出格式 not in GPT_IMAGE_OUTPUT_FORMAT_OPTIONS: raise ValueError("GPT Image 输出格式无效") if 背景 not in GPT_IMAGE_BACKGROUND_OPTIONS: raise ValueError("GPT Image 背景参数无效") if 背景 == "transparent" and 输出格式 == "jpeg": raise ValueError("GPT Image 透明背景仅支持 PNG 或 WebP 输出格式") images_with_references: list[tuple[Image.Image, list[Any]]] = [] try: images_with_references = await _generate_gpt_images( tasks=tasks, model=模型, model_route=模型线路, resolution=分辨率, aspect_ratio=宽高比, quality=质量, seed=int(seed), references=list(references.values()), model_references=list(model_references.values()), mask=_load_mask_tensor(蒙版清单), resize_mode=缩放图片, output_format=输出格式, background=背景, ) except Exception as exc: message = format_o1key_image_error(exc) if message == str(exc): raise raise RuntimeError(message) from None try: output_parts = [] for image, _task_references in images_with_references: output_part = GptImageClient._pil_list_to_tensor([image]) output_parts.append(output_part) output_tensor = torch.cat(output_parts, dim=0) source_metadata = [] for output_part in output_parts: metadata = getattr(output_part, "_o1key_source_metadata", None) if isinstance(metadata, list): source_metadata.extend(metadata) if len(source_metadata) == output_tensor.shape[0]: setattr(output_tensor, "_o1key_source_metadata", source_metadata) _attach_save_settings(output_tensor, references, save_settings) return io.NodeOutput(output_tensor, [], [], "[]") finally: for image, _task_references in images_with_references: image.close() if batch_enabled: try: output_tensor = await _generate_nano_batch( tasks=tasks, model=模型, model_route=模型线路, thinking_level=思考等级, resolution=分辨率, aspect_ratio=宽高比, seed=int(seed), references=list(references.values()), model_references=list(model_references.values()), resize_mode=缩放图片, google_search=google_search, ) except Exception as exc: message = format_o1key_image_error(exc) if message == str(exc): raise raise RuntimeError(message) from None _attach_save_settings(output_tensor, references, save_settings) return io.NodeOutput(output_tensor, [], [], "[]") try: result = await asyncio.to_thread( NanoBanana.execute, prompt=prompt, 模型=模型, 模型线路=模型线路, 思考等级=思考等级, 分辨率=分辨率, 宽高比=宽高比, 生图数量=image_count, seed=int(seed), 参考图组=references, 缩放图片=缩放图片, _o1key_google_search=google_search, _o1key_unlimited_downloads=True, ) except Exception as exc: message = format_o1key_image_error(exc) if message == str(exc): raise raise RuntimeError(message) from None output_tensor = result[0] _attach_save_settings(output_tensor, references, save_settings) return io.NodeOutput(output_tensor, [], [], "[]") class O1keyImageSave(io.ComfyNode): """Save and present images generated by the o1key panel node.""" @classmethod def define_schema(cls): return io.Schema( node_id="O1keyImageSave", display_name="o1key 保存图像", category="o1key/image", description="保存 o1key 图片生成节点的结果，并支持从结果图重新生成。", is_output_node=True, inputs=[ io.Image.Input("images", display_name="图像"), ], hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], outputs=[io.Image.Output("IMAGE", display_name="图像")], ) @classmethod def execute( cls, images, ) -> io.NodeOutput: hidden = getattr(cls, "hidden", None) settings = getattr(images, "_o1key_save_settings", None) if not isinstance(settings, dict): settings = {} descriptors = save_tensor_images( images, str(settings.get("filename_prefix") or "o1key"), settings.get("format", "原始"), folder_paths.get_output_directory(), folder_paths, prompt=getattr(hidden, "prompt", None), extra_pnginfo=getattr(hidden, "extra_pnginfo", None), save_location=settings.get("save_location", ""), naming_rule=settings.get("naming_rule", DEFAULT_SAVE_NAMING_RULE), ) saved = ui.SavedImages([ ui.SavedResult( item["filename"], item["subfolder"], ( io.FolderType.temp if item.get("type") == "temp" else io.FolderType.output ), ) for item in descriptors ]) return io.NodeOutput(images, ui=saved) __all__ = [ "O1keyImageGenerator", "O1keyImageSave", "_load_reference_tensors", "_load_mask_tensor", "_parse_mask_manifest", "_parse_upload_manifest", ]