558 lines
18 KiB
Python
558 lines
18 KiB
Python
"""
|
|
ComfyUI nodes for the LLM.
|
|
Converted to V3 schema
|
|
"""
|
|
|
|
from dataclasses import asdict
|
|
from typing import Literal, Optional
|
|
from datetime import datetime
|
|
from torch import Tensor
|
|
from PIL import Image
|
|
from typing_extensions import override
|
|
from jinja2 import Environment, StrictUndefined
|
|
from comfy_api.latest import ComfyExtension, io
|
|
from .core.chat import chat_completion, grounding, captioning, ExtraParameters
|
|
from .core.utils import tensor_to_pil, pil_to_tensor
|
|
from .core.model_clients import init_client, ClientInfo
|
|
|
|
|
|
class LLMLoader(io.ComfyNode):
|
|
"""
|
|
Load the LLM client.
|
|
|
|
Class methods
|
|
-------------
|
|
define_schema (io.Schema):
|
|
Tell the main program the metadata, input, output parameters of nodes.
|
|
execute:
|
|
Load the LLM client.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""
|
|
Tell the main program the metadata, input, output parameters of nodes.
|
|
"""
|
|
return io.Schema(
|
|
node_id="APILLMLoader",
|
|
display_name="API LLM Loader",
|
|
category="LLM",
|
|
inputs=[
|
|
io.String.Input(
|
|
"base_url",
|
|
multiline=False,
|
|
default="https://api.openai.com/v1",
|
|
tooltip="The API URL of the provider.",
|
|
),
|
|
io.String.Input(
|
|
"api_key",
|
|
multiline=False,
|
|
default="sk-1234567890",
|
|
tooltip="The API key of the provider.",
|
|
),
|
|
io.String.Input(
|
|
"model_name",
|
|
multiline=False,
|
|
default="gpt-5-high",
|
|
tooltip="The name of the model.",
|
|
),
|
|
io.Combo.Input(
|
|
"client_type",
|
|
options=[
|
|
"openai",
|
|
"openai-responses",
|
|
"mistral",
|
|
"ollama",
|
|
"anthropic",
|
|
],
|
|
default="openai",
|
|
tooltip="The type of the LLM client (OpenAI chat comptitable? or Mistral?).",
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Custom("CLIENT_INFO").Output(display_name="client_info"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
base_url: str,
|
|
api_key: str,
|
|
model_name: str,
|
|
client_type: Literal["openai", "openai-responses", "mistral", "ollama"],
|
|
) -> io.NodeOutput:
|
|
"""
|
|
Load the LLM client.
|
|
"""
|
|
client_info: ClientInfo = init_client(
|
|
client_type=client_type,
|
|
base_url=base_url,
|
|
api_key=api_key,
|
|
model=model_name,
|
|
)
|
|
return io.NodeOutput(client_info)
|
|
|
|
|
|
class ExtraParametersComfy(io.ComfyNode):
|
|
"""
|
|
Extra parameters for the chat completion.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""
|
|
Extra parameters for the chat completion.
|
|
"""
|
|
return io.Schema(
|
|
node_id="ExtraParameters",
|
|
display_name="Extra Parameters",
|
|
category="LLM",
|
|
inputs=[
|
|
io.Combo.Input(
|
|
"thinking",
|
|
options=["disabled", "enabled"],
|
|
default="disabled",
|
|
tooltip="Whether to enable the thinking mode for the LLM model.",
|
|
),
|
|
io.Combo.Input(
|
|
"reasoning_effort",
|
|
options=["minimal", "low", "medium", "high"],
|
|
default="medium",
|
|
tooltip="The reasoning effort for the LLM model.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Custom("EXTRA_PARAMETERS").Output(display_name="extra_parameters"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
thinking: Literal["disabled", "enabled"],
|
|
reasoning_effort: Literal["minimal", "low", "medium", "high"],
|
|
) -> io.NodeOutput:
|
|
"""
|
|
Extra parameters for the chat completion.
|
|
"""
|
|
return io.NodeOutput(
|
|
ExtraParameters(thinking=thinking, reasoning_effort=reasoning_effort)
|
|
)
|
|
|
|
|
|
class ChatViaAPI(io.ComfyNode):
|
|
"""
|
|
Chat with the LLM model.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""
|
|
Tell the main program the metadata, input, output parameters of nodes.
|
|
"""
|
|
return io.Schema(
|
|
node_id="ChatViaAPI",
|
|
display_name="API Chat",
|
|
category="LLM",
|
|
inputs=[
|
|
io.Image.Input(
|
|
"images",
|
|
tooltip="The images for the LLM model.",
|
|
optional=True,
|
|
),
|
|
io.Custom("CLIENT_INFO").Input(
|
|
"client_info",
|
|
tooltip="The LLM client info.",
|
|
),
|
|
io.String.Input(
|
|
"system_prompt",
|
|
multiline=True,
|
|
default="You are a helpful assistant.",
|
|
tooltip="The system prompt for the LLM model.",
|
|
),
|
|
io.String.Input(
|
|
"user_prompt",
|
|
multiline=True,
|
|
default="Hello world!",
|
|
tooltip="The system prompt for the LLM model.",
|
|
),
|
|
io.Float.Input(
|
|
"temperature",
|
|
default=1.0,
|
|
min=0.0,
|
|
max=2.0,
|
|
step=0.1,
|
|
tooltip="The temperature parameter for the LLM model.",
|
|
),
|
|
io.Float.Input(
|
|
"top_p",
|
|
default=0.95,
|
|
min=0.0,
|
|
max=2.0,
|
|
step=0.1,
|
|
tooltip="The top_p parameter for the LLM model.",
|
|
),
|
|
io.Int.Input(
|
|
"top_k",
|
|
default=40,
|
|
min=1,
|
|
max=99,
|
|
step=1,
|
|
tooltip="The top_k parameter for the LLM model. Only Ollama is supported.",
|
|
),
|
|
io.Int.Input(
|
|
"max_tokens",
|
|
default=1024,
|
|
min=1,
|
|
max=1_000_000,
|
|
step=1,
|
|
tooltip="The max_tokens for the LLM model.",
|
|
),
|
|
io.Boolean.Input(
|
|
"unload_model_after_chat",
|
|
default=True,
|
|
tooltip=(
|
|
"Whether to unload the LLM model after the chat. "
|
|
+ "Only Ollama is supported."
|
|
),
|
|
),
|
|
io.Custom("EXTRA_PARAMETERS").Input(
|
|
"extra_parameters",
|
|
tooltip="The extra parameters for the LLM model.",
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[
|
|
io.String.Output(display_name="response"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
client_info: ClientInfo,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
temperature: float,
|
|
top_p: float,
|
|
top_k: int,
|
|
max_tokens: int,
|
|
unload_model_after_chat: bool,
|
|
extra_parameters: Optional[ExtraParameters] = None,
|
|
images: Optional[Tensor] = None,
|
|
) -> io.NodeOutput:
|
|
"""
|
|
Chat with the LLM model.
|
|
"""
|
|
if images is not None and images.any():
|
|
pil_images: Optional[list[Image.Image]] = []
|
|
for image in images:
|
|
pil_images.append(tensor_to_pil(image))
|
|
|
|
response: str = chat_completion(
|
|
client_info=client_info,
|
|
system_prompt=system_prompt,
|
|
user_prompt=user_prompt,
|
|
images=pil_images,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
max_tokens=max_tokens,
|
|
unload_after_chat=unload_model_after_chat,
|
|
extra_parameters=extra_parameters,
|
|
)
|
|
return io.NodeOutput(response)
|
|
|
|
|
|
class GenerateBBOX(io.ComfyNode):
|
|
"""
|
|
Generate bboxes based on LLM's grounding abitility.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""
|
|
Tell the main program the metadata, input, output parameters of nodes.
|
|
"""
|
|
return io.Schema(
|
|
node_id="GenerateBBOXViaAPI",
|
|
display_name="Generate BBOXes",
|
|
category="LLM",
|
|
inputs=[
|
|
io.Custom("CLIENT_INFO").Input(
|
|
"client_info",
|
|
tooltip="The LLM client info.",
|
|
),
|
|
io.String.Input(
|
|
"items",
|
|
default="dog",
|
|
tooltip="The system prompt for the LLM model.",
|
|
),
|
|
io.Boolean.Input(
|
|
"unload_model_after_chat",
|
|
default=True,
|
|
tooltip=(
|
|
"Whether to unload the LLM model after the chat. "
|
|
+ "Only Ollama is supported."
|
|
),
|
|
),
|
|
io.Image.Input(
|
|
"image",
|
|
tooltip="The image for the LLM model.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Custom("BBOX").Output(display_name="bbox"),
|
|
io.Image.Output(display_name="BBoxPreviewImage"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
client_info: ClientInfo,
|
|
items: str,
|
|
unload_model_after_chat: bool,
|
|
image: Tensor,
|
|
) -> io.NodeOutput:
|
|
"""
|
|
Generate bboxes based on LLM's grounding abitility.
|
|
"""
|
|
pil_image: Image.Image = tensor_to_pil(image)
|
|
|
|
(response, bbox_image) = grounding(
|
|
client_info=client_info,
|
|
item=items,
|
|
image=pil_image,
|
|
unload_after_chat=unload_model_after_chat,
|
|
)
|
|
# Convert bbox images to Tensor
|
|
result_image: Tensor = pil_to_tensor(bbox_image)
|
|
return io.NodeOutput(response, result_image)
|
|
|
|
|
|
class Caption(io.ComfyNode):
|
|
"""
|
|
Generate bboxes based on LLM's grounding abitility.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""
|
|
Tell the main program the metadata, input, output parameters of nodes.
|
|
"""
|
|
return io.Schema(
|
|
node_id="CaptionViaAPI",
|
|
display_name="Captioner",
|
|
category="LLM",
|
|
inputs=[
|
|
io.Custom("CLIENT_INFO").Input(
|
|
"client_info",
|
|
tooltip="The LLM client info.",
|
|
),
|
|
io.String.Input(
|
|
"language",
|
|
default="English",
|
|
tooltip="The language of the caption.",
|
|
),
|
|
io.Boolean.Input(
|
|
"unload_model_after_chat",
|
|
default=True,
|
|
tooltip=(
|
|
"Whether to unload the LLM model after the chat. "
|
|
+ "Only Ollama is supported."
|
|
),
|
|
),
|
|
io.Int.Input(
|
|
"num_max_sentences",
|
|
default=10,
|
|
min=1,
|
|
max=100,
|
|
step=1,
|
|
tooltip="The maximum number of sentences for the caption. Use 2 for CLIP model. Use 4 for FLUX. Use 10 and more for Qwen-Image.",
|
|
),
|
|
io.Image.Input(
|
|
"image",
|
|
tooltip="The image for the LLM model.",
|
|
),
|
|
io.Custom("EXTRA_PARAMETERS").Input(
|
|
"extra_parameters",
|
|
tooltip="The extra parameters for the LLM model.",
|
|
optional=True,
|
|
),
|
|
],
|
|
outputs=[
|
|
io.String.Output(display_name="Caption"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
client_info: ClientInfo,
|
|
language: str,
|
|
num_max_sentences: int,
|
|
unload_model_after_chat: bool,
|
|
image: Tensor,
|
|
extra_parameters: Optional[ExtraParameters] = None,
|
|
) -> io.NodeOutput:
|
|
"""
|
|
Generate bboxes based on LLM's grounding abitility.
|
|
"""
|
|
pil_image: Image.Image = tensor_to_pil(image)
|
|
|
|
payload = {
|
|
"client_info": client_info,
|
|
"language": language,
|
|
"image": pil_image,
|
|
"unload_after_chat": unload_model_after_chat,
|
|
"num_max_sentences": num_max_sentences,
|
|
}
|
|
|
|
if extra_parameters is not None:
|
|
payload["thinking"] = extra_parameters.thinking
|
|
payload["mode"] = extra_parameters.reasoning_effort
|
|
|
|
response = captioning(**payload)
|
|
|
|
return io.NodeOutput(response)
|
|
|
|
|
|
class ApplyChatTemplate(io.ComfyNode):
|
|
"""
|
|
Apply Chat Template for TextGenerate node. Requires use_default_template closed.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
"""
|
|
images: images
|
|
audio: audio
|
|
video: video
|
|
system_prompt: system's prompt
|
|
user_prompt: user's prompt
|
|
"""
|
|
return io.Schema(
|
|
node_id="ApplyChatTemplate",
|
|
display_name="Apply Chat Template",
|
|
category="LLM",
|
|
inputs=[
|
|
io.Image.Input("image", optional=True),
|
|
io.Image.Input(
|
|
"video", optional=True, tooltip="Video frames as image batch."
|
|
),
|
|
io.Audio.Input("audio", optional=True),
|
|
io.Custom("EXTRA_PARAMETERS").Input(
|
|
"extra_parameters",
|
|
tooltip="The extra parameters for the LLM model.",
|
|
optional=True,
|
|
),
|
|
io.String.Input(
|
|
"system_prompt", default="You are a helpful assistant.", multiline=True
|
|
),
|
|
io.String.Input("user_prompt", multiline=True),
|
|
io.String.Input(
|
|
"chat_template",
|
|
tooltip="Put chat template here. Can be found in original repo of model in huggingface.",
|
|
multiline=True
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output("image"),
|
|
io.Image.Output("video"),
|
|
io.Audio.Output("audio"),
|
|
io.String.Output(display_name="formatted_prompt"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
chat_template: str,
|
|
image: Optional[Tensor] = None,
|
|
video: Optional[Tensor | list[Tensor]] = None,
|
|
audio: Optional[dict | list[dict]] = None,
|
|
extra_parameters: Optional[ExtraParameters] = None,
|
|
) -> io.NodeOutput:
|
|
"""
|
|
Apply Chat Template for TextGenerate node. Requires use_default_template closed.
|
|
"""
|
|
# Check assets number
|
|
_num_images = (
|
|
len([image[i : i + 1] for i in range(image.shape[0])]) if image is not None else 0
|
|
)
|
|
_num_video = len(video) if isinstance(video, list) else 1 if video is not None else 0
|
|
_num_audio = len(audio) if isinstance(audio, list) else 1 if audio is not None else 0
|
|
|
|
# Using dummy data here will not affect the actual resource input.
|
|
_user_contents = (
|
|
[
|
|
{
|
|
"type": "image",
|
|
"image": f"https://dummy.site/image{i}.png",
|
|
}
|
|
for i in range(_num_images)
|
|
]
|
|
+ [
|
|
{
|
|
"type": "video",
|
|
"image": f"https://dummy.site/video{i}.mp4",
|
|
}
|
|
for i in range(_num_video)
|
|
]
|
|
+ [
|
|
{
|
|
"type": "audio",
|
|
"image": f"https://dummy.site/audio{i}.wav",
|
|
}
|
|
for i in range(_num_audio)
|
|
] +
|
|
[{"type": "text", "text": user_prompt}]
|
|
)
|
|
|
|
_base_payload = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": _user_contents},
|
|
]
|
|
|
|
# build up the content by template
|
|
env = Environment(
|
|
undefined=StrictUndefined,
|
|
trim_blocks=True,
|
|
lstrip_blocks=True,
|
|
)
|
|
|
|
template = env.from_string(chat_template)
|
|
|
|
extra = {}
|
|
if extra_parameters is not None:
|
|
extra["reasoning_effort"] = extra_parameters.reasoning_effort
|
|
extra["thinking"] = bool(extra_parameters.thinking == "enabled")
|
|
_formatted_prompt = template.render(messages=_base_payload, tools=None, add_generation_prompt=True, **extra)
|
|
|
|
return io.NodeOutput(image, video, audio, _formatted_prompt)
|
|
|
|
|
|
class PDIDLLMNodes(ComfyExtension):
|
|
"""
|
|
Define ComfyUI nodes.
|
|
"""
|
|
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]: # type: ignore[override]
|
|
"""
|
|
Get the list of nodes.
|
|
"""
|
|
return [LLMLoader, ExtraParametersComfy, ChatViaAPI, GenerateBBOX, Caption, ApplyChatTemplate]
|
|
|
|
|
|
async def comfy_entrypoint() -> ComfyExtension:
|
|
"""
|
|
Register ComfyUI nodes.
|
|
"""
|
|
return PDIDLLMNodes()
|