nodes/Creepybits-ComfyUI-Creepy_nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
Creepybits
2025-05-30 20:33:48 +02:00
committed by GitHub
parent 37d0b75735
commit f340076cfe
1 changed files with 291 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
assets/nodes/Coloring.py
+291
View File
@@ -0,0 +1,291 @@
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 Coloring:
"""
A custom node for ComfyUI that uses the Google Gemini API for text and image generation.
"""
CATEGORY = "Creepybits/Utilities"
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.5-flash-preview-04-17", "gemini-2.5-pro-exp-03-25", "gemini-2.0-flash", "gemini-2.0-flash-exp"],),
"max_output_tokens": ("INT", {"default": 512, "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}),
"top_k": ("INT", {"default": 50, "min": 1, "max": 100}),
},
"optional": {
"api_key_file": ("STRING", {"default": DEFAULT_API_KEY_PATH, "multiline": False}),
"ref_image": ("IMAGE",),
"color_image": ("IMAGE",),
"resize_image_to": (["None", "512", "768", "1024"], {"default": "768"}),
"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, top_k, api_key_file=None, ref_image=None, color_image=None, resize_image_to="None", thinking_mode="disable", safety_threshold="Block None", **kwargs):
"""Generates text using the Google Gemini API."""
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}",)
combined_prompt_text = ""
if system_prompt and system_prompt.strip():
combined_prompt_text += system_prompt.strip() + "\n\n"
contents = []
# --- REFERENCE FUNCTION ---
def process_image(image, image_name="image"): #Pass also the image name
pil_image = 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:
print(f"Warning: Unexpected image data shape {image_data.shape}. Attempting conversion to RGB.")
try:
pil_image = Image.fromarray(image_data).convert('RGB')
except Exception as convert_e:
print(f"Error converting {image_name} format: {convert_e}") #Use the image name in the log
return None #Return none if the image failed to convert
return pil_image
except Exception as e:
print(f"Error processing {image_name} for API request: {e}") #Use the image name in the log
return None #Return none if the image failed to convert
# --- REFERENCE FUNCTION ENDS ---
processed_ref_image = None
if ref_image is not None:
processed_ref_image = process_image(ref_image, "ref_image") # Process ref_image
if processed_ref_image is None:
return ("Error: Failed to process ref_image",) #If ref image fails to convert, error
processed_color_image = None
if color_image is not None:
processed_color_image = process_image(color_image, "color_image") #Process color image
if processed_color_image is None:
return ("Error: Failed to process color_image",) # If color image fails to convert, error
# Resize the ref image, and color image
pil_image = processed_ref_image
pil_image2 = processed_color_image
# Resize Image
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)
print(f"Resizing image from {width}x{height} to {new_width}x{new_height}")
pil_image = pil_image.resize((new_width, new_height))
except ValueError:
print(f"Warning: Invalid resize_image_to value '{resize_image_to}'. Skipping resize.")
except Exception as e:
print(f"Error during image resizing: {e}. Skipping resize.")
if resize_image_to != "None" and pil_image2 is not None:
try:
target_size = int(resize_image_to)
width, height = pil_image2.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)
print(f"Resizing image from {width}x{height} to {new_width}x{new_height}")
pil_image2 = pil_image2.resize((new_width, new_height))
except ValueError:
print(f"Warning: Invalid resize_image_to value '{resize_image_to}'. Skipping resize.")
except Exception as e:
print(f"Error during image resizing: {e}. Skipping resize.")
#Check which pil image to use and add it to contents list
if pil_image is not None:
contents.append(pil_image) # Add processed ref_image to contents
if pil_image2 is not None:
contents.append(pil_image2) # Add processed color_image to contents
# add prompt text to list
if combined_prompt_text or not contents:
contents.append(combined_prompt_text)
#Ensure the list isn't empty to send
if not contents:
return ("Error: No valid content provided. Ensure system prompt, user instructions, or image are present.",)
generation_config = {
"max_output_tokens": max_output_tokens,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
}
if thinking_mode == "enable":
generation_config["thinking_config"] = {"thinking_budget": DEFAULT_THINKING_BUDGET}
print(f"Thinking mode enabled with budget {DEFAULT_THINKING_BUDGET}.")
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
})
print(f"Applying safety threshold '{api_threshold}' to categories: {SAFETY_CATEGORIES}")
else:
print("Safety threshold set to 'Block None'. No safety settings applied.")
try:
model_instance = genai.GenerativeModel(model)
except Exception as e:
print(f"Error getting model '{model}': {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}",)
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):
print("Warning: Could not extract text from response.candidates structure.")
if not generated_text:
print("Generated text is empty or extraction failed.")
if hasattr(response, 'prompt_feedback') and response.prompt_feedback:
print("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':
print(f" Block Reason: {response.prompt_feedback.block_reason.name}")
feedback_info['block_reason'] = response.prompt_feedback.block_reason.name
if hasattr(response.prompt_feedback, 'safety_ratings') and response.prompt_feedback.safety_ratings:
print(" 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}
print(f" Category: {rating_info['category']}, Probability: {rating_info['probability']}")
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:
print(f"An error occurred during Gemini API call: {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']
print(f"API Error Message from response: {api_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}"
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 = {
"Coloring": Coloring,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Coloring": "Coloring Node (Creepybits)",
}