Add Parameter Extractor node #24

This commit is contained in:
Rhys Yang
2023-12-11 19:55:03 +08:00
parent 52882ee594
commit 97c6aa988d
3 changed files with 142 additions and 0 deletions
+1
View File
@@ -17,6 +17,7 @@ files_to_copy = [
"parameterDisplay.js",
"seedGen.js",
"loaderDisplay.js",
"extractorDisplay.js",
]
for file in files_to_copy:
+38
View File
@@ -0,0 +1,38 @@
import {app} from "../../scripts/app.js";
import {ComfyWidgets} from "../../scripts/widgets.js";
// Create a read-only string widget
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;
}
app.registerExtension({
name: "sd_prompt_reader.extractorDisplay",
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (nodeData.name === "SDParameterExtractor") {
const onNodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {
const result = onNodeCreated?.apply(this, arguments);
// Create widgets
const value_display = createWidget(app, this, "value_display", "STRING");
};
// Update widgets
const onExecuted = nodeType.prototype.onExecuted;
nodeType.prototype.onExecuted = function (message) {
onExecuted?.apply(this, arguments);
this.widgets.find(obj => obj.name === "value_display").value = message.text[1]
this.widgets.find(obj => obj.name === "parameter").options.values = message.text[0]
if (this.widgets.find(obj => obj.name === "parameter").value === "parameters not loaded") {
this.widgets.find(obj => obj.name === "parameter").value = message.text[0][0]
}
};
}
},
});
+103
View File
@@ -12,6 +12,7 @@ from itertools import chain
import torch
import json
import re
import numpy as np
from pathlib import Path
from PIL import Image, ImageOps
@@ -936,6 +937,106 @@ class SDBatchLoader:
return True
class SDParameterExtractor:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"settings": (
"STRING",
{"default": "", "multiline": True, "forceInput": True},
)
},
"optional": {
"parameter": (
["parameters not loaded"],
{"default": "parameters not loaded"},
),
"value_type": (["STRING", "INT", "FLOAT"], {"default": "STRING"}),
"parameter_index": (
"INT",
{"default": 0, "min": 0, "max": 255, "step": 1},
),
},
}
RETURN_TYPES = (any_type,)
RETURN_NAMES = ("VALUE",)
FUNCTION = "extract_param"
CATEGORY = "SD Prompt Reader"
def extract_param(
self,
settings: str = "",
parameter: str = "",
value_type: str = "STRING",
parameter_index: int = 0,
):
setting_dict = self.parse_setting(settings)
if not settings or not parameter or parameter == "parameters not loaded":
return {
"ui": {
"text": (list(setting_dict.keys()), ""),
},
"result": ("",),
}
result = setting_dict.get(parameter)
try:
if isinstance(result, tuple):
result = result[parameter_index]
if value_type == "INT":
result = int(result)
elif value_type == "FLOAT":
result = float(result)
except IndexError:
return {
"ui": {
"text": (list(setting_dict.keys()), "Parameter index out of range"),
},
"result": ("",),
}
except (ValueError, TypeError):
return {
"ui": {
"text": (
list(setting_dict.keys()),
f"{parameter}: {result}\n"
f"{result} is not a valid number; it will be output as STRING",
),
},
"result": (result,),
}
return {
"ui": {
"text": (list(setting_dict.keys()), f"{parameter}: {result}"),
},
"result": (result,),
}
@staticmethod
def parse_setting(settings):
pattern = re.compile(r"([^:,]+):\s*\(([^)]+)\)|([^:,]+):\s*([^,]+)")
matches = pattern.findall(settings)
result = {}
for match in matches:
key, value_paren, key_nonparen, value_nonparen = match
if key:
key = key.strip()
value = value_paren.strip()
value = tuple(v.strip() for v in value.split(","))
else:
key = key_nonparen.strip()
value = value_nonparen.strip()
result[key] = value
return result
NODE_CLASS_MAPPINGS = {
"SDPromptReader": SDPromptReader,
"SDPromptSaver": SDPromptSaver,
@@ -943,6 +1044,7 @@ NODE_CLASS_MAPPINGS = {
"SDPromptMerger": SDPromptMerger,
"SDTypeConverter": SDTypeConverter,
"SDBatchLoader": SDBatchLoader,
"SDParameterExtractor": SDParameterExtractor,
}
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -952,4 +1054,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"SDPromptMerger": "SD Prompt Merger",
"SDTypeConverter": "SD Type Converter",
"SDBatchLoader": "SD Batch Loader",
"SDParameterExtractor": "SD Parameter Extractor",
}