144 lines
4.1 KiB
Python
144 lines
4.1 KiB
Python
import numpy as np
|
|
import toml
|
|
|
|
from .utils import ROOT_DIR
|
|
|
|
|
|
class Node:
|
|
"""ComfyUI Node parent class
|
|
Will build and return a prompt according to other sub-components.
|
|
"""
|
|
|
|
def __init__(self, seed=0, data_file=""):
|
|
self.seed = seed
|
|
self.data = self.load_data(data_file)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""ComfyUI node's inputs"""
|
|
required = {
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff
|
|
})
|
|
}
|
|
optional = {}
|
|
inputs = {
|
|
"required": required,
|
|
"optional": optional
|
|
}
|
|
return inputs
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "build_prompt"
|
|
CATEGORY = "🎞️ Scene Composer"
|
|
|
|
def build_prompt(self):
|
|
pass
|
|
|
|
def select_tags(self, tags, p=1, n=1, selected="random", recursive=True):
|
|
"""Return n tags from a string, list or dict
|
|
p: probability of selecting the tags
|
|
n: number of tags to select
|
|
selected: fallback tag to select if not random
|
|
recursive: go deeper randomly to select last children tags"""
|
|
rng = np.random.default_rng(self.seed)
|
|
|
|
if selected != "random":
|
|
return selected
|
|
|
|
if isinstance(tags, dict):
|
|
|
|
p = tags.get("probability", p)
|
|
n = tags.get("repeat", n)
|
|
|
|
if rng.random() > p:
|
|
return ""
|
|
|
|
# If n is a list, choose a random number between the 2 first values
|
|
if isinstance(n, list):
|
|
min_n = n[0]
|
|
max_n = min(n[1], len(tags))
|
|
n = rng.integers(int(min_n), int(max_n))
|
|
|
|
# Handle "list" and "tags" special keys
|
|
if "list" in tags:
|
|
selected_from_list = self.select_tags(tags["list"], p, n)
|
|
tags = tags[selected_from_list]
|
|
|
|
elif "tags" in tags:
|
|
tags = tags.get("tags", [])
|
|
|
|
# Handle recursive tags
|
|
else:
|
|
first_level_keys = list(tags.keys())
|
|
chosen_key = rng.choice(first_level_keys)
|
|
chosen_value = tags[chosen_key]
|
|
|
|
if isinstance(chosen_value, (list, dict)):
|
|
tags = self.select_tags(chosen_value, p, n)
|
|
|
|
if not recursive:
|
|
tags = first_level_keys
|
|
|
|
if isinstance(tags, list):
|
|
tag_names, weights = self.parse_tag_distribution(tags)
|
|
|
|
selected_tags = rng.choice(
|
|
tag_names,
|
|
size=n,
|
|
replace=False,
|
|
p=weights
|
|
)
|
|
tags = self.stringify_tags(selected_tags)
|
|
|
|
return tags
|
|
|
|
@staticmethod
|
|
def parse_tag_distribution(tags):
|
|
"""Choose tags randomly according to defined weights
|
|
If none are defined, tag weight is 1
|
|
Example: "foo?1.5, bar?0.5" -> 75% chance of foo, 25% chance of bar"""
|
|
tag_names = []
|
|
weights = []
|
|
for tag in tags:
|
|
if '?' in tag:
|
|
tag_name, weight = tag.split('?')
|
|
tag_names.append(tag_name)
|
|
weights.append(float(weight))
|
|
else:
|
|
tag_names.append(tag)
|
|
weights.append(1)
|
|
probabilities = np.array(weights) / sum(weights)
|
|
return tag_names, probabilities
|
|
|
|
@staticmethod
|
|
def load_data(filename):
|
|
"""Load data from a TOML file"""
|
|
path = f"{ROOT_DIR}/config/{filename}"
|
|
try:
|
|
with open(path) as file:
|
|
data = toml.load(file)
|
|
except Exception:
|
|
data = {}
|
|
return data
|
|
|
|
@staticmethod
|
|
def build_inputs_list(data):
|
|
"""Build a list with a random option"""
|
|
data = [tag.split('?')[0] for tag in data]
|
|
inputs_list = ["random"] + data
|
|
|
|
return inputs_list
|
|
|
|
@staticmethod
|
|
def stringify_tags(tags):
|
|
"""Return a string from a list of tags"""
|
|
tags = ", ".join(map(str, tags))
|
|
tags = tags.replace(", ,", ",")
|
|
return tags
|
|
|
|
def __str__(self, **kwargs):
|
|
return self.build_prompt(**kwargs)[0]
|