add: StringToolsBalancedChoice(ネスト時その分岐数によって重み付けを行う選択ノード)

__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"

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"]),
+});

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"]),
+})