add files

2023-08-27 22:08:43 +02:00
parent ae0b0a419f
commit 74b97b9316
4 changed files with 86 additions and 0 deletions
@@ -0,0 +1,54 @@
+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"}