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

from typing import Union
from torch import Tensor

from .documentation import short_desc, register_description, coll, DocHelper
from .motion_module_ad import PerBlock, PerBlockId, BlockType, AllPerBlocks
from .utils_model import ModelTypeSD
from .utils_motion import extend_list_to_batch_size


class ADBlockHolder:
    def __init__(self, effect: Union[float, Tensor, None]=None,
                 scales: Union[list[float, Tensor], None]=list()):
        self.effect = effect
        self.scales = scales

    def has_effect(self):
        return self.effect is not None

    def has_scale(self):
        for scale in self.scales:
            if scale is not None:
                return True
        return False

    def is_empty(self):
        has_anything = self.has_effect() or self.has_scale()
        return not has_anything


class ADBlockComboNode:
    NodeID = 'ADE_ADBlockCombo'
    NodeName = 'AD Block 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "effect": ("MULTIVAL",),
                "scale": ("MULTIVAL",),
            },
            "hidden": {
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("AD_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "block_control"

    def block_control(self, effect: Union[float, Tensor, None]=None, scale: Union[float, Tensor, None]=None):
        scales = [scale, scale]
        block = ADBlockHolder(effect=effect, scales=scales)
        if block.is_empty():
            block = None
        return (block,)


class ADBlockIndivNode:
    NodeID = 'ADE_ADBlockIndiv'
    NodeName = 'AD Block+ 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "effect": ("MULTIVAL",),
                "scale_0": ("MULTIVAL",),
                "scale_1": ("MULTIVAL",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("AD_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "block_control"

    def block_control(self, effect: Union[float, Tensor, None]=None,
                      scale_0: Union[float, Tensor, None]=None, scale_1: Union[float, Tensor, None]=None):
        scales = [scale_0, scale_1]
        block = ADBlockHolder(effect=effect, scales=scales)
        if block.is_empty():
            block = None
        return (block,)


class PerBlockHighLevelNode:
    NodeID = 'ADE_PerBlockHighLevel'
    NodeName = 'AD Per Block 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "down": ("AD_BLOCK",),
                "mid": ("AD_BLOCK",),
                "up": ("AD_BLOCK",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("PER_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "create_per_block"

    def create_per_block(self,
                         down: Union[ADBlockHolder, None]=None,
                         mid: Union[ADBlockHolder, None]=None,
                         up: Union[ADBlockHolder, None]=None):
        blocks = []
        d = {
            PerBlockId(block_type=BlockType.DOWN): down,
            PerBlockId(block_type=BlockType.MID): mid,
            PerBlockId(block_type=BlockType.UP): up,
        }
        for id, block in d.items():
            if block is not None:
                blocks.append(PerBlock(id=id, effect=block.effect, scales=block.scales))
        if len(blocks) == 0:
            return (None,)
        return (AllPerBlocks(blocks),)


class PerBlock_SD15_MidLevelNode:
    NodeID = 'ADE_PerBlock_SD15_MidLevel'
    NodeName = 'AD Per Block+ (SD1.5) 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "down_0": ("AD_BLOCK",),
                "down_1": ("AD_BLOCK",),
                "down_2": ("AD_BLOCK",),
                "down_3": ("AD_BLOCK",),
                "mid": ("AD_BLOCK",),
                "up_0": ("AD_BLOCK",),
                "up_1": ("AD_BLOCK",),
                "up_2": ("AD_BLOCK",),
                "up_3": ("AD_BLOCK",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("PER_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "create_per_block"

    def create_per_block(self,
                         down_0: Union[ADBlockHolder, None]=None,
                         down_1: Union[ADBlockHolder, None]=None,
                         down_2: Union[ADBlockHolder, None]=None,
                         down_3: Union[ADBlockHolder, None]=None,
                         mid: Union[ADBlockHolder, None]=None,
                         up_0: Union[ADBlockHolder, None]=None,
                         up_1: Union[ADBlockHolder, None]=None,
                         up_2: Union[ADBlockHolder, None]=None,
                         up_3: Union[ADBlockHolder, None]=None):
        blocks = []
        d = {
            PerBlockId(block_type=BlockType.DOWN, block_idx=0): down_0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=1): down_1,
            PerBlockId(block_type=BlockType.DOWN, block_idx=2): down_2,
            PerBlockId(block_type=BlockType.DOWN, block_idx=3): down_3,
            PerBlockId(block_type=BlockType.MID): mid,
            PerBlockId(block_type=BlockType.UP, block_idx=0): up_0,
            PerBlockId(block_type=BlockType.UP, block_idx=1): up_1,
            PerBlockId(block_type=BlockType.UP, block_idx=2): up_2,
            PerBlockId(block_type=BlockType.UP, block_idx=3): up_3,
        }
        for id, block in d.items():
            if block is not None:
                blocks.append(PerBlock(id=id, effect=block.effect, scales=block.scales))
        if len(blocks) == 0:
            return (None,)
        return (AllPerBlocks(blocks, ModelTypeSD.SD1_5),)


class PerBlock_SD15_LowLevelNode:
    NodeID = 'ADE_PerBlock_SD15_LowLevel'
    NodeName = 'AD Per Block++ (SD1.5) 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "down_0__0": ("AD_BLOCK",),
                "down_0__1": ("AD_BLOCK",),
                "down_1__0": ("AD_BLOCK",),
                "down_1__1": ("AD_BLOCK",),
                "down_2__0": ("AD_BLOCK",),
                "down_2__1": ("AD_BLOCK",),
                "down_3__0": ("AD_BLOCK",),
                "down_3__1": ("AD_BLOCK",),
                "mid": ("AD_BLOCK",),
                "up_0__0": ("AD_BLOCK",),
                "up_0__1": ("AD_BLOCK",),
                "up_0__2": ("AD_BLOCK",),
                "up_1__0": ("AD_BLOCK",),
                "up_1__1": ("AD_BLOCK",),
                "up_1__2": ("AD_BLOCK",),
                "up_2__0": ("AD_BLOCK",),
                "up_2__1": ("AD_BLOCK",),
                "up_2__2": ("AD_BLOCK",),
                "up_3__0": ("AD_BLOCK",),
                "up_3__1": ("AD_BLOCK",),
                "up_3__2": ("AD_BLOCK",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("PER_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "create_per_block"

    def create_per_block(self,
                         down_0__0: Union[ADBlockHolder, None]=None,
                         down_0__1: Union[ADBlockHolder, None]=None,
                         down_1__0: Union[ADBlockHolder, None]=None,
                         down_1__1: Union[ADBlockHolder, None]=None,
                         down_2__0: Union[ADBlockHolder, None]=None,
                         down_2__1: Union[ADBlockHolder, None]=None,
                         down_3__0: Union[ADBlockHolder, None]=None,
                         down_3__1: Union[ADBlockHolder, None]=None,
                         mid: Union[ADBlockHolder, None]=None,
                         up_0__0: Union[ADBlockHolder, None]=None,
                         up_0__1: Union[ADBlockHolder, None]=None,
                         up_0__2: Union[ADBlockHolder, None]=None,
                         up_1__0: Union[ADBlockHolder, None]=None,
                         up_1__1: Union[ADBlockHolder, None]=None,
                         up_1__2: Union[ADBlockHolder, None]=None,
                         up_2__0: Union[ADBlockHolder, None]=None,
                         up_2__1: Union[ADBlockHolder, None]=None,
                         up_2__2: Union[ADBlockHolder, None]=None,
                         up_3__0: Union[ADBlockHolder, None]=None,
                         up_3__1: Union[ADBlockHolder, None]=None,
                         up_3__2: Union[ADBlockHolder, None]=None):
        blocks = []
        d = {
            PerBlockId(block_type=BlockType.DOWN, block_idx=0, module_idx=0): down_0__0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=0, module_idx=1): down_0__1,
            PerBlockId(block_type=BlockType.DOWN, block_idx=1, module_idx=0): down_1__0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=1, module_idx=1): down_1__1,
            PerBlockId(block_type=BlockType.DOWN, block_idx=2, module_idx=0): down_2__0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=2, module_idx=1): down_2__1,
            PerBlockId(block_type=BlockType.DOWN, block_idx=3, module_idx=0): down_3__0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=3, module_idx=1): down_3__1,
            PerBlockId(block_type=BlockType.MID): mid,
            PerBlockId(block_type=BlockType.UP, block_idx=0, module_idx=0): up_0__0,
            PerBlockId(block_type=BlockType.UP, block_idx=0, module_idx=1): up_0__1,
            PerBlockId(block_type=BlockType.UP, block_idx=0, module_idx=2): up_0__2,
            PerBlockId(block_type=BlockType.UP, block_idx=1, module_idx=0): up_1__0,
            PerBlockId(block_type=BlockType.UP, block_idx=1, module_idx=1): up_1__1,
            PerBlockId(block_type=BlockType.UP, block_idx=1, module_idx=2): up_1__2,
            PerBlockId(block_type=BlockType.UP, block_idx=2, module_idx=0): up_2__0,
            PerBlockId(block_type=BlockType.UP, block_idx=2, module_idx=1): up_2__1,
            PerBlockId(block_type=BlockType.UP, block_idx=2, module_idx=2): up_2__2,
            PerBlockId(block_type=BlockType.UP, block_idx=3, module_idx=0): up_3__0,
            PerBlockId(block_type=BlockType.UP, block_idx=3, module_idx=1): up_3__1,
            PerBlockId(block_type=BlockType.UP, block_idx=3, module_idx=2): up_3__2,
        }
        for id, block in d.items():
            if block is not None:
                blocks.append(PerBlock(id=id, effect=block.effect, scales=block.scales))
        if len(blocks) == 0:
            return (None,)
        return (AllPerBlocks(blocks, ModelTypeSD.SD1_5),)


class PerBlock_SD15_FromFloatsNode:
    NodeID = 'ADE_PerBlock_SD15_FromFloats'
    NodeName = 'AD Per Block Floats (SD1.5) 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "effect_21_floats": ("FLOATS",),
                "scale_21_floats": ("FLOATS",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("PER_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "create_per_block"

    Desc = [
        short_desc('Use Floats from Value Schedules to select SD1.5 effect/scale values for blocks.'),
        'SD1.5 Motion Modules contain 21 blocks:',
        'idx 0 - start of down blocks (down_0__0)',
        'idx 7 - end of down blocks   (down_3__1)',
        'idx 8 - mid block            (mid)',
        'idx 9 - start of up blocks   (up_0__0)',
        'idx 20 - end of up blocks    (up_3__2)',
    ]
    register_description(NodeID, Desc)

    def create_per_block(self,
                         effect_21_floats: Union[list[float], None]=None,
                         scale_21_floats: Union[list[float], None]=None):
        if effect_21_floats is None and scale_21_floats is None:
            return (None,)
        # SD1.5 has 21 blocks
        block_total = 21
        holders = [ADBlockHolder() for _ in range(block_total)]
        if effect_21_floats is not None:
            effect_21_floats = extend_list_to_batch_size(effect_21_floats, block_total)
            for effect, holder in zip(effect_21_floats, holders):
                holder.effect = effect
        if scale_21_floats is not None:
            scale_21_floats = extend_list_to_batch_size(scale_21_floats, block_total)
            for scale, holder in zip(scale_21_floats, holders):
                holder.scales = [scale, scale]
        return PerBlock_SD15_LowLevelNode.create_per_block(self, *holders)


class PerBlock_SDXL_MidLevelNode:
    NodeID = 'ADE_PerBlock_SDXL_MidLevel'
    NodeName = 'AD Per Block+ (SDXL) 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "down_0": ("AD_BLOCK",),
                "down_1": ("AD_BLOCK",),
                "down_2": ("AD_BLOCK",),
                "mid": ("AD_BLOCK",),
                "up_0": ("AD_BLOCK",),
                "up_1": ("AD_BLOCK",),
                "up_2": ("AD_BLOCK",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("PER_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "create_per_block"

    def create_per_block(self,
                         down_0: Union[ADBlockHolder, None]=None,
                         down_1: Union[ADBlockHolder, None]=None,
                         down_2: Union[ADBlockHolder, None]=None,
                         mid: Union[ADBlockHolder, None]=None,
                         up_0: Union[ADBlockHolder, None]=None,
                         up_1: Union[ADBlockHolder, None]=None,
                         up_2: Union[ADBlockHolder, None]=None):
        blocks = []
        d = {
            PerBlockId(block_type=BlockType.DOWN, block_idx=0): down_0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=1): down_1,
            PerBlockId(block_type=BlockType.DOWN, block_idx=2): down_2,
            PerBlockId(block_type=BlockType.MID): mid,
            PerBlockId(block_type=BlockType.UP, block_idx=0): up_0,
            PerBlockId(block_type=BlockType.UP, block_idx=1): up_1,
            PerBlockId(block_type=BlockType.UP, block_idx=2): up_2,
        }
        for id, block in d.items():
            if block is not None:
                blocks.append(PerBlock(id=id, effect=block.effect, scales=block.scales))
        if len(blocks) == 0:
            return (None,)
        return (AllPerBlocks(blocks, ModelTypeSD.SDXL),)


class PerBlock_SDXL_LowLevelNode:
    NodeID = 'ADE_PerBlock_SDXL_LowLevel'
    NodeName = 'AD Per Block++ (SDXL) 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "down_0__0": ("AD_BLOCK",),
                "down_0__1": ("AD_BLOCK",),
                "down_1__0": ("AD_BLOCK",),
                "down_1__1": ("AD_BLOCK",),
                "down_2__0": ("AD_BLOCK",),
                "down_2__1": ("AD_BLOCK",),
                "mid": ("AD_BLOCK",),
                "up_0__0": ("AD_BLOCK",),
                "up_0__1": ("AD_BLOCK",),
                "up_0__2": ("AD_BLOCK",),
                "up_1__0": ("AD_BLOCK",),
                "up_1__1": ("AD_BLOCK",),
                "up_1__2": ("AD_BLOCK",),
                "up_2__0": ("AD_BLOCK",),
                "up_2__1": ("AD_BLOCK",),
                "up_2__2": ("AD_BLOCK",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("PER_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "create_per_block"

    def create_per_block(self,
                         down_0__0: Union[ADBlockHolder, None]=None,
                         down_0__1: Union[ADBlockHolder, None]=None,
                         down_1__0: Union[ADBlockHolder, None]=None,
                         down_1__1: Union[ADBlockHolder, None]=None,
                         down_2__0: Union[ADBlockHolder, None]=None,
                         down_2__1: Union[ADBlockHolder, None]=None,
                         mid: Union[ADBlockHolder, None]=None,
                         up_0__0: Union[ADBlockHolder, None]=None,
                         up_0__1: Union[ADBlockHolder, None]=None,
                         up_0__2: Union[ADBlockHolder, None]=None,
                         up_1__0: Union[ADBlockHolder, None]=None,
                         up_1__1: Union[ADBlockHolder, None]=None,
                         up_1__2: Union[ADBlockHolder, None]=None,
                         up_2__0: Union[ADBlockHolder, None]=None,
                         up_2__1: Union[ADBlockHolder, None]=None,
                         up_2__2: Union[ADBlockHolder, None]=None,):
        blocks = []
        d = {
            PerBlockId(block_type=BlockType.DOWN, block_idx=0, module_idx=0): down_0__0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=0, module_idx=1): down_0__1,
            PerBlockId(block_type=BlockType.DOWN, block_idx=1, module_idx=0): down_1__0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=1, module_idx=1): down_1__1,
            PerBlockId(block_type=BlockType.DOWN, block_idx=2, module_idx=0): down_2__0,
            PerBlockId(block_type=BlockType.DOWN, block_idx=2, module_idx=1): down_2__1,
            PerBlockId(block_type=BlockType.MID): mid,
            PerBlockId(block_type=BlockType.UP, block_idx=0, module_idx=0): up_0__0,
            PerBlockId(block_type=BlockType.UP, block_idx=0, module_idx=1): up_0__1,
            PerBlockId(block_type=BlockType.UP, block_idx=0, module_idx=2): up_0__2,
            PerBlockId(block_type=BlockType.UP, block_idx=1, module_idx=0): up_1__0,
            PerBlockId(block_type=BlockType.UP, block_idx=1, module_idx=1): up_1__1,
            PerBlockId(block_type=BlockType.UP, block_idx=1, module_idx=2): up_1__2,
            PerBlockId(block_type=BlockType.UP, block_idx=2, module_idx=0): up_2__0,
            PerBlockId(block_type=BlockType.UP, block_idx=2, module_idx=1): up_2__1,
            PerBlockId(block_type=BlockType.UP, block_idx=2, module_idx=2): up_2__2,
        }
        for id, block in d.items():
            if block is not None:
                blocks.append(PerBlock(id=id, effect=block.effect, scales=block.scales))
        if len(blocks) == 0:
            return (None,)
        return (AllPerBlocks(blocks, ModelTypeSD.SDXL),)


class PerBlock_SDXL_FromFloatsNode:
    NodeID = 'ADE_PerBlock_SDXL_FromFloats'
    NodeName = 'AD Per Block Floats (SDXL) 🎭🅐🅓'
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
            },
            "optional": {
                "effect_16_floats": ("FLOATS",),
                "scale_16_floats": ("FLOATS",),
                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
            }
        }
    
    RETURN_TYPES = ("PER_BLOCK",)
    CATEGORY = "Animate Diff 🎭🅐🅓/per block"
    FUNCTION = "create_per_block"

    Desc = [
        short_desc('Use Floats from Value Schedules to select SDXL effect/scale values for blocks.'),
        'SDXL Motion Modules contain 16 blocks:',
        'idx 0 - start of down blocks (down_0__0)',
        'idx 5 - end of down blocks   (down_2__1)',
        'idx 6 - mid block            (mid)',
        'idx 7 - start of up blocks   (up_0__0)',
        'idx 15 - end of up blocks    (up_2__2)',
    ]
    register_description(NodeID, Desc)

    def create_per_block(self,
                         effect_16_floats: Union[list[float], None]=None,
                         scale_16_floats: Union[list[float], None]=None):
        if effect_16_floats is None and scale_16_floats is None:
            return (None,)
        # SDXL has 16 blocks
        block_total = 16
        holders = [ADBlockHolder() for _ in range(block_total)]
        if effect_16_floats is not None:
            effect_16_floats = extend_list_to_batch_size(effect_16_floats, block_total)
            for effect, holder in zip(effect_16_floats, holders):
                holder.effect = effect
        if scale_16_floats is not None:
            scale_16_floats = extend_list_to_batch_size(scale_16_floats, block_total)
            for scale, holder in zip(scale_16_floats, holders):
                holder.scales = [scale, scale]
        return PerBlock_SDXL_LowLevelNode.create_per_block(self, *holders)