diff --git a/nodes.py b/nodes.py index a2774cb..ce454a0 100644 --- a/nodes.py +++ b/nodes.py @@ -1,7 +1,10 @@ +import base64 +import numpy as np import math import requests from PIL import Image import io +import torch class DrawThingsWrapper: def __init__(self): @@ -22,11 +25,12 @@ class DrawThingsWrapper: } RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("generated_image",) FUNCTION = "generate_image" def generate_image(self, prompt, seed, width, height): # Call the Draw Things API - api_url = "http://127.0.0.1:7860//sdapi/v1/txt2img" + api_url = "http://127.0.0.1:7860/sdapi/v1/txt2img" payload = { "prompt": prompt, @@ -37,15 +41,23 @@ class DrawThingsWrapper: response = requests.post(api_url, json=payload) - if response.status_code == 200: - image_data = response.content # Assuming the API returns raw image data - image = Image.open(io.BytesIO(image_data)) # Convert raw data to an image object - return (image,) # Return as a tuple (ComfyUI expects output in tuple form) - else: - raise Exception(f"Failed to generate image: {response.text}") + # Raise an error if the request failed + response.raise_for_status() -# Register the node (depending on ComfyUI's mechanism, adjust accordingly) + # Parse the JSON response + data = response.json() + # Process the images (assuming they are base64 encoded or raw binary data) + images = [] + for img_data in data['images']: + image_bytes = base64.b64decode(img_data) + # Convert the image data to a Pillow Image object + image = Image.open(io.BytesIO(image_bytes)) + image_np = np.array(image) + # Convert to float32 tensor and normalize + tensor_image = torch.from_numpy(image_np.astype(np.float32) / 255.0) + images.append(tensor_image) + return(torch.stack(images),) NODE_CLASS_MAPPINGS = { "DrawThingsWrapper": DrawThingsWrapper