386 lines
18 KiB
Python
386 lines
18 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io
|
|
import time
|
|
from google import genai
|
|
from google.genai import types
|
|
import comfy.model_management
|
|
import re
|
|
|
|
def parse_retry_delay(error_str):
|
|
"""Extracts suggested wait time from Gemini error message."""
|
|
# Pattern: "Please retry in 21.54s" or similar
|
|
match = re.search(r"Please retry in (\d+\.?\d*)s", error_str)
|
|
if match:
|
|
try:
|
|
return float(match.group(1))
|
|
except:
|
|
pass
|
|
return None
|
|
|
|
# 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."
|
|
)
|
|
|
|
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.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"system_instruction": ("STRING", {"multiline": True, "default": "You are a professional image generator. Maintain high cinematic quality.", "placeholder": "System instructions..."}),
|
|
"prompt": ("STRING", {"multiline": True, "default": "Make this image cyberpunk style", "placeholder": "Enter your prompt here..."}),
|
|
"model": (["gemini-3-pro-image-preview", "gemini-2.0-flash", "gemini-2.0-pro-exp-02-05"], {"default": "gemini-3-pro-image-preview"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
|
|
"aspect_ratio": (["auto", "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"], {"default": "auto"}),
|
|
"resolution": (["1K", "2K", "4K"], {"default": "1K"}),
|
|
"response_modalities": (["IMAGE+TEXT", "IMAGE"], {"default": "IMAGE+TEXT"}),
|
|
|
|
"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}),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("processed_image", "text_response")
|
|
FUNCTION = "process_image"
|
|
CATEGORY = "Gemini AI"
|
|
|
|
def process_image(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):
|
|
if not prompt or not prompt.strip():
|
|
raise ValueError("Prompt 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!")
|
|
|
|
image_parts = []
|
|
if images is not None:
|
|
for i in range(images.shape[0]):
|
|
img_tensor = images[i]
|
|
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")
|
|
image_parts.append(types.Part.from_bytes(data=buffered.getvalue(), mime_type="image/png"))
|
|
|
|
def call_api(current_key):
|
|
# Force REST by disabling HTTP/2 in httpx and setting a long timeout
|
|
client = genai.Client(
|
|
api_key=current_key,
|
|
http_options={
|
|
"timeout": 300000, # 300 seconds (5 minutes)
|
|
"client_args": {"http2": False}
|
|
}
|
|
)
|
|
|
|
# Clean model name mapping
|
|
actual_model = model
|
|
|
|
# Use the official image generation system prompt, appended to user's system instruction
|
|
effective_system = GEMINI_IMAGE_SYS_PROMPT
|
|
if system_instruction and system_instruction.strip():
|
|
effective_system = system_instruction.strip() + "\n\n" + GEMINI_IMAGE_SYS_PROMPT
|
|
|
|
# Match official node: IMAGE -> ["IMAGE"], IMAGE+TEXT -> ["IMAGE", "TEXT"]
|
|
if response_modalities == "IMAGE":
|
|
modalities_list = ["IMAGE"]
|
|
else:
|
|
modalities_list = ["IMAGE", "TEXT"]
|
|
|
|
# Models that do NOT support aspect_ratio (mostly multimodal models generating images vs native image models)
|
|
no_aspect_ratio_models = {"gemini-2.0-flash", "gemini-2.0-pro-exp-02-05"}
|
|
|
|
# Build config (matching working backup: dict-based image_config, no seed)
|
|
config_kwargs = {
|
|
"system_instruction": effective_system,
|
|
"response_modalities": modalities_list,
|
|
"image_config": {"image_size": resolution},
|
|
# Disable AFC (Automatic Function Calling) to reduce overhead during congestion
|
|
"automatic_function_calling": types.AutomaticFunctionCallingConfig(disable=True),
|
|
}
|
|
if aspect_ratio != "auto" and actual_model not in no_aspect_ratio_models:
|
|
config_kwargs["image_config"]["aspect_ratio"] = aspect_ratio
|
|
|
|
# Retry text-only responses per key (model behavior, not key issue)
|
|
TEXT_ONLY_RETRIES = 3
|
|
for text_retry in range(TEXT_ONLY_RETRIES):
|
|
print(f"\n🔵 [Gemini Image Gen] Sending request...")
|
|
print(f" Model: {actual_model} (mapped from {model})")
|
|
print(f" Key: ****{current_key[-4:]}")
|
|
print(f" Config: modalities={modalities_list}, image_config={config_kwargs.get('image_config', 'none')}")
|
|
if text_retry > 0:
|
|
print(f" ↻ Text-only retry {text_retry}/{TEXT_ONLY_RETRIES - 1} (same key)")
|
|
|
|
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:
|
|
# Some chunks might carry finish_reason without 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
|
|
|
|
# Verify if the response was complete
|
|
# STOP (1) or MAX_TOKENS (8) are considered "complete" for our purposes
|
|
if last_finish_reason not in [types.FinishReason.STOP, types.FinishReason.MAX_TOKENS]:
|
|
raise ValueError(f"Stream interrupted or blocked. Reason: {last_finish_reason}. Total chunks: {chunk_count}")
|
|
|
|
if image_found:
|
|
print(f" ✅ Image received ({chunk_count} chunks).")
|
|
return out_tensor, full_text
|
|
|
|
# Text-only response — retry with same key after short delay
|
|
print(f" ⚠️ Model returned text only (attempt {text_retry + 1}/{TEXT_ONLY_RETRIES})")
|
|
if text_retry < TEXT_ONLY_RETRIES - 1:
|
|
# Short sleep before retrying (model non-determinism, not rate limit)
|
|
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: {full_text[:200]}...")
|
|
|
|
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)
|
|
print(f"\n❌ [Gemini Image Gen] Key #{index + 1} (****{key[-4:]}) failed.")
|
|
print(f" Error: {error_str}")
|
|
# Bail immediately on 400 errors (deterministic, will never succeed on retry)
|
|
if "400 INVALID_ARGUMENT" in error_str:
|
|
raise ValueError(f"Gemini Image Gen fatal error (not retryable): {error_str}")
|
|
if index < len(keys) - 1:
|
|
print(f"🔄 [Gemini Image Gen] Rotating to next backup key...")
|
|
|
|
retry_count += 1
|
|
if retry_count < api_max_retries:
|
|
# Use parsed delay if available, otherwise use fixed fallback
|
|
dynamic_delay = parse_retry_delay(error_str) if 'error_str' in locals() else None
|
|
current_wait = int(dynamic_delay) + 1 if dynamic_delay else wait_time
|
|
|
|
print(f"\n⚠️ [Gemini Image Gen] All keys failed in Round {retry_count}/{api_max_retries}.")
|
|
if dynamic_delay:
|
|
print(f"⏳ [Gemini Image Gen] Google requested wait: {dynamic_delay}s. Sleeping {current_wait}s...\n")
|
|
else:
|
|
print(f"⏳ [Gemini Image Gen] Waiting {wait_time}s to avoid rate limits...\n")
|
|
|
|
# Sleep in small increments to allow interrupt
|
|
for _ in range(current_wait):
|
|
if comfy.model_management.processing_interrupted():
|
|
raise Exception("Interrupted by user")
|
|
time.sleep(1)
|
|
|
|
raise ValueError(f"Gemini Image Gen failed after {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-pro-preview", "gemini-3-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-pro-preview"}),
|
|
"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}),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
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, images=None):
|
|
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:
|
|
for i in range(images.shape[0]):
|
|
img_tensor = images[i]
|
|
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")
|
|
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 = []
|
|
|
|
for b in range(batch_count):
|
|
current_seed = (seed + b) % 2147483647
|
|
|
|
def call_api(current_key):
|
|
nonlocal current_seed
|
|
client = genai.Client(
|
|
api_key=current_key,
|
|
http_options={
|
|
"timeout": 300000,
|
|
"client_args": {"http2": False}
|
|
}
|
|
)
|
|
config_kwargs = {
|
|
"system_instruction": system_instruction,
|
|
"seed": current_seed,
|
|
"response_modalities": ["TEXT"],
|
|
"automatic_function_calling": types.AutomaticFunctionCallingConfig(disable=True),
|
|
}
|
|
if max_output_tokens > 0:
|
|
config_kwargs["max_output_tokens"] = max_output_tokens
|
|
|
|
print(f"\n🔵 [Gemini Prompt Gen] Batch {b+1}/{batch_count} - Seed: {current_seed}")
|
|
|
|
response_stream = client.models.generate_content_stream(
|
|
model=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]:
|
|
raise ValueError(f"Stream interrupted. Reason: {last_finish_reason}")
|
|
|
|
if not full_text:
|
|
raise ValueError("Model returned empty text.")
|
|
|
|
print(f" ✅ Done ({chunk_count} chunks, {len(full_text)} chars)")
|
|
return full_text
|
|
|
|
# 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:
|
|
error_str = str(e)
|
|
print(f"\n❌ Item {b+1} failed with Key #{index + 1}")
|
|
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:
|
|
dynamic_delay = parse_retry_delay(error_str)
|
|
current_wait = int(dynamic_delay) + 1 if dynamic_delay else wait_time
|
|
time.sleep(current_wait)
|
|
|
|
if success:
|
|
results.append(item_text)
|
|
else:
|
|
results.append(f"ERROR: Batch {b+1} failed after all retries.")
|
|
|
|
return (results,)
|