From 29d41cd55300c4784c85817b906f2c826c59db5d Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Wed, 14 Aug 2024 23:54:16 -0500
Subject: [PATCH] Add descriptions to scheduling nodes

---
 __init__.py                     |   2 +
 animatediff/documentation.py    |  75 +++++++++++++++++++++++
 animatediff/nodes.py            |  24 ++++----
 animatediff/nodes_scheduling.py | 105 ++++++++++++++++++++++++++++++--
 4 files changed, 190 insertions(+), 16 deletions(-)
 create mode 100644 animatediff/documentation.py

diff --git a/__init__.py b/__init__.py
index a5635e3..b939327 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,9 +2,11 @@ import folder_paths
 from .animatediff.logger import logger
 from .animatediff.utils_model import get_available_motion_models, Folders
 from .animatediff.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+from .animatediff import documentation
 
 if len(get_available_motion_models()) == 0:
     logger.error(f"No motion models found. Please download one and place in: {folder_paths.get_folder_paths(Folders.ANIMATEDIFF_MODELS)}")
 
 WEB_DIRECTORY = "./web"
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
+documentation.format_descriptions(NODE_CLASS_MAPPINGS)
diff --git a/animatediff/documentation.py b/animatediff/documentation.py
new file mode 100644
index 0000000..fa345e9
--- /dev/null
+++ b/animatediff/documentation.py
@@ -0,0 +1,75 @@
+from typing import Union
+
+from .logger import logger
+
+def image(src):
+    return f'<img src={src} style="width: 0px; min-width: 100%">'
+def video(src):
+    return f'<video src={src} autoplay muted loop controls controlslist="nodownload noremoteplayback noplaybackrate" style="width: 0px; min-width: 100%" class="VHS_loopedvideo">'
+def short_desc(desc):
+    return  f'<div id=VHS_shortdesc style="font-size: .8em">{desc}</div>'
+
+def coll(text: str):
+    return f"{text}_collapsed"
+
+descriptions = {
+}
+
+sizes = ['1.4','1.2','1']
+def as_html(entry, depth=0):
+    if isinstance(entry, dict):
+        size = 0.8 if depth < 2 else 1
+        html = ''
+        for k in entry:
+            if k == "collapsed":
+                continue
+            collapse_single = k.endswith("_collapsed")
+            if collapse_single:
+                name = k[:-len("_collapsed")]
+            else:
+                name = k
+            collapse_flag = ' VHS_precollapse' if entry.get("collapsed", False) or collapse_single else ''
+            html += f'<div vhs_title=\"{name}\" style=\"display: flex; font-size: {size}em\" class=\"VHS_collapse{collapse_flag}\"><div style=\"color: #AAA; height: 1.5em;\">[<span style=\"font-family: monospace\">-</span>]</div><div style=\"width: 100%\">{name}: {as_html(entry[k], depth=depth+1)}</div></div>'
+        return html
+    if isinstance(entry, list):
+        html = ''
+        for i in entry:
+            html += f'<div>{as_html(i, depth=depth)}</div>'
+        return html
+    return str(entry)
+
+
+def register_description(node_id: str, desc: Union[list, dict]):
+    descriptions[node_id] = desc
+
+
+def format_descriptions(nodes):
+    for k in descriptions:
+        if k.endswith("_collapsed"):
+            k = k[:-len("_collapsed")]
+        nodes[k].DESCRIPTION = as_html(descriptions[k])
+    # undocumented_nodes = []
+    # for k in nodes:
+    #     if not hasattr(nodes[k], "DESCRIPTION"):
+    #         undocumented_nodes.append(k)
+    # if len(undocumented_nodes) > 0:
+    #     logger.info(f"Undocumented nodes: {undocumented_nodes}")
+
+
+class DocHelper:
+    def __init__(self):
+        self.actual_dict = {}
+    
+    def add(self, add_dict):
+        self.actual_dict.update(add_dict)
+        return self
+
+    def get(self):
+        return self.actual_dict
+    
+    @staticmethod
+    def combine(*args):
+        docs = DocHelper()
+        for doc in args:
+            docs.add(doc)
+        return docs.get()
diff --git a/animatediff/nodes.py b/animatediff/nodes.py
index dc62396..43e718b 100644
--- a/animatediff/nodes.py
+++ b/animatediff/nodes.py
@@ -156,12 +156,12 @@ NODE_CLASS_MAPPINGS = {
     "ADE_NoisedImageInjection": NoisedImageInjectionNode,
     "ADE_NoisedImageInjectOptions": NoisedImageInjectOptionsNode,
     # Scheduling
-    "ADE_PromptScheduling": PromptSchedulingNode,
-    "ADE_PromptSchedulingLatents": PromptSchedulingLatentsNode,
-    "ADE_ValueScheduling": ValueSchedulingNode,
-    "ADE_ValueSchedulingLatents": ValueSchedulingLatentsNode,
-    "ADE_ValuesReplace": AddValuesReplaceNode,
-    "ADE_FloatToFloats": FloatToFloatsNode,
+    PromptSchedulingNode.NodeID: PromptSchedulingNode,
+    PromptSchedulingLatentsNode.NodeID: PromptSchedulingLatentsNode,
+    ValueSchedulingNode.NodeID: ValueSchedulingNode,
+    ValueSchedulingLatentsNode.NodeID: ValueSchedulingLatentsNode,
+    AddValuesReplaceNode.NodeID: AddValuesReplaceNode,
+    FloatToFloatsNode.NodeID: FloatToFloatsNode,
     # Extras Nodes
     "ADE_AnimateDiffUnload": AnimateDiffUnload,
     "ADE_EmptyLatentImageLarge": EmptyLatentImageLarge,
@@ -313,12 +313,12 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ADE_NoisedImageInjection": "Image Injection 🎭🅐🅓",
     "ADE_NoisedImageInjectOptions": "Image Injection Options 🎭🅐🅓",
     # Scheduling
-    "ADE_PromptScheduling": "Prompt Scheduling 🎭🅐🅓",
-    "ADE_PromptSchedulingLatents": "Prompt Scheduling [Latents] 🎭🅐🅓",
-    "ADE_ValueScheduling": "Value Scheduling 🎭🅐🅓",
-    "ADE_ValueSchedulingLatents": "Value Scheduling [Latents] 🎭🅐🅓",
-    "ADE_ValuesReplace": "Add Values Replace 🎭🅐🅓",
-    "ADE_FloatToFloats": "Float to Floats 🎭🅐🅓",
+    PromptSchedulingNode.NodeID: PromptSchedulingNode.NodeName,
+    PromptSchedulingLatentsNode.NodeID: PromptSchedulingLatentsNode.NodeName,
+    ValueSchedulingNode.NodeID: ValueSchedulingNode.NodeName,
+    ValueSchedulingLatentsNode.NodeID: ValueSchedulingLatentsNode.NodeName,
+    AddValuesReplaceNode.NodeID: AddValuesReplaceNode.NodeName,
+    FloatToFloatsNode.NodeID:FloatToFloatsNode.NodeName,
     # Extras Nodes
     "ADE_AnimateDiffUnload": "AnimateDiff Unload 🎭🅐🅓",
     "ADE_EmptyLatentImageLarge": "Empty Latent Image (Big Batch) 🎭🅐🅓",
diff --git a/animatediff/nodes_scheduling.py b/animatediff/nodes_scheduling.py
index 7aaafba..660b09f 100644
--- a/animatediff/nodes_scheduling.py
+++ b/animatediff/nodes_scheduling.py
@@ -1,12 +1,65 @@
 from typing import Union
 
+from .documentation import register_description, short_desc, coll, DocHelper
 from .scheduling import (evaluate_prompt_schedule, evaluate_value_schedule, TensorInterp, PromptOptions,
                          verify_key_value)
 from .utils_model import BIGMAX
 from .logger import logger
 
 
+desc_values = {coll('values'): 'Write your values here.'}
+desc_prompts = {coll('prompts'): 'Write your prompts here.'}
+desc_clip = {'clip': 'CLIP to use for encoding prompts.'}
+desc_latent = {'latent': 'Used to get the amount of frames (max_length) to use for scheduling.'}
+
+desc_prepend_text = {'prepend_text': 'OPTIONAL, adds text before all prompts.'}
+desc_append_text = {'append_text': 'OPTIONAL, adds text after all prompts.'}
+desc_values_replace = {'values_replace': 'OPTIONAL, replaces keys from value_replace keys with provided value schedules. Keys in the prompt are written as `some_key`, surrounded by the ` characters.'}
+desc_tensor_interp = {'tensor_interp': 'Selects method of interpolating prompt conds - defaults to lerp.'}
+desc_print_schedule = {'print_schedule': 'When True, prints output values for each frame.'}
+
+desc_max_length = {'max_length': 'Used to select the intended length of schedule. If set to 0, will use the largest index in the schedule as max_length, but will disable relative indexes (negative and decimal).'}
+desc_floats = {'floats': 'List of floats, likely outputted by a Value Scheduling node.'}
+desc_FLOAT = {'FLOAT': 'Float (or list of floats) to convert to FLOATS type.'}
+desc_value_key = {'value_key': 'Key to use for value schedule in Prompt Scheduling node. Can only contain a-z, A-Z, 0-9, and _ characters. In Prompt Scheduling, keys can be referred to as `some_key`, where the key is surrounded by ` characters.'}
+desc_prev_replace = {'prev_replace': 'OPTIONAL, other values_replace can be chained.'}
+
+desc_output_conditioning = {'CONDITIONING': 'Encoded prompts.'}
+desc_output_latent = {'LATENT': 'Unmodified input latents; can be used as pipe, or can be ignored.'}
+
+desc_format_allowed_idxs = {'allowed idxs':
+        {'single': 'A positive integer (e.g. 0, 2) schedules value for frame. A negative integer (e.g. -1, -5) schedules value for frame from the end (-1 would be the last frame). ' + 
+            'A decimal (e.g. 0.5, 1.0) selects frame based relative location in whole schedule (0.5 would be halfway, 1.0 would be last frame).',
+        'range': 'Using rules above, single:single chooses uninterpolated prompts from start idx (included) to end idx (excluded). Examples -> 0:12, 0:-5, 2:0.5',
+        'hold': 'Putting a colon after a single idx stops interpolation until the next provided index. Examples -> 0:, 0.5:, 16: '}
+    }
+
+desc_format_prompt = [
+    'Scheduling supports two formats: JSON and pythonic.',
+    {'JSON': ['"idx": "your prompt here", ...'],
+     'pythonic': ['idx = "your prompt here", ...']},
+    'The idx is the index of the frame - first frame is 0, last frame is max_frames-1. An idx may be the following:',
+    desc_format_allowed_idxs,
+    'The prompts themselves should be surrounded by double quotes ("your prompt here"). Portions of prompts can use value schedules provided values_replace.',
+    {'JSON': ['"0": "blue rock on mountain",', '"16": "green rock in lake"'],
+     'pythonic': ['0 = "blue rock on mountain",', '16 = "green rock in lake"']}
+]
+
+desc_format_values = [
+    'Scheduling supports two formats: JSON and pythonic.',
+    {'JSON': ['"idx": float/int_value, ...'],
+     'pythonic': ['idx = float/int_value, ...']},
+    'The idx is the index of the frame - first frame is 0, last frame is max_frames-1. An idx may be the following:',
+    desc_format_allowed_idxs,
+    'The values can be written without any special formatting.',
+    {'JSON': ['"0": 1.0,', '"16": 1.3'],
+     'pythonic': ['0 = 1.0,', '16 = 1.3']}
+]
+
+
 class PromptSchedulingLatentsNode:
+    NodeID = 'ADE_PromptSchedulingLatents'
+    NodeName = 'Prompt Scheduling [Latents] 🎭🅐🅓'
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -23,12 +76,19 @@ class PromptSchedulingLatentsNode:
                 "tensor_interp": (TensorInterp._LIST,)
             },
         }
