nodes/silveroxides-ComfyUI_Gemini_Expanded_API
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

improve retry logic

Browse Source
This commit is contained in:
silveroxides
2025-06-11 04:33:30 +02:00
parent a4e7d75664
commit 02094557e1
1 changed files with 183 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gemini_nodes.py
+183 -198
View File
@@ -6,6 +6,7 @@ import numpy as np
from PIL import Image
import io
import folder_paths
from comfy.comfy_types.node_typing import IO, ComfyNodeABC, InputTypeDict
from google import genai
from google.genai import types
import time
@@ -22,7 +23,7 @@ def check_and_install_dependencies():
'requests': 'requests',
'pysocks': 'PySocks',
}
for module_name, package_name in required_packages.items():
try:
importlib.import_module(module_name)
@@ -39,185 +40,180 @@ try:
except Exception as e:
print(f"[WARNING] Error checking dependencies: {str(e)}")
class SSL_GeminiAPIKeyConfig:
class SSL_GeminiAPIKeyConfig(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls):
def INPUT_TYPES(cls) -> InputTypeDict:
return {
"required": {
"api_key": ("STRING", {"multiline": False}),
"api_key": (IO.STRING, {"multiline": False}),
}
}
RETURN_TYPES = ("GEMINI_CONFIG",)
RETURN_NAMES = ("config",)
FUNCTION = "configure"
CATEGORY = "💠SSL/API/Gemini"
def configure(self, api_key):
config = {"api_key": api_key}
return (config,)
class SSL_GeminiTextPrompt:
class SSL_GeminiTextPrompt(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls):
def INPUT_TYPES(cls) -> InputTypeDict:
return {
"required": {
"config": ("GEMINI_CONFIG",),
"prompt": ("STRING", {"multiline": True}),
"system_instruction": ("STRING", {"default": "You are a helpful AI assistant.", "multiline": True}),
"model": (["gemini-1.0-pro", "gemini-exp-1206", "gemini-1.5-flash", "gemini-1.5-pro", "gemini-2.0-flash", "gemini-2.0-flash-lite-001", "gemini-2.0-flash-exp", "gemini-2.0-pro", "gemini-2.0-pro-exp-02-05"], {"default": "gemini-2.0-flash"}),
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_k": ("INT", {"default": 40, "min": 1, "max": 100, "step": 1}),
"max_output_tokens": ("INT", {"default": 8192, "min": 1, "max": 8192, "step": 1}),
"include_images": (["True", "False"], {"default": "True"}),
"prompt": (IO.STRING, {"multiline": True}),
"system_instruction": (IO.STRING, {"default": "You are a helpful AI assistant.", "multiline": True}),
"model": (["gemini-1.0-pro", "gemini-exp-1206", "gemini-1.5-flash", "gemini-1.5-pro", "gemini-2.0-flash", "gemini-2.0-flash-lite-001", "gemini-2.0-flash-exp", "gemini-2.0-pro", "gemini-2.0-flash-live", "gemini-2.5-flash", "gemini-2.5", "gemini-2.5-pro-1p-freebie"], {"default": "gemini-2.0-flash"}),
"temperature": (IO.FLOAT, {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_p": (IO.FLOAT, {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
"top_k": (IO.INT, {"default": 40, "min": 1, "max": 100, "step": 1}),
"max_output_tokens": (IO.INT, {"default": 8192, "min": 1, "max": 8192, "step": 1}),
"include_images": (IO.BOOLEAN, {"default": True}),
},
"optional": {
"input_image": ("IMAGE",),
"input_image_2": ("IMAGE",), # Add second image input
"use_proxy": (["False", "True"], {"default": "False"}),
"proxy_host": ("STRING", {"default": "127.0.0.1"}),
"proxy_port": ("INT", {"default": 7890, "min": 1, "max": 65535}),
"use_seed": (["True", "False"], {"default": "True"}),
"seed": ("INT", {"default": 0, "min": 0, "max": 2147483647}),
"input_image": (IO.IMAGE,),
"input_image_2": (IO.IMAGE,),
"use_proxy": (IO.BOOLEAN, {"default": False}),
"proxy_host": (IO.STRING, {"default": "127.0.0.1"}),
"proxy_port": (IO.INT, {"default": 7890, "min": 1, "max": 65535}),
"use_seed": (IO.BOOLEAN, {"default": True}),
"seed": (IO.INT, {"default": 0, "min": 0, "max": 2147483647}),
}
}
RETURN_TYPES = ("STRING", "IMAGE", "INT")
RETURN_NAMES = ("text", "image", "seed")
FUNCTION = "generate"
CATEGORY = "💠SSL/API/Gemini"
def save_binary_file(self, 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
def generate_empty_image(self, 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
def generate(self, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images,
input_image=None, input_image_2=None, use_proxy="False", proxy_host="127.0.0.1", proxy_port=7890,
use_seed="True", seed=0):
def generate(self, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images,
input_image=None, input_image_2=None, use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890,
use_seed=True, seed=0):
original_http_proxy = os.environ.get('HTTP_PROXY')
original_https_proxy = os.environ.get('HTTPS_PROXY')
original_http_proxy_lower = os.environ.get('http_proxy')
original_https_proxy_lower = os.environ.get('https_proxy')
print(f"[INFO] Starting generation, model: {model}, temperature: {temperature}")
# Improved random seed handling logic
if use_seed == "True":
# If seed is 0, generate a random seed
actual_seed = None
if use_seed == True:
if seed == 0:
# Use current timestamp and random number combination to generate seed, ensuring it's different each run
current_time = int(time.time() * 1000)
random_component = random.randint(0, 1000000)
actual_seed = (current_time + random_component) % 2147483647
print(f"[INFO] Generated random seed: {actual_seed}")
else:
# Use the specified seed value
actual_seed = seed
print(f"[INFO] Using specified seed: {actual_seed}")
# Set random seed
random.seed(actual_seed)
np.random.seed(actual_seed)
torch.manual_seed(actual_seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(actual_seed)
else:
actual_seed = None
print("[INFO] Seed not used")
try:
if use_proxy == "True":
if use_proxy == True:
if not proxy_host.startswith(('http://', 'https://')):
proxy_url = f"http://{proxy_host}:{proxy_port}"
else:
proxy_url = f"{proxy_host}:{proxy_port}"
os.environ['HTTP_PROXY'] = proxy_url
os.environ['HTTPS_PROXY'] = proxy_url
os.environ['http_proxy'] = proxy_url
os.environ['https_proxy'] = proxy_url
os.environ['REQUESTS_CA_BUNDLE'] = ''
print(f"[INFO] Proxy enabled: {proxy_url}")
print(f"[INFO] Initializing Gemini client")
try:
client_options = {}
if use_proxy == "True":
if use_proxy == True:
try:
import google.api_core.http_client
import google.auth.transport.requests
try:
import requests
from requests.adapters import HTTPAdapter
class ProxyAdapter(HTTPAdapter):
def __init__(self, proxy_url, **kwargs):
self.proxy_url = proxy_url
super().__init__(**kwargs)
def add_headers(self, request, **kwargs):
super().add_headers(request, **kwargs)
session = requests.Session()
proxies = {
"http": proxy_url,
"https": proxy_url
}
session.proxies.update(proxies)
adapter = ProxyAdapter(proxy_url, max_retries=3)
adapter = ProxyAdapter(proxy_url, max_retries=2)
session.mount('http://', adapter)
session.mount('https://', adapter)
session.verify = False
http_client = google.api_core.http_client.RequestsHttpClient(session=session)
client_options["http_client"] = http_client
print(f"[INFO] Used requests library to set proxy for HTTP client")
except Exception as proxy_error:
print(f"[WARNING] Failed to set HTTP client proxy: {str(proxy_error)}")
except ImportError as e:
print(f"[WARNING] Failed to import Google API HTTP client library: {str(e)}")
client = genai.Client(api_key=config["api_key"], **client_options)
client = genai.Client(api_key=config["api_key"], http_options=types.HttpOptions(api_version='v1alpha'), **client_options)
print(f"[INFO] Gemini client initialized successfully")
except Exception as e:
print(f"[ERROR] Gemini client initialization failed: {str(e)}")
return (f"Gemini client initialization failed: {str(e)}", self.generate_empty_image(), actual_seed if actual_seed is not None else 0)
try:
import socket
test_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
test_socket.settimeout(5)
test_host = "generativelanguage.googleapis.com"
if use_proxy == "True":
if use_proxy == True:
try:
import socks
test_socket = socks.socksocket()
@@ -225,15 +221,15 @@ class SSL_GeminiTextPrompt:
test_socket.settimeout(5)
except ImportError:
print(f"[WARNING] PySocks library not installed, cannot test connection via proxy")
test_socket.connect((test_host, 443))
test_socket.close()
network_ok = True
except Exception as e:
network_ok = False
print(f"[WARNING] Network connection test failed: {str(e)}")
if use_proxy == "True":
if use_proxy == True:
try:
import subprocess
cmd = f"curl -x {proxy_url} -s -o /dev/null -w '%{{http_code}}' https://{test_host}"
@@ -247,45 +243,42 @@ class SSL_GeminiTextPrompt:
print(f"[WARNING] curl test failed: {str(curl_error)}")
except ImportError:
pass
contents = []
# 处理输入图像
images_to_process = []
if input_image is not None:
images_to_process.append(input_image)
if input_image_2 is not None:
images_to_process.append(input_image_2)
# 如果有图像需要处理
if images_to_process:
try:
img_parts = []
for img in images_to_process:
img_array = img[0].cpu().numpy()
img_array = (img_array * 255).astype(np.uint8)
pil_img = Image.fromarray(img_array)
img_byte_arr = io.BytesIO()
pil_img.save(img_byte_arr, format='PNG')
img_bytes = img_byte_arr.getvalue()
img_part = {"inline_data": {"mime_type": "image/png", "data": img_bytes}}
img_parts.append(img_part)
# Add all image parts and text parts
contents = img_parts + [{"text": prompt}]
except Exception as e:
print(f"[ERROR] Error processing input image: {str(e)}")
return (f"Error processing input image: {str(e)}", self.generate_empty_image(), actual_seed if actual_seed is not None else 0)
else:
contents = prompt
response_modalities = ["text"]
if include_images == "True":
if include_images == True:
response_modalities.append("image")
generate_content_config = types.GenerateContentConfig(
temperature=temperature,
top_p=top_p,
@@ -294,19 +287,19 @@ class SSL_GeminiTextPrompt:
safety_settings=[
types.SafetySetting(
category="HARM_CATEGORY_HARASSMENT",
threshold="BLOCK_NONE", # Block none
threshold="BLOCK_NONE",
),
types.SafetySetting(
category="HARM_CATEGORY_HATE_SPEECH",
threshold="BLOCK_NONE", # Block none
threshold="BLOCK_NONE",
),
types.SafetySetting(
category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
threshold="BLOCK_NONE", # Block none
threshold="BLOCK_NONE",
),
types.SafetySetting(
category="HARM_CATEGORY_DANGEROUS_CONTENT",
threshold="BLOCK_NONE", # Block none
threshold="BLOCK_NONE",
),
],
response_modalities=response_modalities,
@@ -314,166 +307,158 @@ class SSL_GeminiTextPrompt:
types.Part.from_text(text=system_instruction),
],
)
if use_seed == "True" and actual_seed is not None:
if use_seed == True and actual_seed is not None:
try:
generate_content_config.seed = actual_seed
except Exception as seed_error:
print(f"[WARNING] Failed to set seed for API request: {str(seed_error)}")
# text_output = "" # Will be initialized in api_call and retrieved from queue
# image_tensor = None # Will be initialized in api_call and retrieved from queue
try:
print(f"[INFO] Sending API request to Gemini")
start_time = time.time()
timeout = 30 # Timeout for queue.get()
def api_call():
last_exception = None
max_retries = 3
for attempt in range(max_retries):
current_text_output = ""
current_image_tensor = None # Initialized for each attempt
start_time = time.time()
timeout = 30
def api_call():
last_api_exception = None
max_retries = 3
api_response = None
for attempt in range(max_retries):
try:
print(f"[INFO] API call attempt {attempt + 1}/{max_retries}")
api_response = client.models.generate_content(
if use_seed and actual_seed is not None:
current_seed_for_api_call = actual_seed + attempt
try:
generate_content_config.seed = current_seed_for_api_call
if attempt > 0:
print(f"[INFO] Retrying with incremented seed: {current_seed_for_api_call}")
except AttributeError:
print(f"[WARNING] Could not set 'seed' attribute on generate_content_config.")
except Exception as e_set_seed:
print(f"[WARNING] Error setting seed {current_seed_for_api_call} for attempt {attempt + 1}: {str(e_set_seed)}")
response = client.models.generate_content(
model=model,
contents=contents,
config=generate_content_config,
)
print(f"[INFO] Received API response for attempt {attempt + 1}")
if hasattr(api_response, 'candidates') and api_response.candidates:
for part in api_response.candidates[0].content.parts:
if hasattr(part, 'text') and part.text is not None:
current_text_output += part.text
elif hasattr(part, 'inline_data') and part.inline_data is not None:
try:
inline_data = part.inline_data
mime_type = inline_data.mime_type
data = inline_data.data
image_path = self.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 as img_e:
print(f"[ERROR] Image processing error during attempt {attempt + 1}: {str(img_e)}")
current_text_output += f"\nImage processing error: {str(img_e)}"
if current_image_tensor is None:
current_image_tensor = self.generate_empty_image()
if current_image_tensor is None: # Ensure image_tensor is always set
current_image_tensor = self.generate_empty_image()
# --- START OF CORRECTED CODE ---
# This is the robust validation check. It verifies that the necessary parts of the
# response exist and are not None before proceeding.
if not (response.candidates and response.candidates[0].content and response.candidates[0].content.parts):
finish_reason = "UNKNOWN"
if response.candidates and hasattr(response.candidates[0], 'finish_reason'):
# Get the reason if available (e.g., SAFETY, RECITATION, etc.)
finish_reason = response.candidates[0].finish_reason.name
result_queue.put(("success", (current_text_output, current_image_tensor)))
print(f"[INFO] Attempt {attempt + 1}/{max_retries} successful.")
return # Success, exit the function
# Raise a specific error that will be caught by the except block below,
# forcing a retry. This now correctly handles blocked/empty responses.
raise ValueError(f"Response was empty or blocked (Finish Reason: {finish_reason}).")
# --- END OF CORRECTED CODE ---
api_response = response
break # Success! Exit the retry loop.
except Exception as e:
last_exception = e
print(f"[ERROR] Attempt {attempt + 1}/{max_retries} failed: {str(e)}")
last_api_exception = e
print(f"[ERROR] API call attempt {attempt + 1}/{max_retries} failed: {str(e)}")
if attempt < max_retries - 1:
wait_time = 2**attempt # Exponential backoff: 1s, 2s
wait_time = 2**attempt
print(f"[INFO] Waiting {wait_time}s before next attempt...")
time.sleep(wait_time)
# On the last attempt, the loop will end, and we'll put the error in the queue outside the loop.
# If all attempts failed
print(f"[ERROR] All {max_retries} API call and processing attempts failed. Last error: {str(last_exception)}")
result_queue.put(("error", last_exception))
if api_response is None:
print(f"[ERROR] All {max_retries} API call attempts failed. Last error: {str(last_api_exception)}")
result_queue.put(("error", last_api_exception))
return
try:
current_text_output = ""
current_image_tensor = None
for part in api_response.candidates[0].content.parts:
if hasattr(part, 'text') and part.text is not None:
current_text_output += part.text
elif hasattr(part, 'inline_data') and part.inline_data is not None:
try:
inline_data = part.inline_data
mime_type = inline_data.mime_type
data = inline_data.data
image_path = self.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 as img_e:
print(f"[ERROR] Image processing error: {str(img_e)}")
current_text_output += f"\nImage processing error: {str(img_e)}"
if current_image_tensor is None:
current_image_tensor = self.generate_empty_image()
if current_image_tensor is None:
current_image_tensor = self.generate_empty_image()
result_queue.put(("success", (current_text_output, current_image_tensor)))
print("[INFO] API response processing successful.")
except Exception as e_proc:
print(f"[ERROR] A critical error occurred while processing the successful API response: {str(e_proc)}")
result_queue.put(("error", e_proc))
result_queue = queue.Queue()
api_thread = threading.Thread(target=api_call)
api_thread.daemon = True
api_thread.start()
text_output = ""
image_tensor = self.generate_empty_image() # Default to empty image
image_tensor = self.generate_empty_image()
try:
status, result = result_queue.get(timeout=timeout)
elapsed_time = time.time() - start_time
if status == "success":
text_output, image_tensor = result # result is now a tuple (text_output, image_tensor)
text_output, image_tensor = result
print(f"[INFO] API request and processing successfully completed in {elapsed_time:.2f} seconds")
else: # status == "error"
else:
error_exception = result
print(f"[ERROR] Error in API request/processing thread after {elapsed_time:.2f} seconds: {str(error_exception)}")
text_output = f"API call/processing error: {str(error_exception)}"
# image_tensor is already self.generate_empty_image()
error_str = str(error_exception).lower()
if any(term in error_str for term in ["timeout", "connection", "network", "socket", "连接", "网络"]):
if not network_ok and use_proxy != "True":
if not network_ok and use_proxy != True:
text_output += " Network connection test failed, consider enabling proxy."
# return (f"API request failed: {str(error_exception)}. Network connection test failed, consider enabling proxy.", self.generate_empty_image(), actual_seed if actual_seed is not None else 0)
# raise result # No need to raise here, already handled by setting text_output
except queue.Empty:
elapsed_time = time.time() - start_time
print(f"[ERROR] API request/processing timed out in main thread, waited: {elapsed_time:.2f} seconds")
timeout_msg = f"Gemini API request/processing timed out, waited {timeout} seconds."
if not network_ok:
if use_proxy == "True":
if use_proxy == True:
timeout_msg += f" Network connection test failed, the current proxy ({proxy_host}:{proxy_port}) may be invalid, please check proxy settings."
else:
timeout_msg += " Network connection test failed, consider enabling proxy."
else:
timeout_msg += " Network connection test successful, but API request/processing still timed out. This could be due to a busy server or a large request."
text_output = timeout_msg
# image_tensor is already self.generate_empty_image()
# return (timeout_msg, self.generate_empty_image(), actual_seed if actual_seed is not None else 0)
# print(f"[INFO] Received API response") # This line is now part of api_call or success message
# The following block is now inside api_call
# if hasattr(response, 'candidates') and response.candidates:
# for part in response.candidates[0].content.parts:
# if hasattr(part, 'text') and part.text is not None:
# 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 = self.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
# image_tensor = torch.from_numpy(img_array).unsqueeze(0)
# except Exception as e:
# print(f"[ERROR] Image processing error: {str(e)}")
# text_output += f"\nImage processing error: {str(e)}"
except Exception as e: # Catch-all for unforeseen errors in the main try block
except Exception as e:
print(f"[ERROR] Unhandled error in generate method: {str(e)}")
text_output = f"Unhandled error: {str(e)}"
if image_tensor is None: # Ensure image_tensor is set
if image_tensor is None:
image_tensor = self.generate_empty_image()
# if image_tensor is None: # This should be handled by initialization or in api_call
# image_tensor = self.generate_empty_image()
# if use_seed == "True" and actual_seed is not None:
# seed_info = f"\n\n[Seed information: {actual_seed}]"
# text_output += seed_info
return (text_output, image_tensor, actual_seed if actual_seed is not None else 0)
finally:
if original_http_proxy:
@@ -481,28 +466,28 @@ class SSL_GeminiTextPrompt:
else:
if 'HTTP_PROXY' in os.environ:
os.environ.pop('HTTP_PROXY')
if original_https_proxy:
os.environ['HTTPS_PROXY'] = original_https_proxy
else:
if 'HTTPS_PROXY' in os.environ:
os.environ.pop('HTTPS_PROXY')
if original_http_proxy_lower:
os.environ['http_proxy'] = original_http_proxy_lower
else:
if 'http_proxy' in os.environ:
os.environ.pop('http_proxy')
if original_https_proxy_lower:
os.environ['https_proxy'] = original_https_proxy_lower
else:
if 'https_proxy' in os.environ:
os.environ.pop('https_proxy')
if 'REQUESTS_CA_BUNDLE' in os.environ:
os.environ.pop('REQUESTS_CA_BUNDLE')
try:
import requests
if hasattr(requests, 'Session'):
@@ -520,4 +505,4 @@ NODE_CLASS_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = {
"SSL_GeminiAPIKeyConfig": "Configure Gemini API Key",
"SSL_GeminiTextPrompt": "Expanded Gemini Text/Image",
}
}