Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
1035 lines
40 KiB
Python
1035 lines
40 KiB
Python
"""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",
|
||
]
|