nodes/asaddi-ComfyUI-YALLM-node
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial ruff pass

Browse Source
This commit is contained in:
Allan Saddi
2024-11-05 02:19:01 -08:00
parent 8de242d993
commit f1d61967cc
2 changed files with 178 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+3 -3
View File
@@ -1,8 +1,8 @@
# Copyright (c) 2024 Allan Saddi <allan@saddi.com>
from .nodes import NODE_CLASS_MAPPINGS
NODE_DISPLAY_NAME_MAPPINGS = { k: v.TITLE for k,v in NODE_CLASS_MAPPINGS.items() }
NODE_DISPLAY_NAME_MAPPINGS = {k: v.TITLE for k, v in NODE_CLASS_MAPPINGS.items()}
WEB_DIRECTORY = './js'
WEB_DIRECTORY = "./js"
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
nodes.py
+175 -138
View File
@@ -26,11 +26,11 @@ class ProviderDefinition(BaseModel):
def resolve_base_url(self) -> str:
return ProviderDefinition._get_env_or_value(self.base_url)
def resolve_api_key(self) -> str|None:
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':
if not api_key or api_key == "none":
return None
return api_key
@@ -46,9 +46,9 @@ class ProviderDefinition(BaseModel):
Returns:
str: The retrieved value.
"""
PREFIX = 'os.environ/'
PREFIX = "os.environ/"
if value.startswith(PREFIX):
env_var = value[len(PREFIX):]
env_var = value[len(PREFIX) :]
value = os.environ[env_var]
return value
@@ -71,9 +71,9 @@ class DefinitionsConfig:
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 ]
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)
@@ -90,18 +90,18 @@ class DefinitionsConfig:
class Providers(DefinitionsConfig):
_root = 'providers'
_config = 'providers.yaml'
_config_default = 'providers.yaml.example'
_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'
_root = "models"
_config = "models.yaml"
_config_default = "models.yaml.example"
def load(self):
return super().load(ModelDefinition)
@@ -118,19 +118,28 @@ class LLMTextLatch:
@classmethod
def INPUT_TYPES(cls):
return {
'required': {
'text': ('STRING', {
'multiline': True,
}),
'replace': ('BOOLEAN', {
'default': True,
})
"required": {
"text": (
"STRING",
{
"multiline": True,
},
),
"replace": (
"BOOLEAN",
{
"default": True,
},
),
},
'optional': {
'text_input': ('STRING', {
'forceInput': True,
'multiline': True,
}),
"optional": {
"text_input": (
"STRING",
{
"forceInput": True,
"multiline": True,
},
),
},
}
@@ -139,49 +148,47 @@ class LLMTextLatch:
# Always valid, even if empty
return True
TITLE = 'Text Latch'
TITLE = "Text Latch"
RETURN_TYPES = ('STRING',)
RETURN_NAMES = ('text',)
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("text",)
FUNCTION = 'execute'
FUNCTION = "execute"
OUTPUT_NODE = True
CATEGORY = 'YALLM'
CATEGORY = "YALLM"
def execute(self, text, replace, text_input=None):
if text_input is not None and replace:
text = text_input
if text is None:
text = ''
text = ""
return { 'ui': { 'text': text }, 'result': (text,) }
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
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
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)
}
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',
name="temperature",
is_int=False,
min=0.0,
default=0.8,
@@ -190,20 +197,20 @@ class LLMTemperature:
@classmethod
def INPUT_TYPES(cls):
return {
'required': cls._DEF.to_input_type(),
'optional': {
'previous': ('LLMSAMPLER',),
"required": cls._DEF.to_input_type(),
"optional": {
"previous": ("LLMSAMPLER",),
},
}
TITLE = 'LLM Temperature'
TITLE = "LLM Temperature"
RETURN_TYPES = ('LLMSAMPLER',)
RETURN_NAMES = ('llm_sampler',)
RETURN_TYPES = ("LLMSAMPLER",)
RETURN_NAMES = ("llm_sampler",)
FUNCTION = 'execute'
FUNCTION = "execute"
CATEGORY = 'YALLM/samplers'
CATEGORY = "YALLM/samplers"
def execute(self, previous=None, **args):
value = args[self._DEF.name]
@@ -218,57 +225,59 @@ class LLMTemperature:
class LLMTopP(LLMTemperature):
_DEF = SamplerDefinition(
name='top_p',
name="top_p",
is_int=False,
min=0.0,
default=0.95,
max=1.0,
)
TITLE = 'LLM Top-P'
TITLE = "LLM Top-P"
class LLMTopK(LLMTemperature):
_DEF = SamplerDefinition(
name='top_k',
name="top_k",
is_int=True,
min=1,
default=40,
)
TITLE = 'LLM Top-K'
TITLE = "LLM Top-K"
class LLMMinP(LLMTemperature):
_DEF = SamplerDefinition(
name='min_p',
name="min_p",
is_int=False,
min=0.0,
default=0.05,
max=1.0,
)
TITLE = 'LLM Min-P'
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]
ChatMessage = dict[str, str]
ChatHistory = list[ChatMessage]
SamplerSetting = tuple[str,Any]
SamplerSetting = tuple[str, Any]
class LLMModel:
def __init__(self, prov_def: ProviderDefinition, model: str|None=None):
def __init__(self, prov_def: ProviderDefinition, model: str | None = None):
self._llm = openai.OpenAI(
base_url=prov_def.resolve_base_url(),
api_key=(prov_def.resolve_api_key() or 'none'), # openai package requires something for API key
api_key=(
prov_def.resolve_api_key() or "none"
), # openai package requires something for API key
)
if model is None:
self._model = getattr(prov_def, 'model', None)
self._model = getattr(prov_def, "model", None)
else:
self._model = model
@@ -276,32 +285,39 @@ class LLMModel:
models = list(self._llm.models.list())
return [m.id for m in models]
def chat_completion(self, messages: ChatHistory, image: Image|None=None, samplers: list[SamplerSetting]|None=None, seed: int|None=None) -> str:
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 msg.get("role") == "user":
parts = msg.get("content", [])
if isinstance(parts, str):
parts = [{'type': 'text', 'text': parts}]
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')
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}'
}
})
parts.insert(
0,
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{png_b64}"},
},
)
msg['content'] = parts
msg["content"] = parts
break
# print(messages)
@@ -312,26 +328,28 @@ class LLMModel:
sampler_order = []
extra_args = {}
extra_body = {}
for k,v in samplers:
for k, v in samplers:
sampler_order.append(k)
if k in ('temperature', 'top_p'): # The only ones supported directly by openai package
if k in (
"temperature",
"top_p",
): # The only ones supported directly by openai package
extra_args[k] = v
else:
# The rest have to go into "extra_body"
extra_body[k] = v
extra_body['samplers'] = sampler_order # Hopefully won't cause an issue for non-llama.cpp endpoints
extra_body["samplers"] = (
sampler_order # Hopefully won't cause an issue for non-llama.cpp endpoints
)
if seed is not None:
extra_args['seed'] = seed
extra_args["seed"] = seed
# print(f'extra_args = {repr(extra_args)}')
# print(f'extra_body = {repr(extra_body)}')
output = self._llm.chat.completions.create(
messages=messages,
model=self._model,
extra_body=extra_body,
**extra_args
messages=messages, model=self._model, extra_body=extra_body, **extra_args
)
# TODO Since we're not streaming, we'll need some sort of timeout. How to specify that?
@@ -345,9 +363,9 @@ class LLMProvider:
PROVIDERS.refresh()
return {
'required': {
'provider': (PROVIDERS.CHOICES,),
'model': (['fetch.models.first'],)
"required": {
"provider": (PROVIDERS.CHOICES,),
"model": (["fetch.models.first"],),
},
}
@@ -358,20 +376,20 @@ class LLMProvider:
@classmethod
def VALIDATE_INPUTS(cls, model):
if model in ('fetch.models.first', 'failed.to.fetch'):
return 'fetch models and make a selection first'
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)'
TITLE = "LLM Provider (API)"
RETURN_TYPES = ('LLMMODEL',)
RETURN_NAMES = ('llm_model',)
RETURN_TYPES = ("LLMMODEL",)
RETURN_NAMES = ("llm_model",)
FUNCTION = 'execute'
FUNCTION = "execute"
CATEGORY = 'YALLM'
CATEGORY = "YALLM"
def execute(self, provider, model):
llm = LLMModel(PROVIDERS.BY_NAME[provider], model=model)
@@ -384,8 +402,8 @@ class LLMModelNode:
def INPUT_TYPES(cls):
MODELS.refresh()
return {
'required': {
'model': (MODELS.CHOICES,),
"required": {
"model": (MODELS.CHOICES,),
},
}
@@ -394,14 +412,14 @@ class LLMModelNode:
MODELS.refresh()
return MODELS._mtime
TITLE = 'LLM Model (API)'
TITLE = "LLM Model (API)"
RETURN_TYPES = ('LLMMODEL',)
RETURN_NAMES = ('llm_model',)
RETURN_TYPES = ("LLMMODEL",)
RETURN_NAMES = ("llm_model",)
FUNCTION = 'execute'
FUNCTION = "execute"
CATEGORY = 'YALLM'
CATEGORY = "YALLM"
def execute(self, model):
llm = LLMModel(MODELS.BY_NAME[model])
@@ -413,50 +431,67 @@ 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,
}),
"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',),
"optional": {
"llm_sampler": ("LLMSAMPLER",),
"system_prompt": (
"STRING",
{
"multiline": True,
},
),
"image": ("IMAGE",),
},
}
TITLE = 'LLM Chat'
TITLE = "LLM Chat"
RETURN_TYPES = ('STRING',)
RETURN_NAMES = ('completion',)
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("completion",)
OUTPUT_IS_LIST = (True,)
FUNCTION = 'execute'
FUNCTION = "execute"
CATEGORY = 'YALLM'
CATEGORY = "YALLM"
def execute(self, llm_model: LLMModel, seed, user_prompt, llm_sampler=None, system_prompt=None, image=None):
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})
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]
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
@@ -464,7 +499,9 @@ class LLMChat:
# 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)
text = llm_model.chat_completion(
messages, image=pil_image, samplers=llm_sampler, seed=seed
)
result.append(text)
return (result,)
@@ -472,7 +509,7 @@ class LLMChat:
@PromptServer.instance.routes.get("/llm_models")
async def llm_models(request: Request):
name = request.rel_url.query.get('name', None)
name = request.rel_url.query.get("name", None)
if name:
PROVIDERS.refresh()
if (prov_def := PROVIDERS.BY_NAME.get(name, None)) is not None:
@@ -481,20 +518,20 @@ async def llm_models(request: Request):
models = llm.get_models()
return web.json_response(models)
except Exception as e:
print(f'Problem fetching models: {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'])
return web.json_response(["failed.to.fetch"])
NODE_CLASS_MAPPINGS = {
'LLMTextLatch': LLMTextLatch,
'LLMProvider': LLMProvider,
'LLMModel': LLMModelNode,
'LLMChat': LLMChat,
'LLMTemperature': LLMTemperature,
'LLMTopP': LLMTopP,
'LLMTopK': LLMTopK,
'LLMMinP': LLMMinP,
"LLMTextLatch": LLMTextLatch,
"LLMProvider": LLMProvider,
"LLMModel": LLMModelNode,
"LLMChat": LLMChat,
"LLMTemperature": LLMTemperature,
"LLMTopP": LLMTopP,
"LLMTopK": LLMTopK,
"LLMMinP": LLMMinP,
}