Allan Saddi
85423e6600
From experience elsewhere, there seem to be some providers that default to streaming. Also update dependencies in pyproject.toml to match requirements.txt.
616 lines
16 KiB
Python
616 lines
16 KiB
Python
# Copyright (c) 2024 Allan Saddi <allan@saddi.com>
|
|
import base64
|
|
from io import BytesIO
|
|
import os
|
|
from typing import Any, Optional
|
|
|
|
from aiohttp import web
|
|
from aiohttp.web_request import Request
|
|
from PIL.Image import Image
|
|
from pydantic import BaseModel
|
|
import requests
|
|
import torchvision.transforms.functional as F
|
|
import yaml
|
|
|
|
from .utils import PromptUtils, WorkflowUtils
|
|
|
|
from server import PromptServer
|
|
|
|
|
|
BASE_PATH = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
|
|
class ProviderDefinition(BaseModel):
|
|
name: str
|
|
base_url: str
|
|
api_key: Optional[str]
|
|
|
|
def resolve_base_url(self) -> str:
|
|
return ProviderDefinition._get_env_or_value(self.base_url)
|
|
|
|
def resolve_api_key(self) -> str | None:
|
|
api_key = self.api_key
|
|
if api_key:
|
|
api_key = ProviderDefinition._get_env_or_value(api_key)
|
|
if not api_key or api_key == "none":
|
|
return None
|
|
return api_key
|
|
|
|
@staticmethod
|
|
def _get_env_or_value(value: str) -> str:
|
|
"""
|
|
Retrieves the value from either an environment variable or a provided string.
|
|
|
|
Args:
|
|
value (str): The value to retrieve. It can be a string literal or a string
|
|
starting with 'os.environ/' followed by an environment variable name.
|
|
|
|
Returns:
|
|
str: The retrieved value.
|
|
"""
|
|
PREFIX = "os.environ/"
|
|
if value.startswith(PREFIX):
|
|
env_var = value[len(PREFIX) :]
|
|
value = os.environ[env_var]
|
|
return value
|
|
|
|
|
|
class ModelDefinition(ProviderDefinition):
|
|
model: str
|
|
|
|
|
|
class DefinitionsConfig:
|
|
_mtime = None
|
|
|
|
def load(self, def_type) -> None:
|
|
models_file = self._get_models_file()
|
|
|
|
with open(models_file) as inp:
|
|
d = yaml.load(inp, yaml.Loader)
|
|
self._mtime = (models_file, os.path.getmtime(models_file))
|
|
|
|
self.LIST = []
|
|
for value in d[self._root]:
|
|
self.LIST.append(def_type.model_validate(value))
|
|
if not self.LIST:
|
|
raise RuntimeError(f"{models_file}: Need at least one definition")
|
|
self.BY_NAME = {d.name: d for d in self.LIST}
|
|
self.CHOICES = [d.name for d in self.LIST]
|
|
|
|
def _get_models_file(self):
|
|
models_file = os.path.join(BASE_PATH, self._config)
|
|
if not os.path.exists(models_file):
|
|
# Just read from the example for now
|
|
# print(f'{models_file} does not exist; using default')
|
|
models_file = os.path.join(BASE_PATH, self._config_default)
|
|
return models_file
|
|
|
|
def refresh(self):
|
|
models_file = self._get_models_file()
|
|
if self._mtime != (models_file, os.path.getmtime(models_file)):
|
|
self.load()
|
|
|
|
|
|
class Providers(DefinitionsConfig):
|
|
_root = "providers"
|
|
_config = "providers.yaml"
|
|
_config_default = "providers.yaml.example"
|
|
|
|
def load(self):
|
|
return super().load(ProviderDefinition)
|
|
|
|
|
|
class Models(DefinitionsConfig):
|
|
_root = "models"
|
|
_config = "models.yaml"
|
|
_config_default = "models.yaml.example"
|
|
|
|
def load(self):
|
|
return super().load(ModelDefinition)
|
|
|
|
|
|
PROVIDERS = Providers()
|
|
PROVIDERS.load()
|
|
|
|
MODELS = Models()
|
|
MODELS.load()
|
|
|
|
|
|
class LLMPrependAppend:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text_input": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"multiline": True,
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"prefix": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
},
|
|
),
|
|
"suffix": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
TITLE = "Prepend/Append"
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "YALLM"
|
|
|
|
def execute(self, text_input, prefix=None, suffix=None):
|
|
if prefix is None:
|
|
prefix = ""
|
|
if suffix is None:
|
|
suffix = ""
|
|
|
|
text = prefix + text_input + suffix
|
|
|
|
return (text,)
|
|
|
|
|
|
class LLMTextLatch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
},
|
|
),
|
|
"replace": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"text_input": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"multiline": True,
|
|
"lazy": True,
|
|
},
|
|
),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, text):
|
|
# Always valid, even if empty
|
|
# TODO I don't think this is necessary
|
|
return True
|
|
|
|
TITLE = "Text Latch"
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "YALLM"
|
|
|
|
def check_lazy_status(
|
|
self, text, replace, unique_id, prompt, extra_pnginfo, text_input=None
|
|
):
|
|
# Require whatever is connected to text_input to be evaluated
|
|
# contingent on replace.
|
|
if replace and text_input is None:
|
|
# Do we have something connected to text_input?
|
|
pu = PromptUtils(prompt)
|
|
if pu.is_input_connected(unique_id, "text_input"):
|
|
return ["text_input"]
|
|
# All good, nothing else needs to be evaluated
|
|
return []
|
|
|
|
def execute(self, text, replace, unique_id, prompt, extra_pnginfo, text_input=None):
|
|
if text_input is not None and replace:
|
|
text = text_input
|
|
|
|
if text is None:
|
|
text = ""
|
|
|
|
# Save the current value in the outgoing metadata
|
|
wfu = WorkflowUtils(extra_pnginfo)
|
|
wfu.set_widget(unique_id, 0, text)
|
|
|
|
# An oddity: the values of the UI dictionary have to be wrapped in
|
|
# a tuple/list, just like the result?
|
|
return {"ui": {"text": (text,)}, "result": (text,)}
|
|
|
|
|
|
class SamplerDefinition(BaseModel):
|
|
name: str
|
|
is_int: bool
|
|
min: Optional[float | int] = None
|
|
default: float | int
|
|
max: Optional[float | int] = None
|
|
|
|
def to_input_type(self):
|
|
constraints = {}
|
|
if self.min is not None:
|
|
constraints["min"] = self.min
|
|
if self.max is not None:
|
|
constraints["max"] = self.max
|
|
constraints["default"] = self.default
|
|
return {self.name: ("INT" if self.is_int else "FLOAT", constraints)}
|
|
|
|
|
|
class LLMTemperature:
|
|
_DEF = SamplerDefinition(
|
|
name="temperature",
|
|
is_int=False,
|
|
min=0.0,
|
|
default=0.8,
|
|
)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": cls._DEF.to_input_type(),
|
|
"optional": {
|
|
"previous": ("LLMSAMPLER",),
|
|
},
|
|
}
|
|
|
|
TITLE = "LLM Temperature"
|
|
|
|
RETURN_TYPES = ("LLMSAMPLER",)
|
|
RETURN_NAMES = ("llm_sampler",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "YALLM/samplers"
|
|
|
|
def execute(self, previous=None, **args):
|
|
value = args[self._DEF.name]
|
|
|
|
if previous is None:
|
|
previous = []
|
|
|
|
previous.append((self._DEF.name, value))
|
|
|
|
return (previous,)
|
|
|
|
|
|
class LLMTopP(LLMTemperature):
|
|
_DEF = SamplerDefinition(
|
|
name="top_p",
|
|
is_int=False,
|
|
min=0.0,
|
|
default=0.95,
|
|
max=1.0,
|
|
)
|
|
|
|
TITLE = "LLM Top-P"
|
|
|
|
|
|
class LLMTopK(LLMTemperature):
|
|
_DEF = SamplerDefinition(
|
|
name="top_k",
|
|
is_int=True,
|
|
min=1,
|
|
default=40,
|
|
)
|
|
|
|
TITLE = "LLM Top-K"
|
|
|
|
|
|
class LLMMinP(LLMTemperature):
|
|
_DEF = SamplerDefinition(
|
|
name="min_p",
|
|
is_int=False,
|
|
min=0.0,
|
|
default=0.05,
|
|
max=1.0,
|
|
)
|
|
|
|
TITLE = "LLM Min-P"
|
|
|
|
|
|
# TODO typical_p? tfs_z? Any others? mirostat?
|
|
|
|
|
|
# TODO Will we need a more complex definition? (e.g. multi-modal)
|
|
ChatMessage = dict[str, str]
|
|
ChatHistory = list[ChatMessage]
|
|
|
|
SamplerSetting = tuple[str, Any]
|
|
|
|
|
|
class LLMModel:
|
|
def __init__(self, prov_def: ProviderDefinition, model: str | None = None):
|
|
self._base_url = prov_def.resolve_base_url()
|
|
self._headers = {}
|
|
if (api_key := prov_def.resolve_api_key()) and api_key != "none":
|
|
self._headers["authorization"] = f"Bearer {api_key}"
|
|
if model is None:
|
|
self._model = getattr(prov_def, "model", None)
|
|
else:
|
|
self._model = model
|
|
|
|
def get_models(self):
|
|
with requests.get(self._base_url + "/models", headers=self._headers) as r:
|
|
models_json = r.json()
|
|
return [m["id"] for m in models_json["data"]]
|
|
|
|
def chat_completion(
|
|
self,
|
|
messages: ChatHistory,
|
|
image: Image | None = None,
|
|
samplers: list[SamplerSetting] | None = None,
|
|
seed: int | None = None,
|
|
) -> str:
|
|
assert self._model is not None
|
|
|
|
if image is not None:
|
|
# Work backwards through history
|
|
for msg in reversed(messages):
|
|
# And grab the very last user message we see
|
|
if msg.get("role") == "user":
|
|
parts = msg.get("content", [])
|
|
if isinstance(parts, str):
|
|
parts = [{"type": "text", "text": parts}]
|
|
|
|
# Convert image to data URI format
|
|
png_io = BytesIO()
|
|
image.save(png_io, "PNG")
|
|
png_b64 = base64.b64encode(png_io.getbuffer()).decode("ascii")
|
|
|
|
# And insert image just before the user prompt
|
|
parts.insert(
|
|
0,
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:image/png;base64,{png_b64}"},
|
|
},
|
|
)
|
|
|
|
msg["content"] = parts
|
|
break
|
|
|
|
# print(messages)
|
|
|
|
if samplers is None:
|
|
samplers = []
|
|
|
|
sampler_order = []
|
|
extra_args = {}
|
|
for k, v in samplers:
|
|
sampler_order.append(k)
|
|
extra_args[k] = v
|
|
extra_args["samplers"] = (
|
|
sampler_order # Hopefully won't cause an issue for non-llama.cpp endpoints
|
|
)
|
|
|
|
if seed is not None:
|
|
extra_args["seed"] = seed
|
|
|
|
# print(f"extra_args = {repr(extra_args)}")
|
|
|
|
req_json = {
|
|
"model": self._model,
|
|
"messages": messages,
|
|
"stream": False,
|
|
}
|
|
req_json.update(extra_args)
|
|
|
|
# TODO Since we're not streaming, we'll need some sort of timeout. TODO specify that
|
|
with requests.post(
|
|
self._base_url + "/chat/completions", json=req_json, headers=self._headers
|
|
) as r:
|
|
output = r.json()
|
|
|
|
if choices := output.get("choices", []):
|
|
return choices[0]["message"]["content"]
|
|
else:
|
|
return ""
|
|
|
|
|
|
class LLMProvider:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
PROVIDERS.refresh()
|
|
|
|
return {
|
|
"required": {
|
|
"provider": (PROVIDERS.CHOICES,),
|
|
"model": (["fetch.models.first"],),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, provider, model):
|
|
PROVIDERS.refresh()
|
|
return PROVIDERS._mtime
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, model):
|
|
if model in ("fetch.models.first", "failed.to.fetch"):
|
|
return "fetch models and make a selection first"
|
|
|
|
# Ass-u-me the frontend validated the choices
|
|
return True
|
|
|
|
TITLE = "LLM Provider (API)"
|
|
|
|
RETURN_TYPES = ("LLMMODEL",)
|
|
RETURN_NAMES = ("llm_model",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "YALLM"
|
|
|
|
def execute(self, provider, model):
|
|
llm = LLMModel(PROVIDERS.BY_NAME[provider], model=model)
|
|
|
|
return (llm,)
|
|
|
|
|
|
class LLMModelNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
MODELS.refresh()
|
|
return {
|
|
"required": {
|
|
"model": (MODELS.CHOICES,),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, model):
|
|
MODELS.refresh()
|
|
return MODELS._mtime
|
|
|
|
TITLE = "LLM Model (API)"
|
|
|
|
RETURN_TYPES = ("LLMMODEL",)
|
|
RETURN_NAMES = ("llm_model",)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "YALLM"
|
|
|
|
def execute(self, model):
|
|
llm = LLMModel(MODELS.BY_NAME[model])
|
|
|
|
return (llm,)
|
|
|
|
|
|
class LLMChat:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"llm_model": ("LLMMODEL",),
|
|
"seed": (
|
|
"INT",
|
|
{ # Thankfully, frontend has special handling for a widget of this name...
|
|
# min/max/default from the KSampler nodes
|
|
# llama.cpp is only 32-bit (with truncation, no doubt), OpenAI API spec doesn't specify a size
|
|
"min": 0,
|
|
"default": 0,
|
|
"max": 0xFFFFFFFF_FFFFFFFF,
|
|
},
|
|
),
|
|
"user_prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"llm_sampler": ("LLMSAMPLER",),
|
|
"system_prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
},
|
|
),
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
TITLE = "LLM Chat"
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("completion",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "YALLM"
|
|
|
|
def execute(
|
|
self,
|
|
llm_model: LLMModel,
|
|
seed,
|
|
user_prompt,
|
|
llm_sampler=None,
|
|
system_prompt=None,
|
|
image=None,
|
|
):
|
|
messages = []
|
|
if system_prompt:
|
|
messages.append({"role": "system", "content": system_prompt})
|
|
messages.append({"role": "user", "content": user_prompt})
|
|
|
|
result = []
|
|
if image is None:
|
|
text = llm_model.chat_completion(messages, samplers=llm_sampler, seed=seed)
|
|
result.append(text)
|
|
else:
|
|
image = image.permute(0, 3, 1, 2) # Convert to [B,C,H,W]
|
|
for img in image:
|
|
pil_image = F.to_pil_image(img)
|
|
# Fortunately or unfortunately, I'm forced to have this conform
|
|
# to the conventions of my LlamaVision node.
|
|
# I think it would be cleaner to do the conversion here, but
|
|
# I have no idea how that would work for a local
|
|
# Transformers model.
|
|
text = llm_model.chat_completion(
|
|
messages, image=pil_image, samplers=llm_sampler, seed=seed
|
|
)
|
|
result.append(text)
|
|
|
|
return (result,)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/llm_models")
|
|
async def llm_models(request: Request):
|
|
name = request.rel_url.query.get("name", None)
|
|
if name:
|
|
PROVIDERS.refresh()
|
|
if (prov_def := PROVIDERS.BY_NAME.get(name, None)) is not None:
|
|
try:
|
|
llm = LLMModel(prov_def)
|
|
models = llm.get_models()
|
|
return web.json_response(models)
|
|
except Exception as e:
|
|
print(f"Problem fetching models: {e}")
|
|
# Just eat it and fall through.
|
|
|
|
# Fail, but don't error out
|
|
return web.json_response(["failed.to.fetch"])
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LLMPrependAppend": LLMPrependAppend,
|
|
"LLMTextLatch": LLMTextLatch,
|
|
"LLMProvider": LLMProvider,
|
|
"LLMModel": LLMModelNode,
|
|
"LLMChat": LLMChat,
|
|
"LLMTemperature": LLMTemperature,
|
|
"LLMTopP": LLMTopP,
|
|
"LLMTopK": LLMTopK,
|
|
"LLMMinP": LLMMinP,
|
|
}
|