diff --git a/README.md b/README.md index 83efbbd..b800980 100644 --- a/README.md +++ b/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. diff --git a/nodes/flux2_node.py b/nodes/flux2_node.py index f640417..c01c444 100644 --- a/nodes/flux2_node.py +++ b/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 = { diff --git a/nodes/flux2_utils.py b/nodes/flux2_utils.py index 4ca31b3..28ef75d 100644 --- a/nodes/flux2_utils.py +++ b/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()