import io import torch import numpy as np from PIL import Image import paho.mqtt.client as mqtt import ssl import json class MQTTPublishImage: """ A ComfyUI node that publishes images to an MQTT broker. Supports TLS encryption and username/password authentication. """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "broker_url": ("STRING", { "default": "localhost", "multiline": False }), "broker_port": ("INT", { "default": 1883, "min": 1, "max": 65535 }), "topic": ("STRING", { "default": "comfyui/image", "multiline": False }), "use_tls": ("BOOLEAN", { "default": False }), "image_format": (["JPEG", "PNG"],), "quality": ("INT", { "default": 95, "min": 1, "max": 100, "step": 1 }), }, "optional": { "username": ("STRING", { "default": "", "multiline": False }), "password": ("STRING", { "default": "", "multiline": False }), "client_id": ("STRING", { "default": "comfyui_mqtt_client", "multiline": False }), "qos": ("INT", { "default": 1, "min": 0, "max": 2 }), "retain": ("BOOLEAN", { "default": False }), } } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("image", "status") FUNCTION = "publish_image" CATEGORY = "image/mqtt" OUTPUT_NODE = True def tensor_to_pil(self, tensor): """ Convert ComfyUI image tensor to PIL Image. ComfyUI tensors are in format [B, H, W, C] with values in range [0, 1] """ # Handle batch dimension - take first image if batch if len(tensor.shape) == 4: tensor = tensor[0] # Convert to numpy and scale to [0, 255] numpy_image = (tensor.cpu().numpy() * 255).astype(np.uint8) # Convert to PIL Image pil_image = Image.fromarray(numpy_image) return pil_image def image_to_b64(self, pil_image, format="JPEG", quality=95): """ Convert PIL Image to bytes. """ img_byte_arr = io.BytesIO() if format == "JPEG": pil_image.save(img_byte_arr, format='JPEG', quality=quality) else: # PNG pil_image.save(img_byte_arr, format='PNG') img_byte_arr.seek(0) img_byte = img_byte_arr.getvalue() return base64.b64encode(img_byte).decode() def publish_image(self, image, broker_url, broker_port, topic, use_tls, image_format, quality, username="", password="", client_id="comfyui_mqtt_client", qos=1, retain=False): """ Publish image to MQTT broker. """ try: # Convert image tensor to PIL Image pil_image = self.tensor_to_pil(image) # Convert PIL Image to bytes image_b64 = self.image_to_b64(pil_image, image_format, quality) # Create MQTT client client = mqtt.Client(client_id=client_id) # Set up authentication if provided if username and password: client.username_pw_set(username, password) # Set up TLS if requested if use_tls: client.tls_set(cert_reqs=ssl.CERT_REQUIRED, tls_version=ssl.PROTOCOL_TLSv1_2) # Connect to broker client.connect(broker_url, broker_port, keepalive=60) # Start mqtt loop client.loop_start() # Publish image result = client.publish(topic, image_b64, qos=qos, retain=retain) # Wait for publish to complete result.wait_for_publish() # Disconnect client.disconnect() status = f"✓ Published to {broker_url}:{broker_port}/{topic}" print(f"[MQTT] {status}") return (image, status) except Exception as e: error_msg = f"✗ MQTT Error: {str(e)}" print(f"[MQTT] {error_msg}") return (image, error_msg) class MQTTPublishImageJSON: """ A ComfyUI node that publishes image metadata as JSON to an MQTT broker. Useful for sending image information without the full binary payload. """ @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "broker_url": ("STRING", { "default": "localhost", "multiline": False }), "broker_port": ("INT", { "default": 1883, "min": 1, "max": 65535 }), "topic": ("STRING", { "default": "comfyui/metadata", "multiline": False }), "use_tls": ("BOOLEAN", { "default": False }), }, "optional": { "username": ("STRING", { "default": "", "multiline": False }), "password": ("STRING", { "default": "", "multiline": False }), "client_id": ("STRING", { "default": "comfyui_mqtt_json_client", "multiline": False }), "qos": ("INT", { "default": 1, "min": 0, "max": 2 }), "custom_data": ("STRING", { "default": "{}", "multiline": True }), } } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("image", "status") FUNCTION = "publish_metadata" CATEGORY = "image/mqtt" OUTPUT_NODE = True def publish_metadata(self, image, broker_url, broker_port, topic, use_tls, username="", password="", client_id="comfyui_mqtt_json_client", qos=1, custom_data="{}"): """ Publish image metadata as JSON to MQTT broker. """ try: # Extract image metadata if len(image.shape) == 4: batch_size, height, width, channels = image.shape else: batch_size = 1 height, width, channels = image.shape # Parse custom data try: custom_dict = json.loads(custom_data) if custom_data else {} except json.JSONDecodeError: custom_dict = {} # Create metadata dictionary metadata = { "batch_size": int(batch_size), "width": int(width), "height": int(height), "channels": int(channels), "dtype": str(image.dtype), **custom_dict } # Convert to JSON json_payload = json.dumps(metadata) # Create MQTT client client = mqtt.Client(client_id=client_id) # Set up authentication if provided if username and password: client.username_pw_set(username, password) # Set up TLS if requested if use_tls: client.tls_set(cert_reqs=ssl.CERT_REQUIRED, tls_version=ssl.PROTOCOL_TLSv1_2) # Connect to broker client.connect(broker_url, broker_port, keepalive=60) # Publish metadata result = client.publish(topic, json_payload, qos=qos) # Wait for publish to complete result.wait_for_publish() # Disconnect client.disconnect() status = f"✓ Published metadata to {broker_url}:{broker_port}/{topic}" print(f"[MQTT] {status}") return (image, status) except Exception as e: error_msg = f"✗ MQTT Error: {str(e)}" print(f"[MQTT] {error_msg}") return (image, error_msg) # Register nodes NODE_CLASS_MAPPINGS = { "MQTTPublishImage": MQTTPublishImage, "MQTTPublishImageJSON": MQTTPublishImageJSON, } NODE_DISPLAY_NAME_MAPPINGS = { "MQTTPublishImage": "MQTT Publish Image", "MQTTPublishImageJSON": "MQTT Publish Image Metadata (JSON)", }