diff --git a/__init__.py b/__init__.py index b8f8d30..0b8155c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,8 @@ from .nodes.http_serving import HTTPServing from .nodes.all_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .nodes.telegram_serving import TelegramServing +NODE_CLASS_MAPPINGS["TelegramServing"] = TelegramServing +NODE_DISPLAY_NAME_MAPPINGS["TelegramServing"] = "Telegram Serving" NODE_CLASS_MAPPINGS["HTTPServing"] = HTTPServing diff --git a/nodes/all_nodes.py b/nodes/all_nodes.py index 67c0510..7c1e00c 100644 --- a/nodes/all_nodes.py +++ b/nodes/all_nodes.py @@ -32,7 +32,7 @@ class ServingOutput: "image": ("IMAGE",), "frame_duration": ("INT", {"default": 30, "min": 1, "step": 1, "max": 9999999}), }, - } +} RETURN_TYPES = () # RETURN_NAMES = ("image_output_name",) diff --git a/nodes/telegram_serving.py b/nodes/telegram_serving.py new file mode 100644 index 0000000..ad6c8a5 --- /dev/null +++ b/nodes/telegram_serving.py @@ -0,0 +1,95 @@ +import telebot +from telebot import types +from collections import deque +import threading +import io +from PIL import Image +import base64 +import numpy as np +import torch +from .utils import tensorToImageConversion, parse_command_string + +class TelegramServing: + def __init__(self): + self.data_ready = threading.Event() + self.data = deque() + self.telegram_running = False + self.bot = None + self.command_name = None + + def telegram_handler(self): + @self.bot.message_handler(commands=[self.command_name]) + def handle_command(message): + print(f"Received command from {message.chat.id}: {message.text}") + parsed_data = parse_command_string(message.text, self.command_name) + + async def serve_multi_image_function(images): + media_group = [] + for i, img in enumerate(images): + img_np = (img.cpu().numpy() * 255).astype('uint8') + img_pil = Image.fromarray(img_np.squeeze()) + img_bytes = io.BytesIO() + img_pil.save(img_bytes, format='PNG') + img_bytes.seek(0) + media_group.append(types.InputMediaPhoto(img_bytes)) + + self.bot.send_media_group(message.chat.id, media_group) + + def serve_image_function(image, frame_duration): + image_file = tensorToImageConversion(image, frame_duration) + self.bot.send_photo(message.chat.id, image_file) + + parsed_data["serve_image_function"] = serve_image_function + parsed_data["serve_multi_image_function"] = serve_multi_image_function + parsed_data["serve_text_function"] = lambda text: self.bot.reply_to(message, text) + + if message.document: + file_info = self.bot.get_file(message.document.file_id) + downloaded_file = self.bot.download_file(file_info.file_path) + parsed_data["attachment_url_0"] = downloaded_file + + self.data.append(parsed_data) + self.data_ready.set() + + self.bot.polling() + + def get_data(self): + if not self.data: + self.data_ready.wait() + data = self.data.popleft() + self.data_ready.clear() + return data + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "telegram_token": ("STRING", { + "multiline": False, + "default": "" + }), + "command_name": ("STRING", { + "multiline": False, + "default": "generate" + }) + } + } + + RETURN_TYPES = ("SERVING_CONFIG",) + RETURN_NAMES = ("Serving config",) + FUNCTION = "serve" + CATEGORY = "Serving-Toolkit" + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("NaN") + + def serve(self, telegram_token, command_name): + if not self.telegram_running: + self.bot = telebot.TeleBot(telegram_token) + self.command_name = command_name + threading.Thread(target=self.telegram_handler, daemon=True).start() + print(f"Telegram bot running, listening for /{command_name} commands") + self.telegram_running = True + + data = self.get_data() + return (data,) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index be1e6ca..6503bdb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "comfyui-serving-toolkit" description = "This extension adds nodes that allow you to easily serve your workflow (for example using a discord bot) " version = "1.0.0" license = "LICENSE" -dependencies = ["discord.py", "websocket-client", "rel"] +dependencies = ["discord.py", "websocket-client", "rel", "pyTelegramBotAPI"] [project.urls] Repository = "https://github.com/matan1905/ComfyUI-Serving-Toolkit" diff --git a/requirements.txt b/requirements.txt index 661a6a9..aed2a35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ discord.py websocket-client rel +pyTelegramBotAPI opencv-python \ No newline at end of file