From 94dac21d40974f5a7714700f3406afa19fa89705 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Mon, 30 Dec 2024 11:18:42 -0800
Subject: [PATCH] Add comfy_types for node development.

---
 README.md                                      |  6 +++++-
 .../src/{{cookiecutter.project_slug}}/nodes.py | 18 ++++++++++--------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 1182d9d..2f8cd4b 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,10 @@
 # cookiecutter-comfy-extension
 
-## A cookiecutter template for writing a ComfyUI custom node extension
+## Get started writing custom nodes in one step without setting up a ton of Python project config.
+
+An opinionated way to develop ComfyUI custom nodes. Uses cookiecutter to scaffold a python project.
+
+This template helps you 
 
 ## Usage
 
diff --git a/{{cookiecutter.project_slug}}/src/{{cookiecutter.project_slug}}/nodes.py b/{{cookiecutter.project_slug}}/src/{{cookiecutter.project_slug}}/nodes.py
index 5a18093..4572787 100644
--- a/{{cookiecutter.project_slug}}/src/{{cookiecutter.project_slug}}/nodes.py
+++ b/{{cookiecutter.project_slug}}/src/{{cookiecutter.project_slug}}/nodes.py
@@ -1,4 +1,6 @@
-class Example:
+from comfy_types import IO, ComfyNodeABC, InputTypeDict
+from inspect import cleandoc
+class ExampleNode(ComfyNodeABC):
     """
     A example node
 
@@ -31,7 +33,7 @@ def __init__(self):
         pass
 
     @classmethod
-    def INPUT_TYPES(s):
+    def INPUT_TYPES(s) -> InputTypeDict:
         """
             Return a dictionary which contains config for all input fields.
             Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
@@ -48,15 +50,15 @@ def INPUT_TYPES(s):
         """
         return {
             "required": {
-                "image": ("IMAGE", { "tooltip": "This is an image"}),
-                "int_field": ("INT", {
+                "image": (IO.IMAGE, { "tooltip": "This is an image"}),
+                "int_field": (IO.INT, {
                     "default": 0,
                     "min": 0, #Minimum value
                     "max": 4096, #Maximum value
                     "step": 64, #Slider's step
                     "display": "number" # Cosmetic only: display as "number" or "slider"
                 }),
-                "float_field": ("FLOAT", {
+                "float_field": (IO.FLOAT, {
                     "default": 1.0,
                     "min": 0.0,
                     "max": 10.0,
@@ -64,16 +66,16 @@ def INPUT_TYPES(s):
                     "round": 0.001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
                     "display": "number"}),
                 "print_to_screen": (["enable", "disable"],),
-                "string_field": ("STRING", {
+                "string_field": (IO.STRING, {
                     "multiline": False, #True if you want the field to look like the one on the ClipTextEncode node
                     "default": "Hello World!"
                 }),
             },
         }
 
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = (IO.IMAGE,)
     #RETURN_NAMES = ("image_output_name",)
-    DESCRIPTION = "This is an example node"
+    DESCRIPTION = cleandoc(__doc__)
     FUNCTION = "test"
 
     #OUTPUT_NODE = False