290 lines
8.8 KiB
Python
290 lines
8.8 KiB
Python
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)",
|
|
}
|