131 lines
4.1 KiB
Python
131 lines
4.1 KiB
Python
import torch
|
|
import requests
|
|
import base64
|
|
from PIL import Image
|
|
import io
|
|
from requests.adapters import HTTPAdapter
|
|
from urllib3.util.retry import Retry
|
|
|
|
class OneThingAIImageToTextNode:
|
|
def __init__(self):
|
|
self.api_key = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"api_key": ("STRING", {"default": "", "multiline": False}),
|
|
"model": ("STRING", {"default": "gpt-4o", "multiline": False}),
|
|
"system_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"user_content": ("STRING", {"default": "Please describe this image in detail.", "multiline": True}),
|
|
"retries": ("INT", {"default": 3, "min": 0, "max": 5}),
|
|
"timeout": ("INT", {"default": 20, "min": 5, "max": 100}),
|
|
"max_tokens": ("INT", {"default": 500, "min": 100, "max": 10000, "step": 100}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "describe_image"
|
|
CATEGORY = "OneThingAI/CV"
|
|
|
|
def describe_image(self, image, api_key, model, system_prompt, user_content, retries, timeout, max_tokens):
|
|
if not api_key:
|
|
raise ValueError("API key is required")
|
|
|
|
# Convert tensor to PIL Image
|
|
if isinstance(image, torch.Tensor):
|
|
# Ensure image is in the correct format (C,H,W) -> (H,W,C)
|
|
image = image.squeeze(0)
|
|
image = (image * 255).byte()
|
|
image = image.cpu().numpy()
|
|
if image.shape[0] == 3:
|
|
image = image.transpose(1, 2, 0)
|
|
pil_image = Image.fromarray(image)
|
|
else:
|
|
pil_image = image
|
|
|
|
# Convert PIL Image to base64
|
|
buffered = io.BytesIO()
|
|
pil_image.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode()
|
|
|
|
# Prepare API request
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}"
|
|
}
|
|
|
|
# Prepare messages list
|
|
messages = []
|
|
|
|
# Add system message if provided
|
|
if system_prompt:
|
|
messages.append({
|
|
"role": "system",
|
|
"content": system_prompt
|
|
})
|
|
|
|
# Add user message with content and image
|
|
messages.append({
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": user_content
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{img_str}"
|
|
}
|
|
}
|
|
]
|
|
})
|
|
|
|
payload = {
|
|
"messages": messages,
|
|
"model": model,
|
|
"max_tokens": max_tokens
|
|
}
|
|
|
|
try:
|
|
# Configure retry strategy
|
|
retry_strategy = Retry(
|
|
total=retries,
|
|
backoff_factor=1,
|
|
status_forcelist=[429, 500, 502, 503, 504],
|
|
)
|
|
|
|
# Create session with retry strategy
|
|
session = requests.Session()
|
|
adapter = HTTPAdapter(max_retries=retry_strategy)
|
|
session.mount("https://", adapter)
|
|
|
|
# Make the request with retry and timeout
|
|
response = session.post(
|
|
"https://api-model.onethingai.com/v1/chat/completions",
|
|
headers=headers,
|
|
json=payload,
|
|
timeout=timeout
|
|
)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
|
|
# Extract the description from the response
|
|
description = result['choices'][0]['message']['content']
|
|
return (description,)
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
return (f"Error: {str(e)}",)
|
|
finally:
|
|
session.close()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OneThingAI ImageToText": OneThingAIImageToTextNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OneThingAI ImageToText": "OneThingAI Image Understanding"
|
|
}
|