comfyui_o1key/nodes/o1key_video_generator.py

"""Panel-driven parallel video generation and durable result nodes."""

from __future__ import annotations

import json
import os
from typing import Any

import folder_paths
from PIL import Image
from comfy_api.latest import InputImpl, io, ui

from ..utils.image_utils import pil_to_tensor
from ..utils.o1key_video_catalog import (
    SEEDANCE_ASSET_CREATION_MODE_OPTIONS,
    SEEDANCE_MODEL_OPTIONS,
    SEEDANCE_ROUTE_OPTIONS,
    VIDEO_ASPECT_RATIO_OPTIONS,
    VIDEO_DURATION_OPTIONS,
    VIDEO_GENERATION_MODE_OPTIONS,
    VIDEO_PROVIDER_OPTIONS,
    VIDEO_RESOLUTION_OPTIONS,
)


def _parse_result_descriptor(value: str | dict[str, Any] | None) -> dict[str, str] | None:
    if value in (None, "", "{}"):
        return None
    try:
        item = json.loads(value) if isinstance(value, str) else value
    except json.JSONDecodeError:
        raise ValueError("视频结果描述符不是有效 JSON") from None
    if not isinstance(item, dict):
        raise ValueError("视频结果描述符必须是对象")
    filename = os.path.basename(str(item.get("filename") or "").strip())
    subfolder = str(item.get("subfolder") or "").strip().replace("\\", "/")
    folder_type = str(item.get("type") or "output").strip()
    if not filename or folder_type not in {"output", "temp"}:
        raise ValueError("视频结果描述符无效")
    if subfolder.startswith("/") or any(part == ".." for part in subfolder.split("/")):
        raise ValueError("视频结果子目录无效")
    return {"filename": filename, "subfolder": subfolder, "type": folder_type}


def _resolve_result_path(descriptor: dict[str, str]) -> str:
    root = (
        folder_paths.get_output_directory()
        if descriptor["type"] == "output"
        else folder_paths.get_temp_directory()
    )
    root = os.path.realpath(os.path.abspath(root))
    candidate = os.path.realpath(
        os.path.abspath(os.path.join(root, descriptor["subfolder"], descriptor["filename"]))
    )
    try:
        inside = os.path.commonpath((root, candidate)) == root
    except ValueError:
        inside = False
    if not inside or not os.path.isfile(candidate):
        raise ValueError("视频结果文件不存在或已超出允许目录")
    return candidate


def _result_values(
    video_manifest: str | dict[str, Any] | None,
    last_frame_manifest: str | dict[str, Any] | None,
) -> tuple[Any, Any]:
    video_descriptor = _parse_result_descriptor(video_manifest)
    if video_descriptor is None:
        return None, None
    video_path = _resolve_result_path(video_descriptor)

    last_frame_tensor = None
    last_frame_descriptor = _parse_result_descriptor(last_frame_manifest)
    if last_frame_descriptor is not None:
        last_frame_path = _resolve_result_path(last_frame_descriptor)
        with Image.open(last_frame_path) as image:
            image.load()
            last_frame_tensor = pil_to_tensor([image.convert("RGB")])
    return InputImpl.VideoFromFile(video_path), last_frame_tensor


