diff --git a/nodes/ollama_generate_string.py b/nodes/ollama_generate_string.py index ed44b62..27267fc 100644 --- a/nodes/ollama_generate_string.py +++ b/nodes/ollama_generate_string.py @@ -2,7 +2,12 @@ import os import json from collections import deque -from server import PromptServer +from server import PromptServer +from aiohttp import web +import aiohttp + +OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "127.0.0.1:11434") + class OllamaGenerateString : def __init__(self): pass @@ -35,6 +40,13 @@ class OllamaGenerateString : print(" generating new with ", int_value) return { "ui": { "next_int": [int_value], "next_value": [f"super cool: {int_value}"]}, "result": [result] } -@PromptServer.instance.routes.get("/iw/api/ollama_generate") -async def load_drawing(request): - print("/iw/api/ollama_generate") +@PromptServer.instance.routes.get("/iw/api/ollama/models") +async def retrieve_models(request): + print("/iw/api/ollama/models") + async with aiohttp.ClientSession() as session: + async with session.get(f'http://{OLLAMA_HOST}/api/tags') as response: + body = await response.json() + model_simple = [v["model"] for v in body["models"]] + return web.json_response({"status": "success", "models": model_simple}, status=200) + + return web.json_response({"status": "fail"}, status=500) diff --git a/web/ollama_generate_string.js b/web/ollama_generate_string.js index 171de58..62ed9ec 100644 --- a/web/ollama_generate_string.js +++ b/web/ollama_generate_string.js @@ -5,28 +5,45 @@ import { ComfyWidgets } from "../../scripts/widgets.js"; app.registerExtension({ name: "iw.ollama.generate", - async beforeRegisterNodeDef(nodeType, nodeData, app) { + async beforeRegisterNodeDef(nodeType, nodeData, app) { let find_widget = (node, name) => node.widgets.find((w) => w.name === name); + + + if (nodeData.name === "IW_OllamaGenerateString") { + + + const resp = await api.fetchApi('/iw/api/ollama/models', { }); + + if (resp.status !== 200) { + const err = `Error uploading temp file: ${resp.status} - ${resp.statusText}` + useToastStore().addAlert(err) + throw new Error(err) + } + const result_in_json = await resp.json(); + const model_list = result_in_json.models; + console.log("model list:", model_list) + console.log("Something ollama thing"); + nodeType.prototype.initState = function () { console.log("hello"); }; nodeType.prototype.initWidgets = function (app) { - this.setupWidgets(); + this.setupWidgets(model_list); }; - nodeType.prototype.setupWidgets = function (app) { + nodeType.prototype.setupWidgets = function (app, models) { this.addWidget( "combo", "model", - "unknown", + model_list[0], (value) => { this.tool = value; }, { - values: ["a", "b", "c"], - default: "unknown", + values: model_list, + default: model_list[0], } );