From f1d61967cc41585dc61ba4da5c18e5cceb4723a2 Mon Sep 17 00:00:00 2001 From: Allan Saddi Date: Tue, 5 Nov 2024 02:19:01 -0800 Subject: [PATCH] Initial ruff pass --- __init__.py | 6 +- nodes.py | 313 +++++++++++++++++++++++++++++----------------------- 2 files changed, 178 insertions(+), 141 deletions(-) diff --git a/__init__.py b/__init__.py index 1319720..3812780 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,8 @@ # Copyright (c) 2024 Allan Saddi 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"] diff --git a/nodes.py b/nodes.py index 2ecfd2a..078ae09 100644 --- a/nodes.py +++ b/nodes.py @@ -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, }