906 lines
44 KiB
Python
906 lines
44 KiB
Python
import os
|
||
import json
|
||
import uuid
|
||
import re
|
||
import torch
|
||
import numpy as np
|
||
import cv2
|
||
from PIL import Image
|
||
from io import BytesIO
|
||
import folder_paths # type: ignore[reportMissingImports]
|
||
from comfy_api.latest import ComfyExtension, UI, IO # type: ignore[reportMissingImports]
|
||
from google import genai
|
||
from google.genai import types
|
||
import time
|
||
import traceback
|
||
import threading
|
||
import queue
|
||
import sys
|
||
import importlib
|
||
import subprocess
|
||
import random
|
||
import hashlib
|
||
from typing import Any, Tuple
|
||
|
||
def check_and_install_dependencies():
|
||
required_packages = {
|
||
'requests': 'requests',
|
||
'pysocks': 'PySocks',
|
||
}
|
||
|
||
missing_packages = []
|
||
for module_name, package_name in required_packages.items():
|
||
try:
|
||
importlib.import_module(module_name)
|
||
except ImportError:
|
||
missing_packages.append(package_name)
|
||
|
||
if missing_packages:
|
||
print(f"[WARNING] Missing required dependencies: {', '.join(missing_packages)}.")
|
||
print(f"[INFO] Please install them using: pip install {' '.join(missing_packages)}")
|
||
print(f"[INFO] Alternatively, install all requirements: pip install -r requirements.txt")
|
||
|
||
try:
|
||
check_and_install_dependencies()
|
||
except Exception as e:
|
||
print(f"[WARNING] Error checking dependencies: {str(e)}")
|
||
|
||
class GetKeyAPI(IO.ComfyNode):
|
||
@classmethod
|
||
def define_schema(cls) -> IO.Schema:
|
||
return IO.Schema(
|
||
node_id="GetKeyAPI",
|
||
display_name="Get API Key from JSON",
|
||
category="utils/api_keys",
|
||
inputs=[
|
||
IO.String.Input("json_path", default="./input/apikeys.json", multiline=False, tooltip="Path to a .json file with simple top level structure with name as key and api-key as value. See example in custom node folder."),
|
||
IO.Combo.Input("key_id_method", options=["custom", "random_rotate", "increment_rotate"], default="custom", tooltip="custom sets api-key to the api-key with the name set in the key_id widget. random_rotate randomly switches between keys if multiple in the .json and increment_rotate does it in order from first to last, then repeats."),
|
||
IO.Int.Input("rotation_interval", default=0, min=0, tooltip="how many steps to jump when doing rotate."),
|
||
IO.String.Input("key_id", default="placeholder", multiline=False, optional=True, tooltip="Put name of key in the .json here if using custom in key_id_method."),
|
||
],
|
||
outputs=[
|
||
IO.String.Output("API_KEY")
|
||
]
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, json_path: str, key_id_method: str, rotation_interval: int, key_id: str | None = "placeholder") -> IO.NodeOutput:
|
||
api_keys_data = None
|
||
absolute_json_path = os.path.abspath(json_path)
|
||
|
||
try:
|
||
with open(absolute_json_path, 'r') as f:
|
||
api_keys_data = json.load(f)
|
||
except FileNotFoundError:
|
||
raise ValueError(f"RotateKeyAPI Error: JSON file not found at {absolute_json_path}")
|
||
except json.JSONDecodeError:
|
||
raise ValueError(f"RotateKeyAPI Error: Could not decode JSON from {absolute_json_path}. Check file format.")
|
||
except Exception as e:
|
||
raise RuntimeError(f"RotateKeyAPI Error: Unexpected error reading file {absolute_json_path}: {e}")
|
||
|
||
if not isinstance(api_keys_data, dict):
|
||
raise ValueError(f"RotateKeyAPI Error: JSON content is not a dictionary in {absolute_json_path}. Expected format: {{'key_id': 'api_key', ...}}")
|
||
|
||
if not api_keys_data:
|
||
raise ValueError(f"RotateKeyAPI Error: The JSON dictionary in {absolute_json_path} is empty.")
|
||
|
||
selected_key_value = None
|
||
|
||
if key_id_method == "custom":
|
||
if key_id == "placeholder":
|
||
print("RotateKeyAPI Warning: 'custom' method selected but 'key_id' is still the default 'placeholder'. Ensure this is intended or provide a valid key ID.")
|
||
|
||
selected_key_value = api_keys_data.get(key_id)
|
||
|
||
if selected_key_value is None:
|
||
raise ValueError(f"RotateKeyAPI Error: Custom key ID '{key_id}' not found in the JSON dictionary keys.")
|
||
|
||
|
||
elif key_id_method == "random_rotate":
|
||
api_keys_list = list(api_keys_data.values())
|
||
|
||
selected_key_value = random.choice(api_keys_list)
|
||
|
||
elif key_id_method == "increment_rotate":
|
||
api_keys_list = list(api_keys_data.values())
|
||
|
||
index = rotation_interval % len(api_keys_list)
|
||
|
||
try:
|
||
selected_key_value = api_keys_list[index]
|
||
except IndexError:
|
||
raise IndexError(f"RotateKeyAPI Error: Calculated index {index} (from interval {rotation_interval}) is out of bounds for list of size {len(api_keys_list)}.")
|
||
except Exception as e:
|
||
raise RuntimeError(f"RotateKeyAPI Error: Unexpected error accessing item at index {index}: {e}")
|
||
|
||
if not isinstance(selected_key_value, str) or not selected_key_value:
|
||
raise ValueError(f"RotateKeyAPI Error: Retrieved value for selected key is not a valid string. Value: {selected_key_value}")
|
||
|
||
print(f"RotateKeyAPI: Successfully retrieved API key using method '{key_id_method}'.")
|
||
return IO.NodeOutput(selected_key_value)
|
||
|
||
|
||
|
||
class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
|
||
GemConfig = IO.Custom("GEMINI_CONFIG")
|
||
|
||
@classmethod
|
||
def define_schema(cls) -> IO.Schema:
|
||
return IO.Schema(
|
||
node_id="SSL_GeminiAPIKeyConfig",
|
||
display_name="Configure Gemini API Key",
|
||
category="API/Gemini",
|
||
inputs=[
|
||
IO.String.Input("api_key", multiline=False, default=""),
|
||
IO.Combo.Input("api_version", options=["v1", "v1alpha", "v1beta", "v1beta1", "v2beta"], default="v1alpha", tooltip="Select API version to use. v1alpha, v1beta and v2beta are Gemini API specific while v1beta1 is Vertex AI specific. Both can use v1"),
|
||
IO.Boolean.Input("vertexai", default=False, tooltip="Enable Vertex AI / Gemini Enterprise Agent Platform mode. Uses ADC/Service Accounts and ignores api_key."),
|
||
IO.String.Input("vertexai_project", optional=True, tooltip="Google Cloud Project ID. Falls back to GOOGLE_CLOUD_PROJECT env var."),
|
||
IO.String.Input("vertexai_location", optional=True, tooltip="Google Cloud Region (e.g. us-central1). Falls back to GOOGLE_CLOUD_LOCATION env var."),
|
||
],
|
||
outputs=[
|
||
cls.GemConfig.Output("config")
|
||
]
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, api_key: str, api_version: str, vertexai: bool, vertexai_project: str | None = "", vertexai_location: str | None = "") -> IO.NodeOutput:
|
||
config = {
|
||
"api_key": api_key,
|
||
"api_version": api_version,
|
||
"vertexai": vertexai,
|
||
"vertexai_project": vertexai_project,
|
||
"vertexai_location": vertexai_location
|
||
}
|
||
return IO.NodeOutput(config)
|
||
|
||
|
||
|
||
class SSL_GeminiTextPrompt(IO.ComfyNode):
|
||
GemConfig = IO.Custom("GEMINI_CONFIG")
|
||
_cache: dict = {}
|
||
_seed_map_cache: dict = {} # Maps (input_seed, fingerprint_without_seed) -> successful_gemini_seed
|
||
_client_cache: dict = {} # Maps client_key tuple -> genai.Client instance
|
||
|
||
# Define model lists centrally to ensure consistency between cache logic and execution logic
|
||
THINKING_MODELS = [
|
||
"gemini-2.0-flash-thinking-exp", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-flash-thinking-exp-1219",
|
||
"gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-preview-04-17", "gemini-2.5-flash-preview-05-20", "gemini-2.5-flash-preview-09-25",
|
||
"gemini-2.5-pro-preview-06-05", "gemini-2.5-pro-preview-05-06", "gemini-3-pro-preview", "gemini-3-flash-preview",
|
||
"gemini-3.1-pro-preview", "gemini-pro-latest", "gemini-flash-latest"
|
||
]
|
||
GEN3_THINKING_MODELS = [
|
||
"gemini-pro-latest", "gemini-flash-latest", "gemini-3.1-pro-preview",
|
||
"gemini-3-pro-preview", "gemini-3-flash-preview"
|
||
]
|
||
IMAGE_MODELS = ["gemini-2.5-flash-image-preview", "gemini-2.5-flash-image", "gemini-3-pro-image-preview", "nano-banana-pro-preview"]
|
||
MEDIA_RES_MODELS = [
|
||
"gemini-3.1-flash-lite-preview", "gemini-3-pro-preview",
|
||
"gemini-3-flash-preview", "gemini-3.1-pro-preview",
|
||
"gemini-pro-latest", "gemini-flash-latest", "gemini-flash-lite-latest"
|
||
]
|
||
|
||
@classmethod
|
||
def define_schema(cls) -> IO.Schema:
|
||
return IO.Schema(
|
||
node_id="SSL_GeminiTextPrompt",
|
||
display_name="Expanded Gemini Text/Image",
|
||
category="API/Gemini",
|
||
inputs=[
|
||
cls.GemConfig.Input("config"),
|
||
IO.String.Input("prompt", multiline=True),
|
||
IO.String.Input("system_instruction", default="You are a helpful AI assistant.", multiline=True),
|
||
IO.Combo.Input(
|
||
"model",
|
||
options=[
|
||
"gemini-2.0-flash",
|
||
"gemini-2.0-flash-lite",
|
||
"gemini-2.5-pro",
|
||
"gemini-2.5-flash",
|
||
"gemini-2.5-flash-lite",
|
||
"gemini-3-pro-preview",
|
||
"gemini-3-flash-preview",
|
||
"gemini-3.1-flash-lite-preview",
|
||
"gemini-3.1-pro-preview",
|
||
"gemini-2.5-flash-image-preview",
|
||
"nano-banana-pro-preview",
|
||
"gemini-pro-latest",
|
||
"gemini-flash-latest",
|
||
"gemini-flash-lite-latest",
|
||
"gemini-2.5-flash-preview-04-17",
|
||
"gemini-2.5-flash-preview-05-20",
|
||
"gemini-2.5-flash-preview-09-25",
|
||
"gemini-2.5-flash-lite-preview-06-17",
|
||
"gemini-2.5-flash-lite-preview-09-25",
|
||
"gemini-2.5-pro-preview-06-05",
|
||
"gemini-2.5-pro-preview-05-06",
|
||
],
|
||
default="gemini-2.5-flash",
|
||
),
|
||
IO.Float.Input("temperature", default=1.0, min=0.0, max=1.0, step=0.01),
|
||
IO.Float.Input("top_p", default=0.95, min=0.0, max=1.0, step=0.01),
|
||
IO.Int.Input("top_k", default=40, min=1, max=100, step=1),
|
||
IO.Int.Input("max_output_tokens", default=8192, min=1, max=65536, step=1),
|
||
IO.Boolean.Input("include_images", default=False),
|
||
IO.Combo.Input("aspect_ratio", options=["None", "1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5", "21:9"], default="None"),
|
||
IO.Combo.Input("bypass_mode", options=["None", "system_instruction", "prompt", "both"], default="None"),
|
||
IO.Int.Input("thinking_budget", default=0, min=-1, max=24576, step=1, tooltip="0 disables thinking mode, -1 will activate it as default dynamic thinking and anything above 0 sets specific budget"),
|
||
IO.Image.Input("input_image", optional=True),
|
||
IO.Image.Input("input_image_2", optional=True),
|
||
IO.Boolean.Input("use_proxy", default=False),
|
||
IO.String.Input("proxy_host", default="127.0.0.1"),
|
||
IO.Int.Input("proxy_port", default=7890, min=1, max=65535),
|
||
IO.Boolean.Input("use_seed", default=True),
|
||
IO.Int.Input("seed", default=0, min=0, max=2147483647),
|
||
IO.Int.Input("timeout", default=30, min=15, max=300, step=15),
|
||
IO.Boolean.Input("include_thoughts", default=False),
|
||
IO.Combo.Input("thinking_level", options=["None", "low", "medium", "high"], default="None", tooltip="Does not work at the same time as 'thinking_budget'. if this is set, then thinking budget is ignored."),
|
||
IO.Combo.Input("media_resolution", options=["unspecified", "low", "medium", "high"], default="unspecified", tooltip="Set input media resolution for image, video and pdf. This changes tokens consumed."),
|
||
IO.String.Input("retry_pattern", default="", optional=True, multiline=False, tooltip="Regex pattern to match in response text. If matched, retry with new seed. Leave empty to disable."),
|
||
IO.Int.Input("max_retries", default=3, min=0, max=10, step=1, tooltip="Maximum number of retry attempts when pattern matches. 0 disables retry."),
|
||
],
|
||
outputs=[
|
||
IO.String.Output("text"),
|
||
IO.Image.Output("image"),
|
||
IO.Int.Output("final_actual_seed")
|
||
]
|
||
)
|
||
|
||
@classmethod
|
||
def _pad_text_with_joiners(cls, text: str) -> str:
|
||
if not text:
|
||
return ""
|
||
|
||
patternperiod = r"\."
|
||
patternspace = r"\s"
|
||
patterncomma = r","
|
||
patterndash = r"\-"
|
||
patternsingq = r"\'"
|
||
patterndoubq = r'\"'
|
||
patternword = r"(.)(?=.)"
|
||
|
||
replperiod = r"。"
|
||
replspace = r""
|
||
replcomma = r"、"
|
||
repldash = r"‐"
|
||
replsingq = r"ʼ"
|
||
repldoubq = r"ˮ"
|
||
replword = r"\1"
|
||
|
||
joined_textperiod = re.sub(patternperiod, replperiod, text)
|
||
joined_textspace = re.sub(patternspace, replspace, joined_textperiod)
|
||
joined_textcomma = re.sub(patterncomma, replcomma, joined_textspace)
|
||
joined_textdash = re.sub(patterndash, repldash, joined_textcomma)
|
||
joined_textsingq = re.sub(patternsingq, replsingq, joined_textdash)
|
||
joined_textdoubq = re.sub(patterndoubq, repldoubq, joined_textsingq)
|
||
joined_textfinal = re.sub(patternword, replword, joined_textdoubq)
|
||
|
||
print(joined_textfinal)
|
||
|
||
return joined_textfinal
|
||
|
||
@classmethod
|
||
def save_binary_file(cls, data, mime_type):
|
||
ext = ".bin"
|
||
if mime_type == "image/png":
|
||
ext = ".png"
|
||
elif mime_type == "image/jpeg":
|
||
ext = ".jpg"
|
||
|
||
output_dir = folder_paths.get_output_directory()
|
||
gemini_dir = os.path.join(output_dir, "gemini_outputs")
|
||
os.makedirs(gemini_dir, exist_ok=True)
|
||
|
||
file_name = os.path.join(gemini_dir, f"gemini_output_{uuid.uuid4()}{ext}")
|
||
|
||
with open(file_name, "wb") as f:
|
||
f.write(data)
|
||
|
||
return file_name
|
||
|
||
@classmethod
|
||
def generate_empty_image(cls, width=64, height=64):
|
||
empty_image = np.ones((height, width, 3), dtype=np.float32) * 0.2
|
||
tensor = torch.from_numpy(empty_image).unsqueeze(0)
|
||
return tensor
|
||
|
||
@classmethod
|
||
def _compute_fingerprint_and_check_cache(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
|
||
include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed,
|
||
input_image=None, input_image_2=None,
|
||
use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, timeout=30,
|
||
include_thoughts=False, thinking_level=None, media_resolution=None,
|
||
retry_pattern="", max_retries=3):
|
||
|
||
# 1. Hashing Images
|
||
def get_tensor_hash(tensor):
|
||
if tensor is None:
|
||
return "None"
|
||
try:
|
||
return hashlib.sha256(np.ascontiguousarray(tensor.cpu().numpy()).tobytes()).hexdigest()
|
||
except Exception as e:
|
||
print(f"[WARNING] Cache hashing failed for image: {e}")
|
||
return "Error"
|
||
|
||
image_1_hash = get_tensor_hash(input_image)
|
||
image_2_hash = get_tensor_hash(input_image_2)
|
||
|
||
# 2. Determine Effective Parameters based on Model
|
||
# This ensures we don't cache-miss if an irrelevant parameter changes
|
||
|
||
# Defaults
|
||
eff_include_images = False
|
||
eff_aspect_ratio = "None"
|
||
eff_thinking_level = "None"
|
||
eff_thinking_budget = -1
|
||
eff_include_thoughts = False
|
||
|
||
# Logic mirroring _build_generate_content_config exactly
|
||
if include_images and model in cls.IMAGE_MODELS:
|
||
eff_include_images = True
|
||
eff_aspect_ratio = str(aspect_ratio)
|
||
# When generating images, thinking params are ignored
|
||
|
||
elif model in cls.GEN3_THINKING_MODELS and thinking_level is not None and thinking_level != "None":
|
||
eff_thinking_level = thinking_level
|
||
eff_include_thoughts = include_thoughts
|
||
# Budget is ignored in this specific branch
|
||
|
||
elif model in cls.THINKING_MODELS:
|
||
eff_thinking_budget = int(thinking_budget)
|
||
eff_include_thoughts = include_thoughts
|
||
# Thinking level is ignored in this branch
|
||
|
||
# If none of the above, all effective thinking/image params remain default/ignored
|
||
|
||
fingerprint = (
|
||
str(config),
|
||
prompt,
|
||
system_instruction,
|
||
model,
|
||
float(temperature),
|
||
float(top_p),
|
||
int(top_k),
|
||
int(max_output_tokens),
|
||
eff_include_images, # EFFECTIVE include_images
|
||
eff_aspect_ratio, # EFFECTIVE aspect_ratio
|
||
str(bypass_mode),
|
||
eff_thinking_budget, # EFFECTIVE thinking_budget
|
||
use_seed,
|
||
int(seed) if use_seed else 0, # Only use seed in cache if use_seed is True
|
||
image_1_hash,
|
||
image_2_hash,
|
||
bool(use_proxy),
|
||
str(proxy_host),
|
||
int(proxy_port),
|
||
int(timeout),
|
||
eff_include_thoughts, # EFFECTIVE include_thoughts
|
||
eff_thinking_level, # EFFECTIVE thinking_level
|
||
str(media_resolution),
|
||
str(retry_pattern), # Include retry pattern in fingerprint
|
||
int(max_retries) # Include max retries in fingerprint
|
||
)
|
||
|
||
cached = cls._cache.get(fingerprint)
|
||
if use_seed and cached is not None:
|
||
return fingerprint, cached
|
||
return fingerprint, None
|
||
|
||
@classmethod
|
||
def _handle_seed(cls, use_seed, seed):
|
||
actual_seed = None
|
||
if use_seed:
|
||
if seed == 0:
|
||
current_time = int(time.time() * 1000)
|
||
random_component = random.randint(0, 1000000)
|
||
actual_seed = (current_time + random_component) % 2147483647
|
||
print(f"[INFO] Generated random seed: {actual_seed}")
|
||
else:
|
||
actual_seed = seed
|
||
print(f"[INFO] Using specified seed: {actual_seed}")
|
||
|
||
random.seed(actual_seed)
|
||
np.random.seed(actual_seed)
|
||
torch.manual_seed(actual_seed)
|
||
if torch.cuda.is_available():
|
||
torch.cuda.manual_seed_all(actual_seed)
|
||
else:
|
||
print("[INFO] Seed not used")
|
||
|
||
return actual_seed
|
||
|
||
@classmethod
|
||
def _setup_proxy_env(cls, proxy_host, proxy_port):
|
||
if not proxy_host.startswith(('http://', 'https://')):
|
||
proxy_url = f"http://{proxy_host}:{proxy_port}"
|
||
else:
|
||
proxy_url = f"{proxy_host}:{proxy_port}"
|
||
|
||
os.environ['HTTP_PROXY'] = proxy_url
|
||
os.environ['HTTPS_PROXY'] = proxy_url
|
||
os.environ['http_proxy'] = proxy_url
|
||
os.environ['https_proxy'] = proxy_url
|
||
os.environ['REQUESTS_CA_BUNDLE'] = ''
|
||
|
||
print(f"[INFO] Proxy enabled: {proxy_url}")
|
||
return proxy_url
|
||
|
||
@classmethod
|
||
def _build_generate_content_config(cls, model, temperature, top_p, top_k, max_output_tokens, seed,
|
||
include_images, response_modalities, aspect_ratio, padded_system_instruction,
|
||
thinking_level, thinking_budget, include_thoughts, media_resolution):
|
||
# Centralized builder for GenerateContentConfig used by different model/feature branches
|
||
safety = [
|
||
types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"),
|
||
types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"),
|
||
types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
|
||
types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"),
|
||
types.SafetySetting(category="HARM_CATEGORY_CIVIC_INTEGRITY", threshold="BLOCK_NONE"),
|
||
]
|
||
|
||
# Modified: Only trigger image config if include_images is True AND the model is actually an image model
|
||
if include_images and model in cls.IMAGE_MODELS:
|
||
return types.GenerateContentConfig(
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
top_k=top_k,
|
||
seed=seed,
|
||
max_output_tokens=max_output_tokens,
|
||
safety_settings=safety,
|
||
response_modalities=response_modalities,
|
||
image_config=types.ImageConfig(aspect_ratio=aspect_ratio),
|
||
system_instruction=[types.Part.from_text(text=padded_system_instruction)],
|
||
)
|
||
|
||
G3Pro = ["gemini-3.1-pro-preview", "gemini-3-pro-preview"]
|
||
# Modified: Added check for thinking_level != "None"
|
||
if model in G3Pro and thinking_level is not None and thinking_level != "None":
|
||
return types.GenerateContentConfig(
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
top_k=top_k,
|
||
max_output_tokens=max_output_tokens,
|
||
safety_settings=safety,
|
||
thinking_config=types.ThinkingConfig(include_thoughts=include_thoughts, thinking_level=thinking_level),
|
||
response_modalities=response_modalities,
|
||
system_instruction=[types.Part.from_text(text=padded_system_instruction)],
|
||
)
|
||
|
||
if model in cls.THINKING_MODELS:
|
||
return types.GenerateContentConfig(
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
top_k=top_k,
|
||
max_output_tokens=max_output_tokens,
|
||
safety_settings=safety,
|
||
thinking_config=types.ThinkingConfig(include_thoughts=include_thoughts, thinking_budget=thinking_budget),
|
||
response_modalities=response_modalities,
|
||
system_instruction=[types.Part.from_text(text=padded_system_instruction)],
|
||
)
|
||
|
||
return types.GenerateContentConfig(
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
top_k=top_k,
|
||
max_output_tokens=max_output_tokens,
|
||
safety_settings=safety,
|
||
response_modalities=response_modalities,
|
||
system_instruction=[types.Part.from_text(text=padded_system_instruction)],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
|
||
include_images, aspect_ratio, bypass_mode, thinking_budget, input_image=None, input_image_2=None,
|
||
use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, use_seed=False, seed=0, timeout=30,
|
||
include_thoughts=False, thinking_level=None, media_resolution=None,
|
||
retry_pattern="", max_retries=3) -> IO.NodeOutput:
|
||
|
||
print(f"[INFO] SSL_GeminiTextPrompt execute called, model: {model}")
|
||
fingerprint, cached = cls._compute_fingerprint_and_check_cache(
|
||
config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
|
||
include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed,
|
||
input_image, input_image_2,
|
||
use_proxy, proxy_host, proxy_port, timeout,
|
||
include_thoughts, thinking_level, media_resolution,
|
||
retry_pattern, max_retries
|
||
)
|
||
|
||
if cached is not None:
|
||
cached_text, cached_image, cached_seed = cached
|
||
print(f"[INFO] Returning cached result for fingerprint {fingerprint}")
|
||
return IO.NodeOutput(cached_text, cached_image, cached_seed)
|
||
|
||
# --- keep most of the original implementation but converted to classmethod usage ---
|
||
original_http_proxy = os.environ.get('HTTP_PROXY')
|
||
original_https_proxy = os.environ.get('HTTPS_PROXY')
|
||
original_http_proxy_lower = os.environ.get('http_proxy')
|
||
original_https_proxy_lower = os.environ.get('https_proxy')
|
||
|
||
print(f"[INFO] Starting generation, model: {model}, temperature: {temperature}")
|
||
|
||
padded_prompt = prompt
|
||
padded_system_instruction = system_instruction
|
||
|
||
if bypass_mode == "prompt" or bypass_mode == "both":
|
||
padded_prompt = cls._pad_text_with_joiners(prompt)
|
||
print(padded_prompt)
|
||
if bypass_mode == "system_instruction" or bypass_mode == "both":
|
||
padded_system_instruction = cls._pad_text_with_joiners(system_instruction)
|
||
print(padded_system_instruction)
|
||
|
||
actual_seed = cls._handle_seed(use_seed, seed)
|
||
input_seed = seed # Store the original input seed for cache key
|
||
|
||
# Check if we have a cached successful gemini seed for this input seed
|
||
# Build a cache key that excludes the seed itself (to match different gemini seeds for same input)
|
||
if use_seed and retry_pattern and max_retries > 0:
|
||
seed_cache_key = (input_seed, fingerprint[:-4]) # Exclude seed, retry_pattern, max_retries from key
|
||
cached_gemini_seed = cls._seed_map_cache.get(seed_cache_key)
|
||
if cached_gemini_seed is not None:
|
||
print(f"[INFO] Using cached successful gemini seed {cached_gemini_seed} for input seed {input_seed}")
|
||
actual_seed = cached_gemini_seed
|
||
|
||
# Flatten and simplify nested try/except blocks to ensure correct pairing
|
||
original_http_proxy = os.environ.get('HTTP_PROXY')
|
||
original_https_proxy = os.environ.get('HTTPS_PROXY')
|
||
original_http_proxy_lower = os.environ.get('http_proxy')
|
||
original_https_proxy_lower = os.environ.get('https_proxy')
|
||
|
||
text_output = ""
|
||
image_tensor = cls.generate_empty_image()
|
||
proxy_url: str | None = None
|
||
|
||
try:
|
||
if use_proxy:
|
||
proxy_url = cls._setup_proxy_env(proxy_host, proxy_port)
|
||
|
||
# Initialize Gemini client
|
||
client_options = {}
|
||
if use_proxy:
|
||
try:
|
||
import google.api_core.http_client # type: ignore[import]
|
||
import google.auth.transport.requests # type: ignore[import]
|
||
import requests
|
||
from requests.adapters import HTTPAdapter
|
||
|
||
class ProxyAdapter(HTTPAdapter):
|
||
def __init__(self, proxy_url, **kwargs):
|
||
self.proxy_url = proxy_url
|
||
super().__init__(**kwargs)
|
||
|
||
def add_headers(self, request, **kwargs):
|
||
super().add_headers(request, **kwargs)
|
||
|
||
session = requests.Session()
|
||
proxies = {"http": str(proxy_url), "https": str(proxy_url)}
|
||
session.proxies.update(proxies)
|
||
adapter = ProxyAdapter(proxy_url, max_retries=1)
|
||
session.mount('http://', adapter)
|
||
session.mount('https://', adapter)
|
||
session.verify = False
|
||
http_client = google.api_core.http_client.RequestsHttpClient(session=session)
|
||
client_options["http_client"] = http_client
|
||
except Exception:
|
||
# best-effort proxy HTTP client setup; fall back if imports fail
|
||
pass
|
||
|
||
try:
|
||
# 1. Determine Mode (Vertex AI / Gemini Enterprise Agent Platform vs Developer / AI Studio)
|
||
vertexai_flag = config.get("vertexai", False)
|
||
env_vertexai = os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "").lower()
|
||
is_vertexai = vertexai_flag or env_vertexai in ("true", "1")
|
||
|
||
api_version = config.get("api_version")
|
||
|
||
if is_vertexai:
|
||
# Vertex AI Mode (Enterprise Agent Platform) - Uses ADC or Service Accounts
|
||
project = (config.get("vertexai_project") or os.environ.get("GOOGLE_CLOUD_PROJECT", "")).strip()
|
||
location = (config.get("vertexai_location") or os.environ.get("GOOGLE_CLOUD_LOCATION", "")).strip()
|
||
|
||
if not project or not location:
|
||
missing = []
|
||
if not project: missing.append("Project ID (GOOGLE_CLOUD_PROJECT)")
|
||
if not location: missing.append("Location (GOOGLE_CLOUD_LOCATION)")
|
||
err_msg = f"Missing required Vertex AI parameters: {', '.join(missing)}"
|
||
print(f"[ERROR] {err_msg}")
|
||
return IO.NodeOutput(err_msg, cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
|
||
|
||
client_key = ("vertexai", project, location, api_version, proxy_url)
|
||
if client_key not in cls._client_cache:
|
||
cls._client_cache[client_key] = genai.Client(
|
||
vertexai=True,
|
||
project=project,
|
||
location=location,
|
||
http_options=types.HttpOptions(api_version=api_version),
|
||
**client_options
|
||
)
|
||
print(f"[INFO] Created new genai.Client (Vertex AI Mode)")
|
||
else:
|
||
print(f"[INFO] Reusing cached genai.Client (Vertex AI Mode)")
|
||
client = cls._client_cache[client_key]
|
||
|
||
else:
|
||
# Developer Mode (AI Studio) - Uses API Key
|
||
api_key = config.get("api_key")
|
||
client_key = ("developer", api_key, api_version, proxy_url)
|
||
if client_key not in cls._client_cache:
|
||
cls._client_cache[client_key] = genai.Client(
|
||
api_key=api_key,
|
||
http_options=types.HttpOptions(api_version=api_version),
|
||
**client_options
|
||
)
|
||
print(f"[INFO] Created new genai.Client (Developer Mode)")
|
||
else:
|
||
print(f"[INFO] Reusing cached genai.Client (Developer Mode)")
|
||
client = cls._client_cache[client_key]
|
||
|
||
except Exception as e:
|
||
print(f"[ERROR] Gemini client initialization failed: {str(e)}")
|
||
return IO.NodeOutput(f"Gemini client initialization failed: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
|
||
|
||
# Prepare contents (images + prompt)
|
||
images_to_process = []
|
||
if input_image is not None:
|
||
images_to_process.append(input_image)
|
||
if input_image_2 is not None:
|
||
images_to_process.append(input_image_2)
|
||
|
||
if images_to_process:
|
||
try:
|
||
img_parts = []
|
||
for img in images_to_process:
|
||
img_array = img[0].cpu().numpy()
|
||
img_array = (img_array * 255).astype(np.uint8)
|
||
pil_img = Image.fromarray(img_array)
|
||
print(f"[DEBUG] Input image format: {pil_img.mode}")
|
||
img_byte_arr = BytesIO()
|
||
pil_img.save(img_byte_arr, format='PNG')
|
||
img_bytes = img_byte_arr.getvalue()
|
||
if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified":
|
||
img_part = {"inline_data": {"mime_type": "image/png", "data": img_bytes}, "media_resolution": {"level": f"MEDIA_RESOLUTION_{media_resolution.upper()}"}}
|
||
else:
|
||
img_part = {"inline_data": {"mime_type": "image/png", "data": img_bytes}}
|
||
img_parts.append(img_part)
|
||
contents = img_parts + [{"text": padded_prompt}]
|
||
except Exception as e:
|
||
print(f"[ERROR] Error processing input image: {str(e)}")
|
||
return IO.NodeOutput(f"Error processing input image: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
|
||
else:
|
||
contents = padded_prompt
|
||
|
||
# Logic Update: Only request IMAGE modality if include_images is TRUE AND it's the correct model
|
||
# Otherwise we stick to TEXT modality
|
||
response_modalities = ["IMAGE", "TEXT"] if (include_images and model in cls.IMAGE_MODELS) else ["TEXT"]
|
||
|
||
generate_content_config = cls._build_generate_content_config(
|
||
model=model,
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
top_k=top_k,
|
||
max_output_tokens=max_output_tokens,
|
||
seed=seed,
|
||
include_images=include_images,
|
||
response_modalities=response_modalities,
|
||
aspect_ratio=aspect_ratio,
|
||
padded_system_instruction=padded_system_instruction,
|
||
thinking_level=thinking_level,
|
||
thinking_budget=thinking_budget,
|
||
include_thoughts=include_thoughts,
|
||
media_resolution=media_resolution,
|
||
)
|
||
|
||
if use_seed and actual_seed is not None:
|
||
try:
|
||
generate_content_config.seed = actual_seed
|
||
except Exception:
|
||
pass
|
||
|
||
# API call in background thread
|
||
start_time = time.time()
|
||
result_queue: "queue.Queue[Tuple[str, Any]]" = queue.Queue()
|
||
|
||
def api_call():
|
||
last_api_exception = None
|
||
api_response = None
|
||
max_retries = 1
|
||
for attempt in range(max_retries):
|
||
try:
|
||
if use_seed and actual_seed is not None:
|
||
try:
|
||
generate_content_config.seed = actual_seed + attempt
|
||
except Exception:
|
||
pass
|
||
response = client.models.generate_content(model=model, contents=contents, config=generate_content_config)
|
||
if not (response.candidates and getattr(response.candidates[0].content, 'parts', None)):
|
||
finish_reason = "UNKNOWN"
|
||
if response.candidates:
|
||
fr = getattr(response.candidates[0], 'finish_reason', None)
|
||
if fr is not None and hasattr(fr, 'name'):
|
||
finish_reason = fr.name
|
||
raise ValueError(f"Response was empty or blocked (Finish Reason: {finish_reason}).")
|
||
api_response = response
|
||
break
|
||
except Exception as e:
|
||
last_api_exception = e
|
||
if api_response is None:
|
||
result_queue.put(("error", last_api_exception))
|
||
return
|
||
try:
|
||
current_text_output = ""
|
||
current_image_tensor = None
|
||
parts = None
|
||
if api_response.candidates:
|
||
parts = getattr(api_response.candidates[0].content, 'parts', None)
|
||
if parts:
|
||
for part in parts:
|
||
if hasattr(part, 'text') and part.text is not None:
|
||
current_text_output += part.text
|
||
elif hasattr(part, 'inline_data') and part.inline_data is not None:
|
||
try:
|
||
inline_data = part.inline_data
|
||
mime_type = inline_data.mime_type
|
||
data = inline_data.data
|
||
image_path = cls.save_binary_file(data, mime_type)
|
||
img = Image.open(image_path)
|
||
if img.mode != 'RGB':
|
||
img = img.convert('RGB')
|
||
img_array = np.array(img).astype(np.float32) / 255.0
|
||
current_image_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
||
except Exception:
|
||
if current_image_tensor is None:
|
||
current_image_tensor = cls.generate_empty_image()
|
||
if current_image_tensor is None:
|
||
current_image_tensor = cls.generate_empty_image()
|
||
result_queue.put(("success", (current_text_output, current_image_tensor)))
|
||
except Exception as e_proc:
|
||
result_queue.put(("error", e_proc))
|
||
|
||
api_thread = threading.Thread(target=api_call)
|
||
api_thread.daemon = True
|
||
api_thread.start()
|
||
|
||
# Retry loop for pattern matching
|
||
retry_attempt = 0
|
||
retry_needed = True
|
||
while retry_needed:
|
||
retry_needed = False # Will be set to True if pattern matches
|
||
|
||
try:
|
||
status, result = result_queue.get(timeout=timeout)
|
||
elapsed_time = time.time() - start_time
|
||
if status == "success":
|
||
text_output, image_tensor = result
|
||
|
||
# Check if retry pattern matches
|
||
if retry_pattern and max_retries > 0 and retry_attempt < max_retries:
|
||
print(f"[DEBUG] Checking retry pattern '{retry_pattern}' against response (first 200 chars): {text_output[:200]}")
|
||
try:
|
||
compiled_pattern = re.compile(retry_pattern, re.IGNORECASE)
|
||
if compiled_pattern.search(text_output):
|
||
retry_attempt += 1
|
||
print(f"[INFO] Retry pattern matched in response. Retry attempt {retry_attempt}/{max_retries}")
|
||
|
||
# Generate new random gemini seed for retry
|
||
current_time = int(time.time() * 1000)
|
||
random_component = random.randint(0, 1000000)
|
||
actual_seed = (current_time + random_component) % 2147483647
|
||
print(f"[INFO] Retrying with new gemini seed: {actual_seed}")
|
||
|
||
# Update config seed and retry
|
||
if use_seed:
|
||
try:
|
||
generate_content_config.seed = actual_seed
|
||
except Exception:
|
||
pass
|
||
|
||
# Create new queue and thread for retry
|
||
result_queue = queue.Queue()
|
||
start_time = time.time()
|
||
api_thread = threading.Thread(target=api_call)
|
||
api_thread.daemon = True
|
||
api_thread.start()
|
||
retry_needed = True
|
||
continue
|
||
except re.error as regex_err:
|
||
print(f"[WARNING] Invalid retry regex pattern: {regex_err}")
|
||
else:
|
||
error_exception = result
|
||
text_output = f"API call/processing error: {str(error_exception)}"
|
||
|
||
# Check if retry pattern matches error/finish reason
|
||
if retry_pattern and max_retries > 0 and retry_attempt < max_retries:
|
||
print(f"[DEBUG] Checking retry pattern '{retry_pattern}' against error: {text_output}")
|
||
try:
|
||
compiled_pattern = re.compile(retry_pattern, re.IGNORECASE)
|
||
if compiled_pattern.search(text_output):
|
||
retry_attempt += 1
|
||
print(f"[INFO] Retry pattern matched in error. Retry attempt {retry_attempt}/{max_retries}")
|
||
|
||
# Generate new random gemini seed for retry
|
||
current_time = int(time.time() * 1000)
|
||
random_component = random.randint(0, 1000000)
|
||
actual_seed = (current_time + random_component) % 2147483647
|
||
print(f"[INFO] Retrying with new gemini seed: {actual_seed}")
|
||
|
||
# Update config seed and retry
|
||
if use_seed:
|
||
try:
|
||
generate_content_config.seed = actual_seed
|
||
except Exception:
|
||
pass
|
||
|
||
# Create new queue and thread for retry
|
||
result_queue = queue.Queue()
|
||
start_time = time.time()
|
||
api_thread = threading.Thread(target=api_call)
|
||
api_thread.daemon = True
|
||
api_thread.start()
|
||
retry_needed = True
|
||
continue
|
||
except re.error as regex_err:
|
||
print(f"[WARNING] Invalid retry regex pattern: {regex_err}")
|
||
|
||
except queue.Empty:
|
||
text_output = f"Gemini API request/processing timed out, waited {timeout} seconds."
|
||
|
||
except Exception as e:
|
||
print(f"[ERROR] Unhandled error in generate method: {str(e)}")
|
||
text_output = f"Unhandled error: {str(e)}"
|
||
if image_tensor is None:
|
||
image_tensor = cls.generate_empty_image()
|
||
|
||
final_actual_seed = actual_seed if actual_seed is not None else 0
|
||
|
||
# Cache the result
|
||
if use_seed:
|
||
try:
|
||
cls._cache[fingerprint] = (text_output, image_tensor, final_actual_seed)
|
||
except Exception:
|
||
pass
|
||
|
||
# Cache the successful gemini seed for this input seed (if retry was enabled)
|
||
if retry_pattern and max_retries > 0:
|
||
seed_cache_key = (input_seed, fingerprint[:-4])
|
||
cls._seed_map_cache[seed_cache_key] = final_actual_seed
|
||
print(f"[INFO] Cached successful gemini seed {final_actual_seed} for input seed {input_seed}")
|
||
|
||
|
||
try:
|
||
return IO.NodeOutput(text_output, image_tensor, final_actual_seed)
|
||
finally:
|
||
if original_http_proxy:
|
||
os.environ['HTTP_PROXY'] = original_http_proxy
|
||
else:
|
||
if 'HTTP_PROXY' in os.environ:
|
||
os.environ.pop('HTTP_PROXY')
|
||
|
||
if original_https_proxy:
|
||
os.environ['HTTPS_PROXY'] = original_https_proxy
|
||
else:
|
||
if 'HTTPS_PROXY' in os.environ:
|
||
os.environ.pop('HTTPS_PROXY')
|
||
|
||
if original_http_proxy_lower:
|
||
os.environ['http_proxy'] = original_http_proxy_lower
|
||
else:
|
||
if 'http_proxy' in os.environ:
|
||
os.environ.pop('http_proxy')
|
||
|
||
if original_https_proxy_lower:
|
||
os.environ['https_proxy'] = original_https_proxy_lower
|
||
else:
|
||
if 'https_proxy' in os.environ:
|
||
os.environ.pop('https_proxy')
|
||
|
||
if 'REQUESTS_CA_BUNDLE' in os.environ:
|
||
os.environ.pop('REQUESTS_CA_BUNDLE')
|
||
|
||
try:
|
||
import requests
|
||
if hasattr(requests, 'Session'):
|
||
clean_session = requests.Session()
|
||
if hasattr(requests, 'session'):
|
||
requests.session = lambda: clean_session
|
||
except:
|
||
pass
|
||
|
||
# V3 uses ComfyExtension entrypoint in __init__.py to expose nodes |