merge

Add an html input for http serving

.gitignore

@@ -1,4 +1,4 @@
-__pycache__/*
+__pycache__/
 /examples/node_modules/*
 
 .idea

README.md

@@ -2,7 +2,6 @@
 Welcome to the ComfyUI Serving Toolkit, a powerful tool for serving image generation workflows in Discord and other platforms (soon).
 This toolkit is designed to simplify the process of serving your ComfyUI workflow, making image generation bots easier than ever before.
 You can serve on discord, or on websockets.
-* allowed_user - (Optional) The Telegram username of the user allowed to use the bot. If not set, all users can use the bot.
 
 If you need any help, Checkout the [Discord!](https://discord.gg/AyHFchFZuV)
 
@@ -22,11 +21,14 @@ or put all the files inside custom_nodes and run:
 ```
 
 ## The simplest configuration
-Here a simple workflow that will get a !generate \<prompt> and resond with an image
-![image](https://github.com/matan1905/ComfyUI-Serving-Toolkit/assets/24731932/e193be18-7b83-4f44-b119-21230f0b9a16)
+Here a simple workflow that will get a !generate \<prompt> and respond with an existing image, and !help will produce a help message
+![image](https://github.com/user-attachments/assets/bf4bfa6d-3252-40ab-88a4-7ee24147f046)
+
 
 You can copy the workflow json:
-[discordserv.json](https://github.com/matan1905/ComfyUI-Serving-Toolkit/files/13248566/discordserv.json)
+[workflow.json](https://github.com/user-attachments/files/17041080/workflow.4.json)
+
+
 
 
 
@@ -38,6 +40,7 @@ After setting up your workflow, In order for the serving to always be up, you ne
 #### This will require you to keep your ComfyUI and computer running.
 
 ## Nodes
+
 **DiscordServing**
 
 This node is an essencial part of the serving, queueing the prompt it will wait for a single message, process it and optionally return the image.
@@ -65,7 +68,25 @@ Inputs:
 
 Outputs:
 * Serving Config - A basic reference for this serving, used by the other nodes of this toolkit to get arguments and return images.
-  
+
+
+**CommandPickerServing**
+This node allows you to define specific behaviour for specific commands - it will only execute outputs that are connected to it when that specific command is given.
+Note, the command in HTTP Serving is the path after /, and the command in websocket is "command_name" json property
+
+Inputs:
+* serving_config -  a config made by a serving node
+* should_execute - an output of one of the output nodes, will only send that as part of the results if the command received is the one on this node
+* command - the command to filter by
+
+**CommandPickerServing**
+
+**AlwaysExecute**
+For when you don't care about commands, it will always return an output regardless of command given. this will execute for ALL commands, no matter if the commands are registered or not.
+
+Inputs:
+* serving_config -  a config made by a serving node
+* should_execute - an output of one of the output nodes
 
 
 **ServingInputText** 

nodes/all_nodes.py

@@ -2,7 +2,7 @@ import time
 import threading
 from .discord_client import discord_client
 from collections import deque
-from .utils import parse_command_string, tensorToImageConversion
+from .utils import parse_command_string, tensorToImageConversion, CommandRegistry
 import discord
 import asyncio
 import requests
@@ -15,6 +15,7 @@ import json
 import torch
 import torchvision.transforms as transforms
 import cv2
+import nodes
 
 
 
@@ -32,20 +33,23 @@ class ServingOutput:
                 "image": ("IMAGE",),
                 "frame_duration": ("INT", {"default": 30, "min": 1, "step": 1, "max": 9999999}),
             },
+            "optional": {
+                "chained_execution": ("SHOULD_EXECUTE",),
+            }
 }
 
-    RETURN_TYPES = ()
+    RETURN_TYPES = ("SHOULD_EXECUTE",)
     # RETURN_NAMES = ("image_output_name",)
 
     FUNCTION = "out"
 
-    OUTPUT_NODE = True
+    # OUTPUT_NODE = True
 
     CATEGORY = "Serving-Toolkit"
 
-    def out(self, image,serving_config,frame_duration):
+    def out(self, image,serving_config,frame_duration, chained_execution=None):
         serving_config["serve_image_function"](image,frame_duration)
-        return {}
+        return ({}, )
 
 
 
@@ -60,19 +64,22 @@ class ServingTextOutput:
                 "serving_config": ("SERVING_CONFIG",),
                 "text": ("STRING", {"multiline": True, "default": ""}),
             },
+            "optional": {
+                "chained_execution": ("SHOULD_EXECUTE",),
+            }
         }
 
-    RETURN_TYPES = ()
+    RETURN_TYPES = ("SHOULD_EXECUTE",)
     FUNCTION = "out"
-    OUTPUT_NODE = True
+    # OUTPUT_NODE = True
     CATEGORY = "Serving-Toolkit"
 
-    def out(self, serving_config, text):
+    def out(self, serving_config, text, chained_execution=None):
         if "serve_text_function" in serving_config:
             serving_config["serve_text_function"](text)
         else:
             print("Warning: serve_text_function not found in serving_config")
-        return {}
+        return ({}, )
 
 
 class ServingInputText:
@@ -109,6 +116,71 @@ class ServingInputText:
             return (default,)
         return (serving_config[argument],)
 
+class ServingInputTextImage:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "serving_config": ("SERVING_CONFIG",),
+                "argument": ("STRING", {
+                    "multiline": False,
+                    "default": "prompt"
+                }),
+                "default_prompt": ("STRING", {
+                    "multiline": True,
+                    "default": ""
+                }),
+            },
+            "optional":{
+                "default_image": ("IMAGE",)
+            }
+        }
+
+    RETURN_TYPES = ("STRING", "IMAGE",)
+    FUNCTION = "out"
+    CATEGORY = "Serving-Toolkit"
+
+    def convert_color(self, image):
+        if len(image.shape) > 2 and image.shape[2] >= 4:
+            return cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
+        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    def load_image(self, base64_str):
+        nparr = np.frombuffer(base64.b64decode(base64_str), np.uint8)
+        result = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
+        result = self.convert_color(result)
+        result = result.astype(np.float32) / 255.0
+        image = torch.from_numpy(result)[None,]
+        return image
+
+    def out(self, serving_config, argument, default_prompt, default_image = None):
+        attachment_url_key = "attachment_url_0"
+        if attachment_url_key not in serving_config:
+            if default_image is not None:
+                return (default_image,)
+            serving_config["serve_text_function"]("This command requires an image")
+            nodes.interrupt_processing(True)
+            return ('', None)
+
+        attachment_url = serving_config[attachment_url_key]
+        response = requests.get(attachment_url)
+        image = Image.open(io.BytesIO(response.content)).convert("RGB")
+
+        # Convert PIL image to base64 string
+        image_file = io.BytesIO()
+        image.save(image_file, format='PNG')
+        image_file.seek(0)
+        base64_img = base64.b64encode(image_file.read()).decode('utf-8')
+
+        # Use the base64 string to get the image tensor
+        img_out = self.load_image(base64_img)
+        if argument not in serving_config:
+            return (default_prompt, img_out)
+        return (serving_config[argument], img_out)
+
 
 class ServingInputNumber:
     def __init__(self):
@@ -160,16 +232,18 @@ class ServingMultiImageOutput:
                 "serving_config": ("SERVING_CONFIG",),
                 "images": ("IMAGE",),
             },
+            "optional": {
+                "chained_execution": ("SHOULD_EXECUTE",),
+            }
         }
 
-    RETURN_TYPES = ()
+    RETURN_TYPES = ("SHOULD_EXECUTE",)
     FUNCTION = "out"
-    OUTPUT_NODE = True
+    # OUTPUT_NODE = True
     CATEGORY = "Serving-Toolkit"
 
-    def out(self, serving_config, images):
+    def out(self, serving_config, images, chained_execution=None):
         if "serve_multi_image_function" in serving_config:
-            print(f"MultiImageOutput: Received {images.shape[0]} images")
             serve_func = serving_config["serve_multi_image_function"]
             if asyncio.iscoroutinefunction(serve_func):
                 future = asyncio.run_coroutine_threadsafe(serve_func(images), discord_client.loop)
@@ -181,7 +255,7 @@ class ServingMultiImageOutput:
                 serve_func(images)
         else:
             print("Warning: serve_multi_image_function not found in serving_config")
-        return {}
+        return ({}, )
 
 class DiscordServing():
     discord_running = False
@@ -190,7 +264,7 @@ class DiscordServing():
         self.data_ready = threading.Event()
         self.data = deque()
         self.discord_token = None
-        pass
+        self.command_registry = CommandRegistry()
 
     def discord_runner(self):
          discord_client.run(self.discord_token)
@@ -209,10 +283,6 @@ class DiscordServing():
                 "discord_token": ("STRING", {
                     "multiline": True,
                     "default": ""
-                }),
-                "command_name": ("STRING", {
-                    "multiline": False,
-                    "default": "generate"
                 })
             }
         }
@@ -225,74 +295,81 @@ class DiscordServing():
     @classmethod
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
-    # OUTPUT_NODE = False
 
     CATEGORY = "Serving-Toolkit"
 
-    def serve(self, command_name, discord_token):
+    def serve(self, discord_token):
         if not DiscordServing.discord_running:
             self.discord_token = discord_token
             run_discord = threading.Thread(target=self.discord_runner)
             run_discord.start()
             print("Client running")
             DiscordServing.discord_running = True
-        if not self.registered_command: 
-            self.registered_command = True   
-            @discord_client.command(name=command_name)
-            async def execute(ctx):
-                parsed_data = parse_command_string(ctx.message.content,command_name)
+        if not self.registered_command:
+            self.registered_command = True
+            @discord_client.event
+            async def on_message(message):
+                if message.content.startswith('!'):
+                    command_name = message.content.split()[0][1:]
+
+                    if not self.command_registry.has_command(command_name):
+                        return
+
+                    parsed_data = parse_command_string(message.content, command_name)
+
+                    async def serve_multi_image_function(images):
+                        discord_files = []
+                        for i, img in enumerate(images):
+                            try:
+                                if isinstance(img, np.ndarray):
+                                    img_np = img
+                                else:
+                                    img_np = img.cpu().numpy()
+
+                                if img_np.dtype != np.uint8:
+                                    img_np = (img_np * 255).astype(np.uint8)
+
+                                if len(img_np.shape) == 2:
+                                    img_np = np.stack([img_np] * 3, axis=-1)
+                                elif len(img_np.shape) == 3 and img_np.shape[2] == 1:
+                                    img_np = np.concatenate([img_np] * 3, axis=2)
+                                elif img_np.shape[2] == 4:
+                                    img_np = img_np[:, :, :3]
+
+                                img_pil = Image.fromarray(img_np)
+                                img_bytes = io.BytesIO()
+                                img_pil.save(img_bytes, format='PNG')
+                                img_bytes.seek(0)
+                                discord_files.append(discord.File(img_bytes, filename=f'image_{i}.png'))
+                            except Exception as e:
+                                print(f"Error processing image {i + 1}: {str(e)}")
 
-                async def serve_multi_image_function(images):
-                    discord_files = []
-                    for i, img in enumerate(images):
                         try:
-                            if isinstance(img, np.ndarray):
-                                img_np = img
-                            else:
-                                img_np = img.cpu().numpy()
-
-                            if img_np.dtype != np.uint8:
-                                img_np = (img_np * 255).astype(np.uint8)
-
-                            if len(img_np.shape) == 2:
-                                img_np = np.stack([img_np] * 3, axis=-1)
-                            elif len(img_np.shape) == 3 and img_np.shape[2] == 1:
-                                img_np = np.concatenate([img_np] * 3, axis=2)
-                            elif img_np.shape[2] == 4:
-                                img_np = img_np[:, :, :3]
-
-                            img_pil = Image.fromarray(img_np)
-                            img_bytes = io.BytesIO()
-                            img_pil.save(img_bytes, format='PNG')
-                            img_bytes.seek(0)
-                            discord_files.append(discord.File(img_bytes, filename=f'image_{i}.png'))
+                            print("Attempting to send images...")
+                            await message.reply(files=discord_files)
+                            print("Images sent successfully")
                         except Exception as e:
-                            print(f"Error processing image {i + 1}: {str(e)}")
+                            print(f"Error sending images: {str(e)}")
 
-                    try:
-                        print("Attempting to send images...")
-                        await ctx.reply(files=discord_files)
-                        print("Images sent successfully")
-                    except Exception as e:
-                        print(f"Error sending images: {str(e)}")
+                    def serve_image_function(image, frame_duration):
+                        image_file = tensorToImageConversion(image, frame_duration)
+                        asyncio.run_coroutine_threadsafe(message.reply(file=discord.File(image_file, filename='image.webp')), discord_client.loop)
 
-                def serve_image_function(image, frame_duration):
-                    image_file = tensorToImageConversion(image, frame_duration)
-                    asyncio.run_coroutine_threadsafe(ctx.reply(file=discord.File(image_file, filename='image.webp')), discord_client.loop)
-                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: asyncio.run_coroutine_threadsafe(
-                    ctx.reply(content=text), discord_client.loop)
-                parsed_data.update(
-                    {f"attachment_url_{i}": attachment.url for i, attachment in enumerate(ctx.message.attachments)})
-                parsed_data.update({f"attachment_url_{i}": attachment.url for i, attachment in enumerate(ctx.message.attachments)}) # populates all the attachments urls
-                self.data.append(parsed_data)
-                self.data_ready.set()
+                    def is_command(command):
+                        return bool(command_name == command)
 
-        data = self.get_data() 
+                    parsed_data["serve_image_function"] = serve_image_function
+                    parsed_data["serve_multi_image_function"] = serve_multi_image_function
+                    parsed_data["is_command"] = is_command
+                    parsed_data["serve_text_function"] = lambda text: asyncio.run_coroutine_threadsafe(
+                        message.reply(content=text), discord_client.loop)
+                    parsed_data.update({f"attachment_url_{i}": attachment.url for i, attachment in enumerate(message.attachments)})
+                    self.data.append(parsed_data)
+                    self.data_ready.set()
+
+        data = self.get_data()
 
         return (data,)
-
 class WebSocketServing():
     def __init__(self):
         self.data_ready = threading.Event()
@@ -300,10 +377,13 @@ class WebSocketServing():
         self.ws_running = False
         self.websocket_url= None
         self.ws = None
+        self.command_registry = CommandRegistry()
         pass
     def on_message(self,ws,message):
         try:
             parsed = json.loads(message)
+            if "command_name" in parsed and not self.command_registry.has_command(parsed["command_name"]):
+                return
             self.data.append(parsed)
             self.data_ready.set()
         except Exception as e:
@@ -400,10 +480,14 @@ class WebSocketServing():
                     }
                     self.ws.send(json.dumps(response))
 
+        def is_command(command):
+            return bool(data["command_name"] == command)
+
         data["serve_image_function"] = serve_image_function
         data["serve_multi_image_function"] = serve_multi_image_function
         data["serve_text_function"] = lambda text: self.ws.send(
             json.dumps({"text": text, "_requestId": data["_requestId"]}))
+        data["is_command"] = is_command
 
         return (data,)
 
@@ -445,7 +529,9 @@ class ServingInputImage:
         if attachment_url_key not in serving_config:
             if default_image is not None:
                 return (default_image,)
-            raise ValueError("No attachment found in serving_config")
+            serving_config["serve_text_function"]("This command requires an image")
+            nodes.interrupt_processing(True)
+            return (None,)
 
         attachment_url = serving_config[attachment_url_key]
         response = requests.get(attachment_url)
@@ -495,7 +581,9 @@ class ServingInputImageAsLatent:
         if attachment_url_key not in serving_config:
             if default_latent is not None:
                 return (default_latent,)
-            raise ValueError("No attachment found in serving_config")
+            serving_config["serve_text_function"]("This command requires an image")
+            nodes.interrupt_processing(True)
+            return (None,)
 
         attachment_url = serving_config[attachment_url_key]
         response = requests.get(attachment_url)
@@ -512,18 +600,86 @@ class ServingInputImageAsLatent:
         return ({"samples": t},)
 
 
+registered_commands = []
+
+class CommandPickerServing:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "serving_config": ("SERVING_CONFIG", {"lazy": True}),
+                "command_name": ("STRING",),
+                "should_execute": ("SHOULD_EXECUTE", {"lazy": True}),
+            },
+        }
+
+    FUNCTION = "out"
+    CATEGORY = "Serving-Toolkit"
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+
+    def __init__(self):
+        self.commandRegistry = CommandRegistry()
+
+    def check_lazy_status(self, serving_config, command_name, should_execute=None):
+        if serving_config is None:
+            self.commandRegistry.add_command(command_name)
+            return ["serving_config"]
+
+        if serving_config["is_command"](command_name):
+            return ["should_execute", "serving_config"]
+        return None
+
+    def out(self, serving_config, command_name, should_execute=None):
+        if serving_config.get("finalize") and should_execute is not None:
+            serving_config["finalize"]()
+        return {}
+
+
+class AlwaysExecute:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "serving_config": ("SERVING_CONFIG", {"lazy": True}),
+                "should_execute": ("SHOULD_EXECUTE", {"lazy": True}),
+            },
+        }
+
+    FUNCTION = "out"
+    CATEGORY = "Serving-Toolkit"
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+
+    def __init__(self):
+        self.command_registry = CommandRegistry()
+
+    def check_lazy_status(self, serving_config, should_execute):
+        if serving_config is None:
+            self.command_registry.add_catch_all()
+            return ["serving_config"]
+
+        return ["serving_config", "should_execute"]
+
+    def out(self, serving_config, should_execute):
+        if serving_config.get("finalize"):
+            serving_config["finalize"]()
+        return {}
 
 # A dictionary that contains all nodes you want to export with their names
 NODE_CLASS_MAPPINGS = {
     "ServingOutput": ServingOutput,
     "ServingInputText": ServingInputText,
+    "ServingInputTextImage": ServingInputTextImage,
     "ServingInputNumber": ServingInputNumber,
     "DiscordServing": DiscordServing,
     "WebSocketServing": WebSocketServing,
     "ServingInputImage": ServingInputImage,
     "ServingTextOutput": ServingTextOutput,
     "ServingMultiImageOutput": ServingMultiImageOutput,
-    "ServingInputImageAsLatent": ServingInputImageAsLatent
+    "ServingInputImageAsLatent": ServingInputImageAsLatent,
+    "CommandPickerServing": CommandPickerServing,
+    "AlwaysExecute": AlwaysExecute
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -532,10 +688,13 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DiscordServing": "Discord Serving",
     "WebSocketServing": "WebSocket Serving",
     "ServingInputText": "Serving Input Text",
+    "ServingInputTextImage": "Serving Input Text & Image",
     "ServingInputNumber": "Serving Input Number",
     "ServingInputImage": "Serving Input Image",
     "ServingTextOutput": "Serving Text Output",
     "ServingMultiImageOutput": "Serving Multi-Image Output",
-    "ServingInputImageAsLatent": "Serving Input Image as Latent"
+    "ServingInputImageAsLatent": "Serving Input Image as Latent",
+    "CommandPickerServing": "Command Picker Serving",
+    "AlwaysExecute": "Always Execute"
 }
 

nodes/http_serving.py

@@ -4,7 +4,7 @@ from collections import deque
 import json
 import base64
 from io import BytesIO
-from .utils import tensorToImageConversion
+from .utils import tensorToImageConversion, CommandRegistry
 from PIL import Image
 
 
@@ -18,22 +18,34 @@ class HTTPServing:
         self.server = None
         self.output_ready = threading.Event()
         self.output = None
+        self.html_content = None
+        self.path = None
+        self.command_registry = CommandRegistry()
 
     def http_handler(self):
         class RequestHandler(BaseHTTPRequestHandler):
             def do_OPTIONS(self2):
-                if(self.enable_cross_origin_requests):
+                if (self.enable_cross_origin_requests):
                     self2.send_response(200)
                     self2.send_header('Access-Control-Allow-Origin', '*')
                     self2.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
                     self2.send_header('Access-Control-Allow-Headers', '*')
                     self2.end_headers()
 
-
             def do_POST(self2):
+                print(f"Received POST request: {self2.path}")
                 content_length = int(self2.headers['Content-Length'])
                 post_data = self2.rfile.read(content_length)
                 data = json.loads(post_data.decode('utf-8'))
+                self.path = self2.path
+
+                if not self.command_registry.has_command(self.path[1:]):
+                    self2.send_response(404)
+                    self2.send_header('Content-type', 'application/json')
+                    self2.end_headers()
+                    self2.wfile.write(json.dumps({"error": "Command not found"}).encode('utf-8'))
+                    return
+
                 self.data.append(data)
                 self.data_ready.set()
 
@@ -43,7 +55,7 @@ class HTTPServing:
                 self2.send_response(200)
                 self2.send_header('Content-type', 'application/json')
                 # Cors
-                if(self.enable_cross_origin_requests):
+                if (self.enable_cross_origin_requests):
                     self2.send_header('Access-Control-Allow-Origin', '*')
                     self2.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
                     self2.send_header('Access-Control-Allow-Headers', '*')
@@ -51,18 +63,14 @@ class HTTPServing:
                 self2.end_headers()
                 self2.wfile.write(json.dumps(response).encode('utf-8'))
 
-
-
-
             def do_GET(self2):
                 if self2.path == '/':
                     self2.send_response(200)
                     self2.send_header('Content-type', 'text/html')
                     self2.end_headers()
-                    self2.wfile.write(b"HTTP Serving is running, to send data make a POST request to this endpoint.")
+                    self2.wfile.write(self.html_content.encode('utf-8'))
 
         self.server = HTTPServer(('', self.port), RequestHandler)
-        print(f"HTTP Server running on port {self.port}")
         self.server.serve_forever()
 
     def get_data(self):
@@ -78,6 +86,28 @@ class HTTPServing:
             "required": {
                 "port": ("INT", {"default": 8000, "min": 1, "max": 65535}),
                 "enable_cross_origin_requests": ("BOOLEAN", {"default": False}),
+                "html_content": ("STRING", {"multiline": True, "default": """
+                    <!DOCTYPE html>
+                    <html>
+                    <body>
+                        <input type="text" id="prompt" value="">
+                        <button onclick="sendPost()">Send</button>
+                        <img id="resultImage" style="display:none">
+                    
+                        <script>
+                            function sendPost() {
+                                const prompt = document.getElementById('prompt').value;
+                                fetch(window.location.href, { method: 'POST', body: JSON.stringify({ prompt }) })
+                                    .then(response => response.json())
+                                    .then(data => {
+                                        document.getElementById('resultImage').src = 'data:image/png;base64,' + data.base64_img;
+                                        document.getElementById('resultImage').style.display = 'block';
+                                    });
+                            }
+                        </script>
+                    </body>
+                    </html>
+                """})
             }
         }
 
@@ -85,20 +115,22 @@ class HTTPServing:
     RETURN_NAMES = ("Serving config",)
     FUNCTION = "serve"
     CATEGORY = "Serving-Toolkit"
+
     @classmethod
     def IS_CHANGED(cls, **kwargs):
         return float("NaN")
 
-    def serve(self, port,enable_cross_origin_requests):
+    def serve(self, port, enable_cross_origin_requests, html_content):
         self.enable_cross_origin_requests = enable_cross_origin_requests
-
+        self.html_content = html_content
         if not self.http_running:
             self.port = port
             threading.Thread(target=self.http_handler, daemon=True).start()
             print(f"HTTP Server running on port {port}")
             self.http_running = True
 
-        self.output_ready.clear() # Prevent deadlock if failed in previous run
+        self.output = None
+        self.output_ready.clear()  # Prevent deadlock if failed in previous run
         data = self.get_data()
 
         def serve_multi_image_function(images):
@@ -109,31 +141,37 @@ class HTTPServing:
                 Image.fromarray(img_np.squeeze()).save(img_bytes, format='PNG')
                 base64_img = base64.b64encode(img_bytes.getvalue()).decode('utf-8')
                 base64_images.append(base64_img)
-            response = {
-                "base64_images": base64_images,
-            }
+            response = self.output or {}
+            response["base64_images"] = base64_images
             self.output = response
-            self.output_ready.set()
-
 
         def serve_image_function(image, frame_duration):
             image_file = tensorToImageConversion(image, frame_duration)
             base64_img = base64.b64encode(image_file.read()).decode('utf-8')
-            response = {
-                "base64_img": base64_img,
-            }
+            response = self.output or {}
+            response["base64_img"] = base64_img
             self.output = response
-            self.output_ready.set()
 
         def serve_text_function(text):
-            response = {
-                "text": text,
-            }
+            response = self.output or {}
+            response["text"] = text
             self.output = response
+
+        def is_command(command):
+            is_it = bool(self.path) and command == self.path[1:]
+            if is_it:
+                self.path = None
+            return is_it
+
+        def finalize():
             self.output_ready.set()
 
+
+
         data["serve_image_function"] = serve_image_function
         data["serve_multi_image_function"] = serve_multi_image_function
         data["serve_text_function"] = serve_text_function
+        data["is_command"] = is_command
+        data["finalize"] = finalize
 
-        return (data,)
+        return (data,)

nodes/telegram_serving.py

@@ -7,7 +7,8 @@ from PIL import Image
 import base64
 import numpy as np
 import torch
-from .utils import tensorToImageConversion, parse_command_string
+from .utils import tensorToImageConversion, parse_command_string, CommandRegistry
+
 
 class TelegramServing:
     def __init__(self):
@@ -15,19 +16,26 @@ class TelegramServing:
         self.data = deque()
         self.telegram_running = False
         self.bot = None
-        self.command_name = None
         self.allowed_chat_ids = None
+        self.command_registry = CommandRegistry()
 
     def telegram_handler(self):
-        @self.bot.message_handler(commands=[self.command_name])
+        @self.bot.message_handler(func=lambda message: True, content_types=['photo','text'])
         def handle_command(message):
-            chat_id=str(message.chat.id)
-            if self.allowed_chat_ids and not chat_id in self.allowed_chat_ids:
-                print(f"Allowed chatids are: {self.allowed_chat_ids}, but got message from user: {message.from_user.username}, chatid: {chat_id} ! Skipping message.")
+            chat_id = str(message.chat.id)
+            if self.allowed_chat_ids and chat_id not in self.allowed_chat_ids:
+                print(
+                    f"Allowed chatids are: {self.allowed_chat_ids}, but got message from user: {message.from_user.username}, chatid: {chat_id} ! Skipping message.")
                 return  # Silently ignore messages
-            print(f"Received command from {message.chat.id}: {message.text}")
-            parsed_data = parse_command_string(message.text, self.command_name)
-            
+
+            text = message.caption if message.content_type == 'photo' else message.text
+            command_name = text.split()[0][1:] # Extract command name without '/'
+            if not self.command_registry.has_command(command_name):
+                return  # Silently ignore wrong commands
+
+            print(f"Received command from {message.chat.id}: {text}")
+            parsed_data = parse_command_string(text, command_name)
+
             async def serve_multi_image_function(images):
                 media_group = []
                 for i, img in enumerate(images):
@@ -37,21 +45,24 @@ class TelegramServing:
                     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)
+
+                self.bot.send_media_group(message.chat.id, media_group, reply_to_message_id=message.id)
 
             def serve_image_function(image, frame_duration):
                 image_file = tensorToImageConversion(image, frame_duration)
-                self.bot.send_photo(message.chat.id, image_file)
+                self.bot.send_photo(message.chat.id, image_file, reply_to_message_id=message.id)
 
+            def is_command(command):
+                return command == command_name
+
+            parsed_data["is_command"] = is_command
             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
+            if message.photo:
+                file_info = self.bot.get_file(message.photo[2].file_id)
+                parsed_data["attachment_url_0"] = "https://api.telegram.org/file/bot{0}/{1}".format(self.bot.token, file_info.file_path)
 
             self.data.append(parsed_data)
             self.data_ready.set()
@@ -73,11 +84,7 @@ class TelegramServing:
                     "multiline": False,
                     "default": ""
                 }),
-                "command_name": ("STRING", {
-                    "multiline": False,
-                    "default": "generate"
-                }),
-        },
+            },
             "optional": {
                 "allowed_chat_ids": ("STRING", {
                     "multiline": True,
@@ -90,18 +97,18 @@ class TelegramServing:
     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, allowed_chat_ids=""):
+    def serve(self, telegram_token, allowed_chat_ids=""):
         self.allowed_chat_ids = allowed_chat_ids
         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")
+            print("Telegram bot running, listening for all commands")
             self.telegram_running = True
 
         data = self.get_data()
-        return (data,)
+        return (data,)

nodes/utils.py

@@ -38,3 +38,29 @@ def tensorToImageConversion(images, duration):
     return img_byte_array
 
 
+# This class is used to store the commands that are registered by the user.
+class CommandRegistry:
+    _instance = None
+    catch_all = False
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance.commands = []
+        return cls._instance
+
+    def add_command(self, command):
+        self.commands.append(command)
+
+    def get_commands(self):
+        return self.commands
+
+    def clear_commands(self):
+        self.commands.clear()
+
+    def has_command(self, command):
+        return command in self.commands or self.catch_all
+
+    def add_catch_all(self):
+        self.catch_all = True
+

pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 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"
+version = "1.0.1"
 license = "LICENSE"
-dependencies = ["discord.py", "websocket-client", "rel", "pyTelegramBotAPI"]
+dependencies = ["discord.py", "websocket-client", "rel", "pyTelegramBotAPI", "opencv-python"]
 
 [project.urls]
 Repository = "https://github.com/matan1905/ComfyUI-Serving-Toolkit"