ComfyUI-CustomNode/nodes/webhook_forward.py

# -*- coding:utf-8 -*-
"""
File                   webhook_forward.py
Author                 silence
Date                  2025/6/13  13:42
"""
import base64
import io
import os
import random
import string
import time

import numpy as np
import requests
from PIL import Image


class PromptIDGenerator:
    """PromptID生成器：用户可传入或自动生成PromptID"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {},
            "optional": {
                "custom_prompt_id": ("STRING", {"default": "", "placeholder": "留空则自动生成"}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt_id",)
    FUNCTION = "generate_prompt_id"
    OUTPUT_NODE = False
    CATEGORY = "utils"

    def generate_prompt_id(self, custom_prompt_id=""):
        if custom_prompt_id and custom_prompt_id.strip():
            # 用户输入了自定义ID
            prompt_id = custom_prompt_id.strip()
            print(f"📝 使用自定义PromptID: {prompt_id}")
        else:
            # 自动生成ID：时间戳 + 随机字符
            timestamp = str(int(time.time()))
            random_str = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
            prompt_id = f"prompt_{timestamp}_{random_str}"
            print(f"🎲 自动生成PromptID: {prompt_id}")

        return (prompt_id,)


class PlugAndPlayWebhook:
    """即插即用Webhook节点：连上线就能转发数据"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "webhook_url": ("STRING", {"default": "http://127.0.0.1:8010/handler/webhook",
                                           "placeholder": "https://your-api.com/webhook"}),
                "image_url": ("STRING", {"default": "",
                                           "placeholder": "图片的url"}),
            },
            "optional": {
                "prompt_id": ("STRING", {"default": ""}),
            },
            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO", "unique_id": "UNIQUE_ID"},
        }

    RETURN_TYPES = ()
    FUNCTION = "send"
    OUTPUT_NODE = True
    CATEGORY = "utils"

    def send(self, webhook_url, image_url, prompt_id="", prompt=None, extra_pnginfo=None, unique_id=None):
        if not webhook_url:
            raise ValueError("❌ 请填写Webhook URL！")

        # 使用传入的prompt_id，如果没有则用unique_id
        final_prompt_id = prompt_id or unique_id or "unknown"

        # 准备发送的数据
        data = {
            "img_base64": image_url,
            "format": "png",
            "image_url": image_url,
            "prompt_id": final_prompt_id,
            "timestamp": time.time()
        }

        # 发送Webhook
        try:
            response = requests.post(webhook_url, json=data)
            response.raise_for_status()
            print(f'发送的数据:{data}')
        except Exception as e:
            print(f"❌ 发送失败: {str(e)}")

        # 终端节点，无需返回
        return ()


class SaveImageWithOutput:
    """保存图片并输出的节点：既保存图片又能继续传递数据"""

    def __init__(self):
        self.output_dir = "output"
        self.type = "output"
        self.prefix_append = ""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "filename_prefix": ("STRING", {"default": "ComfyUI"}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("images", "file_path")
    FUNCTION = "save_image"
    OUTPUT_NODE = False
    CATEGORY = "image"

    def save_image(self, images, filename_prefix="ComfyUI"):
        filename_prefix += self.prefix_append
        full_output_folder, filename, counter, subfolder, filename_prefix = self.get_save_image_path(filename_prefix,
                                                                                                     self.output_dir)

        file_paths = []

        for i, image in enumerate(images):
            # 转换图片数据
            img_array = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8))

            # 生成文件名
            file = f"{filename_prefix}_{counter + i:05}_.png"
            file_path = os.path.join(full_output_folder, file)

            # 保存图片
            img.save(file_path, compress_level=4)
            file_paths.append(file_path)

        print(f"✅ 图片已保存到: {file_paths[0] if file_paths else '未知路径'}")

        # 直接返回元组：(原始图片数据, 第一个文件路径)
        return (images, file_paths[0] if file_paths else "")

    def get_save_image_path(self, filename_prefix, output_dir):
        def map_filename(filename):
            prefix_len = len(os.path.basename(filename_prefix))
            prefix = filename[:prefix_len + 1]
            try:
                digits = int(filename[prefix_len + 1:].split('_')[0])
            except:
                digits = 0
            return (digits, prefix)

        subfolder = ""
        full_output_folder = os.path.join(output_dir, subfolder)

        if os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) != output_dir:
            print("Saving image outside the output folder is not allowed.")
            return {}

        try:
            counter = max(filter(lambda a: a[1][:-1] == filename_prefix and a[1][-1] == "_",
                                 map(map_filename, os.listdir(full_output_folder))))[0] + 1
        except ValueError:
            counter = 1
        except FileNotFoundError:
            os.makedirs(full_output_folder, exist_ok=True)
            counter = 1

        return full_output_folder, filename_prefix, counter, subfolder, filename_prefix


NODE_CLASS_MAPPINGS = {
    "PlugAndPlayWebhook": PlugAndPlayWebhook,
    "SaveImageWithOutput": SaveImageWithOutput
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "PlugAndPlayWebhook": "Webhook转发器",
    "SaveImageWithOutput": "保存图片(带输出)"
}