122 lines
5.5 KiB
Python
122 lines
5.5 KiB
Python
from collections.abc import Iterable
|
|
from typing import Union
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from comfy_api.latest import io
|
|
|
|
from .utils_motion import create_multival_combo, linear_conversion, normalize_min_max, extend_to_batch_size, extend_list_to_batch_size
|
|
|
|
|
|
class ScaleType:
|
|
ABSOLUTE = "absolute"
|
|
RELATIVE = "relative"
|
|
LIST = [ABSOLUTE, RELATIVE]
|
|
|
|
|
|
class MultivalDynamicNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ADE_MultivalDynamic',
|
|
display_name='Multival 🎭🅐🅓',
|
|
category='Animate Diff 🎭🅐🅓/multival',
|
|
inputs=[io.Float.Input('float_val', default=1.0, min=0.0, step=0.001), io.Mask.Input('mask_optional', optional=True)],
|
|
outputs=[io.Custom('MULTIVAL').Output('MULTIVAL')]
|
|
)
|
|
@classmethod
|
|
def execute(cls, float_val: Union[float, list[float]]=1.0, mask_optional: Tensor=None):
|
|
return io.NodeOutput(create_multival_combo(float_val=float_val, mask_optional=mask_optional))
|
|
|
|
|
|
class MultivalScaledMaskNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ADE_MultivalScaledMask',
|
|
display_name='Multival Scaled Mask 🎭🅐🅓',
|
|
category='Animate Diff 🎭🅐🅓/multival',
|
|
inputs=[io.Float.Input('min_float_val', default=0.0, min=0.0, step=0.001), io.Float.Input('max_float_val', default=1.0, min=0.0, step=0.001), io.Mask.Input('mask'), io.Combo.Input('scaling', options=['absolute', 'relative'] , optional=True)],
|
|
outputs=[io.Custom('MULTIVAL').Output('MULTIVAL')]
|
|
)
|
|
@classmethod
|
|
def execute(cls, min_float_val: float, max_float_val: float, mask: Tensor, scaling: str=ScaleType.ABSOLUTE):
|
|
lengths = [mask.shape[0]]
|
|
iterable_inputs = [False, False]
|
|
val_inputs = [min_float_val, max_float_val]
|
|
if isinstance(min_float_val, Iterable):
|
|
iterable_inputs[0] = True
|
|
val_inputs[0] = list(min_float_val)
|
|
lengths.append(len(min_float_val))
|
|
if isinstance(max_float_val, Iterable):
|
|
iterable_inputs[1] = True
|
|
val_inputs[1] = list(max_float_val)
|
|
lengths.append(len(max_float_val))
|
|
# make sure mask and any iterable float_vals match max length
|
|
max_length = max(lengths)
|
|
mask = extend_to_batch_size(mask, max_length)
|
|
for i in range(len(iterable_inputs)):
|
|
if iterable_inputs[i] == True:
|
|
# make sure tensors will match dimensions of mask
|
|
val_inputs[i] = torch.tensor(extend_list_to_batch_size(val_inputs[i], max_length)).unsqueeze(-1).unsqueeze(-1)
|
|
min_float_val, max_float_val = val_inputs
|
|
if scaling == ScaleType.ABSOLUTE:
|
|
mask = linear_conversion(mask.clone(), new_min=min_float_val, new_max=max_float_val)
|
|
elif scaling == ScaleType.RELATIVE:
|
|
mask = normalize_min_max(mask.clone(), new_min=min_float_val, new_max=max_float_val)
|
|
else:
|
|
raise ValueError(f"scaling '{scaling}' not recognized.")
|
|
return io.NodeOutput(*MultivalDynamicNode.execute(mask_optional=mask).args)
|
|
|
|
|
|
class MultivalDynamicFloatInputNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ADE_MultivalDynamicFloatInput',
|
|
display_name='Multival [Float List] 🎭🅐🅓',
|
|
category='Animate Diff 🎭🅐🅓/multival',
|
|
inputs=[io.Float.Input('float_val', default=1.0, force_input=True, max=10.0, min=0.0, step=0.001), io.Mask.Input('mask_optional', optional=True)],
|
|
outputs=[io.Custom('MULTIVAL').Output('MULTIVAL')]
|
|
)
|
|
@classmethod
|
|
def execute(cls, float_val: Union[float, list[float]]=None, mask_optional: Tensor=None):
|
|
return io.NodeOutput(*MultivalDynamicNode.execute(float_val=float_val, mask_optional=mask_optional).args)
|
|
|
|
|
|
class MultivalDynamicFloatsNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ADE_MultivalDynamicFloats',
|
|
display_name='Multival [Floats] 🎭🅐🅓',
|
|
category='Animate Diff 🎭🅐🅓/multival',
|
|
inputs=[io.Custom('FLOATS').Input('floats', extra_dict={'default': 1.0, 'min': 0.0, 'max': 10.0, 'step': 0.001}), io.Mask.Input('mask_optional', optional=True)],
|
|
outputs=[io.Custom('MULTIVAL').Output('MULTIVAL')]
|
|
)
|
|
@classmethod
|
|
def execute(cls, floats: Union[float, list[float]]=None, mask_optional: Tensor=None):
|
|
return io.NodeOutput(*MultivalDynamicNode.execute(float_val=floats, mask_optional=mask_optional).args)
|
|
|
|
|
|
class MultivalConvertToMaskNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id='ADE_MultivalConvertToMask',
|
|
display_name='Multival to Mask 🎭🅐🅓',
|
|
category='Animate Diff 🎭🅐🅓/multival',
|
|
inputs=[io.Custom('MULTIVAL').Input('multival')],
|
|
outputs=[io.Mask.Output('MASK')]
|
|
)
|
|
@classmethod
|
|
def execute(cls, multival: Union[float, Tensor]):
|
|
# if already tensor, assume is a valid mask
|
|
if type(multival) == Tensor:
|
|
return io.NodeOutput(multival)
|
|
# otherwise, make a single 1x1 mask with the proper value
|
|
shape = (1,1,1)
|
|
converted_multival = torch.ones(shape) * multival
|
|
return io.NodeOutput(converted_multival)
|