89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
import base64
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io as pyio
|
|
from openai import OpenAI
|
|
|
|
class LLM_VLM_API:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"base_url": ("STRING", {"default": ""}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
"model": ("STRING", {"default": ""}),
|
|
"temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
"max_tokens": ("INT", {"default": 4096, "min": 1, "max": 8192, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"system prompt": ("STRING", {"forceInput": True}),
|
|
"text": ("STRING", {"forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("response",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "LLM(VLM) API"
|
|
|
|
def execute(self, base_url, api_key, model, temperature, max_tokens, **kwargs):
|
|
system_prompt = kwargs.get("system prompt", None)
|
|
image = kwargs.get("image", None)
|
|
text = kwargs.get("text", None)
|
|
|
|
if not text and image is None:
|
|
return ("Error: No text or image provided.",)
|
|
|
|
content = []
|
|
if text:
|
|
content.append({"type": "text", "text": text})
|
|
|
|
if image is not None:
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_img = Image.fromarray(img_np)
|
|
buf = pyio.BytesIO()
|
|
pil_img.save(buf, format="PNG")
|
|
img_b64 = base64.b64encode(buf.getvalue()).decode()
|
|
content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{img_b64}"}
|
|
})
|
|
|
|
client = OpenAI(api_key=api_key, base_url=base_url)
|
|
|
|
try:
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
|
|
messages.append({"role": "user", "content": content})
|
|
|
|
resp = client.chat.completions.create(
|
|
model=model,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
messages=messages
|
|
)
|
|
result = resp.choices[0].message.content
|
|
except Exception as e:
|
|
result = str(e)
|
|
|
|
return (result,)
|
|
|
|
class LLM_VLM_Prompt:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "LLM(VLM) API"
|
|
|
|
def execute(self, prompt):
|
|
return (prompt,) |