79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
from .base import BaseNode
|
|
|
|
import numpy as np
|
|
|
|
|
|
class PromptUtilitiesSampleTags(BaseNode):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"tags": ("STRING", {"default": "", "multiline": True}),
|
|
"tags_delimiter": (("new line", ","),),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"max_k": ("INT", {"default": 1, "min": 1, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"min_k": ("INT", {"default": 1, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "sample_tags"
|
|
|
|
def sample_tags(self, tags, tags_delimiter, seed, max_k, min_k):
|
|
sampled_tags = self._sample_tags(tags, tags_delimiter, seed, max_k, min_k)
|
|
return (", ".join(sampled_tags),)
|
|
|
|
def _sample_tags(self, tags_s, tags_delimiter, seed, max_k, min_k):
|
|
assert max_k >= min_k, "max_k must be greater than or equal to min_k."
|
|
|
|
if tags_delimiter == "new line":
|
|
tags = [tag.strip() for tag in tags_s.split("\n")]
|
|
elif tags_delimiter == ",":
|
|
tags = [tag.strip() for tag in tags_s.split(",")]
|
|
|
|
random_gen = np.random.default_rng(seed)
|
|
k = random_gen.integers(min_k, max_k + 1)
|
|
random_gen.shuffle(tags)
|
|
return tags[:k]
|
|
|
|
|
|
class PromptUtilitiesSampleTagsWithWeight(PromptUtilitiesSampleTags):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"tags": ("STRING", {"default": "", "multiline": True}),
|
|
"tags_delimiter": (("new line", ","),),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"max_k": ("INT", {"default": 1, "min": 1, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"min_k": ("INT", {"default": 1, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"max_w": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": -100, "max": 100, "step": 0.01},
|
|
),
|
|
"min_w": (
|
|
"FLOAT",
|
|
{"default": 0.8, "min": -100, "max": 100, "step": 0.01},
|
|
),
|
|
"step_w": (
|
|
"FLOAT",
|
|
{"default": 0.1, "min": -100, "max": 100, "step": 0.01},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "sample_tags"
|
|
|
|
def sample_tags(
|
|
self, tags, tags_delimiter, seed, max_k, min_k, max_w, min_w, step_w
|
|
):
|
|
sampled_tags = self._sample_tags(tags, tags_delimiter, seed, max_k, min_k)
|
|
random_gen = np.random.default_rng(seed)
|
|
weights = random_gen.choice(
|
|
np.arange(min_w, max_w, step_w),
|
|
size=len(sampled_tags),
|
|
)
|
|
sampled_tags = [f"({t}:{round(w, 2)})" for t, w in zip(sampled_tags, weights)]
|
|
return (", ".join(sampled_tags),)
|