1382 lines
69 KiB
Python
1382 lines
69 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, 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_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("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."),
|
||
],
|
||
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) -> 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,
|
||
}
|
||
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",
|
||
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"),
|
||
IO.Image.Output("image"),
|
||
IO.Int.Output("final_actual_seed"),
|
||
IO.String.Output("structured_output", display_name="Structured output"),
|
||
IO.String.Output("thoughts", display_name="Thoughts"),
|
||
]
|
||
)
|
||
|
||
@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 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 _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
|
||
|
||
@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):
|
||
|
||
# 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_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):
|
||
# 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:
|
||
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 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:
|
||
|
||
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)
|
||
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
|
||
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,
|
||
)
|
||
|
||
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)
|
||
image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs)
|
||
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 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")
|
||
|
||
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)),
|
||
)
|
||
|
||
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 = result_queue.get(timeout=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):
|
||
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):
|
||
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)
|
||
|
||
# V3 uses ComfyExtension entrypoint in __init__.py to expose nodes
|