o1keyandClaude Sonnet 4.5
659f94656c
- 新增宽高比（8种）× 分辨率档位（2K/3K/4K）选择，后端自动换算像素 - 5.0支持2K/3K，4.5支持2K/4K，搭配错误时明确报错 - 新增生图数量（1-10），2张以上自动并发请求，加快出图速度 - 移除旧版尺寸预设、宽度、高度输入 Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
421 lines
16 KiB
Python
421 lines
16 KiB
Python
"""
|
||
豆包生图节点
|
||
后端通过 new-api 兼容层调用豆包官方 API
|
||
"""
|
||
|
||
import asyncio
|
||
import time
|
||
import numpy as np
|
||
import torch
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from PIL import Image
|
||
from typing import List, Optional
|
||
|
||
from ..clients.doubao_image_client import DoubaoImageClient
|
||
from ..utils.image_utils import tensor_to_pil
|
||
|
||
|
||
# ── 模型列表 ──────────────────────────────────────────────────────────────────
|
||
_MODELS = [
|
||
"doubao-seedream-5-0-260128",
|
||
"doubao-seedream-4-5-251128",
|
||
]
|
||
|
||
# ── 宽高比列表 ─────────────────────────────────────────────────────────────────
|
||
_ASPECT_RATIOS = ["1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9"]
|
||
|
||
# ── 分辨率档位（每个模型支持的档位不同）──────────────────────────────────────
|
||
# 5.0：2K / 3K
|
||
# 4.5：2K / 4K
|
||
_RESOLUTIONS = ["2K", "3K", "4K"]
|
||
|
||
# ── 像素对照表 ─────────────────────────────────────────────────────────────────
|
||
# 结构：{ 模型版本key: { 分辨率: { 宽高比: (宽, 高) } } }
|
||
_SIZE_TABLE = {
|
||
"5-0": {
|
||
"2K": {
|
||
"1:1": (2048, 2048),
|
||
"4:3": (2304, 1728),
|
||
"3:4": (1728, 2304),
|
||
"16:9": (2848, 1600),
|
||
"9:16": (1600, 2848),
|
||
"3:2": (2496, 1664),
|
||
"2:3": (1664, 2496),
|
||
"21:9": (3136, 1344),
|
||
},
|
||
"3K": {
|
||
"1:1": (3072, 3072),
|
||
"4:3": (3456, 2592),
|
||
"3:4": (2592, 3456),
|
||
"16:9": (4096, 2304),
|
||
"9:16": (2304, 4096),
|
||
"3:2": (3744, 2496),
|
||
"2:3": (2496, 3744),
|
||
"21:9": (4704, 2016),
|
||
},
|
||
},
|
||
"4-5": {
|
||
"2K": {
|
||
"1:1": (2048, 2048),
|
||
"4:3": (2304, 1728),
|
||
"3:4": (1728, 2304),
|
||
"16:9": (2848, 1600),
|
||
"9:16": (1600, 2848),
|
||
"3:2": (2496, 1664),
|
||
"2:3": (1664, 2496),
|
||
"21:9": (3136, 1344),
|
||
},
|
||
"4K": {
|
||
"1:1": (4096, 4096),
|
||
"4:3": (4704, 3520),
|
||
"3:4": (3520, 4704),
|
||
"16:9": (5504, 3040),
|
||
"9:16": (3040, 5504),
|
||
"3:2": (4992, 3328),
|
||
"2:3": (3328, 4992),
|
||
"21:9": (6240, 2656),
|
||
},
|
||
},
|
||
}
|
||
|
||
# 每个模型版本支持的分辨率档位
|
||
_MODEL_RESOLUTIONS = {
|
||
"5-0": ["2K", "3K"],
|
||
"4-5": ["2K", "4K"],
|
||
}
|
||
|
||
# 并发请求超时（秒）
|
||
_CONCURRENT_TIMEOUT = 330
|
||
|
||
|
||
def _model_key(model: str) -> str:
|
||
"""从模型 ID 中提取版本 key（'5-0' 或 '4-5'）。"""
|
||
for key in _SIZE_TABLE:
|
||
if key in model:
|
||
return key
|
||
raise ValueError(f"无法识别模型版本：{model}，支持的模型：{_MODELS}")
|
||
|
||
|
||
def _pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
|
||
"""
|
||
PIL Image 列表 → ComfyUI IMAGE tensor [B, H, W, C]，值域 [0, 1]。
|
||
多张尺寸不同时，以最大尺寸为准，较小图像丢弃。
|
||
"""
|
||
if not images:
|
||
placeholder = Image.new("RGB", (512, 512), color=(128, 128, 128))
|
||
images = [placeholder]
|
||
|
||
base_size = max(images, key=lambda img: img.size[0] * img.size[1]).size
|
||
matched = [img for img in images if img.size == base_size]
|
||
skipped = len(images) - len(matched)
|
||
if skipped:
|
||
print(f"[豆包生图] 丢弃 {skipped} 张非最大尺寸图像，仅输出 {base_size[0]}×{base_size[1]} 的 {len(matched)} 张")
|
||
|
||
tensors = []
|
||
for img in matched:
|
||
arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
||
tensors.append(torch.from_numpy(arr))
|
||
|
||
return torch.stack(tensors, dim=0) # [B, H, W, C]
|
||
|
||
|
||
class DoubaoImage:
|
||
"""豆包生图 —— 通过宽高比 + 分辨率档位选择尺寸，后端自动换算真实像素"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"模型": (
|
||
_MODELS,
|
||
{"default": _MODELS[0]},
|
||
),
|
||
"提示词": (
|
||
"STRING",
|
||
{
|
||
"multiline": True,
|
||
"default": "",
|
||
"tooltip": "用于创建或编辑图像的文本提示",
|
||
},
|
||
),
|
||
"宽高比": (
|
||
_ASPECT_RATIOS,
|
||
{
|
||
"default": "1:1",
|
||
"tooltip": "图像宽高比。所有分辨率档位均支持这些比例",
|
||
},
|
||
),
|
||
"分辨率": (
|
||
_RESOLUTIONS,
|
||
{
|
||
"default": "2K",
|
||
"tooltip": (
|
||
"图像分辨率档位。\n"
|
||
"• Seedream 5.0：支持 2K / 3K\n"
|
||
"• Seedream 4.5：支持 2K / 4K\n"
|
||
"（3K 与 4.5 或 4K 与 5.0 搭配时将报错）"
|
||
),
|
||
},
|
||
),
|
||
"生图数量": (
|
||
"INT",
|
||
{
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 10,
|
||
"step": 1,
|
||
"tooltip": "生成图像的数量。2-10 张时自动并发请求，加快出图速度",
|
||
},
|
||
),
|
||
"种子": (
|
||
"INT",
|
||
{
|
||
"default": 0,
|
||
"min": 0,
|
||
"max": 2147483647,
|
||
"step": 1,
|
||
"control_after_generate": True,
|
||
"tooltip": "用于生成的随机种子",
|
||
},
|
||
),
|
||
"部分失败时停止": (
|
||
"BOOLEAN",
|
||
{
|
||
"default": True,
|
||
"tooltip": (
|
||
"启用时：任意一张失败即抛出错误并中止。\n"
|
||
"禁用时：返回已成功生成的图像，忽略失败项"
|
||
),
|
||
},
|
||
),
|
||
},
|
||
"optional": {
|
||
"图像": (
|
||
"IMAGE",
|
||
{
|
||
"tooltip": (
|
||
"用于图生图的输入图像。"
|
||
"单参考或多参考生成时，可输入1-10张图像列表"
|
||
),
|
||
},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("图像",)
|
||
FUNCTION = "generate"
|
||
CATEGORY = "comfyui_o1key/豆包"
|
||
|
||
# ── 并发核心：在新 event loop 里 gather N 个 _generate_async ─────────────
|
||
|
||
async def _run_concurrent(
|
||
self,
|
||
client: DoubaoImageClient,
|
||
生图数量: int,
|
||
model: str,
|
||
prompt: str,
|
||
size: str,
|
||
seed: int,
|
||
image_tensor,
|
||
pbar,
|
||
) -> List[dict]:
|
||
"""
|
||
并发发起 生图数量 个独立请求，每完成一个推进一格进度条。
|
||
返回结果列表：[{"index": int, "images": [...], "error": str|None}]
|
||
"""
|
||
# 固定参数（顺序生成功能暂时隐藏）
|
||
seq = "disabled"
|
||
max_img = 1
|
||
|
||
async def _one(idx: int) -> dict:
|
||
try:
|
||
imgs = await client._generate_async(
|
||
model=model,
|
||
prompt=prompt,
|
||
size=size,
|
||
seed=seed,
|
||
sequential_image_generation=seq,
|
||
max_images=max_img,
|
||
image_tensor=image_tensor,
|
||
)
|
||
return {"index": idx, "images": imgs, "error": None}
|
||
except Exception as e:
|
||
return {"index": idx, "images": [], "error": str(e)}
|
||
|
||
# 用 as_completed 方式逐个推进进度条
|
||
tasks = [asyncio.create_task(_one(i)) for i in range(生图数量)]
|
||
results = [None] * 生图数量
|
||
completed = 0
|
||
|
||
for coro in asyncio.as_completed(tasks):
|
||
res = await coro
|
||
results[res["index"]] = res
|
||
completed += 1
|
||
status = "✓" if res["error"] is None else f"✗ {res['error']}"
|
||
print(f"[豆包生图] [{completed}/{生图数量}] 第 {res['index'] + 1} 张 → {status}")
|
||
if pbar is not None:
|
||
pbar.update(1)
|
||
|
||
return results
|
||
|
||
# ── 节点主入口 ────────────────────────────────────────────────────────────
|
||
|
||
def generate(
|
||
self,
|
||
模型: str,
|
||
提示词: str,
|
||
宽高比: str,
|
||
分辨率: str,
|
||
生图数量: int,
|
||
种子: int,
|
||
部分失败时停止: bool,
|
||
图像=None,
|
||
):
|
||
start_time = time.time()
|
||
|
||
# 顺序图像生成功能暂时隐藏，固定使用默认值
|
||
顺序图像生成 = "disabled"
|
||
最大图片数 = 1
|
||
|
||
# ── 1. 校验提示词 ─────────────────────────────────────────────────────
|
||
if not 提示词.strip():
|
||
raise ValueError("提示词不能为空，请输入图像描述后重试。")
|
||
|
||
# ── 2. 解析模型版本并校验分辨率兼容性 ────────────────────────────────
|
||
try:
|
||
mkey = _model_key(模型)
|
||
except ValueError as e:
|
||
raise ValueError(str(e)) from None
|
||
|
||
supported = _MODEL_RESOLUTIONS[mkey]
|
||
if 分辨率 not in supported:
|
||
raise ValueError(
|
||
f"模型 {模型} 不支持 {分辨率} 分辨率。\n"
|
||
f"该模型支持：{' / '.join(supported)}"
|
||
)
|
||
|
||
# ── 3. 查表换算真实像素 ───────────────────────────────────────────────
|
||
w, h = _SIZE_TABLE[mkey][分辨率][宽高比]
|
||
size_str = f"{w}x{h}"
|
||
|
||
# ── 4. 打印概要 ───────────────────────────────────────────────────────
|
||
mode_str = "图生图" if 图像 is not None else "文生图"
|
||
print(
|
||
f"[豆包生图] {mode_str} | 模型={模型} | {分辨率} {宽高比} → {size_str}"
|
||
f" | 数量={生图数量} | 种子={种子}"
|
||
)
|
||
|
||
# ── 5. 初始化客户端 ───────────────────────────────────────────────────
|
||
try:
|
||
client = DoubaoImageClient()
|
||
except ValueError as e:
|
||
raise ValueError(str(e)) from None
|
||
|
||
# ── 6. 进度条（按张数计）──────────────────────────────────────────────
|
||
try:
|
||
from comfy.utils import ProgressBar
|
||
pbar = ProgressBar(生图数量)
|
||
except Exception:
|
||
pbar = None
|
||
|
||
# ── 7. 单张 / 多张分支 ────────────────────────────────────────────────
|
||
if 生图数量 == 1:
|
||
# 单张：走原有同步路径
|
||
try:
|
||
pil_images: List[Image.Image] = client.generate_sync(
|
||
model=模型,
|
||
prompt=提示词,
|
||
size=size_str,
|
||
seed=种子,
|
||
sequential_image_generation=顺序图像生成,
|
||
max_images=最大图片数,
|
||
image_tensor=图像,
|
||
)
|
||
except RuntimeError as e:
|
||
raise RuntimeError(str(e)) from None
|
||
except Exception as e:
|
||
raise RuntimeError(f"豆包生图请求失败: {e}") from None
|
||
|
||
if pbar is not None:
|
||
pbar.update(1)
|
||
|
||
else:
|
||
# 多张：并发请求
|
||
def _run_in_thread():
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
return loop.run_until_complete(
|
||
self._run_concurrent(
|
||
client=client,
|
||
生图数量=生图数量,
|
||
model=模型,
|
||
prompt=提示词,
|
||
size=size_str,
|
||
seed=种子,
|
||
image_tensor=图像,
|
||
pbar=pbar,
|
||
)
|
||
)
|
||
finally:
|
||
loop.close()
|
||
|
||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||
future = executor.submit(_run_in_thread)
|
||
try:
|
||
results = future.result(timeout=_CONCURRENT_TIMEOUT)
|
||
except TimeoutError:
|
||
raise RuntimeError(
|
||
f"并发生图超时（>{_CONCURRENT_TIMEOUT}s），请检查网络或减少生图数量"
|
||
)
|
||
|
||
# 统计成功 / 失败
|
||
success_results = [r for r in results if r and r["error"] is None]
|
||
failed_results = [r for r in results if r and r["error"] is not None]
|
||
|
||
if failed_results:
|
||
fail_info = "；".join(
|
||
f"第{r['index']+1}张: {r['error']}" for r in failed_results
|
||
)
|
||
if 部分失败时停止:
|
||
raise RuntimeError(
|
||
f"{len(failed_results)}/{生图数量} 张生成失败：{fail_info}\n"
|
||
"（可将【部分失败时停止】设为 False 以返回已成功的图像）"
|
||
)
|
||
else:
|
||
print(f"[豆包生图] 警告：{len(failed_results)}/{生图数量} 张失败，已忽略：{fail_info}")
|
||
|
||
if not success_results:
|
||
raise RuntimeError("所有图像均生成失败，请检查网络或 API 配置。")
|
||
|
||
# 按原始 index 排序，展平为 PIL 列表
|
||
success_results.sort(key=lambda r: r["index"])
|
||
pil_images = []
|
||
for r in success_results:
|
||
pil_images.extend(r["images"])
|
||
|
||
# ── 8. PIL → tensor ───────────────────────────────────────────────────
|
||
output_tensor = _pil_list_to_tensor(pil_images)
|
||
|
||
# ── 9. 完成日志 ───────────────────────────────────────────────────────
|
||
elapsed = time.time() - start_time
|
||
print(
|
||
f"[豆包生图] 完成！耗时 {elapsed:.1f}s，"
|
||
f"输出 {output_tensor.shape[0]} 张 "
|
||
f"{output_tensor.shape[2]}×{output_tensor.shape[1]}"
|
||
)
|
||
|
||
return (output_tensor,)
|
||
|
||
|
||
# ── 节点注册 ──────────────────────────────────────────────────────────────────
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"DoubaoImage": DoubaoImage,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"DoubaoImage": "豆包生图",
|
||
}
|