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.
410 lines
14 KiB
Python
410 lines
14 KiB
Python
"""Seedream image client for O1Key's asynchronous image API."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import time
|
||
from typing import Any, Awaitable, Callable, Optional, Sequence
|
||
|
||
from PIL import Image
|
||
|
||
from ..utils.nano_banana_async import (
|
||
extract_async_image_result_urls,
|
||
image_to_upload_payload,
|
||
parse_completed_async_image_task,
|
||
poll_async_image_task,
|
||
submit_async_image_task,
|
||
upload_images_to_temp_urls,
|
||
)
|
||
from ..utils.o1key_image_catalog import (
|
||
MAX_UNIFIED_REFERENCE_IMAGES,
|
||
SEEDREAM_MODEL_OPTIONS,
|
||
SEEDREAM_LAYER_RESOLUTION_OPTIONS,
|
||
SEEDREAM_OUTPUT_FORMAT_OPTIONS,
|
||
SEEDREAM_SIZE_MATRIX,
|
||
UNIFIED_IMAGE_ROUTE_OPTIONS,
|
||
)
|
||
|
||
|
||
SEEDREAM_API_MODEL_ID = "dola-seedream-5-0-pro-260628-ep"
|
||
SEEDREAM_REFERENCE_MAX_BYTES = 30 * 1024 * 1024
|
||
SEEDREAM_REFERENCE_MAX_PIXELS = 6000 * 6000
|
||
SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE = 14
|
||
SEEDREAM_REFERENCE_MIN_ASPECT_RATIO = 1 / 16
|
||
SEEDREAM_REFERENCE_MAX_ASPECT_RATIO = 16
|
||
SEEDREAM_LAYER_REFERENCE_MIN_PIXELS = 512 * 512
|
||
|
||
|
||
def validate_seedream_reference_dimensions(
|
||
width: int,
|
||
height: int,
|
||
*,
|
||
label: str = "Seedream 参考图",
|
||
layer_decomposition: bool = False,
|
||
) -> None:
|
||
"""Validate the current Volcengine per-image reference-size contract."""
|
||
width = int(width)
|
||
height = int(height)
|
||
if width <= 0 or height <= 0:
|
||
raise ValueError(f"{label}尺寸无效：{width}×{height}")
|
||
|
||
pixels = width * height
|
||
if (
|
||
not layer_decomposition
|
||
and (
|
||
width <= SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE
|
||
or height <= SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE
|
||
)
|
||
):
|
||
raise ValueError(
|
||
f"{label}宽和高都必须大于 {SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE}px，"
|
||
f"当前为 {width}×{height}"
|
||
)
|
||
ratio = width / height
|
||
if (
|
||
ratio < SEEDREAM_REFERENCE_MIN_ASPECT_RATIO
|
||
or ratio > SEEDREAM_REFERENCE_MAX_ASPECT_RATIO
|
||
):
|
||
raise ValueError(
|
||
f"{label}宽高比必须在 1:16～16:1，当前为 {width}:{height}"
|
||
)
|
||
if layer_decomposition:
|
||
if not (
|
||
SEEDREAM_LAYER_REFERENCE_MIN_PIXELS
|
||
<= pixels
|
||
<= SEEDREAM_REFERENCE_MAX_PIXELS
|
||
):
|
||
raise ValueError(
|
||
f"{label}总像素必须在 512×512（262144）～6000×6000（36000000）之间，"
|
||
f"当前为 {width}×{height}（{pixels}）"
|
||
)
|
||
return
|
||
if pixels > SEEDREAM_REFERENCE_MAX_PIXELS:
|
||
raise ValueError(
|
||
f"{label}总像素不能超过 6000×6000（36000000），"
|
||
f"当前为 {width}×{height}（{pixels}）"
|
||
)
|
||
|
||
|
||
def validate_seedream_reference_image(
|
||
image: Image.Image,
|
||
*,
|
||
label: str = "Seedream 参考图",
|
||
layer_decomposition: bool = False,
|
||
) -> None:
|
||
"""Validate reference dimensions and the exact bytes sent to the uploader."""
|
||
validate_seedream_reference_dimensions(
|
||
image.width,
|
||
image.height,
|
||
label=label,
|
||
layer_decomposition=layer_decomposition,
|
||
)
|
||
payload, _extension, _content_type = image_to_upload_payload(image)
|
||
try:
|
||
payload_size = (
|
||
os.path.getsize(payload)
|
||
if isinstance(payload, (str, os.PathLike))
|
||
else len(payload)
|
||
)
|
||
except (OSError, TypeError) as exc:
|
||
raise ValueError(f"无法读取{label}文件大小") from exc
|
||
if payload_size > SEEDREAM_REFERENCE_MAX_BYTES:
|
||
raise ValueError(
|
||
f"{label}文件不能超过 30MB，当前为 {payload_size / 1024 / 1024:.2f}MB"
|
||
)
|
||
|
||
|
||
def validate_seedream_reference_images(
|
||
images: Sequence[Image.Image],
|
||
*,
|
||
layer_decomposition: bool = False,
|
||
) -> None:
|
||
for index, image in enumerate(images, start=1):
|
||
validate_seedream_reference_image(
|
||
image,
|
||
label=f"Seedream 参考图{index}",
|
||
layer_decomposition=layer_decomposition,
|
||
)
|
||
|
||
|
||
def resolve_seedream_model(model_name: str, route: str) -> str:
|
||
"""Map the stable workflow value to Seedream's API model identifier."""
|
||
if model_name == SEEDREAM_API_MODEL_ID:
|
||
return model_name
|
||
if model_name not in SEEDREAM_MODEL_OPTIONS:
|
||
raise ValueError(f"Seedream 模型无效：{model_name}")
|
||
if route not in UNIFIED_IMAGE_ROUTE_OPTIONS:
|
||
raise ValueError(f"Seedream 模型线路无效：{route}")
|
||
# O1Key currently exposes one Seedream endpoint for every displayed route.
|
||
return SEEDREAM_API_MODEL_ID
|
||
|
||
|
||
def build_seedream_submit_body(
|
||
*,
|
||
model: str,
|
||
prompt: str,
|
||
size: Optional[str],
|
||
output_format: str,
|
||
image_urls: Optional[Sequence[str]] = None,
|
||
layer_decomposition: bool = False,
|
||
) -> dict[str, Any]:
|
||
"""Build and validate the paid Seedream request without logging URLs."""
|
||
normalized_prompt = str(prompt or "").strip()
|
||
if not normalized_prompt and not layer_decomposition:
|
||
raise ValueError("请输入提示词")
|
||
if model != SEEDREAM_API_MODEL_ID:
|
||
raise ValueError(f"Seedream API 模型无效：{model}")
|
||
|
||
normalized_format = str(output_format or "").strip().lower()
|
||
if normalized_format not in SEEDREAM_OUTPUT_FORMAT_OPTIONS:
|
||
raise ValueError("Seedream 输出格式仅支持 png 或 jpeg")
|
||
|
||
if layer_decomposition and normalized_format != "png":
|
||
raise ValueError("Seedream 图层拆分仅支持 png 输出格式")
|
||
|
||
normalized_size = str(size or "").strip().lower().replace("*", "x").replace("×", "x")
|
||
if normalized_size:
|
||
if layer_decomposition:
|
||
normalized_size = "auto" if normalized_size == "auto" else normalized_size.upper()
|
||
if normalized_size not in SEEDREAM_LAYER_RESOLUTION_OPTIONS:
|
||
raise ValueError(f"Seedream 图层拆分分辨率无效：{size}")
|
||
elif normalized_size not in set(SEEDREAM_SIZE_MATRIX.values()):
|
||
raise ValueError(f"Seedream 图片尺寸无效：{size}")
|
||
|
||
urls = [str(url or "").strip() for url in (image_urls or ())]
|
||
if len(urls) > MAX_UNIFIED_REFERENCE_IMAGES:
|
||
raise ValueError(f"Seedream 参考图最多支持 {MAX_UNIFIED_REFERENCE_IMAGES} 张")
|
||
if any(not url.startswith("https://") for url in urls):
|
||
raise ValueError("Seedream 参考图必须使用临时素材 HTTPS URL")
|
||
if layer_decomposition and len(urls) != 1:
|
||
raise ValueError("Seedream 图层拆分必须且只能提供1张参考图")
|
||
|
||
body: dict[str, Any] = {
|
||
"model": model,
|
||
"n": 1,
|
||
"output_format": normalized_format,
|
||
"watermark": False,
|
||
}
|
||
if normalized_size:
|
||
body["size"] = normalized_size
|
||
if normalized_prompt:
|
||
body["prompt"] = normalized_prompt
|
||
if urls:
|
||
body["images"] = urls
|
||
if layer_decomposition:
|
||
body["layer_decomposition"] = True
|
||
return body
|
||
|
||
|
||
def _seedream_result_items(payload: Any) -> list[dict[str, Any]]:
|
||
"""Return the first documented image-item list without exposing its URLs."""
|
||
pending = [payload]
|
||
seen: set[int] = set()
|
||
while pending:
|
||
value = pending.pop(0)
|
||
if not isinstance(value, dict) or id(value) in seen:
|
||
continue
|
||
seen.add(id(value))
|
||
images = value.get("images")
|
||
if isinstance(images, list) and all(isinstance(item, dict) for item in images):
|
||
return images
|
||
for key in ("data", "result", "output"):
|
||
nested = value.get(key)
|
||
if isinstance(nested, dict):
|
||
pending.append(nested)
|
||
return []
|
||
|
||
|
||
def _bounded_int_list(value: Any, *, length: int) -> list[int] | None:
|
||
if not isinstance(value, (list, tuple)) or len(value) != length:
|
||
return None
|
||
try:
|
||
return [int(item) for item in value]
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
def extract_seedream_layer_metadata(payload: Any) -> list[dict[str, Any]]:
|
||
"""Sanitize layer metadata; result URLs are deliberately excluded."""
|
||
metadata: list[dict[str, Any]] = []
|
||
for offset, item in enumerate(_seedream_result_items(payload)):
|
||
try:
|
||
z_index = max(0, min(16, int(item.get("z_index", offset))))
|
||
except (TypeError, ValueError):
|
||
z_index = offset
|
||
safe: dict[str, Any] = {"z_index": z_index}
|
||
for key, limit in (("name", 200), ("description", 1000), ("size", 64), ("output_format", 16)):
|
||
value = item.get(key)
|
||
if isinstance(value, str) and value.strip():
|
||
safe[key] = value.strip()[:limit]
|
||
bounding_box = item.get("bounding_box")
|
||
if isinstance(bounding_box, dict):
|
||
absolute = _bounded_int_list(bounding_box.get("absolute"), length=4)
|
||
normalized = _bounded_int_list(bounding_box.get("normalized"), length=4)
|
||
safe_box = {}
|
||
if absolute is not None:
|
||
safe_box["absolute"] = absolute
|
||
if normalized is not None:
|
||
safe_box["normalized"] = normalized
|
||
if safe_box:
|
||
safe["bounding_box"] = safe_box
|
||
metadata.append(safe)
|
||
return metadata
|
||
|
||
|
||
class SeedreamImageClient:
|
||
"""Upload references, submit one Seedream task, poll it, and decode results."""
|
||
|
||
def __init__(self, *, base_url: str, api_key: str):
|
||
self.base_url = str(base_url).rstrip("/")
|
||
self.api_key = api_key
|
||
|
||
async def generate_async(
|
||
self,
|
||
*,
|
||
session: Any,
|
||
prompt: str,
|
||
model: str,
|
||
size: Optional[str],
|
||
output_format: str,
|
||
images: Optional[Sequence[Image.Image]] = None,
|
||
layer_decomposition: bool = False,
|
||
upload_cache: Optional[dict[int, Awaitable[str]]] = None,
|
||
check_interrupt: Optional[Callable[[], None]] = None,
|
||
progress_callback: Optional[Callable[[float], None]] = None,
|
||
result_url_callback: Optional[Callable[[str], None]] = None,
|
||
log_downloads: bool = True,
|
||
log_task_success: bool = True,
|
||
task_completed_callback: Optional[
|
||
Callable[[str, int, float, list[str]], None]
|
||
] = None,
|
||
) -> tuple[list[Image.Image], dict[str, Any]]:
|
||
if check_interrupt:
|
||
check_interrupt()
|
||
|
||
reference_images = list(images or ())
|
||
validate_seedream_reference_images(
|
||
reference_images,
|
||
layer_decomposition=layer_decomposition,
|
||
)
|
||
|
||
task_started = time.time()
|
||
image_urls = await upload_images_to_temp_urls(
|
||
session=session,
|
||
base_url=self.base_url,
|
||
api_key=self.api_key,
|
||
images=reference_images,
|
||
node_label="Seedream",
|
||
check_interrupt=check_interrupt,
|
||
upload_cache=upload_cache,
|
||
log_success=log_task_success,
|
||
)
|
||
body = build_seedream_submit_body(
|
||
model=model,
|
||
prompt=prompt,
|
||
size=size,
|
||
output_format=output_format,
|
||
image_urls=image_urls,
|
||
layer_decomposition=layer_decomposition,
|
||
)
|
||
|
||
task_id = await submit_async_image_task(
|
||
session,
|
||
self.base_url,
|
||
self.api_key,
|
||
body,
|
||
"Seedream",
|
||
log_body_enabled=False,
|
||
log_success=log_task_success,
|
||
)
|
||
task_payload = await poll_async_image_task(
|
||
session,
|
||
self.base_url,
|
||
self.api_key,
|
||
task_id,
|
||
"Seedream",
|
||
check_interrupt=check_interrupt,
|
||
log_body_enabled=False,
|
||
progress_callback=progress_callback,
|
||
log_success=log_task_success,
|
||
)
|
||
task_done = time.time()
|
||
parse_started = time.time()
|
||
task_payload, parsed = await parse_completed_async_image_task(
|
||
task_payload,
|
||
session,
|
||
self.base_url,
|
||
self.api_key,
|
||
task_id,
|
||
"Seedream",
|
||
check_interrupt=check_interrupt,
|
||
result_url_callback=None,
|
||
log_downloads=log_downloads,
|
||
)
|
||
if isinstance(parsed, tuple) and len(parsed) == 2:
|
||
result_images, metrics = parsed
|
||
else:
|
||
result_images = parsed
|
||
metrics = {
|
||
"download_bytes": 0,
|
||
"download_seconds": 0.0,
|
||
"download_wall_seconds": 0.0,
|
||
"inline_images": 0,
|
||
}
|
||
|
||
result_urls = extract_async_image_result_urls(task_payload)
|
||
result_metadata = extract_seedream_layer_metadata(task_payload)
|
||
if layer_decomposition:
|
||
paired = []
|
||
for index, image in enumerate(result_images):
|
||
metadata = (
|
||
result_metadata[index]
|
||
if index < len(result_metadata)
|
||
else {"z_index": index}
|
||
)
|
||
setattr(image, "_o1key_seedream_layer", metadata)
|
||
paired.append((metadata.get("z_index", index), index, image))
|
||
paired.sort(key=lambda item: (item[0], item[1]))
|
||
result_images = [item[2] for item in paired]
|
||
result_metadata = [
|
||
getattr(image, "_o1key_seedream_layer", {"z_index": index})
|
||
for index, image in enumerate(result_images)
|
||
]
|
||
if result_url_callback:
|
||
for url in result_urls:
|
||
result_url_callback(url)
|
||
if task_completed_callback:
|
||
task_completed_callback(
|
||
task_id,
|
||
len(result_images),
|
||
time.time() - task_started,
|
||
result_urls,
|
||
)
|
||
|
||
return result_images, {
|
||
"task_id": task_id,
|
||
"task_ids": [task_id],
|
||
"task_ms": (task_done - task_started) * 1000,
|
||
"parse_ms": (time.time() - parse_started) * 1000,
|
||
"download_ms": metrics["download_wall_seconds"] * 1000,
|
||
"download_total_ms": metrics["download_seconds"] * 1000,
|
||
"download_bytes": metrics["download_bytes"],
|
||
"inline_images": metrics["inline_images"],
|
||
"result_metadata": result_metadata,
|
||
}
|
||
|
||
|
||
__all__ = [
|
||
"SEEDREAM_API_MODEL_ID",
|
||
"SEEDREAM_LAYER_REFERENCE_MIN_PIXELS",
|
||
"SEEDREAM_REFERENCE_MAX_BYTES",
|
||
"SEEDREAM_REFERENCE_MAX_PIXELS",
|
||
"SeedreamImageClient",
|
||
"build_seedream_submit_body",
|
||
"extract_seedream_layer_metadata",
|
||
"resolve_seedream_model",
|
||
"validate_seedream_reference_dimensions",
|
||
"validate_seedream_reference_image",
|
||
"validate_seedream_reference_images",
|
||
]
|