From b9814a5ba5fe47f701a1f97adb8dfa33eab178ef Mon Sep 17 00:00:00 2001
From: kft334 <93904307+kft334@users.noreply.github.com>
Date: Mon, 8 Jan 2024 07:12:41 -0500
Subject: [PATCH] Update Knodes.py

---
 Knodes.py | 70 +++++++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 55 insertions(+), 15 deletions(-)

diff --git a/Knodes.py b/Knodes.py
index 5faeb14..90d7558 100644
--- a/Knodes.py
+++ b/Knodes.py
@@ -4,6 +4,7 @@ from PIL import Image
 import numpy as np
 import base64
 import torch
+import json
 
 class ImageOutput:
     def __init__(self):
@@ -40,31 +41,64 @@ class ImageOutput:
         PromptServer.instance.send_sync("knodes", {"images": outs, "text": tag})
 
         return (images,)
-    
+
 class LoadImageBase64:
     @classmethod
     def INPUT_TYPES(s):
-        return {
-            "required": {
-            "image": ("STRING", {"multiline": False})}
-            }
+        return {"required": {"image": ("STRING", {"multiline": False})}}
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE", "MASK")
     CATEGORY = "Knodes"
+
     FUNCTION = "Proc"
 
     def Proc(self, image):
+        imgdata = base64.b64decode(image)
+        img = Image.open(BytesIO(imgdata))
 
-        # image is a base64 encoded JObject.
+        if "A" in img.getbands():
+            mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0
+            mask = 1.0 - torch.from_numpy(mask)
+        else:
+            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
 
-        jObject = base64.b64decode(image)
+        img = img.convert("RGB")
+        img = np.array(img).astype(np.float32) / 255.0
+        img = torch.from_numpy(img)[None,]
 
-        # jObject is a JObject with the following structure: [image, image, ...]
-        # where each image is a base64 encoded string.
+        return (img, mask)
+
+class LoadImagesBase64:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"strings": ("STRING", {"multiline": False})}}
+
+    RETURN_TYPES = ("IMAGE",)
+    CATEGORY = "Knodes"
+
+    FUNCTION = "Proc"
+
+    def Proc(self, strings):
+        
+        # Pop the number of images from the string
+        number_of_images = int(strings[:4].lstrip('0'), 16)
+        print("Number Of Images: " + str(number_of_images))
+        strings = strings[4:]
 
         images = []
+        
+        # Structure of the rest of the string: 0xXXXXXXXX (int64) Base64 (string) with 0xXXXXXXXX being the length of the string... repeat number_of_images times
+        for i in range(number_of_images):
+            length = int(strings[:8].lstrip('0'), 16)
+            print("Image #" + str(i) + " Length: " + str(length))
+            strings = strings[8:]
+            single_image = strings[:length]
+            strings = strings[length:]
+            images.append(single_image)
 
-        for single_image in jObject:
+        tensors = []
+
+        for single_image in images:
             imgdata = base64.b64decode(single_image)
             img = Image.open(BytesIO(imgdata))
 
@@ -72,16 +106,22 @@ class LoadImageBase64:
             img = np.array(img).astype(np.float32) / 255.0
             img = torch.from_numpy(img)[None,]
 
-            images.append(img)
+            # shape is [5, 1, 640, 640, 3] = [n, ?, h, w, c] ... n = number of images, h = height, w = width, c = channels.
+            
+            tensors.append(img)
 
-        return (images,)
+        tensors = [img.squeeze(0) for img in tensors]
+
+        return (tensors,)
         
 NODE_CLASS_MAPPINGS = {
     "Image(s) To Websocket (Base64)": ImageOutput,
-    "Load Image(s) From Websocket (Base64)": LoadImageBase64
+    "Load Image From Websocket (Base64)": LoadImageBase64,
+    "Load Images From Websocket (Base64)": LoadImagesBase64
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ImageOutput": "Image(s) To Websocket (Base64)",
-    "LoadImageBase64": "Load Image(s) From Websocket (Base64)"
+    "LoadImageBase64": "Load Image From Websocket (Base64)",
+    "LoadImageBase64": "Load Images From Websocket (Base64)"
 }