141 lines
4.9 KiB
Python
141 lines
4.9 KiB
Python
import base64
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from openai import OpenAI
|
|
|
|
|
|
class LLMImageDescription:
|
|
def __init__(self):
|
|
self.output_dir = "output"
|
|
self.type = "output"
|
|
self._client = None
|
|
self._current_api_key = None
|
|
self._current_api_url = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": (["gpt-4o", "claude-3-5-sonnet-20241022", "gemini-1.5-pro-latest", "gemini-2.0-flash-exp"],),
|
|
"api_url": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"password": True
|
|
}),
|
|
"api_key": ("STRING", {
|
|
"default": "",
|
|
"multiline": False,
|
|
"password": True
|
|
}),
|
|
"prompt_template": ("STRING", {
|
|
"default": "Please describe this image in detail:",
|
|
"multiline": True
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("description",)
|
|
FUNCTION = "process_image"
|
|
CATEGORY = "image/text"
|
|
|
|
def get_client(self, api_key, api_url=None):
|
|
"""Get or create OpenAI client"""
|
|
# 检查参数是否发生变化
|
|
if (self._client and
|
|
(self._current_api_key != api_key or
|
|
self._current_api_url != api_url)):
|
|
# 参数变化，清除旧客户端
|
|
self._client = None
|
|
|
|
if not self._client:
|
|
kwargs = {"api_key": api_key}
|
|
if api_url:
|
|
kwargs["base_url"] = api_url
|
|
# 创建新客户端并保存当前参数
|
|
self._client = OpenAI(**kwargs)
|
|
self._current_api_key = api_key
|
|
self._current_api_url = api_url
|
|
|
|
return self._client
|
|
|
|
@staticmethod
|
|
def convert_image_to_base64(image):
|
|
# Convert PyTorch tensor to PIL Image
|
|
image = image.cpu().numpy()
|
|
image = (image * 255).astype(np.uint8)
|
|
if image.shape[0] == 3: # If image is in CHW format
|
|
image = np.transpose(image, (1, 2, 0))
|
|
pil_image = Image.fromarray(image)
|
|
|
|
# Convert PIL Image to base64
|
|
buffered = BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode()
|
|
return img_str
|
|
|
|
@staticmethod
|
|
def process_with_openai_compatible(prompt_template, base64_image, client, model):
|
|
"""Process image using OpenAI API format"""
|
|
try:
|
|
response = client.chat.completions.create(
|
|
model=model,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": prompt_template
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{base64_image}"
|
|
}
|
|
}
|
|
]
|
|
}
|
|
],
|
|
max_tokens=300
|
|
)
|
|
return response.choices[0].message.content
|
|
except Exception as e:
|
|
print(f"API error: {str(e)}")
|
|
return f"Error: API request failed - {str(e)}"
|
|
|
|
def process_image(self, image, model, api_url, api_key, prompt_template):
|
|
try:
|
|
# Convert the first image in batch to base64
|
|
if len(image.shape) == 4:
|
|
image = image[0]
|
|
base64_image = self.convert_image_to_base64(image)
|
|
|
|
# Get appropriate API URL based on model
|
|
if not api_url:
|
|
api_urls = {
|
|
"gpt-4o": "https://api.openai.com/v1/chat/completions",
|
|
"claude-3-5-sonnet-20241022": "https://api.anthropic.com/v1/messages",
|
|
"gemini-1.5-pro-latest": "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent",
|
|
"gemini-2.0-flash-exp": "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-exp:generateContent"
|
|
}
|
|
api_url = api_urls.get(model, "")
|
|
|
|
# Process using unified OpenAI SDK
|
|
client = self.get_client(api_key, api_url)
|
|
description = self.process_with_openai_compatible(prompt_template, base64_image, client, model)
|
|
|
|
return (description,)
|
|
|
|
except Exception as e:
|
|
print(f"Error in image description: {str(e)}")
|
|
return (f"Error: Failed to generate image description. {str(e)}",)
|
|
|
|
def __del__(self):
|
|
"""Cleanup client on deletion"""
|
|
if self._client:
|
|
self._client.close()
|