Add files via upload
This commit is contained in:
+13
@@ -0,0 +1,13 @@
|
||||
from .llm_vlm_api import LLM_VLM_API, LLM_VLM_Prompt
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LLM_VLM_API": LLM_VLM_API,
|
||||
"LLM_VLM_Prompt": LLM_VLM_Prompt
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LLM_VLM_API": "LLM (VLM) API",
|
||||
"LLM_VLM_Prompt": "LLM (VLM) Prompt"
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,87 @@
|
||||
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": "", "password": True}),
|
||||
"model": ("STRING", {"default": ""}),
|
||||
"temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
},
|
||||
"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, **kwargs):
|
||||
system_prompt = kwargs.get("system prompt", "")
|
||||
image = kwargs.get("image", None)
|
||||
text = kwargs.get("text", "")
|
||||
|
||||
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,
|
||||
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,)
|
||||
@@ -0,0 +1,3 @@
|
||||
openai
|
||||
numpy
|
||||
Pillow
|
||||
Reference in New Issue
Block a user