Add Grok and VEO video workflow support

This commit is contained in:
o1key
2026-06-04 16:06:23 +08:00
parent 5d9aff9ca7
commit 30e9603f77
24 changed files with 3752 additions and 528 deletions
+6 -3
View File
@@ -110,10 +110,13 @@ class KVideoFirstLast:
"mode": mode_api,
"duration": 时长,
}
if 生成音频 == "打开":
body["generate_audio"] = True
metadata = {}
if 尾帧 is not None:
body["metadata"] = {"image_tail": _image_to_base64(尾帧, scale)}
metadata["image_tail"] = _image_to_base64(尾帧, scale)
if 生成音频 == "打开":
metadata["sound"] = "on"
if metadata:
body["metadata"] = metadata
body_str = json.dumps(body, ensure_ascii=False)
body_size = len(body_str.encode("utf-8"))
+4 -1
View File
@@ -82,6 +82,9 @@ class KVideoImage2Video:
CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", seed=0):
if 模式 == "720p" and 生成音频 == "打开":
raise RuntimeError("K26 仅1080p支持音频,请将模式切换为1080p或关闭生成音频。")
api_key = get_api_key_or_raise()
base_url = get_base_url_by_route(网络线路)
headers = {
@@ -109,7 +112,7 @@ class KVideoImage2Video:
"duration": 时长,
}
if 生成音频 == "打开":
body["generate_audio"] = True
body["metadata"] = {"sound": "on"}
body_str = json.dumps(body, ensure_ascii=False)
body_size = len(body_str.encode("utf-8"))
+3 -1
View File
@@ -14,6 +14,7 @@ from .remove_metadata import BatchCleanMetadata
from .video_preview import VideoPreview
from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset
from .veo_video import GoogleVeo
from .newapi_veo_video import Google31Video
from .flux_edit import FluxImageEdit
from .universal_llm import UniversalLLMChat
from .batch_images_o1key import BatchImagesO1key
@@ -22,6 +23,7 @@ from .nano_banana_v2 import NanoBananaV2, NanoBananaV2Batch, AsyncImageGenerator
from .doubao_image import DoubaoImage
from .gpt_image import O1keyGPTImage, O1keyGPTImageBatch
from .grok_image import O1keyGrokImage
from .grok_video import O1keyGrokVideo
from .K_video_firstlast import KVideoFirstLast
from .K_video_image2video import KVideoImage2Video
from .K3_video import K3Video
@@ -33,4 +35,4 @@ from .remove_bg import O1keyRemoveBackground
from .color_remove_bg import O1keyColorRemoveBG
from .grid_splitter import O1keyGridSplitter
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGPTImageBatch', 'O1keyGrokImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG', 'O1keyGridSplitter']
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'Google31Video', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGPTImageBatch', 'O1keyGrokImage', 'O1keyGrokVideo', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG', 'O1keyGridSplitter']
+172 -145
View File
@@ -4,23 +4,21 @@ ComfyUI 自定义节点,用于批量处理图像生成任务
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
"""
import io as _io
import re
import json
import time
import math
import base64
import random
import asyncio
import aiohttp
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Tuple, List
from typing import Callable, Optional, Tuple, List
from PIL import Image
import torch
import numpy as np
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_images_for_request_body_limit
from comfy_api.latest import io
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
from ..utils.file_utils import (
ImageInfo,
load_images_from_folder,
@@ -30,7 +28,7 @@ from ..utils.file_utils import (
save_image,
)
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route, get_api_key_or_raise
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
from ..utils.nano_banana_async import generate_nano_banana_async
from ..clients.gemini_client import GeminiAPIClient
from ..models_config import (
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
@@ -73,64 +71,28 @@ REQUEST_LOG_ENABLED = False
# ============================================================================
_NODE = "Nano Banana"
_ENDPOINT = "/v1/chat/completions"
_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)")
def _get_headers(api_key: str) -> dict:
return {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"X-Accel-Buffering": "no",
"Cache-Control": "no-cache, no-transform",
}
def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
if pbar is None:
return None
last_progress = [0.0]
def _on_progress(progress: float) -> None:
try:
progress = max(0.0, min(float(progress), 1.0))
except (TypeError, ValueError):
return
if progress <= last_progress[0]:
return
pbar.update(progress - last_progress[0])
last_progress[0] = progress
return _on_progress
def _build_request_body(
prompt: str,
model: str,
aspect_ratio: str,
resolution: str,
images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False,
) -> dict:
google_config = {
"image_config": {
"image_size": resolution,
}
}
if aspect_ratio and aspect_ratio != "智能":
google_config["image_config"]["aspect_ratio"] = aspect_ratio
def _make_body(encoded_images: Optional[List[tuple]] = None) -> dict:
content_parts = [{"type": "text", "text": prompt}]
if encoded_images:
for mime_type, b64 in encoded_images:
content_parts.append({
"type": "image_url",
"image_url": {"url": f"data:{mime_type};base64,{b64}"}
})
body = {
"model": model,
"stream": True,
"messages": [{"role": "user", "content": content_parts}],
"extra_body": {"google": google_config},
}
if enable_grounding:
body["extra_body"]["google_search"] = True
return body
encoded_images = None
if images:
encoded_images = encode_images_for_request_body_limit(images, _make_body)
return _make_body(encoded_images)
async def _generate_single_openai(
async def _generate_single_async(
session: aiohttp.ClientSession,
base_url: str,
api_key: str,
@@ -140,82 +102,25 @@ async def _generate_single_openai(
aspect_ratio: str,
images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False,
progress_callback: Optional[Callable[[float], None]] = None,
thinking_level: Optional[str] = None,
) -> List[Image.Image]:
url = f"{base_url}{_ENDPOINT}"
headers = _get_headers(api_key)
body = _build_request_body(
result_images, _ = await generate_nano_banana_async(
session=session,
base_url=base_url,
api_key=api_key,
prompt=prompt,
model=model,
aspect_ratio=aspect_ratio,
resolution=resolution,
aspect_ratio=aspect_ratio,
images=images,
enable_grounding=enable_grounding,
thinking_level=thinking_level,
node_label="BatchNanoBananaPro",
request_log_enabled=REQUEST_LOG_ENABLED,
progress_callback=progress_callback,
)
if REQUEST_LOG_ENABLED:
extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False)
print(f"[请求] POST {url} | model={model} | extra_body={extra}")
last_status = None
for attempt in range(DEFAULT_MAX_RETRIES + 1):
resp = await session.post(url, headers=headers, json=body)
if resp.status == 200:
break
last_status = resp.status
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})")
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
resp.close()
await asyncio.sleep(delay)
continue
error_text = await resp.text()
resp.close()
if resp.status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
try:
err_json = json.loads(error_text)
msg = err_json.get("error", {}).get("message", error_text[:200])
except Exception:
msg = error_text[:200]
raise RuntimeError(f"API 错误 ({resp.status}): {msg}")
else:
if last_status and last_status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
full_content = ""
buffer = ""
async for raw_chunk in resp.content.iter_any():
buffer += raw_chunk.decode("utf-8")
while "\n" in buffer:
line_str, buffer = buffer.split("\n", 1)
line_str = line_str.strip()
if not line_str or not line_str.startswith("data:"):
continue
data_str = line_str[5:].strip()
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
delta = chunk.get("choices", [{}])[0].get("delta", {})
if "content" in delta:
full_content += delta["content"]
except (json.JSONDecodeError, IndexError):
continue
resp.close()
if not full_content:
raise RuntimeError("API 未返回有效内容")
matches = list(_IMAGE_RE.finditer(full_content))
if not matches:
raise RuntimeError(f"响应中未找到图片: {full_content[:100]}")
last_match = matches[-1]
img_data = base64.b64decode(last_match.group(2))
final_image = Image.open(_io.BytesIO(img_data)).convert("RGB")
return [final_image]
return result_images
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
@@ -244,7 +149,7 @@ def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.
return pil_to_tensor(matched)
class BatchNanoBananaPro:
class BatchNanoBananaPro(io.ComfyNode):
"""
批量 Nano Banana 节点
@@ -290,6 +195,116 @@ class BatchNanoBananaPro:
def __init__(self):
pass
@classmethod
def define_schema(cls):
normal_aspect_ratios = [
"智能", "1:1", "2:3", "3:2", "3:4", "4:3",
"4:5", "5:4", "9:16", "16:9", "21:9",
]
nano_banana_2_aspect_ratios = [
"智能", "1:1", "1:4", "1:8", "2:3", "3:2", "3:4",
"4:1", "4:3", "4:5", "5:4", "8:1",
"9:16", "16:9", "21:9",
]
return io.Schema(
node_id="BatchNanoBananaPro",
display_name="批量 Nano Banana",
category="image/batch",
inputs=[
io.String.Input(
"prompt",
default="一个中国女子的OOTD",
multiline=True,
),
io.DynamicCombo.Input("模型", options=[
io.DynamicCombo.Option("Nano Banana Pro", [
io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"),
io.Combo.Input("分辨率", options=["1K", "2K", "4K"], default="2K"),
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
]),
io.DynamicCombo.Option("Nano Banana 2", [
io.Combo.Input("宽高比", options=nano_banana_2_aspect_ratios, default="智能"),
io.Combo.Input("分辨率", options=["512px", "1K", "2K", "4K"], default="2K"),
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
io.Combo.Input("思考深度", options=["", ""], default=""),
]),
io.DynamicCombo.Option("Nano Banana", [
io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"),
io.Combo.Input("分辨率", options=["1K"], default="1K"),
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
]),
]),
io.Combo.Input("图片格式", options=["原始", "JPEG", "PNG", "WebP"], default="原始"),
io.Combo.Input("计费", options=["特价", "官方"], default="特价"),
io.Combo.Input("网络", options=NETWORK_ROUTE_OPTIONS, default="全球加速"),
io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF),
io.String.Input("文件夹1", default="", multiline=False),
io.String.Input("文件夹2", default="", multiline=False),
io.String.Input("文件夹3", default="", multiline=False),
io.String.Input("文件夹4", default="", multiline=False),
io.String.Input("文件夹5", default="", multiline=False),
io.String.Input("保存路径", default="", multiline=False),
io.Combo.Input("图片配对模式", options=cls.PAIRING_MODES, default="不配对"),
io.Image.Input("参考图1", optional=True),
io.Image.Input("参考图2", optional=True),
io.Image.Input("参考图3", optional=True),
io.Image.Input("参考图4", optional=True),
io.Image.Input("参考图5", optional=True),
],
outputs=[
io.Image.Output(display_name="输出图像"),
],
)
@classmethod
def execute(
cls,
prompt,
模型,
图片格式,
计费,
网络,
seed,
文件夹1,
文件夹2,
文件夹3,
文件夹4,
文件夹5,
保存路径,
图片配对模式,
**kwargs,
) -> io.NodeOutput:
model_name = 模型["模型"]
宽高比 = 模型.get("宽高比", "智能")
分辨率 = 模型.get("分辨率", "2K")
思考深度 = 模型.get("思考深度")
谷歌搜索 = 模型.get("谷歌搜索", "关闭")
if 思考深度:
kwargs["思考深度"] = 思考深度
kwargs["谷歌搜索"] = 谷歌搜索
node = cls()
output_tensor, = node.process_batch(
prompt=prompt,
文件夹1=文件夹1,
文件夹2=文件夹2,
文件夹3=文件夹3,
文件夹4=文件夹4,
文件夹5=文件夹5,
seed=seed,
图片配对模式=图片配对模式,
模型=model_name,
计费=计费,
宽高比=宽高比,
分辨率=分辨率,
图片格式=图片格式,
网络=网络,
保存路径=保存路径,
**kwargs,
)
return io.NodeOutput(output_tensor)
def resize_to_megapixels(
self,
@@ -352,7 +367,6 @@ class BatchNanoBananaPro:
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
"default": "不配对"
})
return {
"required": {
"prompt": ("STRING", {
@@ -532,9 +546,11 @@ class BatchNanoBananaPro:
enable_grounding: bool = False,
base_filename: str = None,
image_format: str = "原始",
progress_callback: Optional[Callable[[float], None]] = None,
thinking_level: Optional[str] = None,
) -> dict:
"""
执行单个生成任务(OpenAI 兼容接口)
执行单个生成任务(异步生图接口)
"""
result = {
"task_index": task_index,
@@ -550,10 +566,10 @@ class BatchNanoBananaPro:
# 准备输入图片
input_pil_images = [info.image for info in images]
# 调用 OpenAI 兼容接口生成图片
# 调用异步生图接口生成图片
generated_images = []
try:
gen_images = await _generate_single_openai(
gen_images = await _generate_single_async(
session=session,
base_url=base_url,
api_key=api_key,
@@ -563,6 +579,8 @@ class BatchNanoBananaPro:
aspect_ratio=aspect_ratio,
images=input_pil_images if input_pil_images else None,
enable_grounding=enable_grounding,
progress_callback=progress_callback,
thinking_level=thinking_level,
)
generated_images.extend(gen_images)
except Exception as e:
@@ -661,9 +679,10 @@ class BatchNanoBananaPro:
prompts_per_task: Optional[List[str]] = None,
enable_grounding: bool = False,
image_format: str = "原始",
thinking_level: Optional[str] = None,
) -> List[dict]:
"""
异步批量处理所有任务(OpenAI 兼容接口)
异步批量处理所有任务(异步生图接口)
"""
total_tasks = len(pairs)
@@ -734,6 +753,8 @@ class BatchNanoBananaPro:
enable_grounding=enable_grounding,
base_filename=base_filename,
image_format=image_format,
progress_callback=_make_progress_callback(pbar),
thinking_level=thinking_level,
)
)
tasks.append(task)
@@ -773,13 +794,12 @@ class BatchNanoBananaPro:
else:
fail_count += 1
# 更新 ComfyUI 原生进度条
if pbar is not None:
pbar.update(1)
# 大任务额外显示百分比里程碑
if show_milestone and milestone_index < len(milestones):
progress = completed / total_tasks
if pbar is not None and getattr(pbar, "total", 0):
progress = pbar.current / pbar.total
else:
progress = success_count / total_tasks
if progress >= milestones[milestone_index]:
percentage = int(milestones[milestone_index] * 100)
print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
@@ -856,10 +876,15 @@ class BatchNanoBananaPro:
start_time = time.time()
# 从 kwargs 提取搜索参数(界面显示为「关闭/打开」,转为 bool 供调用)
enable_grounding: bool = False
enable_grounding: bool = kwargs.get("谷歌搜索", "关闭") == "打开"
# 拼接实际模型 ID
base_model_id = self.MODEL_ID_MAP.get(模型, "nano-banana-pro")
思考深度 = kwargs.get("思考深度", "")
thinking_level = None
if base_model_id == "nano-banana-2":
thinking_level = "High" if 思考深度 == "" else "Low"
if base_model_id == "nano-banana":
if 计费 == "官方":
raise ValueError(f"模型 \"{模型}\" 仅支持特价计费")
@@ -961,11 +986,12 @@ class BatchNanoBananaPro:
grounding_str = ""
if enable_grounding:
grounding_str = " | 谷歌搜索接地"
thinking_str = f" | 思考:{thinking_level}" if thinking_level else ""
if batch_prompts:
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}")
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}{thinking_str}")
else:
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}")
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}{thinking_str}")
# 创建 ComfyUI 原生进度条
pbar = None
@@ -1025,6 +1051,7 @@ class BatchNanoBananaPro:
prompts_per_task=prompts_per_task,
enable_grounding=enable_grounding,
image_format=图片格式,
thinking_level=thinking_level,
)
)
except Exception as e:
@@ -1146,7 +1173,7 @@ class BatchNanoBananaPro:
raise RuntimeError(str(e)) from None
except Exception as e:
raise type(e)(str(e)) from None
raise RuntimeError(str(e)) from None
finally:
# 查询余额
+82 -7
View File
@@ -42,6 +42,56 @@ except ImportError:
_FOLDER_PATHS_AVAILABLE = False
def _make_node_progress_callback(progress_bar, task_index: int, total_tasks: int):
if progress_bar is None:
return None
total_units = max(1, total_tasks) * 100
base_units = max(0, task_index - 1) * 100
last_pct = {"value": -1}
def _callback(pct: int):
try:
pct_value = int(round(float(pct)))
except (TypeError, ValueError):
return
pct_value = max(0, min(100, pct_value))
if pct_value < last_pct["value"]:
return
last_pct["value"] = pct_value
progress_bar.update_absolute(
min(total_units, base_units + pct_value),
total_units,
)
return _callback
def _resolve_async_size(value: str) -> str:
value = (value or "").strip()
if not value or value == "智能" or value.lower() == "auto":
return "auto"
first_part = value.split("")[0].strip()
normalized_size = first_part.lower().replace("*", "x").replace("×", "x")
size_parts = [part.strip() for part in normalized_size.split("x")]
if len(size_parts) == 2 and all(part.isdigit() for part in size_parts):
return f"{int(size_parts[0])}x{int(size_parts[1])}"
allowed = {"auto", "1024x1024", "1K", "2K", "4K"}
if first_part in allowed:
return first_part
if "4K" in value:
return "4K"
if "2K" in value:
return "2K"
if "1K" in value:
return "1K"
return "auto"
class O1keyGPTImage:
"""
o1key GPT Image 节点
@@ -123,6 +173,10 @@ class O1keyGPTImage:
"default": "自动",
"tooltip": "Image quality: 高=high, 中=medium, 低=low, 自动=auto",
})
optional_inputs["输出格式"] = (["png", "jpeg", "webp"], {
"default": "jpeg",
"tooltip": "Generated image output format",
})
optional_inputs["seed"] = ("INT", {
"default": 0,
"min": 0,
@@ -160,6 +214,7 @@ class O1keyGPTImage:
网络: str = "全球加速",
分辨率: str = "auto",
质量: str = "自动",
输出格式: str = "jpeg",
生图数量: int = 1,
seed: int = 0,
遮罩=None,
@@ -190,7 +245,7 @@ class O1keyGPTImage:
raise ValueError("提供了遮罩但未提供图片,请同时提供图片和遮罩")
# ── 2. 解析分辨率显示值 → API 参数值 ──────────────────────────────────
size = "auto" if 分辨率 == "智能" else 分辨率.split("")[0].strip()
size = _resolve_async_size(分辨率)
# ── 2b. 解析模型显示值 → API 参数值 ───────────────────────────────────
_model_map = {"gpt-image-2-次卡": "gpt-image-2-c", "gpt-image-2-按量": "gpt-image-2"}
@@ -216,6 +271,8 @@ class O1keyGPTImage:
# ── 5. 调用 API ───────────────────────────────────────────────────
all_pil_images = []
progress_total = len(batch_prompts) if batch_prompts else 1
progress_bar = ProgressBar(progress_total * 100) if _PROGRESS_BAR_AVAILABLE else None
if batch_prompts:
# 批量模式:逐条提示词调用
@@ -226,7 +283,7 @@ class O1keyGPTImage:
print("[o1key GPT Image] 用户取消,已中断批量生成")
raise InterruptProcessingException()
try:
pil_images = client.run_sync(
pil_images = client.generate_image_async_sync(
prompt=p,
model=model,
quality=quality,
@@ -235,6 +292,8 @@ class O1keyGPTImage:
seed=seed,
image_tensor=图片,
mask_tensor=遮罩,
output_format=输出格式,
progress_callback=_make_node_progress_callback(progress_bar, idx, total),
)
all_pil_images.extend(pil_images)
snippet = p[:30] + ("..." if len(p) >= 30 else "")
@@ -245,12 +304,14 @@ class O1keyGPTImage:
error_msg = str(e).split('\n')[0]
snippet = p[:30] + ("..." if len(p) >= 30 else "")
print(f"[o1key GPT Image] [{idx}/{total}] ❌ {snippet}{error_msg}")
if progress_bar is not None:
progress_bar.update_absolute(idx * 100, total * 100)
else:
# 单提示词模式
if not prompt or not prompt.strip():
raise ValueError("提示词不能为空")
try:
pil_images = client.run_sync(
pil_images = client.generate_image_async_sync(
prompt=prompt,
model=model,
quality=quality,
@@ -259,6 +320,8 @@ class O1keyGPTImage:
seed=seed,
image_tensor=图片,
mask_tensor=遮罩,
output_format=输出格式,
progress_callback=_make_node_progress_callback(progress_bar, 1, 1),
)
all_pil_images.extend(pil_images)
except InterruptProcessingException:
@@ -495,7 +558,7 @@ class O1keyGPTImageBatch:
@staticmethod
def _resolve_size(分辨率: str) -> str:
return "auto" if 分辨率 == "智能" else 分辨率.split("")[0].strip()
return _resolve_async_size(分辨率)
@staticmethod
def _resolve_model(模型: str) -> str:
@@ -510,6 +573,15 @@ class O1keyGPTImageBatch:
quality_map = {"": "high", "": "medium", "": "low", "自动": "auto"}
return quality_map.get(质量, "auto")
@staticmethod
def _resolve_output_format(图片格式: str) -> str:
output_format_map = {
"JPEG": "jpeg",
"PNG": "png",
"WebP": "webp",
}
return output_format_map.get(图片格式, "png")
@staticmethod
def _ensure_output_folder(保存路径: str) -> str:
output_folder = (保存路径 or "").strip()
@@ -637,11 +709,12 @@ class O1keyGPTImageBatch:
size = self._resolve_size(分辨率)
model = self._resolve_model(模型)
quality = self._resolve_quality(质量)
output_format = self._resolve_output_format(图片格式)
client = GptImageClient()
client.base_url = get_base_url_by_route(网络)
progress_bar = ProgressBar(total_tasks) if _PROGRESS_BAR_AVAILABLE else None
progress_bar = ProgressBar(total_tasks * 100) if _PROGRESS_BAR_AVAILABLE else None
results = []
all_saved_files = []
@@ -661,7 +734,7 @@ class O1keyGPTImageBatch:
}
try:
pil_images = client.run_sync(
pil_images = client.generate_image_async_sync(
prompt=task_prompt,
model=model,
quality=quality,
@@ -670,6 +743,8 @@ class O1keyGPTImageBatch:
seed=seed,
image_tensor=self._pair_to_tensors(pair),
mask_tensor=遮罩,
output_format=output_format,
progress_callback=_make_node_progress_callback(progress_bar, task_index, total_tasks),
)
saved_files = self._save_images(
images=pil_images,
@@ -691,7 +766,7 @@ class O1keyGPTImageBatch:
results.append(result)
if progress_bar is not None:
progress_bar.update(1)
progress_bar.update_absolute(task_index * 100, total_tasks * 100)
success_count = sum(1 for result in results if result.get("success", False))
total_generated = sum(result.get("generated_count", 0) for result in results)
+283
View File
@@ -0,0 +1,283 @@
"""
Grok Video node.
Submits a /v1/videos task, polls until completion, downloads the mp4,
and returns ComfyUI's native VIDEO object.
"""
import json
import os
from typing import List, Optional
from ..clients.grok_video_client import GrokVideoClient
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
from ..utils.image_utils import encode_images_for_request_body_limit, tensor_to_pil
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
try:
from comfy.utils import ProgressBar
PROGRESS_BAR_AVAILABLE = True
except ImportError:
ProgressBar = None
PROGRESS_BAR_AVAILABLE = False
try:
from comfy_api.input_impl import VideoFromFile
except Exception:
try:
from comfy_api.latest import InputImpl
VideoFromFile = InputImpl.VideoFromFile
except Exception:
VideoFromFile = None
MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"]
ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"]
QUALITY_OPTIONS = ["720p"]
QUALITY_VALUE_MAP = {
"720p": "high",
}
MODEL_SECONDS_OPTIONS = {
"grok-imagine-1.0-video": [6, 10, 12, 16, 20],
}
MAX_REFERENCE_IMAGES = 3
MAX_REQUEST_BODY_BYTES = 20 * 1024 * 1024
def _get_output_dir() -> str:
if FOLDER_PATHS_AVAILABLE:
base = folder_paths.get_output_directory()
else:
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
comfy_root = os.path.dirname(os.path.dirname(plugin_dir))
base = os.path.join(comfy_root, "output")
output_dir = os.path.join(base, "grok_video")
os.makedirs(output_dir, exist_ok=True)
return output_dir
def _format_mb(size_bytes: int) -> str:
return f"{size_bytes / 1024 / 1024:.2f}MB"
def _image_tensor_to_first_pil(image_tensor):
if image_tensor is None:
return None
pil_images = tensor_to_pil(image_tensor)
if not pil_images:
return None
image = pil_images[0]
if image.mode not in ("RGB", "L"):
image = image.convert("RGB")
return image
def _collect_reference_images(**kwargs) -> List[object]:
images = []
for i in range(1, MAX_REFERENCE_IMAGES + 1):
image = _image_tensor_to_first_pil(kwargs.get(f"参考图{i}"))
if image is not None:
images.append(image)
return images
def _to_data_urls(encoded_images) -> List[str]:
return [f"data:{mime};base64,{b64}" for mime, b64 in encoded_images]
def _encode_image_data_urls(
images: List[object],
prompt: str,
model: str,
aspect_ratio: str,
seconds: int,
quality: str,
) -> Optional[List[str]]:
if not images:
return None
def build_body(encoded_images):
return GrokVideoClient.build_video_body(
prompt=prompt,
model=model,
aspect_ratio=aspect_ratio,
seconds=seconds,
quality=quality,
images=_to_data_urls(encoded_images),
)
encoded = encode_images_for_request_body_limit(
images,
build_body=build_body,
max_body_bytes=MAX_REQUEST_BODY_BYTES,
)
data_urls = _to_data_urls(encoded)
return data_urls
def _validate_request_body_size(body: dict) -> None:
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
if body_size > MAX_REQUEST_BODY_BYTES:
raise ValueError(
f"Grok Video 请求体大小 {_format_mb(body_size)} 超过 "
f"{_format_mb(MAX_REQUEST_BODY_BYTES)} 限制,请减少参考图片或降低图片尺寸。"
)
class O1keyGrokVideo:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"提示词": (
"STRING",
{
"default": "",
"multiline": True,
},
),
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": NETWORK_ROUTE_OPTIONS[0]}),
"模型": (MODEL_OPTIONS, {"default": MODEL_OPTIONS[0]}),
"宽高比": (ASPECT_RATIO_OPTIONS, {"default": "16:9"}),
"秒数(按模型限制)": (
"INT",
{
"default": 5,
"min": 5,
"max": 20,
"step": 1,
"display": "number",
},
),
"画质": (QUALITY_OPTIONS, {"default": "720p"}),
},
"optional": {
"参考图1": ("IMAGE",),
"参考图2": ("IMAGE",),
"参考图3": ("IMAGE",),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/Video"
DESCRIPTION = (
"Grok Video /v1/videos task node. Supports prompt plus up to "
"three image references, multiple aspect ratios, model-specific seconds, 720p output."
)
def generate(
self,
**kwargs,
):
if VideoFromFile is None:
raise RuntimeError("当前 ComfyUI 版本不支持原生 VIDEO 输入实现 VideoFromFile。")
提示词 = kwargs.get("提示词", "")
网络线路 = kwargs.get("网络线路", NETWORK_ROUTE_OPTIONS[0])
模型 = kwargs.get("模型", MODEL_OPTIONS[0])
宽高比 = kwargs.get("宽高比", "16:9")
秒数 = kwargs.get("秒数(按模型限制)", kwargs.get("秒数(≤15s", kwargs.get("秒数", 5)))
画质 = kwargs.get("画质", "720p")
prompt = (提示词 or "").strip()
if not prompt:
raise ValueError("提示词不能为空。")
if 模型 not in MODEL_OPTIONS:
raise ValueError(f"模型仅支持: {', '.join(MODEL_OPTIONS)}")
if 宽高比 not in ASPECT_RATIO_OPTIONS:
raise ValueError(f"宽高比仅支持: {', '.join(ASPECT_RATIO_OPTIONS)}")
seconds = int(秒数)
allowed_seconds = MODEL_SECONDS_OPTIONS.get(模型)
if allowed_seconds is not None:
if seconds not in allowed_seconds:
raise ValueError(
f"模型 {模型} 仅支持秒数: "
f"{', '.join(str(s) for s in allowed_seconds)}"
"请修改为正确的秒数后再发起请求。"
)
elif seconds < 5 or seconds > 15:
raise ValueError("秒数仅支持 5 到 15。")
if 画质 not in QUALITY_OPTIONS:
raise ValueError("画质仅支持 720p。")
quality = QUALITY_VALUE_MAP[画质]
reference_images = _collect_reference_images(**kwargs)
image_data_urls = _encode_image_data_urls(
reference_images,
prompt=prompt,
model=模型,
aspect_ratio=宽高比,
seconds=seconds,
quality=quality,
)
request_body = GrokVideoClient.build_video_body(
prompt=prompt,
model=模型,
aspect_ratio=宽高比,
seconds=seconds,
quality=quality,
images=image_data_urls,
)
_validate_request_body_size(request_body)
pbar = ProgressBar(100) if PROGRESS_BAR_AVAILABLE else None
last_progress = [0]
def progress_callback(progress: int, status: str, elapsed: float):
progress_value = max(0, min(100, int(progress or 0)))
if pbar is not None and progress_value > last_progress[0]:
pbar.update(progress_value - last_progress[0])
last_progress[0] = progress_value
client = GrokVideoClient(base_url=get_base_url_by_route(网络线路))
try:
result = client.generate_video_sync(
prompt=prompt,
model=模型,
aspect_ratio=宽高比,
seconds=seconds,
quality=quality,
output_dir=_get_output_dir(),
images=image_data_urls,
poll_interval=5,
timeout=1200,
progress_callback=progress_callback,
)
if pbar is not None and last_progress[0] < 100:
pbar.update(100 - last_progress[0])
video_path = result["video_path"]
print(f"Grok Video:下载完成:{video_path}")
return (VideoFromFile(video_path),)
finally:
try:
balance_data = client.query_balance_sync()
balance_info = client.format_balance_info(balance_data)
print(f"Grok Video{balance_info}")
except Exception:
pass
NODE_CLASS_MAPPINGS = {
"O1keyGrokVideo": O1keyGrokVideo,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"O1keyGrokVideo": "Grok Video",
}
+46 -186
View File
@@ -1,20 +1,16 @@
"""
Nano Banana 节点 (V3)
ComfyUI 自定义节点,用于调用生图模型OpenAI 兼容接口)
ComfyUI 自定义节点,用于调用异步生图模型
使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动
"""
import io as _io
import re
import json
import time
import math
import base64
import random
import asyncio
import aiohttp
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional
from typing import Callable, List, Optional
import torch
import numpy as np
@@ -22,21 +18,15 @@ from PIL import Image
from comfy_api.latest import io
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_images_for_image_size_limit
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
from ..utils.config import (
NETWORK_ROUTE_OPTIONS,
get_base_url_by_route,
get_api_key_or_raise,
)
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
from ..utils.nano_banana_async import generate_nano_banana_async
from ..clients.gemini_client import GeminiAPIClient
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
try:
from comfy.utils import ProgressBar
PROGRESS_BAR_AVAILABLE = True
@@ -51,17 +41,9 @@ except ImportError:
InterruptProcessingException = RuntimeError
processing_interrupted = lambda: False
try:
import psutil
MEMORY_MONITOR_AVAILABLE = True
except ImportError:
MEMORY_MONITOR_AVAILABLE = False
DEBUG_LOG_ENABLED = True
REQUEST_LOG_ENABLED = False
_NODE = "Nano Banana"
_ENDPOINT = "/v1/chat/completions"
_REQUEST_TIMEOUT = 900
_INTERRUPT_CHECK_INTERVAL = 0.2
@@ -112,6 +94,25 @@ def _check_interrupt():
raise InterruptProcessingException()
def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
if pbar is None:
return None
last_progress = [0.0]
def _on_progress(progress: float) -> None:
try:
progress = max(0.0, min(float(progress), 1.0))
except (TypeError, ValueError):
return
if progress <= last_progress[0]:
return
pbar.update(progress - last_progress[0])
last_progress[0] = progress
return _on_progress
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
if not images:
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
@@ -169,66 +170,6 @@ def _build_model_id(model_name: str, resolution: str, billing: str) -> str:
model_id += "-official"
return model_id
_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)")
def _get_headers(api_key: str) -> dict:
return {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"X-Accel-Buffering": "no",
"Cache-Control": "no-cache, no-transform",
}
def _build_request_body(
prompt: str,
model: str,
aspect_ratio: str,
resolution: str,
images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False,
thinking_level: Optional[str] = None,
) -> dict:
google_config = {
"image_config": {
"image_size": resolution,
}
}
if aspect_ratio and aspect_ratio != "智能":
google_config["image_config"]["aspect_ratio"] = aspect_ratio
if thinking_level:
google_config["thinking_config"] = {
"thinking_level": thinking_level.lower(),
"include_thoughts": True,
}
def _make_body(encoded_images: Optional[List[tuple]] = None) -> dict:
content_parts = [{"type": "text", "text": prompt}]
if encoded_images:
for mime_type, b64 in encoded_images:
content_parts.append({
"type": "image_url",
"image_url": {"url": f"data:{mime_type};base64,{b64}"}
})
body = {
"model": model,
"stream": True,
"messages": [{"role": "user", "content": content_parts}],
"extra_body": {"google": google_config},
}
if enable_grounding:
body["extra_body"]["google_search"] = True
return body
encoded_images = None
if images:
encoded_images = encode_images_for_image_size_limit(images)
return _make_body(encoded_images)
async def _generate_single(
session: aiohttp.ClientSession,
@@ -241,105 +182,25 @@ async def _generate_single(
images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False,
thinking_level: Optional[str] = None,
progress_callback: Optional[Callable[[float], None]] = None,
) -> List[Image.Image]:
url = f"{base_url}{_ENDPOINT}"
headers = _get_headers(api_key)
body = _build_request_body(
result_images, timing = await generate_nano_banana_async(
session=session,
base_url=base_url,
api_key=api_key,
prompt=prompt,
model=model,
aspect_ratio=aspect_ratio,
resolution=resolution,
aspect_ratio=aspect_ratio,
images=images,
enable_grounding=enable_grounding,
thinking_level=thinking_level,
node_label="Nano Banana",
request_log_enabled=REQUEST_LOG_ENABLED,
check_interrupt=_check_interrupt,
progress_callback=progress_callback,
)
if REQUEST_LOG_ENABLED:
extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False)
print(f"[请求] POST {url} | model={model} | extra_body={extra}")
last_status = None
resp = None
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT, connect=30, sock_read=_REQUEST_TIMEOUT)
for attempt in range(DEFAULT_MAX_RETRIES + 1):
_check_interrupt()
resp = await session.post(url, headers=headers, json=body, timeout=timeout)
if resp.status == 200:
break
last_status = resp.status
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})")
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
print(f"Nano Banana: {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
resp.close()
await asyncio.sleep(delay)
continue
error_text = await resp.text()
resp.close()
if resp.status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
try:
err_json = json.loads(error_text)
msg = err_json.get("error", {}).get("message", error_text[:200])
except Exception:
msg = error_text[:200]
raise RuntimeError(f"API 错误 ({resp.status}): {msg}")
else:
if last_status and last_status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
full_content = ""
buffer = ""
t_request = time.time()
t_first_token = None
try:
async for raw_chunk in resp.content.iter_any():
_check_interrupt()
if t_first_token is None:
t_first_token = time.time()
buffer += raw_chunk.decode("utf-8")
while "\n" in buffer:
line_str, buffer = buffer.split("\n", 1)
line_str = line_str.strip()
if not line_str or not line_str.startswith("data:"):
continue
data_str = line_str[5:].strip()
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
delta = chunk.get("choices", [{}])[0].get("delta", {})
if "content" in delta:
full_content += delta["content"]
except (json.JSONDecodeError, IndexError):
continue
except aiohttp.ClientPayloadError as e:
if full_content and _IMAGE_RE.search(full_content):
print(f"Nano Banana: 响应流提前结束,但已收到完整图片,继续解析 ({e})")
else:
raise RuntimeError(f"响应流下载中断,请重试或检查网络/代理: {e}") from None
finally:
if resp is not None:
resp.close()
t_done = time.time()
if not full_content:
raise RuntimeError("API 未返回有效内容")
# 思考模型可能输出多张临时图片,最终图片始终是最后一张
matches = list(_IMAGE_RE.finditer(full_content))
if not matches:
raise RuntimeError(f"响应中未找到图片: {full_content[:100]}")
last_match = matches[-1]
img_data = base64.b64decode(last_match.group(2))
final_image = Image.open(_io.BytesIO(img_data)).convert("RGB")
first_token_ms = (t_first_token - t_request) * 1000 if t_first_token else 0
download_ms = (t_done - t_first_token) * 1000 if t_first_token else 0
return [final_image], first_token_ms, download_ms
return result_images, timing["task_ms"], timing["parse_ms"]
async def _generate_single_task(
@@ -354,6 +215,7 @@ async def _generate_single_task(
global_task_index: int,
enable_grounding: bool = False,
thinking_level: Optional[str] = None,
progress_callback: Optional[Callable[[float], None]] = None,
) -> dict:
result = {
"global_task_index": global_task_index,
@@ -364,7 +226,7 @@ async def _generate_single_task(
"error": None,
}
try:
gen_images, first_token_ms, download_ms = await _generate_single(
gen_images, task_ms, parse_ms = await _generate_single(
session=session,
base_url=base_url,
api_key=api_key,
@@ -375,7 +237,9 @@ async def _generate_single_task(
images=images if images else None,
enable_grounding=enable_grounding,
thinking_level=thinking_level,
progress_callback=progress_callback,
)
del task_ms, parse_ms
result["output_images"] = gen_images
result["success"] = True
result["generated_count"] = len(gen_images)
@@ -438,6 +302,7 @@ async def _process_batch_async(
global_task_index=i,
enable_grounding=enable_grounding,
thinking_level=thinking_level,
progress_callback=_make_progress_callback(pbar),
)
)
tasks.append(task)
@@ -473,9 +338,6 @@ async def _process_batch_async(
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
print(f"Nano Banana: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
if pbar is not None:
pbar.update(1)
all_results.extend(batch_results)
import gc; gc.collect()
await asyncio.sleep(0.1)
@@ -576,7 +438,7 @@ class NanoBanana(io.ComfyNode):
pil_imgs = tensor_to_pil(kwargs[key])
input_images.extend(pil_imgs)
if input_images and len(input_images) > 14:
if len(input_images) > 14:
raise ValueError(f"输入图像数量 {len(input_images)} 超过限制 14 张")
batch_prompts = parse_batch_prompts(prompt)
@@ -667,6 +529,7 @@ class NanoBanana(io.ComfyNode):
images=input_images if input_images else None,
enable_grounding=enable_grounding,
thinking_level=thinking_level,
progress_callback=_make_progress_callback(pbar),
)
return loop.run_until_complete(_run_with_interrupt(_do()))
finally:
@@ -674,17 +537,14 @@ class NanoBanana(io.ComfyNode):
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(run_single)
generated_images, first_token_ms, download_ms = future.result(timeout=_REQUEST_TIMEOUT)
if pbar is not None:
pbar.update(1)
generated_images, task_ms, parse_ms = future.result(timeout=_REQUEST_TIMEOUT)
output_tensor = _images_to_tensor_safe(generated_images, _NODE)
elapsed = time.time() - start_time
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
ft_str = f"{first_token_ms/1000:.2f}s"
dl_str = f"{download_ms/1000:.2f}s"
print(f"完成!总耗时 {time_str} | 首字 {ft_str} | 下载 {dl_str} | 成功 {len(generated_images)}")
task_str = f"{task_ms/1000:.2f}s"
parse_str = f"{parse_ms/1000:.2f}s"
print(f"完成!总耗时 {time_str} | 异步任务 {task_str} | 解析 {parse_str} | 成功 {len(generated_images)}")
import gc; gc.collect()
return io.NodeOutput(output_tensor)
@@ -701,7 +561,7 @@ class NanoBanana(io.ComfyNode):
except RuntimeError as e:
raise RuntimeError(str(e)) from None
except Exception as e:
raise type(e)(str(e)) from None
raise RuntimeError(str(e)) from None
finally:
if not was_interrupted:
try:
+12 -17
View File
@@ -31,7 +31,7 @@ from PIL import Image
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
from ..utils.file_utils import load_images_from_folder, pair_images_by_name, pair_images_cartesian
from ..utils.config import get_api_key_or_raise, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
from ..utils.http_error import async_request_with_retry
from ..utils.http_error import async_request_with_retry, extract_structured_error_message, get_friendly_message
from ..models_config import (
get_enabled_async_models,
get_model_provider,
@@ -241,6 +241,10 @@ class NanoBananaV2:
@staticmethod
def _friendly_error(error_msg: str) -> str:
"""将上游错误转化为用户友好的提示"""
structured_message = extract_structured_error_message(error_msg)
if structured_message and structured_message != error_msg:
error_msg = structured_message
if "No available channel for model" in error_msg:
return (
"当前分组下模型不可用,请检查分组是否正确。"
@@ -343,16 +347,16 @@ class NanoBananaV2:
error_text = await response.text()
if not error_text.strip():
error_text = "(服务器未返回错误详情)"
raise RuntimeError(f"查询任务失败 ({response.status}): {error_text}")
raise RuntimeError(f"查询任务失败: {get_friendly_message(response.status, error_text)}")
result = await response.json()
status = provider.extract_status(result)
# 提取进度并回调(封顶 1.0 防止异常值导致进度条溢出)
if on_progress and status in ("SUBMITTED", "IN_PROGRESS"):
# 提取进度并回调;运行中状态不显示 100%,只有 SUCCESS 才补满。
if on_progress and status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
p = provider.extract_progress(result)
if p is not None:
p = min(p, 1.0)
p = min(p, 0.99)
if p > last_progress:
on_progress(p - last_progress)
last_progress = p
@@ -374,7 +378,7 @@ class NanoBananaV2:
print(f"{self.NODE_LABEL}: [轮询] FAILURE 但无错误信息,原始响应: {json.dumps(result, ensure_ascii=False)[:500]}")
friendly_msg = self._friendly_error(error_msg)
raise RuntimeError(f"任务失败: {friendly_msg}")
elif status in ("SUBMITTED", "IN_PROGRESS"):
elif status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
# 分段 sleep,每 0.1 秒检查一次取消信号
sleep_iterations = int(_POLL_INTERVAL / _INTERRUPT_CHECK_INTERVAL)
for _ in range(sleep_iterations):
@@ -406,10 +410,7 @@ class NanoBananaV2:
"error": None,
}
contributed = [0.0] # mutable container,追踪本任务已贡献的 pbar 进度
def _track_progress(delta):
contributed[0] += delta
if on_progress:
on_progress(delta)
@@ -435,15 +436,9 @@ class NanoBananaV2:
result["request_time"] = request_time
result["download_time"] = download_time
except InterruptProcessingException:
# 用户取消:补齐进度后向上传播,不吞掉
if contributed[0] < 1.0 and on_progress:
on_progress(1.0 - contributed[0])
raise
except Exception as e:
result["error"] = str(e) or f"{type(e).__name__}(无错误详情)"
# 失败也补齐 1.0 进度,保证进度条总数正确
if contributed[0] < 1.0 and on_progress:
on_progress(1.0 - contributed[0])
return result
@@ -704,7 +699,7 @@ class NanoBananaV2:
raise RuntimeError(str(e)) from None
except Exception as e:
raise type(e)(str(e)) from None
raise RuntimeError(str(e)) from None
finally:
# 查询并打印余额
@@ -1223,7 +1218,7 @@ class NanoBananaV2Batch(NanoBananaV2):
raise RuntimeError(str(e)) from None
except Exception as e:
raise type(e)(str(e)) from None
raise RuntimeError(str(e)) from None
finally:
try:
+254
View File
@@ -0,0 +1,254 @@
"""
Single-node new-api Veo 3.1 generator.
The node submits a /v1/videos task, waits for completion, downloads the mp4,
and returns ComfyUI's native VIDEO object for the built-in Save Video node.
"""
import os
from io import BytesIO
from typing import Optional, Tuple
from ..clients.newapi_veo_client import NewAPIVeoClient
from ..utils.image_utils import tensor_to_pil
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
try:
from comfy.utils import ProgressBar
PROGRESS_BAR_AVAILABLE = True
except ImportError:
ProgressBar = None
PROGRESS_BAR_AVAILABLE = False
try:
from comfy_api.input_impl import VideoFromFile
except Exception:
VideoFromFile = None
MODEL_OPTIONS = [
"veo-3.1",
]
DURATION_OPTIONS = ["4", "6", "8"]
ASPECT_RATIO_OPTIONS = ["16:9", "9:16"]
RESOLUTION_OPTIONS = ["720p", "1080p"]
TARGET_SIZE_MAP = {
("720p", "16:9"): (1280, 720),
("720p", "9:16"): (720, 1280),
("1080p", "16:9"): (1920, 1080),
("1080p", "9:16"): (1080, 1920),
}
def _get_output_dir() -> str:
if FOLDER_PATHS_AVAILABLE:
return folder_paths.get_output_directory()
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
comfy_root = os.path.dirname(os.path.dirname(plugin_dir))
return os.path.join(comfy_root, "output")
def _get_download_dir() -> str:
output_dir = _get_output_dir()
video_dir = os.path.join(output_dir, "newapi_veo")
os.makedirs(video_dir, exist_ok=True)
return video_dir
def _fit_image_to_target(image, target_size: Tuple[int, int]):
from PIL import Image as PILImage
target_w, target_h = target_size
src_w, src_h = image.size
src_ratio = src_w / src_h
target_ratio = target_w / target_h
if src_w == target_w and src_h == target_h:
return image
resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS
if src_ratio > target_ratio:
scale = target_h / src_h
new_w = round(src_w * scale)
image = image.resize((new_w, target_h), resample=resample)
left = max(0, (new_w - target_w) // 2)
image = image.crop((left, 0, left + target_w, target_h))
else:
scale = target_w / src_w
new_h = round(src_h * scale)
image = image.resize((target_w, new_h), resample=resample)
top = max(0, (new_h - target_h) // 2)
image = image.crop((0, top, target_w, top + target_h))
return image
def _image_to_png_bytes(image_tensor, resolution: str, aspect_ratio: str) -> Optional[bytes]:
if image_tensor is None:
return None
pil_images = tensor_to_pil(image_tensor)
if not pil_images:
return None
image = pil_images[0]
if image.mode != "RGB":
image = image.convert("RGB")
target_size = TARGET_SIZE_MAP.get((resolution, aspect_ratio))
if target_size is not None:
original_size = image.size
image = _fit_image_to_target(image, target_size)
if image.size != original_size:
print(
"NewAPI Veo: input image fitted "
f"{original_size[0]}x{original_size[1]} -> {image.size[0]}x{image.size[1]}"
)
buffer = BytesIO()
image.save(buffer, format="PNG")
image_bytes = buffer.getvalue()
print(
"NewAPI Veo: input_reference PNG "
f"{len(image_bytes) / 1024:.0f} KB ({image.size[0]}x{image.size[1]})"
)
return image_bytes
class Google31Video:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"提示词": (
"STRING",
{
"default": "A cinematic shot of a small robot walking through a rainy neon street.",
"multiline": True,
},
),
"负向提示词": ("STRING", {"default": "", "multiline": True}),
"网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}),
"模型": (MODEL_OPTIONS, {"default": MODEL_OPTIONS[0]}),
"时长": (DURATION_OPTIONS, {"default": "8"}),
"宽高比": (ASPECT_RATIO_OPTIONS, {"default": "16:9"}),
"分辨率": (RESOLUTION_OPTIONS, {"default": "1080p"}),
"生成音频": (["打开", "关闭"], {"default": "打开"}),
"seed": (
"INT",
{
"default": -1,
"min": -1,
"max": 0xFFFFFFFFFFFFFFFF,
"step": 1,
},
),
},
"optional": {
"参考图像": ("IMAGE",),
},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/Video"
DESCRIPTION = (
"Submit a new-api /v1/videos Veo 3.1 task, poll until complete, "
"download the mp4, and output native VIDEO for ComfyUI Save Video."
)
def generate(
self,
提示词: str,
负向提示词: str,
网络线路: str,
模型: str,
时长: str,
宽高比: str,
分辨率: str,
生成音频: str,
seed: int,
参考图像=None,
):
if VideoFromFile is None:
raise RuntimeError("当前 ComfyUI 版本不支持原生 VIDEO 输入实现 VideoFromFile。")
prompt = (提示词 or "").strip()
if not prompt:
raise ValueError("提示词不能为空。")
duration_value = int(时长)
if duration_value not in (4, 6, 8):
raise ValueError("时长仅支持 4、6、8。")
if 宽高比 not in ASPECT_RATIO_OPTIONS:
raise ValueError("宽高比仅支持 16:9 或 9:16。")
if 分辨率 not in RESOLUTION_OPTIONS:
raise ValueError("分辨率仅支持 720p 或 1080p。")
output_dir = _get_download_dir()
image_bytes = _image_to_png_bytes(参考图像, 分辨率, 宽高比)
pbar = ProgressBar(100) if PROGRESS_BAR_AVAILABLE else None
last_progress = [0]
last_status = [""]
def progress_callback(progress: int, status: str, elapsed: float):
if status != last_status[0]:
print(
"NewAPI Veo: polling "
f"status={status} | elapsed={elapsed:.0f}s"
)
last_status[0] = status
progress = max(0, min(100, int(progress or 0)))
if pbar is not None and progress > last_progress[0]:
pbar.update(progress - last_progress[0])
last_progress[0] = progress
client = NewAPIVeoClient(base_url=get_base_url_by_route(网络线路))
result = client.generate_video_sync(
prompt=prompt,
model=模型,
duration=duration_value,
aspect_ratio=宽高比,
resolution=分辨率,
output_dir=output_dir,
negative_prompt=负向提示词,
generate_audio=(生成音频 == "打开"),
image_bytes=image_bytes,
poll_interval=10,
timeout=900,
progress_callback=progress_callback,
)
video_path = result["video_path"]
video = VideoFromFile(video_path)
print(
"NewAPI Veo: completed "
f"| task_id={result['task_id']} | video={video_path}"
)
return (video,)
NODE_CLASS_MAPPINGS = {
"Google31Video": Google31Video,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Google31Video": "Google 3.1 Video",
}
+51 -9
View File
@@ -4,7 +4,9 @@ Seedance 视频生成节点
- Seedance: 文生视频 / 图生视频 / 首尾帧生视频(根据图片输入自动切换模式)
"""
import base64
import io
import json
import os
import tempfile
@@ -13,7 +15,7 @@ import torch
from ..clients.seedance_client import SeedanceClient
from ..clients.gemini_client import GeminiAPIClient
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64, pil_to_tensor
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
from ..utils.r2_uploader import upload_video, upload_audio
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
@@ -28,6 +30,9 @@ _MODELS = [
_RESOLUTIONS = ["720p", "1080p", "480p"]
_MAX_IMAGE_BYTES = 30 * 1024 * 1024
_MAX_REQUEST_BODY_BYTES = 64 * 1024 * 1024
# ── 模型能力判断 ──────────────────────────────────────────────────────────────
@@ -38,13 +43,45 @@ def _supports_camera_fixed(model: str) -> bool:
# ── 工具函数 ──────────────────────────────────────────────────────────────────
def _tensor_to_base64_url(tensor) -> str:
def _format_mb(size_bytes: int) -> str:
return f"{size_bytes / 1024 / 1024:.2f}MB"
def _tensor_to_base64_url(tensor, label: str = "图片") -> str:
"""ComfyUI IMAGE tensor → data:image/png;base64,xxx"""
pil_images = tensor_to_pil(tensor)
b64 = encode_image_to_base64(pil_images[0], format="PNG")
image = pil_images[0]
if image.mode == "RGBA":
image = image.convert("RGB")
buffered = io.BytesIO()
image.save(buffered, format="PNG")
image_bytes = buffered.getvalue()
image_size = len(image_bytes)
if image_size > _MAX_IMAGE_BYTES:
raise ValueError(
f"Seedance {label}大小 {_format_mb(image_size)} 超过单张图片 "
f"{_format_mb(_MAX_IMAGE_BYTES)} 限制,请先压缩或缩小图片。"
)
b64 = base64.b64encode(image_bytes).decode("utf-8")
return f"data:image/png;base64,{b64}"
def _validate_request_body_size(body: dict, tag: str):
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
if body_size > _MAX_REQUEST_BODY_BYTES:
raise ValueError(
f"{tag} 请求体大小 {_format_mb(body_size)} 超过 "
f"{_format_mb(_MAX_REQUEST_BODY_BYTES)} 限制,请减少参考图片数量或降低图片尺寸。"
)
print(
f"[{tag}] 请求体大小: {_format_mb(body_size)} "
f"(限制 {_format_mb(_MAX_REQUEST_BODY_BYTES)})"
)
async def _url_to_tensor(url: str) -> torch.Tensor:
"""从 URL 下载图片并转为 ComfyUI IMAGE tensor,失败时返回 None"""
@@ -201,7 +238,7 @@ class Seedance:
}
elif mode == "i2v":
first_url = _tensor_to_base64_url(first_image)
first_url = _tensor_to_base64_url(first_image, "首帧图片")
metadata["content"] = [
{
"type": "image_url",
@@ -218,8 +255,8 @@ class Seedance:
}
else: # flipflop
first_url = _tensor_to_base64_url(first_image)
last_url = _tensor_to_base64_url(last_image)
first_url = _tensor_to_base64_url(first_image, "首帧图片")
last_url = _tensor_to_base64_url(last_image, "尾帧图片")
metadata["content"] = [
{
"type": "image_url",
@@ -240,6 +277,8 @@ class Seedance:
"metadata": metadata,
}
_validate_request_body_size(body, tag)
# 保存路径(临时文件,避免与下游保存节点重复落盘)
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix=f"{file_prefix}_")
@@ -315,6 +354,7 @@ class SeedanceMultiModal:
web_search = _first(kwargs.get("联网搜索"), "关闭") == "打开"
return_last = _first(kwargs.get("返回末帧图片"), "关闭") == "打开"
seed = _first(kwargs.get("seed"), 0)
network_route = _first(kwargs.get("网络线路"), "全球加速")
# 参考图片:INPUT_IS_LIST 时是 [tensor, tensor, ...] 列表,直接保留
raw_images = kwargs.get("参考图片", None)
@@ -344,11 +384,11 @@ class SeedanceMultiModal:
imgs = ref_images[:9]
if len(ref_images) > 9:
print(f"[SeedanceMultiModal] 参考图片超过9张,仅取前9张(共{len(ref_images)}张)")
for img_tensor in imgs:
for idx, img_tensor in enumerate(imgs, start=1):
# 每个 tensor 可能是 [1,H,W,C] 或 [H,W,C],统一确保有 batch 维
if img_tensor.dim() == 3:
img_tensor = img_tensor.unsqueeze(0)
url = _tensor_to_base64_url(img_tensor)
url = _tensor_to_base64_url(img_tensor, f"参考图片{idx}")
content.append({
"type": "image_url",
"image_url": {"url": url},
@@ -414,11 +454,13 @@ class SeedanceMultiModal:
if first_image_url:
body["image"] = first_image_url
_validate_request_body_size(body, "Seedance多模态")
# ── 保存路径(临时文件,避免与下游保存节点重复落盘)──────────────────
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_")
client = SeedanceClient()
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
client.base_url = get_base_url_by_route(network_route)
pbar = _make_pbar()
on_stage, on_prog = _make_callbacks("Seedance多模态", pbar)