From 9e1ce4e826f7c2c359583fcc6c150e4f95ce9e8b Mon Sep 17 00:00:00 2001
From: Jo Kim <1893942+Jokimbe@users.noreply.github.com>
Date: Sun, 6 Apr 2025 18:02:23 +0200
Subject: [PATCH] Update nodes.py

---
 nodes.py | 41 ++++++++++++++++-------------------------
 1 file changed, 16 insertions(+), 25 deletions(-)

diff --git a/nodes.py b/nodes.py
index b5c63b5..d16d3f2 100644
--- a/nodes.py
+++ b/nodes.py
@@ -22,6 +22,7 @@ from . import LoRA
 from . import GenerationConfiguration
 import hashlib
 import json
+import ctypes
 import pickle
 
 sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
@@ -81,22 +82,14 @@ def convert_response_image(response_image: bytes):
         'channels': channels,
     }
 
-# class TensorParam(ctypes.Structure):
-#     _fields_ = [
-#         ("type", ctypes.c_int),
-#         ("format", ctypes.c_int),
-#         ("datatype", ctypes.c_int),
-#         ("reserved", ctypes.c_int),
-#         ("dim", ctypes.c_array),
-#     ]
-
-class TensorParam:
-    def __init__(self, type, format, datatype, reserved, dim):
-        self.type = type
-        self.format = format
-        self.datatype = datatype
-        self.reserved = reserved
-        self.dim = dim
+class TensorParam(ctypes.Structure):
+    _fields_ = [
+        ("type", ctypes.c_int),
+        ("format", ctypes.c_int),
+        ("datatype", ctypes.c_int),
+        ("reserved", ctypes.c_int),
+        ("dim", ctypes.c_int*4),
+    ]
 
 def convert_image_for_request(img: torch.Tensor):
     # Draw Things: C header + the Float16 blob of -1 to 1 values that represents the image (in RGB order and HWC format, meaning r(0, 0), g(0, 0), b(0, 0), r(1, 0), g(1, 0), b(1, 0) .... (r(x, y) represents the value of red at that particular coordinate). The actual header is a bit more complex, here is the reference: https://github.com/liuliu/s4nnc/blob/main/nnc/Tensor.swift#L1750 the ccv_nnc_tensor_param_t is here: https://github.com/liuliu/ccv/blob/unstable/lib/nnc/ccv_nnc_tfb.h#L79 The type is CCV_TENSOR_CPU_MEMORY, format is CCV_TENSOR_FORMAT_NHWC, datatype is CCV_16F (for Float16), dim is the dimension in N, H, W, C order (in the case it should be 1, actual height, actual width, 3).
@@ -118,21 +111,19 @@ def convert_image_for_request(img: torch.Tensor):
     CCV_16F = 0x20000
     channels = 3
     dimensions = [1, height, width, channels]
-    # dimensions = (ctypes.c_int * len(dimensions))(*dimensions)
+    dimensions = (ctypes.c_int * len(dimensions))(*dimensions)
+
+    header = TensorParam() # ccv_nnc_tensor_param_t()
+    header.type = CCV_TENSOR_CPU_MEMORY
+    header.format = CCV_TENSOR_FORMAT_NHWC
+    header.datatype = CCV_16F
+    header.dim = dimensions
 
-    header = TensorParam(
-    type = CCV_TENSOR_CPU_MEMORY, 
-    format = CCV_TENSOR_FORMAT_NHWC, 
-    datatype = CCV_16F, 
-    reserved = None,
-    dim = dimensions, 
-    ) # ccv_nnc_tensor_param_t()
     print(f"{header.datatype}")
     data = img.to(torch.float16)
 
     # Encode the image as base64
     encoded_string = base64.b64encode(pickle.dumps(header) + tf.io.serialize_tensor(data))
-    # encoded_string = base64.b64encode(data)
     return encoded_string
 
 def get_files(server, port):