138 lines
4.9 KiB
Python
138 lines
4.9 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,)
|