diff --git a/README.md b/README.md index ca5f9f0..83efbbd 100644 --- a/README.md +++ b/README.md @@ -39,14 +39,15 @@ Repo layout (mirrors Floyo Seed API style): - **Output:** `image_url` (STRING, signed URL valid for ~10 minutes) ### FLUX.2 [pro] Image Edit (`Flux2ProImageEdit`) -- **Inputs (required):** `prompt`, `input_image` (URL) -- **Inputs (optional):** `input_image_2` … `input_image_8` (extra refs, keep total <= 8), `width`, `height` (0 = match input), `seed` (`-1` = random), `safety_tolerance`, `output_format` +- **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) ## Usage Notes - The returned URL expires quickly; run your Floyo download/output node immediately after this node. - API statuses handled: `Ready`, `Pending`, and error/moderation states. Failures return an error string in the node output. - Keep total reference images within FLUX.2 [pro] limits (max 8 images, 9MP total). Width/height must be multiples of 16; `0` uses the input image dimensions for edits. Max resolution 4MP (e.g., 2048x2048). +- When connecting ComfyUI image tensors, the node auto-converts to base64. If you prefer URLs, fill the URL fields instead. ## API Reference - Endpoint: `https://api.bfl.ai/v1/flux-2-pro` diff --git a/nodes/flux2_node.py b/nodes/flux2_node.py index f55c1b5..9d024c5 100644 --- a/nodes/flux2_node.py +++ b/nodes/flux2_node.py @@ -1,6 +1,6 @@ from typing import List, Optional -from .flux2_utils import Flux2API, merge_reference_images +from .flux2_utils import Flux2API, image_tensor_to_base64, merge_reference_images def _validate_resolution(width: int, height: int) -> Optional[str]: @@ -85,18 +85,26 @@ class Flux2ProImageEdit: @classmethod def INPUT_TYPES(cls): - optional_refs = { - f"input_image_{idx}": ( - "STRING", - {"default": "", "tooltip": f"Optional reference image URL #{idx} (max 7 extras, total refs <= 8)."}, - ) - for idx in range(2, 9) + base_image_inputs = { + "input_image": ("IMAGE", {"tooltip": "Base image tensor (connect from Load Image). Leave empty if using URL."}), + "input_image_url": ("STRING", {"default": "", "tooltip": "Base image URL (used if provided)."}), } + optional_refs = {} + for idx in range(2, 9): + optional_refs[f"input_image_{idx}"] = ( + "IMAGE", + {"tooltip": f"Optional reference image tensor #{idx}. Leave empty if using URL fields."}, + ) + optional_refs[f"input_image_{idx}_url"] = ( + "STRING", + {"default": "", "tooltip": f"Optional reference image URL #{idx}."}, + ) + return { "required": { "prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "Describe the edit you want (what to change/keep)."}), - "input_image": ("STRING", {"default": "", "tooltip": "Base image URL to edit. Required."}), + **base_image_inputs, }, "optional": { **optional_refs, @@ -116,14 +124,22 @@ class Flux2ProImageEdit: def edit( self, prompt: str, - input_image: str, - input_image_2: str = "", - input_image_3: str = "", - input_image_4: str = "", - input_image_5: str = "", - input_image_6: str = "", - input_image_7: str = "", - input_image_8: str = "", + input_image, + input_image_url: str = "", + input_image_2=None, + input_image_2_url: str = "", + input_image_3=None, + input_image_3_url: str = "", + input_image_4=None, + input_image_4_url: str = "", + input_image_5=None, + input_image_5_url: str = "", + input_image_6=None, + input_image_6_url: str = "", + input_image_7=None, + input_image_7_url: str = "", + input_image_8=None, + input_image_8_url: str = "", width: int = 0, height: int = 0, seed: int = -1, @@ -135,20 +151,37 @@ class Flux2ProImageEdit: if resolution_error: return (f"Error: {resolution_error}",) - additional_images: List[str] = [ - input_image_2, - input_image_3, - input_image_4, - input_image_5, - input_image_6, - input_image_7, - input_image_8, + def resolve_image(source_tensor, source_url: str) -> Optional[str]: + if source_url and source_url.strip(): + return source_url.strip() + if source_tensor is not None: + try: + return image_tensor_to_base64(source_tensor, format="PNG") + except Exception as exc: + print(f"Warning: failed to convert image tensor to base64: {exc}") + return None + return None + + base_image_value = resolve_image(input_image, input_image_url) + if not base_image_value: + return ("Error: Provide a base image tensor or URL.",) + + additional_images: List[str] = [] + ref_pairs = [ + (input_image_2, input_image_2_url), + (input_image_3, input_image_3_url), + (input_image_4, input_image_4_url), + (input_image_5, input_image_5_url), + (input_image_6, input_image_6_url), + (input_image_7, input_image_7_url), + (input_image_8, input_image_8_url), ] - reference_payload = merge_reference_images(input_image, additional_images) + for tensor_val, url_val in ref_pairs: + resolved = resolve_image(tensor_val, url_val) + additional_images.append(resolved if resolved else "") - if not reference_payload.get("input_image"): - return ("Error: input_image URL is required for image editing.",) + reference_payload = merge_reference_images(base_image_value, additional_images) payload = { "prompt": prompt, diff --git a/nodes/flux2_utils.py b/nodes/flux2_utils.py index 32de7a4..4ca31b3 100644 --- a/nodes/flux2_utils.py +++ b/nodes/flux2_utils.py @@ -1,7 +1,12 @@ +import base64 +import io import time from typing import Any, Dict, List, Optional +import numpy as np import requests +import torch +from PIL import Image from .flux2_config import Flux2Config @@ -120,3 +125,47 @@ def merge_reference_images(main_image: str, additional_images: List[str]) -> Dic payload[f"input_image_{idx}"] = image.strip() return payload + + +def image_tensor_to_base64(image: torch.Tensor, format: str = "PNG") -> str: + """ + Convert a ComfyUI image tensor to a base64-encoded string. + + Accepts tensors in BHWC or HWC with float values in [0,1]. + """ + if image is None: + raise ValueError("Image tensor is None") + + # Take first image if batch provided + if image.ndim == 4: + image = image[0] + + if image.ndim != 3: + raise ValueError(f"Expected image tensor with 3 dimensions (HWC), got shape {tuple(image.shape)}") + + # Move to CPU and convert to numpy + if hasattr(image, "detach"): + image = image.detach() + if hasattr(image, "cpu"): + image = image.cpu() + + np_image = image.numpy() + np_image = np.clip(np_image, 0.0, 1.0) + np_image = (np_image * 255).astype(np.uint8) + + # Handle grayscale or missing channel dimension + if np_image.ndim == 2: + np_image = np.stack([np_image] * 3, axis=-1) + if np_image.shape[-1] == 1: + np_image = np.repeat(np_image, 3, axis=-1) + + pil_image = Image.fromarray(np_image) + buffer = io.BytesIO() + + save_params: Dict[str, Any] = {} + if format.upper() == "JPEG": + save_params["quality"] = 95 + + pil_image.save(buffer, format=format.upper(), **save_params) + encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") + return encoded diff --git a/requirements.txt b/requirements.txt index 0eb8cae..2db9150 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,3 @@ requests>=2.31.0 +Pillow>=10.0.0 +numpy>=1.23.0