mcDandy
458901bb2a
should have added them from beginning instead of giving custom names to them. Will not be deleting that.
110 lines
4.6 KiB
Python
110 lines
4.6 KiB
Python
from inspect import cleandoc
|
|
|
|
from comfy_api.latest import io
|
|
from comfy_api.input_impl import VideoFromComponents
|
|
from comfy_api.util import VideoComponents
|
|
|
|
import torch
|
|
|
|
from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, eval_tensor_expr, make_zero_like
|
|
|
|
|
|
from .MathNodeBase import MathNodeBase
|
|
|
|
class VideoMathNode(MathNodeBase):
|
|
"""
|
|
Enables math expressions on Video (images + audio).
|
|
|
|
Inputs:
|
|
a, b, c, d: Video inputs (b, c, d default to zero if not provided)
|
|
w, x, y, z: Float variables for expressions
|
|
Audio: Expression for audio component
|
|
Images: Expression for image component
|
|
|
|
Outputs:
|
|
VIDEO: Result of applying expressions to input videos
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="mrmth_VideoMathNode",
|
|
display_name="Video math",
|
|
category="More math",
|
|
inputs=[
|
|
io.Video.Input(id="a"),
|
|
io.Video.Input(id="b", optional=True),
|
|
io.Video.Input(id="c", optional=True),
|
|
io.Video.Input(id="d", optional=True),
|
|
io.Float.Input(id="w", default=0.0, optional=True, force_input=True),
|
|
io.Float.Input(id="x", default=0.0, optional=True, force_input=True),
|
|
io.Float.Input(id="y", default=0.0, optional=True, lazy=True, force_input=True),
|
|
io.Float.Input(id="z", default=0.0, optional=True, force_input=True),
|
|
io.String.Input(id="Audio", default="a*(1-w)+b*w", tooltip="Expression to apply on audio part of video"),
|
|
io.String.Input(id="Images", default="a*(1-w)+b*w", tooltip="Expression to apply on image part of video"),
|
|
],
|
|
outputs=[
|
|
io.Video.Output(),
|
|
],
|
|
)
|
|
|
|
tooltip = cleandoc(__doc__)
|
|
|
|
@classmethod
|
|
def execute(cls, Audio, Images, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:
|
|
ac = a.get_components()
|
|
|
|
bc = b.get_components() if b is not None else make_zero_like(ac)
|
|
cc = c.get_components() if c is not None else make_zero_like(ac)
|
|
dc = d.get_components() if d is not None else make_zero_like(ac)
|
|
|
|
# Process images (permute to B, C, H, W)
|
|
imgs_a = ac.images.permute(0, 3, 1, 2)
|
|
imgs_b = bc.images.permute(0, 3, 1, 2)
|
|
imgs_c = cc.images.permute(0, 3, 1, 2)
|
|
imgs_d = dc.images.permute(0, 3, 1, 2)
|
|
|
|
img_vars = {
|
|
'a': imgs_a, 'b': imgs_b, 'c': imgs_c, 'd': imgs_d,
|
|
'w': w, 'x': x, 'y': y, 'z': z,
|
|
'X': getIndexTensorAlongDim(imgs_a, 3),
|
|
'Y': getIndexTensorAlongDim(imgs_a, 2),
|
|
'B': getIndexTensorAlongDim(imgs_a, 0), 'frame': getIndexTensorAlongDim(imgs_a, 0),
|
|
'C': getIndexTensorAlongDim(imgs_a, 1), 'channel': getIndexTensorAlongDim(imgs_a, 1),
|
|
'W': imgs_a.shape[3], 'width': imgs_a.shape[3],
|
|
'H': imgs_a.shape[2], 'height': imgs_a.shape[2],
|
|
'R': float(ac.frame_rate), 'frame_rate': float(ac.frame_rate),
|
|
'T': imgs_a.shape[0], 'frame_count': imgs_a.shape[0],
|
|
'N': imgs_a.shape[1], 'channel_count': imgs_a.shape[1],
|
|
} | generate_dim_variables(imgs_a)
|
|
|
|
result_imgs = eval_tensor_expr(Images, img_vars, imgs_a.shape)
|
|
result_imgs = result_imgs.permute(0, 2, 3, 1) # Back to B, H, W, C
|
|
|
|
# Process audio
|
|
audio_a = ac.audio['waveform']
|
|
audio_b = bc.audio['waveform']
|
|
audio_c = cc.audio['waveform']
|
|
audio_d = dc.audio['waveform']
|
|
|
|
audio_vars = {
|
|
'a': audio_a, 'b': audio_b, 'c': audio_c, 'd': audio_d,
|
|
'w': w, 'x': x, 'y': y, 'z': z,
|
|
'B': getIndexTensorAlongDim(audio_a, 0), 'batch': getIndexTensorAlongDim(audio_a, 0),
|
|
'C': getIndexTensorAlongDim(audio_a, 1), 'channel': getIndexTensorAlongDim(audio_a, 1),
|
|
'S': getIndexTensorAlongDim(audio_a, 2), 'sample': getIndexTensorAlongDim(audio_a, 2),
|
|
'R': ac.audio['sample_rate'], 'sample_rate': ac.audio['sample_rate'],
|
|
'T': audio_a.shape[2], 'sample_count': audio_a.shape[2],
|
|
'N': audio_a.shape[1], 'channel_count': audio_a.shape[1],
|
|
} | generate_dim_variables(audio_a)
|
|
|
|
result_audio = eval_tensor_expr(Audio, audio_vars, audio_a.shape)
|
|
|
|
output = VideoFromComponents(VideoComponents(
|
|
images=result_imgs,
|
|
audio={'waveform': result_audio, 'sample_rate': ac.audio['sample_rate']},
|
|
frame_rate=ac.frame_rate,
|
|
metadata=ac.metadata
|
|
))
|
|
return (output,)
|