37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
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,
|
|
SSL_GeminiAPIKeyConfig,
|
|
SSL_GeminiTextPrompt,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> GeminiExtension:
|
|
return GeminiExtension()
|
|
|
|
|
|
__all__ = ['comfy_entrypoint']
|