# IFChatPromptNode.py
import os
import sys
import json
import torch
import shutil
import base64
import platform
import importlib
import subprocess
import numpy as np
import folder_paths
from PIL import Image
import yaml
from io import BytesIO
import asyncio
from typing import List, Union, Dict, Any, Tuple, Optional
from .agent_tool import AgentTool
from .send_request import send_request
#from .transformers_api import TransformersModelManager
import tempfile
import threading
import codecs
from aiohttp import web
from .graphRAG_module import GraphRAGapp
from .colpaliRAG_module import colpaliRAGapp
from .superflorence import FlorenceModule
from .utils import get_api_key, get_models, validate_models, clean_text, process_mask, load_placeholder_image, process_images_for_comfy
#from byaldi import RAGMultiModalModel
# Set up logging
import logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
# Add the ComfyUI directory to the Python path
comfy_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
sys.path.insert(0, comfy_path)
ifchat_prompt_node = None
try:
from server import PromptServer
@PromptServer.instance.routes.post("/IF_ChatPrompt/get_llm_models")
async def get_llm_models_endpoint(request):
data = await request.json()
llm_provider = data.get("llm_provider")
engine = llm_provider
base_ip = data.get("base_ip")
port = data.get("port")
external_api_key = data.get("external_api_key")
logger.debug(f"Received request for LLM models. Provider: {llm_provider}, External API key provided: {bool(external_api_key)}")
if external_api_key:
api_key = external_api_key
logger.debug("Using provided external LLM API key")
else:
api_key_name = f"{llm_provider.upper()}_API_KEY"
try:
api_key = get_api_key(api_key_name, engine)
logger.debug("Using API key from environment or .env file")
except ValueError:
logger.warning(f"No API key found for {llm_provider}. Attempting to proceed without an API key.")
api_key = None
models = get_models(engine, base_ip, port, api_key)
logger.debug(f"Fetched {len(models)} models for {llm_provider}")
return web.json_response(models)
@PromptServer.instance.routes.post("/IF_ChatPrompt/get_embedding_models")
async def get_embedding_models_endpoint(request):
data = await request.json()
embedding_provider = data.get("embedding_provider")
engine = embedding_provider
base_ip = data.get("base_ip")
port = data.get("port")
external_api_key = data.get("external_api_key")
logger.debug(f"Received request for LLM models. Provider: {embedding_provider}, External API key provided: {bool(external_api_key)}")
if external_api_key:
api_key = external_api_key
logger.debug("Using provided external LLM API key")
else:
api_key_name = f"{embedding_provider.upper()}_API_KEY"
try:
api_key = get_api_key(api_key_name, engine)
logger.debug("Using API key from environment or .env file")
except ValueError:
logger.warning(f"No API key found for {embedding_provider}. Attempting to proceed without an API key.")
api_key = None
models = get_models(engine, base_ip, port, api_key)
logger.debug(f"Fetched {len(models)} models for {embedding_provider}")
return web.json_response(models)
@PromptServer.instance.routes.post("/IF_ChatPrompt/upload_file")
async def upload_file_route(request):
try:
reader = await request.multipart()
rag_folder_name = None
file_content = None
filename = None
# Process all parts of the multipart request
while True:
part = await reader.next()
if part is None:
break
if part.name == "rag_root_dir":
rag_folder_name = await part.text()
elif part.filename:
filename = part.filename
file_content = await part.read()
if not filename or not file_content or not rag_folder_name:
return web.json_response({"status": "error", "message": "Missing file, filename, or RAG folder name"})
node = IFChatPrompt()
input_dir = os.path.join(node.rag_dir, rag_folder_name, "input")
if not os.path.exists(input_dir):
os.makedirs(input_dir, exist_ok=True)
file_path = os.path.join(input_dir, filename)
with open(file_path, 'wb') as f:
f.write(file_content)
logger.info(f"File uploaded to: {file_path}")
return web.json_response({"status": "success", "message": f"File uploaded to: {file_path}"})
except Exception as e:
logger.error(f"Error in upload_file_route: {str(e)}")
return web.json_response({"status": "error", "message": f"Error uploading file: {str(e)}"})
@PromptServer.instance.routes.post("/IF_ChatPrompt/setup_and_initialize")
async def setup_and_initialize(request):
global ifchat_prompt_node
data = await request.json()
folder_name = data.get('folder_name', 'rag_data')
if ifchat_prompt_node is None:
ifchat_prompt_node = IFChatPrompt()
init_result = await ifchat_prompt_node.graphrag_app.setup_and_initialize_folder(folder_name, data)
ifchat_prompt_node.rag_folder_name = folder_name
ifchat_prompt_node.colpali_app.set_rag_root_dir(folder_name)
return web.json_response(init_result)
@PromptServer.instance.routes.post("/IF_ChatPrompt/run_indexer")
async def run_indexer_endpoint(request):
try:
data = await request.json()
logger.debug(f"Received indexing request with data: {data}")
global ifchat_prompt_node # Access the global instance
# Set the rag_root_dir in both modules using the global instance
ifchat_prompt_node.graphrag_app.set_rag_root_dir(data.get('rag_folder_name'))
ifchat_prompt_node.colpali_app.set_rag_root_dir(data.get('rag_folder_name'))
query_type = data.get('mode_type')
logger.debug(f"Query type: {query_type}")
logger.debug(f"Starting indexing process for query type: {query_type}")
# Initialize the colpali_model before calling insert, using the global instance
if query_type == 'colpali' or query_type == 'colqwen2' or query_type == 'colpali-v1.2':
_ = ifchat_prompt_node.colpali_app.get_colpali_model(query_type) # This will load or retrieve the cached model
result = await ifchat_prompt_node.colpali_app.insert()
else:
result = await ifchat_prompt_node.graphrag_app.insert()
logger.debug(f"Indexing process completed with result: {result}")
if result:
return web.json_response({"status": "success", "message": f"Indexing complete for {query_type}"})
else:
return web.json_response({"status": "error", "message": "Indexing failed. Check server logs."}, status=500)
except Exception as e:
logger.error(f"Error in run_indexer_endpoint: {str(e)}")
return web.json_response({"status": "error", "message": f"Error during indexing: {str(e)}"}, status=500)
@PromptServer.instance.routes.post("/IF_ChatPrompt/process_chat")
async def process_chat_endpoint(request):
try:
data = await request.json()
# Set default values for required arguments if not provided
defaults = {
"prompt": "",
"assistant": "Cortana", # Default assistant
"neg_prompt": "Default", # Default negative prompt
"embellish_prompt": "Default", # Default embellishment
"style_prompt": "Default", # Default style
"llm_provider": "ollama",
"llm_model": "",
"base_ip": "localhost",
"port": "11434",
"embedding_model": "",
"embedding_provider": "sentence_transformers"
}
# Update data with defaults for missing keys
for key, default_value in defaults.items():
if key not in data:
data[key] = default_value
global ifchat_prompt_node
result = await ifchat_prompt_node.process_chat(**data)
return web.json_response(result)
except Exception as e:
logger.error(f"Error in process_chat_endpoint: {str(e)}")
return web.json_response({
"status": "error",
"message": f"Error processing chat: {str(e)}",
"Question": data.get("prompt", ""),
"Response": f"Error: {str(e)}",
"Negative": "",
"Tool_Output": None,
"Retrieved_Image": None,
"Mask": None
}, status=500)
@PromptServer.instance.routes.post("/IF_ChatPrompt/load_index")
async def load_index_route(request):
try:
data = await request.json()
index_name = data.get('rag_folder_name')
query_type = data.get('query_type')
if not index_name:
logger.error("No index name provided in the request.")
return web.json_response({
"status": "error",
"message": "No index name provided"
})
# Check if index exists in .byaldi directory
byaldi_index_path = os.path.join(".byaldi", index_name)
if not os.path.exists(byaldi_index_path):
logger.error(f"Index not found in .byaldi: {byaldi_index_path}")
return web.json_response({
"status": "error",
"message": f"Index {index_name} does not exist"
})
try:
global ifchat_prompt_node
if ifchat_prompt_node is None:
logger.debug("Initializing IFChatPrompt instance.")
ifchat_prompt_node = IFChatPrompt()
if query_type in ['colpali', 'colqwen2', 'colpali-v1.2']:
logger.debug(f"Loading model for query type: {query_type}")
# Clear any existing cached index
ifchat_prompt_node.colpali_app.cleanup_index()
# First get the base model
colpali_model = ifchat_prompt_node.colpali_app.get_colpali_model(query_type)
if colpali_model:
# Load and cache the new index
model = await ifchat_prompt_node.colpali_app._prepare_model(query_type, index_name)
if not model:
raise ValueError("Failed to load and cache index")
# Set the RAG root directory
ifchat_prompt_node.colpali_app.set_rag_root_dir(index_name)
logger.info(f"Successfully loaded and cached index: {index_name}")
return web.json_response({
"status": "success",
"message": f"Successfully loaded index: {index_name}",
"rag_root_dir": index_name
})
else:
logger.error("Failed to initialize ColPali model.")
raise ValueError("Failed to initialize ColPali model")
else:
logger.error(f"Unsupported query type: {query_type}")
return web.json_response({
"status": "error",
"message": f"Query type {query_type} not supported for loading indexes"
})
except Exception as e:
logger.error(f"Error loading index {index_name}: {str(e)}")
return web.json_response({
"status": "error",
"message": f"Error loading index: {str(e)}"
})
except Exception as e:
logger.error(f"Error in load_index_route: {str(e)}")
return web.json_response({
"status": "error",
"message": f"Error processing request: {str(e)}"
})
# Add this with the other routes
@PromptServer.instance.routes.post("/IF_ChatPrompt/delete_index")
async def delete_index_route(request):
try:
data = await request.json()
index_name = data.get('rag_folder_name')
if not index_name:
return web.json_response({
"status": "error",
"message": "No index name provided"
})
# Path to the index
index_path = os.path.join(".byaldi", index_name)
if not os.path.exists(index_path):
return web.json_response({
"status": "error",
"message": f"Index {index_name} does not exist"
})
# Delete the index directory
try:
shutil.rmtree(index_path)
logger.info(f"Successfully deleted index: {index_name}")
return web.json_response({
"status": "success",
"message": f"Successfully deleted index: {index_name}"
})
except Exception as e:
logger.error(f"Error deleting index {index_name}: {str(e)}")
return web.json_response({
"status": "error",
"message": f"Error deleting index: {str(e)}"
})
except Exception as e:
logger.error(f"Error in delete_index_route: {str(e)}")
return web.json_response({
"status": "error",
"message": f"Error processing request: {str(e)}"
})
except AttributeError:
print("PromptServer.instance not available. Skipping route decoration for IF_ChatPrompt.")
class IFChatPrompt:
def __init__(self):
self.base_ip = "localhost"
self.port = "11434"
self.llm_provider = "ollama"
self.embedding_provider = "sentence_transformers"
self.llm_model = ""
self.embedding_model = ""
self.assistant = "None"
self.random = False
self.comfy_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
self.rag_dir = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI-IF_AI_tools", "IF_AI", "rag")
self.presets_dir = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI-IF_AI_tools", "IF_AI", "presets")
self.stop_file = os.path.join(self.presets_dir, "stop_strings.json")
self.assistants_file = os.path.join(self.presets_dir, "assistants.json")
self.neg_prompts_file = os.path.join(self.presets_dir, "neg_prompts.json")
self.embellish_prompts_file = os.path.join(self.presets_dir, "embellishments.json")
self.style_prompts_file = os.path.join(self.presets_dir, "style_prompts.json")
self.tasks_file = os.path.join(self.presets_dir, "florence_prompts.json")
self.agents_dir = os.path.join(self.presets_dir, "agents")
self.agent_tools = self.load_agent_tools()
self.stop_strings = self.load_presets(self.stop_file)
self.assistants = self.load_presets(self.assistants_file)
self.neg_prompts = self.load_presets(self.neg_prompts_file)
self.embellish_prompts = self.load_presets(self.embellish_prompts_file)
self.style_prompts = self.load_presets(self.style_prompts_file)
self.florence_prompts = self.load_presets(self.tasks_file)
self.keep_alive = False
self.seed = 94687328150
self.messages = []
self.history_steps = 10
self.external_api_key = ""
self.tool_input = ""
self.prime_directives = None
self.rag_folder_name = "rag_data"
self.graphrag_app = GraphRAGapp()
self.colpali_app = colpaliRAGapp()
self.fix_json = True
self.cached_colpali_model = None
self.florence_app = FlorenceModule()
self.florence_models = {}
self.query_type = "global"
self.enable_RAG = False
self.clear_history = False
self.mode = False
self.tool = "None"
self.preset = "Default"
self.precision = "fp16"
self.task = None
self.attention = "sdpa"
self.aspect_ratio = "16:9"
self.top_k_search = 3
self.placeholder_image_path = os.path.join(folder_paths.base_path, "custom_nodes", "ComfyUI-IF_AI_tools", "IF_AI", "placeholder.png")
if not os.path.exists(self.placeholder_image_path):
placeholder = Image.new('RGB', (512, 512), color=(73, 109, 137))
os.makedirs(os.path.dirname(self.placeholder_image_path), exist_ok=True)
placeholder.save(self.placeholder_image_path)
def load_presets(self, file_path: str) -> Dict[str, Any]:
"""
Load JSON presets with support for multiple encodings.
Args:
file_path (str): Path to the JSON preset file
Returns:
Dict[str, Any]: Loaded JSON data or empty dict if loading fails
"""
# List of encodings to try
encodings = ['utf-8', 'utf-8-sig', 'latin1', 'cp1252', 'gbk']
for encoding in encodings:
try:
with codecs.open(file_path, 'r', encoding=encoding) as f:
data = json.load(f)
# If successful, write back with UTF-8 encoding to prevent future issues
try:
with codecs.open(file_path, 'w', encoding='utf-8') as out_f:
json.dump(data, out_f, ensure_ascii=False, indent=2)
except Exception as write_err:
print(f"Warning: Could not write back UTF-8 encoded file: {write_err}")
return data
except UnicodeDecodeError:
continue
except json.JSONDecodeError as e:
print(f"JSON parsing error with {encoding} encoding: {str(e)}")
continue
except Exception as e:
print(f"Error loading presets from {file_path} with {encoding} encoding: {e}")
continue
print(f"Error: Failed to load {file_path} with any supported encoding")
return {}
def load_agent_tools(self):
os.makedirs(self.agents_dir, exist_ok=True)
agent_tools = {}
try:
for filename in os.listdir(self.agents_dir):
if filename.endswith('.json'):
full_path = os.path.join(self.agents_dir, filename)
with open(full_path, 'r') as f:
try:
data = json.load(f)
if 'output_type' not in data:
data['output_type'] = None
agent_tool = AgentTool(**data)
agent_tool.load()
if agent_tool._class_instance is not None:
if agent_tool.python_function:
agent_tools[agent_tool.name] = agent_tool
else:
print(f"Warning: Agent tool {agent_tool.name} in {filename} does not have a python_function defined.")
else:
print(f"Failed to create class instance for {filename}")
except json.JSONDecodeError:
print(f"Error: Invalid JSON in {filename}")
except Exception as e:
print(f"Error loading {filename}: {str(e)}")
return agent_tools
except Exception as e:
print(f"Warning: Error accessing agent tools directory: {str(e)}")
return {}
async def process_chat(
self,
prompt,
llm_provider,
llm_model,
base_ip,
port,
assistant,
neg_prompt,
embellish_prompt,
style_prompt,
embedding_model,
embedding_provider,
external_api_key="",
temperature=0.7,
max_tokens=2048,
seed=0,
random=False,
history_steps=10,
keep_alive=False,
top_k=40,
top_p=0.2,
repeat_penalty=1.1,
stop_string=None,
images=None,
mode=True,
clear_history=False,
text_cleanup=True,
tool=None,
tool_input=None,
prime_directives=None,
enable_RAG=False,
query_type="global",
preset="Default",
rag_folder_name=None,
task=None,
fill_mask=False,
output_mask_select="",
precision="fp16",
attention="sdpa",
aspect_ratio="16:9",
top_k_search=3
):
if external_api_key != "":
llm_api_key = external_api_key
else:
llm_api_key = get_api_key(f"{llm_provider.upper()}_API_KEY", llm_provider)
print(f"LLM API key: {llm_api_key[:5]}...")
if prime_directives is not None:
system_message_str = prime_directives
else:
system_message = self.assistants.get(assistant, "")
system_message_str = json.dumps(system_message)
# Validate LLM model
validate_models(llm_model, llm_provider, "LLM", base_ip, port, llm_api_key)
# Validate embedding model
validate_models(embedding_model, embedding_provider, "embedding", base_ip, port, llm_api_key)
# Handle history
if clear_history:
self.messages = []
elif history_steps > 0:
self.messages = self.messages[-history_steps:]
messages = self.messages
# Handle stop
if stop_string is None or stop_string == "None":
stop_content = None
else:
stop_content = self.stop_strings.get(stop_string, None)
stop = stop_content
if llm_provider not in ["ollama", "llamacpp", "vllm", "lmstudio", "gemeni"]:
if llm_provider == "kobold":
stop = stop_content + \
["\n\n\n\n\n"] if stop_content else ["\n\n\n\n\n"]
elif llm_provider == "mistral":
stop = stop_content + \
["\n\n"] if stop_content else ["\n\n"]
else:
stop = stop_content if stop_content else None
# Handle tools
try:
if tool and tool != "None":
selected_tool = self.agent_tools.get(tool)
if not selected_tool:
raise ValueError(f"Invalid agent tool selected: {tool}")
# Prepare tool execution message
tool_message = f"Execute the {tool} tool with the following input: {prompt}"
system_prompt = json.dumps(selected_tool.system_prompt)
# Send request to LLM for tool execution
generated_text =await send_request(
llm_provider=llm_provider,
base_ip=base_ip,
port=port,
images=images,
model=llm_model,
system_message=system_prompt,
user_message=tool_message,
messages=messages,
seed=seed,
temperature=temperature,
max_tokens=max_tokens,
random=random,
top_k=top_k,
top_p=top_p,
repeat_penalty=repeat_penalty,
stop=stop,
keep_alive=keep_alive,
llm_api_key=llm_api_key,
)
# Parse the generated text for function calls
function_call = None
try:
response_data = json.loads(generated_text)
if 'function_call' in response_data:
function_call = response_data['function_call']
generated_text = response_data['content']
except json.JSONDecodeError:
pass # The response wasn't JSON, so it's just the generated text
# Execute the tool with the LLM's response
tool_args = {
"input": prompt,
"llm_response": generated_text,
"function_call": function_call,
"omni_input": tool_input,
"name": selected_tool.name,
"description": selected_tool.description,
"system_prompt": selected_tool.system_prompt
}
tool_result = selected_tool.execute(tool_args)
# Update messages
messages.append({"role": "user", "content": prompt})
messages.append({
"role": "assistant",
"content": json.dumps(tool_result) if isinstance(tool_result, dict) else str(tool_result)
})
# Process the tool output
if isinstance(tool_result, dict):
if "error" in tool_result:
generated_text = f"Error in {tool}: {tool_result['error']}"
tool_output = None
elif selected_tool.output_type and selected_tool.output_type in tool_result:
tool_output = tool_result[selected_tool.output_type]
generated_text = f"Agent {tool} executed successfully. Output generated."
else:
tool_output = tool_result
generated_text = str(tool_output)
else:
tool_output = tool_result
generated_text = str(tool_output)
return {
"Question": prompt,
"Response": generated_text,
"Negative": self.neg_prompts.get(neg_prompt, ""),
"Tool_Output": tool_output,
"Retrieved_Image": None # No image retrieved in tool execution
}
else:
response = await self.generate_response(
enable_RAG,
query_type,
prompt,
preset,
llm_provider,
base_ip,
port,
images,
llm_model,
system_message_str,
messages,
temperature,
max_tokens,
random,
top_k,
top_p,
repeat_penalty,
stop,
seed,
keep_alive,
llm_api_key,
task,
fill_mask,
output_mask_select,
precision,
attention
)
generated_text = response.get("Response")
selected_neg_prompt_name = neg_prompt
omni = response.get("Tool_Output")
retrieved_image = response.get("Retrieved_Image")
retrieved_mask = response.get("Mask")
# Update messages
messages.append({"role": "user", "content": prompt})
messages.append({"role": "assistant", "content": generated_text})
text_result = str(generated_text).strip()
if mode:
embellish_content = self.embellish_prompts.get(embellish_prompt, "").strip()
style_content = self.style_prompts.get(style_prompt, "").strip()
lines = [line.strip() for line in text_result.split('\n') if line.strip()]
combined_prompts = []
for line in lines:
if text_cleanup:
line = clean_text(line)
formatted_line = f"{embellish_content} {line} {style_content}".strip()
combined_prompts.append(formatted_line)
combined_prompt = "\n".join(formatted_line for formatted_line in combined_prompts)
# Handle negative prompts
if selected_neg_prompt_name == "AI_Fill":
try:
neg_system_message = self.assistants.get("NegativePromptEngineer")
if not neg_system_message:
logger.error("NegativePromptEngineer not found in assistants configuration")
negative_prompt = "Error: NegativePromptEngineer not configured"
else:
user_message = f"Generate negative prompts for the following prompt:\n{text_result}"
system_message_str = json.dumps(neg_system_message)
logger.info(f"Requesting negative prompts for prompt: {text_result[:100]}...")
neg_response = await send_request(
llm_provider=llm_provider,
base_ip=base_ip,
port=port,
images=None,
llm_model=llm_model,
system_message=system_message_str,
user_message=user_message,
messages=[], # Fresh context for negative generation
seed=seed,
temperature=temperature,
max_tokens=max_tokens,
random=random,
top_k=top_k,
top_p=top_p,
repeat_penalty=repeat_penalty,
stop=stop,
keep_alive=keep_alive,
llm_api_key=llm_api_key
)
logger.debug(f"Received negative prompt response: {neg_response}")
if neg_response:
negative_lines = []
for line in neg_response.split('\n'):
line = line.strip()
if line:
negative_lines.append(line)
while len(negative_lines) < len(lines):
negative_lines.append(negative_lines[-1] if negative_lines else "")
negative_lines = negative_lines[:len(lines)]
negative_prompt = "\n".join(negative_lines)
else:
negative_prompt = "Error: Empty response from LLM"
except Exception as e:
logger.error(f"Error generating negative prompts: {str(e)}", exc_info=True)
negative_prompt = f"Error generating negative prompts: {str(e)}"
elif neg_prompt != "None":
neg_content = self.neg_prompts.get(neg_prompt, "").strip()
negative_lines = [neg_content for _ in range(len(lines))]
negative_prompt = "\n".join(negative_lines)
else:
negative_prompt = ""
else:
combined_prompt = text_result
negative_prompt = ""
try:
if isinstance(retrieved_image, torch.Tensor):
# Ensure it's in the correct format (B, C, H, W)
if retrieved_image.dim() == 3: # Single image (C, H, W)
image_tensor = retrieved_image.unsqueeze(0) # Add batch dimension
else:
image_tensor = retrieved_image # Already batched
# Create matching batch masks
batch_size = image_tensor.shape[0]
height = image_tensor.shape[2]
width = image_tensor.shape[3]
# Create white masks (all ones) for each image in batch
mask_tensor = torch.ones((batch_size, 1, height, width),
dtype=torch.float32,
device=image_tensor.device)
if retrieved_mask is not None:
# If we have masks, process them to match the batch
if isinstance(retrieved_mask, torch.Tensor):
if retrieved_mask.dim() == 3: # Single mask
mask_tensor = retrieved_mask.unsqueeze(0)
else:
mask_tensor = retrieved_mask
else:
# Process retrieved_mask if it's not a tensor
mask_tensor = process_mask(retrieved_mask, image_tensor)
else:
image_tensor, default_mask_tensor = process_images_for_comfy(
retrieved_image,
self.placeholder_image_path,
response_key=None,
field_name=None
)
mask_tensor = default_mask_tensor
if retrieved_mask is not None:
mask_tensor = process_mask(retrieved_mask, image_tensor)
return (
prompt,
combined_prompt,
negative_prompt,
omni,
image_tensor,
mask_tensor,
)
except Exception as e:
logger.error(f"Exception in image processing: {str(e)}", exc_info=True)
placeholder_image, placeholder_mask = load_placeholder_image(self.placeholder_image_path)
return (
prompt,
f"Error: {str(e)}",
"",
None,
placeholder_image,
placeholder_mask
)
except Exception as e:
logger.error(f"Exception occurred in process_chat: {str(e)}", exc_info=True)
placeholder_image, placeholder_mask = load_placeholder_image(self.placeholder_image_path)
return (
prompt,
f"Error: {str(e)}",
"",
None,
placeholder_image,
placeholder_mask
)
async def generate_response(
self,
enable_RAG,
query_type,
prompt,
preset,
llm_provider,
base_ip,
port,
images,
llm_model,
system_message_str,
messages,
temperature,
max_tokens,
random,
top_k,
top_p,
repeat_penalty,
stop,
seed,
keep_alive,
llm_api_key,
task=None,
fill_mask=False,
output_mask_select="",
precision="fp16",
attention="sdpa",
):
response_strategies = {
"graphrag": self.graphrag_app.query,
"colpali": self.colpali_app.query,
"florence": self.florence_app.run_florence,
"normal": lambda: send_request(
llm_provider=llm_provider,
base_ip=base_ip,
port=port,
images=images,
llm_model=llm_model,
system_message=system_message_str,
user_message=prompt,
messages=messages,
seed=seed,
temperature=temperature,
max_tokens=max_tokens,
random=random,
top_k=top_k,
top_p=top_p,
repeat_penalty=repeat_penalty,
stop=stop,
keep_alive=keep_alive,
llm_api_key=llm_api_key,
tools=None,
tool_choice=None,
precision=precision,
attention=attention
),
}
florence_tasks = list(self.florence_prompts.keys())
if enable_RAG:
if query_type == "colpali" or query_type == "colpali-v1.2" or query_type == "colqwen2":
strategy = "colpali"
else: # For "global", "local", and "naive" query types
strategy = "graphrag"
elif task and task.lower() != 'none' and task in florence_tasks:
strategy = "florence"
else:
strategy = "normal"
print(f"Strategy: {strategy}")
try:
if strategy == "colpali":
# Ensure the model is loaded before querying
if self.cached_colpali_model is None:
self.cached_colpali_model = self.colpali_app.get_colpali_model(query_type)
response = await response_strategies[strategy](prompt=prompt, query_type=query_type, system_message_str=system_message_str)
return response
elif strategy == "graphrag":
response = await response_strategies[strategy](prompt=prompt, query_type=query_type, preset=preset)
return {
"Question": prompt,
"Response": response[0],
"Negative": "",
"Tool_Output": response[1],
"Retrieved_Image": None,
"Mask": None
}
elif strategy == "florence":
task_content = self.florence_prompts.get(task, "")
response = await response_strategies[strategy](
images=images,
task=task,
task_prompt=task_content,
llm_model=llm_model,
precision=precision,
attention=attention,
fill_mask=fill_mask,
output_mask_select=output_mask_select,
keep_alive=keep_alive,
max_new_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
top_k=top_k,
repetition_penalty=repeat_penalty,
seed=seed,
text_input=prompt,
)
print("Florence response:", response)
return response
else:
response = await response_strategies[strategy]()
print("Normal response:", response)
return {
"Question": prompt,
"Response": response,
"Negative": "",
"Tool_Output": None,
"Retrieved_Image": None,
"Mask": None
}
except Exception as e:
logger.error(f"Error processing strategy: {strategy}")
return {
"Question": prompt,
"Response": f"Error processing task: {str(e)}",
"Negative": "",
"Tool_Output": {"error": str(e)},
"Retrieved_Image": None,
"Mask": None
}
def process_chat_wrapper(self, *args, **kwargs):
try:
loop = asyncio.get_event_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
logger.debug(f"process_chat_wrapper kwargs: {kwargs}")
logger.debug(f"External LLM API Key: {kwargs.get('external_api_key', 'Not provided')}")
return loop.run_until_complete(self.process_chat(*args, **kwargs))
@classmethod
def INPUT_TYPES(cls):
node = cls()
return {
"required": {
"prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "The main text input for the chat or query."}),
"llm_provider": (["xai","llamacpp", "ollama", "kobold", "lmstudio", "textgen", "groq", "gemini", "openai", "anthropic", "mistral", "transformers"], {"default": node.llm_provider, "tooltip": "The provider of the language model to be used."}),
"llm_model": ((), {"tooltip": "The specific language model to be used for processing."}),
"base_ip": ("STRING", {"default": node.base_ip, "tooltip": "IP address of the LLM server."}),
"port": ("STRING", {"default": node.port, "tooltip": "Port number for the LLM server connection."}),
},
"optional": {
"images": ("IMAGE", {"list": True, "tooltip": "Input image(s) for visual processing or context."}),
"precision": (['fp16','bf16','fp32','int8','int4'],{"default": 'bf16', "tooltip": "Select preccision on Transformer models."}),
"attention": (['flash_attention_2','sdpa','xformers', 'Shrek_COT_o1'],{"default": 'sdpa', "tooltip": "Select attention mechanism on Transformer models."}),
"assistant": ([name for name in node.assistants.keys()], {"default": node.assistant, "tooltip": "The pre-defined assistant personality to use for responses."}),
"tool": (["None"] + [name for name in node.agent_tools.keys()], {"default": "None", "tooltip": "Selects a specific tool or agent for task execution."}),
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Controls randomness in output generation. Higher values increase creativity but may reduce coherence."}),
"max_tokens": ("INT", {"default": 2048, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Maximum number of tokens to generate in the response."}),
"top_k": ("INT", {"default": 40, "min": 0, "max": 100, "tooltip": "Limits the next token selection to the K most likely tokens."}),
"top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.1, "tooltip": "Cumulative probability cutoff for token selection."}),
"repeat_penalty": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Penalizes repetition in generated text."}),
"stop_string": ([name for name in node.stop_strings.keys()], {"tooltip": "Specifies a string at which text generation should stop."}),
"seed": ("INT", {"default": 94687328150, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Random seed for reproducible outputs."}),
"random": ("BOOLEAN", {"default": False, "label_on": "Seed", "label_off": "Temperature", "tooltip": "Toggles between using a fixed seed or temperature-based randomness."}),
"history_steps": ("INT", {"default": 10, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Number of previous conversation turns to consider for context."}),
"clear_history": ("BOOLEAN", {"default": False, "label_on": "Clear History", "label_off": "Keep History", "tooltip": "Option to clear or retain conversation history."}),
"keep_alive": ("BOOLEAN", {"default": False, "label_on": "Keeps Model on Memory", "label_off": "Unloads Model from Memory", "tooltip": "Determines whether to keep the model loaded in memory between calls."}),
"text_cleanup": ("BOOLEAN", {"default": True, "label_on": "Clean Response", "label_off": "Raw Text", "tooltip": "Applies text cleaning to the generated output."}),
"mode": ("BOOLEAN", {"default": False, "label_on": "Using SD Mode", "label_off": "Using Chat Mode", "tooltip": "Switches between Stable Diffusion prompt generation and standard chat mode."}),
"embellish_prompt": ([name for name in node.embellish_prompts.keys()], {"tooltip": "Adds pre-defined embellishments to the prompt."}),
"style_prompt": ([name for name in node.style_prompts.keys()], {"tooltip": "Applies a pre-defined style to the prompt."}),
"neg_prompt": ([name for name in node.neg_prompts.keys()], {"tooltip": "Adds a negative prompt to guide what should be avoided in generation."}),
"fill_mask": ("BOOLEAN", {"default": False, "label_on": "Fill Mask", "label_off": "No Fill", "tooltip": "Option to fill masks for Florence tasks."}),
"output_mask_select": ("STRING", {"default": ""}),
"task": ([name for name in node.florence_prompts.keys()], {"default": "None", "tooltip": "Select a Florence task."}),
"embedding_provider": (["llamacpp", "ollama", "kobold", "lmstudio", "textgen", "groq", "gemini", "openai", "anthropic", "mistral", "sentence_transformers"], {"default": node.embedding_provider, "tooltip": "Provider for text embedding model."}),
"embedding_model": ((), {"tooltip": "Specific embedding model to use."}),
"tool_input": ("OMNI", {"default": None, "tooltip": "Additional input for the selected tool."}),
"prime_directives": ("STRING", {"forceInput": True, "tooltip": "System message or prime directive for the AI assistant."}),
"external_api_key":("STRING", {"default": "", "tooltip": "If this is not empty, it will be used instead of the API key from the .env file. Make sure it is empty to use the .env file."}),
"top_k_search": ("INT", {"default": 3, "min": 1, "max": 10, "tooltip": "Find top scored image(s) from RAG."}),
"aspect_ratio": (["1:1", "9:16", "16:9"], {"default": "16:9", "tooltip": "Select the aspect ratio for the image."}),
"enable_RAG": ("BOOLEAN", {"default": False, "label_on": "RAG is Enabled", "label_off": "RAG is Disabled", "tooltip": "Enables Retrieval-Augmented Generation for enhanced context."}),
"query_type": (["global", "local", "naive", "colpali", "colqwen2", "colpali-v1.2"], {"default": "global", "tooltip": "Selects the type of query strategy for RAG."}),
"preset": (["Default", "Detailed", "Quick", "Bullet", "Comprehensive", "High-Level", "Focused"], {"default": "Default"}),
},
"hidden": {
"model": ("STRING", {"default": ""}),
"rag_root_dir": ("STRING", {"default": "rag_data"})
}
}
@classmethod
def IS_CHANGED(cls, **kwargs):
node = cls()
llm_provider = kwargs.get('llm_provider', node.llm_provider)
embedding_provider = kwargs.get('embedding_provider', node.embedding_provider)
base_ip = kwargs.get('base_ip', node.base_ip)
port = kwargs.get('port', node.port)
query_type = kwargs.get('query_type', node.query_type)
external_api_key = kwargs.get('external_api_key', '')
task = kwargs.get('task', node.task)
# Determine which API key to use
def get_api_key_with_fallback(provider, external_api_key):
if external_api_key and external_api_key != '':
return external_api_key
try:
# print(f"Using {provider} API key from .env file")
api_key = get_api_key(f"{provider.upper()}_API_KEY", provider)
# print(f" {api_key} API key for {provider} found in .env file")
return api_key
except ValueError:
return None
api_key = get_api_key_with_fallback(llm_provider, external_api_key)
# Check for changes
llm_provider_changed = llm_provider != node.llm_provider
embedding_provider_changed = embedding_provider != node.embedding_provider
api_key_changed = external_api_key != node.external_api_key
base_ip_changed = base_ip != node.base_ip
port_changed = port != node.port
query_type_changed = query_type != node.query_type
task_changed = task != node.task
# Always fetch new models if the provider, API key, base_ip, or port has changed
if llm_provider_changed or api_key_changed or base_ip_changed or port_changed:
try:
new_llm_models = get_models(llm_provider, base_ip, port, api_key)
except Exception as e:
print(f"Error fetching LLM models: {e}")
new_llm_models = []
llm_model_changed = new_llm_models != node.llm_model
else:
llm_model_changed = False
if embedding_provider_changed or api_key_changed or base_ip_changed or port_changed:
try:
new_embedding_models = get_models(embedding_provider, base_ip, port, api_key)
except Exception as e:
print(f"Error fetching embedding models: {e}")
new_embedding_models = []
embedding_model_changed = new_embedding_models != node.embedding_model
else:
embedding_model_changed = False
if (llm_provider_changed or embedding_provider_changed or llm_model_changed or
embedding_model_changed or query_type_changed or task_changed or api_key_changed or
base_ip_changed or port_changed):
node.llm_provider = llm_provider
node.embedding_provider = embedding_provider
node.base_ip = base_ip
node.port = port
node.external_api_key = external_api_key
node.query_type = query_type
node.task = task
if llm_model_changed:
node.llm_model = new_llm_models
if embedding_model_changed:
node.embedding_model = new_embedding_models
# Update other attributes
for attr in ['seed', 'random', 'history_steps', 'clear_history', 'mode',
'keep_alive', 'tool', 'enable_RAG', 'preset']:
setattr(node, attr, kwargs.get(attr, getattr(node, attr)))
return True
return False
RETURN_TYPES = ("STRING", "STRING", "STRING", "OMNI", "IMAGE", "MASK")
RETURN_NAMES = ("Question", "Response", "Negative", "Tool_Output", "Retrieved_Image", "Mask")
OUTPUT_TOOLTIPS = (
"The original input question or prompt.",
"The generated response from the language model.",
"The negative prompt used (if applicable) for guiding image generation.",
"Output from the selected tool, which can be code or any other data type.",
"An image retrieved by the RAG system, if applicable.",
"Mask image generated by Florence tasks."
)
FUNCTION = "process_chat_wrapper"
OUTPUT_NODE = True
CATEGORY = "ImpactFrames💥🎞️/IF_tools"
DESCRIPTION = "ComfyUI, Support API and Local LLM providers and RAG capabilities. Processes text prompts, handles image inputs, and integrates with different language models and indexing strategies."