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
@@ -0,0 +1,182 @@
|
||||
"""Omni Flash video generation through a normal ComfyUI execution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io as py_io
|
||||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
from comfy_api.latest import InputImpl, io
|
||||
|
||||
from ..clients.omni_flash_client import OmniFlashClient, build_video_body
|
||||
from ..utils.config import get_api_key_or_raise, get_base_url_by_route
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
from ..utils.r2_uploader import upload_image, upload_video
|
||||
from ..utils.video_task import check_interrupt
|
||||
|
||||
|
||||
_MODES = {
|
||||
"文生视频": "text",
|
||||
"参考图视频": "reference",
|
||||
"首尾帧": "first_last_frame",
|
||||
"视频编辑": "edit",
|
||||
}
|
||||
_MAX_EDIT_VIDEO_BYTES = 20 * 1024 * 1024
|
||||
|
||||
|
||||
def _make_progress_callback():
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
bar = ProgressBar(100)
|
||||
except ImportError:
|
||||
bar = None
|
||||
last_progress = -1
|
||||
|
||||
def update(stage: str, value: int, _task_id: str) -> None:
|
||||
nonlocal last_progress
|
||||
# The provider percentage maps directly to the node bar. Reserve the
|
||||
# final point for the local video download and save.
|
||||
current = max(0, min(100, int(value)))
|
||||
if stage != "done":
|
||||
current = min(current, 99)
|
||||
if current <= last_progress:
|
||||
return
|
||||
last_progress = current
|
||||
if bar is not None:
|
||||
bar.update_absolute(current, 100)
|
||||
|
||||
return update
|
||||
|
||||
|
||||
class O1keyOmniFlashVideo(io.ComfyNode):
|
||||
"""One graph node owns validation, submission, polling, and VIDEO output."""
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, **kwargs):
|
||||
"""Do not reuse a completed generation when this node is queued again."""
|
||||
return uuid.uuid4().hex
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="O1keyOmniFlashVideo",
|
||||
display_name="Omni Flash 视频生成",
|
||||
description="连接图片或视频后生成；开始生成按钮运行当前节点。",
|
||||
category="comfyui_o1key/视频",
|
||||
is_output_node=True,
|
||||
not_idempotent=True,
|
||||
inputs=[
|
||||
io.String.Input("提示词", multiline=True, default=""),
|
||||
io.Combo.Input("生成模式", options=list(_MODES), default="文生视频"),
|
||||
io.Combo.Input("分辨率", options=["720p", "1080p"], default="720p"),
|
||||
io.Combo.Input("宽高比", options=["16:9", "9:16"], default="16:9"),
|
||||
io.Autogrow.Input(
|
||||
"参考图片",
|
||||
optional=True,
|
||||
template=io.Autogrow.TemplateNames(
|
||||
input=io.Image.Input("参考图片"),
|
||||
names=[f"参考图片{i}" for i in range(1, 6)],
|
||||
min=0,
|
||||
),
|
||||
),
|
||||
io.Image.Input("首帧图片", optional=True),
|
||||
io.Image.Input("尾帧图片", optional=True, tooltip="可不连接；只连接首帧也能生成。"),
|
||||
io.Video.Input("源视频", optional=True),
|
||||
],
|
||||
outputs=[io.Video.Output("VIDEO", display_name="视频")],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _reference_images(kwargs):
|
||||
group = kwargs.get("参考图片")
|
||||
if isinstance(group, dict):
|
||||
return [value for value in group.values() if value is not None]
|
||||
return [kwargs[f"参考图片{index}"] for index in range(1, 6)
|
||||
if kwargs.get(f"参考图片{index}") is not None]
|
||||
|
||||
@staticmethod
|
||||
def _one_image(value, label):
|
||||
images = tensor_to_pil(value)
|
||||
if len(images) != 1:
|
||||
raise ValueError(f"{label}必须恰好包含 1 张图片")
|
||||
return images[0].convert("RGB")
|
||||
|
||||
@staticmethod
|
||||
def _check_source_video(video):
|
||||
source = video.get_stream_source() if hasattr(video, "get_stream_source") else None
|
||||
if isinstance(source, py_io.BytesIO):
|
||||
size = source.getbuffer().nbytes
|
||||
elif isinstance(source, str) and os.path.isfile(source):
|
||||
size = os.path.getsize(source)
|
||||
if Path(source).suffix.lower() not in {".mp4", ".mov"}:
|
||||
raise ValueError("源视频须为 MP4 或 MOV")
|
||||
else:
|
||||
raise ValueError("无法读取源视频文件")
|
||||
if size > _MAX_EDIT_VIDEO_BYTES:
|
||||
raise ValueError("源视频不能超过 20 MB")
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, **kwargs):
|
||||
mode = _MODES.get(kwargs.get("生成模式"))
|
||||
if mode is None:
|
||||
raise ValueError("生成模式无效")
|
||||
images = cls._reference_images(kwargs)
|
||||
first = kwargs.get("首帧图片")
|
||||
last = kwargs.get("尾帧图片")
|
||||
source = kwargs.get("源视频")
|
||||
if mode == "text" and (images or first is not None or last is not None or source is not None):
|
||||
raise ValueError("文生视频模式不接受媒体输入")
|
||||
if mode == "reference" and (not images or first is not None or last is not None or source is not None):
|
||||
raise ValueError("参考图视频模式只接受参考图片")
|
||||
if mode == "first_last_frame" and (images or first is None or source is not None):
|
||||
raise ValueError("首尾帧模式须连接首帧图片；尾帧图片可选")
|
||||
if mode == "edit" and (source is None or first is not None or last is not None):
|
||||
raise ValueError("视频编辑模式须连接源视频，不能连接首尾帧")
|
||||
if mode == "edit" and len(images) > 5:
|
||||
raise ValueError("视频编辑最多支持 5 张参考图")
|
||||
|
||||
model = "omni_flash_abra_edit" if mode == "edit" else "omni_flash_10s"
|
||||
prompt = kwargs["提示词"]
|
||||
resolution = kwargs["分辨率"]
|
||||
ratio = kwargs["宽高比"]
|
||||
# Validate the scalar contract before any upload or paid request.
|
||||
upload_images = ([first] + ([last] if last is not None else [])) if mode == "first_last_frame" else images
|
||||
build_video_body(model=model, prompt=prompt, resolution=resolution,
|
||||
aspect_ratio=ratio, mode=mode,
|
||||
references=["https://example.invalid/media"] * len(upload_images),
|
||||
source_video_url="https://example.invalid/video" if mode == "edit" else "")
|
||||
pil_images = [cls._one_image(value, "输入") for value in upload_images]
|
||||
if source is not None:
|
||||
cls._check_source_video(source)
|
||||
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
base_url = get_base_url_by_route()
|
||||
urls = []
|
||||
for image in pil_images:
|
||||
check_interrupt()
|
||||
urls.append(await upload_image(image, base_url=base_url))
|
||||
source_url = await upload_video(source, base_url=base_url) if source is not None else ""
|
||||
body = build_video_body(model=model, prompt=prompt, resolution=resolution,
|
||||
aspect_ratio=ratio, mode=mode, references=urls,
|
||||
source_video_url=source_url)
|
||||
|
||||
output_dir = Path(folder_paths.get_output_directory()) / "omni_flash"
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
filename = f"{uuid.uuid4().hex}.mp4"
|
||||
target = output_dir / filename
|
||||
partial = output_dir / f"{filename}.part"
|
||||
report_progress = _make_progress_callback()
|
||||
report_progress("polling", 0, "")
|
||||
try:
|
||||
await OmniFlashClient(base_url=base_url, api_key=api_key).generate(
|
||||
body, str(partial), progress=report_progress,
|
||||
)
|
||||
check_interrupt()
|
||||
os.replace(partial, target)
|
||||
report_progress("done", 100, "")
|
||||
except BaseException:
|
||||
partial.unlink(missing_ok=True)
|
||||
raise
|
||||
return io.NodeOutput(InputImpl.VideoFromFile(str(target)))
|
||||
Reference in New Issue
Block a user