520 lines
18 KiB
Python
520 lines
18 KiB
Python
import io
|
|
import os
|
|
import openai
|
|
import base64
|
|
import numpy as np
|
|
import requests
|
|
import json
|
|
from PIL import Image
|
|
|
|
_CATEGORY = "KYNode/LLM"
|
|
PROTOCOLS = ["openai", "ollama"]
|
|
MODELS = [
|
|
"gpt-4o",
|
|
"gpt-4o-mini",
|
|
"chatgpt-4o-latest",
|
|
"gpt-4-turbo",
|
|
"OpenGVLab/InternVL3-8B",
|
|
"MiniCPM-V-2_6_awq"
|
|
]
|
|
OLLAMA_MODELS = [
|
|
"aha2025/llama-joycaption-beta-one-hf-llava:Q8_0",
|
|
"qwen2.5vl:7b",
|
|
"qwen2.5vl:32b",
|
|
"XiaomiMiMo/MiMo-VL-7B-RL",
|
|
"openbmb/minicpm-o2.6",
|
|
"openbmb/minicpm-v2.6"
|
|
]
|
|
|
|
|
|
class OpenAICaptionImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_in": ("IMAGE", {}),
|
|
"protocol": (PROTOCOLS, {"default": "openai"}),
|
|
"custom_model": ("STRING", {"default": ""}),
|
|
"model": (MODELS,),
|
|
"ollama_model": (OLLAMA_MODELS,),
|
|
"system_prompt": (
|
|
"STRING",
|
|
{"multiline": True, "default": "You are a movie scene director"},
|
|
),
|
|
"caption_prompt": (
|
|
"STRING",
|
|
{"multiline": True, "default": "Describe this image without any speculations"},
|
|
),
|
|
"response_format": (
|
|
["text", "json_object"],
|
|
{"default": "text"},
|
|
),
|
|
"max_tokens": ("INT", {"default": 200}),
|
|
"temperature": ("FLOAT", {"default": 0.5}),
|
|
"top_p": ("FLOAT", {"default": 0.9}),
|
|
"frequency_penalty": ("FLOAT", {"default": 0.0}),
|
|
"presence_penalty": ("FLOAT", {"default": 0.0}),
|
|
"base_url": ("STRING", {"default": "http://127.0.0.1:23333"}),
|
|
"api_key": ("STRING", {"default": "sk-0123456"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text_out",)
|
|
CATEGORY = _CATEGORY
|
|
FUNCTION = "caption"
|
|
|
|
def caption(
|
|
self,
|
|
image_in,
|
|
protocol,
|
|
custom_model,
|
|
model,
|
|
ollama_model,
|
|
system_prompt,
|
|
caption_prompt,
|
|
response_format,
|
|
max_tokens,
|
|
temperature,
|
|
top_p,
|
|
frequency_penalty,
|
|
presence_penalty,
|
|
base_url,
|
|
api_key,
|
|
):
|
|
# image to base64, image is bwhc tensor
|
|
# Convert tensor to PIL Image
|
|
pil_image = Image.fromarray(
|
|
np.clip(255.0 * image_in.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
)
|
|
|
|
# 根据协议选择模型
|
|
if protocol == "ollama":
|
|
if not custom_model:
|
|
custom_model = ollama_model
|
|
return self._caption_ollama(
|
|
pil_image, custom_model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, base_url
|
|
)
|
|
else: # openai
|
|
if not custom_model:
|
|
custom_model = model
|
|
return self._caption_openai(
|
|
pil_image, custom_model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, top_p, frequency_penalty,
|
|
presence_penalty, base_url, api_key, response_format
|
|
)
|
|
|
|
def _caption_openai(self, pil_image, model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, top_p, frequency_penalty,
|
|
presence_penalty, base_url, api_key, response_format="text"):
|
|
# Convert PIL Image to base64
|
|
buffered = io.BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
# Set up OpenAI client
|
|
client = openai.OpenAI(api_key=api_key, base_url=base_url)
|
|
|
|
# Prepare API call parameters
|
|
params = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": caption_prompt},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{img_str}"},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
"timeout": 30,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"frequency_penalty": frequency_penalty,
|
|
"presence_penalty": presence_penalty,
|
|
}
|
|
|
|
# Add response_format if json_object is selected
|
|
if response_format == "json_object":
|
|
params["response_format"] = {"type": "json_object"}
|
|
|
|
# Make API call to OpenAI
|
|
response = client.chat.completions.create(**params)
|
|
|
|
if response.choices[0].message.content is None:
|
|
raise ValueError("No content in response")
|
|
|
|
# Extract and return the caption
|
|
caption = response.choices[0].message.content.strip()
|
|
return (caption,)
|
|
|
|
def _caption_ollama(self, pil_image, model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, base_url):
|
|
# Convert PIL Image to base64
|
|
buffered = io.BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
# Prepare Ollama API request
|
|
url = f"{base_url}/api/chat"
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{
|
|
"role": "user",
|
|
"content": caption_prompt,
|
|
"images": [img_str]
|
|
}
|
|
],
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": temperature,
|
|
"num_predict": max_tokens
|
|
}
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, json=payload, timeout=30)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
|
|
if "message" in result and "content" in result["message"]:
|
|
caption = result["message"]["content"].strip()
|
|
return (caption,)
|
|
else:
|
|
raise ValueError("Invalid response format from Ollama")
|
|
except Exception as e:
|
|
raise ValueError(f"Ollama API error: {str(e)}")
|
|
|
|
|
|
class OpenAICaptionImages:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", {}),
|
|
"protocol": (PROTOCOLS, {"default": "openai"}),
|
|
"custom_model": ("STRING", {"default": ""}),
|
|
"model": (MODELS,),
|
|
"ollama_model": (OLLAMA_MODELS,),
|
|
"system_prompt": (
|
|
"STRING",
|
|
{"multiline": True, "default": "You are a movie scene director"},
|
|
),
|
|
"caption_prompt": (
|
|
"STRING",
|
|
{"default": "Describe this image without any speculations"},
|
|
),
|
|
"max_tokens": ("INT", {"default": 200}),
|
|
"temperature": ("FLOAT", {"default": 0.5}),
|
|
"top_p": ("FLOAT", {"default": 0.9}),
|
|
"frequency_penalty": ("FLOAT", {"default": 0.0}),
|
|
"presence_penalty": ("FLOAT", {"default": 0.0}),
|
|
"base_url": ("STRING", {"default": "http://127.0.0.1:23333"}),
|
|
"api_key": ("STRING", {"default": "sk-0123456"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("captions",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
CATEGORY = _CATEGORY
|
|
FUNCTION = "caption_batch"
|
|
DESCRIPTION = "Caption multiple images using openai-protocol local LLM services (support MiniCPM-V)"
|
|
|
|
def caption_batch(
|
|
self,
|
|
images,
|
|
protocol,
|
|
custom_model,
|
|
model,
|
|
ollama_model,
|
|
system_prompt,
|
|
caption_prompt,
|
|
max_tokens,
|
|
temperature,
|
|
top_p,
|
|
frequency_penalty,
|
|
presence_penalty,
|
|
base_url,
|
|
api_key,
|
|
):
|
|
# 根据协议选择模型
|
|
if protocol == "ollama":
|
|
if not custom_model:
|
|
custom_model = ollama_model
|
|
return self._caption_batch_ollama(
|
|
images, custom_model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, base_url
|
|
)
|
|
else: # openai
|
|
if not custom_model:
|
|
custom_model = model
|
|
return self._caption_batch_openai(
|
|
images, custom_model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, top_p, frequency_penalty,
|
|
presence_penalty, base_url, api_key
|
|
)
|
|
|
|
def _caption_batch_openai(self, images, model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, top_p, frequency_penalty,
|
|
presence_penalty, base_url, api_key):
|
|
# 设置 OpenAI 客户端
|
|
client = openai.OpenAI(api_key=api_key, base_url=base_url)
|
|
|
|
# 处理输入图像
|
|
if len(images.shape) == 4: # 批次输入 (B,H,W,C)
|
|
batch_size = images.shape[0]
|
|
image_list = [images[i:i+1] for i in range(batch_size)]
|
|
else: # 单张图片输入
|
|
image_list = [images]
|
|
|
|
captions = []
|
|
for img in image_list:
|
|
try:
|
|
# 转换图像格式
|
|
pil_image = Image.fromarray(
|
|
np.clip(255.0 * img.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
)
|
|
|
|
# 转换为 base64
|
|
buffered = io.BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
# 准备消息内容
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": caption_prompt},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{img_str}"},
|
|
},
|
|
],
|
|
},
|
|
]
|
|
|
|
# 发送请求
|
|
response = client.chat.completions.create(
|
|
model=model,
|
|
messages=messages,
|
|
timeout=30,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
frequency_penalty=frequency_penalty,
|
|
presence_penalty=presence_penalty,
|
|
)
|
|
|
|
if response.choices[0].message.content is None:
|
|
captions.append("Error: No content in response")
|
|
else:
|
|
caption = response.choices[0].message.content.strip()
|
|
captions.append(caption)
|
|
|
|
except Exception as e:
|
|
captions.append(f"Error: {str(e)}")
|
|
|
|
return (captions,)
|
|
|
|
def _caption_batch_ollama(self, images, model, system_prompt, caption_prompt,
|
|
max_tokens, temperature, base_url):
|
|
# 处理输入图像
|
|
if len(images.shape) == 4: # 批次输入 (B,H,W,C)
|
|
batch_size = images.shape[0]
|
|
image_list = [images[i:i+1] for i in range(batch_size)]
|
|
else: # 单张图片输入
|
|
image_list = [images]
|
|
|
|
captions = []
|
|
for img in image_list:
|
|
try:
|
|
# 转换图像格式
|
|
pil_image = Image.fromarray(
|
|
np.clip(255.0 * img.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)
|
|
)
|
|
|
|
# 转换为 base64
|
|
buffered = io.BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
# Prepare Ollama API request
|
|
url = f"{base_url}/api/chat"
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{
|
|
"role": "user",
|
|
"content": caption_prompt,
|
|
"images": [img_str]
|
|
}
|
|
],
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": temperature,
|
|
"num_predict": max_tokens
|
|
}
|
|
}
|
|
|
|
response = requests.post(url, json=payload, timeout=30)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
|
|
if "message" in result and "content" in result["message"]:
|
|
caption = result["message"]["content"].strip()
|
|
captions.append(caption)
|
|
else:
|
|
captions.append("Error: Invalid response format from Ollama")
|
|
|
|
except Exception as e:
|
|
captions.append(f"Error: {str(e)}")
|
|
|
|
return (captions,)
|
|
|
|
|
|
class OpenAIChat:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"protocol": (PROTOCOLS, {"default": "openai"}),
|
|
"custom_model": ("STRING", {"default": ""}),
|
|
"model": (MODELS,),
|
|
"ollama_model": (OLLAMA_MODELS,),
|
|
"system_prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "You are a professional translator for chinese and english",
|
|
},
|
|
),
|
|
"chat_prompt": ("STRING", {"default": "Hi"}),
|
|
"max_tokens": ("INT", {"default": 200}),
|
|
"temperature": ("FLOAT", {"default": 0.5}),
|
|
"top_p": ("FLOAT", {"default": 0.9}),
|
|
"frequency_penalty": ("FLOAT", {"default": 0.0}),
|
|
"presence_penalty": ("FLOAT", {"default": 0.0}),
|
|
"base_url": ("STRING", {"default": "http://127.0.0.1:23333"}),
|
|
"api_key": ("STRING", {"default": "sk-0123456"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text_out",)
|
|
CATEGORY = _CATEGORY
|
|
FUNCTION = "chat"
|
|
|
|
def chat(
|
|
self,
|
|
protocol,
|
|
custom_model,
|
|
model,
|
|
ollama_model,
|
|
system_prompt,
|
|
chat_prompt,
|
|
max_tokens,
|
|
temperature,
|
|
base_url,
|
|
api_key,
|
|
frequency_penalty,
|
|
top_p,
|
|
presence_penalty,
|
|
):
|
|
# 根据协议选择模型
|
|
if protocol == "ollama":
|
|
if not custom_model:
|
|
custom_model = ollama_model
|
|
return self._chat_ollama(
|
|
custom_model, system_prompt, chat_prompt,
|
|
max_tokens, temperature, base_url
|
|
)
|
|
else: # openai
|
|
if not custom_model:
|
|
custom_model = model
|
|
return self._chat_openai(
|
|
custom_model, system_prompt, chat_prompt,
|
|
max_tokens, temperature, top_p, frequency_penalty,
|
|
presence_penalty, base_url, api_key
|
|
)
|
|
|
|
def _chat_openai(self, model, system_prompt, chat_prompt,
|
|
max_tokens, temperature, top_p, frequency_penalty,
|
|
presence_penalty, base_url, api_key):
|
|
client = openai.OpenAI(api_key=api_key, base_url=base_url)
|
|
|
|
# Make API call to OpenAI
|
|
response = client.chat.completions.create(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": chat_prompt},
|
|
],
|
|
timeout=30,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
frequency_penalty=frequency_penalty,
|
|
presence_penalty=presence_penalty,
|
|
)
|
|
if response.choices[0].message.content is None:
|
|
raise ValueError("No content in response")
|
|
|
|
# Extract and return the caption
|
|
result = response.choices[0].message.content.strip()
|
|
return (result,)
|
|
|
|
def _chat_ollama(self, model, system_prompt, chat_prompt,
|
|
max_tokens, temperature, base_url):
|
|
# Prepare Ollama API request
|
|
url = f"{base_url}/api/chat"
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": chat_prompt}
|
|
],
|
|
"stream": False,
|
|
"options": {
|
|
"temperature": temperature,
|
|
"num_predict": max_tokens
|
|
}
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, json=payload, timeout=30)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
|
|
if "message" in result and "content" in result["message"]:
|
|
result_text = result["message"]["content"].strip()
|
|
return (result_text,)
|
|
else:
|
|
raise ValueError("Invalid response format from Ollama")
|
|
except Exception as e:
|
|
raise ValueError(f"Ollama API error: {str(e)}")
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"KY_OpenAICaptionImage": OpenAICaptionImage,
|
|
"KY_OpenAICaptionImages": OpenAICaptionImages,
|
|
"KY_OpenAIChat": OpenAIChat,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"KY_OpenAICaptionImage": "KY Caption Image by openai-protocol local LLM services",
|
|
"KY_OpenAICaptionImages": "KY Caption Images Batch by openai-protocol local LLM services",
|
|
"KY_OpenAIChat": "KY Chat with openai-protocol local LLM services",
|
|
}
|