154 lines
4.6 KiB
Python
154 lines
4.6 KiB
Python
import base64
|
|
import io
|
|
from PIL import Image
|
|
import numpy as np
|
|
from openai import OpenAI
|
|
|
|
|
|
class ImageToLLMPrompt:
|
|
"""
|
|
A custom node that uploads an image and sends it to an OpenAI-compatible API
|
|
with optional system and user prompts, returning the LLM response as text.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"base_url": ("STRING", {
|
|
"default": "http://localhost:8080/v1",
|
|
"multiline": False
|
|
}),
|
|
"api_key": ("STRING", {
|
|
"default": "sk-test",
|
|
"multiline": False
|
|
}),
|
|
"model": ("STRING", {
|
|
"default": "gpt-4o",
|
|
"multiline": False
|
|
}),
|
|
},
|
|
"optional": {
|
|
"system_prompt": ("STRING", {
|
|
"default": "",
|
|
"multiline": True
|
|
}),
|
|
"user_prompt": ("STRING", {
|
|
"default": "",
|
|
"multiline": True
|
|
}),
|
|
"temperature": ("FLOAT", {
|
|
"default": 0.7,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.1
|
|
}),
|
|
"max_tokens": ("INT", {
|
|
"default": 2048,
|
|
"min": 1,
|
|
"max": 32768,
|
|
"step": 1
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "process_image"
|
|
CATEGORY = "LLM"
|
|
OUTPUT_NODE = True
|
|
|
|
def tensor_to_base64(self, image_tensor):
|
|
"""Convert ComfyUI image tensor to base64 string"""
|
|
# ComfyUI images are in format [B, H, W, C] with values 0-1
|
|
image_np = (image_tensor[0].cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
# Convert to PIL Image
|
|
pil_image = Image.fromarray(image_np)
|
|
|
|
# Convert to base64
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format="PNG")
|
|
buffer.seek(0)
|
|
base64_image = base64.b64encode(buffer.read()).decode('utf-8')
|
|
|
|
return base64_image
|
|
|
|
def process_image(self, image, base_url, api_key, model,
|
|
system_prompt="", user_prompt="", temperature=0.7, max_tokens=2048):
|
|
"""Send image to LLM and return response"""
|
|
|
|
# Convert image to base64
|
|
base64_image = self.tensor_to_base64(image)
|
|
|
|
# Create OpenAI client
|
|
client = OpenAI(
|
|
base_url=base_url,
|
|
api_key=api_key,
|
|
timeout=180
|
|
)
|
|
|
|
# Build messages array
|
|
messages = []
|
|
|
|
# Add optional system prompt
|
|
if system_prompt and system_prompt.strip():
|
|
messages.append({
|
|
"role": "system",
|
|
"content": system_prompt
|
|
})
|
|
|
|
# Build user message content
|
|
user_content = []
|
|
|
|
# Add optional user prompt text
|
|
if user_prompt and user_prompt.strip():
|
|
user_content.append({
|
|
"type": "text",
|
|
"text": user_prompt
|
|
})
|
|
|
|
# Add the image
|
|
user_content.append({
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{base64_image}"
|
|
}
|
|
})
|
|
|
|
messages.append({
|
|
"role": "user",
|
|
"content": user_content
|
|
})
|
|
|
|
# Make API call
|
|
try:
|
|
print(f"[ImageToLLMPrompt] Sending request to {base_url} with model {model}")
|
|
response = client.chat.completions.create(
|
|
model=model,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
messages=messages
|
|
)
|
|
|
|
result = response.choices[0].message.content
|
|
print(f"[ImageToLLMPrompt] LLM Response received (length: {len(result)} chars)")
|
|
print(f"[ImageToLLMPrompt] Response content:\n{result}")
|
|
return {"ui": {"text": (result,)}, "result": (result,)}
|
|
|
|
except Exception as e:
|
|
error_msg = f"LLM API call failed: {str(e)}"
|
|
print(f"[ImageToLLMPrompt] ERROR: {error_msg}")
|
|
return {"ui": {"text": (error_msg,)}, "result": (error_msg,)}
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageToLLMPrompt": ImageToLLMPrompt
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageToLLMPrompt": "Image to LLM Prompt"
|
|
}
|