From 0395dab88cef10a99524aab63ef089c42c3cc92d Mon Sep 17 00:00:00 2001
From: Xingkun <18646154381@163.com>
Date: Thu, 7 Nov 2024 16:55:38 +0800
Subject: [PATCH] init: first commit

---
 .DS_Store                       | Bin 0 -> 6148 bytes
 .gitignore                      |   2 +-
 __init__.py                     |   9 ++++
 mq_client.py                    |  52 ++++++++++++++++++++
 nodes/__init__.py               |   0
 nodes/publish_image_rabbitmq.py |  83 ++++++++++++++++++++++++++++++++
 requirements.txt                |   3 ++
 7 files changed, 148 insertions(+), 1 deletion(-)
 create mode 100644 .DS_Store
 create mode 100644 __init__.py
 create mode 100644 mq_client.py
 create mode 100644 nodes/__init__.py
 create mode 100644 nodes/publish_image_rabbitmq.py
 create mode 100644 requirements.txt

diff --git a/.DS_Store b/.DS_Store
new file mode 100644
index 0000000000000000000000000000000000000000..09f44c293f63ccc693442a19fd52a4a16c67075c
GIT binary patch
literal 6148
zcmeHK!EVz)5S>j^Vk;o|01_OKEOCv3X{!pv#f0?G1D6`X0Z@qR2rL|L8ao6OMe@D<
z7tZ|={)H2~+1;YVAy<UxPBi;w$Mfv$Tifd;BGH>AkBB@Xa^Q>&7tJ4xkF#%B!(>`O
zq0W&~Ndx@4G})-OiGNW6p1Zc1vy^gLoxiWgnLLd%8KZ)4k7w}iPi2URZqk^}z&xc4
zyb;)?n)f>ITaCAiN=XSA#rhS*{q2?UyqK8(av4vuv?zn%yJ+k*cW<<umeY3LMIYrn
zDu?B)=nbcDxc5@ZEWV0|@vC$+A9%N)%e)+>`Dm&LX);2|+t+EH$azoB@}yAQNFQ)q
z#~pb07K^=u!~Kr`I0%*<fAM55==cZwhs&kw+<EZu>B-<zHqPZatAZy1kAkvW2Cv`(
z!OGx2N7Fo$`3Gc~;nE2mLF7KY0Cpj6Qj>?zzd8XPO(EhL%^*KVV>1Z_*d=|Tvc~^W
z<D+ZMDqt13t_qkJmD9ei0@!g|1*`&F72x-S4`&PwRvOi-1C81O0GmiFL(CtG%<&C|
z1}lv?0u!1F)Kp=%7(&yL_f1`Bu+pgMB+TYRm@5mjLlNrgIKQvTNrW10YZb5xtShji
zyB@FqN89)R>m=K=3RngHD+NTeAN6|}l380f2FGiy4}TA5<GM=YZwdmn6r+}x;=6EV
Z$om`sLxYt@WMK9~K*?YmtH57X;1}!JkI4W4

literal 0
HcmV?d00001

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