1816 lines
90 KiB
Python
1816 lines
90 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]
|
||
import comfy.model_management # type: ignore[reportMissingImports]
|
||
from comfy_api.latest import ComfyExtension, UI, IO, Types, VideoFromComponents # type: ignore[reportMissingImports]
|
||
from google import genai
|
||
from google.genai import errors, types
|
||
import time
|
||
import traceback
|
||
import threading
|
||
import queue
|
||
import sys
|
||
import subprocess
|
||
import random
|
||
import hashlib
|
||
import warnings
|
||
from typing import Any, Tuple
|
||
|
||
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_GeminiSafetyConfig(IO.ComfyNode):
|
||
GeminiSafetyConfig = IO.Custom("GEMINI_SAFETY_CONFIG")
|
||
THRESHOLDS = ["off", "block_none", "block_only_high", "block_medium_and_above", "block_low_and_above", "default"]
|
||
|
||
@classmethod
|
||
def define_schema(cls) -> IO.Schema:
|
||
return IO.Schema(
|
||
node_id="SSL_GeminiSafetyConfig",
|
||
display_name="Configure Gemini Safety",
|
||
category="API/Gemini",
|
||
inputs=[
|
||
IO.Combo.Input("harassment", options=cls.THRESHOLDS, default="off", tooltip="Negative or harmful comments targeting identity."),
|
||
IO.Combo.Input("hate_speech", options=cls.THRESHOLDS, default="off", tooltip="Content that promotes violence or incites hatred."),
|
||
IO.Combo.Input("sexually_explicit", options=cls.THRESHOLDS, default="off", tooltip="Contains references to sexual acts or other lewd content."),
|
||
IO.Combo.Input("dangerous_content", options=cls.THRESHOLDS, default="off", tooltip="Promotes, facilitates, or encourages harmful acts."),
|
||
IO.Combo.Input("civic_integrity", options=cls.THRESHOLDS, default="off", tooltip="Content that may mislead or influence elections."),
|
||
IO.Combo.Input("image_harassment", options=cls.THRESHOLDS, default="off", tooltip="Images depicting harassment."),
|
||
IO.Combo.Input("image_hate", options=cls.THRESHOLDS, default="off", tooltip="Images depicting hate speech."),
|
||
IO.Combo.Input("image_sexually_explicit", options=cls.THRESHOLDS, default="off", tooltip="Images containing sexually explicit content."),
|
||
IO.Combo.Input("image_dangerous_content", options=cls.THRESHOLDS, default="off", tooltip="Images depicting dangerous content."),
|
||
IO.Combo.Input("jailbreak", options=cls.THRESHOLDS, default="default", tooltip="Prompts designed to bypass safety filters. 'default' keeps filter disabled."),
|
||
IO.Combo.Input("method", options=["default", "severity", "probability"], default="severity", tooltip="Method for blocking content: severity vs probability score."),
|
||
],
|
||
outputs=[
|
||
cls.GeminiSafetyConfig.Output("safety_config")
|
||
]
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, harassment: str = "off", hate_speech: str = "off",
|
||
sexually_explicit: str = "off", dangerous_content: str = "off",
|
||
civic_integrity: str = "off", image_harassment: str = "off",
|
||
image_hate: str = "off", image_sexually_explicit: str = "off",
|
||
image_dangerous_content: str = "off", jailbreak: str = "default",
|
||
method: str = "severity") -> IO.NodeOutput:
|
||
categories = {
|
||
"harassment": harassment,
|
||
"hate_speech": hate_speech,
|
||
"sexually_explicit": sexually_explicit,
|
||
"dangerous_content": dangerous_content,
|
||
"civic_integrity": civic_integrity,
|
||
"image_harassment": image_harassment,
|
||
"image_hate": image_hate,
|
||
"image_sexually_explicit": image_sexually_explicit,
|
||
"image_dangerous_content": image_dangerous_content,
|
||
"jailbreak": jailbreak,
|
||
}
|
||
settings = []
|
||
for cat, threshold in categories.items():
|
||
if threshold != "default":
|
||
entry: dict[str, Any] = {"type": cat, "threshold": threshold}
|
||
if method != "default":
|
||
entry["method"] = method
|
||
settings.append(entry)
|
||
return IO.NodeOutput(settings)
|
||
|
||
|
||
|
||
class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
|
||
GemConfig = IO.Custom("GEMINI_CONFIG")
|
||
GeminiSafetyConfig = IO.Custom("GEMINI_SAFETY_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("use_vertexai_env", default=False, tooltip="Uses Gemini Enterprise Agent Platform (formerly Vertex AI) with GOOGLE_GENAI_USE_ENTERPRISE or the legacy GOOGLE_GENAI_USE_VERTEXAI environment variable."),
|
||
IO.Boolean.Input("vertexai_express", default=False),
|
||
IO.String.Input("vertexai_project", optional=True),
|
||
IO.String.Input("vertexai_location", optional=True),
|
||
IO.String.Input("google_application_credentials", default="", optional=True, multiline=False, tooltip="Optional absolute path to an application default credentials JSON file for Vertex AI."),
|
||
IO.Boolean.Input("use_cache", default=False, optional=True, tooltip="Reuse unchanged Gemini input context through the API cache after a local result-cache miss."),
|
||
IO.Int.Input("cache_ttl_minutes", default=60, min=1, step=1, optional=True, tooltip="How long Gemini keeps explicitly cached context."),
|
||
IO.Int.Input("cache_seed", default=0, min=0, max=2147483647, control_after_generate=True, optional=True, tooltip="Gemini generation seed used when context caching is enabled."),
|
||
cls.GeminiSafetyConfig.Input("safety_override", optional=True, tooltip="Optional safety override settings from Configure Gemini Safety node."),
|
||
],
|
||
outputs=[
|
||
cls.GemConfig.Output("config")
|
||
]
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, api_key: str, api_version: str, use_vertexai_env: bool, vertexai_express: bool,
|
||
vertexai_project: str | None = "", vertexai_location: str | None = "",
|
||
google_application_credentials: str | None = "", use_cache=False,
|
||
cache_ttl_minutes=60, cache_seed=0,
|
||
safety_override: GeminiSafetyConfig.Type | None = None) -> IO.NodeOutput:
|
||
config = {
|
||
"api_key": api_key,
|
||
"api_version": api_version,
|
||
"use_vertexai_env": use_vertexai_env,
|
||
"vertexai_express": vertexai_express,
|
||
"vertexai_project": vertexai_project,
|
||
"vertexai_location": vertexai_location,
|
||
"google_application_credentials": google_application_credentials,
|
||
"use_cache": use_cache,
|
||
"cache_ttl_minutes": cache_ttl_minutes,
|
||
"cache_seed": cache_seed,
|
||
"safety_override": safety_override,
|
||
}
|
||
return IO.NodeOutput(config)
|
||
|
||
|
||
|
||
class SSL_GeminiVideoConfig(IO.ComfyNode):
|
||
GeminiVideoConfig = IO.Custom("GEMINI_VIDEO_CONFIG")
|
||
|
||
@classmethod
|
||
def define_schema(cls) -> IO.Schema:
|
||
return IO.Schema(
|
||
node_id="SSL_GeminiVideoConfig",
|
||
display_name="Configure Gemini Video Input",
|
||
category="API/Gemini",
|
||
inputs=[
|
||
IO.Video.Input("video"),
|
||
IO.Int.Input("fps", default=1, min=1, max=24, step=1, tooltip="Frames per second sampled by Gemini for video understanding."),
|
||
IO.Boolean.Input("pad_at_start", default=False, tooltip="Prepend black video and silence for half one Gemini sampling interval."),
|
||
IO.Boolean.Input("duration_aware_padding", default=False, tooltip="Use the video's fractional-second duration, rounded down to native video frames, as the start padding duration."),
|
||
],
|
||
outputs=[cls.GeminiVideoConfig.Output("video_config")],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, video: IO.Video.Type, fps: int, pad_at_start: bool = False, duration_aware_padding: bool = False) -> IO.NodeOutput:
|
||
return IO.NodeOutput({
|
||
"video": video,
|
||
"fps": int(fps),
|
||
"pad_at_start": bool(pad_at_start),
|
||
"duration_aware_padding": bool(duration_aware_padding),
|
||
})
|
||
|
||
|
||
class SSL_GeminiTextPrompt(IO.ComfyNode):
|
||
GemConfig = IO.Custom("GEMINI_CONFIG")
|
||
ResponseSchema = IO.Custom("GEMINI_RESPONSE_SCHEMA")
|
||
GeminiVideoConfig = IO.Custom("GEMINI_VIDEO_CONFIG")
|
||
_cache: dict = {}
|
||
_seed_map_cache: dict = {} # Maps (input_seed, fingerprint) -> successful_gemini_seed
|
||
_client_cache: dict = {} # Maps client_key tuple -> genai.Client instance
|
||
_context_cache: dict = {}
|
||
_context_cache_lock = threading.Lock()
|
||
GEMINI_3_7_FLASH = "gemini-3.7-flash"
|
||
GEMINI_3_8_FLASH = "gemini-3.8-flash"
|
||
GEMINI_4_ARGON = "gemini-4-argon"
|
||
GEMINI_4_ARGON_CYBER_PERMISSIVE = "gemini-4-argon-cyber-permissive"
|
||
ARGON_MODELS = (GEMINI_4_ARGON, GEMINI_4_ARGON_CYBER_PERMISSIVE)
|
||
VIDEO_MIME_TYPE = "video/mp4"
|
||
VIDEO_INLINE_LIMIT_BYTES = 100 * 1024 * 1024
|
||
VERTEX_CACHE_INLINE_LIMIT_BYTES = 10 * 1024 * 1024
|
||
|
||
# Placeholder only. Keep this out of the model combo until Google's hosted-model
|
||
# documentation confirms the final ID and gemini-3.6-flash-equivalent capabilities
|
||
# for both Gemini Developer API and Enterprise Agent Platform.
|
||
GEMINI_4_FLASH_PREVIEW = "gemini-4-flash-preview"
|
||
|
||
# Define model lists centrally to ensure consistency between cache logic and execution logic
|
||
THINKING_MODELS = [
|
||
"gemini-1.5-pro-002", "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-pro-exp-03-25",
|
||
"gemini-3-flash-preview", "gemini-3.1-pro-preview", "gemini-3.5-flash", "gemini-3.5-flash-lite", "gemini-3.6-flash", GEMINI_3_7_FLASH, GEMINI_3_8_FLASH, GEMINI_4_ARGON, GEMINI_4_ARGON_CYBER_PERMISSIVE, GEMINI_4_FLASH_PREVIEW, "gemini-pro-latest", "gemini-flash-latest", "gemini-flash-lite-latest"
|
||
]
|
||
GEN3_THINKING_MODELS = [
|
||
"gemini-pro-latest", "gemini-flash-latest", "gemini-3.1-pro-preview",
|
||
"gemini-3-flash-preview", "gemini-3.5-flash", "gemini-3.5-flash-lite", "gemini-3.6-flash", GEMINI_3_7_FLASH, GEMINI_3_8_FLASH, GEMINI_4_ARGON, GEMINI_4_ARGON_CYBER_PERMISSIVE, GEMINI_4_FLASH_PREVIEW
|
||
]
|
||
IMAGE_MODELS = [
|
||
"gemini-2.5-flash-image",
|
||
"gemini-3.1-flash-image", "gemini-3.1-flash-lite-image",
|
||
"gemini-3-pro-image-preview", "gemini-3-pro-image",
|
||
]
|
||
IMAGE_SIZE_BY_SELECTABLE_MODEL = {
|
||
"gemini-3.1-flash-image": ("512", "1K", "2K", "4K"),
|
||
"gemini-3.1-flash-lite-image": ("1K",),
|
||
"gemini-3-pro-image": ("1K", "2K", "4K"),
|
||
}
|
||
MEDIA_RES_MODELS = [
|
||
"gemini-3.1-flash-lite", "gemini-3-flash-preview", "gemini-3.1-pro-preview",
|
||
"gemini-3.5-flash", "gemini-pro-latest", "gemini-flash-latest", "gemini-flash-lite-latest"
|
||
]
|
||
|
||
@classmethod
|
||
def define_schema(cls) -> IO.Schema:
|
||
image_names = [f"image_{index}" for index in range(1, 101)]
|
||
image_template = IO.Autogrow.TemplateNames(
|
||
IO.Image.Input("image_1", optional=True),
|
||
names=image_names,
|
||
min=0,
|
||
)
|
||
return IO.Schema(
|
||
node_id="SSL_GeminiTextPrompt",
|
||
display_name="Expanded Gemini Text/Image",
|
||
category="API/Gemini",
|
||
is_input_list=True,
|
||
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-1.5-pro-002", "gemini-2.0-flash", "gemini-2.0-flash-lite", "gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-exp-03-25", "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite", "gemini-3-flash-preview", "gemini-3.1-flash-lite", "gemini-3.1-pro-preview", "gemini-3.5-flash", "gemini-3.5-flash-lite", "gemini-3.6-flash", cls.GEMINI_3_7_FLASH, cls.GEMINI_3_8_FLASH, cls.GEMINI_4_ARGON, cls.GEMINI_4_ARGON_CYBER_PERMISSIVE, "gemini-3.1-flash-image", "gemini-3.1-flash-lite-image", "gemini-3-pro-image", "gemini-pro-latest", "gemini-flash-latest", "gemini-flash-lite-latest"], 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.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. Gemini 3.7/3.8 Flash always uses this control, defaults to medium, and ignores thinking_budget."),
|
||
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."),
|
||
IO.String.Input("timeout_fallback_text", default="", optional=True, multiline=True, tooltip="Text returned when the Gemini request times out. Leave empty to return the standard timeout message."),
|
||
IO.Combo.Input("image_size", options=["None", "512", "1K", "2K", "4K"], default="None", optional=True, tooltip="Generated image resolution. Gemini 3.1 Flash Lite Image supports only 1K; 512 is supported only by Gemini 3.1 Flash Image."),
|
||
cls.GeminiVideoConfig.Input("video", optional=True, tooltip="Optional configured Gemini video input with embedded audio and sampling FPS."),
|
||
cls.ResponseSchema.Input("response_schema", display_name="Response format", optional=True,
|
||
tooltip="Connect an answer format to request the fields you defined. Leave disconnected for a normal answer."),
|
||
IO.Autogrow.Input(
|
||
"image_inputs",
|
||
template=image_template,
|
||
optional=True,
|
||
tooltip=(
|
||
"Ordered Gemini image parts growing from image_1 through image_100. Images inside a batch "
|
||
"are sent consecutively before the next socket. Provider request-size and model-specific "
|
||
"reference limits still apply."
|
||
),
|
||
),
|
||
],
|
||
outputs=[
|
||
IO.String.Output("text", is_output_list=True),
|
||
IO.Image.Output("image", is_output_list=True),
|
||
IO.Int.Output("final_actual_seed", is_output_list=True),
|
||
IO.String.Output("structured_output", display_name="Structured output", is_output_list=True),
|
||
IO.String.Output("thoughts", display_name="Thoughts", is_output_list=True),
|
||
]
|
||
)
|
||
|
||
@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 _ordered_image_frames(cls, image_inputs: IO.Autogrow.Type | None):
|
||
if not image_inputs:
|
||
return [], []
|
||
|
||
def image_index(name):
|
||
match = re.fullmatch(r"image_(\d+)", name)
|
||
return int(match.group(1)) if match else 101
|
||
|
||
frames = []
|
||
batch_counts = []
|
||
for socket_name, image_batch in sorted(
|
||
image_inputs.items(), key=lambda item: (image_index(item[0]), item[0])
|
||
):
|
||
if image_batch is None:
|
||
continue
|
||
if isinstance(image_batch, list):
|
||
image_batch = image_batch[0] if image_batch else None
|
||
if image_batch is None:
|
||
continue
|
||
if not torch.is_tensor(image_batch) or image_batch.ndim != 4:
|
||
raise ValueError(
|
||
f"Gemini {socket_name} must be a BHWC image tensor."
|
||
)
|
||
batch_size = int(image_batch.shape[0])
|
||
batch_counts.append((socket_name, batch_size))
|
||
frames.extend(
|
||
(socket_name, index, image_batch[index])
|
||
for index in range(batch_size)
|
||
)
|
||
return frames, batch_counts
|
||
|
||
@classmethod
|
||
def _prepare_image_inputs(cls, image_inputs: IO.Autogrow.Type | None) -> dict[str, Any] | None:
|
||
if not image_inputs:
|
||
return None
|
||
image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs)
|
||
if not image_frames:
|
||
return {
|
||
"image_hashes": (),
|
||
"png_bytes_list": [],
|
||
"png_hashes": [],
|
||
"batch_counts": [],
|
||
}
|
||
tensor_hashes = []
|
||
png_bytes_list = []
|
||
png_hashes = []
|
||
for _, _, image_frame in image_frames:
|
||
np_frame = np.ascontiguousarray(image_frame.cpu().numpy())
|
||
tensor_hashes.append(hashlib.sha256(np_frame.tobytes()).hexdigest())
|
||
img_array = (np_frame * 255).astype(np.uint8)
|
||
pil_img = Image.fromarray(img_array)
|
||
img_byte_arr = BytesIO()
|
||
pil_img.save(img_byte_arr, format='PNG')
|
||
img_bytes = img_byte_arr.getvalue()
|
||
png_bytes_list.append(img_bytes)
|
||
png_hashes.append(hashlib.sha256(img_bytes).hexdigest())
|
||
batch_summary = ", ".join(
|
||
f"{socket_name}={batch_size}"
|
||
for socket_name, batch_size in image_batch_counts
|
||
)
|
||
print(
|
||
f"[INFO] Prepared Gemini image inputs: {batch_summary}, "
|
||
f"total={len(png_bytes_list)}"
|
||
)
|
||
return {
|
||
"image_hashes": tuple(tensor_hashes),
|
||
"png_bytes_list": png_bytes_list,
|
||
"png_hashes": png_hashes,
|
||
"batch_counts": image_batch_counts,
|
||
}
|
||
|
||
@classmethod
|
||
def _serialize_video(cls, video: IO.Video.Type | None, fps: int | None = None,
|
||
pad_at_start: bool = False, duration_aware_padding: bool = False):
|
||
if video is None:
|
||
return None, None
|
||
|
||
video_to_save = video
|
||
if pad_at_start:
|
||
if fps is None:
|
||
raise ValueError("Gemini video FPS is required when padding at start.")
|
||
components = video.get_components()
|
||
native_frame_rate = components.frame_rate
|
||
source_frame_rate = float(native_frame_rate)
|
||
source_images = components.images
|
||
source_frame_count = source_images.shape[0]
|
||
if duration_aware_padding:
|
||
source_duration = video.get_duration()
|
||
target_frame_rate = native_frame_rate
|
||
padding_frame_count = int((source_duration % 1.0) * source_frame_rate)
|
||
resampled_images = source_images
|
||
else:
|
||
source_duration = source_frame_count / source_frame_rate
|
||
target_frame_rate = fps * 2
|
||
padding_frame_count = 1
|
||
target_frame_count = max(1, round(source_duration * target_frame_rate))
|
||
source_indices = (
|
||
torch.arange(target_frame_count, device=source_images.device, dtype=torch.float64)
|
||
* source_frame_rate
|
||
/ target_frame_rate
|
||
).floor().long().clamp(max=source_frame_count - 1)
|
||
resampled_images = source_images.index_select(0, source_indices)
|
||
black_frames = torch.zeros_like(resampled_images[:1]).expand(padding_frame_count, -1, -1, -1)
|
||
padded_images = torch.cat((black_frames, resampled_images), dim=0)
|
||
|
||
padded_audio = components.audio
|
||
if padded_audio:
|
||
sample_rate = int(padded_audio["sample_rate"])
|
||
waveform = padded_audio["waveform"]
|
||
silence_samples = round(sample_rate * padding_frame_count / target_frame_rate)
|
||
silence = torch.zeros(
|
||
(*waveform.shape[:-1], silence_samples),
|
||
dtype=waveform.dtype,
|
||
device=waveform.device,
|
||
)
|
||
padded_audio = {
|
||
**padded_audio,
|
||
"waveform": torch.cat((silence, waveform), dim=-1),
|
||
}
|
||
|
||
video_to_save = VideoFromComponents(
|
||
Types.VideoComponents(
|
||
images=padded_images,
|
||
audio=padded_audio,
|
||
frame_rate=target_frame_rate,
|
||
),
|
||
bit_depth=video.get_bit_depth(),
|
||
color_space=video.get_color_space(),
|
||
)
|
||
|
||
buffer = BytesIO()
|
||
video_to_save.save_to(
|
||
buffer,
|
||
format=Types.VideoContainer.MP4,
|
||
codec=Types.VideoCodec.H264,
|
||
)
|
||
video_bytes = buffer.getvalue()
|
||
if len(video_bytes) >= cls.VIDEO_INLINE_LIMIT_BYTES:
|
||
raise ValueError("Gemini video input must be smaller than 100 MB after MP4/H.264 normalization.")
|
||
return video_bytes, cls.VIDEO_MIME_TYPE
|
||
|
||
@classmethod
|
||
def _resolve_gemini_3_7_thinking_level(cls, thinking_level, model_name=None):
|
||
if thinking_level is None or thinking_level == "None":
|
||
if model_name == cls.GEMINI_4_ARGON_CYBER_PERMISSIVE:
|
||
return "high"
|
||
return "medium"
|
||
if thinking_level == "minimal":
|
||
if model_name in (cls.GEMINI_3_8_FLASH, cls.GEMINI_4_ARGON, cls.GEMINI_4_ARGON_CYBER_PERMISSIVE):
|
||
raise ValueError(
|
||
f"{model_name} thinking_level must be low, medium, or high."
|
||
)
|
||
print(
|
||
f"[WARNING] {model_name or 'Gemini 3.7 Flash'} does not support minimal "
|
||
"thinking; using low instead."
|
||
)
|
||
return "low"
|
||
if thinking_level not in {"low", "medium", "high"}:
|
||
raise ValueError(
|
||
f"{model_name or 'Gemini 3.7 Flash'} thinking_level must be low, medium, or high."
|
||
)
|
||
return thinking_level
|
||
|
||
CORE_HARM_CATEGORIES = [
|
||
"harassment",
|
||
"hate_speech",
|
||
"sexually_explicit",
|
||
"dangerous_content",
|
||
"civic_integrity",
|
||
"image_harassment",
|
||
"image_hate",
|
||
"image_sexually_explicit",
|
||
"image_dangerous_content",
|
||
]
|
||
|
||
@classmethod
|
||
def _resolve_safety_settings(cls, config: dict[str, Any]) -> list[types.SafetySetting]:
|
||
override = config.get("safety_override")
|
||
if override and isinstance(override, list) and override:
|
||
settings = []
|
||
for entry in override:
|
||
if isinstance(entry, types.SafetySetting):
|
||
settings.append(entry)
|
||
elif isinstance(entry, dict):
|
||
cat_name = entry.get("type") or entry.get("category", "")
|
||
cat_str = str(cat_name)
|
||
cat_enum = cat_str if cat_str.startswith("HARM_CATEGORY_") else f"HARM_CATEGORY_{cat_str.upper()}"
|
||
thresh = entry.get("threshold", "BLOCK_NONE")
|
||
thresh_enum = str(thresh).upper()
|
||
method = entry.get("method")
|
||
kwargs: dict[str, Any] = {
|
||
"category": cat_enum,
|
||
"threshold": thresh_enum,
|
||
}
|
||
if method and method != "default":
|
||
kwargs["method"] = str(method).upper()
|
||
settings.append(types.SafetySetting(**kwargs))
|
||
if settings:
|
||
return settings
|
||
|
||
is_enterprise = bool(config.get("use_vertexai_env", False) or config.get("vertexai_express", False))
|
||
if is_enterprise:
|
||
return [
|
||
types.SafetySetting(
|
||
category=f"HARM_CATEGORY_{cat.upper()}",
|
||
threshold="OFF",
|
||
method="SEVERITY",
|
||
)
|
||
for cat in cls.CORE_HARM_CATEGORIES
|
||
]
|
||
return [
|
||
types.SafetySetting(
|
||
category=f"HARM_CATEGORY_{cat.upper()}",
|
||
threshold="BLOCK_NONE",
|
||
)
|
||
for cat in cls.CORE_HARM_CATEGORIES
|
||
]
|
||
|
||
@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,
|
||
video_hash=None, video_mime_type=None, video_fps=None,
|
||
image_inputs: IO.Autogrow.Type | None = 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, use_cache=False,
|
||
cache_ttl_minutes=60, cache_seed=0, image_size="None", schema_identity=None,
|
||
precomputed_image_hashes=None):
|
||
|
||
# 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"
|
||
|
||
if precomputed_image_hashes is not None:
|
||
image_hashes = precomputed_image_hashes
|
||
else:
|
||
image_frames, _ = cls._ordered_image_frames(image_inputs)
|
||
image_hashes = tuple(get_tensor_hash(frame) for _, _, frame in image_frames)
|
||
|
||
# 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_image_size = "None"
|
||
eff_thinking_level = "None"
|
||
eff_thinking_budget = -1
|
||
eff_include_thoughts = False
|
||
eff_temperature = None if model in (cls.GEMINI_3_7_FLASH, cls.GEMINI_3_8_FLASH, cls.GEMINI_4_ARGON, cls.GEMINI_4_ARGON_CYBER_PERMISSIVE) else float(temperature)
|
||
eff_top_p = None if model in (cls.GEMINI_3_7_FLASH, cls.GEMINI_3_8_FLASH, cls.GEMINI_4_ARGON, cls.GEMINI_4_ARGON_CYBER_PERMISSIVE) else float(top_p)
|
||
eff_top_k = None if model in (cls.GEMINI_3_7_FLASH, cls.GEMINI_3_8_FLASH, cls.GEMINI_4_ARGON, cls.GEMINI_4_ARGON_CYBER_PERMISSIVE) else int(top_k)
|
||
|
||
# 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)
|
||
eff_image_size = "None" if image_size in (None, "None") else str(image_size)
|
||
# When generating images, thinking params are ignored
|
||
|
||
elif model in (cls.GEMINI_3_7_FLASH, cls.GEMINI_3_8_FLASH, cls.GEMINI_4_ARGON, cls.GEMINI_4_ARGON_CYBER_PERMISSIVE):
|
||
eff_thinking_level = cls._resolve_gemini_3_7_thinking_level(thinking_level, model)
|
||
eff_include_thoughts = include_thoughts
|
||
|
||
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
|
||
|
||
sanitized_config = dict(config)
|
||
for cache_setting in ("use_cache", "cache_ttl_minutes", "cache_seed"):
|
||
sanitized_config.pop(cache_setting, None)
|
||
api_key = sanitized_config.get("api_key")
|
||
if api_key:
|
||
sanitized_config["api_key"] = f"sha256:{hashlib.sha256(str(api_key).encode('utf-8')).hexdigest()}"
|
||
|
||
fingerprint = (
|
||
str(sanitized_config),
|
||
prompt,
|
||
system_instruction,
|
||
model,
|
||
eff_temperature,
|
||
eff_top_p,
|
||
eff_top_k,
|
||
int(max_output_tokens),
|
||
eff_include_images, # EFFECTIVE include_images
|
||
eff_aspect_ratio, # EFFECTIVE aspect_ratio
|
||
eff_image_size, # EFFECTIVE image_size
|
||
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
|
||
bool(use_cache),
|
||
int(cache_ttl_minutes) if use_cache else 0,
|
||
int(cache_seed) if use_cache and use_seed else 0,
|
||
video_hash,
|
||
video_mime_type,
|
||
video_fps,
|
||
image_hashes,
|
||
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
|
||
schema_identity,
|
||
)
|
||
|
||
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):
|
||
if not use_seed:
|
||
print("[INFO] Seed not used")
|
||
return None
|
||
if 0 < seed < 2**31:
|
||
print(f"[INFO] Using specified seed: {seed}")
|
||
return seed
|
||
|
||
generator = torch.Generator(device="cpu")
|
||
if seed == 0:
|
||
generator.seed()
|
||
else:
|
||
generator.manual_seed(seed)
|
||
actual_seed = torch.randint(
|
||
0, 2147483647, (), generator=generator, device="cpu", dtype=torch.int64
|
||
).item()
|
||
print(f"[INFO] Generated API seed: {actual_seed}")
|
||
return actual_seed
|
||
|
||
@classmethod
|
||
def _build_proxy_url(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}"
|
||
|
||
print(f"[INFO] Proxy enabled: {proxy_url}")
|
||
return proxy_url
|
||
|
||
@staticmethod
|
||
def _parse_boolean_environment_variable(name: str) -> bool | None:
|
||
value = os.getenv(name)
|
||
if value is None:
|
||
return None
|
||
|
||
normalized = value.strip().lower()
|
||
if normalized not in {"true", "1", "false", "0"}:
|
||
raise ValueError(f"{name} must be true, false, 1, or 0")
|
||
return normalized in {"true", "1"}
|
||
|
||
@classmethod
|
||
def _resolve_enterprise_environment(cls) -> bool:
|
||
enterprise = cls._parse_boolean_environment_variable("GOOGLE_GENAI_USE_ENTERPRISE")
|
||
legacy_vertex = cls._parse_boolean_environment_variable("GOOGLE_GENAI_USE_VERTEXAI")
|
||
|
||
if enterprise is not None and legacy_vertex is not None and enterprise != legacy_vertex:
|
||
warnings.warn(
|
||
"GOOGLE_GENAI_USE_ENTERPRISE and GOOGLE_GENAI_USE_VERTEXAI conflict; "
|
||
"GOOGLE_GENAI_USE_ENTERPRISE takes precedence.",
|
||
UserWarning,
|
||
stacklevel=2,
|
||
)
|
||
|
||
if enterprise is not None:
|
||
return enterprise
|
||
if legacy_vertex is not None:
|
||
return legacy_vertex
|
||
return True
|
||
|
||
@staticmethod
|
||
def _build_http_options(api_version: str, proxy_url: str | None):
|
||
client_args = {"proxy": proxy_url} if proxy_url else None
|
||
return types.HttpOptions(api_version=api_version, client_args=client_args)
|
||
|
||
@staticmethod
|
||
def _build_context_cache_key(client_key, model, system_instruction, prompt,
|
||
video_hash, video_mime_type, video_fps, image_hashes, media_resolution):
|
||
return (
|
||
client_key,
|
||
model,
|
||
system_instruction,
|
||
prompt,
|
||
video_hash,
|
||
video_mime_type,
|
||
video_fps,
|
||
tuple(image_hashes),
|
||
str(media_resolution),
|
||
)
|
||
|
||
@staticmethod
|
||
def _context_cache_expiry(cache, fallback_ttl_minutes):
|
||
expire_time = getattr(cache, "expire_time", None)
|
||
if expire_time is not None:
|
||
return expire_time.timestamp()
|
||
return time.time() + (int(fallback_ttl_minutes) * 60)
|
||
|
||
@classmethod
|
||
def _get_or_create_context_cache(cls, client, context_key, model, contents,
|
||
system_instruction, cache_ttl_minutes):
|
||
ttl = f"{int(cache_ttl_minutes) * 60}s"
|
||
with cls._context_cache_lock:
|
||
entry = cls._context_cache.get(context_key)
|
||
if entry is not None and entry["expires_at"] > time.time() + 1:
|
||
if entry["ttl_minutes"] == int(cache_ttl_minutes):
|
||
print(f"[INFO] Reusing Gemini context cache {entry['name']}")
|
||
return entry["name"]
|
||
try:
|
||
updated = client.caches.update(
|
||
name=entry["name"],
|
||
config=types.UpdateCachedContentConfig(ttl=ttl),
|
||
)
|
||
entry = {
|
||
"name": updated.name,
|
||
"expires_at": cls._context_cache_expiry(updated, cache_ttl_minutes),
|
||
"ttl_minutes": int(cache_ttl_minutes),
|
||
}
|
||
cls._context_cache[context_key] = entry
|
||
print(f"[INFO] Updated Gemini context cache TTL to {cache_ttl_minutes} minutes")
|
||
return entry["name"]
|
||
except Exception:
|
||
cls._context_cache.pop(context_key, None)
|
||
elif entry is not None:
|
||
cls._context_cache.pop(context_key, None)
|
||
|
||
display_hash = hashlib.sha256(repr(context_key).encode("utf-8")).hexdigest()[:16]
|
||
created = client.caches.create(
|
||
model=model,
|
||
config=types.CreateCachedContentConfig(
|
||
contents=contents,
|
||
system_instruction=system_instruction,
|
||
display_name=f"comfyui-{display_hash}",
|
||
ttl=ttl,
|
||
),
|
||
)
|
||
if not created.name:
|
||
raise ValueError("Gemini did not return a context cache name.")
|
||
cls._context_cache[context_key] = {
|
||
"name": created.name,
|
||
"expires_at": cls._context_cache_expiry(created, cache_ttl_minutes),
|
||
"ttl_minutes": int(cache_ttl_minutes),
|
||
}
|
||
print(f"[INFO] Created Gemini context cache {created.name}")
|
||
return created.name
|
||
|
||
@classmethod
|
||
def _invalidate_context_cache(cls, context_key):
|
||
with cls._context_cache_lock:
|
||
cls._context_cache.pop(context_key, None)
|
||
|
||
@staticmethod
|
||
def _is_missing_context_cache_error(error):
|
||
if isinstance(error, errors.APIError) and error.status == 404:
|
||
return True
|
||
message = str(error).lower()
|
||
cache_reference = any(term in message for term in ("cached content", "cached_content", "cachedcontents"))
|
||
missing_reference = any(term in message for term in ("expired", "invalid", "not found"))
|
||
return cache_reference and missing_reference
|
||
|
||
@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,
|
||
image_size="None", allow_all_people=False, safety_settings=None):
|
||
# Centralized builder for GenerateContentConfig used by different model/feature branches
|
||
safety = safety_settings if safety_settings is not None else [
|
||
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:
|
||
image_config = {}
|
||
if aspect_ratio not in (None, "None"):
|
||
image_config["aspect_ratio"] = aspect_ratio
|
||
if image_size not in (None, "None"):
|
||
supported_image_sizes = cls.IMAGE_SIZE_BY_SELECTABLE_MODEL.get(model)
|
||
if supported_image_sizes is not None and image_size not in supported_image_sizes:
|
||
allowed = ", ".join(supported_image_sizes)
|
||
raise ValueError(f"{model} image_size must be one of: {allowed}.")
|
||
image_config["image_size"] = image_size
|
||
if allow_all_people:
|
||
image_config["person_generation"] = "ALLOW_ALL"
|
||
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(**image_config) if image_config else None,
|
||
system_instruction=[types.Part.from_text(text=padded_system_instruction)],
|
||
)
|
||
|
||
if model in (cls.GEMINI_3_7_FLASH, cls.GEMINI_3_8_FLASH, cls.GEMINI_4_ARGON, cls.GEMINI_4_ARGON_CYBER_PERMISSIVE):
|
||
return types.GenerateContentConfig(
|
||
max_output_tokens=max_output_tokens,
|
||
safety_settings=safety,
|
||
thinking_config=types.ThinkingConfig(
|
||
include_thoughts=include_thoughts,
|
||
thinking_level=cls._resolve_gemini_3_7_thinking_level(thinking_level, model),
|
||
),
|
||
response_modalities=response_modalities,
|
||
system_instruction=[types.Part.from_text(text=padded_system_instruction)],
|
||
)
|
||
|
||
G3Pro = ["gemini-3.1-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)],
|
||
)
|
||
|
||
@staticmethod
|
||
def _reject_json_constant(value):
|
||
raise ValueError(f"Structured response contains invalid JSON constant: {value}.")
|
||
|
||
@classmethod
|
||
def _check_interrupt(cls):
|
||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||
|
||
@classmethod
|
||
def _wait_for_result(cls, result_queue: "queue.Queue[Tuple[str, Any]]", timeout: float) -> Tuple[str, Any]:
|
||
deadline = time.monotonic() + float(timeout)
|
||
while True:
|
||
cls._check_interrupt()
|
||
remaining = deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
raise queue.Empty()
|
||
try:
|
||
return result_queue.get(timeout=min(0.1, remaining))
|
||
except queue.Empty:
|
||
if time.monotonic() >= deadline:
|
||
raise
|
||
|
||
@classmethod
|
||
def _execute_single(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
|
||
include_images, aspect_ratio, bypass_mode, thinking_budget,
|
||
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, timeout_fallback_text="",
|
||
image_size="None",
|
||
video: GeminiVideoConfig.Type | None = None,
|
||
image_inputs: IO.Autogrow.Type | None = None,
|
||
response_schema: ResponseSchema.Type | None = None,
|
||
prepared_images: dict[str, Any] | None = None,
|
||
precomputed_video: Tuple[bytes | None, str | None, str | None, int | None] | None = None) -> IO.NodeOutput:
|
||
|
||
cls._check_interrupt()
|
||
print(f"[INFO] SSL_GeminiTextPrompt execute called, model: {model}")
|
||
use_cache = bool(config.get("use_cache", False))
|
||
cache_ttl_minutes = int(config.get("cache_ttl_minutes", 60))
|
||
cache_seed = int(config.get("cache_seed", 0))
|
||
schema_identity = None
|
||
schema_snapshot = None
|
||
if response_schema is not None:
|
||
if not isinstance(response_schema, dict):
|
||
raise ValueError("Response format must be a schema dictionary from the schema output.")
|
||
if include_images and model in cls.IMAGE_MODELS:
|
||
raise ValueError("Response format cannot be combined with image generation. Disable include_images.")
|
||
schema_identity = json.dumps(response_schema, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
|
||
schema_snapshot = json.loads(schema_identity)
|
||
if precomputed_video is not None:
|
||
video_bytes, video_mime_type, video_hash, video_fps = precomputed_video
|
||
else:
|
||
video_input = video.get("video") if video is not None else None
|
||
video_fps = int(video.get("fps", 1)) if video is not None else None
|
||
video_pad_at_start = bool(video.get("pad_at_start", False)) if video is not None else False
|
||
video_duration_aware_padding = bool(video.get("duration_aware_padding", False)) if video is not None else False
|
||
try:
|
||
video_bytes, video_mime_type = cls._serialize_video(
|
||
video_input,
|
||
video_fps,
|
||
video_pad_at_start,
|
||
video_duration_aware_padding,
|
||
)
|
||
except Exception as e:
|
||
print(f"[ERROR] Error processing input video: {e}")
|
||
output_seed = cache_seed if use_cache and use_seed else seed if use_seed else 0
|
||
return IO.NodeOutput(f"Error processing input video: {e}", cls.generate_empty_image(), output_seed, "", "")
|
||
video_hash = hashlib.sha256(video_bytes).hexdigest() if video_bytes is not None else None
|
||
precomputed_image_hashes = prepared_images["image_hashes"] if prepared_images is not None else None
|
||
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,
|
||
video_hash, video_mime_type, video_fps,
|
||
image_inputs,
|
||
use_proxy, proxy_host, proxy_port, timeout,
|
||
include_thoughts, thinking_level, media_resolution,
|
||
retry_pattern, max_retries, use_cache, cache_ttl_minutes, cache_seed,
|
||
image_size=image_size,
|
||
schema_identity=schema_identity,
|
||
precomputed_image_hashes=precomputed_image_hashes,
|
||
)
|
||
|
||
if cached is not None:
|
||
print(f"[INFO] Returning cached result for fingerprint {fingerprint}")
|
||
return IO.NodeOutput(*cached)
|
||
|
||
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)
|
||
|
||
api_seed = cache_seed if use_cache else seed
|
||
actual_seed = cls._handle_seed(use_seed, api_seed)
|
||
input_seed = api_seed # Store the original API seed for retry cache key
|
||
|
||
# Check if we have a cached successful gemini seed for this input seed
|
||
seed_cache_key = (input_seed, fingerprint)
|
||
if use_seed and retry_pattern and max_retries > 0:
|
||
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 = cls._handle_seed(True, cached_gemini_seed)
|
||
|
||
text_output = ""
|
||
thoughts_output = ""
|
||
structured_output = ""
|
||
is_success = False
|
||
image_tensor = cls.generate_empty_image()
|
||
proxy_url: str | None = None
|
||
|
||
try:
|
||
if use_proxy:
|
||
proxy_url = cls._build_proxy_url(proxy_host, proxy_port)
|
||
|
||
try:
|
||
vertexai_express = config.get("vertexai_express", False)
|
||
use_vertexai_env = config.get("use_vertexai_env", False)
|
||
api_version = config.get("api_version")
|
||
project = config.get("vertexai_project")
|
||
location = config.get("vertexai_location")
|
||
credentials_path = config.get("google_application_credentials")
|
||
http_options = cls._build_http_options(api_version, proxy_url)
|
||
|
||
if use_vertexai_env:
|
||
try:
|
||
env_use = cls._resolve_enterprise_environment()
|
||
if not env_use:
|
||
raise ValueError(
|
||
"GOOGLE_GENAI_USE_ENTERPRISE or GOOGLE_GENAI_USE_VERTEXAI "
|
||
"must be enabled when use_vertexai_env is true"
|
||
)
|
||
|
||
credentials = None
|
||
credentials_cache_key = None
|
||
credential_project = None
|
||
if credentials_path:
|
||
import google.auth
|
||
|
||
absolute_credentials_path = os.path.abspath(os.path.expanduser(credentials_path))
|
||
if not os.path.isfile(absolute_credentials_path):
|
||
raise ValueError(f"Application default credentials file not found: {absolute_credentials_path}")
|
||
credentials, credential_project = google.auth.load_credentials_from_file(
|
||
absolute_credentials_path,
|
||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||
)
|
||
credentials_cache_key = absolute_credentials_path
|
||
|
||
env_project = os.getenv("GOOGLE_CLOUD_PROJECT")
|
||
env_proj = env_project.strip() if env_project else (project or credential_project)
|
||
assert env_proj, "GOOGLE_CLOUD_PROJECT is empty"
|
||
|
||
env_location = os.getenv("GOOGLE_CLOUD_LOCATION")
|
||
env_loc = env_location.strip() if env_location else location
|
||
assert env_loc, "GOOGLE_CLOUD_LOCATION is empty"
|
||
|
||
client_key = ("vertexai_env", env_use, env_proj, env_loc, api_version, proxy_url, credentials_cache_key)
|
||
if client_key not in cls._client_cache:
|
||
cls._client_cache[client_key] = genai.Client(
|
||
enterprise=env_use,
|
||
credentials=credentials,
|
||
project=env_proj,
|
||
location=env_loc,
|
||
http_options=http_options,
|
||
)
|
||
print(f"[INFO] Created new genai.Client (vertexai_env)")
|
||
is_new_client = True
|
||
else:
|
||
print(f"[INFO] Reusing cached genai.Client (vertexai_env)")
|
||
is_new_client = False
|
||
client = cls._client_cache[client_key]
|
||
|
||
except KeyError as e:
|
||
print(f"Missing required environment variable: {e}")
|
||
return IO.NodeOutput(f"Missing environment variable: {e}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0, "", "")
|
||
|
||
except AssertionError as e:
|
||
print(f"Error: {e}")
|
||
return IO.NodeOutput(f"Invalid environment variable: {e}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0, "", "")
|
||
|
||
except ValueError as e:
|
||
print(f"Error: {e}")
|
||
return IO.NodeOutput(f"Invalid Enterprise/Vertex AI configuration: {e}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0, "", "")
|
||
|
||
elif vertexai_express:
|
||
api_key = config.get("api_key")
|
||
if not api_key:
|
||
return IO.NodeOutput(
|
||
"Invalid Enterprise express configuration: API key is required",
|
||
cls.generate_empty_image(),
|
||
actual_seed if actual_seed is not None else 0,
|
||
"", "",
|
||
)
|
||
api_key_hash = hashlib.sha256(str(api_key).encode("utf-8")).hexdigest() if api_key else None
|
||
client_key = ("vertexai_express", api_key_hash, api_version, proxy_url)
|
||
if client_key not in cls._client_cache:
|
||
cls._client_cache[client_key] = genai.Client(
|
||
enterprise=True,
|
||
api_key=api_key,
|
||
http_options=http_options,
|
||
)
|
||
print(f"[INFO] Created new genai.Client (vertexai_express)")
|
||
is_new_client = True
|
||
else:
|
||
print(f"[INFO] Reusing cached genai.Client (vertexai_express)")
|
||
is_new_client = False
|
||
client = cls._client_cache[client_key]
|
||
|
||
else:
|
||
api_key = config.get("api_key")
|
||
api_key_hash = hashlib.sha256(str(api_key).encode("utf-8")).hexdigest() if api_key else None
|
||
client_key = ("standard", api_key_hash, 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=http_options,
|
||
)
|
||
print(f"[INFO] Created new genai.Client (standard)")
|
||
is_new_client = True
|
||
else:
|
||
print(f"[INFO] Reusing cached genai.Client (standard)")
|
||
is_new_client = False
|
||
client = cls._client_cache[client_key]
|
||
|
||
if is_new_client:
|
||
try:
|
||
if hasattr(client, '_api_client') and hasattr(client._api_client, '_access_token'):
|
||
print("[INFO] Pre-fetching auth token to avoid timeout interference...")
|
||
client._api_client._access_token()
|
||
except Exception as auth_e:
|
||
print(f"[WARNING] Pre-auth check failed (will attempt during generation): {auth_e}")
|
||
|
||
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 + video + prompt)
|
||
context_image_hashes = []
|
||
media_parts = []
|
||
context_inline_bytes = (
|
||
(len(video_bytes) if video_bytes is not None else 0)
|
||
+ len(padded_prompt.encode("utf-8"))
|
||
+ len(padded_system_instruction.encode("utf-8"))
|
||
)
|
||
|
||
if prepared_images is not None:
|
||
for img_bytes, png_hash in zip(prepared_images["png_bytes_list"], prepared_images["png_hashes"]):
|
||
context_inline_bytes += len(img_bytes)
|
||
context_image_hashes.append(png_hash)
|
||
if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified":
|
||
img_part = types.Part.from_bytes(
|
||
data=img_bytes,
|
||
mime_type="image/png",
|
||
media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}",
|
||
)
|
||
else:
|
||
img_part = types.Part.from_bytes(
|
||
data=img_bytes,
|
||
mime_type="image/png",
|
||
)
|
||
media_parts.append(img_part)
|
||
else:
|
||
image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs)
|
||
if image_frames:
|
||
try:
|
||
img_parts = []
|
||
for _, _, image_frame in image_frames:
|
||
img_array = image_frame.cpu().numpy()
|
||
img_array = (img_array * 255).astype(np.uint8)
|
||
pil_img = Image.fromarray(img_array)
|
||
img_byte_arr = BytesIO()
|
||
pil_img.save(img_byte_arr, format='PNG')
|
||
img_bytes = img_byte_arr.getvalue()
|
||
context_inline_bytes += len(img_bytes)
|
||
context_image_hashes.append(hashlib.sha256(img_bytes).hexdigest())
|
||
if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified":
|
||
img_part = types.Part.from_bytes(
|
||
data=img_bytes,
|
||
mime_type="image/png",
|
||
media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}",
|
||
)
|
||
else:
|
||
img_part = types.Part.from_bytes(
|
||
data=img_bytes,
|
||
mime_type="image/png",
|
||
)
|
||
img_parts.append(img_part)
|
||
batch_summary = ", ".join(
|
||
f"{socket_name}={batch_size}"
|
||
for socket_name, batch_size in image_batch_counts
|
||
)
|
||
print(
|
||
f"[INFO] Prepared Gemini image inputs: {batch_summary}, "
|
||
f"total={len(img_parts)}"
|
||
)
|
||
media_parts.extend(img_parts)
|
||
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, "", "")
|
||
|
||
if video_bytes is not None:
|
||
if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified":
|
||
video_part = types.Part.from_bytes(
|
||
data=video_bytes,
|
||
mime_type=video_mime_type,
|
||
media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}",
|
||
)
|
||
else:
|
||
video_part = types.Part.from_bytes(data=video_bytes, mime_type=video_mime_type)
|
||
video_part.video_metadata = types.VideoMetadata(fps=video_fps)
|
||
media_parts.append(video_part)
|
||
if media_parts:
|
||
contents = [
|
||
types.UserContent(
|
||
parts=media_parts + [types.Part.from_text(text=padded_prompt)]
|
||
)
|
||
]
|
||
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"]
|
||
context_key = cls._build_context_cache_key(
|
||
client_key,
|
||
model,
|
||
padded_system_instruction,
|
||
padded_prompt,
|
||
video_hash,
|
||
video_mime_type,
|
||
video_fps,
|
||
context_image_hashes,
|
||
media_resolution,
|
||
)
|
||
|
||
skip_context_cache = (
|
||
(
|
||
use_cache
|
||
and video_bytes is not None
|
||
and (config.get("use_vertexai_env", False) or config.get("vertexai_express", False))
|
||
and context_inline_bytes > cls.VERTEX_CACHE_INLINE_LIMIT_BYTES
|
||
)
|
||
or (use_cache and model in cls.ARGON_MODELS)
|
||
)
|
||
if skip_context_cache:
|
||
if use_cache and model in cls.ARGON_MODELS:
|
||
print(f"[WARNING] Explicit context cache skipped for {model} (implicit caching only); using full request")
|
||
else:
|
||
print("[WARNING] Vertex context cache skipped because inline cached media exceeds 10 MB; using full request")
|
||
|
||
resolved_safety_settings = cls._resolve_safety_settings(config)
|
||
|
||
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=actual_seed if use_seed else api_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,
|
||
image_size=image_size,
|
||
allow_all_people=bool(config.get("use_vertexai_env", False) or config.get("vertexai_express", False)),
|
||
safety_settings=resolved_safety_settings,
|
||
)
|
||
|
||
if schema_snapshot is not None:
|
||
generate_content_config.response_mime_type = "application/json"
|
||
generate_content_config.response_json_schema = schema_snapshot
|
||
|
||
if use_seed and actual_seed is not None:
|
||
try:
|
||
generate_content_config.seed = actual_seed
|
||
except Exception:
|
||
pass
|
||
elif use_cache:
|
||
try:
|
||
generate_content_config.seed = None
|
||
except Exception:
|
||
pass
|
||
|
||
# API call in background thread
|
||
start_time = time.time()
|
||
result_queue: "queue.Queue[Tuple[str, Any]]" = queue.Queue()
|
||
|
||
def generate_once():
|
||
# Apply the wire limit even on legacy unseeded image branches.
|
||
if generate_content_config.seed is not None and not -(2**31) <= generate_content_config.seed < 2**31:
|
||
generate_content_config.seed = cls._handle_seed(True, generate_content_config.seed)
|
||
if not use_cache or skip_context_cache:
|
||
return client.models.generate_content(
|
||
model=model,
|
||
contents=contents,
|
||
config=generate_content_config,
|
||
)
|
||
|
||
for cache_attempt in range(2):
|
||
try:
|
||
cache_name = cls._get_or_create_context_cache(
|
||
client,
|
||
context_key,
|
||
model,
|
||
contents,
|
||
padded_system_instruction,
|
||
cache_ttl_minutes,
|
||
)
|
||
except Exception as cache_error:
|
||
print(f"[WARNING] Gemini context cache unavailable; using full request: {cache_error}")
|
||
return client.models.generate_content(
|
||
model=model,
|
||
contents=contents,
|
||
config=generate_content_config,
|
||
)
|
||
|
||
cached_config = generate_content_config.model_copy(deep=True)
|
||
cached_config.cached_content = cache_name
|
||
cached_config.system_instruction = None
|
||
try:
|
||
return client.models.generate_content(
|
||
model=model,
|
||
contents=" ",
|
||
config=cached_config,
|
||
)
|
||
except Exception as cache_error:
|
||
if cache_attempt == 0 and cls._is_missing_context_cache_error(cache_error):
|
||
cls._invalidate_context_cache(context_key)
|
||
print("[INFO] Recreating expired Gemini context cache")
|
||
continue
|
||
raise
|
||
|
||
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 = generate_once()
|
||
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_thoughts_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:
|
||
if getattr(part, "thought", False):
|
||
current_thoughts_output += part.text
|
||
else:
|
||
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()
|
||
if schema_snapshot is not None:
|
||
json.loads(current_text_output, parse_constant=cls._reject_json_constant)
|
||
result_queue.put(("success", (current_text_output, current_image_tensor, current_thoughts_output)))
|
||
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
|
||
is_success = False
|
||
structured_output = ""
|
||
thoughts_output = ""
|
||
|
||
try:
|
||
status, result = cls._wait_for_result(result_queue, timeout)
|
||
elapsed_time = time.time() - start_time
|
||
if status == "success":
|
||
text_output, image_tensor, thoughts_output = result
|
||
structured_output = text_output if schema_snapshot is not None else ""
|
||
is_success = True
|
||
|
||
# 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):
|
||
cls._check_interrupt()
|
||
retry_attempt += 1
|
||
print(f"[INFO] Retry pattern matched in response. Retry attempt {retry_attempt}/{max_retries}")
|
||
|
||
# Generate new random gemini seed for retry
|
||
actual_seed = cls._handle_seed(True, 0)
|
||
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):
|
||
cls._check_interrupt()
|
||
retry_attempt += 1
|
||
print(f"[INFO] Retry pattern matched in error. Retry attempt {retry_attempt}/{max_retries}")
|
||
|
||
# Generate new random gemini seed for retry
|
||
actual_seed = cls._handle_seed(True, 0)
|
||
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 = timeout_fallback_text or f"Gemini API request/processing timed out, waited {timeout} seconds."
|
||
|
||
except Exception as e:
|
||
is_success = False
|
||
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
|
||
|
||
if not is_success:
|
||
structured_output = ""
|
||
thoughts_output = ""
|
||
|
||
# Cache the result
|
||
if use_seed and is_success:
|
||
try:
|
||
cls._cache[fingerprint] = (text_output, image_tensor, final_actual_seed, structured_output, thoughts_output)
|
||
except Exception:
|
||
pass
|
||
|
||
# Cache the successful gemini seed for this input seed (if retry was enabled)
|
||
if retry_pattern and max_retries > 0:
|
||
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}")
|
||
|
||
|
||
return IO.NodeOutput(text_output, image_tensor, final_actual_seed, structured_output, thoughts_output)
|
||
|
||
@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,
|
||
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, timeout_fallback_text="",
|
||
image_size="None",
|
||
video: GeminiVideoConfig.Type | None = None,
|
||
image_inputs: IO.Autogrow.Type | None = None,
|
||
response_schema: ResponseSchema.Type | None = None) -> IO.NodeOutput:
|
||
|
||
def _is_list_arg(val):
|
||
return isinstance(val, list) and not isinstance(val, (str, bytes, dict))
|
||
|
||
def _extract_seed_list(val):
|
||
if not isinstance(val, (list, tuple)):
|
||
return None
|
||
flat = []
|
||
for item in val:
|
||
if isinstance(item, (list, tuple)):
|
||
sub = _extract_seed_list(item)
|
||
if sub is not None:
|
||
flat.extend(sub)
|
||
elif torch.is_tensor(item):
|
||
flat.extend(int(x) for x in item.detach().cpu().reshape(-1).tolist())
|
||
elif isinstance(item, np.ndarray):
|
||
flat.extend(int(x) for x in item.reshape(-1).tolist())
|
||
elif item is not None:
|
||
flat.append(int(item))
|
||
return flat
|
||
|
||
is_prompt_list = _is_list_arg(prompt)
|
||
is_video_list = _is_list_arg(video)
|
||
extracted_seeds = _extract_seed_list(seed)
|
||
is_seed_list = extracted_seeds is not None
|
||
is_config_list = _is_list_arg(config)
|
||
|
||
is_batch = is_prompt_list or is_video_list or is_seed_list or is_config_list
|
||
|
||
if not is_batch:
|
||
return cls._execute_single(
|
||
config=config, prompt=prompt, system_instruction=system_instruction, model=model,
|
||
temperature=temperature, top_p=top_p, top_k=top_k, max_output_tokens=max_output_tokens,
|
||
include_images=include_images, aspect_ratio=aspect_ratio, bypass_mode=bypass_mode,
|
||
thinking_budget=thinking_budget, use_proxy=use_proxy, proxy_host=proxy_host,
|
||
proxy_port=proxy_port, use_seed=use_seed, seed=seed, timeout=timeout,
|
||
include_thoughts=include_thoughts, thinking_level=thinking_level,
|
||
media_resolution=media_resolution, retry_pattern=retry_pattern, max_retries=max_retries,
|
||
timeout_fallback_text=timeout_fallback_text, image_size=image_size,
|
||
video=video, image_inputs=image_inputs, response_schema=response_schema,
|
||
)
|
||
|
||
prompt_list = prompt if is_prompt_list else [prompt]
|
||
video_list = video if is_video_list else [video]
|
||
seed_list = extracted_seeds if is_seed_list else [seed]
|
||
config_list = config if is_config_list else [config]
|
||
|
||
batch_count = max(
|
||
len(prompt_list) if is_prompt_list else 1,
|
||
len(video_list) if is_video_list else 1,
|
||
len(seed_list) if is_seed_list else 1,
|
||
len(config_list) if is_config_list else 1,
|
||
)
|
||
|
||
def _unwrap_scalar(val, default=None):
|
||
while isinstance(val, list) and len(val) > 0 and (isinstance(val[0], list) or not isinstance(val, dict)):
|
||
val = val[0]
|
||
if isinstance(val, list) and len(val) == 1:
|
||
val = val[0]
|
||
return val if val is not None else default
|
||
|
||
def _get_aligned(val, idx, default=None):
|
||
if isinstance(val, list) and val:
|
||
return val[min(idx, len(val) - 1)]
|
||
return val if val is not None else default
|
||
|
||
config_val = _unwrap_scalar(config, {})
|
||
model_val = _unwrap_scalar(model, "gemini-2.5-flash")
|
||
temp_val = float(_unwrap_scalar(temperature, 1.0))
|
||
top_p_val = float(_unwrap_scalar(top_p, 0.95))
|
||
top_k_val = int(_unwrap_scalar(top_k, 40))
|
||
max_tokens_val = int(_unwrap_scalar(max_output_tokens, 8192))
|
||
inc_imgs_val = bool(_unwrap_scalar(include_images, False))
|
||
aspect_val = str(_unwrap_scalar(aspect_ratio, "None"))
|
||
bypass_val = str(_unwrap_scalar(bypass_mode, "None"))
|
||
budget_val = int(_unwrap_scalar(thinking_budget, 0))
|
||
proxy_val = bool(_unwrap_scalar(use_proxy, False))
|
||
phost_val = str(_unwrap_scalar(proxy_host, "127.0.0.1"))
|
||
pport_val = int(_unwrap_scalar(proxy_port, 7890))
|
||
use_seed_val = bool(_unwrap_scalar(use_seed, False))
|
||
timeout_val = int(_unwrap_scalar(timeout, 30))
|
||
thoughts_val = bool(_unwrap_scalar(include_thoughts, False))
|
||
th_level_val = _unwrap_scalar(thinking_level, None)
|
||
media_res_val = _unwrap_scalar(media_resolution, None)
|
||
retry_pat_val = str(_unwrap_scalar(retry_pattern, ""))
|
||
max_retries_val = int(_unwrap_scalar(max_retries, 3))
|
||
timeout_fb_val = str(_unwrap_scalar(timeout_fallback_text, ""))
|
||
img_size_val = str(_unwrap_scalar(image_size, "None"))
|
||
resp_schema_val = _unwrap_scalar(response_schema, None)
|
||
has_multi_seed_input = is_seed_list and len(seed_list) > 1
|
||
|
||
def make_item_kwargs(idx, prepared_imgs=None, precomp_vid=None):
|
||
p = prompt_list[min(idx, len(prompt_list) - 1)] if prompt_list else ""
|
||
v = video_list[min(idx, len(video_list) - 1)] if video_list else None
|
||
|
||
if has_multi_seed_input:
|
||
if idx < len(seed_list):
|
||
s = int(seed_list[idx])
|
||
else:
|
||
base_s = int(seed_list[-1])
|
||
s = (base_s + (idx - len(seed_list) + 1)) if base_s != 0 else 0
|
||
elif is_seed_list and len(seed_list) == 1:
|
||
base_s = int(seed_list[0])
|
||
s = (base_s + idx) if (base_s != 0 and batch_count > 1) else base_s
|
||
elif not is_seed_list and seed != 0:
|
||
s = (int(seed) + idx) if batch_count > 1 else int(seed)
|
||
else:
|
||
s = int(_unwrap_scalar(seed, 0))
|
||
|
||
raw_cfg = _get_aligned(config_list, idx, config_val)
|
||
item_cfg = dict(raw_cfg) if isinstance(raw_cfg, dict) else dict(config_val)
|
||
if item_cfg.get("use_cache"):
|
||
if has_multi_seed_input:
|
||
item_cfg["cache_seed"] = s
|
||
elif batch_count > 1 and not (is_config_list and len(config_list) > 1):
|
||
base_cs = int(item_cfg.get("cache_seed", 0))
|
||
if base_cs != 0:
|
||
item_cfg["cache_seed"] = base_cs + idx
|
||
elif s != 0:
|
||
item_cfg["cache_seed"] = s
|
||
|
||
sys_inst = _get_aligned(system_instruction, idx, "You are a helpful AI assistant.")
|
||
|
||
return {
|
||
"config": item_cfg,
|
||
"prompt": p,
|
||
"system_instruction": sys_inst,
|
||
"model": model_val,
|
||
"temperature": temp_val,
|
||
"top_p": top_p_val,
|
||
"top_k": top_k_val,
|
||
"max_output_tokens": max_tokens_val,
|
||
"include_images": inc_imgs_val,
|
||
"aspect_ratio": aspect_val,
|
||
"bypass_mode": bypass_val,
|
||
"thinking_budget": budget_val,
|
||
"use_proxy": proxy_val,
|
||
"proxy_host": phost_val,
|
||
"proxy_port": pport_val,
|
||
"use_seed": use_seed_val,
|
||
"seed": s,
|
||
"timeout": timeout_val,
|
||
"include_thoughts": thoughts_val,
|
||
"thinking_level": th_level_val,
|
||
"media_resolution": media_res_val,
|
||
"retry_pattern": retry_pat_val,
|
||
"max_retries": max_retries_val,
|
||
"timeout_fallback_text": timeout_fb_val,
|
||
"image_size": img_size_val,
|
||
"video": v,
|
||
"image_inputs": image_inputs,
|
||
"response_schema": resp_schema_val,
|
||
"prepared_images": prepared_imgs,
|
||
"precomputed_video": precomp_vid,
|
||
}
|
||
|
||
if batch_count == 1:
|
||
single_res = cls._execute_single(**make_item_kwargs(0))
|
||
return IO.NodeOutput(
|
||
[single_res[0]],
|
||
[single_res[1]],
|
||
[single_res[2]],
|
||
[single_res[3]],
|
||
[single_res[4]],
|
||
)
|
||
|
||
print(f"[INFO] Preparing sequential list execution for {batch_count} items (model: {model_val})")
|
||
prepared_images = None
|
||
if image_inputs:
|
||
try:
|
||
prepared_images = cls._prepare_image_inputs(image_inputs)
|
||
except Exception:
|
||
prepared_images = None
|
||
|
||
precomputed_video = None
|
||
if len(video_list) == 1 and video_list[0] is not None:
|
||
try:
|
||
v0 = video_list[0]
|
||
v_input = v0.get("video")
|
||
v_fps = int(v0.get("fps", 1))
|
||
v_pad = bool(v0.get("pad_at_start", False))
|
||
v_dur = bool(v0.get("duration_aware_padding", False))
|
||
v_bytes, v_mime = cls._serialize_video(v_input, v_fps, v_pad, v_dur)
|
||
v_hash = hashlib.sha256(v_bytes).hexdigest() if v_bytes is not None else None
|
||
precomputed_video = (v_bytes, v_mime, v_hash, v_fps)
|
||
except Exception:
|
||
precomputed_video = None
|
||
|
||
batch_results = []
|
||
for idx in range(batch_count):
|
||
cls._check_interrupt()
|
||
batch_results.append(
|
||
cls._execute_single(
|
||
**make_item_kwargs(
|
||
idx,
|
||
prepared_imgs=prepared_images,
|
||
precomp_vid=precomputed_video,
|
||
)
|
||
)
|
||
)
|
||
|
||
return IO.NodeOutput(
|
||
[r[0] for r in batch_results],
|
||
[r[1] for r in batch_results],
|
||
[r[2] for r in batch_results],
|
||
[r[3] for r in batch_results],
|
||
[r[4] for r in batch_results],
|
||
)
|
||
|
||
# V3 uses ComfyExtension entrypoint in __init__.py to expose nodes
|