213 lines
7.8 KiB
Python
213 lines
7.8 KiB
Python
import base64
|
|
import folder_paths
|
|
import numpy as np
|
|
import os
|
|
import requests
|
|
import torch
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import subprocess
|
|
|
|
|
|
class UploadToBaseNode:
|
|
def __init__(self):
|
|
self.type = "output"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"api_key": ("STRING", {"default": None}),
|
|
"filename_prefix": ("STRING", {
|
|
"default": "ComfyUI",
|
|
"tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE", {"default": None}),
|
|
"video": ("VIDEO", {"default": None}),
|
|
"folder_id": ("STRING", {"default": None}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO"
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "run"
|
|
CATEGORY = "BASE"
|
|
|
|
def run(self, image=None, video=None, api_key=None, folder_id=None, prompt=None,
|
|
filename_prefix=None, extra_pnginfo=None):
|
|
import json
|
|
from PIL import Image, PngImagePlugin
|
|
from io import BytesIO
|
|
import tempfile
|
|
|
|
# Auto-embed the current prompt as "workflow" if available
|
|
if prompt and isinstance(prompt, dict):
|
|
if extra_pnginfo is None:
|
|
extra_pnginfo = {}
|
|
if "workflow" not in extra_pnginfo:
|
|
extra_pnginfo["workflow"] = prompt
|
|
|
|
if image is None and video is None:
|
|
raise ValueError("Either 'image' or 'video' must be provided.")
|
|
|
|
# Remove api_key from extra_pnginfo if present and sanitize nested workflow nodes
|
|
if extra_pnginfo:
|
|
# Remove any top-level "api_key"
|
|
extra_pnginfo = {
|
|
k: v for k, v in extra_pnginfo.items()
|
|
if k.lower() != "api_key"
|
|
}
|
|
|
|
# Sanitize nested api_key in workflow nodes
|
|
workflow = extra_pnginfo.get("workflow")
|
|
if workflow and isinstance(workflow, dict):
|
|
for node in workflow.get("nodes", []):
|
|
if node.get("type") == "UploadToBaseNode":
|
|
widgets = node.get("widgets_values")
|
|
api_key_index = list(self.INPUT_TYPES()["required"]).index("api_key")
|
|
if isinstance(widgets, list) and len(widgets) > api_key_index:
|
|
widgets[api_key_index] = ""
|
|
|
|
# Determine width and height safely
|
|
if video is not None:
|
|
width, height = video.get_dimensions()
|
|
elif image is not None:
|
|
if torch.is_tensor(image) and image.ndim == 4:
|
|
image_tensor = image[0]
|
|
elif isinstance(image, list) and len(image) > 0:
|
|
image_tensor = image[0]
|
|
else:
|
|
image_tensor = image
|
|
width, height = image_tensor.shape[-1], image_tensor.shape[-2]
|
|
else:
|
|
raise ValueError("Either 'image' or 'video' must be provided.")
|
|
|
|
if not filename_prefix:
|
|
filename_prefix = "ComfyUI"
|
|
|
|
full_output_folder, filename_base, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
|
filename_prefix,
|
|
folder_paths.get_output_directory(),
|
|
width,
|
|
height
|
|
)
|
|
|
|
if video is not None:
|
|
mime_type = "video/mp4"
|
|
filename = f"{filename_base}_{counter:05}_.mp4"
|
|
|
|
# Prepare metadata
|
|
if extra_pnginfo is None:
|
|
extra_pnginfo = {}
|
|
if prompt is not None:
|
|
extra_pnginfo.setdefault("prompt", prompt)
|
|
extra_pnginfo.setdefault("workflow_type", "video")
|
|
|
|
save_path = os.path.join(full_output_folder, filename)
|
|
video.save_to(
|
|
save_path,
|
|
format="mp4",
|
|
codec="h264",
|
|
metadata=extra_pnginfo
|
|
)
|
|
|
|
# Read saved video into buffer for upload
|
|
with open(save_path, "rb") as f:
|
|
buffer = BytesIO(f.read())
|
|
else:
|
|
# Convert batched tensor to list if needed
|
|
if torch.is_tensor(image) and image.ndim == 4:
|
|
image = list(image)
|
|
|
|
def tensor_to_pil(img_tensor):
|
|
img_array = img_tensor.cpu().numpy()
|
|
img_array = np.nan_to_num(img_array, nan=0.0, posinf=255.0, neginf=0.0)
|
|
img_array = (img_array * 255).clip(0, 255).astype("uint8")
|
|
|
|
# If it's a batch of images, select the first one (or raise if unexpected)
|
|
if img_array.ndim == 4:
|
|
if img_array.shape[3] in {1, 3}:
|
|
img_array = img_array[0] # Take first image in batch
|
|
else:
|
|
raise ValueError(
|
|
f"Unexpected channel dimension in batch: {img_array.shape}")
|
|
|
|
if img_array.ndim == 3:
|
|
if img_array.shape[0] in {1, 3}:
|
|
img_array = np.transpose(img_array, (1, 2, 0)) # (C, H, W) → (H, W, C)
|
|
elif img_array.shape[2] not in {1, 3}:
|
|
raise ValueError(f"Unexpected channel dimension: {img_array.shape}")
|
|
elif img_array.ndim == 2:
|
|
img_array = np.expand_dims(img_array, axis=-1)
|
|
img_array = np.repeat(img_array, 3, axis=2)
|
|
else:
|
|
raise ValueError(f"Unsupported image shape: {img_array.shape}")
|
|
|
|
if img_array.shape[2] == 1:
|
|
img_array = np.repeat(img_array, 3, axis=2)
|
|
|
|
return Image.fromarray(img_array)
|
|
|
|
if isinstance(image, list) and len(image) > 1:
|
|
raise ValueError("Animated image export is not supported.")
|
|
|
|
# Static PNG logic
|
|
img_tensor = image[0] if isinstance(image, list) else image
|
|
img = tensor_to_pil(img_tensor)
|
|
|
|
pnginfo = PngImagePlugin.PngInfo()
|
|
if prompt is not None:
|
|
pnginfo.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for k, v in extra_pnginfo.items():
|
|
pnginfo.add_text(k, json.dumps(v))
|
|
|
|
buffer = BytesIO()
|
|
img = img.convert("RGB")
|
|
img.save(buffer, format="PNG", pnginfo=pnginfo)
|
|
mime_type = "image/png"
|
|
filename = f"{filename_base}_{counter:05}_.png"
|
|
|
|
# Save locally for debug
|
|
import tempfile
|
|
verify_path = os.path.join(tempfile.gettempdir(), filename)
|
|
with open(verify_path, "wb") as f:
|
|
f.write(buffer.getvalue())
|
|
|
|
b64_bytes = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
|
|
response = requests.post(
|
|
"https://us-central1-base-14bf3.cloudfunctions.net/uploadToBase", json={
|
|
"bytes": b64_bytes,
|
|
"parent": folder_id,
|
|
"apiKey": api_key,
|
|
"mimeType": mime_type,
|
|
"filename": filename,
|
|
}, timeout=60)
|
|
|
|
if response.status_code != 200:
|
|
raise RuntimeError(f"Upload failed with status {response.status_code}: {response.text}")
|
|
|
|
full_path = os.path.join(full_output_folder, filename)
|
|
with open(full_path, "wb") as f:
|
|
f.write(buffer.getvalue())
|
|
|
|
results = [{
|
|
"filename": filename,
|
|
"subfolder": subfolder,
|
|
"type": self.type,
|
|
}]
|
|
|
|
return {
|
|
"ui": {
|
|
"images": results,
|
|
"animated": (True,) if video is not None else (False,)
|
|
}
|
|
}
|