From 6e29cc41f7e3d0aa2804ad13ebc727e0af67ad8d Mon Sep 17 00:00:00 2001 From: bilt Date: Tue, 8 Jul 2025 09:01:43 +0800 Subject: [PATCH] hf upload --- comfy/cli_args.py | 4 ++++ comfy_extras/nodes_video.py | 23 +++++++++++++++++++++-- requirements.txt | 1 + 3 files changed, 26 insertions(+), 2 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 7234a7ba0..fe144ce88 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -150,6 +150,10 @@ parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build: Enable convenient things that most people using the standalone windows build will probably enjoy (like auto opening the page on startup).") parser.add_argument("--disable-metadata", action="store_true", help="Disable saving prompt metadata in files.") + +parser.add_argument("--hf_token", type=str, help="Hugging Face token for uploading files.") +parser.add_argument("--hf_dataset_name", type=str, help="Hugging Face dataset name for uploading files.") + parser.add_argument("--disable-all-custom-nodes", action="store_true", help="Disable loading all custom nodes.") parser.add_argument("--whitelist-custom-nodes", type=str, nargs='+', default=[], help="Specify custom node folders to load even when --disable-all-custom-nodes is enabled.") parser.add_argument("--disable-api-nodes", action="store_true", help="Disable loading all api nodes.") diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index aae863c6d..468d86178 100644 --- a/comfy_extras/nodes_video.py +++ b/comfy_extras/nodes_video.py @@ -12,6 +12,7 @@ from comfy_api.input import ImageInput, AudioInput, VideoInput from comfy_api.util import VideoContainer, VideoCodec, VideoComponents from comfy_api.input_impl import VideoFromFile, VideoFromComponents from comfy.cli_args import args +from huggingface_hub import HfApi, HfFolder class SaveWEBM: def __init__(self): @@ -94,7 +95,7 @@ class SaveVideo(ComfyNodeABC): "format": (VideoContainer.as_input(), {"default": "auto", "tooltip": "The format to save the video as."}), "codec": (VideoCodec.as_input(), {"default": "auto", "tooltip": "The codec to use for the video."}), "encrypt": ("BOOLEAN", {"default": False, "tooltip": "Enable simple XOR encryption"}), - "encryption_key": ("STRING", {"default": "mysecretkey", "tooltip": "Encryption key for XOR encryption"}), + "encryption_key": ("STRING", {"default": "key4comfy", "tooltip": "Encryption key for XOR encryption"}), }, "hidden": { "prompt": "PROMPT", @@ -110,6 +111,20 @@ class SaveVideo(ComfyNodeABC): CATEGORY = "image/video" DESCRIPTION = "Saves the input images to your ComfyUI output directory." + + def upload(self, hf_token, dataset_name, file_path): + api = HfApi() + HfFolder.save_token(hf_token) + + api.upload_file( + path_or_fileobj=file_path, + path_in_repo=os.path.basename(file_path), + repo_id=dataset_name, + repo_type="dataset", + token=hf_token, + ) + + def simple_xor_encrypt(self, data: bytes, key: str) -> bytes: """简单的XOR加密""" key_bytes = key.encode('utf-8') @@ -173,7 +188,11 @@ class SaveVideo(ComfyNodeABC): # 不加密，直接重命名 os.rename(temp_file_path, final_file_path) - url = f'Click to open Video: {final_file_path}' + # 上传到Hugging Face + self.upload(args.hf_token, args.hf_dataset_name, final_file_path) + final_filename = os.path.basename(final_file_path) + + url = f'Video generated, click to open: {final_filename}' # 只返回url到ui.text return (url,) diff --git a/requirements.txt b/requirements.txt index 27d385389..96597dba6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,6 +15,7 @@ aiohttp>=3.11.8 yarl>=1.18.0 pyyaml Pillow +huggingface_hub scipy tqdm psutil