diff --git a/nodes/ollama_generate_string.py b/nodes/ollama_generate_string.py index 73c56e8..eb7b27c 100644 --- a/nodes/ollama_generate_string.py +++ b/nodes/ollama_generate_string.py @@ -5,10 +5,29 @@ from collections import deque from server import PromptServer from aiohttp import web import aiohttp +import asyncio OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "127.0.0.1:11434") +""" + + To be able to store the generated LLM prompt into the png for reproducibility the flow is as follows: + + This node takes: + generate_seed: This is the seed used for the next generation. + generate_prompt: This is the prompt used for the next generation. + use_seed This is the seed used for the current llm result used as output artifact (generated in previous invocation) + use_result This is the result generated by the previous invocation, used for the image generation. + use_prompt This is the prompt used by the previous invocation, used for the image generation. + + When executed: + generate_seed + generate_prompt is used to generate an LLM response, this is sent to the UI. + The ui stores use_* and sends it to the node for image generation. + + This way the use_* coming in is stored in the generated image. +""" + class OllamaGenerateString : def __init__(self): pass @@ -17,33 +36,57 @@ class OllamaGenerateString : def INPUT_TYPES(s): return { "required": { - "int_value": ("INT",{"defaultInput": True}), - "prompt": ("STRING",{"defaultInput": True, "default": "", "multiline": True, "dynamicPrompts": False}), + "generate": ("BOOLEAN", {"default": True, "label_on": "Generating New", "label_off": "Reusing", "tooltip": "Whether to cycle the result, or keep it fixed to the currently shown value."}), + "generate_seed": ("INT",{"defaultInput": True}), + "generate_prompt": ("STRING",{"defaultInput": True, "default": "", "multiline": True, "dynamicPrompts": False}), + # And the actually used values follow; + "use_seed": ("INT",{"defaultInput": False, "disabled": True, "default":"",}), + "use_result": ("STRING",{"defaultInput": False, "multiline": True, "default":"",}), + "use_prompt": ("STRING",{"defaultInput": False, "multiline": True, "default":"",}), }, "optional": { - }, + }, } RETURN_TYPES = ("STRING", ) FUNCTION = "execute" - OUTPUT_IS_LIST = (True,) + + # Do we need this? + # OUTPUT_IS_LIST = (False,) CATEGORY = "iw" @classmethod - def IS_CHANGED(cls, int_value, prompt ): + def IS_CHANGED(cls, generate_seed, generate_prompt, use_seed=None, use_result=None, use_prompt=None, generate=True): return "foo" - def execute(self, int_value, prompt): - result = (prompt,) - print("Executing with ", prompt) - print(" generating new with ", int_value) - return { "ui": { "next_int": [int_value], "next_value": [f"super cool: {int_value}"]}, "result": [result] } + def execute(self, generate_seed, generate_prompt, use_seed=None, use_result=None, use_prompt=None, generate=True): + result = (use_result,) + print(f"executing with {generate} {generate_seed} {generate_prompt}") + print(f" Using {use_seed} {use_result} {use_prompt}") + + + if generate: + next_seed = generate_seed + next_result = f"generated with {generate_prompt} and seed {generate_seed}" + next_prompt = generate_prompt + else: + next_seed = use_seed + next_result = use_result + next_prompt = use_prompt + + return { + "ui": { + "next_seed": [next_seed,], + "next_result": [next_result,], + "next_prompt": [next_prompt,], + }, + "result": result + } @PromptServer.instance.routes.get("/iw/api/ollama/models") -async def retrieve_models(request): - print("/iw/api/ollama/models") +async def retrieve_models(request): try: async with aiohttp.ClientSession() as session: async with session.get(f'http://{OLLAMA_HOST}/api/tags') as response: diff --git a/web/ollama_generate_string.js b/web/ollama_generate_string.js index 1e311a3..b8f72d4 100644 --- a/web/ollama_generate_string.js +++ b/web/ollama_generate_string.js @@ -2,6 +2,10 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; import { ComfyWidgets } from "../../scripts/widgets.js"; + +// This is helpful: +// https://github.com/Comfy-Org/ComfyUI_frontend/tree/0f0029ca29f4a11b37b84cd4d575280484bdd9e9/src/extensions/core + app.registerExtension({ name: "iw.ollama.generate", @@ -64,10 +68,16 @@ app.registerExtension({ nodeType.prototype.onExecuted = function (ui) { // This is after the fact... so we ready the next value. - console.log("onExecuted next_value", ui.next_value); - let w = find_widget(this, "prompt"); - w.value = ui.next_value; - w.setValue(w.value, false, true); + console.log("onExecuted ui:", ui); + for (let orig_key of ["next_seed", "next_result", "next_prompt"]) { + let next_key = orig_key.replace("next", "use"); + let w = find_widget(this, next_key); + if (w !== undefined) { + w.value = ui[orig_key][0]; + w.setValue(w.value, false, true); + } + + } } }