-    
-    RETURN_TYPES = ("CONDITIONING", "LATENT",)
 
+    RETURN_TYPES = ("CONDITIONING", "LATENT",)
     CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
     FUNCTION = "create_schedule"
 
+    Desc = [
+        short_desc('Encode a schedule of prompts with automatic interpolation, its length matching passed-in latent count.'),
+        {'Format': desc_format_prompt},
+        {coll('Inputs'): DocHelper.combine(desc_prompts, desc_clip, desc_latent, desc_values_replace, desc_prepend_text, desc_append_text, desc_tensor_interp, desc_print_schedule)},
+        {coll('Outputs'): DocHelper.combine(desc_output_conditioning, desc_output_latent)}
+    ]
+    register_description(NodeID, Desc)
+
     def create_schedule(self, prompts: str, clip, latent: dict, print_schedule=False, tensor_interp=TensorInterp.LERP,
                         prepend_text='', append_text='', values_replace=None):
         options = PromptOptions(interp=tensor_interp, prepend_text=prepend_text, append_text=append_text,
@@ -38,6 +98,8 @@ class PromptSchedulingLatentsNode:
 
 
 class PromptSchedulingNode:
+    NodeID = 'ADE_PromptScheduling'
+    NodeName = 'Prompt Scheduling 🎭🅐🅓'
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -56,10 +118,17 @@ class PromptSchedulingNode:
         }
     
     RETURN_TYPES = ("CONDITIONING",)
