publisher/comfyui_o1key
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 新增豆包5.0生图模型节点

Browse Source 
- 新增 DoubaoImage 节点，支持豆包 5.0 文生图模型
- 新增 clients/doubao_image_client.py 封装豆包图像生成 API
- 新增 nodes/doubao_image.py 节点逻辑实现
- 注册节点映射及显示名称「豆包生图」
This commit is contained in:
Jony
2026-04-15 00:38:55 +08:00
parent a5dfebb1eb
commit 0b9d7583c7
4 changed files with 636 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+3 -1
View File
@@ -21,7 +21,7 @@ except Exception:
import ssl
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, SaveCleanImage, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, MultiResPreview, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage
# 报错弹框友好文案（不修改原节点代码，仅在外层统一处理）
_MSG_TIMEOUT = "API 请求超时，请稍后重试或检查网络。"
@@ -78,6 +78,7 @@ NODE_CLASS_MAPPINGS = {
"Seedance": Seedance,
"SeedanceMultiModal": SeedanceMultiModal,
"StreamPreview": StreamPreview,
"DoubaoImage": DoubaoImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -101,6 +102,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"Seedance": "Seedance 视频生成",
"SeedanceMultiModal": "Seedance 多模态参考生视频",
"StreamPreview": "流式文本预览",
"DoubaoImage": "豆包生图",
}
WEB_DIRECTORY = "./web"
clients/doubao_image_client.py
+341
View File
@@ -0,0 +1,341 @@
"""
豆包生图 API 客户端
端点：POST /v1/images/generations/
兼容 new-api 透传格式（OpenAI images/generations 兼容）
设计原则：
- 发送完整正确的请求体，new-api 丢弃字段是其侧问题
- 响应永远是同步 JSON（new-api 强制 stream=false）
- 图像输入以 data:image/png;base64,... 格式内联传递
"""
import asyncio
import json
import time
from io import BytesIO
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional, Union
import aiohttp
from PIL import Image
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
# ── 固定端点 ──────────────────────────────────────────────────────────────────
_ENDPOINT = "/v1/images/generations/"
# ── 轮询 / 请求超时 ───────────────────────────────────────────────────────────
_REQUEST_TIMEOUT = 300 # 单次请求超时秒数（豆包图像生成最长约 60s）
class DoubaoImageClient:
"""
豆包生图客户端（new-api 原生 OpenAI 兼容格式）
new-api 兼容性说明（基于源码分析）：
✅ 透传：model / prompt / size / response_format / watermark / image
❌ 丢弃：seed / sequential_image_generation / sequential_image_generation_options
（进入 Extra map，但 MarshalJSON 中合并代码被注释）
❌ 强制：stream 硬编码 false，图像接口无流式处理
❌ 未实现：/v1/files 文件上传（501）
节点仍发送完整字段，待 new-api 修复后自动生效。
"""
def __init__(self):
self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
self.base_url = get_api_base_url()
# ── 认证头 ────────────────────────────────────────────────────────────────
def _headers(self) -> dict:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
# ── 图像字段构建 ──────────────────────────────────────────────────────────
def _tensor_to_image_field(self, tensor) -> Union[str, List[str]]:
"""
ComfyUI IMAGE tensor → API image 字段值
单张返回字符串，多张返回字符串列表，格式：
data:image/png;base64,<base64数据>
"""
pil_images = tensor_to_pil(tensor)
data_urls = []
for img in pil_images:
b64 = encode_image_to_base64(img, format="PNG")
data_urls.append(f"data:image/png;base64,{b64}")
return data_urls[0] if len(data_urls) == 1 else data_urls
# ── 请求体构建 ────────────────────────────────────────────────────────────
def _build_body(
self,
model: str,
prompt: str,
size: str,
seed: int,
sequential_image_generation: str,
max_images: int,
image_field=None, # str | list[str] | None
) -> dict:
"""
构建完整请求体。
字段说明（对照官方示例）：
- response_format: 固定 "url"（new-api 原样透传给豆包）
- watermark: 固定 False（UI 已移除该参数）
- stream: 固定 False（new-api 强制非流式，此字段不被读取，仅显式注明）
- sequential_image_generation_options: 仅 sequential=auto 时发送
"""
body = {
"model": model,
"prompt": prompt,
"size": size,
"response_format": "url",
"watermark": False,
"seed": seed,
"sequential_image_generation": sequential_image_generation,
}
# 仅 auto 模式才发送 max_images 选项
if sequential_image_generation == "auto":
body["sequential_image_generation_options"] = {
"max_images": max_images
}
# 图像输入（图生图）
if image_field is not None:
body["image"] = image_field
return body
# ── 图像下载 ──────────────────────────────────────────────────────────────
async def _download_image(
self,
url: str,
session: aiohttp.ClientSession,
) -> Image.Image:
"""从 URL 下载图像，返回 PIL.Image。"""
async with session.get(url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(
f"图像下载失败，HTTP {resp.status}，URL: {url}"
)
data = await resp.read()
try:
img = Image.open(BytesIO(data)).convert("RGB")
except Exception as e:
raise RuntimeError(f"图像解码失败: {e}")
return img
# ── 响应解析 ──────────────────────────────────────────────────────────────
async def _parse_response(
self,
resp_json: dict,
session: aiohttp.ClientSession,
) -> List[Image.Image]:
"""
解析 /v1/images/generations 响应，返回 PIL.Image 列表。
期望格式（new-api 原样透传豆包响应）：
{
"created": 1234567890,
"data": [
{"url": "https://..."},
{"url": "https://..."}
]
}
兼容 b64_json 字段（豆包理论上也支持）。
"""
# 检查 API 层级错误
if "error" in resp_json:
err = resp_json["error"]
if isinstance(err, dict):
msg = err.get("message") or err.get("msg") or json.dumps(err, ensure_ascii=False)
else:
msg = str(err)
raise RuntimeError(f"API 返回错误: {msg}")
data_list = resp_json.get("data")
if not data_list:
raise RuntimeError(
f"API 响应中未找到 data 字段，完整响应：\n"
f"{json.dumps(resp_json, ensure_ascii=False, indent=2)}"
)
images: List[Image.Image] = []
for idx, item in enumerate(data_list):
url = item.get("url", "")
b64 = item.get("b64_json", "")
if url and url.startswith("http"):
# 优先使用 URL 模式
img = await self._download_image(url, session)
images.append(img)
print(f"[豆包生图] 第 {idx + 1} 张下载完成 ({img.size[0]}×{img.size[1]})")
elif b64:
# 回退到 base64 模式
import base64 as _b64
try:
img_data = _b64.b64decode(b64)
img = Image.open(BytesIO(img_data)).convert("RGB")
images.append(img)
print(f"[豆包生图] 第 {idx + 1} 张 base64 解码完成 ({img.size[0]}×{img.size[1]})")
except Exception as e:
raise RuntimeError(f"第 {idx + 1} 张 base64 解码失败: {e}")
else:
print(f"[豆包生图] 警告：第 {idx + 1} 条数据既无 url 也无 b64_json，已跳过")
return images
# ── 核心异步生成方法 ──────────────────────────────────────────────────────
async def _generate_async(
self,
model: str,
prompt: str,
size: str,
seed: int,
sequential_image_generation: str,
max_images: int,
image_tensor=None,
) -> List[Image.Image]:
"""
异步完整流程：构建请求 → POST → 解析 → 下载图像。
"""
# 1. 构建 image 字段
image_field = None
if image_tensor is not None:
image_field = self._tensor_to_image_field(image_tensor)
n_imgs = len(image_field) if isinstance(image_field, list) else 1
print(f"[豆包生图] 图生图模式，参考图 {n_imgs} 张")
else:
print(f"[豆包生图] 文生图模式")
# 2. 构建请求体
body = self._build_body(
model=model,
prompt=prompt,
size=size,
seed=seed,
sequential_image_generation=sequential_image_generation,
max_images=max_images,
image_field=image_field,
)
url = f"{self.base_url}{_ENDPOINT}"
print(f"[豆包生图] 提交请求 → {model} | {size}")
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
# 3. 发送 POST 请求
t0 = time.time()
async with session.post(
url,
json=body,
headers=self._headers(),
) as resp:
elapsed_req = time.time() - t0
text = await resp.text()
if resp.status != 200:
# 尝试解析错误信息
try:
err_json = json.loads(text)
err_obj = err_json.get("error", {})
if isinstance(err_obj, dict):
msg = (
err_obj.get("message")
or err_obj.get("msg")
or text
)
else:
msg = str(err_obj) or text
except Exception:
msg = text
raise RuntimeError(
f"请求失败 HTTP {resp.status}: {msg}"
)
try:
resp_json = json.loads(text)
except Exception:
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
print(f"[豆包生图] API 响应耗时 {elapsed_req:.1f}s，开始下载图像...")
# 4. 解析响应 & 下载图像（session 复用）
images = await self._parse_response(resp_json, session)
return images
# ── 同步入口（供 ComfyUI 节点调用）──────────────────────────────────────
def generate_sync(
self,
model: str,
prompt: str,
size: str,
seed: int,
sequential_image_generation: str,
max_images: int,
image_tensor=None,
) -> List[Image.Image]:
"""
同步生成接口（在独立线程中运行事件循环，避免与 ComfyUI 主循环冲突）。
Args:
model: 模型 ID
prompt: 提示词
size: 尺寸字符串，如 "2048x2048"
seed: 随机种子
sequential_image_generation: "disabled" | "auto"
max_images: 最大图片数（auto 模式生效）
image_tensor: ComfyUI IMAGE tensor（可选，图生图用）
Returns:
List[PIL.Image]
"""
coro = self._generate_async(
model=model,
prompt=prompt,
size=size,
seed=seed,
sequential_image_generation=sequential_image_generation,
max_images=max_images,
image_tensor=image_tensor,
)
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return loop.run_until_complete(coro)
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(_run)
try:
return future.result(timeout=_REQUEST_TIMEOUT + 30)
except TimeoutError:
raise RuntimeError(
f"豆包生图超时（>{_REQUEST_TIMEOUT}s），请检查网络或稍后重试"
)
nodes/__init__.py
+2 -1
View File
@@ -18,5 +18,6 @@ from .universal_llm import UniversalLLMChat
from .multi_res_preview import MultiResPreview
from .batch_images_o1key import BatchImagesO1key
from .seedance_video import Seedance, SeedanceMultiModal
from .doubao_image import DoubaoImage
__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'MultiResPreview', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview']
__all__ = ['NanoBananaPro', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'SaveCleanImage', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'MultiResPreview', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage']
nodes/doubao_image.py
+290
View File
@@ -0,0 +1,290 @@
"""
豆包生图节点
1:1 复刻字节跳动 Seedream 4 节点的前端外观（输入/输出/参数/样式）
后端通过 new-api 兼容层调用豆包官方 API
"""
import time
import numpy as np
import torch
from PIL import Image
from typing import List
from ..clients.doubao_image_client import DoubaoImageClient
from ..utils.image_utils import tensor_to_pil
# ── 尺寸预设 ──────────────────────────────────────────────────────────────────
# (显示名, 宽, 高) —— 宽高用于构造 "WxH" size 字符串
RECOMMENDED_PRESETS_SEEDREAM_4 = [
("2048×2048 (1:1)", 2048, 2048),
("2304×1728 (4:3)", 2304, 1728),
("1728×2304 (3:4)", 1728, 2304),
("2560×1440 (16:9)", 2560, 1440),
("1440×2560 (9:16)", 1440, 2560),
("2496×1664 (3:2)", 2496, 1664),
("1664×2496 (2:3)", 1664, 2496),
("3024×1296 (21:9)", 3024, 1296),
("3072×3072 (1:1)", 3072, 3072),
("4096×4096 (1:1)", 4096, 4096),
("自定义", None, None),
]
_PRESET_LABELS = [label for label, _, _ in RECOMMENDED_PRESETS_SEEDREAM_4]
# ── 模型列表 ──────────────────────────────────────────────────────────────────
# 节点下拉选项 = new-api 后台配置的模型 ID（直接透传给 API）
_MODELS = [
"doubao-seedream-5-0-260128",
"doubao-seedream-4-5-251128",
]
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:
"""豆包生图 —— 1:1 复刻字节跳动 Seedream 4 节点前端，后端对接豆包官方 API"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"模型": (
_MODELS,
{"default": _MODELS[0]},
),
"提示词": (
"STRING",
{
"multiline": True,
"default": "",
"tooltip": "用于创建或编辑图像的文本提示",
},
),
"尺寸预设": (
_PRESET_LABELS,
{
"default": _PRESET_LABELS[0],
"tooltip": '选择推荐尺寸。选择"自定义"可使用下方的宽度和高度',
},
),
"宽度": (
"INT",
{
"default": 2048,
"min": 1024,
"max": 6240,
"step": 64,
"tooltip": '图像的自定义宽度。仅当尺寸预设设置为"自定义"时生效',
},
),
"高度": (
"INT",
{
"default": 2048,
"min": 1024,
"max": 4992,
"step": 64,
"tooltip": '图像的自定义高度。仅当尺寸预设设置为"自定义"时生效',
},
),
"顺序图像生成": (
["disabled", "auto"],
{
"default": "disabled",
"tooltip": (
'分组图像生成模式。'
'"disabled"生成单张图像；'
'"auto"让模型决定是否生成多张相关图像（如故事场景、角色变体）'
),
},
),
"最大图片数": (
"INT",
{
"default": 1,
"min": 1,
"max": 15,
"step": 1,
"tooltip": (
"当顺序图像生成='auto'时生成的最大图像数量。"
"总图像数（输入+生成）不能超过15张"
),
},
),
"种子": (
"INT",
{
"default": 0,
"min": 0,
"max": 2147483647,
"step": 1,
"control_after_generate": True,
"tooltip": "用于生成的随机种子",
},
),
"部分失败时停止": (
"BOOLEAN",
{
"default": True,
"tooltip": "如果启用，当任何请求的图像缺失或返回错误时将中止执行",
},
),
},
"optional": {
"图像": (
"IMAGE",
{
"tooltip": (
"用于图生图的输入图像。"
"单参考或多参考生成时，可输入1-10张图像列表"
),
},
),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("图像",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/豆包"
# ── 核心生成方法 ──────────────────────────────────────────────────────────
def generate(
self,
模型: str,
提示词: str,
尺寸预设: str,
宽度: int,
高度: int,
顺序图像生成: str,
最大图片数: int,
种子: int,
部分失败时停止: bool,
图像=None,
):
start_time = time.time()
# ── 1. 校验提示词 ─────────────────────────────────────────────────────
if not 提示词.strip():
raise ValueError("提示词不能为空，请输入图像描述后重试。")
# ── 2. 解析尺寸 ───────────────────────────────────────────────────────
w, h = None, None
for label, tw, th in RECOMMENDED_PRESETS_SEEDREAM_4:
if label == 尺寸预设:
w, h = tw, th
break
if w is None or h is None:
# 自定义尺寸
w, h = 宽度, 高度
print(f"[豆包生图] 自定义尺寸：{w}×{h}")
size_str = f"{w}x{h}"
# ── 3. 打印概要 ───────────────────────────────────────────────────────
mode_str = "图生图" if 图像 is not None else "文生图"
seq_str = f" | 顺序生成=auto(最多{最大图片数}张)" if 顺序图像生成 == "auto" else ""
print(
f"[豆包生图] {mode_str} | 模型={模型} | 尺寸={size_str}"
f" | 种子={种子}{seq_str}"
)
# ── 4. 进度条 ─────────────────────────────────────────────────────────
try:
from comfy.utils import ProgressBar
pbar = ProgressBar(100)
except Exception:
pbar = None
def _pb(pct: int):
if pbar:
pbar.update_absolute(pct, 100)
_pb(0)
# ── 5. 调用客户端 ─────────────────────────────────────────────────────
try:
client = DoubaoImageClient()
except ValueError as e:
raise ValueError(str(e)) from None
_pb(5)
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
_pb(90)
# ── 6. 部分失败判断 ───────────────────────────────────────────────────
if 顺序图像生成 == "auto" and 部分失败时停止:
if len(pil_images) < 最大图片数:
raise RuntimeError(
f"部分图像生成失败：期望 {最大图片数} 张，"
f"实际返回 {len(pil_images)} 张。"
"（可将【部分失败时停止】设为 False 以接受不完整结果）"
)
# ── 7. PIL → tensor ───────────────────────────────────────────────────
output_tensor = _pil_list_to_tensor(pil_images)
_pb(100)
# ── 8. 完成日志 ───────────────────────────────────────────────────────
elapsed = time.time() - start_time
time_str = f"{elapsed:.1f}s"
print(
f"[豆包生图] 完成！耗时 {time_str}，"
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": "豆包生图",
}