diff --git a/__init__.py b/__init__.py index e3adbb7..659bf98 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,11 @@ import random +class StringToolsOptionalDict(dict): + def __contains__(self, key: object) -> bool: + return True + + def sort_kwargs_value(basename, kwargs, separator="_"): sorted_args = {} basename_dict = {} @@ -15,9 +20,7 @@ def sort_kwargs_value(basename, kwargs, separator="_"): basename_dict.items(), key=lambda kv: int(kv[0].split(separator)[-1]) ) - return list(sorted_args.values()) + list( - map(lambda kv: kv[1], sorted_basename_dict) - ) + return list(map(lambda kv: kv[1], sorted_basename_dict)) class StringToolsString: @@ -148,10 +151,31 @@ class StringToolsRandomChoice: return (choice,) +class StringToolsBalancedChoice(StringToolsRandomChoice): + @classmethod + def INPUT_TYPES(s): + input_types = super().INPUT_TYPES() + input_types["optional"] = StringToolsOptionalDict() + return input_types + + def process(self, *args, **kwargs): + seed = 0 + if "seed" in kwargs: + seed = kwargs["seed"] + del kwargs["seed"] + + values = sort_kwargs_value("text", kwargs) + weights = sort_kwargs_value("weight", kwargs) + random.seed(seed) + choice = random.choices(values, weights=weights)[0] + return (choice,) + + NODE_CLASS_MAPPINGS = { "StringToolsString": StringToolsString, "StringToolsText": StringToolsText, "StringToolsConcat": StringToolsConcat, "StringToolsRandomChoice": StringToolsRandomChoice, + "StringToolsBalancedChoice": StringToolsBalancedChoice, } WEB_DIRECTORY = "./js" diff --git a/js/index.js b/js/index.js index dfe34a8..a641d3b 100644 --- a/js/index.js +++ b/js/index.js @@ -1,6 +1,7 @@ import { app } from "/scripts/app.js"; +import { api } from "/scripts/api.js"; import { setWidgetConfig } from "/extensions/core/widgetInputs.js"; -function createCallback(nodename, basename, inputType) { +function createCallback(nodename, basename, inputType, withWeights) { return async function (nodeType, nodeData, app) { if (nodeData.name !== nodename) { return; @@ -8,6 +9,9 @@ function createCallback(nodename, basename, inputType) { const getInputBasename = function (input) { return input.name.split('_')[0]; }; + const getInputExtraname = function (input) { + return input.name.split('_', 2)[1]; + }; const updateInputs = function () { // remove empty inputs for (let index = this.inputs.length; index--;) { @@ -114,6 +118,40 @@ function createCallback(nodename, basename, inputType) { updateInputs.call(this); } }; + if (withWeights !== void 0) { + this.calcNodeInputs = function (prompt, workflow) { + const type = prompt[this.id].class_type; + for (const input of Object.keys(prompt[this.id].inputs)) { + if (getInputBasename({ name: input }) !== basename) { + continue; + } + const extraname = getInputExtraname({ name: input }); + const walkdown = (type, id, sum) => { + for (const input of Object.keys(prompt[id].inputs)) { + const value = prompt[id].inputs[input]; + let start = 0; + if (withWeights.includes(prompt[id].class_type) && + getInputBasename({ name: input }) === basename) { + start = 1; + } + if (Array.isArray(value)) { + sum += walkdown(type, value[0], start); + } + } + return sum; + }; + const value = prompt[this.id].inputs[input]; + if (Array.isArray(value)) { + let sum = walkdown(type, value[0], 0); + if (sum === 0) { + sum = 1; + } + const weightKey = ["weight", extraname].join('_'); + prompt[this.id].inputs[weightKey] = sum; + } + } + }; + } // init inputs if (!this.inputs) { this.inputs = []; @@ -121,6 +159,16 @@ function createCallback(nodename, basename, inputType) { }; }; } +const queuePromptOriginal = api.queuePrompt; +api.queuePrompt = (async function queuePrompt(number, { output, workflow }) { + for (const id of Object.keys(output)) { + const node = app.graph.getNodeById(id); + if (node.calcNodeInputs && typeof node.calcNodeInputs === "function") { + node.calcNodeInputs(output, workflow); + } + } + return await queuePromptOriginal(number, { output, workflow }); +}).bind(api); app.registerExtension({ name: "Taremin.StringToolsConcat", beforeRegisterNodeDef: createCallback("StringToolsConcat", "text", "STRING"), @@ -129,3 +177,7 @@ app.registerExtension({ name: "Taremin.StringToolsRandomChoice", beforeRegisterNodeDef: createCallback("StringToolsRandomChoice", "text", "STRING"), }); +app.registerExtension({ + name: "Taremin.StringToolsBalancedChoice", + beforeRegisterNodeDef: createCallback("StringToolsBalancedChoice", "text", "STRING", ["StringToolsRandomChoice", "StringToolsBalancedChoice"]), +}); diff --git a/src/index.ts b/src/index.ts index 3438202..ea67bbe 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,9 +1,10 @@ import { app } from "/scripts/app.js" +import { api } from "/scripts/api.js" import { setWidgetConfig } from "/extensions/core/widgetInputs.js" declare const LiteGraph: any -function createCallback(nodename: string, basename: string, inputType: string) { +function createCallback(nodename: string, basename: string, inputType: string, withWeights?:string[]) { return async function(nodeType:any, nodeData:any, app:any) { if (nodeData.name !== nodename) { return @@ -13,6 +14,10 @@ function createCallback(nodename: string, basename: string, inputType: string) { return input.name.split('_')[0] } + const getInputExtraname = function(input: any) { + return input.name.split('_', 2)[1] + } + const updateInputs = function(this: any) { // remove empty inputs for (let index = this.inputs.length; index--; ) { @@ -138,6 +143,44 @@ function createCallback(nodename: string, basename: string, inputType: string) { } } + if (withWeights !== void 0) { + this.calcNodeInputs = function(prompt: {[key: string]: {class_type: string, inputs: {[key: string]: any}}}, workflow: {nodes: any[]}) { + const type = prompt[this.id].class_type + + for (const input of Object.keys(prompt[this.id].inputs)) { + if (getInputBasename({name: input}) !== basename) { + continue + } + const extraname = getInputExtraname({name: input}) + const walkdown = (type: string, id: string, sum: number): number => { + for (const input of Object.keys(prompt[id].inputs)) { + const value = prompt[id].inputs[input] + let start = 0 + if ( + withWeights.includes(prompt[id].class_type) && + getInputBasename({name: input}) === basename + ) { + start = 1 + } + if (Array.isArray(value)) { + sum += walkdown(type, value[0], start) + } + } + return sum + } + const value = prompt[this.id].inputs[input] + if (Array.isArray(value)) { + let sum = walkdown(type, value[0], 0) + if (sum === 0) { + sum = 1 + } + const weightKey = ["weight", extraname].join('_') + prompt[this.id].inputs[weightKey] = sum + } + } + } + } + // init inputs if (!this.inputs) { this.inputs = [] @@ -146,6 +189,17 @@ function createCallback(nodename: string, basename: string, inputType: string) { } } +const queuePromptOriginal = api.queuePrompt +api.queuePrompt = (async function queuePrompt(number: number, {output, workflow}: {output: any, workflow: any}) { + for (const id of Object.keys(output)) { + const node = app.graph.getNodeById(id) + if (node.calcNodeInputs && typeof node.calcNodeInputs === "function") { + node.calcNodeInputs(output, workflow) + } + } + return await queuePromptOriginal(number, {output, workflow}) +}).bind(api) + app.registerExtension({ name: "Taremin.StringToolsConcat", beforeRegisterNodeDef: createCallback("StringToolsConcat", "text", "STRING"), @@ -155,3 +209,8 @@ app.registerExtension({ name: "Taremin.StringToolsRandomChoice", beforeRegisterNodeDef: createCallback("StringToolsRandomChoice", "text", "STRING"), }) + +app.registerExtension({ + name: "Taremin.StringToolsBalancedChoice", + beforeRegisterNodeDef: createCallback("StringToolsBalancedChoice", "text", "STRING", ["StringToolsRandomChoice", "StringToolsBalancedChoice"]), +})