Commit

python.py

@@ -4,6 +4,8 @@ import sys
 from typing import Sequence, Mapping, Any, Union
 import torch
 import spaces
+import numpy as np
+from PIL import Image
 # from comfy import model_management
 from nodes import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_1
 from comfy_extras.nodes_custom_sampler import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_2
@@ -23,6 +25,21 @@ hf_hub_download(repo_id="black-forest-labs/FLUX.1-dev", filename="ae.safetensors
 hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="clip_l.safetensors", local_dir="models/text_encoders", token = token)
 hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", filename="t5xxl_fp16.safetensors", local_dir="models/text_encoders", token = token)
 
+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
+
 
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     """Returns the value at the given index of a sequence or mapping.
@@ -266,5 +283,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).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