From dc6fe72cc46108331b84feb9eca412ac7e0b4369 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 11 Feb 2026 23:13:38 +0000 Subject: [PATCH] Add nodes.py with Text-to-Image and Image Edit node implementations Implements the two ComfyUI custom nodes for fal.ai Nano Banana Pro: - NanoBananaProTextToImage: text-to-image generation via fal-ai/nano-banana-pro - NanoBananaProImageEdit: image editing with safety_tolerance (1-6) via fal-ai/nano-banana-pro/edit Uses fal.ai queue API with polling, supports all parameters from README (aspect ratio, resolution, seed, web search, output format, etc). https://claude.ai/code/session_01VtRzQMdLschx9awGyTGGH7 --- nodes.py | 293 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 293 insertions(+) create mode 100644 nodes.py diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..cd6fed0 --- /dev/null +++ b/nodes.py @@ -0,0 +1,293 @@ +import io +import os +import json +import time +import requests +import numpy as np +import torch +from PIL import Image + + +FAL_KEY = os.environ.get("FAL_KEY", "") + +ASPECT_RATIOS = [ + "1:1", + "16:9", + "9:16", + "4:3", + "3:4", + "21:9", + "9:21", + "3:2", + "2:3", + "5:4", +] + +OUTPUT_FORMATS = ["png", "jpeg", "webp"] + +RESOLUTIONS = ["1K", "2K", "4K"] + +SAFETY_TOLERANCE_OPTIONS = ["1", "2", "3", "4", "5", "6"] + + +def _submit_request(endpoint, payload, timeout): + """Submit a request to fal.ai and poll until complete.""" + url = f"https://queue.fal.run/{endpoint}" + headers = { + "Authorization": f"Key {FAL_KEY}", + "Content-Type": "application/json", + } + + resp = requests.post(url, headers=headers, json=payload, timeout=30) + resp.raise_for_status() + queue_data = resp.json() + + request_id = queue_data.get("request_id") + if not request_id: + raise RuntimeError(f"No request_id in queue response: {queue_data}") + + status_url = f"https://queue.fal.run/{endpoint}/requests/{request_id}/status" + result_url = f"https://queue.fal.run/{endpoint}/requests/{request_id}" + + deadline = time.time() + timeout + while time.time() < deadline: + status_resp = requests.get(status_url, headers=headers, timeout=30) + status_resp.raise_for_status() + status_data = status_resp.json() + status = status_data.get("status") + + if status == "COMPLETED": + result_resp = requests.get(result_url, headers=headers, timeout=30) + result_resp.raise_for_status() + return result_resp.json() + elif status in ("FAILED", "CANCELLED"): + raise RuntimeError( + f"fal.ai request {status}: {status_data}" + ) + + time.sleep(2) + + raise TimeoutError(f"fal.ai request timed out after {timeout}s") + + +def _download_image_as_tensor(url): + """Download an image URL and return a ComfyUI IMAGE tensor (B,H,W,C float32 0-1).""" + resp = requests.get(url, timeout=60) + resp.raise_for_status() + img = Image.open(io.BytesIO(resp.content)).convert("RGB") + arr = np.array(img).astype(np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) + + +def _tensor_to_pil(tensor): + """Convert a ComfyUI IMAGE tensor (B,H,W,C) to a PIL Image (first image in batch).""" + if tensor.dim() == 4: + tensor = tensor[0] + arr = (tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8) + return Image.fromarray(arr) + + +def _pil_to_data_uri(img, fmt="png"): + """Convert a PIL image to a data URI string.""" + buf = io.BytesIO() + img.save(buf, format=fmt.upper() if fmt != "jpeg" else "JPEG") + import base64 + encoded = base64.b64encode(buf.getvalue()).decode("ascii") + mime = f"image/{fmt}" if fmt != "jpeg" else "image/jpeg" + return f"data:{mime};base64,{encoded}" + + +class NanoBananaProTextToImage: + """Generate images from text prompts using Nano Banana Pro via fal.ai.""" + + ENDPOINT = "fal-ai/nano-banana-pro" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "num_images": ("INT", {"default": 1, "min": 1, "max": 4}), + "aspect_ratio": (ASPECT_RATIOS, {"default": "1:1"}), + "output_format": (OUTPUT_FORMATS, {"default": "png"}), + "resolution": (RESOLUTIONS, {"default": "1K"}), + "seed": ("INT", {"default": -1, "min": -1, "max": 2**31 - 1}), + "limit_generations": ("BOOLEAN", {"default": False}), + "enable_web_search": ("BOOLEAN", {"default": False}), + "timeout": ("INT", {"default": 300, "min": 30, "max": 600}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("images", "description") + OUTPUT_IS_LIST = (True, False) + FUNCTION = "generate" + CATEGORY = "fal.ai/Nano Banana Pro" + + def generate( + self, + prompt, + num_images=1, + aspect_ratio="1:1", + output_format="png", + resolution="1K", + seed=-1, + limit_generations=False, + enable_web_search=False, + timeout=300, + ): + if not FAL_KEY: + raise ValueError( + "FAL_KEY environment variable is not set. " + "Get a key at https://fal.ai/dashboard/keys" + ) + + payload = { + "prompt": prompt, + "num_images": 1 if limit_generations else num_images, + "aspect_ratio": aspect_ratio, + "output_format": output_format, + "resolution": resolution, + "enable_web_search": enable_web_search, + } + if seed >= 0: + payload["seed"] = seed + + result = _submit_request(self.ENDPOINT, payload, timeout) + + images_data = result.get("images", []) + description = result.get("description", "") + + if not images_data: + raise RuntimeError( + f"No images returned. " + f"Description: {description or 'none'}" + ) + + tensors = [] + for img_info in images_data: + url = img_info.get("url", "") + if url: + tensors.append(_download_image_as_tensor(url)) + + if not tensors: + raise RuntimeError("No valid image URLs in response.") + + return (tensors, description or "") + + +class NanoBananaProImageEdit: + """Edit images with natural language using Nano Banana Pro via fal.ai.""" + + ENDPOINT = "fal-ai/nano-banana-pro/edit" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "image_2": ("IMAGE",), + "image_3": ("IMAGE",), + "num_images": ("INT", {"default": 1, "min": 1, "max": 4}), + "aspect_ratio": (ASPECT_RATIOS, {"default": "1:1"}), + "output_format": (OUTPUT_FORMATS, {"default": "png"}), + "resolution": (RESOLUTIONS, {"default": "1K"}), + "seed": ("INT", {"default": -1, "min": -1, "max": 2**31 - 1}), + "limit_generations": ("BOOLEAN", {"default": False}), + "enable_web_search": ("BOOLEAN", {"default": False}), + "safety_tolerance": (SAFETY_TOLERANCE_OPTIONS, {"default": "4"}), + "timeout": ("INT", {"default": 300, "min": 30, "max": 600}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING") + RETURN_NAMES = ("images", "description") + OUTPUT_IS_LIST = (True, False) + FUNCTION = "edit" + CATEGORY = "fal.ai/Nano Banana Pro" + + def edit( + self, + image, + prompt, + image_2=None, + image_3=None, + num_images=1, + aspect_ratio="1:1", + output_format="png", + resolution="1K", + seed=-1, + limit_generations=False, + enable_web_search=False, + safety_tolerance="4", + timeout=300, + ): + if not FAL_KEY: + raise ValueError( + "FAL_KEY environment variable is not set. " + "Get a key at https://fal.ai/dashboard/keys" + ) + + # Convert primary image to data URI + pil_img = _tensor_to_pil(image) + image_url = _pil_to_data_uri(pil_img, output_format) + + payload = { + "prompt": prompt, + "image_url": image_url, + "num_images": 1 if limit_generations else num_images, + "aspect_ratio": aspect_ratio, + "output_format": output_format, + "resolution": resolution, + "enable_web_search": enable_web_search, + "safety_tolerance": int(safety_tolerance), + } + if seed >= 0: + payload["seed"] = seed + + # Add optional reference images + if image_2 is not None: + pil_2 = _tensor_to_pil(image_2) + payload["image_2_url"] = _pil_to_data_uri(pil_2, output_format) + + if image_3 is not None: + pil_3 = _tensor_to_pil(image_3) + payload["image_3_url"] = _pil_to_data_uri(pil_3, output_format) + + result = _submit_request(self.ENDPOINT, payload, timeout) + + images_data = result.get("images", []) + description = result.get("description", "") + + if not images_data: + raise RuntimeError( + f"No images returned. " + f"Description: {description or 'none'}" + ) + + tensors = [] + for img_info in images_data: + url = img_info.get("url", "") + if url: + tensors.append(_download_image_as_tensor(url)) + + if not tensors: + raise RuntimeError("No valid image URLs in response.") + + return (tensors, description or "") + + +NODE_CLASS_MAPPINGS = { + "NanoBananaProTextToImage": NanoBananaProTextToImage, + "NanoBananaProImageEdit": NanoBananaProImageEdit, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NanoBananaProTextToImage": "Nano Banana Pro (Text to Image) 🍌", + "NanoBananaProImageEdit": "Nano Banana Pro (Image Edit) 🍌", +}