229 lines
7.9 KiB
Python
229 lines
7.9 KiB
Python
"""
|
|
Discord Webhook nodes (safe version: no URL stored in workflows/PNGs)
|
|
|
|
Base: Dayuppy
|
|
Mod: +message input + env/secrets for webhook URL
|
|
"""
|
|
|
|
import os
|
|
import io
|
|
import json
|
|
import tempfile
|
|
import shutil
|
|
import numpy as np
|
|
import asyncio
|
|
from PIL import Image, ImageDraw
|
|
from discord_webhook import AsyncDiscordWebhook
|
|
|
|
# === Secrets handling ===
|
|
BASE_DIR = os.path.dirname(__file__)
|
|
SECRETS_DIR = os.path.join(BASE_DIR, "secrets")
|
|
WEBHOOK_FILE = os.path.join(SECRETS_DIR, "webhook.txt")
|
|
LEGACY_FILE = os.path.join(BASE_DIR, "discord_webhook_url.txt") # backward-compat, deprecated
|
|
|
|
def _read_file(path):
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return f.read().strip()
|
|
except Exception:
|
|
return ""
|
|
|
|
def get_webhook_url():
|
|
# 1) env var (recommended)
|
|
env_url = os.getenv("DISCORD_WEBHOOK_URL", "").strip()
|
|
if env_url:
|
|
return env_url
|
|
# 2) local secret file
|
|
if os.path.exists(WEBHOOK_FILE):
|
|
return _read_file(WEBHOOK_FILE)
|
|
# 3) legacy fallback for old setups
|
|
if os.path.exists(LEGACY_FILE):
|
|
return _read_file(LEGACY_FILE)
|
|
return ""
|
|
|
|
# === Utilities ===
|
|
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
|
|
|
|
def _tensor_or_nd_to_pil_list(image):
|
|
"""
|
|
Accepts:
|
|
- None -> default image
|
|
- numpy ndarray HxWxC or NxHxWxC
|
|
- torch tensor-like with .cpu().numpy()
|
|
- PIL.Image.Image
|
|
Returns: list of PIL images
|
|
"""
|
|
images_to_send = []
|
|
|
|
if image is None:
|
|
images_to_send.append(create_default_image())
|
|
return images_to_send
|
|
|
|
if isinstance(image, Image.Image):
|
|
images_to_send.append(image)
|
|
return images_to_send
|
|
|
|
if isinstance(image, np.ndarray):
|
|
if image.ndim == 4:
|
|
for i in range(image.shape[0]):
|
|
arr = np.clip(image[i] * 255, 0, 255).astype(np.uint8)
|
|
images_to_send.append(Image.fromarray(arr))
|
|
elif image.ndim == 3:
|
|
arr = np.clip(image * 255, 0, 255).astype(np.uint8)
|
|
images_to_send.append(Image.fromarray(arr))
|
|
else:
|
|
raise ValueError("Input numpy array must be 3D or 4D.")
|
|
return images_to_send
|
|
|
|
if hasattr(image, "cpu"): # torch-like
|
|
array = image.cpu().numpy()
|
|
if array.ndim == 4:
|
|
for i in range(array.shape[0]):
|
|
arr = np.clip(array[i] * 255, 0, 255).astype(np.uint8)
|
|
images_to_send.append(Image.fromarray(arr))
|
|
elif array.ndim == 3:
|
|
arr = np.clip(array * 255, 0, 255).astype(np.uint8)
|
|
images_to_send.append(Image.fromarray(arr))
|
|
else:
|
|
raise ValueError("Input tensor must be 3D or 4D.")
|
|
return images_to_send
|
|
|
|
# Fallback: try to treat as PIL
|
|
images_to_send.append(image)
|
|
return images_to_send
|
|
|
|
def _prepare_files_for_discord(images_to_send):
|
|
"""Save PIL images to temp PNGs, enforce <20MB, return list of dicts {data, name}."""
|
|
files = []
|
|
temp_dir = tempfile.mkdtemp()
|
|
try:
|
|
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 >20MB, resize and recompress
|
|
if os.path.getsize(file_path) > 20 * 1024 * 1024:
|
|
w, h = img.width, img.height
|
|
img_small = img.resize((max(1, w // 2), max(1, h // 2)))
|
|
img_small.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"})
|
|
finally:
|
|
shutil.rmtree(temp_dir)
|
|
return files
|
|
|
|
# === Nodes ===
|
|
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. Must start with 'https://discord.com/api/webhooks/'.")
|
|
os.makedirs(SECRETS_DIR, exist_ok=True)
|
|
with open(WEBHOOK_FILE, "w", encoding="utf-8") as f:
|
|
f.write(URL.strip())
|
|
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):
|
|
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):
|
|
imgs = _tensor_or_nd_to_pil_list(image)
|
|
return _prepare_files_for_discord(imgs) if imgs else None
|
|
|
|
async def execute(self, image, send_Message=True, send_Image=True, message="", prepend_message=""):
|
|
webhook_url = get_webhook_url()
|
|
if not webhook_url:
|
|
raise ValueError("Missing Discord webhook URL. Set env DISCORD_WEBHOOK_URL "
|
|
"or create secrets/webhook.txt")
|
|
|
|
final_message = ""
|
|
if send_Message:
|
|
final_message = (prepend_message or "").strip()
|
|
if final_message:
|
|
final_message += "\n"
|
|
final_message += (message or "")
|
|
|
|
files = self._process_image(image) if send_Image else None
|
|
|
|
if files:
|
|
# Discord: max 10 files per message nowadays, mais on reste safe à 4
|
|
batches = [files[i:i + 4] for i in range(0, len(files), 4)]
|
|
for batch in batches:
|
|
await self._send_webhook(webhook_url, final_message, batch)
|
|
else:
|
|
await self._send_webhook(webhook_url, final_message)
|
|
|
|
return (image,)
|
|
|
|
class DiscordPostViaWebhookWithInput(DiscordPostViaWebhook):
|
|
"""
|
|
Variant with connectable inputs for message/prepend (safe: no URL stored in workflow).
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"message": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
|
|
},
|
|
"optional": {
|
|
"send_Message": ("BOOLEAN", {"default": False}),
|
|
"send_Image": ("BOOLEAN", {"default": True}),
|
|
"prepend_message": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
|
|
},
|
|
}
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DiscordSetWebhook": DiscordSetWebhook,
|
|
"DiscordPostViaWebhook": DiscordPostViaWebhook,
|
|
"DiscordPostViaWebhookWithInput": DiscordPostViaWebhookWithInput,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DiscordSetWebhook": "Set Discord Webhook",
|
|
"DiscordPostViaWebhook": "Use Discord Webhook",
|
|
"DiscordPostViaWebhookWithInput": "Use Discord Webhook (+message input)",
|
|
}
|