Merge pull request #25 from ComfyNodePRs/update-publish-yaml

Update Github Action for Publishing to Comfy Registry

.github/workflows/publish.yml

@@ -7,15 +7,19 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'matan1905' }}
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
           ## Add your own personal access token to your Github Repository secrets and reference it here.
-          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

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

README.md

@@ -11,6 +11,7 @@ If you need any help, Checkout the [Discord!](https://discord.gg/AyHFchFZuV)
 * Allows both Images or videos (when in batch mode, such as animatediff - if you return more than one image it will create a video)
 * Add arguments with default values, then allow your users to use them
 * Serve from your own computer, workflow is not inserted into the images so your secrets are 100% safe
+* Support for multiple serving options: Discord, Telegram, HTTP and WebSockets
 
 ## Installation
 [Use ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager)
@@ -20,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)
+
+
 
 
 
@@ -36,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.
@@ -63,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** 
@@ -95,7 +118,7 @@ Inputs that are not in ServingInputText:
 * min_value - the minimum value of this argument
 * step - the steps of this value (setting this to 1 will ensure only whole numbers, 0.5 will allow jumps of half etc)
 
-**ServingOutput**
+**Serving Image/Video Output**
 
 Allows you to return an image/video back to the request
 Inputs:
@@ -103,7 +126,57 @@ Inputs:
 * duration - in the case of a video, what is the time in miliseconds each frame should appear? if you have an FPS number you can use 1000/FPS to calculate the duration value
 
 
+**TelegramServing**
 
+This node allows you to serve your workflow via Telegram. It listens for commands and processes them, enabling interaction with your ComfyUI workflow through a Telegram bot.
+
+Inputs:
+* telegram_token - Your Telegram bot token obtained from the BotFather
+* command_name - The command used to generate, without the '/'. Defaults to 'generate' (so you would have to do /generate <your prompt> --your_argument1 <argument value>)
+
+Outputs:
+* Serving Config - A basic reference for this serving, used by the other nodes of this toolkit to get arguments and return images.
+
+**HTTPServing**
+
+This node enables serving your workflow via HTTP, allowing you to interact with your ComfyUI workflow through HTTP POST requests.
+
+Inputs:
+* port - The port number on which the HTTP server will run. Defaults to 8000.
+
+Outputs:
+* Serving Config - A basic reference for this serving, used by the other nodes of this toolkit to get arguments and return images.
+
+**ServingMultiImageOutput**
+
+Allows you to return multiple images back to the request. This is useful for batch processing or when generating animations.
+
+Inputs:
+* serving_config - A config made by a serving node
+* images - The generated images
+
+**ServingInputImage**
+
+Allows you to input images from the serving platform (e.g., attachments in Discord, Telegram, or files sent via HTTP).
+
+Inputs:
+* serving_config - A config made by a serving node
+* default_image - (Optional) A default image to use if no attachment is found
+
+Outputs:
+* image - The input image as a tensor
+
+**ServingInputImageAsLatent**
+
+Similar to ServingInputImage, but converts the input image directly into a latent representation for use in your workflow.
+
+Inputs:
+* serving_config - A config made by a serving node
+* vae - The VAE model to use for encoding
+* default_latent - (Optional) A default latent to use if no attachment is found
+
+Outputs:
+* latent - The input image encoded as a latent
 
 
 

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", {"multiline": False, "default": ""}),
+                "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,12 +4,13 @@ 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
 
 
 class HTTPServing:
     def __init__(self):
+        self.enable_cross_origin_requests = None
         self.data_ready = threading.Event()
         self.data = deque()
         self.http_running = False
@@ -17,35 +18,59 @@ 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):
+                    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()
 
                 self.output_ready.wait()
+                self.output_ready.clear()
                 response = self.output
                 self2.send_response(200)
                 self2.send_header('Content-type', 'application/json')
+                # Cors
+                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', '*')
+
                 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):
@@ -60,6 +85,29 @@ class HTTPServing:
         return {
             "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>
+                """})
             }
         }
 
@@ -68,13 +116,21 @@ class HTTPServing:
     FUNCTION = "serve"
     CATEGORY = "Serving-Toolkit"
 
-    def serve(self, port):
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        return float("NaN")
+
+    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 = None
+        self.output_ready.clear()  # Prevent deadlock if failed in previous run
         data = self.get_data()
 
         def serve_multi_image_function(images):
@@ -85,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,14 +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):
-            print(f"Received command from {message.chat.id}: {message.text}")
-            parsed_data = parse_command_string(message.text, self.command_name)
-            
+            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
+
+            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):
@@ -32,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()
@@ -68,9 +84,11 @@ class TelegramServing:
                     "multiline": False,
                     "default": ""
                 }),
-                "command_name": ("STRING", {
-                    "multiline": False,
-                    "default": "generate"
+            },
+            "optional": {
+                "allowed_chat_ids": ("STRING", {
+                    "multiline": True,
+                    "default": ""
                 })
             }
         }
@@ -79,16 +97,17 @@ 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):
+    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()

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"