From fe60c8aa2b19720aaa36331df7bc6e21cd4838e4 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Fri, 13 Jun 2025 11:16:41 +0800 Subject: [PATCH] feat: support save video node (#64) Signed-off-by: Frost Ming --- nodes/nodes.py | 8 ++++++++ src/comfy_pack/utils.py | 21 ++++++++------------- uv.lock | 6 +++--- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/nodes/nodes.py b/nodes/nodes.py index 75e749a..fb541b1 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -11,6 +11,7 @@ import node_helpers import numpy as np import torch from comfy_extras.nodes_audio import SaveAudio +from comfy_extras.nodes_video import SaveVideo from PIL import Image, ImageOps, ImageSequence, PngImagePlugin from PIL.PngImagePlugin import PngInfo @@ -465,10 +466,16 @@ class OutputAudio(SaveAudio): } +class OutputVideo(SaveVideo): + CPACK_NODE = True + CATEGORY = "ComfyPack/output" + + NODE_CLASS_MAPPINGS = { "CPackOutputFile": OutputFile, "CPackOutputImage": OutputImage, "CPackOutputAudio": OutputAudio, + "CPackOutputVideo": OutputVideo, "CPackOutputZip": OutputImageWithStringTxt, "CPackInputImage": ImageInput, "CPackInputString": StringInput, @@ -485,6 +492,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CPackInputAny": "Any Input", "CPackOutputImage": "Image Output", "CPackOutputAudio": "Audio Output", + "CPackOutputVideo": "Video Output", "CPackOutputFile": "File Output", "CPackOutputZip": "Zip Output(img + txt file)", } diff --git a/src/comfy_pack/utils.py b/src/comfy_pack/utils.py index f9c401f..a829fae 100644 --- a/src/comfy_pack/utils.py +++ b/src/comfy_pack/utils.py @@ -1,19 +1,14 @@ from __future__ import annotations -import subprocess -from pathlib import Path import re +import subprocess import sys +from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, Union if TYPE_CHECKING: from pydantic import BaseModel -CPACK_OUTPUT_NODES = { - "CPackOutputFile", - "CPackOutputImage", - "CPackOutputZip" -} CPACK_PATH_INPUT_NODES = { "CPackInputFile", @@ -211,7 +206,7 @@ def populate_workflow( for _, node in output_spec.items(): node_id = node["id"] - if node["class_type"] in CPACK_OUTPUT_NODES: + if node["class_type"].startswith("CPackOutput"): workflow[node_id]["inputs"]["filename_prefix"] = ( output_path / f"{session_id}{node_id}_" ).as_posix() @@ -222,7 +217,7 @@ def retrieve_workflow_outputs( workflow: dict, output_path: Path, session_id: str = "", -) -> Union[Path, list[Path], dict[str, Path], dict[str, list[Path]]]: +) -> Union[Path, list[Path], dict[str, Path | list[Path]]]: """ Gets the output file(s) from the workflow. @@ -231,7 +226,7 @@ def retrieve_workflow_outputs( output_path (Path): The path where output files are saved. Returns: - Union[Path, list[Path], dict[str, Path], dict[str, list[Path]]]: + Union[Path, list[Path], dict[str, Path | list[Path]]]: - A single Path if there's only one output file. - A list of Paths if there are multiple files for a single output. - A dictionary mapping output names to Paths or lists of Paths for multiple outputs. @@ -244,7 +239,7 @@ def retrieve_workflow_outputs( value_map = {} for k, node in outputs.items(): node_id = node["id"] - path_strs = list(output_path.glob(f"{node_id}_*")) + path_strs = list(output_path.glob(f"{session_id}{node_id}_*")) if len(path_strs) == 1: value_map[k] = path_strs[0] else: @@ -252,8 +247,8 @@ def retrieve_workflow_outputs( return value_map name, node = next(iter(outputs.items())) - if node["class_type"] not in CPACK_OUTPUT_NODES: - raise ValueError(f"Output node {name} is not of type {CPACK_OUTPUT_NODES}") + if not node["class_type"].startswith("CPackOutput"): + raise ValueError(f"Node {name} is not a comfy-pack output node") node_id = node["id"] outs = list(output_path.glob(f"{session_id}{node_id}_*")) diff --git a/uv.lock b/uv.lock index ab381dd..8e9fcfc 100644 --- a/uv.lock +++ b/uv.lock @@ -209,7 +209,7 @@ wheels = [ [[package]] name = "bentoml" -version = "1.4.13" +version = "1.4.15" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "a2wsgi" }, @@ -255,9 +255,9 @@ dependencies = [ { name = "uvicorn" }, { name = "watchfiles" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fb/15/9fd0bcfb2b4b407142a2362f69a1a5980820a77204405b402830cbdc6424/bentoml-1.4.13.tar.gz", hash = "sha256:9ba3ff502726a1496ebd9733943896797d8d9500a88f88da7fae19d51e833fc5", size = 971767, upload-time = "2025-05-09T07:59:01.918Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/c9/47f188d06196526173e51f1bfa2e461d73bfa9ed5a6f5f18974c9332f184/bentoml-1.4.15.tar.gz", hash = "sha256:e16fd5cbed6cccdafaa8d9c42cacb2d9b39ac2b335eb7b6051930688e6fff4f2", size = 972188, upload-time = "2025-05-23T07:34:23.377Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/01/d64702692bfe11988c4ae0be5b103af71bf9dc7c2c78d1207f47a6e1e71c/bentoml-1.4.13-py3-none-any.whl", hash = "sha256:06859ec789f4f17c3744e1ab1d73da80ee0dc5ce0dd8400040b4b5c725af6567", size = 1151642, upload-time = "2025-05-09T07:58:58.649Z" }, + { url = "https://files.pythonhosted.org/packages/5a/2d/53b0f577abb1d72606b73ac4377b2ce5bafaebc7d7c0a5ee691647790919/bentoml-1.4.15-py3-none-any.whl", hash = "sha256:3e9c411236ddb8a36346e3ab63a0ca742c63dc68140a610501f7b8a9c4ed02a5", size = 1151960, upload-time = "2025-05-23T07:34:20.972Z" }, ] [[package]]