Claude
157b3c9ef3
Prints the queue submission response (to see what URLs fal.ai provides) and the first 5 + every 15th poll status to the ComfyUI console. https://claude.ai/code/session_01VtRzQMdLschx9awGyTGGH7
319 lines
10 KiB
Python
319 lines
10 KiB
Python
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"]
|
|
|
|
|
|
|
|
def _resolve_api_key(fal_api_key=""):
|
|
"""Return the API key from the node field, falling back to the env var."""
|
|
key = fal_api_key.strip() if fal_api_key else ""
|
|
if not key:
|
|
key = FAL_KEY
|
|
if not key:
|
|
raise ValueError(
|
|
"No fal.ai API key provided. Either enter it on the node "
|
|
"or set the FAL_KEY environment variable. "
|
|
"Get a key at https://fal.ai/dashboard/keys"
|
|
)
|
|
return key
|
|
|
|
|
|
def _submit_request(endpoint, payload, timeout, api_key):
|
|
"""Submit a request to fal.ai and poll until complete."""
|
|
url = f"https://queue.fal.run/{endpoint}"
|
|
headers = {
|
|
"Authorization": f"Key {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
resp = requests.post(url, headers=headers, json=payload, timeout=30)
|
|
if not resp.ok:
|
|
raise RuntimeError(
|
|
f"fal.ai submit failed ({resp.status_code}): {resp.text}"
|
|
)
|
|
queue_data = resp.json()
|
|
print(f"[NBP] Queue response keys: {list(queue_data.keys())}")
|
|
print(f"[NBP] Queue response: {queue_data}")
|
|
|
|
request_id = queue_data.get("request_id")
|
|
if not request_id:
|
|
raise RuntimeError(f"No request_id in queue response: {queue_data}")
|
|
|
|
# Use the URLs returned by fal.ai (they include the correct app path).
|
|
status_url = queue_data.get(
|
|
"status_url",
|
|
f"https://queue.fal.run/{endpoint}/requests/{request_id}/status",
|
|
)
|
|
result_url = queue_data.get(
|
|
"response_url",
|
|
f"https://queue.fal.run/{endpoint}/requests/{request_id}",
|
|
)
|
|
print(f"[NBP] status_url: {status_url}")
|
|
print(f"[NBP] result_url: {result_url}")
|
|
|
|
deadline = time.time() + timeout
|
|
poll_count = 0
|
|
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")
|
|
poll_count += 1
|
|
if poll_count <= 5 or poll_count % 15 == 0:
|
|
print(f"[NBP] Poll #{poll_count}: {status_data}")
|
|
|
|
if status == "COMPLETED":
|
|
result_resp = requests.get(result_url, headers=headers, timeout=30)
|
|
if not result_resp.ok:
|
|
raise RuntimeError(
|
|
f"fal.ai result fetch failed ({result_resp.status_code}): "
|
|
f"{result_resp.text}"
|
|
)
|
|
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": {
|
|
"fal_api_key": ("STRING", {"default": "", "multiline": False}),
|
|
"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,
|
|
fal_api_key="",
|
|
num_images=1,
|
|
aspect_ratio="1:1",
|
|
output_format="png",
|
|
resolution="1K",
|
|
seed=-1,
|
|
limit_generations=False,
|
|
enable_web_search=False,
|
|
timeout=300,
|
|
):
|
|
api_key = _resolve_api_key(fal_api_key)
|
|
|
|
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, api_key)
|
|
|
|
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": {
|
|
"fal_api_key": ("STRING", {"default": "", "multiline": False}),
|
|
"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": ("INT", {"default": 4, "min": 1, "max": 6}),
|
|
"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,
|
|
fal_api_key="",
|
|
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,
|
|
):
|
|
api_key = _resolve_api_key(fal_api_key)
|
|
|
|
# Build image_urls array (primary + optional reference images)
|
|
image_urls = [_pil_to_data_uri(_tensor_to_pil(image), output_format)]
|
|
if image_2 is not None:
|
|
image_urls.append(_pil_to_data_uri(_tensor_to_pil(image_2), output_format))
|
|
if image_3 is not None:
|
|
image_urls.append(_pil_to_data_uri(_tensor_to_pil(image_3), output_format))
|
|
|
|
payload = {
|
|
"prompt": prompt,
|
|
"image_urls": image_urls,
|
|
"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": safety_tolerance,
|
|
}
|
|
if seed >= 0:
|
|
payload["seed"] = seed
|
|
|
|
result = _submit_request(self.ENDPOINT, payload, timeout, api_key)
|
|
|
|
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) 🍌",
|
|
}
|