55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
from pathlib import Path
|
|
|
|
import folder_paths
|
|
|
|
|
|
class EmbeddingPicker:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
embeddings = folder_paths.get_filename_list("embeddings")
|
|
|
|
return {
|
|
"required": {
|
|
"prompts": ("STRING", {"forceInput": True}),
|
|
"embedding": ((embeddings),),
|
|
"emphasis": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.05,
|
|
"max": 3.0,
|
|
"step": 0.05,
|
|
},
|
|
),
|
|
"append": (
|
|
"BOOLEAN",
|
|
{"default": False, "label_on": "true ", "label_off": "false "},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("prompts",)
|
|
FUNCTION = "concat_embedding"
|
|
OUTPUT_NODE = False
|
|
|
|
CATEGORY = "utils"
|
|
|
|
def concat_embedding(self, prompts, embedding, emphasis, append):
|
|
emb = "embedding:" + Path(embedding).stem
|
|
|
|
emphasis = f"{emphasis:.3f}"
|
|
if emphasis != "1.000":
|
|
emb = f"({emb}:{emphasis})"
|
|
|
|
output = f"{prompts}, {emb}" if append else f"{emb}, {prompts}"
|
|
|
|
return (output,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"EmbeddingPicker": EmbeddingPicker}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"EmbeddingPicker": "Embedding Picker"}
|