custom_nodes/ComfyUI-AnimateDiff-Evolved/animatediff/nodes_scheduling.py

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 {
            "required": {
                "prompts": ("STRING", {"multiline": True, "default": ''}),
                "clip": ("CLIP",),
                "latent": ("LATENT",),
            },
            "optional": {
                "prepend_text": ("STRING", {"multiline": True, "default": '', "forceInput": True}),
                "append_text": ("STRING", {"multiline": True, "default": '', "forceInput": True}),
                "values_replace": ("VALUES_REPLACE",),
                "print_schedule": ("BOOLEAN", {"default": False}),
                "tensor_interp": (TensorInterp._LIST,)
            },
        }

    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,
                                values_replace=values_replace, print_schedule=print_schedule)
        conditioning = evaluate_prompt_schedule(prompts, latent["samples"].size(0), clip, options)
        return (conditioning, latent)


class PromptSchedulingNode:
    NodeID = 'ADE_PromptScheduling'
    NodeName = 'Prompt Scheduling 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "prompts": ("STRING", {"multiline": True, "default": ''}),
                "clip": ("CLIP",),
            },
            "optional": {
                "prepend_text": ("STRING", {"multiline": True, "default": '', "forceInput": True}),
                "append_text": ("STRING", {"multiline": True, "default": '', "forceInput": True}),
                "values_replace": ("VALUES_REPLACE",),
                "print_schedule": ("BOOLEAN", {"default": False}),
                "max_length": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
                "tensor_interp": (TensorInterp._LIST,)
            },
        }
    
    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,
                                values_replace=values_replace, print_schedule=print_schedule)
        conditioning = evaluate_prompt_schedule(prompts, max_length, clip, options)
        return (conditioning,)


class ValueSchedulingLatentsNode:
    NodeID = 'ADE_ValueSchedulingLatents'
    NodeName = 'Value Scheduling [Latents] 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "values": ("STRING", {"multiline": True, "default": ""}),
                "latent": ("LATENT",),
            },
            "optional": {
                "print_schedule": ("BOOLEAN", {"default": False}),
            },
            "hidden": {
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("FLOAT", "FLOATS", "INT", "INTS")
    CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
    FUNCTION = "create_schedule"

    Desc = [
        short_desc('Create a list of values with automatic interpolation, 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]
        if print_schedule:
            logger.info(f"ValueScheduling ({len(float_vals)} values):")
            for i, val in enumerate(float_vals):
                logger.info(f"{i} = {val}")
        return (float_vals, float_vals, int_vals, int_vals)


class ValueSchedulingNode:
    NodeID = 'ADE_ValueScheduling'
    NodeName = 'Value Scheduling 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "values": ("STRING", {"multiline": True, "default": ""}),
            },
            "optional": {
                "print_schedule": ("BOOLEAN", {"default": False}),
                "max_length": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
            },
            "hidden": {
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }

    RETURN_TYPES = ("FLOAT", "FLOATS", "INT", "INTS")
    CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
    FUNCTION = "create_schedule"

    Desc = [
        short_desc('Create a list of values with automatic interpolation.'),
        {'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]
        if print_schedule:
            logger.info(f"ValueScheduling ({len(float_vals)} values):")
            for i, val in enumerate(float_vals):
                logger.info(f"{i} = {val}")
        return (float_vals, float_vals, int_vals, int_vals)


class AddValuesReplaceNode:
    NodeID = 'ADE_ValuesReplace'
    NodeName = 'Add Values Replace 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "value_key": ("STRING", {"default": ""}),
                "floats": ("FLOATS",)
            },
            "optional": {
                "prev_replace": ("VALUES_REPLACE",),
            },
            "hidden": {
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("VALUES_REPLACE",)
    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)
        # add/replace value floats
        if prev_replace is None:
            prev_replace = {}
        prev_replace = prev_replace.copy()
        if value_key in prev_replace:
            logger.warn(f"Value key '{value_key}' is already present - corresponding floats value will be overriden.")
        prev_replace[value_key] = floats
        return (prev_replace,)


class FloatToFloatsNode:
    NodeID = 'ADE_FloatToFloats'
    NodeName = 'Float to Floats 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "FLOAT": ("FLOAT", {"default": 39, "forceInput": True}),
            },
            "hidden": {
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("FLOATS",)
    CATEGORY = "Animate Diff 🎭🅐🅓/scheduling"
    FUNCTION = "convert_to_floats"

    def convert_to_floats(self, FLOAT: Union[float, list[float]]):
        floats = None
        if isinstance(FLOAT, float):
            floats = [float(FLOAT)]
        else:
            floats = list(FLOAT)
        return (floats,)