diff --git a/.gitignore b/.gitignore index 35ea65a..57f9989 100644 --- a/.gitignore +++ b/.gitignore @@ -161,4 +161,4 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +/.idea diff --git a/__init__.py b/__init__.py index cc4553f..64620b7 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,10 @@ -from .node import EmbeddingPicker +from pathlib import Path +from .node import comfy_entrypoint -NODE_CLASS_MAPPINGS = {"EmbeddingPicker": EmbeddingPicker} -NODE_DISPLAY_NAME_MAPPINGS = {"EmbeddingPicker": "Embedding Picker"} +WEB_DIRECTORY = "./js" +__all__ = ["WEB_DIRECTORY", "comfy_entrypoint"] # ------------------------ Remove copied web extension ------------------------- -from pathlib import Path - node_dir = Path(__file__).resolve().parent comfy_dir = node_dir.parent.parent destination_dir = comfy_dir / "web" / "extensions" / "tropfchen" @@ -13,7 +12,4 @@ destination_path = destination_dir / "epQuickNodes.js" if destination_path.exists(): destination_path.unlink() - # --------------------------- Install web extension ---------------------------- -WEB_DIRECTORY = "./js" -__all__ = ["WEB_DIRECTORY"] diff --git a/node.py b/node.py index 6f917e5..cd852f5 100644 --- a/node.py +++ b/node.py @@ -1,53 +1,57 @@ from pathlib import Path - import folder_paths +from comfy_api.latest import io, ComfyExtension, ui -class EmbeddingPicker: - def __init__(self): - pass - +class EmbeddingPicker(io.ComfyNode): @classmethod - def INPUT_TYPES(self): + def define_schema(cls) -> io.Schema: embeddings = folder_paths.get_filename_list("embeddings") - return { - "required": { - "embedding": ((embeddings),), - "emphasis": ( - "FLOAT", - { - "default": 1.0, - "min": 0.0, - "max": 3.0, - "step": 0.05, - }, + return io.Schema( + node_id="EmbeddingPicker", + display_name="Embedding Picker", + category="utils", + inputs=[ + io.Combo.Input("embedding", options=embeddings), + + io.Float.Input( + "emphasis", + default=1.0, + min=0.0, + max=3.0, + step=0.05 ), - "append": ( - "BOOLEAN", - {"default": False, "label_on": "true ", "label_off": "false "}, - ), - "text": ("STRING", {"multiline": True}), - }, - } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("text",) - FUNCTION = "concat_embedding" - OUTPUT_NODE = False + io.Boolean.Input("append", default=False), - CATEGORY = "utils" + io.String.Input("text", multiline=True) + ], + outputs=[ + io.String.Output("text") + ] + ) - def concat_embedding(self, text, embedding, emphasis, append): + @classmethod + def execute(cls, text, embedding, emphasis, append) -> io.NodeOutput: if emphasis < 0.05: - return (text,) + return io.NodeOutput(text) emb = "embedding:" + Path(embedding).stem - emphasis = f"{emphasis:.3f}" - if emphasis != "1.000": - emb = f"({emb}:{emphasis})" + emphasis_str = f"{emphasis:.3f}" + if emphasis_str != "1.000": + emb = f"({emb}:{emphasis_str})" output = f"{text}, {emb}" if append else f"{emb}, {text}" - return (output,) + return io.NodeOutput(output) + + +class EmbeddingPickerExtension(ComfyExtension): + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [EmbeddingPicker] + + +async def comfy_entrypoint() -> ComfyExtension: + return EmbeddingPickerExtension()