"""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)))