Add Grok and VEO video workflow support
This commit is contained in:
o1key
@@ -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",
|
||||
}
|
||||
Reference in New Issue
Block a user