1025 lines
49 KiB
Python
1025 lines
49 KiB
Python
import subprocess
|
|
import sys
|
|
import json
|
|
import os, glob
|
|
import re
|
|
import base64
|
|
import numpy as np
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import torch
|
|
|
|
# For the Websearch
|
|
from googlesearch import search
|
|
import requests
|
|
from bs4 import BeautifulSoup
|
|
|
|
api_key_oai = os.getenv("OPENAI_API_KEY")
|
|
api_key_anthropic = os.getenv("ANTHROPIC_API_KEY")
|
|
api_key_xai = os.getenv("XAI_API_KEY")
|
|
api_key_deepseek = os.getenv("DEEPSEEK_API_KEY")
|
|
api_key_gemini = os.getenv("GEMINI_API_KEY")
|
|
|
|
# Get the node list
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
json_file_path = os.path.join(script_dir, "../ComfyUI-Manager/custom-node-list.json")
|
|
json_file_path = os.path.normpath(json_file_path)
|
|
|
|
# Get the directory of the current script (Polymath.py)
|
|
script_directory = os.path.dirname(os.path.abspath(__file__))
|
|
custom_instructions_directory = os.path.join(script_directory, 'custom_instructions')
|
|
|
|
# Construct the full path to config.json
|
|
config_path = os.path.join(script_directory, 'config.json')
|
|
|
|
def install_and_import(package):
|
|
try:
|
|
__import__(package)
|
|
except ImportError:
|
|
print(f"Package {package} not found. Installing...")
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
|
|
finally:
|
|
globals()[package] = __import__(package)
|
|
|
|
install_and_import("googlesearch")
|
|
install_and_import("requests")
|
|
install_and_import("bs4")
|
|
|
|
script_directory = os.path.dirname(os.path.abspath(__file__))
|
|
config_path = os.path.join(script_directory, 'config.json')
|
|
ollama_url = "http://127.0.0.1:11434/api/tags"
|
|
|
|
def load_models():
|
|
config_models = {}
|
|
base_urls = {}
|
|
try:
|
|
with open(config_path, 'r') as f:
|
|
config_data = json.load(f)
|
|
# Iterate over all configuration entries
|
|
for conf in config_data:
|
|
models = conf.get('models', {})
|
|
config_models.update(models)
|
|
# Store each baseurl keyed by something, e.g., the first model name
|
|
for name in models.keys():
|
|
base_urls[name] = conf.get('baseurl')
|
|
except FileNotFoundError:
|
|
print(f"Error: config.json not found at {config_path}")
|
|
except json.JSONDecodeError:
|
|
print(f"Error: Could not decode JSON from {config_path}")
|
|
|
|
# Fetch Ollama models
|
|
ollama_models = {}
|
|
try:
|
|
response = requests.get(ollama_url)
|
|
response.raise_for_status()
|
|
models = response.json().get('models', [])
|
|
ollama_models = {
|
|
"Ollama: " + " ".join(model['name'].split(':')): model['name']
|
|
for model in models
|
|
}
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error fetching Ollama models: {e}")
|
|
|
|
# Fetch OpenAI models
|
|
openai_models = {}
|
|
openai_api_key = os.getenv("OPENAI_API_KEY") # Or load from config
|
|
openai_base_url = "https://api.openai.com/v1"
|
|
if openai_api_key:
|
|
try:
|
|
headers = {"Authorization": f"Bearer {openai_api_key}"}
|
|
response = requests.get(f"{openai_base_url}/models", headers=headers)
|
|
response.raise_for_status()
|
|
models = response.json().get('data', [])
|
|
excluded_models = ("tts", "omni", "whisper", "computer", "text", "davinci", "babbage", "codex")
|
|
openai_models = {
|
|
f"OpenAI: {model['id']}": model['id']
|
|
for model in models
|
|
if not model['id'].startswith(excluded_models)
|
|
}
|
|
# Update base_urls for OpenAI models
|
|
for model_name in openai_models.keys():
|
|
base_urls[model_name] = openai_base_url
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error fetching OpenAI models: {e}")
|
|
else:
|
|
print("Warning: OpenAI API key not found. Skipping OpenAI models.")
|
|
|
|
merged_models = { **openai_models, **config_models, **ollama_models}
|
|
models_list = list(merged_models.keys())
|
|
default_model = models_list[0] if models_list else None
|
|
|
|
print("\033[92mPolymath loaded these models: \033[0m", f"\033[38;5;99m{models_list}\033[0m")
|
|
return base_urls, merged_models, models_list, default_model
|
|
|
|
# Load models and set default
|
|
base_url, models_dict, models_list, default_model = load_models()
|
|
|
|
def load_custom_instructions():
|
|
instructions = ["None"]
|
|
if not os.path.exists(custom_instructions_directory):
|
|
os.makedirs(custom_instructions_directory)
|
|
for filename in os.listdir(custom_instructions_directory):
|
|
if filename.endswith(".txt"):
|
|
instructions.append(filename[:-4])
|
|
return instructions
|
|
|
|
custom_instructions_list = load_custom_instructions()
|
|
|
|
def fetch_and_extract_content(url):
|
|
try:
|
|
response = requests.get(url, timeout=10)
|
|
response.raise_for_status()
|
|
soup = BeautifulSoup(response.content, 'html.parser')
|
|
paragraphs = soup.find_all('p')
|
|
text_content = "\n".join([p.text for p in paragraphs])
|
|
return text_content.strip()
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error fetching URL {url}: {e}")
|
|
return None
|
|
except Exception as e:
|
|
print(f"Error processing content from {url}: {e}")
|
|
return None
|
|
|
|
class OllamaAPI:
|
|
def __init__(self, base_url="http://localhost:11434"):
|
|
self.base_url = base_url
|
|
self.chat_history = []
|
|
|
|
def generate_completion(self, model, prompt, stream=False, ollama_chat_mode=False, images=None, seed=42, options=None):
|
|
url = f"{self.base_url}/api/chat" if ollama_chat_mode else f"{self.base_url}/api/generate"
|
|
|
|
payload_options = options if isinstance(options, dict) else {} # Ensure options is a dict
|
|
payload_options['seed'] = seed
|
|
|
|
if ollama_chat_mode:
|
|
# Append the new user message to the chat history
|
|
self.chat_history.append({"role": "user", "content": prompt})
|
|
payload = {
|
|
"model": model,
|
|
"messages": self.chat_history, # Include the entire chat history
|
|
"stream": stream,
|
|
"options": payload_options
|
|
}
|
|
else:
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"stream": stream,
|
|
"options": payload_options
|
|
}
|
|
|
|
if images:
|
|
payload["images"] = images
|
|
if options:
|
|
payload["temperature"] = options.get("temperature", 0.8)
|
|
payload["top_p"] = options.get("top_p", 0.95)
|
|
payload["top_k"] = options.get("top_k", 40)
|
|
payload["num_predict"] = options.get("max_output_tokens", -1)
|
|
payload["seed"] = options.get("seed", 42)
|
|
payload["keep_alive"] = options.get("ollama_keep_alive", 5)
|
|
|
|
response = requests.post(url, json=payload)
|
|
response.raise_for_status()
|
|
response_json = response.json()
|
|
|
|
if ollama_chat_mode:
|
|
# Append the model's response to the chat history
|
|
self.chat_history.append({"role": "assistant", "content": response_json.get('message', {}).get('content', '')})
|
|
|
|
return response_json
|
|
|
|
class PolymathSettings:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"temperature": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 2.0, "step": 0.05}),
|
|
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"top_k": ("INT", {"default": 40, "min": 0, "max": 100}),
|
|
"max_output_tokens": ("INT", {"default": 1024, "min": -1, "max": 65536}),
|
|
"response_format_json": ("BOOLEAN", {"default": False}),
|
|
"ollama_keep_alive": ("INT", {"default": 5, "min": 1, "max": 10}),
|
|
"request_timeout": ("INT", {"default": 120, "min": 0, "max": 600}),
|
|
"dalle_quality": (["standard", "hd"], {"default":"standard"}),
|
|
"dalle_style": (["vivid", "natural"], {"default":"vivid"}),
|
|
"dalle_size": (["1024x1024", "1792x1024", "1024x1792"], {"default":"1024x1024"}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4}),
|
|
"gpt_image_quality": (["low", "medium", "high", "auto"], {"default":"auto"}),
|
|
"gpt_image_background": (["transparent", "opaque", "auto"], {"default":"opaque"}),
|
|
"gpt_image_size": (["1024x1024", "1536x1024", "1024x1536", "auto"], {"default":"auto"}),
|
|
"gemini_aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "2:3", "3:2", "4:5", "5:4", "21:9"], {"default": "1:1"}),
|
|
"gemini_image_size": (["1K", "2K", "4K"], {"default": "1K"})
|
|
}
|
|
}
|
|
# Output a dictionary or a specific tuple/pipe format
|
|
RETURN_TYPES = ("LLM_SETTINGS",) # Or use "BASIC_PIPE" if it fits
|
|
FUNCTION = "get_settings"
|
|
CATEGORY = "Polymath/Settings"
|
|
|
|
def get_settings(self, **kwargs):
|
|
# Package settings into a dictionary for easy use
|
|
return (kwargs,) # Return as a tuple containing the dict
|
|
|
|
class Polymath:
|
|
chat_history = []
|
|
default_model = None
|
|
ollama_api = OllamaAPI() # Class-level instance of OllamaAPI
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": (
|
|
"STRING",
|
|
{
|
|
"default": "Enter your prompt here. Use {additional_text} as a placeholder if needed.",
|
|
"multiline": True
|
|
}
|
|
),
|
|
"additional_text": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True
|
|
}
|
|
),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xfffffff, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"llm_settings": ("LLM_SETTINGS",),
|
|
"model": (models_list, {"default": default_model}),
|
|
"custom_instruction": (custom_instructions_list, {"default": "None"}),
|
|
"enable_web_search": ("BOOLEAN", {"default": False}),
|
|
"list_sources": ("BOOLEAN", {"default": False}),
|
|
"num_search_results": ("INT", {"default": 5, "min": 1, "max": 10}),
|
|
"keep_context": ("BOOLEAN", {"default": True}),
|
|
"ollama_chat_mode": ("BOOLEAN", {"default": False}),
|
|
"compress": ("BOOLEAN", {"default": False}),
|
|
"compression_level": (["soft", "medium", "hard"],),
|
|
"console_log": ("BOOLEAN", {"default": True}),
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING","IMAGE")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "Polymath"
|
|
|
|
def tensor_to_pil(self, img_tensors):
|
|
def to_pil(img_tensor):
|
|
arr = img_tensor.cpu().numpy() if hasattr(img_tensor, "cpu") else np.array(img_tensor)
|
|
arr = np.squeeze(arr)
|
|
if arr.ndim == 3 and arr.shape[0] == 3:
|
|
arr = np.transpose(arr, (1, 2, 0))
|
|
arr = np.clip(255.0 * arr, 0, 255).astype(np.uint8)
|
|
return Image.fromarray(arr)
|
|
|
|
if isinstance(img_tensors, (list, tuple)):
|
|
return [to_pil(t) for t in img_tensors]
|
|
elif hasattr(img_tensors, 'ndim') and img_tensors.ndim == 4: # Batch tensor
|
|
return [to_pil(t) for t in img_tensors]
|
|
else:
|
|
return to_pil(img_tensors)
|
|
|
|
def encode_image(self, image):
|
|
def encode_single(img):
|
|
buffered = BytesIO()
|
|
img.save(buffered, format="PNG")
|
|
return base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
if isinstance(image, list):
|
|
return [encode_single(img) for img in image]
|
|
else:
|
|
return encode_single(image)
|
|
|
|
def execute(
|
|
self, prompt, additional_text, seed, llm_settings=None, model=None, custom_instruction="None", enable_web_search=False, list_sources=False, num_search_results=3, keep_context=True, ollama_chat_mode=False, compress=False, compression_level="soft", console_log=True, image=None
|
|
):
|
|
selected_model_value = models_dict.get(model)
|
|
settings = llm_settings if isinstance(llm_settings, dict) else {}
|
|
|
|
if not selected_model_value:
|
|
print(f"Error: Model '{model}' not found in config.json. Using default.")
|
|
selected_model_value = list(models_dict.values())[0] if models_dict else None
|
|
if not selected_model_value:
|
|
return ("Error: No valid model found.",)
|
|
|
|
selected_base_url = base_url.get(model)
|
|
|
|
if console_log:
|
|
print(
|
|
f"\033[92mPolymath Chat Request: \033[0mModel='{model}', Prompt='{prompt}', Additional Text='{additional_text}'",
|
|
f"Seed='{seed}', Web Search Enabled={enable_web_search}, Keep Context={keep_context},"
|
|
f"Ollama Chat Mode={ollama_chat_mode}, Custom Instruction='{custom_instruction}'"
|
|
)
|
|
|
|
# Process additional_text using the given prompt:
|
|
# If the prompt contains the placeholder "{additional_text}", then replace it with the additional_text input.
|
|
# Otherwise, if additional_text is provided, append it at the end.
|
|
if "{additional_text}" in prompt:
|
|
augmented_prompt = prompt.replace("{additional_text}", additional_text)
|
|
elif additional_text.strip():
|
|
augmented_prompt = prompt + "\n\n" + additional_text
|
|
else:
|
|
augmented_prompt = prompt
|
|
|
|
if enable_web_search:
|
|
urls_in_prompt = re.findall(r'(https?://\S+)', augmented_prompt)
|
|
try:
|
|
search_results_content = []
|
|
searched_urls = set()
|
|
for url in urls_in_prompt:
|
|
if url not in searched_urls:
|
|
print(f"Fetching content from URL in prompt: {url}")
|
|
content = fetch_and_extract_content(url)
|
|
if content:
|
|
search_results_content.append(
|
|
f"Source (from prompt): {url}\nContent:\n{content}\n---\n"
|
|
)
|
|
else:
|
|
search_results_content.append(
|
|
f"Source (from prompt): {url}\nCould not retrieve content.\n---\n"
|
|
)
|
|
searched_urls.add(url)
|
|
for i, url in enumerate(search(augmented_prompt, num_results=num_search_results)):
|
|
if url not in searched_urls:
|
|
print(f"Fetching content from search result {i+1}: {url}")
|
|
content = fetch_and_extract_content(url)
|
|
if content:
|
|
search_results_content.append(
|
|
f"Search Result {i+1}:\nSource: {url}\nContent:\n{content}\n---\n"
|
|
)
|
|
else:
|
|
search_results_content.append(
|
|
f"Search Result {i+1}:\nSource: {url}\nCould not retrieve content.\n---\n"
|
|
)
|
|
searched_urls.add(url)
|
|
if search_results_content:
|
|
augmented_prompt = (
|
|
f"Original query: {augmented_prompt}\n\nWeb search results and linked content:\n\n"
|
|
f"{''.join(search_results_content)}\n\nBased on this information, please provide a response."
|
|
)
|
|
else:
|
|
augmented_prompt = (
|
|
f"Original query: {augmented_prompt}\n\nNo relevant web search results or linked content found. "
|
|
"Please proceed with the original query."
|
|
)
|
|
if console_log:
|
|
print("No relevant web search results or linked content found.")
|
|
except Exception as e:
|
|
print(f"Web Search Error: {e}")
|
|
augmented_prompt = (
|
|
f"Original query: {augmented_prompt}\n\nAn error occurred during web search or fetching linked content. "
|
|
"Please proceed with the original query."
|
|
)
|
|
|
|
if compress:
|
|
compression_chars = {
|
|
"soft": 250,
|
|
"medium": 150,
|
|
"hard": 75,
|
|
}
|
|
char_limit = compression_chars[compression_level]
|
|
augmented_prompt += (
|
|
f" Compress the output to be concise while retaining key visual details. MAX OUTPUT SIZE no more than "
|
|
f"{char_limit} characters."
|
|
)
|
|
if console_log:
|
|
print(f"\033[92mFull Polymath Prompt\033[0m", augmented_prompt)
|
|
|
|
if list_sources:
|
|
augmented_prompt += " Always list all fetched sources."
|
|
|
|
if custom_instruction == "node finder":
|
|
with open(json_file_path, "r", encoding="utf-8") as file:
|
|
json_data = json.load(file)
|
|
filtered_data = [
|
|
f"Title: {node['title']}\nReference: {node['reference']}\nDescription: {node['description']}\n"
|
|
for node in json_data.get("custom_nodes", [])
|
|
]
|
|
data = "\n".join(filtered_data)
|
|
print("Filtered JSON Data:")
|
|
print(data)
|
|
augmented_prompt += data
|
|
|
|
b64 = None
|
|
if image is not None:
|
|
pil_images = self.tensor_to_pil(image)
|
|
b64 = self.encode_image(pil_images)
|
|
if isinstance(b64, str):
|
|
b64 = [b64]
|
|
|
|
|
|
response = self.polymath_interaction(
|
|
selected_base_url,
|
|
api_key_oai,
|
|
api_key_anthropic,
|
|
api_key_xai,
|
|
api_key_deepseek,
|
|
api_key_gemini,
|
|
selected_model_value,
|
|
augmented_prompt,
|
|
console_log,
|
|
keep_context,
|
|
ollama_chat_mode,
|
|
custom_instruction,
|
|
seed,
|
|
llm_settings=settings,
|
|
b64=b64
|
|
)
|
|
return response
|
|
|
|
def polymath_interaction(self, selected_base_url, api_key_oai, api_key_anthropic, api_key_xai, api_key_deepseek, api_key_gemini, model_value, prompt, console_log, keep_context, ollama_chat_mode, custom_instruction, seed=42, llm_settings={}, b64=None):
|
|
settings = llm_settings
|
|
|
|
# GPT Image 1 specific handling
|
|
if model_value.startswith('gpt-image'):
|
|
import requests
|
|
import json
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import base64
|
|
import numpy as np
|
|
import torch
|
|
|
|
image_params = {
|
|
"model": model_value,
|
|
"n": settings.get("batch_size", 1),
|
|
"size": settings.get("gpt_image_size", "1024x1024"),
|
|
"quality": settings.get("gpt_image_quality", "auto"),
|
|
"background": settings.get("gpt_image_background", "opaque"),
|
|
}
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key_oai}"
|
|
}
|
|
|
|
files = {}
|
|
|
|
if console_log:
|
|
print(f"\033[92mGPT-Image-1 Params:\033[0m", {k: v for k, v in image_params.items()})
|
|
|
|
if b64:
|
|
url = f"{selected_base_url}/edits"
|
|
for i, img_str in enumerate(b64):
|
|
try:
|
|
image_bytes = base64.b64decode(img_str)
|
|
files[f"image[{i}]"] = (f"image_{i}.png", BytesIO(image_bytes), "image/png")
|
|
except Exception as e:
|
|
print(f"Error decoding base64 image {i+1}: {e}")
|
|
empty_tensor = torch.zeros((1, 3, 512, 512))
|
|
return ("Error processing input image", empty_tensor)
|
|
files["prompt"] = (None, prompt) # Add prompt to files
|
|
response = requests.post(url, headers=headers, files=files, data=image_params) # Use data for image_params
|
|
else:
|
|
url = f"{selected_base_url}/generations"
|
|
image_params["prompt"] = prompt # Add prompt here for generations
|
|
response = requests.post(url, headers=headers, json=image_params)
|
|
|
|
# 1. Check for basic HTTP errors first
|
|
try:
|
|
response.raise_for_status() # Raises an HTTPError for bad status codes (4xx or 5xx)
|
|
except requests.exceptions.HTTPError as http_err:
|
|
error_detail = f"HTTP Error: {http_err}"
|
|
try:
|
|
# Try to get more specific error detail from response body
|
|
error_detail += f"\nResponse Body: {response.text}"
|
|
except Exception:
|
|
pass # Ignore if response text isn't available
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
# Return error message as string, None for image, matching RETURN_TYPES
|
|
return (error_detail, None)
|
|
except Exception as req_err:
|
|
error_detail = f"Request Exception: {req_err}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
return (error_detail, None)
|
|
|
|
|
|
# 2. Try to parse the JSON response
|
|
try:
|
|
response_data = response.json()
|
|
except json.JSONDecodeError:
|
|
error_detail = f"JSON Decode Error: Failed to parse API response.\nResponse Text: {response.text}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
return (error_detail, None)
|
|
|
|
# 3. Check for "error" key in the parsed JSON (similar to FL_GPT_Image1.py)
|
|
if "error" in response_data:
|
|
error_msg = response_data.get("error", "Unknown API error structure")
|
|
error_detail = f"OpenAI API Error: {json.dumps(error_msg, indent=2)}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
# Special check for organization verification, mentioned in FL_GPT_Image1.py
|
|
if isinstance(error_msg, dict) and "organization verification" in error_msg.get("message", "").lower():
|
|
print("\033[93mNote: This error often indicates OpenAI requires organization verification for GPT-image-1 access.\033[0m")
|
|
error_detail += "\n(Hint: Check OpenAI organization verification requirements)"
|
|
return (error_detail, None)
|
|
|
|
# 4. Check for "data" key existence and ensure it's not empty (similar to FL_GPT_Image1.py)
|
|
if "data" not in response_data or not response_data.get("data"):
|
|
error_detail = f"API Response Error: No 'data' key found or 'data' is empty.\nFull Response: {json.dumps(response_data, indent=2)}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
return (error_detail, None)
|
|
|
|
image_tensors = []
|
|
output_text = f"Generated {len(response_data['data'])} image(s) using these parameters {image_params}"
|
|
|
|
# loop to process images
|
|
for i, img_data in enumerate(response_data["data"]):
|
|
if "b64_json" in img_data:
|
|
try:
|
|
img_bytes = base64.b64decode(img_data["b64_json"])
|
|
pil_image = Image.open(BytesIO(img_bytes))
|
|
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
img_array = np.array(pil_image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
if console_log: print(f"Successfully processed image {i+1} from base64")
|
|
|
|
except Exception as e:
|
|
error_detail = f"Error processing base64 image {i+1}: {str(e)}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
return (f"Image Processing Error: {error_detail}", None)
|
|
|
|
elif "url" in img_data:
|
|
try:
|
|
url_response = requests.get(img_data["url"], timeout=30)
|
|
url_response.raise_for_status()
|
|
|
|
pil_image = Image.open(BytesIO(url_response.content))
|
|
if pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
|
|
img_array = np.array(pil_image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
|
|
if console_log: print(f"Successfully downloaded and processed image {i+1} from URL")
|
|
|
|
except Exception as e:
|
|
error_detail = f"Error downloading/processing image {i+1} from URL {img_data['url']}: {str(e)}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
return (f"Image URL Processing Error: {error_detail}", None)
|
|
|
|
else:
|
|
error_detail = f"No image data ('b64_json' or 'url') found in response item {i+1}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
# This case is unlikely if the 'data' check passed, but good to have
|
|
return (f"API Response Format Error: {error_detail}", None)
|
|
|
|
if not image_tensors:
|
|
error_detail = "No images were successfully processed despite API success."
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
return (error_detail, None)
|
|
elif len(image_tensors) == 1:
|
|
final_image_output = image_tensors[0]
|
|
else:
|
|
# Check tensor shapes before concatenating
|
|
first_shape = image_tensors[0].shape
|
|
if not all(t.shape == first_shape for t in image_tensors):
|
|
error_detail = f"Cannot batch images with different shapes: {[t.shape for t in image_tensors]}"
|
|
print(f"\033[91m!!! Polymath GPT-Image-1 Error !!!\033[0m\n{error_detail}")
|
|
# Maybe just return the first image? Or error out? Let's error.
|
|
return (error_detail, None)
|
|
final_image_output = torch.cat(image_tensors, dim=0)
|
|
|
|
return (output_text, final_image_output)
|
|
|
|
elif model_value.startswith('dall'):
|
|
from openai import OpenAI
|
|
import base64 # Ensure base64 is available
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import numpy as np
|
|
import torch
|
|
|
|
client = OpenAI(api_key=api_key_oai, base_url=selected_base_url)
|
|
dalle_params = {
|
|
"model": model_value,
|
|
"prompt": prompt,
|
|
"n": settings.get("batch_size", 1),
|
|
"size": settings.get("dalle_size", "1024x1024"),
|
|
"response_format": "b64_json"
|
|
}
|
|
|
|
if model_value in ("dall-e-3"):
|
|
dalle_params["quality"] = settings.get("dalle_quality", "standard")
|
|
dalle_params["style"] = settings.get("dalle_style", "vivid")
|
|
elif console_log:
|
|
print("\033[93mWarning:\033[0m `quality` and `style` parameters are only supported for dall-e-3. Ignoring for older DALL-E models.")
|
|
|
|
if console_log:
|
|
print(f"\033[92mDALL-E Params:\033[0m", dalle_params)
|
|
|
|
request_timeout = settings.get("request_timeout", 120) or None
|
|
try:
|
|
response = client.images.generate(**dalle_params, timeout=request_timeout)
|
|
except Exception as e:
|
|
error_message = f"DALL-E API Error: {e}"
|
|
print(f"\033[91m!!! Polymath DALL-E Error !!!\033[0m\n{error_message}")
|
|
return (error_message, None)
|
|
|
|
image_tensors = []
|
|
output_text = response.data[0].revised_prompt if response.data[0].revised_prompt else "(No revised prompt provided)"
|
|
|
|
for i, img_data in enumerate(response.data):
|
|
try:
|
|
decoded_data = base64.b64decode(img_data.b64_json)
|
|
image = Image.open(BytesIO(decoded_data)).convert("RGB")
|
|
img_array = np.array(image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
except Exception as e:
|
|
error_message = f"Error processing DALL-E image data {i+1}: {e}"
|
|
print(f"\033[91m!!! Polymath DALL-E Error !!!\033[0m\n{error_message}")
|
|
return (error_message, None) # Handle image decode errors
|
|
|
|
final_image_output = torch.cat(image_tensors, dim=0) if image_tensors else None # Handle no images
|
|
return (output_text, final_image_output)
|
|
|
|
# OpenAI-based models (e.g., GPT, o1, o3)
|
|
elif model_value.startswith(('gpt', 'o1', 'o3', 'o3', 'chatgpt')):
|
|
from openai import OpenAI
|
|
client = OpenAI(api_key=api_key_oai, base_url=selected_base_url)
|
|
messages = []
|
|
if custom_instruction != "None":
|
|
try:
|
|
with open(os.path.join(custom_instructions_directory, f"{custom_instruction}.txt"), 'r') as f:
|
|
system_instruction = f.read()
|
|
messages.append({"role": "system", "content": system_instruction})
|
|
except Exception as e:
|
|
print("Error reading custom instruction file:", e)
|
|
if keep_context and self.chat_history:
|
|
messages.extend(self.chat_history)
|
|
if b64:
|
|
image_data = {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}
|
|
user_content = [
|
|
{"type": "text", "text": prompt},
|
|
image_data
|
|
]
|
|
messages.append({"role": "user", "content": user_content})
|
|
messages.append({"role": "user", "content": prompt})
|
|
api_params = {
|
|
"model": model_value,
|
|
"messages": messages,
|
|
"stream": False,
|
|
"temperature": settings.get("temperature", 0.8),
|
|
"top_p": settings.get("top_p", 0.95),
|
|
"max_tokens": settings.get("max_output_tokens", 1024),
|
|
"seed": seed,
|
|
#"modalities": modalities,
|
|
}
|
|
stop_seq = settings.get("stop_sequences", "")
|
|
if stop_seq:
|
|
api_params["stop"] = [s.strip() for s in stop_seq.split(',')]
|
|
request_timeout = settings.get("request_timeout", 120) or None
|
|
completion = client.chat.completions.create(
|
|
**api_params,
|
|
timeout=request_timeout
|
|
)
|
|
output_text = completion.choices[0].message.content
|
|
self.chat_history.extend([{"role": "user", "content": prompt},
|
|
{"role": "assistant", "content": output_text}])
|
|
return (output_text,)
|
|
|
|
# Anthropic (e.g., Claude models) branch
|
|
elif model_value.startswith('claude'):
|
|
from anthropic import Anthropic
|
|
client = Anthropic(api_key=api_key_anthropic)
|
|
messages = []
|
|
if custom_instruction != "None":
|
|
try:
|
|
with open(os.path.join(custom_instructions_directory, f"{custom_instruction}.txt"), 'r') as f:
|
|
system_instruction = f.read()
|
|
messages.append({"role": "system", "content": system_instruction})
|
|
except Exception as e:
|
|
print("Error reading custom instruction file:", e)
|
|
if keep_context and self.chat_history:
|
|
messages.extend(self.chat_history)
|
|
messages.append({"role": "user", "content": prompt})
|
|
completion = client.messages.create(
|
|
model=model_value,
|
|
max_tokens=1024,
|
|
messages=messages
|
|
#seed=seed not yet supported
|
|
)
|
|
output_text = completion.content
|
|
self.chat_history.extend([{"role": "user", "content": prompt},
|
|
{"role": "assistant", "content": output_text}])
|
|
return (output_text,)
|
|
|
|
# x.ai / Grok branch
|
|
elif model_value.startswith('grok'):
|
|
from openai import OpenAI
|
|
client = OpenAI(api_key=api_key_xai, base_url=selected_base_url)
|
|
messages = []
|
|
if custom_instruction != "None":
|
|
try:
|
|
with open(os.path.join(custom_instructions_directory, f"{custom_instruction}.txt"), 'r') as f:
|
|
system_instruction = f.read()
|
|
messages.append({"role": "system", "content": system_instruction})
|
|
except Exception as e:
|
|
print("Error reading custom instruction file:", e)
|
|
if keep_context and self.chat_history:
|
|
messages.extend(self.chat_history)
|
|
messages.append({"role": "user", "content": prompt})
|
|
completion = client.chat.completions.create(
|
|
model=model_value,
|
|
messages=messages,
|
|
seed=seed
|
|
)
|
|
output_text = completion.choices[0].message.content
|
|
self.chat_history.extend([{"role": "user", "content": prompt},
|
|
{"role": "assistant", "content": output_text}])
|
|
return (output_text,)
|
|
|
|
# DeepSeek branch (compatible with OpenAI SDK)
|
|
elif model_value in ('deepseek-chat', 'deepseek-reasoner'):
|
|
from openai import OpenAI
|
|
client = OpenAI(api_key=api_key_deepseek, base_url=selected_base_url)
|
|
messages = []
|
|
if custom_instruction != "None":
|
|
try:
|
|
with open(os.path.join(custom_instructions_directory, f"{custom_instruction}.txt"), 'r') as f:
|
|
system_instruction = f.read()
|
|
messages.append({"role": "system", "content": system_instruction})
|
|
except Exception as e:
|
|
print("Error reading custom instruction file:", e)
|
|
if keep_context and self.chat_history:
|
|
messages.extend(self.chat_history)
|
|
messages.append({"role": "user", "content": prompt})
|
|
completion = client.chat.completions.create(
|
|
model=model_value,
|
|
messages=messages,
|
|
stream=False,
|
|
seed=seed
|
|
)
|
|
output_text = completion.choices[0].message.content
|
|
self.chat_history.extend([{"role": "user", "content": prompt},
|
|
{"role": "assistant", "content": output_text}])
|
|
return (output_text,)
|
|
|
|
# Gemini branch
|
|
elif model_value.startswith(('gemini-', 'imagen-', 'nano-', 'banana-')):
|
|
from google import genai
|
|
from google.genai import types
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import os
|
|
import base64
|
|
import numpy as np
|
|
import torch
|
|
|
|
# Initialize the Gemini client (using your Gemini API key)
|
|
client = genai.Client(api_key=api_key_gemini)
|
|
|
|
system_instruction = None
|
|
if custom_instruction != "None":
|
|
try:
|
|
instr_path = os.path.join(custom_instructions_directory, f"{custom_instruction}.txt")
|
|
with open(instr_path, 'r', encoding='utf-8') as f:
|
|
system_instruction = f.read()
|
|
except Exception as e:
|
|
print(f"Warning: Error reading custom instruction file {instr_path}: {e}")
|
|
|
|
prompt_text = (system_instruction + "\n\n" + prompt).strip() if system_instruction else prompt
|
|
|
|
# Build contents (history + current prompt + images) - (keep as is)
|
|
contents = []
|
|
# Add chat history if keep_context is True (convert to Gemini format)
|
|
if keep_context and self.chat_history: # Changed Polymath.chat_history to self.chat_history
|
|
for msg in self.chat_history:
|
|
role = msg.get("role")
|
|
content = msg.get("content")
|
|
gemini_role = "user" if role == "user" else "model"
|
|
# Handle content structure
|
|
if isinstance(content, list): # Already structured (e.g., multimodal parts)
|
|
contents.append({"role": gemini_role, "parts": content})
|
|
elif isinstance(content, str): # Simple text content
|
|
contents.append({"role": gemini_role, "parts": [{"text": content}]})
|
|
|
|
if b64:
|
|
parts = [{"text": prompt_text}]
|
|
for img_str in b64:
|
|
parts.append({"inline_data": {"mime_type": "image/png", "data": img_str}})
|
|
contents.append({"role": "user", "parts": parts})
|
|
else:
|
|
contents.append({"role": "user", "parts": [{"text": prompt_text}]})
|
|
|
|
# Retrieve Gemini specific settings
|
|
aspect_ratio = settings.get("gemini_aspect_ratio", "1:1")
|
|
image_size = settings.get("gemini_image_size", "1K")
|
|
|
|
is_image_model = any(keyword in model_value for keyword in ['image', 'imagen', 'banana'])
|
|
|
|
config_args = {
|
|
"temperature": settings.get("temperature", 0.8),
|
|
"top_p": settings.get("top_p", 0.95),
|
|
"top_k": settings.get("top_k", 40),
|
|
"max_output_tokens": settings.get("max_output_tokens", 8192),
|
|
"seed": seed, # Added seed to config_args
|
|
}
|
|
|
|
if is_image_model:
|
|
config_args["response_modalities"] = ["IMAGE"]
|
|
|
|
allowed_fields = set()
|
|
if hasattr(types.ImageConfig, "model_fields"):
|
|
allowed_fields = set(types.ImageConfig.model_fields.keys())
|
|
|
|
# NOTE: If 'image_size' is ignored, please upgrade google-genai: pip install --upgrade google-genai
|
|
|
|
img_conf_args = {}
|
|
|
|
if "aspect_ratio" in allowed_fields:
|
|
img_conf_args["aspect_ratio"] = aspect_ratio
|
|
if "image_size" in allowed_fields:
|
|
# Only send image_size for Gemini 3 / Nano Banana models
|
|
# Gemini 2.5 and others reject this parameter
|
|
if "3" in model_value or "banana" in model_value or "nano" in model_value:
|
|
img_conf_args["image_size"] = image_size
|
|
elif "sample_image_size" in allowed_fields:
|
|
if "3" in model_value or "banana" in model_value or "nano" in model_value:
|
|
img_conf_args["sample_image_size"] = image_size
|
|
else:
|
|
# Warn if user tries to set non-default size on unsupported setup
|
|
if console_log and image_size != "1K" and ("3" in model_value or "banana" in model_value):
|
|
print(f"Warning: 'image_size' not supported in this google-genai version. Ignoring setting '{image_size}'.")
|
|
|
|
config_args["image_config"] = types.ImageConfig(**img_conf_args)
|
|
|
|
generate_config = types.GenerateContentConfig(**config_args)
|
|
|
|
try:
|
|
if console_log:
|
|
print(f"Gemini Request: Model={model_value}, IsImage={is_image_model}, Config={config_args}")
|
|
|
|
response = client.models.generate_content(
|
|
model=model_value,
|
|
contents=contents,
|
|
config=generate_config
|
|
)
|
|
|
|
output_text = ""
|
|
image_tensors = []
|
|
|
|
# Parse response
|
|
if response.parts:
|
|
for part in response.parts:
|
|
if part.text:
|
|
output_text += part.text + "\n"
|
|
|
|
img_bytes = None
|
|
if hasattr(part, 'as_image') and part.as_image():
|
|
pil_img = part.as_image()
|
|
if not isinstance(pil_img, Image.Image):
|
|
# Try to convert if it's not already an image
|
|
try:
|
|
pil_img = Image.open(BytesIO(pil_img))
|
|
except:
|
|
pass
|
|
|
|
if isinstance(pil_img, Image.Image):
|
|
img_array = np.array(pil_img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
|
|
elif hasattr(part, 'inline_data') and part.inline_data:
|
|
img_data = part.inline_data.data
|
|
try:
|
|
# img_data is likely already bytes
|
|
img_bytes = img_data
|
|
if isinstance(img_data, str):
|
|
try: img_bytes = base64.b64decode(img_data)
|
|
except: img_bytes = img_data.encode('utf-8')
|
|
|
|
pil_image = Image.open(BytesIO(img_bytes)).convert("RGB")
|
|
img_array = np.array(pil_image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
except Exception as e:
|
|
print(f"Error processing inline image data: {e}")
|
|
|
|
# Update chat history if text was generated
|
|
if output_text:
|
|
self.chat_history.extend([{"role": "user", "content": prompt},
|
|
{"role": "assistant", "content": output_text}])
|
|
|
|
# Check for candidates if parts were empty (fallback)
|
|
if not output_text and not image_tensors and hasattr(response, 'candidates') and response.candidates:
|
|
for cand in response.candidates:
|
|
if cand.content and cand.content.parts:
|
|
for part in cand.content.parts:
|
|
if part.text: output_text += part.text + "\n"
|
|
|
|
found_image = False
|
|
|
|
# Try inline_data first as it is seen in the debug dump
|
|
if hasattr(part, 'inline_data') and part.inline_data:
|
|
try:
|
|
img_data = part.inline_data.data
|
|
img_bytes = img_data
|
|
if isinstance(img_data, str):
|
|
try: img_bytes = base64.b64decode(img_data)
|
|
except: img_bytes = img_data.encode('utf-8')
|
|
|
|
pil_image = Image.open(BytesIO(img_bytes)).convert("RGB")
|
|
img_array = np.array(pil_image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
found_image = True
|
|
except Exception as e:
|
|
print(f"Error processing inline image data in candidate: {e}")
|
|
# Fallback: Maybe it WAS base64 encoded bytes?
|
|
try:
|
|
if isinstance(img_data, bytes):
|
|
img_bytes = base64.b64decode(img_data)
|
|
pil_image = Image.open(BytesIO(img_bytes)).convert("RGB")
|
|
img_array = np.array(pil_image).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
found_image = True
|
|
except Exception as e2:
|
|
pass # Fallback failed
|
|
|
|
# Only try as_image if we haven't found one yet for this part
|
|
if not found_image and hasattr(part, 'as_image'):
|
|
try:
|
|
pil_img = part.as_image()
|
|
if pil_img:
|
|
if not isinstance(pil_img, Image.Image):
|
|
try: pil_img = Image.open(BytesIO(pil_img))
|
|
except: pass
|
|
|
|
if isinstance(pil_img, Image.Image):
|
|
pil_img = pil_img.convert("RGB")
|
|
img_array = np.array(pil_img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_array).unsqueeze(0)
|
|
image_tensors.append(img_tensor)
|
|
except Exception as e:
|
|
print(f"Error checking as_image(): {e}")
|
|
|
|
if is_image_model and not image_tensors and not output_text:
|
|
# Return a dummy black image to prevent ComfyUI crash
|
|
dummy_tensor = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
|
|
return ("No image generated. Check console for details.", dummy_tensor)
|
|
|
|
final_image_output = torch.cat(image_tensors, dim=0) if image_tensors else None
|
|
|
|
# Double check return type safety
|
|
if final_image_output is None and is_image_model:
|
|
return (output_text, torch.zeros((1, 512, 512, 3), dtype=torch.float32))
|
|
|
|
if final_image_output is None:
|
|
final_image_output = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
|
|
return (output_text, final_image_output)
|
|
|
|
except Exception as e:
|
|
error_msg = f"Gemini API Error: {e}"
|
|
print(f"\033[91m!!! Polymath Gemini Error !!!\033[0m\n{error_msg}")
|
|
# Return dummy tensor on error
|
|
return (error_msg, torch.zeros((1, 64, 64, 3), dtype=torch.float32))
|
|
|
|
# Fallback to Ollama API
|
|
else:
|
|
images = b64 if b64 else None
|
|
ollama_options = {
|
|
"temperature": settings.get("temperature", 0.8),
|
|
"top_p": settings.get("top_p", 0.95),
|
|
"top_k": settings.get("top_k", 40),
|
|
"num_predict": settings.get("max_output_tokens", -1), # Map max_tokens
|
|
"seed": seed,
|
|
"keep_alive": settings.get("ollama_keep_alive", 5)
|
|
}
|
|
stop_seq = settings.get("stop_sequences", "")
|
|
if stop_seq:
|
|
ollama_options["stop"] = [s.strip() for s in stop_seq.split(',')]
|
|
response = self.ollama_api.generate_completion(
|
|
ollama_chat_mode=ollama_chat_mode,
|
|
model=model_value,
|
|
prompt=prompt,
|
|
stream=False,
|
|
images=images,
|
|
seed=seed,
|
|
options=ollama_options
|
|
)
|
|
if ollama_chat_mode:
|
|
output_text = response.get('message', {}).get('content', '')
|
|
else:
|
|
output_text = response.get('response', '')
|
|
if console_log:
|
|
print("Ollama API Response:", output_text)
|
|
return (output_text,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"polymath_settings": PolymathSettings,
|
|
"polymath_chat": Polymath
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"polymath_settings": "LLM Polymath Chat and API Settings",
|
|
"polymath_chat": "LLM Polymath Chat with Advanced Web and Link Search"
|
|
} |