Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
427 lines
14 KiB
Python
427 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 ..utils.config import get_runtime_config_signature
|
||
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
|
||
self._client_config_signature = 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:
|
||
config_signature = get_runtime_config_signature()
|
||
if self.client is None or config_signature != self._client_config_signature:
|
||
self.client = VeoClient()
|
||
self._client_config_signature = config_signature
|
||
|
||
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",
|
||
}
|