diff --git a/__init__.py b/__init__.py index 468b2b5..d6311e1 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,26 @@ +import os +import asyncio from comfy_api.latest import ComfyExtension, io from .gemini_nodes import GetKeyAPI, SSL_GeminiAPIKeyConfig, SSL_GeminiTextPrompt class GeminiExtension(ComfyExtension): + async def on_load(self) -> None: + # Check if environment variables indicate Vertex AI usage + if os.environ.get("GOOGLE_GENAI_USE_VERTEXAI") == "True" or os.environ.get("GOOGLE_CLOUD_PROJECT"): + asyncio.create_task(self._background_auth()) + + async def _background_auth(self) -> None: + try: + print("[INFO] Attempting background Vertex AI authentication...") + from google import genai + # This will absorb the potential 60s metadata server timeout + client = genai.Client(vertexai=True) + await asyncio.to_thread(client._api_client._access_token) + print("[INFO] Background Vertex AI authentication successful.") + except Exception as e: + print(f"[WARNING] Background Vertex AI auth failed (will retry at inference): {e}") + async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ GetKeyAPI, diff --git a/gemini_nodes.py b/gemini_nodes.py index d715b07..330d9a3 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -581,8 +581,10 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): **client_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: @@ -624,8 +626,10 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): **client_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: @@ -637,10 +641,20 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): **client_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) @@ -858,8 +872,10 @@ 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") + # Cache the result - if use_seed: + if use_seed and is_success: try: cls._cache[fingerprint] = (text_output, image_tensor, final_actual_seed) except Exception: