Bump version: 1.4.0 → 1.5.0

From experience elsewhere, there seem to be some providers that default
to streaming.

Also update dependencies in pyproject.toml to match requirements.txt.

.github/workflows/publish.yml

@@ -8,17 +8,19 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
-    # if this is a forked repository. Skipping the workflow.
-    if: github.event.repository.fork == false 
+    if: ${{ github.repository_owner == 'asaddi' }}
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
           ## Add your own personal access token to your Github Repository secrets and reference it here.
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

README.md

@@ -1,3 +1,5 @@
+
+
 # Yet Another LLM Node (for ComfyUI)
 
 Yet another set of LLM nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI). This one only supports OpenAI-like APIs, but of course can be used with local LLM providers such as [llama.cpp](https://github.com/ggerganov/llama.cpp) (and ollama, LM Studio, among many others).
@@ -69,6 +71,10 @@ The input is optional, meaning you can mute the node leading into it or even dis
 
 The node will output whatever is currently being displayed in the textbox. The node text will also be saved in exported workflows and generated images.
 
+## "Prepend/Append" Node
+
+A simple utility node that takes a text input and optionally adds a prefix and/or suffix to it. This is useful for dynamically injecting instructions or context into prompts without altering the source text, as demonstrated in the `prompt-optimizer` example workflow.
+
 ## My Related Projects
 
 * https://github.com/asaddi/YALLM-LlamaVision Basic ComfyUI node for Llama 3.2 Vision. Sampler/model nodes are cross-compatible with this project.

__init__.py

@@ -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"]

examples/README.md

@@ -16,4 +16,4 @@ The YALLM nodes query a (vision) LLM, asking for a description of the loaded ima
 
 ## prompt-optimizer
 
-Asks the LLM to optimize a prompt by word count while keeping descriptive details. Uses the "String Function" node from <https://github.com/pythongosssss/ComfyUI-Custom-Scripts/> in order to prepend some text to the original prompt in order to use it as a system prompt.
+Asks the LLM to optimize a prompt by word count while keeping descriptive details. Uses the "Prepend/Append" node in order to prepend some text to the original prompt in order to use it as a system prompt.

examples/prompt-optimizer.json

@@ -1,54 +1,7 @@
 {
-  "last_node_id": 5,
-  "last_link_id": 4,
+  "last_node_id": 6,
+  "last_link_id": 7,
   "nodes": [
-    {
-      "id": 4,
-      "type": "StringFunction|pysssss",
-      "pos": {
-        "0": 0,
-        "1": 350
-      },
-      "size": {
-        "0": 400,
-        "1": 244
-      },
-      "flags": {},
-      "order": 2,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "text_b",
-          "type": "STRING",
-          "link": 2,
-          "widget": {
-            "name": "text_b"
-          },
-          "shape": 7
-        }
-      ],
-      "outputs": [
-        {
-          "name": "STRING",
-          "type": "STRING",
-          "links": [
-            3
-          ],
-          "slot_index": 0
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "StringFunction|pysssss"
-      },
-      "widgets_values": [
-        "append",
-        "no",
-        "Use the following text as reference:\n",
-        "",
-        "",
-        "Use the following text as reference:\n\"A beautiful, serene landscape unfolds before us, with rolling hills and lush greenery stretching as far as the eye can see. A delicate, antique glass bottle lies abandoned on a rocky outcropping, its transparent walls glowing with an ethereal light that seems almost otherworldly. The bottle appears to be filled with a deep, rich purple liquid that perfectly captures the color of a distant galaxy, with swirling clouds of gas and dust visible within. The surrounding landscape is bathed in a soft, golden light, with the warm rays of the sun casting a gentle glow across the scene. A few wispy clouds drift lazily across the sky, adding a sense of depth and atmosphere to the already breathtaking view. The overall mood is one of peaceful contemplation, as if the viewer has stumbled upon a secret, hidden world that exists just beyond the edge of reality.\""
-      ]
-    },
     {
       "id": 5,
       "type": "LLMTextLatch",
@@ -90,50 +43,6 @@
         ""
       ]
     },
-    {
-      "id": 2,
-      "type": "LLMTextLatch",
-      "pos": {
-        "0": -450,
-        "1": 350
-      },
-      "size": {
-        "0": 400,
-        "1": 200
-      },
-      "flags": {},
-      "order": 0,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "text_input",
-          "type": "STRING",
-          "link": null,
-          "widget": {
-            "name": "text_input"
-          },
-          "shape": 7
-        }
-      ],
-      "outputs": [
-        {
-          "name": "text",
-          "type": "STRING",
-          "links": [
-            2
-          ],
-          "slot_index": 0
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "LLMTextLatch"
-      },
-      "widgets_values": [
-        "\"A beautiful, serene landscape unfolds before us, with rolling hills and lush greenery stretching as far as the eye can see. A delicate, antique glass bottle lies abandoned on a rocky outcropping, its transparent walls glowing with an ethereal light that seems almost otherworldly. The bottle appears to be filled with a deep, rich purple liquid that perfectly captures the color of a distant galaxy, with swirling clouds of gas and dust visible within. The surrounding landscape is bathed in a soft, golden light, with the warm rays of the sun casting a gentle glow across the scene. A few wispy clouds drift lazily across the sky, adding a sense of depth and atmosphere to the already breathtaking view. The overall mood is one of peaceful contemplation, as if the viewer has stumbled upon a secret, hidden world that exists just beyond the edge of reality.\"",
-        true,
-        ""
-      ]
-    },
     {
       "id": 1,
       "type": "LLMProvider",
@@ -146,7 +55,7 @@
         "1": 106
       },
       "flags": {},
-      "order": 1,
+      "order": 0,
       "mode": 0,
       "inputs": [],
       "outputs": [
@@ -203,7 +112,7 @@
         {
           "name": "system_prompt",
           "type": "STRING",
-          "link": 3,
+          "link": 7,
           "widget": {
             "name": "system_prompt"
           },
@@ -230,6 +139,93 @@
         "Optimize the given text-to-image prompt for word count. Try to keep as many descriptive details as possible.\n\nOUTPUT ONLY THE MODIFIED PROMPT.",
         ""
       ]
+    },
+    {
+      "id": 2,
+      "type": "LLMTextLatch",
+      "pos": {
+        "0": -450,
+        "1": 350
+      },
+      "size": {
+        "0": 400,
+        "1": 200
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "text_input",
+          "type": "STRING",
+          "link": null,
+          "widget": {
+            "name": "text_input"
+          },
+          "shape": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "text",
+          "type": "STRING",
+          "links": [
+            6
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LLMTextLatch"
+      },
+      "widgets_values": [
+        "\"A beautiful, serene landscape unfolds before us, with rolling hills and lush greenery stretching as far as the eye can see. A delicate, antique glass bottle lies abandoned on a rocky outcropping, its transparent walls glowing with an ethereal light that seems almost otherworldly. The bottle appears to be filled with a deep, rich purple liquid that perfectly captures the color of a distant galaxy, with swirling clouds of gas and dust visible within. The surrounding landscape is bathed in a soft, golden light, with the warm rays of the sun casting a gentle glow across the scene. A few wispy clouds drift lazily across the sky, adding a sense of depth and atmosphere to the already breathtaking view. The overall mood is one of peaceful contemplation, as if the viewer has stumbled upon a secret, hidden world that exists just beyond the edge of reality.\"",
+        true,
+        ""
+      ]
+    },
+    {
+      "id": 6,
+      "type": "LLMPrependAppend",
+      "pos": {
+        "0": 0,
+        "1": 350
+      },
+      "size": [
+        400,
+        200
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "text_input",
+          "type": "STRING",
+          "link": 6,
+          "widget": {
+            "name": "text_input"
+          }
+        }
+      ],
+      "outputs": [
+        {
+          "name": "text",
+          "type": "STRING",
+          "links": [
+            7
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LLMPrependAppend"
+      },
+      "widgets_values": [
+        "",
+        "Use the following text as reference:\n",
+        ""
+      ]
     }
   ],
   "links": [
@@ -241,22 +237,6 @@
       0,
       "LLMMODEL"
     ],
-    [
-      2,
-      2,
-      0,
-      4,
-      0,
-      "STRING"
-    ],
-    [
-      3,
-      4,
-      0,
-      3,
-      3,
-      "STRING"
-    ],
     [
       4,
       3,
@@ -264,16 +244,32 @@
       5,
       0,
       "STRING"
+    ],
+    [
+      6,
+      2,
+      0,
+      6,
+      0,
+      "STRING"
+    ],
+    [
+      7,
+      6,
+      0,
+      3,
+      3,
+      "STRING"
     ]
   ],
   "groups": [],
   "config": {},
   "extra": {
     "ds": {
-      "scale": 0.9090909090909091,
+      "scale": 0.7513148009015777,
       "offset": [
-        111.69999999999979,
-        88.6999999999999
+        320.79927272727275,
+        220.6500000000004
       ]
     }
   },

js/yallm-node.js

@@ -2,65 +2,85 @@ import { app } from "../../scripts/app.js";
 import { api } from "../../scripts/api.js";
 app.registerExtension({
 	name: "YALLM.node",
-    async beforeRegisterNodeDef(nodeType, nodeData, app) {
-        if (nodeType.comfyClass==="LLMProvider") {
-            nodeType.prototype.myRefreshModels = function (node, name) {
-                if (!name) name = node.widgets[0].value;
+	async beforeRegisterNodeDef(nodeType, nodeData, app) {
+		if (nodeType.comfyClass === "LLMProvider") {
+			nodeType.prototype.myRefreshModels = function (name) {
+				const provWidget = this.widgets.find(
+					(widget) => widget.name === "provider",
+				);
+				const modelWidget = this.widgets.find(
+					(widget) => widget.name === "model",
+				);
+				const fetchWidget = this.widgets.find(
+					(widget) => widget.name === "fetch models",
+				);
 
-                const myDisableWidgets = (state) => {
-                    node.widgets[0].disabled = state;
-                    node.widgets[1].disabled = state;
-                    // Note: This is the "fetch models" button we add down below
-                    node.widgets[2].disabled = state;
-                };
-                myDisableWidgets(true);
+				const provName = name ?? provWidget.value;
 
-                api.fetchApi(`/llm_models?name=${name}`).then((resp) => {
-                    resp.json().then((data) => {
-                        let widget = node.widgets[1];
-                        widget.options.values = data;
+				const myDisableWidgets = (state) => {
+					provWidget.disabled = state;
+					modelWidget.disabled = state;
+					// Note: This is the "fetch models" button we add down below
+					fetchWidget.disabled = state;
+				};
+				myDisableWidgets(true);
 
-                        if (!widget.options.values.includes(widget.value) && widget.options.values.length > 0) {
-                            widget.value = widget.options.values[0];
-                        }
+				api.fetchApi(`/llm_models?name=${provName}`).then((resp) => {
+					resp.json().then((data) => {
+						const widget = modelWidget;
+						widget.options.values = data;
 
-                        myDisableWidgets(false);
-                    })
-                })
-            }
-        }
-    },
-    async nodeCreated(node) {
-        if (node?.comfyClass==="LLMProvider") {
-            let prov_widget = node.widgets[0];
-            // Whenever provider changes value, fetch models
-            const original_callback = prov_widget.callback;
-            prov_widget.callback = function () {
-                let name = arguments?.[0];
-                node.myRefreshModels(node, name);
-                return original_callback?.apply(this, arguments);
-            }
+						if (
+							!widget.options.values.includes(widget.value) &&
+							widget.options.values.length > 0
+						) {
+							widget.value = widget.options.values[0];
+						}
 
-            // Disable model combo widget until we fetch models
-            node.widgets[1].disabled = true;
+						myDisableWidgets(false);
+					});
+				});
+			};
+		}
+	},
+	async nodeCreated(node) {
+		if (node?.comfyClass === "LLMProvider") {
+			const provWidget = node.widgets.find(
+				(widget) => widget.name === "provider",
+			);
+			const modelWidget = node.widgets.find(
+				(widget) => widget.name === "model",
+			);
 
-            // Add button to fetch models manually
-            const btn = node.addWidget(
-                "button",
-                "fetch models",
-                "models",
-                () => { node.myRefreshModels(node) }
-            )
-            btn.serializeValue = () => void 0;
-        }
-    },
-    async setup() {
-        api.addEventListener("executed", function (event) {
-            var node = app.graph.getNodeById(event.detail.node);
-            if (node?.comfyClass==="LLMTextLatch") {
-                node.widgets[0].value = event.detail.output.text.join("");
-                app.graph.setDirtyCanvas(true, false);
-            }
-        });
-    }
-})
+			// Whenever provider changes value, fetch models
+			const original_callback = provWidget.callback;
+			provWidget.callback = function (...args) {
+				const name = args?.[0];
+				node.myRefreshModels(name);
+				return original_callback?.apply(this, args);
+			};
+
+			// Disable model combo widget until we fetch models
+			modelWidget.disabled = true;
+
+			// Add button to fetch models manually
+			const btn = node.addWidget("button", "fetch models", "models", () => {
+				node.myRefreshModels();
+			});
+			btn.serializeValue = () => void 0;
+		}
+	},
+	async setup() {
+		api.addEventListener("executed", (event) => {
+			const node = app.graph.getNodeById(event.detail.node);
+			if (node?.comfyClass === "LLMTextLatch") {
+				const textWidget = node.widgets.find(
+					(widget) => widget.name === "text",
+				);
+				const texts = event.detail.output.text;
+				textWidget.value = texts[texts.length - 1];
+				app.graph.setDirtyCanvas(true, false);
+			}
+		});
+	},
+});

models.yaml.example

@@ -6,6 +6,16 @@ models:
     api_key: none
     model: whatever
 
+  - name: "[LOCAL] LM Studio"
+    base_url: http://localhost:1234/v1
+    api_key: none
+    model: changeme
+
+  - name: "[LOCAL] Ollama"
+    base_url: http://localhost:11434/v1
+    api_key: none
+    model: changeme
+
 # Example of using a remote endpoint
 # If prefixed with "os.environ/" (similar to LiteLLM), base_url and/or api_key
 # will take their value from the named environment variable.

nodes.py

@@ -6,12 +6,14 @@ from typing import Any, Optional
 
 from aiohttp import web
 from aiohttp.web_request import Request
-import openai
 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
 
 
@@ -26,11 +28,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 +48,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 +73,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 +92,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)
@@ -114,74 +116,156 @@ 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,
-                })
+            "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,
+                        "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'
+    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):
+    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 = ''
+            text = ""
 
-        return { 'ui': { 'text': text }, 'result': (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
+    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 +274,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,90 +302,98 @@ 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):
-        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
-        )
+    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)
+            self._model = getattr(prov_def, "model", None)
         else:
             self._model = model
 
     def get_models(self):
-        models = list(self._llm.models.list())
-        return [m.id for m in models]
+        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:
+    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)
@@ -311,32 +403,35 @@ 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
-                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
-
-        if seed is not None:
-            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
+            extra_args[k] = v
+        extra_args["samplers"] = (
+            sampler_order  # Hopefully won't cause an issue for non-llama.cpp endpoints
         )
 
-        # TODO Since we're not streaming, we'll need some sort of timeout. How to specify that?
+        if seed is not None:
+            extra_args["seed"] = seed
 
