136 lines
4.1 KiB
Python
136 lines
4.1 KiB
Python
import folder_paths
|
|
import os
|
|
import requests
|
|
import json
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import numpy as np
|
|
from comfy.cli_args import args
|
|
|
|
class Webhook:
|
|
"""
|
|
Handles saving generated images and sending webhook notifications.
|
|
Optimized for faster image saving using JPEG and memory-efficient webhook sending.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.save_directory = folder_paths.get_output_directory()
|
|
self.asset_type = "output"
|
|
self.name_prefix = ""
|
|
self.jpeg_quality = 85 # Adjust 0-100 for quality/speed tradeoff
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
|
"webhook_url": ("STRING", {"default": "", "placeholder": "enter webhook URL"}),
|
|
},
|
|
"optional": {
|
|
"metadata": ("STRING", {"default": "", "placeholder": "enter metadata for webhook"}),
|
|
"external_uid": ("STRING", {"default": "", "placeholder": "enter external uid of generated image or session"}),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
CATEGORY = "image/Save Image & Notify"
|
|
FUNCTION = "process_and_notify"
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
|
|
def process_and_notify(
|
|
self,
|
|
images,
|
|
filename_prefix="ComfyUI",
|
|
webhook_url="",
|
|
metadata="",
|
|
external_uid="",
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
custom_metadata = metadata
|
|
|
|
# Determine save path
|
|
output_name = filename_prefix + self.name_prefix
|
|
save_path, base_name, img_counter, sub_path, final_prefix = folder_paths.get_save_image_path(
|
|
output_name,
|
|
self.save_directory,
|
|
images[0].shape[1],
|
|
images[0].shape[0]
|
|
)
|
|
|
|
saved_files = []
|
|
last_output_img = None
|
|
|
|
for idx, img_data in enumerate(images):
|
|
# Efficient tensor -> image conversion
|
|
img_array = (img_data.mul(255).clamp(0, 255).byte().cpu().numpy())
|
|
output_img = Image.fromarray(img_array)
|
|
|
|
# Keep reference to last image for webhook
|
|
if idx == len(images) - 1:
|
|
last_output_img = output_img
|
|
|
|
# Save image to disk
|
|
current_name = base_name.replace("%batch_num%", str(idx))
|
|
img_filename = f"{current_name}_{img_counter:05}_.jpg"
|
|
full_path = os.path.join(save_path, img_filename)
|
|
output_img.save(
|
|
full_path,
|
|
format="JPEG",
|
|
quality=self.jpeg_quality
|
|
)
|
|
|
|
saved_files.append({
|
|
"filename": img_filename,
|
|
"subfolder": sub_path,
|
|
"type": self.asset_type
|
|
})
|
|
img_counter += 1
|
|
|
|
print(f"Successfully saved images: {saved_files}")
|
|
|
|
# Send webhook notification with in-memory image (reuse converted image)
|
|
if webhook_url and last_output_img:
|
|
buffer = BytesIO()
|
|
last_output_img.save(
|
|
buffer,
|
|
format="JPEG",
|
|
quality=self.jpeg_quality
|
|
)
|
|
buffer.seek(0)
|
|
|
|
webhook_data = {
|
|
"metadata": custom_metadata,
|
|
"external_uid": external_uid,
|
|
"images": json.dumps(saved_files)
|
|
}
|
|
|
|
response = requests.post(
|
|
webhook_url,
|
|
data=webhook_data,
|
|
files={"file": ("image.jpg", buffer, "image/jpeg")}
|
|
)
|
|
|
|
if response.status_code >= 300:
|
|
print(f"Webhook failed - Status: {response.status_code}, Response: {response.text}")
|
|
|
|
return {"ui": {"images": saved_files}}
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Webhook": Webhook,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Webhook": "Webhook Image Saver",
|
|
}
|
|
|
|
__all__ = [
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
"Webhook",
|
|
]
|