-
     CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
     FUNCTION = "create_schedule"
 
+    Desc = [
+        short_desc('Encode a schedule of prompts with automatic interpolation.'),
+        {'Format': desc_format_prompt},
+        {coll('Inputs'): DocHelper.combine(desc_prompts, desc_clip, desc_values_replace, desc_prepend_text, desc_append_text, desc_max_length, desc_tensor_interp, desc_print_schedule)},
+        {coll('Outputs'): DocHelper.combine(desc_output_conditioning)}
+    ]
+    register_description(NodeID, Desc)
+
     def create_schedule(self, prompts: str, clip, print_schedule=False, max_length: int=0, tensor_interp=TensorInterp.LERP,
                         prepend_text='', append_text='', values_replace=None):
         options = PromptOptions(interp=tensor_interp, prepend_text=prepend_text, append_text=append_text,
@@ -69,6 +138,8 @@ class PromptSchedulingNode:
 
 
 class ValueSchedulingLatentsNode:
+    NodeID = 'ADE_ValueSchedulingLatents'
+    NodeName = 'Value Scheduling [Latents] 🎭🅐🅓'
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -86,6 +157,13 @@ class ValueSchedulingLatentsNode:
     CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
     FUNCTION = "create_schedule"
 
+    Desc = [
+        short_desc('Create a list of values, its length matching passed-in latent count.'),
+        {'Format': desc_format_values},
+        {coll('Inputs'): DocHelper.combine(desc_values, desc_latent, desc_print_schedule)},
+    ]
+    register_description(NodeID, Desc)
+
     def create_schedule(self, values: str, latent: dict, print_schedule=False):
         float_vals = evaluate_value_schedule(values, latent["samples"].size(0))
         int_vals = [round(x) for x in float_vals]
@@ -97,6 +175,8 @@ class ValueSchedulingLatentsNode:
 
 
 class ValueSchedulingNode:
+    NodeID = 'ADE_ValueScheduling'
+    NodeName = 'Value Scheduling 🎭🅐🅓'
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -114,6 +194,13 @@ class ValueSchedulingNode:
     CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
     FUNCTION = "create_schedule"
 
+    Desc = [
+        short_desc('Create a list of values, its length matching passed-in latent count.'),
+        {'Format': desc_format_values},
+        {coll('Inputs'): DocHelper.combine(desc_values, desc_max_length, desc_print_schedule)},
+    ]
+    register_description(NodeID, Desc)
+
     def create_schedule(self, values: str, max_length: int, print_schedule=False):
         float_vals = evaluate_value_schedule(values, max_length)
         int_vals = [round(x) for x in float_vals]
@@ -125,6 +212,8 @@ class ValueSchedulingNode:
 
 
 class AddValuesReplaceNode:
+    NodeID = 'ADE_ValuesReplace'
+    NodeName = 'Add Values Replace 🎭🅐🅓'
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -142,6 +231,12 @@ class AddValuesReplaceNode:
     CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
     FUNCTION = "add_values_replace"
 
+    Desc = [
+        short_desc('Add a values schedule bound to a key to be used in Prompt Scheduling node.'),
+        {'Inputs': DocHelper.combine(desc_value_key, desc_floats, desc_prev_replace)},
+    ]
+    register_description(NodeID, Desc)
+
     def add_values_replace(self, value_key: str, floats: Union[list[float]], prev_replace: dict=None):
         # key can only have a-z, A-Z, 0-9, and _ characters
         verify_key_value(key=value_key)
@@ -156,6 +251,8 @@ class AddValuesReplaceNode:
 
 
 class FloatToFloatsNode:
+    NodeID = 'ADE_FloatToFloats'
+    NodeName = 'Float to Floats 🎭🅐🅓'
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -177,4 +274,4 @@ class FloatToFloatsNode:
             floats = [float(FLOAT)]
         else:
             floats = list(FLOAT)
-        return (floats,)
+        return (floats,)
\ No newline at end of file