1466 lines
72 KiB
Python
1466 lines
72 KiB
Python
import torch
|
||
import numpy as np
|
||
from PIL import Image
|
||
import io
|
||
import time
|
||
import random
|
||
from google import genai
|
||
from google.genai import types
|
||
import comfy.model_management
|
||
import re
|
||
import os
|
||
import json
|
||
import urllib.request
|
||
import urllib.error
|
||
import concurrent.futures
|
||
import asyncio
|
||
import platform
|
||
import logging
|
||
import base64
|
||
import threading
|
||
import torch.nn.functional as F
|
||
import requests # Added for GAIA proxy support
|
||
|
||
def ensure_same_size(img_list):
|
||
"""Ensures all tensors in the list have the same H, W dimensions by resizing to the first one."""
|
||
if not img_list:
|
||
return []
|
||
|
||
first_img = img_list[0]
|
||
if len(first_img.shape) != 4:
|
||
first_img = first_img.unsqueeze(0)
|
||
|
||
target_h, target_w = first_img.shape[1], first_img.shape[2]
|
||
processed_list = [first_img]
|
||
|
||
for i in range(1, len(img_list)):
|
||
img = img_list[i]
|
||
if len(img.shape) != 4:
|
||
img = img.unsqueeze(0)
|
||
|
||
if img.shape[1] != target_h or img.shape[2] != target_w:
|
||
# Resize using bilinear interpolation
|
||
# Note: ComfyUI images are [B, H, W, C], but F.interpolate expects [B, C, H, W]
|
||
img_permuted = img.permute(0, 3, 1, 2)
|
||
img_resized = F.interpolate(img_permuted, size=(target_h, target_w), mode='bilinear', align_corners=False)
|
||
img_final = img_resized.permute(0, 2, 3, 1)
|
||
processed_list.append(img_final)
|
||
else:
|
||
processed_list.append(img)
|
||
|
||
return processed_list
|
||
|
||
def get_retry_wait_info(error_str, retry_count, default_wait=15):
|
||
"""Calculates backoff wait time and message based on error type and retry count."""
|
||
if not isinstance(error_str, str):
|
||
error_str = str(error_str)
|
||
|
||
err_str_lower = error_str.lower()
|
||
|
||
if "503" in err_str_lower or "overloaded" in err_str_lower:
|
||
wait_time = min(60.0, (2 ** (retry_count - 1)) + random.uniform(0, 1))
|
||
msg = f"⏳ 503 Overload | Exponential backoff: Waiting {wait_time:.2f}s..."
|
||
return wait_time, msg
|
||
|
||
if "429" in err_str_lower or "quota exhausted" in err_str_lower or "rate limit" in err_str_lower:
|
||
wait_time = 60.0
|
||
msg = f"⏳ Rate Limit Cooldown | Waiting 60s for per-minute quota refresh..."
|
||
return wait_time, msg
|
||
|
||
if "500" in err_str_lower or "internal error" in err_str_lower:
|
||
wait_time = random.uniform(5, 10)
|
||
msg = f"⏳ 500 Internal Error | Waiting {wait_time:.2f}s before retry..."
|
||
return wait_time, msg
|
||
|
||
match = re.search(r"Please retry in (\d+\.?\d*)s", error_str)
|
||
if match:
|
||
try:
|
||
wait_time = float(match.group(1))
|
||
msg = f"⏳ Dynamic Cooldown | Waiting {wait_time}s (Google request) before retry..."
|
||
return wait_time, msg
|
||
except:
|
||
pass
|
||
|
||
return float(default_wait), f"⏳ Buffer Wait | Waiting {default_wait}s to clear API congestion..."
|
||
|
||
def clean_error_msg(error_str):
|
||
"""Makes API errors more readable for the console."""
|
||
err_str_lower = str(error_str).lower()
|
||
if "503" in err_str_lower or "overloaded" in err_str_lower:
|
||
return "Model Overloaded (503) - Server capacity reached. Retrying with backoff..."
|
||
if "429" in err_str_lower or "quota exhausted" in err_str_lower or "rate limit" in err_str_lower:
|
||
return "Rate Limit/Quota Exhausted (429). Waiting for rolling 60s refresh..."
|
||
if "404" in err_str_lower or "not found" in err_str_lower:
|
||
return f"Not Found/Model Unavailable (404). Check model name or region. Detail: {str(error_str)[:100]}"
|
||
if "403" in err_str_lower or "permission denied" in err_str_lower:
|
||
return f"Permission Denied (403). Check API key and billing. Detail: {str(error_str)[:100]}"
|
||
if "400" in err_str_lower or "invalid argument" in err_str_lower:
|
||
return f"Invalid Request (400). Check arguments (e.g., inlineData). Detail: {str(error_str)[:100]}"
|
||
if "500" in err_str_lower or "internal error" in err_str_lower:
|
||
return f"Internal Server Error (500). Google side issue. Retrying shortly..."
|
||
|
||
clean = re.sub(r"\{'error': \{.*?\}\}", "", str(error_str)).strip()
|
||
return clean if clean else str(error_str)
|
||
|
||
# ============================================================================
|
||
# VGG PROXY HELPER — Creates a genai.Client from a key OR a full URL
|
||
# ============================================================================
|
||
def _create_genai_client(current_key):
|
||
"""
|
||
Creates a genai.Client OR a GAIA config dict. Supports three modes:
|
||
|
||
1) Standard Google API key (e.g. "AIzaSy..."):
|
||
→ client = genai.Client(api_key=key)
|
||
|
||
2) Full proxy URL (starts with http:// or https://):
|
||
→ client = genai.Client(http_options={"base_url": url})
|
||
|
||
3) GAIA proxy (format "GAIA::base_url::api-key"):
|
||
→ Returns a dict {"mode": "gaia", "base_url": ..., "api_key": ...}
|
||
→ Caller must use manual HTTP call with header "api-key" instead of SDK
|
||
Example: "GAIA::https://int-apim-genaigw.lpl-cloud.com/uat/gaia/vertex::MY_KEY"
|
||
"""
|
||
key = current_key.strip()
|
||
|
||
# Mode 3: GAIA Proxy — raw HTTP with header auth
|
||
if key.startswith("GAIA::"):
|
||
parts = key.split("::", 2)
|
||
gaia_url = parts[1] if len(parts) > 1 else ""
|
||
gaia_key = parts[2] if len(parts) > 2 else ""
|
||
print(f" 🔗 GAIA Proxy Mode | URL: {gaia_url[:50]}...")
|
||
return {
|
||
"mode": "gaia",
|
||
"base_url": gaia_url,
|
||
"api_key": gaia_key,
|
||
}
|
||
|
||
# Mode 2: Full proxy URL
|
||
if key.startswith(("http://", "https://")):
|
||
print(f" 🔗 Proxy Mode | Using custom base URL")
|
||
client = genai.Client(
|
||
api_key="PROXY_ENCAPSULATED_TOKEN",
|
||
http_options={
|
||
"base_url": key,
|
||
"timeout": 300000,
|
||
"client_args": {"http2": False}
|
||
}
|
||
)
|
||
return client
|
||
|
||
# Mode 1: Standard Google API key
|
||
client = genai.Client(
|
||
api_key=key,
|
||
http_options={
|
||
"timeout": 300000,
|
||
"client_args": {"http2": False}
|
||
}
|
||
)
|
||
return client
|
||
|
||
def _gaia_call_image(gaia_cfg, actual_model, prompt, system_instruction, images_tensor, response_modalities, seed, resolution, aspect_ratio, curr_batch, total_batch):
|
||
"""
|
||
GAIA proxy path for IMAGE generation.
|
||
Uses requests.post() with header "api-key" instead of Google SDK.
|
||
"""
|
||
url = f"{gaia_cfg['base_url']}/models/{actual_model}:generateContent"
|
||
|
||
parts = []
|
||
if images_tensor is not None:
|
||
for j in range(images_tensor.shape[0]):
|
||
img_tensor = images_tensor[j]
|
||
np_img = 255. * img_tensor.cpu().numpy()
|
||
pil_img = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
|
||
buf = io.BytesIO()
|
||
pil_img.save(buf, format="PNG")
|
||
parts.append({
|
||
"inlineData": {
|
||
"mimeType": "image/png",
|
||
"data": base64.b64encode(buf.getvalue()).decode("utf-8"),
|
||
}
|
||
})
|
||
parts.append({"text": prompt})
|
||
|
||
modalities_list = ["IMAGE"] if response_modalities == "IMAGE" else ["IMAGE", "TEXT"]
|
||
|
||
# Payload reconstruction for GAIA REST call
|
||
# Added System Instruction, Seed, and Aspect Ratio support
|
||
payload = {
|
||
"contents": [{"role": "user", "parts": parts}],
|
||
"generationConfig": {
|
||
"candidateCount": 1,
|
||
"responseModalities": modalities_list,
|
||
"seed": int(seed) % 2147483647,
|
||
"imageConfig": {"imageSize": resolution}
|
||
},
|
||
"safetySettings": [
|
||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
|
||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
|
||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"},
|
||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_ONLY_HIGH"},
|
||
]
|
||
}
|
||
|
||
if system_instruction and system_instruction.strip():
|
||
payload["systemInstruction"] = {"parts": [{"text": system_instruction.strip() + "\n\n" + GEMINI_IMAGE_SYS_PROMPT}]}
|
||
else:
|
||
payload["systemInstruction"] = {"parts": [{"text": GEMINI_IMAGE_SYS_PROMPT}]}
|
||
|
||
if aspect_ratio not in ("auto", "Aspect Ratio Input"):
|
||
payload["generationConfig"]["imageConfig"]["aspectRatio"] = aspect_ratio
|
||
|
||
batch_info = f" | Batch {curr_batch}/{total_batch}" if total_batch > 1 else ""
|
||
print(f"\n🎨 GAIA Image{batch_info} | Generating image...")
|
||
print(f" Model: {actual_model} | URL: {url[:80]}...")
|
||
|
||
response = requests.post(
|
||
url,
|
||
headers={
|
||
"Content-Type": "application/json",
|
||
"api-key": gaia_cfg["api_key"],
|
||
},
|
||
json=payload,
|
||
timeout=120,
|
||
)
|
||
if response.status_code != 200:
|
||
print(f" ❌ GAIA Response failed ({response.status_code}): {response.text}")
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
if "candidates" not in data or not data["candidates"]:
|
||
raise ValueError(f"GAIA proxy returned no candidates. Detail: {data}")
|
||
|
||
resp_parts = data["candidates"][0]["content"]["parts"]
|
||
full_text = " ".join(p["text"] for p in resp_parts if "text" in p)
|
||
b64_data = next(
|
||
(p["inlineData"]["data"] for p in resp_parts if "inlineData" in p),
|
||
None,
|
||
)
|
||
if b64_data:
|
||
pil_out = Image.open(io.BytesIO(base64.b64decode(b64_data))).convert("RGB")
|
||
out_tensor = torch.from_numpy(
|
||
np.array(pil_out).astype(np.float32) / 255.0
|
||
)[None,]
|
||
print(f" ✅ Success! GAIA image received.")
|
||
return out_tensor[:1], full_text
|
||
|
||
raise ValueError("GAIA proxy returned no image in parts.")
|
||
|
||
|
||
def _gaia_call_text(gaia_cfg, actual_model, user_prompt, system_instruction, images_tensor, seed, max_tokens, batch_idx, batch_count):
|
||
"""
|
||
GAIA proxy path for TEXT generation.
|
||
Uses requests.post() with header "api-key" instead of Google SDK.
|
||
"""
|
||
url = f"{gaia_cfg['base_url']}/models/{actual_model}:generateContent"
|
||
|
||
parts = []
|
||
if images_tensor is not None:
|
||
for j in range(images_tensor.shape[0]):
|
||
img_tensor = images_tensor[j]
|
||
np_img = 255. * img_tensor.cpu().numpy()
|
||
pil_img = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
|
||
buf = io.BytesIO()
|
||
pil_img.save(buf, format="PNG")
|
||
parts.append({
|
||
"inlineData": {
|
||
"mimeType": "image/png",
|
||
"data": base64.b64encode(buf.getvalue()).decode("utf-8"),
|
||
}
|
||
})
|
||
if user_prompt:
|
||
parts.append({"text": user_prompt})
|
||
|
||
payload = {
|
||
"contents": [{"role": "user", "parts": parts}],
|
||
"generationConfig": {
|
||
"candidateCount": 1,
|
||
"responseModalities": ["TEXT"],
|
||
"seed": int(seed) % 2147483647,
|
||
},
|
||
"safetySettings": [
|
||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
|
||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
|
||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"},
|
||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_ONLY_HIGH"},
|
||
]
|
||
}
|
||
|
||
if system_instruction and system_instruction.strip():
|
||
payload["systemInstruction"] = {"parts": [{"text": system_instruction.strip()}]}
|
||
|
||
if max_tokens > 0:
|
||
payload["generationConfig"]["maxOutputTokens"] = max_tokens
|
||
|
||
batch_str = f" | Batch {batch_idx+1}/{batch_count}" if batch_count > 1 else ""
|
||
print(f"\n📝 GAIA Prompt{batch_str} | Processing...")
|
||
print(f" Model: {actual_model} | URL: {url[:80]}...")
|
||
|
||
response = requests.post(
|
||
url,
|
||
headers={
|
||
"Content-Type": "application/json",
|
||
"api-key": gaia_cfg["api_key"],
|
||
},
|
||
json=payload,
|
||
timeout=120,
|
||
)
|
||
if response.status_code != 200:
|
||
print(f" ❌ GAIA Response failed ({response.status_code}): {response.text}")
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
if "candidates" not in data or not data["candidates"]:
|
||
raise ValueError(f"GAIA proxy returned no candidates. Detail: {data}")
|
||
|
||
resp_parts = data["candidates"][0]["content"]["parts"]
|
||
full_text = " ".join(p["text"] for p in resp_parts if "text" in p)
|
||
if not full_text:
|
||
raise ValueError("GAIA proxy returned empty text.")
|
||
|
||
print(f" ✅ Complete | Received {len(full_text)} characters via GAIA.")
|
||
return full_text
|
||
# ============================================================================
|
||
|
||
# Official ComfyUI system prompt for Gemini image generation
|
||
GEMINI_IMAGE_SYS_PROMPT = (
|
||
"You are an expert image-generation engine. You must ALWAYS produce an image.\n"
|
||
"Interpret all user input—regardless of "
|
||
"format, intent, or abstraction—as literal visual directives for image composition.\n"
|
||
"If a prompt is conversational or lacks specific visual details, "
|
||
"you must creatively invent a concrete visual scenario that depicts the concept.\n"
|
||
"Prioritize generating the visual representation above any text, formatting, or conversational requests."
|
||
)
|
||
|
||
# Module-level cache for ComfyUI Org auth token.
|
||
_CACHED_COMFY_ORG_TOKEN = None
|
||
|
||
def _on_prompt_intercept_token(json_data):
|
||
"""Intercept auth tokens from prompt queue requests to cache for nodes."""
|
||
global _CACHED_COMFY_ORG_TOKEN
|
||
try:
|
||
extra_data = json_data.get("extra_data", {})
|
||
token = extra_data.get("auth_token_comfy_org") or extra_data.get("api_key_comfy_org")
|
||
if token:
|
||
_CACHED_COMFY_ORG_TOKEN = token
|
||
except Exception:
|
||
pass
|
||
return json_data
|
||
|
||
try:
|
||
from server import PromptServer
|
||
PromptServer.instance.add_on_prompt_handler(_on_prompt_intercept_token)
|
||
except Exception:
|
||
pass
|
||
|
||
def get_comfy_org_auth(hidden_token=None):
|
||
"""Attempts to get ComfyUI Org authentication token."""
|
||
# We first try to read from comfy-api config if available
|
||
try:
|
||
from comfy_api_nodes.util._helpers import default_base_url
|
||
comfy_api_base = default_base_url()
|
||
except ImportError:
|
||
comfy_api_base = "https://api.comfy.org"
|
||
|
||
auth_header = {}
|
||
auth_token = hidden_token
|
||
if isinstance(auth_token, list):
|
||
auth_token = auth_token[0] if auth_token else ""
|
||
|
||
# Check our own local cached token first
|
||
if not auth_token:
|
||
global _CACHED_COMFY_ORG_TOKEN
|
||
if _CACHED_COMFY_ORG_TOKEN:
|
||
auth_token = _CACHED_COMFY_ORG_TOKEN
|
||
|
||
# Check system modules for any cached token from other nodes (like Chatbot-311)
|
||
if not auth_token:
|
||
try:
|
||
import sys
|
||
for m_name, m in list(sys.modules.items()):
|
||
if hasattr(m, "_CACHED_COMFY_ORG_TOKEN"):
|
||
token = getattr(m, "_CACHED_COMFY_ORG_TOKEN", None)
|
||
if token:
|
||
auth_token = token
|
||
_CACHED_COMFY_ORG_TOKEN = token
|
||
break
|
||
except Exception:
|
||
pass
|
||
|
||
# Try global arguments (CLI) or environment if hidden_token not provided
|
||
if not auth_token:
|
||
try:
|
||
from comfy.cli_args import args
|
||
auth_token = getattr(args, "api_key_comfy_org", None)
|
||
except ImportError:
|
||
pass
|
||
|
||
if not auth_token:
|
||
# Check env vars
|
||
auth_token = os.environ.get("COMFY_API_TOKEN") or os.environ.get("COMFY_ORG_API_KEY")
|
||
|
||
if auth_token:
|
||
auth_header["Authorization"] = f"Bearer {auth_token}"
|
||
# Some versions use X-API-KEY instead, we can provide both or just the Bearer based on typical ComfyUI usage
|
||
auth_header["X-API-KEY"] = auth_token
|
||
|
||
return comfy_api_base, auth_header, auth_token
|
||
|
||
# Shared list of aspect ratio options used by both the selector and the image gen node
|
||
ASPECT_RATIO_OPTIONS = ["Aspect Ratio Input", "auto", "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"]
|
||
RESOLUTION_OPTIONS = ["1K", "2K", "4K"]
|
||
IMAGE_MODEL_OPTIONS = ["gemini-3.1-pro-image-preview", "Nano Banana 2 (Gemini 3.1 Flash Image)"]
|
||
GRID_FORMAT_OPTIONS = ["1x1", "2x2", "3x3", "4x4", "6x2", "2x6", "3x1", "1x3", "4x1", "1x4"]
|
||
|
||
_IMAGE_GEN_CACHE = {}
|
||
_PROMPT_GEN_CACHE = {}
|
||
|
||
class GeminiAspectRatioSelector:
|
||
"""
|
||
ComfyUI Node: Aspect Ratio Selector.
|
||
Outputs the selected aspect ratio as a connectable value.
|
||
Can be connected to one or multiple Gemini 3 Pro Image nodes.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"aspect_ratio": (ASPECT_RATIO_OPTIONS, {"default": "auto"}),
|
||
},
|
||
"optional": {
|
||
"aspect_ratio_in": ("STRING", {"forceInput": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (ASPECT_RATIO_OPTIONS,)
|
||
RETURN_NAMES = ("aspect_ratio",)
|
||
FUNCTION = "select"
|
||
CATEGORY = "Gemini AI"
|
||
|
||
def select(self, aspect_ratio, aspect_ratio_in=None):
|
||
if aspect_ratio == "Aspect Ratio Input":
|
||
if aspect_ratio_in is not None and isinstance(aspect_ratio_in, str):
|
||
clean_ar = aspect_ratio_in.strip()
|
||
if clean_ar in ASPECT_RATIO_OPTIONS and clean_ar != "Aspect Ratio Input":
|
||
return (clean_ar,)
|
||
elif re.match(r"^\d+:\d+$", clean_ar):
|
||
try:
|
||
aw, ah = map(int, clean_ar.split(":"))
|
||
if aw > 0 and ah > 0:
|
||
# Find closest from ASPECT_RATIO_OPTIONS (excluding "auto" and "Aspect Ratio Input")
|
||
actual_ratio = aw / ah
|
||
common_ratios = [
|
||
(1, 1), (2, 3), (3, 2), (3, 4), (4, 3),
|
||
(4, 5), (5, 4), (9, 16), (16, 9), (21, 9)
|
||
]
|
||
best_ratio = None
|
||
min_diff = float('inf')
|
||
for rw, rh in common_ratios:
|
||
diff = abs(actual_ratio - (rw / rh))
|
||
if diff < min_diff:
|
||
min_diff = diff
|
||
best_ratio = f"{rw}:{rh}"
|
||
if best_ratio:
|
||
return (best_ratio,)
|
||
except:
|
||
pass
|
||
return ("auto",)
|
||
return (aspect_ratio,)
|
||
|
||
class GeminiResolutionSelector:
|
||
"""
|
||
ComfyUI Node: Resolution Selector.
|
||
Outputs the selected resolution as a connectable value.
|
||
Can be connected to one or multiple Gemini 3 Pro Image nodes.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"resolution": (RESOLUTION_OPTIONS, {"default": "1K"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (RESOLUTION_OPTIONS,)
|
||
RETURN_NAMES = ("resolution",)
|
||
FUNCTION = "select"
|
||
CATEGORY = "Gemini AI"
|
||
|
||
def select(self, resolution):
|
||
return (resolution,)
|
||
|
||
class GeminiImageModelSelector:
|
||
"""
|
||
ComfyUI Node: Image Model Selector.
|
||
Outputs the selected model as a connectable value.
|
||
Can be connected to one or multiple Gemini 3 Pro Image nodes.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"model": (IMAGE_MODEL_OPTIONS, {"default": "gemini-3-pro-image-preview"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (IMAGE_MODEL_OPTIONS,)
|
||
RETURN_NAMES = ("model",)
|
||
FUNCTION = "select"
|
||
CATEGORY = "Gemini AI"
|
||
|
||
def select(self, model):
|
||
return (model,)
|
||
|
||
class GeminiGridFormatSelector:
|
||
"""
|
||
ComfyUI Node: Grid Format Selector.
|
||
Outputs the selected grid format as a connectable value.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"grid_format": (GRID_FORMAT_OPTIONS, {"default": "1x1"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (GRID_FORMAT_OPTIONS,)
|
||
RETURN_NAMES = ("grid_format",)
|
||
FUNCTION = "select"
|
||
CATEGORY = "Gemini AI"
|
||
|
||
def select(self, grid_format):
|
||
return (grid_format,)
|
||
|
||
# Set Windows Loop Policy once at module level for entire process
|
||
if platform.system() == 'Windows':
|
||
try:
|
||
from asyncio import WindowsProactorEventLoopPolicy
|
||
asyncio.set_event_loop_policy(WindowsProactorEventLoopPolicy())
|
||
except ImportError:
|
||
pass
|
||
|
||
class Gemini3ProImageGenNode:
|
||
"""
|
||
ComfyUI Node for Google Gemini 3 Pro (Image Generation).
|
||
Features:
|
||
- Supports gemini-3-pro-image-preview.
|
||
- 3-Key Iterative Retry System with configurable rounds.
|
||
- Enhanced English logging with Batch counter support.
|
||
"""
|
||
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"system_instruction": ("STRING", {"multiline": True, "default": GEMINI_IMAGE_SYS_PROMPT, "placeholder": "System instructions..."}),
|
||
"prompt": ("STRING", {"multiline": True, "default": "Make this image cyberpunk style", "placeholder": "Enter your prompt here..."}),
|
||
"model": (IMAGE_MODEL_OPTIONS, {"default": "gemini-3-pro-image-preview"}),
|
||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||
|
||
"aspect_ratio": (ASPECT_RATIO_OPTIONS, {"default": "auto"}),
|
||
"resolution": (RESOLUTION_OPTIONS, {"default": "1K"}),
|
||
"response_modalities": (["IMAGE+TEXT", "IMAGE"], {"default": "IMAGE+TEXT"}),
|
||
|
||
"api_key_1": ("STRING", {"multiline": False, "default": "", "placeholder": "API Key or full proxy URL (Required)"}),
|
||
"api_key_2": ("STRING", {"multiline": False, "default": "", "placeholder": "Backup Key/URL 1 (Optional)"}),
|
||
"api_key_3": ("STRING", {"multiline": False, "default": "", "placeholder": "Backup Key/URL 2 (Optional)"}),
|
||
"api_max_retries": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1}),
|
||
"batch_delay": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 60.0, "step": 0.1}),
|
||
"use_comfyui_credits": ("BOOLEAN", {"default": True, "label_on": "Use ComfyUI Credits", "label_off": "Use Custom API Keys"}),
|
||
"bypass": ("BOOLEAN", {"default": False, "label_on": "Use Last Generation", "label_off": "Generate New"}),
|
||
},
|
||
"optional": {
|
||
"images": ("IMAGE",),
|
||
},
|
||
"hidden": {
|
||
"auth_token_comfy_org": "AUTH_TOKEN_COMFY_ORG",
|
||
"unique_id": "UNIQUE_ID",
|
||
}
|
||
}
|
||
|
||
INPUT_IS_LIST = True
|
||
RETURN_TYPES = ("IMAGE", "STRING")
|
||
RETURN_NAMES = ("processed_image", "text_response")
|
||
OUTPUT_IS_LIST = (False, True)
|
||
FUNCTION = "process_batch"
|
||
CATEGORY = "Gemini AI"
|
||
|
||
def process_batch(self, system_instruction, prompt, model, seed, aspect_ratio, resolution, response_modalities, api_key_1, api_key_2, api_key_3, api_max_retries, batch_delay, use_comfyui_credits, bypass=None, images=None, auth_token_comfy_org="", unique_id=None):
|
||
node_id = unique_id[0] if isinstance(unique_id, list) else unique_id
|
||
is_bypass = False
|
||
if bypass is not None:
|
||
is_bypass = bypass[0] if isinstance(bypass, list) else bypass
|
||
|
||
if is_bypass and node_id and node_id in _IMAGE_GEN_CACHE:
|
||
print(f"♻️ Gemini Image Gen | Bypassed. Returning last generation for node {node_id}.")
|
||
return _IMAGE_GEN_CACHE[node_id]
|
||
|
||
# Determine total items to process based on prompts
|
||
batch_size = len(prompt)
|
||
|
||
results_images = []
|
||
results_texts = []
|
||
|
||
available_keys = [k.strip() for k in [api_key_1[0] if api_key_1 else "",
|
||
api_key_2[0] if api_key_2 else "",
|
||
api_key_3[0] if api_key_3 else ""] if k and k.strip()]
|
||
|
||
# --- PRE-PROCESS IMAGES ONCE ---
|
||
pre_images_tensor = None
|
||
pre_image_parts = []
|
||
if images is not None:
|
||
img_list = images if isinstance(images, list) else [images]
|
||
valid_images = [img for img in img_list if img is not None]
|
||
if valid_images:
|
||
try:
|
||
# Robust batching with automatic resizing if mismatch occurs
|
||
synced_images = ensure_same_size(valid_images)
|
||
pre_images_tensor = torch.cat(synced_images, dim=0)
|
||
print(f" 📊 Input Synced | Pre-processing context of {pre_images_tensor.shape[0]} image(s).")
|
||
for j in range(pre_images_tensor.shape[0]):
|
||
img_tensor = pre_images_tensor[j]
|
||
np_img = 255. * img_tensor.cpu().numpy()
|
||
pil_img = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
|
||
buffered = io.BytesIO()
|
||
pil_img.save(buffered, format="PNG")
|
||
pre_image_parts.append(types.Part.from_bytes(data=buffered.getvalue(), mime_type="image/png"))
|
||
except Exception as e:
|
||
print(f" ⚠️ Pre-process Context Warning: {e}")
|
||
|
||
# --- EXECUTION STRATEGY ---
|
||
# Parallel execution for max throughput
|
||
u_credits = use_comfyui_credits[0] if isinstance(use_comfyui_credits, list) else use_comfyui_credits
|
||
max_workers = min(batch_size, 4)
|
||
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
futures = []
|
||
for i in range(batch_size):
|
||
# Intra-batch pacing delay
|
||
curr_delay = batch_delay[i % len(batch_delay)] if isinstance(batch_delay, list) else batch_delay
|
||
if i > 0 and curr_delay > 0:
|
||
time.sleep(curr_delay)
|
||
|
||
# Resolve parameters
|
||
curr_sys = system_instruction[i % len(system_instruction)]
|
||
curr_prompt = prompt[i % len(prompt)]
|
||
curr_model = model[i % len(model)]
|
||
|
||
if len(seed) == 1:
|
||
curr_seed = (seed[0] + i) % 0xffffffffffffffff
|
||
else:
|
||
curr_seed = seed[i % len(seed)]
|
||
|
||
curr_aspect = aspect_ratio[i % len(aspect_ratio)]
|
||
curr_res = resolution[i % len(resolution)]
|
||
curr_modal = response_modalities[i % len(response_modalities)]
|
||
|
||
k1 = api_key_1[0] if api_key_1 else ""
|
||
k2 = api_key_2[0] if api_key_2 else ""
|
||
k3 = api_key_3[0] if api_key_3 else ""
|
||
max_r = api_max_retries[0] if api_max_retries else 10
|
||
round_robin_idx = (i % len(available_keys)) if available_keys else 0
|
||
|
||
futures.append(executor.submit(
|
||
self._exec_single,
|
||
curr_sys, curr_prompt, curr_model, curr_seed, curr_aspect, curr_res, curr_modal,
|
||
k1, k2, k3, max_r,
|
||
images,
|
||
u_credits,
|
||
auth_token=auth_token_comfy_org,
|
||
curr_batch=i+1,
|
||
total_batch=batch_size,
|
||
start_key_idx=round_robin_idx,
|
||
pre_images_tensor=pre_images_tensor,
|
||
pre_image_parts=pre_image_parts
|
||
))
|
||
|
||
for i, future in enumerate(futures):
|
||
try:
|
||
out_img, out_txt = future.result()
|
||
results_images.append(out_img)
|
||
results_texts.append(out_txt)
|
||
except Exception as e:
|
||
is_interrupted = "Interrupted by user" in str(e)
|
||
if is_interrupted:
|
||
print(f" 🛑 Batch Interrupted by User at item {i+1}/{batch_size}")
|
||
else:
|
||
print(f" 🛑 Unexpected Critical Error in batch item {i+1}: {e}")
|
||
|
||
# Fallback to prevent queue interruption and maintain batch size
|
||
print(f" ⚠️ Padding remaining {batch_size - len(results_images)} item(s) with placeholders to preserve generated batch...")
|
||
while len(results_images) < batch_size:
|
||
placeholder = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||
results_images.append(placeholder)
|
||
err_msg = "Generation cancelled by user" if is_interrupted else f"CRITICAL ERROR: {e}"
|
||
results_texts.append(err_msg)
|
||
|
||
break # Stop processing further futures
|
||
|
||
image_batch = torch.cat(ensure_same_size(results_images), dim=0)
|
||
result = (image_batch, results_texts)
|
||
if node_id:
|
||
_IMAGE_GEN_CACHE[node_id] = result
|
||
return result
|
||
|
||
def _exec_single(self, system_instruction, prompt, model, seed, aspect_ratio, resolution, response_modalities, api_key_1, api_key_2, api_key_3, api_max_retries, images=None, use_comfyui_credits=True, auth_token="", curr_batch=1, total_batch=1, start_key_idx=0, pre_images_tensor=None, pre_image_parts=None):
|
||
# Fallback helper to prevent queue interruption
|
||
def get_error_fallback(error_msg):
|
||
print(f" 🛑 Gemini API Error (Item {curr_batch}/{total_batch}): {error_msg}")
|
||
print(f" ⚠️ Returning placeholder to prevent queue interruption.")
|
||
placeholder = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
||
return placeholder, f"ERROR: {error_msg}"
|
||
|
||
if not prompt or not prompt.strip():
|
||
return get_error_fallback("Prompt is required!")
|
||
|
||
all_keys = [k.strip() for k in [api_key_1, api_key_2, api_key_3] if k and k.strip()]
|
||
if not all_keys:
|
||
return get_error_fallback("No API Key provided!")
|
||
|
||
# Re-order keys to start with the load-balanced index first, followed by the others naturally
|
||
keys = all_keys[start_key_idx:] + all_keys[:start_key_idx]
|
||
|
||
image_parts = []
|
||
images_tensor = None
|
||
|
||
if pre_images_tensor is not None:
|
||
images_tensor = pre_images_tensor
|
||
image_parts = pre_image_parts if pre_image_parts is not None else []
|
||
else:
|
||
# Pre-process images to ensure we have a single tensor for Credits path
|
||
# but also can iterate for Direct path
|
||
img_list = []
|
||
if images is not None:
|
||
if isinstance(images, list):
|
||
img_list = [img for img in images if img is not None]
|
||
else:
|
||
img_list = [images]
|
||
|
||
# Merge with robust resizing for Credits path compatibility
|
||
if img_list:
|
||
try:
|
||
synced_images = ensure_same_size(img_list)
|
||
images_tensor = torch.cat(synced_images, dim=0)
|
||
except Exception as e:
|
||
print(f" ⚠️ Input Sync Warning | Could not cat images: {e}. Using first input only.")
|
||
images_tensor = img_list[0] if len(img_list[0].shape) == 4 else img_list[0].unsqueeze(0)
|
||
|
||
# Extract parts for Direct GenAI SDK path
|
||
if images_tensor is not None:
|
||
for j in range(images_tensor.shape[0]):
|
||
img_tensor = images_tensor[j]
|
||
np_img = 255. * img_tensor.cpu().numpy()
|
||
pil_img = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
|
||
|
||
# Log dimensions minimally as requested
|
||
if j == 0 or j == images_tensor.shape[0] - 1:
|
||
print(f" 📸 Image #{j+1} | {pil_img.width}x{pil_img.height}")
|
||
|
||
buffered = io.BytesIO()
|
||
pil_img.save(buffered, format="PNG")
|
||
image_parts.append(types.Part.from_bytes(data=buffered.getvalue(), mime_type="image/png"))
|
||
|
||
def call_api(current_key):
|
||
# --- VGG+GAIA MODIFICATION: Support proxy URL, GAIA proxy, or standard API key ---
|
||
client_or_config = _create_genai_client(current_key)
|
||
|
||
actual_model = model
|
||
if "gemini-3" in actual_model and "image" in actual_model:
|
||
actual_model = "gemini-3-pro-image-preview"
|
||
elif "Nano Banana 2" in actual_model:
|
||
actual_model = "gemini-3-pro-image-preview"
|
||
|
||
# ──── GAIA PROXY PATH ────
|
||
if isinstance(client_or_config, dict) and client_or_config.get("mode") == "gaia":
|
||
return _gaia_call_image(
|
||
client_or_config, actual_model, prompt, system_instruction, images_tensor,
|
||
response_modalities, seed, resolution, aspect_ratio, curr_batch, total_batch
|
||
)
|
||
|
||
# ──── STANDARD GOOGLE SDK PATH ────
|
||
client = client_or_config
|
||
|
||
actual_model = model
|
||
if "gemini-3" in actual_model and "image" in actual_model:
|
||
actual_model = "gemini-3-pro-image-preview"
|
||
elif "Nano Banana 2" in actual_model:
|
||
actual_model = "gemini-3-pro-image-preview"
|
||
|
||
effective_system = GEMINI_IMAGE_SYS_PROMPT
|
||
if system_instruction and system_instruction.strip():
|
||
# VGG logic: Append default prompt to user instructions
|
||
effective_system = system_instruction.strip() + "\n\n" + GEMINI_IMAGE_SYS_PROMPT
|
||
|
||
modalities_list = ["IMAGE"] if response_modalities == "IMAGE" else ["IMAGE", "TEXT"]
|
||
no_aspect_ratio_models = {"gemini-2.0-flash", "gemini-2.0-pro-exp-02-05", "gemini-3.1-pro-preview"}
|
||
|
||
safety_settings = [
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
]
|
||
|
||
config_kwargs = {
|
||
"system_instruction": effective_system,
|
||
"seed": int(seed) % 2147483647,
|
||
"response_modalities": modalities_list,
|
||
"image_config": {"image_size": resolution},
|
||
"automatic_function_calling": types.AutomaticFunctionCallingConfig(disable=True),
|
||
"safety_settings": safety_settings,
|
||
}
|
||
if aspect_ratio not in ("auto", "Aspect Ratio Input") and actual_model not in no_aspect_ratio_models:
|
||
config_kwargs["image_config"]["aspect_ratio"] = aspect_ratio
|
||
|
||
TEXT_ONLY_RETRIES = 3
|
||
for text_retry in range(TEXT_ONLY_RETRIES):
|
||
# Informative logging with Batch counter
|
||
batch_info = f" | Batch {curr_batch}/{total_batch}" if total_batch > 1 else ""
|
||
print(f"\n🎨 Gemini Image{batch_info} | Generating image...")
|
||
print(f" Model: {actual_model} | Resolution: {resolution} | Ratio: {aspect_ratio} | Seed: {seed}")
|
||
# Safe logging: show last 4 chars for keys, or "proxy" for URLs
|
||
key_display = "proxy-url" if current_key.strip().startswith("http") else f"****{current_key[-4:]}"
|
||
key_idx = all_keys.index(current_key) + 1 if current_key in all_keys else 1
|
||
print(f" Key: Using Key #{key_idx} ({key_display})")
|
||
if text_retry > 0:
|
||
print(f" ↻ Safety Retry {text_retry}/{TEXT_ONLY_RETRIES - 1} (Retrying same key...)")
|
||
|
||
# REVERTED TO CORE BEHAVIOR: Just [Prompt] + [Images]
|
||
content_parts = [types.Part.from_text(text=prompt)] + image_parts
|
||
|
||
response_stream = client.models.generate_content_stream(
|
||
model=actual_model,
|
||
contents=[types.Content(role="user", parts=content_parts)],
|
||
config=types.GenerateContentConfig(**config_kwargs),
|
||
)
|
||
|
||
out_tensor = None
|
||
full_text = ""
|
||
image_found = False
|
||
chunk_count = 0
|
||
last_finish_reason = None
|
||
for chunk in response_stream:
|
||
chunk_count += 1
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
if not chunk.candidates or not chunk.candidates[0].content:
|
||
if chunk.candidates and chunk.candidates[0].finish_reason:
|
||
last_finish_reason = chunk.candidates[0].finish_reason
|
||
continue
|
||
|
||
last_finish_reason = chunk.candidates[0].finish_reason
|
||
if not chunk.candidates[0].content.parts:
|
||
continue
|
||
|
||
for part in chunk.candidates[0].content.parts:
|
||
if part.inline_data and part.inline_data.data:
|
||
image_data = part.inline_data.data
|
||
out_pil = Image.open(io.BytesIO(image_data)).convert("RGB")
|
||
out_np = np.array(out_pil).astype(np.float32) / 255.0
|
||
out_tensor = torch.from_numpy(out_np)[None,]
|
||
image_found = True
|
||
if part.text:
|
||
full_text += part.text
|
||
|
||
if last_finish_reason == types.FinishReason.IMAGE_OTHER:
|
||
raise ValueError(f"IMAGE_OTHER: Image generation blocked by Google's policy. Try a more generic prompt.")
|
||
elif last_finish_reason not in [types.FinishReason.STOP, types.FinishReason.MAX_TOKENS, None]:
|
||
if not image_found:
|
||
raise ValueError(f"Stream interrupted. Reason: {last_finish_reason}. Chunks: {chunk_count}")
|
||
elif not image_found:
|
||
raise ValueError(f"Stream interrupted. No image was generated. Chunks: {chunk_count}")
|
||
|
||
if image_found:
|
||
print(f" ✅ Success! Image received ({chunk_count} chunks handled).")
|
||
# Ensure we only return one image even if the model returned multiple candidates/parts
|
||
return out_tensor[:1], full_text
|
||
|
||
print(f" ⚠️ Warning: Model returned text description instead of image.")
|
||
if text_retry < TEXT_ONLY_RETRIES - 1:
|
||
for _ in range(5):
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
time.sleep(1)
|
||
|
||
raise ValueError(f"Model returned text only after {TEXT_ONLY_RETRIES} attempts.")
|
||
|
||
# Try Official ComfyUI API First
|
||
if use_comfyui_credits:
|
||
try:
|
||
# ComfyUI's genai integration uses a direct import or custom endpoint.
|
||
# Since GenAI library defaults to standard Google endpoints, we can override host
|
||
comfy_api_base, auth_headers, actual_token = get_comfy_org_auth(auth_token)
|
||
if not auth_headers:
|
||
print(f"⚠️ ComfyUI Credits toggle is ON, but no ComfyUI API Token found. Falling back to custom keys...")
|
||
else:
|
||
batch_info = f" | Batch {curr_batch}/{total_batch}..." if total_batch > 1 else "..."
|
||
print(f"\n🪙 Using ComfyUI Credits 🎨 Image generation{batch_info}")
|
||
|
||
# We create a client pointing to the proxy
|
||
# Since Google's GenAI SDK uses https://generativelanguage.googleapis.com
|
||
# We can pass http_options to redirect to the comfy proxy
|
||
# But the proxy expects /proxy/vertexai/gemini/... which is Vertex format, not Generative Language format.
|
||
# Actually, the official comfy APIs use a raw urllib call for this. Let's do a direct call to be safe from SDK mismatches.
|
||
|
||
# Ensure we have our ComfyUI API key available under "api_key_comfy_org" or similar
|
||
# It's much simpler to just import the exact comfy_api_nodes class and call it, but
|
||
# we must handle case where user doesn't have comfy_api_nodes installed.
|
||
# We will attempt dynamic import of the comfy_api_nodes package to handle everything including proper chunk parsing.
|
||
try:
|
||
from comfy_api_nodes.nodes_gemini import GeminiImage2
|
||
from comfy_api_nodes.apis.gemini import GeminiPart, GeminiInlineData, GeminiMimeType
|
||
import asyncio
|
||
|
||
f_parts = []
|
||
if images is not None:
|
||
img_list = images if isinstance(images, list) else [images]
|
||
for img_item in img_list:
|
||
img_to_proc = img_item if len(img_item.shape) == 4 else img_item.unsqueeze(0)
|
||
for j in range(img_to_proc.shape[0]):
|
||
img_tensor = img_to_proc[j]
|
||
np_img = 255. * img_tensor.cpu().numpy()
|
||
pil_img = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
|
||
buffered = io.BytesIO()
|
||
pil_img.save(buffered, format="PNG")
|
||
import base64
|
||
b64_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||
f_parts.append(GeminiPart(inlineData=GeminiInlineData(mimeType=GeminiMimeType.image_png, data=b64_str)))
|
||
|
||
# Prepare thread-safe async execution
|
||
|
||
# Note: GeminiImage2 needs auth to be stored magically in class state via the mixin if running alone,
|
||
# but if run from our node, it might fail unless we mock the class.
|
||
# Since ComfyUI's standard `sync_op` uses `node_cls.hidden.auth_token_comfy_org`.
|
||
# We must inject it into the class temporarily.
|
||
from comfy.cli_args import args as comfy_args
|
||
auth_val = actual_token
|
||
|
||
if not auth_val:
|
||
raise ValueError("No authorization token found.")
|
||
|
||
if not hasattr(GeminiImage2, "hidden") or GeminiImage2.hidden is None:
|
||
class DummyHidden: pass
|
||
GeminiImage2.hidden = DummyHidden()
|
||
|
||
GeminiImage2.hidden.auth_token_comfy_org = auth_val
|
||
GeminiImage2.hidden.api_key_comfy_org = auth_val
|
||
GeminiImage2.hidden.unique_id = "Fallback_Generated_Node"
|
||
|
||
def _run_sync_op():
|
||
# Optimized loop creation
|
||
new_loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(new_loop)
|
||
|
||
if hasattr(new_loop, 'set_exception_handler'):
|
||
def silence_connection_reset(loop, context):
|
||
if "exception" in context:
|
||
exc = context["exception"]
|
||
if isinstance(exc, (ConnectionResetError, ConnectionAbortedError)):
|
||
return
|
||
loop.default_exception_handler(context)
|
||
new_loop.set_exception_handler(silence_connection_reset)
|
||
|
||
try:
|
||
comfy_model = model
|
||
if "gemini-3.1-pro-image" in model or "gemini-3-pro" in model:
|
||
comfy_model = "gemini-3-pro-image-preview"
|
||
elif "Nano Banana 2" in model:
|
||
comfy_model = "gemini-3.1-flash-image-preview"
|
||
|
||
safe_seed = int(seed) % 2147483647
|
||
|
||
# Process images_tensor directly from pre_images_tensor if available
|
||
# (already handled in _exec_single parameters)
|
||
comfy_aspect_ratio = aspect_ratio
|
||
if comfy_aspect_ratio == "Aspect Ratio Input":
|
||
comfy_aspect_ratio = "auto"
|
||
|
||
return new_loop.run_until_complete(
|
||
asyncio.wait_for(
|
||
GeminiImage2.execute(
|
||
prompt=prompt,
|
||
model=comfy_model,
|
||
seed=safe_seed,
|
||
aspect_ratio=comfy_aspect_ratio,
|
||
resolution=resolution,
|
||
response_modalities=response_modalities,
|
||
images=images_tensor,
|
||
files=None,
|
||
system_prompt=system_instruction
|
||
),
|
||
timeout=300.0
|
||
)
|
||
)
|
||
finally:
|
||
try:
|
||
new_loop.run_until_complete(new_loop.shutdown_asyncgens())
|
||
except: pass
|
||
new_loop.close()
|
||
|
||
# Retry logic for ComfyUI Credits
|
||
CREDIT_RETRIES = 2
|
||
current_credit_attempt = 0
|
||
|
||
while current_credit_attempt < CREDIT_RETRIES:
|
||
current_credit_attempt += 1
|
||
try:
|
||
final_res = _run_sync_op()
|
||
if final_res:
|
||
print(f" ✅ Success! Prompt {curr_batch} via Credits.")
|
||
return final_res[0][:1], final_res[1]
|
||
break
|
||
except Exception as e:
|
||
err_str = str(e)
|
||
if ("10054" in err_str or "ConnectionResetError" in err_str) and current_credit_attempt < CREDIT_RETRIES:
|
||
time.sleep(1)
|
||
continue
|
||
raise e
|
||
|
||
except Exception as e:
|
||
if "Interrupted by user" in str(e): raise e
|
||
print(f" ❌ ComfyUI Credits failed. Reason: {e}")
|
||
print(f" 🔄 Falling back to custom API Keys System...")
|
||
except Exception as e:
|
||
if "Interrupted by user" in str(e): raise e
|
||
print(f" ❌ ComfyUI Credits failed. Reason: {e}")
|
||
print(f" 🔄 Falling back to custom API Keys System...")
|
||
|
||
retry_count = 0
|
||
wait_time = 15
|
||
while retry_count < api_max_retries:
|
||
for index, key in enumerate(keys):
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
try:
|
||
return call_api(key)
|
||
except Exception as e:
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
error_str = str(e)
|
||
orig_idx = all_keys.index(key) + 1 if key in all_keys else (index + 1)
|
||
print(f"\n❌ Key Failure | Key #{orig_idx} (****{key[-4:]}) failed.")
|
||
print(f" Reason: {clean_error_msg(error_str)}")
|
||
|
||
if "disconnected" in error_str.lower() or "remote protocol" in error_str.lower() or "10054" in error_str:
|
||
print(f" ℹ️ Notice: Transient network hiccup. Automatic retry in progress...")
|
||
|
||
if "400 INVALID_ARGUMENT" in error_str:
|
||
raise ValueError(f"Fatal error (not retryable): {error_str}")
|
||
if index < len(keys) - 1:
|
||
print(f"🔄 Rotating API keys to find an available worker...")
|
||
|
||
retry_count += 1
|
||
if retry_count < api_max_retries:
|
||
current_wait_sec, wait_msg = get_retry_wait_info(error_str if 'error_str' in locals() else "", retry_count, default_wait=wait_time)
|
||
|
||
print(f"⚠️ Round {retry_count}/{api_max_retries} Exhausted | All keys temporarily unavailable.")
|
||
print(wait_msg)
|
||
|
||
for _ in range(int(current_wait_sec) + 1):
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
time.sleep(1)
|
||
|
||
# Collect the last error
|
||
last_failure_reason = "No successful response from any key."
|
||
if 'error_str' in locals():
|
||
last_failure_reason = clean_error_msg(error_str)
|
||
|
||
# Instead of raising and stopping the queue, return a placeholder
|
||
return get_error_fallback(f"Gemini API Fail: {last_failure_reason} (Tried {api_max_retries} rounds)")
|
||
|
||
class GeminiPromptGenerator:
|
||
"""
|
||
ComfyUI Node for Google Gemini (Text/Prompt Generation).
|
||
Features:
|
||
- Supports multiple text models.
|
||
- 3-Key Iterative Retry System with configurable rounds.
|
||
- Enhanced English logging.
|
||
"""
|
||
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"system_instruction": ("STRING", {"multiline": True, "default": "You are a creative writer. Describe the image in detail.", "placeholder": "System instructions..."}),
|
||
"user_prompt": ("STRING", {"multiline": True, "default": "Describe this image", "placeholder": "User prompt..."}),
|
||
"batch_count": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}),
|
||
"model": (["gemini-3.5-flash", "gemini-3.1-pro-preview", "gemini-3.1-flash-preview", "gemini-2.0-pro-exp-02-05", "gemini-2.0-flash", "gemini-2.0-flash-lite-preview-02-05", "gemini-1.5-pro", "gemini-1.5-flash"], {"default": "gemini-3.5-flash"}),
|
||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||
"max_output_tokens": ("INT", {"default": 0, "min": 0, "max": 128000, "step": 64}),
|
||
"api_key_1": ("STRING", {"multiline": False, "default": "", "placeholder": "Primary API Key (Required)"}),
|
||
"api_key_2": ("STRING", {"multiline": False, "default": "", "placeholder": "Backup API Key 1 (Optional)"}),
|
||
"api_key_3": ("STRING", {"multiline": False, "default": "", "placeholder": "Backup API Key 2 (Optional)"}),
|
||
"api_max_retries": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1}),
|
||
"batch_delay": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 60.0, "step": 0.1}),
|
||
"use_comfyui_credits": ("BOOLEAN", {"default": True, "label_on": "Use ComfyUI Credits", "label_off": "Use Custom API Keys"}),
|
||
"bypass": ("BOOLEAN", {"default": False, "label_on": "Use Last Generation", "label_off": "Generate New"}),
|
||
},
|
||
"optional": {
|
||
"images": ("IMAGE",),
|
||
},
|
||
"hidden": {
|
||
"auth_token_comfy_org": "AUTH_TOKEN_COMFY_ORG",
|
||
"unique_id": "UNIQUE_ID",
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("text_output",)
|
||
OUTPUT_IS_LIST = (True,)
|
||
FUNCTION = "generate_text"
|
||
CATEGORY = "Gemini API"
|
||
|
||
def generate_text(self, system_instruction, user_prompt, model, seed, max_output_tokens, api_key_1, api_key_2, api_key_3, api_max_retries, batch_count, batch_delay, use_comfyui_credits, bypass=None, images=None, auth_token_comfy_org="", unique_id=None):
|
||
if bypass and unique_id and unique_id in _PROMPT_GEN_CACHE:
|
||
print(f"♻️ Gemini Prompt Gen | Bypassed. Returning last generation for node {unique_id}.")
|
||
return _PROMPT_GEN_CACHE[unique_id]
|
||
|
||
if not user_prompt and images is None:
|
||
raise ValueError("At least a prompt or an image is required.")
|
||
|
||
keys = [k.strip() for k in [api_key_1, api_key_2, api_key_3] if k and k.strip()]
|
||
if not keys:
|
||
raise ValueError("No API Key provided!")
|
||
|
||
input_parts = []
|
||
if images is not None:
|
||
# Standard ComfyUI input can be a single tensor or list
|
||
img_list = images if isinstance(images, list) else [images]
|
||
|
||
# Filter out None and handle list of tensors
|
||
valid_images = [img for img in img_list if img is not None]
|
||
|
||
if valid_images:
|
||
print(f" 📊 Prompt Debug | Processing {len(valid_images)} image input(s).")
|
||
# Ensure same size for internal consistency if needed,
|
||
# though GenAI Parts can technically handle mixed sizes.
|
||
# Still, we'll keep it consistent with the Image node.
|
||
synced_images = ensure_same_size(valid_images)
|
||
for i, img_to_proc in enumerate(synced_images):
|
||
for j in range(img_to_proc.shape[0]):
|
||
img_tensor = img_to_proc[j]
|
||
np_img = 255. * img_tensor.cpu().numpy()
|
||
pil_img = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
|
||
|
||
if j == 0:
|
||
print(f" 📸 Image Ref | Dim: {pil_img.width}x{pil_img.height}")
|
||
|
||
buffered = io.BytesIO()
|
||
pil_img.save(buffered, format="PNG")
|
||
input_parts.append(types.Part.from_bytes(data=buffered.getvalue(), mime_type="image/png"))
|
||
|
||
if user_prompt:
|
||
input_parts.append(types.Part.from_text(text=user_prompt))
|
||
|
||
results = []
|
||
|
||
# Build available keys list for load balancing
|
||
available_keys = [k.strip() for k in [api_key_1, api_key_2, api_key_3] if k and k.strip()]
|
||
|
||
for b in range(batch_count):
|
||
try:
|
||
if b > 0 and batch_delay > 0:
|
||
print(f" ⏳ Batch delay: Waiting {batch_delay}s")
|
||
time.sleep(batch_delay)
|
||
|
||
current_seed = (seed + b) % 2147483647
|
||
# Load balancing offset for API keys
|
||
start_key_idx = b % len(available_keys)
|
||
keys = available_keys[start_key_idx:] + available_keys[:start_key_idx]
|
||
|
||
def call_api(current_key):
|
||
nonlocal current_seed
|
||
# --- VGG+GAIA MODIFICATION: Support proxy URL or standard API key ---
|
||
client_or_config = _create_genai_client(current_key)
|
||
|
||
# model mapping for Google GenAI SDK
|
||
actual_model = model
|
||
if "gemini-3.5" in model:
|
||
actual_model = model
|
||
elif "gemini-3.1-pro" in model or "gemini-3-pro" in model:
|
||
actual_model = "gemini-3.1-pro-preview"
|
||
elif "gemini-3" in model and "flash" in model:
|
||
actual_model = "gemini-2.0-flash"
|
||
|
||
# ──── GAIA PROXY PATH ────
|
||
if isinstance(client_or_config, dict) and client_or_config.get("mode") == "gaia":
|
||
# Prepare images tensor for GAIA (already handled via input_parts but _gaia needs tensor)
|
||
img_tensor_for_gaia = None
|
||
if images is not None:
|
||
img_list = images if isinstance(images, list) else [images]
|
||
valid = [img for img in img_list if img is not None]
|
||
if valid:
|
||
synced = ensure_same_size(valid)
|
||
img_tensor_for_gaia = torch.cat(synced, dim=0)
|
||
return _gaia_call_text(
|
||
client_or_config, actual_model, user_prompt, system_instruction,
|
||
img_tensor_for_gaia, current_seed, max_output_tokens, b, batch_count
|
||
)
|
||
|
||
# ──── STANDARD GOOGLE SDK PATH ────
|
||
client = client_or_config
|
||
|
||
config_kwargs = {
|
||
"system_instruction": system_instruction,
|
||
"seed": current_seed,
|
||
"response_modalities": ["TEXT"],
|
||
"automatic_function_calling": types.AutomaticFunctionCallingConfig(disable=True),
|
||
"safety_settings": [
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=types.HarmBlockThreshold.BLOCK_ONLY_HIGH),
|
||
],
|
||
}
|
||
if max_output_tokens > 0:
|
||
config_kwargs["max_output_tokens"] = max_output_tokens
|
||
|
||
batch_str = f" | Batch {b+1}/{batch_count}" if batch_count > 1 else ""
|
||
print(f"\n📝 Gemini Prompt{batch_info} | Processing..." if 'batch_info' in locals() else f"\n📝 Gemini Prompt{batch_str} | Processing...")
|
||
print(f" Model: {model} | Seed: {current_seed}")
|
||
key_display = "proxy-url" if current_key.strip().startswith("http") else f"****{current_key[-4:]}"
|
||
key_idx = available_keys.index(current_key) + 1 if current_key in available_keys else 1
|
||
print(f" Key: Using Key #{key_idx} ({key_display})")
|
||
|
||
response_stream = client.models.generate_content_stream(
|
||
model=actual_model,
|
||
contents=[types.Content(role="user", parts=input_parts)],
|
||
config=types.GenerateContentConfig(**config_kwargs),
|
||
)
|
||
|
||
full_text = ""
|
||
chunk_count = 0
|
||
last_finish_reason = None
|
||
for chunk in response_stream:
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
if not chunk.candidates:
|
||
continue
|
||
|
||
last_finish_reason = chunk.candidates[0].finish_reason
|
||
if not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
|
||
continue
|
||
|
||
for part in chunk.candidates[0].content.parts:
|
||
if part.text:
|
||
full_text += part.text
|
||
chunk_count += 1
|
||
|
||
if last_finish_reason not in [types.FinishReason.STOP, types.FinishReason.MAX_TOKENS, None]:
|
||
if not full_text:
|
||
raise ValueError(f"Stream interrupted. Reason: {last_finish_reason}")
|
||
elif not full_text:
|
||
raise ValueError(f"Stream interrupted. No text was generated.")
|
||
|
||
if not full_text:
|
||
raise ValueError("Model returned empty text.")
|
||
|
||
print(f" ✅ Complete | Received {len(full_text)} characters in {chunk_count} chunks.")
|
||
return full_text
|
||
|
||
# Try Official ComfyUI API First
|
||
u_credits = use_comfyui_credits[0] if isinstance(use_comfyui_credits, list) else use_comfyui_credits
|
||
if u_credits:
|
||
try:
|
||
comfy_api_base, auth_headers, actual_token = get_comfy_org_auth(auth_token_comfy_org)
|
||
if not auth_headers:
|
||
print(f"⚠️ ComfyUI Credits toggle is ON, but no ComfyUI API Token found. Falling back to custom keys...")
|
||
else:
|
||
batch_str = f" | Batch {b+1}/{batch_count}..." if batch_count > 1 else "..."
|
||
print(f"\n🪙 Using ComfyUI Credits 📝 Prompt generation{batch_str}")
|
||
import asyncio
|
||
|
||
try:
|
||
from comfy_api_nodes.nodes_gemini import GeminiNode
|
||
|
||
from comfy.cli_args import args as comfy_args
|
||
auth_val = actual_token
|
||
|
||
if not auth_val:
|
||
raise ValueError("No authorization token found.")
|
||
|
||
if not hasattr(GeminiNode, "hidden") or GeminiNode.hidden is None:
|
||
class DummyHidden: pass
|
||
GeminiNode.hidden = DummyHidden()
|
||
|
||
GeminiNode.hidden.auth_token_comfy_org = auth_val
|
||
GeminiNode.hidden.api_key_comfy_org = auth_val
|
||
GeminiNode.hidden.unique_id = "Fallback_Generated_Node"
|
||
|
||
import threading
|
||
thread_res = []
|
||
thread_err = []
|
||
def _run_async_prompt():
|
||
import asyncio
|
||
import platform
|
||
import logging
|
||
|
||
# SILENCE ALL ASYNCIO/PROACTOR NOISE GLOBALLY IN THIS THREAD
|
||
logging.getLogger("asyncio").setLevel(logging.CRITICAL)
|
||
|
||
if platform.system() == 'Windows':
|
||
try:
|
||
from asyncio import WindowsProactorEventLoopPolicy
|
||
asyncio.set_event_loop_policy(WindowsProactorEventLoopPolicy())
|
||
except ImportError:
|
||
pass
|
||
|
||
new_loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(new_loop)
|
||
|
||
# Patch the loop to ignore the connection reset error log on Windows
|
||
if hasattr(new_loop, 'set_exception_handler'):
|
||
def silence_connection_reset(loop, context):
|
||
if "exception" in context:
|
||
exc = context["exception"]
|
||
if isinstance(exc, (ConnectionResetError, ConnectionAbortedError)):
|
||
return # Silence it
|
||
loop.default_exception_handler(context)
|
||
new_loop.set_exception_handler(silence_connection_reset)
|
||
|
||
try:
|
||
# model mapping for ComfyUI Credits internal names
|
||
comfy_model = model
|
||
if "gemini-3.5" in model:
|
||
comfy_model = model
|
||
elif "gemini-3.1-pro" in model or "gemini-3-pro" in model:
|
||
comfy_model = "gemini-3.1-pro-preview"
|
||
elif "gemini-3" in model and "flash" in model:
|
||
comfy_model = "gemini-3.1-flash-lite-preview"
|
||
|
||
result = new_loop.run_until_complete(
|
||
asyncio.wait_for(
|
||
GeminiNode.execute(
|
||
prompt=user_prompt,
|
||
model=comfy_model,
|
||
seed=current_seed,
|
||
images=images,
|
||
system_prompt=system_instruction
|
||
),
|
||
timeout=120.0
|
||
)
|
||
)
|
||
thread_res.append(result)
|
||
except Exception as e:
|
||
thread_err.append(e)
|
||
finally:
|
||
# Graceful shutdown
|
||
try:
|
||
new_loop.run_until_complete(new_loop.shutdown_asyncgens())
|
||
if hasattr(new_loop, "shutdown_default_executor"):
|
||
new_loop.run_until_complete(new_loop.shutdown_default_executor())
|
||
except:
|
||
pass
|
||
new_loop.close()
|
||
|
||
# Retry logic for ComfyUI Credits specifically for transient errors
|
||
CREDIT_RETRIES = 2
|
||
current_credit_attempt = 0
|
||
item_success = False
|
||
|
||
while current_credit_attempt < CREDIT_RETRIES:
|
||
current_credit_attempt += 1
|
||
thread_res = []
|
||
thread_err = []
|
||
|
||
t = threading.Thread(target=_run_async_prompt)
|
||
t.start()
|
||
t.join()
|
||
|
||
if thread_err:
|
||
err_str = str(thread_err[0])
|
||
if "10054" in err_str or "ConnectionResetError" in err_str or "forcibly closed" in err_str:
|
||
if current_credit_attempt < CREDIT_RETRIES:
|
||
print(f" ⏳ Connection Reset (10054). Retrying Credits ({current_credit_attempt}/{CREDIT_RETRIES})...")
|
||
time.sleep(2)
|
||
continue
|
||
raise thread_err[0]
|
||
|
||
res = thread_res[0]
|
||
print(" ✅ Success! Prompt generated via ComfyUI Credits.")
|
||
results.append(res[0])
|
||
item_success = True
|
||
break
|
||
|
||
if item_success:
|
||
continue # Move to next batch item
|
||
|
||
except Exception as e:
|
||
if "Interrupted by user" in str(e): raise e
|
||
print(f" ❌ ComfyUI Credits failed. Reason: {e}")
|
||
print(f" 🔄 Falling back to custom API Keys System...")
|
||
except Exception as e:
|
||
if "Interrupted by user" in str(e): raise e
|
||
print(f" ❌ ComfyUI Credits failed. Reason: {e}")
|
||
print(f" 🔄 Falling back to custom API Keys System...")
|
||
|
||
|
||
# Strategy: execute this batch item with full key rotation/retry
|
||
retry_count = 0
|
||
wait_time = 5
|
||
success = False
|
||
item_text = ""
|
||
|
||
while retry_count < api_max_retries:
|
||
for index, key in enumerate(keys):
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
try:
|
||
item_text = call_api(key)
|
||
success = True
|
||
break
|
||
except Exception as e:
|
||
if "Interrupted by user" in str(e): raise e
|
||
error_str = str(e)
|
||
orig_idx = available_keys.index(key) + 1 if key in available_keys else (index + 1)
|
||
print(f"\n❌ Key Failure | Key #{orig_idx} (****{key[-4:]}) failed.")
|
||
print(f" Reason: {clean_error_msg(error_str)}")
|
||
if "400 INVALID_ARGUMENT" in error_str:
|
||
raise ValueError(f"Fatal: {error_str}")
|
||
|
||
if success:
|
||
break
|
||
|
||
retry_count += 1
|
||
if retry_count < api_max_retries:
|
||
current_wait_sec, wait_msg = get_retry_wait_info(error_str if 'error_str' in locals() else "", retry_count, default_wait=wait_time)
|
||
print(f"⚠️ Round {retry_count}/{api_max_retries} Exhausted | {wait_msg}")
|
||
for _ in range(int(current_wait_sec) + 1):
|
||
if comfy.model_management.processing_interrupted():
|
||
raise Exception("Interrupted by user")
|
||
time.sleep(1)
|
||
|
||
if success:
|
||
results.append(item_text)
|
||
else:
|
||
results.append(f"ERROR: Batch {b+1} failed after all retries.")
|
||
except Exception as e:
|
||
is_interrupted = "Interrupted by user" in str(e)
|
||
if is_interrupted:
|
||
print(f" 🛑 Batch Interrupted by User at item {b+1}/{batch_count}")
|
||
else:
|
||
print(f" 🛑 Prompt Error (Item {b+1}/{batch_count}): {e}")
|
||
|
||
# Fallback to prevent queue interruption and maintain batch size
|
||
print(f" ⚠️ Padding remaining {batch_count - len(results)} item(s) to preserve generated batch...")
|
||
while len(results) < batch_count:
|
||
err_msg = "Generation cancelled by user" if is_interrupted else f"ERROR: {e}"
|
||
results.append(err_msg)
|
||
|
||
break # Stop processing further batch items
|
||
|
||
if unique_id:
|
||
_PROMPT_GEN_CACHE[unique_id] = (results,)
|
||
return (results,)
|