feat(gemini): expand image and auth support

This commit is contained in:
silveroxides
2026-08-09 15:29:08 +02:00
parent 0e9efbc7f6
commit a9a589d11e
3 changed files with 437 additions and 29 deletions
+112 -29
View File
@@ -137,6 +137,7 @@ class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
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."),
],
outputs=[
cls.GemConfig.Output("config")
@@ -144,8 +145,18 @@ class SSL_GeminiAPIKeyConfig(IO.ComfyNode):
)
@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 = "") -> 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}
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 = "") -> 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,
}
return IO.NodeOutput(config)
@@ -174,6 +185,12 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
@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",
@@ -191,8 +208,16 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
IO.Combo.Input("aspect_ratio", options=["None", "1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5", "21:9"], default="None"),
IO.Combo.Input("bypass_mode", options=["None", "system_instruction", "prompt", "both"], default="None"),
IO.Int.Input("thinking_budget", default=0, min=-1, max=24576, step=1, tooltip="0 disables thinking mode, -1 will activate it as default dynamic thinking and anything above 0 sets specific budget"),
IO.Image.Input("input_image", optional=True),
IO.Image.Input("input_image_2", optional=True),
IO.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."
),
),
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),
@@ -204,6 +229,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
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."),
],
outputs=[
IO.String.Output("text"),
@@ -270,10 +296,33 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
tensor = torch.from_numpy(empty_image).unsqueeze(0)
return tensor
@classmethod
def _flatten_image_inputs(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
flattened = []
for _, 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("Gemini image inputs must be BHWC image tensors.")
flattened.extend(
image_batch[index : index + 1]
for index in range(image_batch.shape[0])
)
return flattened
@classmethod
def _compute_fingerprint_and_check_cache(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed,
input_image=None, input_image_2=None,
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):
@@ -288,8 +337,10 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
print(f"[WARNING] Cache hashing failed for image: {e}")
return "Error"
image_1_hash = get_tensor_hash(input_image)
image_2_hash = get_tensor_hash(input_image_2)
image_hashes = tuple(
get_tensor_hash(image)
for image in cls._flatten_image_inputs(image_inputs)
)
# 2. Determine Effective Parameters based on Model
# This ensures we don't cache-miss if an irrelevant parameter changes
@@ -319,8 +370,13 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
# If none of the above, all effective thinking/image params remain default/ignored
sanitized_config = dict(config)
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(config),
str(sanitized_config),
prompt,
system_instruction,
model,
@@ -334,8 +390,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
eff_thinking_budget, # EFFECTIVE thinking_budget
use_seed,
int(seed) if use_seed else 0, # Only use seed in cache if use_seed is True
image_1_hash,
image_2_hash,
image_hashes,
bool(use_proxy),
str(proxy_host),
int(proxy_port),
@@ -456,16 +511,17 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
@classmethod
def execute(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget, input_image=None, input_image_2=None,
include_images, aspect_ratio, bypass_mode, thinking_budget,
image_inputs: IO.Autogrow.Type | None = None,
use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, use_seed=False, seed=0, timeout=30,
include_thoughts=False, thinking_level=None, media_resolution=None,
retry_pattern="", max_retries=3) -> IO.NodeOutput:
retry_pattern="", max_retries=3, timeout_fallback_text="") -> IO.NodeOutput:
print(f"[INFO] SSL_GeminiTextPrompt execute called, model: {model}")
fingerprint, cached = cls._compute_fingerprint_and_check_cache(
config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens,
include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed,
input_image, input_image_2,
image_inputs,
use_proxy, proxy_host, proxy_port, timeout,
include_thoughts, thinking_level, media_resolution,
retry_pattern, max_retries
@@ -515,6 +571,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
text_output = ""
image_tensor = cls.generate_empty_image()
proxy_url: str | None = None
timed_out = False
try:
if use_proxy:
@@ -556,22 +613,43 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
api_version = config.get("api_version")
project = config.get("vertexai_project")
location = config.get("vertexai_location")
credentials_path = config.get("google_application_credentials")
if use_vertexai_env:
try:
env_use = os.environ["GOOGLE_GENAI_USE_VERTEXAI"].strip() if "GOOGLE_GENAI_USE_VERTEXAI" in os.environ else "True"
assert env_use, "GOOGLE_GENAI_USE_VERTEXAI is empty"
env_use_value = os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "true").strip().lower()
if env_use_value not in {"true", "1", "false", "0"}:
raise ValueError("GOOGLE_GENAI_USE_VERTEXAI must be true, false, 1, or 0")
env_use = env_use_value in {"true", "1"}
if not env_use:
raise ValueError("GOOGLE_GENAI_USE_VERTEXAI must be enabled when use_vertexai_env is true")
env_proj = os.environ["GOOGLE_CLOUD_PROJECT"].strip() if "GOOGLE_CLOUD_PROJECT" in os.environ else project
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_proj = os.environ["GOOGLE_CLOUD_PROJECT"].strip() if "GOOGLE_CLOUD_PROJECT" in os.environ else (project or credential_project)
assert env_proj, "GOOGLE_CLOUD_PROJECT is empty"
env_loc = os.environ["GOOGLE_CLOUD_LOCATION"].strip() if "GOOGLE_CLOUD_LOCATION" in os.environ else location
assert env_loc, "GOOGLE_CLOUD_LOCATION is empty"
client_key = ("vertexai_env", env_use, env_proj, env_loc, api_version, proxy_url)
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(
vertexai=env_use,
credentials=credentials,
project=env_proj,
location=env_loc,
http_options=types.HttpOptions(api_version=api_version),
@@ -592,6 +670,10 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
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 Vertex AI configuration: {e}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
elif vertexai_express:
if not project:
@@ -614,11 +696,13 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
else:
os.environ.setdefault("GOOGLE_CLOUD_LOCATION", location)
client_key = ("vertexai_express", config.get("api_key"), project, location, api_version, proxy_url)
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 = ("vertexai_express", api_key_hash, project, location, api_version, proxy_url)
if client_key not in cls._client_cache:
cls._client_cache[client_key] = genai.Client(
vertexai=True,
api_key=config.get("api_key"),
api_key=api_key,
http_options=types.HttpOptions(api_version=api_version),
**client_options
)
@@ -630,10 +714,12 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
client = cls._client_cache[client_key]
else:
client_key = ("standard", config.get("api_key"), api_version, proxy_url)
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=config.get("api_key"),
api_key=api_key,
http_options=types.HttpOptions(api_version=api_version),
**client_options
)
@@ -657,11 +743,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
return IO.NodeOutput(f"Gemini client initialization failed: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0)
# Prepare contents (images + prompt)
images_to_process = []
if input_image is not None:
images_to_process.append(input_image)
if input_image_2 is not None:
images_to_process.append(input_image_2)
images_to_process = cls._flatten_image_inputs(image_inputs)
if images_to_process:
try:
@@ -859,7 +941,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
print(f"[WARNING] Invalid retry regex pattern: {regex_err}")
except queue.Empty:
text_output = f"Gemini API request/processing timed out, waited {timeout} seconds."
timed_out = True
text_output = timeout_fallback_text or f"Gemini API request/processing timed out, waited {timeout} seconds."
except Exception as e:
print(f"[ERROR] Unhandled error in generate method: {str(e)}")
@@ -869,7 +952,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
final_actual_seed = actual_seed if actual_seed is not None else 0
is_success = not text_output.startswith("API call/processing error:") and not text_output.startswith("Gemini API request/processing timed out")
is_success = not timed_out and not text_output.startswith("API call/processing error:") and not text_output.startswith("Gemini API request/processing timed out")
# Cache the result
if use_seed and is_success:
@@ -924,4 +1007,4 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
except:
pass
# V3 uses ComfyExtension entrypoint in __init__.py to expose nodes
# V3 uses ComfyExtension entrypoint in __init__.py to expose nodes