342 lines
14 KiB
Python
342 lines
14 KiB
Python
"""
|
|
ComfyUI BananaForge Batch Node
|
|
Process multiple prompts at once, generating one image per prompt.
|
|
Supports both text-to-image and image-to-image modes.
|
|
"""
|
|
|
|
import io
|
|
import re
|
|
import base64
|
|
import json
|
|
import time
|
|
import urllib.request
|
|
import urllib.error
|
|
from pathlib import Path
|
|
from typing import Optional, List
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from comfy_api.latest import io as comfy_io
|
|
|
|
|
|
def load_env_file() -> dict:
|
|
"""Load environment variables from .env file."""
|
|
env_vars = {}
|
|
node_dir = Path(__file__).parent
|
|
env_path = node_dir / ".env"
|
|
|
|
if env_path.exists():
|
|
try:
|
|
with open(env_path, 'r', encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line and not line.startswith('#') and '=' in line:
|
|
key, _, value = line.partition('=')
|
|
key = key.strip()
|
|
value = value.strip()
|
|
if (value.startswith('"') and value.endswith('"')) or \
|
|
(value.startswith("'") and value.endswith("'")):
|
|
value = value[1:-1]
|
|
env_vars[key] = value
|
|
except Exception:
|
|
pass
|
|
return env_vars
|
|
|
|
|
|
class GeminiBatchNode(comfy_io.ComfyNode):
|
|
"""
|
|
Process multiple prompts - one image per line.
|
|
Supports both text-to-image (no input image) and image-to-image (with reference).
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> comfy_io.Schema:
|
|
return comfy_io.Schema(
|
|
node_id="BananaForgeBatch",
|
|
display_name="🍌 BananaForge Batch (Multi-Prompt)",
|
|
description="Generate multiple images from multiple prompts (one per line)",
|
|
category="Banana Forge",
|
|
inputs=[
|
|
# API Configuration
|
|
comfy_io.String.Input(
|
|
"api_url",
|
|
default="https://generativelanguage.googleapis.com/v1beta",
|
|
multiline=False,
|
|
tooltip="API base URL"
|
|
),
|
|
comfy_io.String.Input(
|
|
"api_key",
|
|
default="",
|
|
multiline=False,
|
|
tooltip="API key (leave empty to use .env file)"
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"use_env_file",
|
|
options=["no", "yes"],
|
|
default="no",
|
|
tooltip="Load API key from .env file"
|
|
),
|
|
comfy_io.String.Input(
|
|
"model_name",
|
|
default="gemini-3-pro-image-preview",
|
|
multiline=False,
|
|
tooltip="Model name"
|
|
),
|
|
|
|
# Multi-line Prompts (one per line)
|
|
comfy_io.String.Input(
|
|
"prompts",
|
|
default="A girl sitting on a chair\nA girl eating an apple\nA girl standing by the window",
|
|
multiline=True,
|
|
tooltip="Multiple prompts - ONE PER LINE. Each line generates one image."
|
|
),
|
|
|
|
# Optional Reference Image for I2I mode
|
|
comfy_io.Image.Input(
|
|
"reference_image",
|
|
optional=True,
|
|
tooltip="Optional: Reference image for image-to-image mode (used for ALL prompts)"
|
|
),
|
|
|
|
# Image Settings
|
|
comfy_io.Combo.Input(
|
|
"aspect_ratio",
|
|
options=["1:1", "16:9", "9:16", "4:3", "3:4"],
|
|
default="1:1",
|
|
tooltip="Output aspect ratio"
|
|
),
|
|
comfy_io.Combo.Input(
|
|
"image_size",
|
|
options=["default", "2K", "4K"],
|
|
default="default",
|
|
tooltip="Output resolution"
|
|
),
|
|
|
|
# Rate Limit
|
|
comfy_io.Float.Input(
|
|
"wait_between",
|
|
default=2.0,
|
|
min=0.0,
|
|
max=60.0,
|
|
step=0.5,
|
|
tooltip="Wait time between API calls (to avoid rate limits)"
|
|
),
|
|
],
|
|
outputs=[
|
|
comfy_io.Image.Output(display_name="IMAGES"),
|
|
comfy_io.Int.Output(display_name="COUNT"),
|
|
],
|
|
hidden=[
|
|
comfy_io.Hidden.prompt,
|
|
comfy_io.Hidden.unique_id,
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
api_url: str,
|
|
api_key: str,
|
|
use_env_file: str,
|
|
model_name: str,
|
|
prompts: str,
|
|
aspect_ratio: str,
|
|
image_size: str,
|
|
wait_between: float,
|
|
reference_image: Optional[torch.Tensor] = None,
|
|
hidden_prompt: Optional[dict] = None,
|
|
hidden_unique_id: Optional[str] = None,
|
|
) -> comfy_io.NodeOutput:
|
|
"""Execute batch image generation."""
|
|
|
|
# Load from .env if requested (always override when enabled)
|
|
if use_env_file == "yes":
|
|
env_vars = load_env_file()
|
|
# Always use .env values if they exist
|
|
env_api_key = env_vars.get("API_KEY", env_vars.get("GEMINI_API_KEY", ""))
|
|
env_api_url = env_vars.get("API_URL", env_vars.get("GEMINI_API_URL", ""))
|
|
if env_api_key:
|
|
api_key = env_api_key
|
|
if env_api_url:
|
|
api_url = env_api_url
|
|
|
|
# Validate inputs
|
|
if not api_url or not api_url.strip():
|
|
raise ValueError("❌ API URL cannot be empty. Set it in the node or .env file.")
|
|
if not api_key or not api_key.strip():
|
|
raise ValueError("❌ API key cannot be empty. Set it in the node or .env file.")
|
|
if not model_name or not model_name.strip():
|
|
raise ValueError("❌ Model name cannot be empty.")
|
|
if not prompts or not prompts.strip():
|
|
raise ValueError("❌ Prompts cannot be empty.")
|
|
|
|
actual_model = model_name.strip()
|
|
is_pro_model = "pro" in actual_model.lower()
|
|
|
|
# Parse prompts - one per line, skip empty lines
|
|
prompt_list = [p.strip() for p in prompts.strip().split('\n') if p.strip()]
|
|
|
|
if len(prompt_list) == 0:
|
|
raise ValueError("❌ No valid prompts found. Add prompts one per line.")
|
|
|
|
# Determine mode
|
|
mode = "img2img" if reference_image is not None else "txt2img"
|
|
print(f"🍌 [Batch] Mode: {mode.upper()}, Prompts: {len(prompt_list)}")
|
|
|
|
# Prepare reference image if provided
|
|
ref_b64 = None
|
|
if reference_image is not None:
|
|
ref_b64 = cls._tensor_to_base64(reference_image)
|
|
print(f"🍌 [Batch] Reference image prepared")
|
|
|
|
# Build endpoint
|
|
api_url = api_url.strip().rstrip('/')
|
|
endpoint = f"{api_url}/models/{actual_model}:generateContent"
|
|
|
|
# Process each prompt
|
|
generated_images = []
|
|
|
|
for i, prompt in enumerate(prompt_list):
|
|
print(f"🍌 [Batch] Processing {i+1}/{len(prompt_list)}: {prompt[:50]}...")
|
|
|
|
# Wait between calls (except first)
|
|
if i > 0 and wait_between > 0:
|
|
print(f"🍌 [Batch] Waiting {wait_between}s...")
|
|
time.sleep(wait_between)
|
|
|
|
# Add ratio/size info to prompt
|
|
enhanced_prompt = prompt
|
|
enhanced_prompt += f"\n[Image settings: aspect ratio {aspect_ratio}"
|
|
if image_size != "default":
|
|
enhanced_prompt += f", resolution {image_size}"
|
|
enhanced_prompt += "]"
|
|
|
|
# Build content parts
|
|
if ref_b64:
|
|
# Image-to-image mode
|
|
parts = [
|
|
{"text": enhanced_prompt},
|
|
{"inlineData": {"mimeType": "image/png", "data": ref_b64}}
|
|
]
|
|
else:
|
|
# Text-to-image mode
|
|
parts = [{"text": enhanced_prompt}]
|
|
|
|
# Build payload
|
|
payload = {
|
|
"contents": [{"parts": parts}],
|
|
"generationConfig": {
|
|
"responseModalities": ["IMAGE"],
|
|
"imageConfig": {"aspectRatio": aspect_ratio}
|
|
}
|
|
}
|
|
|
|
if image_size != "default" and is_pro_model:
|
|
payload["generationConfig"]["imageConfig"]["imageSize"] = image_size
|
|
|
|
try:
|
|
result = cls._call_api(endpoint, api_key.strip(), payload)
|
|
image_tensor = cls._decode_image(result["image"])
|
|
generated_images.append(image_tensor)
|
|
print(f"🍌 [Batch] ✓ Image {i+1} generated")
|
|
except Exception as e:
|
|
print(f"🍌 [Batch] ✗ Image {i+1} failed: {str(e)[:100]}")
|
|
# Continue with other prompts
|
|
|
|
if len(generated_images) == 0:
|
|
raise ValueError("❌ All prompts failed. Check API key and prompts.")
|
|
|
|
# Combine all images into a batch tensor
|
|
combined = torch.cat(generated_images, dim=0)
|
|
print(f"🍌 [Batch] Complete: {len(generated_images)}/{len(prompt_list)} images generated")
|
|
|
|
return comfy_io.NodeOutput(combined, len(generated_images))
|
|
|
|
@classmethod
|
|
def _tensor_to_base64(cls, tensor: torch.Tensor) -> str:
|
|
"""Convert tensor to base64 PNG."""
|
|
if len(tensor.shape) == 4:
|
|
tensor = tensor[0]
|
|
np_img = (tensor.numpy() * 255).astype(np.uint8)
|
|
pil_img = Image.fromarray(np_img, mode='RGB')
|
|
buffer = io.BytesIO()
|
|
pil_img.save(buffer, format='PNG')
|
|
return base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
@classmethod
|
|
def _call_api(cls, endpoint: str, api_key: str, payload: dict, max_retries: int = 3) -> dict:
|
|
"""Call the Gemini API with retry logic."""
|
|
headers = {"Content-Type": "application/json", "x-goog-api-key": api_key}
|
|
data = json.dumps(payload).encode('utf-8')
|
|
|
|
last_error = None
|
|
for attempt in range(max_retries):
|
|
try:
|
|
request = urllib.request.Request(endpoint, data=data, headers=headers, method='POST')
|
|
with urllib.request.urlopen(request, timeout=300) as response:
|
|
response_data = json.loads(response.read().decode('utf-8'))
|
|
result = {"image": None}
|
|
|
|
if 'candidates' not in response_data or len(response_data['candidates']) == 0:
|
|
if 'promptFeedback' in response_data:
|
|
raise ValueError(f"Content blocked: {response_data['promptFeedback'].get('blockReason', 'Unknown')}")
|
|
raise ValueError(f"No candidates in response")
|
|
|
|
candidate = response_data['candidates'][0]
|
|
if candidate.get('finishReason') == 'SAFETY':
|
|
raise ValueError("Content blocked by safety filters")
|
|
|
|
if 'content' not in candidate or 'parts' not in candidate['content']:
|
|
raise ValueError("No content in response")
|
|
|
|
for part in candidate['content']['parts']:
|
|
if 'inlineData' in part and 'data' in part['inlineData']:
|
|
result["image"] = part['inlineData']['data']
|
|
break
|
|
elif 'text' in part:
|
|
# Check for image URL in markdown
|
|
text = part['text']
|
|
match = re.search(r'!\[.*?\]\((https?://[^\s\)]+)\)', text)
|
|
if match:
|
|
result["image"] = cls._download_image(match.group(1))
|
|
break
|
|
|
|
if result["image"] is None:
|
|
raise ValueError("No image in response")
|
|
|
|
return result
|
|
|
|
except urllib.error.HTTPError as e:
|
|
error_msg = e.read().decode('utf-8') if e.fp else str(e)
|
|
raise ValueError(f"API Error ({e.code}): {error_msg[:200]}")
|
|
except urllib.error.URLError as e:
|
|
last_error = f"Connection error: {str(e.reason)}"
|
|
print(f"🍌 [Batch] Retry {attempt + 1}/{max_retries}: {last_error}")
|
|
if attempt < max_retries - 1:
|
|
time.sleep(2 ** attempt) # Exponential backoff
|
|
except Exception as e:
|
|
# Catch RemoteDisconnected and other transient errors
|
|
last_error = str(e)
|
|
print(f"🍌 [Batch] Retry {attempt + 1}/{max_retries}: {last_error}")
|
|
if attempt < max_retries - 1:
|
|
time.sleep(2 ** attempt)
|
|
|
|
raise ValueError(f"Failed after {max_retries} attempts: {last_error}")
|
|
|
|
@classmethod
|
|
def _download_image(cls, url: str) -> str:
|
|
"""Download image from URL and return as base64."""
|
|
req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
|
|
with urllib.request.urlopen(req, timeout=60) as response:
|
|
return base64.b64encode(response.read()).decode('utf-8')
|
|
|
|
@classmethod
|
|
def _decode_image(cls, b64_data: str) -> torch.Tensor:
|
|
"""Decode base64 to tensor."""
|
|
pil_image = Image.open(io.BytesIO(base64.b64decode(b64_data)))
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
return torch.from_numpy(np.array(pil_image).astype(np.float32) / 255.0).unsqueeze(0)
|