Files
HM-RunningHub-ComfyUI_RH_LL…/node.py
T

164 lines
5.5 KiB
Python

from openai import OpenAI
import time
from PIL import Image
import numpy as np
import base64
import os
import io
def encode_image_b64(ref_image):
"""
Encode ComfyUI IMAGE tensor to base64 JPEG without resizing.
Notes:
- Keep original resolution (no resize).
- Avoid temporary files (in-memory encoding).
"""
i = 255.0 * ref_image.cpu().numpy()[0]
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
buf = io.BytesIO()
# Use JPEG to match the existing OpenAI-compatible payload mime label.
img.save(buf, format="JPEG", quality=95, optimize=True)
return base64.b64encode(buf.getvalue()).decode("utf-8")
def _get_video_file_path(video):
"""
Try to extract a filesystem path from a ComfyUI VIDEO object.
Returns None if it cannot be resolved.
"""
# VideoFromFile type (private attribute)
if hasattr(video, "_VideoFromFile__file"):
path = getattr(video, "_VideoFromFile__file", None)
if isinstance(path, str) and os.path.exists(path):
return path
# Stream-like sources
if hasattr(video, "get_stream_source"):
try:
stream_source = video.get_stream_source()
if isinstance(stream_source, str) and os.path.exists(stream_source):
return stream_source
except Exception:
pass
# Common attributes
for attr in ("path", "file"):
if hasattr(video, attr):
path = getattr(video, attr, None)
if isinstance(path, str) and os.path.exists(path):
return path
return None
def encode_video_b64(video):
"""
Encode ComfyUI VIDEO object to base64 MP4 bytes.
Notes:
- No ffmpeg processing, no compression, no resizing.
- If a file path is available, read it directly.
- Otherwise, try saving via save_to() to a temp mp4 and read back.
"""
video_path = _get_video_file_path(video)
if video_path:
with open(video_path, "rb") as f:
return base64.b64encode(f.read()).decode("utf-8")
if hasattr(video, "save_to"):
temp_path = f"temp_video_{time.time()}.mp4"
try:
video.save_to(temp_path)
with open(temp_path, "rb") as f:
return base64.b64encode(f.read()).decode("utf-8")
finally:
try:
if os.path.exists(temp_path):
os.remove(temp_path)
except Exception:
pass
raise ValueError(f"Unable to read video data from object type: {type(video)}")
class RH_LLMAPI_Node():
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"api_baseurl": ("STRING", {"multiline": True}),
"api_key": ("STRING", {"default": ""}),
"model": ("STRING", {"default": ""}),
"role": ("STRING", {"multiline": True, "default": "You are a helpful assistant"}),
"prompt": ("STRING", {"multiline": True, "default": "Hello"}),
"temperature": ("FLOAT", {"default": 0.6}),
"seed": ("INT", {"default": 100}),
},
"optional": {
"ref_image": ("IMAGE",),
"video": ("VIDEO",),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("describe",)
FUNCTION = "rh_run_llmapi"
CATEGORY = "Runninghub"
def rh_run_llmapi(self, api_baseurl, api_key, model, role, prompt, temperature, seed, ref_image=None, video=None):
client = OpenAI(api_key=api_key, base_url=api_baseurl)
# Priority: video > image > text (align with reference node behavior)
if video is not None:
base64_video = encode_video_b64(video)
messages = [
{'role': 'system', 'content': f'{role}'},
{'role': 'user',
'content': [
{
"type": "text",
"text": f"{prompt}"
},
{
"type": "video_url",
"video_url": {
"url": f"data:video/mp4;base64,{base64_video}"
}
},
]},
]
elif ref_image is None:
messages = [
{'role': 'system', 'content': f'{role}'},
{'role': 'user', 'content': f'{prompt}'},
]
else:
base64_image = encode_image_b64(ref_image)
messages = [
{'role': 'system', 'content': f'{role}'},
{'role': 'user',
'content': [
{
"type": "text",
"text": f"{prompt}"
},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{base64_image}"
}
},
]},
]
completion = client.chat.completions.create(model=model, messages=messages, temperature=temperature)
if completion is not None and hasattr(completion, 'choices'):
prompt = completion.choices[0].message.content
else:
prompt = 'Error'
return (prompt,)