504 lines
18 KiB
Python
504 lines
18 KiB
Python
import io
|
|
import os
|
|
import json
|
|
import openai
|
|
from openai import OpenAI
|
|
import base64
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Optional, List, Dict, Any
|
|
from .modelverse_api.client import ModelverseClient
|
|
from comfy.comfy_types.node_typing import IO
|
|
from server import PromptServer
|
|
import folder_paths
|
|
|
|
# Hardcoded models list
|
|
DEFAULT_MODELS = [
|
|
'ByteDance/doubao-seed-1.6',
|
|
'Qwen/QwQ-32B',
|
|
'Qwen/Qwen-Image',
|
|
'Qwen/Qwen-Image-Edit',
|
|
'Qwen/Qwen3-32B',
|
|
'Qwen/Qwen3-Coder',
|
|
'doubao-seedance-2-0-260128',
|
|
'kling-v3',
|
|
'kling-v3-omni',
|
|
'happyhorse-1.0-t2v',
|
|
'happyhorse-1.0-i2v',
|
|
'happyhorse-1.0-r2v',
|
|
'veo-3.1-generate-001',
|
|
'veo-3.1-fast-generate-001',
|
|
'baidu/ernie-4.5-turbo-128k',
|
|
'baidu/ernie-4.5-turbo-vl-32k',
|
|
'baidu/ernie-x1-turbo-32k',
|
|
'black-forest-labs/flux-kontext-max',
|
|
'black-forest-labs/flux-kontext-max/multi',
|
|
'black-forest-labs/flux-kontext-max/text-to-image',
|
|
'black-forest-labs/flux-kontext-pro',
|
|
'black-forest-labs/flux-kontext-pro/multi',
|
|
'black-forest-labs/flux-kontext-pro/text-to-image',
|
|
'black-forest-labs/flux.1-dev',
|
|
'claude-4-sonnet',
|
|
'deepseek-ai/DeepSeek-R1',
|
|
'deepseek-ai/DeepSeek-R1-0528',
|
|
'deepseek-ai/DeepSeek-V3-0324',
|
|
'deepseek-ai/DeepSeek-V3.1',
|
|
'gemini-2.5-flash',
|
|
'gemini-2.5-flash-image',
|
|
'gemini-2.5-pro',
|
|
'gemini-3.1-flash-image',
|
|
'gemini-3-pro-image',
|
|
'gpt-4.1-mini',
|
|
'grok-4',
|
|
'moonshotai/Kimi-K2-Instruct',
|
|
'openai/gpt-4.1',
|
|
'openai/gpt-5',
|
|
'openai/gpt-5-mini',
|
|
'viduq2',
|
|
'viduq2-pro',
|
|
'viduq2-turbo',
|
|
'viduq2-pro-fast',
|
|
'viduq3-pro',
|
|
'viduq3-turbo',
|
|
'zai-org/glm-5',
|
|
'stepfun-ai/step1x-edit'
|
|
]
|
|
|
|
# Default selected model
|
|
DEFAULT_MODEL = "zai-org/glm-5"
|
|
|
|
|
|
class ModelverseChat:
|
|
"""OpenAI Chat node with support for text, images, and files"""
|
|
|
|
# Class-level cache for models list
|
|
_cached_models = None
|
|
_cache_initialized = False
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def get_models_list(cls):
|
|
"""Try to fetch models from ModelVerse API, fallback to default models if failed"""
|
|
# Return cached models if already fetched
|
|
if cls._cache_initialized:
|
|
return cls._cached_models
|
|
|
|
try:
|
|
# Try to create a client and fetch models
|
|
# Note: This will work if there's a valid API key available, otherwise fallback
|
|
client = OpenAI(base_url="https://api.modelverse.cn/v1",api_key="xxxsu")
|
|
models_response = client.models.list()
|
|
|
|
# Extract model IDs from the response
|
|
if hasattr(models_response, 'data') and models_response.data:
|
|
model_ids = []
|
|
for model in models_response.data:
|
|
if hasattr(model, 'id') and model.id:
|
|
model_ids.append(model.id)
|
|
|
|
if model_ids:
|
|
print(f"ModelverseChat: Successfully fetched {len(model_ids)} models from API")
|
|
cls._cached_models = sorted(model_ids) # Sort for better user experience
|
|
cls._cache_initialized = True
|
|
return cls._cached_models
|
|
|
|
# If no models found, fall back to default
|
|
print("ModelverseChat: No models found in API response, using default models")
|
|
cls._cached_models = DEFAULT_MODELS
|
|
cls._cache_initialized = True
|
|
return cls._cached_models
|
|
|
|
except Exception as e:
|
|
print(f"ModelverseChat: Failed to fetch models from API ({str(e)}), using default models")
|
|
cls._cached_models = DEFAULT_MODELS
|
|
cls._cache_initialized = True
|
|
return cls._cached_models
|
|
|
|
@classmethod
|
|
def clear_models_cache(cls):
|
|
"""Clear the cached models list to force a refresh on next call"""
|
|
cls._cached_models = None
|
|
cls._cache_initialized = False
|
|
print("ModelverseChat: Models cache cleared")
|
|
|
|
def display_message_on_node(self, message: str, node_id: str) -> None:
|
|
"""Display the current response message on the node UI."""
|
|
render_spec = {
|
|
"node_id": node_id,
|
|
"component": "MessageDisplayWidget",
|
|
"props": {
|
|
"message": message,
|
|
},
|
|
}
|
|
PromptServer.instance.send_sync(
|
|
"display_component",
|
|
render_spec,
|
|
)
|
|
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Get models list dynamically
|
|
available_models = cls.get_models_list()
|
|
|
|
return {
|
|
"required": {
|
|
"client": ("MODELVERSE_API_CLIENT",),
|
|
"model": (available_models, {"default": DEFAULT_MODEL if DEFAULT_MODEL in available_models else available_models[0] if available_models else DEFAULT_MODEL}),
|
|
"user_prompt": (IO.STRING, {
|
|
"multiline": True,
|
|
"default": "What can you tell me about this?"
|
|
}),
|
|
"temperature": (IO.FLOAT, {
|
|
"default": 0.7,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.1
|
|
}),
|
|
"max_tokens": (IO.INT, {
|
|
"default": 6000,
|
|
"min": 1,
|
|
"max": 128000
|
|
}),
|
|
"top_p": (IO.FLOAT, {
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
},
|
|
"optional": {
|
|
"system_prompt": (IO.STRING, {
|
|
"multiline": True,
|
|
"default": "You are a helpful assistant."
|
|
}),
|
|
"image_in": (IO.IMAGE, {}),
|
|
"files": ("OPENAI_INPUT_FILES", {
|
|
"default": None,
|
|
"tooltip": "Optional file(s) to use as context for the model. Accepts inputs from the OpenAI Input Files node."
|
|
}),
|
|
# "response_format": (["text", "json_object"], {"default": "text"}),
|
|
"response_format": (["text"], {"default": "text"}),
|
|
"presence_penalty": (IO.FLOAT, {
|
|
"default": 0.0,
|
|
"min": -2.0,
|
|
"max": 2.0,
|
|
"step": 0.1
|
|
}),
|
|
"frequency_penalty": (IO.FLOAT, {
|
|
"default": 0.0,
|
|
"min": -2.0,
|
|
"max": 2.0,
|
|
"step": 0.1
|
|
}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (IO.STRING,)
|
|
RETURN_NAMES = ("response",)
|
|
CATEGORY = "UCLOUD_MODELVERSE"
|
|
FUNCTION = "chat"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
"""Tell ComfyUI that this node's inputs may change"""
|
|
return "static"
|
|
|
|
def chat(self,
|
|
client: Dict[str, str],
|
|
model: str,
|
|
user_prompt: str,
|
|
temperature: float,
|
|
max_tokens: int,
|
|
top_p: float,
|
|
unique_id: Optional[str] = None,
|
|
system_prompt: Optional[str] = "You are a helpful assistant.",
|
|
image_in: Optional[Any] = None,
|
|
files: Optional[List[Any]] = None,
|
|
response_format: str = "text",
|
|
presence_penalty: float = 0.0,
|
|
frequency_penalty: float = 0.0) -> tuple:
|
|
|
|
# Create ModelverseClient and get API key
|
|
api_key = client.get("api_key")
|
|
if not api_key:
|
|
raise ValueError("No API key found in the client")
|
|
|
|
# Create ModelverseClient instance to get the actual API key
|
|
modelverse_client = ModelverseClient(api_key)
|
|
|
|
# Initialize OpenAI client with the API key from ModelverseClient
|
|
openai_client = openai.OpenAI(api_key=modelverse_client.api_key,base_url="https://api.modelverse.cn/v1")
|
|
|
|
# Build messages
|
|
messages = []
|
|
|
|
# Add system message if provided
|
|
if system_prompt:
|
|
messages.append({
|
|
"role": "system",
|
|
"content": system_prompt
|
|
})
|
|
|
|
# Build user message content
|
|
user_content = []
|
|
|
|
# Add main prompt
|
|
user_content.append({
|
|
"type": "text",
|
|
"text": user_prompt
|
|
})
|
|
|
|
# Add file content if provided
|
|
if files:
|
|
print(f"OpenAIChat: Processing {len(files)} files")
|
|
for file_info in files:
|
|
# Handle the actual file format from OpenAIInputFiles
|
|
if hasattr(file_info, 'file_data') and hasattr(file_info, 'filename'):
|
|
# Decode base64 content
|
|
try:
|
|
# Extract base64 data (remove data:text/plain;base64, prefix)
|
|
if file_info.file_data.startswith('data:'):
|
|
base64_data = file_info.file_data.split(',', 1)[1]
|
|
else:
|
|
base64_data = file_info.file_data
|
|
|
|
# Decode base64 to text
|
|
content = base64.b64decode(base64_data).decode('utf-8')
|
|
|
|
file_text = f"\n\nFile: {file_info.filename}\nContent:\n{content}"
|
|
user_content.append({
|
|
"type": "text",
|
|
"text": file_text
|
|
})
|
|
print(f"OpenAIChat: Added file {file_info.filename} with {len(content)} characters")
|
|
except Exception as e:
|
|
print(f"OpenAIChat: Error decoding file {file_info.filename}: {str(e)}")
|
|
elif isinstance(file_info, dict) and "content" in file_info and "filename" in file_info:
|
|
# Handle simple dict format (backup)
|
|
file_text = f"\n\nFile: {file_info['filename']}\nContent:\n{file_info['content']}"
|
|
user_content.append({
|
|
"type": "text",
|
|
"text": file_text
|
|
})
|
|
print(f"OpenAIChat: Added file {file_info['filename']} with {len(file_info['content'])} characters")
|
|
else:
|
|
print(f"OpenAIChat: Invalid file format: {file_info}")
|
|
else:
|
|
print("OpenAIChat: No files provided")
|
|
|
|
# Add image if provided
|
|
if image_in is not None:
|
|
# Convert tensor to PIL Image
|
|
# Handle both single images and batches
|
|
if len(image_in.shape) == 4:
|
|
# Take first image from batch
|
|
image_in = image_in[0]
|
|
|
|
# Convert from tensor format (H, W, C) to numpy array
|
|
image_array = image_in.cpu().numpy()
|
|
|
|
# Ensure values are in 0-255 range
|
|
image_array = np.clip(255. * image_array, 0, 255).astype(np.uint8)
|
|
|
|
# Create PIL Image
|
|
pil_image = Image.fromarray(image_array)
|
|
|
|
# Convert PIL Image to base64
|
|
buffered = io.BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
# Add image to content
|
|
user_content.append({
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{img_str}"
|
|
}
|
|
})
|
|
|
|
# Add user message with all content
|
|
if len(user_content) == 1:
|
|
# If only text, use simple format
|
|
messages.append({
|
|
"role": "user",
|
|
"content": user_content[0]["text"]
|
|
})
|
|
else:
|
|
# If multiple content types, use array format
|
|
messages.append({
|
|
"role": "user",
|
|
"content": user_content
|
|
})
|
|
|
|
# Prepare API call parameters
|
|
api_params = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
"top_p": top_p,
|
|
"presence_penalty": presence_penalty,
|
|
"frequency_penalty": frequency_penalty,
|
|
}
|
|
|
|
# Add response format if JSON is requested
|
|
if response_format == "json_object":
|
|
api_params["response_format"] = {"type": "json_object"}
|
|
|
|
try:
|
|
# Make API call to OpenAI
|
|
response = openai_client.chat.completions.create(**api_params)
|
|
|
|
# Extract response content
|
|
if response.choices and len(response.choices) > 0:
|
|
content = response.choices[0].message.content
|
|
if content is None:
|
|
raise ValueError("No content in response")
|
|
|
|
# Display the response message on the node UI
|
|
if unique_id:
|
|
self.display_message_on_node(content.strip(), unique_id)
|
|
|
|
# Return the response as a tuple
|
|
return (content.strip(),)
|
|
else:
|
|
raise ValueError("No choices in response")
|
|
|
|
except Exception as e:
|
|
error_msg = f"OpenAI API Error: {str(e)}"
|
|
print(error_msg)
|
|
return (error_msg,)
|
|
|
|
|
|
class ModelverseInputFiles:
|
|
"""
|
|
Loads and formats input files for OpenAI API.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""
|
|
For details about the supported file input types, see:
|
|
https://platform.openai.com/docs/guides/pdf-files?api-mode=responses
|
|
"""
|
|
input_dir = folder_paths.get_input_directory()
|
|
input_files = [
|
|
f
|
|
for f in os.scandir(input_dir)
|
|
if f.is_file()
|
|
and (f.name.endswith(".txt") or f.name.endswith(".pdf") or f.name.endswith(".md"))
|
|
and f.stat().st_size < 32 * 1024 * 1024
|
|
]
|
|
input_files = sorted(input_files, key=lambda x: x.name)
|
|
input_files = [f.name for f in input_files]
|
|
return {
|
|
"required": {},
|
|
"optional": {
|
|
"file": (
|
|
IO.COMBO,
|
|
{
|
|
"options": input_files,
|
|
"default": input_files[0] if input_files else None,
|
|
"tooltip": "Input files to include as context for the model. Only accepts text (.txt), markdown (.md) and PDF (.pdf) files for now.",
|
|
},
|
|
),
|
|
"OPENAI_INPUT_FILES": (
|
|
"OPENAI_INPUT_FILES",
|
|
{
|
|
"tooltip": "An optional additional file(s) to batch together with the file loaded from this node. Allows chaining of input files so that a single message can include multiple input files.",
|
|
"default": None,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
DESCRIPTION = "Loads and prepares input files (text, markdown, pdf, etc.) to include as inputs for the OpenAI Chat Node. The files will be read by the OpenAI model when generating a response. 🛈 TIP: Can be chained together with other OpenAI Input File nodes."
|
|
RETURN_TYPES = ("OPENAI_INPUT_FILES",)
|
|
FUNCTION = "prepare_files"
|
|
CATEGORY = "UCLOUD_MODELVERSE"
|
|
|
|
def read_file_content(self, file_path: str) -> str:
|
|
"""Read content from a file."""
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
return f.read()
|
|
except UnicodeDecodeError:
|
|
# Try with different encoding if UTF-8 fails
|
|
try:
|
|
with open(file_path, 'r', encoding='gbk') as f:
|
|
return f.read()
|
|
except:
|
|
return f"Error: Unable to read file {file_path} - encoding issue"
|
|
except Exception as e:
|
|
return f"Error reading file {file_path}: {str(e)}"
|
|
|
|
def text_to_data_uri(self, content: str) -> str:
|
|
"""Convert text content to data URI."""
|
|
encoded_content = base64.b64encode(content.encode('utf-8')).decode('ascii')
|
|
return f"data:text/plain;base64,{encoded_content}"
|
|
|
|
def create_input_file_content(self, file_path: str):
|
|
"""Create a file content object compatible with the reference implementation."""
|
|
content = self.read_file_content(file_path)
|
|
|
|
# Create a simple object that mimics the expected structure
|
|
class InputFileContent:
|
|
def __init__(self, content, filename):
|
|
self.file_data = self.text_to_data_uri(content)
|
|
self.filename = filename
|
|
self.type = "input_file"
|
|
self.file_id = None
|
|
|
|
def text_to_data_uri(self, content: str) -> str:
|
|
encoded_content = base64.b64encode(content.encode('utf-8')).decode('ascii')
|
|
return f"data:text/plain;base64,{encoded_content}"
|
|
|
|
return InputFileContent(content, os.path.basename(file_path))
|
|
|
|
def prepare_files(
|
|
self, file: str = None, OPENAI_INPUT_FILES: List[Any] = None
|
|
) -> tuple:
|
|
"""
|
|
Loads and formats input files for OpenAI API.
|
|
"""
|
|
files = []
|
|
|
|
# Add current file if provided
|
|
if file is not None and file != "":
|
|
try:
|
|
file_path = folder_paths.get_annotated_filepath(file)
|
|
input_file_content = self.create_input_file_content(file_path)
|
|
files.append(input_file_content)
|
|
print(f"OpenAIInputFiles: Successfully loaded file {file}")
|
|
except Exception as e:
|
|
print(f"OpenAIInputFiles: Error loading file {file}: {str(e)}")
|
|
|
|
# Add previous files if provided
|
|
if OPENAI_INPUT_FILES is not None:
|
|
files.extend(OPENAI_INPUT_FILES)
|
|
print(f"OpenAIInputFiles: Added {len(OPENAI_INPUT_FILES)} previous files")
|
|
|
|
print(f"OpenAIInputFiles: Returning {len(files)} files total")
|
|
return (files,)
|
|
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ModelverseChat": ModelverseChat,
|
|
"ModelverseInputFiles": ModelverseInputFiles,
|
|
# "OpenAICaptionImage": OpenAICaptionImage,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ModelverseChat": "Modelverse Chat",
|
|
"ModelverseInputFiles": "Modelverse Input Files",
|
|
# "OpenAICaptionImage": "OpenAI Caption Image",
|
|
}
|