Return image tensors so Save Image works

README.md

@@ -36,12 +36,12 @@ Repo layout (mirrors Floyo Seed API style):
 ### FLUX.2 [pro] Text-to-Image (`Flux2ProTextToImage`)
 - **Inputs (required):** `prompt`, `width`, `height`
 - **Inputs (optional):** `seed` (`-1` = random), `safety_tolerance` (0-6), `output_format` (`jpeg`/`png`)
-- **Output:** `image_url` (STRING, signed URL valid for ~10 minutes)
+- **Outputs:** `image_url` (STRING, signed URL valid for ~10 minutes), `image` (IMAGE tensor ready for Save Image)
 
 ### FLUX.2 [pro] Image Edit (`Flux2ProImageEdit`)
 - **Inputs (required):** `prompt`, `input_image` (IMAGE tensor) or `input_image_url` (URL). If both are set, URL wins.
 - **Inputs (optional):** `input_image_2` … `input_image_8` (IMAGE tensors) and matching `input_image_2_url` … `input_image_8_url`. URLs win over tensors. Keep total refs <= 8. `width`, `height` (0 = match input), `seed` (`-1` = random), `safety_tolerance`, `output_format`
-- **Output:** `image_url` (STRING, signed URL valid for ~10 minutes)
+- **Outputs:** `image_url` (STRING, signed URL valid for ~10 minutes), `image` (IMAGE tensor ready for Save Image)
 
 ## Usage Notes
 - The returned URL expires quickly; run your Floyo download/output node immediately after this node.

nodes/flux2_node.py

@@ -1,6 +1,12 @@
 from typing import List, Optional
 
-from .flux2_utils import Flux2API, image_tensor_to_base64, merge_reference_images
+from .flux2_utils import (
+    Flux2API,
+    download_image_to_tensor,
+    image_tensor_to_base64,
+    merge_reference_images,
+)
+from .flux2_utils import _blank_image_tensor as _blank_image
 
 
 def _validate_resolution(width: int, height: int) -> Optional[str]:
@@ -37,8 +43,8 @@ class Flux2ProTextToImage:
             },
         }
 
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("image_url",)
+    RETURN_TYPES = ("STRING", "IMAGE")
+    RETURN_NAMES = ("image_url", "image")
     FUNCTION = "generate"
     CATEGORY = "Floyo/Flux2 Pro"
 
@@ -69,11 +75,13 @@ class Flux2ProTextToImage:
             run_result = client.run(payload)
             image_url = run_result.get("sample")
             if not image_url:
-                return ("Error: FLUX.2 response did not include an image URL.",)
+                return ("Error: FLUX.2 response did not include an image URL.", _blank_image())
 
-            return (image_url,)
+            image_tensor = download_image_to_tensor(image_url)
+
+            return (image_url, image_tensor)
         except Exception as exc:  # noqa: BLE001 - ComfyUI expects string errors
-            return (f"Error generating with FLUX.2: {exc}",)
+            return (f"Error generating with FLUX.2: {exc}", _blank_image())
 
 
 class Flux2ProImageEdit:
@@ -116,8 +124,8 @@ class Flux2ProImageEdit:
             },
         }
 
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("image_url",)
+    RETURN_TYPES = ("STRING", "IMAGE")
+    RETURN_NAMES = ("image_url", "image")
     FUNCTION = "edit"
     CATEGORY = "Floyo/Flux2 Pro"
 
@@ -197,11 +205,13 @@ class Flux2ProImageEdit:
             run_result = client.run(payload)
             image_url = run_result.get("sample")
             if not image_url:
-                return ("Error: FLUX.2 response did not include an image URL.",)
+                return ("Error: FLUX.2 response did not include an image URL.", _blank_image())
 
-            return (image_url,)
+            image_tensor = download_image_to_tensor(image_url)
+
+            return (image_url, image_tensor)
         except Exception as exc:  # noqa: BLE001 - ComfyUI expects string errors
-            return (f"Error editing with FLUX.2: {exc}",)
+            return (f"Error editing with FLUX.2: {exc}", _blank_image())
 
 
 NODE_CLASS_MAPPINGS = {

nodes/flux2_utils.py

@@ -169,3 +169,30 @@ def image_tensor_to_base64(image: torch.Tensor, format: str = "PNG") -> str:
     pil_image.save(buffer, format=format.upper(), **save_params)
     encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
     return encoded
+
+
+def _blank_image_tensor(width: int = 512, height: int = 512) -> torch.Tensor:
+    """Create a blank black image tensor for fallback."""
+    blank = np.zeros((height, width, 3), dtype=np.uint8)
+    return torch.from_numpy(blank.astype(np.float32) / 255.0)[None, ...]
+
+
+def download_image_to_tensor(url: str) -> torch.Tensor:
+    """
+    Download an image from URL and convert it to a ComfyUI-style tensor.
+
+    Returns a 4D tensor (1, H, W, C) with float values in [0,1].
+    """
+    if not url:
+        return _blank_image_tensor()
+
+    try:
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()
+        img = Image.open(io.BytesIO(response.content)).convert("RGB")
+        img_np = np.array(img).astype(np.float32) / 255.0
+        img_tensor = torch.from_numpy(img_np)[None, ...]
+        return img_tensor
+    except Exception as exc:
+        print(f"[Flux2] Failed to download image from {url}: {exc}")
+        return _blank_image_tensor()