-        return output.choices[0].message.content
+        # 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:
@@ -345,9 +440,9 @@ class LLMProvider:
         PROVIDERS.refresh()
 
         return {
-            'required': {
-                'provider': (PROVIDERS.CHOICES,),
-                'model': (['fetch.models.first'],)
+            "required": {
+                "provider": (PROVIDERS.CHOICES,),
+                "model": (["fetch.models.first"],),
             },
         }
 
@@ -358,20 +453,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 +479,8 @@ class LLMModelNode:
     def INPUT_TYPES(cls):
         MODELS.refresh()
         return {
-            'required': {
-                'model': (MODELS.CHOICES,),
+            "required": {
+                "model": (MODELS.CHOICES,),
             },
         }
 
@@ -394,14 +489,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 +508,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 +576,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 +586,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 +595,21 @@ 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,
+    "LLMPrependAppend": LLMPrependAppend,
+    "LLMTextLatch": LLMTextLatch,
+    "LLMProvider": LLMProvider,
+    "LLMModel": LLMModelNode,
+    "LLMChat": LLMChat,
+    "LLMTemperature": LLMTemperature,
+    "LLMTopP": LLMTopP,
+    "LLMTopK": LLMTopK,
+    "LLMMinP": LLMMinP,
 }

providers.yaml.example

@@ -5,6 +5,14 @@ providers:
     base_url: http://localhost:8080/v1
     api_key: none
 
+  - name: "[LOCAL] LM Studio"
+    base_url: http://localhost:1234/v1
+    api_key: none
+
+  - name: "[LOCAL] Ollama"
+    base_url: http://localhost:11434/v1
+    api_key: none
+
 # Example of using a remote endpoint
 # If prefixed with "os.environ/" (similar to LiteLLM), base_url and/or api_key
 # will take their value from the named environment variable.

pyproject.toml

@@ -1,12 +1,14 @@
 [project]
 name = "comfyui-yallm-node"
 description = "Yet another set of LLM nodes for ComfyUI (for local/remote OpenAI-like APIs, multi-modal models supported)"
-version = "1.3.0"
+version = "1.5.0"
 license = {file = "LICENSE.txt"}
 dependencies = [
-    "openai",
+    "pillow",
     "pydantic",
-    "pyyaml"
+    "pyyaml",
+    "requests",
+    "torchvision"
 ]
 
 [project.urls]
@@ -19,7 +21,7 @@ DisplayName = "ComfyUI-YALLM-node"
 Icon = ""
 
 [tool.bumpversion]
-current_version = "1.3.0"
+current_version = "1.5.0"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

requirements.txt

@@ -1,7 +1,6 @@
-openai
-pydantic
-
 # These should already be installed by ComfyUI
 pillow
+pydantic
 pyyaml
+requests
 torchvision

utils.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2024 Allan Saddi <allan@saddi.com>
+class PromptUtils:
+    def __init__(self, prompt):
+        self.prompt = prompt
+
+    def get_node_info(self, node_id: int | str) -> dict:
+        node_id = str(node_id)
+        return self.prompt[node_id]
+
+    def is_input_connected(self, node_id: int | str, name: str) -> bool:
+        node_info = self.get_node_info(node_id)
+        input = node_info.get("inputs", {}).get(name)
+        # In reality, input will be a 2-elem list of source node ID + slot
+        # if it is connected.
+        return input is not None
+
+    @property
+    def output_map(self):
+        if hasattr(self, "_output_map"):
+            return self._output_map
+        # A map of (src node ID, slot) -> (dest node ID, "input name")
+        self._output_map: dict[tuple[str, int], list[tuple[str, str]]] = {}
+        for dest_id, dest_info in self.prompt.items():
+            for input_name, input in dest_info.get("inputs", {}).items():
+                if isinstance(input, list) and len(input) == 2:
+                    # This is apparently a link
+                    src_id, src_slot = input
+                    out_key = (src_id, src_slot)
+                    outputs = self._output_map.get(out_key)
+                    if outputs is None:
+                        self._output_map[out_key] = outputs = []
+                    outputs.append((dest_id, input_name))
+        return self._output_map
+
+    def is_output_connected(self, node_id: int | str, slot: int) -> bool:
+        out_key = (str(node_id), slot)
+        return out_key in self.output_map
+
+    # The assumption is that the given output slot can only be connected
+    # to nodes of the same type as the one being checked.
+    def get_downstream_nodes(self, node_id: int | str, slot: int) -> list[str]:
+        downstream: list[str] = []
+
+        to_check = [str(node_id)]
+        while to_check:
+            check_id = to_check.pop(0)
+
+            # Yes, consider the starting node downstream from itself as well
+            downstream.append(check_id)
+
+            outputs = self.output_map.get((check_id, slot))
+            if outputs is not None:
+                # NB There can be multiple outputs
+                out_ids: list[str] = [link_info[0] for link_info in outputs]
+                to_check.extend(out_ids)
+
+        return downstream
+
+
+class WorkflowUtils:
+    def __init__(self, extra_pnginfo):
+        # Note: API mode does not have extra_pnginfo (it is None)
+        self.extra_pnginfo = extra_pnginfo
+
+    def get_node_info(self, node_id: str | int) -> dict | None:
+        # Be somewhat paranoid since API mode offers much less access
+        if (
+            self.extra_pnginfo is not None
+            and (workflow := self.extra_pnginfo.get("workflow")) is not None
+            and (nodes := workflow.get("nodes")) is not None
+        ):
+            node_id = str(node_id)
+            # TODO For the time being, we aren't caching anything.
+            # Concievably, if there are a large number of nodes, a lot of time
+            # could be wasted searching.
+            my_node = [node_info for node_info in nodes if node_info["id"] == node_id]
+            if my_node:
+                return my_node[0]
+        return None
+
+    def set_property(self, node_id: str | int, property_name: str, value):
+        if (my_node := self.get_node_info(node_id)) is not None:
+            my_node["properties"][property_name] = value
+
+    def set_widget(self, node_id: str | int, slot: int, value):
+        if (my_node := self.get_node_info(node_id)) is not None:
+            values = my_node.get("widgets_values", [])
+            # TODO When would there be a mismatch?
+            if slot < len(values):
+                values[slot] = value
+
+
+class AnyType(str):
+    def __ne__(self, other):
+        if self == "*":
+            # If we're the wildcard, match anything
+            return False
+        return super().__ne__(other)
+
+
+class ComboType(str):
+    def __ne__(self, other):
+        if self == "*" and isinstance(other, list):
+            return False
+        return True