Files
comfyui_o1key/utils/o1key_image_jobs.py
Jony ba920f2b66 Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
2026-09-24 19:56:48 +08:00

2236 lines
92 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Independent concurrent job scheduler for the o1key image creation panel."""
from __future__ import annotations
import asyncio
from collections import deque
from dataclasses import dataclass, field
import json
import os
import re
import shutil
import threading
import time
import uuid
from typing import Any, Awaitable, Callable, Optional
from PIL import Image, ImageOps
from .o1key_image_catalog import (
BANANA_ASPECT_RATIO_OPTIONS,
BANANA_RESOLUTION_OPTIONS,
GPT_IMAGE_ASPECT_RATIO_OPTIONS,
GPT_IMAGE_BACKGROUND_OPTIONS,
GPT_IMAGE_EXACT_SIZE_OPTIONS,
GPT_IMAGE_OUTPUT_FORMAT_OPTIONS,
GPT_IMAGE_COUNTS,
GPT_IMAGE_RESOLUTION_OPTIONS,
SEEDREAM_ASPECT_RATIO_OPTIONS,
SEEDREAM_OUTPUT_FORMAT_OPTIONS,
SEEDREAM_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,
is_gpt_image_model,
resolve_gpt_image_quality,
is_seedream_model,
resolve_gpt_image_size,
resolve_seedream_layer_size,
resolve_seedream_size,
)
from .http_error import format_o1key_image_error
from .image_utils import (
IMAGE_BATCH_MODE_GROUP_TO_MODELS,
IMAGE_BATCH_MODE_SINGLE_REFERENCES,
IMAGE_BATCH_MODES,
expand_image_generation_tasks,
)
from .o1key_image_save import (
DEFAULT_SAVE_NAMING_RULE,
detect_image_format,
normalize_naming_rule,
normalize_save_format,
normalize_save_location,
save_temp_images,
)
MAX_CONCURRENT_BATCHES = 10
MAX_REFERENCE_IMAGES = MAX_UNIFIED_REFERENCE_IMAGES
MAX_BATCH_IMAGES = MAX_UNIFIED_BATCH_IMAGES
MAX_RETAINED_JOBS = 200
JOB_HISTORY_FILENAME = "image_job_history.json"
IMAGE_COUNTS = set(UNIFIED_IMAGE_COUNTS)
MAX_CONCURRENT_REQUESTS_PER_JOB = max(UNIFIED_IMAGE_COUNTS)
OUTPUT_FILENAME_PATTERN = re.compile(
r"^image_(\d+)_(\d+)\.(png|jpe?g|webp)$",
re.IGNORECASE,
)
RESOLUTIONS = {UNIFIED_IMAGE_SMART_RESOLUTION, *BANANA_RESOLUTION_OPTIONS}
ASPECT_RATIOS = set(BANANA_ASPECT_RATIO_OPTIONS)
THINKING_LEVELS = {"低": "minimal", "高": "high"}
MODELS = set(UNIFIED_IMAGE_MODEL_OPTIONS)
MODEL_ROUTES = set(UNIFIED_IMAGE_ROUTE_OPTIONS)
EventSender = Callable[[str, dict[str, Any]], Awaitable[None]]
ProgressCallback = Callable[[float], None]
JobExecutor = Callable[
[dict[str, Any], ProgressCallback],
Awaitable[dict[str, Any]],
]
def _publish_partial_images(
progress_callback: Optional[ProgressCallback],
images: list[dict[str, Any]],
) -> None:
"""Publish fully written temp descriptors when the scheduler supports it."""
publish = getattr(progress_callback, "publish_images", None)
if callable(publish) and images:
publish([dict(item) for item in images])
def _safe_layer_metadata(value: Any) -> dict[str, Any] | None:
"""Copy the bounded Seedream metadata subset; URLs are never accepted."""
if not isinstance(value, dict):
return None
try:
z_index = max(0, min(16, int(value.get("z_index", 0))))
except (TypeError, ValueError):
z_index = 0
result: dict[str, Any] = {"z_index": z_index}
for key, limit in (("name", 200), ("description", 1000), ("size", 64), ("output_format", 16)):
item = value.get(key)
if isinstance(item, str) and item.strip():
result[key] = item.strip()[:limit]
box = value.get("bounding_box")
if isinstance(box, dict):
safe_box = {}
for key in ("absolute", "normalized"):
item = box.get(key)
if not isinstance(item, (list, tuple)) or len(item) != 4:
continue
try:
safe_box[key] = [int(number) for number in item]
except (TypeError, ValueError):
continue
if safe_box:
result["bounding_box"] = safe_box
return result
def _canonical_batch_id(value: Any) -> str:
try:
parsed = uuid.UUID(str(value))
except (TypeError, ValueError, AttributeError):
raise ValueError("批次 ID 无效") from None
canonical = str(parsed)
if canonical != str(value).strip().lower():
raise ValueError("批次 ID 必须使用标准 UUID 格式")
return canonical
def _history_image_descriptor(value: Any, batch_id: str) -> dict[str, Any] | None:
if not isinstance(value, dict):
return None
filename = str(value.get("filename") or "")[:1024]
folder_type = str(value.get("type") or "")
if not filename or folder_type not in {"output", "temp"}:
return None
descriptor: dict[str, Any] = {
"filename": filename,
"subfolder": str(value.get("subfolder") or "")[:1024],
"type": folder_type,
"batch_id": batch_id,
}
if value.get("external_saved") is True:
descriptor["external_saved"] = True
for key in ("request_index", "result_index"):
try:
descriptor[key] = max(1, int(value.get(key) or 1))
except (TypeError, ValueError):
descriptor[key] = 1
layer = _safe_layer_metadata(value.get("layer"))
if layer is not None:
descriptor["layer"] = layer
return descriptor
def _history_item(value: Any) -> dict[str, Any] | None:
"""Return the bounded, credential-free subset used by the task-history UI."""
if not isinstance(value, dict):
return None
try:
batch_id = _canonical_batch_id(value.get("batch_id"))
generator_node_id = _safe_node_id(value.get("generator_node_id"), "生成节点 ID")
save_node_id = _safe_node_id(value.get("save_node_id"), "保存节点 ID")
except ValueError:
return None
state = str(value.get("state") or "")
if state not in {"completed", "failed", "cancelled"}:
return None
raw_images = value.get("images")
if not isinstance(raw_images, list):
raw_images = []
images = [
descriptor
for descriptor in (
_history_image_descriptor(item, batch_id)
for item in raw_images[:MAX_UNIFIED_IMAGE_TASKS]
)
if descriptor is not None
]
raw_failed_items = value.get("failed_items")
if not isinstance(raw_failed_items, list):
raw_failed_items = []
failed_items = []
for item in raw_failed_items[:MAX_UNIFIED_IMAGE_TASKS]:
if not isinstance(item, dict):
continue
try:
request_index = max(1, int(item.get("request_index") or 1))
except (TypeError, ValueError):
request_index = 1
failed_items.append({
"request_index": request_index,
"error": str(item.get("error") or "")[:600],
})
raw_warnings = value.get("warnings")
if not isinstance(raw_warnings, list):
raw_warnings = []
warnings = [str(item)[:600] for item in raw_warnings[:50]]
def _timestamp(name: str, fallback: float = 0) -> int:
try:
return max(0, int(float(value.get(name) or fallback)))
except (TypeError, ValueError):
return max(0, int(fallback))
create_time = _timestamp("create_time")
execution_end_time = _timestamp("execution_end_time", create_time)
try:
total_count = max(
0,
min(MAX_UNIFIED_IMAGE_TASKS, int(value.get("total_count") or 0)),
)
except (TypeError, ValueError):
total_count = 0
return {
"batch_id": batch_id,
"state": state,
"generator_node_id": generator_node_id,
"save_node_id": save_node_id,
"images": images,
"warnings": warnings,
"failed_items": failed_items,
"error": str(value.get("error") or "")[:600],
"total_count": total_count or len(images) + len(failed_items),
"succeeded_count": len(images),
"failed_count": len(failed_items),
"progress": 1.0,
"progress_percent": 100,
"create_time": create_time,
"execution_start_time": _timestamp("execution_start_time"),
"execution_end_time": execution_end_time,
}
class PersistentJobHistory:
"""Atomic, bounded history index containing only safe terminal summaries."""
def __init__(self, path: str, max_items: int = MAX_RETAINED_JOBS):
self.path = os.path.abspath(path)
self.max_items = max(1, int(max_items))
self._lock = threading.Lock()
def _read_unlocked(self) -> list[dict[str, Any]]:
try:
with open(self.path, "r", encoding="utf-8") as handle:
payload = json.load(handle)
except (OSError, ValueError, TypeError):
return []
raw_items = payload.get("items") if isinstance(payload, dict) else None
if not isinstance(raw_items, list):
return []
items = [item for item in (_history_item(value) for value in raw_items) if item]
items.sort(key=lambda item: item["execution_end_time"], reverse=True)
return items[: self.max_items]
def _write_unlocked(self, items: list[dict[str, Any]]) -> None:
directory = os.path.dirname(self.path)
os.makedirs(directory, exist_ok=True)
temporary = f"{self.path}.{uuid.uuid4().hex}.tmp"
try:
with open(temporary, "w", encoding="utf-8", newline="\n") as handle:
json.dump(
{"version": 1, "items": items[: self.max_items]},
handle,
ensure_ascii=False,
separators=(",", ":"),
)
os.replace(temporary, self.path)
finally:
if os.path.exists(temporary):
try:
os.remove(temporary)
except OSError:
pass
def list(self, limit: int = MAX_RETAINED_JOBS) -> list[dict[str, Any]]:
with self._lock:
return [dict(item) for item in self._read_unlocked()[: max(0, int(limit))]]
def upsert(self, value: Any) -> None:
item = _history_item(value)
if item is None:
return
with self._lock:
items = [
existing for existing in self._read_unlocked()
if existing["batch_id"] != item["batch_id"]
]
items.append(item)
items.sort(key=lambda existing: existing["execution_end_time"], reverse=True)
self._write_unlocked(items)
def delete(self, batch_id: str) -> None:
canonical = _canonical_batch_id(batch_id)
with self._lock:
items = [
item for item in self._read_unlocked()
if item["batch_id"] != canonical
]
self._write_unlocked(items)
def clear(self) -> None:
with self._lock:
self._write_unlocked([])
def _safe_node_id(value: Any, field_name: str) -> int:
try:
result = int(value)
except (TypeError, ValueError):
raise ValueError(f"{field_name} 无效") from None
if result < 0:
raise ValueError(f"{field_name} 无效")
return result
def normalize_job_payload(payload: Any) -> dict[str, Any]:
"""Validate and copy all user-controlled fields into an immutable job snapshot."""
if not isinstance(payload, dict):
raise ValueError("请求体必须是对象")
model = str(payload.get("model") or "").strip()
layer_decomposition = (
is_seedream_model(model) and payload.get("layer_decomposition") is True
)
prompt = str(payload.get("prompt") or "").strip()
if not prompt and not layer_decomposition:
raise ValueError("请输入提示词")
model_route = str(payload.get("model_route") or "").strip()
resolution = str(
payload.get("resolution") or UNIFIED_IMAGE_SMART_RESOLUTION
).strip()
aspect_ratio = str(payload.get("aspect_ratio") or "智能").strip()
thinking = str(payload.get("thinking_level") or "低").strip()
if model not in MODELS:
raise ValueError("模型无效")
if model_route not in MODEL_ROUTES:
raise ValueError("模型线路无效")
google_search = model == "Nano Banana 2" and payload.get("google_search") is True
quality = str(payload.get("quality") or "自动").strip()
resize_mode = str(
payload.get("resize_mode") or ("智能缩放" if is_gpt_image_model(model) else "不缩放")
).strip()
if resize_mode not in {"不缩放", "智能缩放"}:
raise ValueError("缩放图片参数无效")
output_format = None
background = None
if is_gpt_image_model(model):
output_format = str(payload.get("output_format") or "png").strip().lower()
background = str(payload.get("background") or "auto").strip().lower()
if resolution not in {
UNIFIED_IMAGE_SMART_RESOLUTION,
*GPT_IMAGE_RESOLUTION_OPTIONS,
*GPT_IMAGE_EXACT_SIZE_OPTIONS,
}:
raise ValueError("GPT Image 分辨率无效")
if aspect_ratio not in GPT_IMAGE_ASPECT_RATIO_OPTIONS:
raise ValueError("GPT Image 宽高比无效")
resolve_gpt_image_quality(model, quality)
if output_format not in GPT_IMAGE_OUTPUT_FORMAT_OPTIONS:
raise ValueError("GPT Image 输出格式无效")
if background not in GPT_IMAGE_BACKGROUND_OPTIONS:
raise ValueError("GPT Image 背景参数无效")
if background == "transparent" and output_format == "jpeg":
raise ValueError("GPT Image 透明背景仅支持 PNG 或 WebP 输出格式")
elif is_seedream_model(model):
output_format = str(
payload.get("output_format") or ("png" if layer_decomposition else "jpeg")
).strip().lower()
if layer_decomposition:
if resolution != UNIFIED_IMAGE_SMART_RESOLUTION:
resolve_seedream_layer_size(resolution)
if output_format != "png":
raise ValueError("Seedream 图层拆分仅支持 png 输出格式")
else:
if resolution not in {
UNIFIED_IMAGE_SMART_RESOLUTION,
*SEEDREAM_RESOLUTION_OPTIONS,
}:
raise ValueError("Seedream 分辨率无效")
if aspect_ratio not in SEEDREAM_ASPECT_RATIO_OPTIONS:
raise ValueError("Seedream 宽高比无效")
if output_format not in SEEDREAM_OUTPUT_FORMAT_OPTIONS:
raise ValueError("Seedream 输出格式仅支持 png 或 jpeg")
else:
if resolution not in RESOLUTIONS:
raise ValueError("分辨率无效")
if aspect_ratio not in ASPECT_RATIOS:
raise ValueError("宽高比无效")
if thinking not in THINKING_LEVELS:
raise ValueError("思考等级无效")
if model not in {"Nano Banana 2", "Nano Banana 2 Lite"} and (
resolution == "512" or aspect_ratio in {"1:4", "1:8", "4:1", "8:1"}
):
raise ValueError("当前分辨率或宽高比仅支持 Nano Banana 2 系列")
if model == "Nano Banana" and resolution not in {
UNIFIED_IMAGE_SMART_RESOLUTION,
"1K",
}:
raise ValueError("Nano Banana 模型仅支持 1K 分辨率")
naming_rule = normalize_naming_rule(payload.get("naming_rule"))
filename_prefix = str(payload.get("filename_prefix") or "o1key").strip()
if naming_rule == "自定义前缀" and (
not filename_prefix or len(filename_prefix) > 512
):
raise ValueError("文件名前缀无效")
save_format = (
"原始"
if is_gpt_image_model(model) or is_seedream_model(model)
else normalize_save_format(payload.get("save_format"))
)
save_location = normalize_save_location(payload.get("save_location"))
try:
image_count = int(payload.get("image_count", 1))
except (TypeError, ValueError):
raise ValueError("生图数量无效") from None
if is_gpt_image_model(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 IMAGE_COUNTS:
raise ValueError("生图数量仅支持:1、2、4、9")
batch_enabled = payload.get("batch_enabled") is True or str(
payload.get("batch_enabled") or ""
).strip().lower() in {"1", "true", "开启"}
batch_mode = str(
payload.get("batch_mode") or IMAGE_BATCH_MODE_GROUP_TO_MODELS
).strip()
if batch_enabled and batch_mode not in IMAGE_BATCH_MODES:
raise ValueError("批量模式无效")
if not batch_enabled:
batch_mode = IMAGE_BATCH_MODE_GROUP_TO_MODELS
if layer_decomposition and batch_enabled:
raise ValueError("Seedream 图层拆分不支持批量出图")
if layer_decomposition and image_count != 1:
raise ValueError("Seedream 图层拆分的生图数量必须为1")
references = payload.get("references") or []
if not isinstance(references, list):
raise ValueError("参考图清单必须是数组")
reference_limit = MAX_BATCH_IMAGES if batch_enabled else MAX_REFERENCE_IMAGES
if len(references) > reference_limit:
raise ValueError(f"参考图最多支持 {reference_limit} 张")
safe_references = []
for item in references:
if not isinstance(item, dict):
raise ValueError("参考图清单包含无效项目")
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 or folder_type != "input":
raise ValueError("参考图必须来自 ComfyUI input 目录")
safe_references.append({"name": name, "subfolder": subfolder, "type": "input"})
if layer_decomposition and len(safe_references) != 1:
raise ValueError("Seedream 图层拆分必须且只能上传1张参考图")
model_references = payload.get("model_references") or []
if not isinstance(model_references, list):
raise ValueError("目标图清单必须是数组")
safe_model_references = []
if batch_enabled and batch_mode != IMAGE_BATCH_MODE_SINGLE_REFERENCES:
if len(model_references) > MAX_BATCH_IMAGES:
raise ValueError(f"目标图最多支持 {MAX_BATCH_IMAGES} 张")
for item in model_references:
if not isinstance(item, dict):
raise ValueError("目标图清单包含无效项目")
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 or folder_type != "input":
raise ValueError("目标图必须来自 ComfyUI input 目录")
safe_model_references.append({
"name": name,
"subfolder": subfolder,
"type": "input",
})
tasks = expand_image_generation_tasks(
prompt,
image_count,
batch_enabled=batch_enabled,
batch_mode=batch_mode,
reference_count=len(safe_references),
model_reference_count=len(safe_model_references),
)
if len(tasks) > MAX_UNIFIED_IMAGE_TASKS:
raise ValueError(f"单次生成任务最多支持 {MAX_UNIFIED_IMAGE_TASKS} 个")
if (
batch_enabled
and batch_mode == IMAGE_BATCH_MODE_GROUP_TO_MODELS
and len(safe_references) + 1 > MAX_REFERENCE_IMAGES
):
raise ValueError(
f"整组素材模式每次请求最多支持 {MAX_REFERENCE_IMAGES - 1} 张素材图和1张目标图"
)
mask = payload.get("mask") or None
safe_mask = None
if mask is not None:
if batch_enabled:
raise ValueError("批量出图暂不支持蒙版,请先移除蒙版")
if not is_gpt_image_model(model):
raise ValueError("当前模型不支持蒙版")
if not isinstance(mask, dict):
raise ValueError("蒙版清单必须是对象")
name = str(mask.get("name") or "").strip()
subfolder = str(mask.get("subfolder") or "").strip()
folder_type = str(mask.get("type") or "input").strip()
if not name or folder_type != "input":
raise ValueError("蒙版必须来自 ComfyUI input 目录")
if not safe_references:
raise ValueError("提供了蒙版但未提供参考图")
safe_mask = {"name": name, "subfolder": subfolder, "type": "input"}
try:
seed = max(0, int(payload.get("seed", 0)))
except (TypeError, ValueError):
seed = 0
job = {
"batch_id": _canonical_batch_id(payload.get("batch_id")),
"generator_node_id": _safe_node_id(payload.get("generator_node_id"), "生成节点 ID"),
"save_node_id": _safe_node_id(payload.get("save_node_id"), "保存节点 ID"),
"prompt": prompt,
"model": model,
"model_route": model_route,
"thinking_level": thinking,
"resolution": resolution,
"aspect_ratio": aspect_ratio,
"image_count": image_count,
"task_prompts": [task["prompt"] for task in tasks],
"tasks": tasks,
"total_task_count": len(tasks),
"seed": seed,
"references": safe_references,
"batch_enabled": batch_enabled,
"batch_mode": batch_mode,
"model_references": safe_model_references,
"quality": quality,
"resize_mode": resize_mode,
"filename_prefix": filename_prefix or "o1key",
"save_format": save_format,
"save_location": save_location,
"naming_rule": naming_rule,
"mask": safe_mask,
"submitted_at": time.time(),
"layer_decomposition": layer_decomposition,
}
if is_gpt_image_model(model):
job["output_format"] = output_format
job["background"] = background
elif is_seedream_model(model):
job["output_format"] = output_format
elif google_search:
job["google_search"] = True
return job
def _is_within(root: str, candidate: str) -> bool:
try:
return os.path.commonpath([root, candidate]) == root
except ValueError:
return False
def snapshot_reference_files(
job: dict[str, Any],
input_directory: str,
temp_directory: str,
) -> str:
"""Copy references into a batch-owned directory and preserve their exact order."""
input_root = os.path.abspath(input_directory)
jobs_root = os.path.abspath(os.path.join(temp_directory, "o1key_image_jobs"))
batch_root = os.path.abspath(os.path.join(jobs_root, job["batch_id"]))
reference_root = os.path.join(batch_root, "references")
model_reference_root = os.path.join(batch_root, "models")
if not _is_within(jobs_root, batch_root):
raise ValueError("批次临时目录无效")
if is_seedream_model(job.get("model")):
from ..clients.seedream_image_client import validate_seedream_reference_image
sources: list[tuple[str, str]] = []
for index, item in enumerate(job["references"], start=1):
sources.append((
os.path.abspath(os.path.join(input_root, item.get("subfolder", ""), item["name"])),
f"Seedream 参考图{index}",
))
reference_offset = len(sources)
for index, item in enumerate(job.get("model_references", ()), start=1):
sources.append((
os.path.abspath(os.path.join(input_root, item.get("subfolder", ""), item["name"])),
f"Seedream 参考图{reference_offset + index}",
))
for source, label in sources:
if not _is_within(input_root, source) or not os.path.isfile(source):
raise ValueError(f"{label}不存在或路径不安全")
image = _load_reference_image(source)
try:
validate_seedream_reference_image(
image,
label=label,
layer_decomposition=bool(job.get("layer_decomposition")),
)
finally:
image.close()
os.makedirs(jobs_root, exist_ok=True)
os.makedirs(reference_root, exist_ok=False)
if job.get("model_references"):
os.makedirs(model_reference_root, exist_ok=False)
snapshot_paths: list[str] = []
model_snapshot_paths: list[str] = []
try:
for index, item in enumerate(job["references"], start=1):
source = os.path.abspath(
os.path.join(input_root, item.get("subfolder", ""), item["name"])
)
if not _is_within(input_root, source) or not os.path.isfile(source):
raise ValueError(f"参考图不存在或路径不安全:{item['name']}")
extension = os.path.splitext(source)[1].lower()
if extension not in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}:
extension = ".img"
target = os.path.join(reference_root, f"reference_{index:02d}{extension}")
shutil.copy2(source, target)
snapshot_paths.append(target)
for index, item in enumerate(job.get("model_references", ()), start=1):
source = os.path.abspath(
os.path.join(input_root, item.get("subfolder", ""), item["name"])
)
if not _is_within(input_root, source) or not os.path.isfile(source):
raise ValueError(f"目标图不存在或路径不安全:{item['name']}")
extension = os.path.splitext(source)[1].lower()
if extension not in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}:
extension = ".img"
target = os.path.join(model_reference_root, f"model_{index:02d}{extension}")
shutil.copy2(source, target)
model_snapshot_paths.append(target)
mask_path = None
if job.get("mask"):
item = job["mask"]
source = os.path.abspath(
os.path.join(input_root, item.get("subfolder", ""), item["name"])
)
if not _is_within(input_root, source) or not os.path.isfile(source):
raise ValueError(f"蒙版不存在或路径不安全:{item['name']}")
extension = os.path.splitext(source)[1].lower()
if extension not in {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}:
extension = ".img"
mask_path = os.path.join(batch_root, f"mask{extension}")
shutil.copy2(source, mask_path)
except BaseException:
shutil.rmtree(batch_root, ignore_errors=True)
raise
job["reference_paths"] = tuple(snapshot_paths)
job["model_reference_paths"] = tuple(model_snapshot_paths)
job["mask_path"] = mask_path
job["snapshot_root"] = batch_root
return batch_root
def cleanup_job_snapshot(job: dict[str, Any], temp_directory: str) -> None:
jobs_root = os.path.abspath(os.path.join(temp_directory, "o1key_image_jobs"))
snapshot_root = os.path.abspath(str(job.get("snapshot_root") or ""))
if snapshot_root and snapshot_root != jobs_root and _is_within(jobs_root, snapshot_root):
shutil.rmtree(snapshot_root, ignore_errors=True)
def _load_reference_image(path: str) -> Image.Image:
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", os.path.basename(path))
if orientation in (None, 1):
setattr(image, "_o1key_original_path", path)
return image
async def _load_reference_images(paths) -> list[Image.Image]:
"""Load a bounded path group and close partial results on failure."""
images: list[Image.Image] = []
try:
for path in paths:
images.append(await asyncio.to_thread(_load_reference_image, path))
return images
except BaseException:
for image in images:
image.close()
raise
def _task_reference_paths(job: dict[str, Any], request_index: int) -> tuple[str, ...]:
tasks = job.get("tasks") or ()
if request_index >= len(tasks):
return tuple(job.get("reference_paths", ()))
task = tasks[request_index]
reference_paths = tuple(job.get("reference_paths", ()))
model_paths = tuple(job.get("model_reference_paths", ()))
return tuple(
reference_paths[index]
for index in task.get("reference_indices", ())
) + tuple(
model_paths[index]
for index in task.get("model_reference_indices", ())
)
def _load_mask_tensor(path: str | None):
if not path:
return None
from .image_utils import pil_to_tensor
with Image.open(path) as opened:
image = ImageOps.exif_transpose(opened).convert("RGB").copy()
try:
return pil_to_tensor([image])[..., :3].mean(dim=3)
finally:
image.close()
def _prepare_output_directory(job: dict[str, Any], output_directory: str) -> tuple[str, str]:
"""Use ComfyUI's temp root for provider results awaiting the save node."""
del job
output_root = os.path.abspath(output_directory)
if not os.path.isdir(output_root):
raise ValueError("ComfyUI 临时目录不存在")
return output_root, ""
def _result_filename(
job: dict[str, Any],
request_index: int,
result_index: int,
extension: str = "png",
) -> str:
"""Build a collision-proof filename for a result in the shared temp root."""
normalized_extension = "jpg" if extension == "jpeg" else extension
return (
f"o1key_{job['batch_id']}_image_"
f"{request_index + 1:02d}_{result_index + 1:02d}.{normalized_extension}"
)
def _provider_result_format(image: Image.Image) -> str:
raw = getattr(image, "_o1key_original_bytes", None)
detected = detect_image_format(raw) if isinstance(raw, bytes) else None
if detected == "JPEG":
return "jpg"
if detected in {"PNG", "WEBP"}:
return detected.lower()
return "png"
def _write_provider_result(image: Image.Image, output_path: str) -> None:
"""Write provider bytes unchanged; modified/unknown pixels fall back to lossless PNG."""
raw = getattr(image, "_o1key_original_bytes", None)
detected = detect_image_format(raw) if isinstance(raw, bytes) else None
temporary_path = f"{output_path}.tmp"
try:
if detected in {"PNG", "JPEG", "WEBP"}:
with open(temporary_path, "wb") as handle:
handle.write(raw)
else:
image.save(temporary_path, format="PNG", compress_level=4)
os.replace(temporary_path, output_path)
finally:
if os.path.exists(temporary_path):
try:
os.remove(temporary_path)
except OSError:
pass
async def _execute_nano_banana_job(
job: dict[str, Any],
output_directory: str,
progress_callback: Optional[ProgressCallback] = None,
) -> dict[str, Any]:
"""Generate and atomically save one Nano Banana batch."""
from .config import get_api_key_or_raise, get_base_url_by_route
from .http2_client import create_http_client
from .nano_banana_async import (
VERBOSE_LOG_ENABLED,
generate_nano_banana_async,
prepare_nano_banana_inline_images,
)
from .nano_banana_models import resolve_nano_banana_model
actual_model = resolve_nano_banana_model(job["model"], job["model_route"])
batch_tag = f"[o1key {job['batch_id'][:8]}]"
batch_started_at = time.perf_counter()
thinking_level = (
THINKING_LEVELS[job["thinking_level"]]
if job["model"] == "Nano Banana 2"
else None
)
api_key = get_api_key_or_raise("O1KEY_API_KEY")
base_url = get_base_url_by_route()
output_path, relative_subfolder = await asyncio.to_thread(
_prepare_output_directory,
job,
output_directory,
)
print(
f"{batch_tag} 开始 | {job['model']} | {job['resolution']} {job['aspect_ratio']} "
f"| 提示词={len(job['task_prompts']) // job['image_count']} "
f"| 每条={job['image_count']} | 总任务={job['total_task_count']} "
f"| 参考图={len(job.get('reference_paths', ()))}"
)
batch_enabled = bool(job.get("batch_enabled"))
group_batch = (
batch_enabled
and job.get("batch_mode") == IMAGE_BATCH_MODE_GROUP_TO_MODELS
)
reference_images = await _load_reference_images(
job.get("reference_paths", ()) if not batch_enabled or group_batch else ()
)
model_reference_images = await _load_reference_images(
job.get("model_reference_paths", ()) if not batch_enabled else ()
)
reference_inline_images: list[dict[str, Any]] = []
request_progress = [0.0] * job["total_task_count"]
request_semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS_PER_JOB)
reference_prepare_lock = asyncio.Lock()
def _update_request_progress(request_index: int, value: float) -> None:
normalized = max(0.0, min(float(value), 1.0))
if normalized <= request_progress[request_index]:
return
request_progress[request_index] = normalized
if progress_callback:
progress_callback(sum(request_progress) / len(request_progress))
try:
if reference_images and not job.get("batch_enabled"):
encode_started_at = time.perf_counter()
reference_inline_images = await prepare_nano_banana_inline_images(
reference_images,
model=actual_model,
prompt=max(job["task_prompts"], key=lambda value: len(value.encode("utf-8"))),
resolution=job["resolution"],
aspect_ratio=job["aspect_ratio"],
thinking_level=thinking_level,
google_search=job.get("google_search") is True,
resize_mode=job["resize_mode"],
node_label=f"o1key批次 {job['batch_id'][:8]}",
)
print(
f"{batch_tag} 参考图已内联编码 | {len(reference_inline_images)} 张 "
f"| {time.perf_counter() - encode_started_at:.1f}s"
)
for image in reference_images:
try:
image.close()
except Exception:
pass
reference_images.clear()
async with create_http_client(
http2=True,
max_connections=32,
max_keepalive_connections=16,
) as session:
async def _generate_and_save(request_index: int) -> list[dict[str, Any]]:
async def _heartbeat() -> None:
while True:
await asyncio.sleep(1.5)
elapsed = time.monotonic() - started_at
estimate = min(0.72, 0.04 + (elapsed / 60.0) * 0.66)
_update_request_progress(request_index, estimate)
images: list[Image.Image] = []
try:
async with request_semaphore:
task_inline_images = reference_inline_images
if batch_enabled:
task = job["tasks"][request_index]
async with reference_prepare_lock:
task_images: list[Image.Image] = []
try:
if group_batch:
task_images.extend(
reference_images[index].copy()
for index in task["reference_indices"]
)
task_images.extend(await _load_reference_images(
job["model_reference_paths"][index]
for index in task["model_reference_indices"]
))
else:
task_images = await _load_reference_images(
_task_reference_paths(job, request_index)
)
task_inline_images = await prepare_nano_banana_inline_images(
task_images,
model=actual_model,
prompt=job["task_prompts"][request_index],
resolution=job["resolution"],
aspect_ratio=job["aspect_ratio"],
thinking_level=thinking_level,
google_search=job.get("google_search") is True,
resize_mode=job["resize_mode"],
node_label=(
f"o1key批次 {job['batch_id'][:8]}"
f"#{request_index + 1}"
),
)
finally:
for task_image in task_images:
task_image.close()
started_at = time.monotonic()
heartbeat = asyncio.create_task(_heartbeat())
try:
images, timing = await generate_nano_banana_async(
session=session,
base_url=base_url,
api_key=api_key,
prompt=job["task_prompts"][request_index],
model=actual_model,
resolution=job["resolution"],
aspect_ratio=job["aspect_ratio"],
inline_images=task_inline_images or None,
thinking_level=thinking_level,
google_search=job.get("google_search") is True,
resize_mode=job["resize_mode"],
node_label=f"o1key批次 {job['batch_id'][:8]}#{request_index + 1}",
request_log_enabled=VERBOSE_LOG_ENABLED,
progress_callback=lambda value: _update_request_progress(
request_index,
float(value) * 0.86,
),
download_semaphore=None,
log_task_success=False,
log_downloads=False,
)
finally:
# Progress reporting is auxiliary. Never await its
# cancellation on the critical download/save path.
heartbeat.cancel()
# generate_nano_banana_async returns after result downloads.
_update_request_progress(request_index, 0.96)
dimensions = ",".join(f"{image.width}×{image.height}" for image in images)
downloaded_bytes = int(timing.get("download_bytes") or 0)
download_seconds = float(timing.get("download_ms") or 0) / 1000
size_text = (
f"{downloaded_bytes / (1024 * 1024):.2f} MiB"
if downloaded_bytes
else "内联结果"
)
speed_text = (
f" | 速度={downloaded_bytes / (1024 * 1024) / download_seconds:.2f} MiB/s"
if downloaded_bytes and download_seconds > 0
else ""
)
print(
f"{batch_tag} #{request_index + 1} 结果就绪 "
f"| task={timing.get('task_id', '<unknown>')} "
f"| 生成={float(timing.get('task_ms') or 0) / 1000:.1f}s "
f"| 下载={download_seconds:.1f}s{speed_text} "
f"| {len(images)}{dimensions} {size_text}"
)
descriptors: list[dict[str, Any]] = []
for result_index, image in enumerate(images):
filename = _result_filename(
job,
request_index,
result_index,
_provider_result_format(image),
)
target = os.path.join(output_path, filename)
await asyncio.to_thread(
_write_provider_result,
image,
target,
)
descriptors.append({
"filename": filename,
"subfolder": relative_subfolder,
"type": "temp",
"batch_id": job["batch_id"],
"request_index": request_index + 1,
"result_index": result_index + 1,
})
_publish_partial_images(progress_callback, descriptors)
_update_request_progress(request_index, 1.0)
return descriptors
finally:
for image in images:
try:
image.close()
except Exception:
pass
images.clear()
results = await asyncio.gather(
*(
_generate_and_save(index)
for index in range(job["total_task_count"])
),
return_exceptions=True,
)
finally:
for image in reference_images:
try:
image.close()
except Exception:
pass
reference_images.clear()
for image in model_reference_images:
try:
image.close()
except Exception:
pass
model_reference_images.clear()
images: list[dict[str, Any]] = []
warnings: list[str] = []
failed_items: list[dict[str, Any]] = []
for request_index, result in enumerate(results, start=1):
if isinstance(result, BaseException):
message = str(result)
warnings.append(message)
failed_items.append({"request_index": request_index, "error": message})
else:
images.extend(result)
if not images:
detail = warnings[0] if warnings else "生成完成但没有可保存的图片"
raise RuntimeError(detail)
warning_text = f" | 警告={len(warnings)}" if warnings else ""
print(
f"{batch_tag} 完成 | 原始结果={len(images)} 张 "
f"| 总耗时={time.perf_counter() - batch_started_at:.1f}s{warning_text} "
f"| {relative_subfolder or 'temp 根目录'}"
)
return {"images": images, "warnings": warnings, "failed_items": failed_items}
async def _execute_gpt_image_job(
job: dict[str, Any],
output_directory: str,
progress_callback: Optional[ProgressCallback] = None,
) -> dict[str, Any]:
"""Generate a GPT Image batch as concurrent single-image requests (n=1)."""
from ..clients.gpt_image_client import GptImageClient, resolve_gpt_image_model
from .config import get_base_url_by_route
from .image_utils import pil_to_tensor
actual_model = resolve_gpt_image_model(job["model"], job["model_route"])
quality = resolve_gpt_image_quality(job["model"], job["quality"])
client = GptImageClient()
client.base_url = get_base_url_by_route()
client.response_log_enabled = False
client.poll_log_enabled = False
batch_tag = f"[o1key {job['batch_id'][:8]}]"
batch_started_at = time.perf_counter()
output_path, relative_subfolder = await asyncio.to_thread(
_prepare_output_directory,
job,
output_directory,
)
batch_enabled = bool(job.get("batch_enabled"))
group_batch = (
batch_enabled
and job.get("batch_mode") == IMAGE_BATCH_MODE_GROUP_TO_MODELS
)
# Normal requests need every reference. Group batches reuse only their
# bounded source set; large target/cartesian manifests are loaded per task.
reference_images = await _load_reference_images(
job.get("reference_paths", ()) if not batch_enabled or group_batch else ()
)
model_reference_images = await _load_reference_images(
job.get("model_reference_paths", ()) if not batch_enabled else ()
)
try:
reference_tensors = [
await asyncio.to_thread(pil_to_tensor, [image])
for image in reference_images
]
model_reference_tensors = [
await asyncio.to_thread(pil_to_tensor, [image])
for image in model_reference_images
]
finally:
for image in [*reference_images, *model_reference_images]:
image.close()
reference_images.clear()
model_reference_images.clear()
mask_tensor = await asyncio.to_thread(_load_mask_tensor, job.get("mask_path"))
request_progress = [0.0] * job["total_task_count"]
request_semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS_PER_JOB)
reference_prepare_lock = asyncio.Lock()
print(
f"{batch_tag} 开始 | {job['model']} | {job['resolution']} "
f"| 提示词={len(job['task_prompts']) // job['image_count']} "
f"| 每条={job['image_count']} | 总任务={job['total_task_count']} | 每次 n=1 "
f"| 参考图={len(reference_tensors)}"
)
def _update_request_progress(request_index: int, value: float) -> None:
normalized = max(0.0, min(float(value), 1.0))
if normalized <= request_progress[request_index]:
return
request_progress[request_index] = normalized
if progress_callback:
progress_callback(sum(request_progress) / len(request_progress))
async def _generate_and_save(request_index: int) -> list[dict[str, Any]]:
async def _heartbeat() -> None:
while True:
await asyncio.sleep(1.5)
elapsed = time.monotonic() - started_at
estimate = min(0.72, 0.04 + (elapsed / 60.0) * 0.66)
_update_request_progress(request_index, estimate)
images: list[Image.Image] = []
local_reference_images: list[Image.Image] = []
local_reference_tensors: list[Any] = []
try:
async with request_semaphore:
task_reference_tensors = reference_tensors
if batch_enabled:
task = job["tasks"][request_index]
async with reference_prepare_lock:
if group_batch:
task_reference_tensors = [
reference_tensors[index]
for index in task["reference_indices"]
]
local_paths = [
job["model_reference_paths"][index]
for index in task["model_reference_indices"]
]
else:
task_reference_tensors = []
local_paths = list(_task_reference_paths(job, request_index))
local_reference_images = await _load_reference_images(local_paths)
local_reference_tensors = [
await asyncio.to_thread(pil_to_tensor, [image])
for image in local_reference_images
]
task_reference_tensors = [
*task_reference_tensors,
*local_reference_tensors,
]
for reference_image in local_reference_images:
reference_image.close()
local_reference_images.clear()
started_at = time.monotonic()
heartbeat = asyncio.create_task(_heartbeat())
try:
images = await client.generate_image_async(
prompt=job["task_prompts"][request_index],
model=actual_model,
quality=quality,
size=(
None
if job["resolution"] == UNIFIED_IMAGE_SMART_RESOLUTION
else resolve_gpt_image_size(
job["resolution"], job["aspect_ratio"]
)
),
n=1,
seed=job["seed"],
image_tensor=task_reference_tensors or None,
mask_tensor=mask_tensor,
output_format=job["output_format"],
background=job["background"],
resize_mode=job["resize_mode"],
progress_callback=lambda value: _update_request_progress(
request_index,
float(value) / 100.0 * 0.86,
),
special_price_parallel=False,
log_downloads=False,
log_request_start=False,
log_prefix=f"{batch_tag} #{request_index + 1}",
)
finally:
heartbeat.cancel()
# Provider submission/result retrieval no longer needs the
# request tensors. Release them before another queued task is
# allowed to prepare its full-resolution references.
if batch_enabled:
task_reference_tensors.clear()
local_reference_tensors.clear()
if not images:
raise RuntimeError(f"第 {request_index + 1} 次请求没有返回图片")
_update_request_progress(request_index, 0.96)
selected = images[0]
filename = _result_filename(
job,
request_index,
0,
_provider_result_format(selected),
)
target = os.path.join(output_path, filename)
await asyncio.to_thread(
_write_provider_result,
selected,
target,
)
descriptor = {
"filename": filename,
"subfolder": relative_subfolder,
"type": "temp",
"batch_id": job["batch_id"],
"request_index": request_index + 1,
"result_index": 1,
}
_publish_partial_images(progress_callback, [descriptor])
_update_request_progress(request_index, 1.0)
return [descriptor]
finally:
for image in images:
image.close()
images.clear()
for image in local_reference_images:
image.close()
local_reference_images.clear()
local_reference_tensors.clear()
try:
results = await asyncio.gather(
*(
_generate_and_save(index)
for index in range(job["total_task_count"])
),
return_exceptions=True,
)
finally:
reference_tensors.clear()
model_reference_tensors.clear()
descriptors: list[dict[str, Any]] = []
warnings: list[str] = []
failed_items: list[dict[str, Any]] = []
for request_index, result in enumerate(results, start=1):
if isinstance(result, BaseException):
message = str(result)
warnings.append(message)
failed_items.append({"request_index": request_index, "error": message})
else:
descriptors.extend(result)
if not descriptors:
raise RuntimeError(warnings[0] if warnings else "生成完成但没有可保存的图片")
warning_text = f" | 警告={len(warnings)}" if warnings else ""
print(
f"{batch_tag} 完成 | 原始结果={len(descriptors)} 张 "
f"| 总耗时={time.perf_counter() - batch_started_at:.1f}s{warning_text} "
f"| {relative_subfolder or 'temp 根目录'}"
)
return {
"images": descriptors,
"warnings": warnings,
"failed_items": failed_items,
}
async def _execute_seedream_job(
job: dict[str, Any],
output_directory: str,
progress_callback: Optional[ProgressCallback] = None,
) -> dict[str, Any]:
"""Generate and atomically save one Seedream batch."""
from ..clients.seedream_image_client import SeedreamImageClient, resolve_seedream_model
from .config import get_api_key_or_raise, get_base_url_by_route
from .http2_client import create_http_client
api_key = get_api_key_or_raise("O1KEY_API_KEY")
base_url = get_base_url_by_route()
client = SeedreamImageClient(base_url=base_url, api_key=api_key)
actual_model = resolve_seedream_model(job["model"], job["model_route"])
output_path, relative_subfolder = await asyncio.to_thread(
_prepare_output_directory,
job,
output_directory,
)
batch_tag = f"[o1key {job['batch_id'][:8]}]"
batch_started_at = time.perf_counter()
batch_enabled = bool(job.get("batch_enabled"))
group_batch = (
batch_enabled
and job.get("batch_mode") == IMAGE_BATCH_MODE_GROUP_TO_MODELS
)
shared_references = await _load_reference_images(
job.get("reference_paths", ()) if not batch_enabled or group_batch else ()
)
shared_model_references = await _load_reference_images(
job.get("model_reference_paths", ()) if not batch_enabled else ()
)
request_progress = [0.0] * job["total_task_count"]
request_semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS_PER_JOB)
upload_cache: dict[int, Awaitable[str]] = {}
def _update_request_progress(request_index: int, value: float) -> None:
normalized = max(0.0, min(float(value), 1.0))
if normalized <= request_progress[request_index]:
return
request_progress[request_index] = normalized
if progress_callback:
progress_callback(sum(request_progress) / len(request_progress))
print(
f"{batch_tag} 开始 | {job['model']} | {job['resolution']} {job['aspect_ratio']} "
f"| 每条={job['image_count']} | 总任务={job['total_task_count']} "
f"| 参考图={len(job.get('reference_paths', ()))}"
)
async with create_http_client(
http2=True,
max_connections=32,
max_keepalive_connections=16,
) as session:
async def _generate_and_save(request_index: int) -> list[dict[str, Any]]:
images: list[Image.Image] = []
local_reference_images: list[Image.Image] = []
try:
async with request_semaphore:
task = job["tasks"][request_index]
if not batch_enabled:
task_reference_images = [
shared_references[index]
for index in task["reference_indices"]
] + [
shared_model_references[index]
for index in task["model_reference_indices"]
]
elif group_batch:
model_paths = [
job["model_reference_paths"][index]
for index in task["model_reference_indices"]
]
local_reference_images = await _load_reference_images(model_paths)
task_reference_images = [
shared_references[index]
for index in task["reference_indices"]
] + local_reference_images
else:
local_reference_images = await _load_reference_images(
_task_reference_paths(job, request_index)
)
task_reference_images = local_reference_images
_update_request_progress(request_index, 0.02)
images, _timing = await client.generate_async(
session=session,
prompt=job["task_prompts"][request_index],
model=actual_model,
size=(
None
if job["resolution"] == UNIFIED_IMAGE_SMART_RESOLUTION
else (
resolve_seedream_size(
job["resolution"], job["aspect_ratio"]
)
if not job.get("layer_decomposition")
else resolve_seedream_layer_size(job["resolution"])
)
),
output_format=job["output_format"],
images=task_reference_images,
layer_decomposition=bool(job.get("layer_decomposition")),
# Batch-local PIL objects are closed after each task;
# do not retain id-based cache entries that could be
# reused by a later, different image object.
upload_cache=upload_cache if not batch_enabled else None,
progress_callback=lambda value: _update_request_progress(
request_index,
0.08 + float(value) * 0.82,
),
log_downloads=False,
log_task_success=False,
)
if not images:
raise RuntimeError(f"第 {request_index + 1} 次请求没有返回图片")
_update_request_progress(request_index, 0.96)
selected = images[0]
selected_images = images if job.get("layer_decomposition") else [selected]
descriptors = []
for result_index, result_image in enumerate(selected_images):
filename = _result_filename(
job,
request_index,
result_index,
_provider_result_format(result_image),
)
await asyncio.to_thread(
_write_provider_result,
result_image,
os.path.join(output_path, filename),
)
descriptor = {
"filename": filename,
"subfolder": relative_subfolder,
"type": "temp",
"batch_id": job["batch_id"],
"request_index": request_index + 1,
"result_index": result_index + 1,
}
layer = _safe_layer_metadata(
getattr(result_image, "_o1key_seedream_layer", None)
)
if layer is not None:
descriptor["layer"] = layer
descriptors.append(descriptor)
_publish_partial_images(progress_callback, descriptors)
_update_request_progress(request_index, 1.0)
return descriptors
finally:
for image in images:
image.close()
for image in local_reference_images:
image.close()
try:
results = await asyncio.gather(
*(
_generate_and_save(index)
for index in range(job["total_task_count"])
),
return_exceptions=True,
)
finally:
for image in [*shared_references, *shared_model_references]:
image.close()
descriptors: list[dict[str, Any]] = []
warnings: list[str] = []
failed_items: list[dict[str, Any]] = []
for request_index, result in enumerate(results, start=1):
if isinstance(result, BaseException):
message = str(result)
warnings.append(message)
failed_items.append({"request_index": request_index, "error": message})
else:
descriptors.extend(result)
if not descriptors:
raise RuntimeError(warnings[0] if warnings else "生成完成但没有可保存的图片")
warning_text = f" | 警告={len(warnings)}" if warnings else ""
print(
f"{batch_tag} 完成 | 原始结果={len(descriptors)} 张 "
f"| 总耗时={time.perf_counter() - batch_started_at:.1f}s{warning_text} "
f"| {relative_subfolder or 'temp 根目录'}"
)
return {
"images": descriptors,
"warnings": warnings,
"failed_items": failed_items,
}
async def execute_image_job(
job: dict[str, Any],
output_directory: str,
progress_callback: Optional[ProgressCallback] = None,
) -> dict[str, Any]:
"""Dispatch one isolated batch to its model-family executor."""
try:
if is_gpt_image_model(job["model"]):
return await _execute_gpt_image_job(
job,
output_directory,
progress_callback,
)
if is_seedream_model(job["model"]):
return await _execute_seedream_job(
job,
output_directory,
progress_callback,
)
return await _execute_nano_banana_job(
job,
output_directory,
progress_callback,
)
except Exception as exc:
message = format_o1key_image_error(exc)
if message == str(exc):
raise
raise RuntimeError(message) from None
def find_completed_job_outputs(
batch_id: str,
output_directory: str,
generator_node_id: int,
save_node_id: int,
temp_directory: str | None = None,
) -> Optional[dict[str, Any]]:
"""Rebuild a completed result after a frontend or ComfyUI restart."""
output_root = os.path.abspath(output_directory)
direct_prefix = f"o1key_{batch_id}_"
direct_roots = [(output_root, "output")]
if temp_directory:
direct_roots.append((os.path.abspath(temp_directory), "temp"))
for direct_root, folder_type in direct_roots:
descriptors: list[dict[str, Any]] = []
if not os.path.isdir(direct_root):
continue
try:
root_entries = sorted(os.scandir(direct_root), key=lambda item: item.name)
except OSError:
root_entries = []
for entry in root_entries:
if not entry.is_file(follow_symlinks=False) or not entry.name.startswith(direct_prefix):
continue
match = OUTPUT_FILENAME_PATTERN.match(entry.name[len(direct_prefix):])
if not match:
continue
descriptors.append({
"filename": entry.name,
"subfolder": "",
"type": folder_type,
"batch_id": batch_id,
"request_index": int(match.group(1)),
"result_index": int(match.group(2)),
})
if descriptors:
return {
"batch_id": batch_id,
"state": "completed",
"generator_node_id": generator_node_id,
"save_node_id": save_node_id,
"images": descriptors,
"warnings": [],
"failed_items": [],
"error": "",
"total_count": len(descriptors),
"succeeded_count": len(descriptors),
"failed_count": 0,
"progress": 1.0,
"progress_percent": 100,
"active_batches": 0,
"queued_batches": 0,
"max_concurrent_batches": MAX_CONCURRENT_BATCHES,
"recovered_from_disk": True,
}
# Compatibility: results created by older releases remain recoverable from
# output/o1key_parallel/<date>/<batch-id>/.
parallel_root = os.path.abspath(os.path.join(output_root, "o1key_parallel"))
if not os.path.isdir(parallel_root):
return None
try:
date_entries = sorted(os.scandir(parallel_root), key=lambda item: item.name, reverse=True)
except OSError:
return None
for date_entry in date_entries:
if not date_entry.is_dir(follow_symlinks=False):
continue
batch_output = os.path.abspath(os.path.join(date_entry.path, batch_id))
if not _is_within(parallel_root, batch_output) or not os.path.isdir(batch_output):
continue
descriptors: list[dict[str, Any]] = []
try:
filenames = sorted(os.listdir(batch_output))
except OSError:
return None
for filename in filenames:
match = OUTPUT_FILENAME_PATTERN.match(filename)
if not match or not os.path.isfile(os.path.join(batch_output, filename)):
continue
descriptors.append({
"filename": filename,
"subfolder": f"o1key_parallel/{date_entry.name}/{batch_id}",
"type": "output",
"batch_id": batch_id,
"request_index": int(match.group(1)),
"result_index": int(match.group(2)),
})
if descriptors:
return {
"batch_id": batch_id,
"state": "completed",
"generator_node_id": generator_node_id,
"save_node_id": save_node_id,
"images": descriptors,
"warnings": [],
"failed_items": [],
"error": "",
"total_count": len(descriptors),
"succeeded_count": len(descriptors),
"failed_count": 0,
"progress": 1.0,
"progress_percent": 100,
"active_batches": 0,
"queued_batches": 0,
"max_concurrent_batches": MAX_CONCURRENT_BATCHES,
"recovered_from_disk": True,
}
return None
def _resolve_temp_result_paths(
batch_id: str,
descriptors: Any,
temp_directory: str,
) -> list[str]:
if not isinstance(descriptors, list) or not descriptors:
raise ValueError("没有可保存的临时图片")
if len(descriptors) > MAX_UNIFIED_IMAGE_TASKS:
raise ValueError("单次保存的图片数量过多")
temp_root = os.path.abspath(temp_directory)
prefix = f"o1key_{batch_id}_"
paths: list[str] = []
for item in descriptors:
if not isinstance(item, dict) or item.get("type") != "temp":
raise ValueError("图片结果必须来自 ComfyUI temp 目录")
if str(item.get("batch_id") or "") != batch_id:
raise ValueError("图片结果与批次 ID 不匹配")
filename = str(item.get("filename") or "")
if (
not filename
or os.path.basename(filename) != filename
or not filename.startswith(prefix)
or not OUTPUT_FILENAME_PATTERN.match(filename[len(prefix):])
):
raise ValueError("临时图片文件名无效")
candidate = os.path.abspath(os.path.join(temp_root, filename))
if not _is_within(temp_root, candidate) or not os.path.isfile(candidate):
raise ValueError(f"临时图片不存在:{filename}")
paths.append(candidate)
return paths
@dataclass
class JobRecord:
job: dict[str, Any]
state: str = "queued"
images: list[dict[str, Any]] = field(default_factory=list)
warnings: list[str] = field(default_factory=list)
failed_items: list[dict[str, Any]] = field(default_factory=list)
error: str = ""
progress: float = 0.0
started_at: float = 0.0
ended_at: float = 0.0
last_emitted_percent: int = -1
task: Optional[asyncio.Task] = None
save_lock: Optional[asyncio.Lock] = None
def _saved_record_images(record: JobRecord, batch_id: str) -> list[dict[str, Any]]:
images = list(record.images)
if not images or any(
not isinstance(item, dict)
or not (
item.get("type") == "output"
or (item.get("type") == "temp" and item.get("external_saved") is True)
)
or str(item.get("batch_id") or "") != batch_id
for item in images
):
return []
return [dict(item) for item in images]
class ParallelImageJobManager:
"""FIFO scheduler with ten active batch slots and isolated result events."""
def __init__(
self,
executor: JobExecutor,
event_sender: EventSender,
max_concurrent: int = MAX_CONCURRENT_BATCHES,
history_store: PersistentJobHistory | None = None,
):
self.executor = executor
self.event_sender = event_sender
self.max_concurrent = max_concurrent
self.history_store = history_store
self.jobs: dict[str, JobRecord] = {}
self.completed_order: deque[str] = deque()
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._semaphore: Optional[asyncio.Semaphore] = None
self._lock: Optional[asyncio.Lock] = None
def _ensure_loop_resources(self) -> None:
loop = asyncio.get_running_loop()
if self._loop is loop:
return
if self._loop is not None and any(
record.state in {"queued", "running"} for record in self.jobs.values()
):
raise RuntimeError("后台任务调度器事件循环发生变化")
self._loop = loop
self._semaphore = asyncio.Semaphore(self.max_concurrent)
self._lock = asyncio.Lock()
def _counts(self) -> tuple[int, int]:
active = sum(record.state == "running" for record in self.jobs.values())
queued = sum(record.state == "queued" for record in self.jobs.values())
return active, queued
def _queue_position(self, record: JobRecord) -> int:
if record.state != "queued":
return 0
queued = [item for item in self.jobs.values() if item.state == "queued"]
try:
return queued.index(record) + 1
except ValueError:
return 0
async def _emit(self, record: JobRecord) -> None:
active, queued = self._counts()
payload = {
"batch_id": record.job["batch_id"],
"generator_node_id": record.job["generator_node_id"],
"save_node_id": record.job["save_node_id"],
"state": record.state,
"images": list(record.images),
"warnings": list(record.warnings),
"failed_items": list(record.failed_items),
"error": record.error,
"total_count": (
len(record.images)
if record.state == "completed" and record.job.get("layer_decomposition")
else record.job["total_task_count"]
),
"request_count": record.job["total_task_count"],
"result_count": len(record.images),
"succeeded_count": len(record.images),
"failed_count": len(record.failed_items),
"progress": record.progress,
"progress_percent": round(record.progress * 100),
"active_batches": active,
"queued_batches": queued,
"queue_position": self._queue_position(record),
"max_concurrent_batches": self.max_concurrent,
"create_time": round(float(record.job.get("submitted_at") or 0) * 1000),
"execution_start_time": round(record.started_at * 1000) or 0,
"execution_end_time": round(record.ended_at * 1000) or 0,
}
try:
await self.event_sender("o1key.image_job", payload)
except Exception:
# A disconnected browser must never cancel a paid generation task.
pass
async def _emit_running_progress(self, record: JobRecord) -> None:
# Progress callbacks are scheduled tasks. If the executor finishes in
# the same event-loop turn, do not turn a stale progress tick into a
# duplicate terminal event.
if record.state == "running":
await self._emit(record)
async def submit(self, job: dict[str, Any]) -> dict[str, Any]:
self._ensure_loop_resources()
async with self._lock:
batch_id = job["batch_id"]
if batch_id in self.jobs:
raise ValueError("批次 ID 已存在")
record = JobRecord(job=dict(job))
self.jobs[batch_id] = record
record.task = asyncio.create_task(self._run(record))
active, queued = self._counts()
return {
"batch_id": batch_id,
"state": record.state,
"active_batches": active,
"queued_batches": queued,
"queue_position": self._queue_position(record),
"max_concurrent_batches": self.max_concurrent,
"total_count": record.job["total_task_count"],
"succeeded_count": 0,
"failed_count": 0,
"create_time": round(float(record.job.get("submitted_at") or 0) * 1000),
}
async def _run(self, record: JobRecord) -> None:
try:
await self._emit(record)
async with self._semaphore:
record.state = "running"
record.started_at = time.time()
await self._emit(record)
def _on_progress(value: float) -> None:
normalized = max(record.progress, min(float(value), 1.0))
percent = int(normalized * 100)
record.progress = normalized
if percent == record.last_emitted_percent:
return
record.last_emitted_percent = percent
asyncio.create_task(self._emit_running_progress(record))
def _on_partial_images(images: Any) -> None:
if record.state != "running" or not isinstance(images, list):
return
merged = {
(
item.get("request_index"),
item.get("result_index"),
item.get("filename"),
item.get("subfolder"),
item.get("type"),
): item
for item in record.images
if isinstance(item, dict)
}
changed = False
for value in images[:MAX_UNIFIED_IMAGE_TASKS]:
descriptor = _history_image_descriptor(
value,
record.job["batch_id"],
)
if descriptor is None:
continue
key = (
descriptor.get("request_index"),
descriptor.get("result_index"),
descriptor.get("filename"),
descriptor.get("subfolder"),
descriptor.get("type"),
)
if merged.get(key) != descriptor:
merged[key] = descriptor
changed = True
if not changed:
return
record.images = sorted(
merged.values(),
key=lambda item: (
int(item.get("request_index") or 1),
int(item.get("result_index") or 1),
str(item.get("filename") or ""),
),
)[:MAX_UNIFIED_IMAGE_TASKS]
asyncio.create_task(self._emit_running_progress(record))
_on_progress.publish_images = _on_partial_images
result = await self.executor(
record.job,
_on_progress,
)
record.images = list(result.get("images") or [])
record.warnings = list(result.get("warnings") or [])
record.failed_items = list(result.get("failed_items") or [])
record.progress = 1.0
record.state = "completed"
record.ended_at = time.time()
await self._emit(record)
except asyncio.CancelledError:
record.state = "cancelled"
record.error = "任务已取消"
record.ended_at = time.time()
print(f"[o1key {record.job['batch_id'][:8]}] 已取消")
await self._emit(record)
except Exception as exc:
record.state = "failed"
record.error = str(exc)
record.ended_at = time.time()
record.failed_items = [
{"request_index": index + 1, "error": record.error}
for index in range(record.job["total_task_count"])
]
error_text = " ".join(record.error.split())
if len(error_text) > 600:
error_text = error_text[:597] + "..."
print(f"[o1key {record.job['batch_id'][:8]}] 失败 | {error_text}")
await self._emit(record)
finally:
record.ended_at = record.ended_at or time.time()
if self.history_store is not None:
try:
await asyncio.to_thread(
self.history_store.upsert,
self.status(record.job["batch_id"]),
)
except OSError:
# A history-index failure must not change a paid job result.
pass
self.completed_order.append(record.job["batch_id"])
while len(self.completed_order) > MAX_RETAINED_JOBS:
expired = self.completed_order.popleft()
self.jobs.pop(expired, None)
async def cancel(self, batch_id: str) -> Optional[dict[str, Any]]:
self._ensure_loop_resources()
async with self._lock:
record = self.jobs.get(batch_id)
if record is None:
return None
task = record.task
if record.state in {"queued", "running"} and task is not None and not task.done():
task.cancel()
if task is not None and not task.done():
try:
await task
except asyncio.CancelledError:
# Defensive compatibility if a future _run implementation
# chooses to propagate cancellation again.
pass
if record.state in {"queued", "running"}:
record.state = "cancelled"
record.error = "任务已取消"
await self._emit(record)
return self.status(batch_id)
def status(self, batch_id: str) -> Optional[dict[str, Any]]:
record = self.jobs.get(batch_id)
if record is None:
return None
active, queued = self._counts()
return {
"batch_id": batch_id,
"state": record.state,
"generator_node_id": record.job["generator_node_id"],
"save_node_id": record.job["save_node_id"],
"images": list(record.images),
"warnings": list(record.warnings),
"failed_items": list(record.failed_items),
"error": record.error,
"total_count": (
len(record.images)
if record.state == "completed" and record.job.get("layer_decomposition")
else record.job["total_task_count"]
),
"request_count": record.job["total_task_count"],
"result_count": len(record.images),
"succeeded_count": len(record.images),
"failed_count": len(record.failed_items),
"progress": record.progress,
"progress_percent": round(record.progress * 100),
"active_batches": active,
"queued_batches": queued,
"queue_position": self._queue_position(record),
"max_concurrent_batches": self.max_concurrent,
"create_time": round(float(record.job.get("submitted_at") or 0) * 1000),
"execution_start_time": round(record.started_at * 1000) or 0,
"execution_end_time": round(record.ended_at * 1000) or 0,
}
def register_o1key_image_job_routes(PromptServer, web, folder_paths):
"""Register REST endpoints once and return the process-wide scheduler."""
async def _send_event(event: str, payload: dict[str, Any]) -> None:
await PromptServer.instance.send_json(event, payload)
async def _executor(
job: dict[str, Any],
progress_callback: ProgressCallback,
):
try:
return await execute_image_job(
job,
folder_paths.get_temp_directory(),
progress_callback,
)
finally:
await asyncio.to_thread(
cleanup_job_snapshot,
job,
folder_paths.get_temp_directory(),
)
get_user_directory = getattr(folder_paths, "get_user_directory", None)
history_root = (
get_user_directory()
if callable(get_user_directory)
else folder_paths.get_output_directory()
)
history_store = PersistentJobHistory(
os.path.join(history_root, "o1key", JOB_HISTORY_FILENAME)
)
manager = ParallelImageJobManager(
_executor,
_send_event,
history_store=history_store,
)
@PromptServer.instance.routes.post("/o1key/image/jobs")
async def submit_o1key_image_job(request):
job = None
try:
payload = await request.json()
job = normalize_job_payload(payload)
await asyncio.to_thread(
snapshot_reference_files,
job,
folder_paths.get_input_directory(),
folder_paths.get_temp_directory(),
)
result = await manager.submit(job)
return web.json_response(result, status=202, headers={"Cache-Control": "no-store"})
except ValueError as exc:
if job is not None:
await asyncio.to_thread(
cleanup_job_snapshot,
job,
folder_paths.get_temp_directory(),
)
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
if job is not None:
await asyncio.to_thread(
cleanup_job_snapshot,
job,
folder_paths.get_temp_directory(),
)
return web.json_response({"error": str(exc)}, status=500)
@PromptServer.instance.routes.get("/o1key/image/jobs/history")
async def get_o1key_image_job_history(request):
try:
limit = max(0, min(MAX_RETAINED_JOBS, int(request.query.get("limit", 64))))
except (TypeError, ValueError):
limit = 64
items = await asyncio.to_thread(history_store.list, limit)
return web.json_response({"items": items}, headers={"Cache-Control": "no-store"})
@PromptServer.instance.routes.post("/o1key/image/jobs/history")
async def update_o1key_image_job_history(request):
try:
payload = await request.json()
if not isinstance(payload, dict):
raise ValueError("请求体必须是对象")
if payload.get("clear") is True:
await asyncio.to_thread(history_store.clear)
else:
batch_id = _canonical_batch_id(payload.get("batch_id"))
await asyncio.to_thread(history_store.delete, batch_id)
return web.json_response({"ok": True}, headers={"Cache-Control": "no-store"})
except ValueError as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
@PromptServer.instance.routes.get("/o1key/image/jobs/{batch_id}")
async def get_o1key_image_job(request):
try:
batch_id = _canonical_batch_id(request.match_info.get("batch_id"))
except ValueError as exc:
return web.json_response({"error": str(exc)}, status=400)
result = manager.status(batch_id)
if result is None:
try:
generator_node_id = _safe_node_id(
request.query.get("generator_node_id"),
"生成节点 ID",
)
save_node_id = _safe_node_id(
request.query.get("save_node_id"),
"保存节点 ID",
)
except ValueError:
return web.json_response({"error": "任务不存在"}, status=404)
result = await asyncio.to_thread(
find_completed_job_outputs,
batch_id,
folder_paths.get_output_directory(),
generator_node_id,
save_node_id,
folder_paths.get_temp_directory(),
)
if result is None:
return web.json_response({"error": "任务不存在"}, status=404)
return web.json_response(result, headers={"Cache-Control": "no-store"})
@PromptServer.instance.routes.post("/o1key/image/jobs/{batch_id}/cancel")
async def cancel_o1key_image_job(request):
try:
batch_id = _canonical_batch_id(request.match_info.get("batch_id"))
except ValueError as exc:
return web.json_response({"error": str(exc)}, status=400)
result = await manager.cancel(batch_id)
if result is None:
return web.json_response({"error": "任务不存在"}, status=404)
record = manager.jobs.get(batch_id)
if record is not None:
await asyncio.to_thread(
cleanup_job_snapshot,
record.job,
folder_paths.get_temp_directory(),
)
return web.json_response(result, headers={"Cache-Control": "no-store"})
@PromptServer.instance.routes.post("/o1key/image/save")
async def save_o1key_image_results(request):
try:
payload = await request.json()
if not isinstance(payload, dict):
raise ValueError("请求体必须是对象")
batch_id = _canonical_batch_id(payload.get("batch_id"))
save_node_id = _safe_node_id(payload.get("save_node_id"), "保存节点 ID")
record = manager.jobs.get(batch_id)
if record is not None and int(record.job["save_node_id"]) != save_node_id:
raise ValueError("保存节点 ID 与批次不匹配")
job_settings = record.job if record is not None else None
save_format = normalize_save_format(
job_settings.get("save_format")
if job_settings is not None
else payload.get("format")
)
if record is not None and (
is_gpt_image_model(record.job.get("model", ""))
or is_seedream_model(record.job.get("model", ""))
):
save_format = "原始"
naming_rule = normalize_naming_rule(
job_settings.get("naming_rule")
if job_settings is not None
else payload.get("naming_rule")
)
save_location = normalize_save_location(
job_settings.get("save_location")
if job_settings is not None
else payload.get("save_location")
)
filename_prefix = str(
(
job_settings.get("filename_prefix")
if job_settings is not None
else payload.get("filename_prefix")
)
or "o1key"
).strip()
if naming_rule == "自定义前缀" and (
not filename_prefix or len(filename_prefix) > 512
):
raise ValueError("文件名前缀无效")
prompt = payload.get("prompt")
extra_pnginfo = payload.get("extra_pnginfo")
if extra_pnginfo is not None and not isinstance(extra_pnginfo, dict):
raise ValueError("工作流元数据无效")
async def _save_results():
if record is not None:
saved = _saved_record_images(record, batch_id)
if saved:
return saved
source_paths = _resolve_temp_result_paths(
batch_id,
payload.get("images"),
folder_paths.get_temp_directory(),
)
references = record.job.get("references", []) if record is not None else []
main_filename = references[0].get("name") if references else None
descriptors = await asyncio.to_thread(
save_temp_images,
source_paths,
filename_prefix,
save_format,
folder_paths.get_output_directory(),
folder_paths,
prompt,
extra_pnginfo,
save_location,
naming_rule,
main_filename,
)
source_descriptors = payload.get("images") or []
for index, descriptor in enumerate(descriptors):
descriptor["batch_id"] = batch_id
if index < len(source_descriptors):
source = source_descriptors[index]
if isinstance(source, dict):
try:
request_index = int(source.get("request_index") or index + 1)
result_index = int(source.get("result_index") or 1)
except (TypeError, ValueError):
request_index = index + 1
result_index = 1
descriptor["request_index"] = max(1, request_index)
descriptor["result_index"] = max(1, result_index)
layer = _safe_layer_metadata(source.get("layer"))
if layer is not None:
descriptor["layer"] = layer
if record is not None:
record.images = [dict(item) for item in descriptors]
try:
await asyncio.to_thread(
history_store.upsert,
manager.status(batch_id),
)
except OSError:
# Saving pixels remains successful if history storage is unavailable.
pass
return descriptors
if record is not None:
if record.save_lock is None:
record.save_lock = asyncio.Lock()
async with record.save_lock:
descriptors = await _save_results()
else:
descriptors = await _save_results()
return web.json_response({"images": descriptors}, headers={"Cache-Control": "no-store"})
except ValueError as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
return manager
__all__ = [
"JOB_HISTORY_FILENAME",
"MAX_CONCURRENT_BATCHES",
"MAX_CONCURRENT_REQUESTS_PER_JOB",
"ParallelImageJobManager",
"PersistentJobHistory",
"cleanup_job_snapshot",
"execute_image_job",
"find_completed_job_outputs",
"normalize_job_payload",
"register_o1key_image_job_routes",
"snapshot_reference_files",
]