275 lines
12 KiB
Python
275 lines
12 KiB
Python
import os
|
|
import json
|
|
import google.generativeai as genai
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import torch
|
|
import numpy as np
|
|
|
|
import comfy.utils
|
|
|
|
|
|
DEFAULT_API_KEY_PATH = r"C:\AI\Comfy\ComfyUI\custom_nodes\Creepy_nodes\assets\scripts\gemini_api_key.txt"
|
|
DEFAULT_THINKING_BUDGET = 4096
|
|
|
|
SAFETY_THRESHOLDS = ["Block None", "Block Low", "Block Medium", "Block High"]
|
|
|
|
SAFETY_THRESHOLD_MAP = {
|
|
"Block None": "BLOCK_NONE",
|
|
"Block Low": "BLOCK_LOW_AND_ABOVE",
|
|
"Block Medium": "BLOCK_MEDIUM_AND_ABOVE",
|
|
"Block High": "BLOCK_HIGH_AND_ABOVE",
|
|
}
|
|
|
|
SAFETY_CATEGORIES = [
|
|
"HARM_CATEGORY_HARASSMENT",
|
|
"HARM_CATEGORY_HATE_SPEECH",
|
|
"HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
|
"HARM_CATEGORY_DANGEROUS_CONTENT",
|
|
]
|
|
|
|
|
|
class GeminiAPI:
|
|
"""
|
|
A custom node for ComfyUI that uses the Google Gemini API for text and image generation
|
|
via the official google-generativeai library.
|
|
Includes optional image input (with resizing), system prompt, user instructions,
|
|
thinking mode, and safety settings control.
|
|
"""
|
|
|
|
CATEGORY = "Creepybits/API"
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
|
|
FUNCTION = "generate_text"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"system_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"model": (["gemini-2.0-flash", "gemini-2.0-flash-lite", "gemini-2.5-flash-lite-preview-09-2025", "gemini-2.5-flash-preview-09-2025", "gemini-2.0-flash-exp"],),
|
|
"max_output_tokens": ("INT", {"default": 1024, "min": 1, "max": 4096}),
|
|
"temperature": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
"top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"top_k": ("INT", {"default": 50, "min": 1, "max": 100}),
|
|
},
|
|
"optional": {
|
|
"user_instructions": ("STRING", {"multiline": True, "default": ""}),
|
|
"api_key_file": ("STRING", {"default": DEFAULT_API_KEY_PATH, "multiline": False}),
|
|
"image": ("IMAGE",),
|
|
"resize_image_to": (["None", "512", "768", "1024"], {"default": "None"}),
|
|
"thinking_mode": (["disable", "enable"], {"default": "disable"}),
|
|
"safety_threshold": (SAFETY_THRESHOLDS, {"default": "Block None"}),
|
|
}
|
|
}
|
|
|
|
|
|
def generate_text(self, system_prompt, model, max_output_tokens, temperature, top_p, seed, top_k, user_instructions="", api_key_file=None, image=None, resize_image_to="None", thinking_mode="disable", safety_threshold="Block None"):
|
|
"""
|
|
Generates text using the Google Gemini API via the google-generativeai library.
|
|
Handles optional image input (with resizing), system prompt, user instructions,
|
|
thinking mode, and safety settings.
|
|
"""
|
|
api_key = None
|
|
|
|
# --- API Key Handling ---
|
|
api_key = os.environ.get("GOOGLE_API_KEY")
|
|
|
|
if not api_key and api_key_file and os.path.exists(api_key_file):
|
|
try:
|
|
with open(api_key_file, 'r') as f:
|
|
api_key = f.read().strip()
|
|
except Exception as e:
|
|
print(f"Warning: Error reading Gemini API key file '{api_key_file}': {e}.")
|
|
elif not api_key and api_key_file and not os.path.exists(api_key_file):
|
|
print(f"Warning: GOOGLE_API_KEY environment variable not set and API key file not found at '{api_key_file}'.")
|
|
|
|
|
|
if not api_key:
|
|
return ("Error: Gemini API key not found. Please set the GOOGLE_API_KEY environment variable or provide a valid path to an API key file.",)
|
|
|
|
try:
|
|
genai.configure(api_key=api_key)
|
|
except Exception as e:
|
|
return (f"Error configuring Gemini API with key: {e}",)
|
|
|
|
|
|
# --- Prepare the 'contents' payload for the genai library ---
|
|
combined_prompt_text = ""
|
|
if system_prompt and system_prompt.strip():
|
|
combined_prompt_text += system_prompt.strip() + "\n\n"
|
|
|
|
if user_instructions and user_instructions.strip():
|
|
combined_prompt_text += user_instructions.strip()
|
|
|
|
contents = []
|
|
|
|
if combined_prompt_text or image is None:
|
|
contents.append(combined_prompt_text)
|
|
|
|
|
|
pil_image = None
|
|
if image is not None:
|
|
try:
|
|
image_data = image.numpy()
|
|
image_data = image_data[0]
|
|
image_data = (image_data * 255).astype(np.uint8)
|
|
|
|
if len(image_data.shape) == 2:
|
|
pil_image = Image.fromarray(image_data, 'L').convert('RGB')
|
|
elif len(image_data.shape) == 3 and image_data.shape[2] == 3:
|
|
pil_image = Image.fromarray(image_data, 'RGB')
|
|
elif len(image_data.shape) == 3 and image_data.shape[2] == 4:
|
|
pil_image = Image.fromarray(image_data, 'RGBA').convert('RGB')
|
|
else:
|
|
return ("Error: Could not process image format for API request.",)
|
|
|
|
|
|
# --- Resize Image if specified ---
|
|
if resize_image_to != "None" and pil_image is not None:
|
|
try:
|
|
target_size = int(resize_image_to)
|
|
width, height = pil_image.size
|
|
if height > width:
|
|
new_height = target_size
|
|
new_width = int(width * (target_size / height))
|
|
else:
|
|
new_width = target_size
|
|
new_height = int(height * (target_size / width))
|
|
|
|
new_height = max(1, new_height)
|
|
new_width = max(1, new_width)
|
|
|
|
pil_image = pil_image.resize((new_width, new_height))
|
|
except ValueError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
if pil_image is not None:
|
|
contents.append(pil_image)
|
|
|
|
|
|
except Exception as e:
|
|
return (f"Error: Failed to process image: {e}",)
|
|
|
|
if not contents:
|
|
return ("Error: No valid content provided. Ensure system prompt, user instructions, or image are present.",)
|
|
|
|
|
|
# --- Construct the generationConfig payload ---
|
|
generation_config = {
|
|
"max_output_tokens": max_output_tokens,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
# Note: Seed parameter removed as it caused a 400 error previously.
|
|
# The Gemini API documentation doesn't explicitly mention a seed parameter
|
|
# for this endpoint at the time of writing.
|
|
}
|
|
|
|
# --- Add thinkingConfig if thinking_mode is enabled ---
|
|
if thinking_mode == "enable":
|
|
generation_config["thinking_config"] = {"thinking_budget": DEFAULT_THINKING_BUDGET}
|
|
|
|
|
|
# --- Construct the safety_settings payload ---
|
|
safety_settings = []
|
|
api_threshold = SAFETY_THRESHOLD_MAP.get(safety_threshold, "BLOCK_NONE")
|
|
|
|
if api_threshold != "BLOCK_NONE":
|
|
for category_name in SAFETY_CATEGORIES:
|
|
safety_settings.append({
|
|
"category": category_name,
|
|
"threshold": api_threshold
|
|
})
|
|
|
|
|
|
# --- Get the Generative Model ---
|
|
try:
|
|
model_instance = genai.GenerativeModel(model)
|
|
except Exception as e:
|
|
if "Invalid model" in str(e):
|
|
return (f"Error: Invalid model name '{model}'. Check model availability.",)
|
|
return (f"Error: Could not get Gemini model '{model}'. Check model name, API status, or internet connection: {e}",)
|
|
|
|
|
|
# --- Send API Request using genai library ---
|
|
try:
|
|
response = model_instance.generate_content(
|
|
contents,
|
|
generation_config=generation_config,
|
|
safety_settings=safety_settings,
|
|
)
|
|
|
|
generated_text = ""
|
|
if hasattr(response, 'text'):
|
|
generated_text = response.text
|
|
elif hasattr(response, 'candidates') and response.candidates:
|
|
try:
|
|
if hasattr(response.candidates[0], 'content') and response.candidates[0].content and \
|
|
hasattr(response.candidates[0].content, 'parts') and response.candidates[0].content.parts:
|
|
generated_text = response.candidates[0].content.parts[0].text
|
|
except (AttributeError, IndexError, KeyError):
|
|
pass
|
|
|
|
|
|
if not generated_text:
|
|
if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
|
|
feedback_info = {}
|
|
if hasattr(response.prompt_feedback, 'block_reason') and response.prompt_feedback.block_reason and hasattr(response.prompt_feedback.block_reason, 'name') and response.prompt_feedback.block_reason.name != 'UNASSIGNED':
|
|
feedback_info['block_reason'] = response.prompt_feedback.block_reason.name
|
|
|
|
if hasattr(response.prompt_feedback, 'safety_ratings') and response.prompt_feedback.safety_ratings:
|
|
feedback_info['safety_ratings'] = []
|
|
for rating in response.prompt_feedback.safety_ratings:
|
|
if hasattr(rating, 'probability') and hasattr(rating.probability, 'name') and rating.probability.name != 'UNSPECIFIED':
|
|
rating_info = {"category": rating.category.name, "probability": rating.probability.name}
|
|
feedback_info['safety_ratings'].append(rating_info)
|
|
|
|
if 'block_reason' in feedback_info:
|
|
return (f"Error: Content generation blocked due to {feedback_info['block_reason']}. Check console for safety feedback.",)
|
|
elif 'safety_ratings' in feedback_info and feedback_info['safety_ratings']:
|
|
return ("Error: No text generated. High safety ratings. Check console for details.",)
|
|
return ("Error: No text generated. API returned empty content/candidates. Check Prompt Feedback in console.",)
|
|
else:
|
|
if hasattr(response, 'candidates') and not response.candidates:
|
|
return ("Error: API returned no candidates (likely blocked). No prompt feedback details available.",)
|
|
return ("Error: No text generated from the API and no specific feedback provided. Check console.",)
|
|
|
|
|
|
return (generated_text,)
|
|
|
|
|
|
except Exception as e:
|
|
error_message = str(e)
|
|
if hasattr(e, 'response') and e.response:
|
|
if hasattr(e.response, 'text'):
|
|
try:
|
|
error_json = json.loads(e.response.text)
|
|
if 'error' in error_json and 'message' in error_json['error']:
|
|
api_error_message = error_json['error']['message']
|
|
error_message = f"API Error: {api_error_message}"
|
|
elif hasattr(e.response, 'status_code') and hasattr(e.response, 'reason'):
|
|
error_message = f"HTTP Error: {e.response.status_code} {e.response.reason}"
|
|
except (json.JSONDecodeError, AttributeError):
|
|
if hasattr(e.response, 'status_code') and hasattr(e.response, 'reason'):
|
|
error_message = f"HTTP Error: {e.response.status_code} {e.response.reason}"
|
|
pass
|
|
|
|
elif hasattr(e.response, 'status_code') and hasattr(e.response, 'reason'):
|
|
error_message = f"HTTP Error: {e.response.status_code} {e.response.reason}"
|
|
|
|
return (f"Error: An error occurred during Gemini API call: {error_message}",)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GeminiAPI": GeminiAPI,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GeminiAPI": "Gemini 2.5 Flash/Pro API (Creepybits)",
|
|
}
|