diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..4de22b2 --- /dev/null +++ b/__init__.py @@ -0,0 +1,27 @@ +""" +Copyright (C) 2024 Dayuppy + + This library is free software; you can redistribute it and/or + modify it under the terms of the GNU Lesser General Public + License as published by the Free Software Foundation; either + version 2.1 of the License, or (at your option) any later version. + + This library is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + Lesser General Public License for more details. + + You should have received a copy of the GNU Lesser General Public + License along with this library; if not, write to the Free Software + Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 + USA + +@author: Dayuppy +@title: Discord Webhook +@nickname: DiscordWebhook +@description: A very simple Discord webhook integration node for ComfyUI that lets you post images and text. +""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..67d8bac --- /dev/null +++ b/nodes.py @@ -0,0 +1,217 @@ +""" +Copyright (C) 2024 Dayuppy + + This library is free software; you can redistribute it and/or + modify it under the terms of the GNU Lesser General Public + License as published by the Free Software Foundation; either + version 2.1 of the License, or (at your option) any later version. + + This library is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + Lesser General Public License for more details. + + You should have received a copy of the GNU Lesser General Public + License along with this library; if not, write to the Free Software + Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 + USA + +@author: Dayuppy +@title: Discord Webhook +@nickname: DiscordWebhook +@description: A very simple Discord webhook integration node for ComfyUI that lets you post images and text. +""" + +import os +import tempfile +import shutil +import numpy as np +import asyncio +from PIL import Image, ImageDraw +from discord_webhook import AsyncDiscordWebhook + +def create_default_image(): + """Create a simple TV test pattern image.""" + image = Image.new("RGB", (128, 128), "black") + colors = ["white", "yellow", "cyan", "green", "magenta", "red", "blue", "black"] + bar_width = 128 // len(colors) + draw = ImageDraw.Draw(image) + + for i, color in enumerate(colors): + draw.rectangle([i * bar_width, 0, (i + 1) * bar_width, 128], fill=color) + + return image + +class DiscordSetWebhook: + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("DUMMY IMAGE",) + OUTPUT_NODE = True + FUNCTION = "execute" + CATEGORY = "Discord" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"URL": ("STRING",)}} + + def execute(self, URL): + if not URL.startswith("https://discord.com/api/webhooks/"): + raise ValueError("Invalid URL format. URL should start with 'https://discord.com/api/webhooks/' Please reference https://support.discord.com/hc/en-us/articles/228383668-Intro-to-Webhooks") + + with open("discord_webhook_url.txt", "w") as f: + f.write(URL) + + return (create_default_image(),) + +class DiscordPostViaWebhook: + RETURN_TYPES = ("IMAGE",) + OUTPUT_NODE = True + FUNCTION = "execute" + CATEGORY = "Discord" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",) + }, + "optional": { + "send_Message": ("BOOLEAN", {"default": False}), + "send_Image": ("BOOLEAN", {"default": True}), + "message": ("STRING", {"default": "", "multiline": True}), + "prepend_message": ("STRING", {"default": "", "multiline": True}), + }, + } + + async def send_webhook(self, url, message, files=None): + """Send the webhook with the given message and up to 4 files.""" + webhook = AsyncDiscordWebhook(url=url, content=message[:2000], timeout=30.0) + if files: + for file in files: + webhook.add_file(file=file["data"], filename=file["name"]) + await webhook.execute() + + def process_image(self, image): + """Process the image (or batch of images) and return them in a format suitable for Discord.""" + images_to_send = [] + + if image is None: + image = create_default_image() + + # Check if it's a batched image (4D array: [batch_size, height, width, channels]) + if isinstance(image, np.ndarray): + if image.ndim == 4: + # Batch of images, process each image separately + for i in range(image.shape[0]): + img = image[i] # Don't squeeze unless we know the axis has size 1 + img = Image.fromarray(np.clip(img * 255, 0, 255).astype(np.uint8)) + images_to_send.append(img) + elif image.ndim == 3: + # Single image (3D array: [height, width, channels]) + image = Image.fromarray(np.clip(image * 255, 0, 255).astype(np.uint8)) + images_to_send.append(image) + else: + raise ValueError("Input image array must be 3D or 4D (batch of images).") + + elif hasattr(image, "cpu"): + array = image.cpu().numpy() + + # Handle batched images + if array.ndim == 4: # Batch of images + for i in range(array.shape[0]): + img_array = array[i] # Only access the i-th image, no squeeze needed + img = Image.fromarray(np.clip(img_array * 255, 0, 255).astype(np.uint8)) + images_to_send.append(img) + elif array.ndim == 3: + # Single image + array = np.clip(array * 255, 0, 255).astype(np.uint8) + images_to_send.append(Image.fromarray(array)) + else: + raise ValueError("Input tensor must be 3D or 4D (batch of images).") + + else: + images_to_send.append(image) + + # Save each image to a temporary file and collect file data + files = [] + temp_dir = tempfile.mkdtemp() + + for idx, img in enumerate(images_to_send): + file_path = os.path.join(temp_dir, f"image_{idx}.png") + img.save(file_path, format="PNG", compress_level=1) + + # If the file size exceeds 20MB, resize and save again + if os.path.getsize(file_path) > 20 * 1024 * 1024: + img = img.resize((img.width // 2, img.height // 2)) + img.save(file_path, format="PNG", compress_level=9) + + with open(file_path, "rb") as f: + files.append({"data": f.read(), "name": f"image_{idx}.png"}) + + # Clean up the temporary directory + shutil.rmtree(temp_dir) + + return files + + async def execute(self, image, send_Message=True, send_Image=True, message="", prepend_message=""): + with open("discord_webhook_url.txt", "r") as f: + webhook_url = f.read().strip() + + if not webhook_url: + raise ValueError("Webhook URL is empty.") + + if send_Message: + message = f"{prepend_message}\n{message}" + + files = self.process_image(image) if send_Image else None + + if files: + # Split files into batches of 4 (Discord limit) + batches = [files[i:i + 4] for i in range(0, len(files), 4)] + + # Send multiple webhooks if necessary + for batch in batches: + await self.send_webhook(webhook_url, message, batch) + else: + # No images to send, just send the message + await self.send_webhook(webhook_url, message) + + return (image,) + +NODE_CLASS_MAPPINGS = { + "DiscordSetWebhook": DiscordSetWebhook, + "DiscordPostViaWebhook": DiscordPostViaWebhook +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "DiscordSetWebhook": "Set Discord Webhook", + "DiscordPostViaWebhook": "Use Discord Webhook" +} +class DiscordPostViaWebhookWithInput(DiscordPostViaWebhook): + """ + Variante du node avec entrée connectable pour le message (prompt). + Basée sur DiscordPostViaWebhook pour ne rien casser. + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + "optional": { + "send_Message": ("BOOLEAN", {"default": False}), + "send_Image": ("BOOLEAN", {"default": True}), + # 👇 Ajout important: on force une ENTRÉE (socket) pour le texte + "message": ("STRING", {"default": "", "multiline": True, "forceInput": True}), + # Optionnel: on peut aussi forcer une entrée pour le prepend + "prepend_message": ("STRING", {"default": "", "multiline": True, "forceInput": True}), + }, + } + +# Ajoute ces lignes à la fin des mappings existants +NODE_CLASS_MAPPINGS.update({ + "DiscordPostViaWebhookWithInput": DiscordPostViaWebhookWithInput +}) + +NODE_DISPLAY_NAME_MAPPINGS.update({ + "DiscordPostViaWebhookWithInput": "Use Discord Webhook (+message input)" +})