Working txt2img generator

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