87 lines
3.6 KiB
Python
87 lines
3.6 KiB
Python
import torch
|
|
from antlr4 import CommonTokenStream, InputStream
|
|
from .Parser.MathExprParser import MathExprParser
|
|
from .Parser.MathExprLexer import MathExprLexer
|
|
from .Parser.TensorEvalVisitor import TensorEvalVisitor
|
|
from .helper_functions import getIndexTensorAlongDim
|
|
|
|
from comfy_api.latest import io
|
|
|
|
class AudioMathNode(io.ComfyNode):
|
|
"""
|
|
This node enables the use of math expressions on AUDIO tensors.
|
|
inputs:
|
|
a, b, c, d:
|
|
AUDIO, bound to variables with the same name. Defaults to zero AUDIO if not provided.
|
|
w, x, y, z:
|
|
Floats, bound to variables of the expression. Defaults to 0.0 if not provided.
|
|
Audio expression:
|
|
String, describing expression to apply to audio tensors.
|
|
|
|
outputs:
|
|
AUDIO:
|
|
Returns an AUDIO object that contains the result of the math expression applied to the input audio tensors.
|
|
"""
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="mrmth_AudioMathNode",
|
|
category="More math",
|
|
display_name="Audio math",
|
|
inputs=[
|
|
io.Audio.Input(id="a", tooltip="Input audio tensor"),
|
|
io.Audio.Input(id="b", optional=True, tooltip="Second input audio tensor"),
|
|
io.Audio.Input(id="c", optional=True, tooltip="Third input audio tensor"),
|
|
io.Audio.Input(id="d", optional=True, tooltip="Fourth input audio tensor"),
|
|
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, force_input=True),
|
|
io.Float.Input(id="z", default=0.0, optional=True, force_input=True),
|
|
io.String.Input(id="AudioExpr", default="a*(1-w)+b*w", tooltip="Expression to apply on input audio tensors"),
|
|
],
|
|
outputs=[
|
|
io.Audio.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, a, AudioExpr, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
|
|
|
|
|
|
bv = b if b else {'waveform':torch.zeros_like(a['waveform']),'sample_rate':a['sample_rate']}
|
|
cv = c if c else {'waveform':torch.zeros_like(a['waveform']),'sample_rate':a['sample_rate']}
|
|
dv = d if d else {'waveform':torch.zeros_like(a['waveform']),'sample_rate':a['sample_rate']}
|
|
|
|
B = getIndexTensorAlongDim(a['waveform'], 0)
|
|
C = getIndexTensorAlongDim(a['waveform'], 1)
|
|
S = getIndexTensorAlongDim(a['waveform'], 2)
|
|
R = torch.full_like(S, a['sample_rate'], dtype=torch.float32)
|
|
T = torch.full_like(S, a['waveform'].shape[2], dtype=torch.float32)
|
|
|
|
variables = {
|
|
'a': a['waveform'], 'b': bv['waveform'], 'c': cv['waveform'], 'd': dv['waveform'],
|
|
'w': w, 'x': x, 'y': y, 'z': z,
|
|
'B': B, 'C': C, 'S': S,'R': R, 'T' : T, 'N': a['waveform'].shape[1],
|
|
'batch': B, 'channel': C, 'sample': S, 'sample_rate': R, 'sample_count': T,'channel_count': a['waveform'].shape[1]
|
|
}
|
|
|
|
input_stream = InputStream(AudioExpr)
|
|
lexer = MathExprLexer(input_stream)
|
|
stream = CommonTokenStream(lexer)
|
|
parser = MathExprParser(stream)
|
|
tree = parser.expr()
|
|
|
|
visitor = TensorEvalVisitor(variables, a['waveform'].shape)
|
|
result_tensor = visitor.visit(tree)
|
|
|
|
# Create output dictionary with the same sample rate
|
|
output = {
|
|
'waveform': result_tensor,
|
|
'sample_rate': a['sample_rate']
|
|
}
|
|
|
|
return (output,)
|