Publish current ComfyUI O1Key code baseline

Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
2026-09-24 19:56:48 +08:00
parent 3e337722ab
commit ba920f2b66
183 changed files with 49496 additions and 9934 deletions
+227 -199
View File
@@ -1,30 +1,25 @@
"""
Grok Video node.
"""Lean ComfyUI nodes for O1Key Grok Imagine Video."""
Submits a /v1/videos task, polls until completion, downloads the mp4,
and returns ComfyUI's native VIDEO object.
"""
import json
import asyncio
import math
import os
from typing import List, Optional
import re
from typing import Dict
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
from ..utils.config import get_base_url_by_route
from ..utils.image_utils import tensor_to_pil
from ..utils.r2_uploader import upload_audio, upload_image, upload_video
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
folder_paths = None
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
@@ -36,135 +31,94 @@ 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
MODEL_OPTIONS = list(GrokVideoClient.MODEL_OPTIONS)
ASPECT_RATIO_OPTIONS = list(GrokVideoClient.ASPECT_RATIO_OPTIONS)
RESOLUTION_OPTIONS = list(GrokVideoClient.RESOLUTION_OPTIONS)
GENERATION_MODE_OPTIONS = ["文生视频", "图生视频", "参考生视频"]
EDIT_MODE_OPTIONS = ["编辑视频", "续写视频"]
IMAGE_INPUT_NAMES = [f"图片{i}" for i in range(1, 8)]
AUDIO_INPUT_NAMES = ["音频素材", "音频素材2", "音频素材3"]
def _get_output_dir() -> str:
if FOLDER_PATHS_AVAILABLE:
base = folder_paths.get_output_directory()
if folder_paths is not None:
base_dir = folder_paths.get_temp_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")
base_dir = os.path.join(os.path.dirname(os.path.dirname(plugin_dir)), "temp")
output_dir = os.path.join(base_dir, "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):
def _single_pil_image(image_tensor, input_name: str):
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
images = tensor_to_pil(image_tensor)
if len(images) != 1:
raise ValueError(f"{input_name} 只能连接 1 张图片,请拆分批次后再连接。")
return images[0].convert("RGB")
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 _parse_voice_ids(value: object) -> list[str]:
voice_ids = [item.strip() for item in re.split(r"[,\n]", str(value or "")) if item.strip()]
if len(voice_ids) > 3:
raise ValueError("参考音色 ID 最多填写 3 个。")
return voice_ids
def _to_data_urls(encoded_images) -> List[str]:
return [f"data:{mime};base64,{b64}" for mime, b64 in encoded_images]
def _video_duration_seconds(video) -> float:
getter = getattr(video, "get_duration", None)
if not callable(getter):
raise ValueError("无法读取输入视频时长;请连接 ComfyUI 原生 VIDEO 输出。")
try:
duration = float(getter())
except Exception as exc:
raise ValueError("无法读取输入视频时长;请确认视频文件可以正常解码。") from exc
if not math.isfinite(duration) or duration <= 0:
raise ValueError("输入视频时长无效;请确认视频文件可以正常解码。")
return duration
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 _progress_callback():
progress_bar = ProgressBar(100) if ProgressBar is not None else None
progress_value = [0]
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),
)
def callback(progress: int, _status: str, _elapsed: float) -> None:
current = max(0, min(100, int(progress or 0)))
if progress_bar is not None and current > progress_value[0]:
progress_bar.update(current - progress_value[0])
progress_value[0] = current
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
return progress_bar, progress_value, callback
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)} 限制,请减少参考图片或降低图片尺寸。"
)
def _finish_video(result: Dict[str, object], progress_bar, progress_value):
if progress_bar is not None and progress_value[0] < 100:
progress_bar.update(100 - progress_value[0])
video_path = result["video_path"]
print(f"Grok Video:下载完成:{video_path}")
return (VideoFromFile(video_path),)
class O1keyGrokVideo:
"""文生、图生或参考素材生 Grok 视频。素材会自动上传为 URL。"""
@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]}),
"生成模式": (GENERATION_MODE_OPTIONS, {"default": "文生视频"}),
"提示词": ("STRING", {"default": "", "multiline": True}),
"模型": (MODEL_OPTIONS, {"default": GrokVideoClient.DEFAULT_MODEL}),
"时长(秒)": ("INT", {"default": 8, "min": 1, "max": 15, "step": 1}),
"宽高比": (ASPECT_RATIO_OPTIONS, {"default": "16:9"}),
"秒数(按模型限制)": (
"INT",
{
"default": 5,
"min": 5,
"max": 20,
"step": 1,
"display": "number",
},
),
"画质": (QUALITY_OPTIONS, {"default": "720p"}),
"分辨率": (RESOLUTION_OPTIONS, {"default": "480p"}),
"参考音色ID(逗号分隔)": ("STRING", {"default": ""}),
},
"optional": {
"参考图1": ("IMAGE",),
"参考图2": ("IMAGE",),
"参考图3": ("IMAGE",),
**{input_name: ("IMAGE",) for input_name in IMAGE_INPUT_NAMES},
**{input_name: ("AUDIO",) for input_name in AUDIO_INPUT_NAMES},
},
}
@@ -172,112 +126,186 @@ class O1keyGrokVideo:
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."
"支持文生、图生和多参考素材生成。图生视频只连接图片 1;参考生视频最多使用 7 张图和 "
"3 个参考音频(AUDIO 或 voice_id 合计)。Grok 1.5 的文生/图生可选 1080p,多参考最高 720p。"
)
def generate(
self,
**kwargs,
):
def generate(self, **kwargs):
if VideoFromFile is None:
raise RuntimeError("当前 ComfyUI 版本不支持原生 VIDEO 输入实现 VideoFromFile。")
raise RuntimeError("当前 ComfyUI 版本不支持 VideoFromFile,无法输出 VIDEO")
提示词 = 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")
mode = kwargs.get("生成模式", "文生视频")
model = kwargs.get("模型", GrokVideoClient.DEFAULT_MODEL)
prompt = (kwargs.get("提示词") or "").strip()
connected_images = [
(input_name, image)
for input_name in IMAGE_INPUT_NAMES
if (image := _single_pil_image(kwargs.get(input_name), input_name)) is not None
]
images = [image for _, image in connected_images]
audios = [kwargs.get(input_name) for input_name in AUDIO_INPUT_NAMES if kwargs.get(input_name) is not None]
voice_ids = _parse_voice_ids(kwargs.get("参考音色ID(逗号分隔)"))
duration = kwargs.get("时长(秒)", 8)
aspect_ratio = kwargs.get("宽高比", "16:9")
resolution = kwargs.get("分辨率", "480p")
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。")
if mode not in GENERATION_MODE_OPTIONS:
raise ValueError(f"不支持的生成模式:{mode}")
if len(audios) + len(voice_ids) > 3:
raise ValueError("参考音频与参考音色 ID 合计最多 3 个。")
quality = QUALITY_VALUE_MAP[画质]
reference_images = _collect_reference_images(**kwargs)
image_data_urls = _encode_image_data_urls(
reference_images,
if mode == "文生视频":
if images or audios or voice_ids:
raise ValueError("文生视频不需要连接图像或音频素材。")
elif mode == "图生视频":
if len(images) != 1 or connected_images[0][0] != "图片1":
raise ValueError("图生视频需要在“图片 1”连接 1 张图片,其他图片端口请留空。")
if audios or voice_ids:
raise ValueError("图生视频不支持音频素材,请使用参考生视频。")
else:
if not images and not audios and not voice_ids:
raise ValueError("参考生视频至少需要连接图像素材或音频素材。")
placeholder_image = {"url": "https://example.invalid/image"} if mode == "图生视频" else None
placeholder_references = (
[{"url": f"https://example.invalid/reference-{index}"} for index in range(len(images))]
if mode == "参考生视频"
else []
)
placeholder_audios = (
[{"url": f"https://example.invalid/audio-{index}"} for index in range(len(audios))]
+ [{"voice_id": voice_id} for voice_id in voice_ids]
if mode == "参考生视频"
else []
)
# Validate every user-controlled field before temporary uploads or paid generation calls.
GrokVideoClient.build_video_body(
operation="generate",
prompt=prompt,
model=模型,
aspect_ratio=宽高比,
seconds=seconds,
quality=quality,
model=model,
duration=duration,
aspect_ratio=aspect_ratio,
resolution=resolution,
image=placeholder_image,
reference_images=placeholder_references,
reference_audios=placeholder_audios,
)
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)
base_url = get_base_url_by_route()
client = GrokVideoClient(base_url=base_url)
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,
async def upload_materials():
image_urls, audio_urls = await asyncio.gather(
asyncio.gather(*(upload_image(image, base_url=base_url) for image in images)),
asyncio.gather(*(upload_audio(audio, base_url=base_url) for audio in audios)),
)
return list(image_urls), list(audio_urls)
if pbar is not None and last_progress[0] < 100:
pbar.update(100 - last_progress[0])
image_urls, audio_urls = client.run_async_in_thread(upload_materials())
if mode == "文生视频":
image = None
reference_images = []
reference_audios = []
elif mode == "图生视频":
image = {"url": image_urls[0]}
reference_images = []
reference_audios = []
else:
image = None
reference_images = [{"url": url} for url in image_urls]
reference_audios = [
*({"url": url} for url in audio_urls),
*({"voice_id": voice_id} for voice_id in voice_ids),
]
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
progress_bar, progress_value, callback = _progress_callback()
result = client.run_video_sync(
operation="generate",
prompt=prompt,
model=model,
duration=duration,
aspect_ratio=aspect_ratio,
resolution=resolution,
image=image,
reference_images=reference_images,
reference_audios=reference_audios,
output_dir=_get_output_dir(),
progress_callback=callback,
)
return _finish_video(result, progress_bar, progress_value)
class O1keyGrokVideoEdit:
"""编辑或续写 Grok 视频。输入 VIDEO 会自动上传为 URL。"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"操作": (EDIT_MODE_OPTIONS, {"default": "编辑视频"}),
"提示词": ("STRING", {"default": "", "multiline": True}),
"续写时长(秒)": ("INT", {"default": 6, "min": 2, "max": 10, "step": 1}),
"模型": (MODEL_OPTIONS, {"default": GrokVideoClient.DEFAULT_MODEL}),
},
"optional": {"视频素材": ("VIDEO",)},
}
RETURN_TYPES = ("VIDEO",)
RETURN_NAMES = ("视频",)
FUNCTION = "generate"
CATEGORY = "comfyui_o1key/Video"
DESCRIPTION = (
"编辑或续写视频。编辑输入最长 8.7 秒,并保留原时长和宽高比,输出最高 720p;"
"续写时长为 2–10 秒,输出总时长等于输入时长加续写时长。"
)
def generate(self, **kwargs):
if VideoFromFile is None:
raise RuntimeError("当前 ComfyUI 版本不支持 VideoFromFile,无法输出 VIDEO。")
video = kwargs.get("视频素材")
if video is None:
raise ValueError("请连接一个 VIDEO 类型的视频素材。")
selected_operation = kwargs.get("操作", "编辑视频")
if selected_operation not in EDIT_MODE_OPTIONS:
raise ValueError(f"不支持的 Grok 视频操作:{selected_operation}")
operation = "edit" if selected_operation == "编辑视频" else "extend"
prompt = (kwargs.get("提示词") or "").strip()
model = kwargs.get("模型", GrokVideoClient.DEFAULT_MODEL)
duration = kwargs.get("续写时长(秒)", 6)
if operation == "edit" and _video_duration_seconds(video) > 8.7:
raise ValueError("Grok 视频编辑的输入视频不能超过 8.7 秒。")
GrokVideoClient.build_video_body(
operation=operation,
prompt=prompt,
model=model,
duration=duration,
video={"url": "https://example.invalid/video"},
)
base_url = get_base_url_by_route()
client = GrokVideoClient(base_url=base_url)
video_url = client.run_async_in_thread(upload_video(video, base_url=base_url))
progress_bar, progress_value, callback = _progress_callback()
result = client.run_video_sync(
operation=operation,
prompt=prompt,
model=model,
duration=duration,
video={"url": video_url},
output_dir=_get_output_dir(),
progress_callback=callback,
)
return _finish_video(result, progress_bar, progress_value)
NODE_CLASS_MAPPINGS = {
"O1keyGrokVideo": O1keyGrokVideo,
"O1keyGrokVideoEdit": O1keyGrokVideoEdit,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"O1keyGrokVideo": "Grok Video",
"O1keyGrokVideoEdit": "Grok Video Edit",
}