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,)