107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
# Shared Gemini helper (google-genai)
|
|
#
|
|
# These nodes used to talk to the legacy `google-generativeai` SDK, which hard
|
|
# pinned google-ai-generativelanguage==0.6.15 and therefore dragged protobuf<6
|
|
# plus the whole grpcio / google-api stack into the ComfyUI environment.
|
|
# `google-genai` is the current SDK and needs none of that.
|
|
#
|
|
# Legacy -> current mapping used across the pack:
|
|
# genai.configure(api_key=...) -> genai.Client(api_key=...)
|
|
# genai.GenerativeModel(m, safety=...) -> config=types.GenerateContentConfig(...)
|
|
# model.generate_content(contents) -> client.models.generate_content(
|
|
# model=m, contents=..., config=...)
|
|
|
|
import os
|
|
|
|
try:
|
|
from google import genai
|
|
from google.genai import types as genai_types
|
|
|
|
GEMINI_AVAILABLE = True
|
|
except ImportError:
|
|
genai = None
|
|
genai_types = None
|
|
GEMINI_AVAILABLE = False
|
|
|
|
|
|
GEMINI_INSTALL_HINT = (
|
|
"google-genai is not installed. Install it with: pip install google-genai"
|
|
)
|
|
|
|
# The permissive safety posture these nodes have always shipped with.
|
|
_SAFETY_CATEGORIES = (
|
|
"HARM_CATEGORY_HARASSMENT",
|
|
"HARM_CATEGORY_HATE_SPEECH",
|
|
"HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
|
"HARM_CATEGORY_DANGEROUS_CONTENT",
|
|
)
|
|
|
|
|
|
def get_gemini_client(api_key=None):
|
|
"""
|
|
Build a google-genai client.
|
|
|
|
Falls back to the GEMINI_API_KEY environment variable and raises the same
|
|
errors the nodes raised before the migration.
|
|
"""
|
|
if not GEMINI_AVAILABLE:
|
|
raise ImportError(GEMINI_INSTALL_HINT)
|
|
|
|
key = api_key or os.environ.get("GEMINI_API_KEY")
|
|
if not key:
|
|
raise ValueError("GEMINI_API_KEY environment variable is not set")
|
|
|
|
return genai.Client(api_key=key)
|
|
|
|
|
|
def build_gemini_config(**kwargs):
|
|
"""
|
|
GenerateContentConfig with every safety filter set to BLOCK_NONE.
|
|
|
|
Any extra keyword (temperature, max_output_tokens, system_instruction, ...)
|
|
is forwarded; None values are dropped so callers can pass optionals
|
|
unconditionally.
|
|
"""
|
|
if not GEMINI_AVAILABLE:
|
|
raise ImportError(GEMINI_INSTALL_HINT)
|
|
|
|
return genai_types.GenerateContentConfig(
|
|
safety_settings=[
|
|
genai_types.SafetySetting(category=category, threshold="BLOCK_NONE")
|
|
for category in _SAFETY_CATEGORIES
|
|
],
|
|
**{k: v for k, v in kwargs.items() if v is not None},
|
|
)
|
|
|
|
|
|
def gemini_generate(client, model, contents, **config_kwargs):
|
|
"""
|
|
client.models.generate_content with the shared safety config applied.
|
|
|
|
`contents` may be a string, a PIL image, or a list mixing both - google-genai
|
|
converts PIL images to inline image parts on its own.
|
|
"""
|
|
if not isinstance(contents, list):
|
|
contents = [contents]
|
|
|
|
return client.models.generate_content(
|
|
model=model,
|
|
contents=contents,
|
|
config=build_gemini_config(**config_kwargs),
|
|
)
|
|
|
|
|
|
def gemini_finished_normally(candidate):
|
|
"""
|
|
True when a candidate completed with FinishReason.STOP.
|
|
|
|
The legacy SDK exposed finish_reason as the int 1; google-genai returns a
|
|
FinishReason string enum, so compare by name instead.
|
|
"""
|
|
reason = getattr(candidate, "finish_reason", None)
|
|
if reason is None:
|
|
return True
|
|
|
|
name = getattr(reason, "name", None) or str(reason).rsplit(".", 1)[-1]
|
|
return str(name).upper() == "STOP"
|