nodes/Aryan185-ComfyUI-ExternalAPI-Helpers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Updated gemini models and added new params

Browse Source
This commit is contained in:
Aryan185
2026-05-04 07:03:06 +00:00
parent e5f6323192
commit 40cde5caea
7 changed files with 203 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gemini_node.py
+27 -17
View File
@@ -12,12 +12,15 @@ class GeminiChatNode:
def INPUT_TYPES(cls):
return {
"required": {
"api_key": ("STRING", {"default": "", "multiline": False, "tooltip": "Directly put Gemini API key or .env variable name (GEMINI_API_KEY)"}),
"model": (["gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.5-flash-lite", "gemini-3.1-pro-preview", "gemini-3.1-flash-lite-preview", "gemini-3-flash-preview", "gemini-flash-latest", "gemini-flash-lite-latest", "gemini-2.0-flash", "gemini-2.0-flash-lite"], {"default": "gemini-2.5-flash"}),
"prompt": ("STRING", {"multiline": True}),
"model": (["gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.5-flash-lite", "gemini-3-pro-preview", "gemini-3.1-pro-preview", "gemini-3-flash-preview", "gemini-3.1-flash-lite-preview", "gemini-flash-latest", "gemini-flash-lite-latest", "gemini-2.0-flash", "gemini-2.0-flash-lite"], {"default": "gemini-2.5-flash"}),
"temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 2.0, "step": 0.1}),
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
"thinking": ("BOOLEAN", {"default": False}),
"google_search": ("BOOLEAN", {"default": False}),
"url_context": ("BOOLEAN", {"default": False}),
"seed": ("INT", {"default": 69, "min": -1, "max": 2147483646, "step": 1}),
"api_key": ("STRING", {"default": "", "multiline": False, "tooltip": "Directly put Gemini API key or .env variable name (GEMINI_API_KEY)"})
},
"optional": {
"system_instruction": ("STRING", {"multiline": True, "default": ""}),
@@ -26,27 +29,28 @@ class GeminiChatNode:
"audio": ("AUDIO",),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("response",)
FUNCTION = "generate"
CATEGORY = "text/generation"
def generate(self, prompt, model, temperature, thinking, seed, api_key,
system_instruction=None, thinking_budget=-1, image=None, audio=None):
def generate(self, prompt, model, temperature, top_p, thinking, google_search, url_context, seed, api_key,
system_instruction=None, thinking_budget=0, image=None, audio=None):
key = os.environ.get(api_key.strip(), api_key.strip()) or os.environ.get("GEMINI_API_KEY")
if not key: raise ValueError("Error: No API key provided.")
client = genai.Client(api_key=key, http_options={'api_version': 'v1beta'})
parts = [types.Part.from_text(text=prompt)]
if image is not None:
arr = (image[0].cpu().numpy() * 255).astype(np.uint8)
buf = io.BytesIO()
Image.fromarray(arr).save(buf, format="PNG")
parts.append(types.Part.from_bytes(mime_type="image/png", data=buf.getvalue()))
for i in range(image.shape[0]):
arr = (image[i].cpu().numpy() * 255).astype(np.uint8)
buf = io.BytesIO()
Image.fromarray(arr).save(buf, format="PNG")
parts.append(types.Part.from_bytes(mime_type="image/png", data=buf.getvalue()))
if audio is not None:
wf = audio.get("waveform") if isinstance(audio, dict) else audio[0]
sr = audio.get("sample_rate", 44100) if isinstance(audio, dict) else audio[1]
@@ -67,25 +71,31 @@ class GeminiChatNode:
if "gemini-2.0" in model_lower:
print("Gemini-2.0 models do not support thinking - disabling thinking config")
else:
final_budget = 0 # Default disabled
final_budget = 0
if not thinking:
if "gemini-2.5-pro" in model_lower or "gemini-3-pro-preview" in model_lower:
if "gemini-2.5-pro" in model_lower or "gemini-3.1-pro-preview" in model_lower:
print("Pro models cannot have thinking turned off - defaulting thinking budget to -1")
final_budget = -1
else:
final_budget = thinking_budget
if ("gemini-2.5-pro" in model_lower or "gemini-3-pro-preview" in model_lower) and final_budget == 0:
if ("gemini-2.5-pro" in model_lower or "gemini-3.1-pro-preview" in model_lower) and final_budget == 0:
print("Pro models cannot have thinking turned off - defaulting thinking budget to -1")
final_budget = -1
t_config = types.ThinkingConfig(thinking_budget=final_budget)
tools = []
if google_search: tools.append(types.Tool(googleSearch=types.GoogleSearch()))
if url_context: tools.append(types.Tool(url_context=types.UrlContext()))
config = types.GenerateContentConfig(
temperature=temperature,
top_p=top_p,
seed=seed,
system_instruction=system_instruction.strip() if system_instruction else None,
thinking_config=t_config
thinking_config=t_config,
tools=tools if tools else None
)
response = client.models.generate_content(