diff --git a/web/ComfyUI-DrawThings-gRPC.js b/web/ComfyUI-DrawThings-gRPC.js index 843b30a..9263097 100644 --- a/web/ComfyUI-DrawThings-gRPC.js +++ b/web/ComfyUI-DrawThings-gRPC.js @@ -20,6 +20,17 @@ app.registerExtension({ } }, + beforeConfigureGraph(graph) { + for (const node of graph.nodes) { + if (node.type === "DrawThingsPositive" || node.type === "DrawThingsNegative") { + node.type = "DrawThingsPrompt" + node.properties['Node name for S&R'] = "DrawThingsPrompt" + delete node.properties.ver + node.widgets_values.unshift("...") + } + } + }, + async beforeRegisterNodeDef(nodeType, nodeData, app) { if (nodeType.comfyClass === "DrawThingsSampler") { updateProto(nodeType, samplerProto) diff --git a/web/configProperties.js b/web/configProperties.js index f4e2451..3886711 100644 --- a/web/configProperties.js +++ b/web/configProperties.js @@ -1,4 +1,4 @@ -import { updateNodeModels } from './models.js' +import { updateNodeModelsX } from './models.js' export const samplers = [ "DPM++ 2M Karras", @@ -149,7 +149,7 @@ function roundBy64(value) { /** @type {Record void>} */ const importers = { model: async (k, v, w, n, c) => { - await updateNodeModels(n) + await updateNodeModelsX(n) const matchingOption = w?.options?.values?.find(ov => ov.value?.file === v) if (matchingOption) w.value = matchingOption }, diff --git a/web/dtModelNodes.js b/web/dtModelNodes.js index 68aa150..748cbaa 100644 --- a/web/dtModelNodes.js +++ b/web/dtModelNodes.js @@ -1,6 +1,6 @@ import { app } from "../../scripts/app.js" import { setCallback } from "./dynamicInputs.js" -import { updateNodeModels } from "./models.js" +import { updateNodeModelsX, modelService, getMenuItem } from "./models.js" import { updateProto } from "./util.js" export const dtModelNodeTypes = ["DrawThingsSampler", "DrawThingsControlNet", "DrawThingsLoRA", "DrawThingsUpscaler", "DrawThingsRefiner", "DrawThingsPrompt"] @@ -18,6 +18,11 @@ app.registerExtension({ } } }, + + afterConfigureGraph() { + console.log('configure?') + modelService.updateNodes() + } }) /** @type {import("@comfyorg/litegraph").LGraphNode} */ @@ -25,7 +30,7 @@ const dtModelNodeProto = { saveSelectedModels() { const modelWidgets = this.widgets.filter((w) => w.options?.modelType) const selections = modelWidgets.reduce((acc, w) => { - if (w.value?.value?.version !== "fail") acc[w.name] = w.value + if (typeof(w.value) === "object") acc[w.name] = w.value else acc[w.name] = this._lastSelectedModel?.[w.name] return acc }, {}) @@ -62,7 +67,12 @@ const dtModelNodeProto = { } } }, - + onAdded() { + console.log('added', this.id) + }, + onRemoved() { + console.log('removed', this.id) + }, } /** @type {import("@comfyorg/litegraph").LGraphNode} */ @@ -70,31 +80,65 @@ const dtServerNodeProto = { onNodeCreated() { // update when server or port changes const serverWidget = this.widgets.find((w) => w.name === "server") - if (serverWidget) setCallback(serverWidget, "callback", () => updateNodeModels(this)) + if (serverWidget) setCallback(serverWidget, "callback", () => modelService.updateNodes()) const portWidget = this.widgets.find((w) => w.name === "port") - if (portWidget) setCallback(portWidget, "callback", () => updateNodeModels(this)) + if (portWidget) setCallback(portWidget, "callback", () => modelService.updateNodes()) + + const tlsWidget = this.widgets.find((w) => w.name === "use_tls") + if (tlsWidget) setCallback(tlsWidget, "callback", () => modelService.updateNodes()) }, - onConfigure() { - updateNodeModels(this) - }, + // onConfigure() { + // updateNodeModelsX(this) + // }, getServer() { const server = this.widgets.find((w) => w.name === "server")?.value const port = this.widgets.find((w) => w.name === "port")?.value - return { server, port } + const useTls = this.widgets.find((w) => w.name === "use_tls")?.value + return { server, port, useTls } }, getModelVersion() { return this.widgets.find((w) => w.options?.modelType === "models")?.value?.value?.version }, + + /** + * @param {import("./models.js").ModelInfo} models + * @param {string} version + */ + updateModels(models, version) { + const widget = this.widgets.find(w => w.name === "model") + if (models === null) { + widget.options.values = ["Not connected", "Click to retry"] + widget.value = "Not connected" + return + } + + widget.options.values = ["(None selected)", ...models.models.map(m => getMenuItem(m)).sort((a, b) => a.content.localeCompare(b.content))] + + if (widget.value === "Click to retry" || widget.value === "Not connected") { + if (this._lastSelectedModel?.model) widget.value = this._lastSelectedModel.model + else widget.value = "(None selected)" + } + + if (widget.value?.toString() === "[object Object]") { + const value = { + ...widget.value, + toString() { + return this.value.name + }, + } + widget.value = value + } + } } /** @type {import("@comfyorg/litegraph").LGraphNode} */ const dtExtraNodeProto = { onConnectionsChange(...args) { const isConnected = args[2] - if (isConnected) updateNodeModels(this) + if (isConnected) updateNodeModelsX(this) }, } diff --git a/web/models.js b/web/models.js index b00f63e..547b66b 100644 --- a/web/models.js +++ b/web/models.js @@ -1,4 +1,61 @@ import { dtModelNodeTypes } from './dtModelNodes.js' +import * as App from "../../scripts/app.js" + +/** @type {import("@comfyorg/comfyui-frontend-types").ComfyApp} */ +const app = App.app + +class ModelService { + constructor() { + + } + + async updateNodes() { + console.debug('updating models') + const dtModelNodes = app.graph.nodes.filter(n => n.isDtServerNode !== undefined) + const serverNodes = dtModelNodes.filter(n => n.isDtServerNode) + + const nodesUpdated = new Map(dtModelNodes.map(n => ([n, false]))) + + const serverModels = new Map() + + for (const sn of serverNodes) { + nodesUpdated[sn] = true + const { server, port, useTls } = sn.getServer() + + if (!server || !port || useTls === undefined) continue + const key = modelInfoStoreKey(server, port, useTls) + if (!serverModels.has(key)) { + serverModels.set(key, await getModels(server, port, useTls)) + } + + const models = serverModels.get(key) + + /** @param {import("@comfyorg/litegraph").LGraphNode} node */ + function updateInputs(node) { + nodesUpdated[nodesUpdated] = true + node?.updateModels(models) + for (let slot = 0; slot < node.inputs.length; slot++) { + const input = node.getInputNode(slot) + if (input && nodesUpdated[input] === false) updateInputs(input) + } + } + + updateInputs(sn) + } + + // these nodes are not connected to a sampler node + // if there's a single server, update them with those models + // otherwise, they are sol + const models = serverModels.size === 1 ? serverModels.values().next().value : null + for (const node of nodesUpdated.keys()) { + if (nodesUpdated[node] === false) { + node.updateModels?.(null) + } + } + } +} + +export const modelService = new ModelService() /** * @param node {LGraphNode} @@ -12,8 +69,8 @@ export function DtModelTypeHandler(node, inputName, inputData, app) { "(None selected)", /** @type WidgetCallback> */ (value, graph, node) => { - if (node.saveSelectedModels && value.value?.version !== "fail") node.saveSelectedModels() - updateNodeModels(node) + node.saveSelectedModels?.() + modelService.updateNodes() }, { values: failedConnectionOptions.map((o) => getMenuItem(o, false)), @@ -44,12 +101,25 @@ export async function getFiles(server, port, useTls) { return filesInfoResponse } -/** @typedef {{ models: any[], controlNets: any[], loras: any[], upscalers: any[]}} ModelInfo */ +/* @typedef {{ models: any[], controlNets: any[], loras: any[], upscalers: any[]}} ModelInfo */ + +/** + * @typedef {Object} Model + * @property {string} name + * @property {string} file + * @property {string} version + */ + +/** + * @typedef {Object} ModelInfo + * @property {Model[]} models + */ + /** @type Map */ const modelInfoStore = new Map() /** @type Map> */ const modelInfoRequests = new Map() -const modelInfoStoreKey = (server, port) => `${server}:${port}` +const modelInfoStoreKey = (server, port, useTls) => `${server}:${port}${useTls ? ":tls" : ""}` // yes this is kind of hacky :) const failedConnectionOptions = ["Couldn't connect to server", "Check server and click to retry"].map((c, i) => ({ @@ -64,6 +134,35 @@ const notConnectedOptions = ["Not connected to sampler node", "Connect to a samp }) ) +async function getModels(server, port, useTls) { + if (!server || !port || useTls === undefined) return + const key = modelInfoStoreKey(server, port, useTls) + + if (modelInfoRequests.has(key)) { + const request = modelInfoRequests.get(key) + await request + } + else { + const promise = new Promise((resolve) => { + console.debug("checking DT server", key, " (", fetches++, ")") + getFiles(server, port, useTls).then(async (response) => { + if (!response.ok) { + modelInfoStore.set(key, null) + } else { + const data = await response.json() + modelInfoStore.set(key, data) + } + modelInfoRequests.delete(key) + resolve() + }) + }) + modelInfoRequests.set(key, promise) + await promise + } + + return modelInfoStore.get(key) +} + const failedConnectionInfo = { models: failedConnectionOptions, controlNets: notConnectedOptions, @@ -75,9 +174,10 @@ const failedConnectionInfo = { modelInfoStore.set(modelInfoStoreKey(), failedConnectionInfo) let fetches = 0 -let updaates = 0 +let updates = 0 /** @param node {LGraphNode} */ -export async function updateNodeModels(node, updateDisconnected = false) { +export async function updateNodeModelsX(node, updateDisconnected = false) { + return // find the sampler node let root = findRoot(node) if (!root) { @@ -202,10 +302,10 @@ function getModelOptions(modelInfo, version) { return { models, controlNets, loras, upscalers, textualInversions, isNotConnected } } -function getMenuItem(model, disabled) { +export function getMenuItem(model, disabled) { return { value: model, - content: model.version && model.version !== "fail" ? `${model.name} (${model.version})` : model.name, + content: model.version && model.version !== "fail" ? `${model.name} (${getVersionAbbrev(model.version)})` : model.name, toString() { return model.name }, @@ -225,9 +325,27 @@ function getMenuItem(model, disabled) { const modelComparator = (a, b) => a.version?.localeCompare(b.version) || a.name.localeCompare(b.name) const versionNames = { - v1: "SD", - "sdxl_base_v0.9": "SDXL", - flux1: "Flux", + v1: 'SD', + v2: 'SD2', + 'kandinsky2.1': 'Kan', + 'sdxl_base_v0.9': 'SDXL', + 'sdxl_refiner_v0.9': 'SDXL', + ssd_1b: 'SSD', + svd_i2v: 'SVD', + 'wurstchen_v3.0_stage_c': 'Wur', + 'wurstchen_v3.0_stage_b': 'Wur', + sd3: 'SD3', + pixart: 'Pix', + auraflow: 'AF', + flux1: 'F1', + sd3_large: 'SD3L', + hunyuan_video: 'Hun', + 'wan_v2.1_1.3b': 'Wan', + 'wan_v2.1_14b': 'Wan', + hidream_i1: 'HiD', +} +function getVersionAbbrev(version) { + return versionNames[version] ?? version } function setValidOption(widget, node, isNotConnected) {