class O1keyVideoGenerator(io.ComfyNode):
    """A frontend-operated generator; each click creates an independent job."""

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="O1keyVideoGenerator",
            display_name="o1key 视频生成",
            category="o1key/video",
            description="点击节点内按钮提交独立后台视频任务，并自动连接原生保存节点。",
            inputs=[
                io.String.Input("prompt", default="", multiline=True, socketless=True),
                io.Combo.Input("provider", options=VIDEO_PROVIDER_OPTIONS, default="seedance", socketless=True),
                io.Combo.Input("model", options=SEEDANCE_MODEL_OPTIONS, default="seedance-2.0", socketless=True),
                io.Combo.Input("route", options=SEEDANCE_ROUTE_OPTIONS, default="domestic", socketless=True),
                io.Combo.Input("generation_mode", options=VIDEO_GENERATION_MODE_OPTIONS, default="multimodal", socketless=True),
                io.Combo.Input("resolution", options=VIDEO_RESOLUTION_OPTIONS, default="720p", socketless=True),
                io.Combo.Input("aspect_ratio", options=VIDEO_ASPECT_RATIO_OPTIONS, default="auto", socketless=True),
                io.Combo.Input("duration", options=VIDEO_DURATION_OPTIONS, default="5", socketless=True),
                io.Boolean.Input("generate_audio", default=False, socketless=True),
                io.Boolean.Input("return_last_frame", default=False, socketless=True),
                io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, socketless=True),
                io.String.Input("media_manifest", default="{}", multiline=True, socketless=True),
                io.String.Input("asset_manifest", default="{}", multiline=True, socketless=True),
                io.String.Input("provider_options", default="{}", multiline=True, socketless=True),
                io.String.Input("filename_prefix", default="o1key_video", socketless=True),
                io.String.Input("save_location", default="video", socketless=True),
                # Append-only: keep every released widgets_values position stable.
                io.Combo.Input(
                    "asset_creation_mode",
                    options=SEEDANCE_ASSET_CREATION_MODE_OPTIONS,
                    default="auto",
                    socketless=True,
                ),
                io.String.Input("video_manifest", default="{}", multiline=True, socketless=True),
                io.String.Input("last_frame_manifest", default="{}", multiline=True, socketless=True),
            ],
            outputs=[
                io.Video.Output("VIDEO", display_name="VIDEO"),
                io.Image.Output("LAST_FRAME", display_name="LAST_FRAME"),
            ],
            not_idempotent=True,
        )

    @classmethod
    def execute(
        cls,
        video_manifest: str = "{}",
        last_frame_manifest: str = "{}",
        **_kwargs,
    ) -> io.NodeOutput:
        # Generation remains owned by /o1key/video/jobs. Native execution only
        # resolves the latest completed local descriptors and never spends again.
        video, last_frame = _result_values(video_manifest, last_frame_manifest)
        if video is None:
            return io.NodeOutput(block_execution="请先在 o1key 视频生成节点中完成一次生成")
        return io.NodeOutput(video, last_frame)


class O1keyVideoResult(io.ComfyNode):
    """A completed background result that can later feed native VIDEO workflows."""

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="O1keyVideoResult",
            display_name="o1key 视频结果",
            category="o1key/video",
            description="显示独立后台任务状态；完成后可向下游输出原生 VIDEO。",
            is_deprecated=True,
            inputs=[
                io.String.Input("batch_id", default="", socketless=True),
                io.String.Input("video_manifest", default="{}", multiline=True, socketless=True),
                io.String.Input("last_frame_manifest", default="{}", multiline=True, socketless=True),
            ],
            outputs=[
                io.Video.Output("VIDEO", display_name="VIDEO"),
                io.Image.Output("LAST_FRAME", display_name="LAST_FRAME"),
            ],
        )

    @classmethod
    def execute(
        cls,
        batch_id: str = "",
        video_manifest: str = "{}",
        last_frame_manifest: str = "{}",
    ) -> io.NodeOutput:
        del batch_id
        video_descriptor = _parse_result_descriptor(video_manifest)
        if video_descriptor is None:
            raise ValueError("视频任务尚未完成，没有可输出的视频")
        video, last_frame_tensor = _result_values(video_manifest, last_frame_manifest)

        preview = ui.PreviewVideo([video_descriptor])
        return io.NodeOutput(
            video,
            last_frame_tensor,
            ui=preview,
        )


NODE_CLASS_MAPPINGS = {
    "O1keyVideoGenerator": O1keyVideoGenerator,
    "O1keyVideoResult": O1keyVideoResult,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "O1keyVideoGenerator": "o1key 视频生成",
    "O1keyVideoResult": "o1key 视频结果",
}


__all__ = ["O1keyVideoGenerator", "O1keyVideoResult"]