Add image tensor support and base64 upload for Flux2 edit

This commit is contained in:
RITIK
2025-11-26 13:57:58 +05:30
parent 11d769c584
commit 2cdc026d11
4 changed files with 114 additions and 29 deletions
+3 -2
View File
@@ -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`
+60 -27
View File
@@ -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,
+49
View File
@@ -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
+2
View File
@@ -1 +1,3 @@
requests>=2.31.0
Pillow>=10.0.0
numpy>=1.23.0