o1keyandClaude Sonnet 4.5
9ab209b2b7
Complete rewrite/sync of comfyui_o1key custom nodes. Treat this commit as the current canonical version. Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
423 lines
14 KiB
Python
423 lines
14 KiB
Python
"""
|
||
Google Veo 视频生成节点
|
||
ComfyUI 自定义节点，调用 Veo API 生成视频
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import time
|
||
from typing import Optional, Tuple
|
||
|
||
import torch
|
||
|
||
from ..utils.image_utils import tensor_to_pil
|
||
from ..clients.veo_client import VeoClient
|
||
from ..models_config import (
|
||
get_enabled_veo_models,
|
||
VEO_MODELS,
|
||
VEO_RESOLUTION_MAP,
|
||
)
|
||
|
||
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:
|
||
PROGRESS_BAR_AVAILABLE = False
|
||
print("⚠️ GoogleVeo: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||
|
||
|
||
def _get_video_output_dir() -> str:
|
||
"""获取视频输出目录: ComfyUI/output/video"""
|
||
if FOLDER_PATHS_AVAILABLE:
|
||
base = folder_paths.get_output_directory()
|
||
else:
|
||
plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
base = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "output")
|
||
video_dir = os.path.join(base, "video")
|
||
os.makedirs(video_dir, exist_ok=True)
|
||
return video_dir
|
||
|
||
|
||
def _get_next_counter(directory: str, prefix: str) -> int:
|
||
"""扫描目录，获取下一个可用的文件计数器"""
|
||
if not os.path.exists(directory):
|
||
return 1
|
||
pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)")
|
||
max_counter = 0
|
||
for f in os.listdir(directory):
|
||
m = pattern.match(f)
|
||
if m:
|
||
max_counter = max(max_counter, int(m.group(1)))
|
||
return max_counter + 1
|
||
|
||
|
||
def _fit_image_to_target(image, target_size: str):
|
||
"""
|
||
将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。
|
||
"""
|
||
from PIL import Image as PILImage
|
||
|
||
parts = target_size.lower().split("x")
|
||
target_w, target_h = int(parts[0]), int(parts[1])
|
||
|
||
src_w, src_h = image.size
|
||
src_ratio = src_w / src_h
|
||
target_ratio = target_w / target_h
|
||
|
||
if abs(src_ratio - target_ratio) < 0.01 and src_w <= target_w and src_h <= target_h:
|
||
return image
|
||
|
||
print(f"Veo: 参考图片 {src_w}x{src_h} (比例 {src_ratio:.2f}) → 目标 {target_w}x{target_h} (比例 {target_ratio:.2f})")
|
||
|
||
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)
|
||
new_h = target_h
|
||
image = image.resize((new_w, new_h), resample=resample)
|
||
left = (new_w - target_w) // 2
|
||
image = image.crop((left, 0, left + target_w, target_h))
|
||
else:
|
||
scale = target_w / src_w
|
||
new_w = target_w
|
||
new_h = round(src_h * scale)
|
||
image = image.resize((new_w, new_h), resample=resample)
|
||
top = (new_h - target_h) // 2
|
||
image = image.crop((0, top, target_w, top + target_h))
|
||
|
||
print(f"Veo: 参考图片已适配为 {image.size[0]}x{image.size[1]}")
|
||
return image
|
||
|
||
|
||
def _compress_image_to_bytes(image, target_size: Optional[str] = None) -> bytes:
|
||
"""
|
||
将 PIL Image 适配目标分辨率并编码为 PNG 字节
|
||
"""
|
||
from io import BytesIO
|
||
|
||
if image.mode != "RGB":
|
||
image = image.convert("RGB")
|
||
|
||
if target_size:
|
||
image = _fit_image_to_target(image, target_size)
|
||
|
||
buffered = BytesIO()
|
||
image.save(buffered, format="PNG")
|
||
size_kb = buffered.tell() / 1024
|
||
print(f"Veo: 参考图片编码为 PNG，{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})")
|
||
return buffered.getvalue()
|
||
|
||
|
||
class GoogleVeo:
|
||
"""
|
||
Google Veo 视频生成节点
|
||
|
||
功能：
|
||
- 文生视频：基于提示词生成视频
|
||
- 图生视频：基于首帧/尾帧/参考图生成视频
|
||
- 异步轮询：自动等待生成完成并下载
|
||
"""
|
||
|
||
def __init__(self):
|
||
self.client = None
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
enabled_models = get_enabled_veo_models()
|
||
if not enabled_models:
|
||
enabled_models = ["请在 models_config.py 中启用 Veo 模型"]
|
||
|
||
# 分辨率选项
|
||
resolution_options = ["720p", "1080p", "4K"]
|
||
|
||
# 宽高比选项
|
||
aspect_ratio_options = ["16:9", "9:16"]
|
||
|
||
# 视频秒数选项
|
||
seconds_options = ["4", "6", "8"]
|
||
|
||
return {
|
||
"required": {
|
||
"prompt": ("STRING", {
|
||
"default": "A calico cat playing a piano on stage",
|
||
"multiline": True,
|
||
}),
|
||
"模型": (enabled_models, {
|
||
"default": enabled_models[0] if enabled_models else "Veo3.1",
|
||
}),
|
||
"分辨率": (resolution_options, {
|
||
"default": "720p",
|
||
}),
|
||
"宽高比": (aspect_ratio_options, {
|
||
"default": "9:16",
|
||
}),
|
||
"视频时长": (seconds_options, {
|
||
"default": "8",
|
||
}),
|
||
"seed": ("INT", {
|
||
"default": 0,
|
||
"min": 0,
|
||
"max": 0xffffffffffffffff,
|
||
}),
|
||
"生成数量": ("INT", {
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 10,
|
||
"step": 1,
|
||
}),
|
||
},
|
||
"optional": {
|
||
"首帧": ("IMAGE",),
|
||
"尾帧": ("IMAGE",),
|
||
"参考图": ("IMAGE",),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("预览视频",)
|
||
FUNCTION = "generate_video"
|
||
CATEGORY = "video/generation"
|
||
|
||
DESCRIPTION = (
|
||
"Google Veo 视频生成节点。\n"
|
||
"支持文生视频和图生视频（图生视频支持首帧、尾帧、参考图）。\n"
|
||
"视频保存到 ComfyUI/output/video/ 目录。\n\n"
|
||
"【模型说明】\n"
|
||
"• Veo3.1：Google 最新视频生成模型\n\n"
|
||
"【分辨率说明】\n"
|
||
"• 720p：标清\n"
|
||
"• 1080p：高清\n"
|
||
"• 4K：超高清\n\n"
|
||
"【时长说明】\n"
|
||
"• 4秒：短视频\n"
|
||
"• 6秒：标准\n"
|
||
"• 8秒：长视频（默认）\n\n"
|
||
"【图生视频说明】\n"
|
||
"• 首帧：视频开始的第一帧图像\n"
|
||
"• 尾帧：视频结束时的最后一帧图像\n"
|
||
"• 参考图：参考图像（与首帧/尾帧配合使用）\n"
|
||
"• 至少需要提供首帧或参考图之一"
|
||
)
|
||
|
||
def generate_video(
|
||
self,
|
||
prompt: str,
|
||
模型: str,
|
||
**kwargs,
|
||
) -> Tuple[str]:
|
||
分辨率 = kwargs.pop("分辨率", "720p")
|
||
宽高比 = kwargs.pop("宽高比", "9:16")
|
||
视频时长 = kwargs.pop("视频时长", "8")
|
||
seed = kwargs.pop("seed", 0)
|
||
生成数量 = kwargs.pop("生成数量", 1)
|
||
start_time = time.time()
|
||
|
||
# 解析视频时长
|
||
seconds = int(视频时长)
|
||
|
||
# 解析分辨率和宽高比，映射到模型名称
|
||
size_key = f"{分辨率}_{宽高比}"
|
||
actual_size = VEO_RESOLUTION_MAP.get(size_key)
|
||
if not actual_size:
|
||
# 默认值
|
||
actual_size = "720x1280" # 720p 9:16
|
||
|
||
# 检查是否有参考图输入
|
||
首帧 = kwargs.get("首帧")
|
||
尾帧 = kwargs.get("尾帧")
|
||
参考图 = kwargs.get("参考图")
|
||
|
||
has_image = 首帧 is not None or 尾帧 is not None or 参考图 is not None
|
||
|
||
# 根据是否有图片选择模型前缀
|
||
if has_image:
|
||
model_prefix = "veo3.1"
|
||
else:
|
||
model_prefix = "veo3.1"
|
||
|
||
# 构建完整模型名称
|
||
# 格式: veo3.1-portrait / veo3.1-landscape / veo3.1-portrait-fl / veo3.1-landscape-fl 等
|
||
if 分辨率 == "720p":
|
||
res_suffix = ""
|
||
if 宽高比 == "9:16":
|
||
orientation = "portrait"
|
||
else:
|
||
orientation = "landscape"
|
||
elif 分辨率 == "1080p":
|
||
res_suffix = "-hd"
|
||
if 宽高比 == "9:16":
|
||
orientation = "portrait"
|
||
else:
|
||
orientation = "landscape"
|
||
else: # 4K
|
||
res_suffix = "-4k"
|
||
if 宽高比 == "9:16":
|
||
orientation = "portrait"
|
||
else:
|
||
orientation = "landscape"
|
||
|
||
# 图生视频添加 -fl 后缀
|
||
if has_image:
|
||
model_suffix = f"-{orientation}-fl{res_suffix}"
|
||
else:
|
||
model_suffix = f"-{orientation}{res_suffix}"
|
||
|
||
model = f"{model_prefix}{model_suffix}"
|
||
|
||
# 准备图片字节
|
||
first_frame_bytes = None
|
||
last_frame_bytes = None
|
||
reference_bytes = None
|
||
|
||
if 首帧 is not None:
|
||
pil_images = tensor_to_pil(首帧)
|
||
if pil_images:
|
||
first_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||
|
||
if 尾帧 is not None:
|
||
pil_images = tensor_to_pil(尾帧)
|
||
if pil_images:
|
||
last_frame_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||
|
||
if 参考图 is not None:
|
||
pil_images = tensor_to_pil(参考图)
|
||
if pil_images:
|
||
reference_bytes = _compress_image_to_bytes(pil_images[0], target_size=actual_size)
|
||
|
||
mode_str = "图生视频" if has_image else "文生视频"
|
||
print(f"Veo: {mode_str} | 并发{生成数量}个 | 模型: {model} | {seconds}秒 | {分辨率} {宽高比}")
|
||
|
||
# 准备保存路径
|
||
video_dir = _get_video_output_dir()
|
||
counter = _get_next_counter(video_dir, "veo")
|
||
|
||
# ProgressBar
|
||
pbar = None
|
||
if PROGRESS_BAR_AVAILABLE:
|
||
pbar = ProgressBar(生成数量 if 生成数量 > 1 else 100)
|
||
|
||
try:
|
||
if self.client is None:
|
||
self.client = VeoClient()
|
||
|
||
if 生成数量 == 1:
|
||
save_path = os.path.join(video_dir, f"veo_{counter:05d}.mp4")
|
||
last_progress = [0]
|
||
|
||
def progress_callback(progress_pct: int):
|
||
print(
|
||
f"\rVeo: 生成中... 进度: {progress_pct}%",
|
||
end="", flush=True
|
||
)
|
||
if pbar is not None and progress_pct > last_progress[0]:
|
||
pbar.update(progress_pct - last_progress[0])
|
||
last_progress[0] = progress_pct
|
||
|
||
def on_stage(stage: str):
|
||
if stage == "submitting":
|
||
print("Veo: 正在提交视频生成任务...")
|
||
elif stage.startswith("submitted:"):
|
||
vid = stage.split(":", 1)[1]
|
||
print(f"Veo: 视频任务已提交，ID: {vid}")
|
||
elif stage == "polling":
|
||
print("Veo: 等待视频生成...")
|
||
elif stage == "downloading":
|
||
print("")
|
||
print("Veo: 视频生成完成，正在下载...")
|
||
|
||
result_path = self.client.generate_video_sync(
|
||
prompt=prompt,
|
||
model=model,
|
||
seconds=seconds,
|
||
size=actual_size,
|
||
save_path=save_path,
|
||
first_frame_bytes=first_frame_bytes,
|
||
last_frame_bytes=last_frame_bytes,
|
||
reference_bytes=reference_bytes,
|
||
seed=seed,
|
||
progress_callback=progress_callback,
|
||
on_stage=on_stage,
|
||
)
|
||
result_paths = [result_path]
|
||
|
||
else:
|
||
save_paths = [
|
||
os.path.join(video_dir, f"veo_{counter + i:05d}.mp4")
|
||
for i in range(生成数量)
|
||
]
|
||
success_count = [0]
|
||
|
||
def batch_progress_callback(current: int, total: int, success: bool, error_msg):
|
||
if success:
|
||
success_count[0] += 1
|
||
print(f"Veo: 第 {current}/{total} 个视频完成 ✓")
|
||
else:
|
||
print(f"Veo: 第 {current}/{total} 个视频失败 ✗")
|
||
if error_msg:
|
||
print(f"原始错误详情:\n{error_msg}")
|
||
if pbar is not None:
|
||
pbar.update(1)
|
||
|
||
print(f"Veo: 正在并发提交 {生成数量} 个视频任务，请耐心等待...")
|
||
result_paths = self.client.generate_batch_videos_sync(
|
||
prompt=prompt,
|
||
model=model,
|
||
seconds=seconds,
|
||
size=actual_size,
|
||
save_paths=save_paths,
|
||
first_frame_bytes=first_frame_bytes,
|
||
last_frame_bytes=last_frame_bytes,
|
||
reference_bytes=reference_bytes,
|
||
seed=seed,
|
||
progress_callback=batch_progress_callback,
|
||
)
|
||
|
||
elapsed = time.time() - start_time
|
||
time_str = f"{elapsed:.2f}s" if elapsed >= 1 else f"{elapsed:.3f}s"
|
||
print(f"Veo: 完成！总耗时 {time_str} | 已生成 {len(result_paths)} 个视频")
|
||
for p in result_paths:
|
||
print(f" → {p}")
|
||
|
||
output_path = "\n".join(result_paths)
|
||
return (output_path,)
|
||
|
||
except ValueError as e:
|
||
error_msg = str(e)
|
||
print(f"\nVeo: ❌ {error_msg}")
|
||
raise ValueError(error_msg) from None
|
||
|
||
except RuntimeError as e:
|
||
error_msg = str(e)
|
||
print(f"\nVeo: ❌ {error_msg}")
|
||
raise RuntimeError(error_msg) from None
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
print(f"\nVeo: ❌ {error_msg}")
|
||
raise type(e)(error_msg) from None
|
||
|
||
finally:
|
||
if self.client is not None:
|
||
try:
|
||
balance_data = self.client.query_balance_sync()
|
||
balance_info = self.client.format_balance_info(balance_data)
|
||
print(f"Veo: {balance_info}")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"GoogleVeo": GoogleVeo,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"GoogleVeo": "Google Veo - ab",
|
||
}
|