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.
531 lines
20 KiB
Python
531 lines
20 KiB
Python
"""
|
||
Sora 视频生成节点
|
||
ComfyUI 自定义节点，调用 Sora API 生成视频
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import time
|
||
from math import gcd
|
||
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.sora_client import SoraClient
|
||
from ..models_config import (
|
||
get_enabled_sora_models,
|
||
get_all_sora_seconds,
|
||
get_all_sora_sizes,
|
||
get_sora_supported_seconds,
|
||
get_sora_supported_sizes,
|
||
get_sora_seconds_with_labels,
|
||
get_sora_sizes_with_labels,
|
||
SORA_MODELS,
|
||
)
|
||
|
||
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("⚠️ SoraVideo: comfy.utils.ProgressBar 不可用，将只使用终端进度显示")
|
||
|
||
|
||
def _size_to_display(size: str) -> str:
|
||
"""
|
||
将 'WxH' 格式的分辨率转换为友好显示名。
|
||
|
||
例如：
|
||
"720x1280" → "720P 9:16"
|
||
"1280x720" → "720P 16:9"
|
||
"1024x1792" → "1K 4:7"
|
||
"1792x1024" → "1K 7:4"
|
||
|
||
Args:
|
||
size: 分辨率字符串，格式 "WxH"
|
||
|
||
Returns:
|
||
友好显示名字符串
|
||
"""
|
||
parts = size.lower().split("x")
|
||
w, h = int(parts[0]), int(parts[1])
|
||
short_side = min(w, h)
|
||
if short_side >= 3840:
|
||
res = "4K"
|
||
elif short_side >= 1920:
|
||
res = "2K"
|
||
elif short_side >= 1080:
|
||
res = "1K"
|
||
elif short_side >= 720:
|
||
res = "720P"
|
||
elif short_side >= 480:
|
||
res = "480P"
|
||
else:
|
||
res = f"{short_side}P"
|
||
g = gcd(w, h)
|
||
ratio = f"{w // g}:{h // g}"
|
||
return f"{res} {ratio} ({size})"
|
||
|
||
|
||
def _build_size_display_map(sizes: list) -> dict:
|
||
"""
|
||
构建 显示名 → 实际值 映射字典。
|
||
|
||
Args:
|
||
sizes: 实际分辨率列表，如 ["720x1280", "1280x720"]
|
||
|
||
Returns:
|
||
字典，key 为显示名，value 为实际分辨率字符串
|
||
"""
|
||
mapping = {}
|
||
for size in sizes:
|
||
display = _size_to_display(size)
|
||
if display in mapping:
|
||
# 极少数情况下防止重名
|
||
display = f"{display} ({size})"
|
||
mapping[display] = size
|
||
return mapping
|
||
|
||
|
||
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):
|
||
"""
|
||
将参考图片按 "等比缩放覆盖 + 居中裁剪" 策略适配到目标分辨率。
|
||
|
||
策略 (Cover Crop)：
|
||
1. 比较图片宽高比和目标宽高比
|
||
2. 等比缩放，使图片最短边刚好覆盖目标对应边（图片完全覆盖目标区域）
|
||
3. 居中裁剪多余部分，得到精确目标尺寸
|
||
|
||
Args:
|
||
image: PIL Image 对象
|
||
target_size: 目标分辨率字符串，格式 "WxH"（如 "720x1280"）
|
||
|
||
Returns:
|
||
适配后的 PIL Image 对象
|
||
"""
|
||
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"Sora: 参考图片 {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
|
||
|
||
# Cover Crop: 缩放使图片完全覆盖目标区域，然后居中裁剪
|
||
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"Sora: 参考图片已适配为 {image.size[0]}x{image.size[1]}")
|
||
return image
|
||
|
||
|
||
def _compress_image_for_upload(
|
||
image,
|
||
target_size: Optional[str] = None,
|
||
) -> bytes:
|
||
"""
|
||
将 PIL Image 适配目标分辨率并编码为 PNG 字节，用于上传。
|
||
|
||
============================================================
|
||
⚠️ 已验证可用的标准做法，请勿随意修改以下编码逻辑！
|
||
============================================================
|
||
经过多轮调试（2026-02-28），以下参数组合为唯一验证成功的方案：
|
||
|
||
1. 图片格式：PNG（format="PNG"）
|
||
- 不可改为 JPEG —— API 会校验 Content-Type，抓包确认服务端使用 image/png
|
||
- 不可使用 base64 字符串 —— 会报 "expected a file, got a string"
|
||
- 不可使用 data URI —— 服务端不识别，返回 500
|
||
|
||
2. 图片尺寸：必须与视频分辨率完全一致（target_size）
|
||
- 不可缩放降采样 —— 会报 "Inpaint image must match the requested width and height"
|
||
- 尺寸由 _fit_image_to_target() 保证（等比缩放 + 居中裁剪）
|
||
|
||
3. 上传方式：由调用方（sora_client.py）以 multipart/form-data 文件字段上传
|
||
- filename="reference.png", content_type="image/png"
|
||
- 不可改回 application/json —— 服务端校验 input_reference 必须为 file 类型
|
||
============================================================
|
||
|
||
Args:
|
||
image: PIL Image 对象
|
||
target_size: 目标分辨率字符串 "WxH"（如 "720x1280"）
|
||
|
||
Returns:
|
||
PNG 格式的二进制字节
|
||
"""
|
||
from io import BytesIO
|
||
|
||
# 统一转换为 RGB（去除透明通道及其他模式）
|
||
if image.mode != "RGB":
|
||
image = image.convert("RGB")
|
||
|
||
# 适配到目标分辨率（等比缩放 + 居中裁剪）
|
||
# ⚠️ 必须保持此尺寸不变，API 强制要求参考图片与视频分辨率完全一致
|
||
if target_size:
|
||
image = _fit_image_to_target(image, target_size)
|
||
|
||
# ⚠️ 必须使用 PNG 格式，不可改为 JPEG 或其他格式
|
||
buffered = BytesIO()
|
||
image.save(buffered, format="PNG")
|
||
size_kb = buffered.tell() / 1024
|
||
print(f"Sora: 参考图片编码为 PNG，{size_kb:.0f} KB ({image.size[0]}x{image.size[1]})")
|
||
return buffered.getvalue()
|
||
|
||
|
||
class SoraVideo:
|
||
"""
|
||
Sora 视频生成节点
|
||
|
||
功能：
|
||
- 文生视频：基于提示词生成视频
|
||
- 图生视频：基于参考图片和提示词生成视频
|
||
- 异步轮询：自动等待生成完成并下载
|
||
"""
|
||
|
||
def __init__(self):
|
||
self.client = None
|
||
self._client_config_signature = None
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP
|
||
|
||
enabled_models = get_enabled_sora_models()
|
||
if not enabled_models:
|
||
enabled_models = ["请在 models_config.py 中启用至少一个 Sora 模型"]
|
||
|
||
# 构建秒数选项列表（按数字顺序排序）
|
||
# 格式: ["4", "8", "10", "12", "15", "25(pro)"]
|
||
all_seconds_display = []
|
||
seen_seconds = set()
|
||
for model_id in enabled_models:
|
||
supported = get_sora_supported_seconds(model_id)
|
||
for s in supported:
|
||
if s not in seen_seconds:
|
||
seen_seconds.add(s)
|
||
display = SECONDS_DISPLAY_MAP.get(s, str(s))
|
||
all_seconds_display.append((s, display))
|
||
# 按秒数数值排序
|
||
all_seconds_display = sorted(all_seconds_display, key=lambda x: x[0])
|
||
seconds_options = [d for _, d in all_seconds_display] if all_seconds_display else ["4", "8", "12"]
|
||
|
||
# 构建分辨率选项列表（去重）
|
||
# 格式: ["720P", "1080P"]
|
||
seen_resolutions = set()
|
||
for model_id in enabled_models:
|
||
supported = get_sora_supported_sizes(model_id)
|
||
for size in supported:
|
||
if size in RESOLUTION_DISPLAY_MAP:
|
||
res_name, _ = RESOLUTION_DISPLAY_MAP[size]
|
||
seen_resolutions.add(res_name)
|
||
resolution_options = sorted(list(seen_resolutions)) if seen_resolutions else ["720P"]
|
||
|
||
return {
|
||
"required": {
|
||
"prompt": ("STRING", {
|
||
"default": "A calico cat playing a piano on stage",
|
||
"multiline": True,
|
||
}),
|
||
"模型": (enabled_models, {
|
||
"default": enabled_models[0],
|
||
}),
|
||
"分辨率": (resolution_options, {
|
||
"default": resolution_options[0] if resolution_options else "720P",
|
||
}),
|
||
"宽高比": (["竖屏", "横屏"], {
|
||
"default": "竖屏",
|
||
}),
|
||
"视频时长": (seconds_options, {
|
||
"default": seconds_options[0] if seconds_options else "4",
|
||
}),
|
||
"生成数量": ("INT", {
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 10,
|
||
"step": 1,
|
||
}),
|
||
"seed": ("INT", {
|
||
"default": 0,
|
||
"min": 0,
|
||
"max": 0xffffffffffffffff
|
||
}),
|
||
},
|
||
"optional": {
|
||
"参考图片": ("IMAGE",),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("预览视频",)
|
||
FUNCTION = "generate_video"
|
||
CATEGORY = "video/generation"
|
||
|
||
DESCRIPTION = (
|
||
"Sora 视频生成节点。\n"
|
||
"支持文生视频和图生视频，自动轮询任务状态并下载视频。\n"
|
||
"视频保存到 ComfyUI/output/video/ 目录。\n\n"
|
||
"【模型说明】\n"
|
||
"• sora-2：官方模型，支持 4/8/12秒、720P 分辨率\n"
|
||
"• sora-2-pro：增强模型，支持全时长（含25秒）、1080P 分辨率\n\n"
|
||
"【时长说明】\n"
|
||
"• 25(pro)：仅 sora-2-pro 支持的25秒时长\n\n"
|
||
"【分辨率说明】\n"
|
||
"• 720P：sora-2 和 sora-2-pro 均支持\n"
|
||
"• 1080P：仅 sora-2-pro 支持的高清分辨率"
|
||
)
|
||
|
||
def generate_video(
|
||
self,
|
||
prompt: str,
|
||
模型: str,
|
||
**kwargs,
|
||
) -> Tuple[str]:
|
||
from ..models_config import SECONDS_DISPLAY_MAP, RESOLUTION_DISPLAY_MAP
|
||
|
||
视频时长_display = kwargs.pop("视频时长", "4")
|
||
分辨率_display = kwargs.pop("分辨率", "720P")
|
||
宽高比 = kwargs.pop("宽高比", "竖屏")
|
||
生成数量 = kwargs.pop("生成数量", 1)
|
||
seed = kwargs.pop("seed", 0)
|
||
start_time = time.time()
|
||
|
||
# 解析秒数显示值（如 "25(pro)" → 25）
|
||
seconds = 4 # 默认
|
||
for actual, display in SECONDS_DISPLAY_MAP.items():
|
||
if display == 视频时长_display:
|
||
seconds = actual
|
||
break
|
||
# 如果找不到映射，尝试直接解析数字
|
||
if seconds == 4 and 视频时长_display != "4":
|
||
try:
|
||
seconds = int(视频时长_display.replace("(pro)", ""))
|
||
except ValueError:
|
||
seconds = 4
|
||
|
||
# 根据分辨率和宽高比确定实际分辨率值
|
||
分辨率 = "720x1280" # 默认
|
||
for actual, (res_name, orientation) in RESOLUTION_DISPLAY_MAP.items():
|
||
if res_name == 分辨率_display and orientation == 宽高比:
|
||
分辨率 = actual
|
||
break
|
||
|
||
# 检查参考图片
|
||
ref_image = kwargs.get("参考图片")
|
||
ref_image_bytes = None
|
||
if ref_image is not None:
|
||
pil_images = tensor_to_pil(ref_image)
|
||
if pil_images:
|
||
ref_image_bytes = _compress_image_for_upload(pil_images[0], target_size=分辨率)
|
||
|
||
mode_str = "图生视频 (含参考图)" if ref_image_bytes else "文生视频"
|
||
# 获取用户友好的显示值用于日志
|
||
seconds_display = SECONDS_DISPLAY_MAP.get(seconds, str(seconds))
|
||
res_display = f"{分辨率_display} {宽高比}"
|
||
if 生成数量 > 1:
|
||
print(f"Sora: {mode_str} | 并发{生成数量}个 | {模型} | {seconds_display} | {res_display}")
|
||
else:
|
||
print(f"Sora: {mode_str} | {模型} | {seconds_display} | {res_display}")
|
||
|
||
# 校验参数兼容性
|
||
supported_seconds = get_sora_supported_seconds(模型)
|
||
if supported_seconds and seconds not in supported_seconds:
|
||
# 构建带标签的支持时长列表
|
||
supported_labels = []
|
||
for s in supported_seconds:
|
||
display = SECONDS_DISPLAY_MAP.get(s, str(s))
|
||
supported_labels.append(display)
|
||
raise ValueError(
|
||
f"时长 {SECONDS_DISPLAY_MAP.get(seconds, str(seconds))} 与模型 \"{模型}\" 不兼容！\n"
|
||
f"该模型支持的时长: {', '.join(supported_labels)}"
|
||
)
|
||
|
||
supported_sizes = get_sora_supported_sizes(模型)
|
||
if supported_sizes and 分辨率 not in supported_sizes:
|
||
# 检查该分辨率是否为Pro独占
|
||
pro_only_sizes = ["1024x1792", "1792x1024"]
|
||
_, orientation = RESOLUTION_DISPLAY_MAP.get(分辨率, (分辨率, ""))
|
||
extra_hint = f"\n提示：1080P {orientation} 为 sora-2-pro 独占，请切换模型或选择720P。" if 分辨率 in pro_only_sizes else ""
|
||
raise ValueError(
|
||
f"分辨率 \"{分辨率_display} {宽高比}\" 与模型 \"{模型}\" 不兼容！"
|
||
f"支持的分辨率: {', '.join(supported_sizes)}" + extra_hint
|
||
)
|
||
|
||
# 准备保存路径
|
||
video_dir = _get_video_output_dir()
|
||
counter = _get_next_counter(video_dir, "sora")
|
||
|
||
# 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 = SoraClient()
|
||
self._client_config_signature = config_signature
|
||
|
||
if 生成数量 == 1:
|
||
# ── 单个视频：保留详细进度（提交→轮询→下载）
|
||
save_path = os.path.join(video_dir, f"sora_{counter:05d}.mp4")
|
||
last_progress = [0]
|
||
|
||
def progress_callback(progress_pct: int):
|
||
print(
|
||
f"\rSora: 生成中... 进度: {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("Sora: 正在提交视频生成任务...")
|
||
elif stage.startswith("submitted:"):
|
||
vid = stage.split(":", 1)[1]
|
||
print(f"Sora: 视频任务已提交，ID: {vid}")
|
||
elif stage == "polling":
|
||
print("Sora: 等待视频生成...")
|
||
elif stage == "downloading":
|
||
print("") # 换行（结束 \r 行）
|
||
print("Sora: 视频生成完成，正在下载...")
|
||
|
||
result_path = self.client.generate_video_sync(
|
||
prompt=prompt,
|
||
model=模型,
|
||
seconds=seconds,
|
||
size=分辨率,
|
||
save_path=save_path,
|
||
input_reference_bytes=ref_image_bytes,
|
||
seed=seed,
|
||
progress_callback=progress_callback,
|
||
on_stage=on_stage,
|
||
)
|
||
result_paths = [result_path]
|
||
|
||
else:
|
||
# ── 批量并发：同时提交多个任务
|
||
save_paths = [
|
||
os.path.join(video_dir, f"sora_{counter + i:05d}.mp4")
|
||
for i in range(生成数量)
|
||
]
|
||
success_count = [0]
|
||
fail_count = [0]
|
||
|
||
def batch_progress_callback(current: int, total: int, success: bool, error_msg):
|
||
if success:
|
||
success_count[0] += 1
|
||
print(f"Sora: 第 {current}/{total} 个视频完成 ✓")
|
||
else:
|
||
fail_count[0] += 1
|
||
print(f"Sora: 第 {current}/{total} 个视频失败 ✗")
|
||
if error_msg:
|
||
print(f"原始错误详情:\n{error_msg}")
|
||
if pbar is not None:
|
||
pbar.update(1)
|
||
|
||
print(f"Sora: 正在并发提交 {生成数量} 个视频任务，请耐心等待...")
|
||
result_paths = self.client.generate_batch_videos_sync(
|
||
prompt=prompt,
|
||
model=模型,
|
||
seconds=seconds,
|
||
size=分辨率,
|
||
save_paths=save_paths,
|
||
input_reference_bytes=ref_image_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"Sora: 完成！总耗时 {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"\nSora: ❌ {error_msg}")
|
||
raise ValueError(error_msg) from None
|
||
|
||
except RuntimeError as e:
|
||
error_msg = str(e)
|
||
print(f"\nSora: ❌ {error_msg}")
|
||
raise RuntimeError(error_msg) from None
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
print(f"\nSora: ❌ {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"Sora: {balance_info}")
|
||
except Exception:
|
||
pass
|