diff --git a/__init__.py b/__init__.py index 964aabb..70f1481 100644 --- a/__init__.py +++ b/__init__.py @@ -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" diff --git a/clients/doubao_image_client.py b/clients/doubao_image_client.py new file mode 100644 index 0000000..0f870a1 --- /dev/null +++ b/clients/doubao_image_client.py @@ -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, + """ + 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），请检查网络或稍后重试" + ) diff --git a/nodes/__init__.py b/nodes/__init__.py index 5799075..831568f 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -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'] diff --git a/nodes/doubao_image.py b/nodes/doubao_image.py new file mode 100644 index 0000000..9ddefb7 --- /dev/null +++ b/nodes/doubao_image.py @@ -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": "豆包生图", +}