1206 lines
55 KiB
Python
1206 lines
55 KiB
Python
import sys
|
|
import os
|
|
import warnings
|
|
|
|
# Suppress warnings related to IMAGE_SAFETY finish reason which is normal
|
|
warnings.filterwarnings("ignore", message="IMAGE_SAFETY is not a valid FinishReason")
|
|
|
|
# Add site-packages directory to Python's sys.path
|
|
'''
|
|
site_packages_path = os.path.join(sys.prefix, 'Lib', 'site-packages')
|
|
if site_packages_path not in sys.path:
|
|
sys.path.insert(0, site_packages_path)
|
|
'''
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import io
|
|
import logging
|
|
import random
|
|
import base64
|
|
import time
|
|
import uuid
|
|
import hashlib
|
|
import json
|
|
import re
|
|
|
|
# Import PromptServer for real-time UI updates
|
|
try:
|
|
from server import PromptServer
|
|
HAS_PROMPTSERVER = True
|
|
except ImportError:
|
|
logger.warning("PromptServer could not be imported. Real-time UI updates will be disabled.")
|
|
HAS_PROMPTSERVER = False
|
|
PromptServer = None
|
|
|
|
from .env_utils import get_api_key
|
|
from .utils import ChatHistory
|
|
from .image_utils import (
|
|
create_placeholder_image,
|
|
prepare_batch_images,
|
|
process_images_for_comfy,
|
|
tensor_to_pil,
|
|
resize_image_to_dimensions,
|
|
pil_to_tensor,
|
|
)
|
|
from .response_utils import prepare_response
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Helper function for safely sending WebSocket messages
|
|
def send_ws_message_safe(event_name, data, node_id=None):
|
|
"""Safely send a WebSocket message with error handling."""
|
|
if not HAS_PROMPTSERVER or PromptServer is None:
|
|
return False
|
|
|
|
try:
|
|
# Ensure node_id is in the data if provided
|
|
if node_id and isinstance(data, dict) and "node_id" not in data:
|
|
data["node_id"] = node_id
|
|
|
|
PromptServer.instance.send_sync(event_name, data)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error sending WebSocket message '{event_name}': {e}")
|
|
return False
|
|
|
|
def generate_consistent_seed(input_seed=0, use_random=False):
|
|
"""
|
|
Generate a consistent seed for the Gemini API.
|
|
|
|
This function uses a more reliable approach to seed generation:
|
|
- If input_seed is non-zero and use_random is False, use the input_seed
|
|
- Otherwise, generate a high-quality random seed based on uuid and time
|
|
|
|
Returns:
|
|
int: A seed value within the INT32 range (0 to 2^31-1)
|
|
"""
|
|
max_int32 = 2**31 - 1
|
|
|
|
if input_seed != 0 and not use_random:
|
|
# Use the provided seed, but ensure it's within INT32 range
|
|
adjusted_seed = input_seed % max_int32
|
|
logger.info(f"Using provided seed (adjusted to INT32 range): {adjusted_seed}")
|
|
return adjusted_seed
|
|
|
|
# For random seeds, use a more robust method that won't collide with ComfyUI's seed generation
|
|
# Create a unique identifier by combining:
|
|
# 1. A UUID (universally unique)
|
|
# 2. Current high-precision time
|
|
# 3. ComfyUI's random seed (if we wanted to use it)
|
|
|
|
unique_id = str(uuid.uuid4())
|
|
current_time = str(time.time_ns()) # Nanosecond precision
|
|
random_component = str(random.randint(0, max_int32))
|
|
|
|
# Combine and hash all components to get a deterministic but high-quality random value
|
|
combined = unique_id + current_time + random_component
|
|
hash_hex = hashlib.md5(combined.encode()).hexdigest()
|
|
|
|
# Convert first 8 characters of hash to integer and ensure within INT32 range
|
|
hash_int = int(hash_hex[:8], 16) % max_int32
|
|
|
|
logger.info(f"Generated random seed: {hash_int}")
|
|
return hash_int
|
|
|
|
|
|
class GeminiNode:
|
|
def __init__(self):
|
|
self.api_key = ""
|
|
self.chat_history = ChatHistory()
|
|
self.last_external_api_key = "" # Track the last external API key
|
|
self.api_key_source = None # Track where the API key came from
|
|
|
|
# First check system environment variables
|
|
system_api_key = os.environ.get("GEMINI_API_KEY", "")
|
|
if system_api_key:
|
|
self.api_key = system_api_key
|
|
self.api_key_source = "system environment variables"
|
|
logger.info("Successfully loaded Gemini API key from system environment")
|
|
else:
|
|
# Next, try to directly check shell configuration files
|
|
home_dir = os.path.expanduser("~")
|
|
shell_config_files = [
|
|
os.path.join(home_dir, ".zshrc"),
|
|
os.path.join(home_dir, ".bashrc"),
|
|
os.path.join(home_dir, ".bash_profile")
|
|
]
|
|
|
|
import re
|
|
shell_key = None
|
|
shell_source = None
|
|
|
|
for config_file in shell_config_files:
|
|
if os.path.exists(config_file):
|
|
logger.debug(f"Checking {config_file} for API key...")
|
|
try:
|
|
with open(config_file, 'r') as f:
|
|
content = f.read()
|
|
# Look for export VAR=value or VAR=value patterns
|
|
patterns = [
|
|
r'export\s+GEMINI_API_KEY=[\'\"]?([^\s\'\"]+)[\'\"]?',
|
|
r'GEMINI_API_KEY=[\'\"]?([^\s\'\"]+)[\'\"]?'
|
|
]
|
|
for pattern in patterns:
|
|
matches = re.findall(pattern, content)
|
|
if matches:
|
|
shell_key = matches[0]
|
|
shell_source = os.path.basename(config_file)
|
|
logger.info(f"Found Gemini API key in {shell_source}")
|
|
# Also set in environment for future use
|
|
os.environ["GEMINI_API_KEY"] = shell_key
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error reading {config_file}: {str(e)}")
|
|
if shell_key:
|
|
break
|
|
|
|
if shell_key:
|
|
self.api_key = shell_key
|
|
self.api_key_source = shell_source
|
|
logger.info(f"Successfully loaded Gemini API key from {shell_source}")
|
|
else:
|
|
# Last resort: check .env files
|
|
env_api_key = get_api_key("GEMINI_API_KEY", "Gemini")
|
|
if env_api_key:
|
|
self.api_key = env_api_key
|
|
self.api_key_source = ".env file"
|
|
logger.info("Successfully loaded Gemini API key from .env file")
|
|
else:
|
|
logger.warning("No Gemini API key found in any location (system env, shell configs, .env). You'll need to provide it in the node.")
|
|
|
|
# Log key information (masked for security)
|
|
if self.api_key:
|
|
masked_key = self.api_key[:4] + "..." + self.api_key[-4:] if len(self.api_key) > 8 else "****"
|
|
logger.info(f"Using Gemini API key ({masked_key}) from {self.api_key_source}")
|
|
|
|
# Check for Google Generative AI SDK
|
|
self.genai_available = self._check_genai_availability()
|
|
|
|
def _check_genai_availability(self):
|
|
"""Check if Google Generative AI SDK is available"""
|
|
try:
|
|
# Import just to check availability
|
|
from google import genai
|
|
|
|
return True
|
|
except ImportError:
|
|
logger.error(
|
|
"Google Generative AI SDK not installed. Install with: pip install google-generativeai"
|
|
)
|
|
return False
|
|
|
|
def _get_active_api_key(self, external_api_key=""):
|
|
"""
|
|
Determine the active API key to use based on priority order:
|
|
1. External API key passed to the method (if valid)
|
|
2. System environment variable
|
|
3. Previously stored API key in the instance
|
|
|
|
Args:
|
|
external_api_key (str): External API key provided in the node UI
|
|
|
|
Returns:
|
|
tuple: (api_key, api_key_source) where:
|
|
api_key (str): The API key to use
|
|
api_key_source (str): Description of where the key came from
|
|
"""
|
|
# Clean and validate external API key
|
|
cleaned_external_key = external_api_key.strip() if external_api_key else ""
|
|
|
|
# Determine API key to use based on priority
|
|
if cleaned_external_key:
|
|
logger.info("Using API key provided in the node")
|
|
# Save it for future reference
|
|
self.last_external_api_key = cleaned_external_key
|
|
return cleaned_external_key, "external"
|
|
|
|
# Next check system environment variables
|
|
system_key = os.environ.get("GEMINI_API_KEY", "")
|
|
if system_key:
|
|
logger.info("Using API key from system environment variable")
|
|
return system_key, "system"
|
|
|
|
# Next try the API key already in the instance
|
|
if self.api_key:
|
|
logger.info(f"Using API key from {self.api_key_source}")
|
|
return self.api_key, self.api_key_source
|
|
|
|
# Last resort: try previously provided external key
|
|
if self.last_external_api_key:
|
|
logger.info("Using previously provided external API key")
|
|
return self.last_external_api_key, "cached"
|
|
|
|
# No API key found
|
|
logger.warning("No API key found")
|
|
return "", ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": "Create a vivid word-picture representation of this image include elements that characterize the subject, costume, prop elemts, the action, the background, layout and composition elements present on the scene, be sure to mention the style and mood of the scene. Like it would a film director or director of photography"}),
|
|
"operation_mode": (
|
|
["analysis", "generate_text", "generate_images", "generate_sequence"],
|
|
{"default": "generate_images"},
|
|
),
|
|
"model_name": (
|
|
[
|
|
"gemini-2.0-flash-exp",
|
|
"gemini-2.0-pro",
|
|
"gemini-2.0-flash",
|
|
"gemini-2.0-flash-exp-image-generation",
|
|
"gemini-1.5-pro-latest",
|
|
"gemini-1.5-pro",
|
|
"gemini-1.5-pro-001",
|
|
"gemini-1.5-pro-vision-latest",
|
|
],
|
|
{"default": "gemini-2.0-flash-exp"},
|
|
),
|
|
"temperature": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
"video": ("IMAGE",),
|
|
"audio": ("AUDIO",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
|
|
"batch_count": ("INT", {"default": 4, "min": 1, "max": 8}),
|
|
"aspect_ratio": (
|
|
["none", "1:1", "16:9", "9:16", "4:3", "3:4", "5:4", "4:5"],
|
|
{"default": "none"},
|
|
),
|
|
"external_api_key": ("STRING", {"default": ""}),
|
|
"chat_mode": ("BOOLEAN", {"default": False}),
|
|
"clear_history": ("BOOLEAN", {"default": False}),
|
|
"structured_output": ("BOOLEAN", {"default": False}),
|
|
"max_images": ("INT", {"default": 6, "min": 1, "max": 16}),
|
|
"max_output_tokens": ("INT", {"default": 8192, "min": 1, "max": 32768}),
|
|
"use_random_seed": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"unique_id": "UNIQUE_ID"
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE")
|
|
RETURN_NAMES = ("text", "image")
|
|
FUNCTION = "generate_content"
|
|
CATEGORY = "ImpactFrames💥🎞️/LLM"
|
|
|
|
def generate_content(
|
|
self,
|
|
prompt,
|
|
operation_mode="analysis",
|
|
chat_mode=False,
|
|
clear_history=False,
|
|
images=None,
|
|
video=None,
|
|
audio=None,
|
|
external_api_key="",
|
|
max_images=6,
|
|
batch_count=1,
|
|
seed=0,
|
|
max_output_tokens=8192,
|
|
temperature=0.4,
|
|
structured_output=False,
|
|
aspect_ratio="none",
|
|
use_random_seed=False,
|
|
model_name="gemini-2.0-flash-exp",
|
|
unique_id=None,
|
|
extra_pnginfo=None,
|
|
prompt_=None,
|
|
):
|
|
"""Generate content using Gemini model with various input types."""
|
|
|
|
# Check if Google Generative AI SDK is available
|
|
if not self.genai_available:
|
|
return (
|
|
"ERROR: Google Generative AI SDK not installed. Install with: pip install google-generativeai",
|
|
create_placeholder_image(),
|
|
)
|
|
|
|
# Import here to avoid ImportError during ComfyUI startup
|
|
try:
|
|
from google import genai
|
|
from google.genai import types
|
|
logger.info(f"Google Generative AI SDK path: {genai.__file__}")
|
|
except ImportError:
|
|
return ("ERROR: Failed to import Google Generative AI SDK", create_placeholder_image())
|
|
|
|
# Get API key using the helper method
|
|
api_key, api_key_source = self._get_active_api_key(external_api_key)
|
|
|
|
if not api_key:
|
|
return (
|
|
"ERROR: No API key provided. Please set GEMINI_API_KEY in your environment or"
|
|
" provide it in the external_api_key field.",
|
|
create_placeholder_image(),
|
|
)
|
|
|
|
# --- Robust Prompt Extraction ---
|
|
prompt_text = ""
|
|
actual_prompt_input = prompt # Use the main 'prompt' input first
|
|
|
|
# Check if the main prompt input is a dict with nested structure
|
|
if isinstance(actual_prompt_input, dict):
|
|
# If prompt is a dict, try to extract text from common patterns
|
|
if "text" in actual_prompt_input:
|
|
prompt_text = str(actual_prompt_input["text"])
|
|
elif "prompt" in actual_prompt_input:
|
|
prompt_text = str(actual_prompt_input["prompt"])
|
|
elif "inputs" in actual_prompt_input and isinstance(actual_prompt_input["inputs"], dict):
|
|
# Try to extract from a nested inputs structure
|
|
inputs = actual_prompt_input["inputs"]
|
|
if "prompt" in inputs:
|
|
prompt_text = str(inputs["prompt"])
|
|
elif "text" in inputs:
|
|
prompt_text = str(inputs["text"])
|
|
else:
|
|
# If no standard keys found, convert the whole dict to string
|
|
prompt_text = str(actual_prompt_input)
|
|
else:
|
|
# If no known structure, convert the whole dict to string
|
|
prompt_text = str(actual_prompt_input)
|
|
logger.info(f"Extracted prompt from dictionary: {prompt_text[:50]}...")
|
|
elif isinstance(actual_prompt_input, (list, tuple)):
|
|
# If prompt is a list or tuple, join elements
|
|
prompt_text = " ".join(str(item) for item in actual_prompt_input)
|
|
logger.info(f"Converted list/tuple prompt to string: {prompt_text[:50]}...")
|
|
else:
|
|
# Otherwise, ensure it's a string
|
|
prompt_text = str(actual_prompt_input)
|
|
|
|
# Check if prompt_text is empty or just whitespace, but we have prompt_ (workflow data)
|
|
if not prompt_text.strip() and prompt_ and isinstance(prompt_, dict):
|
|
try:
|
|
logger.debug("Extracted prompt is empty, trying workflow data (prompt_)")
|
|
# Try to find the node data in the workflow
|
|
for node_id, node_data in prompt_.items():
|
|
if isinstance(node_data, dict) and node_data.get("class_type") == "GeminiNode":
|
|
if "inputs" in node_data and "prompt" in node_data["inputs"]:
|
|
prompt_text = str(node_data["inputs"]["prompt"])
|
|
logger.info(f"Extracted prompt from workflow data: {prompt_text[:50]}...")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error extracting prompt from workflow data: {e}")
|
|
|
|
logger.info(f"Using final prompt text: {prompt_text[:100]}...")
|
|
# --- End of Robust Prompt Extraction ---
|
|
|
|
if clear_history:
|
|
self.chat_history.clear()
|
|
|
|
# Generate a consistent seed for this operation
|
|
operation_seed = generate_consistent_seed(seed, use_random_seed)
|
|
|
|
# Handle different operation modes
|
|
if operation_mode == "generate_images":
|
|
return self.generate_images(
|
|
prompt=prompt_text,
|
|
model_name=model_name
|
|
if "image-generation" in model_name
|
|
else "gemini-2.0-flash-exp-image-generation",
|
|
images=images,
|
|
batch_count=batch_count,
|
|
temperature=temperature,
|
|
seed=operation_seed,
|
|
max_images=max_images,
|
|
aspect_ratio=aspect_ratio,
|
|
use_random_seed=use_random_seed,
|
|
external_api_key=api_key,
|
|
api_key_source=api_key_source,
|
|
unique_id=unique_id,
|
|
extra_pnginfo=extra_pnginfo,
|
|
prompt_=prompt_,
|
|
)
|
|
elif operation_mode == "generate_sequence":
|
|
# Call new sequence generation method
|
|
return self.generate_sequence(
|
|
prompt=prompt_text,
|
|
model_name=model_name,
|
|
images=images,
|
|
video=video,
|
|
audio=audio,
|
|
temperature=temperature,
|
|
seed=operation_seed,
|
|
max_output_tokens=max_output_tokens,
|
|
external_api_key=api_key,
|
|
unique_id=unique_id,
|
|
extra_pnginfo=extra_pnginfo,
|
|
prompt_=prompt_,
|
|
)
|
|
|
|
# Initialize the API client with the API key
|
|
try:
|
|
client = genai.Client(api_key=api_key)
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
logger.error(f"Error initializing Gemini client: {error_msg}", exc_info=True)
|
|
|
|
# Provide more helpful messages for common errors
|
|
if "invalid api key" in error_msg.lower() or "unauthorized" in error_msg.lower():
|
|
return (
|
|
"ERROR: Invalid Gemini API key. Please check your API key and try again.",
|
|
create_placeholder_image(),
|
|
)
|
|
elif "quota" in error_msg.lower() or "exceeded" in error_msg.lower():
|
|
return (
|
|
"ERROR: API quota exceeded. You've reached your usage limit for the Gemini API.",
|
|
create_placeholder_image(),
|
|
)
|
|
|
|
return (f"Error initializing Gemini client: {error_msg}", create_placeholder_image())
|
|
|
|
# Configure safety settings and generation parameters
|
|
safety_settings = [
|
|
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
|
]
|
|
|
|
generation_config = types.GenerateContentConfig(
|
|
max_output_tokens=max_output_tokens,
|
|
temperature=temperature,
|
|
seed=operation_seed, # Use our consistent seed
|
|
safety_settings=safety_settings
|
|
)
|
|
|
|
try:
|
|
if chat_mode:
|
|
# Handle chat mode with proper history
|
|
history = self.chat_history.get_messages_for_api()
|
|
|
|
# Create chat session
|
|
chat_session = client.chats.create(
|
|
model=model_name,
|
|
history=history
|
|
)
|
|
|
|
# Create appropriate content parts based on input type
|
|
contents = prepare_response(
|
|
prompt,
|
|
"image" if images is not None else "text",
|
|
None,
|
|
images,
|
|
video,
|
|
audio,
|
|
max_images,
|
|
)
|
|
# Extract content for chat format
|
|
if (
|
|
isinstance(contents, list)
|
|
and len(contents) == 1
|
|
and isinstance(contents[0], dict)
|
|
and "parts" in contents[0]
|
|
):
|
|
contents = contents[0]["parts"]
|
|
|
|
# Send message to chat and get response
|
|
response = chat_session.send_message(
|
|
content=contents,
|
|
config=generation_config,
|
|
)
|
|
|
|
# Add to history and format response
|
|
self.chat_history.add_message("user", prompt)
|
|
self.chat_history.add_message("assistant", response.text)
|
|
|
|
# Return the chat history
|
|
generated_content = self.chat_history.get_formatted_history()
|
|
|
|
else:
|
|
# Standard non-chat mode - prepare content for each input type
|
|
contents = prepare_response(
|
|
prompt,
|
|
"image" if images is not None else "text",
|
|
None,
|
|
images,
|
|
video,
|
|
audio,
|
|
max_images,
|
|
)
|
|
|
|
# Add structured output instruction if requested
|
|
if structured_output:
|
|
if (
|
|
isinstance(contents, list)
|
|
and len(contents) > 0
|
|
and isinstance(contents[0], dict)
|
|
and "parts" in contents[0]
|
|
and len(contents[0]["parts"]) > 0
|
|
):
|
|
if (
|
|
isinstance(contents[0]["parts"][0], dict)
|
|
and "text" in contents[0]["parts"][0]
|
|
):
|
|
contents[0]["parts"][0][
|
|
"text"
|
|
] = f"Please provide the response in a structured format. {contents[0]['parts'][0]['text']}"
|
|
|
|
# Generate content using the model
|
|
response = client.models.generate_content(
|
|
model=model_name,
|
|
contents=contents,
|
|
config=generation_config,
|
|
)
|
|
|
|
generated_content = response.text
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating content: {str(e)}", exc_info=True)
|
|
generated_content = f"Error: {str(e)}"
|
|
|
|
# For analysis mode, return the text response and an empty placeholder image
|
|
return (generated_content, create_placeholder_image())
|
|
|
|
def generate_images(
|
|
self,
|
|
prompt,
|
|
model_name,
|
|
images=None,
|
|
batch_count=1,
|
|
temperature=0.4,
|
|
seed=0,
|
|
max_images=6,
|
|
aspect_ratio="none",
|
|
use_random_seed=False,
|
|
external_api_key="",
|
|
api_key_source=None,
|
|
unique_id=None,
|
|
extra_pnginfo=None,
|
|
prompt_=None,
|
|
):
|
|
"""Generate images using Gemini models with image generation capabilities.
|
|
NOTE: Makes ONE API call. Batch_count > 1 doesn't multiply API calls here.
|
|
It might influence the prompt if we adapt it later, but currently ignored.
|
|
The API itself might return multiple images (up to its limit).
|
|
"""
|
|
logger.info(f"Starting generate_images for node {unique_id}")
|
|
|
|
if not self.genai_available:
|
|
return ("ERROR: Google Generative AI SDK not installed.", create_placeholder_image())
|
|
|
|
# --- Force the Correct Model ---
|
|
# Use the specific model designed for this integrated generation
|
|
forced_model_name = "gemini-2.0-flash-exp-image-generation"
|
|
if model_name != forced_model_name:
|
|
logger.warning(f"Operation mode is generate_images, overriding model to {forced_model_name}")
|
|
model_name = forced_model_name # Override
|
|
|
|
try:
|
|
from google import genai
|
|
from google.genai import types
|
|
except ImportError:
|
|
return ("ERROR: Failed to import Google Generative AI SDK", create_placeholder_image())
|
|
|
|
# --- API Key ---
|
|
# Get API key using the helper method
|
|
api_key, api_key_source = self._get_active_api_key(external_api_key)
|
|
|
|
if not api_key:
|
|
return (
|
|
"ERROR: No API key available for image generation.",
|
|
create_placeholder_image(),
|
|
)
|
|
|
|
# Handle prompt - ensure it's a string before processing
|
|
if isinstance(prompt, dict):
|
|
# If prompt is a dict, try to extract text from it
|
|
if "text" in prompt:
|
|
prompt_text = str(prompt["text"])
|
|
else:
|
|
# If no text key, convert the whole dict to string
|
|
prompt_text = str(prompt)
|
|
logger.info(f"Converted dict prompt to string for image generation: {prompt_text[:50]}...")
|
|
elif isinstance(prompt, (list, tuple)):
|
|
# If prompt is a list or tuple, join elements
|
|
prompt_text = " ".join(str(item) for item in prompt)
|
|
logger.info(f"Converted list/tuple prompt to string for image generation: {prompt_text[:50]}...")
|
|
else:
|
|
# Otherwise, ensure it's a string
|
|
prompt_text = str(prompt)
|
|
|
|
# --- Initialize Client ---
|
|
try:
|
|
client = genai.Client(api_key=api_key)
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
logger.error(f"Error initializing Gemini client for image generation: {error_msg}", exc_info=True)
|
|
if "invalid api key" in error_msg.lower():
|
|
return ("ERROR: Invalid Gemini API key.", create_placeholder_image())
|
|
if "quota" in error_msg.lower():
|
|
return ("ERROR: API quota exceeded.", create_placeholder_image())
|
|
return (f"Error initializing Gemini client: {error_msg}", create_placeholder_image())
|
|
|
|
# --- Seed for this single call ---
|
|
# Use the consistent seed generated earlier
|
|
current_seed = seed
|
|
logger.info(f"Using seed for image generation call: {current_seed}")
|
|
|
|
# --- Aspect Ratio / Dimensions ---
|
|
aspect_ratio_dimensions = {
|
|
"none": (1024, 1024), "1:1": (1024, 1024), "16:9": (1408, 768),
|
|
"9:16": (768, 1408), "4:3": (1280, 896), "3:4": (896, 1280),
|
|
"5:4": (1024, 819), "4:5": (819, 1024),
|
|
}
|
|
target_width, target_height = aspect_ratio_dimensions.get(aspect_ratio, (1024, 1024))
|
|
logger.info(f"Target resolution {target_width}x{target_height} for aspect ratio {aspect_ratio}")
|
|
dimension_prompt_part = f" Ensure the image dimensions are {target_width}x{target_height}." if aspect_ratio != "none" else ""
|
|
|
|
# --- Generation Config ---
|
|
generation_config = types.GenerateContentConfig(
|
|
temperature=temperature,
|
|
response_modalities=["Text", "Image"], # MUST include both
|
|
seed=current_seed,
|
|
safety_settings=[ # Keep safety settings relaxed for now
|
|
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
|
],
|
|
# candidate_count is NOT directly applicable here like in text generation for multiple choices.
|
|
# The API might return multiple images based on its internal logic, not this count.
|
|
)
|
|
|
|
# --- Prepare Content ---
|
|
content_parts = []
|
|
# Add prompt text, including dimension request
|
|
content_parts.append({"text": prompt_text + dimension_prompt_part})
|
|
|
|
# Process and add reference images if provided
|
|
if images is not None and isinstance(images, torch.Tensor) and images.nelement() > 0:
|
|
pil_images = []
|
|
# Handle batch dim
|
|
num_ref_images = min(images.shape[0], 6) if images.dim() == 4 else 1
|
|
image_source = images if images.dim() == 4 else images.unsqueeze(0) # Add batch dim if single
|
|
|
|
for i in range(num_ref_images):
|
|
try:
|
|
pil_img = tensor_to_pil(image_source[i])
|
|
# Resize reference images reasonably
|
|
pil_img = resize_image_to_dimensions(pil_img, 1024, 1024, crop_to_fit=False)
|
|
# Convert PIL to bytes
|
|
buffer = BytesIO()
|
|
pil_img.save(buffer, format="PNG")
|
|
image_bytes = buffer.getvalue()
|
|
# Add image part
|
|
content_parts.append({
|
|
"inline_data": { "mime_type": "image/png", "data": image_bytes }
|
|
})
|
|
except Exception as ref_img_e:
|
|
logger.error(f"Error processing reference image {i}: {ref_img_e}")
|
|
|
|
api_content = {"parts": content_parts}
|
|
|
|
# --- Make ONE API Call ---
|
|
all_generated_images = []
|
|
response_text_parts = []
|
|
status_text = ""
|
|
finish_reason = "unknown"
|
|
final_text = "No text returned by the model." # Initialize to default value
|
|
|
|
try:
|
|
logger.info(f"Calling Gemini model {model_name} for image generation (Node ID: {unique_id}, Seed: {current_seed})")
|
|
send_ws_message_safe("if-gemini-sequence-text", {"node_id": unique_id, "text": f"🔄 Calling {model_name} (Seed: {current_seed})..."})
|
|
|
|
response = client.models.generate_content(
|
|
model=model_name, contents=api_content, config=generation_config
|
|
)
|
|
|
|
# Process the response and extract parts
|
|
all_generated_images = []
|
|
response_text_parts = []
|
|
finish_reason = "UNKNOWN"
|
|
|
|
if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
|
|
if hasattr(response.candidates[0], "finish_reason"):
|
|
finish_reason = response.candidates[0].finish_reason
|
|
|
|
for part_idx, part in enumerate(response.candidates[0].content.parts):
|
|
if hasattr(part, 'text') and part.text:
|
|
response_text_parts.append(part.text)
|
|
send_ws_message_safe("if-gemini-sequence-text", {"node_id": unique_id, "text": part.text, "part_index": f"text_{part_idx}"})
|
|
|
|
elif hasattr(part, 'inline_data') and part.inline_data and 'image' in part.inline_data.mime_type:
|
|
try:
|
|
# Extract and add the image to the collection
|
|
image_bytes = part.inline_data.data
|
|
all_generated_images.append(image_bytes) # Store raw bytes
|
|
# Send image part to UI if needed
|
|
img_b64 = base64.b64encode(image_bytes).decode('utf-8')
|
|
send_ws_message_safe("if-gemini-sequence-image", {"node_id": unique_id, "image_b64": img_b64, "mime_type": part.inline_data.mime_type, "part_index": f"img_{part_idx}"})
|
|
except Exception as img_e:
|
|
logger.error(f"Error processing image part: {img_e}")
|
|
send_ws_message_safe("if-gemini-sequence-error", {"node_id": unique_id, "message": f"Error processing image: {img_e}"})
|
|
else:
|
|
if response.candidates and hasattr(response.candidates[0], "finish_reason"):
|
|
finish_reason = response.candidates[0].finish_reason
|
|
logger.warning(f"No valid parts found. Finish reason: {finish_reason}")
|
|
status_text = f"No content parts received. Reason: {finish_reason}"
|
|
send_ws_message_safe("if-gemini-sequence-text", {"node_id": unique_id, "text": status_text})
|
|
|
|
# Update final_text based on response_text_parts
|
|
final_text = "\n".join(response_text_parts) if response_text_parts else "No text returned by the model."
|
|
|
|
except Exception as e:
|
|
error_msg = f"Error calling Gemini API: {str(e)}"
|
|
logger.error(error_msg, exc_info=True)
|
|
status_text = error_msg
|
|
send_ws_message_safe("if-gemini-sequence-error", {"node_id": unique_id, "message": error_msg})
|
|
|
|
|
|
# Process images if any were generated
|
|
num_images_generated = len(all_generated_images)
|
|
|
|
if num_images_generated > 0:
|
|
logger.info(f"Generated {num_images_generated} images, processing for ComfyUI...")
|
|
image_tensors_list = []
|
|
|
|
for i, img_bytes in enumerate(all_generated_images[:max_images]):
|
|
try:
|
|
# Convert bytes to PIL image and then to tensor
|
|
pil_image = Image.open(BytesIO(img_bytes)).convert("RGB")
|
|
# Match aspect ratio if requested
|
|
if aspect_ratio != "none":
|
|
target_width, target_height = aspect_ratio_dimensions.get(aspect_ratio, (1024, 1024))
|
|
pil_image = resize_image_to_dimensions(pil_image, target_width, target_height)
|
|
# Convert to tensor
|
|
img_tensor = pil_to_tensor(pil_image)
|
|
image_tensors_list.append(img_tensor)
|
|
except Exception as tensor_e:
|
|
logger.error(f"Error processing image {i} to tensor: {tensor_e}")
|
|
|
|
# If we have processed images, batch them; otherwise return placeholder
|
|
if not image_tensors_list:
|
|
result_text = f"Generated {num_images_generated} image(s) but failed to process them.\nReason: {finish_reason}\n{status_text}\nModel Text: {final_text}"
|
|
send_ws_message_safe("if-gemini-sequence-complete", {"node_id": unique_id, "message": "Processing failed."})
|
|
return (result_text, create_placeholder_image())
|
|
|
|
# Process multiple images
|
|
if len(image_tensors_list) > 1:
|
|
# Ensure all tensors have the same dimensions
|
|
height, width = image_tensors_list[0].shape[1:3]
|
|
for i in range(1, len(image_tensors_list)):
|
|
if image_tensors_list[i].shape[1:3] != (height, width):
|
|
# Resize to match the first image's dimensions
|
|
pil_img = tensor_to_pil(image_tensors_list[i])
|
|
pil_img = pil_img.resize((width, height), Image.LANCZOS)
|
|
image_tensors_list[i] = pil_to_tensor(pil_img)
|
|
|
|
# Now batch all tensors
|
|
final_image_batch = torch.cat(image_tensors_list, dim=0)
|
|
else:
|
|
# Single image case
|
|
final_image_batch = image_tensors_list[0]
|
|
|
|
num_images_returned = len(image_tensors_list)
|
|
logger.info(f"Returning {num_images_returned} processed images")
|
|
|
|
# Check for requested image count from prompt
|
|
requested_count = None
|
|
prompt_lower = prompt_text.lower()
|
|
count_match = re.search(r'generate\s+(\d+)\s+images', prompt_lower)
|
|
if count_match:
|
|
try:
|
|
requested_count = int(count_match.group(1))
|
|
except ValueError:
|
|
requested_count = None
|
|
|
|
# Add explanation about batch_count vs. API generation
|
|
api_note = ""
|
|
if batch_count > 1:
|
|
api_note = f"\nNote: batch_count={batch_count} was provided, but the Gemini API makes a single API call that may return multiple images based on its internal logic."
|
|
|
|
if requested_count and requested_count > num_images_generated and num_images_generated > 0:
|
|
api_note += f"\nNote: You requested {requested_count} images in the prompt, but the API returned {num_images_generated}. The Gemini API determines the appropriate number of images based on the context and content requirements."
|
|
|
|
# Format final text output
|
|
result_text = (
|
|
f"Generated {num_images_generated} images, returned {num_images_returned}.\n"
|
|
f"Reason: {finish_reason}\n"
|
|
f"Model Text: {final_text}"
|
|
f"{api_note}"
|
|
)
|
|
|
|
# Create more informative completion message
|
|
completion_message = f"Complete: {num_images_returned} images."
|
|
if requested_count and requested_count > num_images_generated:
|
|
completion_message = f"Complete: {num_images_returned} images (API returned fewer images than the {requested_count} requested)."
|
|
|
|
send_ws_message_safe("if-gemini-sequence-complete", {
|
|
"node_id": unique_id,
|
|
"message": completion_message,
|
|
"requested_count": requested_count if requested_count else None,
|
|
"api_count": num_images_generated,
|
|
"processed_count": num_images_returned
|
|
})
|
|
|
|
return (result_text, final_image_batch)
|
|
|
|
else:
|
|
# No images generated
|
|
result_text = f"No images generated by {model_name}. Reason: {finish_reason}\n{status_text}\nModel Text: {final_text}"
|
|
send_ws_message_safe("if-gemini-sequence-complete", {"node_id": unique_id, "message": "Complete: No images generated."})
|
|
return (result_text, create_placeholder_image())
|
|
|
|
def generate_sequence(
|
|
self,
|
|
prompt,
|
|
model_name,
|
|
images=None,
|
|
video=None,
|
|
audio=None,
|
|
temperature=0.4,
|
|
seed=0,
|
|
max_output_tokens=8192,
|
|
max_images=4,
|
|
aspect_ratio="none",
|
|
external_api_key="",
|
|
unique_id=None,
|
|
extra_pnginfo=None,
|
|
prompt_=None,
|
|
):
|
|
logger.info(f"Starting generate_sequence for node {unique_id}")
|
|
# Check for SDK availability
|
|
if not self.genai_available:
|
|
return ("ERROR: Google Generative AI SDK not installed.", create_placeholder_image())
|
|
|
|
# Send initialization message to UI if available
|
|
send_ws_message_safe("if-gemini-sequence-init", {"node_id": unique_id})
|
|
|
|
if not unique_id:
|
|
logger.warning("Node unique_id is missing, UI updates may be incomplete")
|
|
|
|
# Force the correct model for generation regardless of user selection
|
|
enforced_model = "gemini-2.0-flash-exp-image-generation"
|
|
if model_name != enforced_model:
|
|
logger.warning(f"Sequence generation requires {enforced_model}. Overriding selected model {model_name}.")
|
|
model_name = enforced_model
|
|
|
|
# Get API key with proper logic
|
|
api_key, api_key_source = self._get_active_api_key(external_api_key)
|
|
if not api_key:
|
|
return ("ERROR: No API key found. Please provide a valid Gemini API key.", create_placeholder_image())
|
|
|
|
try:
|
|
# Import SDK
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
# Initialize client
|
|
try:
|
|
client = genai.Client(api_key=api_key)
|
|
except Exception as client_e:
|
|
error_msg = f"Failed to initialize API client: {str(client_e)}"
|
|
logger.error(error_msg)
|
|
send_ws_message_safe("if-gemini-sequence-error", {"node_id": unique_id, "message": error_msg})
|
|
return (error_msg, create_placeholder_image())
|
|
|
|
# Prepare content for API call
|
|
aspect_ratio_dimensions = {
|
|
"none": (1024, 1024),
|
|
"1:1": (1024, 1024),
|
|
"16:9": (1408, 768),
|
|
"9:16": (768, 1408),
|
|
"4:3": (1280, 896),
|
|
"3:4": (896, 1280),
|
|
"5:4": (1024, 819),
|
|
"4:5": (819, 1024),
|
|
}
|
|
target_width, target_height = aspect_ratio_dimensions.get(aspect_ratio, (1024, 1024))
|
|
|
|
# Add aspect ratio request to prompt if specified
|
|
dimension_prompt = ""
|
|
if aspect_ratio != "none":
|
|
dimension_prompt = f" Please create images with aspect ratio {aspect_ratio} ({target_width}x{target_height})."
|
|
|
|
# Handle prompt - ensure it's a string before processing
|
|
if isinstance(prompt, dict):
|
|
# If prompt is a dict, try to extract text from it
|
|
if "text" in prompt:
|
|
prompt_text = str(prompt["text"])
|
|
else:
|
|
# If no text key, convert the whole dict to string
|
|
prompt_text = str(prompt)
|
|
logger.info(f"Converted dict prompt to string: {prompt_text[:50]}...")
|
|
elif isinstance(prompt, (list, tuple)):
|
|
# If prompt is a list or tuple, join elements
|
|
prompt_text = " ".join(str(item) for item in prompt)
|
|
logger.info(f"Converted list/tuple prompt to string: {prompt_text[:50]}...")
|
|
else:
|
|
# Otherwise, ensure it's a string
|
|
prompt_text = str(prompt)
|
|
|
|
# Enhance prompt for sequence generation
|
|
prompt_lower = prompt_text.lower()
|
|
if "generate" in prompt_lower and "images" in prompt_lower:
|
|
# Try to extract number if user asked for specific count
|
|
count_match = re.search(r'generate\s+(\d+)\s+images', prompt_lower)
|
|
count_str = count_match.group(1) if count_match else "multiple"
|
|
enhanced_prompt = f"Create a sequence of text descriptions and corresponding images for the following request, aiming for around {count_str} steps if possible: {prompt_text}{dimension_prompt}"
|
|
logger.info(f"Enhanced sequence prompt: {enhanced_prompt[:100]}...")
|
|
final_prompt = enhanced_prompt
|
|
else:
|
|
final_prompt = f"{prompt_text}{dimension_prompt}"
|
|
|
|
# Prepare content parts
|
|
content_parts = [{"text": final_prompt}]
|
|
|
|
# Add reference images if provided
|
|
if images is not None and isinstance(images, torch.Tensor) and images.nelement() > 0:
|
|
# Handle single image
|
|
if images.dim() == 3:
|
|
images = images.unsqueeze(0) # Add batch dimension [1, H, W, C]
|
|
|
|
# Limit number of reference images to process
|
|
num_ref_images = min(images.shape[0], 6) # Process up to 6 reference images
|
|
logger.info(f"Adding {num_ref_images} reference images to request")
|
|
|
|
for i in range(num_ref_images):
|
|
try:
|
|
# Convert tensor to PIL image
|
|
pil_img = tensor_to_pil(images[i])
|
|
# Resize image for API (max 4MB)
|
|
pil_img = resize_image_to_dimensions(pil_img, 1024, 1024, crop_to_fit=False)
|
|
# Convert to bytes
|
|
img_buf = BytesIO()
|
|
pil_img.save(img_buf, format="PNG")
|
|
img_bytes = img_buf.getvalue()
|
|
# Add to content parts
|
|
content_parts.append({"inline_data": {"mime_type": "image/png", "data": img_bytes}})
|
|
except Exception as img_e:
|
|
logger.error(f"Error processing reference image {i}: {str(img_e)}")
|
|
|
|
# Configure generation settings
|
|
safety_settings = [
|
|
{"category": cat, "threshold": "BLOCK_NONE"}
|
|
for cat in ["HARM_CATEGORY_HARASSMENT", "HARM_CATEGORY_HATE_SPEECH",
|
|
"HARM_CATEGORY_SEXUALLY_EXPLICIT", "HARM_CATEGORY_DANGEROUS_CONTENT"]
|
|
]
|
|
|
|
generation_config = types.GenerateContentConfig(
|
|
max_output_tokens=max_output_tokens,
|
|
temperature=temperature,
|
|
seed=seed,
|
|
response_modalities=['Text', 'Image'], # This is key for mixed text+image output
|
|
safety_settings=safety_settings
|
|
)
|
|
|
|
# Make API call
|
|
logger.info(f"Calling Gemini model {model_name} for sequence generation (Node ID: {unique_id}, Seed: {seed})")
|
|
send_ws_message_safe("if-gemini-sequence-text", {"node_id": unique_id, "text": f"🔄 Calling {model_name} (Seed: {seed})..."})
|
|
|
|
response = client.models.generate_content(
|
|
model=model_name,
|
|
contents={"parts": content_parts},
|
|
config=generation_config
|
|
)
|
|
|
|
# Process response
|
|
all_text_parts = []
|
|
all_image_tensors = []
|
|
finish_reason = "UNKNOWN"
|
|
|
|
if response.candidates and response.candidates[0].content:
|
|
if hasattr(response.candidates[0], "finish_reason"):
|
|
finish_reason = response.candidates[0].finish_reason.name
|
|
logger.info(f"Response finish reason: {finish_reason}")
|
|
|
|
# Process each part in the response
|
|
for part_idx, part in enumerate(response.candidates[0].content.parts):
|
|
if hasattr(part, 'text') and part.text:
|
|
logger.debug(f"Node {unique_id}: Received text part")
|
|
all_text_parts.append(part.text)
|
|
send_ws_message_safe("if-gemini-sequence-text", {"node_id": unique_id, "text": part.text, "part_index": f"text_{part_idx}"})
|
|
|
|
elif hasattr(part, 'inline_data') and part.inline_data and 'image' in part.inline_data.mime_type:
|
|
try:
|
|
logger.debug(f"Node {unique_id}: Received image part {part_idx}")
|
|
# Process the image
|
|
image_bytes = part.inline_data.data
|
|
img_b64 = base64.b64encode(image_bytes).decode('utf-8')
|
|
send_ws_message_safe("if-gemini-sequence-image", {"node_id": unique_id, "image_b64": img_b64, "mime_type": part.inline_data.mime_type, "part_index": f"img_{part_idx}"})
|
|
|
|
# Convert to tensor for final output batch
|
|
if len(all_image_tensors) < max_images: # Only process up to max_images
|
|
pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
|
|
# Apply aspect ratio if specified
|
|
if aspect_ratio != "none":
|
|
pil_image = resize_image_to_dimensions(pil_image, target_width, target_height)
|
|
# Convert to tensor
|
|
img_tensor = pil_to_tensor(pil_image)
|
|
all_image_tensors.append(img_tensor)
|
|
else:
|
|
logger.warning(f"Node {unique_id}: Skipping image {part_idx} as max_images limit reached")
|
|
except Exception as img_e:
|
|
logger.error(f"Node {unique_id}: Error processing image part: {str(img_e)}", exc_info=True)
|
|
send_ws_message_safe("if-gemini-sequence-error", {"node_id": unique_id, "message": f"Error processing image: {str(img_e)}", "part_index": f"img_{part_idx}"})
|
|
else:
|
|
logger.warning(f"Node {unique_id}: Received unsupported part type or empty part")
|
|
else:
|
|
logger.warning(f"No valid parts found in response. Finish reason: {finish_reason}")
|
|
status_text = f"Generation finished but no content received. (Reason: {finish_reason})"
|
|
send_ws_message_safe("if-gemini-sequence-text", {"node_id": unique_id, "text": status_text})
|
|
all_text_parts.append(status_text)
|
|
|
|
# Prepare final outputs
|
|
# Create batched tensor from all images
|
|
if all_image_tensors:
|
|
try:
|
|
# Ensure all tensors are the same size for batching
|
|
target_h, target_w = all_image_tensors[0].shape[1:3]
|
|
resized_tensors = []
|
|
|
|
for tensor in all_image_tensors:
|
|
if tensor.shape[1] != target_h or tensor.shape[2] != target_w:
|
|
logger.debug(f"Resizing image from {tensor.shape[1]}x{tensor.shape[2]} to {target_h}x{target_w}")
|
|
pil_img = tensor_to_pil(tensor.squeeze(0))
|
|
resized = resize_image_to_dimensions(pil_img, target_w, target_h)
|
|
resized_tensor = pil_to_tensor(resized)
|
|
resized_tensors.append(resized_tensor)
|
|
else:
|
|
resized_tensors.append(tensor)
|
|
|
|
final_image_batch = torch.cat(resized_tensors, dim=0)
|
|
except Exception as batch_e:
|
|
logger.error(f"Error creating image batch: {str(batch_e)}")
|
|
final_image_batch = create_placeholder_image()
|
|
else:
|
|
# No images in response
|
|
final_image_batch = create_placeholder_image()
|
|
|
|
# Create final text output with improved information
|
|
final_text = "\n\n".join(all_text_parts) if all_text_parts else "No text generated."
|
|
|
|
# Check if user requested specific number of images
|
|
requested_count = None
|
|
count_match = re.search(r'generate\s+(\d+)\s+images', prompt_lower)
|
|
if count_match:
|
|
try:
|
|
requested_count = int(count_match.group(1))
|
|
except ValueError:
|
|
requested_count = None
|
|
|
|
# Add information about image generation limitations if applicable
|
|
num_images_generated = len(all_image_tensors)
|
|
api_note = ""
|
|
if requested_count and requested_count > num_images_generated and num_images_generated > 0:
|
|
api_note = (f"\n\nNote: You requested {requested_count} images, but the API returned {num_images_generated}. "
|
|
f"The Gemini API often returns fewer images than requested as it determines "
|
|
f"the appropriate number based on the context and content requirements.")
|
|
|
|
# Create result text with enhanced information
|
|
result_text = f"Sequence generation with {model_name}\n"
|
|
result_text += f"Finish reason: {finish_reason}\n"
|
|
result_text += f"Generated {len(all_text_parts)} text parts and {num_images_generated} images.\n"
|
|
if api_note:
|
|
result_text += api_note
|
|
result_text += f"\n\n{final_text}"
|
|
|
|
# Send final completion message
|
|
completion_message = f"✅ Sequence generation complete: {len(all_text_parts)} text parts, {num_images_generated} images. Reason: {finish_reason}"
|
|
if requested_count and requested_count > num_images_generated:
|
|
completion_message += f" (Note: API returned fewer images than the {requested_count} requested)"
|
|
|
|
send_ws_message_safe("if-gemini-sequence-complete", {
|
|
"node_id": unique_id,
|
|
"message": completion_message,
|
|
"text_count": len(all_text_parts),
|
|
"image_count": num_images_generated
|
|
})
|
|
|
|
logger.info(f"Node {unique_id}: Sequence finished. Returning {len(all_text_parts)} text parts and {num_images_generated} images.")
|
|
return (result_text, final_image_batch)
|
|
|
|
except Exception as e:
|
|
error_msg = f"Error during sequence generation: {str(e)}"
|
|
logger.error(f"Node {unique_id}: {error_msg}", exc_info=True)
|
|
send_ws_message_safe("if-gemini-sequence-error", {"node_id": unique_id, "message": error_msg})
|
|
return (error_msg, create_placeholder_image())
|
|
|
|
def get_available_models(api_key):
|
|
"""Get available Gemini models for a given API key"""
|
|
try:
|
|
from google import genai
|
|
|
|
# Initialize client with the provided API key
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
# List available models
|
|
models_response = client.models.list()
|
|
|
|
# Filter for Gemini models only
|
|
gemini_models = []
|
|
for model in models_response:
|
|
if "gemini" in model.name.lower():
|
|
# Extract just the model name from the full path
|
|
model_name = model.name.split('/')[-1]
|
|
gemini_models.append(model_name)
|
|
|
|
# Ensure we always have the default models available
|
|
default_models = [
|
|
"gemini-2.0-flash-exp",
|
|
"gemini-2.0-pro",
|
|
"gemini-2.0-flash",
|
|
"gemini-2.0-flash-exp-image-generation"
|
|
]
|
|
|
|
for model in default_models:
|
|
if model not in gemini_models:
|
|
gemini_models.append(model)
|
|
|
|
return gemini_models
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error retrieving models: {str(e)}")
|
|
# Return default models on error
|
|
return [
|
|
"gemini-2.0-flash-exp",
|
|
"gemini-2.0-pro",
|
|
"gemini-2.0-flash",
|
|
"gemini-2.0-flash-exp-image-generation"
|
|
]
|
|
|
|
def check_gemini_api_key(api_key):
|
|
"""Check if a Gemini API key is valid by attempting to list models"""
|
|
try:
|
|
from google import genai
|
|
|
|
# Initialize client with the provided API key
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
# Try to list models as a simple API test
|
|
models = client.models.list()
|
|
|
|
# If we get here, the API key is valid
|
|
return True, "API key is valid. Successfully connected to Gemini API."
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
logger.error(f"API key validation error: {error_msg}")
|
|
|
|
# Provide more helpful error messages
|
|
if "invalid api key" in error_msg.lower() or "unauthorized" in error_msg.lower():
|
|
return False, "Invalid API key. Please check your API key and try again."
|
|
elif "quota" in error_msg.lower() or "exceeded" in error_msg.lower():
|
|
return False, "API quota exceeded. You've reached your usage limit for the Gemini API."
|
|
else:
|
|
return False, f"Error validating API key: {error_msg}"
|
|
|
|
# Ensure pil_to_tensor function is available
|
|
def pil_to_tensor(pil_image):
|
|
"""Convert PIL image to tensor with shape [1, H, W, 3]"""
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
img_array = np.array(pil_image).astype(np.float32) / 255.0
|
|
# Ensure the array has 3 dimensions [H, W, C] before adding batch
|
|
if img_array.ndim == 2: # Handle grayscale
|
|
img_array = np.stack((img_array,)*3, axis=-1)
|
|
img_tensor = torch.from_numpy(img_array)[None,] # Add batch dimension [1, H, W, 3]
|
|
return img_tensor
|