Author SHA1 Message Date
GitHub Actions 09fc4fbb87 Update repository 2024-09-05 21:59:54 +00:00
8 changed files with 126 additions and 468 deletions
+2 -6
View File
@@ -7,19 +7,15 @@ 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@v1
uses: Comfy-Org/publish-node-action@main
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 }}
+1 -1
View File
@@ -1,4 +1,4 @@
__pycache__/
__pycache__/*
/examples/node_modules/*
.idea
+6 -78
View File
@@ -2,6 +2,7 @@
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.
**ServingTextOutput**
If you need any help, Checkout the [Discord!](https://discord.gg/AyHFchFZuV)
@@ -11,7 +12,6 @@ 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)
@@ -21,14 +21,11 @@ or put all the files inside custom_nodes and run:
```
## The simplest configuration
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)
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)
You can copy the workflow json:
[workflow.json](https://github.com/user-attachments/files/17041080/workflow.4.json)
[discordserv.json](https://github.com/matan1905/ComfyUI-Serving-Toolkit/files/13248566/discordserv.json)
@@ -40,7 +37,6 @@ 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.
@@ -68,25 +64,7 @@ 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**
@@ -118,7 +96,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)
**Serving Image/Video Output**
**ServingOutput**
Allows you to return an image/video back to the request
Inputs:
@@ -126,57 +104,7 @@ 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
+74 -233
View File
@@ -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, CommandRegistry
from .utils import parse_command_string, tensorToImageConversion
import discord
import asyncio
import requests
@@ -15,7 +15,6 @@ import json
import torch
import torchvision.transforms as transforms
import cv2
import nodes
@@ -33,23 +32,20 @@ class ServingOutput:
"image": ("IMAGE",),
"frame_duration": ("INT", {"default": 30, "min": 1, "step": 1, "max": 9999999}),
},
"optional": {
"chained_execution": ("SHOULD_EXECUTE",),
}
}
RETURN_TYPES = ("SHOULD_EXECUTE",)
RETURN_TYPES = ()
# RETURN_NAMES = ("image_output_name",)
FUNCTION = "out"
# OUTPUT_NODE = True
OUTPUT_NODE = True
CATEGORY = "Serving-Toolkit"
def out(self, image,serving_config,frame_duration, chained_execution=None):
def out(self, image,serving_config,frame_duration):
serving_config["serve_image_function"](image,frame_duration)
return ({}, )
return {}
@@ -64,22 +60,19 @@ class ServingTextOutput:
"serving_config": ("SERVING_CONFIG",),
"text": ("STRING", {"multiline": True, "default": ""}),
},
"optional": {
"chained_execution": ("SHOULD_EXECUTE",),
}
}
RETURN_TYPES = ("SHOULD_EXECUTE",)
RETURN_TYPES = ()
FUNCTION = "out"
# OUTPUT_NODE = True
OUTPUT_NODE = True
CATEGORY = "Serving-Toolkit"
def out(self, serving_config, text, chained_execution=None):
def out(self, serving_config, text):
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:
@@ -116,71 +109,6 @@ 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):
@@ -232,18 +160,16 @@ class ServingMultiImageOutput:
"serving_config": ("SERVING_CONFIG",),
"images": ("IMAGE",),
},
"optional": {
"chained_execution": ("SHOULD_EXECUTE",),
}
}
RETURN_TYPES = ("SHOULD_EXECUTE",)
RETURN_TYPES = ()
FUNCTION = "out"
# OUTPUT_NODE = True
OUTPUT_NODE = True
CATEGORY = "Serving-Toolkit"
def out(self, serving_config, images, chained_execution=None):
def out(self, serving_config, images):
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)
@@ -255,7 +181,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
@@ -264,7 +190,7 @@ class DiscordServing():
self.data_ready = threading.Event()
self.data = deque()
self.discord_token = None
self.command_registry = CommandRegistry()
pass
def discord_runner(self):
discord_client.run(self.discord_token)
@@ -283,6 +209,10 @@ class DiscordServing():
"discord_token": ("STRING", {
"multiline": True,
"default": ""
}),
"command_name": ("STRING", {
"multiline": False,
"default": "generate"
})
}
}
@@ -295,81 +225,74 @@ class DiscordServing():
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("NaN")
# OUTPUT_NODE = False
CATEGORY = "Serving-Toolkit"
def serve(self, discord_token):
def serve(self, command_name, 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.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)}")
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)
async def serve_multi_image_function(images):
discord_files = []
for i, img in enumerate(images):
try:
print("Attempting to send images...")
await message.reply(files=discord_files)
print("Images sent successfully")
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 sending images: {str(e)}")
print(f"Error processing image {i + 1}: {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)
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 is_command(command):
return bool(command_name == command)
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()
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()
data = self.get_data()
return (data,)
class WebSocketServing():
def __init__(self):
self.data_ready = threading.Event()
@@ -377,13 +300,10 @@ 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:
@@ -480,14 +400,10 @@ 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,)
@@ -529,9 +445,7 @@ class ServingInputImage:
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,)
raise ValueError("No attachment found in serving_config")
attachment_url = serving_config[attachment_url_key]
response = requests.get(attachment_url)
@@ -581,9 +495,7 @@ class ServingInputImageAsLatent:
if attachment_url_key not in serving_config:
if default_latent is not None:
return (default_latent,)
serving_config["serve_text_function"]("This command requires an image")
nodes.interrupt_processing(True)
return (None,)
raise ValueError("No attachment found in serving_config")
attachment_url = serving_config[attachment_url_key]
response = requests.get(attachment_url)
@@ -600,86 +512,18 @@ 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,
"CommandPickerServing": CommandPickerServing,
"AlwaysExecute": AlwaysExecute
"ServingInputImageAsLatent": ServingInputImageAsLatent
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -688,13 +532,10 @@ 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",
"CommandPickerServing": "Command Picker Serving",
"AlwaysExecute": "Always Execute"
"ServingInputImageAsLatent": "Serving Input Image as Latent"
}
+22 -84
View File
@@ -4,13 +4,12 @@ from collections import deque
import json
import base64
from io import BytesIO
from .utils import tensorToImageConversion, CommandRegistry
from .utils import tensorToImageConversion
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
@@ -18,59 +17,35 @@ 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(self.html_content.encode('utf-8'))
self2.wfile.write(b"HTTP Serving is running, to send data make a POST request to this endpoint.")
self.server = HTTPServer(('', self.port), RequestHandler)
print(f"HTTP Server running on port {self.port}")
self.server.serve_forever()
def get_data(self):
@@ -85,29 +60,6 @@ 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>
"""})
}
}
@@ -116,21 +68,13 @@ class HTTPServing:
FUNCTION = "serve"
CATEGORY = "Serving-Toolkit"
@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
def serve(self, port):
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):
@@ -141,37 +85,31 @@ 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 = self.output or {}
response["base64_images"] = base64_images
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 = self.output or {}
response["base64_img"] = base64_img
response = {
"base64_img": base64_img,
}
self.output = response
def serve_text_function(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()
def serve_text_function(text):
response = {
"text": text,
}
self.output = response
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,)
+19 -38
View File
@@ -7,8 +7,7 @@ from PIL import Image
import base64
import numpy as np
import torch
from .utils import tensorToImageConversion, parse_command_string, CommandRegistry
from .utils import tensorToImageConversion, parse_command_string
class TelegramServing:
def __init__(self):
@@ -16,26 +15,14 @@ class TelegramServing:
self.data = deque()
self.telegram_running = False
self.bot = None
self.allowed_chat_ids = None
self.command_registry = CommandRegistry()
self.command_name = None
def telegram_handler(self):
@self.bot.message_handler(func=lambda message: True, content_types=['photo','text'])
@self.bot.message_handler(commands=[self.command_name])
def handle_command(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
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)
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):
@@ -45,24 +32,21 @@ 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, reply_to_message_id=message.id)
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, reply_to_message_id=message.id)
self.bot.send_photo(message.chat.id, image_file)
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.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)
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()
@@ -84,11 +68,9 @@ class TelegramServing:
"multiline": False,
"default": ""
}),
},
"optional": {
"allowed_chat_ids": ("STRING", {
"multiline": True,
"default": ""
"command_name": ("STRING", {
"multiline": False,
"default": "generate"
})
}
}
@@ -97,17 +79,16 @@ 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, allowed_chat_ids=""):
self.allowed_chat_ids = allowed_chat_ids
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("Telegram bot running, listening for all commands")
print(f"Telegram bot running, listening for /{command_name} commands")
self.telegram_running = True
data = self.get_data()
-26
View File
@@ -38,29 +38,3 @@ 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
+2 -2
View File
@@ -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.1"
version = "1.0.0"
license = "LICENSE"
dependencies = ["discord.py", "websocket-client", "rel", "pyTelegramBotAPI", "opencv-python"]
dependencies = ["discord.py", "websocket-client", "rel", "pyTelegramBotAPI"]
[project.urls]
Repository = "https://github.com/matan1905/ComfyUI-Serving-Toolkit"