From 300669de661a23230acfa569cd2bfaa1e367478e Mon Sep 17 00:00:00 2001 From: Rhys Yang Date: Fri, 13 Oct 2023 23:55:08 +0800 Subject: [PATCH] Add seedGen.js --- js/parameterDisplay.js | 59 ++++++++++++++++ js/seedGen.js | 155 +++++++++++++++++++++++++++++++++++++++++ nodes.py | 61 +++++++++++----- 3 files changed, 256 insertions(+), 19 deletions(-) create mode 100644 js/parameterDisplay.js create mode 100644 js/seedGen.js diff --git a/js/parameterDisplay.js b/js/parameterDisplay.js new file mode 100644 index 0000000..9310e14 --- /dev/null +++ b/js/parameterDisplay.js @@ -0,0 +1,59 @@ +import {app} from "/scripts/app.js"; +import {ComfyWidgets} from "/scripts/widgets.js"; + +// Create a read-only string widget with opacity set +function createWidget(app, node, widgetName, type) { + const widget = ComfyWidgets[type](node, widgetName, ["STRING", {multiline: true}], app).widget; + widget.inputEl.readOnly = true; + widget.inputEl.style.textAlign = "center"; + widget.inputEl.style.fontSize = "0.75rem"; + + return widget; +} + +// Displays prompt and setting on the node +app.registerExtension({ + name: "sd_prompt_reader.parameterDisplay", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "SDParameterGenerator") { + const onNodeCreated = nodeType.prototype.onNodeCreated; + + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreated?.apply(this, arguments); + + // Create prompt and setting widgets + const aspect_ratio_display = createWidget(app, this, "aspect_ratio_display", "STRING"); + const steps_display = createWidget(app, this, "steps_display", "STRING"); + // Resize the node + const nodeWidth = this.size[0]; + const nodeHeight = this.size[1]; + this.setSize([nodeWidth*2, nodeHeight * 1.2]); + return result; + + }; + + // Update widgets + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, arguments); + let ar_message; + if (message.text[0] === "custom") { + ar_message = "Custom aspect ratio: " + message.text[2] + " x " + message.text[3]; + } + else{ + ar_message = `Optimal resolution for ${message.text[1]} model +with aspect ratio ${message.text[0]}: ${message.text[2]} x ${message.text[3]}`; + } + + const start_at_float = parseFloat(message.text[5]) + const base_percentage = Math.round(start_at_float* 100) + "%"; + const refiner_percentage = Math.round((1-start_at_float)* 100) + "%"; + const step_message = `Total steps: ${message.text[4]}, refiner start at ${base_percentage}, +Base steps: ${message.text[6]} (${base_percentage}), Refiner steps: ${message.text[7]} (${refiner_percentage})`; + + this.widgets.find(obj => obj.name === "aspect_ratio_display").value = ar_message; + this.widgets.find(obj => obj.name === "steps_display").value = step_message; + }; + } + }, +}); \ No newline at end of file diff --git a/js/seedGen.js b/js/seedGen.js new file mode 100644 index 0000000..321ac0d --- /dev/null +++ b/js/seedGen.js @@ -0,0 +1,155 @@ +/** + * Modified from: https://github.com/rgthree/rgthree-comfy/blob/main/web/seed.js + * Modified by: receyuki + */ +import {app} from "../../scripts/app.js"; +import {ComfyWidgets} from "../../scripts/widgets.js"; +const LAST_SEED_BUTTON_LABEL = "(Use Last Queued Seed)"; +const SPECIAL_SEED_RANDOM = -1; +const SPECIAL_SEED_INCREMENT = -2; +const SPECIAL_SEED_DECREMENT = -3; +const SPECIAL_SEEDS = [SPECIAL_SEED_RANDOM, SPECIAL_SEED_INCREMENT, SPECIAL_SEED_DECREMENT]; +class SeedControl { + constructor(node) { + this.lastSeed = undefined; + this.serializedCtx = {}; + this.lastSeedValue = null; + this.node = node; + this.node.constructor.exposedActions = ["Randomize seed each time", "Use last queued seed"]; + const handleAction = this.node.handleAction; + this.node.handleAction = async (action) => { + handleAction && handleAction.call(this.node, action); + if (action === "Randomize each time") { + this.seedWidget.value = SPECIAL_SEED_RANDOM; + } + else if (action === "Use last queued seed") { + this.seedWidget.value = this.lastSeed != null ? this.lastSeed : this.seedWidget.value; + this.lastSeedButton.name = LAST_SEED_BUTTON_LABEL; + this.lastSeedButton.disabled = true; + } + }; + this.node.properties = this.node.properties || {}; + for (const [i, w] of this.node.widgets.entries()) { + if (w.name === "seed") { + this.seedWidget = w; + } + else if (w.name === "control_after_generate") { + this.node.widgets.splice(i, 1); + } + } + if (!this.seedWidget) { + throw new Error("Something's wrong; expected seed widget"); + } + const randMax = Math.min(1125899906842624, this.seedWidget.options.max); + const randMin = Math.max(0, this.seedWidget.options.min); + const randomRange = (randMax - Math.max(0, randMin)) / (this.seedWidget.options.step / 10); + this.node.addWidget("button", "Randomize seed each time", null, () => { + this.seedWidget.value = SPECIAL_SEED_RANDOM; + }, { serialize: false }); + this.node.addWidget("button", "New fixed random seed", null, () => { + this.seedWidget.value = + Math.floor(Math.random() * randomRange) * (this.seedWidget.options.step / 10) + randMin; + }, { serialize: false }); + this.lastSeedButton = this.node.addWidget("button", LAST_SEED_BUTTON_LABEL, null, () => { + this.seedWidget.value = this.lastSeed != null ? this.lastSeed : this.seedWidget.value; + this.lastSeedButton.name = LAST_SEED_BUTTON_LABEL; + this.lastSeedButton.disabled = true; + }, { serialize: false }); + this.lastSeedButton.disabled = true; + this.seedWidget.serializeValue = async (node, index) => { + const inputSeed = this.seedWidget.value; + this.serializedCtx = { + inputSeed: this.seedWidget.value, + }; + if (SPECIAL_SEEDS.includes(this.serializedCtx.inputSeed)) { + if (typeof this.lastSeed === "number" && !SPECIAL_SEEDS.includes(this.lastSeed)) { + if (inputSeed === SPECIAL_SEED_INCREMENT) { + this.serializedCtx.seedUsed = this.lastSeed + 1; + } + else if (inputSeed === SPECIAL_SEED_INCREMENT) { + this.serializedCtx.seedUsed = this.lastSeed - 1; + } + } + if (!this.serializedCtx.seedUsed || SPECIAL_SEEDS.includes(this.serializedCtx.seedUsed)) { + this.serializedCtx.seedUsed = + Math.floor(Math.random() * randomRange) * (this.seedWidget.options.step / 10) + randMin; + } + } + else { + this.serializedCtx.seedUsed = this.seedWidget.value; + } + node.widgets_values[index] = this.serializedCtx.seedUsed; + this.seedWidget.value = this.serializedCtx.seedUsed; + this.lastSeed = this.serializedCtx.seedUsed; + if (SPECIAL_SEEDS.includes(this.serializedCtx.inputSeed)) { + this.lastSeedButton.name = `${this.serializedCtx.seedUsed}`; + this.lastSeedButton.disabled = false; + if (this.lastSeedValue) { + this.lastSeedValue.value = `Last Seed: ${this.serializedCtx.seedUsed}`; + } + } + else { + this.lastSeedButton.name = LAST_SEED_BUTTON_LABEL; + this.lastSeedButton.disabled = true; + } + return this.serializedCtx.seedUsed; + }; + this.seedWidget.afterQueued = () => { + if (this.serializedCtx.inputSeed) { + this.seedWidget.value = this.serializedCtx.inputSeed; + } + this.serializedCtx = {}; + }; + this.node.getExtraMenuOptions = (_, options) => { + options.splice(options.length - 1, 0, { + content: "Show/Hide Last Seed Value", + callback: (_value, _options, _event, _parentMenu, _node) => { + this.node.properties["showLastSeed"] = !this.node.properties["showLastSeed"]; + if (this.node.properties["showLastSeed"]) { + this.addLastSeedValue(); + } + else { + this.removeLastSeedValue(); + } + }, + }); + }; + } + addLastSeedValue() { + if (this.lastSeedValue) + return; + this.lastSeedValue = ComfyWidgets["STRING"](this.node, "last_seed", ["STRING", { multiline: true }], app).widget; + this.lastSeedValue.inputEl.readOnly = true; + this.lastSeedValue.inputEl.style.fontSize = "0.75rem"; + this.lastSeedValue.inputEl.style.textAlign = "center"; + this.lastSeedValue.serializeValue = async (node, index) => { + node.widgets_values[index] = ""; + return ""; + }; + // this.node.computeSize(); + } + removeLastSeedValue() { + if (!this.lastSeedValue) + return; + this.lastSeedValue.inputEl.remove(); + this.node.widgets.splice(this.node.widgets.indexOf(this.lastSeedValue), 1); + this.lastSeedValue = null; + // this.node.computeSize(); + } +} +app.registerExtension({ + name: "sd_prompt_reader.seedGen", + async beforeRegisterNodeDef(nodeType, nodeData, _app) { + if (nodeData.name === "SDParameterGenerator") { + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + onNodeCreated ? onNodeCreated.apply(this, []) : undefined; + this.seedControl = new SeedControl(this); + // Resize the node + const nodeWidth = this.size[0]; + const nodeHeight = this.size[1]; + this.setSize([nodeWidth*1.5, nodeHeight * 1.1]); + }; + } + }, +}); diff --git a/nodes.py b/nodes.py index dd0efd2..5609bfc 100644 --- a/nodes.py +++ b/nodes.py @@ -203,7 +203,7 @@ class SDPromptSaver: "seed": ( "INT", { - "default": 0, + "default": -1, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, }, @@ -586,21 +586,38 @@ class SDParameterGenerator: base_steps = int(steps * refiner_start) refiner_steps = steps - base_steps - return checkpoint + ( - ckpt_name, - seed, - steps, - cfg, - sampler_name, - scheduler, - width, - height, - batch_size, - positive_ascore, - negative_ascore, - base_steps, - refiner_steps, - ) + return { + "ui": { + "text": ( + aspect_ratio, + model_version, + width, + height, + steps, + refiner_start, + base_steps, + refiner_steps, + ) + }, + "result": ( + checkpoint + + ( + ckpt_name, + seed, + steps, + cfg, + sampler_name, + scheduler, + width, + height, + batch_size, + positive_ascore, + negative_ascore, + base_steps, + refiner_steps, + ) + ), + } class SDPromptMerger: @@ -635,9 +652,15 @@ class SDTypeConverter: return { "required": {}, "optional": { - "model_name": (folder_paths.get_filename_list("checkpoints"),), - "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), - "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), + "model_name": ( + folder_paths.get_filename_list("checkpoints"), + {"forceInput": True}, + ), + "sampler_name": ( + comfy.samplers.KSampler.SAMPLERS, + {"forceInput": True}, + ), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"forceInput": True}), }, }