commit

python.py

@@ -12,6 +12,7 @@ from custom_nodes.ComfyUI_Comfyroll_CustomNodes.node_mappings import NODE_CLASS_
 from comfy_extras.nodes_model_advanced import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_5
 from comfy_extras.nodes_flux import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_6
 from torchvision.transforms.functional import to_pil_image
+from PIL import Image
 import time
 
 from huggingface_hub import hf_hub_download
@@ -126,6 +127,21 @@ def import_custom_nodes() -> None:
     # Initializing custom nodes
     init_extra_nodes()
 
+def preprocess_image_tensor(image):
+    # If image has a batch dimension (shape: [1, C, H, W]), remove it.
+    if image.ndim == 4 and image.shape[0] == 1:
+        image = image.squeeze(0)
+    # If image is in channels-first format (i.e. [C, H, W]) and has 3 or 4 channels,
+    # convert it to channels-last format ([H, W, C]).
+    if image.ndim == 3 and image.shape[0] in [1, 3, 4]:
+        image = image.permute(1, 2, 0)
+    # Ensure the image values are between 0 and 1. Then scale them to [0, 255].
+    image = image.detach().cpu().numpy()
+    image = np.clip(image, 0, 1) * 255
+    # Convert to unsigned 8-bit integer type.
+    image = image.astype(np.uint8)
+    return image
+
 
 
 add_comfyui_directory_to_sys_path()
@@ -266,5 +282,7 @@ def generate_image(prompt,
             # saveimage_9 = saveimage.save_images(
             #     filename_prefix="image", images=get_value_at_index(vaedecode_8, 0)
             # )
-            image = to_pil_image(get_value_at_index(vaedecode_8, 0).squeeze(0).cpu())
-            return image, seed
+            image_tensor = get_value_at_index(vaedecode_8, 0)
+            preprocessed_image = preprocess_image_tensor(image_tensor)
+            pil_image = Image.fromarray(preprocessed_image)
+            return pil_image, seed