375 lines
16 KiB
Python
375 lines
16 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import base64
|
|
from io import BytesIO
|
|
import json
|
|
import re
|
|
import urllib.request
|
|
|
|
# Try to import openai, give a clear console error if missing
|
|
try:
|
|
from openai import OpenAI
|
|
except ImportError:
|
|
print("\n[Universal LLM] ERROR: 'openai' Python package is not installed. Please run: pip install openai\n")
|
|
OpenAI = None
|
|
|
|
# --- Fetch models at ComfyUI startup ---
|
|
def get_available_models():
|
|
"""Pings local API ports at startup to build a dynamic dropdown list."""
|
|
models = ["default_model", "gpt-3.5-turbo"] # Fallbacks
|
|
|
|
# Common local LLM ports: 5000 (TabbyAPI/YALS), 8000 (vLLM), 11434 (Ollama), 1234 (LM Studio)
|
|
ports_to_check = [5000, 8000, 11434, 1234]
|
|
|
|
for port in ports_to_check:
|
|
url = f"http://127.0.0.1:{port}/v1/models"
|
|
try:
|
|
req = urllib.request.Request(url)
|
|
with urllib.request.urlopen(req, timeout=0.5) as response:
|
|
data = json.loads(response.read())
|
|
if "data" in data:
|
|
models.extend([m["id"] for m in data["data"]])
|
|
except Exception:
|
|
pass # Ignore if server isn't running on this port
|
|
|
|
# Remove duplicates and return
|
|
return list(set(models))
|
|
|
|
# Fetch the list once when the module loads
|
|
AVAILABLE_MODELS = get_available_models()
|
|
|
|
|
|
# ======================================================================
|
|
# HUB NODE: Pure OpenAI Compatible Client
|
|
# ======================================================================
|
|
|
|
class UniversalOpenAIHubNode:
|
|
def __init__(self):
|
|
self.client = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": "Describe this image or answer the question."}),
|
|
"base_url": ("STRING", {"default": "http://127.0.0.1:5000/v1"}),
|
|
"api_key": ("STRING", {"default": "EMPTY"}),
|
|
"model": (AVAILABLE_MODELS,),
|
|
"system_prompt": ("STRING", {"multiline": True, "default": "You are a helpful AI assistant."}),
|
|
|
|
# Strict OpenAI standard parameters
|
|
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}),
|
|
"top_p": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"max_tokens": ("INT", {"default": 8192, "min": 1, "max": 128000, "step": 1}),
|
|
"presence_penalty": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.05}),
|
|
"frequency_penalty": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.05}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
"optional": {
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
"image_4": ("IMAGE",),
|
|
"extra_parameters": ("LLM_PARAMS",), # Magic socket for backend-specific nodes
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("reasoning_text", "response_text") # Reasoning on top!
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Universal LLM"
|
|
|
|
def process_image(self, img_tensor):
|
|
"""Converts ComfyUI Image Tensor to Base64 String"""
|
|
if len(img_tensor.shape) == 3:
|
|
img_tensor = img_tensor.unsqueeze(0)
|
|
|
|
b64_images = []
|
|
for i in range(img_tensor.shape[0]):
|
|
single_img = img_tensor[i]
|
|
if single_img.shape[-1] != 3:
|
|
single_img = single_img.permute(1, 2, 0)
|
|
|
|
image_np = (single_img.cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(image_np)
|
|
buffered = BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
b64_images.append(img_b64)
|
|
|
|
return b64_images
|
|
|
|
def extract_reasoning(self, content, reasoning):
|
|
"""Robust parser for DeepSeek-R1 and Gemma-style reasoning tags."""
|
|
# If the backend natively split it for us (vLLM, TabbyAPI w/ reasoning: true)
|
|
if reasoning:
|
|
return content.strip(), reasoning.strip()
|
|
|
|
# Define known reasoning tag pairs (Start Tag, End Tag)
|
|
tag_pairs = [
|
|
("<think>", "</think>"), # DeepSeek / Standard
|
|
("<|channel>thought", "<channel|>") # Gemma
|
|
]
|
|
|
|
for start_tag, end_tag in tag_pairs:
|
|
# Check if the end tag exists in the content
|
|
if end_tag in content:
|
|
parts = content.split(end_tag, 1)
|
|
reasoning_raw = parts[0]
|
|
content = parts[1]
|
|
|
|
# Handle case where both start and end tags exist
|
|
if start_tag in reasoning_raw:
|
|
reasoning = reasoning_raw.split(start_tag, 1)[-1]
|
|
else:
|
|
# Handle case where only end tag exists (model started reasoning immediately)
|
|
reasoning = reasoning_raw
|
|
|
|
return content.strip(), reasoning.strip()
|
|
|
|
# If no end tag, check if the start tag exists (generation cut off by max_tokens)
|
|
elif start_tag in content:
|
|
parts = content.split(start_tag, 1)
|
|
content = parts[0] # Captures anything before the start tag
|
|
reasoning = parts[1]
|
|
|
|
return content.strip(), reasoning.strip()
|
|
|
|
# If no tags are found at all, return the content as-is
|
|
return content.strip(), ""
|
|
|
|
def generate(self, prompt, base_url, api_key, model, system_prompt,
|
|
temperature, top_p, max_tokens, presence_penalty, frequency_penalty, seed,
|
|
image_1=None, image_2=None, image_3=None, image_4=None,
|
|
extra_parameters=None):
|
|
|
|
if OpenAI is None:
|
|
return ("Error: 'openai' python package is missing. Check ComfyUI console.", "")
|
|
|
|
# 1. Initialize Client
|
|
if self.client is None or self.client.base_url != base_url or self.client.api_key != api_key:
|
|
self.client = OpenAI(api_key=api_key, base_url=base_url)
|
|
|
|
# 2. Build Messages
|
|
messages = []
|
|
if system_prompt.strip():
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
|
|
user_content = []
|
|
if prompt.strip():
|
|
user_content.append({"type": "text", "text": prompt})
|
|
|
|
# 3. Process Images
|
|
all_images = [img for img in [image_1, image_2, image_3, image_4] if img is not None]
|
|
for img_batch in all_images:
|
|
b64_images = self.process_image(img_batch)
|
|
for b64 in b64_images:
|
|
user_content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{b64}"}
|
|
})
|
|
|
|
messages.append({"role": "user", "content": user_content})
|
|
|
|
# 4. Standard OpenAI Payload
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_tokens": max_tokens,
|
|
"presence_penalty": presence_penalty,
|
|
"frequency_penalty": frequency_penalty,
|
|
"seed": seed
|
|
}
|
|
|
|
# 5. Inject Spoke Node Parameters
|
|
if isinstance(extra_parameters, dict):
|
|
payload["extra_body"] = extra_parameters
|
|
|
|
# 6. Execute Request
|
|
try:
|
|
response = self.client.chat.completions.create(**payload)
|
|
msg = response.choices[0].message
|
|
|
|
raw_content = msg.content or ""
|
|
raw_reasoning = getattr(msg, "reasoning_content", "")
|
|
|
|
# 7. Parse Reasoning
|
|
final_content, final_reasoning = self.extract_reasoning(raw_content, raw_reasoning)
|
|
|
|
# Return swapped to match the new RETURN_NAMES order
|
|
return (final_reasoning, final_content)
|
|
|
|
except Exception as e:
|
|
return (f"API Error: {str(e)}", "")
|
|
|
|
|
|
# ======================================================================
|
|
# SPOKE NODES: Backend-Specific Parameter Generators
|
|
# ======================================================================
|
|
|
|
class TabbyAPISamplersNode:
|
|
"""TabbyAPI specific samplers (ExLlamaV2 engine)"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"top_k": ("INT", {"default": 0, "min": -1, "max": 1000}),
|
|
"min_p": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"repetition_penalty": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
|
|
"xtc_probability": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"xtc_threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"dry_multiplier": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 5.0, "step": 0.1}),
|
|
"dry_base": ("FLOAT", {"default": 1.75, "min": 0.0, "max": 5.0, "step": 0.05}),
|
|
"dry_allowed_length": ("INT", {"default": 2, "min": 0, "max": 100}),
|
|
"mirostat_mode": ("INT", {"default": 0, "min": 0, "max": 2}),
|
|
"max_temp": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}),
|
|
"min_temp": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.05}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLM_PARAMS",)
|
|
RETURN_NAMES = ("tabby_parameters",)
|
|
FUNCTION = "get_params"
|
|
CATEGORY = "Universal LLM/Samplers"
|
|
|
|
def get_params(self, **kwargs):
|
|
params = {}
|
|
if kwargs.get("top_k", 0) > 0: params["top_k"] = kwargs["top_k"]
|
|
if kwargs.get("min_p", 0.0) > 0.0: params["min_p"] = kwargs["min_p"]
|
|
if kwargs.get("repetition_penalty", 1.0) != 1.0: params["repetition_penalty"] = kwargs["repetition_penalty"]
|
|
|
|
if kwargs.get("xtc_probability", 0.0) > 0.0:
|
|
params["xtc_probability"] = kwargs["xtc_probability"]
|
|
params["xtc_threshold"] = kwargs["xtc_threshold"]
|
|
|
|
if kwargs.get("dry_multiplier", 0.0) > 0.0:
|
|
params["dry_multiplier"] = kwargs["dry_multiplier"]
|
|
params["dry_base"] = kwargs["dry_base"]
|
|
params["dry_allowed_length"] = kwargs["dry_allowed_length"]
|
|
|
|
if kwargs.get("mirostat_mode", 0) > 0:
|
|
params["mirostat_mode"] = kwargs["mirostat_mode"]
|
|
|
|
if kwargs.get("max_temp", 1.0) != 1.0 or kwargs.get("min_temp", 1.0) != 1.0:
|
|
params["max_temp"] = kwargs["max_temp"]
|
|
params["min_temp"] = kwargs["min_temp"]
|
|
|
|
return (params,)
|
|
|
|
|
|
class YALSSamplersNode:
|
|
"""YALS specific samplers (llama.cpp engine)"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"top_k": ("INT", {"default": 0, "min": -1, "max": 1000}),
|
|
"min_p": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"repetition_penalty": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
|
|
"xtc_probability": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"xtc_threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"dry_multiplier": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 5.0, "step": 0.1}),
|
|
"dry_base": ("FLOAT", {"default": 1.75, "min": 0.0, "max": 5.0, "step": 0.05}),
|
|
"dry_allowed_length": ("INT", {"default": 2, "min": 0, "max": 100}),
|
|
"mirostat_mode": ("INT", {"default": 0, "min": 0, "max": 2}),
|
|
},
|
|
"optional": {
|
|
"json_schema": ("STRING", {"multiline": True, "default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLM_PARAMS",)
|
|
RETURN_NAMES = ("yals_parameters",)
|
|
FUNCTION = "get_params"
|
|
CATEGORY = "Universal LLM/Samplers"
|
|
|
|
def get_params(self, **kwargs):
|
|
params = {}
|
|
if kwargs.get("top_k", 0) > 0: params["top_k"] = kwargs["top_k"]
|
|
if kwargs.get("min_p", 0.0) > 0.0: params["min_p"] = kwargs["min_p"]
|
|
if kwargs.get("repetition_penalty", 1.0) != 1.0: params["repetition_penalty"] = kwargs["repetition_penalty"]
|
|
|
|
if kwargs.get("xtc_probability", 0.0) > 0.0:
|
|
params["xtc_probability"] = kwargs["xtc_probability"]
|
|
params["xtc_threshold"] = kwargs["xtc_threshold"]
|
|
|
|
if kwargs.get("dry_multiplier", 0.0) > 0.0:
|
|
params["dry_multiplier"] = kwargs["dry_multiplier"]
|
|
params["dry_base"] = kwargs["dry_base"]
|
|
params["dry_allowed_length"] = kwargs["dry_allowed_length"]
|
|
|
|
if kwargs.get("mirostat_mode", 0) > 0:
|
|
params["mirostat_mode"] = kwargs["mirostat_mode"]
|
|
|
|
schema_str = kwargs.get("json_schema", "").strip()
|
|
if schema_str:
|
|
try:
|
|
params["json_schema"] = json.loads(schema_str)
|
|
except json.JSONDecodeError:
|
|
print("[YALS Sampler] Warning: Invalid JSON schema provided. Ignored.")
|
|
|
|
return (params,)
|
|
|
|
|
|
class vLLMSamplersNode:
|
|
"""vLLM specific samplers and engine features"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"top_k": ("INT", {"default": -1, "min": -1, "max": 1000}),
|
|
"min_p": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"repetition_penalty": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05}),
|
|
"use_beam_search": ("BOOLEAN", {"default": False}),
|
|
"best_of": ("INT", {"default": 1, "min": 1, "max": 20}),
|
|
"ignore_eos": ("BOOLEAN", {"default": False}),
|
|
"skip_special_tokens": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"guided_json": ("STRING", {"multiline": True, "default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLM_PARAMS",)
|
|
RETURN_NAMES = ("vllm_parameters",)
|
|
FUNCTION = "get_params"
|
|
CATEGORY = "Universal LLM/Samplers"
|
|
|
|
def get_params(self, **kwargs):
|
|
params = {}
|
|
if kwargs.get("top_k", -1) != -1: params["top_k"] = kwargs["top_k"]
|
|
if kwargs.get("min_p", 0.0) > 0.0: params["min_p"] = kwargs["min_p"]
|
|
if kwargs.get("repetition_penalty", 1.0) != 1.0: params["repetition_penalty"] = kwargs["repetition_penalty"]
|
|
|
|
if kwargs.get("use_beam_search"): params["use_beam_search"] = True
|
|
if kwargs.get("best_of", 1) > 1: params["best_of"] = kwargs["best_of"]
|
|
if kwargs.get("ignore_eos"): params["ignore_eos"] = True
|
|
if not kwargs.get("skip_special_tokens", True): params["skip_special_tokens"] = False
|
|
|
|
schema_str = kwargs.get("guided_json", "").strip()
|
|
if schema_str:
|
|
params["guided_json"] = schema_str
|
|
|
|
return (params,)
|
|
|
|
|
|
# ======================================================================
|
|
# NODE REGISTRATION
|
|
# ======================================================================
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UniversalOpenAIHubNode": UniversalOpenAIHubNode,
|
|
"TabbyAPISamplersNode": TabbyAPISamplersNode,
|
|
"YALSSamplersNode": YALSSamplersNode,
|
|
"vLLMSamplersNode": vLLMSamplersNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UniversalOpenAIHubNode": "Universal OpenAI Hub",
|
|
"TabbyAPISamplersNode": "TabbyAPI Samplers",
|
|
"YALSSamplersNode": "YALS Samplers",
|
|
"vLLMSamplersNode": "vLLM Samplers",
|
|
} |