File size: 5,914 Bytes
91286ad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 |
# ComfyUI/custom_nodes/zip_output_to_hf.py
#
# Node: ZipOutputToHuggingFace
#
# - Zips the ComfyUI output folder
# - Uploads the zip to a Hugging Face repo
# - Returns a direct download URL as STRING
import os
import time
import zipfile
import tempfile
import folder_paths # ComfyUI's helper for paths
try:
from huggingface_hub import HfApi
except ImportError:
HfApi = None
class ZipOutputToHuggingFace:
"""
Zip the ComfyUI output folder and upload the archive to Hugging Face.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
# Where to upload; default to your repo
"hf_repo_id": (
"STRING",
{"default": "saliacoel/MyCustomNodes"},
),
},
"optional": {
# Hugging Face token; can be left empty if you use env vars / hf auth login
"hf_token": (
"STRING",
{
"default": "",
"multiline": False,
},
),
# Prefix for the zip filename
"zip_name_prefix": (
"STRING",
{
"default": "comfy_output",
"multiline": False,
},
),
# Subfolder in the repo to put the zip into
"remote_dir": (
"STRING",
{
"default": "exports",
"multiline": False,
},
),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("download_url",)
FUNCTION = "zip_and_upload"
CATEGORY = "utils/huggingface"
OUTPUT_NODE = True # can act as a terminal node in your graph
def zip_and_upload(
self,
hf_repo_id: str,
hf_token: str = "",
zip_name_prefix: str = "comfy_output",
remote_dir: str = "exports",
):
# 1. Check huggingface_hub availability
if HfApi is None:
raise RuntimeError(
"huggingface_hub is not installed.\n"
"Install it in your ComfyUI environment, for example:\n"
" pip install huggingface_hub"
)
# 2. Get ComfyUI output directory
output_dir = folder_paths.get_output_directory()
if not os.path.isdir(output_dir):
raise RuntimeError(f"Output directory does not exist: {output_dir}")
# 3. Ensure there is at least one file to zip
has_files = False
for _, _, files in os.walk(output_dir):
if files:
has_files = True
break
if not has_files:
raise RuntimeError(
f"No files found in output directory: {output_dir}"
)
# 4. Create zip in a temp directory (so we don't zip the zip)
timestamp = time.strftime("%Y%m%d-%H%M%S")
safe_prefix = (zip_name_prefix or "").strip() or "comfy_output"
zip_filename = f"{safe_prefix}_{timestamp}.zip"
tmp_dir = tempfile.gettempdir()
zip_path = os.path.join(tmp_dir, zip_filename)
zip_abs = os.path.abspath(zip_path)
# 5. Build the zip archive of everything under output/
with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
for root, _, files in os.walk(output_dir):
for fname in files:
file_path = os.path.join(root, fname)
# Just in case tempdir == output_dir, avoid zipping ourselves
if os.path.abspath(file_path) == zip_abs:
continue
rel_path = os.path.relpath(file_path, output_dir)
zipf.write(file_path, arcname=rel_path)
# 6. Prepare Hugging Face upload
repo_id = (hf_repo_id or "").strip()
if not repo_id:
# Clean up zip if we created it but can't use it
try:
os.remove(zip_path)
except OSError:
pass
raise RuntimeError("hf_repo_id cannot be empty")
remote_dir = (remote_dir or "").strip().strip("/")
if remote_dir:
path_in_repo = f"{remote_dir}/{zip_filename}"
else:
path_in_repo = zip_filename
api = HfApi()
# 7. Upload to Hugging Face
try:
api.upload_file(
path_or_fileobj=zip_path,
path_in_repo=path_in_repo,
repo_id=repo_id,
# token is optional; if empty, hf_hub will use env/CLI token if present
token=hf_token or None,
)
except Exception as e:
# Clean up local zip on failure
try:
os.remove(zip_path)
except OSError:
pass
raise RuntimeError(f"Failed to upload to Hugging Face: {e}")
# 8. Build a direct download URL
# This gives the raw file: https://huggingface.co/{repo_id}/resolve/main/{path_in_repo}
download_url = f"https://huggingface.co/{repo_id}/resolve/main/{path_in_repo}"
# 9. Clean up temp zip (we only keep the file on HF)
try:
os.remove(zip_path)
except OSError:
pass
# Return as a single STRING output
return (download_url,)
# Register node with ComfyUI
NODE_CLASS_MAPPINGS = {
"ZipOutputToHuggingFace": ZipOutputToHuggingFace,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ZipOutputToHuggingFace": "Zip Output → Hugging Face",
}
|