diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..09f44c2 Binary files /dev/null and b/.DS_Store differ diff --git a/.gitignore b/.gitignore index 82f9275..7b6caf3 100644 --- a/.gitignore +++ b/.gitignore @@ -159,4 +159,4 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +.idea/ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..7ece095 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ +from .nodes.publish_image_rabbitmq import * + +NODE_CLASS_MAPPINGS = { + "Publish Image To RabbitMQ": PublishImageToRabbitMQ, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Publish Image To RabbitMQ": "💾 Publish Your Image to RabbitMQ", +} \ No newline at end of file diff --git a/mq_client.py b/mq_client.py new file mode 100644 index 0000000..827cd50 --- /dev/null +++ b/mq_client.py @@ -0,0 +1,52 @@ +import json + +import pika + + +class RabbitMQClient: + def __init__( + self, + exchange: str, + routing_key: str = 'image', + host: str = 'localhost', + port: int = 5672, + username: str = 'guest', + password: str = 'guest' + ): + self.connection = pika.BlockingConnection(pika.ConnectionParameters( + host=host, + port=port, + credentials=pika.PlainCredentials(username, password) + )) + self.channel = self.connection.channel() + self.exchange = exchange + self.routing_key = routing_key + + def declare_exchange( + self, + exchange_type: str = 'direct', + durable: bool = True, + auto_delete: bool = False, + internal: bool = False + ): + """test exchange declaration""" + self.channel.exchange_declare( + exchange=self.exchange, + exchange_type=exchange_type, + durable=durable, + auto_delete=auto_delete, + internal=internal + ) + + def publish(self, message: dict): + self.channel.basic_publish( + exchange=self.exchange, + routing_key=self.routing_key, + body=json.dumps(message) + ) + + def close(self): + self.channel.close() + self.connection.close() + + diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/publish_image_rabbitmq.py b/nodes/publish_image_rabbitmq.py new file mode 100644 index 0000000..431d740 --- /dev/null +++ b/nodes/publish_image_rabbitmq.py @@ -0,0 +1,83 @@ +import base64 +from io import BytesIO + +import numpy as np +from PIL import Image + +from server import PromptServer +from ..mq_client import RabbitMQClient + + +class PublishImageToRabbitMQ: + @classmethod + def INPUT_TYPES(s): + return { + 'required': { + 'exchange': ('STRING', {'multiline': False, 'default': 'comfy'}), + 'routing_key': ('STRING', {'multiline': False, 'default': 'image'}), + 'host': ('STRING', {'multiline': False, 'default': 'localhost'}), + 'port': ('INT', {'default': 5672}), + 'username': ('STRING', {'multiline': False, 'default': 'guest'}), + 'password': ('STRING', {'multiline': False, 'default': 'guest'}), + 'exchange_type': ('STRING', { + 'multiline': False, + 'default': 'direct', + 'options': ['fanout', 'direct', 'topic'] # TODO support header type + }), + 'durable': ('BOOLEAN', {'default': True}), + 'auto_delete': ('BOOLEAN', {'default': False}), + 'internal': ('BOOLEAN', {'default': False}), + 'images': ('IMAGE',), + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = ('STRING',) + # RETURN_NAMES = () + FUNCTION = 'publish_images' + OUTPUT_NODE = True + OUTPUT_IS_LIST = (True,) + CATEGORY = 'ComfyRabbitMQ' + + def publish_images( + self, + exchange, + routing_key, + host, + port, + username, + password, + exchange_type, + durable, + auto_delete, + internal, + images, + prompt=None, + extra_pnginfo=None + ): + prompt_id = PromptServer.instance.last_prompt_id # get current prompt id + + results = [] + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + img_byte_array = BytesIO() + img.save(img_byte_array, format='PNG') + img_data = img_byte_array.getvalue() + img_base64 = base64.b64encode(img_data).decode('utf-8') + + results.append({ + 'batch_number': batch_number, + 'base64_data': img_base64, + }) + + mq_client = RabbitMQClient(exchange, routing_key, host, port, username, password) + mq_client.declare_exchange(exchange_type, durable, auto_delete, internal) + mq_client.publish({ + 'images': results, + 'prompt_id': prompt_id, + }) + mq_client.close() + + return { "ui": { "images": prompt_id } } + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..59efa26 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +numpy==2.1.3 +pika==1.3.2 +pillow==11.0.0