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