91 lines
2.5 KiB
Python
91 lines
2.5 KiB
Python
import base64
|
|
import io
|
|
from PIL import Image
|
|
import numpy as np
|
|
import anthropic
|
|
import os
|
|
|
|
api_key = os.environ.get("ANTHROPIC_API_KEY")
|
|
|
|
|
|
class DataSet_ClaudeAIChatImage:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
|
|
api_models = ["claude-3-5-sonnet-latest",
|
|
"claude-3-5-haiku-latest", "claude-3-opus-latest"]
|
|
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": (api_models, {"default": api_models[0]}),
|
|
# "system_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"user_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"max_tokens": ("INT", {"default": 1024})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "🔶DATASET🔶"
|
|
|
|
def to_base64(self, image):
|
|
image = image[0]
|
|
i = 255. * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
buffered = io.BytesIO()
|
|
img.save(buffered, format="PNG")
|
|
return base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
|
|
def generate(self, image, model, user_prompt, max_tokens):
|
|
try:
|
|
|
|
base64img = self.to_base64(image)
|
|
api_client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
|
|
|
|
chat_completion = api_client.messages.create(
|
|
model=model,
|
|
max_tokens=max_tokens,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": "image/png",
|
|
"data": base64img,
|
|
},
|
|
},
|
|
{"type": "text", "text": user_prompt}
|
|
],
|
|
}
|
|
],
|
|
)
|
|
|
|
msg = ""
|
|
|
|
for message in chat_completion.content:
|
|
if message.type == 'text':
|
|
msg = message.text
|
|
|
|
return (msg,)
|
|
|
|
|
|
except Exception as e:
|
|
return (f"Error: {str(e)}",)
|
|
|
|
|
|
N_CLASS_MAPPINGS = {
|
|
"DataSet_ClaudeAIChatImage": DataSet_ClaudeAIChatImage,
|
|
}
|
|
|
|
N_DISPLAY_NAME_MAPPINGS = {
|
|
"DataSet_ClaudeAIChatImage": "DataSet_ClaudeAIChatImage",
|
|
}
|